retriever.py

"""
Pinecone retriever — connects to the stem-embed index,
encodes queries with BGE-large, returns relevant NCERT chunks.
"""

from sentence_transformers import SentenceTransformer
from pinecone import Pinecone
from config import PINECONE_API_KEY, PINECONE_INDEX, EMBED_MODEL_NAME, BGE_QUERY_PREFIX, RETRIEVAL_TOP_K


# --- Load embedding model (CPU, ~1.3GB, loads once at startup) ---
_embed_model = SentenceTransformer(EMBED_MODEL_NAME, device="cpu")

# --- Connect to Pinecone ---
_pc    = Pinecone(api_key=PINECONE_API_KEY)
_index = _pc.Index(PINECONE_INDEX)


def search(query: str, top_k: int = RETRIEVAL_TOP_K) -> list[dict]:
    """
    Encode query with BGE prefix, retrieve top-k chunks from Pinecone.
    Returns list of dicts: {text, subject, class_level, chapter, source, score}
    """
    q_vec = _embed_model.encode(
        [BGE_QUERY_PREFIX + query],
        normalize_embeddings=True,
        convert_to_numpy=True,
    )[0].tolist()

    response = _index.query(vector=q_vec, top_k=top_k, include_metadata=True)

    results = []
    for m in response.matches:
        meta = m.metadata
        results.append({
            "text":        meta.get("text", ""),
            "subject":     meta.get("subject", ""),
            "class_level": meta.get("class_level", ""),
            "chapter":     meta.get("chapter", ""),
            "source":      meta.get("source", ""),
            "score":       round(m.score, 4),
        })
    return results


def format_context(results: list[dict]) -> str:
    """
    Format search results into a context string
    suitable for injection into the system prompt.
    """
    if not results:
        return "No relevant context found."

    blocks = []
    for i, r in enumerate(results, 1):
        label = f"Class {r['class_level']} {r['subject']} Ch.{r['chapter']}"
        blocks.append(f"[{i}. {label} | score={r['score']}]\n{r['text']}")

    return "\n\n".join(blocks)