RAG Step 3: Set Up pgvector

pgvector adds vector search to PostgreSQL. Instead of SQLite, we'll use PostgreSQL with the pgvector extension.

Why PostgreSQL + pgvector?

Step 1: Start PostgreSQL with Docker Compose

Create advanced/docker-compose.rag.yml:

services:
  db:
    image: pgvector/pgvector:pg16
    environment:
      POSTGRES_USER: aiproject
      POSTGRES_PASSWORD: aiproject
      POSTGRES_DB: aiproject
    ports:
      - "5432:5432"
    volumes:
      - pgvector_data:/var/lib/postgresql/data
    healthcheck:
      test: ["CMD-SHELL", "pg_isready -U aiproject"]
      interval: 5s
      timeout: 5s
      retries: 5

  app:
    build: ../backend
    ports:
      - "5000:5000"
    volumes:
      - ../backend:/app
      - sqlite_data:/app/instance
    env_file:
      - ../backend/app/.env
    environment:
      DATABASE_URL: postgresql://aiproject:aiproject@db:5432/aiproject
    depends_on:
      db:
        condition: service_healthy

volumes:
  pgvector_data:
  sqlite_data:

This file adds a PostgreSQL service alongside your Flask app. The app service connects to the database using the DATABASE_URL environment variable.

Start both services:

cd aiproject/advanced
docker compose -f docker-compose.rag.yml up

Step 2: Install Dependencies

pip install psycopg2-binary sentence-transformers

Step 3: Add the Document Model

Create a new model for storing document chunks and their embeddings. Add to backend/app/models.py:

from pgvector.sqlalchemy import Vector


class Document(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    title = db.Column(db.String(500), nullable=False)
    content = db.Column(db.Text, nullable=False)
    source = db.Column(db.String(500), default="")
    category = db.Column(db.String(100), default="general")
    embedding = db.Column(Vector(384), nullable=False)

    def to_dict(self):
        return {
            "id": self.id,
            "title": self.title,
            "content": self.content[:200],
            "source": self.source,
            "category": self.category,
        }

The Vector(384) column stores embeddings. The <=> operator will search by cosine distance.

Also add pgvector to requirements.txt:

pgvector
psycopg2-binary
sentence-transformers

Step 4: Create the Database Tables

The embeddings table needs to be created. Run the app once with PostgreSQL connected:

cd aiproject
docker compose -f advanced/docker-compose.rag.yml up --build

The app will create tables automatically (via db.create_all()). You can verify:

docker compose -f advanced/docker-compose.rag.yml exec db psql -U aiproject -c "\dt"

You should see tables including document with an embedding column.

Step 5: Test the Connection and Model

Run this quick check:

# test_pgvector.py
from app import create_app, db
from app.models import Document

app = create_app()
with app.app_context():
    db.create_all()
    count = Document.query.count()
    print(f"Connection OK. Documents in DB: {count}")

Run it:

DATABASE_URL="postgresql://aiproject:aiproject@localhost:5432/aiproject" python test_pgvector.py
# → Connection OK. Documents in DB: 0

Write an Automated Test

Add a test that verifies the Document model works with pgvector. Create advanced/rag/tests/test_knowledge_base.py:

import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from sample_data import DOCUMENTS
from knowledge_base import chunk_text, get_all_chunks


def test_chunk_text_creates_chunks():
    text = "word " * 500
    chunks = chunk_text(text, chunk_size=100, overlap=20)
    assert len(chunks) >= 5
    assert all(isinstance(c, str) for c in chunks)


def test_chunk_text_short_text():
    text = "short text"
    chunks = chunk_text(text, chunk_size=100, overlap=20)
    assert len(chunks) == 1
    assert chunks[0] == "short text"


def test_get_all_chunks_returns_list():
    chunks = get_all_chunks(DOCUMENTS[:3])
    assert len(chunks) > 0
    assert all("title" in c for c in chunks)
    assert all("content" in c for c in chunks)

Run them:

cd aiproject
python -m pytest advanced/rag/tests/ -v

These tests run without a database — they test the chunking logic in isolation.

Write a quick test:

# test_pgvector.py
from app import create_app, db
from app.models import Document

app = create_app()
with app.app_context():
    db.create_all()
    count = Document.query.count()
    print(f"Connection OK. Documents in DB: {count}")

Run it:

DATABASE_URL="postgresql://aiproject:aiproject@localhost:5432/aiproject" python test_pgvector.py
# → Connection OK. Documents in DB: 0

How It Works

┌─────────────────────────────────────────────────────────┐
│                    pgvector                              │
│                                                         │
│  ┌─────────────────────────────────────────────────┐    │
│  │  document table                                  │    │
│  │  id │ title │ content │ source │ embedding (384) │    │
│  ├─────┼───────┼─────────┼────────┼─────────────────┤    │
│  │  1  │ Tut…  │ Clark…  │ url…   │ [0.04, -0.13,…] │    │
│  │  2  │ Fin…  │ Reso…   │ url…   │ [0.21, 0.07,…]  │    │
│  │  3  │ Dis…  │ The…    │ url…   │ [-0.09, 0.33,…] │    │
│  └─────────────────────────────────────────────────┘    │
│                                                         │
│  Query: SELECT * FROM document                          │
│         ORDER BY embedding <=> :query_vec               │
│         LIMIT 5                                         │
└─────────────────────────────────────────────────────────┘