src/data/SemanticCache.py

import logging
import hashlib

from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings

from src.data.chroma_config import COLLECTION_NAME, PERSIST_DIRECTORY, ensure_persist_dir
from src.core.settings import get_settings

logger = logging.getLogger(__name__)


class SemanticCache:
    def __init__(
        self,
        persist_directory: str = PERSIST_DIRECTORY,
        collection_name: str = COLLECTION_NAME,
    ):
        """
        Initializes the Semantic Caching layer utilizing ChromaDB.
        """
        self.persist_directory = persist_directory
        settings = get_settings()
        self.embeddings = OpenAIEmbeddings(model=settings.models.embedding_model)

        ensure_persist_dir(self.persist_directory)

        self.vector_store = Chroma(
            collection_name=collection_name,
            embedding_function=self.embeddings,
            persist_directory=self.persist_directory,
        )

    def _doc_id(self, normalized_query: str) -> str:
        h = hashlib.sha256()
        h.update(normalized_query.encode("utf-8"))
        return h.hexdigest()

    def check_cache(self, query: str, threshold: float = 0.70):
        """
        Check the vector store for a semantically similar query.
        Returns the cached response if a match is found above the similarity threshold.
        """
        # Normalize the query to prevent case/whitespace from skewing the embeddings
        query_normalized = query.lower().strip()
        where = {"namespace": "semantic_cache"}
        try:
            results = self.vector_store.similarity_search_with_relevance_scores(
                query_normalized, k=1, filter=where
            )
        except TypeError:
            # Older wrappers use `where` not `filter`.
            results = self.vector_store.similarity_search_with_relevance_scores(
                query_normalized, k=1, where=where
            )  # type: ignore[call-arg]
        except Exception:
            logger.exception("SemanticCache: check_cache failed")
            return None

        if results:
            doc, score = results[0]
            if score >= threshold:
                logger.info(f"Semantic Match Score: {score:.3f} greater than threshold {threshold:.3f}")
                meta = getattr(doc, "metadata", {}) or {}
                return meta.get("response")
        return None

    def save_to_cache(self, query: str, response: str):
        """
        Save the given query and response combo to the semantic cache.
        """
        query_normalized = query.lower().strip()
        doc_id = self._doc_id(query_normalized)

        # Best-effort de-dupe: keep the first cached answer for a given normalized query.
        try:
            existing = set(self.vector_store.get(ids=[doc_id]).get("ids", []))
        except Exception:
            existing = set()
        if doc_id in existing:
            return

        self.vector_store.add_texts(
            texts=[query_normalized],
            metadatas=[{"namespace": "semantic_cache", "response": response}],
            ids=[doc_id],
        )