06_retrieve_context.py

"""
06_retrieve_context.py
------------------------
Stage 6 of the RAG pipeline: CONTEXT RETRIEVAL.

Given a user question, embeds the question with the same embedding model
used to build the vector store, then queries ChromaDB for the most
semantically similar chunks. These chunks become the "context" that stage
7 (prompting) feeds to the LLM.

Run directly to try a retrieval query against the existing store:
    python 06_retrieve_context.py
"""

from __future__ import annotations

import importlib.util
import sys
import os
from dataclasses import dataclass
from typing import List

_THIS_DIR = os.path.dirname(os.path.abspath(__file__))

DEFAULT_TOP_K = 4


def _import_module_from_file(module_name: str, filename: str):
    path = os.path.join(_THIS_DIR, filename)
    spec = importlib.util.spec_from_file_location(module_name, path)
    module = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = module  # required so dataclasses can resolve cls.__module__
    spec.loader.exec_module(module)
    return module


_store_module = _import_module_from_file("store_stage", "05_create_chroma_store.py")
_vector_module = _import_module_from_file("vector_stage", "04_vector_representation.py")


@dataclass
class RetrievedChunk:
    """A single chunk retrieved from the vector store for a query."""

    chunk_id: str
    text: str
    metadata: dict
    distance: float


def retrieve_context(
    query: str,
    top_k: int = DEFAULT_TOP_K,
    persist_dir: str = _store_module.CHROMA_PERSIST_DIR,
    collection_name: str = _store_module.CHROMA_COLLECTION_NAME,
) -> List[RetrievedChunk]:
    """
    Retrieve the `top_k` most relevant chunks for `query` from the Chroma store.

    Returns an empty list if the store has no data yet, so callers can
    detect retrieval failure and avoid hallucinating an answer.
    """
    if not query or not query.strip():
        return []

    collection = _store_module.load_collection(persist_dir, collection_name)

    if collection.count() == 0:
        return []

    query_embedding = _vector_module.embed_texts([query])[0]

    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=min(top_k, collection.count()),
        include=["documents", "metadatas", "distances"],
    )

    retrieved: List[RetrievedChunk] = []
    ids = results.get("ids", [[]])[0]
    documents = results.get("documents", [[]])[0]
    metadatas = results.get("metadatas", [[]])[0]
    distances = results.get("distances", [[]])[0]

    for chunk_id, text, metadata, distance in zip(ids, documents, metadatas, distances):
        retrieved.append(
            RetrievedChunk(
                chunk_id=chunk_id,
                text=text,
                metadata=metadata or {},
                distance=distance,
            )
        )

    return retrieved


if __name__ == "__main__":
    sample_query = "What is this project about?"
    results = retrieve_context(sample_query)

    if not results:
        print(
            "No results found. Make sure you've built the Chroma store first:\n"
            "    python 05_create_chroma_store.py"
        )
    else:
        print(f"Top {len(results)} chunk(s) retrieved for query: '{sample_query}'\n")
        for i, chunk in enumerate(results, start=1):
            preview = chunk.text[:150].replace("\n", " ")
            source = chunk.metadata.get("filename", "unknown")
            print(f"{i}. [{source}] (distance={chunk.distance:.4f})")
            print(f"   {preview}...\n")