fix(otel): stamp the Langfuse root observation from the post-guardrail request and response

litellm/integrations/otel/langfuse_logger.py

@@ -8,41 +8,26 @@ from litellm.integrations.otel.model.request_io import request_input, response_o
 from litellm.integrations.otel.plumbing.context import request_root_span
 
 if TYPE_CHECKING:
-    from litellm.caching.dual_cache import DualCache
     from litellm.proxy._types import UserAPIKeyAuth
-    from litellm.types.utils import CallTypesLiteral, ModelResponseStream
-
-ROOT_OBSERVATION_IO_CALL_TYPES: Final = frozenset(
-    {"completion", "acompletion", "responses", "aresponses", "anthropic_messages", "aanthropic_messages"}
-)
+    from litellm.types.utils import ModelResponseStream
 
 
 class LangfuseOpenTelemetryV2(OpenTelemetryV2):
     """Stamps the request's input and output on the root observation while it is still recording.
 
     Langfuse shows a trace's input and output from its root observation. The proxy's root span ends
-    when the response is sent, before the success callback runs, so the stamps have to come from the
-    request-task hooks: input at pre-call, output at post-call success or at the end of the stream.
+    when the response is sent, before the success callback runs, so both stamps come from the
+    post-call hooks in the request task: the request as it stands after the pre-call chain and the
+    response as it is returned, for the call types whose response renders as a message.
     """
 
-    async def async_pre_call_hook(
-        self,
-        user_api_key_dict: "UserAPIKeyAuth",
-        cache: "DualCache",
-        data: Mapping[str, object],
-        call_type: "CallTypesLiteral",
-    ) -> None:
-        await super().async_pre_call_hook(user_api_key_dict, cache, data, call_type)
-        if call_type in ROOT_OBSERVATION_IO_CALL_TYPES:
-            self._stamp_root(LANGFUSE_OBSERVATION_INPUT, lambda: request_input(data))
-
     async def async_post_call_success_hook(
         self,
         data: Mapping[str, object],
         user_api_key_dict: "UserAPIKeyAuth",
         response: object,
     ) -> None:
-        self._stamp_root(LANGFUSE_OBSERVATION_OUTPUT, lambda: response_output(response))
+        self._stamp_root_io(data, lambda: response_output(response))
 
     async def async_post_call_streaming_iterator_hook(
         self,
@@ -54,16 +39,22 @@ class LangfuseOpenTelemetryV2(OpenTelemetryV2):
         async for chunk in response:
             relayed.append(chunk)
             yield chunk
-        self._stamp_root(LANGFUSE_OBSERVATION_OUTPUT, lambda: stream_output(tuple(relayed), request_data))
+        self._stamp_root_io(request_data, lambda: stream_output(tuple(relayed), request_data))
 
-    def _stamp_root(self, key: str, render: Callable[[], str | None]) -> None:
+    def _stamp_root_io(self, data: Mapping[str, object], render_output: Callable[[], str | None]) -> None:
         root: Final = request_root_span()
         if root is None or not root.is_recording():
             return
         try:
-            value: Final = render()
+            output: Final = render_output()
+            if output is None:
+                return
+            root.set_attribute(LANGFUSE_OBSERVATION_OUTPUT, output)
+            rendered_input: Final = request_input(data)
         except Exception:  # noqa: BLE001  # telemetry must never fail the request it describes
-            verbose_logger.debug("otel v2 langfuse: could not render %s for the root observation", key, exc_info=True)
+            verbose_logger.debug(
+                "otel v2 langfuse: could not render the root observation input or output", exc_info=True
+            )
             return
-        if value is not None:
-            root.set_attribute(key, value)
+        if rendered_input is not None:
+            root.set_attribute(LANGFUSE_OBSERVATION_INPUT, rendered_input)

litellm/integrations/otel/logger.py

@@ -723,13 +723,14 @@ class OpenTelemetryV2(CustomLogger):
         self,
         user_api_key_dict: "UserAPIKeyAuth",
         cache: "DualCache",
-        data: Mapping[str, object],
+        data: dict,
         call_type: "CallTypesLiteral",
-    ) -> None:
+    ) -> dict:
         self.seed_request_identity(
             user_api_key_dict,
             model=model_from_request_data(data),
         )
