RAG

RAG search example. This demo allows you to ask questions about an October 2024 snapshot of the Logfire documentation.

Demonstrates:

This is done by creating a database containing each section of the markdown documentation, then registering the search tool with the Pydantic AI agent.

Logic for extracting sections from markdown files and a JSON file with that data is available in this gist.

PostgreSQL with pgvector is used as the search database, the easiest way to download and run pgvector is using Docker:

bash

mkdir postgres-data
docker run --rm \
  -e POSTGRES_PASSWORD=postgres \
  -p 54320:5432 \
  -v `pwd`/postgres-data:/var/lib/postgresql/data \
  pgvector/pgvector:pg17

As with the SQL gen example, we run postgres on port 54320 to avoid conflicts with any other postgres instances you may have running. We also mount the PostgreSQL data directory locally to persist the data if you need to stop and restart the container.

With that running and dependencies installed and environment variables set, we can build the search database with (WARNING: this requires the OPENAI_API_KEY env variable and will calling the OpenAI embedding API around 300 times to generate embeddings for each section of the documentation):

bash

python/uv-run -m pydantic_ai_examples.rag build

(Note building the database doesn't use Pydantic AI right now, instead it uses the OpenAI SDK directly.)

!!! note "Embedding model and index schema" This example uses text-embedding-3-small for documents and queries, storing its 1,536-dimensional output in a vector(1536) column. If you change the model or dimensions, stop PostgreSQL, delete the example's postgres-data directory (this removes the entire local example database), update DB_SCHEMA if needed, then restart PostgreSQL and rerun build. pgvector's HNSW vector index supports up to 2,000 dimensions.

You can then ask the agent a question with:

bash

python/uv-run -m pydantic_ai_examples.rag search "How do I configure logfire to work with FastAPI?"

Example Code

snippet {path="/examples/pydantic_ai_examples/rag.py"}