litellm/tests/unit/integrations/otel/test_langfuse_logger.py

"""Tests for the Langfuse OTel v2 loggers: the trace name and the root observation's input and output are
stamped from the request task while the root span is still recording, so Langfuse can show them on the trace."""

import asyncio
import json
from collections.abc import AsyncIterator, Mapping, Sequence
from typing import Final

import pytest

pytest.importorskip("opentelemetry")

from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter  # noqa: E402

import litellm  # noqa: E402
from litellm.caching.dual_cache import DualCache  # 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.proxy.utils import ProxyLogging  # noqa: E402
from litellm.types.llms.openai import (  # noqa: E402
    ResponseCompletedEvent,
    ResponsesAPIResponse,
    ResponsesAPIStreamEvents,
)
from litellm.types.utils import (  # noqa: E402
    Choices,
    Delta,
    Embedding,
    EmbeddingResponse,
    Message,
    ModelResponse,
    ModelResponseStream,
    StreamingChoices,
)

INPUT_ATTR: Final = "langfuse.observation.input"
OUTPUT_ATTR: Final = "langfuse.observation.output"
TRACE_NAME_ATTR: Final = "langfuse.trace.name"
TRACE_CONTROL_ATTRS: Final = (TRACE_NAME_ATTR, "user.id", "session.id", "langfuse.trace.tags")
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_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_post_call_success_hook(data=CHAT_DATA, user_api_key_dict=UserAPIKeyAuth(), response=response)
    )

    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():
    logger, exporter = _logger()

    _run_request(logger, CHAT_DATA, "acompletion", object())

    attrs = _root_attrs(exporter)
    assert INPUT_ATTR not in attrs and OUTPUT_ATTR not in attrs


def _run_named_request(
    logger: OpenTelemetryV2, exporter: InMemorySpanExporter, litellm_params: Mapping[str, object]
) -> tuple[Mapping[str, object], Mapping[str, object]]:
    response: Final = ModelResponse(choices=[Choices(message=Message(role="assistant", content="pong"))])
    root: Final = _start_root(logger)
    logger.log_pre_api_call(
        model="gpt-5.4-mini", messages=[], kwargs={"litellm_call_id": "call_1", "litellm_params": litellm_params}
    )
    root.end()
    payload: Final = {
        "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": litellm_params}, response, None, None
        )
    )
    generation: Final = next(
        span for span in exporter.get_finished_spans() if span.name != LITELLM_PROXY_REQUEST_SPAN_NAME
    )
    return _root_attrs(exporter), dict(generation.attributes or {})


@pytest.mark.parametrize("capture", ["span_only", "no_content"])
def test_langfuse_trace_name_header_names_the_root_and_the_generation_over_body_metadata(capture):
    logger, exporter = _logger(capture=capture)

    root_attrs, generation_attrs = _run_named_request(
        logger,
        exporter,
        {
            "metadata": {"trace_name": "from-body"},
            "proxy_server_request": {"headers": {"langfuse_trace_name": "from-header"}},
        },
    )

    assert root_attrs[TRACE_NAME_ATTR] == "from-header"
    assert generation_attrs[TRACE_NAME_ATTR] == "from-header"


def test_body_metadata_trace_name_names_the_root_and_the_generation():
    logger, exporter = _logger()

    root_attrs, generation_attrs = _run_named_request(
        logger, exporter, {"metadata": {"trace_name": "from-body"}, "proxy_server_request": {"headers": {}}}
    )

    assert root_attrs[TRACE_NAME_ATTR] == "from-body"
    assert generation_attrs[TRACE_NAME_ATTR] == "from-body"


def test_unnamed_request_leaves_the_trace_name_off_both_spans():
    logger, exporter = _logger()

    root_attrs, generation_attrs = _run_named_request(logger, exporter, {"proxy_server_request": {"headers": {}}})

    assert TRACE_NAME_ATTR not in root_attrs and TRACE_NAME_ATTR not in generation_attrs


