feat(prompts): add fallback to get_prompts (langfuse#792)

hassiebp · web-flow · commit 697382488c68 · 2024-07-10T19:10:13.000+02:00
diff --git a/langfuse/client.py b/langfuse/client.py
@@ -49,6 +49,8 @@
 from langfuse.api.resources.prompts.types import (
     CreatePromptRequest_Chat,
     CreatePromptRequest_Text,
+    Prompt_Text,
+    Prompt_Chat,
 )
 from langfuse.api.resources.trace.types.traces import Traces
 from langfuse.api.resources.utils.resources.pagination.types.meta_response import (
@@ -876,6 +878,7 @@ def get_prompt(
         label: Optional[str] = None,
         type: Literal["chat"],
         cache_ttl_seconds: Optional[int] = None,
+        fallback: Optional[List[ChatMessageDict]] = None,
     ) -> ChatPromptClient: ...
 
     @overload
@@ -887,6 +890,7 @@ def get_prompt(
         label: Optional[str] = None,
         type: Literal["text"] = "text",
         cache_ttl_seconds: Optional[int] = None,
+        fallback: Optional[str] = None,
     ) -> TextPromptClient: ...
 
     def get_prompt(
@@ -897,6 +901,7 @@ def get_prompt(
         label: Optional[str] = None,
         type: Literal["chat", "text"] = "text",
         cache_ttl_seconds: Optional[int] = None,
+        fallback: Union[Optional[List[ChatMessageDict]], Optional[str]] = None,
     ) -> PromptClient:
         """Get a prompt.
 
@@ -914,6 +919,7 @@ def get_prompt(
             cache_ttl_seconds: Optional[int]: Time-to-live in seconds for caching the prompt. Must be specified as a
             keyword argument. If not set, defaults to 60 seconds.
             type: Literal["chat", "text"]: The type of the prompt to retrieve. Defaults to "text".
+            fallback: Union[Optional[List[ChatMessageDict]], Optional[str]]: The prompt string to return if fetching the prompt fails. Important on the first call where no cached prompt is available. Follows Langfuse prompt formatting with double curly braces for variables. Defaults to None.
 
         Returns:
             The prompt object retrieved from the cache or directly fetched if not cached or expired of type
@@ -936,9 +942,40 @@ def get_prompt(
         cached_prompt = self.prompt_cache.get(cache_key)
 
         if cached_prompt is None:
-            return self._fetch_prompt_and_update_cache(
-                name, version=version, label=label, ttl_seconds=cache_ttl_seconds
-            )
+            self.log.debug(f"Prompt '{cache_key}' not found in cache.")
+            try:
+                return self._fetch_prompt_and_update_cache(
+                    name, version=version, label=label, ttl_seconds=cache_ttl_seconds
+                )
+            except Exception as e:
+                if fallback:
+                    self.log.warn(
+                        f"Returning fallback prompt for '{cache_key}' due to fetch error: {e}"
+                    )
+
+                    fallback_client_args = {
+                        "name": name,
+                        "prompt": fallback,
+                        "type": type,
+                        "version": version or 0,
+                        "config": {},
+                        "labels": [label] if label else [],
+                        "tags": [],
+                    }
+
+                    if type == "text":
+                        return TextPromptClient(
+                            prompt=Prompt_Text(**fallback_client_args),
+                            is_fallback=True,
+                        )
+
+                    if type == "chat":
+                        return ChatPromptClient(
+                            prompt=Prompt_Chat(**fallback_client_args),
+                            is_fallback=True,
+                        )
+
+                raise e
 
         if cached_prompt.is_expired():
             try:
@@ -973,7 +1010,10 @@ def _fetch_prompt_and_update_cache(
 
             self.log.debug(f"Fetching prompt '{cache_key}' from server...")
             promptResponse = self.client.prompts.get(
-                self._url_encode(name), version=version, label=label
+                self._url_encode(name),
+                version=version,
+                label=label,
+                request_options={"max_retries": 2},
             )
 
             if promptResponse.type == "chat":
diff --git a/langfuse/model.py b/langfuse/model.py
@@ -61,12 +61,13 @@ class BasePromptClient(ABC):
     labels: List[str]
     tags: List[str]
 
-    def __init__(self, prompt: Prompt):
+    def __init__(self, prompt: Prompt, is_fallback: bool = False):
         self.name = prompt.name
         self.version = prompt.version
         self.config = prompt.config
         self.labels = prompt.labels
         self.tags = prompt.tags
+        self.is_fallback = is_fallback
 
     @abstractmethod
     def compile(self, **kwargs) -> Union[str, List[ChatMessage]]:
@@ -127,8 +128,8 @@ def _compile_template_string(content: str, data: Dict[str, Any] = {}) -> str:
 
 
 class TextPromptClient(BasePromptClient):
-    def __init__(self, prompt: Prompt_Text):
-        super().__init__(prompt)
+    def __init__(self, prompt: Prompt_Text, is_fallback: bool = False):
+        super().__init__(prompt, is_fallback)
         self.prompt = prompt.prompt
 
     def compile(self, **kwargs) -> str:
@@ -158,8 +159,8 @@ def get_langchain_prompt(self):
 
 
 class ChatPromptClient(BasePromptClient):
-    def __init__(self, prompt: Prompt_Chat):
-        super().__init__(prompt)
+    def __init__(self, prompt: Prompt_Chat, is_fallback: bool = False):
+        super().__init__(prompt, is_fallback)
         self.prompt = [
             ChatMessageDict(role=p.role, content=p.content) for p in prompt.prompt
         ]
diff --git a/langfuse/utils/__init__.py b/langfuse/utils/__init__.py
@@ -21,7 +21,7 @@ def _get_timestamp():
 def _create_prompt_context(
     prompt: typing.Optional[PromptClient] = None,
 ):
-    if prompt is not None:
+    if prompt is not None and not prompt.is_fallback:
         return {"prompt_version": prompt.version, "prompt_name": prompt.name}
 
     return {"prompt_version": None, "prompt_name": None}
diff --git a/tests/test_prompt.py b/tests/test_prompt.py
@@ -421,8 +421,10 @@ def test_get_fresh_prompt(langfuse):
     mock_server_call = langfuse.client.prompts.get
     mock_server_call.return_value = prompt
 
-    result = langfuse.get_prompt(prompt_name)
-    mock_server_call.assert_called_once_with(prompt_name, version=None, label=None)
+    result = langfuse.get_prompt(prompt_name, fallback="fallback")
+    mock_server_call.assert_called_once_with(
+        prompt_name, version=None, label=None, request_options={"max_retries": 2}
+    )
 
     assert result == TextPromptClient(prompt)
 
@@ -480,7 +482,7 @@ def test_get_valid_cached_prompt(langfuse):
     mock_server_call = langfuse.client.prompts.get
     mock_server_call.return_value = prompt
 
-    result_call_1 = langfuse.get_prompt(prompt_name)
+    result_call_1 = langfuse.get_prompt(prompt_name, fallback="fallback")
     assert mock_server_call.call_count == 1
     assert result_call_1 == prompt_client
 
@@ -742,3 +744,82 @@ def test_get_fresh_prompt_when_version_changes(langfuse):
     result_call_2 = langfuse.get_prompt(prompt_name, version=2)
     assert mock_server_call.call_count == 2
     assert result_call_2 == version_changed_prompt_client
+
+
+def test_do_not_return_fallback_if_fetch_success():
+    langfuse = Langfuse()
+    prompt_name = create_uuid()
+    prompt_client = langfuse.create_prompt(
+        name=prompt_name,
+        prompt="test prompt",
+        labels=["production"],
+    )
+
+    second_prompt_client = langfuse.get_prompt(prompt_name, fallback="fallback")
+
+    assert prompt_client.name == second_prompt_client.name
+    assert prompt_client.version == second_prompt_client.version
+    assert prompt_client.prompt == second_prompt_client.prompt
+    assert prompt_client.config == second_prompt_client.config
+    assert prompt_client.config == {}
+
+
+def test_fallback_text_prompt():
+    langfuse = Langfuse()
+
+    fallback_text_prompt = "this is a fallback text prompt with {{variable}}"
+
+    # Should throw an error if prompt not found and no fallback provided
+    with pytest.raises(Exception):
+        langfuse.get_prompt("nonexistent_prompt")
+
+    prompt = langfuse.get_prompt("nonexistent_prompt", fallback=fallback_text_prompt)
+
+    assert prompt.prompt == fallback_text_prompt
+    assert (
+        prompt.compile(variable="value") == "this is a fallback text prompt with value"
+    )
+
+
+def test_fallback_chat_prompt():
+    langfuse = Langfuse()
+    fallback_chat_prompt = [
+        {"role": "system", "content": "fallback system"},
+        {"role": "user", "content": "fallback user name {{name}}"},
+    ]
+
+    # Should throw an error if prompt not found and no fallback provided
+    with pytest.raises(Exception):
+        langfuse.get_prompt("nonexistent_chat_prompt", type="chat")
+
+    prompt = langfuse.get_prompt(
+        "nonexistent_chat_prompt", type="chat", fallback=fallback_chat_prompt
+    )
+
+    assert prompt.prompt == fallback_chat_prompt
+    assert prompt.compile(name="Jane") == [
+        {"role": "system", "content": "fallback system"},
+        {"role": "user", "content": "fallback user name Jane"},
+    ]
+
+
+def test_do_not_link_observation_if_fallback():
+    langfuse = Langfuse()
+    trace_id = create_uuid()
+
+    fallback_text_prompt = "this is a fallback text prompt with {{variable}}"
+
+    # Should throw an error if prompt not found and no fallback provided
+    with pytest.raises(Exception):
+        langfuse.get_prompt("nonexistent_prompt")
+
+    prompt = langfuse.get_prompt("nonexistent_prompt", fallback=fallback_text_prompt)
+
+    langfuse.trace(id=trace_id).generation(prompt=prompt, input="this is a test input")
+    langfuse.flush()
+
+    api = get_api()
+    trace = api.trace.get(trace_id)
+
+    assert len(trace.observations) == 1
+    assert trace.observations[0].prompt_id is None
