Building a RAG chatbot with pgvector on StackBlaze

CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE documents (
  id          BIGSERIAL PRIMARY KEY,
  source      TEXT NOT NULL,
  chunk_index INTEGER NOT NULL,
  content     TEXT NOT NULL,
  embedding   VECTOR(1536) NOT NULL,
  created_at  TIMESTAMPTZ DEFAULT NOW()
);

-- IVFFlat index for approximate nearest-neighbor search
CREATE INDEX ON documents USING ivfflat (embedding vector_cosine_ops)
  WITH (lists = 100);

import os
import psycopg2
from openai import OpenAI
from pathlib import Path

client = OpenAI(api_key=os.environ['OPENAI_API_KEY'])
conn = psycopg2.connect(os.environ['DATABASE_URL'])

CHUNK_SIZE = 800
CHUNK_OVERLAP = 100

def chunk_text(text: str) -> list[str]:
    chunks = []
    start = 0
    while start < len(text):
        end = start + CHUNK_SIZE
        chunks.append(text[start:end])
        start += CHUNK_SIZE - CHUNK_OVERLAP
    return chunks

def embed(texts: list[str]) -> list[list[float]]:
    response = client.embeddings.create(model='text-embedding-3-small', input=texts)
    return [item.embedding for item in response.data]

def ingest_file(path: Path) -> None:
    text = path.read_text(encoding='utf-8')
    chunks = chunk_text(text)
    embeddings = embed(chunks)
    with conn.cursor() as cur:
        for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)):
            cur.execute(
                'INSERT INTO documents (source, chunk_index, content, embedding) VALUES (%s, %s, %s, %s)',
                (str(path), i, chunk, embedding),
            )
    conn.commit()
    print(f'Ingested {len(chunks)} chunks from {path}')

for doc in Path('docs').rglob('*.md'):
    ingest_file(doc)

conn.close()

def retrieve(question: str, top_k: int = 5) -> list[dict]:
    embedding = client.embeddings.create(
        model='text-embedding-3-small', input=question
    ).data[0].embedding
    with conn.cursor() as cur:
        cur.execute(
            'SELECT source, chunk_index, content, 1 - (embedding <=> %s::vector) AS similarity FROM documents ORDER BY embedding <=> %s::vector LIMIT %s',
            (embedding, embedding, top_k),
        )
        rows = cur.fetchall()
    return [{'source': r[0], 'chunk_index': r[1], 'content': r[2], 'similarity': r[3]} for r in rows]

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from openai import OpenAI
import os
from .retrieval import retrieve

app = FastAPI()
client = OpenAI(api_key=os.environ['OPENAI_API_KEY'])

class QuestionRequest(BaseModel):
    question: str
    top_k: int = 5

@app.post('/ask')
def ask(req: QuestionRequest):
    chunks = retrieve(req.question, top_k=req.top_k)
    context = '\n\n---\n\n'.join(c['content'] for c in chunks)
    sources = list(dict.fromkeys(c['source'] for c in chunks))
    response = client.chat.completions.create(
        model='gpt-4o',
        messages=[
            {'role': 'system', 'content': 'You are a helpful assistant that answers questions based on the provided documentation. If the answer is not in the documentation, say so.'},
            {'role': 'user', 'content': f'Documentation:\n\n{context}\n\nQuestion: {req.question}'},
        ],
        temperature=0.2,
    )
    return {'answer': response.choices[0].message.content, 'sources': sources}