From a3e4fa5b81e82f71c74fb9e7dc859c6cb40495f5 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Wed, 22 Oct 2025 14:20:39 -0700
Subject: [PATCH] fix: initial commit fixing langfuse request/response logging
 with OTEL

---
 litellm/integrations/arize/_utils.py   | 68 ++++++++++++++++++--------
 tests/local_testing/test_completion.py | 37 ++++++++++++--
 2 files changed, 82 insertions(+), 23 deletions(-)

diff --git a/litellm/integrations/arize/_utils.py b/litellm/integrations/arize/_utils.py
index e93ef128b4a..693f0bc176f 100644
--- a/litellm/integrations/arize/_utils.py
+++ b/litellm/integrations/arize/_utils.py
@@ -35,6 +35,47 @@ def safe_set_attribute(span: Span, key: str, value: Any):
     span.set_attribute(key, primitive_value)
 
 
+def set_response_output_messages(span: Span, response_obj):
+    """
+    Sets output message attributes on the span from the LLM response.
+
+    Args:
+        span: The OpenTelemetry span to set attributes on
+        response_obj: The response object containing choices with messages
+    """
+    from litellm.integrations._types.open_inference import (
+        MessageAttributes,
+        SpanAttributes,
+    )
+
+    safe_set_attribute(
+        span,
+        "langfuse.observation.output",
+        response_obj.model_dump_json(),
+    )
+
+    for idx, choice in enumerate(response_obj.get("choices", [])):
+        response_message = choice.get("message", {})
+        safe_set_attribute(
+            span,
+            SpanAttributes.OUTPUT_VALUE,
+            response_message.get("content", ""),
+        )
+
+        # This shows up under `output_messages` tab on the span page.
+        prefix = f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.{idx}"
+        safe_set_attribute(
+            span,
+            f"{prefix}.{MessageAttributes.MESSAGE_ROLE}",
+            response_message.get("role"),
+        )
+        safe_set_attribute(
+            span,
+            f"{prefix}.{MessageAttributes.MESSAGE_CONTENT}",
+            response_message.get("content", ""),
+        )
+
+
 def set_attributes(span: Span, kwargs, response_obj):  # noqa: PLR0915
     """
     Populates span with OpenInference-compliant LLM attributes for Arize and Phoenix tracing.
@@ -153,6 +194,12 @@ def set_attributes(span: Span, kwargs, response_obj):  # noqa: PLR0915
         )
         messages = kwargs.get("messages")
 
+        safe_set_attribute(
+            span,
+            "langfuse.observation.input",
+            json.dumps(messages),
+        )
+
         # for /chat/completions
         # https://docs.arize.com/arize/large-language-models/tracing/semantic-conventions
         if messages:
@@ -235,26 +282,7 @@ def set_attributes(span: Span, kwargs, response_obj):  # noqa: PLR0915
 
         # Captures response tokens, message, and content.
         if hasattr(response_obj, "get"):
-            for idx, choice in enumerate(response_obj.get("choices", [])):
-                response_message = choice.get("message", {})
-                safe_set_attribute(
-                    span,
-                    SpanAttributes.OUTPUT_VALUE,
-                    response_message.get("content", ""),
-                )
-
-                # This shows up under `output_messages` tab on the span page.
-                prefix = f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.{idx}"
-                safe_set_attribute(
-                    span,
-                    f"{prefix}.{MessageAttributes.MESSAGE_ROLE}",
-                    response_message.get("role"),
-                )
-                safe_set_attribute(
-                    span,
-                    f"{prefix}.{MessageAttributes.MESSAGE_CONTENT}",
-                    response_message.get("content", ""),
-                )
+            set_response_output_messages(span, response_obj)
 
             # Token usage info.
             usage = response_obj and response_obj.get("usage")
diff --git a/tests/local_testing/test_completion.py b/tests/local_testing/test_completion.py
index d91306b0208..50baa9c8032 100644
--- a/tests/local_testing/test_completion.py
+++ b/tests/local_testing/test_completion.py
@@ -2514,8 +2514,6 @@ def test_completion_azure_key_completion_arg():
         pytest.fail(f"Error occurred: {e}")
 
 
-
-
 async def test_re_use_azure_async_client():
     try:
         print("azure gpt-3.5 ASYNC with clie nttest\n\n")
@@ -4357,7 +4355,6 @@ def test_deepseek_reasoning_content_completion():
         pytest.skip("Model is timing out")
 
 
-
 def test_qwen_text_completion():
     # litellm._turn_on_debug()
     resp = litellm.completion(
@@ -4492,3 +4489,37 @@ def test_completion_gpt_4o_empty_str():
             messages=[{"role": "user", "content": ""}],
         )
         assert resp.choices[0].message.content is not None
+
+
+def test_edit_note():
+    litellm.callbacks = ["langfuse_otel"]
+    response = completion(
+        model="gpt-4o",
+        messages=[
+            {
+                "role": "system",
+                "content": "Your only job is to call the edit_note tool with the content specified in the user's message.",
+            },
+            {
+                "role": "user",
+                "content": "Edit the note with the content: 'This is a test note.'",
+            },
+        ],
+        tools=[
+            {
+                "type": "function",
+                "function": {
+                    "name": "edit_note",
+                    "description": "Edit the note with the content specified in the user's message.",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "content": {"type": "string"},
+                        },
+                    },
+                },
+            },
+        ],
+    )
+
+    return response