feat - _add_prompt_to_generation_params for langfuse

litellm/integrations/langfuse.py

@@ -1,11 +1,13 @@
 #### What this does ####
 #    On success, logs events to Langfuse
-import os
 import copy
+import os
 import traceback
+
 from packaging.version import Version
-from litellm._logging import verbose_logger
+
 import litellm
+from litellm._logging import verbose_logger
 
 
 class LangFuseLogger:
@@ -14,8 +16,8 @@ class LangFuseLogger:
         self, langfuse_public_key=None, langfuse_secret=None, flush_interval=1
     ):
         try:
-            from langfuse import Langfuse
             import langfuse
+            from langfuse import Langfuse
         except Exception as e:
             raise Exception(
                 f"\033[91mLangfuse not installed, try running 'pip install langfuse' to fix this error: {e}\n{traceback.format_exc()}\033[0m"
@@ -251,7 +253,7 @@ class LangFuseLogger:
         input,
         response_obj,
     ):
-        from langfuse.model import CreateTrace, CreateGeneration
+        from langfuse.model import CreateGeneration, CreateTrace
 
         verbose_logger.warning(
             "Please upgrade langfuse to v2.0.0 or higher: https://github.com/langfuse/langfuse-python/releases/tag/v2.0.1"
@@ -533,30 +535,9 @@ class LangFuseLogger:
                 generation_params["parent_observation_id"] = parent_observation_id
 
             if supports_prompt:
-                user_prompt = clean_metadata.pop("prompt", None)
-                if user_prompt is None:
-                    pass
-                elif isinstance(user_prompt, dict):
-                    from langfuse.model import (
-                        TextPromptClient,
-                        ChatPromptClient,
-                        Prompt_Text,
-                        Prompt_Chat,
+                generation_params = _add_prompt_to_generation_params(
+                    generation_params=generation_params, clean_metadata=clean_metadata
                 )
-
-                    if user_prompt.get("type", "") == "chat":
-                        _prompt_chat = Prompt_Chat(**user_prompt)
-                        generation_params["prompt"] = ChatPromptClient(
-                            prompt=_prompt_chat
-                        )
-                    elif user_prompt.get("type", "") == "text":
-                        _prompt_text = Prompt_Text(**user_prompt)
-                        generation_params["prompt"] = TextPromptClient(
-                            prompt=_prompt_text
-                        )
-                else:
-                    generation_params["prompt"] = user_prompt
-
             if output is not None and isinstance(output, str) and level == "ERROR":
                 generation_params["status_message"] = output
 
@@ -569,5 +550,58 @@ class LangFuseLogger:
 
             return generation_client.trace_id, generation_id
         except Exception as e:
-            verbose_logger.debug(f"Langfuse Layer Error - {traceback.format_exc()}")
+            verbose_logger.error(f"Langfuse Layer Error - {traceback.format_exc()}")
             return None, None
+
+
+def _add_prompt_to_generation_params(
+    generation_params: dict, clean_metadata: dict
+) -> dict:
+    from langfuse.model import (
+        ChatPromptClient,
+        Prompt_Chat,
+        Prompt_Text,
+        TextPromptClient,
+    )
+
+    user_prompt = clean_metadata.pop("prompt", None)
+    if user_prompt is None:
+        pass
+    elif isinstance(user_prompt, dict):
+        if user_prompt.get("type", "") == "chat":
+            _prompt_chat = Prompt_Chat(**user_prompt)
+            generation_params["prompt"] = ChatPromptClient(prompt=_prompt_chat)
+        elif user_prompt.get("type", "") == "text":
+            _prompt_text = Prompt_Text(**user_prompt)
+            generation_params["prompt"] = TextPromptClient(prompt=_prompt_text)
+        elif "version" in user_prompt and "prompt" in user_prompt:
+            # prompts
+            if isinstance(user_prompt["prompt"], str):
+                _prompt_obj = Prompt_Text(
+                    name=user_prompt["name"],
+                    prompt=user_prompt["prompt"],
+                    version=user_prompt["version"],
+                    config=user_prompt.get("config", None),
+                )
+                generation_params["prompt"] = TextPromptClient(prompt=_prompt_obj)
+
+            elif isinstance(user_prompt["prompt"], list):
+                _prompt_obj = Prompt_Chat(
+                    name=user_prompt["name"],
+                    prompt=user_prompt["prompt"],
+                    version=user_prompt["version"],
+                    config=user_prompt.get("config", None),
+                )
+                generation_params["prompt"] = ChatPromptClient(prompt=_prompt_obj)
+            else:
+                verbose_logger.error(
+                    "[Non-blocking] Langfuse Logger: Invalid prompt format"
+                )
+        else:
+            verbose_logger.error(
+                "[Non-blocking] Langfuse Logger: Invalid prompt format. No prompt logged to Langfuse"
+            )
+    else:
+        generation_params["prompt"] = user_prompt
+
+    return generation_params