Cómo construir un sistema RAG con PostgreSQL

• Con un sistema de Retrieval-Augmented Generation (RAG) se puede crear un asistente de IA que responda preguntas con base en una base de conocimiento interna existente (wiki, manuales, materiales de capacitación y referencia, etc.)
• Es posible construir un sistema RAG usando solo PostgreSQL, pgvector, ollama y menos de 200 líneas de código en Go

Overview

• Se usan algunos párrafos de una historia como "corpus de documentos" y, para cada documento, se generan embeddings del documento con Llama3 de Meta (hospedado localmente con ollama)
• Los documentos y embeddings se almacenan en una tabla de PostgreSQL, y la extensión pgvector permite guardar y acceder a los embeddings
• Para una consulta del usuario, se recupera de la tabla el documento más relevante y se genera una respuesta con Llama3
• ollama ofrece una API HTTP similar a OpenAI para generar embeddings y respuestas de chat
• El código en Go se comunica con Postgres usando jackc/pgx y pgvector-go, y usa el paquete cliente de API de ollama para manejar las llamadas a la API HTTP

Ejecutar modelos con Ollama

• Ollama es una herramienta que permite ejecutar modelos open source de forma local y ofrece una API REST estilo OpenAI
• Ejecutar el modelo llama3 con el comando ollama pull llama3
• El servidor HTTP de ollama está disponible por defecto en 127.0.0.1:11434

Instalar pgvector

• pgvector es una extensión para PostgreSQL versiones 12~16; si se usa el repositorio pgdg APT, puede instalarse con sudo apt install postgresql-16-pgvector
• Después de la instalación, configurar la extensión en la base de datos con create extension vector;
• Crear una tabla para guardar documentos y embeddings con create table items (id serial primary key, doc text, embedding vector(4096));

Datos de documentos

• Se usan 4 párrafos de la historia de Sherlock Holmes "The Boscombe Valley Mystery" (dominio público - Proyecto Gutenberg)

Código

• Se puede usar el código demo publicado en GitHub bajo licencia MIT
• La inserción de documentos se hace con INSERT INTO items (doc, embedding) VALUES ($1, $2)
• La búsqueda del documento más relevante se hace con SELECT doc FROM items ORDER BY embedding <-> $1 LIMIT 1 (<-> es un operador proporcionado por pgvector)
• Las llamadas a la API de Ollama usan el paquete de Go de ollama
  • Para generar embeddings se usa api.EmbeddingRequest
  • Para generar respuestas de chat se usa api.ChatRequest (incluyendo en el prompt el documento recuperado)

Interfaz de línea de comandos

• ragdemo -insert {path-to-doc-file} guarda documentos en la base de datos
• ragdemo -query {query-text} ingresa un prompt y genera una respuesta

Proceso completo

1. Al guardar un documento con la opción -insert, se lee el contenido del archivo, se generan embeddings con Llama3 y se guardan en PostgreSQL
2. Al usar la opción -query, se generan embeddings del prompt y luego se comparan con los demás embeddings de la tabla items para recuperar el documento "más cercano" (calculando distancia L2 con el operador <->)
3. El documento recuperado se incluye en el prompt y se envía a Llama3 para generar y mostrar una respuesta de chat

Tips adicionales

• Considerar usar un modelo especializado en generación de embeddings (en lugar de llama3)
• Para idiomas distintos del inglés, puede ser necesario buscar un modelo más adecuado
• Además de la distancia L2, se pueden probar otros métodos de distancia (pgvector soporta otros métodos)
• Un escaneo completo de la tabla escala mal, por lo que conviene usar índices de pgvector, entre otros
• También puede ayudar usar más documentos en la etapa de generación, o traer documentos adicionales con coincidencia de palabras clave, etc.
• Ajustar el prompt de generación y probar distintos LLM puede mejorar la calidad de salida