mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-28 06:23:07 +00:00
refactor(chat): log bridged chat streams on the Responses API as chat rows
The chat stream that wraps the native Responses stream now logs the request, so the spend row carries the chat completion the client received instead of the Responses object. The Responses iterator hands its completed response to the bridge instead of dispatching success logging itself.
This commit is contained in:
parent
e3a4ec7387
commit
8edddf5506
3 changed files with 146 additions and 1 deletions
|
|
@ -70,6 +70,18 @@ def _executor_logging_obj(
|
|||
return native_logging_obj(logging_obj, request_data, request_data)
|
||||
|
||||
|
||||
def _logged_by_client_stream(_completed: object) -> None:
|
||||
return None
|
||||
|
||||
|
||||
def _hand_completion_to_client_stream(result: object, executor_logging: "LiteLLMLoggingObj | None") -> None:
|
||||
"""The chat stream wrapping the executor's Responses stream logs the request once, in the shape the client got."""
|
||||
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
|
||||
|
||||
if executor_logging is not None and isinstance(result, BaseResponsesAPIStreamingIterator):
|
||||
result.hand_completion_to(_logged_by_client_stream)
|
||||
|
||||
|
||||
def _retagged_as_responses(
|
||||
request_data: Mapping[str, object], logging_obj: "LiteLLMLoggingObj"
|
||||
) -> Mapping[str, object]:
|
||||
|
|
@ -303,6 +315,7 @@ class ResponsesToCompletionBridgeHandler:
|
|||
return self._apply_post_stream_processing(result, model, custom_llm_provider)
|
||||
from .transformation import is_structured_output_request
|
||||
|
||||
_hand_completion_to_client_stream(result, executor_logging)
|
||||
completion_stream: Final = self.transformation_handler.get_model_response_iterator(
|
||||
streaming_response=result,
|
||||
sync_stream=True,
|
||||
|
|
@ -411,6 +424,7 @@ class ResponsesToCompletionBridgeHandler:
|
|||
return self._apply_post_stream_processing(result, model, custom_llm_provider)
|
||||
from .transformation import is_structured_output_request
|
||||
|
||||
_hand_completion_to_client_stream(result, executor_logging)
|
||||
completion_stream: Final = self.transformation_handler.get_model_response_iterator(
|
||||
streaming_response=result,
|
||||
sync_stream=False,
|
||||
|
|
|
|||
|
|
@ -363,6 +363,7 @@ class BaseResponsesAPIStreamingIterator:
|
|||
"""
|
||||
|
||||
_pending_logging_tasks: tuple[asyncio.Task[object], ...] = ()
|
||||
_on_complete: Callable[[object], None] | None = None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -416,6 +417,10 @@ class BaseResponsesAPIStreamingIterator:
|
|||
) # GUARANTEE OPENAI HEADERS IN RESPONSE
|
||||
self._raw_response_headers: Mapping[str, str] = MappingProxyType(dict(self.response.headers or {}))
|
||||
|
||||
def hand_completion_to(self, on_complete: Callable[[object], None]) -> None:
|
||||
"""A bridge whose own client stream logs the request takes the completed response instead of the logger."""
|
||||
self._on_complete = on_complete
|
||||
|
||||
def _check_max_streaming_duration(self) -> None:
|
||||
"""Raise litellm.Timeout if the stream has exceeded LITELLM_MAX_STREAMING_DURATION_SECONDS."""
|
||||
if LITELLM_MAX_STREAMING_DURATION_SECONDS is None:
|
||||
|
|
@ -615,7 +620,9 @@ class BaseResponsesAPIStreamingIterator:
|
|||
self._restore_provider_response_headers(logging_response)
|
||||
|
||||
end_time: Final = datetime.now()
|
||||
if is_async:
|
||||
if self._on_complete is not None:
|
||||
self._on_complete(logging_response)
|
||||
elif is_async:
|
||||
logging_coroutine: Final = self.logging_obj.dispatch_success_handlers(
|
||||
logging_response,
|
||||
start_time=self.start_time,
|
||||
|
|
|
|||
|
|
@ -1,9 +1,16 @@
|
|||
import asyncio
|
||||
import json
|
||||
import math
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
from openai.types.chat import ChatCompletion
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.caching.caching import Cache
|
||||
|
|
@ -428,3 +435,120 @@ def test_completion_keeps_responses_surface_for_what_only_it_serves(bridge_model
|
|||
assert result is cached
|
||||
assert responses_call.call_count == 1
|
||||
assert executor_call.call_count == 0
|
||||
|
||||
|
||||
_LOGGED_ROW: Final = TypeAdapter(dict[str, object])
|
||||
_LOGGED_CHAT: Final = TypeAdapter(ChatCompletion)
|
||||
|
||||
|
||||
class _SuccessRowRecorder(CustomLogger):
|
||||
rows: tuple[dict[str, object], ...] = ()
|
||||
|
||||
async def async_log_success_event(
|
||||
self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
self.rows = (*self.rows, _LOGGED_ROW.validate_python(kwargs["standard_logging_object"]))
|
||||
|
||||
|
||||
def _responses_sse(events: Sequence[Mapping[str, object]]) -> bytes:
|
||||
return "".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events).encode()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bridged_chat_stream_logs_one_chat_row_billed_from_provider_usage(monkeypatch: pytest.MonkeyPatch):
|
||||
recorder: Final = _SuccessRowRecorder()
|
||||
monkeypatch.setattr(litellm, "callbacks", [recorder])
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
prices: Final = {"input_cost_per_token": 1e-6, "output_cost_per_token": 4e-6, "cache_read_input_token_cost": 1e-7}
|
||||
litellm.register_model(
|
||||
{
|
||||
"openai/bridge-stream-model": {
|
||||
"litellm_provider": "openai",
|
||||
"mode": "responses",
|
||||
"max_tokens": 100,
|
||||
**prices,
|
||||
}
|
||||
}
|
||||
)
|
||||
usage: Final[dict[str, object]] = {
|
||||
"input_tokens": 100,
|
||||
"input_tokens_details": {"cached_tokens": 40},
|
||||
"output_tokens": 20,
|
||||
"output_tokens_details": {"reasoning_tokens": 5},
|
||||
"total_tokens": 120,
|
||||
}
|
||||
message: Final[dict[str, object]] = {
|
||||
"type": "message",
|
||||
"id": "msg_stream",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "streamed", "annotations": []}],
|
||||
}
|
||||
response: Final[dict[str, object]] = {
|
||||
"id": "resp_stream",
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"model": "bridge-stream-model",
|
||||
"output": [],
|
||||
"parallel_tool_calls": True,
|
||||
"tool_choice": "auto",
|
||||
"tools": [],
|
||||
}
|
||||
sse: Final = _responses_sse(
|
||||
[
|
||||
{"type": "response.created", "sequence_number": 0, "response": {**response, "status": "in_progress"}},
|
||||
{"type": "response.output_item.added", "sequence_number": 1, "output_index": 0, "item": message},
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"sequence_number": 2,
|
||||
"item_id": "msg_stream",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": "streamed",
|
||||
},
|
||||
{"type": "response.output_item.done", "sequence_number": 3, "output_index": 0, "item": message},
|
||||
{
|
||||
"type": "response.completed",
|
||||
"sequence_number": 4,
|
||||
"response": {**response, "status": "completed", "output": [message], "usage": usage},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
with respx.mock(assert_all_called=False) as upstream:
|
||||
responses_route = upstream.post("https://example.invalid/v1/responses").mock(
|
||||
return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, content=sse)
|
||||
)
|
||||
stream = await litellm.acompletion(
|
||||
model="openai/bridge-stream-model",
|
||||
api_base="https://example.invalid/v1",
|
||||
api_key="x",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=True,
|
||||
)
|
||||
assert isinstance(stream, CustomStreamWrapper)
|
||||
text = "".join([chunk.choices[0].delta.content or "" async for chunk in stream if chunk.choices])
|
||||
for _ in range(100):
|
||||
if recorder.rows:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
assert text == "streamed"
|
||||
assert responses_route.call_count == 1
|
||||
assert len(recorder.rows) == 1
|
||||
row: Final = recorder.rows[0]
|
||||
logged: Final = _LOGGED_CHAT.validate_python(row["response"])
|
||||
assert logged.object == "chat.completion"
|
||||
assert logged.choices[0].message.content == "streamed"
|
||||
assert row["prompt_tokens"] == 100
|
||||
assert row["completion_tokens"] == 20
|
||||
cost: Final = row["response_cost"]
|
||||
assert isinstance(cost, float)
|
||||
assert math.isclose(
|
||||
cost,
|
||||
60 * prices["input_cost_per_token"]
|
||||
+ 40 * prices["cache_read_input_token_cost"]
|
||||
+ 20 * prices["output_cost_per_token"],
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue