mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
Merge branch 'litellm_internal_staging' of https://github.com/BerriAI/litellm into litellm_together_structured_outputs
This commit is contained in:
mateo-berri
commit
7aa8efcf47
27 changed files with 1421 additions and 62 deletions
1
.github/workflows/test-unit.yml
vendored
1
.github/workflows/test-unit.yml
vendored
|
|
@ -164,6 +164,7 @@ jobs:
|
|||
tests/test_litellm/proxy/public_endpoints
|
||||
tests/test_litellm/proxy/prompts
|
||||
tests/test_litellm/proxy/rag_endpoints
|
||||
tests/test_litellm/proxy/rerank_endpoints
|
||||
tests/test_litellm/proxy/realtime_endpoints
|
||||
tests/test_litellm/proxy/ui_crud_endpoints
|
||||
tests/test_litellm/proxy/config_resolvers
|
||||
|
|
|
|||
97
litellm/litellm_core_utils/json_fragment_accumulator.py
Normal file
97
litellm/litellm_core_utils/json_fragment_accumulator.py
Normal file
|
|
@ -0,0 +1,97 @@
|
|||
import json
|
||||
from typing import Final, cast # noqa: TID251 # raw_decode returns tuple[Any, int]; no cast-free unpack
|
||||
|
||||
|
||||
class JSONFragmentAccumulator:
|
||||
"""
|
||||
Buffers a JSON value that arrives piecemeal over a stream (SSE data split
|
||||
across TCP packets, one shard per network read, etc) without the O(n^2)
|
||||
cost of repeated `buffer += fragment` string concatenation, and without
|
||||
the O(n^2) cost of re-copying the unconsumed remainder on every peeled
|
||||
value when one payload holds many concatenated JSON values.
|
||||
|
||||
Fragments are appended to a list in O(1). The buffer is only rebuilt into
|
||||
a single string, and only decoded, when a caller asks for a value via
|
||||
`pop_next_value`, and `could_close_json` lets callers skip that rebuild
|
||||
entirely for fragments that plainly cannot close a JSON value yet. Once
|
||||
rebuilt, consumed values are dropped by advancing a cursor rather than
|
||||
slicing a new string, so draining N concatenated values already sitting
|
||||
in the buffer costs O(n) total, not O(n^2).
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._chunks: list[str] = [] # mutable-ok: O(1) append; string concat would copy the buffer each time
|
||||
self._buffer: str = (
|
||||
"" # mutable-ok: lazily materialized join of _chunks, rebuilt only when _chunks is non-empty
|
||||
)
|
||||
self._offset: int = 0 # mutable-ok: cursor past already-consumed values; avoids re-slicing on every pop
|
||||
self._could_close: bool = False # mutable-ok: cached heuristic; rescanning past fragments was itself O(n^2)
|
||||
|
||||
def __bool__(self) -> bool:
|
||||
return bool(self._chunks) or self._offset < len(self._buffer)
|
||||
|
||||
def append(self, fragment: str) -> None:
|
||||
self._chunks.append(fragment) # mutable-ok: see __init__
|
||||
stripped: Final = fragment.rstrip()
|
||||
if stripped:
|
||||
self._could_close = stripped[-1] in ("}", "]") # mutable-ok: see __init__
|
||||
|
||||
def could_close_json(self) -> bool:
|
||||
"""
|
||||
Whether the buffer's logical last non-whitespace byte is "}" or "]",
|
||||
i.e. whether a JSON value could plausibly be complete. Tracked
|
||||
incrementally in `append` rather than rescanned here, so a run of
|
||||
blank keepalive fragments (e.g. from a malformed upstream stream)
|
||||
can't make this, or the join+parse it gates, cost O(n^2).
|
||||
"""
|
||||
return self._could_close
|
||||
|
||||
def _materialize(self) -> None:
|
||||
if not self._chunks:
|
||||
return
|
||||
unconsumed: Final = self._buffer[self._offset :]
|
||||
self._buffer = unconsumed + "".join(self._chunks) # mutable-ok: merge pending fragments, once per append batch
|
||||
self._offset = 0 # mutable-ok: see __init__
|
||||
self._chunks = [] # mutable-ok: see __init__
|
||||
|
||||
def pop_next_value(self) -> tuple[bool, object]:
|
||||
"""
|
||||
Attempt to decode one complete JSON value from the front of the
|
||||
buffer. On success, advances a cursor past that value (keeping any
|
||||
unconsumed tail, e.g. a second concatenated value, in place rather
|
||||
than copying it) and returns (True, value). If the buffer is empty
|
||||
or holds no complete value yet, it is left untouched and this
|
||||
returns (False, None).
|
||||
"""
|
||||
self._materialize()
|
||||
length: Final = len(self._buffer)
|
||||
start = self._offset
|
||||
while start < length and self._buffer[start].isspace():
|
||||
start += 1
|
||||
if start >= length:
|
||||
self._offset = start # mutable-ok: see __init__
|
||||
return False, None
|
||||
decoder: Final = json.JSONDecoder()
|
||||
try:
|
||||
raw_value: Final = decoder.raw_decode(self._buffer, start)
|
||||
except json.JSONDecodeError:
|
||||
return False, None
|
||||
decoded, end_index = cast("tuple[object, int]", raw_value) # cast-ok: raw_decode returns tuple[Any, int]
|
||||
self._offset = end_index # mutable-ok: see __init__
|
||||
if self._offset >= len(self._buffer):
|
||||
self._buffer = "" # mutable-ok: see __init__
|
||||
self._offset = 0 # mutable-ok: see __init__
|
||||
self._could_close = False # mutable-ok: buffer is empty, nothing can close
|
||||
return True, decoded
|
||||
|
||||
def snapshot(self) -> str:
|
||||
self._materialize()
|
||||
return self._buffer[self._offset :]
|
||||
|
||||
def set(self, value: str) -> None:
|
||||
"""Replace the buffer's contents with a single fragment."""
|
||||
self._chunks = [] # mutable-ok: see __init__
|
||||
self._buffer = value # mutable-ok: see __init__
|
||||
self._offset = 0 # mutable-ok: see __init__
|
||||
stripped: Final = value.rstrip()
|
||||
self._could_close = bool(stripped) and stripped[-1] in ("}", "]") # mutable-ok: see __init__
|
||||
|
|
@ -18,6 +18,7 @@ from litellm.anthropic_beta_headers_manager import (
|
|||
)
|
||||
from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
|
||||
from litellm.litellm_core_utils.core_helpers import map_finish_reason
|
||||
from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAccumulator
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
HTTPHandler,
|
||||
|
|
@ -654,7 +655,7 @@ class ModelResponseIterator:
|
|||
|
||||
# For handling partial JSON chunks from fragmentation
|
||||
# See: https://github.com/BerriAI/litellm/issues/17473
|
||||
self.accumulated_json: str = ""
|
||||
self._json_buffer = JSONFragmentAccumulator()
|
||||
self.chunk_type: Literal["valid_json", "accumulated_json"] = "valid_json"
|
||||
|
||||
# Track current content block type to avoid emitting tool calls for non-tool blocks
|
||||
|
|
@ -678,6 +679,14 @@ class ModelResponseIterator:
|
|||
self._current_server_tool_id: str | None = None
|
||||
self._container_id: str | None = None
|
||||
|
||||
@property
|
||||
def accumulated_json(self) -> str:
|
||||
return self._json_buffer.snapshot()
|
||||
|
||||
@accumulated_json.setter
|
||||
def accumulated_json(self, value: str) -> None:
|
||||
self._json_buffer.set(value)
|
||||
|
||||
def check_empty_tool_call_args(self) -> bool:
|
||||
"""
|
||||
Check if the tool call block so far has been an empty string
|
||||
|
|
@ -1149,31 +1158,39 @@ class ModelResponseIterator:
|
|||
container: Final = message_delta["delta"].get("container")
|
||||
return finish_reason, usage, container
|
||||
|
||||
def _handle_accumulated_json_chunk(self, data_str: str) -> ModelResponseStream | None:
|
||||
def _handle_accumulated_json_chunk(self, data_str: str, is_final: bool = False) -> ModelResponseStream | None:
|
||||
"""
|
||||
Handle partial JSON chunks by accumulating them until valid JSON is received.
|
||||
|
||||
This fixes network fragmentation issues where SSE data chunks may be split
|
||||
across TCP packets. See: https://github.com/BerriAI/litellm/issues/17473
|
||||
|
||||
Mid-stream, defer parsing until the buffer's last byte can close a value:
|
||||
attempting a parse after every fragment of one large object is O(n^2) and
|
||||
holds the GIL, freezing the event loop. At end of stream (is_final) no more
|
||||
data is coming, so drain whatever complete values remain regardless of the
|
||||
trailing byte.
|
||||
|
||||
Args:
|
||||
data_str: The JSON string to parse (without "data:" prefix)
|
||||
is_final: True when called from the end-of-stream drain, where the
|
||||
trailing-byte heuristic no longer applies
|
||||
|
||||
Returns:
|
||||
ModelResponseStream if JSON is complete, None if still accumulating
|
||||
"""
|
||||
# Accumulate JSON data
|
||||
self.accumulated_json += data_str
|
||||
self._json_buffer.append(data_str)
|
||||
|
||||
# Try to parse the accumulated JSON
|
||||
try:
|
||||
data_json: Final = json.loads(self.accumulated_json)
|
||||
self.accumulated_json = "" # Reset after successful parsing
|
||||
return self.chunk_parser(chunk=data_json)
|
||||
except json.JSONDecodeError:
|
||||
# If it's not valid JSON yet, continue to the next chunk
|
||||
if not is_final and not self._json_buffer.could_close_json():
|
||||
return None
|
||||
|
||||
while True:
|
||||
found, decoded = self._json_buffer.pop_next_value()
|
||||
if not found:
|
||||
return None
|
||||
if isinstance(decoded, dict):
|
||||
return self.chunk_parser(chunk=decoded)
|
||||
|
||||
def _parse_sse_data(self, str_line: str) -> ModelResponseStream | None:
|
||||
"""
|
||||
Parse SSE data line, handling both complete and partial JSON chunks.
|
||||
|
|
@ -1209,13 +1226,10 @@ class ModelResponseIterator:
|
|||
chunk = self.response_iterator.__next__()
|
||||
except StopIteration:
|
||||
# If we have accumulated JSON when stream ends, try to parse it
|
||||
if self.accumulated_json:
|
||||
try:
|
||||
data_json = json.loads(self.accumulated_json)
|
||||
self.accumulated_json = ""
|
||||
return self.chunk_parser(chunk=data_json)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
if self._json_buffer:
|
||||
result = self._handle_accumulated_json_chunk(data_str="", is_final=True)
|
||||
if result is not None:
|
||||
return result
|
||||
raise StopIteration
|
||||
except ValueError as e:
|
||||
raise RuntimeError(f"Error receiving chunk from stream: {e}")
|
||||
|
|
@ -1258,13 +1272,10 @@ class ModelResponseIterator:
|
|||
chunk = await self.async_response_iterator.__anext__()
|
||||
except StopAsyncIteration:
|
||||
# If we have accumulated JSON when stream ends, try to parse it
|
||||
if self.accumulated_json:
|
||||
try:
|
||||
data_json = json.loads(self.accumulated_json)
|
||||
self.accumulated_json = ""
|
||||
return self.chunk_parser(chunk=data_json)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
if self._json_buffer:
|
||||
result = self._handle_accumulated_json_chunk(data_str="", is_final=True)
|
||||
if result is not None:
|
||||
return result
|
||||
raise StopAsyncIteration
|
||||
except ValueError as e:
|
||||
raise RuntimeError(f"Error receiving chunk from stream: {e}")
|
||||
|
|
|
|||
|
|
@ -214,6 +214,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
|
|||
system text -> message(role=system, input_text)
|
||||
user text -> message(role=user, input_text)
|
||||
user image -> message(role=user, input_image)
|
||||
user document -> message(role=user, input_file)
|
||||
user tool_result -> function_call_output
|
||||
assistant text -> message(role=assistant, output_text)
|
||||
assistant thinking -> reasoning
|
||||
|
|
@ -268,6 +269,12 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
|
|||
{"type": "input_image", "image_url": url}, block.get("prompt_cache_breakpoint")
|
||||
)
|
||||
)
|
||||
elif btype == "document":
|
||||
file_part = self._translate_anthropic_document_block_to_file_part(block)
|
||||
if file_part:
|
||||
user_parts.append(
|
||||
with_prompt_cache_breakpoint(file_part, block.get("prompt_cache_breakpoint"))
|
||||
)
|
||||
elif btype == "tool_result":
|
||||
tool_use_id = block.get("tool_use_id", "")
|
||||
inner = block.get("content")
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ class BedrockRerankHandler(BaseAWSLLM):
|
|||
async def arerank(
|
||||
self,
|
||||
prepared_request: BedrockPreparedRequest,
|
||||
logging_obj: LitellmLogging,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
client: AsyncHTTPHandler | None = None,
|
||||
):
|
||||
|
|
@ -40,6 +41,7 @@ class BedrockRerankHandler(BaseAWSLLM):
|
|||
headers=dict(prepared_request["prepped"].headers),
|
||||
data=prepared_request["body"],
|
||||
timeout=timeout,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as err:
|
||||
|
|
@ -98,6 +100,7 @@ class BedrockRerankHandler(BaseAWSLLM):
|
|||
if _is_async:
|
||||
return self.arerank(
|
||||
prepared_request,
|
||||
logging_obj=logging_obj,
|
||||
timeout=timeout,
|
||||
client=client if client is not None and isinstance(client, AsyncHTTPHandler) else None,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1203,6 +1203,7 @@ class BaseLLMHTTPHandler:
|
|||
headers=headers,
|
||||
data=json.dumps(request_data),
|
||||
timeout=timeout,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
except Exception as e:
|
||||
raise self._handle_error(e=e, provider_config=provider_config)
|
||||
|
|
|
|||
|
|
@ -4,17 +4,24 @@ Translates from OpenAI's `/v1/chat/completions` to Together AI's `/v1/chat/compl
|
|||
Docs: https://docs.together.ai/docs/chat-overview
|
||||
"""
|
||||
|
||||
from collections.abc import Callable, Container
|
||||
from typing import Final
|
||||
from collections.abc import Callable, Container, Coroutine
|
||||
from typing import (
|
||||
Final,
|
||||
Literal,
|
||||
cast, # noqa: TID251 # rebuilding a TypedDict minus keys has no checked spelling
|
||||
overload,
|
||||
)
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.exceptions import UnsupportedParamsError
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.utils import supports_function_calling, supports_response_schema
|
||||
|
||||
from ...openai.chat.gpt_transformation import OpenAIGPTConfig
|
||||
|
||||
TOOL_CALLING_PARAMS: Final = ("tools", "tool_choice", "function_call")
|
||||
LITELLM_INTERNAL_ASSISTANT_FIELDS: Final = frozenset({"thinking_blocks", "provider_specific_fields"})
|
||||
FUNCTION_CALLING_DOCS_URL: Final = "https://docs.together.ai/docs/function-calling"
|
||||
STRUCTURED_OUTPUTS_DOCS_URL: Final = "https://docs.together.ai/docs/inference/chat/structured-outputs"
|
||||
|
||||
|
|
@ -102,7 +109,50 @@ def _drop_response_format(passed_params: Container[str], model: str, drop_params
|
|||
)
|
||||
|
||||
|
||||
def _without_litellm_internal_fields(message: AllMessageValues) -> AllMessageValues:
|
||||
if message["role"] != "assistant" or LITELLM_INTERNAL_ASSISTANT_FIELDS.isdisjoint(message):
|
||||
return message
|
||||
return cast( # cast-ok: rebuilding the same TypedDict minus internal keys loses the narrowed type
|
||||
"AllMessageValues",
|
||||
{ # mutable-ok: TypedDict rebuild minus internal keys
|
||||
key: value for key, value in message.items() if key not in LITELLM_INTERNAL_ASSISTANT_FIELDS
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
class TogetherAIChatConfig(OpenAIGPTConfig):
|
||||
@overload
|
||||
def _transform_messages(
|
||||
self,
|
||||
messages: list[AllMessageValues], # mutable-ok: inherited contract
|
||||
model: str,
|
||||
is_async: Literal[True],
|
||||
) -> Coroutine[object, object, list[AllMessageValues]]: ... # mutable-ok: inherited contract
|
||||
|
||||
@overload
|
||||
def _transform_messages(
|
||||
self,
|
||||
messages: list[AllMessageValues], # mutable-ok: inherited contract
|
||||
model: str,
|
||||
is_async: Literal[False] = False,
|
||||
) -> list[AllMessageValues]: ... # mutable-ok: inherited contract
|
||||
|
||||
def _transform_messages(
|
||||
self,
|
||||
messages: list[AllMessageValues], # mutable-ok: inherited contract
|
||||
model: str,
|
||||
is_async: bool = False,
|
||||
) -> list[AllMessageValues] | Coroutine[object, object, list[AllMessageValues]]: # mutable-ok: inherited contract
|
||||
"""Together consumes replayed assistant `reasoning_content` (preserved thinking via
|
||||
`chat_template_kwargs: {"clear_thinking": false}`), so it must stay in the payload;
|
||||
only litellm-internal fields are stripped before sending."""
|
||||
stripped: Final = [ # mutable-ok: super() requires a list
|
||||
_without_litellm_internal_fields(message) for message in messages
|
||||
]
|
||||
if is_async:
|
||||
return super()._transform_messages(stripped, model, is_async=True)
|
||||
return super()._transform_messages(stripped, model, is_async=False)
|
||||
|
||||
def map_openai_params(
|
||||
self,
|
||||
non_default_params: dict,
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ from litellm.constants import (
|
|||
DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH_LITE,
|
||||
DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_PRO,
|
||||
)
|
||||
from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAccumulator
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
_encode_tool_call_id_with_signature,
|
||||
)
|
||||
|
|
@ -3087,7 +3088,7 @@ class ModelResponseIterator:
|
|||
self.streaming_response = streaming_response
|
||||
self.response = response
|
||||
self.chunk_type: Literal["valid_json", "accumulated_json"] = "valid_json"
|
||||
self.accumulated_json = ""
|
||||
self._json_buffer = JSONFragmentAccumulator()
|
||||
self.sent_first_chunk = False
|
||||
self.logging_obj = logging_obj
|
||||
self.response_headers = response_headers or {}
|
||||
|
|
@ -3095,6 +3096,14 @@ class ModelResponseIterator:
|
|||
self.cumulative_tool_call_index: int = 0
|
||||
self.has_seen_tool_calls: bool = False
|
||||
|
||||
@property
|
||||
def accumulated_json(self) -> str:
|
||||
return self._json_buffer.snapshot()
|
||||
|
||||
@accumulated_json.setter
|
||||
def accumulated_json(self, value: str) -> None:
|
||||
self._json_buffer.set(value)
|
||||
|
||||
@staticmethod
|
||||
def _check_streaming_error(chunk: dict) -> None:
|
||||
"""Detect embedded errors (e.g. 429 RESOURCE_EXHAUSTED) in streaming chunks and raise VertexAIError."""
|
||||
|
|
@ -3298,8 +3307,8 @@ class ModelResponseIterator:
|
|||
return self.chunk_parser(chunk=json_chunk)
|
||||
|
||||
def handle_accumulated_json_chunk(self, chunk: str, is_final: bool = False) -> Optional["ModelResponseStream"]:
|
||||
message: Final = litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or ""
|
||||
self.accumulated_json = (self.accumulated_json + message.replace("\n\n", "")).strip()
|
||||
message: Final = (litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or "").replace("\n\n", "")
|
||||
self._json_buffer.append(message)
|
||||
|
||||
# Mid-stream, defer parsing until the buffer's last byte can close a value:
|
||||
# attempting a parse after every fragment of one large object is O(n^2) and
|
||||
|
|
@ -3307,27 +3316,23 @@ class ModelResponseIterator:
|
|||
# data is coming, so drain whatever complete values remain regardless of the
|
||||
# trailing byte, otherwise a complete leading value sitting behind a truncated
|
||||
# trailing one would be silently dropped.
|
||||
if not is_final and (not self.accumulated_json or self.accumulated_json[-1] not in "}]"):
|
||||
if not is_final and not self._json_buffer.could_close_json():
|
||||
return None
|
||||
|
||||
# Peel one complete JSON value from the front of the buffer and keep the
|
||||
# unconsumed tail. Running json.loads over the whole buffer would fail
|
||||
# forever once it held more than one concatenated value ("Extra data") while
|
||||
# never resetting the buffer, so the buffer grew without bound and pinned the
|
||||
# core. raw_decode reports where the value ended, so concatenated values drain
|
||||
# one call at a time. A leading non-dict value (never emitted by Gemini in
|
||||
# practice) is consumed and skipped so it cannot block the dict values behind it.
|
||||
decoder: Final = json.JSONDecoder()
|
||||
while self.accumulated_json:
|
||||
try:
|
||||
raw_value = decoder.raw_decode(self.accumulated_json)
|
||||
except json.JSONDecodeError:
|
||||
# core. pop_next_value reports where the value ended, so concatenated values
|
||||
# drain one call at a time. A leading non-dict value (never emitted by Gemini
|
||||
# in practice) is consumed and skipped so it cannot block the dict values
|
||||
# behind it.
|
||||
while True:
|
||||
found, decoded = self._json_buffer.pop_next_value()
|
||||
if not found:
|
||||
return None
|
||||
decoded, end_index = cast("tuple[object, int]", raw_value) # cast-ok: raw_decode -> tuple[Any,int]
|
||||
self.accumulated_json = self.accumulated_json[end_index:].strip()
|
||||
if isinstance(decoded, dict):
|
||||
return self.chunk_parser(chunk=decoded)
|
||||
return None
|
||||
|
||||
def _common_chunk_parsing_logic(self, chunk: str) -> Optional["ModelResponseStream"]:
|
||||
try:
|
||||
|
|
@ -3351,7 +3356,7 @@ class ModelResponseIterator:
|
|||
try:
|
||||
chunk: Final = self.response_iterator.__next__()
|
||||
except StopIteration:
|
||||
if self.chunk_type == "accumulated_json" and self.accumulated_json:
|
||||
if self.chunk_type == "accumulated_json" and self._json_buffer:
|
||||
result: Final = self.handle_accumulated_json_chunk(chunk="", is_final=True)
|
||||
if result is not None:
|
||||
return result
|
||||
|
|
@ -3375,7 +3380,7 @@ class ModelResponseIterator:
|
|||
try:
|
||||
chunk: Final = await self.async_response_iterator.__anext__()
|
||||
except StopAsyncIteration:
|
||||
if self.chunk_type == "accumulated_json" and self.accumulated_json:
|
||||
if self.chunk_type == "accumulated_json" and self._json_buffer:
|
||||
result: Final = self.handle_accumulated_json_chunk(chunk="", is_final=True)
|
||||
if result is not None:
|
||||
return result
|
||||
|
|
|
|||
|
|
@ -3546,6 +3546,8 @@ class SpendLogsMetadata(TypedDict):
|
|||
litellm_overhead_time_ms: float | None # LiteLLM overhead time in milliseconds
|
||||
attempted_retries: int | None # Number of retries attempted (0 = first attempt succeeded)
|
||||
max_retries: int | None # Max retries configured for this request
|
||||
attempted_fallbacks: ReadOnly[int | None] # Number of fallbacks attempted (0 = primary model group served)
|
||||
original_model_group: ReadOnly[str | None] # Model group requested before any fallbacks
|
||||
cost_breakdown: CostBreakdown | None # Detailed cost breakdown (input_cost, output_cost, margin, discount, etc.)
|
||||
compression_savings: CompressionSavingsMetadata | None
|
||||
autorouter_savings: ReadOnly[float | None] # stamped by the logging payload; None = not auto-routed
|
||||
|
|
|
|||
|
|
@ -90,12 +90,15 @@ async def rerank(
|
|||
fastapi_response.headers.update(
|
||||
ProxyBaseLLMRequestProcessing.get_custom_headers(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
call_id=hidden_params.get("litellm_call_id", None) or data.get("litellm_call_id", None),
|
||||
model_id=model_id,
|
||||
cache_key=cache_key,
|
||||
api_base=api_base,
|
||||
version=version,
|
||||
response_cost=hidden_params.get("response_cost", None),
|
||||
model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
|
||||
request_data=data,
|
||||
hidden_params=hidden_params,
|
||||
**additional_headers,
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -131,6 +131,8 @@ def _get_spend_logs_metadata(
|
|||
litellm_overhead_time_ms=None,
|
||||
attempted_retries=None,
|
||||
max_retries=None,
|
||||
attempted_fallbacks=None,
|
||||
original_model_group=None,
|
||||
cost_breakdown=None,
|
||||
compression_savings=None,
|
||||
autorouter_savings=autorouter_savings,
|
||||
|
|
|
|||
|
|
@ -141,6 +141,7 @@ from litellm.router_utils.cooldown_handlers import (
|
|||
is_advisor_orchestration_failure,
|
||||
)
|
||||
from litellm.router_utils.fallback_event_handlers import (
|
||||
AttemptedFallbackTargets,
|
||||
_check_non_standard_fallback_format,
|
||||
get_fallback_model_group,
|
||||
run_async_fallback,
|
||||
|
|
@ -6930,6 +6931,20 @@ class Router:
|
|||
If it fails after num_retries, fall back to another model group
|
||||
"""
|
||||
model_group: Final[str | None] = kwargs.get("model")
|
||||
if not isinstance(kwargs.get("attempted_targets"), AttemptedFallbackTargets):
|
||||
_fallback_metadata_key: Final = _get_router_metadata_variable_name(
|
||||
function_name=getattr(kwargs.get("original_function"), "__name__", None)
|
||||
)
|
||||
_sibling_metadata_key: Final = (
|
||||
"metadata" if _fallback_metadata_key == "litellm_metadata" else "litellm_metadata"
|
||||
)
|
||||
if isinstance(_sibling_metadata := kwargs.get(_sibling_metadata_key), dict):
|
||||
_sibling_metadata.pop("attempted_fallbacks", None)
|
||||
_sibling_metadata.pop("original_model_group", None)
|
||||
if isinstance(_fallback_metadata := kwargs.get(_fallback_metadata_key), dict):
|
||||
_fallback_metadata["attempted_fallbacks"] = 0
|
||||
if model_group is not None:
|
||||
_fallback_metadata["original_model_group"] = model_group
|
||||
include_fallback_errors: Final = kwargs.get("include_fallback_errors", False) is True
|
||||
disable_fallbacks: Final[bool | None] = kwargs.pop("disable_fallbacks", False)
|
||||
fallbacks: Final[list | None] = kwargs.get("fallbacks", self.fallbacks)
|
||||
|
|
|
|||
|
|
@ -374,11 +374,13 @@ async def run_async_fallback(
|
|||
kwargs["model"] = mg
|
||||
elif isinstance(mg, dict):
|
||||
kwargs.update(mg)
|
||||
fallback_depth = fallback_depth + 1
|
||||
kwargs[metadata_variable_name] = {
|
||||
"original_model_group": original_model_group,
|
||||
**(kwargs.get(metadata_variable_name) or {}),
|
||||
"model_group": kwargs.get("model", None),
|
||||
"attempted_fallbacks": fallback_depth,
|
||||
}
|
||||
fallback_depth = fallback_depth + 1
|
||||
kwargs["fallback_depth"] = fallback_depth
|
||||
kwargs["max_fallbacks"] = max_fallbacks
|
||||
kwargs["attempted_targets"] = attempted
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@
|
|||
"user": "",
|
||||
"team_id": "",
|
||||
"organization_id": "",
|
||||
"metadata": "{\"applied_guardrails\": [], \"batch_models\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"guardrail_information\": null, \"compression_savings\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}",
|
||||
"metadata": "{\"applied_guardrails\": [], \"attempted_fallbacks\": null, \"original_model_group\": null, \"batch_models\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"guardrail_information\": null, \"compression_savings\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}",
|
||||
"cache_key": "Cache OFF",
|
||||
"spend": 0.00022500000000000002,
|
||||
"total_tokens": 30,
|
||||
|
|
|
|||
|
|
@ -0,0 +1,236 @@
|
|||
import json
|
||||
import time
|
||||
from unittest.mock import patch
|
||||
|
||||
from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAccumulator
|
||||
|
||||
|
||||
def test_initial_state_is_empty():
|
||||
accumulator = JSONFragmentAccumulator()
|
||||
assert not accumulator
|
||||
assert accumulator.could_close_json() is False
|
||||
assert accumulator.snapshot() == ""
|
||||
|
||||
|
||||
def test_could_close_json_true_only_when_last_fragment_closes_a_value():
|
||||
accumulator = JSONFragmentAccumulator()
|
||||
accumulator.append('{"a": ')
|
||||
assert accumulator.could_close_json() is False
|
||||
|
||||
accumulator.append("1}")
|
||||
assert accumulator.could_close_json() is True
|
||||
|
||||
|
||||
def test_could_close_json_looks_past_trailing_blank_fragments():
|
||||
"""A whitespace-only or empty fragment (e.g. the flush call at end of
|
||||
stream) must not mask a real closing byte in an earlier fragment."""
|
||||
accumulator = JSONFragmentAccumulator()
|
||||
accumulator.append('{"a": 1}')
|
||||
accumulator.append("")
|
||||
accumulator.append(" \n")
|
||||
assert accumulator.could_close_json() is True
|
||||
|
||||
|
||||
def test_pop_next_value_on_empty_buffer_returns_false_without_touching_state():
|
||||
accumulator = JSONFragmentAccumulator()
|
||||
found, value = accumulator.pop_next_value()
|
||||
assert found is False
|
||||
assert value is None
|
||||
|
||||
|
||||
def test_pop_next_value_on_incomplete_buffer_leaves_buffer_untouched():
|
||||
accumulator = JSONFragmentAccumulator()
|
||||
accumulator.append('{"candidates": [{"content":')
|
||||
|
||||
found, value = accumulator.pop_next_value()
|
||||
|
||||
assert found is False
|
||||
assert value is None
|
||||
assert accumulator.snapshot() == '{"candidates": [{"content":'
|
||||
|
||||
|
||||
def test_pop_next_value_decodes_single_complete_object_and_clears_buffer():
|
||||
accumulator = JSONFragmentAccumulator()
|
||||
accumulator.append('{"candidates": [{"content": {"parts": [{"text": "hi"}]}}]}')
|
||||
|
||||
found, value = accumulator.pop_next_value()
|
||||
|
||||
assert found is True
|
||||
assert value == {"candidates": [{"content": {"parts": [{"text": "hi"}]}}]}
|
||||
assert accumulator.snapshot() == ""
|
||||
assert not accumulator
|
||||
|
||||
|
||||
def test_pop_next_value_reassembles_a_value_split_across_many_fragments():
|
||||
obj = {"candidates": [{"content": {"parts": [{"text": "x" * 5000}]}}]}
|
||||
blob = json.dumps(obj)
|
||||
fragments = [blob[i : i + 37] for i in range(0, len(blob), 37)]
|
||||
assert len(fragments) > 10, "need a genuinely multi-fragment payload"
|
||||
|
||||
accumulator = JSONFragmentAccumulator()
|
||||
found = False
|
||||
value = None
|
||||
for fragment in fragments:
|
||||
accumulator.append(fragment)
|
||||
if accumulator.could_close_json():
|
||||
found, value = accumulator.pop_next_value()
|
||||
|
||||
assert found is True
|
||||
assert value == obj
|
||||
|
||||
|
||||
def test_pop_next_value_peels_one_value_and_keeps_remainder():
|
||||
"""Two concatenated envelopes in the buffer must both surface, one per
|
||||
call, instead of json.loads's "Extra data" failure wedging the buffer."""
|
||||
obj = '{"a": 1}'
|
||||
accumulator = JSONFragmentAccumulator()
|
||||
accumulator.append(obj + obj)
|
||||
|
||||
first_found, first_value = accumulator.pop_next_value()
|
||||
assert first_found is True
|
||||
assert first_value == {"a": 1}
|
||||
assert accumulator.snapshot() == obj, "second value must remain buffered"
|
||||
|
||||
second_found, second_value = accumulator.pop_next_value()
|
||||
assert second_found is True
|
||||
assert second_value == {"a": 1}
|
||||
assert not accumulator
|
||||
|
||||
|
||||
def test_pop_next_value_skips_non_ascii_whitespace_between_concatenated_values():
|
||||
"""A separator like U+00A0 (non-breaking space) between two concatenated
|
||||
values must not strand the second value forever. `raw_decode` only skips
|
||||
the narrow `json.decoder.WHITESPACE` set, so the accumulator's own
|
||||
whitespace skip must be as tolerant as `str.strip()` was before this
|
||||
class replaced it, not merely match `raw_decode`'s narrower set."""
|
||||
accumulator = JSONFragmentAccumulator()
|
||||
accumulator.append('{"a": 1}' + "\xa0" + '{"a": 2}')
|
||||
|
||||
first_found, first_value = accumulator.pop_next_value()
|
||||
assert first_found is True
|
||||
assert first_value == {"a": 1}
|
||||
|
||||
second_found, second_value = accumulator.pop_next_value()
|
||||
assert second_found is True, "the second value must not be permanently stranded"
|
||||
assert second_value == {"a": 2}
|
||||
assert not accumulator
|
||||
|
||||
|
||||
def test_pop_next_value_advances_past_a_non_dict_leading_value():
|
||||
accumulator = JSONFragmentAccumulator()
|
||||
accumulator.append("[1, 2]" + '{"a": 1}')
|
||||
|
||||
first_found, first_value = accumulator.pop_next_value()
|
||||
assert first_found is True
|
||||
assert first_value == [1, 2]
|
||||
|
||||
second_found, second_value = accumulator.pop_next_value()
|
||||
assert second_found is True
|
||||
assert second_value == {"a": 1}
|
||||
|
||||
|
||||
def test_set_and_snapshot_roundtrip():
|
||||
accumulator = JSONFragmentAccumulator()
|
||||
accumulator.set('{"a": 1}')
|
||||
assert accumulator.snapshot() == '{"a": 1}'
|
||||
assert accumulator
|
||||
|
||||
accumulator.set("")
|
||||
assert accumulator.snapshot() == ""
|
||||
assert not accumulator
|
||||
|
||||
|
||||
def test_append_never_calls_raw_decode(): # test-quality-ok: TQ002 - laziness contract has no caller-observable proxy other than spying on the stdlib call it must defer
|
||||
"""Appending must be O(1) bookkeeping only; the O(n) join+decode is
|
||||
deferred entirely to pop_next_value."""
|
||||
accumulator = JSONFragmentAccumulator()
|
||||
with patch.object(json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode) as spy:
|
||||
for fragment in ['{"a":', " 1", "}"]:
|
||||
accumulator.append(fragment)
|
||||
assert spy.call_count == 0
|
||||
|
||||
|
||||
def test_pop_next_value_calls_raw_decode_at_most_once_per_value():
|
||||
accumulator = JSONFragmentAccumulator()
|
||||
accumulator.append('{"a": 1}' * 3)
|
||||
|
||||
with patch.object(json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode) as spy:
|
||||
for _ in range(3):
|
||||
found, _ = accumulator.pop_next_value()
|
||||
assert found is True
|
||||
assert spy.call_count == 3
|
||||
|
||||
|
||||
def test_accumulation_of_many_fragments_is_not_quadratic():
|
||||
"""Regression guard: appending 1000 shards must stay O(n) total, not the
|
||||
O(n^2) cost of repeated `buffer += fragment` string concatenation."""
|
||||
accumulator = JSONFragmentAccumulator()
|
||||
shard = "x" * 2048
|
||||
|
||||
start = time.perf_counter()
|
||||
for _ in range(1000):
|
||||
accumulator.append(shard)
|
||||
elapsed_ms = (time.perf_counter() - start) * 1000
|
||||
|
||||
assert elapsed_ms < 50, f"1000-fragment append took {elapsed_ms:.1f} ms (expected < 50 ms); O(n^2) regression?"
|
||||
|
||||
|
||||
def test_draining_many_concatenated_values_is_not_quadratic():
|
||||
"""
|
||||
Regression guard: peeling N JSON values already sitting in one buffer,
|
||||
one pop_next_value() call per value with no new fragments in between,
|
||||
must be O(n) total. Re-copying the shrinking remainder on every pop
|
||||
(slicing a new string instead of advancing a cursor) makes total drain
|
||||
time scale with the square of the buffer size.
|
||||
|
||||
Uses a doubling ratio rather than an absolute ms budget so it isn't
|
||||
flaky on a slower or busier CI runner: doubling the input should
|
||||
roughly double an O(n) drain's time but roughly quadruple an O(n^2)
|
||||
drain's time, and that ratio holds regardless of machine speed.
|
||||
"""
|
||||
|
||||
def drain_time_ms(n: int) -> float:
|
||||
accumulator = JSONFragmentAccumulator()
|
||||
accumulator.append('{"a": 1}' * n)
|
||||
start = time.perf_counter()
|
||||
drained = 0
|
||||
while True:
|
||||
found, _ = accumulator.pop_next_value()
|
||||
if not found:
|
||||
break
|
||||
drained += 1
|
||||
assert drained == n
|
||||
return (time.perf_counter() - start) * 1000
|
||||
|
||||
small_ms = drain_time_ms(40_000)
|
||||
large_ms = drain_time_ms(80_000)
|
||||
|
||||
ratio = large_ms / max(small_ms, 0.001)
|
||||
assert ratio < 3.0, (
|
||||
f"doubling drained values scaled time by {ratio:.2f}x ({small_ms:.1f} ms -> {large_ms:.1f} ms); "
|
||||
"expected roughly 2x for O(n); O(n^2) regression?"
|
||||
)
|
||||
|
||||
|
||||
def test_could_close_json_after_many_blank_fragments_is_not_quadratic():
|
||||
"""
|
||||
Regression test: a hostile upstream can send malformed JSON that never
|
||||
closes, followed by thousands of blank keepalive fragments. Rescanning
|
||||
every blank fragment on each could_close_json() call would make N calls
|
||||
cost O(n^2) total; it must be O(1) regardless of how many blank
|
||||
fragments preceded it.
|
||||
"""
|
||||
accumulator = JSONFragmentAccumulator()
|
||||
accumulator.append('{"a": ') # never closes
|
||||
|
||||
start = time.perf_counter()
|
||||
for _ in range(20_000):
|
||||
accumulator.append("")
|
||||
accumulator.could_close_json()
|
||||
elapsed_ms = (time.perf_counter() - start) * 1000
|
||||
|
||||
assert accumulator.could_close_json() is False
|
||||
assert elapsed_ms < 300, (
|
||||
f"20000 blank-fragment could_close_json() calls took {elapsed_ms:.1f} ms "
|
||||
"(expected < 300 ms); quadratic rescan regression?"
|
||||
)
|
||||
|
|
@ -1008,6 +1008,143 @@ def test_multiple_partial_chunks_accumulation():
|
|||
assert result3.choices[0].delta.content == "Hello"
|
||||
|
||||
|
||||
def test_accumulated_json_partial_fragment_returns_none_without_parsing():
|
||||
"""
|
||||
Regression test: before the shared JSONFragmentAccumulator, every partial
|
||||
fragment triggered a `json.loads` attempt over the whole growing buffer,
|
||||
unlike Vertex which already deferred parsing until the buffer could close.
|
||||
A fragment that can't close a JSON value must not trigger a decode attempt.
|
||||
"""
|
||||
iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(), sync_stream=True, json_mode=False
|
||||
)
|
||||
iterator.chunk_type = "accumulated_json"
|
||||
|
||||
with patch.object(
|
||||
json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode
|
||||
) as spy:
|
||||
result = iterator._handle_accumulated_json_chunk(
|
||||
'{"type":"content_block_delta","index":0,"delta":'
|
||||
)
|
||||
assert result is None
|
||||
assert spy.call_count == 0, "incomplete buffer should not be parsed"
|
||||
|
||||
|
||||
def test_accumulated_json_does_not_reparse_every_fragment():
|
||||
"""
|
||||
Regression test for the O(n^2) json.loads-per-fragment anti-pattern: a
|
||||
payload split across many fragments must be parsed ~once, not once per
|
||||
fragment.
|
||||
"""
|
||||
text = "x" * 200_000
|
||||
blob = json.dumps(
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}
|
||||
)
|
||||
fragments = [blob[i : i + 4096] for i in range(0, len(blob), 4096)]
|
||||
assert len(fragments) > 10, "need a multi-fragment payload to exercise the bug"
|
||||
|
||||
iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(), sync_stream=True, json_mode=False
|
||||
)
|
||||
iterator.chunk_type = "accumulated_json"
|
||||
|
||||
parsed = None
|
||||
with patch.object(
|
||||
json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode
|
||||
) as spy:
|
||||
for fragment in fragments:
|
||||
out = iterator._handle_accumulated_json_chunk(fragment)
|
||||
if out is not None:
|
||||
parsed = out
|
||||
parse_calls = spy.call_count
|
||||
|
||||
assert parsed is not None, "the reassembled chunk must still parse"
|
||||
assert parsed.choices[0].delta.content == text
|
||||
assert parse_calls <= 2, (
|
||||
f"raw_decode was called {parse_calls} times for {len(fragments)} fragments; "
|
||||
"the O(n^2) per-fragment re-parse has regressed"
|
||||
)
|
||||
|
||||
|
||||
def test_accumulated_json_concatenated_envelopes_do_not_wedge():
|
||||
"""
|
||||
Regression test: Anthropic's single `json.loads(self.accumulated_json)`
|
||||
call raised "Extra data" on two concatenated envelopes and, since the
|
||||
buffer was never reset on that failure, returned None forever while
|
||||
growing without bound. The shared accumulator peels one value at a time
|
||||
and keeps the remainder, so both values surface across two calls.
|
||||
"""
|
||||
obj = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"a"}}'
|
||||
iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(), sync_stream=True, json_mode=False
|
||||
)
|
||||
iterator.chunk_type = "accumulated_json"
|
||||
|
||||
first = iterator._handle_accumulated_json_chunk(obj + obj)
|
||||
assert first is not None
|
||||
assert first.choices[0].delta.content == "a"
|
||||
|
||||
second = iterator._handle_accumulated_json_chunk("")
|
||||
assert second is not None
|
||||
assert second.choices[0].delta.content == "a"
|
||||
|
||||
assert iterator.accumulated_json == ""
|
||||
|
||||
|
||||
def test_accumulated_json_heuristic_passes_but_value_still_incomplete():
|
||||
"""
|
||||
A buffer whose newest fragment ends in '}' can still be genuinely
|
||||
incomplete (an inner object closed, the outer one didn't). The
|
||||
heuristic must let the parse attempt through, and pop_next_value
|
||||
finding nothing must propagate as None rather than raising.
|
||||
"""
|
||||
iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(), sync_stream=True, json_mode=False
|
||||
)
|
||||
iterator.chunk_type = "accumulated_json"
|
||||
|
||||
result = iterator._handle_accumulated_json_chunk('{"type": {"nested": 1}')
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_accumulated_json_setter_and_sync_end_of_stream_drain():
|
||||
"""
|
||||
The accumulated_json setter and __next__'s StopIteration drain branch:
|
||||
a buffered partial JSON must still parse and return when the
|
||||
underlying stream ends, instead of being silently dropped.
|
||||
"""
|
||||
obj = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"a"}}'
|
||||
iterator = ModelResponseIterator(
|
||||
streaming_response=iter([]), sync_stream=True, json_mode=False
|
||||
)
|
||||
iterator.chunk_type = "accumulated_json"
|
||||
iterator.accumulated_json = obj # exercises the setter
|
||||
|
||||
result = iterator.__next__()
|
||||
assert result is not None
|
||||
assert result.choices[0].delta.content == "a"
|
||||
|
||||
|
||||
def test_accumulated_json_async_end_of_stream_drain():
|
||||
"""Async twin of the sync end-of-stream drain test: __anext__'s
|
||||
StopAsyncIteration branch must also parse a buffered value."""
|
||||
import asyncio
|
||||
|
||||
obj = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"a"}}'
|
||||
iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(), sync_stream=False, json_mode=False
|
||||
)
|
||||
iterator.chunk_type = "accumulated_json"
|
||||
iterator.accumulated_json = obj
|
||||
mock_async_iterator = MagicMock()
|
||||
mock_async_iterator.__anext__ = AsyncMock(side_effect=StopAsyncIteration)
|
||||
iterator.async_response_iterator = mock_async_iterator
|
||||
|
||||
result = asyncio.run(iterator.__anext__())
|
||||
assert result is not None
|
||||
assert result.choices[0].delta.content == "a"
|
||||
|
||||
|
||||
def test_web_search_tool_result_no_extra_tool_calls():
|
||||
"""
|
||||
Test that web_search_tool_result blocks don't emit tool call chunks.
|
||||
|
|
|
|||
|
|
@ -1688,6 +1688,87 @@ class TestToolResultDocuments:
|
|||
]
|
||||
|
||||
|
||||
class TestUserContentDocuments:
|
||||
"""Documents in plain user content must survive translation (LIT-6144): each
|
||||
document block becomes an input_file part of the user message, in block order,
|
||||
exactly like image blocks become input_image parts. Untranslatable documents
|
||||
are dropped without disturbing the surrounding parts."""
|
||||
|
||||
PDF_B64 = "JVBERi0xLjQKJSBQT05H"
|
||||
PDF_DATA_URI = "data:application/pdf;base64,JVBERi0xLjQKJSBQT05H"
|
||||
PDF_URL = "https://example.com/report.pdf"
|
||||
EXPLICIT = {"mode": "explicit"}
|
||||
|
||||
def _translate(self, user_content):
|
||||
return _ADAPTER.translate_messages_to_responses_input([{"role": "user", "content": user_content}])
|
||||
|
||||
@staticmethod
|
||||
def _user_content(items):
|
||||
return next(item for item in items if item.get("type") == "message" and item.get("role") == "user")["content"]
|
||||
|
||||
def _base64_document(self, **extra):
|
||||
return {
|
||||
"type": "document",
|
||||
"source": {"type": "base64", "media_type": "application/pdf", "data": self.PDF_B64},
|
||||
**extra,
|
||||
}
|
||||
|
||||
def test_document_then_text_keeps_block_order(self):
|
||||
content = self._user_content(
|
||||
self._translate([self._base64_document(), {"type": "text", "text": "what does the pdf say?"}])
|
||||
)
|
||||
assert content == [
|
||||
{"type": "input_file", "filename": "document.pdf", "file_data": self.PDF_DATA_URI},
|
||||
{"type": "input_text", "text": "what does the pdf say?"},
|
||||
]
|
||||
|
||||
def test_document_title_becomes_filename(self):
|
||||
content = self._user_content(self._translate([self._base64_document(title="quarterly-report.pdf")]))
|
||||
assert content == [
|
||||
{"type": "input_file", "filename": "quarterly-report.pdf", "file_data": self.PDF_DATA_URI}
|
||||
]
|
||||
|
||||
def test_url_document_becomes_file_url_part(self):
|
||||
content = self._user_content(
|
||||
self._translate([{"type": "document", "source": {"type": "url", "url": self.PDF_URL}}])
|
||||
)
|
||||
assert content == [{"type": "input_file", "file_url": self.PDF_URL}]
|
||||
|
||||
def test_document_only_content_still_produces_user_message(self):
|
||||
content = self._user_content(self._translate([self._base64_document()]))
|
||||
assert content == [{"type": "input_file", "filename": "document.pdf", "file_data": self.PDF_DATA_URI}]
|
||||
|
||||
def test_empty_base64_data_drops_only_the_document_part(self):
|
||||
content = self._user_content(
|
||||
self._translate(
|
||||
[
|
||||
{"type": "text", "text": "still here"},
|
||||
{"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": ""}},
|
||||
]
|
||||
)
|
||||
)
|
||||
assert content == [{"type": "input_text", "text": "still here"}]
|
||||
|
||||
def test_non_dict_source_drops_only_the_document_part(self):
|
||||
content = self._user_content(
|
||||
self._translate([{"type": "text", "text": "still here"}, {"type": "document", "source": self.PDF_URL}])
|
||||
)
|
||||
assert content == [{"type": "input_text", "text": "still here"}]
|
||||
|
||||
def test_document_breakpoint_rides_on_the_file_part(self):
|
||||
content = self._user_content(
|
||||
self._translate([self._base64_document(prompt_cache_breakpoint=self.EXPLICIT)])
|
||||
)
|
||||
assert content == [
|
||||
{
|
||||
"type": "input_file",
|
||||
"filename": "document.pdf",
|
||||
"file_data": self.PDF_DATA_URI,
|
||||
"prompt_cache_breakpoint": self.EXPLICIT,
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def _contains_key(value, key) -> bool:
|
||||
if isinstance(value, dict):
|
||||
return key in value or any(_contains_key(v, key) for v in value.values())
|
||||
|
|
|
|||
|
|
@ -77,7 +77,7 @@ def test_bedrock_rerank_header_forwarding_sync(model):
|
|||
|
||||
with (
|
||||
patch.object(client, "post") as mock_post,
|
||||
patch(
|
||||
patch( # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport
|
||||
"litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params",
|
||||
return_value=mock_credentials_info,
|
||||
),
|
||||
|
|
@ -170,7 +170,7 @@ async def test_bedrock_rerank_header_forwarding_async(model):
|
|||
|
||||
with (
|
||||
patch.object(client, "post", new_callable=AsyncMock) as mock_post,
|
||||
patch(
|
||||
patch( # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport
|
||||
"litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params",
|
||||
return_value=mock_credentials_info,
|
||||
),
|
||||
|
|
@ -241,7 +241,7 @@ def test_bedrock_rerank_timeout_sync():
|
|||
|
||||
with (
|
||||
patch.object(client, "post") as mock_post,
|
||||
patch(
|
||||
patch( # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport
|
||||
"litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params",
|
||||
return_value=mock_credentials_info,
|
||||
),
|
||||
|
|
@ -285,7 +285,7 @@ async def test_bedrock_rerank_timeout_async():
|
|||
|
||||
with (
|
||||
patch.object(client, "post", new_callable=AsyncMock) as mock_post,
|
||||
patch(
|
||||
patch( # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport
|
||||
"litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params",
|
||||
return_value=mock_credentials_info,
|
||||
),
|
||||
|
|
@ -340,7 +340,7 @@ def test_bedrock_rerank_extra_headers_and_headers_merge():
|
|||
|
||||
with (
|
||||
patch.object(client, "post") as mock_post,
|
||||
patch(
|
||||
patch( # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport
|
||||
"litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params",
|
||||
return_value=mock_credentials_info,
|
||||
),
|
||||
|
|
@ -400,3 +400,32 @@ def test_bedrock_rerank_extra_headers_and_headers_merge():
|
|||
|
||||
except Exception as e:
|
||||
pytest.fail(f"Failed to merge and forward headers: {str(e)}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bedrock_rerank_records_llm_api_duration():
|
||||
"""The bedrock rerank handler must feed httpx timing into the logging obj, so the
|
||||
proxy can emit x-litellm-overhead-duration-ms / x-litellm-timing-* on /rerank."""
|
||||
import httpx
|
||||
|
||||
def handle(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(200, json=bedrock_rerank_response)
|
||||
|
||||
client = AsyncHTTPHandler()
|
||||
client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
|
||||
|
||||
with patch( # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport
|
||||
"litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params",
|
||||
return_value=create_mock_credentials(),
|
||||
):
|
||||
response = await litellm.arerank(
|
||||
model="bedrock/arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0",
|
||||
query=test_query,
|
||||
documents=test_documents,
|
||||
top_n=3,
|
||||
client=client,
|
||||
aws_region_name="us-east-1",
|
||||
)
|
||||
|
||||
assert response._hidden_params["litellm_overhead_time_ms"] is not None
|
||||
assert response._hidden_params["_response_ms"] >= response._hidden_params["litellm_overhead_time_ms"]
|
||||
|
|
|
|||
|
|
@ -2528,6 +2528,37 @@ def test_only_callbacks_that_can_charge_a_frame_are_collected_for_ws_quota(monke
|
|||
assert _collect_ws_project_quota_callbacks() == (quota,)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_rerank_records_llm_api_duration():
|
||||
"""arerank must feed the httpx timing into the logging obj, so the proxy can emit
|
||||
x-litellm-overhead-duration-ms / x-litellm-timing-* on /rerank."""
|
||||
|
||||
def handle(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "rerank-1",
|
||||
"results": [{"index": 0, "relevance_score": 0.9}],
|
||||
"meta": {"api_version": {"version": "2"}, "billed_units": {"search_units": 1}},
|
||||
},
|
||||
)
|
||||
|
||||
client = AsyncHTTPHandler()
|
||||
client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
|
||||
|
||||
response = await litellm.arerank(
|
||||
model="cohere/rerank-v3.5",
|
||||
query="what is the capital of france",
|
||||
documents=["paris", "berlin"],
|
||||
top_n=1,
|
||||
api_key="fake-key",
|
||||
client=client,
|
||||
)
|
||||
|
||||
assert response._hidden_params["litellm_overhead_time_ms"] is not None
|
||||
assert response._hidden_params["_response_ms"] >= response._hidden_params["litellm_overhead_time_ms"]
|
||||
|
||||
|
||||
class _JSONBodyVideoConfig(OpenAIVideoConfig):
|
||||
def use_multipart_form_data(self) -> bool:
|
||||
return False
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import json
|
||||
import logging
|
||||
from collections.abc import Mapping, Sequence
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import httpx
|
||||
|
|
@ -326,6 +327,96 @@ def test_streaming_chunk_preserves_tool_call_index_and_id():
|
|||
assert continuation["function"]["arguments"] == '{"city": "San'
|
||||
|
||||
|
||||
REPLAYED_ASSISTANT_MESSAGE = {
|
||||
"role": "assistant",
|
||||
"content": "The digit sum is 11.",
|
||||
"reasoning_content": "The secret number is 47. 4 + 7 = 11.",
|
||||
"thinking_blocks": [{"type": "thinking", "thinking": "The secret number is 47.", "signature": ""}],
|
||||
"provider_specific_fields": {"thinking_blocks": [{"type": "thinking", "thinking": "The secret number is 47."}]},
|
||||
}
|
||||
|
||||
PRESERVED_THINKING_MESSAGES = [
|
||||
{"role": "user", "content": "Pick a secret two-digit number and tell me only its digit sum."},
|
||||
REPLAYED_ASSISTANT_MESSAGE,
|
||||
{"role": "user", "content": "What was the secret number?"},
|
||||
]
|
||||
|
||||
|
||||
def _assert_internal_fields_stripped_reasoning_kept(transformed_messages: Sequence[Mapping[str, object]]):
|
||||
assistant_message = transformed_messages[1]
|
||||
assert assistant_message["reasoning_content"] == REPLAYED_ASSISTANT_MESSAGE["reasoning_content"]
|
||||
assert "thinking_blocks" not in assistant_message
|
||||
assert "provider_specific_fields" not in assistant_message
|
||||
assert assistant_message["content"] == REPLAYED_ASSISTANT_MESSAGE["content"]
|
||||
assert transformed_messages[0] == PRESERVED_THINKING_MESSAGES[0]
|
||||
assert transformed_messages[2] == PRESERVED_THINKING_MESSAGES[2]
|
||||
|
||||
|
||||
def test_transform_request_keeps_reasoning_content_strips_internal_fields():
|
||||
request = TogetherAIChatConfig().transform_request(
|
||||
model=REASONING_MODEL,
|
||||
messages=[dict(message) for message in PRESERVED_THINKING_MESSAGES],
|
||||
optional_params={},
|
||||
litellm_params={"custom_llm_provider": "together_ai"},
|
||||
headers={},
|
||||
)
|
||||
|
||||
_assert_internal_fields_stripped_reasoning_kept(request["messages"])
|
||||
|
||||
|
||||
async def test_async_transform_request_keeps_reasoning_content_strips_internal_fields():
|
||||
request = await TogetherAIChatConfig().async_transform_request(
|
||||
model=REASONING_MODEL,
|
||||
messages=[dict(message) for message in PRESERVED_THINKING_MESSAGES],
|
||||
optional_params={},
|
||||
litellm_params={"custom_llm_provider": "together_ai"},
|
||||
headers={},
|
||||
)
|
||||
|
||||
_assert_internal_fields_stripped_reasoning_kept(request["messages"])
|
||||
|
||||
|
||||
def test_completion_sends_chat_template_kwargs_and_preserved_reasoning():
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
captured_requests: list[httpx.Request] = []
|
||||
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
captured_requests.append(request)
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "chatcmpl-together-preserved",
|
||||
"object": "chat.completion",
|
||||
"created": 1234567890,
|
||||
"model": REASONING_MODEL,
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": "47"},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
|
||||
},
|
||||
)
|
||||
|
||||
client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond)))
|
||||
|
||||
litellm.completion(
|
||||
model=f"together_ai/{REASONING_MODEL}",
|
||||
messages=[dict(message) for message in PRESERVED_THINKING_MESSAGES],
|
||||
chat_template_kwargs={"clear_thinking": False},
|
||||
api_key="fake-key",
|
||||
client=client,
|
||||
)
|
||||
|
||||
request_body = json.loads(captured_requests[0].content)
|
||||
assert request_body["chat_template_kwargs"] == {"clear_thinking": False}
|
||||
assert "extra_body" not in request_body
|
||||
_assert_internal_fields_stripped_reasoning_kept(request_body["messages"])
|
||||
|
||||
|
||||
def test_together_ai_config_alias_points_at_chat_config():
|
||||
assert litellm.TogetherAIConfig is litellm.TogetherAIChatConfig
|
||||
config = litellm.TogetherAIConfig(max_tokens=10)
|
||||
|
|
|
|||
|
|
@ -3002,8 +3002,11 @@ def test_accumulated_json_does_not_reparse_every_fragment():
|
|||
|
||||
The buffer only becomes a complete JSON object on the final fragment, so a
|
||||
correct implementation parses it ~once, not once per fragment. We assert the
|
||||
full chunk still parses correctly AND that json.loads is not called on every
|
||||
fragment (which is what made it quadratic).
|
||||
full chunk still parses correctly AND that the buffer is not decoded on
|
||||
every fragment (which is what made it quadratic). Post-migration to the
|
||||
shared JSONFragmentAccumulator, decoding goes through
|
||||
`json.JSONDecoder.raw_decode`, not `json.loads` (see the equivalent
|
||||
Anthropic tests) so the spy targets that call, not `json.loads`.
|
||||
"""
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
ModelResponseIterator,
|
||||
|
|
@ -3024,7 +3027,9 @@ def test_accumulated_json_does_not_reparse_every_fragment():
|
|||
assert len(fragments) > 10, "need a multi-fragment payload to exercise the bug"
|
||||
|
||||
parsed = None
|
||||
with patch("json.loads", wraps=json.loads) as spy:
|
||||
with patch.object(
|
||||
json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode
|
||||
) as spy:
|
||||
for fragment in fragments:
|
||||
out = iterator.handle_accumulated_json_chunk(chunk=fragment)
|
||||
if out is not None:
|
||||
|
|
@ -3035,14 +3040,16 @@ def test_accumulated_json_does_not_reparse_every_fragment():
|
|||
assert parsed.choices[0].delta.content == text, "content must be preserved intact"
|
||||
|
||||
assert parse_calls <= 2, (
|
||||
f"json.loads was called {parse_calls} times for {len(fragments)} "
|
||||
f"raw_decode was called {parse_calls} times for {len(fragments)} "
|
||||
"fragments; the O(n^2) per-fragment re-parse has regressed"
|
||||
)
|
||||
|
||||
|
||||
def test_accumulated_json_partial_fragment_returns_none_without_parsing():
|
||||
"""A fragment that cannot complete the JSON must not trigger a json.loads
|
||||
parse of the whole growing buffer (issue #26181)."""
|
||||
"""A fragment that cannot complete the JSON must not trigger a decode
|
||||
attempt over the whole growing buffer (issue #26181). Decoding goes
|
||||
through `json.JSONDecoder.raw_decode` post-JSONFragmentAccumulator
|
||||
migration, not `json.loads`."""
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
ModelResponseIterator,
|
||||
)
|
||||
|
|
@ -3054,7 +3061,9 @@ def test_accumulated_json_partial_fragment_returns_none_without_parsing():
|
|||
)
|
||||
iterator.chunk_type = "accumulated_json"
|
||||
|
||||
with patch("json.loads", wraps=json.loads) as spy:
|
||||
with patch.object(
|
||||
json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode
|
||||
) as spy:
|
||||
result = iterator.handle_accumulated_json_chunk(
|
||||
chunk='{"candidates": [{"content": {"parts": [{"text": "partial'
|
||||
)
|
||||
|
|
@ -5552,3 +5561,21 @@ def test_accumulated_json_skips_non_dict_leading_value():
|
|||
|
||||
assert len(out) == 1
|
||||
assert out[0].choices[0].delta.content == "a"
|
||||
|
||||
|
||||
def test_accumulated_json_async_end_of_stream_drains_buffered_value():
|
||||
"""Async twin of test_accumulated_json_end_of_stream_drains_all_buffered_values:
|
||||
__anext__'s StopAsyncIteration branch must also parse a buffered value."""
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
obj = '{"candidates":[{"content":{"parts":[{"text":"a"}]}}],"usageMetadata":{}}'
|
||||
iterator = _accumulating_gemini_iterator()
|
||||
iterator.accumulated_json = obj
|
||||
mock_async_iterator = MagicMock()
|
||||
mock_async_iterator.__anext__ = AsyncMock(side_effect=StopAsyncIteration)
|
||||
iterator.async_response_iterator = mock_async_iterator
|
||||
|
||||
result = asyncio.run(iterator.__anext__())
|
||||
assert result is not None
|
||||
assert result.choices[0].delta.content == "a"
|
||||
|
|
|
|||
0
tests/test_litellm/proxy/rerank_endpoints/__init__.py
Normal file
0
tests/test_litellm/proxy/rerank_endpoints/__init__.py
Normal file
120
tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py
Normal file
120
tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py
Normal file
|
|
@ -0,0 +1,120 @@
|
|||
"""
|
||||
Tests for rerank_endpoints/endpoints.py response headers.
|
||||
"""
|
||||
|
||||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import Request, Response
|
||||
|
||||
import litellm.proxy.common_request_processing as common_request_processing_mod
|
||||
import litellm.proxy.proxy_server as proxy_server_mod
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.rerank_endpoints.endpoints import rerank
|
||||
from litellm.types.utils import RerankResponse
|
||||
|
||||
HIDDEN_PARAMS = {
|
||||
"model_id": "deployment-1",
|
||||
"api_base": "https://bedrock-agent-runtime.us-east-1.amazonaws.com",
|
||||
"response_cost": 0.002,
|
||||
"_response_ms": 1500.5,
|
||||
"litellm_overhead_time_ms": 12.5,
|
||||
"callback_duration_ms": 1.25,
|
||||
"timing_llm_api_ms": 1488.0,
|
||||
"timing_pre_processing_ms": 10.0,
|
||||
"timing_post_processing_ms": 2.5,
|
||||
"timing_message_copy_ms": 0.01,
|
||||
}
|
||||
|
||||
|
||||
def _build_request() -> Request:
|
||||
body = json.dumps({"model": "rerank-model", "query": "q", "documents": ["a", "b"]}).encode()
|
||||
|
||||
async def receive():
|
||||
return {"type": "http.request", "body": body, "more_body": False}
|
||||
|
||||
return Request(
|
||||
scope={
|
||||
"type": "http",
|
||||
"method": "POST",
|
||||
"path": "/rerank",
|
||||
"headers": [(b"content-type", b"application/json")],
|
||||
"query_string": b"",
|
||||
},
|
||||
receive=receive,
|
||||
)
|
||||
|
||||
|
||||
async def _call_rerank(hidden_params: dict = HIDDEN_PARAMS) -> Response:
|
||||
response = RerankResponse(id="rerank-1", results=[{"index": 0, "relevance_score": 0.9}])
|
||||
response._hidden_params = dict(hidden_params)
|
||||
|
||||
fastapi_response = Response()
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"])
|
||||
proxy_logging_obj.update_request_status = AsyncMock()
|
||||
|
||||
async def fake_add_litellm_data_to_request(**kwargs):
|
||||
return {**kwargs["data"], "litellm_call_id": "call-123"}
|
||||
|
||||
async def fake_route_request(**kwargs):
|
||||
async def _call():
|
||||
return response
|
||||
|
||||
return _call()
|
||||
|
||||
with (
|
||||
patch.object(proxy_server_mod, "add_litellm_data_to_request", fake_add_litellm_data_to_request), # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler
|
||||
patch.object(proxy_server_mod, "route_request", fake_route_request), # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler
|
||||
patch.object(proxy_server_mod, "proxy_logging_obj", proxy_logging_obj), # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler
|
||||
patch.object(proxy_server_mod, "llm_router", MagicMock()), # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler
|
||||
patch.object(proxy_server_mod, "version", "1.2.3"), # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler
|
||||
):
|
||||
await rerank(
|
||||
request=_build_request(),
|
||||
fastapi_response=fastapi_response,
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
|
||||
)
|
||||
|
||||
return fastapi_response
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rerank_emits_latency_and_cost_headers():
|
||||
"""/rerank must surface the same hidden_params-derived headers as /chat/completions."""
|
||||
fastapi_response = await _call_rerank()
|
||||
|
||||
assert fastapi_response.headers["x-litellm-call-id"] == "call-123"
|
||||
assert fastapi_response.headers["x-litellm-response-duration-ms"] == "1500.5"
|
||||
assert fastapi_response.headers["x-litellm-overhead-duration-ms"] == "12.5"
|
||||
assert fastapi_response.headers["x-litellm-callback-duration-ms"] == "1.25"
|
||||
assert fastapi_response.headers["x-litellm-response-cost"] == "0.002"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rerank_emits_detailed_timing_headers_when_enabled():
|
||||
"""LITELLM_DETAILED_TIMING must also work on /rerank, not just /chat/completions."""
|
||||
with patch.object(common_request_processing_mod, "LITELLM_DETAILED_TIMING", True): # test-quality-ok: LITELLM_DETAILED_TIMING is a module constant; toggling it is the behavior under test
|
||||
fastapi_response = await _call_rerank()
|
||||
|
||||
assert fastapi_response.headers["x-litellm-timing-llm-api-ms"] == "1488.0"
|
||||
assert fastapi_response.headers["x-litellm-timing-pre-processing-ms"] == "10.0"
|
||||
assert fastapi_response.headers["x-litellm-timing-post-processing-ms"] == "2.5"
|
||||
assert fastapi_response.headers["x-litellm-timing-message-copy-ms"] == "0.01"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rerank_emits_zero_response_cost_header():
|
||||
"""A free deployment costs 0.0, which is a real cost and must not be dropped."""
|
||||
fastapi_response = await _call_rerank({**HIDDEN_PARAMS, "response_cost": 0.0})
|
||||
|
||||
assert fastapi_response.headers["x-litellm-response-cost"] == "0.0"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rerank_omits_detailed_timing_headers_when_disabled():
|
||||
with patch.object(common_request_processing_mod, "LITELLM_DETAILED_TIMING", False): # test-quality-ok: LITELLM_DETAILED_TIMING is a module constant; toggling it is the behavior under test
|
||||
fastapi_response = await _call_rerank()
|
||||
|
||||
assert "x-litellm-timing-llm-api-ms" not in fastapi_response.headers
|
||||
|
|
@ -2629,7 +2629,7 @@ class TestSpendLogsPayload:
|
|||
"model": "gpt-4o",
|
||||
"user": "",
|
||||
"team_id": "",
|
||||
"metadata": '{"applied_guardrails": [], "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}',
|
||||
"metadata": '{"applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}',
|
||||
"cache_key": "Cache OFF",
|
||||
"spend": 0.00022500000000000002,
|
||||
"total_tokens": 30,
|
||||
|
|
@ -2725,7 +2725,7 @@ class TestSpendLogsPayload:
|
|||
"model": "claude-4-sonnet-20250514",
|
||||
"user": "",
|
||||
"team_id": "",
|
||||
"metadata": '{"applied_guardrails": [], "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}',
|
||||
"metadata": '{"applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}',
|
||||
"cache_key": "Cache OFF",
|
||||
"spend": 0.01383,
|
||||
"total_tokens": 2598,
|
||||
|
|
@ -2819,7 +2819,7 @@ class TestSpendLogsPayload:
|
|||
"model": "claude-4-sonnet-20250514",
|
||||
"user": "",
|
||||
"team_id": "",
|
||||
"metadata": '{"applied_guardrails": [], "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}',
|
||||
"metadata": '{"applied_guardrails": [], "attempted_fallbacks": 0, "original_model_group": "my-anthropic-model-group", "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}',
|
||||
"cache_key": "Cache OFF",
|
||||
"spend": 0.01383,
|
||||
"total_tokens": 2598,
|
||||
|
|
|
|||
|
|
@ -3642,3 +3642,200 @@ def test_caller_forged_autorouter_savings_is_discarded(bucket):
|
|||
)
|
||||
metadata = json.loads(payload["metadata"])
|
||||
assert metadata["autorouter_savings"] is None
|
||||
|
||||
|
||||
def test_get_logging_payload_includes_fallback_info_in_spend_logs_metadata():
|
||||
"""
|
||||
Test that fallback info (attempted_fallbacks, original_model_group) from metadata
|
||||
is included in the spend logs metadata JSON.
|
||||
"""
|
||||
kwargs = {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"user_api_key": "sk-test-key",
|
||||
"attempted_fallbacks": 2,
|
||||
"original_model_group": "azure-gpt-fallback",
|
||||
}
|
||||
},
|
||||
"standard_logging_object": StandardLoggingPayload(
|
||||
id="test-fallback-123",
|
||||
call_type="completion",
|
||||
stream=False,
|
||||
response_cost=0.001,
|
||||
status="success",
|
||||
total_tokens=100,
|
||||
prompt_tokens=50,
|
||||
completion_tokens=50,
|
||||
startTime=1234567890.0,
|
||||
endTime=1234567891.0,
|
||||
completionStartTime=None,
|
||||
model_map_information=StandardLoggingModelInformation(
|
||||
model_map_key="gpt-3.5-turbo", model_map_value=None
|
||||
),
|
||||
model="gpt-3.5-turbo",
|
||||
model_id="model-123",
|
||||
model_group="openai",
|
||||
custom_llm_provider="openai",
|
||||
api_base="https://api.openai.com",
|
||||
metadata=StandardLoggingMetadata(
|
||||
user_api_key_hash="test_hash",
|
||||
user_api_key_alias=None,
|
||||
user_api_key_team_id=None,
|
||||
user_api_key_org_id=None,
|
||||
user_api_key_user_id=None,
|
||||
user_api_key_team_alias=None,
|
||||
spend_logs_metadata=None,
|
||||
requester_ip_address=None,
|
||||
requester_metadata=None,
|
||||
user_api_key_end_user_id=None,
|
||||
),
|
||||
cache_hit=False,
|
||||
cache_key=None,
|
||||
saved_cache_cost=0.0,
|
||||
request_tags=[],
|
||||
end_user=None,
|
||||
requester_ip_address=None,
|
||||
messages=[],
|
||||
response={},
|
||||
error_str=None,
|
||||
model_parameters={},
|
||||
hidden_params=StandardLoggingHiddenParams(
|
||||
model_id="model-123",
|
||||
cache_key=None,
|
||||
api_base="https://api.openai.com",
|
||||
response_cost="0.001",
|
||||
litellm_overhead_time_ms=None,
|
||||
additional_headers=None,
|
||||
batch_models=None,
|
||||
litellm_model_name=None,
|
||||
usage_object=None,
|
||||
),
|
||||
),
|
||||
}
|
||||
|
||||
response_obj = {
|
||||
"id": "test-response-retry",
|
||||
"choices": [{"message": {"content": "Hello!"}}],
|
||||
"usage": {
|
||||
"total_tokens": 100,
|
||||
"prompt_tokens": 50,
|
||||
"completion_tokens": 50,
|
||||
},
|
||||
}
|
||||
|
||||
start_time = datetime.datetime.now(timezone.utc)
|
||||
end_time = datetime.datetime.now(timezone.utc)
|
||||
|
||||
payload = get_logging_payload(
|
||||
kwargs=kwargs,
|
||||
response_obj=response_obj,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
|
||||
metadata = json.loads(payload["metadata"])
|
||||
|
||||
assert (
|
||||
metadata.get("attempted_fallbacks") == 2
|
||||
), f"Expected attempted_fallbacks=2, got {metadata.get('attempted_fallbacks')}"
|
||||
assert (
|
||||
metadata.get("original_model_group") == "azure-gpt-fallback"
|
||||
), f"Expected original_model_group=azure-gpt-fallback, got {metadata.get('original_model_group')}"
|
||||
|
||||
|
||||
def test_get_logging_payload_handles_missing_fallback_info_gracefully():
|
||||
"""
|
||||
Test that fallback fields are None when not present in metadata (backward compatibility).
|
||||
"""
|
||||
kwargs = {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"user_api_key": "sk-test-key",
|
||||
}
|
||||
},
|
||||
"standard_logging_object": StandardLoggingPayload(
|
||||
id="test-no-fallback-456",
|
||||
call_type="completion",
|
||||
stream=False,
|
||||
response_cost=0.001,
|
||||
status="success",
|
||||
total_tokens=100,
|
||||
prompt_tokens=50,
|
||||
completion_tokens=50,
|
||||
startTime=1234567890.0,
|
||||
endTime=1234567891.0,
|
||||
completionStartTime=None,
|
||||
model_map_information=StandardLoggingModelInformation(
|
||||
model_map_key="gpt-3.5-turbo", model_map_value=None
|
||||
),
|
||||
model="gpt-3.5-turbo",
|
||||
model_id="model-123",
|
||||
model_group="openai",
|
||||
custom_llm_provider="openai",
|
||||
api_base="https://api.openai.com",
|
||||
metadata=StandardLoggingMetadata(
|
||||
user_api_key_hash="test_hash",
|
||||
user_api_key_alias=None,
|
||||
user_api_key_team_id=None,
|
||||
user_api_key_org_id=None,
|
||||
user_api_key_user_id=None,
|
||||
user_api_key_team_alias=None,
|
||||
spend_logs_metadata=None,
|
||||
requester_ip_address=None,
|
||||
requester_metadata=None,
|
||||
user_api_key_end_user_id=None,
|
||||
),
|
||||
cache_hit=False,
|
||||
cache_key=None,
|
||||
saved_cache_cost=0.0,
|
||||
request_tags=[],
|
||||
end_user=None,
|
||||
requester_ip_address=None,
|
||||
messages=[],
|
||||
response={},
|
||||
error_str=None,
|
||||
model_parameters={},
|
||||
hidden_params=StandardLoggingHiddenParams(
|
||||
model_id="model-123",
|
||||
cache_key=None,
|
||||
api_base="https://api.openai.com",
|
||||
response_cost="0.001",
|
||||
litellm_overhead_time_ms=None,
|
||||
additional_headers=None,
|
||||
batch_models=None,
|
||||
litellm_model_name=None,
|
||||
usage_object=None,
|
||||
),
|
||||
),
|
||||
}
|
||||
|
||||
response_obj = {
|
||||
"id": "test-response-no-fallback",
|
||||
"choices": [{"message": {"content": "Hello!"}}],
|
||||
"usage": {
|
||||
"total_tokens": 100,
|
||||
"prompt_tokens": 50,
|
||||
"completion_tokens": 50,
|
||||
},
|
||||
}
|
||||
|
||||
start_time = datetime.datetime.now(timezone.utc)
|
||||
end_time = datetime.datetime.now(timezone.utc)
|
||||
|
||||
payload = get_logging_payload(
|
||||
kwargs=kwargs,
|
||||
response_obj=response_obj,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
|
||||
metadata = json.loads(payload["metadata"])
|
||||
|
||||
assert (
|
||||
metadata.get("attempted_fallbacks") is None
|
||||
), "attempted_fallbacks should be None when not provided"
|
||||
assert (
|
||||
metadata.get("original_model_group") is None
|
||||
), "original_model_group should be None when not provided"
|
||||
|
|
|
|||
|
|
@ -308,7 +308,11 @@ async def test_run_async_fallback_handles_explicitly_none_metadata():
|
|||
metadata=None,
|
||||
)
|
||||
|
||||
assert router.received_kwargs["metadata"] == {"model_group": "azure-group"}
|
||||
assert router.received_kwargs["metadata"] == {
|
||||
"model_group": "azure-group",
|
||||
"attempted_fallbacks": 1,
|
||||
"original_model_group": "openai-group",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -843,3 +847,45 @@ class TestRunAsyncFallbackTriggersCooldown:
|
|||
)
|
||||
|
||||
mock_trigger.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async_fallback_stamps_fallback_info_into_metadata():
|
||||
"""Spend logs are built from the request metadata of the nested call, so the
|
||||
fallback signal has to be stamped there before recursing."""
|
||||
router = RecordingRouter()
|
||||
|
||||
await run_async_fallback(
|
||||
litellm_router=router,
|
||||
fallback_model_group=["fallback-model"],
|
||||
original_model_group="primary-model",
|
||||
original_exception=RuntimeError("original failed"),
|
||||
max_fallbacks=3,
|
||||
fallback_depth=0,
|
||||
)
|
||||
|
||||
metadata = router.received_kwargs["metadata"]
|
||||
assert metadata["attempted_fallbacks"] == 1
|
||||
assert metadata["original_model_group"] == "primary-model"
|
||||
assert metadata["model_group"] == "fallback-model"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async_fallback_preserves_original_model_group_on_nested_fallback():
|
||||
"""A second-level fallback receives the first fallback target as its
|
||||
original_model_group argument, so the first-stamped value must survive the hop."""
|
||||
router = RecordingRouter()
|
||||
|
||||
await run_async_fallback(
|
||||
litellm_router=router,
|
||||
fallback_model_group=["second-fallback"],
|
||||
original_model_group="first-fallback",
|
||||
original_exception=RuntimeError("first fallback failed"),
|
||||
max_fallbacks=3,
|
||||
fallback_depth=1,
|
||||
metadata={"attempted_fallbacks": 1, "original_model_group": "primary-model"},
|
||||
)
|
||||
|
||||
metadata = router.received_kwargs["metadata"]
|
||||
assert metadata["attempted_fallbacks"] == 2
|
||||
assert metadata["original_model_group"] == "primary-model"
|
||||
|
|
|
|||
|
|
@ -10325,3 +10325,168 @@ async def test_factory_function_anthropic_messages_uses_streaming_fallback_dispa
|
|||
result = await wrapped(model="primary")
|
||||
assert result == "ok"
|
||||
mock_anthropic.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_function_with_fallbacks_stamps_zero_attempted_fallbacks():
|
||||
"""A request served by the primary model group records attempted_fallbacks=0 and
|
||||
the requested model group in metadata, mirroring the x-litellm-attempted-fallbacks header."""
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"},
|
||||
}
|
||||
]
|
||||
)
|
||||
metadata = {}
|
||||
|
||||
await router.acompletion(
|
||||
model="gpt-3.5-turbo",
|
||||
messages=[{"role": "user", "content": "hey"}],
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
assert metadata["attempted_fallbacks"] == 0
|
||||
assert metadata["original_model_group"] == "gpt-3.5-turbo"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_function_with_fallbacks_stamps_route_bucket_not_litellm_metadata():
|
||||
"""A chat completion carrying both metadata buckets gets stamped in the route's bucket
|
||||
(metadata), matching where run_async_fallback rewrites, so the two never diverge."""
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"},
|
||||
}
|
||||
]
|
||||
)
|
||||
metadata = {}
|
||||
litellm_metadata = {"client_key": "client_value"}
|
||||
|
||||
await router.acompletion(
|
||||
model="gpt-3.5-turbo",
|
||||
messages=[{"role": "user", "content": "hey"}],
|
||||
metadata=metadata,
|
||||
litellm_metadata=litellm_metadata,
|
||||
)
|
||||
|
||||
assert metadata["attempted_fallbacks"] == 0
|
||||
assert metadata["original_model_group"] == "gpt-3.5-turbo"
|
||||
assert litellm_metadata["client_key"] == "client_value"
|
||||
assert "attempted_fallbacks" not in litellm_metadata
|
||||
assert "original_model_group" not in litellm_metadata
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_function_with_fallbacks_overrides_client_supplied_stamp_values():
|
||||
"""Client-supplied attempted_fallbacks and original_model_group are replaced on entry,
|
||||
so a reused metadata dict or a spoofed value cannot leak stale attribution into logs."""
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"},
|
||||
}
|
||||
]
|
||||
)
|
||||
metadata = {"attempted_fallbacks": 99, "original_model_group": "stale-group"}
|
||||
|
||||
await router.acompletion(
|
||||
model="gpt-3.5-turbo",
|
||||
messages=[{"role": "user", "content": "hey"}],
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
assert metadata["attempted_fallbacks"] == 0
|
||||
assert metadata["original_model_group"] == "gpt-3.5-turbo"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_function_with_fallbacks_stamps_despite_forged_reentry_params():
|
||||
"""A client injecting fallback_depth or a JSON-shaped attempted_targets via request
|
||||
litellm params cannot skip the entry stamp; only the router's own in-process
|
||||
AttemptedFallbackTargets instance marks a genuine re-entrant hop."""
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"},
|
||||
}
|
||||
]
|
||||
)
|
||||
metadata = {"attempted_fallbacks": 99, "original_model_group": "spoofed-group"}
|
||||
|
||||
await router.acompletion(
|
||||
model="gpt-3.5-turbo",
|
||||
messages=[{"role": "user", "content": "hey"}],
|
||||
metadata=metadata,
|
||||
fallback_depth=3,
|
||||
attempted_targets={"keys": ["spoofed-group"]},
|
||||
)
|
||||
|
||||
assert metadata["attempted_fallbacks"] == 0
|
||||
assert metadata["original_model_group"] == "gpt-3.5-turbo"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_function_with_fallbacks_skips_stamp_on_genuine_reentrant_hop():
|
||||
"""A re-entrant hop carrying the router's own AttemptedFallbackTargets instance keeps
|
||||
the per-hop metadata that run_async_fallback wrote instead of resetting it to zero."""
|
||||
from litellm.router_utils.fallback_event_handlers import AttemptedFallbackTargets
|
||||
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"},
|
||||
}
|
||||
]
|
||||
)
|
||||
metadata = {"attempted_fallbacks": 1, "original_model_group": "prod-chat"}
|
||||
|
||||
await router.acompletion(
|
||||
model="gpt-3.5-turbo",
|
||||
messages=[{"role": "user", "content": "hey"}],
|
||||
metadata=metadata,
|
||||
attempted_targets=AttemptedFallbackTargets(keys=frozenset(("prod-chat",))),
|
||||
)
|
||||
|
||||
assert metadata["attempted_fallbacks"] == 1
|
||||
assert metadata["original_model_group"] == "prod-chat"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_function_with_fallbacks_scrubs_spoofed_values_from_sibling_bucket():
|
||||
"""Spend logs read a truthy litellm_metadata dict in preference to metadata, so spoofed
|
||||
stamp keys planted in the bucket the route does not own are removed on entry instead of
|
||||
flowing into the spend log row."""
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"},
|
||||
}
|
||||
]
|
||||
)
|
||||
metadata = {}
|
||||
litellm_metadata = {
|
||||
"attempted_fallbacks": 99,
|
||||
"original_model_group": "spoofed-group",
|
||||
"client_key": "client_value",
|
||||
}
|
||||
|
||||
await router.acompletion(
|
||||
model="gpt-3.5-turbo",
|
||||
messages=[{"role": "user", "content": "hey"}],
|
||||
metadata=metadata,
|
||||
litellm_metadata=litellm_metadata,
|
||||
)
|
||||
|
||||
assert "attempted_fallbacks" not in litellm_metadata
|
||||
assert "original_model_group" not in litellm_metadata
|
||||
assert litellm_metadata["client_key"] == "client_value"
|
||||
assert metadata["attempted_fallbacks"] == 0
|
||||
assert metadata["original_model_group"] == "gpt-3.5-turbo"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue