scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-03 02:22:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

feat: propagate team_id and team_alias to all child OTEL spans (#28273)

Browse source 
- 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 <yassinkortam@g.ucla.edu>
This commit is contained in:
Yassin Kortam 2026-05-19 15:31:25 -07:00 • committed by lorenzbaraldi
parent da835b78ee
commit 86f0ffd429
3 changed files with 510 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

82
litellm/integrations/opentelemetry.py
View file

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

143
tests/test_litellm/integrations/test_opentelemetry.py
View file

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

285
tests/test_litellm/integrations/test_otel_team_attributes_matrix.py Normal file
View file

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