mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-05 18:24:49 +00:00
Fix 3 OpenTelemetry tracing bugs in proxy integration (#27757)
1. Missing litellm_request child span when proxy parent in metadata: _get_span_context now returns (ctx, None) for the metadata-injected proxy parent so the primary span is always emitted as a child of ctx. Proxy span lifecycle managed by new _end_proxy_span_from_kwargs. 2. open_telemetry_logger overwrite by later handlers: _init_otel_logger_on_litellm_proxy now uses first-registered-wins — only assigns proxy_server.open_telemetry_logger when currently None. 3. Duplicate litellm_request success spans in streaming paths: Added _mark_success_span_once with per-handler dedupe key stored in kwargs metadata, suppressing the second span when both sync and async success callbacks fire for the same request. Co-authored-by: Yassin Kortam <yassinkortam@g.ucla.edu> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
co-authored by
Yassin Kortam
Claude Opus 4.7
parent
8f25942ecf
commit
1c4e4d4a60
@@ -888,6 +888,15 @@ def log_guardrail_information(func):
|
||||
- pre_call
|
||||
- during_call
|
||||
- post_call
|
||||
|
||||
Some guardrails (e.g. ``block_code_execution``) call
|
||||
``add_standard_logging_guardrail_information_to_request_data`` directly
|
||||
from inside the wrapped function so they can record a richer payload
|
||||
(structured detections, tracing detail) than this decorator's
|
||||
"allow"/"mask"/raw-response default. To avoid double-recording in that
|
||||
case (which would emit two spans, two Datadog records, two spend-log
|
||||
entries, etc.), snapshot the entry count before invocation: if the
|
||||
wrapped function already appended its own entry, skip the auto-record.
|
||||
"""
|
||||
import functools
|
||||
import inspect
|
||||
@@ -907,6 +916,16 @@ def log_guardrail_information(func):
|
||||
return GuardrailEventHooks.post_call
|
||||
return None
|
||||
|
||||
def _count_recorded_guardrail_entries(request_data: dict) -> int:
|
||||
total = 0
|
||||
for container_key in ("metadata", "litellm_metadata"):
|
||||
container = request_data.get(container_key)
|
||||
if isinstance(container, dict):
|
||||
entries = container.get("standard_logging_guardrail_information")
|
||||
if isinstance(entries, list):
|
||||
total += len(entries)
|
||||
return total
|
||||
|
||||
@functools.wraps(func)
|
||||
async def async_wrapper(*args, **kwargs):
|
||||
start_time = datetime.now() # Move start_time inside the wrapper
|
||||
@@ -919,8 +938,11 @@ def log_guardrail_information(func):
|
||||
if func.__name__ == "apply_guardrail" and "inputs" in kwargs:
|
||||
original_inputs = kwargs.get("inputs")
|
||||
|
||||
entries_before = _count_recorded_guardrail_entries(request_data)
|
||||
try:
|
||||
response = await func(*args, **kwargs)
|
||||
if _count_recorded_guardrail_entries(request_data) > entries_before:
|
||||
return response
|
||||
return self._process_response(
|
||||
response=response,
|
||||
request_data=request_data,
|
||||
@@ -931,6 +953,8 @@ def log_guardrail_information(func):
|
||||
original_inputs=original_inputs,
|
||||
)
|
||||
except Exception as e:
|
||||
if _count_recorded_guardrail_entries(request_data) > entries_before:
|
||||
raise
|
||||
return self._process_error(
|
||||
e=e,
|
||||
request_data=request_data,
|
||||
@@ -952,8 +976,11 @@ def log_guardrail_information(func):
|
||||
if func.__name__ == "apply_guardrail" and "inputs" in kwargs:
|
||||
original_inputs = kwargs.get("inputs")
|
||||
|
||||
entries_before = _count_recorded_guardrail_entries(request_data)
|
||||
try:
|
||||
response = func(*args, **kwargs)
|
||||
if _count_recorded_guardrail_entries(request_data) > entries_before:
|
||||
return response
|
||||
return self._process_response(
|
||||
response=response,
|
||||
request_data=request_data,
|
||||
@@ -962,6 +989,8 @@ def log_guardrail_information(func):
|
||||
original_inputs=original_inputs,
|
||||
)
|
||||
except Exception as e:
|
||||
if _count_recorded_guardrail_entries(request_data) > entries_before:
|
||||
raise
|
||||
return self._process_error(
|
||||
e=e,
|
||||
request_data=request_data,
|
||||
|
||||
@@ -237,7 +237,14 @@ class OpenTelemetry(CustomLogger):
|
||||
not isinstance(cb, OpenTelemetry) for cb in litellm.service_callback
|
||||
):
|
||||
litellm.service_callback.append(self)
|
||||
setattr(proxy_server, "open_telemetry_logger", self)
|
||||
# avoid proxy logger ownership being overwritten by later
|
||||
# handlers. Multiple integrations (default OTEL, Langfuse OTEL,
|
||||
# Arize OTEL, etc.) may initialize in sequence; without this guard,
|
||||
# the last one silently replaces the first and breaks expected
|
||||
# routing for proxy_server.open_telemetry_logger consumers.
|
||||
# Behavior: first-registered wins.
|
||||
if getattr(proxy_server, "open_telemetry_logger", None) is None:
|
||||
setattr(proxy_server, "open_telemetry_logger", self)
|
||||
|
||||
def _get_or_create_provider(
|
||||
self,
|
||||
@@ -794,12 +801,100 @@ class OpenTelemetry(CustomLogger):
|
||||
# End of Team/Key Based Logging Control Flow
|
||||
#########################################################
|
||||
|
||||
def _emit_once(self, kwargs: dict, *scope: object) -> bool:
|
||||
"""Return True the first time this handler is asked to emit a span
|
||||
for the given (handler, scope) on this kwargs; False on repeats.
|
||||
|
||||
Used to suppress duplicate span emission for two distinct patterns:
|
||||
|
||||
1. **Handler-level dual-fire**: streaming code paths trigger both
|
||||
the sync and async callback for one request, so ``_handle_success``
|
||||
/ ``_handle_failure`` would otherwise produce two
|
||||
``litellm_request`` spans. Scope: ``("success",)`` / ``("failure",)``.
|
||||
2. **Payload-driven multi-entrypoint emission**: a span loop that
|
||||
reads entries from ``standard_logging_payload`` (currently only
|
||||
guardrails) is invoked from multiple lifecycle points
|
||||
(post-call hooks, success callback, failure callback). The list
|
||||
can be re-read with mutated entries between calls, so dedupe
|
||||
must be at entry granularity. Scope: the entry's stable identity.
|
||||
|
||||
``scope`` parts can be any hashable identity. The marker is stored
|
||||
in ``kwargs["litellm_params"]["metadata"]["_otel_internal"]`` so it
|
||||
is request-local (kwargs is shared across the sync/async callbacks
|
||||
and lifecycle hooks for one request).
|
||||
"""
|
||||
litellm_params = kwargs.get("litellm_params")
|
||||
if not isinstance(litellm_params, dict):
|
||||
litellm_params = {}
|
||||
kwargs["litellm_params"] = litellm_params
|
||||
|
||||
_metadata = litellm_params.get("metadata")
|
||||
if not isinstance(_metadata, dict):
|
||||
_metadata = {}
|
||||
litellm_params["metadata"] = _metadata
|
||||
|
||||
_otel_internal = _metadata.get("_otel_internal")
|
||||
if not isinstance(_otel_internal, dict):
|
||||
_otel_internal = {}
|
||||
_metadata["_otel_internal"] = _otel_internal
|
||||
|
||||
spans_logged = _otel_internal.get("spans_logged")
|
||||
if not isinstance(spans_logged, dict):
|
||||
spans_logged = {}
|
||||
_otel_internal["spans_logged"] = spans_logged
|
||||
|
||||
dedupe_key = (self.__class__.__name__, id(self), *scope)
|
||||
if spans_logged.get(dedupe_key) is True:
|
||||
return False
|
||||
|
||||
spans_logged[dedupe_key] = True
|
||||
return True
|
||||
|
||||
def _end_proxy_span_from_kwargs(self, kwargs: dict, end_time) -> None:
|
||||
"""Close the proxy-level parent span if it is still recording.
|
||||
|
||||
This helper retrieves the proxy span directly from kwargs metadata
|
||||
and closes it after all child spans have been recorded.
|
||||
|
||||
Only called from the success path. The failure path deliberately
|
||||
leaves the proxy span open so ``async_post_call_failure_hook`` can
|
||||
append the ``"Failed Proxy Server Request"`` child span before
|
||||
closing it.
|
||||
|
||||
Only spans named ``LITELLM_PROXY_REQUEST_SPAN_NAME`` are closed —
|
||||
externally provided spans must not be closed by LiteLLM.
|
||||
"""
|
||||
litellm_params = kwargs.get("litellm_params", {}) or {}
|
||||
_metadata = litellm_params.get("metadata", {}) or {}
|
||||
proxy_span = _metadata.get("litellm_parent_otel_span", None)
|
||||
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()
|
||||
):
|
||||
proxy_span.end(end_time=self._to_ns(end_time))
|
||||
|
||||
def _handle_success(self, kwargs, response_obj, start_time, end_time):
|
||||
"""Create the litellm_request span then close the proxy span."""
|
||||
verbose_logger.debug(
|
||||
"OpenTelemetry Logger: Logging kwargs: %s, OTEL config settings=%s",
|
||||
kwargs,
|
||||
self.config,
|
||||
)
|
||||
|
||||
# sync + async success handlers can both fire for one
|
||||
# request (notably in streaming code paths). Guard against duplicate
|
||||
# span writes — but still close the proxy span on the skip path so
|
||||
# the trace doesn't leak an open root span.
|
||||
if not self._emit_once(kwargs, "success"):
|
||||
verbose_logger.debug(
|
||||
"OpenTelemetry: skipping duplicate success span for handler=%s",
|
||||
self.__class__.__name__,
|
||||
)
|
||||
self._end_proxy_span_from_kwargs(kwargs, end_time)
|
||||
return
|
||||
|
||||
ctx, parent_span = self._get_span_context(kwargs)
|
||||
|
||||
if self.config.ignore_context_propagation:
|
||||
@@ -859,7 +954,7 @@ class OpenTelemetry(CustomLogger):
|
||||
|
||||
# 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
|
||||
# However, proxy-created spans should be closed here.
|
||||
if (
|
||||
parent_span is not None
|
||||
and hasattr(parent_span, "name")
|
||||
@@ -867,6 +962,11 @@ class OpenTelemetry(CustomLogger):
|
||||
):
|
||||
parent_span.end(end_time=self._to_ns(end_time))
|
||||
|
||||
# close the proxy span explicitly from kwargs metadata
|
||||
# after all child spans (litellm_request, guardrail, raw_request)
|
||||
# have been fully recorded and exported.
|
||||
self._end_proxy_span_from_kwargs(kwargs, end_time)
|
||||
|
||||
def _start_primary_span(
|
||||
self,
|
||||
kwargs,
|
||||
@@ -1296,6 +1396,21 @@ class OpenTelemetry(CustomLogger):
|
||||
for guardrail_information in guardrail_information_list:
|
||||
start_time_float = guardrail_information.get("start_time")
|
||||
end_time_float = guardrail_information.get("end_time")
|
||||
|
||||
# ``_create_guardrail_span`` is called from three lifecycle
|
||||
# points (``async_post_call_success_hook``, ``_handle_success``,
|
||||
# ``_handle_failure``) and re-reads the (mutating) entry list
|
||||
# each time. Dedupe at entry granularity so a single real
|
||||
# guardrail invocation produces exactly one span per handler.
|
||||
if not self._emit_once(
|
||||
kwargs,
|
||||
"guardrail",
|
||||
guardrail_information.get("guardrail_name"),
|
||||
start_time_float,
|
||||
guardrail_information.get("guardrail_mode"),
|
||||
):
|
||||
continue
|
||||
|
||||
start_time_datetime = datetime.now()
|
||||
if start_time_float is not None:
|
||||
start_time_datetime = datetime.fromtimestamp(start_time_float)
|
||||
@@ -1349,6 +1464,21 @@ class OpenTelemetry(CustomLogger):
|
||||
kwargs,
|
||||
self.config,
|
||||
)
|
||||
|
||||
# sync + async failure handlers can both fire for one
|
||||
# request (notably in streaming code paths), producing two
|
||||
# semantically identical ERROR spans. Unlike the success path, the
|
||||
# proxy span is intentionally left open here so that
|
||||
# ``async_post_call_failure_hook`` can append the
|
||||
# "Failed Proxy Server Request" child span before closing it —
|
||||
# there is no proxy-span side-effect to preserve on the skip path.
|
||||
if not self._emit_once(kwargs, "failure"):
|
||||
verbose_logger.debug(
|
||||
"OpenTelemetry: skipping duplicate failure span for handler=%s",
|
||||
self.__class__.__name__,
|
||||
)
|
||||
return
|
||||
|
||||
_parent_context, parent_otel_span = self._get_span_context(kwargs)
|
||||
|
||||
if self.config.ignore_context_propagation:
|
||||
@@ -2188,7 +2318,7 @@ class OpenTelemetry(CustomLogger):
|
||||
verbose_logger.debug(
|
||||
"OpenTelemetry: Using explicit parent span from metadata"
|
||||
)
|
||||
return trace.set_span_in_context(parent_otel_span), parent_otel_span
|
||||
return trace.set_span_in_context(parent_otel_span), None
|
||||
|
||||
# Priority 2: HTTP traceparent header
|
||||
if traceparent is not None:
|
||||
|
||||
@@ -929,6 +929,91 @@ class TestEventTypeLogging:
|
||||
assert len(logged_info) == 1
|
||||
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.post_call
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_log_guardrail_information_skips_auto_record_if_function_already_recorded(
|
||||
self,
|
||||
):
|
||||
"""When a wrapped guardrail function records its own entry directly
|
||||
(e.g. block_code_execution.apply_guardrail records a rich
|
||||
``[detections...]`` payload), the decorator must NOT also append its
|
||||
own ``"allow"``/raw-response entry — otherwise every backend
|
||||
(OTEL spans, Datadog, Langfuse, spend logs) double-records one
|
||||
logical guardrail invocation."""
|
||||
from litellm.integrations.custom_guardrail import log_guardrail_information
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
class TestGuardrail(CustomGuardrail):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
guardrail_name="block-code",
|
||||
event_hook=GuardrailEventHooks.pre_call,
|
||||
)
|
||||
|
||||
@log_guardrail_information
|
||||
async def apply_guardrail(self, inputs, request_data, **kwargs):
|
||||
self.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_json_response=[{"action_taken": "block"}],
|
||||
request_data=request_data,
|
||||
guardrail_status="success",
|
||||
event_type=GuardrailEventHooks.pre_call,
|
||||
)
|
||||
return inputs
|
||||
|
||||
guardrail = TestGuardrail()
|
||||
request_data = {"metadata": {}}
|
||||
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["x"]}, request_data=request_data
|
||||
)
|
||||
|
||||
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
|
||||
assert len(logged_info) == 1, (
|
||||
f"Decorator must not double-record when the wrapped function "
|
||||
f"already appended its own entry; got {len(logged_info)} entries"
|
||||
)
|
||||
assert logged_info[0]["guardrail_response"] == [{"action_taken": "block"}]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_log_guardrail_information_skips_auto_record_on_exception_if_function_already_recorded(
|
||||
self,
|
||||
):
|
||||
"""Same as above on the failure path: if the wrapped function
|
||||
appended an entry in its ``finally`` block before re-raising, the
|
||||
decorator must just re-raise without auto-recording on top."""
|
||||
from litellm.integrations.custom_guardrail import log_guardrail_information
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
class TestGuardrail(CustomGuardrail):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
guardrail_name="block-code",
|
||||
event_hook=GuardrailEventHooks.pre_call,
|
||||
)
|
||||
|
||||
@log_guardrail_information
|
||||
async def apply_guardrail(self, inputs, request_data, **kwargs):
|
||||
try:
|
||||
raise ValueError("blocked")
|
||||
finally:
|
||||
self.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_json_response=[{"action_taken": "block"}],
|
||||
request_data=request_data,
|
||||
guardrail_status="guardrail_intervened",
|
||||
event_type=GuardrailEventHooks.pre_call,
|
||||
)
|
||||
|
||||
guardrail = TestGuardrail()
|
||||
request_data = {"metadata": {}}
|
||||
|
||||
with pytest.raises(ValueError, match="blocked"):
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["x"]}, request_data=request_data
|
||||
)
|
||||
|
||||
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
|
||||
assert len(logged_info) == 1
|
||||
assert logged_info[0]["guardrail_status"] == "guardrail_intervened"
|
||||
|
||||
def test_add_standard_logging_falls_back_to_event_hook_when_event_type_is_none(
|
||||
self,
|
||||
):
|
||||
@@ -1086,9 +1171,12 @@ class TestCustomGuardrailSpendLogMatchRedaction:
|
||||
][0]["match"]
|
||||
== "[REDACTED]"
|
||||
)
|
||||
assert raw["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][
|
||||
"match"
|
||||
] == "GG"
|
||||
assert (
|
||||
raw["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][
|
||||
"match"
|
||||
]
|
||||
== "GG"
|
||||
)
|
||||
|
||||
def test_add_standard_logging_redacts_regex_field(self):
|
||||
cg = CustomGuardrail(guardrail_name="test-rail")
|
||||
|
||||
@@ -1982,11 +1982,13 @@ class TestOpenTelemetryExternalSpan(unittest.TestCase):
|
||||
- raw_gen_ai_request spans are children of litellm_request spans
|
||||
- Correct hierarchy: external_parent → litellm_request → raw_gen_ai_request
|
||||
"""
|
||||
import copy
|
||||
|
||||
# Initialize OpenTelemetry
|
||||
otel = OpenTelemetry(tracer_provider=self.tracer_provider)
|
||||
|
||||
# Load test data
|
||||
kwargs, response_obj = self._create_test_kwargs_and_response()
|
||||
kwargs1, response_obj = self._create_test_kwargs_and_response()
|
||||
kwargs2 = copy.deepcopy(kwargs1)
|
||||
|
||||
# Create external parent span using our test TracerProvider
|
||||
tracer = self.tracer_provider.get_tracer(__name__)
|
||||
@@ -1999,7 +2001,7 @@ class TestOpenTelemetryExternalSpan(unittest.TestCase):
|
||||
# First completion call
|
||||
start_time = datetime.utcnow()
|
||||
end_time = start_time + timedelta(seconds=1)
|
||||
otel._handle_success(kwargs, response_obj, start_time, end_time)
|
||||
otel._handle_success(kwargs1, response_obj, start_time, end_time)
|
||||
|
||||
# Verify parent span is still recording
|
||||
self.assertTrue(
|
||||
@@ -2010,7 +2012,7 @@ class TestOpenTelemetryExternalSpan(unittest.TestCase):
|
||||
# Second completion call
|
||||
start_time2 = end_time
|
||||
end_time2 = start_time2 + timedelta(seconds=1)
|
||||
otel._handle_success(kwargs, response_obj, start_time2, end_time2)
|
||||
otel._handle_success(kwargs2, response_obj, start_time2, end_time2)
|
||||
|
||||
# Verify parent span is still recording
|
||||
self.assertTrue(
|
||||
@@ -3161,7 +3163,6 @@ class TestResponseIdFallback(unittest.TestCase):
|
||||
mock_span.set_attribute.assert_any_call("litellm.call_id", call_id)
|
||||
|
||||
|
||||
|
||||
class TestOpenTelemetryResponsesAPI(unittest.TestCase):
|
||||
"""
|
||||
Tests for Responses API (/v1/responses) OTel span attributes.
|
||||
@@ -3374,7 +3375,9 @@ class TestOpenTelemetryResponsesAPI(unittest.TestCase):
|
||||
|
||||
# No output messages should be set since the text is empty
|
||||
raw = self._get_attr(mock_span, "gen_ai.output.messages")
|
||||
self.assertIsNone(raw, "Empty output text should not produce gen_ai.output.messages")
|
||||
self.assertIsNone(
|
||||
raw, "Empty output text should not produce gen_ai.output.messages"
|
||||
)
|
||||
|
||||
def test_choices_still_work(self):
|
||||
"""Existing choices-based responses must still work (no regression)."""
|
||||
@@ -3616,7 +3619,6 @@ class TestTransformResponsesAPIOutput(unittest.TestCase):
|
||||
result = otel._transform_responses_api_output_to_otel(output)
|
||||
self.assertEqual(result[0]["role"], "assistant")
|
||||
|
||||
|
||||
def test_pydantic_like_objects_accepted(self):
|
||||
"""Items with .get() but not isinstance(dict) should be accepted."""
|
||||
|
||||
@@ -3738,12 +3740,16 @@ class TestResponsesAPIToolCallSpanAttributes(unittest.TestCase):
|
||||
],
|
||||
}
|
||||
|
||||
otel.set_attributes(span=mock_span, kwargs=self._base_kwargs(), response_obj=response_obj)
|
||||
otel.set_attributes(
|
||||
span=mock_span, kwargs=self._base_kwargs(), response_obj=response_obj
|
||||
)
|
||||
|
||||
# Verify per-tool-call attributes were set (same format as choices branch)
|
||||
attr_names = [call[0][0] for call in mock_span.set_attribute.call_args_list]
|
||||
tool_call_attrs = [a for a in attr_names if "function_call" in a]
|
||||
self.assertTrue(len(tool_call_attrs) > 0, "Per-tool-call span attributes should be emitted")
|
||||
self.assertTrue(
|
||||
len(tool_call_attrs) > 0, "Per-tool-call span attributes should be emitted"
|
||||
)
|
||||
|
||||
# Verify the name attribute specifically
|
||||
mock_span.set_attribute.assert_any_call(
|
||||
@@ -3778,7 +3784,9 @@ class TestResponsesAPIToolCallSpanAttributes(unittest.TestCase):
|
||||
],
|
||||
}
|
||||
|
||||
otel.set_attributes(span=mock_span, kwargs=self._base_kwargs(), response_obj=response_obj)
|
||||
otel.set_attributes(
|
||||
span=mock_span, kwargs=self._base_kwargs(), response_obj=response_obj
|
||||
)
|
||||
|
||||
mock_span.set_attribute.assert_any_call(
|
||||
"gen_ai.completion.0.function_call.name", "get_weather"
|
||||
@@ -3786,3 +3794,488 @@ class TestResponsesAPIToolCallSpanAttributes(unittest.TestCase):
|
||||
mock_span.set_attribute.assert_any_call(
|
||||
"gen_ai.completion.1.function_call.name", "get_time"
|
||||
)
|
||||
|
||||
|
||||
class TestOpenTelemetryProxyParentSpanChildEmission(unittest.TestCase):
|
||||
"""When metadata includes litellm_parent_otel_span (the proxy
|
||||
span), the primary litellm_request span must still be created as a child
|
||||
so the trace hierarchy is complete."""
|
||||
|
||||
def _build_kwargs(self, parent_span):
|
||||
return {
|
||||
"model": "gpt-4",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"optional_params": {},
|
||||
"litellm_params": {
|
||||
"custom_llm_provider": "openai",
|
||||
"metadata": {"litellm_parent_otel_span": parent_span},
|
||||
},
|
||||
"standard_logging_object": {
|
||||
"id": "test-id",
|
||||
"call_type": "completion",
|
||||
"metadata": {},
|
||||
"hidden_params": {},
|
||||
},
|
||||
}
|
||||
|
||||
def test_get_span_context_returns_none_parent_for_metadata_span(self):
|
||||
"""_get_span_context Priority 1 must return (ctx, None) — never the
|
||||
parent span object — so callers always create litellm_request as a
|
||||
child of ctx."""
|
||||
tracer_provider = TracerProvider()
|
||||
otel = OpenTelemetry(tracer_provider=tracer_provider)
|
||||
otel.tracer = tracer_provider.get_tracer(__name__)
|
||||
|
||||
parent_span = otel.tracer.start_span("some_external_parent")
|
||||
kwargs = self._build_kwargs(parent_span)
|
||||
|
||||
ctx, returned_parent = otel._get_span_context(kwargs)
|
||||
|
||||
self.assertIsNotNone(ctx, "ctx should carry the parent for child spans")
|
||||
self.assertIsNone(
|
||||
returned_parent,
|
||||
"parent_span return slot must be None so callers create litellm_request",
|
||||
)
|
||||
parent_span.end()
|
||||
|
||||
def test_litellm_request_emitted_as_child_of_proxy_parent_span(self):
|
||||
"""End-to-end: proxy span in metadata should yield exactly one
|
||||
litellm_request span parented to it, with no extra root span."""
|
||||
from litellm.integrations.opentelemetry import (
|
||||
LITELLM_PROXY_REQUEST_SPAN_NAME,
|
||||
LITELLM_REQUEST_SPAN_NAME,
|
||||
)
|
||||
|
||||
span_exporter = InMemorySpanExporter()
|
||||
tracer_provider = TracerProvider()
|
||||
tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter))
|
||||
|
||||
otel = OpenTelemetry(tracer_provider=tracer_provider)
|
||||
otel.tracer = tracer_provider.get_tracer(__name__)
|
||||
|
||||
proxy_span = otel.tracer.start_span(LITELLM_PROXY_REQUEST_SPAN_NAME)
|
||||
kwargs = self._build_kwargs(proxy_span)
|
||||
|
||||
start = datetime.utcnow()
|
||||
end = start + timedelta(seconds=1)
|
||||
otel._handle_success(kwargs, response_obj=None, start_time=start, end_time=end)
|
||||
|
||||
spans = span_exporter.get_finished_spans()
|
||||
litellm_spans = [s for s in spans if s.name == LITELLM_REQUEST_SPAN_NAME]
|
||||
proxy_spans = [s for s in spans if s.name == LITELLM_PROXY_REQUEST_SPAN_NAME]
|
||||
|
||||
self.assertEqual(
|
||||
len(litellm_spans), 1, "Exactly one litellm_request span must be emitted"
|
||||
)
|
||||
self.assertEqual(
|
||||
len(proxy_spans), 1, "Proxy span should be closed exactly once"
|
||||
)
|
||||
|
||||
litellm_span = litellm_spans[0]
|
||||
self.assertIsNotNone(
|
||||
litellm_span.parent, "litellm_request must have a parent (not root)"
|
||||
)
|
||||
self.assertEqual(
|
||||
litellm_span.parent.span_id,
|
||||
proxy_spans[0].context.span_id,
|
||||
"litellm_request must be a child of the proxy span",
|
||||
)
|
||||
|
||||
def test_end_proxy_span_from_kwargs_closes_recording_proxy_span(self):
|
||||
from litellm.integrations.opentelemetry import LITELLM_PROXY_REQUEST_SPAN_NAME
|
||||
|
||||
span_exporter = InMemorySpanExporter()
|
||||
tracer_provider = TracerProvider()
|
||||
tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter))
|
||||
|
||||
otel = OpenTelemetry(tracer_provider=tracer_provider)
|
||||
otel.tracer = tracer_provider.get_tracer(__name__)
|
||||
|
||||
proxy_span = otel.tracer.start_span(LITELLM_PROXY_REQUEST_SPAN_NAME)
|
||||
self.assertTrue(proxy_span.is_recording())
|
||||
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {"litellm_parent_otel_span": proxy_span},
|
||||
}
|
||||
}
|
||||
otel._end_proxy_span_from_kwargs(kwargs, end_time=datetime.utcnow())
|
||||
|
||||
self.assertFalse(
|
||||
proxy_span.is_recording(), "Proxy span should be closed by helper"
|
||||
)
|
||||
|
||||
def test_end_proxy_span_from_kwargs_does_not_close_external_span(self):
|
||||
"""Spans not named LITELLM_PROXY_REQUEST_SPAN_NAME must not be closed —
|
||||
they may belong to external owners (Langfuse SDK, user code, etc.)."""
|
||||
tracer_provider = TracerProvider()
|
||||
otel = OpenTelemetry(tracer_provider=tracer_provider)
|
||||
otel.tracer = tracer_provider.get_tracer(__name__)
|
||||
|
||||
external = otel.tracer.start_span("external_caller_span")
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {"litellm_parent_otel_span": external},
|
||||
}
|
||||
}
|
||||
otel._end_proxy_span_from_kwargs(kwargs, end_time=datetime.utcnow())
|
||||
|
||||
self.assertTrue(
|
||||
external.is_recording(),
|
||||
"External (non-proxy) parent span must not be closed by LiteLLM",
|
||||
)
|
||||
external.end()
|
||||
|
||||
|
||||
class TestOpenTelemetryProxyLoggerFirstRegisteredWins(unittest.TestCase):
|
||||
"""open_telemetry_logger ownership must not be silently
|
||||
overwritten by later handlers. First-registered wins."""
|
||||
|
||||
def _install_fake_proxy_server(self):
|
||||
"""Install a stub ``litellm.proxy.proxy_server`` so the test does
|
||||
not depend on optional proxy dependencies (websockets, etc.).
|
||||
Returns (fake_module, cleanup_fn)."""
|
||||
import importlib
|
||||
import types
|
||||
|
||||
proxy_pkg_name = "litellm.proxy"
|
||||
proxy_server_name = "litellm.proxy.proxy_server"
|
||||
|
||||
previous_pkg = sys.modules.get(proxy_pkg_name)
|
||||
previous_mod = sys.modules.get(proxy_server_name)
|
||||
|
||||
# Ensure litellm.proxy package object exists
|
||||
if previous_pkg is None:
|
||||
try:
|
||||
pkg = importlib.import_module(proxy_pkg_name)
|
||||
except Exception:
|
||||
pkg = types.ModuleType(proxy_pkg_name)
|
||||
sys.modules[proxy_pkg_name] = pkg
|
||||
else:
|
||||
pkg = previous_pkg
|
||||
|
||||
fake = types.ModuleType(proxy_server_name)
|
||||
fake.open_telemetry_logger = None
|
||||
sys.modules[proxy_server_name] = fake
|
||||
setattr(pkg, "proxy_server", fake)
|
||||
|
||||
def cleanup():
|
||||
if previous_mod is not None:
|
||||
sys.modules[proxy_server_name] = previous_mod
|
||||
setattr(pkg, "proxy_server", previous_mod)
|
||||
else:
|
||||
sys.modules.pop(proxy_server_name, None)
|
||||
if hasattr(pkg, "proxy_server"):
|
||||
try:
|
||||
delattr(pkg, "proxy_server")
|
||||
except AttributeError:
|
||||
pass
|
||||
if previous_pkg is None and proxy_pkg_name in sys.modules:
|
||||
if sys.modules[proxy_pkg_name] is pkg:
|
||||
# Leave it in place — removing it would break later imports
|
||||
pass
|
||||
|
||||
return fake, cleanup
|
||||
|
||||
def test_first_registered_handler_keeps_ownership(self):
|
||||
fake_proxy_server, cleanup = self._install_fake_proxy_server()
|
||||
try:
|
||||
first = OpenTelemetry()
|
||||
self.assertIs(
|
||||
fake_proxy_server.open_telemetry_logger,
|
||||
first,
|
||||
"First registered handler must own the proxy logger slot",
|
||||
)
|
||||
|
||||
second = OpenTelemetry()
|
||||
self.assertIs(
|
||||
fake_proxy_server.open_telemetry_logger,
|
||||
first,
|
||||
"Second handler must NOT overwrite the first-registered logger",
|
||||
)
|
||||
self.assertIsNot(
|
||||
fake_proxy_server.open_telemetry_logger,
|
||||
second,
|
||||
"Proxy logger must remain pointed at the first handler",
|
||||
)
|
||||
finally:
|
||||
cleanup()
|
||||
|
||||
def test_assignment_happens_when_slot_is_unset(self):
|
||||
fake_proxy_server, cleanup = self._install_fake_proxy_server()
|
||||
try:
|
||||
handler = OpenTelemetry()
|
||||
self.assertIs(fake_proxy_server.open_telemetry_logger, handler)
|
||||
finally:
|
||||
cleanup()
|
||||
|
||||
def test_existing_non_none_logger_is_preserved(self):
|
||||
"""If ``proxy_server.open_telemetry_logger`` is already set to any
|
||||
non-None value, a new handler must not overwrite it — even if the
|
||||
existing value is not an OpenTelemetry instance."""
|
||||
fake_proxy_server, cleanup = self._install_fake_proxy_server()
|
||||
try:
|
||||
sentinel = object()
|
||||
fake_proxy_server.open_telemetry_logger = sentinel
|
||||
OpenTelemetry()
|
||||
self.assertIs(
|
||||
fake_proxy_server.open_telemetry_logger,
|
||||
sentinel,
|
||||
"Existing non-None logger must not be overwritten",
|
||||
)
|
||||
finally:
|
||||
cleanup()
|
||||
|
||||
|
||||
class TestOpenTelemetrySpanDedupe(unittest.TestCase):
|
||||
"""``_emit_once`` is a per-request, per-handler idempotency guard that
|
||||
prevents duplicate span emission across two distinct dual-fire patterns:
|
||||
|
||||
1. Handler-level: streaming triggers both sync and async success/failure
|
||||
callbacks for one request — the second call would otherwise produce a
|
||||
duplicate ``litellm_request`` span.
|
||||
2. Payload-driven entry-level: ``_create_guardrail_span`` is invoked
|
||||
from three lifecycle points (post-call hook, success, failure) and
|
||||
re-reads a mutating list — the same logical guardrail invocation
|
||||
would otherwise be emitted up to three times.
|
||||
"""
|
||||
|
||||
def _build_kwargs(self, *, exception: bool = False):
|
||||
kwargs = {
|
||||
"model": "gpt-4",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"optional_params": {},
|
||||
"litellm_params": {
|
||||
"custom_llm_provider": "openai",
|
||||
"metadata": {},
|
||||
},
|
||||
"standard_logging_object": {
|
||||
"id": "test-id",
|
||||
"call_type": "completion",
|
||||
"metadata": {},
|
||||
"hidden_params": {},
|
||||
},
|
||||
}
|
||||
if exception:
|
||||
kwargs["exception"] = Exception("test error")
|
||||
return kwargs
|
||||
|
||||
def test_emit_once_first_call_returns_true_then_false(self):
|
||||
otel = OpenTelemetry()
|
||||
kwargs = self._build_kwargs()
|
||||
self.assertTrue(otel._emit_once(kwargs, "success"))
|
||||
self.assertFalse(
|
||||
otel._emit_once(kwargs, "success"),
|
||||
"Repeat call for same handler+scope+kwargs must be deduped",
|
||||
)
|
||||
|
||||
def test_emit_once_distinct_scopes_dont_collide(self):
|
||||
"""Different scopes on the same handler+kwargs must each emit once."""
|
||||
otel = OpenTelemetry()
|
||||
kwargs = self._build_kwargs()
|
||||
self.assertTrue(otel._emit_once(kwargs, "success"))
|
||||
self.assertTrue(
|
||||
otel._emit_once(kwargs, "failure"),
|
||||
"Failure scope must be independent of success scope",
|
||||
)
|
||||
self.assertTrue(
|
||||
otel._emit_once(kwargs, "guardrail", "block-code", 1.0, "pre_call"),
|
||||
"Guardrail entry scope must be independent of success/failure scopes",
|
||||
)
|
||||
self.assertFalse(otel._emit_once(kwargs, "success"))
|
||||
self.assertFalse(otel._emit_once(kwargs, "failure"))
|
||||
self.assertFalse(
|
||||
otel._emit_once(kwargs, "guardrail", "block-code", 1.0, "pre_call")
|
||||
)
|
||||
|
||||
def test_emit_once_separate_handlers_each_emit(self):
|
||||
"""Two distinct handler instances must each emit exactly once for the
|
||||
same scope."""
|
||||
otel_a = OpenTelemetry()
|
||||
otel_b = OpenTelemetry()
|
||||
kwargs = self._build_kwargs()
|
||||
self.assertTrue(otel_a._emit_once(kwargs, "success"))
|
||||
self.assertTrue(
|
||||
otel_b._emit_once(kwargs, "success"),
|
||||
"Different handler instance must not share the first handler's marker",
|
||||
)
|
||||
self.assertFalse(otel_a._emit_once(kwargs, "success"))
|
||||
self.assertFalse(otel_b._emit_once(kwargs, "success"))
|
||||
|
||||
def test_emit_once_handles_missing_metadata(self):
|
||||
otel = OpenTelemetry()
|
||||
kwargs = {"litellm_params": {}}
|
||||
self.assertTrue(otel._emit_once(kwargs, "success"))
|
||||
self.assertFalse(otel._emit_once(kwargs, "success"))
|
||||
|
||||
def test_emit_once_handles_missing_litellm_params(self):
|
||||
otel = OpenTelemetry()
|
||||
kwargs = {}
|
||||
self.assertTrue(otel._emit_once(kwargs, "success"))
|
||||
self.assertFalse(otel._emit_once(kwargs, "success"))
|
||||
|
||||
def test_handle_success_emits_single_litellm_request_span_on_double_call(self):
|
||||
"""Sync + async callback paths firing for the same kwargs must
|
||||
result in exactly one litellm_request span."""
|
||||
from litellm.integrations.opentelemetry import LITELLM_REQUEST_SPAN_NAME
|
||||
|
||||
span_exporter = InMemorySpanExporter()
|
||||
tracer_provider = TracerProvider()
|
||||
tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter))
|
||||
|
||||
otel = OpenTelemetry(tracer_provider=tracer_provider)
|
||||
otel.tracer = tracer_provider.get_tracer(__name__)
|
||||
|
||||
kwargs = self._build_kwargs()
|
||||
start = datetime.utcnow()
|
||||
end = start + timedelta(seconds=1)
|
||||
|
||||
otel._handle_success(kwargs, response_obj=None, start_time=start, end_time=end)
|
||||
otel._handle_success(kwargs, response_obj=None, start_time=start, end_time=end)
|
||||
|
||||
spans = span_exporter.get_finished_spans()
|
||||
litellm_spans = [s for s in spans if s.name == LITELLM_REQUEST_SPAN_NAME]
|
||||
self.assertEqual(
|
||||
len(litellm_spans),
|
||||
1,
|
||||
f"Exactly one litellm_request span expected, got {len(litellm_spans)}",
|
||||
)
|
||||
|
||||
def test_handle_success_dedupe_skip_still_closes_proxy_span(self):
|
||||
"""When the success path is short-circuited as a duplicate, the
|
||||
proxy span must still be closed so traces don't leak."""
|
||||
from litellm.integrations.opentelemetry import LITELLM_PROXY_REQUEST_SPAN_NAME
|
||||
|
||||
span_exporter = InMemorySpanExporter()
|
||||
tracer_provider = TracerProvider()
|
||||
tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter))
|
||||
|
||||
otel = OpenTelemetry(tracer_provider=tracer_provider)
|
||||
otel.tracer = tracer_provider.get_tracer(__name__)
|
||||
|
||||
proxy_span = otel.tracer.start_span(LITELLM_PROXY_REQUEST_SPAN_NAME)
|
||||
kwargs = self._build_kwargs()
|
||||
kwargs["litellm_params"]["metadata"]["litellm_parent_otel_span"] = proxy_span
|
||||
|
||||
otel._emit_once(kwargs, "success") # pre-mark to force dedupe-skip branch
|
||||
self.assertTrue(proxy_span.is_recording())
|
||||
|
||||
start = datetime.utcnow()
|
||||
end = start + timedelta(seconds=1)
|
||||
otel._handle_success(kwargs, response_obj=None, start_time=start, end_time=end)
|
||||
|
||||
self.assertFalse(
|
||||
proxy_span.is_recording(),
|
||||
"Dedupe-skip path must still close the proxy span via _end_proxy_span_from_kwargs",
|
||||
)
|
||||
|
||||
def test_handle_failure_emits_single_error_span_on_double_call(self):
|
||||
"""Sync + async failure callback paths firing for the same kwargs
|
||||
must result in exactly one ERROR litellm_request span."""
|
||||
from opentelemetry.trace import StatusCode
|
||||
|
||||
from litellm.integrations.opentelemetry import LITELLM_REQUEST_SPAN_NAME
|
||||
|
||||
span_exporter = InMemorySpanExporter()
|
||||
tracer_provider = TracerProvider()
|
||||
tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter))
|
||||
|
||||
otel = OpenTelemetry(tracer_provider=tracer_provider)
|
||||
otel.tracer = tracer_provider.get_tracer(__name__)
|
||||
|
||||
kwargs = self._build_kwargs(exception=True)
|
||||
start = datetime.utcnow()
|
||||
end = start + timedelta(seconds=1)
|
||||
|
||||
otel._handle_failure(kwargs, response_obj=None, start_time=start, end_time=end)
|
||||
otel._handle_failure(kwargs, response_obj=None, start_time=start, end_time=end)
|
||||
|
||||
spans = span_exporter.get_finished_spans()
|
||||
litellm_spans = [s for s in spans if s.name == LITELLM_REQUEST_SPAN_NAME]
|
||||
self.assertEqual(
|
||||
len(litellm_spans),
|
||||
1,
|
||||
f"Exactly one litellm_request ERROR span expected, got {len(litellm_spans)}",
|
||||
)
|
||||
self.assertEqual(litellm_spans[0].status.status_code, StatusCode.ERROR)
|
||||
|
||||
def test_create_guardrail_span_dedupes_across_lifecycle_entrypoints(self):
|
||||
"""``_create_guardrail_span`` is called from post-call-success hook,
|
||||
``_handle_success``, and ``_handle_failure``. A single guardrail
|
||||
invocation (identified by ``(name, start_time, mode)``) must produce
|
||||
exactly one span per handler even when the underlying entry is
|
||||
mutated between calls (e.g. proxy enriches ``guardrail_response``)."""
|
||||
span_exporter = InMemorySpanExporter()
|
||||
tracer_provider = TracerProvider()
|
||||
tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter))
|
||||
|
||||
otel = OpenTelemetry(tracer_provider=tracer_provider)
|
||||
otel.tracer = tracer_provider.get_tracer(__name__)
|
||||
|
||||
kwargs = self._build_kwargs()
|
||||
guardrail_entry = {
|
||||
"guardrail_name": "block-code",
|
||||
"guardrail_mode": "pre_call",
|
||||
"guardrail_response": "allow",
|
||||
"start_time": 1.0,
|
||||
"end_time": 2.0,
|
||||
}
|
||||
kwargs["standard_logging_object"]["guardrail_information"] = [guardrail_entry]
|
||||
|
||||
otel._create_guardrail_span(kwargs=kwargs, context=None)
|
||||
# Mutate the entry between calls — proxy enriches the response.
|
||||
guardrail_entry["guardrail_response"] = [
|
||||
{"type": "code_block", "action_taken": "block"}
|
||||
]
|
||||
guardrail_entry["end_time"] = 3.0
|
||||
otel._create_guardrail_span(kwargs=kwargs, context=None)
|
||||
otel._create_guardrail_span(kwargs=kwargs, context=None)
|
||||
|
||||
guardrail_spans = [
|
||||
s for s in span_exporter.get_finished_spans() if s.name == "guardrail"
|
||||
]
|
||||
self.assertEqual(
|
||||
len(guardrail_spans),
|
||||
1,
|
||||
f"Exactly one guardrail span expected per logical invocation, got {len(guardrail_spans)}",
|
||||
)
|
||||
|
||||
def test_create_guardrail_span_emits_distinct_entries(self):
|
||||
"""Two real guardrail invocations (different ``start_time``) must
|
||||
each emit a span — entry-level dedupe must not collapse them."""
|
||||
span_exporter = InMemorySpanExporter()
|
||||
tracer_provider = TracerProvider()
|
||||
tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter))
|
||||
|
||||
otel = OpenTelemetry(tracer_provider=tracer_provider)
|
||||
otel.tracer = tracer_provider.get_tracer(__name__)
|
||||
|
||||
kwargs = self._build_kwargs()
|
||||
kwargs["standard_logging_object"]["guardrail_information"] = [
|
||||
{
|
||||
"guardrail_name": "block-code",
|
||||
"guardrail_mode": "pre_call",
|
||||
"guardrail_response": "allow",
|
||||
"start_time": 1.0,
|
||||
"end_time": 2.0,
|
||||
},
|
||||
{
|
||||
"guardrail_name": "block-code",
|
||||
"guardrail_mode": "post_call",
|
||||
"guardrail_response": "allow",
|
||||
"start_time": 5.0,
|
||||
"end_time": 6.0,
|
||||
},
|
||||
]
|
||||
|
||||
otel._create_guardrail_span(kwargs=kwargs, context=None)
|
||||
otel._create_guardrail_span(kwargs=kwargs, context=None)
|
||||
|
||||
guardrail_spans = [
|
||||
s for s in span_exporter.get_finished_spans() if s.name == "guardrail"
|
||||
]
|
||||
self.assertEqual(
|
||||
len(guardrail_spans),
|
||||
2,
|
||||
f"Two distinct guardrail invocations expected, got {len(guardrail_spans)}",
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user