src/retriever.py

"""
retriever.py
------------
Lightweight TF-IDF + cosine-similarity retriever used to ground the LLM's
answers in the warehouse knowledge base (simple RAG pipeline). Kept
dependency-light (scikit-learn only) so it trains instantly and runs fast
on the free CPU tier of Hugging Face Spaces.
"""

from dataclasses import dataclass
from typing import List

import joblib
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

from src.knowledge_base import KNOWLEDGE_BASE


@dataclass
class RetrievedDoc:
    id: str
    title: str
    text: str
    score: float


class KBRetriever:
    def __init__(self):
        self.vectorizer = TfidfVectorizer(stop_words="english", ngram_range=(1, 2))
        self.doc_ids = [d["id"] for d in KNOWLEDGE_BASE]
        self.docs = {d["id"]: d for d in KNOWLEDGE_BASE}
        corpus = [d["title"] + ". " + d["text"] for d in KNOWLEDGE_BASE]
        self.doc_matrix = self.vectorizer.fit_transform(corpus)

    def retrieve(self, query: str, k: int = 2) -> List[RetrievedDoc]:
        q_vec = self.vectorizer.transform([query])
        sims = cosine_similarity(q_vec, self.doc_matrix).flatten()
        top_idx = np.argsort(sims)[::-1][:k]
        results = []
        for idx in top_idx:
            doc_id = self.doc_ids[idx]
            d = self.docs[doc_id]
            results.append(RetrievedDoc(id=doc_id, title=d["title"], text=d["text"], score=float(sims[idx])))
        return results

    def save(self, path: str):
        joblib.dump(self.vectorizer, path)

    @staticmethod
    def top1_id(query: str, retriever: "KBRetriever") -> str:
        return retriever.retrieve(query, k=1)[0].id