diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 97266096ef91..82f37cf1deb7 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1612,6 +1612,83 @@ async def _response_cost_calculator_async( ) -> 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).""" + if self.stream is not True: + return False + if result is not None: + 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 +2111,7 @@ def success_handler( # noqa: PLR0915 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 +2567,11 @@ async def async_success_handler( # noqa: PLR0915 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 +3021,7 @@ def failure_handler( # noqa: PLR0915 ): # 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( diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index fa7faf3035d3..29c0d0629e87 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1808,8 +1808,10 @@ def run_success_logging_and_cache_storage(self, processed_chunk, cache_hit: bool 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() @@ -2206,23 +2208,19 @@ async def __anext__(self) -> "ModelResponseStream": # noqa: PLR0915 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 diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index ef1d64335b40..094d8548a665 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1290,7 +1290,7 @@ async def base_process_llm_request( # noqa: PLR0915 # (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 ( @@ -1411,7 +1411,7 @@ async def _on_deferred_stream_complete( 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, @@ -1422,22 +1422,6 @@ async def _on_deferred_stream_complete( 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. @@ -1639,7 +1623,7 @@ async def _run_deferred_stream_guardrails( ) -> 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) @@ -1655,8 +1639,6 @@ async def _run_deferred_stream_guardrails( 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 @@ -1716,7 +1698,7 @@ async def _run_deferred_stream_guardrails( finally: try: asyncio.create_task( - captured_logging_obj.async_success_handler( + captured_logging_obj.dispatch_success_handlers( _response, cache_hit=cache_hit, start_time=None, @@ -1725,21 +1707,7 @@ async def _run_deferred_stream_guardrails( ) 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, ) diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 235a38b75f9e..915134705721 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -7,7 +7,6 @@ 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,27 +144,13 @@ async def _route_streaming_logging_to_handler( end_time=end_time, model=model, ) - await litellm_logging_obj.async_success_handler( + 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, - **kwargs, - ) except Exception as e: verbose_proxy_logger.error( f"Error in _route_streaming_logging_to_handler: {str(e)}" diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index 0bc0183aa7c6..6282f5072fc7 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -11,7 +11,6 @@ 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,8 @@ async def _handle_logging( 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.""" + await logging_obj.dispatch_success_handlers( result=( json.dumps(result) if isinstance(result, dict) diff --git a/tests/proxy_unit_tests/test_proxy_reject_logging.py b/tests/proxy_unit_tests/test_proxy_reject_logging.py index 51a92fa3b4b6..e0b575f4a71c 100644 --- a/tests/proxy_unit_tests/test_proxy_reject_logging.py +++ b/tests/proxy_unit_tests/test_proxy_reject_logging.py @@ -95,6 +95,21 @@ def log_failure_event(self, kwargs, response_obj, start_time, end_time): ) +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 @@ def log_failure_event(self, kwargs, response_obj, start_time, end_time): "/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 diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 07ab29c5231d..b64cb7c69056 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -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 @@ class DummyLogger(CustomLogger): 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 diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 49d3c51e340f..63e2cb7f35c1 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -569,8 +569,6 @@ class MockCallback(CustomLogger): == 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]) ) diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py index e10258c0829d..f6803c0a35b2 100644 --- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -38,6 +38,24 @@ # --------------------------------------------------------------------------- +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 @@ async def async_post_call_success_hook( 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 @@ async def track_async_success(*args, **kwargs): 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 @@ async def track_async_success(*args, **kwargs): 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 @@ async def async_post_call_success_hook( 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 @@ async def track_async_success(*args, **kwargs): 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 @@ async def apply_guardrail( 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 @@ async def async_post_call_success_hook( 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 @@ async def async_post_call_success_hook( 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 @@ async def async_post_call_success_hook( 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 @@ async def track_async_success(*args, **kwargs): 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") @@ -1054,7 +1072,7 @@ async def track_async_success(*args, **kwargs): 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):