Construire un système RAG avec PostgreSQL

• Un système de Retrieval-Augmented Generation (RAG) permet de créer un assistant IA capable de répondre à des questions à partir d’une base de connaissances interne existante (wiki, manuels, supports de formation et documents de référence, etc.)
• Il est possible de construire un système RAG uniquement avec PostgreSQL, pgvector, ollama et moins de 200 lignes de code Go

Overview

• Utilise quelques paragraphes comme « corpus de documents » et génère un embedding pour chaque document avec Llama3 de Meta (hébergé localement via ollama)
• Stocke les documents et leurs embeddings dans une table PostgreSQL, avec l’extension pgvector pour enregistrer et accéder aux embeddings
• Pour une requête utilisateur, récupère dans la table le document le plus pertinent, puis génère une réponse avec Llama3
• ollama fournit une API HTTP similaire à OpenAI pour générer des embeddings et des réponses de chat
• Le code Go communique avec Postgres via jackc/pgx et pgvector-go, et utilise le package API client d’ollama pour gérer les appels à l’API HTTP

Exécuter les modèles avec Ollama

• Ollama est un outil qui permet d’exécuter des modèles open source en local et fournit une API REST au style OpenAI
• Exécuter le modèle llama3 avec la commande ollama pull llama3
• Le serveur HTTP d’ollama est disponible par défaut sur 127.0.0.1:11434

Installation de pgvector

• pgvector est une extension pour PostgreSQL versions 12 à 16 ; avec le dépôt APT pgdg, il peut être installé via sudo apt install postgresql-16-pgvector
• Après installation, activer l’extension dans la base de données avec create extension vector;
• Crée une table pour stocker les documents et embeddings avec create table items (id serial primary key, doc text, embedding vector(4096));

Données de documents

• Utilise 4 paragraphes de l’histoire de Sherlock Holmes « The Boscombe Valley Mystery » (domaine public - Project Gutenberg)

Code

• Le code de démonstration, publié sur GitHub sous licence MIT, peut être utilisé
• L’insertion des documents se fait avec INSERT INTO items (doc, embedding) VALUES ($1, $2)
• La recherche du document le plus pertinent se fait avec SELECT doc FROM items ORDER BY embedding <-> $1 LIMIT 1 (l’opérateur <-> est fourni par pgvector)
• Les appels à l’API Ollama utilisent le package Go ollama
  • La génération d’embeddings utilise api.EmbeddingRequest
  • La génération de réponses de chat utilise api.ChatRequest (en incluant le document récupéré dans le prompt)

Interface en ligne de commande

• ragdemo -insert {path-to-doc-file} enregistre un document dans la base de données
• ragdemo -query {query-text} envoie un prompt et génère une réponse

Processus complet

1. Lors de l’enregistrement d’un document avec l’option -insert, le contenu du fichier est lu, un embedding est généré avec Llama3, puis stocké dans PostgreSQL
2. Lors de l’utilisation de l’option -query, un embedding du prompt est généré puis comparé aux autres embeddings de la table items afin de retrouver le document « nearest neighbor » (<-> calcule la distance L2)
3. Le document récupéré est inclus dans le prompt envoyé à Llama3, qui génère ensuite une réponse de chat affichée en sortie

Conseils supplémentaires

• Envisager l’utilisation d’un modèle spécialisé dans la génération d’embeddings (plutôt que llama3)
• Pour les langues autres que l’anglais, il peut être nécessaire de chercher un modèle plus adapté
• Il est possible d’essayer d’autres méthodes de calcul de distance que la distance L2 (pgvector en prend en charge d’autres)
• Un scan complet de la table passe mal à l’échelle, donc mieux vaut utiliser des index pgvector ou équivalents
• À l’étape de génération, il peut être utile d’exploiter davantage de documents ou d’en récupérer d’autres via du keyword matching
• Ajuster le prompt de génération et tester différents LLM peut améliorer la qualité de sortie