src/retrieval.py

"""
retrieval.py — End-to-end retrieval pipeline.

Embeds the query using the same model as the documents,
searches FAISS for top-k nearest neighbors,
and applies a relevance threshold before passing to generation.
"""

import logging
import numpy as np
from src.embeddings import embed_texts, search_index
from src.utils import Timer

logger = logging.getLogger("enterprise-rag.retrieval")

# Below this cosine similarity, retrieved chunks are considered too weak.
# The system will return a fallback response instead of risking hallucination.
RELEVANCE_THRESHOLD = 0.35


def retrieve_relevant_chunks(
    query: str,
    chunks: list,
    faiss_index,
    top_k: int = 5,
) -> dict:
    """
    Retrieve the most relevant document chunks for a user query.

    Args:
        query       — user's natural language question
        chunks      — list of chunk text strings
        faiss_index — built FAISS index from embeddings.py
        top_k       — number of chunks to retrieve

    Returns dict:
        retrieved_chunks     — list of chunk text strings
        scores               — cosine similarity scores
        retrieval_latency_ms — time taken in ms
        is_relevant          — bool: top score above threshold
        warning              — message if quality is low, else None
    """
    result = {
        "retrieved_chunks": [],
        "scores": [],
        "retrieval_latency_ms": 0,
        "is_relevant": False,
        "warning": None,
    }

    if not chunks or faiss_index is None:
        result["warning"] = "No documents indexed. Please upload a PDF first."
        return result

    if not query or not query.strip():
        result["warning"] = "Empty query received."
        return result

    with Timer() as t:
        query_embedding = embed_texts([query.strip()])[0]
        scores_raw, indices = search_index(faiss_index, query_embedding, top_k)

    result["retrieval_latency_ms"] = round(t.elapsed_ms, 2)

    if len(indices) == 0:
        result["warning"] = "FAISS returned no results."
        return result

    retrieved = []
    scores_out = []

    for idx, score in zip(indices, scores_raw):
        if 0 <= idx < len(chunks):
            chunk_text = chunks[idx]["text"] if isinstance(chunks[idx], dict) else chunks[idx]
            retrieved.append(chunk_text)
            scores_out.append(float(score))

    result["retrieved_chunks"] = retrieved
    result["scores"] = scores_out
    result["is_relevant"] = bool(scores_out and scores_out[0] >= RELEVANCE_THRESHOLD)

    if not result["is_relevant"] and scores_out:
        result["warning"] = (
            f"Top similarity score is {scores_out[0]:.3f} — below threshold "
            f"({RELEVANCE_THRESHOLD}). The document may not contain an answer "
            f"to this question."
        )

    logger.info(
        f"Retrieved {len(retrieved)} chunks in {t.elapsed_ms:.1f}ms | "
        f"Top score: {scores_out[0]:.4f if scores_out else 'N/A'}"
    )

    return result