From 86f0ffd42971cd9525e44827ec5d7f124efd4299 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Tue, 19 May 2026 15:31:25 -0700 Subject: [PATCH] feat: propagate team_id and team_alias to all child OTEL spans (#28273) - Add `_set_team_attributes_on_span` helper to stamp team_id/team_alias onto any span, ensuring these attributes are not limited to the root litellm_request span - Add `_set_team_attributes_from_kwargs` helper to extract team metadata from the standard_logging_object in kwargs and apply them to a span - Apply team attributes to raw request spans via `_maybe_log_raw_request` so downstream consumers can filter traces by team without needing the root span - Apply team attributes to guardrail spans so guardrail activity can be correlated to teams in tracing backends - Apply team attributes to exception logging spans to preserve team context during failure paths - Add comprehensive unit tests covering all new helpers, including edge cases where metadata or standard_logging_object is absent Co-authored-by: Yassin Kortam --- litellm/integrations/opentelemetry.py | 82 +++++ .../integrations/test_opentelemetry.py | 143 +++++++++ .../test_otel_team_attributes_matrix.py | 285 ++++++++++++++++++ 3 files changed, 510 insertions(+) create mode 100644 tests/test_litellm/integrations/test_otel_team_attributes_matrix.py diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index ced15a01660..a70574952b8 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -673,6 +673,15 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): if parent_otel_span is not None: parent_otel_span.set_status(Status(StatusCode.ERROR)) + # Stamp team attributes onto the SERVER (root) span too, so the + # trace root is team-filterable on the failure path like the + # child exception span below. + self._set_team_attributes_on_span( + span=parent_otel_span, + team_id=user_api_key_dict.team_id, + team_alias=user_api_key_dict.team_alias, + ) + # Stamp structured error attrs on the SERVER span itself; the # failure path otherwise only sets its status (_handle_failure # records on the litellm_request child span). Inline import: @@ -709,6 +718,11 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): key="exception", value=str(original_exception), ) + self._set_team_attributes_on_span( + span=exception_logging_span, + team_id=user_api_key_dict.team_id, + team_alias=user_api_key_dict.team_alias, + ) exception_logging_span.set_status(Status(StatusCode.ERROR)) exception_logging_span.end(end_time=self._to_ns(datetime.now())) @@ -1012,6 +1026,10 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): ): parent_span.end(end_time=self._to_ns(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. @@ -1070,8 +1088,70 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): ) raw_span.set_status(Status(StatusCode.OK)) self.set_raw_request_attributes(raw_span, kwargs, response_obj) + self._set_team_attributes_from_kwargs(raw_span, kwargs) raw_span.end(end_time=self._to_ns(end_time)) + def _set_team_attributes_on_span( + self, + span: Span, + team_id: Optional[str], + team_alias: Optional[str], + ) -> None: + """Stamp team_id / team_alias onto a span so every child span of a + litellm_request trace carries them, not just the root span. + + Empty strings are treated as absent: a request made with the master + key or a team-less virtual key carries ``user_api_key_team_id=""`` + in ``standard_logging_object.metadata``; propagating that to every + span only adds noise that makes traces look mis-instrumented. + """ + if team_id: + self.safe_set_attribute( + span=span, + key="metadata.user_api_key_team_id", + value=team_id, + ) + if team_alias: + self.safe_set_attribute( + span=span, + key="metadata.user_api_key_team_alias", + value=team_alias, + ) + + def _set_team_attributes_from_kwargs(self, span: Span, kwargs: dict) -> None: + """Pull team_id / team_alias from the standard logging metadata in kwargs and stamp them onto span.""" + std_log = kwargs.get("standard_logging_object") + md: dict = {} + if isinstance(std_log, dict): + md = std_log.get("metadata") or {} + elif std_log is not None: + md = getattr(std_log, "metadata", None) or {} + self._set_team_attributes_on_span( + span=span, + team_id=md.get("user_api_key_team_id"), + team_alias=md.get("user_api_key_team_alias"), + ) + + def _set_team_attributes_on_proxy_span_from_kwargs(self, kwargs: dict) -> None: + """Stamp team attributes onto the proxy SERVER (root) span so the + trace root is filterable by team, not just its children. The root + span is created in auth before the team is resolved and is + otherwise only closed (never re-attributed) on the success path. + + Guarded to the LiteLLM-created proxy span (by name + recording) so + externally provided parent spans are never mutated. + """ + litellm_params = kwargs.get("litellm_params") or {} + metadata = litellm_params.get("metadata") or {} + proxy_span = metadata.get("litellm_parent_otel_span") + if ( + proxy_span is not None + and getattr(proxy_span, "name", None) == LITELLM_PROXY_REQUEST_SPAN_NAME + and hasattr(proxy_span, "is_recording") + and proxy_span.is_recording() + ): + self._set_team_attributes_from_kwargs(proxy_span, kwargs) + def _record_metrics(self, kwargs, response_obj, start_time, end_time): duration_s = (end_time - start_time).total_seconds() params = kwargs.get("litellm_params") or {} @@ -1537,6 +1617,8 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): value=guardrail_information.get("guardrail_response"), ) + self._set_team_attributes_from_kwargs(guardrail_span, kwargs) + guardrail_span.end(end_time=self._to_ns(end_time_datetime)) def _handle_failure(self, kwargs, response_obj, start_time, end_time): diff --git a/tests/test_litellm/integrations/test_opentelemetry.py b/tests/test_litellm/integrations/test_opentelemetry.py index 87298b5a7a3..6de855262bd 100644 --- a/tests/test_litellm/integrations/test_opentelemetry.py +++ b/tests/test_litellm/integrations/test_opentelemetry.py @@ -88,6 +88,149 @@ class TestOpenTelemetryGuardrails(unittest.TestCase): otel.tracer.start_span.assert_not_called() +class TestOpenTelemetryTeamAttributesOnChildSpans(unittest.TestCase): + """team_id / team_alias must land on every child span of a + litellm_request trace, not only the root litellm_request span.""" + + def _slo_metadata(self): + return { + "user_api_key_team_id": "team-123", + "user_api_key_team_alias": "my-team", + } + + @patch("litellm.integrations.opentelemetry.datetime") + def test_guardrail_span_has_team_attributes(self, mock_datetime): + otel = OpenTelemetry() + otel.tracer = MagicMock() + mock_span = MagicMock() + otel.tracer.start_span.return_value = mock_span + + guardrail_info = { + "guardrail_name": "test_guardrail", + "guardrail_mode": "input", + "guardrail_response": "filtered_content", + "start_time": 1609459200.0, + "end_time": 1609459201.0, + } + kwargs = { + "standard_logging_object": { + "guardrail_information": [guardrail_info], + "metadata": self._slo_metadata(), + } + } + + otel._create_guardrail_span(kwargs=kwargs, context=None) + + mock_span.set_attribute.assert_any_call( + "metadata.user_api_key_team_id", "team-123" + ) + mock_span.set_attribute.assert_any_call( + "metadata.user_api_key_team_alias", "my-team" + ) + + @patch.dict(os.environ, {"OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT": ""}) + @patch("litellm.turn_off_message_logging", False) + def test_raw_request_span_has_team_attributes(self): + otel = OpenTelemetry() + otel.message_logging = True + + mock_tracer = MagicMock() + mock_span = MagicMock() + mock_tracer.start_span.return_value = mock_span + otel.get_tracer_to_use_for_request = MagicMock(return_value=mock_tracer) + otel.set_raw_request_attributes = MagicMock() + otel._to_ns = MagicMock(return_value=1234567890) + + kwargs = { + "litellm_params": {"metadata": {}}, + "standard_logging_object": {"metadata": self._slo_metadata()}, + } + otel._maybe_log_raw_request( + kwargs, {}, datetime.now(), datetime.now(), MagicMock() + ) + + mock_span.set_attribute.assert_any_call( + "metadata.user_api_key_team_id", "team-123" + ) + mock_span.set_attribute.assert_any_call( + "metadata.user_api_key_team_alias", "my-team" + ) + + def test_helper_skips_when_team_values_missing(self): + otel = OpenTelemetry() + mock_span = MagicMock() + + otel._set_team_attributes_on_span(span=mock_span, team_id=None, team_alias=None) + + mock_span.set_attribute.assert_not_called() + + def test_helper_skips_when_team_values_are_empty_strings(self): + """A master-key / team-less request carries user_api_key_team_id='' + in metadata. Propagating '' to every span is noise that makes + traces look mis-instrumented; treat empty as absent.""" + otel = OpenTelemetry() + mock_span = MagicMock() + + otel._set_team_attributes_on_span(span=mock_span, team_id="", team_alias="") + + mock_span.set_attribute.assert_not_called() + + def test_helper_reads_metadata_from_kwargs(self): + otel = OpenTelemetry() + mock_span = MagicMock() + + otel._set_team_attributes_from_kwargs( + mock_span, + {"standard_logging_object": {"metadata": self._slo_metadata()}}, + ) + + mock_span.set_attribute.assert_any_call( + "metadata.user_api_key_team_id", "team-123" + ) + mock_span.set_attribute.assert_any_call( + "metadata.user_api_key_team_alias", "my-team" + ) + + def test_helper_handles_missing_standard_logging_object(self): + otel = OpenTelemetry() + mock_span = MagicMock() + + otel._set_team_attributes_from_kwargs(mock_span, {}) + + mock_span.set_attribute.assert_not_called() + + def test_failure_hook_exception_span_has_team_attributes(self): + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + tracer = provider.get_tracer(__name__) + + otel = OpenTelemetry() + otel.tracer = tracer + server_span = tracer.start_span("Received Proxy Server Request") + + user_api_key_dict = MagicMock() + user_api_key_dict.parent_otel_span = server_span + user_api_key_dict.team_id = "team-123" + user_api_key_dict.team_alias = "my-team" + + asyncio.run( + otel.async_post_call_failure_hook( + request_data={}, + original_exception=ValueError("boom"), + user_api_key_dict=user_api_key_dict, + traceback_str="trace", + ) + ) + + finished = {s.name: s for s in exporter.get_finished_spans()} + exception_span = finished["Failed Proxy Server Request"] + assert exception_span.attributes["metadata.user_api_key_team_id"] == "team-123" + assert ( + exception_span.attributes["metadata.user_api_key_team_alias"] == "my-team" + ) + + class TestOpenTelemetryCostBreakdown(unittest.TestCase): def test_cost_breakdown_emitted_to_otel_span(self): """ diff --git a/tests/test_litellm/integrations/test_otel_team_attributes_matrix.py b/tests/test_litellm/integrations/test_otel_team_attributes_matrix.py new file mode 100644 index 00000000000..1ce55fa7a58 --- /dev/null +++ b/tests/test_litellm/integrations/test_otel_team_attributes_matrix.py @@ -0,0 +1,285 @@ +""" +Matrix test: team_id / team_alias must land on EVERY span of a proxy +request trace, for a representative set of endpoints x HTTP outcomes. + +Endpoints + - /v1/chat/completions (OpenAI-format LLM path) + - /v1/messages (Anthropic-format LLM path) + - /team/info (management/admin path) + +Outcomes + - 2xx success + - 3xx redirect (LLM endpoints never 3xx -> N/A; admin too) + - 4xx client error (auth / validation failure) + - 5xx server error (upstream / DB failure) + +Strategy + These assertions exercise the real OpenTelemetry callback the proxy + invokes for each path, with a SERVER parent span (as + ``user_api_key_auth`` creates) and an in-memory exporter. Each cell + drives the path, then asserts team attributes on every span that path + actually emits. + + - success path -> ``log_success_event`` -> litellm_request + + raw_gen_ai_request + guardrail child spans. + - failure path -> ``async_post_call_failure_hook`` -> Failed Proxy + Server Request exception child span. + + Admin endpoints do not run the LLM success callback, so their only + trace surface is the SERVER span (success) or the exception child span + (failure) -- the cells below assert exactly that. +""" + +import asyncio +import os +import sys +import unittest +from datetime import datetime +from unittest.mock import MagicMock + +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) + +sys.path.insert(0, os.path.abspath("../..")) + +from litellm.integrations.opentelemetry import ( + LITELLM_PROXY_REQUEST_SPAN_NAME, + OpenTelemetry, +) + +TEAM_ID = "team-123" +TEAM_ALIAS = "my-team" +TEAM_ID_ATTR = "metadata.user_api_key_team_id" +TEAM_ALIAS_ATTR = "metadata.user_api_key_team_alias" + + +def _make_otel(): + """OTel callback whose every span lands in an in-memory exporter.""" + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + + otel = OpenTelemetry() + otel.tracer = provider.get_tracer(__name__) + # raw_gen_ai_request sub-span is gated on message logging. + otel.message_logging = True + return otel, exporter + + +def _server_span(otel): + """Mirror the SERVER span user_api_key_auth opens per request.""" + return otel.create_litellm_proxy_request_started_span( + start_time=datetime.now(), headers={} + ) + + +def _slo(call_type, with_guardrail=False): + """standard_logging_object the proxy attaches, carrying team metadata.""" + md = { + "user_api_key_team_id": TEAM_ID, + "user_api_key_team_alias": TEAM_ALIAS, + } + slo = {"metadata": md, "call_type": call_type} + if with_guardrail: + slo["guardrail_information"] = [ + { + "guardrail_name": "test_guardrail", + "guardrail_mode": "input", + "guardrail_response": "ok", + "start_time": 1609459200.0, + "end_time": 1609459201.0, + } + ] + return slo + + +def _success_kwargs(call_type, server_span, with_guardrail=True): + """kwargs the success callback receives for an LLM proxy request.""" + return { + "model": "gpt-4.1-mini", + "litellm_call_id": "call-abc", + "call_type": call_type, + "litellm_params": { + "metadata": { + "litellm_parent_otel_span": server_span, + "user_api_key_team_id": TEAM_ID, + "user_api_key_team_alias": TEAM_ALIAS, + } + }, + "standard_logging_object": _slo(call_type, with_guardrail=with_guardrail), + "messages": [{"role": "user", "content": "hi"}], + } + + +def _team_user_api_key_dict(server_span): + d = MagicMock() + d.parent_otel_span = server_span + d.team_id = TEAM_ID + d.team_alias = TEAM_ALIAS + return d + + +def _spans_by_name(exporter): + return {s.name: s for s in exporter.get_finished_spans()} + + +def _assert_team_attrs(span, where): + assert span.attributes.get(TEAM_ID_ATTR) == TEAM_ID, ( + f"{where}: missing/blank {TEAM_ID_ATTR} " + f"(got {span.attributes.get(TEAM_ID_ATTR)!r})" + ) + assert span.attributes.get(TEAM_ALIAS_ATTR) == TEAM_ALIAS, ( + f"{where}: missing/blank {TEAM_ALIAS_ATTR} " + f"(got {span.attributes.get(TEAM_ALIAS_ATTR)!r})" + ) + + +class _Boom(Exception): + """Upstream/DB style 5xx.""" + + status_code = 500 + + +class _ClientErr(Exception): + """Auth/validation style 4xx.""" + + status_code = 401 + + +# --------------------------------------------------------------------------- +# LLM success cells: litellm_request + raw_gen_ai_request + guardrail spans +# --------------------------------------------------------------------------- +class TestLLMSuccessCells(unittest.TestCase): + def _run_success(self, call_type): + otel, exporter = _make_otel() + server_span = _server_span(otel) + kwargs = _success_kwargs(call_type, server_span) + now = datetime.now() + otel.log_success_event(kwargs, {"id": "resp-1"}, now, now) + return _spans_by_name(exporter) + + def test_chat_completions_2xx(self): + spans = self._run_success("completion") + for name in ( + LITELLM_PROXY_REQUEST_SPAN_NAME, + "litellm_request", + "raw_gen_ai_request", + "guardrail", + ): + assert name in spans, f"chat/completions 2xx: missing span {name}" + _assert_team_attrs(spans[name], f"chat/completions 2xx [{name}]") + + def test_v1_messages_2xx(self): + spans = self._run_success("anthropic_messages") + for name in ( + LITELLM_PROXY_REQUEST_SPAN_NAME, + "litellm_request", + "raw_gen_ai_request", + "guardrail", + ): + assert name in spans, f"v1/messages 2xx: missing span {name}" + _assert_team_attrs(spans[name], f"v1/messages 2xx [{name}]") + + +# --------------------------------------------------------------------------- +# LLM failure cells: Failed Proxy Server Request exception child span +# --------------------------------------------------------------------------- +class TestLLMFailureCells(unittest.TestCase): + def _run_failure(self, exc): + """Drive the failure hook, then close the SERVER span (the proxy + closes it after the hook in real flow) so both the exception child + span and the SERVER root span are asserted.""" + otel, exporter = _make_otel() + server_span = _server_span(otel) + asyncio.run( + otel.async_post_call_failure_hook( + request_data={}, + original_exception=exc, + user_api_key_dict=_team_user_api_key_dict(server_span), + traceback_str="tb", + ) + ) + server_span.end() + return _spans_by_name(exporter) + + def _assert_all(self, spans, where): + for name in ("Failed Proxy Server Request", LITELLM_PROXY_REQUEST_SPAN_NAME): + assert name in spans, f"{where}: missing span {name}" + _assert_team_attrs(spans[name], f"{where} [{name}]") + + def test_chat_completions_4xx(self): + self._assert_all( + self._run_failure(_ClientErr("bad key")), "chat/completions 4xx" + ) + + def test_chat_completions_5xx(self): + self._assert_all( + self._run_failure(_Boom("upstream blew up")), "chat/completions 5xx" + ) + + def test_v1_messages_4xx(self): + self._assert_all( + self._run_failure(_ClientErr("bad anthropic key")), "v1/messages 4xx" + ) + + def test_v1_messages_5xx(self): + self._assert_all( + self._run_failure(_Boom("anthropic upstream timeout")), "v1/messages 5xx" + ) + + +# --------------------------------------------------------------------------- +# Admin /team/info cells. +# 2xx: admin path never runs the LLM success callback -> its only trace +# surface is the SERVER span; no child spans are emitted. +# 3xx: management endpoints do not redirect -> N/A (documented, no run). +# 4xx/5xx: proxy_logging post_call_failure_hook -> exception child span. +# --------------------------------------------------------------------------- +class TestAdminTeamInfoCells(unittest.TestCase): + def _run_admin_failure(self, exc): + otel, exporter = _make_otel() + server_span = _server_span(otel) + asyncio.run( + otel.async_post_call_failure_hook( + request_data={}, + original_exception=exc, + user_api_key_dict=_team_user_api_key_dict(server_span), + traceback_str="tb", + ) + ) + server_span.end() + return _spans_by_name(exporter) + + def test_team_info_4xx(self): + spans = self._run_admin_failure(_ClientErr("team not found")) + for name in ("Failed Proxy Server Request", LITELLM_PROXY_REQUEST_SPAN_NAME): + _assert_team_attrs(spans[name], f"/team/info 4xx [{name}]") + + def test_team_info_5xx(self): + spans = self._run_admin_failure(_Boom("db connection lost")) + for name in ("Failed Proxy Server Request", LITELLM_PROXY_REQUEST_SPAN_NAME): + _assert_team_attrs(spans[name], f"/team/info 5xx [{name}]") + + def test_team_info_2xx_only_server_span_no_orphan_children(self): + """Admin success path emits no LLM child spans; nothing to stamp + beyond the SERVER span. This pins that contract so a future + regression that starts emitting child spans here without team + attrs is caught.""" + otel, exporter = _make_otel() + server_span = _server_span(otel) + server_span.end() + spans = _spans_by_name(exporter) + assert set(spans) == { + LITELLM_PROXY_REQUEST_SPAN_NAME + }, f"/team/info 2xx: unexpected child spans {set(spans)}" + + def test_team_info_3xx_not_applicable(self): + """Management endpoints return JSON, never a 3xx redirect.""" + self.skipTest("/team/info has no 3xx redirect path (N/A)") + + +if __name__ == "__main__": + unittest.main()