Document retriever

The retrieval half of RAG. A set of document chunks is indexed in the vector

README Try it Source

Showcase — Document retriever

The retrieval half of RAG. A set of document chunks is indexed in the vector store once; your question returns the top passages by cosine similarity, each tagged with its source id so the result is traceable. There's no generation here — just the part that finds the right context to hand a model next chapter.

Run

bash bootstrap-secrets.sh              # reads ../../../../.env, writes secrets/
docker compose up --build              # default: PROVIDER=openai

Open http://localhost:3000. Run against Gemini with PROVIDER=gemini docker compose up --build.

What's where

backend/store.py — the shared in-memory VectorStore (add + search).
backend/ai_openai.py / backend/ai_gemini.py — the doc chunks, the embed function, and the top-k retrieval.
frontend/app/page.tsx — question box + ranked passages.

Stop

docker compose down

Run locally

Download the project as a ZIP and run it with Docker. Brings up a FastAPI backend + Next.js frontend on localhost:3000.

Download doc-retriever.zip

unzip doc-retriever.zip
cd doc-retriever
bash bootstrap-secrets.sh   # one-time: pulls API keys into ./secrets
docker compose up --build   # default provider: openai
# or:  PROVIDER=gemini docker compose up --build

Type some input, pick a provider, and run the same code shown in Source against the live API. Sign-in required.

Input

Provider

Output

The same modules the Run button hits. The whole project (frontend, Dockerfile, compose) is in the ZIP under README.

`backend/ai_openai.py`

"""Showcase 1 (OpenAI): a document retriever over a chunked knowledge base.

This is the retrieval half of RAG (next chapter): index a set of document
chunks in the vector store, then return the top passages for a question — with
their source ids, so the answer is traceable. No generation yet; just the part
that finds the right context.
"""
from openai import OpenAI

from store import VectorStore

_client = OpenAI()

_MODEL = "text-embedding-3-small"

# A handful of chunks from a fictional product's docs. In a real system these
# come from splitting real documents; the store doesn't care where they're from.
_KB = [
    {"id": "plans", "text": "The Free plan includes 3 projects. Pro is $12/month for unlimited projects and priority support."},
    {"id": "export", "text": "Export your data any time as CSV or JSON from Settings > Data. Exports include all projects."},
    {"id": "sso", "text": "Single sign-on (SAML) is available on the Enterprise plan and is configured under Admin > Security."},
    {"id": "api", "text": "The REST API is rate-limited to 100 requests per minute. Keys are created under Settings > API."},
    {"id": "delete", "text": "Deleting a project is permanent and cannot be undone. Deleted data is purged after 30 days."},
    {"id": "invite", "text": "Invite teammates from Members > Invite. Invited users get Editor access by default."},
    {"id": "billing", "text": "Plan changes are prorated. Downgrades take effect at the end of the current billing cycle."},
    {"id": "mobile", "text": "The mobile app supports viewing and commenting; editing requires the web app."},
]

_store: VectorStore | None = None


def _embed(texts: list[str]) -> list[list[float]]:
    return [item.embedding for item in _client.embeddings.create(model=_MODEL, input=texts).data]


def _get_store() -> VectorStore:
    global _store
    if _store is None:  # build the index once
        _store = VectorStore(_embed)
        _store.add(_KB)
    return _store


def run(question: str) -> str:
    store = _get_store()
    hits = store.search(question.strip(), k=4)
    rows = [f"{score:.3f}  [{item['id']}]  {item['text']}" for score, item in hits]
    return f"Top passages for your question (index of {store.size()} chunks):\n\n" + "\n".join(rows)

`backend/ai_gemini.py`

"""Showcase 1 (Gemini): a document retriever over a chunked knowledge base.

Same store, same chunks, Gemini's embedding model.
"""
import os

from google import genai

from store import VectorStore

_client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])

_MODEL = "gemini-embedding-001"

_KB = [
    {"id": "plans", "text": "The Free plan includes 3 projects. Pro is $12/month for unlimited projects and priority support."},
    {"id": "export", "text": "Export your data any time as CSV or JSON from Settings > Data. Exports include all projects."},
    {"id": "sso", "text": "Single sign-on (SAML) is available on the Enterprise plan and is configured under Admin > Security."},
    {"id": "api", "text": "The REST API is rate-limited to 100 requests per minute. Keys are created under Settings > API."},
    {"id": "delete", "text": "Deleting a project is permanent and cannot be undone. Deleted data is purged after 30 days."},
    {"id": "invite", "text": "Invite teammates from Members > Invite. Invited users get Editor access by default."},
    {"id": "billing", "text": "Plan changes are prorated. Downgrades take effect at the end of the current billing cycle."},
    {"id": "mobile", "text": "The mobile app supports viewing and commenting; editing requires the web app."},
]

_store: VectorStore | None = None


def _embed(texts: list[str]) -> list[list[float]]:
    return [e.values for e in _client.models.embed_content(model=_MODEL, contents=texts).embeddings]


def _get_store() -> VectorStore:
    global _store
    if _store is None:
        _store = VectorStore(_embed)
        _store.add(_KB)
    return _store


def run(question: str) -> str:
    store = _get_store()
    hits = store.search(question.strip(), k=4)
    rows = [f"{score:.3f}  [{item['id']}]  {item['text']}" for score, item in hits]
    return f"Top passages for your question (index of {store.size()} chunks):\n\n" + "\n".join(rows)