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

Align responses API streaming hooks with chat pipeline

Browse Source
This commit is contained in:
LingXuanYin 2025-12-29 14:08:30 +08:00
parent 4cdcdc98e8
commit 499fc9f48a
3 changed files with 393 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
litellm/llms/custom_httpx/llm_http_handler.py
View File

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

217
litellm/responses/streaming_iterator.py
View File

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

165
tests/llm_responses_api_testing/test_responses_hooks.py Normal file
View File

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