mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
fix(responses): unwrap completed response for cost tracking on bridge path
When a non-streaming caller drains a Responses API streaming iterator (e.g. the Anthropic /v1/messages -> Responses API bridge against a force-streaming provider like chatgpt/*), _log_completed_response passed the ResponseCompletedEvent wrapper to dispatch_success_handlers. The success handlers only unwrap it in their assembled-stream branch, which a non-streaming caller never reaches — so standard_logging_object was never built, and _PROXY_track_cost_callback raised "Cost tracking failed: standard_logging_object not found". Fix: add _response_for_success_logging() that copies and unwraps .completed_response to the inner ResponsesAPIResponse before passing it to the success handlers. This lets _is_recognized_call_type_for_logging match in the non-streaming branch so standard_logging_object is built. Related: #36426, #36450
This commit is contained in:
parent
ca0b951a43
commit
6a4e558f61
2 changed files with 83 additions and 7 deletions
|
|
@ -412,13 +412,7 @@ class BaseResponsesAPIStreamingIterator:
|
|||
# to chat completion format (prompt_tokens/completion_tokens) for internal logging
|
||||
# Use model_dump + model_validate instead of deepcopy to avoid pickle errors with
|
||||
# Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192)
|
||||
logging_response = self.completed_response
|
||||
if self.completed_response is not None and hasattr(self.completed_response, "model_dump"):
|
||||
try:
|
||||
logging_response = type(self.completed_response).model_validate(self.completed_response.model_dump())
|
||||
except Exception:
|
||||
# Fallback to original if serialization fails
|
||||
pass
|
||||
logging_response = self._response_for_success_logging()
|
||||
|
||||
end_time: Final = datetime.now()
|
||||
if is_async:
|
||||
|
|
@ -448,6 +442,29 @@ class BaseResponsesAPIStreamingIterator:
|
|||
)
|
||||
self._run_post_success_hooks(end_time=end_time)
|
||||
|
||||
def _response_for_success_logging(
|
||||
self,
|
||||
) -> ResponsesAPIStreamingResponse | ResponsesAPIResponse | None:
|
||||
"""Build the response object to pass to success handlers.
|
||||
|
||||
``self.completed_response`` is a ``ResponseCompletedEvent`` wrapper.
|
||||
The success handlers only unwrap it in their assembled-stream branch,
|
||||
which a non-streaming caller draining this iterator never reaches.
|
||||
Unwrap it here so the inner ``ResponsesAPIResponse`` reaches the
|
||||
handlers and ``standard_logging_object`` is built for cost tracking.
|
||||
"""
|
||||
completed = self.completed_response
|
||||
if completed is None or not hasattr(completed, "model_dump"):
|
||||
return completed
|
||||
try:
|
||||
copied = type(completed).model_validate(completed.model_dump())
|
||||
except Exception:
|
||||
return completed
|
||||
unwrapped = getattr(copied, "response", None)
|
||||
if isinstance(unwrapped, ResponsesAPIResponse):
|
||||
return unwrapped
|
||||
return copied
|
||||
|
||||
def _handle_logging_completed_response(self):
|
||||
"""Base implementation - should be overridden by subclasses"""
|
||||
|
||||
|
|
|
|||
|
|
@ -305,3 +305,62 @@ def test_stream_cache_write_completes_when_asyncio_run_closes_the_loop(monkeypat
|
|||
asyncio.run(_short_lived_script())
|
||||
|
||||
assert len(writes) == 1
|
||||
async def test_completed_response_unwrapped_for_success_logging():
|
||||
"""The iterator must pass the unwrapped ResponsesAPIResponse (not the
|
||||
ResponseCompletedEvent wrapper) to dispatch_success_handlers so that
|
||||
_success_handler_helper_fn recognises it and builds
|
||||
standard_logging_object for cost tracking.
|
||||
|
||||
Without this fix, non-streaming callers that drain a force-streaming
|
||||
provider (e.g. the Anthropic /v1/messages -> Responses API bridge) get
|
||||
``standard_logging_object not found`` because the wrapper event is not
|
||||
recognised by _is_recognized_call_type_for_logging in the non-streaming
|
||||
branch of _success_handler_helper_fn.
|
||||
"""
|
||||
logging_obj = _logging_obj_stub()
|
||||
dispatched_args: list = []
|
||||
|
||||
async def _dispatch(result, **kwargs):
|
||||
dispatched_args.append(result)
|
||||
|
||||
logging_obj.dispatch_success_handlers.side_effect = _dispatch
|
||||
|
||||
iterator = _make_iterator(
|
||||
sse_events=_COMPLETE_STREAM_EVENTS,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
async for _ in iterator:
|
||||
pass
|
||||
|
||||
assert len(dispatched_args) == 1, (
|
||||
f"Expected dispatch_success_handlers called once; got {len(dispatched_args)}"
|
||||
)
|
||||
logged = dispatched_args[0]
|
||||
assert isinstance(logged, ResponsesAPIResponse), (
|
||||
f"Expected ResponsesAPIResponse unwrapped for logging; got {type(logged).__name__}"
|
||||
)
|
||||
|
||||
|
||||
def test_sync_completed_response_unwrapped_for_success_logging():
|
||||
"""Sync counterpart of the unwrapping test."""
|
||||
logging_obj = _logging_obj_stub()
|
||||
success_args: list = []
|
||||
|
||||
logging_obj.async_success_handler.side_effect = lambda result, **kw: success_args.append(result)
|
||||
|
||||
iterator = _make_sync_iterator(
|
||||
sse_events=_COMPLETE_STREAM_EVENTS,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
for _ in iterator:
|
||||
pass
|
||||
|
||||
assert len(success_args) == 1, (
|
||||
f"Expected async_success_handler called once; got {len(success_args)}"
|
||||
)
|
||||
logged = success_args[0]
|
||||
assert isinstance(logged, ResponsesAPIResponse), (
|
||||
f"Expected ResponsesAPIResponse unwrapped for logging; got {type(logged).__name__}"
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue