fix(otel): stamp Langfuse root observation input and output from the request task

This commit is contained in:
mateo-berri 2026-09-02 11:57:45 -07:00
parent 4049a075bd
commit cc2cbb36f3
6 changed files with 519 additions and 11 deletions

View file

@ -0,0 +1,69 @@
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Mapping
from typing import TYPE_CHECKING, Final
from litellm._logging import verbose_logger
from litellm.integrations.otel.logger import OpenTelemetryV2
from litellm.integrations.otel.mappers.langfuse import LANGFUSE_OBSERVATION_INPUT, LANGFUSE_OBSERVATION_OUTPUT
from litellm.integrations.otel.model.request_io import request_input, response_output, stream_output
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"}
)
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.
"""
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))
async def async_post_call_streaming_iterator_hook(
self,
user_api_key_dict: "UserAPIKeyAuth",
response: "AsyncIterator[ModelResponseStream]",
request_data: Mapping[str, object],
) -> "AsyncGenerator[ModelResponseStream, None]":
relayed: Final[list[ModelResponseStream]] = [] # mutable-ok: relayed as they arrive, assembled at end of stream
async for chunk in response:
relayed.append(chunk)
yield chunk
self._stamp_root(LANGFUSE_OBSERVATION_OUTPUT, lambda: stream_output(tuple(relayed), request_data))
def _stamp_root(self, key: str, render: Callable[[], str | None]) -> None:
root: Final = request_root_span()
if root is None or not root.is_recording():
return
try:
value: Final = render()
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)
return
if value is not None:
root.set_attribute(key, value)

View file

@ -4,6 +4,7 @@ from collections import OrderedDict
from collections.abc import Callable, Iterator, Mapping, Sequence
from contextlib import contextmanager
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, cast
from opentelemetry.context import Context, attach, get_current
@ -722,14 +723,13 @@ class OpenTelemetryV2(CustomLogger):
self,
user_api_key_dict: "UserAPIKeyAuth",
cache: "DualCache",
data: dict,
data: Mapping[str, object],
call_type: "CallTypesLiteral",
) -> dict:
) -> None:
self.seed_request_identity(
user_api_key_dict,
model=model_from_request_data(data),
)
return data
def record_error_attributes_on_span(
self,
@ -909,3 +909,29 @@ def phase_span(name: str) -> "Iterator[Span | None]":
return
with logger.start_phase_span(name) as span:
yield span
def build_otel_v2_logger(
config: OpenTelemetryV2Config,
callback_name: str | None = None,
tracer_provider: TracerProvider | None = None,
logger_provider: LoggerProvider | None = None,
meter_provider: "MeterProvider | None" = None,
settings: Mapping[str, object] = MappingProxyType({}),
) -> OpenTelemetryV2:
return _logger_class(config)(
config=config,
callback_name=callback_name,
tracer_provider=tracer_provider,
logger_provider=logger_provider,
meter_provider=meter_provider,
**settings,
)
def _logger_class(config: OpenTelemetryV2Config) -> type[OpenTelemetryV2]:
if "langfuse" not in config.mapper_names or not config.capture_span_content:
return OpenTelemetryV2
from litellm.integrations.otel.langfuse_logger import LangfuseOpenTelemetryV2
return LangfuseOpenTelemetryV2

View file

@ -11,6 +11,7 @@ the JSON-serialized payloads. ``_llm_call`` just applies both tables.
import json
from collections.abc import Callable
from typing import Final
from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue, SpanData
from litellm.integrations.otel.mappers.utils import (
@ -25,6 +26,9 @@ from litellm.integrations.otel.model.payloads import (
LLMUsage,
)
LANGFUSE_OBSERVATION_INPUT: Final = "langfuse.observation.input"
LANGFUSE_OBSERVATION_OUTPUT: Final = "langfuse.observation.output"
class LangfuseMapper:
_LLM_CALL_ATTRS: dict[str, Callable[[LLMCallSpanData], AttrValue | None]] = {
@ -56,8 +60,8 @@ class LangfuseMapper:
"langfuse.observation.model.parameters": lambda d: json_if(
collect(LangfuseMapper._MODEL_PARAMS, d.request_params)
),
"langfuse.observation.input": lambda d: serialize_messages(d.messages_in),
"langfuse.observation.output": lambda d: serialize_messages(output_messages(d)),
LANGFUSE_OBSERVATION_INPUT: lambda d: serialize_messages(d.messages_in),
LANGFUSE_OBSERVATION_OUTPUT: lambda d: serialize_messages(output_messages(d)),
"langfuse.observation.usage_details": lambda d: json_if(collect(LangfuseMapper._USAGE_FIELDS, d.usage)),
"langfuse.observation.cost_details": lambda d: (
json.dumps({"total": d.response_cost}) if d.response_cost is not None else None

View file

@ -0,0 +1,90 @@
from collections.abc import Mapping, Sequence
from typing import Final, Literal
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm.integrations.otel.mappers.utils import json_or_none
from litellm.proxy.guardrails.anthropic_sse import assemble_anthropic_sse_stream, is_raw_sse_stream
from litellm.types.llms.openai import ResponseCompletedEvent, ResponsesAPIResponse
from litellm.types.utils import ModelResponse, ModelResponseStream
_SYSTEM_KEYS: Final = ("system", "instructions")
_TURNS: Final = TypeAdapter(tuple[object, ...])
_MESSAGES: Final = TypeAdapter(list[object] | None)
class _Turn(TypedDict):
role: ReadOnly[str]
content: ReadOnly[object]
class _AnthropicMessage(BaseModel):
model_config = ConfigDict(frozen=True)
type: Literal["message"] = Field(exclude=True)
role: str = "assistant"
content: object = None
def request_input(data: Mapping[str, object]) -> str | None:
turns: Final = data.get("messages", data.get("input"))
if turns is None:
return None
return json_or_none((*_system_turns(data), *_user_turns(turns)))
def _system_turns(data: Mapping[str, object]) -> tuple[_Turn, ...]:
return tuple(_Turn(role="system", content=data[key]) for key in _SYSTEM_KEYS if data.get(key) is not None)
def _user_turns(turns: object) -> tuple[object, ...]:
if isinstance(turns, str):
return (_Turn(role="user", content=turns),)
try:
return _TURNS.validate_python(turns)
except ValidationError:
return (_Turn(role="user", content=turns),)
def response_output(response: object) -> str | None:
match response:
case ModelResponse():
return json_or_none(tuple(choice.message.model_dump(exclude_none=True) for choice in response.choices))
case ResponsesAPIResponse():
return json_or_none(response.model_dump(exclude_none=True).get("output"))
case _:
return _anthropic_message_output(response)
def _anthropic_message_output(message: object) -> str | None:
try:
parsed: Final = _AnthropicMessage.model_validate(message)
except ValidationError:
return None
return json_or_none((parsed.model_dump(),))
def stream_output(chunks: Sequence[object], data: Mapping[str, object]) -> str | None:
if not chunks:
return None
if is_raw_sse_stream(chunks):
return response_output(assemble_anthropic_sse_stream(chunks))
if all(isinstance(chunk, ModelResponseStream) for chunk in chunks):
return response_output(_assembled_chat_stream(chunks, data))
return response_output(_completed_response(chunks))
def _assembled_chat_stream(chunks: Sequence[object], data: Mapping[str, object]) -> object:
try:
return litellm.stream_chunk_builder( # pyright: ignore[reportUnknownMemberType] # upstream types chunks as a bare list
chunks=list(chunks), # mutable-ok: stream_chunk_builder takes a list
messages=_MESSAGES.validate_python(data.get("messages")),
)
except (litellm.APIError, ValidationError):
return None
def _completed_response(chunks: Sequence[object]) -> ResponsesAPIResponse | None:
return next((chunk.response for chunk in reversed(chunks) if isinstance(chunk, ResponseCompletedEvent)), None)

View file

@ -4390,13 +4390,15 @@ def _init_custom_logger_compatible_class(
from litellm.integrations.otel.model.config import is_otel_v2_enabled
if is_otel_v2_enabled():
from litellm.integrations.otel.logger import OpenTelemetryV2
from litellm.integrations.otel.logger import OpenTelemetryV2, build_otel_v2_logger
from litellm.integrations.otel.model.config import OpenTelemetryV2Config
for callback in _in_memory_loggers:
if type(callback) is OpenTelemetryV2:
if isinstance(callback, OpenTelemetryV2):
return callback
otel_logger_v2: Final = OpenTelemetryV2(
**_get_custom_logger_settings_from_proxy_server(callback_name=logging_integration)
otel_settings: Final = _get_custom_logger_settings_from_proxy_server(callback_name=logging_integration)
otel_logger_v2: Final = build_otel_v2_logger(
config=OpenTelemetryV2Config(**otel_settings), settings=otel_settings
)
_in_memory_loggers.append(otel_logger_v2)
_maybe_auto_initialize_arize_phoenix(_in_memory_loggers)
@ -4759,7 +4761,7 @@ def _maybe_construct_otel_v2(callback_name: str, _in_memory_loggers: list[Custom
if not is_otel_v2_enabled():
return None
from litellm.integrations.otel.logger import OpenTelemetryV2
from litellm.integrations.otel.logger import OpenTelemetryV2, build_otel_v2_logger
from litellm.integrations.otel.presets import PRESET_BY_CALLBACK
preset_fn: Final = PRESET_BY_CALLBACK.get(callback_name)
@ -4774,7 +4776,7 @@ def _maybe_construct_otel_v2(callback_name: str, _in_memory_loggers: list[Custom
# If env vars are missing or the preset raises, defer to the legacy path
# so customers get the same error story they had before V2 landed.
return None
v2_logger: Final = OpenTelemetryV2(config=config, callback_name=callback_name)
v2_logger: Final = build_otel_v2_logger(config=config, callback_name=callback_name)
_in_memory_loggers.append(v2_logger)
return v2_logger

View file

@ -0,0 +1,317 @@
"""Tests for ``LangfuseOpenTelemetryV2``: the root observation's input and output are stamped from the
request-task hooks, while the root span is still recording, so Langfuse can show them on the trace."""
import asyncio
import json
from collections.abc import AsyncIterator, Sequence
from typing import Final
import pytest
pytest.importorskip("opentelemetry")
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter # noqa: E402
from litellm.caching.dual_cache import DualCache # noqa: E402
from litellm.integrations.otel.langfuse_logger import LangfuseOpenTelemetryV2 # noqa: E402
from litellm.integrations.otel.logger import OpenTelemetryV2, build_otel_v2_logger # noqa: E402
from litellm.integrations.otel.model.config import OpenTelemetryV2Config, is_otel_v2_enabled # noqa: E402
from litellm.integrations.otel.model.spans import LITELLM_PROXY_REQUEST_SPAN_NAME, SpanRole # noqa: E402
from litellm.integrations.otel.plumbing import context as otel_context # noqa: E402
from litellm.integrations.otel.plumbing import providers # noqa: E402
from litellm.integrations.otel.plumbing.context import set_request_root_span # noqa: E402
from litellm.litellm_core_utils.litellm_logging import _maybe_construct_otel_v2 # noqa: E402
from litellm.proxy._types import UserAPIKeyAuth # noqa: E402
from litellm.types.llms.openai import ( # noqa: E402
ResponseCompletedEvent,
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
)
from litellm.types.utils import ( # noqa: E402
Choices,
Delta,
Message,
ModelResponse,
ModelResponseStream,
StreamingChoices,
)
INPUT_ATTR: Final = "langfuse.observation.input"
OUTPUT_ATTR: Final = "langfuse.observation.output"
CHAT_DATA: Final = {"model": "gpt-5.4-mini", "messages": [{"role": "user", "content": "ping"}]}
@pytest.fixture(autouse=True)
def _reset_request_root_span():
otel_context._request_root_span.set(None)
yield
otel_context._request_root_span.set(None)
def _logger(*, capture: str = "span_only", mappers: Sequence[str] = ("genai", "langfuse")):
cfg = OpenTelemetryV2Config(exporter="in_memory", mapper_names=list(mappers), capture_message_content=capture)
exporter = InMemorySpanExporter()
tracer_provider = providers.build_tracer_provider(cfg, exporter=exporter)
return build_otel_v2_logger(config=cfg, tracer_provider=tracer_provider), exporter
def _start_root(logger):
root = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME)
set_request_root_span(root)
return root
def _root_attrs(exporter):
by_name = {span.name: span for span in exporter.get_finished_spans()}
return dict(by_name[LITELLM_PROXY_REQUEST_SPAN_NAME].attributes or {})
def _run_request(logger, data: dict, call_type: str, response: object):
root = _start_root(logger)
asyncio.run(logger.async_pre_call_hook(UserAPIKeyAuth(), DualCache(), data, call_type))
asyncio.run(logger.async_post_call_success_hook(data=data, user_api_key_dict=UserAPIKeyAuth(), response=response))
root.end()
async def _relay(logger, chunks: Sequence[object], data: dict) -> list[object]:
async def source() -> AsyncIterator[object]:
for chunk in chunks:
yield chunk
return [chunk async for chunk in logger.async_post_call_streaming_iterator_hook(UserAPIKeyAuth(), source(), data)]
def _run_stream(logger, data: dict, chunks: Sequence[object]) -> list[object]:
root = _start_root(logger)
asyncio.run(logger.async_pre_call_hook(UserAPIKeyAuth(), DualCache(), data, "acompletion"))
relayed = asyncio.run(_relay(logger, chunks, data))
root.end()
return relayed
def _chat_chunk(content: str | None, finish_reason: str | None = None) -> ModelResponseStream:
return ModelResponseStream(
id="chatcmpl-1",
created=1,
model="gpt-5.4-mini",
choices=[StreamingChoices(index=0, delta=Delta(content=content), finish_reason=finish_reason)],
)
def _responses_api_response() -> ResponsesAPIResponse:
return ResponsesAPIResponse(
id="resp_1",
created_at=1,
output=[
{
"type": "message",
"id": "msg_1",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "pong", "annotations": []}],
}
],
)
def _anthropic_sse_frames() -> tuple[bytes, ...]:
events = (
{
"type": "message_start",
"message": {
"id": "msg_1",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-5",
"content": [],
"stop_reason": None,
"usage": {"input_tokens": 1, "output_tokens": 0},
},
},
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "po"}},
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ng"}},
{"type": "content_block_stop", "index": 0},
{"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 2}},
{"type": "message_stop"},
)
return tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events)
def test_chat_request_stamps_root_observation_input_and_output():
logger, exporter = _logger()
response = ModelResponse(choices=[Choices(message=Message(role="assistant", content="pong"))])
_run_request(logger, CHAT_DATA, "acompletion", response)
attrs = _root_attrs(exporter)
assert json.loads(attrs[INPUT_ATTR]) == [{"role": "user", "content": "ping"}]
output = json.loads(attrs[OUTPUT_ATTR])
assert [(turn["role"], turn["content"]) for turn in output] == [("assistant", "pong")]
def test_responses_request_folds_instructions_into_input_and_stamps_output_items():
logger, exporter = _logger()
data = {"model": "gpt-5.4-mini", "instructions": "be terse", "input": "ping"}
_run_request(logger, data, "aresponses", _responses_api_response())
attrs = _root_attrs(exporter)
assert json.loads(attrs[INPUT_ATTR]) == [
{"role": "system", "content": "be terse"},
{"role": "user", "content": "ping"},
]
output = json.loads(attrs[OUTPUT_ATTR])
assert output[0]["role"] == "assistant"
assert output[0]["content"][0]["text"] == "pong"
def test_anthropic_messages_request_folds_system_into_input_and_stamps_content_blocks():
logger, exporter = _logger()
data = {"model": "claude-sonnet-4-5", "system": "be terse", "messages": [{"role": "user", "content": "ping"}]}
response = {"type": "message", "role": "assistant", "content": [{"type": "text", "text": "pong"}]}
_run_request(logger, data, "aanthropic_messages", response)
attrs = _root_attrs(exporter)
assert json.loads(attrs[INPUT_ATTR]) == [
{"role": "system", "content": "be terse"},
{"role": "user", "content": "ping"},
]
assert json.loads(attrs[OUTPUT_ATTR]) == [{"role": "assistant", "content": [{"type": "text", "text": "pong"}]}]
def test_chat_stream_relays_chunks_untouched_and_stamps_assembled_output():
logger, exporter = _logger()
chunks = (_chat_chunk("po"), _chat_chunk("ng"), _chat_chunk(None, finish_reason="stop"))
relayed = _run_stream(logger, CHAT_DATA, chunks)
assert [id(chunk) for chunk in relayed] == [id(chunk) for chunk in chunks]
output = json.loads(_root_attrs(exporter)[OUTPUT_ATTR])
assert [(turn["role"], turn["content"]) for turn in output] == [("assistant", "pong")]
def test_responses_stream_stamps_output_from_the_completed_event():
logger, exporter = _logger()
completed = ResponseCompletedEvent(
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=_responses_api_response()
)
chunks = ({"type": "response.created"}, {"type": "response.output_text.delta", "delta": "pong"}, completed)
relayed = _run_stream(logger, {"model": "gpt-5.4-mini", "input": "ping"}, chunks)
assert relayed == list(chunks)
output = json.loads(_root_attrs(exporter)[OUTPUT_ATTR])
assert output[0]["content"][0]["text"] == "pong"
def test_anthropic_sse_stream_stamps_output_from_the_assembled_frames():
logger, exporter = _logger()
frames = _anthropic_sse_frames()
relayed = _run_stream(
logger, {"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "ping"}]}, frames
)
assert relayed == list(frames)
output = json.loads(_root_attrs(exporter)[OUTPUT_ATTR])
assert [(turn["role"], turn["content"]) for turn in output] == [("assistant", "pong")]
def test_root_observation_io_survives_the_root_ending_before_the_success_callback():
logger, exporter = _logger()
response = ModelResponse(choices=[Choices(message=Message(role="assistant", content="pong"))])
root = _start_root(logger)
asyncio.run(logger.async_pre_call_hook(UserAPIKeyAuth(), DualCache(), CHAT_DATA, "acompletion"))
logger.log_pre_api_call(
model="gpt-5.4-mini",
messages=[],
kwargs={"litellm_call_id": "call_1", "litellm_params": {"metadata": {}}},
)
asyncio.run(
logger.async_post_call_success_hook(data=CHAT_DATA, user_api_key_dict=UserAPIKeyAuth(), response=response)
)
root.end()
payload = {
"call_type": "acompletion",
"custom_llm_provider": "openai",
"model": "gpt-5.4-mini",
"messages": CHAT_DATA["messages"],
"response": response.model_dump(),
"status": "success",
"litellm_call_id": "call_1",
"metadata": {},
"hidden_params": {},
}
asyncio.run(
logger.async_log_success_event(
{"standard_logging_object": payload, "litellm_params": {"metadata": {}}}, response, None, None
)
)
attrs = _root_attrs(exporter)
assert INPUT_ATTR in attrs and OUTPUT_ATTR in attrs
generation = next(span for span in exporter.get_finished_spans() if span.name != LITELLM_PROXY_REQUEST_SPAN_NAME)
assert OUTPUT_ATTR in dict(generation.attributes or {})
def test_root_already_ended_is_left_alone():
logger, exporter = _logger()
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")
)
root.end()
assert INPUT_ATTR not in _root_attrs(exporter)
def test_unrenderable_output_never_raises_into_the_request():
logger, exporter = _logger()
_run_request(logger, CHAT_DATA, "acompletion", object())
assert OUTPUT_ATTR not in _root_attrs(exporter)
@pytest.mark.parametrize(
("capture", "mappers"),
[("no_content", ("genai", "langfuse")), ("span_only", ("genai",))],
)
def test_factory_keeps_the_base_logger_unless_langfuse_content_capture_is_on(capture, mappers):
logger, exporter = _logger(capture=capture, mappers=mappers)
assert type(logger) is OpenTelemetryV2
_run_request(logger, CHAT_DATA, "acompletion", ModelResponse())
attrs = _root_attrs(exporter)
assert INPUT_ATTR not in attrs and OUTPUT_ATTR not in attrs
def test_langfuse_otel_preset_builds_the_langfuse_logger(monkeypatch):
monkeypatch.setenv("LITELLM_OTEL_V2", "true")
monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "pk")
monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk")
monkeypatch.setenv("LANGFUSE_HOST", "https://cloud.langfuse.com")
monkeypatch.setenv("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", "span_only")
is_otel_v2_enabled.cache_clear()
loggers: list = []
try:
built = _maybe_construct_otel_v2("langfuse_otel", loggers)
assert isinstance(built, LangfuseOpenTelemetryV2)
assert _maybe_construct_otel_v2("langfuse_otel", loggers) is built
finally:
is_otel_v2_enabled.cache_clear()