Skip to content
Search

Retrieval-Augmented Generation

lastmod 2026-09-04created 2024-04-2618 edits

Retrieval-augmented generation (RAG) is an architectural approach that improves relevance by bringing data into the context of a large language model (LLM).

Categories

Building a RAG system with PostgreSQL

With a Retrieval-Augmented Generation (RAG) system, you can build an AI assistant that answers questions based on an existing internal knowledge base (wikis, manuals, training and reference materials, etc.)

A RAG system can be built with just PostgreSQL, pgvector, ollama, and under 200 lines of Go code

Overview

  • Uses a few paragraphs of a story as the "document corpus" and generates a document embedding for each document using Meta's Llama3 (hosted locally with ollama)
  • Stores documents and embeddings in a PostgreSQL table, storing and accessing embeddings through the pgvector extension
  • For a user query, retrieves the single most relevant document from the table and generates a response with Llama3
  • ollama provides an OpenAI-like HTTP API that generates embeddings and chat responses
  • The Go code talks to Postgres using jackc/pgx and pgvector-go, and handles HTTP API calls using the ollama client API package

Running models with Ollama

  • Ollama is a tool that lets you run open-source models locally, and provides an OpenAI-style REST API
  • Run the llama3 model with the ollama pull llama3 command
  • ollama's HTTP server is available at 127.0.0.1:11434 by default

Installing pgvector

  • pgvector is an extension for PostgreSQL versions 12–16; when using the pgdg APT repository it can be installed with sudo apt install postgresql-16-pgvector
  • After installing, set up the extension in the database with create extension vector;
  • Create a table for storing documents and embeddings with create table items (id serial primary key, doc text, embedding vector(4096));

Document data

  • Uses 4 paragraphs from the Sherlock Holmes story "The Boscombe Valley Mystery" (public domain - Project Gutenberg)

Code

  • Demo code released on GitHub under the MIT license is available
  • Inserting a document is INSERT INTO items (doc, embedding) VALUES ($1, $2)
  • Retrieving the most relevant document is SELECT doc FROM items ORDER BY embedding <-> $1 LIMIT 1 (the <-> operator is provided by pgvector)
  • Ollama API calls use the ollama Go package
    • Embedding generation uses api.EmbeddingRequest
    • Chat response generation uses api.ChatRequest (with the retrieved document included in the prompt)

Command-line interface

  • Store a document in the database with ragdemo -insert {path-to-doc-file}
  • Enter a prompt and generate a response with ragdemo -query {query-text}

Overall process

  • When storing a document with the -insert option, it reads the file contents, generates an embedding with Llama3, and stores it in PostgreSQL
  • When using the -query option, it generates the prompt embedding and then compares it with the other embeddings in the items table to retrieve the "nearest" document (L2 distance computed with the <-> operator)
  • The retrieved document is included in the prompt and passed to Llama3, and the chat response is generated and printed

Additional tips

  • Consider using a model specialized for embedding generation (instead of llama3)
  • For languages other than English, look for more suitable models
  • You can try other distance measures besides L2 distance (pgvector supports other methods)
  • Full table scans scale poorly, so use pgvector indexes and the like
  • Using more documents in the generation step, or fetching additional documents via keyword matching and so on, also helps
  • Output quality can be improved by tuning the generation prompt and trying various LLMs

See also

  • LangChain
  • Pykoi - A UI library for collecting data & feedback for LLMs
  • PEFT
  • MUVERA - Makes multi-vector search as fast as single-vector search

Favorite site

Articles

References


  1. Guidebook_to_the_State-of-the-Art_Embeddings_and_Information_Retrieval_-_Ko.pdf ↩

author Zer0 sitemap