Guarde e pesquise embeddings do Ollama no PostgreSQL com o pgvector

O README do pgvector o descreve como busca por similaridade de vetores de código aberto para o Postgres. Ele oferece suporte a busca exata e aproximada de vizinhos mais próximos; vetores de precisão simples, meia precisão, binários e esparsos; e distância L2, produto interno, distância de cosseno, distância L1, distância de Hamming e distância de Jaccard. Como os vetores são colunas comuns do PostgreSQL, eles também contam com conformidade ACID, recuperação em um ponto no tempo e JOINs.

Para uma busca nos seus documentos, você guarda uma linha por trecho com o texto e o vetor, gera o embedding da pergunta com o mesmo modelo e ordena as linhas pela distância até o vetor da pergunta.

O pgvector é compatível com o PostgreSQL 13 e posteriores. O README lista várias formas de instalá-lo, entre elas estas:

Uma coluna do tipo vector(n) guarda vetores com exatamente n dimensões, e n deve ser o comprimento dos vetores do seu modelo de embeddings. Um vetor com outro comprimento é rejeitado: quando inserimos um vetor com 2 dimensões em uma coluna vector(3), o PostgreSQL retornou "expected 3 dimensions, not 2". O all-minilm retorna 384 dimensões, então sua coluna é vector(384).

O README informa que cada vetor ocupa 4 × dimensões + 8 bytes, com cada elemento sendo um número de ponto flutuante de precisão simples, e permite até 16,000 dimensões em uma coluna. Um índice aceita menos dimensões, como mostrado abaixo. O tipo halfvec guarda valores de meia precisão em 2 × dimensões + 8 bytes.

O pgvector lê um vetor como texto entre colchetes, como '[1,2,3]', e o README diz que todos os elementos devem ser finitos (sem NaN, Infinity nem -Infinity). Cada array do campo embeddings de uma resposta de /api/embed já tem essa forma, e a resposta mantém a ordem da entrada, então o primeiro texto corresponde ao primeiro vetor.

O README mostra INSERT para poucas linhas, INSERT … ON CONFLICT para upserts e COPY para carga em massa. O conversor abaixo emparelha seus textos com a resposta e escreve instruções INSERT com as aspas escapadas. Executamos a saída dele com textos contendo aspas, barras invertidas, quebras de linha e emoji, e eles voltaram sem alterações.

A documentação do Ollama recomenda a similaridade de cosseno para a maioria dos casos de busca semântica e diz que o /api/embed retorna vetores com comprimento 1. Para esses vetores, a distância de cosseno, a distância L2 e o produto interno colocam as linhas na mesma ordem, porque a distância L2 ao quadrado é 2 menos 2 vezes a similaridade de cosseno. No nosso teste com o all-minilm, as três primeiras linhas foram as mesmas com as três, e a distância L2 de 0.931 coincidiu com a similaridade de cosseno de 0.566.

Cada distância tem o seu próprio operador. Uma busca ordena pelo operador e pega as primeiras linhas com LIMIT, e um índice é criado para uma classe de operador.

Sem índice, o pgvector faz uma busca exata de vizinhos mais próximos, que segundo o README oferece recall perfeito. Um índice aproximado troca um pouco de recall por velocidade, e o README avisa que as consultas podem retornar resultados diferentes depois que você adiciona um.

Um índice só é usado quando a consulta tem ORDER BY e LIMIT e o ORDER BY é o próprio operador de distância em ordem crescente. ORDER BY embedding <=> $1 LIMIT 5 pode usar o índice; ORDER BY 1 - (embedding <=> $1) DESC não pode. O README também observa que uma varredura da tabela pode ser mais rápida quando a tabela é pequena.

Com HNSW, os resultados são limitados por hnsw.ef_search, que é 40 por padrão. Em uma tabela de teste de 200 linhas com um índice HNSW, LIMIT 100 devolveu 40 linhas; depois de SET hnsw.ef_search = 100, devolveu 100. Defina hnsw.ef_search com pelo menos o número de linhas de que você precisa. SET LOCAL dentro de uma transação o altera apenas para uma consulta.

A filtragem também reduz os resultados. Com índices aproximados, o README diz que a condição WHERE é aplicada depois que o índice é percorrido: se uma condição corresponde a 10% das linhas, o HNSW com o hnsw.ef_search padrão de 40 devolve apenas 4 linhas em média. Desde a versão 0.8.0, as varreduras iterativas do índice (SET hnsw.iterative_scan = strict_order, ou relaxed_order) percorrem mais do índice até encontrar linhas suficientes, até hnsw.max_scan_tuples (20,000 por padrão).

