diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 90383ddb60d..1204ad51016 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -1427,6 +1427,10 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): if log_span is not None: self._emit_semantic_logs(kwargs, response_obj, log_span) + # Stamp team attributes onto the SERVER (root) span before it is + # closed, so the trace root carries them like every child span. + self._set_team_attributes_on_proxy_span_from_kwargs(kwargs) + # 6. Do NOT end parent span - it should be managed by its creator # External spans (from Langfuse, user code, HTTP headers, global context) must not be closed by LiteLLM # However, proxy-created spans should be closed here. @@ -1439,10 +1443,6 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): ): self._close_proxy_span_ok(parent_span, end_time) - # Stamp team attributes onto the SERVER (root) span before it is - # closed, so the trace root carries them like every child span. - self._set_team_attributes_on_proxy_span_from_kwargs(kwargs) - # close the proxy span explicitly from kwargs metadata # after all child spans (litellm_request, guardrail, raw_request) # have been fully recorded and exported. @@ -1552,9 +1552,19 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): Guarded to the LiteLLM-created proxy span (by name + recording) so externally provided parent spans are never mutated. """ - litellm_params: Final = kwargs.get("litellm_params") or {} - metadata: Final = litellm_params.get("metadata") or {} - proxy_span: Final = metadata.get("litellm_parent_otel_span") + litellm_params_value: Final[Any] = cast(Any, kwargs.get("litellm_params")) # cast-ok: dynamic logging kwargs + litellm_params: Final[Mapping[str, Any]] = ( + litellm_params_value if isinstance(litellm_params_value, Mapping) else {} + ) + metadata_value: Final[Any] = cast(Any, litellm_params.get("metadata")) # cast-ok: dynamic metadata payload + metadata: Final[Mapping[str, Any]] = metadata_value if isinstance(metadata_value, Mapping) else {} + litellm_metadata_value: Final[Any] = cast(Any, litellm_params.get("litellm_metadata")) # cast-ok: dynamic field + litellm_metadata: Final[Mapping[str, Any]] = ( + litellm_metadata_value if isinstance(litellm_metadata_value, Mapping) else {} + ) + proxy_span: Final[Any] = metadata.get("litellm_parent_otel_span") or litellm_metadata.get( + "litellm_parent_otel_span" + ) if ( proxy_span is not None and getattr(proxy_span, "name", None) == LITELLM_PROXY_REQUEST_SPAN_NAME diff --git a/tests/unit/integrations/test_opentelemetry.py b/tests/unit/integrations/test_opentelemetry.py index 9ba09085ce8..8f653443802 100644 --- a/tests/unit/integrations/test_opentelemetry.py +++ b/tests/unit/integrations/test_opentelemetry.py @@ -3,6 +3,7 @@ import concurrent.futures import contextlib import gc import json +import logging import os import sys import threading @@ -19,6 +20,7 @@ import pytest # Adds the grandparent directory to sys.path to allow importing project modules from opentelemetry import trace +from opentelemetry.trace import Span from opentelemetry._logs.severity import SeverityNumber from opentelemetry.sdk._logs import LogData, LogRecord from opentelemetry.sdk._logs import LoggerProvider as OTLoggerProvider @@ -38,6 +40,7 @@ import litellm from litellm.integrations import opentelemetry as otel_module from litellm.integrations.arize.arize_phoenix import ArizePhoenixLogger from litellm.integrations.opentelemetry import ( + LITELLM_PROXY_REQUEST_SPAN_NAME, LITELLM_REQUEST_SPAN_NAME, OpenTelemetry, OpenTelemetryConfig, @@ -7338,3 +7341,55 @@ class TestOpentelemetryUnitTests: # Assert: error.message should be set from error_str using ErrorAttributes constant mock_span.set_attribute.assert_called_with(ErrorAttributes.ERROR_MESSAGE, "Fallback error message") + + +@pytest.mark.parametrize("parent_container", ("metadata", "litellm_metadata")) +def test_success_stamps_team_attributes_before_proxy_span_end( + caplog: pytest.LogCaptureFixture, parent_container: str +) -> None: + exporter: Final[InMemorySpanExporter] = InMemorySpanExporter() + provider: Final[TracerProvider] = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + otel: Final[OpenTelemetry] = OpenTelemetry( + tracer_provider=provider, message_logging=False + ) + server_span: Final[Span] = otel.create_litellm_proxy_request_started_span( + start_time=datetime(2026, 1, 1, tzinfo=timezone.utc), headers={} + ) + parent_metadata: Final[dict[str, object]] = { + "litellm_parent_otel_span": server_span, + "user_api_key_team_id": "test-team-id", + "user_api_key_team_alias": "test-team", + } + litellm_params: Final[dict[str, object]] = {parent_container: parent_metadata} + kwargs: Final[dict[str, object]] = { + "model": "test-model", + "litellm_call_id": "test-call", + "call_type": "completion", + "litellm_params": litellm_params, + "standard_logging_object": { + "metadata": { + "user_api_key_team_id": "test-team-id", + "user_api_key_team_alias": "test-team", + }, + "call_type": "completion", + }, + } + + with caplog.at_level(logging.WARNING, logger="opentelemetry.sdk.trace"): + end_time: Final[datetime] = datetime(2026, 1, 1, 0, 1, tzinfo=timezone.utc) + otel.log_success_event(kwargs, {"id": "test-response"}, end_time, end_time) + + warnings: Final[list[logging.LogRecord]] = [ + record + for record in caplog.records + if record.getMessage() == "Setting attribute on ended span." + ] + assert warnings == [] + + spans: Final[dict[str, ReadableSpan]] = { + span.name: span for span in exporter.get_finished_spans() + } + root: Final[ReadableSpan] = spans[LITELLM_PROXY_REQUEST_SPAN_NAME] + assert root.attributes["metadata.user_api_key_team_id"] == "test-team-id" + assert root.attributes["metadata.user_api_key_team_alias"] == "test-team"