semantic cache to support other tools

app.py

@@ -2,7 +2,7 @@ import dspy
 import gradio as gr
 from dotenv import load_dotenv
 
-from medirag.cache.local import SemanticCaching
+from medirag.cache.local import LocalSemanticCache
 from medirag.index.kdbai import KDBAIDailyMedIndexer
 from medirag.rag.qa import RAG, DailyMedRetrieve
 from medirag.rag.wf import RAGWorkflow
@@ -21,7 +21,7 @@ dspy.settings.configure(lm=turbo, rm=rm)
 # Set the LLM model
 Settings.llm = OpenAI(model="gpt-3.5-turbo")
 
-sm = SemanticCaching(
+sm = LocalSemanticCache(
     model_name="sentence-transformers/all-mpnet-base-v2", dimension=768, json_file="rag_test_cache.json"
 )
 

medirag/cache/abc.py

@@ -0,0 +1,21 @@
+from abc import ABC, abstractmethod
+
+
+class SemanticCache(ABC):
+    """
+    Abstract base class for semantic caching mechanisms.
+    """
+
+    @abstractmethod
+    def lookup(self, question: str, cosine_threshold: float):
+        """
+        Retrieve a response from the cache based on the question and cosine similarity threshold.
+        """
+        pass
+
+    @abstractmethod
+    def save(self, question: str, answer: str):
+        """
+        Save a question-answer pair to the cache.
+        """
+        pass

medirag/cache/local.py

@@ -5,14 +5,16 @@ from pydantic import BaseModel, ValidationError
 from sentence_transformers import SentenceTransformer
 from loguru import logger
 
+from medirag.cache.abc import SemanticCache
 
-class SemanticCache(BaseModel):
+
+class SemanticCacheModel(BaseModel):
     questions: list[str] = []
     embeddings: list[list[float]] = []
     response_text: list[str] = []
 
 
-class SemanticCaching:
+class LocalSemanticCache(SemanticCache):
     def __init__(
         self,
         model_name: str = "sentence-transformers/all-mpnet-base-v2",
@@ -24,14 +26,14 @@ class SemanticCaching:
         self.json_file = json_file
         self.vector_index = faiss.IndexFlatIP(self.dimension)
         self.encoder = SentenceTransformer(model_name)
-        self._cache = SemanticCache()  # Initialize with a default SemanticCache to avoid NoneType issues
+        self._cache = SemanticCacheModel()  # Initialize with a default SemanticCache to avoid NoneType issues
         self.load_cache()
 
     def load_cache(self) -> None:
         try:
             with open(self.json_file, "r") as file:
                 data = json.load(file)
-            self._cache = SemanticCache(**data)  # Use unpacking to handle Pydantic validation
+            self._cache = SemanticCacheModel(**data)  # Use unpacking to handle Pydantic validation
             for emb in self._cache.embeddings:
                 np_emb = np.array(emb, dtype=np.float32)
                 faiss.normalize_L2(np_emb.reshape(1, -1))
@@ -71,7 +73,7 @@ class SemanticCaching:
         logger.info("New response saved to cache.")
 
     def clear(self):
-        self._cache = SemanticCache()
+        self._cache = SemanticCacheModel()
         self.vector_index.reset()
         self.save_cache()
         logger.info("Cache cleared.")

tests/cache/test_semantic_cache.py

@@ -1,31 +1,31 @@
 import pytest
-from medirag.cache.local import SemanticCaching
+from medirag.cache.local import LocalSemanticCache
 
 
 # Fixture to initialize the SemanticCaching object
 @pytest.fixture(scope="module")
-def semantic_caching():
+def semantic_cache():
     # Initialize the SemanticCaching class with a test cache file
-    return SemanticCaching(
+    return LocalSemanticCache(
         model_name="sentence-transformers/all-mpnet-base-v2", dimension=768, json_file="real_test_cache.json"
     )
 
 
-def test_save_and_lookup_in_cache(semantic_caching):
+def test_save_and_lookup_in_cache(semantic_cache):
     # Clear any existing cache data
-    semantic_caching.clear()
+    semantic_cache.clear()
 
     # Step 1: Lookup should return None for a question not in the cache
-    initial_lookup = semantic_caching.lookup("What is the capital of France?")
+    initial_lookup = semantic_cache.lookup("What is the capital of France?")
     assert initial_lookup is None
 
     # Step 2: Save a response to the cache
-    semantic_caching.save("What is the capital of France?", "Paris")
+    semantic_cache.save("What is the capital of France?", "Paris")
 
     # Step 3: Lookup the same question; it should now return the cached response
-    cached_response = semantic_caching.lookup("What is the capital of France?")
+    cached_response = semantic_cache.lookup("What is the capital of France?")
     assert cached_response is not None
     assert cached_response == "Paris"
 
     # Cleanup: Clear the cache after test
-    semantic_caching.clear()
+    semantic_cache.clear()

tests/rag/test_rag.py

@@ -1,4 +1,4 @@
-from medirag.cache.local import SemanticCaching
+from medirag.cache.local import LocalSemanticCache
 from medirag.index.local import LocalIndexer
 
 # from medirag.index.kdbai import KDBAIDailyMedIndexer
@@ -36,10 +36,9 @@ def test_rag_with_example(data_dir):
 
     rag = RAG(k=3)
 
-    sm = SemanticCaching(
+    sm = LocalSemanticCache(
         model_name="sentence-transformers/all-mpnet-base-v2", dimension=768, json_file="rag_test_cache.json"
     )
-    # sm.load_cache()
 
     result1 = ask_med_question(sm, rag, query)
     print(result1)