Ein RAG-System mit PostgreSQL aufbauen

• Mit einem Retrieval-Augmented-Generation-(RAG)-System lässt sich ein KI-Assistent erstellen, der Fragen auf Basis einer bestehenden internen Wissensbasis beantworten kann, etwa Wikis, Handbücher sowie Schulungs- und Referenzmaterialien
• Ein RAG-System lässt sich allein mit PostgreSQL, pgvector, ollama und weniger als 200 Zeilen Go-Code aufbauen

Überblick

• Einige Absätze einer Geschichte werden als „Dokumentkorpus“ verwendet, und für jedes Dokument werden mit Metas Llama3 Dokument-Embeddings erzeugt, lokal gehostet über ollama
• Dokumente und Embeddings werden in einer PostgreSQL-Tabelle gespeichert; die Erweiterung pgvector übernimmt das Speichern und den Zugriff auf die Embeddings
• Für eine Nutzerabfrage wird das relevanteste Dokument aus der Tabelle abgerufen und anschließend mit Llama3 eine Antwort erzeugt
• ollama stellt eine OpenAI-ähnliche HTTP-API bereit, um Embeddings und Chat-Antworten zu erzeugen
• Der Go-Code kommuniziert mit Postgres über jackc/pgx und pgvector-go und nutzt das ollama-Client-API-Paket für die Aufrufe der HTTP-API

Modelle mit Ollama ausführen

• Ollama ist ein Tool, mit dem sich Open-Source-Modelle lokal ausführen lassen, und es bietet eine REST-API im OpenAI-Stil
• Mit dem Befehl ollama pull llama3 lässt sich das llama3-Modell ausführen
• Der HTTP-Server von ollama ist standardmäßig unter 127.0.0.1:11434 erreichbar

pgvector installieren

• pgvector ist eine Erweiterung für PostgreSQL 12 bis 16 und kann bei Nutzung des pgdg-APT-Repositories mit sudo apt install postgresql-16-pgvector installiert werden
• Nach der Installation wird die Erweiterung mit create extension vector; in der Datenbank eingerichtet
• Mit create table items (id serial primary key, doc text, embedding vector(4096)); wird eine Tabelle zum Speichern von Dokumenten und Embeddings erstellt

Dokumentdaten

• Verwendet werden vier Absätze aus der Sherlock-Holmes-Geschichte „The Boscombe Valley Mystery“ (Public Domain, Projekt Gutenberg)

Code

• Es steht Demo-Code unter MIT-Lizenz auf GitHub zur Verfügung
• Dokumente werden mit INSERT INTO items (doc, embedding) VALUES ($1, $2) eingefügt
• Das relevanteste Dokument wird mit SELECT doc FROM items ORDER BY embedding <-> $1 LIMIT 1 gesucht (der Operator <-> wird von pgvector bereitgestellt)
• Für Aufrufe der Ollama-API wird das Go-Paket von ollama verwendet
  • Zur Erzeugung von Embeddings wird api.EmbeddingRequest verwendet
  • Zur Erzeugung von Chat-Antworten wird api.ChatRequest verwendet, wobei das abgerufene Dokument in den Prompt aufgenommen wird

Kommandozeilenoberfläche

• Mit ragdemo -insert {path-to-doc-file} werden Dokumente in der Datenbank gespeichert
• Mit ragdemo -query {query-text} wird ein Prompt eingegeben und eine Antwort erzeugt

Gesamtablauf

1. Beim Speichern eines Dokuments mit der Option -insert wird der Dateiinhalt gelesen, mit Llama3 ein Embedding erzeugt und in PostgreSQL gespeichert
2. Bei Verwendung der Option -query wird ein Prompt-Embedding erzeugt und mit den übrigen Embeddings in der Tabelle items verglichen, um das „nächste“ Dokument zu finden (L2-Distanzberechnung über den Operator <->)
3. Das gefundene Dokument wird in den Prompt aufgenommen, an Llama3 übergeben, und die erzeugte Chat-Antwort wird ausgegeben

Zusätzliche Tipps

• Es kann sinnvoll sein, ein speziell für Embeddings ausgelegtes Modell zu verwenden, statt llama3
• Für andere Sprachen als Englisch sollte nach besser geeigneten Modellen gesucht werden
• Neben der L2-Distanz lassen sich auch andere Distanzmethoden ausprobieren, da pgvector weitere Verfahren unterstützt
• Ein vollständiger Tabellenscan ist nicht gut skalierbar; daher sollten etwa pgvector-Indizes genutzt werden
• In der Generierungsphase kann es helfen, mehr Dokumente einzubeziehen oder zusätzliche Dokumente per Keyword-Matching nachzuladen
• Durch Anpassung des Generierungs-Prompts und Tests mit verschiedenen LLMs lässt sich die Ausgabequalität verbessern