Spaces:
Sleeping
Sleeping
| """ | |
| src/generation/semantic_cache.py | |
| ================================= | |
| In-memory Semantic Cache for RAG query results. | |
| Uses cosine similarity over NumPy matrices to find semantically identical questions. | |
| """ | |
| import numpy as np | |
| from typing import Tuple, List, Dict, Optional | |
| from src.utils.logger import logger | |
| class SemanticCache: | |
| """ | |
| Singleton Semantic Cache. | |
| Max 500 overall queries. FIFO eviction. | |
| """ | |
| _instance = None | |
| def __new__(cls, *args, **kwargs): | |
| if cls._instance is None: | |
| cls._instance = super().__new__(cls) | |
| cls._instance._initialized = False | |
| return cls._instance | |
| def __init__(self, max_size: int = 500, threshold: float = 0.95): | |
| if self._initialized: | |
| return | |
| self.max_size = max_size | |
| self.threshold = threshold | |
| # We store queries in a parallel array structure | |
| # embeddings: np.ndarray of shape (N, 768) | |
| self.embeddings: np.ndarray = np.empty((0, 768), dtype=np.float32) | |
| # payloads parallel to embeddings rows | |
| # [{"company": str, "fy": str, "text": str, "sources": list}, ...] | |
| self.payloads: List[Dict] = [] | |
| self._initialized = True | |
| logger.info(f"Semantic Cache initialized (max_size={max_size}, threshold={threshold})") | |
| def find_match(self, query_emb: np.ndarray, company: str, fiscal_year: Optional[str]) -> Optional[Tuple[str, list]]: | |
| """ | |
| Finds a cached response using cosine similarity. | |
| Query embedding must be L2 normalized (BGE embedder does this). | |
| """ | |
| if len(self.payloads) == 0: | |
| return None | |
| # query_emb shape: (768,) | |
| # self.embeddings shape: (N, 768) | |
| # Cosine similarity is just the dot product since vectors are L2 normalized | |
| similarities = np.dot(self.embeddings, query_emb) | |
| # Get the index of the highest similarity | |
| best_idx = int(np.argmax(similarities)) | |
| best_score = similarities[best_idx] | |
| if best_score >= self.threshold: | |
| # Check hard filters (company and FY must match exactly) | |
| p = self.payloads[best_idx] | |
| if p["company"] == company and p["fy"] == fiscal_year: | |
| logger.info(f"Semantic Cache HIT (score={best_score:.4f})") | |
| return p["text"], p["sources"] | |
| return None | |
| def add(self, query_emb: np.ndarray, company: str, fiscal_year: Optional[str], text: str, sources: list): | |
| """Adds a new response to the cache, evicting the oldest if full.""" | |
| # Check size and apply FIFO eviction | |
| if len(self.payloads) >= self.max_size: | |
| # Remove oldest (index 0) | |
| self.embeddings = self.embeddings[1:] | |
| self.payloads.pop(0) | |
| # Append new embedding | |
| query_emb_2d = query_emb.reshape(1, -1) | |
| if self.embeddings.shape[0] == 0: | |
| self.embeddings = query_emb_2d | |
| else: | |
| self.embeddings = np.vstack([self.embeddings, query_emb_2d]) | |
| # Append payload | |
| self.payloads.append({ | |
| "company": company, | |
| "fy": fiscal_year, | |
| "text": text, | |
| "sources": sources | |
| }) | |
| logger.info(f"Added to Semantic Cache (size={len(self.payloads)}/{self.max_size})") | |