src/retriever.py

# ============================================================
# FILE: src/retriever.py
# ============================================================
# PURPOSE:
# Combine embedding generation and vector search.
#
# RAG retrieval flow:
#
# user question
# -> query embedding
# -> vector database search
# -> top matching chunks
#
# AI ENGINEER PRODUCTION TIP:
# Always inspect retrieved chunks when debugging RAG quality.
# If retrieval is wrong, the LLM cannot produce a grounded answer.
# ============================================================

from typing import Any, Dict, List

from src.embeddings import EmbeddingModel
from src.vector_store import ChromaVectorStore


class Retriever:
    """
    Retrieves relevant chunks for a user question.
    """

    def __init__(
        self,
        embedding_model: EmbeddingModel,
        vector_store: ChromaVectorStore,
    ) -> None:
        self.embedding_model = embedding_model
        self.vector_store = vector_store

    def retrieve(
        self,
        question: str,
        top_k: int,
    ) -> List[Dict[str, Any]]:
        """
        Retrieve top-k relevant chunks.

        top_k:
        - number of chunks to return
        - common values: 3, 4, 5, 8
        """

        question = question.strip()

        if not question:
            raise ValueError("Question cannot be empty.")

        query_embedding = self.embedding_model.embed_query(question)

        return self.vector_store.query(
            query_embedding=query_embedding,
            top_k=top_k,
        )