feat: add caching if prompt request fails

literalai/api/__init__.py

@@ -1,8 +1,8 @@
+from threading import Lock
 import logging
 import os
 import uuid
 from typing import (
-    TYPE_CHECKING,
     Any,
     Callable,
     Dict,
@@ -102,9 +102,6 @@
 )
 from literalai.prompt_engineering.prompt import Prompt, ProviderSettings
 
-if TYPE_CHECKING:
-    from typing import Tuple  # noqa: F401
-
 import httpx
 
 from literalai.my_types import Environment, PaginatedResponse
@@ -145,6 +142,59 @@ def handle_bytes(item):
     return handle_bytes(variables)
 
 
+class SharedPromptCache:
+    """
+    Thread-safe singleton cache for storing prompts.
+    Only one instance will exist regardless of how many times it's instantiated.
+    """
+    _instance = None
+    _lock = Lock()
+
+    def __new__(cls):
+        with cls._lock:
+            if cls._instance is None:
+                cls._instance = super().__new__(cls)
+
+                cls._instance._prompts: dict[str, Prompt] = {}
+                cls._instance._name_index: dict[str, str] = {}
+                cls._instance._name_version_index: dict[tuple[str, int], str] = {}
+        return cls._instance
+
+    def get(
+        self,
+        id: Optional[str] = None,
+        name: Optional[str] = None,
+        version: Optional[int] = None
+    ) -> Optional[Prompt]:
+        """
+        Retrieves a prompt using the most specific criteria provided.
+        Lookup priority: id, name-version, name
+        """
+        if id:
+            prompt_id = id
+        elif name and version:
+            prompt_id = self._name_version_index.get((name, version))
+        elif name:
+            prompt_id = self._name_index.get(name)
+
+        return self._prompts.get(prompt_id) if prompt_id else None
+
+    def put(self, prompt: Prompt):
+        with self._lock: 
+            self._prompts[prompt.id] = prompt
+            self._name_index[prompt.name] = prompt.id
+            self._name_version_index[(prompt.name, prompt.version)] = prompt.id
+
+    def clear(self) -> None:
+        """
+        Clears all cached promopts and indices.
+        """
+        with self._lock:
+            self._prompts.clear()
+            self._name_index.clear()
+            self._name_version_index.clear()
+
+
 class BaseLiteralAPI:
     def __init__(
         self,
@@ -169,6 +219,8 @@ def __init__(
         self.graphql_endpoint = self.url + "/api/graphql"
         self.rest_endpoint = self.url + "/api"
 
+        self.prompt_cache = SharedPromptCache()
+
     @property
     def headers(self):
         from literalai.version import __version__
@@ -186,6 +238,7 @@ def headers(self):
         return h
 
 
+
 class LiteralAPI(BaseLiteralAPI):
     """
     ```python
@@ -200,8 +253,8 @@ class LiteralAPI(BaseLiteralAPI):
     R = TypeVar("R")
 
     def make_gql_call(
-        self, description: str, query: str, variables: Dict[str, Any]
-    ) -> Dict:
+        self, description: str, query: str, variables: dict[str, Any], timeout: Optional[int] = 10
+    ) -> dict:
         def raise_error(error):
             logger.error(f"Failed to {description}: {error}")
             raise Exception(error)
@@ -212,7 +265,7 @@ def raise_error(error):
                 self.graphql_endpoint,
                 json={"query": query, "variables": variables},
                 headers=self.headers,
-                timeout=10,
+                timeout=timeout,
             )
 
             try:
@@ -233,7 +286,7 @@ def raise_error(error):
 
             if json.get("data"):
                 if isinstance(json["data"], dict):
-                    for key, value in json["data"].items():
+                    for value in json["data"].values():
                         if value and value.get("ok") is False:
                             raise_error(
                                 f"""Failed to {description}: {
@@ -242,10 +295,6 @@ def raise_error(error):
 
             return json
 
-        # This should not be reached, exceptions should be thrown beforehands
-        # Added because of mypy
-        raise Exception("Unknown error")
-
     def make_rest_call(self, subpath: str, body: Dict[str, Any]) -> Dict:
         with httpx.Client(follow_redirects=True) as client:
             response = client.post(
@@ -276,8 +325,9 @@ def gql_helper(
         description: str,
         variables: Dict,
         process_response: Callable[..., R],
+        timeout: Optional[int] = None,
     ) -> R:
-        response = self.make_gql_call(description, query, variables)
+        response = self.make_gql_call(description, query, variables, timeout)
         return process_response(response)
 
     # User API
@@ -684,7 +734,7 @@ def upload_file(
         # Prepare form data
         form_data = (
             {}
-        )  # type: Dict[str, Union[Tuple[Union[str, None], Any], Tuple[Union[str, None], Any, Any]]]
+        )  # type: Dict[str, Union[tuple[Union[str, None], Any], tuple[Union[str, None], Any, Any]]]
         for field_name, field_value in fields.items():
             form_data[field_name] = (None, field_value)
 
@@ -1362,13 +1412,28 @@ def get_prompt(
         Returns:
             Prompt: The prompt with the given identifier or name.
         """
-        if id:
-            return self.gql_helper(*get_prompt_helper(self, id=id))
-        elif name:
-            return self.gql_helper(*get_prompt_helper(self, name=name, version=version))
-        else:
+        if not (id or name):
             raise ValueError("Either the `id` or the `name` must be provided.")
 
+        cached_prompt = self.prompt_cache.get(id, name, version)
+        timeout = 1 if cached_prompt else None
+
+        try:
+            if id:
+                prompt = self.gql_helper(*get_prompt_helper(self, id=id, timeout=timeout))
+            elif name:
+                prompt = self.gql_helper(*get_prompt_helper(self, name=name, version=version, timeout=timeout))
+
+            self.prompt_cache.put(prompt)
+            return prompt
+
+        except Exception as e:
+            if cached_prompt:
+                logger.warning("Failed to get prompt from API, returning cached prompt")
+                return cached_prompt
+
+            raise e
+
     def create_prompt_variant(
         self,
         name: str,
@@ -1451,7 +1516,7 @@ class AsyncLiteralAPI(BaseLiteralAPI):
     R = TypeVar("R")
 
     async def make_gql_call(
-        self, description: str, query: str, variables: Dict[str, Any]
+        self, description: str, query: str, variables: Dict[str, Any], timeout: Optional[int] = 10
     ) -> Dict:
         def raise_error(error):
             logger.error(f"Failed to {description}: {error}")
@@ -1464,7 +1529,7 @@ def raise_error(error):
                 self.graphql_endpoint,
                 json={"query": query, "variables": variables},
                 headers=self.headers,
-                timeout=10,
+                timeout=timeout,
             )
 
             try:
@@ -1485,7 +1550,7 @@ def raise_error(error):
 
             if json.get("data"):
                 if isinstance(json["data"], dict):
-                    for key, value in json["data"].items():
+                    for value in json["data"].values():
                         if value and value.get("ok") is False:
                             raise_error(
                                 f"""Failed to {description}: {
@@ -1494,10 +1559,6 @@ def raise_error(error):
 
             return json
 
-        # This should not be reached, exceptions should be thrown beforehands
-        # Added because of mypy
-        raise Exception("Unkown error")
-
     async def make_rest_call(self, subpath: str, body: Dict[str, Any]) -> Dict:
         async with httpx.AsyncClient(follow_redirects=True) as client:
             response = await client.post(
@@ -1528,8 +1589,9 @@ async def gql_helper(
         description: str,
         variables: Dict,
         process_response: Callable[..., R],
+        timeout: Optional[int] = 10,
     ) -> R:
-        response = await self.make_gql_call(description, query, variables)
+        response = await self.make_gql_call(description, query, variables, timeout)
         return process_response(response)
 
     async def get_users(
@@ -1963,7 +2025,7 @@ async def upload_file(
         # Prepare form data
         form_data = (
             {}
-        )  # type: Dict[str, Union[Tuple[Union[str, None], Any], Tuple[Union[str, None], Any, Any]]]
+        )  # type: dict[str, Union[tuple[Union[str, None], Any], tuple[Union[str, None], Any, Any]]]
         for field_name, field_value in fields.items():
             form_data[field_name] = (None, field_value)
 
@@ -2593,16 +2655,34 @@ async def get_prompt(
         name: Optional[str] = None,
         version: Optional[int] = None,
     ) -> Prompt:
-        sync_api = LiteralAPI(self.api_key, self.url)
-        if id:
-            return await self.gql_helper(*get_prompt_helper(sync_api, id=id))
-        elif name:
-            return await self.gql_helper(
-                *get_prompt_helper(sync_api, name=name, version=version)
-            )
-        else:
+        if not (id or name):
             raise ValueError("Either the `id` or the `name` must be provided.")
 
+        sync_api = LiteralAPI(self.api_key, self.url)
+        cached_prompt = self.prompt_cache.get(id, name, version)
+        timeout = 1 if cached_prompt else None
+
+        try:
+            if id:
+                prompt = await self.gql_helper(
+                    *get_prompt_helper(sync_api, id=id, timeout=timeout)
+                )
+            elif name:
+                prompt = await self.gql_helper(
+                    *get_prompt_helper(
+                        sync_api, name=name, version=version, timeout=timeout
+                    )
+                )
+
+            self.prompt_cache.put(prompt)
+            return prompt
+
+        except Exception as e:
+            if cached_prompt:
+                logger.warning("Failed to get prompt from API, returning cached prompt")
+                return cached_prompt
+            raise e
+
     get_prompt.__doc__ = LiteralAPI.get_prompt.__doc__
 
     async def update_prompt_ab_testing(

literalai/api/prompt_helpers.py

@@ -61,6 +61,7 @@ def get_prompt_helper(
     id: Optional[str] = None,
     name: Optional[str] = None,
     version: Optional[int] = 0,
+    timeout: Optional[int] = None,
 ):
     variables = {"id": id, "name": name, "version": version}
 
@@ -70,7 +71,7 @@ def process_response(response):
 
     description = "get prompt"
 
-    return gql.GET_PROMPT_VERSION, description, variables, process_response
+    return gql.GET_PROMPT_VERSION, description, variables, process_response, timeout
 
 
 def create_prompt_variant_helper(