Harshavard21/FinRAG

1"""2src/generation/semantic_cache.py3=================================4In-memory Semantic Cache for RAG query results.5Uses cosine similarity over NumPy matrices to find semantically identical questions.6"""7 8import numpy as np9from typing import Tuple, List, Dict, Optional10from src.utils.logger import logger11 12class SemanticCache:13    """14    Singleton Semantic Cache.15    Max 500 overall queries. FIFO eviction.16    """17    _instance = None18 19    def __new__(cls, *args, **kwargs):20        if cls._instance is None:21            cls._instance = super().__new__(cls)22            cls._instance._initialized = False23        return cls._instance24 25    def __init__(self, max_size: int = 500, threshold: float = 0.95):26        if self._initialized:27            return28            29        self.max_size = max_size30        self.threshold = threshold31        32        # We store queries in a parallel array structure33        # embeddings: np.ndarray of shape (N, 768)34        self.embeddings: np.ndarray = np.empty((0, 768), dtype=np.float32)35        36        # payloads parallel to embeddings rows37        # [{"company": str, "fy": str, "text": str, "sources": list}, ...]38        self.payloads: List[Dict] = []39        40        self._initialized = True41        logger.info(f"Semantic Cache initialized (max_size={max_size}, threshold={threshold})")42 43    def find_match(self, query_emb: np.ndarray, company: str, fiscal_year: Optional[str]) -> Optional[Tuple[str, list]]:44        """45        Finds a cached response using cosine similarity.46        Query embedding must be L2 normalized (BGE embedder does this).47        """48        if len(self.payloads) == 0:49            return None50 51        # query_emb shape: (768,)52        # self.embeddings shape: (N, 768)53        # Cosine similarity is just the dot product since vectors are L2 normalized54        similarities = np.dot(self.embeddings, query_emb)55        56        # Get the index of the highest similarity57        best_idx = int(np.argmax(similarities))58        best_score = similarities[best_idx]59        60        if best_score >= self.threshold:61            # Check hard filters (company and FY must match exactly)62            p = self.payloads[best_idx]63            if p["company"] == company and p["fy"] == fiscal_year:64                logger.info(f"Semantic Cache HIT (score={best_score:.4f})")65                return p["text"], p["sources"]66                67        return None68 69    def add(self, query_emb: np.ndarray, company: str, fiscal_year: Optional[str], text: str, sources: list):70        """Adds a new response to the cache, evicting the oldest if full."""71        # Check size and apply FIFO eviction72        if len(self.payloads) >= self.max_size:73            # Remove oldest (index 0)74            self.embeddings = self.embeddings[1:]75            self.payloads.pop(0)76            77        # Append new embedding78        query_emb_2d = query_emb.reshape(1, -1)79        if self.embeddings.shape[0] == 0:80            self.embeddings = query_emb_2d81        else:82            self.embeddings = np.vstack([self.embeddings, query_emb_2d])83            84        # Append payload85        self.payloads.append({86            "company": company,87            "fy": fiscal_year,88            "text": text,89            "sources": sources90        })91        logger.info(f"Added to Semantic Cache (size={len(self.payloads)}/{self.max_size})")92