mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
Merge pull request #41337 from BerriAI/litellm_fix_responses_stream_absent_usage_recount
fix(responses): recount tokens when a streamed response completes without usage
This commit is contained in:
kerry-berri
GitHub
commit
5960881640
2 changed files with 437 additions and 11 deletions
|
|
@ -32,6 +32,9 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
|
|||
)
|
||||
from litellm.litellm_core_utils.thread_pool_executor import executor
|
||||
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
|
||||
from litellm.responses.litellm_completion_transformation.transformation import (
|
||||
LiteLLMCompletionResponsesConfig,
|
||||
)
|
||||
from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils
|
||||
from litellm.types.integrations.custom_logger import converted_stream_requested
|
||||
from litellm.types.llms.openai import (
|
||||
|
|
@ -257,6 +260,7 @@ class BaseResponsesAPIStreamingIterator:
|
|||
self._failure_handled = False # Track if failure handler has been called
|
||||
self._yielded_first_chunk = False
|
||||
self._generated_content = ""
|
||||
self._generated_tool_arguments = ""
|
||||
self._completed_response_cached = False
|
||||
self._completed_response_logged = False
|
||||
self._completed_response_cache_hit: bool | None = None
|
||||
|
|
@ -352,6 +356,10 @@ class BaseResponsesAPIStreamingIterator:
|
|||
_delta: Final = getattr(openai_responses_api_chunk, "delta", None)
|
||||
if isinstance(_delta, str):
|
||||
self._generated_content += _delta
|
||||
elif _event_type in _TOOL_ARGUMENTS_DELTA_EVENTS:
|
||||
_args_delta: Final = getattr(openai_responses_api_chunk, "delta", None)
|
||||
if isinstance(_args_delta, str):
|
||||
self._generated_tool_arguments += _args_delta
|
||||
_stream_model_id: Final = _model_id_from_metadata(self.litellm_metadata)
|
||||
if _event_type in (
|
||||
ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
|
||||
|
|
@ -419,14 +427,41 @@ class BaseResponsesAPIStreamingIterator:
|
|||
openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
|
||||
openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED,
|
||||
):
|
||||
self.completed_response = openai_responses_api_chunk
|
||||
_stamp_responses_usage_cost(getattr(openai_responses_api_chunk, "response", None), self.logging_obj)
|
||||
_response_obj: Final[object] = getattr(openai_responses_api_chunk, "response", None)
|
||||
_estimate_wanted: Final[bool] = _chunk_type in (
|
||||
openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
||||
openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
|
||||
)
|
||||
_billed_response: Final[ResponsesAPIResponse | None] = _billed_terminal_response(
|
||||
_response_obj,
|
||||
(
|
||||
lambda: (
|
||||
_estimate_usage_safely(
|
||||
self.model or "",
|
||||
self.request_data.get("input"),
|
||||
self.request_data,
|
||||
self._generated_content + self._generated_tool_arguments,
|
||||
)
|
||||
if _estimate_wanted
|
||||
else None
|
||||
)
|
||||
),
|
||||
)
|
||||
_terminal_chunk: Final = (
|
||||
openai_responses_api_chunk
|
||||
if _billed_response is None or _billed_response is _response_obj
|
||||
else openai_responses_api_chunk.model_copy(update={"response": _billed_response})
|
||||
)
|
||||
self.completed_response = _terminal_chunk
|
||||
_stamp_responses_usage_cost(_billed_response, self.logging_obj)
|
||||
|
||||
if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED:
|
||||
self._handle_logging_failed_response()
|
||||
else:
|
||||
self._handle_logging_completed_response()
|
||||
|
||||
return _terminal_chunk
|
||||
|
||||
return openai_responses_api_chunk
|
||||
|
||||
return None
|
||||
|
|
@ -655,7 +690,9 @@ class BaseResponsesAPIStreamingIterator:
|
|||
if cache is None:
|
||||
return
|
||||
|
||||
cached_response: Final = response_obj.model_dump_json()
|
||||
cached_response: Final = _dump_json_safely(response_obj)
|
||||
if cached_response is None:
|
||||
return
|
||||
if is_async:
|
||||
from litellm.caching.caching_handler import create_cache_write_task
|
||||
|
||||
|
|
@ -1301,6 +1338,31 @@ def _add_text_like_part_events(
|
|||
)
|
||||
|
||||
|
||||
def _billed_terminal_response(
|
||||
response_obj: object, estimate: Callable[[], ResponseAPIUsage | None] | None
|
||||
) -> ResponsesAPIResponse | None:
|
||||
if isinstance(response_obj, ResponsesAPIResponse):
|
||||
return (
|
||||
response_obj
|
||||
if response_obj.usage is not None or estimate is None
|
||||
else response_obj.model_copy(update={"usage": estimate()})
|
||||
)
|
||||
if not isinstance(response_obj, dict):
|
||||
return None
|
||||
usage: Final[object] = response_obj.get("usage") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # a model_constructed terminal event leaves response as an untyped dict
|
||||
return ResponsesAPIResponse.model_construct(
|
||||
**{**response_obj, "usage": usage if usage is not None or estimate is None else estimate()} # pyright: ignore[reportUnknownArgumentType, reportArgumentType] # same untyped dict spread
|
||||
)
|
||||
|
||||
|
||||
def _dump_json_safely(response: BaseModel) -> str | None:
|
||||
try:
|
||||
return response.model_dump_json()
|
||||
except Exception as exc:
|
||||
verbose_logger.debug("could not serialize completed response for cache: %s", exc)
|
||||
return None
|
||||
|
||||
|
||||
def _logging_copy(event: object) -> object:
|
||||
"""Hand logging callbacks a copy, so their usage rewrite (Responses shape to chat shape) never
|
||||
reaches the event the caller is iterating. The round trip through ``model_dump`` sidesteps the
|
||||
|
|
@ -1332,6 +1394,56 @@ def _usage_as_model(usage: object) -> ResponseAPIUsage | None:
|
|||
return None
|
||||
|
||||
|
||||
_TOOL_ARGUMENTS_DELTA_EVENTS: Final = frozenset(
|
||||
{
|
||||
ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA,
|
||||
ResponsesAPIStreamEvents.CUSTOM_TOOL_CALL_INPUT_DELTA,
|
||||
ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DELTA,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _estimate_usage_from_text(
|
||||
model: str,
|
||||
request_input: object,
|
||||
responses_api_request: Mapping[str, object],
|
||||
generated_text: str,
|
||||
) -> ResponseAPIUsage:
|
||||
messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # the transformer's signature is partially untyped
|
||||
input=request_input, # pyright: ignore[reportArgumentType] # the raw Responses API input is a str or ResponseInputParam list, matching the helper's declared union
|
||||
responses_api_request=dict(responses_api_request),
|
||||
)
|
||||
input_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped
|
||||
model=model, messages=messages
|
||||
)
|
||||
output_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped
|
||||
model=model, text=generated_text, count_response_tokens=True
|
||||
)
|
||||
return ResponseAPIUsage(
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
total_tokens=input_tokens + output_tokens,
|
||||
)
|
||||
|
||||
|
||||
def _estimate_usage_safely(
|
||||
model: str,
|
||||
request_input: object,
|
||||
responses_api_request: Mapping[str, object],
|
||||
generated_text: str,
|
||||
) -> ResponseAPIUsage | None:
|
||||
try:
|
||||
return _estimate_usage_from_text(
|
||||
model=model,
|
||||
request_input=request_input,
|
||||
responses_api_request=responses_api_request,
|
||||
generated_text=generated_text,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.debug("Could not estimate usage from stream text, billing $0: %s", e)
|
||||
return None
|
||||
|
||||
|
||||
def _stamp_responses_usage_cost(
|
||||
response_obj: ResponsesAPIResponse | None, logging_obj: LiteLLMLoggingObj | None
|
||||
) -> None:
|
||||
|
|
|
|||
|
|
@ -5,17 +5,20 @@ completion_start_time = end_time."""
|
|||
|
||||
import json
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
from typing import Final, Optional
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from pydantic_core import PydanticSerializationError
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
|
||||
from litellm.responses.streaming_iterator import (
|
||||
ResponsesAPIStreamingIterator,
|
||||
SyncResponsesAPIStreamingIterator,
|
||||
_estimate_usage_from_text,
|
||||
)
|
||||
from litellm.types.llms.openai import (
|
||||
ResponseAPIUsage,
|
||||
|
|
@ -31,16 +34,23 @@ def _sse_event(payload: dict) -> bytes:
|
|||
|
||||
def _mock_config() -> Mock:
|
||||
mock_config = Mock(spec=BaseResponsesAPIConfig)
|
||||
mock_responses_api_response = Mock(spec=ResponsesAPIResponse)
|
||||
mock_responses_api_response.id = "resp_ttft"
|
||||
mock_responses_api_response = ResponsesAPIResponse(
|
||||
id="resp_ttft",
|
||||
created_at=0,
|
||||
status="completed",
|
||||
model="gpt-4o-mini",
|
||||
object="response",
|
||||
output=[],
|
||||
usage=ResponseAPIUsage(input_tokens=1, output_tokens=1, total_tokens=2),
|
||||
)
|
||||
|
||||
def _transform(model, parsed_chunk, logging_obj):
|
||||
evt_type = parsed_chunk.get("type")
|
||||
if evt_type == "response.completed":
|
||||
completed = Mock(spec=ResponseCompletedEvent)
|
||||
completed.type = ResponsesAPIStreamEvents.RESPONSE_COMPLETED
|
||||
completed.response = mock_responses_api_response
|
||||
return completed
|
||||
return ResponseCompletedEvent(
|
||||
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
||||
response=mock_responses_api_response,
|
||||
)
|
||||
stub = Mock()
|
||||
stub.type = evt_type
|
||||
return stub
|
||||
|
|
@ -54,6 +64,8 @@ def _make_iterator(
|
|||
sse_events: list[bytes],
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
trailing_error: Optional[Exception] = None,
|
||||
config: Mock | None = None,
|
||||
request_data: dict | None = None,
|
||||
) -> ResponsesAPIStreamingIterator:
|
||||
async def aiter_bytes():
|
||||
for evt in sse_events:
|
||||
|
|
@ -68,10 +80,11 @@ def _make_iterator(
|
|||
return ResponsesAPIStreamingIterator(
|
||||
response=mock_response,
|
||||
model="gpt-4o-mini",
|
||||
responses_api_provider_config=_mock_config(),
|
||||
responses_api_provider_config=config or _mock_config(),
|
||||
logging_obj=logging_obj,
|
||||
litellm_metadata={},
|
||||
custom_llm_provider="openai",
|
||||
request_data=request_data,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -329,6 +342,88 @@ def test_run_post_success_hooks_does_not_report_generation_time_as_overhead():
|
|||
assert "litellm_overhead_time_ms" not in iterator.completed_response._hidden_params
|
||||
|
||||
|
||||
def _mock_config_with_completed_response(response: ResponsesAPIResponse) -> Mock:
|
||||
mock_config = Mock(spec=BaseResponsesAPIConfig)
|
||||
|
||||
def _transform(model, parsed_chunk, logging_obj):
|
||||
evt_type = parsed_chunk.get("type")
|
||||
if evt_type == "response.completed":
|
||||
return ResponseCompletedEvent(
|
||||
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
||||
response=response,
|
||||
)
|
||||
stub = Mock()
|
||||
stub.type = evt_type
|
||||
if "delta" in parsed_chunk:
|
||||
stub.delta = parsed_chunk.get("delta")
|
||||
if "item" in parsed_chunk:
|
||||
stub.item = parsed_chunk.get("item")
|
||||
return stub
|
||||
|
||||
mock_config.transform_streaming_response.side_effect = _transform
|
||||
return mock_config
|
||||
|
||||
|
||||
def _responses_api_response_without_usage() -> ResponsesAPIResponse:
|
||||
return ResponsesAPIResponse(
|
||||
id="resp_no_usage",
|
||||
created_at=int(datetime(2025, 1, 1).timestamp()),
|
||||
status="completed",
|
||||
model="gpt-4o-mini",
|
||||
object="response",
|
||||
output=[],
|
||||
usage=None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_completed_event_without_usage_gets_text_estimate():
|
||||
"""A response.completed event carrying usage: null still bills: the
|
||||
iterator estimates usage from the request input and generated text."""
|
||||
response = _responses_api_response_without_usage()
|
||||
iterator = _make_iterator(
|
||||
sse_events=[
|
||||
_sse_event({"type": "response.output_text.delta", "delta": "hello world"}),
|
||||
_sse_event({"type": "response.completed", "response": {}}),
|
||||
],
|
||||
logging_obj=_logging_obj_stub(),
|
||||
config=_mock_config_with_completed_response(response),
|
||||
request_data={"input": "count these input tokens please"},
|
||||
)
|
||||
|
||||
async for _ in iterator:
|
||||
pass
|
||||
|
||||
usage = iterator.completed_response.response.usage
|
||||
assert usage is not None
|
||||
assert usage.input_tokens > 0
|
||||
assert usage.output_tokens > 0
|
||||
assert usage.total_tokens == usage.input_tokens + usage.output_tokens
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_completed_event_with_usage_is_left_untouched():
|
||||
"""Provider-reported usage on response.completed wins over the estimate."""
|
||||
response = _responses_api_response_with_usage()
|
||||
iterator = _make_iterator(
|
||||
sse_events=[
|
||||
_sse_event({"type": "response.output_text.delta", "delta": "hello world"}),
|
||||
_sse_event({"type": "response.completed", "response": {}}),
|
||||
],
|
||||
logging_obj=_logging_obj_stub(),
|
||||
config=_mock_config_with_completed_response(response),
|
||||
request_data={"input": "count these input tokens please"},
|
||||
)
|
||||
|
||||
async for _ in iterator:
|
||||
pass
|
||||
|
||||
usage = iterator.completed_response.response.usage
|
||||
assert usage.input_tokens == 20
|
||||
assert usage.output_tokens == 60
|
||||
assert usage.total_tokens == 80
|
||||
|
||||
|
||||
def _responses_api_response_with_usage() -> ResponsesAPIResponse:
|
||||
return ResponsesAPIResponse(
|
||||
id="resp_lit6427",
|
||||
|
|
@ -628,3 +723,222 @@ async def test_streaming_logging_copy_keeps_client_usage_when_response_fails_val
|
|||
assert isinstance(client_usage, ResponseAPIUsage)
|
||||
assert client_usage.input_tokens == 29
|
||||
assert client_usage.cost == pytest.approx(0.0001)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_completed_event_without_usage_counts_tool_call_arguments():
|
||||
"""A function-call-only stream still bills output tokens: streamed
|
||||
function_call_arguments deltas feed the text estimate."""
|
||||
response = _responses_api_response_without_usage()
|
||||
iterator = _make_iterator(
|
||||
sse_events=[
|
||||
_sse_event(
|
||||
{
|
||||
"type": "response.output_item.added",
|
||||
"item": {"type": "function_call", "name": "get_weather", "call_id": "call_1"},
|
||||
}
|
||||
),
|
||||
_sse_event(
|
||||
{
|
||||
"type": "response.function_call_arguments.delta",
|
||||
"delta": '{"location": "San Francisco", "unit": "celsius"}',
|
||||
}
|
||||
),
|
||||
_sse_event({"type": "response.completed", "response": {}}),
|
||||
],
|
||||
logging_obj=_logging_obj_stub(),
|
||||
config=_mock_config_with_completed_response(response),
|
||||
request_data={"input": "what is the weather in san francisco"},
|
||||
)
|
||||
|
||||
async for _ in iterator:
|
||||
pass
|
||||
|
||||
usage = iterator.completed_response.response.usage
|
||||
assert usage is not None
|
||||
assert usage.output_tokens > 0
|
||||
assert usage.total_tokens == usage.input_tokens + usage.output_tokens
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_completed_event_without_usage_counts_multimodal_input_as_messages():
|
||||
"""Multimodal request input is counted as chat messages, not as a JSON blob:
|
||||
a huge base64 image must not inflate the estimated input tokens."""
|
||||
image_input: Final = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "input_text", "text": "what is in this image"},
|
||||
{
|
||||
"type": "input_image",
|
||||
"image_url": "data:image/png;base64," + "A" * 4000,
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
json_count: Final = litellm.token_counter(model="gpt-4o-mini", text=json.dumps(image_input))
|
||||
response = _responses_api_response_without_usage()
|
||||
iterator = _make_iterator(
|
||||
sse_events=[
|
||||
_sse_event({"type": "response.output_text.delta", "delta": "it is a cat"}),
|
||||
_sse_event({"type": "response.completed", "response": {}}),
|
||||
],
|
||||
logging_obj=_logging_obj_stub(),
|
||||
config=_mock_config_with_completed_response(response),
|
||||
request_data={"input": image_input},
|
||||
)
|
||||
|
||||
async for _ in iterator:
|
||||
pass
|
||||
|
||||
usage = iterator.completed_response.response.usage
|
||||
assert usage is not None
|
||||
assert usage.input_tokens < json_count / 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_completed_event_survives_a_failing_usage_estimate():
|
||||
"""A malformed request input that makes the message transformer raise must not
|
||||
break a stream that previously completed: the estimate is best-effort and
|
||||
falls back to usage None."""
|
||||
malformed_input: Final = [{"type": "message", "role": "user", "content": 42}]
|
||||
with pytest.raises(ValueError, match="Invalid content type"):
|
||||
_estimate_usage_from_text("gpt-4o-mini", malformed_input, {"input": malformed_input}, "hello world")
|
||||
|
||||
response = _responses_api_response_without_usage()
|
||||
iterator = _make_iterator(
|
||||
sse_events=[
|
||||
_sse_event({"type": "response.output_text.delta", "delta": "hello world"}),
|
||||
_sse_event({"type": "response.completed", "response": {}}),
|
||||
],
|
||||
logging_obj=_logging_obj_stub(),
|
||||
config=_mock_config_with_completed_response(response),
|
||||
request_data={"input": malformed_input},
|
||||
)
|
||||
|
||||
yielded: list = []
|
||||
async for chunk in iterator:
|
||||
yielded.append(chunk)
|
||||
|
||||
assert yielded
|
||||
assert iterator.completed_response.response.usage is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"tool_delta_event_type",
|
||||
["response.custom_tool_call_input.delta", "response.mcp_call_arguments.delta"],
|
||||
)
|
||||
async def test_completed_event_without_usage_counts_tool_input_deltas(tool_delta_event_type):
|
||||
"""Custom-tool and MCP argument deltas feed the streamed usage fallback the
|
||||
same way function_call_arguments deltas do."""
|
||||
response = _responses_api_response_without_usage()
|
||||
iterator = _make_iterator(
|
||||
sse_events=[
|
||||
_sse_event({"type": tool_delta_event_type, "delta": '{"query": "weather in sf"}'}),
|
||||
_sse_event({"type": "response.completed", "response": {}}),
|
||||
],
|
||||
logging_obj=_logging_obj_stub(),
|
||||
config=_mock_config_with_completed_response(response),
|
||||
request_data={"input": "what is the weather in san francisco"},
|
||||
)
|
||||
|
||||
async for _ in iterator:
|
||||
pass
|
||||
|
||||
usage = iterator.completed_response.response.usage
|
||||
assert usage is not None
|
||||
assert usage.output_tokens > 0
|
||||
assert usage.total_tokens == usage.input_tokens + usage.output_tokens
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_completed_event_with_a_dict_response_is_typed_and_billed():
|
||||
"""transform_streaming_response can model_construct a terminal event whose
|
||||
response stays a plain dict; the iterator must type it so the estimated
|
||||
usage reaches the cost stamping path."""
|
||||
dict_response: Final = {
|
||||
"id": "resp_dict",
|
||||
"model": "gpt-4o-mini",
|
||||
"object": "response",
|
||||
"output": [],
|
||||
"usage": None,
|
||||
}
|
||||
|
||||
def _transform(model, parsed_chunk, logging_obj):
|
||||
if parsed_chunk.get("type") == "response.completed":
|
||||
return ResponseCompletedEvent.model_construct(type="response.completed", response=dict_response)
|
||||
stub: Final = Mock()
|
||||
stub.type = parsed_chunk.get("type")
|
||||
if "delta" in parsed_chunk:
|
||||
stub.delta = parsed_chunk.get("delta")
|
||||
return stub
|
||||
|
||||
config: Final = Mock(spec=BaseResponsesAPIConfig)
|
||||
config.transform_streaming_response.side_effect = _transform
|
||||
logging_obj: Final = _logging_obj_stub()
|
||||
logging_obj._response_cost_calculator.return_value = 0.000704
|
||||
iterator: Final = _make_iterator(
|
||||
sse_events=[
|
||||
_sse_event({"type": "response.output_text.delta", "delta": "hello world"}),
|
||||
_sse_event({"type": "response.completed", "response": {}}),
|
||||
],
|
||||
logging_obj=logging_obj,
|
||||
config=config,
|
||||
request_data={"input": "count these input tokens please"},
|
||||
)
|
||||
|
||||
yielded: Final = [chunk async for chunk in iterator]
|
||||
|
||||
terminal_event: Final = iterator.completed_response
|
||||
assert yielded[-1] is terminal_event
|
||||
completed_response: Final = terminal_event.response
|
||||
assert isinstance(completed_response, ResponsesAPIResponse)
|
||||
usage: Final = completed_response.usage
|
||||
assert usage is not None
|
||||
assert usage.input_tokens > 0
|
||||
assert usage.output_tokens > 0
|
||||
assert usage.cost == pytest.approx(0.000704)
|
||||
logging_obj._response_cost_calculator.assert_any_call(result=completed_response)
|
||||
|
||||
|
||||
def test_billed_terminal_response_keeps_a_response_that_already_has_usage():
|
||||
from litellm.responses.streaming_iterator import _billed_terminal_response
|
||||
|
||||
response: Final = _responses_api_response_with_usage()
|
||||
|
||||
assert _billed_terminal_response(response, None) is response
|
||||
|
||||
|
||||
def test_billed_terminal_response_copies_when_estimating_and_leaves_the_original_untouched():
|
||||
from litellm.responses.streaming_iterator import _billed_terminal_response
|
||||
|
||||
response: Final = _responses_api_response_without_usage()
|
||||
estimated: Final = ResponseAPIUsage(input_tokens=3, output_tokens=4, total_tokens=7)
|
||||
|
||||
billed: Final = _billed_terminal_response(response, lambda: estimated)
|
||||
|
||||
assert billed is not response
|
||||
assert billed.usage is estimated
|
||||
assert response.usage is None
|
||||
|
||||
|
||||
def test_persist_completed_response_to_cache_survives_an_unserializable_response(monkeypatch):
|
||||
bad_response: Final = ResponsesAPIResponse.model_construct(id="r", output=[object()], usage=None)
|
||||
with pytest.raises(PydanticSerializationError):
|
||||
bad_response.model_dump_json()
|
||||
|
||||
logging_obj: Final = _logging_obj_stub()
|
||||
caching_handler: Final = Mock()
|
||||
caching_handler.request_kwargs = {"stream": True}
|
||||
logging_obj._llm_caching_handler = caching_handler
|
||||
iterator: Final = _make_iterator(sse_events=[], logging_obj=logging_obj)
|
||||
iterator.completed_response = ResponseCompletedEvent.model_construct(
|
||||
type="response.completed", response=bad_response
|
||||
)
|
||||
cache: Final = Mock()
|
||||
monkeypatch.setattr(litellm, "cache", cache)
|
||||
|
||||
iterator._persist_completed_response_to_cache(is_async=False)
|
||||
|
||||
cache.add_cache.assert_not_called()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue