ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[internal copy of #29089] fix: duplicate claude code traces (#29311)

Browse Source
This commit is contained in:
Mateo Wang 2026-05-29 22:23:24 -07:00 committed by GitHub
parent c754c560dd
commit 581c30f1e8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
10 changed files with 566 additions and 122 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

106
litellm/litellm_core_utils/litellm_logging.py
View File

@ -1612,6 +1612,90 @@ class Logging(LiteLLMLoggingBaseClass):
) -> Optional[float]:
return self._response_cost_calculator(result=result, cache_hit=cache_hit)
@staticmethod
def _is_sync_litellm_request(litellm_params: dict) -> bool:
"""True for sync SDK entrypoints (``completion``), false for async (``acompletion``, etc.)."""
return (
litellm_params.get(CallTypes.acompletion.value, False) is not True
and litellm_params.get(CallTypes.aresponses.value, False) is not True
and litellm_params.get(CallTypes.aembedding.value, False) is not True
and litellm_params.get(CallTypes.aimage_generation.value, False) is not True
and litellm_params.get(CallTypes.atranscription.value, False) is not True
)
def _is_assembled_stream_success(self, result=None) -> bool:
"""Final assembled stream export (not a per-chunk success call).
Per-chunk callers pass a ``ModelResponseStream`` (or ``None``); the
final assembled response is any other non-``None`` value (typically a
``ModelResponse``). Treating a chunk as the assembled response would
prematurely set the ``has_dispatched_final_stream_success`` dedup
guard and silently suppress the real final stream log.
"""
if self.stream is not True:
return False
if result is not None and not isinstance(result, ModelResponseStream):
return True
return (
"async_complete_streaming_response" in self.model_call_details
or self.model_call_details.get("complete_streaming_response") is not None
)
async def dispatch_success_handlers(
self,
result=None,
start_time=None,
end_time=None,
cache_hit=None,
prefer_async_handlers: bool = False,
**kwargs,
) -> None:
"""Route success logging to async and/or sync handlers for this request.
``prefer_async_handlers`` only bypasses the sync-SDK-only shortcut (e.g.
``async for`` on a stream from ``completion()``). Legacy string callbacks
still run via ``executor.submit(success_handler)`` when configured.
"""
from litellm.litellm_core_utils.thread_pool_executor import executor
if self._is_assembled_stream_success(result):
if self.model_call_details.get("has_dispatched_final_stream_success"):
return
self.model_call_details["has_dispatched_final_stream_success"] = True
litellm_params = self.model_call_details.get("litellm_params", {}) or {}
sync_sdk = self._is_sync_litellm_request(litellm_params)
passthrough = self.call_type == CallTypes.pass_through.value
if sync_sdk and not prefer_async_handlers and not passthrough:
self.success_handler(
result,
start_time=start_time,
end_time=end_time,
cache_hit=cache_hit,
**kwargs,
)
return
await self.async_success_handler(
result,
start_time=start_time,
end_time=end_time,
cache_hit=cache_hit,
**kwargs,
)
if not self._should_run_sync_callbacks_for_async_calls():
return
executor.submit(
self.success_handler,
result,
start_time=start_time,
end_time=end_time,
cache_hit=cache_hit,
**kwargs,
)
def should_run_logging(
self,
event_type: Literal[
@ -2034,13 +2118,7 @@ class Logging(LiteLLMLoggingBaseClass):
standard_logging_object=kwargs.get("standard_logging_object", None),
)
litellm_params = self.model_call_details.get("litellm_params", {})
is_sync_request = (
litellm_params.get(CallTypes.acompletion.value, False) is not True
and litellm_params.get(CallTypes.aresponses.value, False) is not True
and litellm_params.get(CallTypes.aembedding.value, False) is not True
and litellm_params.get(CallTypes.aimage_generation.value, False) is not True
and litellm_params.get(CallTypes.atranscription.value, False) is not True
)
is_sync_request = self._is_sync_litellm_request(litellm_params)
try:
## BUILD COMPLETE STREAMED RESPONSE
complete_streaming_response: Optional[
@ -2496,9 +2574,11 @@ class Logging(LiteLLMLoggingBaseClass):
print_verbose(
"Logging Details LiteLLM-Async Success Call, cache_hit={}".format(cache_hit)
)
if not self.should_run_logging(
if not self._is_assembled_stream_success(
result
) and not self.should_run_logging(
event_type="async_success"
): # prevent double logging
): # prevent double logging (non-streaming)
return
## CALCULATE COST FOR BATCH JOBS
@ -2948,13 +3028,7 @@ class Logging(LiteLLMLoggingBaseClass):
): # prevent double logging
return
litellm_params = self.model_call_details.get("litellm_params", {})
is_sync_request = (
litellm_params.get(CallTypes.acompletion.value, False) is not True
and litellm_params.get(CallTypes.aresponses.value, False) is not True
and litellm_params.get(CallTypes.aembedding.value, False) is not True
and litellm_params.get(CallTypes.aimage_generation.value, False) is not True
and litellm_params.get(CallTypes.atranscription.value, False) is not True
)
is_sync_request = self._is_sync_litellm_request(litellm_params)
try:
start_time, end_time = self._failure_handler_helper_fn(

20
litellm/litellm_core_utils/streaming_handler.py
View File

@ -1835,8 +1835,10 @@ class CustomStreamWrapper:
processed_chunk, None, None, cache_hit
)
)
## SYNC LOGGING
self.logging_obj.success_handler(processed_chunk, None, None, cache_hit)
## SYNC LOGGING — only for sync SDK entrypoints; async proxy paths export via async_success_handler
litellm_params = self.logging_obj.model_call_details.get("litellm_params", {})
if self.logging_obj._is_sync_litellm_request(litellm_params):
self.logging_obj.success_handler(processed_chunk, None, None, cache_hit)
def finish_reason_handler(self):
model_response = self.model_response_creator()
@ -2231,23 +2233,19 @@ class CustomStreamWrapper:
cache_hit,
)
else:
# prefer_async_handlers routes CustomLogger to async_success_handler
# when consumers use ``async for`` on sync-SDK streams. Legacy string
# callbacks still run via executor.submit inside dispatch_success_handlers.
asyncio.create_task(
self.logging_obj.async_success_handler(
self.logging_obj.dispatch_success_handlers(
complete_streaming_response,
cache_hit=cache_hit,
start_time=None,
end_time=None,
prefer_async_handlers=True,
)
)
executor.submit(
self.logging_obj.success_handler,
complete_streaming_response,
cache_hit=cache_hit,
start_time=None,
end_time=None,
)
raise StopAsyncIteration # Re-raise StopIteration
else:
self.sent_last_chunk = True

49
litellm/proxy/common_request_processing.py
View File

@ -1291,7 +1291,7 @@ class ProxyBaseLLMRequestProcessing:
# (ProxyLogging._fire_deferred_stream_logging) fires the
# closure after the full streaming pipeline finishes.
# The closure runs non-apply_guardrail hooks on the
# assembled response, then fires both logging handlers.
# assembled response, then fires success logging.
# Only for CustomStreamWrapper — raw async generators from
# passthrough routes bypass CSW and would orphan the closure.
from litellm.litellm_core_utils.streaming_handler import (
@ -1427,33 +1427,18 @@ class ProxyBaseLLMRequestProcessing:
logging_obj._on_deferred_stream_complete = None # type: ignore[union-attr]
try:
asyncio.create_task(
logging_obj.async_success_handler(
logging_obj.dispatch_success_handlers(
response,
cache_hit=None,
start_time=None,
end_time=None,
prefer_async_handlers=True,
)
)
except Exception as e:
verbose_proxy_logger.exception(
"Error in orphaned streaming async logging: %s", e
)
try:
from litellm.litellm_core_utils.thread_pool_executor import (
executor as _exc,
)
_exc.submit(
logging_obj.success_handler,
response,
cache_hit=None,
start_time=None,
end_time=None,
)
except Exception as e:
verbose_proxy_logger.exception(
"Error in orphaned streaming sync logging: %s", e
)
# Always return the client-requested model name (not provider-prefixed internal identifiers)
# for OpenAI-compatible responses.
@ -1740,7 +1725,7 @@ class ProxyBaseLLMRequestProcessing:
) -> None:
"""
Run non-streaming post-call guardrail hooks on an assembled streaming
response, then fire both async and sync logging handlers.
response, then fire success logging via ``dispatch_success_handlers``.
Called by ProxyLogging._fire_deferred_stream_logging after the full
streaming pipeline (including unified_guardrail end-of-stream blocks)
@ -1756,8 +1741,6 @@ class ProxyBaseLLMRequestProcessing:
Extracted as a static method so tests can call the production
implementation directly rather than reimplementing the closure.
"""
from litellm.litellm_core_utils.thread_pool_executor import executor
_response = assembled_response
try:
from litellm.proxy.proxy_server import llm_router as _global_llm_router
@ -1816,31 +1799,23 @@ class ProxyBaseLLMRequestProcessing:
)
finally:
try:
# Proxy streaming always runs in async context and proxy spend
# logging is async-only; force async dispatch so DB/spend
# callbacks fire regardless of the call-type heuristic in
# _is_sync_litellm_request (which only recognizes a subset of
# async markers stored in litellm_params).
asyncio.create_task(
captured_logging_obj.async_success_handler(
captured_logging_obj.dispatch_success_handlers(
_response,
cache_hit=cache_hit,
start_time=None,
end_time=None,
prefer_async_handlers=True,
)
)
except Exception as e:
verbose_proxy_logger.exception(
"Error in deferred streaming async logging: %s",
e,
)
try:
executor.submit(
captured_logging_obj.success_handler,
_response,
cache_hit=cache_hit,
start_time=None,
end_time=None,
)
except Exception as e:
verbose_proxy_logger.exception(
"Error in deferred streaming sync logging: %s",
"Error in deferred streaming success logging: %s",
e,
)

22
litellm/proxy/pass_through_endpoints/streaming_handler.py
View File

@ -7,7 +7,6 @@ import httpx
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.thread_pool_executor import executor
from litellm.proxy._types import PassThroughEndpointLoggingResultValues
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
@ -145,25 +144,16 @@ class PassThroughStreamingHandler:
end_time=end_time,
model=model,
)
await litellm_logging_obj.async_success_handler(
# Always reached from an async context (anthropic_messages,
# google_genai, and proxy pass-through stream tasks). prefer_async_handlers
# keeps async-only loggers running even when call_type isn't pass_through
# and litellm_params lacks an async flag (e.g. aanthropic_messages).
await litellm_logging_obj.dispatch_success_handlers(
result=standard_logging_response_object,
start_time=start_time,
end_time=end_time,
cache_hit=False,
**kwargs,
)
if (
litellm_logging_obj._should_run_sync_callbacks_for_async_calls()
is False
):
return
executor.submit(
litellm_logging_obj.success_handler,
result=standard_logging_response_object,
end_time=end_time,
cache_hit=False,
start_time=start_time,
prefer_async_handlers=True,
**kwargs,
)
except Exception as e:

24
litellm/proxy/pass_through_endpoints/success_handler.py
View File

@ -11,7 +11,6 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
PassthroughStandardLoggingPayload,
)
from litellm.types.utils import StandardPassThroughResponseObject
from litellm.utils import executor as thread_pool_executor
from .llm_provider_handlers.anthropic_passthrough_logging_handler import (
AnthropicPassthroughLoggingHandler,
@ -94,19 +93,15 @@ class PassThroughEndpointLogging:
cache_hit: bool,
**kwargs,
):
"""Helper function to handle both sync and async logging operations"""
# Submit to thread pool for sync logging
thread_pool_executor.submit(
logging_obj.success_handler,
standard_logging_response_object,
start_time,
end_time,
cache_hit,
**kwargs,
)
# Handle async logging
await logging_obj.async_success_handler(
"""Log pass-through success via the shared async dispatch path."""
# Always reached from pass_through_async_success_handler, which runs in
# an async context. call_type is "pass_through_endpoint" here, so the
# passthrough guard in dispatch_success_handlers already forces the
# async handler to run; pass prefer_async_handlers explicitly to match
# the streaming sibling (_route_streaming_logging_to_handler) and keep
# async-only loggers (e.g. the proxy spend logger) firing regardless of
# how the call-type classification evolves.
await logging_obj.dispatch_success_handlers(
result=(
json.dumps(result)
if isinstance(result, dict)
@ -115,6 +110,7 @@ class PassThroughEndpointLogging:
start_time=start_time,
end_time=end_time,
cache_hit=False,
prefer_async_handlers=True,
**kwargs,
)

117
tests/pass_through_unit_tests/test_unit_test_streaming.py
View File

@ -97,6 +97,123 @@ async def test_chunk_processor_yields_raw_bytes(endpoint_type, url_route):
), "Collected chunks do not match raw chunks"
@pytest.mark.asyncio
async def test_route_streaming_logging_runs_async_handler_for_sdk_passthrough():
"""
SDK pass-through streaming (anthropic_messages, google generate_content) must run
the async success handler so async-only loggers record the assembled stream.
Regression for duplicate-trace dedupe: dispatch_success_handlers treated these as
sync SDK requests because call_type is not ``pass_through_endpoint`` and
litellm_params carries no ``acompletion`` flag, so only the sync success_handler
ran and CustomLogger.async_log_success_event never fired.
"""
import time
from litellm.types.utils import CallTypes
logging_obj = LiteLLMLoggingObj(
model="claude-sonnet-4-5",
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type=CallTypes.anthropic_messages.value,
start_time=time.time(),
litellm_call_id="test-id",
function_id="fn",
)
logging_obj.model_call_details["litellm_params"] = {"anthropic_messages": True}
with (
patch.object(
PassThroughStreamingHandler,
"_build_passthrough_logging_result",
return_value=({"id": "slp"}, {}),
),
patch.object(
logging_obj, "async_success_handler", new_callable=AsyncMock
) as mock_async,
patch.object(
logging_obj, "success_handler", new_callable=MagicMock
) as mock_sync,
patch.object(
logging_obj,
"_should_run_sync_callbacks_for_async_calls",
return_value=False,
),
):
await PassThroughStreamingHandler._route_streaming_logging_to_handler(
litellm_logging_obj=logging_obj,
passthrough_success_handler_obj=MagicMock(),
url_route="/v1/messages",
request_body={},
endpoint_type=EndpointType.ANTHROPIC,
start_time=datetime.now(),
raw_bytes=[],
end_time=datetime.now(),
)
mock_async.assert_awaited_once()
mock_sync.assert_not_called()
@pytest.mark.asyncio
async def test_handle_logging_runs_async_handler_for_passthrough():
"""
Non-streaming pass-through logging (_handle_logging) must always run the
async success handler so async-only loggers (e.g. the proxy spend logger)
record the request.
_handle_logging is only ever reached from pass_through_async_success_handler
(an async context), so it forces async dispatch via prefer_async_handlers.
This pins that contract independent of the call-type classification: even a
call_type that _is_sync_litellm_request would classify as sync (here
"completion" with no async marker in litellm_params) must still reach
async_success_handler. Without prefer_async_handlers=True the sync-only
branch would return early and async_log_success_event would never fire.
"""
import time
from litellm.types.utils import CallTypes
logging_obj = LiteLLMLoggingObj(
model="claude-sonnet-4-5",
messages=[{"role": "user", "content": "hi"}],
stream=False,
call_type=CallTypes.completion.value,
start_time=time.time(),
litellm_call_id="test-id",
function_id="fn",
)
logging_obj.model_call_details["litellm_params"] = {}
handler = PassThroughEndpointLogging()
with (
patch.object(
logging_obj, "async_success_handler", new_callable=AsyncMock
) as mock_async,
patch.object(
logging_obj, "success_handler", new_callable=MagicMock
) as mock_sync,
patch.object(
logging_obj,
"_should_run_sync_callbacks_for_async_calls",
return_value=False,
),
):
await handler._handle_logging(
logging_obj=logging_obj,
standard_logging_response_object={"id": "slp"},
result="",
start_time=datetime.now(),
end_time=datetime.now(),
cache_hit=False,
)
mock_async.assert_awaited_once()
mock_sync.assert_not_called()
def test_convert_raw_bytes_to_str_lines():
"""
Test that the _convert_raw_bytes_to_str_lines method correctly converts raw bytes to a list of strings

19
tests/proxy_unit_tests/test_proxy_reject_logging.py
View File

@ -95,6 +95,21 @@ router = Router(
)
def _register_proxy_test_logger(callback_logger: testLogger) -> None:
"""
Register the test logger on global callback lists.
``function_setup`` dedupes by object identity; each parametrized case
constructs a new ``testLogger`` and must replace the global lists, not
only ``litellm.callbacks``.
"""
litellm.callbacks = [callback_logger]
litellm.success_callback = [callback_logger]
litellm.failure_callback = [callback_logger]
litellm._async_success_callback = [callback_logger]
litellm._async_failure_callback = [callback_logger]
@pytest.mark.parametrize(
"route, body",
[
@ -115,7 +130,7 @@ router = Router(
"/v1/embeddings",
{
"input": "The food was delicious and the waiter...",
"model": "text-embedding-ada-002",
"model": "fake-model",
"encoding_format": "float",
},
),
@ -133,7 +148,7 @@ async def test_chat_completion_request_with_redaction(route, body):
setattr(proxy_server, "llm_router", router)
_test_logger = testLogger()
litellm.callbacks = [_test_logger]
_register_proxy_test_logger(_test_logger)
litellm.set_verbose = True
# Prepare the query string

212
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
View File

@ -1,6 +1,7 @@
import os
import sys
from unittest.mock import MagicMock, patch
import asyncio
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -786,6 +787,211 @@ def test_success_handler_runs_sync_callbacks_for_sync_requests(logging_obj, call
dummy_logger.log_stream_event.assert_not_called()
def test_is_sync_litellm_request():
assert LitellmLogging._is_sync_litellm_request({}) is True
assert LitellmLogging._is_sync_litellm_request({"acompletion": True}) is False
@pytest.mark.asyncio
async def test_dispatch_success_handlers_invokes_callbacks_once_for_final_stream(
logging_obj,
):
"""Second final-stream dispatch must not re-export (CSW + deferred guardrail paths)."""
import litellm
from litellm.integrations.custom_logger import CustomLogger
class MockCallback(CustomLogger):
pass
mock_callback = MockCallback()
original_async_callbacks = list(litellm._async_success_callback or [])
litellm._async_success_callback = [mock_callback]
result = ModelResponse(
id="resp-dedupe",
model="gpt-4o-mini",
choices=[
{
"message": {"role": "assistant", "content": "hi"},
"finish_reason": "stop",
"index": 0,
}
],
usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
)
try:
logging_obj.stream = True
logging_obj.model_call_details["litellm_params"] = {"acompletion": True}
with (
patch.object(
mock_callback, "async_log_success_event", new_callable=AsyncMock
) as mock_async_log,
patch.object(mock_callback, "log_success_event") as mock_sync_log,
patch.object(
logging_obj,
"_success_handler_helper_fn",
return_value=(time.time(), time.time(), result),
),
patch.object(
logging_obj,
"_get_assembled_streaming_response",
return_value=result,
),
patch.object(
logging_obj,
"_should_run_sync_callbacks_for_async_calls",
return_value=True,
),
):
await logging_obj.dispatch_success_handlers(result=result)
await logging_obj.dispatch_success_handlers(result=result)
mock_async_log.assert_awaited_once()
mock_sync_log.assert_not_called()
finally:
litellm._async_success_callback = original_async_callbacks
@pytest.mark.asyncio
async def test_dispatch_success_handlers_sync_path_invokes_callback_once_for_final_stream(
logging_obj,
):
"""Sync dispatch path must also dedupe when dispatch is called twice."""
import litellm
from litellm.integrations.custom_logger import CustomLogger
class MockCallback(CustomLogger):
pass
mock_callback = MockCallback()
original_success_callbacks = list(litellm.success_callback or [])
litellm.success_callback = [mock_callback]
result = ModelResponse(
id="resp-sync-dedupe",
model="gpt-4o-mini",
choices=[
{
"message": {"role": "assistant", "content": "hi"},
"finish_reason": "stop",
"index": 0,
}
],
usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
)
try:
logging_obj.stream = True
logging_obj.model_call_details["litellm_params"] = {}
with (
patch.object(mock_callback, "log_success_event") as mock_sync_log,
patch.object(
mock_callback, "async_log_success_event", new_callable=AsyncMock
) as mock_async_log,
patch.object(
logging_obj,
"_success_handler_helper_fn",
return_value=(time.time(), time.time(), result),
),
patch.object(
logging_obj,
"_get_assembled_streaming_response",
return_value=result,
),
):
await logging_obj.dispatch_success_handlers(result=result)
await logging_obj.dispatch_success_handlers(result=result)
mock_sync_log.assert_called_once()
mock_async_log.assert_not_awaited()
finally:
litellm.success_callback = original_success_callbacks
@pytest.mark.asyncio
async def test_dispatch_prefer_async_handlers_runs_legacy_callbacks(
logging_obj,
):
"""``prefer_async_handlers`` must not skip executor.submit for string callbacks."""
result = ModelResponse(
id="resp-prefer-async",
model="gpt-4o-mini",
choices=[
{
"message": {"role": "assistant", "content": "hi"},
"finish_reason": "stop",
"index": 0,
}
],
)
logging_obj.stream = True
logging_obj.model_call_details["litellm_params"] = {}
with (
patch.object(
logging_obj, "async_success_handler", new_callable=AsyncMock
) as mock_async,
patch.object(
logging_obj, "success_handler", new_callable=MagicMock
) as mock_sync,
patch.object(
logging_obj,
"_should_run_sync_callbacks_for_async_calls",
return_value=True,
),
patch(
"litellm.litellm_core_utils.litellm_logging.executor.submit"
) as mock_submit,
):
await logging_obj.dispatch_success_handlers(
result=result,
prefer_async_handlers=True,
)
mock_async.assert_awaited_once()
mock_sync.assert_not_called()
mock_submit.assert_called_once()
@pytest.mark.asyncio
async def test_dispatch_success_handlers_invokes_async_callback_for_pass_through(
logging_obj,
):
"""Pass-through must use async_success_handler (CustomLogger skips sync success_handler)."""
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.utils import CallTypes
class MockCallback(CustomLogger):
pass
mock_callback = MockCallback()
original_async_callbacks = list(litellm._async_success_callback or [])
litellm._async_success_callback = [mock_callback]
logging_obj.call_type = CallTypes.pass_through.value
logging_obj.stream = False
logging_obj.model_call_details["litellm_params"] = {}
try:
with (
patch.object(
mock_callback, "async_log_success_event", new_callable=AsyncMock
) as mock_async_log,
patch.object(mock_callback, "log_success_event") as mock_sync_log,
):
await logging_obj.dispatch_success_handlers(result={"id": "pt-1"})
mock_async_log.assert_awaited_once()
mock_sync_log.assert_not_called()
finally:
litellm._async_success_callback = original_async_callbacks
def test_success_handler_skips_guardrail_logging_hook_when_disabled(logging_obj):
"""Ensure CustomGuardrail logging_hook is skipped when should_run_guardrail is False."""
import datetime
@ -1351,7 +1557,7 @@ async def test_e2e_generate_cold_storage_object_key_with_custom_logger_s3_path()
Test that _generate_cold_storage_object_key uses s3_path from custom logger instance.
"""
from datetime import datetime, timezone
from unittest.mock import MagicMock, patch
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
@ -1404,7 +1610,7 @@ async def test_e2e_generate_cold_storage_object_key_with_logger_no_s3_path():
Test that _generate_cold_storage_object_key falls back to empty s3_path when logger has no s3_path.
"""
from datetime import datetime, timezone
from unittest.mock import MagicMock, patch
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup

16
tests/test_litellm/litellm_core_utils/test_streaming_handler.py
View File

@ -569,8 +569,6 @@ async def test_streaming_with_usage_and_logging(sync_mode: bool):
== final_usage_block
)
print(mock_log_success_event.call_args.kwargs.keys())
def test_streaming_handler_with_stop_chunk(
initialized_custom_stream_wrapper: CustomStreamWrapper,
@ -2036,23 +2034,19 @@ async def test_azure_streaming_role_preserved_with_include_usage(sync_mode: bool
chunks.append(chunk)
# The prompt_filter chunk should be forwarded with choices=[]
assert len(chunks[0].choices) == 0, (
f"Expected prompt_filter chunk with choices=[], got {len(chunks[0].choices)} choices"
)
assert (
len(chunks[0].choices) == 0
), f"Expected prompt_filter chunk with choices=[], got {len(chunks[0].choices)} choices"
# At least one chunk must have role='assistant' in its delta
has_role = any(
len(c.choices) > 0
and getattr(c.choices[0].delta, "role", None) == "assistant"
len(c.choices) > 0 and getattr(c.choices[0].delta, "role", None) == "assistant"
for c in chunks
)
assert has_role, (
"No chunk contained role='assistant' in delta (issue #24221). "
"Chunk deltas: "
+ str([
c.choices[0].delta if c.choices else "no choices"
for c in chunks
])
+ str([c.choices[0].delta if c.choices else "no choices" for c in chunks])
)

103
tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py
View File

@ -18,7 +18,7 @@ import asyncio
import os
import sys
from typing import Any
from unittest.mock import MagicMock, patch
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -38,6 +38,24 @@ from litellm.types.guardrails import GuardrailEventHooks
# ---------------------------------------------------------------------------
def _attach_mock_success_dispatch(mock_logging_obj, async_success_fn):
"""Match production entrypoint: ``_run_deferred_stream_guardrails`` uses dispatch."""
async def dispatch_success_handlers(
result=None, start_time=None, end_time=None, cache_hit=None, **kwargs
):
await async_success_fn(
result,
start_time=start_time,
end_time=end_time,
cache_hit=cache_hit,
**kwargs,
)
mock_logging_obj.dispatch_success_handlers = dispatch_success_handlers
mock_logging_obj.async_success_handler = async_success_fn
class PostCallGuardrail(CustomGuardrail):
"""A post-call guardrail."""
@ -454,7 +472,7 @@ class TestDeferredStreamingClosure:
async def track_async_success(*args, **kwargs):
pass
mock_logging_obj.async_success_handler = track_async_success
_attach_mock_success_dispatch(mock_logging_obj, track_async_success)
tracking_guardrail = TrackingGuardrail()
tracking_logger = TrackingLogger()
@ -511,7 +529,7 @@ class TestDeferredStreamingClosure:
nonlocal logged_response
logged_response = args[0] if args else None
mock_logging_obj.async_success_handler = track_async_success
_attach_mock_success_dispatch(mock_logging_obj, track_async_success)
class ModifyingGuardrail(CustomGuardrail):
def __init__(self):
@ -573,7 +591,7 @@ class TestDeferredStreamingClosure:
nonlocal logging_called
logging_called = True
mock_logging_obj.async_success_handler = track_async_success
_attach_mock_success_dispatch(mock_logging_obj, track_async_success)
guardrail = BlockingGuardrail()
@ -621,7 +639,7 @@ class TestDeferredStreamingClosure:
async def track_async_success(*args, **kwargs):
pass
mock_logging_obj.async_success_handler = track_async_success
_attach_mock_success_dispatch(mock_logging_obj, track_async_success)
guardrail = TransientErrorGuardrail()
@ -656,7 +674,7 @@ class TestDeferredStreamingClosure:
nonlocal logged_response
logged_response = args[0] if args else None
mock_logging_obj.async_success_handler = track_async_success
_attach_mock_success_dispatch(mock_logging_obj, track_async_success)
class TestGuardrail(CustomGuardrail):
def __init__(self):
@ -739,7 +757,7 @@ class TestDeferredStreamingClosure:
async def track_async_success(*args, **kwargs):
pass
mock_logging_obj.async_success_handler = track_async_success
_attach_mock_success_dispatch(mock_logging_obj, track_async_success)
guardrail = ApplyGuardrailType()
@ -792,7 +810,7 @@ class TestDeferredStreamingClosure:
async def track_async_success(*args, **kwargs):
pass
mock_logging_obj.async_success_handler = track_async_success
_attach_mock_success_dispatch(mock_logging_obj, track_async_success)
guardrail = IteratorHookGuardrail()
@ -847,7 +865,7 @@ class TestDeferredStreamingClosure:
async def track_async_success(*args, **kwargs):
pass
mock_logging_obj.async_success_handler = track_async_success
_attach_mock_success_dispatch(mock_logging_obj, track_async_success)
guardrail = InspectingGuardrail()
@ -914,7 +932,7 @@ class TestDeferredStreamingClosure:
async def track_async_success(*args, **kwargs):
pass
mock_logging_obj.async_success_handler = track_async_success
_attach_mock_success_dispatch(mock_logging_obj, track_async_success)
guardrail_a = TaggedGuardrail("guardrail-a")
guardrail_b = TaggedGuardrail("guardrail-b")
@ -962,7 +980,7 @@ class TestDeferredStreamingClosure:
nonlocal logging_called
logging_called = True
mock_logging_obj.async_success_handler = track_async_success
_attach_mock_success_dispatch(mock_logging_obj, track_async_success)
def exploding_merge(data, llm_router):
raise RuntimeError("Simulated init failure")
@ -986,6 +1004,67 @@ class TestDeferredStreamingClosure:
logging_called is True
), "Logging must fire even when guardrail initialization raises"
@pytest.mark.asyncio
async def test_deferred_logging_forces_async_for_sync_classified_call_type(self):
"""
Regression: proxy deferred streaming logging must reach the async success
handler (which runs the async-only DB/spend logger) even when the call
type is classified as a sync SDK request by _is_sync_litellm_request.
Without prefer_async_handlers=True, an async proxy stream whose
litellm_params lacks a recognized async marker would enter the sync
branch of dispatch_success_handlers and silently skip spend tracking.
Uses the real dispatch_success_handlers via the production
_run_deferred_stream_guardrails entrypoint.
"""
import time
from litellm.litellm_core_utils.litellm_logging import (
Logging as LiteLLMLoggingObj,
)
logging_obj = LiteLLMLoggingObj(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type="completion", # not pass_through_endpoint
start_time=time.time(),
litellm_call_id="test-id",
function_id="fn",
)
# litellm_params with no recognized async marker -> classified sync.
logging_obj.model_call_details["litellm_params"] = {}
assert LiteLLMLoggingObj._is_sync_litellm_request({}) is True
with (
patch.object(
logging_obj, "async_success_handler", new_callable=AsyncMock
) as mock_async,
patch.object(
logging_obj, "success_handler", new_callable=MagicMock
) as mock_sync,
patch.object(
logging_obj,
"_should_run_sync_callbacks_for_async_calls",
return_value=False,
),
patch("litellm.callbacks", [PostCallGuardrail()]),
):
await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails(
captured_data={"model": "gpt-4o-mini", "metadata": {}},
captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"),
captured_logging_obj=logging_obj,
assembled_response=MagicMock(),
cache_hit=False,
)
await asyncio.sleep(0)
await asyncio.sleep(0)
mock_async.assert_awaited_once()
mock_sync.assert_not_called()
# ---------------------------------------------------------------------------
# 7. _fire_deferred_stream_logging
@ -1054,7 +1133,7 @@ class TestFireDeferredStreamLogging:
nonlocal logged_response
logged_response = args[0] if args else None
mock_logging_obj.async_success_handler = track_async_success
_attach_mock_success_dispatch(mock_logging_obj, track_async_success)
class InfoWritingGuardrail(CustomGuardrail):
def __init__(self):