retriever.py

"""Retrieve the most relevant chunks for a user query."""

from config import TOP_K
from embeddings import embed_query
from vector_store import load_index, search_index

# Cached index loaded once at first call
_index = None
_chunks = None


def get_index():
    global _index, _chunks
    if _index is None:
        _index, _chunks = load_index()
    return _index, _chunks


def retrieve(query: str, top_k: int = TOP_K) -> list[dict]:
    """Return top_k relevant chunks for the given query."""
    index, chunks = get_index()
    query_vec = embed_query(query)
    results = search_index(index, chunks, query_vec, top_k)
    return results


def reset_index_cache():
    """Force reload of the index on next retrieval (useful after rebuilding)."""
    global _index, _chunks
    _index = None
    _chunks = None