RAG Basics

uv venv && source .venv/bin/activate
uv pip install "openai>=1.60" "chromadb>=0.5" "tiktoken>=0.8"
export OPENAI_API_KEY="sk-..."

def chunk_text(text: str, size: int = 500, overlap: int = 50) -> list[str]:
    chunks = []
    start = 0
    while start < len(text):
        chunks.append(text[start:start + size])
        start += size - overlap
    return chunks

import chromadb
from openai import OpenAI
 
client = OpenAI()
chroma = chromadb.Client()
collection = chroma.create_collection("docs")
 
texts = ["FastAPI is a Python web framework.", "pytest runs Python tests."]
emb = client.embeddings.create(model="text-embedding-3-small", input=texts)
collection.add(ids=["1","2"], documents=texts, embeddings=[e.embedding for e in emb.data])

query = "How do I test Python?"
q_emb = client.embeddings.create(model="text-embedding-3-small", input=[query])
results = collection.query(query_embeddings=[q_emb.data[0].embedding], n_results=2)
print(results["documents"])

context = "\n".join(results["documents"][0])
prompt = f"Answer using ONLY this context:\n{context}\n\nQuestion: {query}"

response = client.chat.completions.create(
    model="gpt-4o-mini",
    messages=[{"role": "user", "content": prompt}],
    temperature=0,
)
print(response.choices[0].message.content)

def rag_answer(query: str, collection, k: int = 3) -> str:
    q_emb = client.embeddings.create(model="text-embedding-3-small", input=[query])
    hits = collection.query(query_embeddings=[q_emb.data[0].embedding], n_results=k)
    context = "\n---\n".join(hits["documents"][0])
    resp = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": "Answer from context only. Say 'I don't know' if not found."},
            {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {query}"},
        ],
        temperature=0,
    )
    return resp.choices[0].message.content or ""

for i, (doc, meta) in enumerate(zip(hits["documents"][0], hits["metadatas"][0])):
    print(f"[{i+1}] {doc[:80]}...")

collection.query(
    query_embeddings=[q_emb.data[0].embedding],
    n_results=3,
    where={"source": "docs"},
)

def recall_at_k(query: str, expected_text: str, k: int = 5) -> bool:
    hits = collection.query(query_embeddings=[embed(query)], n_results=k)
    return any(expected_text in doc for doc in hits["documents"][0])