Como construir um sistema RAG com PostgreSQL

• Com um sistema de Retrieval-Augmented Generation (RAG), é possível criar um assistente de IA que responde perguntas com base em uma base interna de conhecimento já existente (wiki, manuais, materiais de treinamento e referência etc.)
• É possível construir um sistema RAG usando apenas PostgreSQL, pgvector, ollama e menos de 200 linhas de código em Go

Visão geral

• Usa alguns parágrafos de texto como um "corpus de documentos" e, para cada documento, gera embeddings com o Llama3 da Meta (hospedado localmente via ollama)
• Armazena os documentos e os embeddings em uma tabela do PostgreSQL, usando a extensão pgvector para salvar e acessar os embeddings
• Para uma consulta do usuário, busca na tabela o 1 documento mais relevante e gera uma resposta com o Llama3
• O ollama fornece uma API HTTP semelhante à da OpenAI para gerar embeddings e respostas de chat
• O código em Go usa jackc/pgx e pgvector-go para se comunicar com o Postgres, e usa o pacote de API cliente do ollama para lidar com chamadas à API HTTP

Executando modelos com Ollama

• Ollama é uma ferramenta que permite executar modelos open source localmente e fornece uma API REST no estilo OpenAI
• Execute o modelo llama3 com o comando ollama pull llama3
• O servidor HTTP do ollama fica disponível por padrão em 127.0.0.1:11434

Instalando o pgvector

• pgvector é uma extensão para PostgreSQL nas versões 12 a 16; ao usar o repositório pgdg APT, é possível instalar com sudo apt install postgresql-16-pgvector
• Após a instalação, configure a extensão no banco com create extension vector;
• Crie uma tabela para armazenar documentos e embeddings com create table items (id serial primary key, doc text, embedding vector(4096));

Dados dos documentos

• São usados 4 parágrafos do conto de Sherlock Holmes "The Boscombe Valley Mystery" (domínio público - Project Gutenberg)

Código

• O código de demonstração publicado no GitHub sob licença MIT pode ser usado
• A inserção de documentos é feita com INSERT INTO items (doc, embedding) VALUES ($1, $2)
• A busca do documento mais relevante é feita com SELECT doc FROM items ORDER BY embedding <-> $1 LIMIT 1 (o operador <-> é fornecido pelo pgvector)
• As chamadas à API do Ollama usam o pacote Go do ollama
  • A geração de embeddings usa api.EmbeddingRequest
  • A geração de respostas de chat usa api.ChatRequest (incluindo no prompt o documento recuperado)

Interface de linha de comando

• Use ragdemo -insert {path-to-doc-file} para salvar um documento no banco de dados
• Use ragdemo -query {query-text} para inserir um prompt e gerar uma resposta

Processo completo

1. Ao salvar um documento com a opção -insert, o conteúdo do arquivo é lido, os embeddings são gerados com Llama3 e armazenados no PostgreSQL
2. Ao usar a opção -query, o embedding do prompt é gerado e comparado com os demais embeddings da tabela items para buscar o documento "vizinho mais próximo" (calculando a distância L2 com o operador <->)
3. O documento recuperado é incluído no prompt e enviado ao Llama3, que gera e exibe a resposta de chat

Dicas extras

• Considere usar um modelo especializado em geração de embeddings (em vez de llama3)
• Para idiomas além do inglês, pode ser necessário procurar um modelo mais adequado
• Também é possível testar outros métodos de cálculo de distância além da distância L2 (o pgvector oferece suporte a outros métodos)
• Como fazer varredura completa da tabela não escala bem, vale usar índices do pgvector e outros recursos
• Na etapa de geração, também pode ajudar usar mais documentos ou buscar documentos adicionais com correspondência por palavras-chave
• Ajustar o prompt de geração e testar diferentes LLMs pode melhorar a qualidade da saída