@pytest.mark.parametrize("capture", ["span_only", "no_content"])
def test_body_metadata_user_session_and_tags_land_on_the_root_and_the_generation(capture):
    logger, exporter = _logger(capture=capture)

    root_attrs, generation_attrs = _run_named_request(
        logger,
        exporter,
        {
            "metadata": {
                "trace_user_id": "user-42",
                "session_id": "session-7",
                "tags": ["prod", "eval", "nightly"],
                "user_api_key_team_id": "team-from-proxy",
            },
            "proxy_server_request": {"headers": {}},
        },
    )

    for attrs in (root_attrs, generation_attrs):
        assert attrs["user.id"] == "user-42"
        assert attrs["session.id"] == "session-7"
        assert tuple(attrs["langfuse.trace.tags"]) == ("prod", "eval", "nightly")
        assert TRACE_NAME_ATTR not in attrs


def test_langfuse_user_and_session_headers_beat_body_metadata_on_both_spans():
    logger, exporter = _logger()

    root_attrs, generation_attrs = _run_named_request(
        logger,
        exporter,
        {
            "metadata": {"trace_user_id": "from-body", "session_id": "from-body"},
            "proxy_server_request": {
                "headers": {"langfuse_trace_user_id": "from-header", "langfuse_session_id": "from-header-s"}
            },
        },
    )

    for attrs in (root_attrs, generation_attrs):
        assert attrs["user.id"] == "from-header"
        assert attrs["session.id"] == "from-header-s"


def test_caller_metadata_cannot_override_the_proxy_team_identity():
    logger, exporter = _logger()
    response: Final = ModelResponse(choices=[Choices(message=Message(role="assistant", content="pong"))])
    litellm_params: Final = {
        "metadata": {"trace_user_id": "u", "trace_metadata": {"team_id": "spoofed"}, "team_id": "spoofed"}
    }
    logger.log_pre_api_call(
        model="gpt-5.4-mini", messages=[], kwargs={"litellm_call_id": "call_1", "litellm_params": litellm_params}
    )
    payload: Final = {
        "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": {
            "user_api_key_team_id": "real-team",
            "user_api_key_team_alias": "real-alias",
            "team_id": "spoofed",
            "team_alias": "spoofed",
        },
        "hidden_params": {},
    }
    asyncio.run(
        logger.async_log_success_event(
            {"standard_logging_object": payload, "litellm_params": litellm_params}, response, None, None
        )
    )

    attrs: Final = dict(exporter.get_finished_spans()[0].attributes or {})
    assert attrs["user.id"] == "u"
    assert attrs["langfuse.trace.metadata.team_id"] == "real-team"
    assert attrs["langfuse.trace.metadata.team_alias"] == "real-alias"
    assert "langfuse.trace.metadata" not in attrs and "langfuse.trace.id" not in attrs


def test_a_request_without_trace_controls_stamps_none_of_them():
    logger, exporter = _logger()

    root_attrs, generation_attrs = _run_named_request(
        logger, exporter, {"metadata": {"user_api_key_team_id": "t1", "tags": []}, "proxy_server_request": {"headers": {}}}
    )

    assert set(TRACE_CONTROL_ATTRS).isdisjoint(root_attrs)
    assert set(TRACE_CONTROL_ATTRS).isdisjoint(generation_attrs)


@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)

    _run_request(logger, CHAT_DATA, "acompletion", ModelResponse())
    attrs = _root_attrs(exporter)
    assert INPUT_ATTR not in attrs and OUTPUT_ATTR not in attrs


@pytest.mark.parametrize(
    ("capture", "mappers", "relays_streams"),
    [
        ("span_only", ("genai", "langfuse"), True),
        ("no_content", ("genai", "langfuse"), False),
        ("span_only", ("genai",), False),
    ],
)
def test_only_langfuse_content_capture_takes_proxy_streams_off_the_fast_path(
    monkeypatch, capture, mappers, relays_streams
):
    logger, _ = _logger(capture=capture, mappers=mappers)
    monkeypatch.setattr(litellm, "callbacks", [logger])

    assert ProxyLogging._callback_capabilities().has_iterator_override is relays_streams


def test_langfuse_otel_preset_builds_a_logger_that_stamps_the_root(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 built is not None
        assert _maybe_construct_otel_v2("langfuse_otel", loggers) is built
        root = _start_root(built)
        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()