diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 34ea598a65..f571da3cd1 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -65,6 +65,7 @@ from litellm.responses.streaming_iterator import ( ResponsesAPIStreamingIterator, SyncResponsesAPIStreamingIterator, ) +from litellm.types.utils import CallTypes from litellm.types.containers.main import ( ContainerFileListResponse, ContainerListResponse, @@ -2060,6 +2061,13 @@ class BaseLLMHTTPHandler: if extra_body: data.update(extra_body) + request_context: Dict[str, Any] = {"input": input} + try: + request_context.update(response_api_optional_request_params) + except Exception: + pass + request_context["litellm_params"] = dict(litellm_params) + ## LOGGING logging_obj.pre_call( input=input, @@ -2097,6 +2105,8 @@ class BaseLLMHTTPHandler: responses_api_provider_config=responses_api_provider_config, litellm_metadata=litellm_metadata, custom_llm_provider=custom_llm_provider, + request_data=request_context, + call_type=CallTypes.responses.value, ) return SyncResponsesAPIStreamingIterator( @@ -2106,6 +2116,8 @@ class BaseLLMHTTPHandler: responses_api_provider_config=responses_api_provider_config, litellm_metadata=litellm_metadata, custom_llm_provider=custom_llm_provider, + request_data=request_context, + call_type=CallTypes.responses.value, ) else: # For non-streaming requests @@ -2189,6 +2201,13 @@ class BaseLLMHTTPHandler: if extra_body: data.update(extra_body) + request_context: Dict[str, Any] = {"input": input} + try: + request_context.update(response_api_optional_request_params) + except Exception: + pass + request_context["litellm_params"] = dict(litellm_params) + ## LOGGING logging_obj.pre_call( input=input, @@ -2227,6 +2246,8 @@ class BaseLLMHTTPHandler: responses_api_provider_config=responses_api_provider_config, litellm_metadata=litellm_metadata, custom_llm_provider=custom_llm_provider, + request_data=request_context, + call_type=CallTypes.responses.value, ) # Return the streaming iterator @@ -2237,6 +2258,8 @@ class BaseLLMHTTPHandler: responses_api_provider_config=responses_api_provider_config, litellm_metadata=litellm_metadata, custom_llm_provider=custom_llm_provider, + request_data=request_context, + call_type=CallTypes.responses.value, ) else: # For non-streaming, proceed as before @@ -8288,4 +8311,4 @@ class BaseLLMHTTPHandler: return skills_api_provider_config.transform_delete_skill_response( raw_response=response, logging_obj=logging_obj, - ) \ No newline at end of file + ) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 0407776029..0b838f916e 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1,5 +1,6 @@ import asyncio import json +import traceback from datetime import datetime from typing import Any, Dict, Optional @@ -11,6 +12,9 @@ from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base +from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( + update_response_metadata, +) from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.responses.utils import ResponsesAPIRequestUtils @@ -22,7 +26,8 @@ from litellm.types.llms.openai import ( ResponsesAPIStreamEvents, ResponsesAPIStreamingResponse, ) -from litellm.utils import CustomStreamWrapper +from litellm.types.utils import CallTypes +from litellm.utils import CustomStreamWrapper, async_post_call_success_deployment_hook class BaseResponsesAPIStreamingIterator: @@ -40,6 +45,8 @@ class BaseResponsesAPIStreamingIterator: logging_obj: LiteLLMLoggingObj, litellm_metadata: Optional[Dict[str, Any]] = None, custom_llm_provider: Optional[str] = None, + request_data: Optional[Dict[str, Any]] = None, + call_type: Optional[str] = None, ): self.response = response self.model = model @@ -47,21 +54,25 @@ class BaseResponsesAPIStreamingIterator: self.finished = False self.responses_api_provider_config = responses_api_provider_config self.completed_response: Optional[ResponsesAPIStreamingResponse] = None - self.start_time = datetime.now() + self.start_time = getattr(logging_obj, "start_time", datetime.now()) - # set request kwargs + # track request context for hooks self.litellm_metadata = litellm_metadata self.custom_llm_provider = custom_llm_provider + self.request_data: Dict[str, Any] = request_data or {} + self.call_type: Optional[str] = call_type # set hidden params for response headers (e.g., x-litellm-model-id) - # This matches ths stream wrapper in litellm/litellm_core_utils/streaming_handler.py + # This matches the stream wrapper in litellm/litellm_core_utils/streaming_handler.py _api_base = get_api_base( model=model or "", optional_params=self.logging_obj.model_call_details.get( "litellm_params", {} ), ) - _model_info: Dict = litellm_metadata.get("model_info", {}) if litellm_metadata else {} + _model_info: Dict = ( + litellm_metadata.get("model_info", {}) if litellm_metadata else {} + ) self._hidden_params = { "model_id": _model_info.get("id", None), "api_base": _api_base, @@ -102,13 +113,21 @@ class BaseResponsesAPIStreamingIterator: # if "response" in parsed_chunk, then encode litellm specific information like custom_llm_provider response_object = getattr(openai_responses_api_chunk, "response", None) if response_object: - response = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( - responses_api_response=response_object, - litellm_metadata=self.litellm_metadata, - custom_llm_provider=self.custom_llm_provider, + response = ( + ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( + responses_api_response=response_object, + litellm_metadata=self.litellm_metadata, + custom_llm_provider=self.custom_llm_provider, + ) ) setattr(openai_responses_api_chunk, "response", response) + # Allow callbacks to modify chunk before returning + openai_responses_api_chunk = run_async_function( + async_function=self._call_post_streaming_deployment_hook, + chunk=openai_responses_api_chunk, + ) + # Store the completed response if ( openai_responses_api_chunk @@ -149,11 +168,159 @@ class BaseResponsesAPIStreamingIterator: except json.JSONDecodeError: # If we can't parse the chunk, continue return None + except Exception as e: + # Ensure failures trigger failure hooks + self._handle_failure(e) + raise def _handle_logging_completed_response(self): """Base implementation - should be overridden by subclasses""" pass + async def _call_post_streaming_deployment_hook(self, chunk): + """ + Allow callbacks to modify streaming chunks before returning (parity with chat). + """ + try: + # Align with chat pipeline: use logging_obj model_call_details + call_type + typed_call_type: Optional[CallTypes] = None + if self.call_type is not None: + try: + typed_call_type = CallTypes(self.call_type) + except ValueError: + typed_call_type = None + if typed_call_type is None: + try: + typed_call_type = CallTypes(getattr(self.logging_obj, "call_type", None)) + except Exception: + typed_call_type = None + + request_data = self.request_data or getattr( + self.logging_obj, "model_call_details", {} + ) + callbacks = getattr(litellm, "callbacks", None) or [] + hooks_ran = False + for callback in callbacks: + if hasattr(callback, "async_post_call_streaming_deployment_hook"): + hooks_ran = True + result = await callback.async_post_call_streaming_deployment_hook( + request_data=request_data, + response_chunk=chunk, + call_type=typed_call_type, + ) + if result is not None: + chunk = result + if hooks_ran: + setattr(chunk, "_post_streaming_hooks_ran", True) + return chunk + except Exception: + return chunk + + async def call_post_streaming_hooks_for_testing(self, chunk): + """ + Helper to invoke streaming deployment hooks explicitly (used in tests). + """ + return await self._call_post_streaming_deployment_hook(chunk) + + def _run_post_success_hooks(self, end_time: datetime): + """ + Run post-call deployment hooks and update metadata similar to chat pipeline. + """ + if self.completed_response is None: + return + + request_payload: Dict[str, Any] = {} + if isinstance(self.request_data, dict): + request_payload.update(self.request_data) + try: + if hasattr(self.logging_obj, "model_call_details"): + request_payload.update(self.logging_obj.model_call_details) + except Exception: + pass + if "litellm_params" not in request_payload: + try: + request_payload["litellm_params"] = getattr( + self.logging_obj, "model_call_details", {} + ).get("litellm_params", {}) + except Exception: + request_payload["litellm_params"] = {} + + try: + update_response_metadata( + result=self.completed_response, + logging_obj=self.logging_obj, + model=self.model, + kwargs=request_payload, + start_time=self.start_time, + end_time=end_time, + ) + except Exception: + # Non-blocking + pass + + try: + typed_call_type: Optional[CallTypes] = None + if self.call_type is not None: + try: + typed_call_type = CallTypes(self.call_type) + except ValueError: + typed_call_type = None + except Exception: + typed_call_type = None + if typed_call_type is None: + try: + typed_call_type = CallTypes.responses + except Exception: + typed_call_type = None + + try: + # Call synchronously; async hook will be executed via asyncio.run in a new loop + run_async_function( + async_function=async_post_call_success_deployment_hook, + request_data=request_payload, + response=self.completed_response, + call_type=typed_call_type, + ) + except Exception: + pass + + def _handle_failure(self, exception: Exception): + """ + Trigger failure handlers before bubbling the exception. + """ + traceback_exception = traceback.format_exc() + try: + run_async_function( + async_function=self.logging_obj.async_failure_handler, + exception=exception, + traceback_exception=traceback_exception, + start_time=self.start_time, + end_time=datetime.now(), + ) + except Exception: + pass + + try: + executor.submit( + self.logging_obj.failure_handler, + exception, + traceback_exception, + self.start_time, + datetime.now(), + ) + except Exception: + pass + + +async def call_post_streaming_hooks_for_testing(iterator, chunk): + """ + Module-level helper for tests to ensure hooks can be invoked even if the iterator is wrapped. + """ + hook_fn = getattr(iterator, "_call_post_streaming_deployment_hook", None) + if hook_fn is None: + return chunk + return await hook_fn(chunk) + class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): """ @@ -168,6 +335,8 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): logging_obj: LiteLLMLoggingObj, litellm_metadata: Optional[Dict[str, Any]] = None, custom_llm_provider: Optional[str] = None, + request_data: Optional[Dict[str, Any]] = None, + call_type: Optional[str] = None, ): super().__init__( response, @@ -176,6 +345,8 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): logging_obj, litellm_metadata, custom_llm_provider, + request_data, + call_type, ) self.stream_iterator = response.aiter_lines() @@ -203,16 +374,21 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): except httpx.HTTPError as e: # Handle HTTP errors self.finished = True + self._handle_failure(e) + raise e + except Exception as e: + self.finished = True + self._handle_failure(e) raise e def _handle_logging_completed_response(self): """Handle logging for completed responses in async context""" # Create a deep copy for logging to avoid modifying the response object that will be returned to the user - # The logging handlers may transform usage from Responses API format (input_tokens/output_tokens) + # The logging handlers may transform usage from Responses API format (input_tokens/output_tokens) # to chat completion format (prompt_tokens/completion_tokens) for internal logging import copy logging_response = copy.deepcopy(self.completed_response) - + asyncio.create_task( self.logging_obj.async_success_handler( result=logging_response, @@ -229,6 +405,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): start_time=self.start_time, end_time=datetime.now(), ) + self._run_post_success_hooks(end_time=datetime.now()) class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): @@ -244,6 +421,8 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): logging_obj: LiteLLMLoggingObj, litellm_metadata: Optional[Dict[str, Any]] = None, custom_llm_provider: Optional[str] = None, + request_data: Optional[Dict[str, Any]] = None, + call_type: Optional[str] = None, ): super().__init__( response, @@ -252,6 +431,8 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): logging_obj, litellm_metadata, custom_llm_provider, + request_data, + call_type, ) self.stream_iterator = response.iter_lines() @@ -279,16 +460,21 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): except httpx.HTTPError as e: # Handle HTTP errors self.finished = True + self._handle_failure(e) + raise e + except Exception as e: + self.finished = True + self._handle_failure(e) raise e def _handle_logging_completed_response(self): """Handle logging for completed responses in sync context""" # Create a deep copy for logging to avoid modifying the response object that will be returned to the user - # The logging handlers may transform usage from Responses API format (input_tokens/output_tokens) + # The logging handlers may transform usage from Responses API format (input_tokens/output_tokens) # to chat completion format (prompt_tokens/completion_tokens) for internal logging import copy logging_response = copy.deepcopy(self.completed_response) - + run_async_function( async_function=self.logging_obj.async_success_handler, result=logging_response, @@ -304,6 +490,7 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): start_time=self.start_time, end_time=datetime.now(), ) + self._run_post_success_hooks(end_time=datetime.now()) class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): @@ -324,6 +511,8 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): logging_obj: LiteLLMLoggingObj, litellm_metadata: Optional[Dict[str, Any]] = None, custom_llm_provider: Optional[str] = None, + request_data: Optional[Dict[str, Any]] = None, + call_type: Optional[str] = None, ): super().__init__( response=response, @@ -332,6 +521,8 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): logging_obj=logging_obj, litellm_metadata=litellm_metadata, custom_llm_provider=custom_llm_provider, + request_data=request_data, + call_type=call_type, ) # one-time transform diff --git a/tests/llm_responses_api_testing/test_responses_hooks.py b/tests/llm_responses_api_testing/test_responses_hooks.py new file mode 100644 index 0000000000..8c0f7dab2a --- /dev/null +++ b/tests/llm_responses_api_testing/test_responses_hooks.py @@ -0,0 +1,165 @@ +import asyncio +from datetime import datetime +from types import SimpleNamespace + +import httpx +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.responses import streaming_iterator as streaming_module +from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator +from litellm.types.llms.openai import ResponsesAPIStreamEvents +from litellm.types.utils import CallTypes + + +class _FakeLoggingObj: + def __init__(self): + self.success_calls = 0 + self.async_success_calls = 0 + self.failure_calls = 0 + self.async_failure_calls = 0 + self.start_time = datetime.now() + self.model_call_details = {"litellm_params": {}} + + # Signature alignment with Logging handlers + def success_handler(self, *args, **kwargs): + self.success_calls += 1 + + async def async_success_handler(self, *args, **kwargs): + self.async_success_calls += 1 + + def failure_handler(self, *args, **kwargs): + self.failure_calls += 1 + + async def async_failure_handler(self, *args, **kwargs): + self.async_failure_calls += 1 + + +@pytest.mark.asyncio +async def test_responses_streaming_triggers_hooks(monkeypatch): + """ + Ensure streaming iterator fires success + post-call hooks for responses API. + """ + hook_calls = {"post_call": 0, "metadata": 0} + seen = {} + + async def fake_post_call(request_data, response, call_type): + hook_calls["post_call"] += 1 + seen["request_data"] = request_data + seen["call_type"] = call_type + + def fake_update_metadata(**kwargs): + hook_calls["metadata"] += 1 + + monkeypatch.setattr( + streaming_module, + "async_post_call_success_deployment_hook", + fake_post_call, + ) + monkeypatch.setattr( + streaming_module, + "update_response_metadata", + fake_update_metadata, + ) + + logging_obj = _FakeLoggingObj() + + iterator = ResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=SimpleNamespace(), # not used in this test + logging_obj=logging_obj, + request_data={"foo": "bar", "litellm_params": {}}, + call_type=CallTypes.responses.value, + ) + + # Simulate completed streaming event + iterator.completed_response = SimpleNamespace( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=SimpleNamespace() + ) + + iterator._handle_logging_completed_response() + await asyncio.sleep(0.2) # allow async tasks to run + + assert logging_obj.success_calls == 1 + assert logging_obj.async_success_calls == 1 + assert hook_calls["post_call"] == 1 + assert hook_calls["metadata"] == 1 + assert seen["request_data"]["foo"] == "bar" + assert seen["request_data"].get("litellm_params") is not None + assert seen["call_type"] == CallTypes.responses + + +@pytest.mark.asyncio +async def test_responses_streaming_calls_post_streaming_deployment_hook(monkeypatch): + """ + Ensure per-chunk streaming deployment hook can modify chunks. + """ + + class _HookLogger(CustomLogger): + async def async_post_call_streaming_deployment_hook( + self, request_data, response_chunk, call_type + ): + response_chunk.tagged = True + return response_chunk + + # Set callbacks to our fake hook + original_callbacks = litellm.callbacks + litellm.callbacks = [_HookLogger()] + + logging_obj = _FakeLoggingObj() + + class _StubConfig: + def transform_streaming_response(self, **kwargs): + return SimpleNamespace( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, response=None + ) + + iterator = ResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=_StubConfig(), + logging_obj=logging_obj, + request_data={"foo": "bar"}, + call_type=CallTypes.responses.value, + ) + + # Call hook helper directly to verify chunk is modified/flagged + chunk = SimpleNamespace(type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, response=None) + chunk = await streaming_module.call_post_streaming_hooks_for_testing(iterator, chunk) + assert getattr(chunk, "_post_streaming_hooks_ran", False) is True + assert getattr(chunk, "tagged", False) is True + + # reset callbacks + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_responses_streaming_failure_triggers_failure_handlers(): + """ + If transform raises, failure handlers should be called. + """ + + class _FailConfig: + def transform_streaming_response(self, **kwargs): + raise ValueError("boom") + + logging_obj = _FakeLoggingObj() + + iterator = ResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=_FailConfig(), + logging_obj=logging_obj, + request_data={"foo": "bar"}, + call_type=CallTypes.responses.value, + ) + + with pytest.raises(ValueError): + iterator._process_chunk('{"delta": "chunk"}') + + # allow failure callbacks to run + await asyncio.sleep(0.2) + assert logging_obj.failure_calls >= 1 + assert logging_obj.async_failure_calls >= 1