Quando a busca devolve as linhas mais próximas, o modelo precisa delas no prompt. O /api/chat do Ollama recebe uma lista de mensagens, cada uma com um papel (system, user, assistant ou tool) e o seu content. Uma organização simples é uma mensagem system com as regras e uma mensagem user que lista as linhas, cada uma com um número, como [1] e [2], seguida da pergunta. Com os números, a resposta pode dizer quais linhas usou, e você pode associar cada número à sua linha.

Para obter as citações em um formato fixo, dê ao campo format um JSON Schema com a resposta e uma lista de citações, cada uma com o número da fonte e um trecho citado curto. A página de saídas estruturadas do Ollama recomenda escrever o esquema também no prompt e baixar a temperatura (temperature), por exemplo para 0, para obter uma saída mais determinística. No psql, a função json_agg reúne as linhas de uma consulta em um array JSON, que a ferramenta abaixo lê diretamente.

As regras, as linhas e a pergunta precisam caber no tamanho do contexto. A página sobre o tamanho do contexto do Ollama indica um padrão de 4k tokens com menos de 24 GiB de VRAM, e o FAQ indica 4096 tokens; ollama ps na nossa máquina de teste mostrou 4096. num_ctx em options ou a variável de ambiente OLLAMA_CONTEXT_LENGTH o altera.

Quando o prompt é mais longo, o /api/chat não devolve um erro por padrão. No código-fonte do Ollama (0.40.11), o servidor primeiro descarta mensagens anteriores inteiras, mantendo as mensagens system e a última mensagem. Se o prompt restante ainda for longo demais, ele mantém os primeiros tokens (num_keep, 4 por padrão) e o final do prompt e descarta a parte do meio. No nosso teste com qwen2.5:7b e num_ctx 1024, um prompt de 2,068 tokens foi cortado para 514. A mensagem system e as primeiras fontes foram perdidas: com a resposta na fonte [1], o modelo respondeu com uma palavra-código que não estava em nenhuma fonte, e com a resposta na última fonte, respondeu corretamente, mas sem o número da fonte que a mensagem system pedia. A resposta não continha nenhum erro; o log do servidor tinha o aviso "truncating input prompt".

O código-fonte do Ollama também tem um campo truncate para requisições de chat, que não consta na página do /api/chat. Com "truncate": false, a mesma requisição devolveu HTTP 400 com o erro "request (2068 tokens) exceeds the available context size (1024 tokens), try increasing it". O campo prompt_eval_count de uma resposta é o número de tokens do prompt; no teste com o corte, foi 514, não 2,068.

Uma citação pode ser verificada de forma mecânica: o número deve se referir a uma fonte que foi enviada, e o trecho citado deveria aparecer nessa fonte. No nosso teste com quatro fontes e qwen2.5:7b, os trechos citados nas respostas em JSON coincidiam palavra por palavra com as fontes. Isso não prova que a resposta esteja certa: ao perguntar o tamanho de um vetor de 384 dimensões, o modelo citou a regra correta (4 vezes as dimensões mais 8 bytes) e mesmo assim respondeu 1552 bytes em vez de 1,544.

Olhe também as respostas sem citações. Quando as fontes não continham a resposta, o modelo uma vez respondeu com o próprio conhecimento e sem citações, e outra vez disse que não sabia, mas ainda assim citou uma fonte sem relação. Se done_reason na resposta for "length", a resposta foi cortada por num_predict e pode estar sem as citações.

O README diz que um índice HNSW é construído de forma significativamente mais rápida quando o grafo cabe em maintenance_work_mem, mostra um aviso quando ele deixa de caber e alerta para não definir maintenance_work_mem tão alto a ponto de esgotar a memória do servidor. Como com outros índices, é mais rápido criar o índice depois de carregar os dados iniciais. Os índices não precisam caber na memória, mas provavelmente você verá um desempenho melhor se couberem; SELECT pg_size_pretty(pg_relation_size('index_name')) mostra o tamanho de um índice.

Para um conjunto de trabalho menor e índices menores, o README sugere halfvec, indexação de meia precisão ou quantização binária. Se o Ollama roda no mesmo servidor, o PostgreSQL e os modelos compartilham a mesma memória; nosso artigo sobre como acompanhar a memória do Ollama explica o que determina quanto os modelos usam.