Ollamaの埋め込み(embedding)をPostgreSQLのpgvectorで保存して検索する

4 次元の 3 行。テキストは /api/embed のリクエストボディ、ベクトルは Ollama /api/embed のレスポンス として読みました。

psqlなどのクライアントで実行します。行は1つのトランザクションで挿入されます。

BEGIN;

INSERT INTO documents (content, embedding) VALUES
  ('Embeddings turn text into numeric vectors.', '[0.620186,0.410123,0.550165,0.380114]'),
  ('pgvector stores vectors in PostgreSQL, next to the rest of your data.', '[0.578815,0.469039,0.518937,0.419142]'),
  ('Ollama''s /api/embed returns one vector per input, in the same order.', '[0.119982,0.809879,-0.329951,0.46993]');

COMMIT;

CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE documents (
  id bigserial PRIMARY KEY,
  content text NOT NULL,
  embedding vector(384) NOT NULL
);

CREATE INDEX ON documents USING hnsw (embedding vector_cosine_ops) WITH (m = 16, ef_construction = 64);

PREPARE は、/api/embed で得た質問のベクトルを $1 として受け取ります。'[0.01,-0.02,…]' のようにテキストとして渡してください。

PREPARE search_documents(vector(384)) AS
  SELECT id, content, 1 - (embedding <=> $1) AS cosine_similarity
  FROM documents
  ORDER BY embedding <=> $1
  LIMIT 5;

-- EXECUTE search_documents('[0.0123, -0.0456, ...]');

出典3件をオブジェクトのJSON配列として読み取りました。

このマシン上の Ollama(localhost:11434)にリクエストを送ります。

curl http://localhost:11434/api/chat -d '{
  "model": "qwen2.5:7b",
  "messages": [
    {
      "role": "system",
      "content": "Answer the question using only the numbered sources in the user message. Cite the number of each source you use. If the sources do not contain the answer, say that you do not know and cite nothing. Answer in the same language as the question. For each citation, give the source number and copy a short quote from that source exactly as it is written. Return JSON that matches this schema: {\"type\":\"object\",\"properties\":{\"answer\":{\"type\":\"string\"},\"citations\":{\"type\":\"array\",\"items\":{\"type\":\"object\",\"properties\":{\"source\":{\"type\":\"integer\"},\"quote\":{\"type\":\"string\"}},\"required\":[\"source\",\"quote\"]}}},\"required\":[\"answer\",\"citations\"]}"
    },
    {
      "role": "user",
      "content": "Sources:\n\n[1] (id: 12)\nWith HNSW, hnsw.ef_search is 40 by default, so a query returns at most 40 rows unless you raise it.\n\n[2] (id: 7)\nSET LOCAL inside a transaction changes hnsw.ef_search for a single query.\n\n[3] (id: 3)\npgvector supports PostgreSQL 13 and later.\n\nQuestion: Why does my HNSW search return only 40 rows with LIMIT 100?"
    }
  ],
  "stream": false,
  "format": {
    "type": "object",
    "properties": {
      "answer": {
        "type": "string"
      },
      "citations": {
        "type": "array",
        "items": {
          "type": "object",
          "properties": {
            "source": {
              "type": "integer"
            },
            "quote": {
              "type": "string"
            }
          },
          "required": [
            "source",
            "quote"
          ]
        }
      }
    },
    "required": [
      "answer",
      "citations"
    ]
  },
  "options": {
    "temperature": 0
  },
  "truncate": false
}'

{
  "model": "qwen2.5:7b",
  "messages": [
    {
      "role": "system",
      "content": "Answer the question using only the numbered sources in the user message. Cite the number of each source you use. If the sources do not contain the answer, say that you do not know and cite nothing. Answer in the same language as the question. For each citation, give the source number and copy a short quote from that source exactly as it is written. Return JSON that matches this schema: {\"type\":\"object\",\"properties\":{\"answer\":{\"type\":\"string\"},\"citations\":{\"type\":\"array\",\"items\":{\"type\":\"object\",\"properties\":{\"source\":{\"type\":\"integer\"},\"quote\":{\"type\":\"string\"}},\"required\":[\"source\",\"quote\"]}}},\"required\":[\"answer\",\"citations\"]}"
    },
    {
      "role": "user",
      "content": "Sources:\n\n[1] (id: 12)\nWith HNSW, hnsw.ef_search is 40 by default, so a query returns at most 40 rows unless you raise it.\n\n[2] (id: 7)\nSET LOCAL inside a transaction changes hnsw.ef_search for a single query.\n\n[3] (id: 3)\npgvector supports PostgreSQL 13 and later.\n\nQuestion: Why does my HNSW search return only 40 rows with LIMIT 100?"
    }
  ],
  "stream": false,
  "format": {
    "type": "object",
    "properties": {
      "answer": {
        "type": "string"
      },
      "citations": {
        "type": "array",
        "items": {
          "type": "object",
          "properties": {
            "source": {
              "type": "integer"
            },
            "quote": {
              "type": "string"
            }
          },
          "required": [
            "source",
            "quote"
          ]
        }
      }
    },
    "required": [
      "answer",
      "citations"
    ]
  },
  "options": {
    "temperature": 0
  },
  "truncate": false
}

レスポンスでは、回答は message.content にあります。prompt_eval_count が思っていたよりかなり少ない場合は、プロンプトが切り詰められた可能性があります。

引用