Merge pull request #1722 from BerriAI/litellm_langfuse_log_embedding

[Feat] Langfuse log embeddings

litellm/integrations/langfuse.py

@@ -9,6 +9,7 @@ dotenv.load_dotenv()  # Loading env variables using dotenv
 import traceback
 from packaging.version import Version
 from litellm._logging import verbose_logger
+import litellm
 
 
 class LangFuseLogger:
@@ -63,11 +64,15 @@ class LangFuseLogger:
                         pass
 
             # end of processing langfuse ########################
-            input = prompt
-            output = response_obj["choices"][0]["message"].json()
-            print_verbose(
-                f"OUTPUT IN LANGFUSE: {output}; original: {response_obj['choices'][0]['message']}"
-            )
+            if kwargs.get("call_type", None) == "embedding" or isinstance(
+                response_obj, litellm.EmbeddingResponse
+            ):
+                input = prompt
+                output = response_obj["data"]
+            else:
+                input = prompt
+                output = response_obj["choices"][0]["message"].json()
+            print_verbose(f"OUTPUT IN LANGFUSE: {output}; original: {response_obj}")
             self._log_langfuse_v2(
                 user_id,
                 metadata,
@@ -173,8 +178,12 @@ class LangFuseLogger:
         tags = []
         supports_tags = Version(langfuse.version.__version__) >= Version("2.6.3")
 
+        generation_name = metadata.get("generation_name", None)
+        if generation_name is None:
+            # just log `litellm-{call_type}` as the generation name
+            generation_name = f"litellm-{kwargs.get('call_type', 'completion')}"
         trace_params = {
-            "name": metadata.get("generation_name", "litellm-completion"),
+            "name": generation_name,
             "input": input,
             "output": output,
             "user_id": metadata.get("trace_user_id", user_id),
@@ -190,7 +199,7 @@ class LangFuseLogger:
         trace = self.Langfuse.trace(**trace_params)
 
         trace.generation(
-            name=metadata.get("generation_name", "litellm-completion"),
+            name=generation_name,
             id=metadata.get("generation_id", None),
             startTime=start_time,
             endTime=end_time,

litellm/tests/test_alangfuse.py

@@ -277,7 +277,24 @@ def test_langfuse_logging_custom_generation_name():
         print(e)
 
 
-test_langfuse_logging_custom_generation_name()
+# test_langfuse_logging_custom_generation_name()
+
+
+@pytest.mark.skip(reason="beta test - checking langfuse output")
+def test_langfuse_logging_embedding():
+    try:
+        litellm.set_verbose = True
+        litellm.success_callback = ["langfuse"]
+        response = litellm.embedding(
+            model="text-embedding-ada-002",
+            input=["gm", "ishaan"],
+        )
+        print(response)
+    except litellm.Timeout as e:
+        pass
+    except Exception as e:
+        pytest.fail(f"An exception occurred - {e}")
+        print(e)
 
 
 @pytest.mark.skip(reason="beta test - checking langfuse output")