+        return data
 
     def record_error_attributes_on_span(
         self,

tests/test_litellm/integrations/otel/test_langfuse_logger.py

@@ -31,6 +31,8 @@ from litellm.types.llms.openai import (  # noqa: E402
 from litellm.types.utils import (  # noqa: E402
     Choices,
     Delta,
+    Embedding,
+    EmbeddingResponse,
     Message,
     ModelResponse,
     ModelResponseStream,
@@ -258,26 +260,41 @@ def test_root_observation_io_survives_the_root_ending_before_the_success_callbac
     assert OUTPUT_ATTR in dict(generation.attributes or {})
 
 
+def test_root_input_is_the_request_as_the_pre_call_chain_left_it():
+    logger, exporter = _logger()
+    raw = {"model": "gpt-5.4-mini", "messages": [{"role": "user", "content": "my ssn is 123-45-6789"}]}
+    masked = {"model": "gpt-5.4-mini", "messages": [{"role": "user", "content": "my ssn is [REDACTED]"}]}
+    response = ModelResponse(choices=[Choices(message=Message(role="assistant", content="noted"))])
+    root = _start_root(logger)
+    asyncio.run(logger.async_pre_call_hook(UserAPIKeyAuth(), DualCache(), raw, "acompletion"))
+    asyncio.run(logger.async_post_call_success_hook(data=masked, user_api_key_dict=UserAPIKeyAuth(), response=response))
+    root.end()
+
+    assert json.loads(_root_attrs(exporter)[INPUT_ATTR]) == masked["messages"]
+
+
 def test_root_already_ended_is_left_alone():
     logger, exporter = _logger()
+    response = ModelResponse(choices=[Choices(message=Message(role="assistant", content="pong"))])
     root = _start_root(logger)
     root.end()
 
-    asyncio.run(logger.async_pre_call_hook(UserAPIKeyAuth(), DualCache(), CHAT_DATA, "acompletion"))
-
-    assert INPUT_ATTR not in _root_attrs(exporter)
-
-
-def test_non_chat_call_types_do_not_stamp_input():
-    logger, exporter = _logger()
-    root = _start_root(logger)
-
     asyncio.run(
-        logger.async_pre_call_hook(UserAPIKeyAuth(), DualCache(), {"model": "e", "input": "ping"}, "aembedding")
+        logger.async_post_call_success_hook(data=CHAT_DATA, user_api_key_dict=UserAPIKeyAuth(), response=response)
     )
-    root.end()
 
-    assert INPUT_ATTR not in _root_attrs(exporter)
+    attrs = _root_attrs(exporter)
+    assert INPUT_ATTR not in attrs and OUTPUT_ATTR not in attrs
+
+
+def test_responses_without_a_message_body_stamp_neither_input_nor_output():
+    logger, exporter = _logger()
+    embedding = EmbeddingResponse(model="e", data=[Embedding(embedding=[0.1], index=0, object="embedding")])
+
+    _run_request(logger, {"model": "e", "input": "ping"}, "aembedding", embedding)
+
+    attrs = _root_attrs(exporter)
+    assert INPUT_ATTR not in attrs and OUTPUT_ATTR not in attrs
 
 
 def test_unrenderable_output_never_raises_into_the_request():
@@ -285,7 +302,8 @@ def test_unrenderable_output_never_raises_into_the_request():
 
     _run_request(logger, CHAT_DATA, "acompletion", object())
 
-    assert OUTPUT_ATTR not in _root_attrs(exporter)
+    attrs = _root_attrs(exporter)
+    assert INPUT_ATTR not in attrs and OUTPUT_ATTR not in attrs
 
 
 @pytest.mark.parametrize(
@@ -331,7 +349,11 @@ def test_langfuse_otel_preset_builds_a_logger_that_stamps_the_root(monkeypatch):
         assert built is not None
         assert _maybe_construct_otel_v2("langfuse_otel", loggers) is built
         root = _start_root(built)
-        asyncio.run(built.async_pre_call_hook(UserAPIKeyAuth(), DualCache(), CHAT_DATA, "acompletion"))
-        assert INPUT_ATTR in dict(root.attributes or {})
+        response = ModelResponse(choices=[Choices(message=Message(role="assistant", content="pong"))])
+        asyncio.run(
+            built.async_post_call_success_hook(data=CHAT_DATA, user_api_key_dict=UserAPIKeyAuth(), response=response)
+        )
+        attrs = dict(root.attributes or {})
+        assert INPUT_ATTR in attrs and OUTPUT_ATTR in attrs
     finally:
         is_otel_v2_enabled.cache_clear()