Step 6: Query — Similarity Search + Gemini Answer

# src/retrieve.py
"""Find the K most similar chunks to a question."""
 
from .db import get_conn
from .embeddings import embed
 
 
def retrieve(question: str, k: int = 5) -> list[dict]:
    """Return the K chunks closest to the question."""
    [query_vec] = embed([question], task="RETRIEVAL_QUERY")
    query_str = "[" + ",".join(f"{x:.7f}" for x in query_vec) + "]"
 
    with get_conn() as conn:
        cur = conn.cursor()
        cur.execute(
            """
            SELECT
              id,
              source,
              chunk_index,
              content,
              embedding <=> %s AS distance
            FROM chunks
            ORDER BY embedding <=> %s
            LIMIT %s
            """,
            (query_str, query_str, k),
        )
        rows = cur.fetchall()
 
    return [
        {
            "id": r[0],
            "source": r[1],
            "chunk_index": r[2],
            "content": r[3],
            "distance": float(r[4]),
        }
        for r in rows
    ]

# scripts/test_retrieve.py
from src.retrieve import retrieve
 
results = retrieve("What is retrieval-augmented generation?", k=3)
for r in results:
    print(f"[{r['distance']:.3f}] {r['source']}#{r['chunk_index']}")
    print(r['content'][:200], "\n")

uv run python -m scripts.test_retrieve

# src/generate.py
"""Ask Gemini to answer a question using only the provided context."""
 
from google import genai
from google.genai.types import GenerateContentConfig
 
from .config import load_config
 
_cfg = load_config()
_client = genai.Client(vertexai=True, project=_cfg.project, location=_cfg.region)
 
MODEL = "gemini-flash-latest"
 
SYSTEM_PROMPT = """You answer questions from the provided context.
Rules:
- Use only the context. Do not use outside knowledge.
- If the context does not contain the answer, reply exactly: "I don't know based on the provided documents."
- Quote short phrases from the context when helpful.
- Cite sources at the end as: Sources: <filename>#<chunk_index>, ...
"""
 
 
def generate_answer(question: str, chunks: list[dict]) -> str:
    context = "\n\n---\n\n".join(
        f"[{c['source']}#{c['chunk_index']}]\n{c['content']}" for c in chunks
    )
 
    prompt = (
        f"Context:\n{context}\n\n"
        f"Question: {question}\n\n"
        "Answer:"
    )
 
    response = _client.models.generate_content(
        model=MODEL,
        contents=prompt,
        config=GenerateContentConfig(
            system_instruction=SYSTEM_PROMPT,
            temperature=0.1,
        ),
    )
    return response.text or "I don't know based on the provided documents."

# src/ask.py
"""End-to-end: question → retrieved chunks → grounded answer."""
 
from .generate import generate_answer
from .retrieve import retrieve
 
 
def ask(question: str, k: int = 5) -> dict:
    chunks = retrieve(question, k=k)
    answer = generate_answer(question, chunks)
    return {
        "question": question,
        "answer": answer,
        "sources": [
            {"source": c["source"], "chunk_index": c["chunk_index"], "distance": c["distance"]}
            for c in chunks
        ],
    }

# scripts/test_ask.py
import json
from src.ask import ask
 
print(json.dumps(ask("What is retrieval-augmented generation?"), indent=2))

uv run python -m scripts.test_ask

{
  "question": "What is retrieval-augmented generation?",
  "answer": "Retrieval-augmented generation (RAG) is a technique that... [...] Sources: rag-wikipedia.txt#0, rag-wikipedia.txt#1",
  "sources": [
    {"source": "rag-wikipedia.txt", "chunk_index": 0, "distance": 0.142},
    {"source": "rag-wikipedia.txt", "chunk_index": 1, "distance": 0.218}
  ]
}