tools/retrieval_tool.py

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

from dotenv import load_dotenv

load_dotenv()

from chroma import get_collection  # noqa: E402


def retrieve_chunks(query: str, top_k: int = 5) -> list[str]:
    """Search the vector store for document chunks most relevant to the query.

    Embeds the query using the same OpenAI embedding model used during ingestion,
    performs a cosine-similarity search against the ChromaDB collection, and
    returns the top-k matching text chunks formatted with source metadata so
    the caller can cite the origin of each passage.

    Args:
        query:  The natural-language question or search string to look up.
        top_k:  Maximum number of chunks to return (default 5).

    Returns:
        A list of strings, each formatted as:
            "[source: <filename>, page: <n>, chunk: <i>]\\n<chunk text>"
        Page number is omitted for DOCX files where it is unavailable.
    """
    collection = get_collection()
    count = collection.count()
    if count == 0:
        return ["No documents have been indexed yet. Please ingest documents first."]

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

    chunks: list[str] = []
    for doc, meta in zip(results["documents"][0], results["metadatas"][0]):
        source = meta.get("source_file", "unknown")
        chunk_idx = meta.get("chunk_index", "?")
        page = meta.get("page_number")
        if page is not None:
            header = f"[source: {source}, page: {page}, chunk: {chunk_idx}]"
        else:
            header = f"[source: {source}, chunk: {chunk_idx}]"
        chunks.append(f"{header}\n{doc}")

    return chunks