fix: test and implementation

Author: Matthieu-OD

literalai/api/__init__.py

@@ -143,16 +143,15 @@ def handle_bytes(item):
 
 class SharedCache:
     """
-    Thread-safe singleton cache for storing prompts with memory leak prevention.
+    Singleton cache for storing data.
     Only one instance will exist regardless of how many times it's instantiated.
     """
     _instance = None
-    _cache: dict[str, Any]
 
     def __new__(cls):
         if cls._instance is None:
             cls._instance = super().__new__(cls)
-            cls._instance.cache = {}
+            cls._instance._cache = {}
         return cls._instance
 
     def get_cache(self) -> dict[str, Any]:
@@ -162,14 +161,23 @@ def get(self, key: str) -> Optional[Any]:
         """
         Retrieves a value from the cache using the provided key.
         """
+        if not isinstance(key, str):
+            raise TypeError("Key must be a string")
         return self._cache.get(key)
 
     def put(self, key: str, value: Any):
         """
         Stores a value in the cache.
         """
+        if not isinstance(key, str):
+            raise TypeError("Key must be a string")
         self._cache[key] = value
 
+    def put_prompt(self, prompt: Prompt):
+        self.put(prompt.id, prompt)
+        self.put(prompt.name, prompt)
+        self.put(f"{prompt.name}-{prompt.version}", prompt)
+
     def clear(self) -> None:
         """
         Clears all cached values.
@@ -1398,7 +1406,7 @@ def get_prompt(
             raise ValueError("Either the `id` or the `name` must be provided.")
 
         get_prompt_query, description, variables, process_response, timeout, cached_prompt = get_prompt_helper(
-            api=self,id=id, name=name, version=version, prompt_cache=self.prompt_cache
+            api=self,id=id, name=name, version=version, cache=self.cache
         )
 
         try:
@@ -1407,9 +1415,8 @@ def get_prompt(
             elif name:
                 prompt = self.gql_helper(get_prompt_query, description, variables, process_response, timeout)
 
-            self.cache.put(prompt.id, prompt)
-            self.cache.put(prompt.name, prompt)
-            self.cache.put(f"{prompt.name}-{prompt.version}", prompt)
+            self.cache.put_prompt(prompt)
+
             return prompt
 
         except Exception as e:
@@ -2645,7 +2652,7 @@ async def get_prompt(
 
         sync_api = LiteralAPI(self.api_key, self.url)
         get_prompt_query, description, variables, process_response, timeout, cached_prompt = get_prompt_helper(
-            api=sync_api, id=id, name=name, version=version, prompt_cache=self.prompt_cache
+            api=sync_api, id=id, name=name, version=version, cache=self.cache
         )
 
         try:
@@ -2658,7 +2665,6 @@ async def get_prompt(
                     get_prompt_query, description, variables, process_response, timeout
                 )
 
-            self.prompt_cache.put(prompt)
             return prompt
 
         except Exception as e:

literalai/api/prompt_helpers.py

@@ -5,7 +5,7 @@
 
 if TYPE_CHECKING:
     from literalai.api import LiteralAPI
-    from literalai.api import SharedPromptCache
+    from literalai.api import SharedCache
 
 from literalai.api import gql
 
@@ -73,24 +73,24 @@ def get_prompt_helper(
     id: Optional[str] = None,
     name: Optional[str] = None,
     version: Optional[int] = 0,
-    prompt_cache: Optional["SharedPromptCache"] = None,
+    cache: Optional["SharedCache"] = None,
 ) -> tuple[str, str, dict, Callable, int, Optional[Prompt]]:
     """Helper function for getting prompts with caching logic"""
 
     cached_prompt = None
     timeout = 10
 
-    if prompt_cache:
-        cached_prompt = prompt_cache.get(get_prompt_cache_key(id, name, version))
+    if cache:
+        cached_prompt = cache.get(get_prompt_cache_key(id, name, version))
         timeout = 1 if cached_prompt else timeout
 
     variables = {"id": id, "name": name, "version": version}
 
     def process_response(response):
         prompt_version = response["data"]["promptVersion"]
         prompt = Prompt.from_dict(api, prompt_version) if prompt_version else None
-        if prompt_cache:
-            prompt_cache.put(prompt)
+        if cache and prompt:
+            cache.put_prompt(prompt)
         return prompt
 
     description = "get prompt"

tests/unit/test_cache.py

@@ -1,10 +1,7 @@
 import pytest
-from threading import Thread
-import time
-import random
 
 from literalai.prompt_engineering.prompt import Prompt
-from literalai.api import SharedPromptCache, LiteralAPI
+from literalai.api import SharedCache, LiteralAPI
 
 def default_prompt(id: str = "1", name: str = "test", version: int = 1) -> Prompt:
     return Prompt(
@@ -26,156 +23,69 @@ def default_prompt(id: str = "1", name: str = "test", version: int = 1) -> Promp
     )
 
 def test_singleton_instance():
-    """Test that SharedPromptCache maintains singleton pattern"""
-    cache1 = SharedPromptCache()
-    cache2 = SharedPromptCache()
+    """Test that SharedCache maintains singleton pattern"""
+    cache1 = SharedCache()
+    cache2 = SharedCache()
     assert cache1 is cache2
 
 def test_get_empty_cache():
     """Test getting from empty cache returns None"""
-    cache = SharedPromptCache()
+    cache = SharedCache()
     cache.clear() 
 
-    assert cache._prompts == {}
-    assert cache._name_index == {}
-    assert cache._name_version_index == {}
+    assert cache.get_cache() == {}
 
-def test_put_and_get_by_id():
-    """Test storing and retrieving prompt by ID"""
-    cache = SharedPromptCache()
+def test_put_and_get_prompt_by_id_by_name_version_by_name():
+    """Test storing and retrieving prompt by ID by name-version by name"""
+    cache = SharedCache()
     cache.clear()
 
     prompt = default_prompt()
-    cache.put(prompt)
+    cache.put_prompt(prompt)
 
-    retrieved = cache.get(id="1")
-    assert retrieved is prompt
-    assert retrieved.id == "1"
-    assert retrieved.name == "test"
-    assert retrieved.version == 1
-
-def test_put_and_get_by_name():
-    """Test storing and retrieving prompt by name"""
-    cache = SharedPromptCache()
-    cache.clear()
+    retrieved_by_id = cache.get(id="1")
+    assert retrieved_by_id is prompt
 
-    prompt = default_prompt()
-    cache.put(prompt)
+    retrieved_by_name_version = cache.get(name="test", version=1)
+    assert retrieved_by_name_version is prompt
 
-    retrieved = cache.get(name="test")
-    assert retrieved is prompt
-    assert retrieved.name == "test"
-
-def test_put_and_get_by_name_version():
-    """Test storing and retrieving prompt by name and version"""
-    cache = SharedPromptCache()
-    cache.clear()
-    
-    prompt = default_prompt()
-    cache.put(prompt)
-    
-    retrieved = cache.get(name="test", version=1)
-    assert retrieved is prompt
-    assert retrieved.name == "test"
-    assert retrieved.version == 1
-
-def test_multiple_versions():
-    """Test handling multiple versions of the same prompt"""
-    cache = SharedPromptCache()
-    cache.clear()
-    
-    prompt1 = default_prompt()
-    prompt2 = default_prompt(id="2", version=2)
-    
-    cache.put(prompt1)
-    cache.put(prompt2)
-    
-    assert cache.get(name="test", version=1) is prompt1
-    assert cache.get(name="test", version=2) is prompt2
-    
-    assert cache.get(name="test") is prompt2 
+    retrieved_by_name = cache.get(name="test")
+    assert retrieved_by_name is prompt
 
 def test_clear_cache():
     """Test clearing the cache"""
-    cache = SharedPromptCache()
+    cache = SharedCache()
     prompt = default_prompt()
-    cache.put(prompt)
+    cache.put_prompt(prompt)
 
     cache.clear()
-    assert cache._prompts == {}
-    assert cache._name_index == {}
-    assert cache._name_version_index == {}
+    assert cache.get_cache() == {}
 
 def test_update_existing_prompt():
     """Test updating an existing prompt"""
-    cache = SharedPromptCache()
+    cache = SharedCache()
     cache.clear()
 
     prompt1 = default_prompt()
     prompt2 = default_prompt(id="1", version=2)
 
-    cache.put(prompt1)
-    cache.put(prompt2)
+    cache.put_prompt(prompt1)
+    cache.put_prompt(prompt2)
 
     retrieved = cache.get(id="1")
     assert retrieved is prompt2
     assert retrieved.version == 2
 
-def test_lookup_priority():
-    """Test that lookup priority is id > name-version > name"""
-    cache = SharedPromptCache()
-    cache.clear()
-    
-    prompt1 = default_prompt()
-    prompt2 = default_prompt(id="2", name="test", version=2)
-    
-    cache.put(prompt1)
-    cache.put(prompt2)
-    
-    assert cache.get(id="1", name="test", version=2) is prompt1
-    
-    assert cache.get(name="test", version=2) is prompt2
-
-def test_thread_safety():
-    """Test thread safety of the cache"""
-    cache = SharedPromptCache()
-    cache.clear()
-    
-    def worker(worker_id: int):
-        for i in range(100):
-            prompt = default_prompt(
-                id=f"{worker_id}-{i}",
-                name=f"test-{worker_id}",
-                version=i
-            )
-            cache.put(prompt)
-            time.sleep(random.uniform(0, 0.001))
-            
-            retrieved = cache.get(id=prompt.id)
-            assert retrieved is prompt
-    
-    threads = [Thread(target=worker, args=(i,)) for i in range(10)]
-    
-    for t in threads:
-        t.start()
-    for t in threads:
-        t.join()
-    
-    for worker_id in range(10):
-        for i in range(100):
-            prompt_id = f"{worker_id}-{i}"
-            assert cache.get(id=prompt_id) is not None
-
 def test_error_handling():
     """Test error handling for invalid inputs"""
-    cache = SharedPromptCache()
+    cache = SharedCache()
     cache.clear()
 
-    assert cache.get() is None
-    assert cache.get(id=None, name=None, version=None) is None
+    assert cache.get_cache() == {}
+    assert cache.get(key="") is None
 
     with pytest.raises(TypeError):
-        cache.get(version="invalid")  # type: ignore
+        cache.get(5)  # type: ignore
 
     with pytest.raises(TypeError):
-        cache.put("not a prompt")  # type: ignore
+        cache.put(5, "test")  # type: ignore