From f5d6e89ad8341e4df5400a0bd0e99509b8ae62c8 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 10:42:43 +0000 Subject: [PATCH 01/33] test(integration): follow #47441 metadata-routed counting in count_tokens wire tests (#47525) * test(integration): follow #47441 metadata-routed Mantle counting in count_tokens wire tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): build the runtime count body without mutation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_bedrock_mantle_count_tokens_wire.py | 205 ++++++++++++------ .../providers/test_count_tokens_model_wire.py | 81 ++++++- 2 files changed, 207 insertions(+), 79 deletions(-) diff --git a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py index ab0c35d0ca4..4fe29fa49d8 100644 --- a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py @@ -70,6 +70,18 @@ _SCRIPTED_STATUS: Final = re.compile(r"status=(\d{3})") _STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _JSON_LIST: Final = TypeAdapter(list[JsonValue]) +_COST_MAP_PATH: Final = Path(__file__).resolve().parents[3] / "model_prices_and_context_window.json" +_COST_MAP: Final = TypeAdapter(dict[str, dict[str, JsonValue]]) + + +def _count_model_source(model: str) -> dict[str, JsonValue]: + cost_map: Final = _COST_MAP.validate_json(_COST_MAP_PATH.read_bytes()) + token_count_model: Final = cost_map[model].get("token_count_model") + assert isinstance(token_count_model, dict), model + count_model: Final = token_count_model.get("model") + assert isinstance(count_model, str), token_count_model + return {"litellm_count_source": "token_count_model", "litellm_count_model": count_model} + _SDK_MESSAGES: Final[list[MessageParam]] = [{"role": "user", "content": "Count this message"}] _MESSAGES: Final = _JSON_LIST.validate_json(json.dumps(_SDK_MESSAGES)) @@ -94,6 +106,32 @@ _GEMINI_BODY: Final[dict[str, JsonValue]] = {"contents": [{"role": "user", "part _GEMINI_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": "Count this"}] +def _runtime_count_body(prompt: str, include_system_and_tools: bool = False) -> dict[str, JsonValue]: + tool: Final = _SDK_TOOLS[0] + converse: Final[dict[str, JsonValue]] = { + "messages": [{"role": "user", "content": [{"text": prompt}]}], + **( + { + "system": [{"text": _SYSTEM}], + "toolConfig": { + "tools": [ + { + "toolSpec": { + "name": tool["name"], + "description": tool["description"], + "inputSchema": {"json": tool["input_schema"]}, + } + } + ] + }, + } + if include_system_and_tools + else {} + ), + } + return {"input": {"converse": converse}} + + def _mantle_body(**fields: JsonValue) -> dict[str, JsonValue]: return {"model": _OPUS_BASE, "messages": _MESSAGES, **fields} @@ -223,7 +261,11 @@ def _owned_config(path: Path, runtime_url: str, settings: Mapping[str, JsonValue { **config, "model_list": [ - {"model_name": name, "litellm_params": _litellm_params(model, runtime_url)} + { + "model_name": name, + "litellm_params": _litellm_params(model, runtime_url), + "model_info": {"token_count_model": {}}, + } for name, model in _OWNED_MODELS.items() ], "litellm_settings": {**litellm_settings, **settings}, @@ -233,7 +275,11 @@ def _owned_config(path: Path, runtime_url: str, settings: Mapping[str, JsonValue return path -def _deployment(scenario: Scenario, api_base: str, model: str = _OPUS) -> str: +def _fallback_deployment(scenario: Scenario, api_base: str, model: str = _OPUS) -> str: + return scenario.model(model_info={"token_count_model": {}}, api_key=None, **_litellm_params(model, api_base)) + + +def _metadata_deployment(scenario: Scenario, api_base: str, model: str = _OPUS) -> str: return scenario.model(model_info=None, api_key=None, **_litellm_params(model, api_base)) @@ -290,6 +336,14 @@ def _runtime_count_targets(wire: Wire) -> tuple[str, ...]: return tuple(target_of(request) for request in counts) +def _runtime_count_requests(wire: Wire) -> tuple[tuple[str, dict[str, JsonValue]], ...]: + counts: Final = tuple(request for request in wire.drain() if _COUNT_TARGET.match(target_of(request))) + for request in counts: + assert request.method == "POST", request.method + _assert_sigv4(request) + return tuple((target_of(request), _JSON_OBJECT.validate_json(request.body)) for request in counts) + + def _clients(stack: ExitStack, base_url: str, count: int, timeout: float = 30) -> tuple[httpx.Client, ...]: return tuple( stack.enter_context(httpx.Client(base_url=base_url, timeout=timeout, trust_env=False)) for _ in range(count) @@ -376,7 +430,7 @@ def test_messages_count_tokens_counts_through_mantle_when_the_runtime_cannot( wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS}) assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",) @@ -392,7 +446,7 @@ def test_messages_count_tokens_without_system_or_tools_sends_a_bare_body_to_mant wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, _bare(model)) assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 @@ -408,7 +462,7 @@ def test_utils_token_counter_call_endpoint_reports_the_mantle_tokenizer( wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = gateway.request( "POST", "/utils/token_counter", _full(model), params={"call_endpoint": "true"} ) @@ -427,14 +481,14 @@ def test_gemini_count_tokens_route_reports_the_mantle_total(counting_proxy: Owne wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _metadata_deployment(scenario, runtime.url) response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", _GEMINI_BODY) assert _payload(response) == { "totalTokens": _MANTLE_COUNT, "promptTokensDetails": [], - "litellm_count_source": "provider", + **_count_model_source(_OPUS), }, response.text - assert len(_runtime_count_targets(runtime)) == 1 + assert _runtime_count_requests(runtime) == () assert _mantle_bodies(mantle) == ({"model": _OPUS_BASE, "messages": _GEMINI_MESSAGES},) @@ -447,14 +501,16 @@ def test_gemini_count_tokens_route_reports_the_runtime_total_for_a_model_bedrock wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url, _SONNET) + model: Final = _metadata_deployment(scenario, runtime.url, _SONNET) response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", _GEMINI_BODY) assert _payload(response) == { "totalTokens": _RUNTIME_COUNT, "promptTokensDetails": [], - "litellm_count_source": "provider", + **_count_model_source(_SONNET), }, response.text - assert _runtime_count_targets(runtime) == (f"/model/{_SONNET}/count-tokens",) + assert _runtime_count_requests(runtime) == ( + (f"/model/{_SONNET}/count-tokens", _runtime_count_body("Count this")), + ) assert mantle.drain() == () @@ -466,7 +522,7 @@ def test_runtime_404_on_claude_falls_back_to_mantle(counting_proxy: OwnedProxy, wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS}) assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 @@ -481,7 +537,7 @@ def test_runtime_404_and_mantle_404_surface_provider_error(counting_proxy: Owned wire_server(_mantle(_rejecting(404)), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS}) assert response.status_code == 404, response.text assert _REJECTION in response.text @@ -496,7 +552,7 @@ def test_runtime_404_on_nova_does_not_call_mantle(counting_proxy: OwnedProxy, ma wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url, _NOVA) + model: Final = _fallback_deployment(scenario, runtime.url, _NOVA) body: Final = {**_bare(model), "messages": [{"role": "user", "content": "status=404"}]} response: Final = _count(gateway, body) assert response.status_code == 404, response.text @@ -504,9 +560,7 @@ def test_runtime_404_on_nova_does_not_call_mantle(counting_proxy: OwnedProxy, ma assert mantle.drain() == () -def test_google_count_tokens_forwards_system_and_tools_to_mantle( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: +def test_google_count_tokens_forwards_system_and_tools_to_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: gateway: Final = counting_proxy.gateway system_text: Final = "Use the provided tool" tool_schema: Final = {"type": "object", "properties": {"location": {"type": "string"}}} @@ -526,7 +580,7 @@ def test_google_count_tokens_forwards_system_and_tools_to_mantle( wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", body) assert _payload(response) == { "totalTokens": _MANTLE_COUNT, @@ -575,7 +629,7 @@ def test_google_count_tokens_normalizes_uppercase_parameters_for_mantle( wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", body) assert _payload(response) == { "totalTokens": _MANTLE_COUNT, @@ -607,7 +661,7 @@ def test_responses_input_tokens_counts_through_mantle(counting_proxy: OwnedProxy wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = gateway.request( "POST", "/v1/responses/input_tokens", {"model": model, "input": "Count this message"} ) @@ -627,7 +681,7 @@ def test_anthropic_sdk_count_tokens_counts_through_mantle(counting_proxy: OwnedP wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) counted: Final = _anthropic_client(gateway).messages.count_tokens( model=model, messages=_SDK_MESSAGES, system=_SYSTEM, tools=_SDK_TOOLS ) @@ -643,7 +697,7 @@ def test_async_anthropic_sdk_count_tokens_counts_through_mantle(counting_proxy: wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) counted: Final = asyncio.run( _async_anthropic_client(gateway).messages.count_tokens(model=model, messages=_SDK_MESSAGES, system=_SYSTEM) ) @@ -659,17 +713,36 @@ def test_model_the_runtime_counts_never_reaches_mantle(counting_proxy: OwnedProx wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url, _SONNET) + model: Final = _metadata_deployment(scenario, runtime.url, _SONNET) assert _payload(_count(gateway, _full(model))) == { "input_tokens": _RUNTIME_COUNT, - "litellm_count_source": "provider", + **_count_model_source(_SONNET), } detailed: Final = gateway.request( "POST", "/utils/token_counter", _full(model), params={"call_endpoint": "true"} ) payload: Final = _payload(detailed) - assert (payload["total_tokens"], payload["tokenizer_type"]) == (_RUNTIME_COUNT, "bedrock_api"), detailed.text - assert _runtime_count_targets(runtime) == (f"/model/{_SONNET}/count-tokens",) * 2 + assert payload == { + "total_tokens": _RUNTIME_COUNT, + "request_model": model, + "model_used": _count_model_source(_SONNET)["litellm_count_model"], + "tokenizer_type": "bedrock_api", + "count_source": "token_count_model", + "original_response": {"input_tokens": _RUNTIME_COUNT}, + "error": False, + "error_message": None, + "status_code": None, + }, detailed.text + assert _runtime_count_requests(runtime) == ( + ( + f"/model/{_SONNET}/count-tokens", + _runtime_count_body("Count this message", include_system_and_tools=True), + ), + ( + f"/model/{_SONNET}/count-tokens", + _runtime_count_body("Count this message", include_system_and_tools=True), + ), + ) assert mantle.drain() == () @@ -682,7 +755,7 @@ def test_non_claude_model_the_runtime_cannot_count_surfaces_the_error_without_ma wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url, _NOVA) + model: Final = _fallback_deployment(scenario, runtime.url, _NOVA) response: Final = _count(gateway, _bare(model)) assert response.status_code == 400, response.text assert _UNSUPPORTED in response.text and "input_tokens" not in response.text, response.text @@ -697,7 +770,7 @@ def test_runtime_403_on_a_claude_model_never_reaches_mantle(counting_proxy: Owne wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) body: Final[dict[str, JsonValue]] = { "model": model, "messages": [{"role": "user", "content": "status=403 Count this message"}], @@ -712,7 +785,7 @@ def test_runtime_403_on_a_claude_model_never_reaches_mantle(counting_proxy: Owne def test_unreachable_runtime_surfaces_the_error_without_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: gateway: Final = counting_proxy.gateway with wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario: - model: Final = _deployment(scenario, _closed_port_url()) + model: Final = _fallback_deployment(scenario, _closed_port_url()) response: Final = _count(gateway, _full(model)) assert 500 <= response.status_code < 600, response.text assert "input_tokens" not in response.text, response.text @@ -728,7 +801,7 @@ def test_bedrock_passthrough_count_tokens_still_answers_the_runtime_rejection( wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = gateway.request("POST", "/bedrock/v1/messages/count_tokens", _full(model)) assert response.status_code == 400, response.text assert _UNSUPPORTED in response.text and "input_tokens" not in response.text, response.text @@ -745,7 +818,7 @@ def test_responses_input_tokens_with_instructions_counts_through_mantle( wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = gateway.request( "POST", "/v1/responses/input_tokens", @@ -770,7 +843,7 @@ def test_messages_count_tokens_lifts_the_leading_system_run_for_mantle( wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, count_request(model, case)) assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 @@ -786,7 +859,7 @@ def test_anthropic_sdk_count_tokens_lifts_the_leading_system_through_mantle( wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) counted: Final = _anthropic_client(gateway).messages.count_tokens(model=model, messages=_LIFT_SDK_MESSAGES) assert counted.input_tokens == _MANTLE_COUNT, counted assert len(_runtime_count_targets(runtime)) == 1 @@ -802,7 +875,7 @@ def test_async_anthropic_sdk_count_tokens_lifts_the_leading_system_through_mantl wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) counted: Final = asyncio.run( _async_anthropic_client(gateway).messages.count_tokens(model=model, messages=_LIFT_SDK_MESSAGES) ) @@ -820,7 +893,7 @@ def test_utils_token_counter_call_endpoint_counts_a_leading_system_through_mantl wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = gateway.request( "POST", "/utils/token_counter", count_request(model, STRING_CASE), params={"call_endpoint": "true"} ) @@ -841,7 +914,7 @@ def test_openai_sdk_input_tokens_lifts_instructions_through_mantle( wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) counted: Final = openai_client(gateway).responses.input_tokens.count( model=model, input=USER_TEXT, instructions=INSTRUCTION ) @@ -859,7 +932,7 @@ def test_async_openai_sdk_input_tokens_lifts_instructions_through_mantle( wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) counted: Final = asyncio.run( async_openai_client(gateway).responses.input_tokens.count( model=model, input=USER_TEXT, instructions=INSTRUCTION @@ -879,7 +952,7 @@ def test_responses_input_tokens_lifts_instructions_ahead_of_a_leading_system_ite wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = gateway.request( "POST", "/v1/responses/input_tokens", @@ -903,7 +976,7 @@ def test_messages_count_tokens_forwards_tools_beside_the_lifted_system_to_mantle wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, {**count_request(model, STRING_CASE), "tools": _TOOLS}) assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 @@ -919,7 +992,7 @@ def test_messages_count_tokens_keeps_a_mid_conversation_system_in_place_for_mant wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) body: Final[dict[str, JsonValue]] = { "model": model, "messages": [LEADING, USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP], @@ -941,7 +1014,7 @@ def test_messages_count_tokens_surfaces_the_mantle_error_when_every_message_is_s wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) body: Final[dict[str, JsonValue]] = {"model": model, "messages": [LEADING]} response: Final = _count(gateway, body) assert response.status_code == 400, response.text @@ -959,7 +1032,7 @@ def test_messages_count_tokens_surfaces_the_mantle_error_beside_a_non_text_syste wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, {**count_request(model, STRING_CASE), "system": 5}) assert response.status_code == 400, response.text assert "input_tokens" not in response.text, response.text @@ -976,7 +1049,7 @@ def test_messages_count_tokens_repeated_request_lifts_the_leading_system_each_ti wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) answers: Final = tuple(_payload(_count(gateway, count_request(model, STRING_CASE))) for _ in range(2)) assert answers == ({"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"},) * 2 assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",) * 2 @@ -992,7 +1065,7 @@ def test_messages_count_tokens_surfaces_the_counting_error_for_a_leading_system_ wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) body: Final[dict[str, JsonValue]] = {"model": model, "messages": [{"role": "system"}, USER]} response: Final = _count(gateway, body) assert 500 <= response.status_code < 600, response.text @@ -1015,7 +1088,7 @@ def test_messages_count_tokens_duplicate_messages_key_lifts_the_last_value_for_m wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) first: Final = json.dumps([USER]) last: Final = json.dumps([LEADING, USER]) response: Final = gateway.client.post( @@ -1064,7 +1137,7 @@ def test_mantle_outage_between_concurrent_lifted_waves_errors_then_recovers( pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) runtime: Final = stack.enter_context(wire_server(_runtime)) scenario: Final = stack.enter_context(gateway.scenario()) - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) body: Final = count_request(model, STRING_CASE) def count(client: httpx.Client) -> tuple[int, JsonValue]: @@ -1092,7 +1165,7 @@ def test_messages_count_tokens_surfaces_the_mantle_status_when_mantle_errors( wire_server(_mantle(_rejecting(status)), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, _full(model)) assert response.status_code == status, response.text assert "input_tokens" not in response.text, response.text @@ -1114,7 +1187,7 @@ def test_messages_count_tokens_errors_when_mantle_answers_without_input_tokens( wire_server(_mantle(lambda _request: Reply(body=body)), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, _bare(model)) assert 500 <= response.status_code < 600, response.text assert len(_runtime_count_targets(runtime)) == 1 @@ -1124,16 +1197,14 @@ def test_messages_count_tokens_errors_when_mantle_answers_without_input_tokens( def test_messages_count_tokens_errors_when_mantle_is_unreachable(counting_proxy: OwnedProxy) -> None: gateway: Final = counting_proxy.gateway with wire_server(_runtime) as runtime, gateway.scenario() as scenario: - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, _full(model)) assert 500 <= response.status_code < 600, response.text assert "input_tokens" not in response.text, response.text assert len(_runtime_count_targets(runtime)) == 1 -def test_messages_count_tokens_errors_when_mantle_never_answers( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: +def test_messages_count_tokens_errors_when_mantle_never_answers(counting_proxy: OwnedProxy, mantle_port: int) -> None: gateway: Final = counting_proxy.gateway held: Final[SimpleQueue[str]] = SimpleQueue() release: Final = threading.Event() @@ -1150,7 +1221,7 @@ def test_messages_count_tokens_errors_when_mantle_never_answers( mantle: Final = stack.enter_context(wire_server(_mantle(hold), port=mantle_port)) stack.callback(release.set) scenario: Final = stack.enter_context(gateway.scenario()) - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) body: Final = _full(model) future: Final = pool.submit(_counted_on, client, gateway.key, body) eventually(held.qsize, lambda size: size == 1, seconds=30) @@ -1172,7 +1243,7 @@ def test_messages_count_tokens_surfaces_the_mantle_error_for_a_non_text_system( wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, {**_bare(model), "system": 5}) assert response.status_code == 400, response.text assert "input_tokens" not in response.text, response.text @@ -1194,7 +1265,7 @@ def test_messages_count_tokens_rejects_malformed_tools_without_calling_either_pe wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) refused: Final = _count(gateway, {**_bare(model), "tools": tools}) assert refused.status_code == 400, refused.text assert object_value(_JSON_OBJECT.validate_json(refused.content)["error"])["type"] == "invalid_request_error" @@ -1213,7 +1284,7 @@ def test_messages_count_tokens_without_messages_is_rejected( wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, {"model": model, "system": _SYSTEM, "tools": _TOOLS, **fields}) assert response.status_code == 400, response.text assert "messages parameter is required" in response.text, response.text @@ -1234,7 +1305,7 @@ def test_messages_count_tokens_unauthenticated_request_never_reaches_either_peer wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = gateway.request( "POST", "/v1/messages/count_tokens", {**_full(model), **fields}, key="sk-not-a-key-this-proxy-issued" ) @@ -1253,7 +1324,7 @@ def test_messages_count_tokens_forwards_a_5kb_system_verbatim_to_mantle( wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, {**_bare(model), "system": system}) assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 @@ -1269,7 +1340,7 @@ def test_messages_count_tokens_duplicate_system_and_tools_keys_forward_one_value wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) fields: Final = f'"system": {json.dumps(_SYSTEM)}, "tools": {json.dumps(_TOOLS)}' response: Final = gateway.client.post( "/v1/messages/count_tokens", @@ -1292,7 +1363,7 @@ def test_messages_count_tokens_leaves_empty_tools_and_system_out_of_the_mantle_b wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, {**_bare(model), "tools": [], "system": ""}) assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 @@ -1308,7 +1379,7 @@ def test_messages_count_tokens_repeated_request_reaches_both_peers_each_time( wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) answers: Final = tuple(_payload(_count(gateway, _bare(model))) for _ in range(2)) assert answers == ({"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"},) * 2 assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",) * 2 @@ -1364,7 +1435,7 @@ def test_mantle_outage_between_concurrent_waves_errors_then_recovers( pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) runtime: Final = stack.enter_context(wire_server(_runtime)) scenario: Final = stack.enter_context(gateway.scenario()) - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) body: Final = _full(model) def generate_then_count(client: httpx.Client) -> tuple[int, int, JsonValue]: @@ -1398,7 +1469,7 @@ def test_slow_mantle_holds_concurrent_counts_without_stalling_the_proxy( scenario: Final = stack.enter_context(gateway.scenario()) pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) stack.callback(release.set) - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) futures: Final = tuple( pool.submit(_generated_then_counted, client, gateway.key, model, _full(model)) for client in clients ) @@ -1495,7 +1566,7 @@ def test_chat_completions_on_the_same_deployment_still_generate( wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = gateway.request( "POST", "/v1/chat/completions", @@ -1523,7 +1594,7 @@ def test_messages_endpoint_on_the_same_deployment_still_generates(counting_proxy wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = gateway.request( "POST", "/v1/messages", @@ -1546,7 +1617,7 @@ def test_responses_endpoint_on_the_same_deployment_still_generates( wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = gateway.request("POST", "/v1/responses", {"model": model, "input": f"marker-{marker}"}) assert response.status_code == 200, response.text assert answer(marker) in response.text, response.text @@ -1562,6 +1633,6 @@ def test_utils_token_counter_local_mode_never_calls_either_peer(counting_proxy: wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) assert _local_count(gateway, _full(model)) > 0 assert runtime.drain() == () and mantle.drain() == () diff --git a/tests/integration/providers/test_count_tokens_model_wire.py b/tests/integration/providers/test_count_tokens_model_wire.py index 46e5396ef64..b6ace7ddef3 100644 --- a/tests/integration/providers/test_count_tokens_model_wire.py +++ b/tests/integration/providers/test_count_tokens_model_wire.py @@ -31,7 +31,6 @@ _PROFILE_ID: Final = "lit9563count" _PROFILE_ARN: Final = f"arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/{_PROFILE_ID}" _COST_MAP_PATH: Final = Path(__file__).resolve().parents[3] / "model_prices_and_context_window.json" _REGION: Final = "us-west-2" -_COUNT_REGION: Final = "us-east-1" _ACCESS_KEY: Final = "AKIAINTEGRATIONCOUNT" _SECRET_KEY: Final = "integration-count-secret" _PROVIDER_COUNT: Final = 1345 @@ -39,12 +38,37 @@ _COUNT_MODEL_COUNT: Final = 2718 _FRAMING: Final = 50 _UNSUPPORTED: Final = "The provided model doesn't support counting tokens." _COUNT_TARGET: Final = re.compile(r"^/model/(.+)/count-tokens$") +_MANTLE_COUNT_TARGET: Final = "/anthropic/v1/messages/count_tokens" +_MANTLE_VERSION: Final = "2023-06-01" _SCRIPTED_STATUS: Final = re.compile(r"status=(\d{3})") _SCRIPTED_COUNT_MODEL_STATUS: Final = re.compile(r"count_model=(\d{3})") _CREDENTIAL: Final = re.compile(r"Credential=([^/]+)/\d+/([^/]+)/") _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _COST_MAP: Final = TypeAdapter(dict[str, dict[str, JsonValue]]) + +def _shipped_token_count_model(model: str) -> dict[str, JsonValue]: + token_count_model: Final = _COST_MAP.validate_json(_COST_MAP_PATH.read_bytes())[model].get("token_count_model") + assert isinstance(token_count_model, dict), model + return token_count_model + + +def _shipped_token_count_model_field(model: str, field: str) -> str: + value: Final = _shipped_token_count_model(model).get(field) + assert isinstance(value, str), (model, field, value) + return value + + +def _count_model_source(model: str) -> dict[str, JsonValue]: + return { + "litellm_count_source": "token_count_model", + "litellm_count_model": _shipped_token_count_model_field(model, "model"), + } + + +_COUNT_REGION: Final = _shipped_token_count_model_field(_SONNET, "region") +_COUNT_MODEL_SOURCE: Final = _count_model_source(_SONNET) + _SYSTEM: Final = "You are a terse assistant that answers in one sentence" _WEATHER_SCHEMA: Final[dict[str, JsonValue]] = { "type": "object", @@ -62,7 +86,6 @@ _SDK_TOOLS: Final[list[MessageCountTokensToolParam]] = [ }, } ] -_COUNT_MODEL_SOURCE: Final = {"litellm_count_source": "token_count_model", "litellm_count_model": _OPUS} _OPENAI_PRIMARY: Final = "gpt-count-primary" _OPENAI_COUNT_MODEL: Final = "gpt-count-model" _OPENAI_COUNT_TARGET: Final = "/v1/responses/input_tokens" @@ -140,8 +163,17 @@ def _runtime(request: Request) -> Reply: def _mantle(request: Request) -> Reply: model: Final = _JSON_OBJECT.validate_json(request.body).get("model") + status: Final = _scripted(_SCRIPTED_COUNT_MODEL_STATUS, request, 200) + if status == 200: + return _json_reply(200, {"input_tokens": _COUNT_MODEL_COUNT}) + error_type: Final = { + 401: "authentication_error", + 404: "not_found_error", + 429: "rate_limit_error", + }[status] return _json_reply( - 404, {"type": "error", "error": {"type": "not_found_error", "message": f"scripted mantle 404 for {model}"}} + status, + {"type": "error", "error": {"type": error_type, "message": f"scripted mantle {status} for {model}"}}, ) @@ -165,6 +197,24 @@ def _mantle_models(wire: Wire) -> tuple[JsonValue, ...]: return tuple(_JSON_OBJECT.validate_json(request.body).get("model") for request in wire.drain()) +def _mantle_count_requests(wire: Wire) -> tuple[tuple[tuple[str, str], dict[str, JsonValue]], ...]: + requests: Final = tuple(wire.drain()) + for request in requests: + assert (request.method, request.target) == ("POST", _MANTLE_COUNT_TARGET), request.target + assert request.headers["anthropic-version"] == _MANTLE_VERSION, request.headers + assert request.headers["content-type"] == "application/json", request.headers + return tuple((_signer(request), _JSON_OBJECT.validate_json(request.body)) for request in requests) + + +def _mantle_body(prompt: str) -> dict[str, JsonValue]: + return { + "model": _OPUS_BASE, + "messages": [{"role": "user", "content": prompt}], + "system": _SYSTEM, + "tools": _SDK_TOOLS, + } + + def _url(port: int) -> str: return f"http://127.0.0.1:{port}" @@ -255,15 +305,15 @@ def test_shipped_token_count_model_counts_the_request_in_its_region_without_tryi ): counted: Final = _count(gateway, _deployment(scenario, runtime_port), 404) assert (counted.input_tokens, counted.model_extra) == (_COUNT_MODEL_COUNT, _COUNT_MODEL_SOURCE), counted - assert _runtime_counts(runtime) == ((_OPUS_BASE, (_ACCESS_KEY, _COUNT_REGION), _converse_body(_prompt(404))),) - assert _mantle_models(mantle) == () + assert _runtime_counts(runtime) == () + assert _mantle_count_requests(mantle) == (((_ACCESS_KEY, _COUNT_REGION), _mantle_body(_prompt(404))),) @pytest.mark.parametrize( ("count_model_status", "error", "message"), [ - (401, anthropic.AuthenticationError, f"scripted runtime 401 for {_OPUS_BASE}"), - (429, anthropic.RateLimitError, f"scripted runtime 429 for {_OPUS_BASE}"), + (401, anthropic.AuthenticationError, f"scripted mantle 401 for {_OPUS_BASE}"), + (429, anthropic.RateLimitError, f"scripted mantle 429 for {_OPUS_BASE}"), (404, anthropic.NotFoundError, f"scripted mantle 404 for {_OPUS_BASE}"), ], ) @@ -283,9 +333,16 @@ def test_token_count_model_failure_answers_its_error_and_names_the_region_to_ena ): with pytest.raises(error) as refused: _count(gateway, _deployment(scenario, runtime_port), 200, count_model_status=count_model_status) - assert message in str(refused.value) and f"{_OPUS} in {_COUNT_REGION}" in str(refused.value), refused.value - assert _runtime_models(runtime) == (_OPUS_BASE,) - assert _mantle_models(mantle) == ((_OPUS_BASE,) if count_model_status == 404 else ()) + assert message in str(refused.value) and ( + f"{_COUNT_MODEL_SOURCE['litellm_count_model']} in {_COUNT_REGION}" in str(refused.value) + ), refused.value + assert _runtime_counts(runtime) == () + assert _mantle_count_requests(mantle) == ( + ( + (_ACCESS_KEY, _COUNT_REGION), + _mantle_body(_prompt(200, count_model_status)), + ), + ) def test_framing_is_added_to_the_local_estimate_when_the_runtime_cannot_count( @@ -359,8 +416,8 @@ def test_inference_profile_arn_counts_with_the_token_count_model_of_its_base_mod ) counted: Final = _count(gateway, model, 404) assert (counted.input_tokens, counted.model_extra) == (_COUNT_MODEL_COUNT, _COUNT_MODEL_SOURCE), counted - assert _runtime_models(runtime) == (_OPUS_BASE,) - assert _mantle_models(mantle) == () + assert _runtime_counts(runtime) == () + assert _mantle_count_requests(mantle) == (((_ACCESS_KEY, _COUNT_REGION), _mantle_body(_prompt(404))),) @pytest.mark.parametrize( From 963333d654d3622f74b0bff0ecb41e9d62397135 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 10:58:42 +0000 Subject: [PATCH 02/33] fix(proxy)!: let proxy admin keys read raw cloud storage file ids on GET /files/{id}/content (#40621) * fix(proxy): let proxy admin keys read raw cloud storage file ids on GET /files/{id}/content Since #30595 the endpoint answered 400 for every raw gs:// or s3:// id, which broke the documented unmanaged Vertex batch flow at its last step, while a double-encoded id slipped past the guard entirely. Raw ids are now readable by proxy admin keys only, every other key gets 403, and the guard decodes the id fully before deciding, so encoding depth no longer matters * fix(proxy): decode nested percent-encoded file ids iteratively An id percent-encoded a thousand or more times sent the recursive decoder past the interpreter's recursion limit, so a proxy admin key got a 500 instead of a verdict on the id. The decoder now walks the encodings in a loop, so any depth resolves * fix(proxy): cap the decode passes the raw cloud storage file id guard makes Any authenticated key could send a file id wrapped in tens of thousands of percent-encoding layers and the guard would unwrap every one of them, each pass walking the whole id, so a 120 KB request line held a worker for about five seconds and a 2 MB one for minutes. The guard now unwraps at most eight layers, and an id still changing after that is treated as a raw cloud storage id, so it stays admin-only instead of being decoded any further * refactor(proxy): type the decode step the raw cloud storage file id guard folds over * test(proxy): return a streaming result from the admin raw id test mock now that vertex_ai file content streams * test(proxy): send raw cloud ids to the content route the way uvicorn delivers them Starlette's TestClient decodes the path twice, so the twice-encoded cases reached the route fully decoded and never hit the decode loop. The cases now name the id the route receives, and the helper tests take a typed quote helper and Final constants. * test(proxy): mark the raw cloud id test locals Final --- .../cloud_storage_security.py | 23 +++- .../openai_files_endpoints/files_endpoints.py | 10 +- .../test_cloud_storage_security.py | 42 +++++++ .../test_files_endpoint.py | 118 +++++++++++++++--- 4 files changed, 170 insertions(+), 23 deletions(-) diff --git a/litellm/litellm_core_utils/cloud_storage_security.py b/litellm/litellm_core_utils/cloud_storage_security.py index c913e4de02e..e9689bd7d8a 100644 --- a/litellm/litellm_core_utils/cloud_storage_security.py +++ b/litellm/litellm_core_utils/cloud_storage_security.py @@ -1,6 +1,7 @@ import posixpath import re from collections.abc import Mapping, Sequence +from itertools import accumulate, repeat from types import MappingProxyType from typing import Any, Final, cast from urllib.parse import quote, unquote @@ -20,15 +21,33 @@ MANAGED_CLOUD_STORAGE_SCHEMES: Final = ("s3://", "gs://") _MAPPING_PROXY_TYPE: Final[type] = type(MappingProxyType({})) +MAX_FILE_ID_DECODE_PASSES: Final = 8 + + +def _unquote_once(current: str, _: str) -> str: + return unquote(current) + + +def _decodings(value: str) -> tuple[str, ...]: + return tuple(accumulate(repeat(value, MAX_FILE_ID_DECODE_PASSES + 2), _unquote_once)) + + def is_managed_cloud_storage_uri(file_id: str) -> bool: """ - True if file_id is a raw cloud-storage object URI (e.g. ``s3://bucket/key``). + True if file_id is a raw cloud-storage object URI (e.g. ``s3://bucket/key``), + up to ``MAX_FILE_ID_DECODE_PASSES`` layers of percent-encoding deep, or an id + encoded deeper than that, which no client produces and which is treated as raw + rather than decoded any further. These are internal provider artifacts. On the multi-tenant proxy they must be retrieved through their managed unified file id so owner/team access is enforced; a raw URI supplied by a caller bypasses that check. """ - return isinstance(file_id, str) and file_id.startswith(MANAGED_CLOUD_STORAGE_SCHEMES) + if not isinstance(file_id, str): + return False + decodings: Final = _decodings(file_id) + settled: Final = decodings[-1] == decodings[-2] + return decodings[-1].startswith(MANAGED_CLOUD_STORAGE_SCHEMES) or not settled _SAFE_OBJECT_COMPONENT_PATTERN: Final = re.compile(r"[^A-Za-z0-9._-]+") diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index bdc71908c7c..8823f01d781 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -1102,14 +1102,10 @@ async def get_file_content( } ) else: - # A raw cloud-storage URI (s3://, gs://) supplied here would skip the - # managed-file owner/team check that only runs for unified ids, letting - # a caller read another tenant's object by its key. Such objects are only - # reachable through their managed unified id. - if is_managed_cloud_storage_uri(file_id): + if is_managed_cloud_storage_uri(file_id) and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException( - status_code=400, - detail="Raw cloud storage file ids cannot be retrieved directly. Use the LiteLLM managed file id returned when the file was created.", + status_code=403, + detail="Raw cloud storage file ids can only be retrieved by a proxy admin key. Use the LiteLLM managed file id returned when the file was created.", ) # Check for model-based credential routing ( diff --git a/tests/unit/litellm_core_utils/test_cloud_storage_security.py b/tests/unit/litellm_core_utils/test_cloud_storage_security.py index c3a2511a263..46131aa64e5 100644 --- a/tests/unit/litellm_core_utils/test_cloud_storage_security.py +++ b/tests/unit/litellm_core_utils/test_cloud_storage_security.py @@ -1,4 +1,11 @@ +from functools import reduce +from typing import Final +from urllib.parse import quote + +import pytest + from litellm.litellm_core_utils.cloud_storage_security import ( + MAX_FILE_ID_DECODE_PASSES, is_managed_cloud_storage_uri, ) @@ -13,3 +20,38 @@ def test_is_managed_cloud_storage_uri_ignores_provider_and_unified_ids(): assert not is_managed_cloud_storage_uri("file-abc123") assert not is_managed_cloud_storage_uri("bGl0ZWxsbV9wcm94eQ==") assert not is_managed_cloud_storage_uri("") + + +@pytest.mark.parametrize( + "file_id", + ( + "gs%3A%2F%2Fbucket%2Flitellm-vertex-files%2Fprediction-model%2Fpredictions.jsonl", + "gs%253A%252F%252Fbucket%252Flitellm-vertex-files%252Fprediction-model%252Fpredictions.jsonl", + "s3%3A%2F%2Fbucket%2Flitellm-batch-outputs%2Fx.jsonl.out", + "s3%253A%252F%252Fbucket%252Flitellm-batch-outputs%252Fx.jsonl.out", + ), +) +def test_is_managed_cloud_storage_uri_sees_through_percent_encoding(file_id: str): + assert is_managed_cloud_storage_uri(file_id) + + +def _quote_once(current: str, _: int) -> str: + return quote(current, safe="") + + +def _quoted_times(value: str, times: int) -> str: + return reduce(_quote_once, range(times), value) + + +def test_is_managed_cloud_storage_uri_decodes_up_to_the_pass_cap(): + assert is_managed_cloud_storage_uri(_quoted_times("gs://bucket/x", MAX_FILE_ID_DECODE_PASSES)) + assert not is_managed_cloud_storage_uri(_quoted_times("file-abc/x", MAX_FILE_ID_DECODE_PASSES)) + + +def test_is_managed_cloud_storage_uri_fails_closed_on_an_id_encoded_past_the_pass_cap(): + nested_gs_id: Final = "gs" + "%" + "25" * 5000 + "3A//bucket/x" + nested_plain_id: Final = "%" + "25" * 5000 + + assert is_managed_cloud_storage_uri(_quoted_times("file-abc/x", MAX_FILE_ID_DECODE_PASSES + 1)) + assert is_managed_cloud_storage_uri(nested_gs_id) + assert is_managed_cloud_storage_uri(nested_plain_id) diff --git a/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py index 793921255ec..eb222366812 100644 --- a/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py @@ -1,5 +1,6 @@ import json from typing import Final, List, Optional +from urllib.parse import quote from unittest.mock import ANY, AsyncMock, MagicMock import pytest @@ -12,7 +13,7 @@ from pytest_mock import MockerFixture import litellm from litellm import Router from litellm.files.types import FileContentStreamingResult -from litellm.proxy._types import LiteLLM_UserTableFiltered, UserAPIKeyAuth +from litellm.proxy._types import LiteLLM_UserTableFiltered, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.hooks import get_proxy_hook from litellm.proxy.management_endpoints.internal_user_endpoints import ui_view_users from litellm.proxy.openai_files_endpoints.file_content_streaming_handler import ( @@ -235,23 +236,112 @@ def test_invalid_purpose(mocker: MockerFixture, monkeypatch, llm_router: Router) assert error["param"] == "purpose" -def test_get_file_content_rejects_raw_cloud_storage_uri(llm_router: Router): - """A raw s3:// file id must be rejected on the proxy content endpoint. +RAW_GCS_OUTPUT_FILE_ID: Final = ( + "gs://my-gcs-bucket/litellm-vertex-files/publishers/google/models/gemini-3.8-flash/" + "prediction-model-2026-09-10T20:27:18.178864Z/predictions.jsonl" +) +RAW_S3_OUTPUT_FILE_ID: Final = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out" - Such an id is not a managed unified id, so it would otherwise skip the - owner/team access check and let a caller read another tenant's batch output - object by its key. Callers must use the managed unified file id. - """ - from urllib.parse import quote - s3_file_id = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out" - response = client.get( - f"/v1/files/{quote(s3_file_id, safe='')}/content?provider=bedrock", - headers={"Authorization": "Bearer test-key"}, +def _ids_a_server_hands_the_route_for(raw_id: str, provider: str) -> tuple[tuple[str, str], ...]: + encoded_once: Final = quote(raw_id, safe="") + return ((raw_id, provider), (encoded_once, provider), (quote(encoded_once, safe=""), provider)) + + +RAW_CLOUD_FILE_IDS_AS_THE_ROUTE_RECEIVES_THEM: Final = _ids_a_server_hands_the_route_for( + RAW_GCS_OUTPUT_FILE_ID, "vertex_ai" +) + _ids_a_server_hands_the_route_for(RAW_S3_OUTPUT_FILE_ID, "bedrock") + + +def _testclient_path_segment_delivering(route_file_id: str) -> str: + return quote(quote(route_file_id, safe=""), safe="") + + +def _override_auth( + monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, llm_router: Router, user_role: LitellmUserRoles +) -> ProxyLogging: + import litellm.proxy.proxy_server as ps + + proxy_logging_obj: Final = setup_proxy_logging_object(monkeypatch, llm_router) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + proxy_logging_obj.update_request_status = mocker.AsyncMock() + proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock() + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="test-key", user_role=user_role, user_id="test-user" ) + return proxy_logging_obj - assert response.status_code == 400 - assert "managed file id" in response.json()["error"]["message"].lower() + +@pytest.mark.parametrize(("route_file_id", "provider"), RAW_CLOUD_FILE_IDS_AS_THE_ROUTE_RECEIVES_THEM) +def test_get_file_content_answers_403_for_a_raw_cloud_id_from_a_non_admin_key( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + llm_router: Router, + route_file_id: str, + provider: str, +): + import litellm.proxy.proxy_server as ps + + _override_auth(monkeypatch, mocker, llm_router, LitellmUserRoles.INTERNAL_USER) + afile_content: Final = mocker.AsyncMock() + monkeypatch.setattr(litellm, "afile_content", afile_content) + + try: + response: Final = client.get( + f"/v1/files/{_testclient_path_segment_delivering(route_file_id)}/content?provider={provider}", + headers={"Authorization": "Bearer test-key"}, + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 403, response.text + assert "proxy admin" in response.json()["error"]["message"] + afile_content.assert_not_called() + + +@pytest.mark.parametrize(("route_file_id", "provider"), RAW_CLOUD_FILE_IDS_AS_THE_ROUTE_RECEIVES_THEM) +def test_get_file_content_forwards_a_raw_cloud_id_from_a_proxy_admin_key( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + llm_router: Router, + route_file_id: str, + provider: str, +): + import litellm.proxy.proxy_server as ps + + _override_auth(monkeypatch, mocker, llm_router, LitellmUserRoles.PROXY_ADMIN) + content: Final = b'{"custom_id": "request-1"}\n' + + async def _stream(): + yield content + + async def _provider_file_content(**kwargs: object) -> FileContentStreamingResult | HttpxBinaryResponseContent: + if kwargs.get("stream"): + return FileContentStreamingResult(stream_iterator=_stream(), headers={"content-length": str(len(content))}) + return HttpxBinaryResponseContent( + response=httpx.Response( + status_code=200, content=content, headers={"content-type": "application/octet-stream"} + ) + ) + + afile_content: Final = mocker.AsyncMock(side_effect=_provider_file_content) + monkeypatch.setattr(litellm, "afile_content", afile_content) + + try: + response: Final = client.get( + f"/v1/files/{_testclient_path_segment_delivering(route_file_id)}/content?provider={provider}", + headers={"Authorization": "Bearer test-key"}, + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 200, response.text + assert response.content == content + afile_content.assert_called_once() + assert afile_content.call_args.kwargs["custom_llm_provider"] == provider + assert afile_content.call_args.kwargs["file_id"] == route_file_id def test_mock_create_audio_file(mocker: MockerFixture, monkeypatch, llm_router: Router): From 808333682ed2f0a5fff5277d862fd224935accc5 Mon Sep 17 00:00:00 2001 From: Vincent Gimenes <147169146+VincentG1234@users.noreply.github.com> Date: Tue, 20 Oct 2026 13:35:18 +0200 Subject: [PATCH 03/33] feat(embeddings): enforce per-request input batch limit (#44574) * feat(embeddings): add max_batch_inputs pre-call validation Signed-off-by: Vincent Gimenes * add unit tests for max embedding batch size Signed-off-by: Vincent Gimenes * feat(embeddings): enforce batch limit on extra_body input override Signed-off-by: Vincent Gimenes * chore(dashboard): update schema.d.ts for embedding batch size model field Signed-off-by: Vincent Gimenes * fix(embeddings): close Bedrock batch bypass and satisfy ruff PT012 Count max(top-level input, extra_body.input) for batch limits so providers that ignore extra_body cannot bypass max_embedding_batch_size. Add router tests and refactor pytest.raises blocks for ruff-tests. Signed-off-by: Vincent Gimenes * fix(embeddings): drop unchecked casts in router batch-size check Signed-off-by: Vincent Gimenes * fix(embeddings): route oversized batches to a group deployment that admits them * fix(embeddings): count extra_body input for any top-level input and prefer admitting deployments before cooldown * fix(embeddings): log cap rejections to failure callbacks and scope batch routing to embedding calls * fix(embeddings): apply the batch cap after tag and pre-call filters and fall back on cap rejections * fix(router): keep embedding cap rejections out of deployment health and hold only for routable cooling deployments * fix(router): clear the embedding cap rejection mark before a fallback's provider call A fallback or retry shares the request's logging object, so the mark a rejected attempt set made telemetry skip the fallback deployment's real call. Unit tests now pass a fixed logging start time instead of reading the clock. * fix(router): log embedding batch rejections once, before a fallback runs * refactor(router): drop docstrings that restate the code --------- Signed-off-by: Vincent Gimenes Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/router.py | 116 +- .../embedding_batch_size_check.py | 199 ++++ litellm/types/router.py | 2 + .../test_embedding_batch_size_check.py | 1024 +++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 + 5 files changed, 1327 insertions(+), 16 deletions(-) create mode 100644 litellm/router_utils/embedding_batch_size_check.py create mode 100644 tests/unit/router_utils/test_embedding_batch_size_check.py diff --git a/litellm/router.py b/litellm/router.py index aadbdaa8520..ac0933e158b 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -242,6 +242,15 @@ from litellm.router_utils.cooldown_handlers import ( is_video_lookup_not_found, set_cooldown_deployments, ) +from litellm.router_utils.embedding_batch_size_check import ( + async_log_rejection_before_provider_call, + clear_rejected_before_provider_call, + deployments_admitting_embedding_batch, + deployments_when_none_admits_embedding_batch, + embedding_batch_too_large_error, + embedding_deployment_selection, + log_rejection_before_provider_call, +) from litellm.router_utils.fallback_event_handlers import ( MID_STREAM_FALLBACK_CONTROLS_KEY, AttemptedFallbackTargets, @@ -6861,19 +6870,33 @@ class Router: except Exception as e: raise e - def _embedding(self, input: str | list, model: str, **kwargs): + def _embedding(self, input: str | list[object], model: str, **kwargs): model_name = None try: verbose_router_logger.debug("Inside embedding()- model: %s; kwargs: %s", model, kwargs) - deployment: Final = self.get_available_deployment( - model=model, - input=input, - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) + with embedding_deployment_selection(): + deployment: Final = self.get_available_deployment( + model=model, + input=input, + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + ) self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) data: Final = deployment["litellm_params"].copy() model_name = data["model"] + batch_error: Final = embedding_batch_too_large_error( + deployment=_OPAQUE_VALUE_ADAPTER.validate_python(deployment), + input=input, + request_kwargs=_OPAQUE_VALUE_ADAPTER.validate_python(kwargs), + model=model, + ) + logging_obj: Final = _REQUEST_KWARGS_ADAPTER.validate_python(kwargs).get("litellm_logging_obj") + if batch_error is not None: + if isinstance(logging_obj, LiteLLMLogging): + log_rejection_before_provider_call(logging_obj, batch_error) + raise batch_error + if isinstance(logging_obj, LiteLLMLogging): + clear_rejected_before_provider_call(logging_obj) potential_model_client: Final = self._get_client(deployment=deployment, kwargs=kwargs, client_type="sync") # check if provided keys == client keys # @@ -6936,20 +6959,34 @@ class Router: ) raise e - async def _aembedding(self, input: str | list, model: str, **kwargs): + async def _aembedding(self, input: str | list[object], model: str, **kwargs): model_name = None try: verbose_router_logger.debug("Inside _aembedding()- model: %s; kwargs: %s", model, kwargs) parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - deployment: Final = await self.async_get_available_deployment( - model=model, - input=input, - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) + with embedding_deployment_selection(): + deployment: Final = await self.async_get_available_deployment( + model=model, + input=input, + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + ) self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) data: Final = deployment["litellm_params"].copy() model_name = data["model"] + batch_error: Final = embedding_batch_too_large_error( + deployment=_OPAQUE_VALUE_ADAPTER.validate_python(deployment), + input=input, + request_kwargs=_OPAQUE_VALUE_ADAPTER.validate_python(kwargs), + model=model, + ) + logging_obj: Final = _REQUEST_KWARGS_ADAPTER.validate_python(kwargs).get("litellm_logging_obj") + if batch_error is not None: + if isinstance(logging_obj, LiteLLMLogging): + await async_log_rejection_before_provider_call(logging_obj, batch_error) + raise batch_error + if isinstance(logging_obj, LiteLLMLogging): + clear_rejected_before_provider_call(logging_obj) model_client: Final = self._get_async_openai_model_client( deployment=deployment, kwargs=kwargs, @@ -14680,10 +14717,34 @@ class Router: batch_candidates: Final = cast( # cast-ok: router deployments are plain dicts list[dict[str, object]], healthy_deployments ) + embedding_batch_admitting: Final = deployments_admitting_embedding_batch( + batch_candidates, + _OPAQUE_VALUE_ADAPTER.validate_python(input), + _OPAQUE_VALUE_ADAPTER.validate_python(request_kwargs), + ) + embedding_batch_routed: Final = ( + embedding_batch_admitting + if embedding_batch_admitting is not None + else deployments_when_none_admits_embedding_batch( + batch_candidates, + cast(list[dict[str, object]], _pre_cooldown_deployments), # cast-ok: router deployments are plain dicts + await self._async_routable_cooled_deployment_ids( + model=model, + request_deployments=request_deployments, + cooldown_deployments=cooldown_deployments, + request_kwargs=cast(Mapping[str, object], request_kwargs), # cast-ok: str-keyed request kwargs + has_request_input=True, + ), + _OPAQUE_VALUE_ADAPTER.validate_python(input), + _OPAQUE_VALUE_ADAPTER.validate_python(request_kwargs), + ) + ) healthy_deployments = litellm.utils.get_order_filtered_deployments( list( self._narrow_to_video_deployment( - without_deployments_that_cannot_batch(batch_candidates, caller_deployments=lambda: batch_candidates) + without_deployments_that_cannot_batch( + embedding_batch_routed, caller_deployments=lambda: batch_candidates + ) ) ), target_order=_target_order, @@ -15890,11 +15951,34 @@ class Router: batch_candidates: Final = cast( # cast-ok: router deployments are plain dicts list[dict[str, object]], healthy_deployments ) + embedding_batch_admitting: Final = deployments_admitting_embedding_batch( + batch_candidates, + _OPAQUE_VALUE_ADAPTER.validate_python(input), + routing_kwargs, + ) + embedding_batch_routed: Final = ( + embedding_batch_admitting + if embedding_batch_admitting is not None + else deployments_when_none_admits_embedding_batch( + batch_candidates, + cast(list[dict[str, object]], _pre_cooldown_deployments), # cast-ok: router deployments are plain dicts + self._deployment_ids( + self._region_routable_deployments( + request_deployments, + cast(Mapping[str, object] | None, request_kwargs), # cast-ok: str-keyed request kwargs + True, + ) + ) + & frozenset(cooldown_deployments), + _OPAQUE_VALUE_ADAPTER.validate_python(input), + routing_kwargs, + ) + ) healthy_deployments = litellm.utils.get_order_filtered_deployments( list( self._narrow_to_video_deployment( without_deployments_that_cannot_batch( - batch_candidates, + embedding_batch_routed, caller_deployments=lambda: deployments_the_caller_may_use(batch_candidates, routing_kwargs), ) ) diff --git a/litellm/router_utils/embedding_batch_size_check.py b/litellm/router_utils/embedding_batch_size_check.py new file mode 100644 index 00000000000..6000bc9766d --- /dev/null +++ b/litellm/router_utils/embedding_batch_size_check.py @@ -0,0 +1,199 @@ +import traceback +from collections.abc import Generator, Iterable, Mapping, Sequence +from contextlib import contextmanager +from contextvars import ContextVar +from typing import TYPE_CHECKING, Final + +from typing_extensions import TypeIs # noqa: TID251 # TypeIs reaches typing only on 3.13 + +import litellm +from litellm.constants import PROXY_REJECTED_BEFORE_ROUTING_KEY +from litellm.types.utils import CallTypes + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +_SELECTING_FOR_EMBEDDING: Final[ContextVar[bool]] = ContextVar("selecting_for_embedding", default=False) + + +def _is_str_keyed_mapping( + value: object, +) -> TypeIs[Mapping[str, object]]: # guard-ok: model_info, litellm_params and extra_body are str-keyed config objects + return isinstance(value, Mapping) + + +def _is_embedding_input( + value: object, +) -> TypeIs[str | Sequence[object]]: # guard-ok: only the length and first item are read + return isinstance(value, Sequence) + + +def count_embedding_inputs(input: str | Sequence[object]) -> int: + if isinstance(input, str): + return 1 + if input and isinstance(input[0], int): + return 1 + return len(input) + + +def effective_embedding_input( + input: object, + extra_body: object, +) -> str | Sequence[object] | None: + """Return the input shape used for batch-size checks, or None when neither is an embedding input. + + Some providers merge ``extra_body`` over the request body (OpenAI-compatible paths); + others ignore ``extra_body.input`` and use the top-level input (e.g. Bedrock Titan). + Use whichever of the two yields the larger input count so limits cannot be bypassed. + """ + override: Final = extra_body.get("input") if _is_str_keyed_mapping(extra_body) else None + candidates: Final = tuple(value for value in (input, override) if _is_embedding_input(value)) + return max(candidates, key=count_embedding_inputs, default=None) + + +def _litellm_param(deployment: Mapping[str, object], key: str) -> object: + litellm_params: Final = deployment.get("litellm_params") + return litellm_params.get(key) if _is_str_keyed_mapping(litellm_params) else None + + +def _embedding_batch_limit(deployment: Mapping[str, object]) -> int | None: + model_info: Final = deployment.get("model_info") + limit: Final = model_info.get("max_embedding_batch_size") if _is_str_keyed_mapping(model_info) else None + return limit if isinstance(limit, int) else None + + +def _requested_embedding_count( + deployment: Mapping[str, object], + input: object, + request_kwargs: Mapping[str, object], +) -> int: + extra_body: Final = request_kwargs.get("extra_body", _litellm_param(deployment, "extra_body")) + effective_input: Final = effective_embedding_input(input, extra_body) + return 0 if effective_input is None else count_embedding_inputs(effective_input) + + +def _deployment_admits_embedding_batch( + deployment: Mapping[str, object], + input: object, + request_kwargs: Mapping[str, object], +) -> bool: + limit: Final = _embedding_batch_limit(deployment) + return limit is None or _requested_embedding_count(deployment, input, request_kwargs) <= limit + + +def _str_keyed_or_empty(value: object) -> Mapping[str, object]: + return value if _is_str_keyed_mapping(value) else {} + + +def _deployment_id(deployment: Mapping[str, object]) -> str | None: + model_info: Final = deployment.get("model_info") + deployment_id: Final = model_info.get("id") if _is_str_keyed_mapping(model_info) else None + return None if deployment_id is None else str(deployment_id) + + +def _is_paused(deployment: Mapping[str, object]) -> bool: + model_info: Final = deployment.get("model_info") + return _is_str_keyed_mapping(model_info) and model_info.get("blocked") is True + + +@contextmanager +def embedding_deployment_selection() -> Generator[None]: + """Apply the input limit to deployment selection inside this block, which only embedding calls enter.""" + token: Final = _SELECTING_FOR_EMBEDDING.set(True) + try: + yield + finally: + _SELECTING_FOR_EMBEDDING.reset(token) + + +def deployments_admitting_embedding_batch( + deployments: list[dict[str, object]], # mutable-ok: router deployment-list contract + input: object, + request_kwargs: object, +) -> list[dict[str, object]] | None: # mutable-ok: returns a mutable deployment list + """Keep the deployments whose input limit admits the request, all of them outside embedding selection, else None""" + if not _SELECTING_FOR_EMBEDDING.get(): + return deployments + kwargs: Final = _str_keyed_or_empty(request_kwargs) + admitting: Final = [ + deployment for deployment in deployments if _deployment_admits_embedding_batch(deployment, input, kwargs) + ] + return admitting or None + + +def deployments_when_none_admits_embedding_batch( + deployments: list[dict[str, object]], # mutable-ok: router deployment-list contract + before_cooldown: Iterable[Mapping[str, object]], + routable_cooling_ids: frozenset[str], + input: object, + request_kwargs: object, +) -> list[dict[str, object]]: # mutable-ok: returns a mutable deployment list + """None while an unpaused cooling deployment the request routes to admits it, so it waits out the cooldown, + else the deployments with the largest limit, so the 400 names the most the group accepts""" + kwargs: Final = _str_keyed_or_empty(request_kwargs) + routable_cooling: Final = ( + deployment + for deployment in before_cooldown + if _deployment_id(deployment) in routable_cooling_ids and not _is_paused(deployment) + ) + if any(_deployment_admits_embedding_batch(deployment, input, kwargs) for deployment in routable_cooling): + return [] + largest_limit: Final = max((_embedding_batch_limit(deployment) or 0 for deployment in deployments), default=0) + return [deployment for deployment in deployments if _embedding_batch_limit(deployment) == largest_limit] + + +def _mark_litellm_params(logging_obj: "LiteLLMLoggingObj", *keys: str) -> None: + litellm_params: Final = logging_obj.model_call_details.get("litellm_params") + if not isinstance(litellm_params, dict): + return + for key in keys: + litellm_params[key] = True + + +def log_rejection_before_provider_call(logging_obj: "LiteLLMLoggingObj", error: Exception) -> None: + """Mark the request as never reaching a provider, so deployment health metrics skip its failure, then log it""" + _mark_litellm_params(logging_obj, PROXY_REJECTED_BEFORE_ROUTING_KEY) + logging_obj.failure_handler(error, "".join(traceback.format_exception(error))) + + +async def async_log_rejection_before_provider_call(logging_obj: "LiteLLMLoggingObj", error: Exception) -> None: + """Log it like an async provider failure: sync handlers inline, before a fallback swaps the deployment they read, + and the call marked async, so a logger with both hooks records it once""" + _mark_litellm_params(logging_obj, PROXY_REJECTED_BEFORE_ROUTING_KEY, CallTypes.aembedding.value) + traceback_text: Final = "".join(traceback.format_exception(error)) + logging_obj.failure_handler(error, traceback_text) + await logging_obj.async_failure_handler(error, traceback_text) + + +def clear_rejected_before_provider_call(logging_obj: "LiteLLMLoggingObj") -> None: + """Drop the mark an earlier rejected attempt left, since a fallback or retry shares the logging object""" + litellm_params: Final = logging_obj.model_call_details.get("litellm_params") + if isinstance(litellm_params, dict) and PROXY_REJECTED_BEFORE_ROUTING_KEY in litellm_params: + del litellm_params[PROXY_REJECTED_BEFORE_ROUTING_KEY] + + +def embedding_batch_too_large_error( + deployment: object, + input: object, + request_kwargs: object, + model: str, +) -> litellm.BadRequestError | None: + if not _is_str_keyed_mapping(deployment): + return None + limit: Final = _embedding_batch_limit(deployment) + if limit is None: + return None + + input_count: Final = _requested_embedding_count(deployment, input, _str_keyed_or_empty(request_kwargs)) + if input_count <= limit: + return None + + return litellm.BadRequestError( + message=( + f"Embedding request has {input_count} inputs, which exceeds " + f"max_embedding_batch_size={limit} for model={model}. " + "Split the input into smaller batches." + ), + model=model, + llm_provider="", + ) diff --git a/litellm/types/router.py b/litellm/types/router.py index 6cd758d864e..1f159cb8b16 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -323,6 +323,8 @@ class ModelInfo(MirroredPricingParams): # in the spend log row's metadata. Set it on every deployment of the group. internal_router_model: bool | None = None + max_embedding_batch_size: Annotated[int, Field(gt=0)] | None = None + def __init__(self, id: str | int | None = None, **params) -> None: if id is None: id = str(uuid.uuid4()) # Generate a UUID if id is None or not provided diff --git a/tests/unit/router_utils/test_embedding_batch_size_check.py b/tests/unit/router_utils/test_embedding_batch_size_check.py new file mode 100644 index 00000000000..d5986d9e6f7 --- /dev/null +++ b/tests/unit/router_utils/test_embedding_batch_size_check.py @@ -0,0 +1,1024 @@ +import datetime +import json +import threading +from collections.abc import Iterator +from types import MappingProxyType +from typing import Final + +import httpx +import pytest +import respx +from prometheus_client import REGISTRY + +import litellm +from litellm import Router +from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo +from litellm.constants import PROXY_REJECTED_BEFORE_ROUTING_KEY +from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.prometheus import PrometheusLogger +from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging +from litellm.types.router import RouterRateLimitError +from litellm.router_utils.embedding_batch_size_check import ( + count_embedding_inputs, + deployments_admitting_embedding_batch, + deployments_when_none_admits_embedding_batch, + effective_embedding_input, + embedding_batch_too_large_error, + embedding_deployment_selection, +) + +_MODEL: Final = "embed-group" +_PROVIDER: Final = "openai" +_OPENAI_EMBED_BASE: Final = "https://batch-limit-embed.local/v1" +_LOGGING_START: Final = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc) +_VLLM_EMBED_BASE: Final = "https://batch-limit-vllm.local/v1" +_UNCAPPED_EMBED_BASE: Final = "https://batch-limit-uncapped.local/v1" + + +def _model_info(max_embedding_batch_size: int | None) -> dict[str, object]: + if max_embedding_batch_size is None: + return {} + return {"max_embedding_batch_size": max_embedding_batch_size} + + +def _deployment( + limit: int | None, + extra_body: dict[str, object] | None = None, + deployment_id: str = "embed-deployment", +) -> dict[str, object]: + return { + "model_info": {"id": deployment_id, **_model_info(limit)}, + "litellm_params": {} if extra_body is None else {"extra_body": extra_body}, + } + + +def _admitted( + deployments: list[dict[str, object]], + input_value: object, + request_kwargs: dict[str, object] | None = None, + cooling: tuple[dict[str, object], ...] = (), + routable_cooling_ids: frozenset[str] = frozenset(), +) -> list[dict[str, object]]: + admitting: Final = deployments_admitting_embedding_batch(deployments, input_value, request_kwargs) + if admitting is not None: + return admitting + return deployments_when_none_admits_embedding_batch( + deployments, (*deployments, *cooling), routable_cooling_ids, input_value, request_kwargs + ) + + +def _batch_error( + input_value: object, + limit: int | None, + request_kwargs: dict[str, object] | None = None, +) -> litellm.BadRequestError | None: + return embedding_batch_too_large_error( + deployment={"model_info": _model_info(limit), "litellm_params": {"custom_llm_provider": _PROVIDER}}, + input=input_value, + request_kwargs=request_kwargs or {}, + model=_MODEL, + ) + + +def _embedding_router( + *, + provider_model: str, + api_base: str, + max_embedding_batch_size: int | None, + deployment_id: str, + litellm_params_extra: dict[str, object] | None = None, +) -> Router: + litellm_params: dict[str, object] = { + "model": provider_model, + "api_key": "sk-fake", + "api_base": api_base, + } + if litellm_params_extra: + litellm_params.update(litellm_params_extra) + return Router( + model_list=[ + { + "model_name": "embed", + "litellm_params": litellm_params, + "model_info": {"id": deployment_id, **_model_info(max_embedding_batch_size)}, + } + ], + num_retries=0, + ) + + +async def _call_embedding(router: Router, *, sync: bool, **kwargs: object) -> None: + if sync: + router.embedding(**kwargs) + else: + await router.aembedding(**kwargs) + + +def _vllm_embedding_response(input_count: int) -> httpx.Response: + return httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": i, "embedding": [0.1, 0.2]} for i in range(input_count)], + "model": "bge-m3", + "usage": {"prompt_tokens": input_count, "total_tokens": input_count}, + }, + ) + + +@pytest.mark.parametrize( + ("input_value", "expected_count"), + [ + ("hello", 1), + (["a", "b"], 2), + ([1, 2, 3, 4, 5], 1), + ([[1, 2, 3], [4, 5, 6]], 2), + ], +) +def test_openai_input_shapes(input_value: object, expected_count: int) -> None: + assert count_embedding_inputs(input_value) == expected_count + + +def test_batch_count_is_independent_of_string_length() -> None: + long_text: Final = "word " * 500 + assert count_embedding_inputs([long_text, long_text]) == 2 + + +_MIXED_EXTRA_BODY: Final = { + "truncate": "END", + "encoding_format": "float", + "input": ["a", "b", "c"], + "user": "batch-job", +} + + +@pytest.mark.parametrize( + ("top_level", "extra_body", "expected"), + [ + ("ok", None, "ok"), + ("ok", {"input": ["a", "b", "c"]}, ["a", "b", "c"]), + (["a", "b"], {"truncate": "END"}, ["a", "b"]), + ("ignored-top-level", _MIXED_EXTRA_BODY, ["a", "b", "c"]), + ( + ["would", "be", "three"], + {"dimensions": 256, "input": "single-string", "model": "ignored"}, + ["would", "be", "three"], + ), + ("ok", {"truncate": "NONE", "input": ("a", "b")}, ("a", "b")), + (["x", "y", "z"], {"input": [[1, 2], [3, 4]], "foo": "bar"}, ["x", "y", "z"]), + ( + ["a", "b", "c"], + {"input": "ok"}, + ["a", "b", "c"], + ), + (["a", "b"], {}, ["a", "b"]), + (["a", "b", "c"], {"input": None, "truncate": "END"}, ["a", "b", "c"]), + (["a", "b"], {"input": {"nested": "dict"}, "truncate": "END"}, ["a", "b"]), + (["a", "b"], {"input": 42, "dimensions": 128}, ["a", "b"]), + (None, {"input": ["a", "b"]}, ["a", "b"]), + (7, {"input": ["a", "b"]}, ["a", "b"]), + (None, None, None), + ], +) +def test_resolves_provider_bound_input( + top_level: object, + extra_body: dict[str, object] | None, + expected: object, +) -> None: + assert effective_embedding_input(top_level, extra_body) == expected + + +@pytest.mark.parametrize( + "extra_body", + ["not-a-mapping", ["input", "list"]], +) +def test_non_mapping_extra_body_keeps_top_level(extra_body: object) -> None: + top_level: Final = ["a", "b", "c"] + assert effective_embedding_input(top_level, extra_body) == top_level + + +def test_readonly_mapping_extra_body_with_input() -> None: + extra_body: Final = MappingProxyType({"input": ["a", "b"], "encoding_format": "float"}) + assert effective_embedding_input("ok", extra_body) == ["a", "b"] + + +class TestDeploymentsAdmittingEmbeddingBatch: + @pytest.fixture(autouse=True) + def _embedding_selection(self) -> Iterator[None]: + with embedding_deployment_selection(): + yield + + @pytest.mark.parametrize( + ("input_value", "limit"), + [ + ("hello", 1), + (["a", "b"], 2), + (["a", "b", "c"], 3), + ([1, 2, 3, 4, 5], 1), + ([[1, 2, 3], [4, 5, 6]], 2), + ], + ) + def test_keeps_deployments_whose_limit_admits_the_batch(self, input_value: str | list, limit: int) -> None: + capped: Final = _deployment(limit) + uncapped: Final = _deployment(None) + assert _admitted([capped, uncapped], input_value) == [capped, uncapped] + + def test_drops_deployments_whose_limit_is_below_the_batch(self) -> None: + capped: Final = _deployment(2) + uncapped: Final = _deployment(None) + assert _admitted([capped, uncapped], ["a", "b", "c"]) == [uncapped] + + def test_keeps_the_largest_limit_deployments_when_none_admits_the_batch(self) -> None: + smaller: Final = _deployment(1, deployment_id="smaller") + small: Final = _deployment(2, deployment_id="small") + other_small: Final = _deployment(2, deployment_id="other-small") + assert _admitted([small, smaller, other_small], ["a", "b", "c"]) == [small, other_small] + + def test_request_without_embedding_input_keeps_every_deployment(self) -> None: + capped: Final = _deployment(1) + uncapped: Final = _deployment(None) + assert _admitted([capped, uncapped], None) == [capped, uncapped] + + def test_request_extra_body_replaces_the_deployment_extra_body(self) -> None: + capped: Final = _deployment(2, extra_body={"input": ["a", "b", "c"]}) + uncapped: Final = _deployment(None) + assert _admitted([capped, uncapped], "ok") == [uncapped] + assert _admitted([capped, uncapped], "ok", {"extra_body": {}}) == [capped, uncapped] + + @pytest.mark.parametrize("top_level", [None, 7, {"text": "a"}]) + def test_extra_body_input_routes_when_top_level_input_is_not_a_list(self, top_level: object) -> None: + capped: Final = _deployment(2) + uncapped: Final = _deployment(None) + request_kwargs: Final = {"extra_body": {"input": ["a", "b", "c"]}} + assert _admitted([capped, uncapped], top_level, request_kwargs) == [uncapped] + + def test_keeps_no_deployment_while_only_a_cooling_one_admits_the_batch(self) -> None: + capped: Final = _deployment(2, deployment_id="capped") + cooling: Final = _deployment(None, deployment_id="cooling") + assert _admitted([capped], ["a", "b", "c"], cooling=(cooling,), routable_cooling_ids=frozenset({"cooling"})) == [] + + def test_cooling_deployment_the_request_cannot_route_to_does_not_hold_the_batch(self) -> None: + capped: Final = _deployment(2, deployment_id="capped") + cooling: Final = _deployment(None, deployment_id="cooling") + assert _admitted([capped], ["a", "b", "c"], cooling=(cooling,)) == [capped] + + def test_paused_cooling_deployment_does_not_hold_the_batch(self) -> None: + capped: Final = _deployment(2, deployment_id="capped") + paused: Final = {"model_info": {"id": "cooling", "blocked": True}, "litellm_params": {}} + assert _admitted([capped], ["a", "b", "c"], cooling=(paused,), routable_cooling_ids=frozenset({"cooling"})) == [ + capped + ] + + def test_keeps_every_deployment_when_the_cooling_one_also_rejects_the_batch(self) -> None: + capped: Final = _deployment(2, deployment_id="capped") + cooling: Final = _deployment(1, deployment_id="cooling") + assert _admitted([capped], ["a", "b", "c"], cooling=(cooling,), routable_cooling_ids=frozenset({"cooling"})) == [ + capped + ] + + +@pytest.mark.parametrize( + ("input_value", "limit", "message_match"), + [ + (["a", "b", "c", "d"], 3, "4 inputs"), + ([[1], [2], [3]], 2, None), + ( + effective_embedding_input( + "ok", + {"truncate": "END", "input": ["a", "b", "c", "d"], "user": "x"}, + ), + 3, + "4 inputs", + ), + ( + effective_embedding_input( + ["a", "b", "c"], + {"input": {"bad": "shape"}, "encoding_format": "float"}, + ), + 2, + None, + ), + ], +) +def test_rejects_inputs_above_limit( + input_value: str | list, + limit: int, + message_match: str | None, +) -> None: + error: Final = _batch_error(input_value, limit) + assert isinstance(error, litellm.BadRequestError) + assert error.status_code == 400 + if message_match is not None: + assert message_match in error.message + assert f"max_embedding_batch_size={limit}" in error.message + + +@pytest.mark.parametrize("top_level", [None, 7, {"text": "a"}]) +def test_extra_body_override_counts_when_top_level_input_is_not_a_list(top_level: object) -> None: + error: Final = _batch_error(top_level, 2, {"extra_body": {"input": ["a", "b", "c"]}}) + assert error is not None + assert "3 inputs" in error.message + + +class TestRouterEmbeddingBatchLimit: + @pytest.fixture(autouse=True) + def _disable_aiohttp(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + + @pytest.mark.parametrize("sync", [False, True], ids=["async", "sync"]) + @pytest.mark.asyncio + async def test_oversized_top_level_input_rejects_before_provider_call(self, sync: bool) -> None: + router: Final = _embedding_router( + provider_model="openai/text-embedding-3-small", + api_base=_OPENAI_EMBED_BASE, + max_embedding_batch_size=2, + deployment_id="embed-batch-capped", + ) + embeddings_url: Final = f"{_OPENAI_EMBED_BASE}/embeddings" + empty_ok: Final = httpx.Response(200, json={"object": "list", "data": [], "model": "x", "usage": {}}) + + with respx.mock(assert_all_called=False) as respx_mock: + route: Final = respx_mock.post(embeddings_url).mock(return_value=empty_ok) + with pytest.raises(litellm.BadRequestError, match="max_embedding_batch_size=2"): + await _call_embedding( + router, + sync=sync, + model="embed", + input=["one", "two", "three"], + ) + assert route.call_count == 0 + + @pytest.mark.parametrize("sync", [False, True], ids=["async", "sync"]) + @pytest.mark.asyncio + async def test_oversized_batch_rejection_reaches_failure_callbacks(self, sync: bool) -> None: + failures: Final[list[tuple[Exception, object]]] = [] + + class FailureRecorder(CustomLogger): + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + failures.append((kwargs["exception"], kwargs["litellm_params"].get(PROXY_REJECTED_BEFORE_ROUTING_KEY))) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + failures.append((kwargs["exception"], kwargs["litellm_params"].get(PROXY_REJECTED_BEFORE_ROUTING_KEY))) + + recorder: Final = FailureRecorder() + logging_obj: Final = LitellmLogging( + model="embed", + messages=[], + stream=False, + call_type="embedding" if sync else "aembedding", + start_time=_LOGGING_START, + litellm_call_id="embed-batch-rejected", + function_id="embed-batch-rejected", + dynamic_failure_callbacks=[recorder] if sync else None, + dynamic_async_failure_callbacks=None if sync else [recorder], + ) + logging_obj.model_call_details["litellm_params"] = {"metadata": {}, "proxy_server_request": {}} + router: Final = _embedding_router( + provider_model="openai/text-embedding-3-small", + api_base=_OPENAI_EMBED_BASE, + max_embedding_batch_size=2, + deployment_id="embed-batch-capped", + ) + + with respx.mock(assert_all_called=False) as respx_mock: + route: Final = respx_mock.post(f"{_OPENAI_EMBED_BASE}/embeddings").mock( + return_value=httpx.Response(200, json={"object": "list", "data": [], "model": "x", "usage": {}}) + ) + with pytest.raises(litellm.BadRequestError, match="3 inputs") as excinfo: + await _call_embedding( + router, + sync=sync, + model="embed", + input=["one", "two", "three"], + litellm_logging_obj=logging_obj, + ) + assert route.call_count == 0 + assert failures == [(excinfo.value, True)] + + @pytest.mark.asyncio + async def test_async_rejection_logs_each_failure_hook_once_before_raising(self) -> None: + hook_events: Final[list[str]] = [] + legacy_callback_threads: Final[list[int]] = [] + legacy_callback_ran: Final = threading.Event() + + class BothHooks(CustomLogger): + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + hook_events.append("sync") + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + hook_events.append("async") + + def legacy_failure_callback(kwargs, completion_response, start_time, end_time) -> None: + legacy_callback_threads.append(threading.get_ident()) + legacy_callback_ran.set() + + both_hooks: Final = BothHooks() + logging_obj: Final = LitellmLogging( + model="embed", + messages=[], + stream=False, + call_type="aembedding", + start_time=_LOGGING_START, + litellm_call_id="embed-batch-logged-once", + function_id="embed-batch-logged-once", + dynamic_failure_callbacks=[legacy_failure_callback, both_hooks], + dynamic_async_failure_callbacks=[both_hooks], + ) + logging_obj.model_call_details["litellm_params"] = {"metadata": {}, "proxy_server_request": {}} + router: Final = _embedding_router( + provider_model="openai/text-embedding-3-small", + api_base=_OPENAI_EMBED_BASE, + max_embedding_batch_size=2, + deployment_id="embed-batch-capped", + ) + + with pytest.raises(litellm.BadRequestError, match="3 inputs"): + await router.aembedding(model="embed", input=["one", "two", "three"], litellm_logging_obj=logging_obj) + + assert legacy_callback_ran.wait(timeout=5) + assert legacy_callback_threads == [threading.get_ident()] + assert hook_events == ["async"] + + @pytest.mark.parametrize("sync", [False, True], ids=["async", "sync"]) + @pytest.mark.asyncio + async def test_fallback_after_a_rejection_logs_its_provider_call_as_reached(self, sync: bool) -> None: + litellm_params: Final[dict[str, object]] = {"model": "openai/text-embedding-3-small", "api_key": "sk-fake"} + router: Final = Router( + model_list=[ + { + "model_name": "embed-capped", + "litellm_params": {**litellm_params, "api_base": _OPENAI_EMBED_BASE}, + "model_info": {"id": "capped-box", "max_embedding_batch_size": 2}, + }, + { + "model_name": "embed-fallback", + "litellm_params": {**litellm_params, "api_base": _VLLM_EMBED_BASE}, + "model_info": {"id": "fallback-box"}, + }, + ], + fallbacks=[{"embed-capped": ["embed-fallback"]}], + num_retries=0, + ) + logging_obj: Final = LitellmLogging( + model="embed-capped", + messages=[], + stream=False, + call_type="embedding" if sync else "aembedding", + start_time=_LOGGING_START, + litellm_call_id="embed-batch-fallback", + function_id="embed-batch-fallback", + ) + logging_obj.update_environment_variables( + litellm_params={"metadata": {}, "proxy_server_request": {}}, optional_params={} + ) + + with respx.mock(assert_all_called=False) as respx_mock: + capped_route: Final = respx_mock.post(f"{_OPENAI_EMBED_BASE}/embeddings").mock( + return_value=_vllm_embedding_response(3) + ) + fallback_route: Final = respx_mock.post(f"{_VLLM_EMBED_BASE}/embeddings").mock( + return_value=_vllm_embedding_response(3) + ) + await _call_embedding( + router, + sync=sync, + model="embed-capped", + input=["one", "two", "three"], + litellm_logging_obj=logging_obj, + ) + assert (capped_route.call_count, fallback_route.call_count) == (0, 1) + assert PROXY_REJECTED_BEFORE_ROUTING_KEY not in logging_obj.model_call_details["litellm_params"] + + @pytest.fixture + def _isolated_prometheus_registry(self) -> Iterator[None]: + def unregister_all() -> None: + for collector in tuple(REGISTRY._collector_to_names): + REGISTRY.unregister(collector) + + unregister_all() + yield + unregister_all() + + @pytest.mark.asyncio + @pytest.mark.usefixtures("_isolated_prometheus_registry") + async def test_oversized_batch_rejection_leaves_the_deployment_healthy_in_prometheus(self) -> None: + prometheus: Final = PrometheusLogger() + metadata: Final[dict[str, object]] = {} + logging_obj: Final = LitellmLogging( + model="embed", + messages=[], + stream=False, + call_type="aembedding", + start_time=_LOGGING_START, + litellm_call_id="embed-batch-prometheus", + function_id="embed-batch-prometheus", + dynamic_async_failure_callbacks=[prometheus], + ) + logging_obj.model_call_details["litellm_params"] = {"metadata": metadata, "proxy_server_request": {}} + router: Final = _embedding_router( + provider_model="openai/text-embedding-3-small", + api_base=_OPENAI_EMBED_BASE, + max_embedding_batch_size=2, + deployment_id="embed-batch-capped", + ) + + with respx.mock(assert_all_called=False) as respx_mock: + respx_mock.post(f"{_OPENAI_EMBED_BASE}/embeddings").mock( + return_value=httpx.Response(200, json={"object": "list", "data": [], "model": "x", "usage": {}}) + ) + with pytest.raises(litellm.BadRequestError, match="3 inputs"): + await router.aembedding( + model="embed", input=["one", "two", "three"], metadata=metadata, litellm_logging_obj=logging_obj + ) + assert prometheus.litellm_deployment_state._metrics == {} + + @pytest.mark.parametrize("sync", [False, True], ids=["async", "sync"]) + @pytest.mark.asyncio + async def test_no_limit_configured_forwards_any_batch(self, sync: bool) -> None: + router: Final = _embedding_router( + provider_model="hosted_vllm/bge-m3", + api_base=_VLLM_EMBED_BASE, + max_embedding_batch_size=None, + deployment_id="vllm-embed-uncapped", + ) + + with respx.mock(assert_all_called=False) as respx_mock: + route: Final = respx_mock.post(f"{_VLLM_EMBED_BASE}/embeddings").mock( + return_value=_vllm_embedding_response(100) + ) + await _call_embedding(router, sync=sync, model="embed", input=["a"] * 100) + assert len(json.loads(route.calls.last.request.content)["input"]) == 100 + + @pytest.mark.parametrize("sync", [False, True], ids=["async", "sync"]) + @pytest.mark.asyncio + async def test_oversized_batch_goes_to_the_group_deployment_that_admits_it(self, sync: bool) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "embed", + "litellm_params": { + "model": "hosted_vllm/bge-m3", + "api_key": "sk-fake", + "api_base": _VLLM_EMBED_BASE, + "weight": 1000, + }, + "model_info": {"id": "vllm-embed-capped", "max_embedding_batch_size": 2}, + }, + { + "model_name": "embed", + "litellm_params": { + "model": "hosted_vllm/bge-m3", + "api_key": "sk-fake", + "api_base": _UNCAPPED_EMBED_BASE, + "weight": 1, + }, + "model_info": {"id": "vllm-embed-uncapped"}, + }, + ], + num_retries=0, + ) + + with respx.mock(assert_all_called=False) as respx_mock: + capped_route: Final = respx_mock.post(f"{_VLLM_EMBED_BASE}/embeddings").mock( + return_value=_vllm_embedding_response(3) + ) + uncapped_route: Final = respx_mock.post(f"{_UNCAPPED_EMBED_BASE}/embeddings").mock( + return_value=_vllm_embedding_response(3) + ) + for _ in range(5): + await _call_embedding(router, sync=sync, model="embed", input=["a", "b", "c"]) + assert (capped_route.call_count, uncapped_route.call_count) == (0, 5) + + @pytest.mark.asyncio + async def test_non_embedding_call_with_input_ignores_embedding_batch_limit(self) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "embed", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_key": "sk-fake", + "api_base": _OPENAI_EMBED_BASE, + "order": 1, + }, + "model_info": {"id": "capped", "max_embedding_batch_size": 2}, + }, + { + "model_name": "embed", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_key": "sk-fake", + "api_base": _UNCAPPED_EMBED_BASE, + "order": 2, + }, + "model_info": {"id": "uncapped"}, + }, + ], + num_retries=0, + ) + responses_input: Final = [{"role": "user", "content": str(turn)} for turn in range(3)] + + async_pick: Final = await router.async_get_available_deployment( + model="embed", request_kwargs={}, input=responses_input + ) + sync_pick: Final = router.get_available_deployment(model="embed", request_kwargs={}, input=responses_input) + + assert (async_pick["model_info"]["id"], sync_pick["model_info"]["id"]) == ("capped", "capped") + + @pytest.mark.parametrize("sync", [False, True], ids=["async", "sync"]) + @pytest.mark.asyncio + async def test_oversized_batch_waits_for_a_cooling_down_deployment_that_admits_it(self, sync: bool) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "embed", + "litellm_params": { + "model": "hosted_vllm/bge-m3", + "api_key": "sk-fake", + "api_base": _VLLM_EMBED_BASE, + }, + "model_info": {"id": "vllm-embed-capped", "max_embedding_batch_size": 2}, + }, + { + "model_name": "embed", + "litellm_params": { + "model": "hosted_vllm/bge-m3", + "api_key": "sk-fake", + "api_base": _UNCAPPED_EMBED_BASE, + }, + "model_info": {"id": "vllm-embed-uncapped"}, + }, + ], + num_retries=0, + ) + + with respx.mock(assert_all_called=False) as respx_mock: + capped_route: Final = respx_mock.post(f"{_VLLM_EMBED_BASE}/embeddings").mock( + return_value=_vllm_embedding_response(3) + ) + uncapped_route: Final = respx_mock.post(f"{_UNCAPPED_EMBED_BASE}/embeddings").mock( + return_value=httpx.Response(429, json={"error": {"message": "rate limited"}}) + ) + with pytest.raises(litellm.RateLimitError): + await _call_embedding(router, sync=sync, model="embed", input=["a", "b", "c"]) + with pytest.raises(RouterRateLimitError): + await _call_embedding(router, sync=sync, model="embed", input=["a", "b", "c"]) + assert (capped_route.call_count, uncapped_route.call_count) == (0, 1) + + @pytest.mark.parametrize("sync", [False, True], ids=["async", "sync"]) + @pytest.mark.asyncio + async def test_oversized_batch_names_the_cap_while_a_paused_admitting_deployment_cools_down( + self, sync: bool + ) -> None: + uncapped_params: Final = {"model": "hosted_vllm/bge-m3", "api_key": "sk-fake", "api_base": _UNCAPPED_EMBED_BASE} + router: Final = Router( + model_list=[ + { + "model_name": "embed", + "litellm_params": { + "model": "hosted_vllm/bge-m3", + "api_key": "sk-fake", + "api_base": _VLLM_EMBED_BASE, + }, + "model_info": {"id": "vllm-embed-capped", "max_embedding_batch_size": 2}, + }, + {"model_name": "embed", "litellm_params": uncapped_params, "model_info": {"id": "vllm-embed-uncapped"}}, + ], + num_retries=0, + ) + + with respx.mock(assert_all_called=False) as respx_mock: + capped_route: Final = respx_mock.post(f"{_VLLM_EMBED_BASE}/embeddings").mock( + return_value=_vllm_embedding_response(3) + ) + uncapped_route: Final = respx_mock.post(f"{_UNCAPPED_EMBED_BASE}/embeddings").mock( + return_value=httpx.Response(429, json={"error": {"message": "rate limited"}}) + ) + with pytest.raises(litellm.RateLimitError): + await _call_embedding(router, sync=sync, model="embed", input=["a", "b", "c"]) + router.upsert_deployment( + Deployment( + model_name="embed", + litellm_params=LiteLLM_Params(**uncapped_params), + model_info=ModelInfo(id="vllm-embed-uncapped", blocked=True), + ) + ) + with pytest.raises(litellm.BadRequestError, match="max_embedding_batch_size=2"): + await _call_embedding(router, sync=sync, model="embed", input=["a", "b", "c"]) + assert (capped_route.call_count, uncapped_route.call_count) == (0, 1) + + @pytest.mark.parametrize("sync", [False, True], ids=["async", "sync"]) + @pytest.mark.asyncio + async def test_extra_body_input_override_rejects_before_provider_call(self, sync: bool) -> None: + router: Final = _embedding_router( + provider_model="hosted_vllm/bge-m3", + api_base=_VLLM_EMBED_BASE, + max_embedding_batch_size=2, + deployment_id="vllm-embed-batch-capped", + ) + embeddings_url: Final = f"{_VLLM_EMBED_BASE}/embeddings" + extra_body: Final = {"truncate": "END", "input": ["a", "b", "c"]} + + with respx.mock(assert_all_called=False) as respx_mock: + route: Final = respx_mock.post(embeddings_url).mock(return_value=_vllm_embedding_response(3)) + with pytest.raises(litellm.BadRequestError, match="3 inputs"): + await _call_embedding( + router, + sync=sync, + model="embed", + input="ok", + extra_body=extra_body, + ) + assert route.call_count == 0 + + @pytest.mark.asyncio + async def test_deployment_extra_body_input_counts_toward_batch_limit(self) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "embed", + "litellm_params": { + "model": "hosted_vllm/bge-m3", + "api_key": "sk-fake", + "api_base": _VLLM_EMBED_BASE, + "extra_body": {"input": ["a", "b", "c"]}, + }, + "model_info": { + "id": "vllm-embed-deployment-extra-body", + "max_embedding_batch_size": 2, + }, + } + ], + num_retries=0, + ) + embeddings_url: Final = f"{_VLLM_EMBED_BASE}/embeddings" + + with respx.mock(assert_all_called=False) as respx_mock: + route: Final = respx_mock.post(embeddings_url).mock(return_value=_vllm_embedding_response(3)) + with pytest.raises(litellm.BadRequestError, match="3 inputs"): + await router.aembedding(model="embed", input="ok") + assert route.call_count == 0 + + @pytest.mark.asyncio + async def test_extra_body_input_at_limit_reaches_provider(self) -> None: + router: Final = _embedding_router( + provider_model="hosted_vllm/bge-m3", + api_base=_VLLM_EMBED_BASE, + max_embedding_batch_size=2, + deployment_id="vllm-embed-batch-capped", + ) + embeddings_url: Final = f"{_VLLM_EMBED_BASE}/embeddings" + + with respx.mock(assert_all_called=False) as respx_mock: + route: Final = respx_mock.post(embeddings_url).mock(return_value=_vllm_embedding_response(2)) + await router.aembedding(model="embed", input="ok", extra_body={"input": ["a", "b"]}) + assert route.call_count == 1 + assert json.loads(route.calls.last.request.content)["input"] == ["a", "b"] + + @pytest.mark.asyncio + async def test_unrelated_extra_body_does_not_bypass_top_level_batch_limit(self) -> None: + router: Final = _embedding_router( + provider_model="hosted_vllm/bge-m3", + api_base=_VLLM_EMBED_BASE, + max_embedding_batch_size=2, + deployment_id="vllm-embed-batch-capped", + ) + embeddings_url: Final = f"{_VLLM_EMBED_BASE}/embeddings" + + with respx.mock(assert_all_called=False) as respx_mock: + route: Final = respx_mock.post(embeddings_url).mock(return_value=_vllm_embedding_response(3)) + with pytest.raises(litellm.BadRequestError, match="3 inputs"): + await _call_embedding( + router, + sync=False, + model="embed", + input=["a", "b", "c"], + extra_body={"truncate": "END", "encoding_format": "float"}, + ) + assert route.call_count == 0 + + @pytest.mark.parametrize("sync", [False, True], ids=["async", "sync"]) + @pytest.mark.asyncio + async def test_bedrock_ignores_extra_body_input_override_for_provider_but_guard_uses_top_level( + self, + sync: bool, + ) -> None: + router: Final = _embedding_router( + provider_model="bedrock/amazon.titan-embed-text-v1", + api_base="", + max_embedding_batch_size=2, + deployment_id="bedrock-titan-embed-batch-capped", + litellm_params_extra={ + "aws_access_key_id": "AKIAFAKE", + "aws_secret_access_key": "fake-secret", + "aws_region_name": "us-east-1", + "custom_llm_provider": "bedrock", + }, + ) + + with pytest.raises(litellm.BadRequestError, match="3 inputs"): + await _call_embedding( + router, + sync=sync, + model="embed", + input=["a", "b", "c"], + extra_body={"input": "ok"}, + ) + + @pytest.mark.parametrize("sync", [False, True], ids=["async", "sync"]) + @pytest.mark.asyncio + async def test_explicit_custom_llm_provider_in_deployment_rejects_oversized_batch( + self, + sync: bool, + ) -> None: + router: Final = _embedding_router( + provider_model="openai/text-embedding-3-small", + api_base=_OPENAI_EMBED_BASE, + max_embedding_batch_size=2, + deployment_id="embed-batch-capped-explicit-provider", + litellm_params_extra={"custom_llm_provider": "openai"}, + ) + embeddings_url: Final = f"{_OPENAI_EMBED_BASE}/embeddings" + empty_ok: Final = httpx.Response(200, json={"object": "list", "data": [], "model": "x", "usage": {}}) + + with respx.mock(assert_all_called=False) as respx_mock: + route: Final = respx_mock.post(embeddings_url).mock(return_value=empty_ok) + with pytest.raises(litellm.BadRequestError, match="max_embedding_batch_size=2"): + await _call_embedding( + router, + sync=sync, + model="embed", + input=["one", "two", "three"], + ) + assert route.call_count == 0 + + @pytest.mark.asyncio + async def test_oversized_batch_names_the_cap_when_tags_exclude_the_admitting_deployment(self) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "embed", + "litellm_params": { + "model": "hosted_vllm/bge-m3", + "api_key": "sk-fake", + "api_base": _VLLM_EMBED_BASE, + "tags": ["team-a"], + }, + "model_info": {"id": "team-a-capped", "max_embedding_batch_size": 2}, + }, + { + "model_name": "embed", + "litellm_params": { + "model": "hosted_vllm/bge-m3", + "api_key": "sk-fake", + "api_base": _UNCAPPED_EMBED_BASE, + "tags": ["team-b"], + }, + "model_info": {"id": "team-b-uncapped"}, + }, + ], + enable_tag_filtering=True, + num_retries=0, + ) + + with respx.mock(assert_all_called=False) as respx_mock: + capped_route: Final = respx_mock.post(f"{_VLLM_EMBED_BASE}/embeddings").mock( + return_value=_vllm_embedding_response(3) + ) + uncapped_route: Final = respx_mock.post(f"{_UNCAPPED_EMBED_BASE}/embeddings").mock( + return_value=_vllm_embedding_response(3) + ) + with pytest.raises(litellm.BadRequestError, match="max_embedding_batch_size=2"): + await router.aembedding(model="embed", input=["a", "b", "c"], metadata={"tags": ["team-a"]}) + assert (capped_route.call_count, uncapped_route.call_count) == (0, 0) + + @pytest.mark.asyncio + async def test_oversized_batch_names_the_cap_while_another_tags_admitting_deployment_cools_down(self) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "embed", + "litellm_params": { + "model": "hosted_vllm/bge-m3", + "api_key": "sk-fake", + "api_base": _VLLM_EMBED_BASE, + "tags": ["team-a"], + }, + "model_info": {"id": "team-a-capped", "max_embedding_batch_size": 2}, + }, + { + "model_name": "embed", + "litellm_params": { + "model": "hosted_vllm/bge-m3", + "api_key": "sk-fake", + "api_base": _UNCAPPED_EMBED_BASE, + "tags": ["team-b"], + }, + "model_info": {"id": "team-b-uncapped"}, + }, + ], + enable_tag_filtering=True, + num_retries=0, + ) + + with respx.mock(assert_all_called=False) as respx_mock: + capped_route: Final = respx_mock.post(f"{_VLLM_EMBED_BASE}/embeddings").mock( + return_value=_vllm_embedding_response(3) + ) + uncapped_route: Final = respx_mock.post(f"{_UNCAPPED_EMBED_BASE}/embeddings").mock( + return_value=httpx.Response(429, json={"error": {"message": "rate limited"}}) + ) + with pytest.raises(litellm.RateLimitError): + await router.aembedding(model="embed", input=["a", "b", "c"], metadata={"tags": ["team-b"]}) + with pytest.raises(litellm.BadRequestError, match="max_embedding_batch_size=2"): + await router.aembedding(model="embed", input=["a", "b", "c"], metadata={"tags": ["team-a"]}) + assert (capped_route.call_count, uncapped_route.call_count) == (0, 1) + + @pytest.mark.parametrize("sync", [False, True], ids=["async", "sync"]) + @pytest.mark.asyncio + async def test_oversized_batch_names_the_cap_while_another_regions_admitting_deployment_cools_down( + self, sync: bool + ) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "embed", + "litellm_params": { + "model": "hosted_vllm/bge-m3", + "api_key": "sk-fake", + "api_base": _VLLM_EMBED_BASE, + "region_name": "eu", + }, + "model_info": {"id": "eu-capped", "max_embedding_batch_size": 2}, + }, + { + "model_name": "embed", + "litellm_params": { + "model": "hosted_vllm/bge-m3", + "api_key": "sk-fake", + "api_base": _UNCAPPED_EMBED_BASE, + "region_name": "us", + }, + "model_info": {"id": "us-uncapped"}, + }, + ], + enable_pre_call_checks=True, + num_retries=0, + ) + + with respx.mock(assert_all_called=False) as respx_mock: + capped_route: Final = respx_mock.post(f"{_VLLM_EMBED_BASE}/embeddings").mock( + return_value=_vllm_embedding_response(3) + ) + uncapped_route: Final = respx_mock.post(f"{_UNCAPPED_EMBED_BASE}/embeddings").mock( + return_value=httpx.Response(429, json={"error": {"message": "rate limited"}}) + ) + with pytest.raises(litellm.RateLimitError): + await _call_embedding( + router, sync=sync, model="embed", input=["a", "b", "c"], allowed_model_region="us" + ) + with pytest.raises(litellm.BadRequestError, match="max_embedding_batch_size=2"): + await _call_embedding( + router, sync=sync, model="embed", input=["a", "b", "c"], allowed_model_region="eu" + ) + assert (capped_route.call_count, uncapped_route.call_count) == (0, 1) + + @pytest.mark.parametrize("sync", [False, True], ids=["async", "sync"]) + @pytest.mark.asyncio + async def test_oversized_batch_rejection_falls_back_with_explicit_provider(self, sync: bool) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "embed", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_key": "sk-fake", + "api_base": _OPENAI_EMBED_BASE, + "custom_llm_provider": "openai", + }, + "model_info": {"id": "explicit-provider-capped", "max_embedding_batch_size": 2}, + }, + { + "model_name": "embed-fallback", + "litellm_params": { + "model": "hosted_vllm/bge-m3", + "api_key": "sk-fake", + "api_base": _UNCAPPED_EMBED_BASE, + }, + "model_info": {"id": "fallback-uncapped"}, + }, + ], + fallbacks=[{"embed": ["embed-fallback"]}], + fallback_on_unsupported_params=False, + num_retries=0, + ) + + with respx.mock(assert_all_called=False) as respx_mock: + capped_route: Final = respx_mock.post(f"{_OPENAI_EMBED_BASE}/embeddings").mock( + return_value=_vllm_embedding_response(3) + ) + fallback_route: Final = respx_mock.post(f"{_UNCAPPED_EMBED_BASE}/embeddings").mock( + return_value=_vllm_embedding_response(3) + ) + await _call_embedding(router, sync=sync, model="embed", input=["a", "b", "c"]) + assert (capped_route.call_count, fallback_route.call_count) == (0, 1) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 992a7d91b0d..2b6898e39d9 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -50107,6 +50107,8 @@ export interface components { input_cost_per_token?: number | null; /** Internal Router Model */ internal_router_model?: boolean | null; + /** Max Embedding Batch Size */ + max_embedding_batch_size?: number | null; /** * Member Auto Router * @default false From 06e3e32c88326930645e3aabc2e454d7dfc4bfdd Mon Sep 17 00:00:00 2001 From: Ondra Zahradnik Date: Tue, 20 Oct 2026 13:37:46 +0200 Subject: [PATCH 04/33] fix(vertex_ai): reuse known context caches instead of listing on every request (#46857) Vertex AI and Gemini context caching now remember each cache they create or find until 10s before its expireTime, so a repeat cache-marked request goes straight to generateContent instead of listing every cachedContents first. Concurrent async requests for one prefix share one lookup or create When generation rejects a remembered cache with a 400, 403 or 404 that names the cached content, the request rebuilds the cache and sends once more. Any other 400, 403 or 404 drops the record without a retry Fixes #36395 --- .../context_caching/known_cached_contents.py | 156 +++++ .../context_caching/rejected_cache_retry.py | 74 ++ .../vertex_ai_context_caching.py | 179 +++-- .../vertex_and_google_ai_studio_gemini.py | 251 ++++--- .../vertex_ai/context_caching/conftest.py | 10 + .../test_known_cached_contents.py | 663 ++++++++++++++++++ 6 files changed, 1173 insertions(+), 160 deletions(-) create mode 100644 litellm/llms/vertex_ai/context_caching/known_cached_contents.py create mode 100644 litellm/llms/vertex_ai/context_caching/rejected_cache_retry.py create mode 100644 tests/unit/llms/vertex_ai/context_caching/conftest.py create mode 100644 tests/unit/llms/vertex_ai/context_caching/test_known_cached_contents.py diff --git a/litellm/llms/vertex_ai/context_caching/known_cached_contents.py b/litellm/llms/vertex_ai/context_caching/known_cached_contents.py new file mode 100644 index 00000000000..6b7f098d304 --- /dev/null +++ b/litellm/llms/vertex_ai/context_caching/known_cached_contents.py @@ -0,0 +1,156 @@ +import asyncio +import contextlib +import re +import threading +import time +import weakref +from collections.abc import Awaitable, Callable, Mapping +from datetime import datetime, timezone +from typing import Final, Literal + +KNOWN_CACHE_EXPIRY_MARGIN_SECONDS: Final = 10.0 +MAX_KNOWN_CACHED_CONTENTS: Final = 10_000 +_CACHE_REJECTED_STATUS_CODES: Final = frozenset({400, 403, 404}) +_NAMES_CACHED_CONTENT: Final = re.compile(r"cached?\s*content", re.IGNORECASE) +_FORBIDS_FIELDS_BESIDE_A_CACHE: Final = re.compile(r"system[_\s]instruction", re.IGNORECASE) + +_RFC3339_FRACTION: Final = re.compile(r"\.(\d{1,6})\d*") + + +def names_rejected_cache(status_code: int, message: str) -> bool: + return ( + status_code in _CACHE_REJECTED_STATUS_CODES + and _NAMES_CACHED_CONTENT.search(message) is not None + and _FORBIDS_FIELDS_BESIDE_A_CACHE.search(message) is None + ) + + +def parse_expire_time(expire_time: object) -> float | None: + """Parse a Vertex RFC 3339 ``expireTime`` (nanosecond precision, ``Z`` suffix) to epoch seconds.""" + if not isinstance(expire_time, str) or not expire_time: + return None + normalized: Final = _RFC3339_FRACTION.sub(lambda m: "." + m.group(1), expire_time.replace("Z", "+00:00"), count=1) + try: + parsed: Final = datetime.fromisoformat(normalized) + except ValueError: + return None + if parsed.tzinfo is None: + return parsed.replace(tzinfo=timezone.utc).timestamp() + return parsed.timestamp() + + +class KnownCachedContents: + """Process-local record of the cachedContents this process created or found. + + Vertex meters every cachedContents list page against a per-minute, per-location quota, + so listing on every cache-marked request exhausts it under load. A cache is remembered + until ``KNOWN_CACHE_EXPIRY_MARGIN_SECONDS`` before the ``expireTime`` Google reported, + and after that requests list again, as they did before the record existed. + """ + + def __init__(self, clock: Callable[[], float] = time.time) -> None: + self._clock: Final = clock + self._entries: dict[tuple[tuple[str, str], str], tuple[str, float]] = {} # mutable-ok: guarded by _lock + self._lock: Final = threading.Lock() + self._inflight: weakref.WeakKeyDictionary[ # mutable-ok: private, guarded by _lock + asyncio.AbstractEventLoop, dict[tuple[tuple[str, str], str], asyncio.Future[str | None]] + ] = weakref.WeakKeyDictionary() + + @staticmethod + def scope( + base_url: str, + custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"], + headers: Mapping[str, str], + vertex_project: str | None, + vertex_location: str | None, + ) -> tuple[str, str]: + """Where a cache name is valid: a Vertex project and location, or the project a Gemini key belongs to. + + A custom ``api_base`` can replace the whole Vertex path, so the project and location are kept beside the URL. + The Gemini key is read from the headers actually sent, so an ``extra_headers`` override scopes the + cache to the caller's project. It is kept as-is in this in-memory scope and never logged or persisted. + """ + if custom_llm_provider != "gemini": + return (base_url, f"{vertex_project}/{vertex_location}") + sent_keys: Final = [value for name, value in headers.items() if name.lower() == "x-goog-api-key"] + return (base_url, sent_keys[-1] if sent_keys else "") + + def has_not_expired(self, expire_time: object) -> bool: + expires_at: Final = parse_expire_time(expire_time) + return expires_at is None or expires_at > self._clock() + + def get(self, scope: tuple[str, str], cache_key: str) -> str | None: + with self._lock: + entry: Final = self._entries.get((scope, cache_key)) + if entry is None: + return None + name, reuse_until = entry + if reuse_until > self._clock(): + return name + del self._entries[(scope, cache_key)] + return None + + def remember(self, scope: tuple[str, str], cache_key: str, name: str | None, expire_time: object) -> None: + expires_at: Final = parse_expire_time(expire_time) + if not name or expires_at is None: + return + reuse_until: Final = expires_at - KNOWN_CACHE_EXPIRY_MARGIN_SECONDS + now: Final = self._clock() + if reuse_until <= now: + return + with self._lock: + if len(self._entries) >= MAX_KNOWN_CACHED_CONTENTS: + self._entries = {k: v for k, v in self._entries.items() if v[1] > now} + while len(self._entries) >= MAX_KNOWN_CACHED_CONTENTS: + del self._entries[next(iter(self._entries))] + self._entries[(scope, cache_key)] = (name, reuse_until) + + def forget_rejected(self, cached_content: str | None, status_code: int) -> None: + if cached_content is None or status_code not in _CACHE_REJECTED_STATUS_CODES: + return + with self._lock: + self._entries = {key: entry for key, entry in self._entries.items() if entry[0] != cached_content} + + async def single_flight( + self, + scope: tuple[str, str], + cache_key: str, + resolve: Callable[[], Awaitable[str | None]], + wait_seconds: float | None, + ) -> str | None: + """Run ``resolve`` once per (scope, cache_key) at a time within this event loop. + + Concurrent requests for the same prefix wait for the first one's lookup-or-create + instead of each listing (and creating a duplicate cache). A waiter resolves on its + own when that first attempt fails or has not finished within ``wait_seconds``, and + never cancels it. + """ + loop: Final = asyncio.get_running_loop() + own: Final[asyncio.Future[str | None]] = loop.create_future() + with self._lock: + inflight: Final = self._inflight.setdefault(loop, {}) + leader: Final = inflight.setdefault((scope, cache_key), own) + if leader is not own: + with contextlib.suppress(asyncio.TimeoutError): + if name := await asyncio.wait_for(asyncio.shield(leader), wait_seconds): + return name + return await resolve() + try: + resolved: Final = await resolve() + except BaseException: + own.set_result(None) + raise + else: + own.set_result(resolved) + return resolved + finally: + with self._lock: + inflight.pop((scope, cache_key), None) + + def clear(self) -> None: + with self._lock: + self._entries.clear() + self._inflight.clear() + + +known_cached_contents: Final = KnownCachedContents() diff --git a/litellm/llms/vertex_ai/context_caching/rejected_cache_retry.py b/litellm/llms/vertex_ai/context_caching/rejected_cache_retry.py new file mode 100644 index 00000000000..43e4ee8dc84 --- /dev/null +++ b/litellm/llms/vertex_ai/context_caching/rejected_cache_retry.py @@ -0,0 +1,74 @@ +from collections.abc import Awaitable, Callable +from copy import deepcopy +from functools import partial +from typing import Final, TypeVar + +import httpx + +from litellm.litellm_core_utils.core_helpers import independent_snapshot +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.vertex_ai import RequestBody +from litellm.utils import is_cached_message + +from .known_cached_contents import names_rejected_cache + +Built: Final = TypeVar("Built") +Sent: Final = TypeVar("Sent") + + +def cache_rebuild( + build: Callable[..., Built], messages: list[AllMessageValues], optional_params: dict +) -> Callable[[], Built] | None: + """Copies the inputs before the first transform, which pops keys from ``optional_params`` and ``messages``.""" + if not any(map(is_cached_message, messages)): + return None + return partial(build, messages=deepcopy(messages), optional_params=independent_snapshot(optional_params)) + + +def _status_and_message(error: httpx.HTTPStatusError | BaseLLMException) -> tuple[int, str]: + match error: + case httpx.HTTPStatusError(): + return error.response.status_code, error.response.text + case BaseLLMException(): + return error.status_code, error.message + + +def _cache_rejected(body: RequestBody, error: httpx.HTTPStatusError | BaseLLMException) -> bool: + return body.get("cachedContent") is not None and names_rejected_cache(*_status_and_message(error)) + + +async def post_with_fresh_cache_async( + body: RequestBody, + rebuild: Callable[[], Awaitable[RequestBody]] | None, + post: Callable[[RequestBody], Awaitable[Sent]], + log_request: Callable[[RequestBody], None], +) -> tuple[RequestBody, Sent]: + try: + return body, await post(body) + except (httpx.HTTPStatusError, BaseLLMException) as error: + if rebuild is None or not _cache_rejected(body, error): + raise + fresh_body: Final = await rebuild() + log_request(fresh_body) + if fresh_body.get("cachedContent") == body.get("cachedContent"): + raise + return fresh_body, await post(fresh_body) + + +def post_with_fresh_cache_sync( + body: RequestBody, + rebuild: Callable[[], RequestBody] | None, + post: Callable[[RequestBody], Sent], + log_request: Callable[[RequestBody], None], +) -> tuple[RequestBody, Sent]: + try: + return body, post(body) + except (httpx.HTTPStatusError, BaseLLMException) as error: + if rebuild is None or not _cache_rejected(body, error): + raise + fresh_body: Final = rebuild() + log_request(fresh_body) + if fresh_body.get("cachedContent") == body.get("cachedContent"): + raise + return fresh_body, post(fresh_body) diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index 06b8f3a3fbb..5bff43828d6 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py @@ -20,6 +20,7 @@ from litellm.utils import is_prompt_caching_valid_prompt from ..common_utils import VertexAIError, get_vertex_base_url from ..vertex_llm_base import VertexBase +from .known_cached_contents import KnownCachedContents, known_cached_contents from .transformation import ( cached_messages_end_on_supported_turn, separate_cached_messages, @@ -35,6 +36,7 @@ MAX_PAGINATION_PAGES: Final = 100 # Reasonable upper bound for pagination class _CachedContentItem(TypedDict): displayName: ReadOnly[NotRequired[object]] name: ReadOnly[NotRequired[str | None]] + expireTime: ReadOnly[NotRequired[str | None]] @with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) @@ -51,6 +53,7 @@ class _CachedContentUsageMetadata(TypedDict): @with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) class _CreatedCachedContent(TypedDict): name: ReadOnly[NotRequired[str | None]] + expireTime: ReadOnly[NotRequired[str | None]] usageMetadata: ReadOnly[NotRequired[_CachedContentUsageMetadata | None]] @@ -67,6 +70,14 @@ def _record_cache_creation_tokens(logging_obj: Logging, created: _CreatedCachedC logging_obj.vertex_ai_cache_creation_input_tokens = max(tokens, 0) # rebind-ok: set provider count +def _wait_seconds(timeout: float | httpx.Timeout | None) -> float | None: + match timeout: + case httpx.Timeout(): + return timeout.read + case _: + return timeout + + class ContextCachingEndpoints(VertexBase): """ Covers context caching endpoints for Vertex AI + Google AI Studio @@ -74,8 +85,9 @@ class ContextCachingEndpoints(VertexBase): v0: covers Google AI Studio """ - def __init__(self) -> None: + def __init__(self, known_caches: KnownCachedContents = known_cached_contents) -> None: super().__init__() + self._known_caches: Final = known_caches def _get_token_and_url_context_caching( self, @@ -129,7 +141,7 @@ class ContextCachingEndpoints(VertexBase): self, cache_key: str, client: HTTPHandler, - headers: dict, + headers: dict[str, str], # mutable-ok: passed through to the HTTP client, which takes a dict api_key: str, api_base: str | None, logging_obj: Logging, @@ -159,6 +171,12 @@ class ContextCachingEndpoints(VertexBase): vertex_auth_header=vertex_auth_header, model=model, ) + scope: Final = KnownCachedContents.scope( + base_url, custom_llm_provider, headers, vertex_project=vertex_project, vertex_location=vertex_location + ) + known_name: Final = self._known_caches.get(scope, cache_key) + if known_name is not None: + return known_name page_token: str | None = None @@ -201,8 +219,14 @@ class ContextCachingEndpoints(VertexBase): # Check current page for matching cache_key for cached_item in raw_response["cachedContents"]: display_name = cached_item.get("displayName") - if display_name is not None and display_name == cache_key: - return cached_item.get("name") + if display_name is None or display_name != cache_key: + continue + expire_time = cached_item.get("expireTime") + if not self._known_caches.has_not_expired(expire_time): + continue + name = cached_item.get("name") + self._known_caches.remember(scope, cache_key, name, expire_time) + return name # Check if there are more pages page_token = raw_response.get("nextPageToken") @@ -216,7 +240,7 @@ class ContextCachingEndpoints(VertexBase): self, cache_key: str, client: AsyncHTTPHandler, - headers: dict, + headers: dict[str, str], # mutable-ok: passed through to the HTTP client, which takes a dict api_key: str, api_base: str | None, logging_obj: Logging, @@ -246,6 +270,12 @@ class ContextCachingEndpoints(VertexBase): vertex_auth_header=vertex_auth_header, model=model, ) + scope: Final = KnownCachedContents.scope( + base_url, custom_llm_provider, headers, vertex_project=vertex_project, vertex_location=vertex_location + ) + known_name: Final = self._known_caches.get(scope, cache_key) + if known_name is not None: + return known_name page_token: str | None = None @@ -288,8 +318,14 @@ class ContextCachingEndpoints(VertexBase): # Check current page for matching cache_key for cached_item in raw_response["cachedContents"]: display_name = cached_item.get("displayName") - if display_name is not None and display_name == cache_key: - return cached_item.get("name") + if display_name is None or display_name != cache_key: + continue + expire_time = cached_item.get("expireTime") + if not self._known_caches.has_not_expired(expire_time): + continue + name = cached_item.get("name") + self._known_caches.remember(scope, cache_key, name, expire_time) + return name # Check if there are more pages page_token = raw_response.get("nextPageToken") @@ -451,6 +487,14 @@ class ContextCachingEndpoints(VertexBase): raw_response_cached: Final = _CreatedCachedContentRoot.model_validate(response.json()).root _record_cache_creation_tokens(logging_obj, raw_response_cached) + self._known_caches.remember( + KnownCachedContents.scope( + url, custom_llm_provider, headers, vertex_project=vertex_project, vertex_location=vertex_location + ), + generated_cache_key, + raw_response_cached.get("name"), + raw_response_cached.get("expireTime"), + ) return ( non_cached_messages, optional_params, @@ -548,68 +592,77 @@ class ContextCachingEndpoints(VertexBase): generated_cache_key: Final = local_cache_obj.get_cache_key( messages=cached_messages, tools=tools, tool_choice=tool_choice, model=model ) - google_cache_name: Final = await self.async_check_cache( - cache_key=generated_cache_key, - client=client, - headers=headers, - api_key=api_key, - api_base=api_base, - logging_obj=logging_obj, - custom_llm_provider=custom_llm_provider, - vertex_project=vertex_project, - vertex_location=vertex_location, - vertex_auth_header=vertex_auth_header, - model=model, + scope: Final = KnownCachedContents.scope( + url, custom_llm_provider, headers, vertex_project=vertex_project, vertex_location=vertex_location ) - if google_cache_name: - return non_cached_messages, optional_params, google_cache_name - - ## TRANSFORM REQUEST - cached_content_request_body: Final = transform_openai_messages_to_gemini_context_caching( - model=model, - messages=cached_messages, - cache_key=generated_cache_key, - custom_llm_provider=custom_llm_provider, - vertex_project=vertex_project, - vertex_location=vertex_location, - ) - - cached_content_request_body["tools"] = tools - if tool_choice is not None: - cached_content_request_body["toolConfig"] = tool_choice - - ## LOGGING - logging_obj.pre_call( - input=messages, - api_key="", - additional_args={ - "complete_input_dict": cached_content_request_body, - "api_base": url, - "headers": headers, - }, - ) - - try: - response: Final = await client.post( - url=url, + async def lookup_or_create() -> str | None: + google_cache_name: Final = await self.async_check_cache( + cache_key=generated_cache_key, + client=client, headers=headers, - json=cached_content_request_body, + api_key=api_key, + api_base=api_base, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_auth_header=vertex_auth_header, + model=model, ) - response.raise_for_status() - except httpx.HTTPStatusError as err: - error_code: Final = err.response.status_code - raise VertexAIError(status_code=error_code, message=err.response.text) - except httpx.TimeoutException: - raise VertexAIError(status_code=408, message="Timeout error occurred.") - raw_response_cached: Final = _CreatedCachedContentRoot.model_validate(response.json()).root - _record_cache_creation_tokens(logging_obj, raw_response_cached) - return ( - non_cached_messages, - optional_params, - raw_response_cached.get("name"), + if google_cache_name: + return google_cache_name + + ## TRANSFORM REQUEST + cached_content_request_body: Final = transform_openai_messages_to_gemini_context_caching( + model=model, + messages=cached_messages, + cache_key=generated_cache_key, + custom_llm_provider=custom_llm_provider, + vertex_project=vertex_project, + vertex_location=vertex_location, + ) + + cached_content_request_body["tools"] = tools + if tool_choice is not None: + cached_content_request_body["toolConfig"] = tool_choice + + ## LOGGING + logging_obj.pre_call( + input=messages, + api_key="", + additional_args={ + "complete_input_dict": cached_content_request_body, + "api_base": url, + "headers": headers, + }, + ) + + try: + response: Final = await client.post( + url=url, + headers=headers, + json=cached_content_request_body, + ) + response.raise_for_status() + except httpx.HTTPStatusError as err: + error_code: Final = err.response.status_code + raise VertexAIError(status_code=error_code, message=err.response.text) + except httpx.TimeoutException: + raise VertexAIError(status_code=408, message="Timeout error occurred.") + + raw_response_cached: Final = _CreatedCachedContentRoot.model_validate(response.json()).root + _record_cache_creation_tokens(logging_obj, raw_response_cached) + self._known_caches.remember( + scope, generated_cache_key, raw_response_cached.get("name"), raw_response_cached.get("expireTime") + ) + return raw_response_cached.get("name") + + cache_name: Final = await self._known_caches.single_flight( + scope, generated_cache_key, lookup_or_create, _wait_seconds(timeout) ) + return non_cached_messages, optional_params, cache_name def get_cache(self): pass diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 257fafbb6b8..114b9cd80ae 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -96,6 +96,12 @@ from ..common_utils import ( build_vertex_schema, supports_response_json_schema, ) +from ..context_caching.known_cached_contents import known_cached_contents +from ..context_caching.rejected_cache_retry import ( + cache_rebuild, + post_with_fresh_cache_async, + post_with_fresh_cache_sync, +) from ..vertex_llm_base import VertexBase from .grounding_requests import calculate_grounding_requests from .long_decode import ( @@ -2786,7 +2792,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): litellm_params: dict, api_key: str | dict | None = None, api_base: str | None = None, - ) -> dict: + ) -> dict[str, str]: default_headers: Final = { "Content-Type": "application/json", } @@ -2800,6 +2806,20 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return default_headers +def _log_generate_request( + logging_obj: LoggingClass, + messages: Sequence[AllMessageValues], + api_base: str, + headers: Mapping[str, str], + body: RequestBody, +) -> None: + logging_obj.pre_call( + input=messages, + api_key="", + additional_args={"complete_input_dict": body, "api_base": api_base, "headers": headers}, + ) + + async def make_call( client: AsyncHTTPHandler | None, # module-level client gemini_client: AsyncHTTPHandler | None, # if passed by user @@ -2832,6 +2852,7 @@ async def make_call( response.raise_for_status() except httpx.HTTPStatusError as e: exception_string: Final = str(await e.response.aread()) + known_cached_contents.forget_rejected(request_body.get("cachedContent"), e.response.status_code) raise VertexAIError( status_code=e.response.status_code, message=VertexGeminiConfig().translate_exception_str(exception_string), @@ -2845,14 +2866,18 @@ async def make_call( ) async def open_hop(body: RequestBody) -> httpx.Response: - return await client.post( - url_for(api_base, body), - headers=dict(await hop_headers()), - data=json.dumps(body), - stream=True, - logging_obj=logging_obj, - timeout=timeout, - ) + try: + return await client.post( + url_for(api_base, body), + headers=dict(await hop_headers()), + data=json.dumps(body), + stream=True, + logging_obj=logging_obj, + timeout=timeout, + ) + except httpx.HTTPStatusError as e: + known_cached_contents.forget_rejected(body.get("cachedContent"), e.response.status_code) + raise completion_stream: Final = ModelResponseIterator( streaming_response=response.aiter_lines(), @@ -2891,14 +2916,18 @@ def make_sync_call( if client is None: client = HTTPHandler() # Create a new client if none provided - response: Final = client.post( - url_for(api_base, request_body), - headers=headers, - data=data, - stream=True, - logging_obj=logging_obj, - timeout=timeout, - ) + try: + response: Final = client.post( + url_for(api_base, request_body), + headers=headers, + data=data, + stream=True, + logging_obj=logging_obj, + timeout=timeout, + ) + except httpx.HTTPStatusError as e: + known_cached_contents.forget_rejected(request_body.get("cachedContent"), e.response.status_code) + raise if response.status_code != 200 and response.status_code != 201: raise VertexAIError( @@ -2908,14 +2937,18 @@ def make_sync_call( ) def open_hop(body: RequestBody) -> httpx.Response: - return client.post( - url_for(api_base, body), - headers=dict(hop_headers()), - data=json.dumps(body), - stream=True, - logging_obj=logging_obj, - timeout=timeout, - ) + try: + return client.post( + url_for(api_base, body), + headers=dict(hop_headers()), + data=json.dumps(body), + stream=True, + logging_obj=logging_obj, + timeout=timeout, + ) + except httpx.HTTPStatusError as e: + known_cached_contents.forget_rejected(body.get("cachedContent"), e.response.status_code) + raise completion_stream: Final = ModelResponseIterator( streaming_response=response.iter_lines(), @@ -2983,14 +3016,14 @@ class VertexLLM(VertexBase): custom_llm_provider: Literal[ "vertex_ai", "vertex_ai_beta", "gemini" ], # if it's vertex_ai or gemini (google ai studio) - messages: list, + messages: list[AllMessageValues], model_response: ModelResponse, print_verbose: Callable, data: dict, timeout: float | httpx.Timeout | None, encoding: object, - logging_obj, - stream, + logging_obj: LoggingClass, + stream: bool | None, optional_params: dict, litellm_params: dict, logger_fn=None, @@ -3033,6 +3066,17 @@ class VertexLLM(VertexBase): litellm_params=litellm_params, ) + rebuild: Final = cache_rebuild( + partial( + async_transform_request_body, + **data, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_auth_header=auth_header, + ), + messages, + optional_params, + ) request_body: Final = await async_transform_request_body( **data, vertex_project=vertex_project, @@ -3041,39 +3085,37 @@ class VertexLLM(VertexBase): ) ## LOGGING - logging_obj.pre_call( - input=messages, - api_key="", - additional_args={ - "complete_input_dict": request_body, - "api_base": api_base, - "headers": headers, - }, + log_request: Final = partial(_log_generate_request, logging_obj, messages, api_base, headers) + log_request(request_body) + + open_stream: Final = partial( + make_call, + gemini_client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None), + api_base=api_base, + headers=headers, + model=model, + messages=messages, + logging_obj=logging_obj, + hop_headers=partial( + self._fresh_headers_async, + headers, + extra_headers, + vertex_credentials, + vertex_project, + custom_llm_provider, + ), + timeout=timeout, ) - request_body_str: Final = json.dumps(request_body) + async def open_with_fresh_cache(client: AsyncHTTPHandler | None) -> ModelResponseIterator: + async def post(body: RequestBody) -> ModelResponseIterator: + return await open_stream(client=client, data=json.dumps(body), request_body=body) + + return (await post_with_fresh_cache_async(request_body, rebuild, post, log_request))[1] + streaming_response: Final = CustomStreamWrapper( completion_stream=None, - make_call=partial( - make_call, - gemini_client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None), - api_base=api_base, - headers=headers, - data=request_body_str, - model=model, - messages=messages, - logging_obj=logging_obj, - request_body=request_body, - hop_headers=partial( - self._fresh_headers_async, - headers, - extra_headers, - vertex_credentials, - vertex_project, - custom_llm_provider, - ), - timeout=timeout, - ), + make_call=open_with_fresh_cache, model=model, custom_llm_provider="vertex_ai_beta", logging_obj=logging_obj, @@ -3083,7 +3125,7 @@ class VertexLLM(VertexBase): async def async_completion( self, model: str, - messages: list, + messages: list[AllMessageValues], model_response: ModelResponse, print_verbose: Callable, data: dict, @@ -3093,7 +3135,7 @@ class VertexLLM(VertexBase): timeout: float | httpx.Timeout | None, encoding: object, logging_obj: LoggingClass, - stream, + stream: bool | None, optional_params: dict, litellm_params: dict, logger_fn=None, @@ -3136,6 +3178,17 @@ class VertexLLM(VertexBase): litellm_params=litellm_params, ) + rebuild: Final = cache_rebuild( + partial( + async_transform_request_body, + **data, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_auth_header=auth_header, + ), + messages, + optional_params, + ) request_body: Final = await async_transform_request_body( **data, vertex_project=vertex_project, @@ -3151,15 +3204,8 @@ class VertexLLM(VertexBase): else: client = client ## LOGGING - logging_obj.pre_call( - input=messages, - api_key="", - additional_args={ - "complete_input_dict": request_body, - "api_base": api_base, - "headers": headers, - }, - ) + log_request: Final = partial(_log_generate_request, logging_obj, messages, api_base, headers) + log_request(request_body) fresh_headers: Final = partial( self._fresh_headers_async, headers, extra_headers, vertex_credentials, vertex_project, custom_llm_provider @@ -3175,6 +3221,7 @@ class VertexLLM(VertexBase): ) response.raise_for_status() except httpx.HTTPStatusError as err: + known_cached_contents.forget_rejected(body.get("cachedContent"), err.response.status_code) raise VertexAIError( status_code=err.response.status_code, message=err.response.text, @@ -3188,14 +3235,14 @@ class VertexLLM(VertexBase): ) return response - first_response: Final = await post(request_body) + sent_body, first_response = await post_with_fresh_cache_async(request_body, rebuild, post, log_request) return _transform_long_decode_hops( model=model, - hops=await collect_hops_async(parse_body(first_response), request_body, post), + hops=await collect_hops_async(parse_body(first_response), sent_body, post), first_response=first_response, model_response=model_response, logging_obj=logging_obj, - request_body=request_body, + request_body=sent_body, messages=messages, ) @@ -3320,6 +3367,17 @@ class VertexLLM(VertexBase): ) ## TRANSFORMATION ## + rebuild: Final = cache_rebuild( + partial( + sync_transform_request_body, + **transform_request_params, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_auth_header=auth_header, + ), + messages, + optional_params, + ) data: Final = sync_transform_request_body( **transform_request_params, vertex_project=vertex_project, @@ -3328,15 +3386,8 @@ class VertexLLM(VertexBase): ) ## LOGGING - logging_obj.pre_call( - input=messages, - api_key="", - additional_args={ - "complete_input_dict": data, - "api_base": url, - "headers": headers, - }, - ) + log_request: Final = partial(_log_generate_request, logging_obj, messages, url, headers) + log_request(data) fresh_headers: Final = partial( self._fresh_headers, headers, extra_headers, vertex_credentials, vertex_project, custom_llm_provider @@ -3344,22 +3395,27 @@ class VertexLLM(VertexBase): ## SYNC STREAMING CALL ## if stream is True: - request_data_str: Final = json.dumps(data) + open_stream: Final = partial( + make_sync_call, + gemini_client=(client if client is not None and isinstance(client, HTTPHandler) else None), + api_base=url, + model=model, + messages=messages, + logging_obj=logging_obj, + headers=headers, + hop_headers=fresh_headers, + timeout=timeout, + ) + + def open_with_fresh_cache(client: HTTPHandler | None) -> ModelResponseIterator: + def post(body: RequestBody) -> ModelResponseIterator: + return open_stream(client=client, data=json.dumps(body), request_body=body) + + return post_with_fresh_cache_sync(data, rebuild, post, log_request)[1] + streaming_response: Final = CustomStreamWrapper( completion_stream=None, - make_call=partial( - make_sync_call, - gemini_client=(client if client is not None and isinstance(client, HTTPHandler) else None), - api_base=url, - data=request_data_str, - model=model, - messages=messages, - logging_obj=logging_obj, - headers=headers, - request_body=data, - hop_headers=fresh_headers, - timeout=timeout, - ), + make_call=open_with_fresh_cache, model=model, custom_llm_provider="vertex_ai_beta", logging_obj=logging_obj, @@ -3385,6 +3441,7 @@ class VertexLLM(VertexBase): ) response.raise_for_status() except httpx.HTTPStatusError as err: + known_cached_contents.forget_rejected(body.get("cachedContent"), err.response.status_code) raise VertexAIError( status_code=err.response.status_code, message=err.response.text, @@ -3398,14 +3455,14 @@ class VertexLLM(VertexBase): ) return response - first_response: Final = post(data) + sent_body, first_response = post_with_fresh_cache_sync(data, rebuild, post, log_request) return _transform_long_decode_hops( model=model, - hops=collect_hops_sync(parse_body(first_response), data, post), + hops=collect_hops_sync(parse_body(first_response), sent_body, post), first_response=first_response, model_response=model_response, logging_obj=logging_obj, - request_body=data, + request_body=sent_body, messages=messages, ) diff --git a/tests/unit/llms/vertex_ai/context_caching/conftest.py b/tests/unit/llms/vertex_ai/context_caching/conftest.py new file mode 100644 index 00000000000..9c80c964608 --- /dev/null +++ b/tests/unit/llms/vertex_ai/context_caching/conftest.py @@ -0,0 +1,10 @@ +import pytest + +from litellm.llms.vertex_ai.context_caching.known_cached_contents import known_cached_contents + + +@pytest.fixture(autouse=True) +def _forget_known_cached_contents(): + known_cached_contents.clear() + yield + known_cached_contents.clear() diff --git a/tests/unit/llms/vertex_ai/context_caching/test_known_cached_contents.py b/tests/unit/llms/vertex_ai/context_caching/test_known_cached_contents.py new file mode 100644 index 00000000000..0c5c3facb56 --- /dev/null +++ b/tests/unit/llms/vertex_ai/context_caching/test_known_cached_contents.py @@ -0,0 +1,663 @@ +import asyncio +import json +import threading +from datetime import datetime, timezone +from typing import Final +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import openai +import pytest +import respx + +import litellm + +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.vertex_ai.common_utils import VertexAIError +from litellm.llms.vertex_ai.context_caching import vertex_ai_context_caching +from litellm.llms.vertex_ai.context_caching.known_cached_contents import KnownCachedContents, parse_expire_time +from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ContextCachingEndpoints +from litellm.utils import get_prompt_cache_min_tokens + +_REGISTRY_MODULE: Final = "litellm.llms.vertex_ai.context_caching.known_cached_contents" +_SCOPE: Final = ("https://vertex.test/v1/projects/p/locations/us/cachedContents", "") +_NOW: Final = datetime(2026, 10, 14, 18, 0, tzinfo=timezone.utc).timestamp() +_FAR_FUTURE: Final = "2100-01-01T00:00:00.123456789Z" +_GATEWAY: Final = "https://gateway.test/vertex" +_PREFIX_REPEATS: Final = max( + get_prompt_cache_min_tokens(model=f"{provider}/gemini-3.1-pro") for provider in ("vertex_ai", "gemini") +) + + +@pytest.fixture +def clock() -> MagicMock: + return MagicMock(return_value=_NOW) + + +@pytest.fixture +def endpoints(clock: MagicMock) -> ContextCachingEndpoints: + return ContextCachingEndpoints(known_caches=KnownCachedContents(clock=clock)) + + +def _expire_time(seconds_from_now: float) -> str: + expires: Final = datetime.fromtimestamp(_NOW + seconds_from_now, timezone.utc) + return expires.strftime("%Y-%m-%dT%H:%M:%S.%f") + "123Z" + + +def _messages(label: str = "shared") -> list: + text: Final = f"{label} cacheable prefix " * _PREFIX_REPEATS + return [ + {"role": "system", "content": [{"type": "text", "text": text, "cache_control": {"type": "ephemeral"}}]}, + {"role": "user", "content": "question"}, + ] + + +def _response(body: dict) -> MagicMock: + return MagicMock(json=lambda: body) + + +def _created(name: str, expire_time: str | None) -> dict: + body: Final = {"name": name, "usageMetadata": {"totalTokenCount": 2048}} + return body if expire_time is None else {**body, "expireTime": expire_time} + + +def _sync_call( + endpoints: ContextCachingEndpoints, + client: HTTPHandler, + api_key: str = "k", + custom_llm_provider: str = "vertex_ai", + messages=None, + extra_headers: dict | None = None, + api_base: str | None = None, + vertex_project: str = "p", +): + return endpoints.check_and_create_cache( + extra_headers=extra_headers, + messages=messages or _messages(), + optional_params={}, + api_key=api_key, + api_base=api_base, + model="gemini-3.1-pro", + client=client, + timeout=30.0, + logging_obj=MagicMock(model_call_details={}), + custom_llm_provider=custom_llm_provider, + vertex_project=vertex_project, + vertex_location="us", + vertex_auth_header="token", + ) + + +async def _async_call( + endpoints: ContextCachingEndpoints, + client: AsyncHTTPHandler, + timeout: float | httpx.Timeout = 30.0, + api_base: str | None = None, + vertex_project: str = "p", +): + return await endpoints.async_check_and_create_cache( + messages=_messages(), + optional_params={}, + api_key="k", + api_base=api_base, + model="gemini-3.1-pro", + client=client, + timeout=timeout, + logging_obj=MagicMock(model_call_details={}), + custom_llm_provider="vertex_ai", + vertex_project=vertex_project, + vertex_location="us", + vertex_auth_header="token", + ) + + +def test_created_cache_is_reused_without_listing(endpoints: ContextCachingEndpoints): + client: Final = MagicMock(spec=HTTPHandler) + client.get.return_value = _response({}) + client.post.return_value = _response(_created("cachedContents/1", _expire_time(60))) + + assert _sync_call(endpoints, client)[2] == "cachedContents/1" + assert _sync_call(endpoints, client)[2] == "cachedContents/1" + + assert client.get.call_count == 1 + assert client.post.call_count == 1 + + +def test_listed_cache_is_reused_without_listing_again(endpoints: ContextCachingEndpoints): + client: Final = MagicMock(spec=HTTPHandler) + key: Final = vertex_ai_context_caching.local_cache_obj.get_cache_key( + messages=_messages()[:1], tools=None, tool_choice=None, model="gemini-3.1-pro" + ) + client.get.return_value = _response( + {"cachedContents": [{"name": "cachedContents/listed", "displayName": key, "expireTime": _expire_time(60)}]} + ) + + assert _sync_call(endpoints, client)[2] == "cachedContents/listed" + assert _sync_call(endpoints, client)[2] == "cachedContents/listed" + + assert client.get.call_count == 1 + client.post.assert_not_called() + + +def test_listed_cache_past_its_expire_time_is_not_used(endpoints: ContextCachingEndpoints): + client: Final = MagicMock(spec=HTTPHandler) + key: Final = vertex_ai_context_caching.local_cache_obj.get_cache_key( + messages=_messages()[:1], tools=None, tool_choice=None, model="gemini-3.1-pro" + ) + client.get.return_value = _response( + {"cachedContents": [{"name": "cachedContents/stale", "displayName": key, "expireTime": _expire_time(-1)}]} + ) + client.post.return_value = _response(_created("cachedContents/fresh", _expire_time(60))) + + assert _sync_call(endpoints, client)[2] == "cachedContents/fresh" + + +def test_short_lived_cache_is_reused_from_the_listing(endpoints: ContextCachingEndpoints, clock: MagicMock): + client: Final = MagicMock(spec=HTTPHandler) + key: Final = vertex_ai_context_caching.local_cache_obj.get_cache_key( + messages=_messages()[:1], tools=None, tool_choice=None, model="gemini-3.1-pro" + ) + client.get.side_effect = [ + _response({}), + _response( + {"cachedContents": [{"name": "cachedContents/1", "displayName": key, "expireTime": _expire_time(8)}]} + ), + ] + client.post.return_value = _response(_created("cachedContents/1", _expire_time(8))) + + assert _sync_call(endpoints, client)[2] == "cachedContents/1" + clock.return_value = _NOW + 1 + assert _sync_call(endpoints, client)[2] == "cachedContents/1" + + assert client.post.call_count == 1 + + +def test_cache_without_expire_time_is_not_remembered(endpoints: ContextCachingEndpoints): + client: Final = MagicMock(spec=HTTPHandler) + client.get.return_value = _response({}) + client.post.return_value = _response(_created("cachedContents/1", None)) + + _sync_call(endpoints, client) + _sync_call(endpoints, client) + + assert client.get.call_count == 2 + + +def test_remembered_cache_is_listed_again_once_its_reuse_window_closes( + endpoints: ContextCachingEndpoints, clock: MagicMock +): + client: Final = MagicMock(spec=HTTPHandler) + key: Final = vertex_ai_context_caching.local_cache_obj.get_cache_key( + messages=_messages()[:1], tools=None, tool_choice=None, model="gemini-3.1-pro" + ) + client.get.side_effect = [ + _response({}), + _response( + {"cachedContents": [{"name": "cachedContents/listed", "displayName": key, "expireTime": _expire_time(600)}]} + ), + ] + client.post.return_value = _response(_created("cachedContents/1", _expire_time(60))) + assert _sync_call(endpoints, client)[2] == "cachedContents/1" + + clock.return_value = _NOW + 55 + assert _sync_call(endpoints, client)[2] == "cachedContents/listed" + + +def test_gemini_caches_are_not_shared_across_api_keys(endpoints: ContextCachingEndpoints): + client: Final = MagicMock(spec=HTTPHandler) + client.get.return_value = _response({}) + client.post.return_value = _response(_created("cachedContents/1", _expire_time(60))) + + _sync_call(endpoints, client, api_key="key-a", custom_llm_provider="gemini") + _sync_call(endpoints, client, api_key="key-b", custom_llm_provider="gemini") + + assert client.post.call_count == 2 + + +def test_gemini_cache_created_with_an_overridden_api_key_is_not_reused_with_the_configured_key( + endpoints: ContextCachingEndpoints, +): + client: Final = MagicMock(spec=HTTPHandler) + client.get.return_value = _response({}) + client.post.side_effect = [ + _response(_created("cachedContents/caller-project", _expire_time(60))), + _response(_created("cachedContents/deployment-project", _expire_time(60))), + ] + + caller: Final = _sync_call( + endpoints, + client, + api_key="deployment-key", + custom_llm_provider="gemini", + extra_headers={"X-Goog-Api-Key": "caller-key"}, + ) + configured: Final = _sync_call(endpoints, client, api_key="deployment-key", custom_llm_provider="gemini") + + assert (caller[2], configured[2]) == ("cachedContents/caller-project", "cachedContents/deployment-project") + + +def test_vertex_caches_behind_one_custom_api_base_are_not_shared_across_projects(endpoints: ContextCachingEndpoints): + client: Final = MagicMock(spec=HTTPHandler) + client.get.return_value = _response({}) + client.post.side_effect = [ + _response(_created("cachedContents/project-a", _expire_time(60))), + _response(_created("cachedContents/project-b", _expire_time(60))), + ] + + sent: Final = [ + _sync_call(endpoints, client, api_base=_GATEWAY, vertex_project=project)[2] + for project in ("project-a", "project-b", "project-a") + ] + + assert sent == ["cachedContents/project-a", "cachedContents/project-b", "cachedContents/project-a"] + + +@pytest.mark.asyncio +async def test_async_vertex_caches_behind_one_custom_api_base_are_not_shared_across_projects( + endpoints: ContextCachingEndpoints, +): + client: Final = MagicMock(spec=AsyncHTTPHandler) + client.get = AsyncMock(return_value=_response({})) + client.post = AsyncMock( + side_effect=[ + _response(_created("cachedContents/project-a", _expire_time(60))), + _response(_created("cachedContents/project-b", _expire_time(60))), + ] + ) + + sent: Final = [ + (await _async_call(endpoints, client, api_base=_GATEWAY, vertex_project=project))[2] + for project in ("project-a", "project-b", "project-a") + ] + + assert sent == ["cachedContents/project-a", "cachedContents/project-b", "cachedContents/project-a"] + + +def test_different_prefixes_do_not_share_a_cache(endpoints: ContextCachingEndpoints): + client: Final = MagicMock(spec=HTTPHandler) + client.get.return_value = _response({}) + client.post.side_effect = [ + _response(_created("cachedContents/a", _expire_time(60))), + _response(_created("cachedContents/b", _expire_time(60))), + ] + + assert _sync_call(endpoints, client, messages=_messages("first"))[2] == "cachedContents/a" + assert _sync_call(endpoints, client, messages=_messages("second"))[2] == "cachedContents/b" + + +@pytest.mark.asyncio +async def test_concurrent_requests_for_one_prefix_list_and_create_once(endpoints: ContextCachingEndpoints): + client: Final = MagicMock(spec=AsyncHTTPHandler) + + async def slow_list(**_): + await asyncio.sleep(0.01) + return _response({}) + + client.get = AsyncMock(side_effect=slow_list) + client.post = AsyncMock(return_value=_response(_created("cachedContents/1", _expire_time(60)))) + + results: Final = await asyncio.gather(*(_async_call(endpoints, client) for _ in range(5))) + + assert {result[2] for result in results} == {"cachedContents/1"} + assert client.get.await_count == 1 + assert client.post.await_count == 1 + + +@pytest.mark.asyncio +async def test_async_created_cache_is_reused_by_a_later_request(endpoints: ContextCachingEndpoints): + client: Final = MagicMock(spec=AsyncHTTPHandler) + client.get = AsyncMock(return_value=_response({})) + client.post = AsyncMock(return_value=_response(_created("cachedContents/1", _expire_time(60)))) + + first: Final = await _async_call(endpoints, client) + client.get = AsyncMock(side_effect=AssertionError("a known cache must not be listed again")) + second: Final = await _async_call(endpoints, client) + + assert (first[2], second[2]) == ("cachedContents/1", "cachedContents/1") + + +@pytest.mark.asyncio +async def test_async_listed_cache_past_its_expire_time_is_not_used(endpoints: ContextCachingEndpoints): + client: Final = MagicMock(spec=AsyncHTTPHandler) + key: Final = vertex_ai_context_caching.local_cache_obj.get_cache_key( + messages=_messages()[:1], tools=None, tool_choice=None, model="gemini-3.1-pro" + ) + client.get = AsyncMock( + return_value=_response( + {"cachedContents": [{"name": "cachedContents/stale", "displayName": key, "expireTime": _expire_time(-1)}]} + ) + ) + client.post = AsyncMock(return_value=_response(_created("cachedContents/fresh", _expire_time(60)))) + + assert (await _async_call(endpoints, client))[2] == "cachedContents/fresh" + + +@pytest.mark.asyncio +async def test_async_create_http_error_surfaces_its_status_code(endpoints: ContextCachingEndpoints): + client: Final = MagicMock(spec=AsyncHTTPHandler) + client.get = AsyncMock(return_value=_response({})) + request: Final = httpx.Request("POST", "https://vertex.test/v1/projects/p/locations/us/cachedContents") + failed: Final = httpx.Response(429, request=request, text="RESOURCE_EXHAUSTED") + client.post = AsyncMock(return_value=failed) + + with pytest.raises(VertexAIError) as raised: + await _async_call(endpoints, client) + + assert raised.value.status_code == 429 + + +@pytest.mark.asyncio +async def test_waiters_resolve_on_their_own_when_the_first_request_fails(endpoints: ContextCachingEndpoints): + client: Final = MagicMock(spec=AsyncHTTPHandler) + + async def slow_list(**_): + await asyncio.sleep(0.01) + return _response({}) + + client.get = AsyncMock(side_effect=slow_list) + client.post = AsyncMock( + side_effect=[ + VertexAIError(status_code=500, message="create failed"), + _response(_created("cachedContents/2", _expire_time(60))), + ] + ) + + results: Final = await asyncio.gather( + _async_call(endpoints, client), _async_call(endpoints, client), return_exceptions=True + ) + + assert isinstance(results[0], VertexAIError) + assert results[1][2] == "cachedContents/2" + + +@pytest.mark.parametrize("waiter_timeout", [0.05, httpx.Timeout(30.0, read=0.05)]) +@pytest.mark.asyncio +async def test_waiter_resolves_on_its_own_once_its_timeout_passes( + endpoints: ContextCachingEndpoints, waiter_timeout: float | httpx.Timeout +): + client: Final = MagicMock(spec=AsyncHTTPHandler) + + async def slow_list(**_): + await asyncio.sleep(0.5) + return _response({}) + + client.get = AsyncMock(side_effect=slow_list) + client.post = AsyncMock(return_value=_response(_created("cachedContents/1", _expire_time(60)))) + + first, waiter = await asyncio.gather(_async_call(endpoints, client), _async_call(endpoints, client, waiter_timeout)) + + assert (first[2], waiter[2]) == ("cachedContents/1", "cachedContents/1") + assert client.get.await_count == 2 + + +def test_parse_expire_time_accepts_nanosecond_precision(): + assert parse_expire_time("2026-10-14T18:06:21.123456789Z") == pytest.approx( + datetime(2026, 10, 14, 18, 6, 21, 123456, tzinfo=timezone.utc).timestamp() + ) + assert ( + parse_expire_time("2026-10-14T18:06:21Z") == datetime(2026, 10, 14, 18, 6, 21, tzinfo=timezone.utc).timestamp() + ) + assert parse_expire_time("not a time") is None + assert parse_expire_time(None) is None + + +def test_parse_expire_time_reads_a_missing_offset_as_utc(): + assert ( + parse_expire_time("2026-10-14T18:06:21") == datetime(2026, 10, 14, 18, 6, 21, tzinfo=timezone.utc).timestamp() + ) + + +def test_cache_expiring_within_the_margin_is_not_remembered(): + known: Final = KnownCachedContents(clock=lambda: _NOW) + known.remember(_SCOPE, "key", "cachedContents/soon", _expire_time(3)) + + assert known.get(_SCOPE, "key") is None + + +def test_registry_is_bounded(): + known: Final = KnownCachedContents(clock=lambda: _NOW) + with patch(f"{_REGISTRY_MODULE}.MAX_KNOWN_CACHED_CONTENTS", 3): + for index in range(5): + known.remember(_SCOPE, f"key-{index}", f"cachedContents/{index}", _expire_time(60)) + + assert known.get(_SCOPE, "key-0") is None + assert known.get(_SCOPE, "key-4") == "cachedContents/4" + + +def _generated(stream: bool) -> httpx.Response: + body: Final = { + "candidates": [{"content": {"role": "model", "parts": [{"text": "ok"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 1, "totalTokenCount": 11}, + } + if stream: + return httpx.Response(200, text=f"data: {json.dumps(body)}\n\n", headers={"content-type": "text/event-stream"}) + return httpx.Response(200, json=body) + + +def _rejected(status_code: int, message: str) -> httpx.Response: + return httpx.Response(status_code, json={"error": {"code": status_code, "message": message}}) + + +async def _complete(sync_mode: bool, stream: bool, **kwargs) -> str: + request: Final = {"model": "gemini/gemini-3.1-pro", "messages": _messages(), "stream": stream, **kwargs} + if sync_mode: + response = litellm.completion(**request) + if stream: + return "".join(chunk.choices[0].delta.content or "" for chunk in response) + return response.choices[0].message.content + response = await litellm.acompletion(**request) + if stream: + return "".join([chunk.choices[0].delta.content or "" async for chunk in response]) + return response.choices[0].message.content + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.asyncio +async def test_rejected_cache_is_replaced_when_a_request_param_cannot_be_deep_copied( + respx_mock: respx.MockRouter, + fake_provider_credentials: None, + monkeypatch: pytest.MonkeyPatch, + sync_mode: bool, + stream: bool, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.get(url__regex=r".*/cachedContents.*").mock(return_value=httpx.Response(200, json={})) + respx_mock.post(url__regex=r".*/cachedContents$").mock( + side_effect=[ + httpx.Response(200, json=_created("cachedContents/deleted", _FAR_FUTURE)), + httpx.Response(200, json=_created("cachedContents/fresh", _FAR_FUTURE)), + ] + ) + generate_route: Final = respx_mock.post(url__regex=r".*:(stream)?[gG]enerateContent.*").mock( + side_effect=[_generated(stream), _rejected(404, "CachedContent not found"), _generated(stream)] + ) + lock: Final = threading.Lock() + + assert await _complete(sync_mode, stream, lock_holder=lock) == "ok" + assert await _complete(sync_mode, stream, lock_holder=lock) == "ok" + + sent: Final = [json.loads(call.request.content)["cachedContent"] for call in generate_route.calls] + assert sent == ["cachedContents/deleted", "cachedContents/deleted", "cachedContents/fresh"] + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.parametrize( + ("status_code", "message", "replaced"), + [ + (400, "Invalid resource state for cache content 4484391867902853120.", True), + (403, "CachedContent not found (or permission denied)", True), + (404, "CachedContent not found", True), + (400, "Invalid value at 'generation_config.temperature'", False), + (500, "CachedContent not found", False), + ], +) +@pytest.mark.asyncio +async def test_remembered_cache_rejected_by_generate_content_is_replaced_within_the_same_call( + respx_mock: respx.MockRouter, + fake_provider_credentials: None, + monkeypatch: pytest.MonkeyPatch, + sync_mode: bool, + stream: bool, + status_code: int, + message: str, + replaced: bool, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + list_route: Final = respx_mock.get(url__regex=r".*/cachedContents.*").mock( + return_value=httpx.Response(200, json={}) + ) + respx_mock.post(url__regex=r".*/cachedContents$").mock( + side_effect=[ + httpx.Response(200, json=_created("cachedContents/deleted", _FAR_FUTURE)), + httpx.Response(200, json=_created("cachedContents/fresh", _FAR_FUTURE)), + ] + ) + generate_route: Final = respx_mock.post(url__regex=r".*:(stream)?[gG]enerateContent.*").mock( + side_effect=[_generated(stream), _rejected(status_code, message), _generated(stream)] + ) + logged: Final[list[dict]] = [] + + assert await _complete(sync_mode, stream) == "ok" + if replaced: + assert await _complete(sync_mode, stream, logger_fn=logged.append) == "ok" + else: + with pytest.raises(openai.APIStatusError): + await _complete(sync_mode, stream, logger_fn=logged.append) + + sent: Final = [json.loads(call.request.content)["cachedContent"] for call in generate_route.calls] + assert sent == ["cachedContents/deleted", "cachedContents/deleted"] + (["cachedContents/fresh"] if replaced else []) + assert list_route.call_count == (2 if replaced else 1) + assert "enerateContent" in logged[-1]["litellm_params"]["api_base"] + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.parametrize( + ("status_code", "message", "forgotten"), + [ + (400, "Invalid value at 'generation_config.temperature'", True), + ( + 400, + "Tool config, tools and system instruction should not be set in therequest when using cached content.", + True, + ), + ( + 400, + "CachedContent can not be used with GenerateContent request setting system_instruction, tools or tool_config.", + True, + ), + (404, "Requested entity was not found.", True), + (429, "Resource has been exhausted (e.g. check quota).", False), + (500, "Internal error encountered.", False), + ], +) +@pytest.mark.asyncio +async def test_error_that_does_not_name_a_missing_cache_is_not_retried( + respx_mock: respx.MockRouter, + fake_provider_credentials: None, + monkeypatch: pytest.MonkeyPatch, + sync_mode: bool, + stream: bool, + status_code: int, + message: str, + forgotten: bool, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + list_route: Final = respx_mock.get(url__regex=r".*/cachedContents.*").mock( + return_value=httpx.Response(200, json={}) + ) + respx_mock.post(url__regex=r".*/cachedContents$").mock( + side_effect=[ + httpx.Response(200, json=_created("cachedContents/first", _FAR_FUTURE)), + httpx.Response(200, json=_created("cachedContents/second", _FAR_FUTURE)), + ] + ) + generate_route: Final = respx_mock.post(url__regex=r".*:(stream)?[gG]enerateContent.*").mock( + side_effect=[_generated(stream), _rejected(status_code, message), _generated(stream)] + ) + + assert await _complete(sync_mode, stream) == "ok" + with pytest.raises(openai.APIStatusError): + await _complete(sync_mode, stream) + assert await _complete(sync_mode, stream) == "ok" + + third: Final = "cachedContents/second" if forgotten else "cachedContents/first" + assert [json.loads(call.request.content)["cachedContent"] for call in generate_route.calls] == [ + "cachedContents/first", + "cachedContents/first", + third, + ] + assert list_route.call_count == (2 if forgotten else 1) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.asyncio +async def test_rejected_caller_supplied_cached_content_is_not_sent_again( + respx_mock: respx.MockRouter, + fake_provider_credentials: None, + monkeypatch: pytest.MonkeyPatch, + sync_mode: bool, + stream: bool, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + list_route: Final = respx_mock.get(url__regex=r".*/cachedContents.*").mock( + return_value=httpx.Response(200, json={}) + ) + generate_route: Final = respx_mock.post(url__regex=r".*:(stream)?[gG]enerateContent.*").mock( + side_effect=[_rejected(404, "CachedContent not found"), _generated(stream)] + ) + + with pytest.raises(openai.APIStatusError): + await _complete(sync_mode, stream, cached_content="cachedContents/callers") + + assert [json.loads(call.request.content)["cachedContent"] for call in generate_route.calls] == [ + "cachedContents/callers" + ] + assert list_route.call_count == 0 + + +def _cut_off_stream() -> httpx.Response: + body: Final = { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "ok"}]}, + "finishReason": "CONTINUATION", + "continuationToken": "dG9rLTE=", + } + ], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 1, "totalTokenCount": 11}, + } + return httpx.Response(200, text=f"data: {json.dumps(body)}\n\n", headers={"content-type": "text/event-stream"}) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_cache_rejected_on_a_resume_hop_is_not_sent_again( + respx_mock: respx.MockRouter, + fake_provider_credentials: None, + monkeypatch: pytest.MonkeyPatch, + sync_mode: bool, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + list_route: Final = respx_mock.get(url__regex=r".*/cachedContents.*").mock( + return_value=httpx.Response(200, json={}) + ) + respx_mock.post(url__regex=r".*/cachedContents$").mock( + side_effect=[ + httpx.Response(200, json=_created("cachedContents/deleted", _FAR_FUTURE)), + httpx.Response(200, json=_created("cachedContents/fresh", _FAR_FUTURE)), + ] + ) + generate_route: Final = respx_mock.post(url__regex=r".*:streamGenerateContent.*").mock( + side_effect=[_cut_off_stream(), _rejected(404, "CachedContent not found"), _generated(True)] + ) + + await _complete(sync_mode, True, max_tokens=100_000) + await _complete(sync_mode, True, max_tokens=100_000) + + sent: Final = [json.loads(call.request.content)["cachedContent"] for call in generate_route.calls] + assert sent == ["cachedContents/deleted", "cachedContents/deleted", "cachedContents/fresh"] + assert list_route.call_count == 2 From d776277647c4c6f2ab7f576222d00e14e2927e95 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 12:59:37 +0000 Subject: [PATCH 05/33] fix(responses): count billed background responses toward the creator's project and per-model budgets (#47504) * test(responses): billed background responses count toward the creator's project and model budgets * fix(responses): carry the creator's project and per-model budgets on the billed background response * fix(responses): count a key's budget tier toward background per-model budgets * fix(responses): bill a background response's project only while the key stays on its team * revert(responses): keep billing the creating key's project after a team move A key can only keep or detach its project, so the billed project always pairs with the team it was created under. This reverts commit ded4f084c5. * fix(responses): bill a background response under the model name the client sent, with its project alias * fix(responses): name the team's public model as a billed background response's model group The model group lookup now passes the creating key's team, so a team model or team wildcard bills under the public name a normal call logs and per-model budgets on that name move --- .../common_utils/check_responses_cost.py | 26 +++ .../managed_object_attribution.py | 55 ++++- .../test_background_response_billed_once.py | 54 +++++ .../common_utils/test_check_responses_cost.py | 209 +++++++++++++++++- 4 files changed, 330 insertions(+), 14 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py index 00779bfe8ec..1fa461d7ab7 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py @@ -59,6 +59,19 @@ if TYPE_CHECKING: from litellm.router import Router _STORED_RESPONSE_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object]) + + +class _RoutedModelInfo(BaseModel): + model_config = ConfigDict(frozen=True) + id: str | None = None + + +class _RoutedDeployment(BaseModel): + model_config = ConfigDict(frozen=True) + model_info: _RoutedModelInfo = _RoutedModelInfo() + + +_ROUTED_DEPLOYMENTS: Final = TypeAdapter(tuple[_RoutedDeployment, ...]) _RETRIEVE_RESPONSE_CALL_TYPE: Final = "aget_responses" @@ -314,10 +327,23 @@ class CheckResponsesCost: INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, **target.model_metadata, **(_deployment_metadata(deployment) if deployment is not None else {}), + **self._requested_model_group(target, job.team_id), } ), ) + def _requested_model_group(self, target: _PollTarget, team_id: str | None) -> Mapping[str, object]: + requested: Final = target.model_metadata.get("model_group") + if target.deployment is None or not isinstance(requested, str): + return MappingProxyType({}) + routed: Final = _ROUTED_DEPLOYMENTS.validate_python( + self.llm_router.get_model_list(model_name=requested, team_id=team_id) or () + ) + serving_id: Final = target.deployment.model_info.id + return MappingProxyType( + {"model_group": requested} if any(row.model_info.id == serving_id for row in routed) else {} + ) + async def _bill_once( self, job: ManagedObjectRow, diff --git a/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py b/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py index 3c7735b53c0..16b1b0e973c 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py @@ -107,8 +107,10 @@ class ManagedObjectCreatorAttribution: user's alias, because a resolvable name is more useful on the spend row than a null. user_api_key_org_id is resolved too: the spend update writer reads it off this metadata to increment organization spend, so leaving it out drops the cost from org accounting. - A background response also carries its end user and the key, team, and user budgets, the - way the read that billed it before did, and never names a key-scoped creator as its user. + A background response also carries its end user, the creating key's project and its alias, + and the key, team, and user budgets and per-model budgets, the way the read that billed it + before did, and never names a key-scoped creator as its user. A key that sets no budget of + its own takes its budget tier's, the same rule auth applies. """ def __init__(self, prisma_client: PrismaClient, poller_name: str) -> None: @@ -126,9 +128,23 @@ class ManagedObjectCreatorAttribution: "user_api_key_user_id": _creator_user_id(job.created_by), **_shared_metadata(job, rows), **({"user_api_key_end_user_id": end_user_id} if end_user_id else {}), - **_budget_fields(rows.key, "user_api_key_spend", "user_api_key_max_budget"), - **_budget_fields(rows.team, "user_api_key_team_spend", "user_api_key_team_max_budget"), - **_budget_fields(rows.user, "user_api_key_user_spend", "user_api_key_user_max_budget"), + "user_api_key_project_id": getattr(rows.key, "project_id", None), + "user_api_key_project_alias": getattr( + getattr(rows.key, "litellm_project_table", None), "project_alias", None + ), + **_key_budget_fields(rows.key), + **_budget_fields( + rows.team, + "user_api_key_team_spend", + "user_api_key_team_max_budget", + "user_api_key_team_model_max_budget", + ), + **_budget_fields( + rows.user, + "user_api_key_user_spend", + "user_api_key_user_max_budget", + "user_api_key_user_model_max_budget", + ), } async def _rows(self, job: ManagedObjectRow, object_id: str) -> _CreatorRows: @@ -155,7 +171,9 @@ class ManagedObjectCreatorAttribution: if not api_key or _is_cli_session_key(api_key, created_by): return None try: - return await _token_table(self._prisma_client).find_unique(where={"token": api_key}) + return await _token_table(self._prisma_client).find_unique( + where={"token": api_key}, include={"litellm_budget_table": True, "litellm_project_table": True} + ) except Exception as e: verbose_proxy_logger.error(f"{self._poller_name}: could not look up the creating key for {object_id}: {e}") return None @@ -212,7 +230,28 @@ def _organization_id(row: object) -> str | None: return cast("str | None", getattr(row, "organization_id", None)) if row is not None else None -def _budget_fields(row: object, spend_key: str, max_budget_key: str) -> Mapping[str, object]: +def _budget_fields(row: object, spend_key: str, max_budget_key: str, model_max_budget_key: str) -> Mapping[str, object]: if row is None: return {} - return {spend_key: getattr(row, "spend", None), max_budget_key: getattr(row, "max_budget", None)} + return { + spend_key: getattr(row, "spend", None), + max_budget_key: getattr(row, "max_budget", None), + model_max_budget_key: getattr(row, "model_max_budget", None), + } + + +def _key_budget_fields(key: object) -> Mapping[str, object]: + if key is None: + return {} + tier: Final = getattr(key, "litellm_budget_table", None) + return { + "user_api_key_spend": getattr(key, "spend", None), + "user_api_key_max_budget": _own_or_tier_budget(key, tier, "max_budget"), + "user_api_key_model_max_budget": _own_or_tier_budget(key, tier, "model_max_budget"), + } + + +def _own_or_tier_budget(key: object, tier: object, field: str) -> object: + own: Final = getattr(key, field, None) + tier_value: Final = getattr(tier, field, None) + return own if tier_value is None or own not in (None, {}) else tier_value diff --git a/tests/integration/spend/test_background_response_billed_once.py b/tests/integration/spend/test_background_response_billed_once.py index 8a9818a7532..6d481a98fd0 100644 --- a/tests/integration/spend/test_background_response_billed_once.py +++ b/tests/integration/spend/test_background_response_billed_once.py @@ -5,6 +5,7 @@ from collections.abc import Callable, Mapping from concurrent.futures import ThreadPoolExecutor from typing import Final +import httpx import openai import pytest from integration._support.client import Gateway, Scenario, eventually, held, list_value, object_value, string_value @@ -611,6 +612,59 @@ def test_a_background_response_bills_the_end_user_who_created_it_once(gateway: G assert _charged_end_users(upstream.model_id) == [end_user] +def _project_spend(project_id: str) -> float: + rows: Final = read_rows('SELECT spend FROM "LiteLLM_ProjectTable" WHERE project_id = %s', (project_id,)) + return number(rows[0]["spend"]) + + +def _next_create(scenario: Scenario, key: str, model: str) -> httpx.Response: + response: Final = scenario.gateway.request( + "POST", "/v1/responses", {"model": model, "input": PROMPT, "background": True, "store": False}, key=key + ) + if response.status_code == 200: + scenario.cleanups.callback(forget_row, string_value(object_value(response.json())["id"])) + return response + + +def _key_budget(scenario: Scenario, budget_owner: str, model_budget: JsonValue) -> Mapping[str, JsonValue]: + if budget_owner == "key": + return {"model_max_budget": model_budget} + if budget_owner == "user": + return {"user_id": scenario.user(model_max_budget=model_budget)} + if budget_owner == "budget_tier": + return {"budget_id": scenario.budget(model_max_budget=model_budget)} + return {} + + +@pytest.mark.timeout(240) +@pytest.mark.parametrize("budget_owner", ["key", "team", "user", "budget_tier"]) +def test_a_billed_background_response_counts_toward_its_creators_project_and_model_budget( + gateway: Gateway, budget_owner: str +) -> None: + with gateway.scenario() as scenario: + upstream: Final = scripted_upstream(scenario, background_routes("completed")) + model_budget: Final[JsonValue] = {upstream.model: {"max_budget": PRICE / 2, "budget_duration": "1d"}} + team_id: Final = scenario.team(**({"model_max_budget": model_budget} if budget_owner == "team" else {})) + project_id: Final = scenario.project(team_id) + key: Final = scenario.key( + team_id=team_id, project_id=project_id, **_key_budget(scenario, budget_owner, model_budget) + ) + response_id: Final = create_background_response(scenario, key, upstream.model) + read: Final = read_response(gateway, key, response_id) + assert (read,) in (settling_reads("completed", 1), free_reads("completed", 1)) + eventually( + lambda: (key_spend(key), _project_spend(project_id)), + lambda spends: spends[0] > 0 and costs(spends[0])(spends[1]), + seconds=SPEND_LANDING_SECONDS, + ) + refused: Final = eventually( + lambda: _next_create(scenario, key, upstream.model), + lambda response: response.status_code != 200, + seconds=SPEND_LANDING_SECONDS, + ) + assert f"exceeded budget for model={upstream.model}" in refused.text, refused.text + + @pytest.mark.timeout(240) def test_the_read_that_settles_a_background_response_splits_its_charge_into_input_and_output( gateway: Gateway, diff --git a/tests/unit/proxy/common_utils/test_check_responses_cost.py b/tests/unit/proxy/common_utils/test_check_responses_cost.py index f606658fa09..d78c3ef6185 100644 --- a/tests/unit/proxy/common_utils/test_check_responses_cost.py +++ b/tests/unit/proxy/common_utils/test_check_responses_cost.py @@ -7,7 +7,7 @@ import dataclasses from collections.abc import Mapping from dataclasses import dataclass from datetime import datetime -from types import MappingProxyType +from types import MappingProxyType, SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock, Mock, patch @@ -1217,8 +1217,6 @@ class TestCheckResponsesCost: """The poll's billing event is the one billed event of a background job, so its spend row must name the key, team, org, user, and tags of the create it settles, read off the managed object row, never the poller itself.""" - from types import SimpleNamespace - mock_job = MagicMock() mock_job.unified_object_id = "resp_test_attributed" mock_job.id = "job-attributed" @@ -1494,6 +1492,122 @@ async def test_the_billing_event_carries_the_creators_end_user(): assert billed.metadata["user_api_key_end_user_id"] == "customer-42" +@pytest.mark.asyncio +async def test_the_billing_event_carries_the_creating_keys_project_and_every_creator_model_budget(): + key_budget: Final = {"gpt-5": {"max_budget": 1.0, "budget_duration": "1d"}} + team_budget: Final = {"gpt-5": {"max_budget": 2.0, "budget_duration": "7d"}} + user_budget: Final = {"gpt-5": {"max_budget": 3.0, "budget_duration": "30d"}} + table: Final = _InMemoryManagedObjectTable( + rows=( + _StoredResponse( + id="job-budgets", + unified_object_id="resp_budgets", + status="queued", + created_by="alice", + api_key="hash-alice", + team_id="team-alpha", + ), + ) + ) + billed_logger: Final = _RecordingBilledLogger() + poller: Final = _poller_over(table, billed_logger) + poller.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=SimpleNamespace( + key_alias=None, + organization_id=None, + spend=0.0, + max_budget=None, + project_id="project-1", + litellm_project_table=SimpleNamespace(project_alias="project-one"), + model_max_budget=key_budget, + ) + ) + poller.prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=SimpleNamespace( + team_alias=None, organization_id=None, spend=0.0, max_budget=None, model_max_budget=team_budget + ) + ) + poller.prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=SimpleNamespace( + user_email=None, user_alias=None, spend=0.0, max_budget=None, model_max_budget=user_budget + ) + ) + + await poller.bill_if_finished(response_id="resp_budgets", response=_terminal_response()) + + (billed,) = billed_logger.billed + assert { + field: billed.metadata.get(field) + for field in ( + "user_api_key_project_id", + "user_api_key_project_alias", + "user_api_key_model_max_budget", + "user_api_key_team_model_max_budget", + "user_api_key_user_model_max_budget", + ) + } == { + "user_api_key_project_id": "project-1", + "user_api_key_project_alias": "project-one", + "user_api_key_model_max_budget": key_budget, + "user_api_key_team_model_max_budget": team_budget, + "user_api_key_user_model_max_budget": user_budget, + } + + +_TIER_MODEL_BUDGET: Final = {"gpt-5": {"max_budget": 4.0, "budget_duration": "1d"}} +_KEY_MODEL_BUDGET: Final = {"gpt-5": {"max_budget": 1.0, "budget_duration": "1d"}} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("key_max_budget", "key_model_max_budget", "billed_max_budget", "billed_model_max_budget"), + [ + (None, {}, 5.0, _TIER_MODEL_BUDGET), + (None, None, 5.0, _TIER_MODEL_BUDGET), + (2.0, _KEY_MODEL_BUDGET, 2.0, _KEY_MODEL_BUDGET), + ], + ids=["key_with_an_empty_budget_map", "key_with_no_budget_map", "key_with_its_own_budgets"], +) +async def test_the_billing_event_takes_the_keys_budget_tier_only_where_the_key_sets_no_budget_of_its_own( + key_max_budget: float | None, + key_model_max_budget: dict[str, object] | None, + billed_max_budget: float, + billed_model_max_budget: dict[str, object], +): + table: Final = _InMemoryManagedObjectTable( + rows=( + _StoredResponse( + id="job-tier", + unified_object_id="resp_tier", + status="queued", + created_by="key:hash-tier", + api_key="hash-tier", + ), + ) + ) + billed_logger: Final = _RecordingBilledLogger() + poller: Final = _poller_over(table, billed_logger) + poller.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=SimpleNamespace( + key_alias=None, + organization_id=None, + spend=0.0, + max_budget=key_max_budget, + project_id=None, + model_max_budget=key_model_max_budget, + litellm_budget_table=SimpleNamespace(max_budget=5.0, model_max_budget=_TIER_MODEL_BUDGET), + ) + ) + + await poller.bill_if_finished(response_id="resp_tier", response=_terminal_response()) + + (billed,) = billed_logger.billed + assert (billed.metadata.get("user_api_key_max_budget"), billed.metadata.get("user_api_key_model_max_budget")) == ( + billed_max_budget, + billed_model_max_budget, + ) + + @pytest.mark.asyncio @pytest.mark.parametrize( ("created_by", "billed_user_id"), @@ -1599,6 +1713,10 @@ async def _bill_through_a_priced_deployment( read: ResponsesAPIResponse, litellm_params: Mapping[str, object] = _PRICED_LITELLM_PARAMS, model_info: Mapping[str, object] = MappingProxyType({}), + model_name: str = "priced-gpt", + requested_model: str = "priced-gpt", + model_group_alias: Mapping[str, str] = MappingProxyType({}), + team_id: str | None = None, ) -> _PricedBilling: import litellm from litellm_enterprise.proxy.common_utils.check_responses_cost import CheckResponsesCost @@ -1610,7 +1728,16 @@ async def _bill_through_a_priced_deployment( custom_llm_provider="openai", model_id="deployment-custom", response_id="resp_upstream" ) table: Final = _InMemoryManagedObjectTable( - rows=(_StoredResponse(id="job-priced", unified_object_id=response_id, status="queued", api_key="hash-alice"),) + rows=( + _StoredResponse( + id="job-priced", + unified_object_id=response_id, + status="queued", + api_key="hash-alice", + file_object={"model": requested_model}, + team_id=team_id, + ), + ) ) prisma_client: Final = MagicMock() prisma_client.db.litellm_managedobjecttable = table @@ -1620,11 +1747,12 @@ async def _bill_through_a_priced_deployment( router: Final = litellm.Router( model_list=[ { - "model_name": "priced-gpt", + "model_name": model_name, "litellm_params": dict(litellm_params), "model_info": {"id": "deployment-custom", **model_info}, } - ] + ], + model_group_alias=dict(model_group_alias), ) charge: Final = await CheckResponsesCost( proxy_logging_obj=MagicMock(), prisma_client=prisma_client, llm_router=router @@ -1657,6 +1785,75 @@ async def test_the_billing_event_prices_the_fetched_usage_at_the_deployment_rate assert billing.table.statuses == {"job-priced": "completed"} +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model_name", "deployment_model", "requested_model", "model_group_alias", "model_info", "model_group"), + [ + ("openai/*", "openai/*", "openai/gpt-5.6", {}, {}, "openai/gpt-5.6"), + ("priced-gpt", "openai/gpt-5.6", "priced-alias", {"priced-alias": "priced-gpt"}, {}, "priced-alias"), + ("priced-gpt", "openai/gpt-5.6", "gpt-5.6-2026-01-01", {}, {}, "priced-gpt"), + ( + "model_name_team-1_abc", + "openai/gpt-5.6", + "team-gpt", + {}, + {"team_id": "team-1", "team_public_model_name": "team-gpt"}, + "team-gpt", + ), + ( + "model_name_team-1_wild", + "openai/*", + "openai/gpt-5.6", + {}, + {"team_id": "team-1", "team_public_model_name": "openai/*"}, + "openai/gpt-5.6", + ), + ( + "model_name_team-2_abc", + "openai/gpt-5.6", + "team-gpt", + {}, + {"team_id": "team-2", "team_public_model_name": "team-gpt"}, + "model_name_team-2_abc", + ), + ], + ids=[ + "wildcard_route", + "model_group_alias", + "fallback_kept_the_provider_model_name", + "team_model", + "team_wildcard", + "another_teams_model", + ], +) +async def test_the_billing_event_names_the_requested_model_as_its_model_group_only_while_it_routes_to_the_deployment( + monkeypatch: pytest.MonkeyPatch, + model_name: str, + deployment_model: str, + requested_model: str, + model_group_alias: Mapping[str, str], + model_info: Mapping[str, object], + model_group: str, +): + import litellm + + capture: Final = _spend_capture_logger() + monkeypatch.setattr(litellm, "_async_success_callback", [capture]) + + await _bill_through_a_priced_deployment( + _terminal_response(), + litellm_params={**_PRICED_LITELLM_PARAMS, "model": deployment_model}, + model_info=model_info, + model_name=model_name, + requested_model=requested_model, + model_group_alias=model_group_alias, + team_id="team-1", + ) + + (payload,) = capture.payloads + assert payload["model_group"] == model_group + + @pytest.mark.asyncio @pytest.mark.parametrize("base_model_in", ["litellm_params", "model_info"]) async def test_the_billing_event_prices_the_usage_at_the_deployments_base_model( From dd20bb0ae9bd8f59ce4d34aadeac62c4cbda1b9f Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 20 Oct 2026 13:01:45 +0000 Subject: [PATCH 06/33] fix(bedrock): apply the 16 maxTokens floor to Kimi K3 on Converse (#43875) * fix(bedrock): apply the 16 maxTokens floor to Kimi K3 on Converse Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(e2e): use lane AWS credentials for the Kimi K3 messages test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(e2e): bypass the proxy cache and declare test metadata for the Kimi K3 probe --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../bedrock/chat/converse_transformation.py | 2 +- .../coverage_registry/llm_conversational.yaml | 1 + .../test_bedrock_kimi_k3_messages_e2e.py | 71 +++++++++++++++++++ .../chat/test_converse_transformation.py | 10 +++ 4 files changed, 83 insertions(+), 1 deletion(-) create mode 100644 tests/e2e/llm_translation/test_bedrock_kimi_k3_messages_e2e.py diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 9deed66d647..b505ae47c61 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -385,7 +385,7 @@ class AmazonConverseConfig(BaseConfig): @staticmethod def _requires_min_max_tokens(model: str) -> bool: - return re.search(r"openai\.gpt-\d|xai\.grok-", model) is not None + return re.search(r"openai\.gpt-\d|xai\.grok-|moonshotai\.kimi-k3", model) is not None def _is_nova_2_model(self, model: str) -> bool: """ diff --git a/tests/e2e/coverage_registry/llm_conversational.yaml b/tests/e2e/coverage_registry/llm_conversational.yaml index b11ae907cfb..c40aeade1c1 100644 --- a/tests/e2e/coverage_registry/llm_conversational.yaml +++ b/tests/e2e/coverage_registry/llm_conversational.yaml @@ -68,6 +68,7 @@ - {id: llm.messages.anthropic.vision.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: anthropic, capability: vision, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Vision via Messages API"} - {id: llm.messages.anthropic.prompt_cache_5m.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: anthropic, capability: prompt_cache_5m, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Prompt caching via Messages API"} - {id: llm.messages.anthropic.thinking.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: anthropic, capability: thinking, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Extended thinking via Messages API"} +- {id: llm.messages.bedrock_converse.basic.nonstream.works_for_kimi_k3_max_tokens_1, module: llm, tier: P0, subject_endpoint: messages, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works_for_kimi_k3_max_tokens_1], source: "llm_translation/test_bedrock_kimi_k3_messages_e2e.py", rationale: "Bedrock Converse accepts max_tokens=1 on the Anthropic Messages surface after LiteLLM applies the model minimum", fail_before_fix: proven} - {id: llm.messages.bedrock_invoke.mid_conversation_system.nonstream.cache_hit, module: llm, tier: P0, subject_endpoint: messages, route: bedrock_invoke, capability: mid_conversation_system, streaming: nonstream, assertions: [works, cache_hit], source: "llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py", rationale: "Flagged Claude 4.8+/5 must keep mid-conversation system reminders in messages; hoisting mutates the system prefix and collapses the prompt cache (#32578/#32831/#32882)", fail_before_fix: proven} - {id: llm.messages.bedrock_invoke.mid_conversation_system.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: bedrock_invoke, capability: mid_conversation_system, streaming: nonstream, assertions: [works], source: "llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py", rationale: "Claude <= 4.7 rejects role system inside messages; unflagged models must convert reminders to user turns in place (hoisting collapses the prompt cache) or every Claude Code session 400s (#32831)", fail_before_fix: proven} - {id: llm.messages.bedrock_invoke.web_search_server_tool.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: bedrock_invoke, capability: web_search_server_tool, streaming: nonstream, assertions: [works], source: "llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py", rationale: "Bedrock hosts no web_search server tool, so this only works because interception rewrites it before the upstream call and the agentic loop feeds the results back in native shape; a regression that short-circuits or forwards it instead yields raw text or AWS's 400", fail_before_fix: unproven} diff --git a/tests/e2e/llm_translation/test_bedrock_kimi_k3_messages_e2e.py b/tests/e2e/llm_translation/test_bedrock_kimi_k3_messages_e2e.py new file mode 100644 index 00000000000..b381fd0d9e1 --- /dev/null +++ b/tests/e2e/llm_translation/test_bedrock_kimi_k3_messages_e2e.py @@ -0,0 +1,71 @@ +from __future__ import annotations + +from typing import Final + +import pytest +from anthropic.types import Message, MessageParam +from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta +from lifecycle import ResourceManager +from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import NO_PROXY_CACHE, SdkClients + +pytestmark = pytest.mark.e2e + +BEDROCK_MODEL: Final = "bedrock/us.moonshotai.kimi-k3" +MODEL_NAME: Final = "claude-kimi-k3" +MODEL_SWITCH_BETA: Final = ( + "claude-code-20250219,context-1m-2025-08-07,interleaved-thinking-2025-05-14," + "redact-thinking-2026-02-12,context-management-2025-06-27," + "prompt-caching-scope-2026-01-05,mid-conversation-system-2026-04-07" +) + + +def _register_model(proxy: ProxyClient, resources: ResourceManager) -> tuple[str, str]: + model_id: Final = proxy.create_model( + MODEL_NAME, + LiteLLMParamsBody( + model=BEDROCK_MODEL, + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name="os.environ/AWS_REGION", + drop_params=True, + ), + ) + resources.defer(lambda: proxy.delete_model(model_id)) + return MODEL_NAME, resources.key(models=[MODEL_NAME]) + + +class TestBedrockKimiK3Messages: + @pytest.mark.covers("llm.messages.bedrock_converse.basic.nonstream.works_for_kimi_k3_max_tokens_1") + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.MESSAGES, + providers=(Provider.BEDROCK,), + models=(BEDROCK_MODEL,), + mode=Mode.NONSTREAM, + ) + ) + def test_model_switch_probe_max_tokens_1_returns_message( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients + ) -> None: + model, key = _register_model(proxy, resources) + client: Final = sdk.anthropic(key) + message_param: Final[MessageParam] = { + "role": "user", + "content": [{"type": "text", "text": "Hi", "cache_control": {"type": "ephemeral"}}], + } + message: Final = client.messages.create( + model=model, + max_tokens=1, + stream=False, + system="You are Claude Code, an assistant for software development.", + messages=[message_param], + extra_headers={"anthropic-beta": MODEL_SWITCH_BETA}, + extra_query={"beta": "true"}, + extra_body=NO_PROXY_CACHE, + ) + + assert isinstance(message, Message) + assert message.role == "assistant" diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index f0dd40d4bc9..1fb5542c964 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -896,6 +896,16 @@ def test_reasoning_with_forced_tool_choice_switches_to_auto(): ("us.xai.grok-4.6", "max_tokens", 1, 16), ("global.xai.grok-4.6", "max_completion_tokens", 1, 16), ("us.xai.grok-4.6", "max_tokens", 32, 32), + ("us.moonshotai.kimi-k3", "max_tokens", 1, 16), + ("global.moonshotai.kimi-k3", "max_completion_tokens", 1, 16), + ("moonshotai.kimi-k3", "max_tokens", 64, 64), + ( + "arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.moonshotai.kimi-k3", + "max_tokens", + 1, + 16, + ), + ("moonshotai.kimi-k2.5", "max_tokens", 1, 1), ("anthropic.claude-sonnet-4-5-20250929-v1:0", "max_tokens", 1, 1), ("arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.openai.gpt-6-astra", "max_tokens", 1, 16), ("arn:aws:bedrock:us-east-1:123456789012:inference-profile/global.xai.grok-4.6", "max_tokens", 1, 16), From 4818b767d8b4c4e5fa9885ab953d6f5cfe7c7ca5 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 20 Oct 2026 06:51:05 -0700 Subject: [PATCH 07/33] fix(cost-map): registry audit 2026-10-20, add volcengine doubao-seed-2-1 260915 rows and undated azure data zone rows (#47561) * fix(cost-map): add volcengine doubao-seed-2-1-pro-260915 and doubao-seed-2-1-lite-260915 Pro row absorbed from #47372, lite row added from the same Volcengine Ark pricing and model list docs Co-authored-by: donghongfei <1247543758@qq.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(cost-map): cover volcengine doubao-seed-2-1-pro-260915 row Co-authored-by: donghongfei <1247543758@qq.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(cost-map): add undated azure/eu and azure/us rows for o1, o3, o3-mini, o4-mini, gpt-4o, gpt-4o-mini Absorbed from #47499 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: donghongfei <1247543758@qq.com> --- ...odel_prices_and_context_window_backup.json | 269 ++++++++++++++++++ model_prices_and_context_window.json | 269 ++++++++++++++++++ tests/unit/test_cost_calculator.py | 34 +++ .../test_model_prices_and_context_window.py | 23 ++ 4 files changed, 595 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 8b3771033a5..903210c57e5 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -4916,6 +4916,26 @@ "supports_tool_choice": true, "supports_vision": true }, + "azure/eu/gpt-4o": { + "deprecation_date": "2027-04-14", + "cache_creation_input_token_cost": 1.38e-06, + "cache_read_input_token_cost": 1.375e-06, + "input_cost_per_token": 2.75e-06, + "input_cost_per_token_batches": 1.375e-06, + "litellm_provider": "azure", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.1e-05, + "output_cost_per_token_batches": 5.5e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "azure/eu/gpt-4o-2024-11-20": { "deprecation_date": "2027-04-14", "cache_creation_input_token_cost": 1.38e-06, @@ -4936,6 +4956,25 @@ "supports_tool_choice": true, "supports_vision": true }, + "azure/eu/gpt-4o-mini": { + "cache_read_input_token_cost": 8.3e-08, + "deprecation_date": "2027-04-14", + "input_cost_per_token": 1.65e-07, + "input_cost_per_token_batches": 8.3e-08, + "litellm_provider": "azure", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 6.6e-07, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "azure/eu/gpt-4o-mini-2024-07-18": { "cache_read_input_token_cost": 8.3e-08, "deprecation_date": "2027-04-14", @@ -5247,6 +5286,25 @@ "supports_vision": true, "cache_read_input_token_cost_batches": 2.75e-09 }, + "azure/eu/o1": { + "cache_read_input_token_cost": 8.25e-06, + "deprecation_date": "2026-11-19", + "input_cost_per_token": 1.65e-05, + "input_cost_per_token_batches": 8.25e-06, + "litellm_provider": "azure", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "chat", + "output_cost_per_token": 6.6e-05, + "output_cost_per_token_batches": 3.3e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_tool_choice": true, + "supports_vision": true + }, "azure/eu/o1-2024-12-17": { "cache_read_input_token_cost": 8.25e-06, "deprecation_date": "2026-11-19", @@ -5283,6 +5341,24 @@ "supports_prompt_caching": true, "supports_vision": false }, + "azure/eu/o3-mini": { + "cache_read_input_token_cost": 6.05e-07, + "deprecation_date": "2026-11-19", + "input_cost_per_token": 1.21e-06, + "input_cost_per_token_batches": 6.05e-07, + "litellm_provider": "azure", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "chat", + "output_cost_per_token": 4.84e-06, + "output_cost_per_token_batches": 2.42e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": false + }, "azure/eu/o3-mini-2025-01-31": { "cache_read_input_token_cost": 6.05e-07, "deprecation_date": "2026-11-19", @@ -11540,6 +11616,26 @@ "supports_tool_choice": true, "supports_vision": true }, + "azure/us/gpt-4o": { + "deprecation_date": "2027-04-14", + "cache_creation_input_token_cost": 1.38e-06, + "cache_read_input_token_cost": 1.375e-06, + "input_cost_per_token": 2.75e-06, + "input_cost_per_token_batches": 1.375e-06, + "litellm_provider": "azure", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.1e-05, + "output_cost_per_token_batches": 5.5e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "azure/us/gpt-4o-2024-11-20": { "deprecation_date": "2027-04-14", "cache_creation_input_token_cost": 1.38e-06, @@ -11560,6 +11656,25 @@ "supports_tool_choice": true, "supports_vision": true }, + "azure/us/gpt-4o-mini": { + "cache_read_input_token_cost": 8.3e-08, + "deprecation_date": "2027-04-14", + "input_cost_per_token": 1.65e-07, + "input_cost_per_token_batches": 8.3e-08, + "litellm_provider": "azure", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 6.6e-07, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "azure/us/gpt-4o-mini-2024-07-18": { "cache_read_input_token_cost": 8.3e-08, "deprecation_date": "2027-04-14", @@ -11871,6 +11986,25 @@ "supports_tool_choice": true, "supports_vision": true }, + "azure/us/o1": { + "cache_read_input_token_cost": 8.25e-06, + "deprecation_date": "2026-11-19", + "input_cost_per_token": 1.65e-05, + "input_cost_per_token_batches": 8.25e-06, + "litellm_provider": "azure", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "chat", + "output_cost_per_token": 6.6e-05, + "output_cost_per_token_batches": 3.3e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_tool_choice": true, + "supports_vision": true + }, "azure/us/o1-2024-12-17": { "cache_read_input_token_cost": 8.25e-06, "deprecation_date": "2026-11-19", @@ -11907,6 +12041,39 @@ "supports_prompt_caching": true, "supports_vision": false }, + "azure/us/o3": { + "deprecation_date": "2026-11-19", + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 2.2e-06, + "input_cost_per_token_batches": 1.1e-06, + "litellm_provider": "azure", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "chat", + "output_cost_per_token": 8.8e-06, + "output_cost_per_token_batches": 4.4e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "azure/us/o3-2025-04-16": { "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5.5e-07, @@ -11940,6 +12107,24 @@ "supports_tool_choice": true, "supports_vision": true }, + "azure/us/o3-mini": { + "cache_read_input_token_cost": 6.05e-07, + "deprecation_date": "2026-11-19", + "input_cost_per_token": 1.21e-06, + "input_cost_per_token_batches": 6.05e-07, + "litellm_provider": "azure", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "chat", + "output_cost_per_token": 4.84e-06, + "output_cost_per_token_batches": 2.42e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": false + }, "azure/us/o3-mini-2025-01-31": { "cache_read_input_token_cost": 6.05e-07, "deprecation_date": "2026-11-19", @@ -11958,6 +12143,27 @@ "supports_tool_choice": true, "supports_vision": false }, + "azure/us/o4-mini": { + "cache_read_input_token_cost": 3.03e-07, + "deprecation_date": "2026-11-19", + "input_cost_per_token": 1.21e-06, + "input_cost_per_token_batches": 6.05e-07, + "litellm_provider": "azure", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "chat", + "output_cost_per_token": 4.84e-06, + "output_cost_per_token_batches": 2.42e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_parallel_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "azure/us/o4-mini-2025-04-16": { "cache_read_input_token_cost": 3.03e-07, "deprecation_date": "2026-11-19", @@ -61148,6 +61354,41 @@ "supports_tool_choice": false, "supports_vision": true }, + "volcengine/doubao-seed-2-1-pro-260915": { + "cache_read_input_token_cost": 1.725e-07, + "input_cost_per_token": 8.625e-07, + "litellm_provider": "volcengine", + "max_input_tokens": 1024000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 4.3125e-06, + "source": "https://www.volcengine.com/docs/82379/1544106", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": false, + "supports_vision": true + }, + "volcengine/doubao-seed-2-1-lite-260915": { + "cache_read_input_audio_token_cost": 3.45e-07, + "cache_read_input_token_cost": 2.3e-08, + "input_cost_per_audio_token": 1.725e-06, + "input_cost_per_token": 1.15e-07, + "litellm_provider": "volcengine", + "max_input_tokens": 1024000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 3.88126e-07, + "source": "https://www.volcengine.com/docs/82379/1544106", + "supports_audio_input": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": false, + "supports_vision": true + }, "volcengine/doubao-seed-2-1-turbo-260628": { "cache_read_input_token_cost": 8.625e-08, "input_cost_per_token": 4.3125e-07, @@ -70141,6 +70382,20 @@ "output_cost_per_token_batches": 2.42e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, + "azure/eu/o3": { + "deprecation_date": "2026-11-19", + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 2.2e-06, + "input_cost_per_token_batches": 1.1e-06, + "litellm_provider": "azure", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "chat", + "output_cost_per_token": 8.8e-06, + "output_cost_per_token_batches": 4.4e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + }, "azure/eu/o3-2025-04-16": { "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5.5e-07, @@ -70164,6 +70419,20 @@ "output_cost_per_token": 4.4e-05, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, + "azure/eu/o4-mini": { + "deprecation_date": "2026-11-19", + "cache_read_input_token_cost": 3.03e-07, + "input_cost_per_token": 1.21e-06, + "input_cost_per_token_batches": 6.05e-07, + "litellm_provider": "azure", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "chat", + "output_cost_per_token": 4.84e-06, + "output_cost_per_token_batches": 2.42e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + }, "azure/eu/o4-mini-2025-04-16": { "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 3.03e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8b3771033a5..903210c57e5 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -4916,6 +4916,26 @@ "supports_tool_choice": true, "supports_vision": true }, + "azure/eu/gpt-4o": { + "deprecation_date": "2027-04-14", + "cache_creation_input_token_cost": 1.38e-06, + "cache_read_input_token_cost": 1.375e-06, + "input_cost_per_token": 2.75e-06, + "input_cost_per_token_batches": 1.375e-06, + "litellm_provider": "azure", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.1e-05, + "output_cost_per_token_batches": 5.5e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "azure/eu/gpt-4o-2024-11-20": { "deprecation_date": "2027-04-14", "cache_creation_input_token_cost": 1.38e-06, @@ -4936,6 +4956,25 @@ "supports_tool_choice": true, "supports_vision": true }, + "azure/eu/gpt-4o-mini": { + "cache_read_input_token_cost": 8.3e-08, + "deprecation_date": "2027-04-14", + "input_cost_per_token": 1.65e-07, + "input_cost_per_token_batches": 8.3e-08, + "litellm_provider": "azure", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 6.6e-07, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "azure/eu/gpt-4o-mini-2024-07-18": { "cache_read_input_token_cost": 8.3e-08, "deprecation_date": "2027-04-14", @@ -5247,6 +5286,25 @@ "supports_vision": true, "cache_read_input_token_cost_batches": 2.75e-09 }, + "azure/eu/o1": { + "cache_read_input_token_cost": 8.25e-06, + "deprecation_date": "2026-11-19", + "input_cost_per_token": 1.65e-05, + "input_cost_per_token_batches": 8.25e-06, + "litellm_provider": "azure", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "chat", + "output_cost_per_token": 6.6e-05, + "output_cost_per_token_batches": 3.3e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_tool_choice": true, + "supports_vision": true + }, "azure/eu/o1-2024-12-17": { "cache_read_input_token_cost": 8.25e-06, "deprecation_date": "2026-11-19", @@ -5283,6 +5341,24 @@ "supports_prompt_caching": true, "supports_vision": false }, + "azure/eu/o3-mini": { + "cache_read_input_token_cost": 6.05e-07, + "deprecation_date": "2026-11-19", + "input_cost_per_token": 1.21e-06, + "input_cost_per_token_batches": 6.05e-07, + "litellm_provider": "azure", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "chat", + "output_cost_per_token": 4.84e-06, + "output_cost_per_token_batches": 2.42e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": false + }, "azure/eu/o3-mini-2025-01-31": { "cache_read_input_token_cost": 6.05e-07, "deprecation_date": "2026-11-19", @@ -11540,6 +11616,26 @@ "supports_tool_choice": true, "supports_vision": true }, + "azure/us/gpt-4o": { + "deprecation_date": "2027-04-14", + "cache_creation_input_token_cost": 1.38e-06, + "cache_read_input_token_cost": 1.375e-06, + "input_cost_per_token": 2.75e-06, + "input_cost_per_token_batches": 1.375e-06, + "litellm_provider": "azure", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.1e-05, + "output_cost_per_token_batches": 5.5e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "azure/us/gpt-4o-2024-11-20": { "deprecation_date": "2027-04-14", "cache_creation_input_token_cost": 1.38e-06, @@ -11560,6 +11656,25 @@ "supports_tool_choice": true, "supports_vision": true }, + "azure/us/gpt-4o-mini": { + "cache_read_input_token_cost": 8.3e-08, + "deprecation_date": "2027-04-14", + "input_cost_per_token": 1.65e-07, + "input_cost_per_token_batches": 8.3e-08, + "litellm_provider": "azure", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 6.6e-07, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "azure/us/gpt-4o-mini-2024-07-18": { "cache_read_input_token_cost": 8.3e-08, "deprecation_date": "2027-04-14", @@ -11871,6 +11986,25 @@ "supports_tool_choice": true, "supports_vision": true }, + "azure/us/o1": { + "cache_read_input_token_cost": 8.25e-06, + "deprecation_date": "2026-11-19", + "input_cost_per_token": 1.65e-05, + "input_cost_per_token_batches": 8.25e-06, + "litellm_provider": "azure", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "chat", + "output_cost_per_token": 6.6e-05, + "output_cost_per_token_batches": 3.3e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_tool_choice": true, + "supports_vision": true + }, "azure/us/o1-2024-12-17": { "cache_read_input_token_cost": 8.25e-06, "deprecation_date": "2026-11-19", @@ -11907,6 +12041,39 @@ "supports_prompt_caching": true, "supports_vision": false }, + "azure/us/o3": { + "deprecation_date": "2026-11-19", + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 2.2e-06, + "input_cost_per_token_batches": 1.1e-06, + "litellm_provider": "azure", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "chat", + "output_cost_per_token": 8.8e-06, + "output_cost_per_token_batches": 4.4e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "azure/us/o3-2025-04-16": { "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5.5e-07, @@ -11940,6 +12107,24 @@ "supports_tool_choice": true, "supports_vision": true }, + "azure/us/o3-mini": { + "cache_read_input_token_cost": 6.05e-07, + "deprecation_date": "2026-11-19", + "input_cost_per_token": 1.21e-06, + "input_cost_per_token_batches": 6.05e-07, + "litellm_provider": "azure", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "chat", + "output_cost_per_token": 4.84e-06, + "output_cost_per_token_batches": 2.42e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": false + }, "azure/us/o3-mini-2025-01-31": { "cache_read_input_token_cost": 6.05e-07, "deprecation_date": "2026-11-19", @@ -11958,6 +12143,27 @@ "supports_tool_choice": true, "supports_vision": false }, + "azure/us/o4-mini": { + "cache_read_input_token_cost": 3.03e-07, + "deprecation_date": "2026-11-19", + "input_cost_per_token": 1.21e-06, + "input_cost_per_token_batches": 6.05e-07, + "litellm_provider": "azure", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "chat", + "output_cost_per_token": 4.84e-06, + "output_cost_per_token_batches": 2.42e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_parallel_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "azure/us/o4-mini-2025-04-16": { "cache_read_input_token_cost": 3.03e-07, "deprecation_date": "2026-11-19", @@ -61148,6 +61354,41 @@ "supports_tool_choice": false, "supports_vision": true }, + "volcengine/doubao-seed-2-1-pro-260915": { + "cache_read_input_token_cost": 1.725e-07, + "input_cost_per_token": 8.625e-07, + "litellm_provider": "volcengine", + "max_input_tokens": 1024000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 4.3125e-06, + "source": "https://www.volcengine.com/docs/82379/1544106", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": false, + "supports_vision": true + }, + "volcengine/doubao-seed-2-1-lite-260915": { + "cache_read_input_audio_token_cost": 3.45e-07, + "cache_read_input_token_cost": 2.3e-08, + "input_cost_per_audio_token": 1.725e-06, + "input_cost_per_token": 1.15e-07, + "litellm_provider": "volcengine", + "max_input_tokens": 1024000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 3.88126e-07, + "source": "https://www.volcengine.com/docs/82379/1544106", + "supports_audio_input": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": false, + "supports_vision": true + }, "volcengine/doubao-seed-2-1-turbo-260628": { "cache_read_input_token_cost": 8.625e-08, "input_cost_per_token": 4.3125e-07, @@ -70141,6 +70382,20 @@ "output_cost_per_token_batches": 2.42e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, + "azure/eu/o3": { + "deprecation_date": "2026-11-19", + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 2.2e-06, + "input_cost_per_token_batches": 1.1e-06, + "litellm_provider": "azure", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "chat", + "output_cost_per_token": 8.8e-06, + "output_cost_per_token_batches": 4.4e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + }, "azure/eu/o3-2025-04-16": { "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5.5e-07, @@ -70164,6 +70419,20 @@ "output_cost_per_token": 4.4e-05, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, + "azure/eu/o4-mini": { + "deprecation_date": "2026-11-19", + "cache_read_input_token_cost": 3.03e-07, + "input_cost_per_token": 1.21e-06, + "input_cost_per_token_batches": 6.05e-07, + "litellm_provider": "azure", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "chat", + "output_cost_per_token": 4.84e-06, + "output_cost_per_token_batches": 2.42e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + }, "azure/eu/o4-mini-2025-04-16": { "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 3.03e-07, diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 99119799ac0..fee4db586cd 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -357,6 +357,40 @@ def test_completion_cost_prices_rule_only_response_model_at_the_azure_deployment assert cost == pytest.approx(usage.prompt_tokens * 0.0000111 + usage.completion_tokens * 0.0000222) +@pytest.mark.parametrize("region", ["eu", "us"]) +@pytest.mark.parametrize( + ("alias", "dated"), + [ + ("o1", "o1-2024-12-17"), + ("o3", "o3-2025-04-16"), + ("o3-mini", "o3-mini-2025-01-31"), + ("o4-mini", "o4-mini-2025-04-16"), + ("gpt-4o", "gpt-4o-2024-11-20"), + ("gpt-4o-mini", "gpt-4o-mini-2024-07-18"), + ], +) +def test_completion_cost_bills_undated_azure_data_zone_base_model_at_its_dated_rate( + _local_model_cost_map: None, region: str, alias: str, dated: str +) -> None: + """LIT-9812: a deployment with base_model azure/eu/o3 (undated, like the global azure/o3) must bill the EU + Data Zone rate of its latest dated row, not $0 and not the global rate""" + dated_info: Final = litellm.model_cost[f"azure/{region}/{dated}"] + usage: Final = Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26) + response: Final = ModelResponse(model=dated, choices=[], usage=usage) + + cost: Final = completion_cost( + completion_response=response, + model="azure/my-data-zone-deployment", + custom_llm_provider="azure", + base_model=f"azure/{region}/{alias}", + ) + + assert cost == pytest.approx( + usage.prompt_tokens * dated_info["input_cost_per_token"] + + usage.completion_tokens * dated_info["output_cost_per_token"] + ) + + def test_completion_cost_orders_candidates_without_extra_huggingface_config_requests( _local_model_cost_map: None, ) -> None: diff --git a/tests/unit/test_model_prices_and_context_window.py b/tests/unit/test_model_prices_and_context_window.py index 395fe8c9750..20b46e27213 100644 --- a/tests/unit/test_model_prices_and_context_window.py +++ b/tests/unit/test_model_prices_and_context_window.py @@ -719,3 +719,26 @@ def test_unregistered_provider_guard_flags_only_labels_nobody_registered(): "unknown_root-new_family_models", "vertex_ai-new_family_models", ] + + +@pytest.mark.parametrize("path", (PRICES_PATH, BACKUP_PRICES_PATH), ids=("main", "backup")) +def test_volcengine_doubao_seed_2_1_pro_260915_row_matches_published_metadata(path: Path): + """Volcengine Ark lists doubao-seed-2-1-pro-260915 at 6 / 30 CNY per million input / output tokens with + cache hits at 1.2 (ark.volcengine.com model detail and docs.volcengine.com/docs/82379/1544106, read + 2026-10-20), a 1024k context window and a 256k output cap. USD rows use the 6.9565 CNY/USD rate of the + sibling 260628 row.""" + rows: Mapping[str, object] = json.loads(path.read_text()) + entry = rows.get("volcengine/doubao-seed-2-1-pro-260915") + assert isinstance(entry, dict) + assert entry["litellm_provider"] == "volcengine" + assert entry["mode"] == "chat" + assert entry["max_input_tokens"] == 1024000 + assert entry["max_output_tokens"] == 256000 + assert entry["max_tokens"] == 256000 + # Volcengine Ark pricing https://www.volcengine.com/docs/82379/1544106 checked 2026-10-20 + assert entry["input_cost_per_token"] == pytest.approx(6.0 / 6.9565 / 1_000_000, rel=1e-4) + assert entry["output_cost_per_token"] == pytest.approx(30.0 / 6.9565 / 1_000_000, rel=1e-4) + assert entry["cache_read_input_token_cost"] == pytest.approx(1.2 / 6.9565 / 1_000_000, rel=1e-4) + sibling = rows["volcengine/doubao-seed-2-1-pro-260628"] + assert isinstance(sibling, dict) + assert all(entry[key] == sibling[key] for key in ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost")) From ea2a842d7d675bfd6387f438b02d09760eabb24a Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 20 Oct 2026 06:51:53 -0700 Subject: [PATCH 08/33] test(perf): set the seven new perf ceilings from 124 lane runs across v1.103.0 to v1.106.0 (#47527) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/perf/cases/mixed_size_requests_2rps.py | 14 +++++++------- tests/e2e/perf/cases/small_bodies_5rps.py | 14 +++++++------- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/tests/e2e/perf/cases/mixed_size_requests_2rps.py b/tests/e2e/perf/cases/mixed_size_requests_2rps.py index 850658a15c5..d4f1fdddc9c 100644 --- a/tests/e2e/perf/cases/mixed_size_requests_2rps.py +++ b/tests/e2e/perf/cases/mixed_size_requests_2rps.py @@ -64,14 +64,14 @@ CASE: Final = PerformanceTestCase( Limit("unexpected_status_count", max=3), Limit("gateway_added_seconds", percentile=50, max=0.58), Limit("gateway_added_seconds", percentile=99, max=6.1), - Limit("litellm_overhead_seconds", percentile=50, max=0.05), - Limit("litellm_overhead_seconds", percentile=99, max=2.2), - Limit("litellm_response_seconds", percentile=50, max=1.2), + Limit("litellm_overhead_seconds", percentile=50, max=0.1), + Limit("litellm_overhead_seconds", percentile=99, max=3.4), + Limit("litellm_response_seconds", percentile=50, max=1.0), Limit("litellm_response_seconds", percentile=99, max=3.8), - Limit("prometheus_redis_latency_mean_seconds", max=0.003), - Limit("postgres_active_sessions", max=0.004), - Limit("postgres_slowest_statement_seconds", max=0.045), - Limit("postgres_peak_connections", max=30), + Limit("prometheus_redis_latency_mean_seconds", max=0.005), + Limit("postgres_active_sessions", max=0.016), + Limit("postgres_slowest_statement_seconds", max=0.25), + Limit("postgres_peak_connections", max=45), Limit("time_to_first_token_seconds", percentile=50, max=2.7), Limit("time_to_first_token_seconds", percentile=99, max=8.8), Limit("memory_growth_megabytes", max=600), diff --git a/tests/e2e/perf/cases/small_bodies_5rps.py b/tests/e2e/perf/cases/small_bodies_5rps.py index bc22a9a96c8..f93c64afba4 100644 --- a/tests/e2e/perf/cases/small_bodies_5rps.py +++ b/tests/e2e/perf/cases/small_bodies_5rps.py @@ -65,14 +65,14 @@ CASE: Final = PerformanceTestCase( limits=( Limit("unexpected_status_count", max=3), Limit("gateway_added_seconds", percentile=50, max=0.29), - Limit("litellm_overhead_seconds", percentile=50, max=0.02), - Limit("litellm_overhead_seconds", percentile=99, max=0.8), - Limit("litellm_response_seconds", percentile=50, max=0.8), - Limit("litellm_response_seconds", percentile=99, max=2.8), + Limit("litellm_overhead_seconds", percentile=50, max=0.012), + Limit("litellm_overhead_seconds", percentile=99, max=0.35), + Limit("litellm_response_seconds", percentile=50, max=1.0), + Limit("litellm_response_seconds", percentile=99, max=2.0), Limit("prometheus_redis_latency_mean_seconds", max=0.006), - Limit("postgres_active_sessions", max=0.006), - Limit("postgres_slowest_statement_seconds", max=0.035), - Limit("postgres_peak_connections", max=30), + Limit("postgres_active_sessions", max=0.01), + Limit("postgres_slowest_statement_seconds", max=0.025), + Limit("postgres_peak_connections", max=45), Limit("time_to_first_token_seconds", percentile=99, max=3.7), Limit("memory_growth_megabytes", max=100), Limit("postgres_calls_per_request", max=1.8), From 5145250d83dbf8e7113f09e626517b8f184fea7f Mon Sep 17 00:00:00 2001 From: Max Azatian <66359507+HardMax71@users.noreply.github.com> Date: Tue, 20 Oct 2026 15:52:26 +0200 Subject: [PATCH 09/33] fix(langfuse_otel)!: send the completion start time (#45597) * fix(langfuse_otel): send the completion start time langfuse_otel never set langfuse.observation.completion_start_time, so Langfuse showed no time to first token for any call, streamed or not. The legacy langfuse callback already sends kwargs["completion_start_time"]. Send the same value as an ISO timestamp with an explicit offset. * fix(langfuse_otel): type kwargs and test against a real span Type the kwargs of _set_langfuse_specific_attributes as dict[str, Any], the declared type of model_call_details, so the new completion start time read adds no unknown-type errors. The tests now check a real SDK span instead of patching safe_set_attribute, and the parameter is typed. * fix(otel): send the completion start time from the v2 Langfuse mapper With LITELLM_OTEL_V2 the langfuse_otel spans come from LangfuseMapper, which had no completion start time either. Carry the payload's completionStartTime on LLMCallSpanData for streamed calls, the same calls that get time_to_first_chunk_seconds, and map it to langfuse.observation.completion_start_time. * fix(langfuse_otel): take the logging kwargs as a read-only Mapping The new kwargs annotation tripped LIT001. Both helpers only read the logging kwargs, so type them Mapping[str, Any] instead of dict. * fix(langfuse_otel): time /v1/messages streams from the call start and send the V2 completion start time like the default emitter * fix(langfuse_otel): time Gemini, Interactions and A2A streams from their first chunk and send the start time as JSON The native Gemini, Interactions and A2A streams never recorded a first chunk, so with completion_start_time now sent both emitters showed TTFT equal to latency on them. They now record it, and Gemini and Interactions take the call's start time like /v1/messages. Both emitters send the completion start time as a JSON string, the form the Langfuse SDK and the langfuse callback send, through one serializer. * fix(langfuse_otel): time A2A streams from the call start and record a cached replay's first event A2A streams took their start time when the iterator was built, so their latency left out the work before it. They now use the call's start time like every other route. Cached /v1/messages stream replays recorded no first event, so their time to first token fell back to the end time. The replay now records it on the first event it yields. The tests use the real Interactions config and typed logger hooks. * fix(a2a): time streamed A2A calls from before the agent card fetch and make the cached replay test clock-free * test(a2a): type the card fetch test and stop rebinding the seed helper's recorder * test(otel): expect completion_start_time and cache_hit in the B3 span keys --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/a2a_protocol/main.py | 4 +- litellm/a2a_protocol/streaming_iterator.py | 4 +- litellm/google_genai/streaming_iterator.py | 4 +- .../integrations/langfuse/langfuse_otel.py | 15 +++- .../langfuse/langfuse_otel_attributes.py | 5 ++ litellm/integrations/langfuse/langfuse_sdk.py | 3 +- litellm/integrations/otel/logger.py | 1 + litellm/integrations/otel/mappers/langfuse.py | 7 ++ litellm/integrations/otel/model/metadata.py | 2 + litellm/integrations/otel/model/payloads.py | 3 + litellm/interactions/streaming_iterator.py | 4 +- .../pass_through/messages/response_cache.py | 3 + .../messages/streaming_iterator.py | 2 +- litellm/types/integrations/langfuse_otel.py | 1 + ...test_arize_otel_v2_openinference_family.py | 3 + .../test_a2a_streaming_iterator.py | 54 ++++++++++++ tests/unit/a2a_protocol/test_main.py | 52 +++++++++++- .../test_google_genai_streaming_iterator.py | 84 +++++++++++++++++-- .../integrations/otel/test_langfuse_logger.py | 53 ++++++++++++ .../otel/test_otel_v2_vendor_mappers.py | 27 ++++++ tests/unit/integrations/test_langfuse_otel.py | 25 ++++++ .../test_interactions_streaming_iterator.py | 51 +++++++++++ .../messages/test_response_cache.py | 57 ++++++++++++- .../messages/test_streaming_iterator.py | 54 +++++++++++- 24 files changed, 497 insertions(+), 21 deletions(-) diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index e5bd38bb7ed..f44e3f6a81b 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -602,9 +602,9 @@ def _build_streaming_logging_obj( litellm_params: dict[str, object] | None, metadata: dict[str, object] | None, proxy_server_request: dict[str, object] | None, + start_time: datetime.datetime, ) -> Logging: """Build logging object for streaming A2A requests.""" - start_time: Final = datetime.datetime.now() model: Final = f"a2a_agent/{agent_name}" logging_obj: Final = Logging( @@ -718,6 +718,7 @@ async def asend_message_streaming( if request is None: raise ValueError("request is required") + called_at: Final = datetime.datetime.now() _raw_logging_obj: Final = kwargs.get("litellm_logging_obj") logging_obj: Logging | None = _raw_logging_obj if isinstance(_raw_logging_obj, Logging) else None @@ -751,6 +752,7 @@ async def asend_message_streaming( litellm_params=litellm_params, metadata=metadata, proxy_server_request=proxy_server_request, + start_time=called_at, ) verbose_logger.info("A2A send_message_streaming request_id=%s, agent=%s", request.id, agent_name) diff --git a/litellm/a2a_protocol/streaming_iterator.py b/litellm/a2a_protocol/streaming_iterator.py index 0b649f62a66..328681b8dbe 100644 --- a/litellm/a2a_protocol/streaming_iterator.py +++ b/litellm/a2a_protocol/streaming_iterator.py @@ -41,7 +41,7 @@ class A2AStreamingIterator: self.request = request self.logging_obj = logging_obj self.agent_name = agent_name - self.start_time = datetime.now() + self.start_time = logging_obj.start_time # Collect chunks for token counting self.chunks: list[SendStreamingMessageResponse] = [] @@ -54,6 +54,8 @@ class A2AStreamingIterator: async def __anext__(self) -> "SendStreamingMessageResponse": try: chunk: Final = await self.stream.__anext__() + if self.logging_obj.completion_start_time is None: + self.logging_obj.update_completion_start_time(completion_start_time=datetime.now()) # Store chunk self.chunks.append(chunk) diff --git a/litellm/google_genai/streaming_iterator.py b/litellm/google_genai/streaming_iterator.py index 37d08dfb59a..a22c47cb1e7 100644 --- a/litellm/google_genai/streaming_iterator.py +++ b/litellm/google_genai/streaming_iterator.py @@ -71,7 +71,7 @@ class BaseGoogleGenAIGenerateContentStreamingIterator: ): self.litellm_logging_obj = litellm_logging_obj self.request_body = request_body - self.start_time = datetime.now() + self.start_time = litellm_logging_obj.start_time self.collected_chunks: list[bytes] = [] self.model = model self.custom_llm_provider = custom_llm_provider @@ -192,6 +192,8 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateCo async def __anext__(self): try: chunk: Final = await _anext_google_genai_sse_chunk(self.stream_iterator) + if self.litellm_logging_obj.completion_start_time is None: + self.litellm_logging_obj.update_completion_start_time(completion_start_time=datetime.now()) self.collected_chunks.append(chunk) return chunk except StopAsyncIteration: diff --git a/litellm/integrations/langfuse/langfuse_otel.py b/litellm/integrations/langfuse/langfuse_otel.py index ae63ab99929..f195cd97e3c 100644 --- a/litellm/integrations/langfuse/langfuse_otel.py +++ b/litellm/integrations/langfuse/langfuse_otel.py @@ -1,7 +1,7 @@ import base64 import json import os -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Optional @@ -9,6 +9,7 @@ from litellm._logging import verbose_logger from litellm.integrations.arize import _utils from litellm.integrations.langfuse.langfuse_otel_attributes import ( LangfuseLLMObsOTELAttributes, + serialize_langfuse_datetime, ) from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig from litellm.litellm_core_utils.safe_json_loads import safe_json_loads @@ -55,7 +56,7 @@ class LangfuseOtelLogger(OpenTelemetry): LangfuseOtelLogger._set_langfuse_specific_attributes(span=span, kwargs=kwargs, response_obj=response_obj) @staticmethod - def _extract_langfuse_metadata(kwargs: dict) -> dict: + def _extract_langfuse_metadata(kwargs: Mapping[str, Any]) -> dict: """ Extracts Langfuse metadata from the standard LiteLLM kwargs structure. @@ -221,7 +222,7 @@ class LangfuseOtelLogger(OpenTelemetry): ) @staticmethod - def _set_langfuse_specific_attributes(span: Span, kwargs, response_obj): + def _set_langfuse_specific_attributes(span: Span, kwargs: Mapping[str, Any], response_obj): """ Sets Langfuse specific metadata attributes onto the OTEL span. @@ -255,6 +256,14 @@ class LangfuseOtelLogger(OpenTelemetry): safe_dumps(messages), ) + completion_start_time: Final = kwargs.get("completion_start_time") + if isinstance(completion_start_time, datetime): + safe_set_attribute( + span, + LangfuseSpanAttributes.OBSERVATION_COMPLETION_START_TIME.value, + serialize_langfuse_datetime(completion_start_time), + ) + LangfuseOtelLogger._set_observation_output(span=span, response_obj=response_obj) @staticmethod diff --git a/litellm/integrations/langfuse/langfuse_otel_attributes.py b/litellm/integrations/langfuse/langfuse_otel_attributes.py index 1eb3ce1a9c2..2873dba93ed 100644 --- a/litellm/integrations/langfuse/langfuse_otel_attributes.py +++ b/litellm/integrations/langfuse/langfuse_otel_attributes.py @@ -6,6 +6,7 @@ Relevant Issue: https://github.com/BerriAI/litellm/issues/13764 import json from collections.abc import Mapping +from datetime import datetime from typing import TYPE_CHECKING, Any, Final from pydantic import BaseModel @@ -29,6 +30,10 @@ if TYPE_CHECKING: from opentelemetry.trace import Span +def serialize_langfuse_datetime(value: datetime) -> str: + return json.dumps(value.astimezone().isoformat()) + + def get_output_content_by_type( response_obj: None | dict diff --git a/litellm/integrations/langfuse/langfuse_sdk.py b/litellm/integrations/langfuse/langfuse_sdk.py index 864c0697eb3..51b08cb72df 100644 --- a/litellm/integrations/langfuse/langfuse_sdk.py +++ b/litellm/integrations/langfuse/langfuse_sdk.py @@ -39,6 +39,7 @@ from pydantic import ConfigDict import litellm from litellm._logging import verbose_logger from litellm.integrations.langfuse.langfuse import PROMPT_CACHE_TTL_ENV, parse_langfuse_debug, whole_number +from litellm.integrations.langfuse.langfuse_otel_attributes import serialize_langfuse_datetime from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.custom_httpx.http_handler import HTTPHandler, get_httpx_client from litellm.types.llms.base import LiteLLMBaseModel @@ -168,7 +169,7 @@ def _string_or_none(value: object) -> str | None: def _serialize_datetime(value: object) -> str | None: """A datetime the way the SDK's ``EventSerializer`` sends one: a JSON string, naive values read as local time.""" if isinstance(value, datetime): - return safe_dumps(value.astimezone().isoformat()) + return serialize_langfuse_datetime(value) return _serialize(value) diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index 73448221acb..5092b248eac 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -594,6 +594,7 @@ class OpenTelemetryV2(CustomLogger): payload, capture_content=self._capture_span_content(), time_to_first_chunk_seconds=call.time_to_first_chunk_seconds, + completion_start_seconds=call.completion_start_seconds, request_route=request_root_http_route(), request_purpose=call.purpose, trace=call.trace, diff --git a/litellm/integrations/otel/mappers/langfuse.py b/litellm/integrations/otel/mappers/langfuse.py index 68860931b76..e60ea539188 100644 --- a/litellm/integrations/otel/mappers/langfuse.py +++ b/litellm/integrations/otel/mappers/langfuse.py @@ -12,8 +12,10 @@ the JSON-serialized payloads. ``trace_attributes`` maps the caller's trace contr import json from collections.abc import Callable +from datetime import datetime, timezone from typing import Final +from litellm.integrations.langfuse.langfuse_otel_attributes import serialize_langfuse_datetime from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue, SpanData from litellm.integrations.otel.mappers.utils import ( collect, @@ -45,6 +47,11 @@ class LangfuseMapper: "langfuse.observation.id": lambda d: d.identity.call_id or None, "langfuse.trace.metadata.team_id": lambda d: d.identity.team_id or None, "langfuse.trace.metadata.team_alias": lambda d: d.identity.team_alias or None, + "langfuse.observation.completion_start_time": lambda d: ( + serialize_langfuse_datetime(datetime.fromtimestamp(d.completion_start_seconds, tz=timezone.utc)) + if d.completion_start_seconds is not None + else None + ), } # Sub-tables folded into their respective JSON blobs. diff --git a/litellm/integrations/otel/model/metadata.py b/litellm/integrations/otel/model/metadata.py index a3151c24519..5d37894bb8f 100644 --- a/litellm/integrations/otel/model/metadata.py +++ b/litellm/integrations/otel/model/metadata.py @@ -258,6 +258,7 @@ class LLMCallEvent: # needs to be reasonable for a span that never gets closed (a leak). provisional_span_name: str time_to_first_chunk_seconds: float | None + completion_start_seconds: float | None trace: TraceControls session_id: str | None @@ -280,6 +281,7 @@ class LLMCallEvent: cache_hit=bool(payload.get("cache_hit")) if payload else False, # pyright: ignore[reportUnknownMemberType] # TypedDict .get on a str key is partially unknown provisional_span_name=f"{operation.value} {model}".strip(), time_to_first_chunk_seconds=time_to_first_chunk_seconds(kwargs), + completion_start_seconds=_epoch_seconds(kwargs.get("completion_start_time")), trace=trace, session_id=caller_session_id(kwargs, trace), ) diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index e15072c009b..d11505ce026 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -429,6 +429,7 @@ class LLMCallSpanData: choices_out: tuple[Mapping[str, object], ...] = () system_fingerprint: str | None = None time_to_first_chunk_seconds: float | None = None + completion_start_seconds: float | None = None # The requested output modality, set only on the routes that pin one (image # generation, speech, transcription, OCR), and the litellm route itself, which # keeps routes the convention folds into one operation distinguishable. @@ -450,6 +451,7 @@ class LLMCallSpanData: payload: StandardLoggingPayload, capture_content: bool = False, time_to_first_chunk_seconds: float | None = None, + completion_start_seconds: float | None = None, request_route: str | None = None, request_purpose: str | None = None, trace: TraceControls | None = None, @@ -503,6 +505,7 @@ class LLMCallSpanData: choices_out=choices_out if capture_content else (), system_fingerprint=as_str(response.get("system_fingerprint")), time_to_first_chunk_seconds=time_to_first_chunk_seconds, + completion_start_seconds=completion_start_seconds, output_type=resolve_output_type(call_type), call_type=call_type or None, request_route=request_route or context.identity.request_route, diff --git a/litellm/interactions/streaming_iterator.py b/litellm/interactions/streaming_iterator.py index 64b599edfe2..a2468cd894f 100644 --- a/litellm/interactions/streaming_iterator.py +++ b/litellm/interactions/streaming_iterator.py @@ -48,7 +48,7 @@ class BaseInteractionsAPIStreamingIterator: self.finished = False self.interactions_api_config = interactions_api_config self.completed_response: InteractionsAPIStreamingResponse | None = None - self.start_time = datetime.now() + self.start_time = logging_obj.start_time # set request kwargs self.litellm_metadata = litellm_metadata @@ -87,6 +87,8 @@ class BaseInteractionsAPIStreamingIterator: # Format as InteractionsAPIStreamingResponse if isinstance(parsed_chunk, dict): + if self.logging_obj.completion_start_time is None: + self.logging_obj.update_completion_start_time(completion_start_time=datetime.now()) streaming_response: Final = self.interactions_api_config.transform_streaming_response( model=self.model, parsed_chunk=parsed_chunk, diff --git a/litellm/llms/anthropic/pass_through/messages/response_cache.py b/litellm/llms/anthropic/pass_through/messages/response_cache.py index 4c5f47bc650..81d01f1437e 100644 --- a/litellm/llms/anthropic/pass_through/messages/response_cache.py +++ b/litellm/llms/anthropic/pass_through/messages/response_cache.py @@ -1,5 +1,6 @@ import re from collections.abc import AsyncIterator, Mapping, Sequence +from datetime import datetime from types import MappingProxyType from typing import TYPE_CHECKING, Final, cast @@ -153,6 +154,8 @@ class CachedAnthropicMessagesStreamIterator(BaseAnthropicMessagesStreamingIterat await self._handle_streaming_logging(chunks) raise StopAsyncIteration chunk: Final = self.chunks[self.current_index] + if self.completion_start_time is None: + self.completion_start_time = datetime.now() self.current_index += 1 return chunk diff --git a/litellm/llms/anthropic/pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/pass_through/messages/streaming_iterator.py index 44db193ea83..bfa309fad83 100644 --- a/litellm/llms/anthropic/pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/pass_through/messages/streaming_iterator.py @@ -506,7 +506,7 @@ class BaseAnthropicMessagesStreamingIterator: ): self.litellm_logging_obj = litellm_logging_obj self.request_body = request_body - self.start_time = datetime.now() + self.start_time = litellm_logging_obj.start_time self.completion_start_time: datetime | None = None async def _handle_streaming_logging(self, collected_chunks: list[bytes], *, stream_teardown: bool = False): diff --git a/litellm/types/integrations/langfuse_otel.py b/litellm/types/integrations/langfuse_otel.py index 4bcd0448c3a..5647f28c757 100644 --- a/litellm/types/integrations/langfuse_otel.py +++ b/litellm/types/integrations/langfuse_otel.py @@ -29,6 +29,7 @@ class LangfuseSpanAttributes(str, Enum): # ---- Observation input/output ---- OBSERVATION_INPUT = "langfuse.observation.input" OBSERVATION_OUTPUT = "langfuse.observation.output" + OBSERVATION_COMPLETION_START_TIME = "langfuse.observation.completion_start_time" # ---- Trace-level metadata ---- TRACE_USER_ID = "user.id" diff --git a/tests/integration/observability/test_arize_otel_v2_openinference_family.py b/tests/integration/observability/test_arize_otel_v2_openinference_family.py index 242b7bfe8c1..0b0c077cbee 100644 --- a/tests/integration/observability/test_arize_otel_v2_openinference_family.py +++ b/tests/integration/observability/test_arize_otel_v2_openinference_family.py @@ -43,6 +43,7 @@ _GENAI_B3_KEYS_WITHOUT_BAGGAGE: Final = frozenset( "gen_ai.usage.prompt_tokens", "gen_ai.usage.total_tokens", "litellm.api_key.hash", + "litellm.cache_hit", "litellm.call_id", "litellm.call_type", "litellm.cost.discount_amount", @@ -85,6 +86,7 @@ _LANGFUSE_B3_KEYS: Final = frozenset( "gen_ai.usage.output_tokens", "gen_ai.usage.prompt_tokens", "gen_ai.usage.total_tokens", + "langfuse.observation.completion_start_time", "langfuse.observation.cost_details", "langfuse.observation.id", "langfuse.observation.input", @@ -94,6 +96,7 @@ _LANGFUSE_B3_KEYS: Final = frozenset( "langfuse.observation.type", "langfuse.observation.usage_details", "litellm.api_key.hash", + "litellm.cache_hit", "litellm.call_id", "litellm.call_type", "litellm.cost.discount_amount", diff --git a/tests/unit/a2a_protocol/test_a2a_streaming_iterator.py b/tests/unit/a2a_protocol/test_a2a_streaming_iterator.py index 2e883e91fda..6d685674984 100644 --- a/tests/unit/a2a_protocol/test_a2a_streaming_iterator.py +++ b/tests/unit/a2a_protocol/test_a2a_streaming_iterator.py @@ -6,6 +6,8 @@ async_success_handler (cross-thread pydantic mutation segfaults pydantic-core). import asyncio import time +from datetime import datetime +from typing import Final from types import SimpleNamespace import pytest @@ -135,3 +137,55 @@ async def test_stream_completion_counts_tokens_off_the_event_loop(monkeypatch): assert usage.prompt_tokens > 100_000 assert usage.completion_tokens > 100_000 assert_loop_stayed_free(took, lags) + + +class _SuccessRecorder(CustomLogger): + def __init__(self): + super().__init__() + self.logged: list[tuple[datetime, object]] = [] + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.logged.append((start_time, kwargs.get("completion_start_time"))) + + async def async_log_stream_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.logged.append((start_time, kwargs.get("completion_start_time"))) + + +@pytest.mark.asyncio +async def test_stream_logs_its_first_chunk_time_measured_from_the_call_start(): + call_start: Final = datetime(2026, 10, 9, 12, 0, 0) + recorder: Final = _SuccessRecorder() + logging_obj: Final = LitellmLogging( + model="a2a/test-agent", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="a2a_send_message_streaming", + start_time=call_start, + litellm_call_id="a2a-first-chunk", + function_id="a2a-first-chunk", + dynamic_async_success_callbacks=[recorder], + ) + + async def _stream(): + yield _AgentChunk("Hel") + yield _AgentChunk("Hello") + + iterator: Final = A2AStreamingIterator( + stream=_stream(), + request=SimpleNamespace(params=SimpleNamespace(message={"role": "user", "parts": [{"kind": "text", "text": "hi"}]})), + logging_obj=logging_obj, + agent_name="test-agent", + ) + tasks_before: Final = asyncio.all_tasks() + + _ = await iterator.__anext__() + first_chunk_at: Final = logging_obj.completion_start_time + _ = [chunk async for chunk in iterator] + await asyncio.gather(*(asyncio.all_tasks() - tasks_before)) + + assert first_chunk_at is not None + assert recorder.logged == [(call_start, first_chunk_at)] diff --git a/tests/unit/a2a_protocol/test_main.py b/tests/unit/a2a_protocol/test_main.py index bab14dd26d6..75fd9eda75b 100644 --- a/tests/unit/a2a_protocol/test_main.py +++ b/tests/unit/a2a_protocol/test_main.py @@ -1,7 +1,9 @@ """Tests for litellm/a2a_protocol/main.py non-streaming send behavior.""" import asyncio +import datetime import json +from typing import Final import httpx import pytest @@ -155,6 +157,7 @@ def test_streaming_logging_obj_carries_call_type_into_model_call_details(): litellm_params=None, metadata=None, proxy_server_request=None, + start_time=datetime.datetime(2026, 10, 20, 12, 0, 0), ) assert logging_obj.model_call_details["call_type"] == "asend_message_streaming" @@ -352,6 +355,7 @@ async def _seed_shared_a2a_client( rpc_reply=_RPC_REPLY, rpc_stream_events=None, provider: str = httpxSpecialProvider.A2AProvider, + recorder: _RequestRecorder | None = None, ) -> _RequestRecorder: """Put the one A2A client the cache will hand out behind a mock transport. @@ -359,10 +363,12 @@ async def _seed_shared_a2a_client( it. The injected client is a real httpx.AsyncClient, so the merge of per-request headers over client defaults, which is what these tests are about, stays real. """ - recorder = _RequestRecorder(card=card, rpc_reply=rpc_reply, rpc_stream_events=rpc_stream_events) + active_recorder: Final = recorder or _RequestRecorder( + card=card, rpc_reply=rpc_reply, rpc_stream_events=rpc_stream_events + ) handler = AsyncHTTPHandler(timeout=DEFAULT_A2A_AGENT_TIMEOUT) owned_client = handler.client - handler.client = httpx.AsyncClient(transport=httpx.MockTransport(recorder)) + handler.client = httpx.AsyncClient(transport=httpx.MockTransport(active_recorder)) await owned_client.aclose() litellm.in_memory_llm_clients_cache.set_cache( @@ -374,8 +380,8 @@ async def _seed_shared_a2a_client( ) assert seeded is handler, "cache key drifted from get_async_httpx_client; these tests would test nothing" - recorder.client = handler.client - return recorder + active_recorder.client = handler.client + return active_recorder @pytest.fixture @@ -679,6 +685,43 @@ class _UsageRecorder(CustomLogger): self.logged.set() +class _CardServedAtRecorder(_RequestRecorder): + def __init__(self) -> None: + super().__init__(rpc_stream_events=_LANGGRAPH_STREAM_EVENTS) + self.card_served_at: datetime.datetime | None = None + + def __call__(self, request: httpx.Request) -> httpx.Response: + if request.method == "GET" and self.card_served_at is None: + self.card_served_at = datetime.datetime.now() + return super().__call__(request) + + +@pytest.mark.asyncio +async def test_streamed_send_is_timed_from_before_the_agent_card_fetch( + isolated_client_cache: LLMClientCache, monkeypatch: pytest.MonkeyPatch +): + from litellm.a2a_protocol import main as a2a_main + + card_recorder: Final = _CardServedAtRecorder() + await _seed_shared_a2a_client(recorder=card_recorder) + logger: Final = _UsageRecorder() + monkeypatch.setattr(litellm, "callbacks", [logger]) + monkeypatch.setattr(litellm, "success_callback", [logger]) + monkeypatch.setattr(litellm, "_async_success_callback", [logger]) + request: Final = SendStreamingMessageRequest( + id="reply", + params=MessageSendParams( + message={"messageId": "m-user", "role": "user", "parts": [{"kind": "text", "text": "hi"}]} + ), + ) + + _ = [chunk async for chunk in a2a_main.asend_message_streaming(request=request, api_base="http://127.0.0.1:9")] + await asyncio.wait_for(logger.logged.wait(), timeout=10) + + assert card_recorder.card_served_at is not None + assert logger.payload["startTime"] <= card_recorder.card_served_at.timestamp() + + @pytest.mark.asyncio async def test_asend_message_counts_usage_off_the_event_loop(monkeypatch): from tests.large_text import text @@ -742,6 +785,7 @@ def test_streaming_logging_obj_keeps_agent_credentials_out_of_logging_params(): }, metadata={"user_api_key": "hashed"}, proxy_server_request={"url": "http://localhost:4000"}, + start_time=datetime.datetime(2026, 10, 20, 12, 0, 0), ) expected = { diff --git a/tests/unit/google_genai/test_google_genai_streaming_iterator.py b/tests/unit/google_genai/test_google_genai_streaming_iterator.py index e8ec2848233..120c99b6e42 100644 --- a/tests/unit/google_genai/test_google_genai_streaming_iterator.py +++ b/tests/unit/google_genai/test_google_genai_streaming_iterator.py @@ -1,16 +1,55 @@ +import asyncio import json +from datetime import datetime +from typing import Final from unittest.mock import MagicMock +import httpx import pytest from litellm.google_genai.streaming_iterator import ( AsyncGoogleGenAIGenerateContentStreamingIterator, GoogleGenAIGenerateContentStreamingIterator, ) +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType +class _SuccessRecorder(CustomLogger): + def __init__(self): + super().__init__() + self.logged: list[tuple[datetime, object]] = [] + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.logged.append((start_time, kwargs.get("completion_start_time"))) + + +def _logging_obj( + start_time: datetime = datetime(2026, 10, 9, 12, 0, 0), recorder: _SuccessRecorder | None = None +) -> LiteLLMLoggingObj: + logging_obj: Final = LiteLLMLoggingObj( + model="gemini-3.1-flash", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="agenerate_content_stream", + start_time=start_time, + litellm_call_id="google-genai-stream", + function_id="google-genai-stream", + dynamic_async_success_callbacks=[recorder] if recorder else None, + ) + logging_obj.update_environment_variables( + model="gemini-3.1-flash", + user="", + optional_params={}, + litellm_params={"custom_llm_provider": "gemini"}, + custom_llm_provider="gemini", + ) + return logging_obj + + @pytest.mark.parametrize( "custom_llm_provider, expected_endpoint_type", [("gemini", EndpointType.GEMINI), ("vertex_ai", EndpointType.VERTEX_AI)], @@ -31,7 +70,7 @@ def test_streaming_logging_targets_the_provider_that_served_the_request( iterator = iterator_cls( response=MagicMock(), model="gemini-3.1-flash-image", - logging_obj=MagicMock(spec=LiteLLMLoggingObj), + logging_obj=_logging_obj(), generate_content_provider_config=MagicMock(), litellm_metadata={}, custom_llm_provider=custom_llm_provider, @@ -73,7 +112,7 @@ async def test_async_streaming_iterator_yields_complete_sse_events(): iterator = AsyncGoogleGenAIGenerateContentStreamingIterator( response=mock_response, model="gemini-3.1-flash-image-preview", - logging_obj=MagicMock(spec=LiteLLMLoggingObj), + logging_obj=_logging_obj(), generate_content_provider_config=MagicMock(), litellm_metadata={}, custom_llm_provider="gemini", @@ -95,7 +134,7 @@ def test_sync_streaming_iterator_yields_complete_sse_events(): iterator = GoogleGenAIGenerateContentStreamingIterator( response=mock_response, model="gemini-3.1-flash-image-preview", - logging_obj=MagicMock(spec=LiteLLMLoggingObj), + logging_obj=_logging_obj(), generate_content_provider_config=MagicMock(), litellm_metadata={}, custom_llm_provider="gemini", @@ -123,7 +162,7 @@ async def test_async_streaming_iterator_preserves_multi_field_sse_event(): iterator = AsyncGoogleGenAIGenerateContentStreamingIterator( response=mock_response, model="gemini-test", - logging_obj=MagicMock(spec=LiteLLMLoggingObj), + logging_obj=_logging_obj(), generate_content_provider_config=MagicMock(), litellm_metadata={}, custom_llm_provider="gemini", @@ -146,7 +185,7 @@ async def test_async_streaming_iterator_forwards_sse_comment_events(): iterator = AsyncGoogleGenAIGenerateContentStreamingIterator( response=mock_response, model="gemini-test", - logging_obj=MagicMock(spec=LiteLLMLoggingObj), + logging_obj=_logging_obj(), generate_content_provider_config=MagicMock(), litellm_metadata={}, custom_llm_provider="gemini", @@ -154,3 +193,38 @@ async def test_async_streaming_iterator_forwards_sse_comment_events(): chunk = await iterator.__anext__() assert chunk == b": keepalive\n\n" + + +def _gemini_event(text: str, finish_reason: str | None = None) -> str: + candidate: Final = {"content": {"role": "model", "parts": [{"text": text}]}, "index": 0} + payload: Final = { + "candidates": [{**candidate, "finishReason": finish_reason} if finish_reason else candidate], + "usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 2, "totalTokenCount": 5}, + "modelVersion": "gemini-3.1-flash", + } + return f"data: {json.dumps(payload)}\n\n" + + +@pytest.mark.asyncio +async def test_async_stream_logs_its_first_chunk_time_measured_from_the_call_start(): + call_start: Final = datetime(2026, 10, 9, 12, 0, 0) + recorder: Final = _SuccessRecorder() + logging_obj: Final = _logging_obj(call_start, recorder) + response: Final = httpx.Response(200, content=(_gemini_event("Hel") + _gemini_event("lo", "STOP")).encode()) + iterator: Final = AsyncGoogleGenAIGenerateContentStreamingIterator( + response=response, + model="gemini-3.1-flash", + logging_obj=logging_obj, + generate_content_provider_config=MagicMock(), + litellm_metadata={}, + custom_llm_provider="gemini", + ) + tasks_before: Final = asyncio.all_tasks() + + _ = await iterator.__anext__() + first_chunk_at: Final = logging_obj.completion_start_time + _ = [chunk async for chunk in iterator] + await asyncio.gather(*(asyncio.all_tasks() - tasks_before)) + + assert first_chunk_at is not None + assert recorder.logged == [(call_start, first_chunk_at)] diff --git a/tests/unit/integrations/otel/test_langfuse_logger.py b/tests/unit/integrations/otel/test_langfuse_logger.py index ff2060d91c2..d55f0577e03 100644 --- a/tests/unit/integrations/otel/test_langfuse_logger.py +++ b/tests/unit/integrations/otel/test_langfuse_logger.py @@ -4,6 +4,7 @@ stamped from the request task while the root span is still recording, so Langfus import asyncio import json from collections.abc import AsyncIterator, Mapping, Sequence +from datetime import datetime, timezone from typing import Final import pytest @@ -11,6 +12,7 @@ import pytest pytest.importorskip("opentelemetry") from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter # noqa: E402 +from opentelemetry.trace import SpanKind # noqa: E402 import litellm # noqa: E402 from litellm.caching.dual_cache import DualCache # noqa: E402 @@ -43,6 +45,7 @@ from litellm.types.utils import ( # noqa: E402 INPUT_ATTR: Final = "langfuse.observation.input" OUTPUT_ATTR: Final = "langfuse.observation.output" TRACE_NAME_ATTR: Final = "langfuse.trace.name" +COMPLETION_START_ATTR: Final = "langfuse.observation.completion_start_time" TRACE_CONTROL_ATTRS: Final = (TRACE_NAME_ATTR, "user.id", "session.id", "langfuse.trace.tags") CHAT_DATA: Final = {"model": "gpt-5.4-mini", "messages": [{"role": "user", "content": "ping"}]} @@ -541,3 +544,53 @@ def test_langfuse_otel_preset_builds_a_logger_that_stamps_the_root(monkeypatch): assert INPUT_ATTR in attrs and OUTPUT_ATTR in attrs finally: is_otel_v2_enabled.cache_clear() + + +def _llm_call_kwargs(optional_params: Mapping[str, object], completion_start: datetime | None) -> dict[str, object]: + return { + "standard_logging_object": { + "litellm_call_id": "call_1", + "call_type": "acompletion", + "custom_llm_provider": "openai", + "model": "gpt-5.4-mini", + "stream": bool(optional_params.get("stream")), + "status": "success", + "completionStartTime": datetime(2026, 10, 9, 12, 0, 9, tzinfo=timezone.utc).timestamp(), + }, + "litellm_params": {"metadata": {}}, + "optional_params": optional_params, + "api_call_start_time": datetime(2026, 10, 9, 12, 0, 0, tzinfo=timezone.utc), + "completion_start_time": completion_start, + } + + +def _generation_attrs(kwargs: dict[str, object], *, fail: bool = False) -> Mapping[str, object]: + logger, exporter = _logger() + logger.log_pre_api_call(model="gpt-5.4-mini", messages=[], kwargs=kwargs) + hook: Final = logger.async_log_failure_event if fail else logger.async_log_success_event + asyncio.run(hook(kwargs, None, None, None)) + (generation,) = [span for span in exporter.get_finished_spans() if span.kind is SpanKind.CLIENT] + return dict(generation.attributes or {}) + + +@pytest.mark.parametrize( + "optional_params", + [ + pytest.param({"stream": True}, id="streamed"), + pytest.param({}, id="not_streamed_or_pass_through_stream"), + ], +) +def test_generation_carries_litellms_completion_start_time_like_the_default_emitter( + optional_params: Mapping[str, object], +) -> None: + completion_start: Final = datetime(2026, 10, 9, 12, 0, 0, 250000, tzinfo=timezone.utc) + + attrs: Final = _generation_attrs(_llm_call_kwargs(optional_params, completion_start)) + + assert datetime.fromisoformat(json.loads(attrs[COMPLETION_START_ATTR])) == completion_start + + +def test_failed_generation_without_a_first_chunk_has_no_completion_start_time() -> None: + attrs: Final = _generation_attrs(_llm_call_kwargs({"stream": True}, None), fail=True) + + assert COMPLETION_START_ATTR not in attrs diff --git a/tests/unit/integrations/otel/test_otel_v2_vendor_mappers.py b/tests/unit/integrations/otel/test_otel_v2_vendor_mappers.py index ca7a9b424ad..7e112fcf6a4 100644 --- a/tests/unit/integrations/otel/test_otel_v2_vendor_mappers.py +++ b/tests/unit/integrations/otel/test_otel_v2_vendor_mappers.py @@ -7,6 +7,7 @@ backends, so one trace lights up every configured destination. import json from collections.abc import Mapping +from datetime import datetime, timezone from itertools import chain from typing import Final @@ -580,6 +581,32 @@ def test_langfuse_usage_details_omit_zero_cache_and_reasoning_counts(): assert usage == {"input": 12, "output": 8, "total": 20} +def _langfuse_payload() -> dict[str, object]: + return { + "call_type": "acompletion", + "custom_llm_provider": "openai", + "model": "gpt-4o", + "completionStartTime": 1791547299.0, + } + + +def test_langfuse_mapper_sends_the_completion_start_time_as_the_sdks_json_string(): + data: Final = LLMCallSpanData.from_standard_logging_payload( + _langfuse_payload(), completion_start_seconds=1791547200.25 + ) + + attrs: Final = LangfuseMapper().map(data) + + sent: Final = datetime.fromisoformat(json.loads(attrs["langfuse.observation.completion_start_time"])) + assert sent == datetime(2026, 10, 9, 12, 0, 0, 250000, tzinfo=timezone.utc) + + +def test_langfuse_mapper_ignores_the_payloads_end_time_fallback_when_litellm_recorded_no_completion_start(): + attrs: Final = LangfuseMapper().map(LLMCallSpanData.from_standard_logging_payload(_langfuse_payload())) + + assert "langfuse.observation.completion_start_time" not in attrs + + def test_langfuse_mapper_names_the_trace_from_the_caller(): named = LangfuseMapper().map(_llm_call(trace=TraceControls(name="nightly-eval"))) assert named["langfuse.trace.name"] == "nightly-eval" diff --git a/tests/unit/integrations/test_langfuse_otel.py b/tests/unit/integrations/test_langfuse_otel.py index 13a49f40d7e..9f536ab0dac 100644 --- a/tests/unit/integrations/test_langfuse_otel.py +++ b/tests/unit/integrations/test_langfuse_otel.py @@ -168,6 +168,31 @@ class TestLangfuseOtelIntegration: assert span.attributes["langfuse.environment"] == "team-a-env" + @pytest.mark.parametrize( + "completion_start_time", + [ + datetime(2026, 10, 9, 12, 0, 0, 250000, tzinfo=timezone.utc), + datetime(2026, 10, 9, 12, 0, 0, 250000), + ], + ids=["aware", "naive_local_time"], + ) + def test_set_langfuse_specific_attributes_completion_start_time(self, completion_start_time: datetime): + """The first-token time goes out the way the Langfuse SDK sends it: a JSON string, with an explicit offset.""" + span: Final = TracerProvider().get_tracer(__name__).start_span("litellm_request") + + LangfuseOtelLogger._set_langfuse_specific_attributes(span, {"completion_start_time": completion_start_time}, {}) + + sent: Final = datetime.fromisoformat(json.loads(span.attributes["langfuse.observation.completion_start_time"])) + assert sent.tzinfo is not None + assert sent.timestamp() == completion_start_time.timestamp() + + def test_set_langfuse_specific_attributes_without_completion_start_time(self): + span: Final = TracerProvider().get_tracer(__name__).start_span("litellm_request") + + LangfuseOtelLogger._set_langfuse_specific_attributes(span, {}, {}) + + assert "langfuse.observation.completion_start_time" not in span.attributes + def test_extract_langfuse_metadata_basic(self): """Ensure metadata is correctly pulled from litellm_params.""" metadata_in = {"generation_name": "my-gen", "custom": "data"} diff --git a/tests/unit/interactions/test_interactions_streaming_iterator.py b/tests/unit/interactions/test_interactions_streaming_iterator.py index c52ab36deaa..afcd0d559d7 100644 --- a/tests/unit/interactions/test_interactions_streaming_iterator.py +++ b/tests/unit/interactions/test_interactions_streaming_iterator.py @@ -5,7 +5,10 @@ async_success_handler (cross-thread pydantic mutation segfaults pydantic-core). """ import asyncio +import json import time +from datetime import datetime +from typing import Final import httpx import pytest @@ -16,6 +19,7 @@ from litellm.interactions import streaming_iterator as interactions_streaming_it from litellm.interactions.streaming_iterator import InteractionsAPIStreamingIterator from litellm.litellm_core_utils import thread_pool_executor as thread_pool_executor_module from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging +from litellm.llms.gemini.interactions.transformation import GoogleAIStudioInteractionsConfig from litellm.types.interactions import InteractionsAPIStreamingResponse @@ -97,3 +101,50 @@ async def test_custom_logger_only_never_submits_sync_success_handler(monkeypatch assert recorder.async_hook_fired is True assert recording_executor.submitted_for(logging_obj) == [] + + +class _SuccessRecorder(CustomLogger): + def __init__(self): + super().__init__() + self.logged: list[tuple[datetime, object]] = [] + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.logged.append((start_time, kwargs.get("completion_start_time"))) + + +@pytest.mark.asyncio +async def test_stream_logs_its_first_event_time_measured_from_the_call_start(): + call_start: Final = datetime(2026, 10, 9, 12, 0, 0) + recorder: Final = _SuccessRecorder() + logging_obj: Final = LitellmLogging( + model="gemini/gemini-3-pro-preview", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="ainteraction", + start_time=call_start, + litellm_call_id="interactions-stream", + function_id="interactions-stream", + dynamic_async_success_callbacks=[recorder], + ) + events: Final = ( + {"event_type": "step.delta", "id": "interaction-1"}, + {"event_type": "interaction.completed", "id": "interaction-1", "status": "completed"}, + ) + response: Final = httpx.Response(200, content="".join(f"data: {json.dumps(e)}\n\n" for e in events).encode()) + iterator: Final = InteractionsAPIStreamingIterator( + response=response, + model="gemini/gemini-3-pro-preview", + interactions_api_config=GoogleAIStudioInteractionsConfig(), + logging_obj=logging_obj, + ) + tasks_before: Final = asyncio.all_tasks() + + _ = await iterator.__anext__() + first_event_at: Final = logging_obj.completion_start_time + _ = [event async for event in iterator] + await asyncio.gather(*(asyncio.all_tasks() - tasks_before)) + + assert first_event_at is not None + assert recorder.logged == [(call_start, first_event_at)] diff --git a/tests/unit/llms/anthropic/pass_through/messages/test_response_cache.py b/tests/unit/llms/anthropic/pass_through/messages/test_response_cache.py index b88a57f440d..1df4ccec02c 100644 --- a/tests/unit/llms/anthropic/pass_through/messages/test_response_cache.py +++ b/tests/unit/llms/anthropic/pass_through/messages/test_response_cache.py @@ -1,5 +1,5 @@ import asyncio -from typing import Any, AsyncIterator, Dict, List +from typing import Any, AsyncIterator, Dict, Final, List import pytest @@ -10,10 +10,13 @@ import litellm from litellm._internal_context import in_post_response_phase from litellm.caching.caching import Cache, LiteLLMCacheType from litellm.caching.caching_handler import LLMCachingHandler +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.anthropic.pass_through.messages import handler from litellm.llms.anthropic.pass_through.messages.response_cache import ( AnthropicMessagesStreamCacheWriter, + CachedAnthropicMessagesStreamIterator, ) STREAM_EVENTS: List[bytes] = [ @@ -299,6 +302,58 @@ async def test_cached_stream_replay_logs_once_when_polled_after_exhaustion(): await asyncio.sleep(0) mock_route.assert_called_once() + + +class _ReplayRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged: list[object] = [] + + async def async_log_success_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + self.logged.append(kwargs.get("completion_start_time")) + + +@pytest.mark.asyncio +async def test_cached_stream_replay_logs_when_its_first_event_was_replayed(): + recorder: Final = _ReplayRecorder() + logging_obj: Final = LiteLLMLoggingObj( + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "which greek letter?"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.datetime(2026, 10, 20, 12, 0, 0), + litellm_call_id="cached-replay", + function_id="cached-replay", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model="claude-sonnet-4-5", + user="", + optional_params={}, + litellm_params={"custom_llm_provider": "anthropic"}, + custom_llm_provider="anthropic", + ) + iterator: Final = CachedAnthropicMessagesStreamIterator( + events=[event.decode("utf-8") for event in STREAM_EVENTS], + litellm_logging_obj=logging_obj, + request_body={"model": "claude-sonnet-4-5", "stream": True}, + ) + + _ = await iterator.__anext__() + first_event_at: Final = iterator.completion_start_time + _ = await _collect(iterator) + await GLOBAL_LOGGING_WORKER.flush() + + assert first_event_at is not None + assert recorder.logged == [first_event_at] + + class _HeldBackStream: has_buffered_provider_output = True diff --git a/tests/unit/llms/anthropic/pass_through/messages/test_streaming_iterator.py b/tests/unit/llms/anthropic/pass_through/messages/test_streaming_iterator.py index b995c6dba9b..fd9aa75be74 100644 --- a/tests/unit/llms/anthropic/pass_through/messages/test_streaming_iterator.py +++ b/tests/unit/llms/anthropic/pass_through/messages/test_streaming_iterator.py @@ -5,6 +5,7 @@ from datetime import datetime from typing import Final, Literal from unittest.mock import patch +import httpx import pytest @@ -859,13 +860,15 @@ class _SuccessRecorder(CustomLogger): self.success_kwargs.append(kwargs) -def _make_priced_logging_obj(call_id: str, recorder: _SuccessRecorder, model: str) -> LiteLLMLoggingObj: +def _make_priced_logging_obj( + call_id: str, recorder: _SuccessRecorder, model: str, start_time: datetime | None = None +) -> LiteLLMLoggingObj: logging_obj = LiteLLMLoggingObj( model=model, messages=[{"role": "user", "content": "hi"}], stream=True, call_type="anthropic_messages", - start_time=datetime.now(), + start_time=start_time or datetime.now(), litellm_call_id=call_id, function_id=call_id, dynamic_async_success_callbacks=[recorder], @@ -880,6 +883,53 @@ def _make_priced_logging_obj(call_id: str, recorder: _SuccessRecorder, model: st return logging_obj +@pytest.mark.asyncio +async def test_streamed_messages_log_the_call_start_instead_of_when_the_upstream_answered(): + model: Final = "claude-sonnet-5" + call_start: Final = datetime(2026, 10, 9, 12, 0, 0) + recorder: Final = _SuccessRecorder() + logging_obj: Final = _make_priced_logging_obj("logs_call_start", recorder, model, start_time=call_start) + logging_obj._on_deferred_stream_complete = lambda logging_coroutine: logging_coroutine + iterator: Final = BaseAnthropicMessagesStreamingIterator( + litellm_logging_obj=logging_obj, request_body={"model": model, "stream": True} + ) + + async def _full_stream(): + for event in (*_STREAM_PREFIX, *_STREAM_TAIL): + yield event + + _ = [chunk async for chunk in iterator.async_sse_wrapper(_full_stream())] + (logging_coroutine,) = logging_obj._deferred_stream_complete_args + await logging_coroutine + + assert len(recorder.success_kwargs) == 1 + assert recorder.success_kwargs[0]["standard_logging_object"]["startTime"] == call_start.timestamp() + + +@pytest.mark.asyncio +async def test_streamed_messages_from_anthropic_log_the_call_start_instead_of_when_the_upstream_answered(): + model: Final = "claude-sonnet-5" + call_start: Final = datetime(2026, 10, 9, 12, 0, 0) + recorder: Final = _SuccessRecorder() + logging_obj: Final = _make_priced_logging_obj("anthropic_logs_call_start", recorder, model, start_time=call_start) + logging_obj._on_deferred_stream_complete = lambda logging_coroutine: logging_coroutine + request_body: Final = {"model": model, "stream": True} + iterator: Final = BaseAnthropicMessagesStreamingIterator(litellm_logging_obj=logging_obj, request_body=request_body) + sse: Final = "".join(f"event: {e['type']}\ndata: {json.dumps(e)}\n\n" for e in (*_STREAM_PREFIX, *_STREAM_TAIL)) + + _ = [ + chunk + async for chunk in iterator.get_async_streaming_response_iterator( + httpx.Response(200, content=sse.encode()), request_body, logging_obj + ) + ] + (logging_coroutine,) = logging_obj._deferred_stream_complete_args + await logging_coroutine + + assert len(recorder.success_kwargs) == 1 + assert recorder.success_kwargs[0]["standard_logging_object"]["startTime"] == call_start.timestamp() + + class _UpstreamClosedOnDetach: """Upstream that yields its events and then, like a socket read, waits until it is closed.""" From b88f64e65f29a34ee7363a3e2cd8189e850e4387 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 14:17:07 +0000 Subject: [PATCH 10/33] refactor(types): replace Any with proven types in 6 files (#47268) * refactor(types): replace Any with proven types in 39 files Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(anthropic): use valid response fixtures for strict parsing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(types): revert anthropic chat response seam Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(types): revert validation seams that changed tolerated inputs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(types): revert remaining validation seams to keep behavior unchanged Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(types): drop no-op search router annotation and helper docstring Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(types): keep ip address utils at base Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../integrations/datadog/datadog_llm_obs.py | 24 ++++++++++--- litellm/llms/base_llm/chat/transformation.py | 2 +- litellm/models/organization_membership.py | 3 +- litellm/proxy/auth/auth_checks.py | 2 +- .../workflow_management_endpoints.py | 35 +++++++++++++++---- .../complexity_router/complexity_router.py | 19 +++++++--- 6 files changed, 66 insertions(+), 19 deletions(-) diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py index 270d799742c..a72023eaedf 100644 --- a/litellm/integrations/datadog/datadog_llm_obs.py +++ b/litellm/integrations/datadog/datadog_llm_obs.py @@ -12,7 +12,7 @@ import os from collections.abc import Mapping, Sequence from datetime import datetime from types import MappingProxyType -from typing import Any, Final, Literal +from typing import Final, Literal import httpx from pydantic import ConfigDict, TypeAdapter, ValidationError @@ -54,6 +54,7 @@ from litellm.types.utils import ( PROMPT_CARRYING_GUARDRAIL_FIELDS, PROMPT_QUOTING_ROUTING_DECISION_FIELDS, CallTypes, + StandardLoggingMetadata, StandardLoggingPayload, StandardLoggingPayloadErrorInformation, ) @@ -93,8 +94,12 @@ _ROUTER_DIMENSIONS: Final[tuple[str, ...]] = ("router_tier", "router_cause", "ro _COST_DIMENSIONS: Final[tuple[str, ...]] = ("team", "user", "key_alias", "model_group", *_ROUTER_DIMENSIONS) -def _metadata_of(standard_logging_payload: StandardLoggingPayload) -> Mapping[str, Any]: - metadata: Final = standard_logging_payload.get("metadata") +def _logging_metadata_of(standard_logging_payload: StandardLoggingPayload) -> StandardLoggingMetadata | None: + return standard_logging_payload.get("metadata") + + +def _metadata_of(standard_logging_payload: StandardLoggingPayload) -> Mapping[str, object]: + metadata: Final = _logging_metadata_of(standard_logging_payload) return metadata or _EMPTY_MAPPING @@ -1057,14 +1062,23 @@ class DataDogLLMObsLogger(CustomBatchLogger): # Get budget information from metadata metadata: Final = _metadata_of(standard_logging_payload) + logging_metadata: Final = _logging_metadata_of(standard_logging_payload) # API key max budget - user_api_key_max_budget: Final = metadata.get("user_api_key_max_budget") + user_api_key_max_budget: Final = ( + logging_metadata["user_api_key_max_budget"] + if logging_metadata and "user_api_key_max_budget" in logging_metadata + else None + ) if user_api_key_max_budget is not None: spend_metrics["user_api_key_max_budget"] = float(user_api_key_max_budget) # API key spend - user_api_key_spend: Final = metadata.get("user_api_key_spend") + user_api_key_spend: Final = ( + logging_metadata["user_api_key_spend"] + if logging_metadata and "user_api_key_spend" in logging_metadata + else None + ) if user_api_key_spend is not None: try: spend_metrics["user_api_key_spend"] = float(user_api_key_spend) diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index 51c8e4b8bef..b2eb4cde5ce 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -428,7 +428,7 @@ class BaseConfig(ABC): streaming_response: Union[Iterator[str], AsyncIterator[str], "ModelResponse"], sync_stream: bool, json_mode: bool | None = False, - ) -> Any: + ) -> object: pass def get_model_response_iterator_for_request( diff --git a/litellm/models/organization_membership.py b/litellm/models/organization_membership.py index e9697cf57b6..6a0d22cb63b 100644 --- a/litellm/models/organization_membership.py +++ b/litellm/models/organization_membership.py @@ -6,7 +6,6 @@ Canonical definition for ``litellm_organizationmembership``. Re-exported from """ from datetime import datetime -from typing import Any from pydantic import ConfigDict, model_validator @@ -24,7 +23,7 @@ class LiteLLM_OrganizationMembershipTable(LiteLLMPydanticObjectBase): budget_id: str | None = None created_at: datetime updated_at: datetime - user: Any | None = None + user: object | None = None litellm_budget_table: LiteLLM_BudgetTable | None = None user_email: str | None = None diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 5a09a0efbf6..8887d86570a 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2663,7 +2663,7 @@ async def _load_team_membership_on_cache_miss( parent_otel_span: Span | None, proxy_logging_obj: ProxyLogging | None, ) -> LiteLLM_TeamMembership | None: - redis_cached: Final[object] = await user_api_key_cache.async_get_cache(key=cache_key) + redis_cached: Final = await _raw_cache(user_api_key_cache).async_get_cache(key=cache_key) redis_membership: Final = _membership_from_cached_payload(redis_cached) if not isinstance(redis_membership, _TeamMembershipCacheMiss): return redis_membership diff --git a/litellm/proxy/management_endpoints/workflow_management_endpoints.py b/litellm/proxy/management_endpoints/workflow_management_endpoints.py index 890a344dfe0..649941bc233 100644 --- a/litellm/proxy/management_endpoints/workflow_management_endpoints.py +++ b/litellm/proxy/management_endpoints/workflow_management_endpoints.py @@ -147,6 +147,34 @@ class _MessageCreateData(TypedDict, total=False): session_id: str +class _EventTable(Protocol): + async def create(self, *, data: _EventCreateData) -> object: ... + + +class _RunTable(Protocol): + async def update(self, *, where: Mapping[str, str], data: Mapping[str, str]) -> object: ... + + +class _WorkflowTx(Protocol): + @property + def litellm_workflowevent(self) -> _EventTable: ... + + @property + def litellm_workflowrun(self) -> _RunTable: ... + + +async def _append_event_in_tx( + tx: _WorkflowTx, run_id: str, event_data: _EventCreateData, new_status: str | None +) -> object: + event: Final = await tx.litellm_workflowevent.create(data=event_data) + if new_status: + await tx.litellm_workflowrun.update( + where={"run_id": run_id}, + data={"status": new_status}, + ) + return event + + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -387,12 +415,7 @@ async def append_workflow_event( event_data["data"] = _json(data.data) async with prisma_client.db.tx() as tx: - event: object = await tx.litellm_workflowevent.create(data=event_data) - if new_status: - await tx.litellm_workflowrun.update( - where={"run_id": run_id}, - data={"status": new_status}, - ) + event: object = await _append_event_in_tx(tx, run_id, event_data, new_status) return event diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 1638f20c8e3..a475358e1f0 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -22,7 +22,7 @@ import json import random import re import time -from collections.abc import Callable, Iterator, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from itertools import accumulate, chain, islice, takewhile from threading import Lock from types import MappingProxyType @@ -30,6 +30,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast from pydantic import BaseModel, TypeAdapter, ValidationError, create_model from pydantic_core import ErrorDetails +from typing_extensions import ReadOnly, TypedDict from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger @@ -411,9 +412,19 @@ def _append_custom_keywords(base_keywords: list[str], custom_keywords: list[str] return [*base_keywords, *deduped_custom.values()] -def _parent_session_kwargs(request_kwargs: Mapping[str, object] | None) -> Mapping[str, Any]: +class _ParentSessionKwargs(TypedDict, total=False): + litellm_session_id: ReadOnly[object] + litellm_trace_id: ReadOnly[object] + + +def _parent_session_kwargs(request_kwargs: Mapping[str, object] | None) -> _ParentSessionKwargs: kwargs: Final = request_kwargs or {} - return {k: kwargs[k] for k in ("litellm_session_id", "litellm_trace_id") if kwargs.get(k) is not None} + session_id: Final = kwargs.get("litellm_session_id") + trace_id: Final = kwargs.get("litellm_trace_id") + return _ParentSessionKwargs( + **({"litellm_session_id": session_id} if session_id is not None else {}), + **({"litellm_trace_id": trace_id} if trace_id is not None else {}), + ) def _response_cost_or_none(response: object) -> float | None: @@ -2828,7 +2839,7 @@ class ComplexityRouter(CustomLogger): "originating_request_masked": masked_originating_request(request_kwargs), "body": {"model": llm_config.model, **payload}, } - classify: Final = ( + classify: Final[Callable[..., Awaitable[ModelResponse | ResponsesAPIResponse]]] = ( self.litellm_router_instance.aresponses if encrypted_task is not None else self.litellm_router_instance.acompletion From 4963ba4d9906170427a82731ba1b26c4f82be1b8 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 20 Oct 2026 07:56:29 -0700 Subject: [PATCH 11/33] test(translation): add sonnet 5.5 anthropic direct cross-model mid row (#47566) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../chat_completions/cross_model/README.md | 4 +- .../chat_completions/cross_model/anthropic.py | 350 ++++++++++++++++++ .../chat_completions/cross_model/case.py | 5 + 3 files changed, 357 insertions(+), 2 deletions(-) diff --git a/tests/integration/translation/chat_completions/cross_model/README.md b/tests/integration/translation/chat_completions/cross_model/README.md index 90710d5597c..99a4225480e 100644 --- a/tests/integration/translation/chat_completions/cross_model/README.md +++ b/tests/integration/translation/chat_completions/cross_model/README.md @@ -8,7 +8,7 @@ Every row sends `CODING_AGENT_REQUEST` to `/v1/chat/completions`, changing only | Mid | Opus 5 and 5.5, Sonnet 5.5, Gemini 3.8 Flash and 3.1 Pro Preview, GPT 5.6 Sol, 6 Sol and 6.1 Sol, GPT 5.6 Terra | | High | Fable 5 and 5.1, Gemini 4 Argon, GPT 6 Astra | -The 51 rows cover Anthropic direct, OpenAI direct, Vertex, Gemini direct, Azure OpenAI, Azure AI and Bedrock Converse where accessible. Vertex Haiku 5.5 was unavailable during capture and is not represented +The 52 rows cover Anthropic direct, OpenAI direct, Vertex, Gemini direct, Azure OpenAI, Azure AI and Bedrock Converse where accessible. Vertex Haiku 5.5 was unavailable during capture and is not represented The request contains system instructions, conversation history, previous assistant tool calls and results, four function definitions, automatic parallel tools, low reasoning effort and an output limit. Each row exercises a tool-call response, raw streaming, SDK stream assembly and a follow-up that replays the returned assistant message and tool-call IDs @@ -33,7 +33,7 @@ Streaming cases store `expected_litellm_message` and `expected_litellm_usage`. M Usage requires nonnegative integer counts, positive prompt and completion counts, and `total_tokens == prompt_tokens + completion_tokens`. The complete normalized usage must equal the independently expected counters, including cached and reasoning tokens and derived text/cache details. Missing, changed or unexpected counters fail. Only null optional fields, the explicitly listed provider metadata/duplicate counters and computed `cost` are excluded. Thus a wrong prompt count with an adjusted total still fails -The raw stream check assembles tool names and argument fragments by index. The SDK check separately uses `OpenAI.chat.completions.stream()` without removing reasoning fields. All 51 SDK checks run, including the seven Haiku rows. [LIT-9653](https://linear.app/litellm-ai/issue/LIT-9653) was fixed in [#46611](https://github.com/BerriAI/litellm/pull/46611), so there are no Haiku skips +The raw stream check assembles tool names and argument fragments by index. The SDK check separately uses `OpenAI.chat.completions.stream()` without removing reasoning fields. All 52 SDK checks run, including the seven Haiku rows. [LIT-9653](https://linear.app/litellm-ai/issue/LIT-9653) was fixed in [#46611](https://github.com/BerriAI/litellm/pull/46611), so there are no Haiku skips ## Recorded upstreams diff --git a/tests/integration/translation/chat_completions/cross_model/anthropic.py b/tests/integration/translation/chat_completions/cross_model/anthropic.py index 3df7da17ef6..9302103f369 100644 --- a/tests/integration/translation/chat_completions/cross_model/anthropic.py +++ b/tests/integration/translation/chat_completions/cross_model/anthropic.py @@ -11,6 +11,7 @@ from integration.translation.chat_completions.bases.anthropic import ( CLAUDE_HAIKU_5_5_TEST_CASE, CLAUDE_OPUS_5_5_TEST_CASE, CLAUDE_OPUS_5_TEST_CASE, + CLAUDE_SONNET_5_5_TEST_CASE, ) from integration.translation.chat_completions.templates import ( CODING_AGENT_REQUEST, @@ -2098,3 +2099,352 @@ CLAUDE_OPUS_5_AGENT_STREAMING_TEST_CASE: Final = StreamingTranslationTestCase( "completion_tokens_details": {"reasoning_tokens": 0, "text_tokens": 105}, }, ) + + +CLAUDE_SONNET_5_5_AGENT_FOLLOWUP_TEXT: Final = ( + "The bug is that `retries_remaining` in `src/retry.py` returns `limit + attempts` instead of " + "`max(0, limit - attempts)`, so it can return too many retries and can't be clamped at zero. The only other " + "caller is `src/worker.py` line 12 (`retries_remaining(3, attempts)`), and `tests/test_retry.py::test_retries_remaining` " + "fails with `assert 4 == 2` for `retries_remaining(3, 1)`." +) + + +CLAUDE_SONNET_5_5_AGENT_TEST_CASE: Final = replace( + CLAUDE_SONNET_5_5_TEST_CASE, + scenario="agent_tools", + litellm_request={**CODING_AGENT_REQUEST, "model": "anthropic/claude-sonnet-5-5"}, + expected_provider_endpoint="/v1/messages", + expected_provider_headers=CLAUDE_OPUS_5_5_AGENT_TEST_CASE.expected_provider_headers, + expected_provider_request={ + **ANTHROPIC_AGENT_REQUEST, + "model": "claude-sonnet-5-5", + "thinking": {"type": "adaptive", "display": "summarized"}, + "output_config": {"effort": "low"}, + }, + mock_provider_response={ + **ANTHROPIC_AGENT_RESPONSE, + "model": "claude-sonnet-5-5", + "id": "msg_011CgDcA1vnTwnTPTcJv9Lb7", + "content": [ + { + "type": "tool_use", + "id": "toolu_01RK8oPY5AnH3K4R8LfPXPTi", + "name": "search_code", + "input": {"query": "retries_remaining"}, + "caller": {"type": "direct"}, + }, + { + "type": "tool_use", + "id": "toolu_013REgHECA7Rceip3FrejMEy", + "name": "run_tests", + "input": {"path": "tests/test_retry.py"}, + "caller": {"type": "direct"}, + }, + ], + "stop_reason": "tool_use", + "usage": { + **object_value(ANTHROPIC_AGENT_RESPONSE["usage"]), + "input_tokens": 1184, + "output_tokens": 105, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + "speed": "standard", + }, + "diagnostics": None, + }, + expected_litellm_response={ + **CODING_AGENT_RESPONSE, + "model": "anthropic/claude-sonnet-5-5", + "choices": [ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "role": "assistant", + "tool_calls": [ + { + "index": 0, + "caller": {"type": "direct"}, + "function": { + "arguments": '{"query": "retries_remaining"}', + "name": "search_code", + }, + "id": "toolu_01RK8oPY5AnH3K4R8LfPXPTi", + "type": "function", + }, + { + "index": 1, + "caller": {"type": "direct"}, + "function": { + "arguments": '{"path": "tests/test_retry.py"}', + "name": "run_tests", + }, + "id": "toolu_013REgHECA7Rceip3FrejMEy", + "type": "function", + }, + ], + "provider_specific_fields": {"citations": None, "thinking_blocks": None}, + "content": None, + }, + } + ], + "usage": { + "completion_tokens": 105, + "prompt_tokens": 1184, + "total_tokens": 1289, + "completion_tokens_details": {"reasoning_tokens": 0, "text_tokens": 105}, + "prompt_tokens_details": { + "cached_tokens": 0, + "text_tokens": 1184, + "cache_write_tokens": 0, + "cache_creation_tokens": 0, + "cache_creation_token_details": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + }, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "inference_geo": "global", + "speed": "standard", + "service_tier": "standard", + }, + }, +) + + +CLAUDE_SONNET_5_5_AGENT_FOLLOWUP_TEST_CASE: Final = replace( + CLAUDE_SONNET_5_5_AGENT_TEST_CASE, + scenario="agent_followup", + litellm_request=followup_request( + CLAUDE_SONNET_5_5_AGENT_TEST_CASE.litellm_request, + CLAUDE_SONNET_5_5_AGENT_TEST_CASE.expected_litellm_response, + ), + expected_provider_endpoint="/v1/messages", + expected_provider_request={ + **ANTHROPIC_AGENT_REQUEST, + "model": "claude-sonnet-5-5", + "messages": [ + *list_value(ANTHROPIC_AGENT_REQUEST["messages"]), + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_01RK8oPY5AnH3K4R8LfPXPTi", + "name": "search_code", + "input": {"query": "retries_remaining"}, + }, + { + "type": "tool_use", + "id": "toolu_013REgHECA7Rceip3FrejMEy", + "name": "run_tests", + "input": {"path": "tests/test_retry.py"}, + }, + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_01RK8oPY5AnH3K4R8LfPXPTi", + "content": '{"matches": [{"path": "src/worker.py", "line": 12, "text": "remaining = ' + 'retries_remaining(3, attempts)"}]}', + }, + { + "type": "tool_result", + "tool_use_id": "toolu_013REgHECA7Rceip3FrejMEy", + "content": '{"exit_code": 1, "output": "FAILED tests/test_retry.py::test_retries_remaining - ' + 'assert 4 == 2"}', + }, + { + "type": "text", + "text": "Stop after these diagnostics. Do not call more tools or change files. Summarize the " + "bug, the caller found and the failed test in two short sentences.", + }, + ], + }, + ], + "thinking": {"type": "adaptive", "display": "summarized"}, + "output_config": {"effort": "low"}, + }, + mock_provider_response={ + **ANTHROPIC_AGENT_RESPONSE, + "model": "claude-sonnet-5-5", + "id": "msg_011CgDcZdb2JGywm8DnwxFvn", + "content": [{"type": "text", "text": CLAUDE_SONNET_5_5_AGENT_FOLLOWUP_TEXT}], + "stop_reason": "end_turn", + "usage": { + **object_value(ANTHROPIC_AGENT_RESPONSE["usage"]), + "input_tokens": 1492, + "output_tokens": 155, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + "speed": "standard", + }, + "diagnostics": None, + }, + expected_litellm_response={ + **CODING_AGENT_RESPONSE, + "model": "anthropic/claude-sonnet-5-5", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": CLAUDE_SONNET_5_5_AGENT_FOLLOWUP_TEXT, + "role": "assistant", + "provider_specific_fields": {"citations": None, "thinking_blocks": None}, + }, + } + ], + "usage": { + "completion_tokens": 155, + "prompt_tokens": 1492, + "total_tokens": 1647, + "completion_tokens_details": {"reasoning_tokens": 0, "text_tokens": 155}, + "prompt_tokens_details": { + "cached_tokens": 0, + "text_tokens": 1492, + "cache_write_tokens": 0, + "cache_creation_tokens": 0, + "cache_creation_token_details": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + }, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "inference_geo": "global", + "speed": "standard", + "service_tier": "standard", + }, + }, +) + + +CLAUDE_SONNET_5_5_AGENT_STREAMING_TEST_CASE: Final = StreamingTranslationTestCase( + scenario="agent_streaming", + litellm_endpoint=CLAUDE_SONNET_5_5_AGENT_TEST_CASE.litellm_endpoint, + litellm_request={ + **CLAUDE_SONNET_5_5_AGENT_TEST_CASE.litellm_request, + "stream": True, + "stream_options": {"include_usage": True}, + }, + expected_provider_headers=CLAUDE_SONNET_5_5_AGENT_TEST_CASE.expected_provider_headers, + expected_provider_endpoint="/v1/messages", + expected_provider_request={ + **CLAUDE_SONNET_5_5_AGENT_TEST_CASE.expected_provider_request, + "stream": True, + }, + mock_provider_frames=( + SSEFrame( + data={ + "type": "message_start", + "message": { + **ANTHROPIC_AGENT_RESPONSE, + "model": "claude-sonnet-5-5", + "id": "msg_011CgDcC8agu11ruKbps6qFH", + "content": [], + "stop_reason": None, + "usage": { + **object_value(ANTHROPIC_AGENT_RESPONSE["usage"]), + "input_tokens": 1184, + "output_tokens": 16, + "service_tier": "standard", + "inference_geo": "global", + "speed": "standard", + }, + "diagnostics": None, + }, + }, + event="message_start", + ), + SSEFrame( + data={ + "type": "content_block_start", + "index": 0, + "content_block": { + "type": "tool_use", + "id": "toolu_01LmCacjLUDXcY139vPTCceR", + "name": "search_code", + "input": {}, + "caller": {"type": "direct"}, + }, + }, + event="content_block_start", + ), + SSEFrame(data={"type": "ping"}, event="ping"), + *( + SSEFrame( + data={ + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": partial_json}, + }, + event="content_block_delta", + ) + for partial_json in ("", '{"query', '": "ret', "ries_re", "maining", '"}') + ), + SSEFrame(data={"type": "content_block_stop", "index": 0}, event="content_block_stop"), + SSEFrame( + data={ + "type": "content_block_start", + "index": 1, + "content_block": { + "type": "tool_use", + "id": "toolu_011cyQH4o1bcDJYsWQoLSzGT", + "name": "run_tests", + "input": {}, + "caller": {"type": "direct"}, + }, + }, + event="content_block_start", + ), + *( + SSEFrame( + data={ + "type": "content_block_delta", + "index": 1, + "delta": {"type": "input_json_delta", "partial_json": partial_json}, + }, + event="content_block_delta", + ) + for partial_json in ("", '{"path"', ': "test', "s/test_", "retry.p", 'y"}') + ), + SSEFrame(data={"type": "content_block_stop", "index": 1}, event="content_block_stop"), + SSEFrame( + data={ + "type": "message_delta", + "delta": { + "stop_reason": "tool_use", + "stop_sequence": None, + "stop_details": None, + "container": None, + }, + "usage": { + "input_tokens": 1184, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "output_tokens": 105, + "output_tokens_details": {"thinking_tokens": 0}, + }, + }, + event="message_delta", + ), + SSEFrame(data={"type": "message_stop"}, event="message_stop"), + ), + expected_litellm_message=coding_agent_stream_message( + ("toolu_01LmCacjLUDXcY139vPTCceR", "toolu_011cyQH4o1bcDJYsWQoLSzGT") + ), + expected_litellm_usage={ + "prompt_tokens": 1184, + "completion_tokens": 105, + "total_tokens": 1289, + "prompt_tokens_details": { + "cached_tokens": 0, + "text_tokens": 1184, + "cache_write_tokens": 0, + "cache_creation_tokens": 0, + "cache_creation_token_details": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + }, + "completion_tokens_details": {"reasoning_tokens": 0, "text_tokens": 105}, + }, +) diff --git a/tests/integration/translation/chat_completions/cross_model/case.py b/tests/integration/translation/chat_completions/cross_model/case.py index 189592b26c2..8dee0fae2ab 100644 --- a/tests/integration/translation/chat_completions/cross_model/case.py +++ b/tests/integration/translation/chat_completions/cross_model/case.py @@ -132,6 +132,11 @@ MID: Final[tuple[AgentCases, ...]] = ( bedrock_converse.CLAUDE_SONNET_5_5_AGENT_FOLLOWUP_TEST_CASE, bedrock_converse.CLAUDE_SONNET_5_5_AGENT_STREAMING_TEST_CASE, ), + ( + anthropic.CLAUDE_SONNET_5_5_AGENT_TEST_CASE, + anthropic.CLAUDE_SONNET_5_5_AGENT_FOLLOWUP_TEST_CASE, + anthropic.CLAUDE_SONNET_5_5_AGENT_STREAMING_TEST_CASE, + ), ( bedrock_converse.CLAUDE_OPUS_5_TEST_CASE, bedrock_converse.CLAUDE_OPUS_5_AGENT_FOLLOWUP_TEST_CASE, From 9e3d3a87f4016a621d14e808c916f61d59d9e8af Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 15:00:31 +0000 Subject: [PATCH 12/33] fix(anthropic, bedrock)!: keep a redacted_thinking content item on every Claude route (#47489) * test(anthropic, bedrock): cover a redacted_thinking content item on every Claude route * fix(anthropic, bedrock): keep a redacted_thinking content item on every Claude route * fix(anthropic, bedrock): send each replayed thinking block once, without cache_control or blank data * fix(anthropic, bedrock): keep thinking order and gate replay key drops on drop_params A thinking block that only one of content and thinking_blocks carries now goes back in its original place instead of being prepended. cache_control on a thinking block and blank redacted data are dropped only under drop_params and otherwise reach Claude as sent. Converse drops a cachePoint after reasoning under drop_params, skips a thinking_blocks entry that is not a block, and the cache control hook no longer marks a thinking part. * fix(anthropic, snowflake): match a signed thinking block on its signature and send each block once on Snowflake Cortex A signed thinking block in both content and thinking_blocks is the same block when the signatures match, even if the thinking text differs, so it is sent once from content instead of twice. Snowflake Cortex now merges thinking_blocks into content the same way the Anthropic routes do. A redacted_thinking item with blank data no longer counts as sent for system message placement. * fix(anthropic, bedrock): send a thinking block once when consecutive assistant messages both carry it Consecutive assistant messages merge into one Claude turn, so a block one message holds in content and the next holds in thinking_blocks went out twice and Claude answered 400. The merged turn now keeps the first copy of each redacted or signed block on the Anthropic routes and both Bedrock Converse builders. Snowflake Cortex goes back to its previous message building, since its double send cannot be proven live without a Cortex account and is tracked on its own. * test(anthropic, bedrock): cover a redacted block with null data and a non-block part beside a cache breakpoint * fix(anthropic, bedrock): match only text-valued thinking items when merging replayed blocks --- .../anthropic_cache_control_hook.py | 21 +- .../prompt_templates/common_utils.py | 10 + .../prompt_templates/factory.py | 256 +++++++---- .../mid_conversation_system.py | 19 +- litellm/llms/anthropic/chat/transformation.py | 5 +- .../bedrock/chat/converse_transformation.py | 2 + .../test_anthropic_cache_control_hook.py | 33 ++ ...llm_core_utils_prompt_templates_factory.py | 409 +++++++++++++++++- ...rompt_templates_mid_conversation_system.py | 28 +- 9 files changed, 660 insertions(+), 123 deletions(-) diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index abdbfc79d41..b0deec24c1a 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -342,6 +342,20 @@ EXTERNAL_BREAKPOINTS_STAMP: Final = "_litellm_external_breakpoints" CUSTOM_LLM_PROVIDER_STAMP: Final = "_litellm_custom_llm_provider" +def _accepts_cache_control(part: object) -> bool: + match part: + case {"type": "thinking" | "redacted_thinking"}: + return False + case dict(): + return True + case _: + return False + + +def _last_index_accepting_cache_control(parts: Sequence[object]) -> int | None: + return next((index for index in reversed(range(len(parts))) if _accepts_cache_control(parts[index])), None) + + class AnthropicCacheControlHook(CustomPromptManagement): @staticmethod def _request_value(request_kwargs: object, key: str) -> object: @@ -740,10 +754,11 @@ class AnthropicCacheControlHook(CustomPromptManagement): # 1. if string, insert cache control in the message if isinstance(message_content, str): message["cache_control"] = control - # 2. list of objects - only apply to last item per Anthropic spec + # 2. list of objects - only apply to the last item Claude accepts a breakpoint on elif isinstance(message_content, list): - if len(message_content) > 0 and isinstance(message_content[-1], dict): - message_content[-1]["cache_control"] = control # pyright: ignore[reportGeneralTypeIssues] # loose runtime dict + markable_index: Final = _last_index_accepting_cache_control(message_content) + if markable_index is not None: + message_content[markable_index]["cache_control"] = control # pyright: ignore[reportGeneralTypeIssues] # loose runtime dict return message @staticmethod diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index b704dfa2675..b89e4d609ad 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -2325,6 +2325,16 @@ def is_unsignable_thinking_block(block: object) -> bool: return not (isinstance(signature, str) and len(signature) > 0) +def has_blank_redacted_data(block: object) -> bool: + match block: + case {"type": "redacted_thinking", "data": str() as data}: + return not data.strip() + case {"type": "redacted_thinking"}: + return True + case _: + return False + + def native_redacted_thinking_data(data: str) -> str: """Claude's own `data` string for a redacted_thinking block. diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 481ec6a5076..3d865bd1819 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -6,10 +6,10 @@ import json import mimetypes import re import xml.etree.ElementTree as ET -from collections.abc import Container, Iterable, Iterator, Mapping, Sequence +from collections.abc import Callable, Container, Iterable, Iterator, Mapping, Sequence from enum import Enum from types import MappingProxyType -from typing import Annotated, Any, Final, TypeAlias, TypedDict, cast, overload +from typing import Annotated, Any, Final, TypeAlias, TypedDict, TypeVar, cast, overload from jinja2.sandbox import ImmutableSandboxedEnvironment from pydantic import BaseModel, BeforeValidator, ConfigDict, JsonValue, TypeAdapter, ValidationError, model_validator @@ -61,6 +61,7 @@ from litellm.types.utils import GenericImageParsingChunk from .common_utils import ( converse_redacted_content, convert_content_list_to_str, + has_blank_redacted_data, infer_content_type_from_url_and_content, is_encrypted_reasoning_block, is_non_content_values_set, @@ -2550,22 +2551,125 @@ def _pair_tool_calls_with_results(messages: Sequence[AllMessageValues]) -> tuple def _replayable_thinking_blocks( thinking_blocks: list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock], + *, + drop_params: bool = False, ) -> list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock]: + keys_to_drop: Final = _UNSUPPORTED_THINKING_KEYS if drop_params else _STREAM_ONLY_THINKING_KEYS return [ - _without_stream_index(_with_native_redacted_data(block)) + _without_keys(_with_native_redacted_data(block), keys_to_drop) for block in thinking_blocks - if not is_unsignable_thinking_block(block) + if not (is_unsignable_thinking_block(block) or (drop_params and has_blank_redacted_data(block))) ] -def _without_stream_index( +def _as_thinking_block(item: object) -> ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock | None: + match item: + case {"type": "thinking" | "redacted_thinking"}: + return cast( # cast-ok: the mapping pattern matched a thinking block's type tag + ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, item + ) + case _: + return None + + +def _thinking_identity(item: object) -> tuple[str | None, ...] | None: + match item: + case {"type": "redacted_thinking", "data": str() as data}: + return ("redacted_thinking", native_redacted_thinking_data(data)) + case {"type": "thinking", "signature": str() as signature} if signature: + return ("thinking", signature) + case {"type": str() as kind} if kind in ("thinking", "redacted_thinking"): + fields: Final = cast(Mapping[str, object], item) # cast-ok: the mapping pattern matched + return _unsigned_thinking_identity(kind, fields.get("thinking"), fields.get("signature")) + case _: + return None + + +def _unsigned_thinking_identity(kind: str, thinking: object, signature: object) -> tuple[str | None, ...] | None: + match thinking, signature: + case (str() | None as text, str() | None as signed): + return (kind, text, signed) + case _: + return None + + +_Replayed: Final = TypeVar("_Replayed") + + +def _with_missing_thinking_blocks( + content: Iterable[_Replayed], thinking_blocks: Sequence[_Replayed] +) -> tuple[_Replayed, ...]: + content_items: Final = tuple(content) + block_identities: Final = tuple(map(_thinking_identity, thinking_blocks)) + content_positions: Final[Mapping[tuple[object, ...] | None, int]] = MappingProxyType( + { + identity: index + for index, identity in reversed(tuple(enumerate(map(_thinking_identity, content_items)))) + if identity is not None + } + ) + anchors: Final = tuple(index for index, identity in enumerate(block_identities) if identity in content_positions) + if not anchors: + return (*thinking_blocks, *content_items) + missing_before: Final = MappingProxyType( + { + content_positions[block_identities[end]]: thinking_blocks[start + 1 : end] + for start, end in zip((-1, *anchors), anchors) + } + ) + missing_after: Final = MappingProxyType( + {content_positions[block_identities[anchors[-1]]]: thinking_blocks[anchors[-1] + 1 :]} + ) + return tuple(_interleave(content_items, missing_before, missing_after)) + + +def _without_repeated_thinking( + items: Sequence[_Replayed], identity: Callable[[_Replayed], tuple[object, ...] | None] +) -> tuple[_Replayed, ...]: + identities: Final = tuple(map(identity, items)) + first_positions: Final[Mapping[tuple[object, ...], int]] = MappingProxyType( + {key: index for index, key in reversed(tuple(enumerate(identities))) if key is not None} + ) + return tuple( + item for index, (item, key) in enumerate(zip(items, identities)) if key is None or first_positions[key] == index + ) + + +def _interleave( + content: Sequence[_Replayed], + before: Mapping[int, Sequence[_Replayed]], + after: Mapping[int, Sequence[_Replayed]], +) -> Iterator[_Replayed]: + for index, item in enumerate(content): + yield from before.get(index, ()) + yield item + yield from after.get(index, ()) + + +def _thinking_in_replay_order( + content: str | Iterable[object] | None, thinking_blocks: Sequence[object] | None +) -> tuple[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, ...]: + content_items: Final = content if isinstance(content, list) else () + return tuple( + block + for item in _with_missing_thinking_blocks(content_items, thinking_blocks or ()) + if (block := _as_thinking_block(item)) is not None + ) + + +_STREAM_ONLY_THINKING_KEYS: Final = frozenset({"index"}) +_UNSUPPORTED_THINKING_KEYS: Final = frozenset({"index", "cache_control"}) + + +def _without_keys( block: ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, + keys: frozenset[str], ) -> ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock: - if "index" not in block: + if keys.isdisjoint(block): return block - return cast( # cast-ok: dropping one key keeps the block's TypedDict shape + return cast( # cast-ok: dropping keys keeps the block's TypedDict shape ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, - {key: value for key, value in block.items() if key != "index"}, + {key: value for key, value in block.items() if key not in keys}, ) @@ -2629,6 +2733,7 @@ def anthropic_messages_pt( messages: list[AllMessageValues], model: str, llm_provider: str, + drop_params: bool = False, ) -> _AnthropicMessageList: """ format messages for anthropic @@ -2816,7 +2921,9 @@ def anthropic_messages_pt( _raw_thinking_blocks = assistant_content_block.get("thinking_blocks", None) thinking_blocks = ( - _replayable_thinking_blocks(_raw_thinking_blocks) if _raw_thinking_blocks is not None else None + _replayable_thinking_blocks(_raw_thinking_blocks, drop_params=drop_params) + if _raw_thinking_blocks is not None + else None ) # Check if tool_calls contain server tool calls (web search, etc.) @@ -2951,42 +3058,28 @@ def anthropic_messages_pt( # SEQUENTIAL MODE: No server tool calls, or no thinking blocks, # or content is a list. Use the original sequential approach. - # When content is a list, check if it already contains thinking - # blocks inline. If so, skip prepending thinking_blocks to avoid - # duplication and preserve the original interleaved order. - # Fixes the gap where list-content messages bypass INTERLEAVED - # MODE and still get thinking blocks prepended out of order. _content_is_list = "content" in assistant_content_block and isinstance( assistant_content_block["content"], list ) _content_list = assistant_content_block.get("content") if _content_is_list else None - _list_has_thinking = False - if _content_is_list and _content_list is not None: - for _item in _content_list: - if isinstance(_item, dict) and _item.get("type") in ( - "thinking", - "redacted_thinking", - ): - _list_has_thinking = True - break if ( - thinking_blocks is not None and not _list_has_thinking + thinking_blocks is not None and not _content_is_list ): # IMPORTANT: ADD THIS FIRST, ELSE ANTHROPIC WILL RAISE AN ERROR assistant_content.extend(thinking_blocks) if _content_is_list and _content_list is not None: - for m in _content_list: + for m in _with_missing_thinking_blocks(_content_list, thinking_blocks or ()): if not isinstance(m, dict): continue # handle thinking blocks - text_block = cast(str, m.get("text", "")) - if m.get("type", "") == "thinking" and not is_unsignable_thinking_block(m): - anthropic_message: ChatCompletionThinkingBlock | AnthropicMessagesTextParam = cast( - ChatCompletionThinkingBlock, m + if (thinking_block := _as_thinking_block(m)) is not None: + assistant_content.extend( + _replayable_thinking_blocks([thinking_block], drop_params=drop_params) ) - assistant_content.append(anthropic_message) + continue + text_block = cast(str, m.get("text", "")) # handle text - elif ( + if ( m.get("type", "") == "text" and len(text_block) > 0 ): # don't pass empty text blocks. anthropic api raises errors. anthropic_message = AnthropicMessagesTextParam(type="text", text=text_block) @@ -3058,7 +3151,12 @@ def anthropic_messages_pt( msg_i += 1 if assistant_content: - new_messages.append({"role": "assistant", "content": assistant_content}) + new_messages.append( + { + "role": "assistant", + "content": list(_without_repeated_thinking(assistant_content, _thinking_identity)), + } + ) ## MID-CONVERSATION SYSTEM MESSAGES (placement is the caller's job) ## while msg_i < len(messages) and messages[msg_i]["role"] == "system": @@ -4190,6 +4288,16 @@ def _convert_to_bedrock_tool_call_result( return content_block +def _reasoning_identity(block: BedrockContentBlock) -> tuple[object, ...] | None: + match block: + case {"reasoningContent": {"redactedContent": redacted}}: + return ("redacted_thinking", redacted) + case {"reasoningContent": {"reasoningText": {"signature": str() as signature}}} if signature: + return ("thinking", signature) + case _: + return None + + def _deduplicate_bedrock_content_blocks( blocks: list[BedrockContentBlock], block_key: str, @@ -4673,6 +4781,7 @@ class BedrockConverseMessagesProcessor: assistant_continue_message: str | ChatCompletionAssistantMessage | None = None, deployment_model: str | None = None, base_model: str | None = None, + drop_params: bool = False, ) -> list[BedrockMessageBlock]: from litellm.llms.bedrock.common_utils import bedrock_converse_accepts_redacted_thinking @@ -4826,33 +4935,21 @@ class BedrockConverseMessagesProcessor: ) _assistant_content = assistant_message_block.get("content", None) thinking_blocks = assistant_message_block.get("thinking_blocks") - - if thinking_blocks is not None: - converted_thinking_blocks = ( - BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks( - thinking_blocks, to_claude=to_claude - ) - ) - assistant_content = BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content( - thinking_blocks=converted_thinking_blocks, - assistant_parts=assistant_content, + converted_thinking_blocks = ( + BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks( + _thinking_in_replay_order(_assistant_content, thinking_blocks), to_claude=to_claude ) + ) + assistant_content = BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content( + thinking_blocks=converted_thinking_blocks, + assistant_parts=assistant_content, + ) if _assistant_content is not None and isinstance(_assistant_content, list): assistants_parts: list[BedrockContentBlock] = [] for element in _assistant_content: if isinstance(element, dict): - if element["type"] == "thinking" or element["type"] == "redacted_thinking": - thinking_block = BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks( - thinking_blocks=[element], to_claude=to_claude - ) - assistants_parts = ( - BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content( - thinking_blocks=thinking_block, - assistant_parts=assistants_parts, - ) - ) - elif element["type"] == "text": + if element["type"] == "text": # Skip completely empty strings to avoid blank content blocks if element.get("text", "").strip(): assistants_part = BedrockContentBlock(text=element["text"]) @@ -4871,7 +4968,9 @@ class BedrockConverseMessagesProcessor: model=model, deployment_model=deployment_model, ) - if _cache_point_block is not None: + if _cache_point_block is not None and not ( + drop_params and _as_thinking_block(element) is not None + ): assistants_parts.append(_cache_point_block) assistant_content.extend(assistants_parts) elif _assistant_content is not None and isinstance(_assistant_content, str): @@ -4899,7 +4998,11 @@ class BedrockConverseMessagesProcessor: msg_i += 1 - assistant_content = _deduplicate_bedrock_content_blocks(assistant_content, "toolUse") + assistant_content = list( + _without_repeated_thinking( + _deduplicate_bedrock_content_blocks(assistant_content, "toolUse"), _reasoning_identity + ) + ) assistant_content = _sort_bedrock_assistant_content_blocks(assistant_content) if assistant_content: @@ -5083,6 +5186,7 @@ def _bedrock_converse_messages_pt( assistant_continue_message: str | ChatCompletionAssistantMessage | None = None, deployment_model: str | None = None, base_model: str | None = None, + drop_params: bool = False, ) -> list[BedrockMessageBlock]: """ Converts given messages from OpenAI format to Bedrock format @@ -5244,35 +5348,21 @@ def _bedrock_converse_messages_pt( ) _assistant_content = assistant_message_block.get("content", None) thinking_blocks = assistant_message_block.get("thinking_blocks") - - if thinking_blocks is not None: - converted_thinking_blocks = ( - BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks( - thinking_blocks, to_claude=to_claude - ) - ) - assistant_content = BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content( - thinking_blocks=converted_thinking_blocks, - assistant_parts=assistant_content, + converted_thinking_blocks = ( + BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks( + _thinking_in_replay_order(_assistant_content, thinking_blocks), to_claude=to_claude ) + ) + assistant_content = BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content( + thinking_blocks=converted_thinking_blocks, + assistant_parts=assistant_content, + ) if _assistant_content is not None and isinstance(_assistant_content, list): assistants_parts: list[BedrockContentBlock] = [] for element in _assistant_content: if isinstance(element, dict): - if element["type"] == "thinking" or element["type"] == "redacted_thinking": - thinking_block = ( - BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks( - thinking_blocks=[element], to_claude=to_claude - ) - ) - assistants_parts = ( - BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content( - thinking_blocks=thinking_block, - assistant_parts=assistants_parts, - ) - ) - elif element["type"] == "text": + if element["type"] == "text": # AWS Bedrock doesn't allow empty or whitespace-only text content # Skip completely empty strings to avoid blank content blocks if element.get("text", "").strip(): @@ -5292,7 +5382,9 @@ def _bedrock_converse_messages_pt( model=model, deployment_model=deployment_model, ) - if _cache_point_block is not None: + if _cache_point_block is not None and not ( + drop_params and _as_thinking_block(element) is not None + ): assistants_parts.append(_cache_point_block) assistant_content.extend(assistants_parts) elif _assistant_content is not None and isinstance(_assistant_content, str): @@ -5313,7 +5405,11 @@ def _bedrock_converse_messages_pt( msg_i += 1 - assistant_content = _deduplicate_bedrock_content_blocks(assistant_content, "toolUse") + assistant_content = list( + _without_repeated_thinking( + _deduplicate_bedrock_content_blocks(assistant_content, "toolUse"), _reasoning_identity + ) + ) assistant_content = _sort_bedrock_assistant_content_blocks(assistant_content) if assistant_content: diff --git a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py index fd0096f9f4d..ecf419695d0 100644 --- a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py +++ b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py @@ -49,7 +49,7 @@ from litellm.types.llms.openai import ( ChatCompletionUserMessage, ) -from .common_utils import is_unsignable_thinking_block +from .common_utils import has_blank_redacted_data, is_unsignable_thinking_block CONVERTED_SYSTEM_NOTE: Final = ( "Operator note (not from the user): the following was originally a mid-conversation system-role reminder." @@ -267,27 +267,22 @@ def _block_index_by_message(blocks: Sequence[tuple[bool, tuple[int, ...]]]) -> M def _thinking_block_renders(block: object) -> bool: """A thinking block the converter keeps: one Anthropic can verify, so never bridged encrypted reasoning.""" - return message_field(block, "type") in _THINKING_BLOCK_TYPES and not is_unsignable_thinking_block(block) + return message_field(block, "type") in _THINKING_BLOCK_TYPES and not ( + is_unsignable_thinking_block(block) or has_blank_redacted_data(block) + ) def _assistant_part_renders(part: object) -> bool: """A text part always renders: the converter pads empty text with a placeholder.""" part_type: Final = message_field(part, "type") - if part_type == "thinking": + if part_type in _THINKING_BLOCK_TYPES: return _thinking_block_renders(part) return part_type in _RENDERED_ASSISTANT_PART_TYPES or ( isinstance(part_type, str) and part_type.endswith("_tool_result") ) -def _separate_thinking_blocks_render(message: object, parts: Sequence[object]) -> bool: - """``thinking_blocks`` reach the wire only when no inline thinking part claims the slot. - - The converter skips the separate blocks as soon as the content list carries a - ``thinking`` or ``redacted_thinking`` part, whether or not that part itself renders. - """ - if any(message_field(part, "type") in _THINKING_BLOCK_TYPES for part in parts): - return False +def _separate_thinking_blocks_render(message: object) -> bool: return any(_thinking_block_renders(block) for block in parts_of(message_field(message, "thinking_blocks"))) @@ -404,7 +399,7 @@ def _assistant_renders(message: object, answered: frozenset[str]) -> bool: parts: Final = parts_of(content) return ( any(_assistant_part_renders(part) for part in parts) - or _separate_thinking_blocks_render(message, parts) + or _separate_thinking_blocks_render(message) or _tool_calls_render(message, answered) or bool(message_field(message, "function_call")) or bool(message_field(message_field(message, "provider_specific_fields"), "compaction_blocks")) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 867dd935a03..9a63afd5cc0 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -2267,6 +2267,8 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): anthropic_messages_pt, ) + drop_params: Final = litellm.drop_params is True or litellm_params.get("drop_params") is True + if "tools" not in optional_params and messages is not None and has_tool_call_blocks(messages): optional_params["tools"], _ = self._map_tools(add_dummy_tool(custom_llm_provider="anthropic")) @@ -2293,7 +2295,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): AnthropicConfig._maybe_drop_speed_param( model=model, optional_params=optional_params, - drop_params=litellm.drop_params or litellm_params.get("drop_params") is True, + drop_params=drop_params, custom_llm_provider=self.custom_llm_provider, ) @@ -2351,6 +2353,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): model=model, messages=list(conversation), llm_provider=self._resolved_provider, + drop_params=drop_params, ) except Exception as e: raise AnthropicError( diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index b505ae47c61..cd3727401ef 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -2127,6 +2127,7 @@ class AmazonConverseConfig(BaseConfig): llm_provider="bedrock_converse", user_continue_message=litellm_params.pop("user_continue_message", None), base_model=self._deployment_base_model(_DEPLOYMENT_PARAMS.validate_python(litellm_params)), + drop_params=litellm.drop_params is True or litellm_params.get("drop_params") is True, ), capability_model, ) @@ -2207,6 +2208,7 @@ class AmazonConverseConfig(BaseConfig): llm_provider="bedrock_converse", user_continue_message=litellm_params.pop("user_continue_message", None), base_model=self._deployment_base_model(_DEPLOYMENT_PARAMS.validate_python(litellm_params)), + drop_params=litellm.drop_params is True or litellm_params.get("drop_params") is True, ), capability_model, ) diff --git a/tests/unit/integrations/test_anthropic_cache_control_hook.py b/tests/unit/integrations/test_anthropic_cache_control_hook.py index 22db24f6ffd..c70c7080ab8 100644 --- a/tests/unit/integrations/test_anthropic_cache_control_hook.py +++ b/tests/unit/integrations/test_anthropic_cache_control_hook.py @@ -1705,6 +1705,39 @@ def test_cache_control_hook_does_not_overwrite_existing_cache_control(): assert _count_cache_control(processed) == 1 +@pytest.mark.parametrize( + "content, marked_parts", + [ + ([{"type": "text", "text": "Checking."}, {"type": "redacted_thinking", "data": "abc"}], [0]), + ([{"type": "redacted_thinking", "data": "abc"}], []), + ([{"type": "thinking", "thinking": "hm", "signature": "s"}], []), + ([{"type": "text", "text": "Checking."}, "not a content block"], [0]), + ], + ids=["text-before-trailing-redacted-thinking", "redacted-thinking-only", "signed-thinking-only", "text-before-a-string"], +) +def test_assistant_injection_never_marks_a_thinking_part( + content: list[dict[str, str] | str], marked_parts: list[int] +): + """Claude 400s a cache_control on a thinking or redacted_thinking block, and Bedrock Converse 400s a cache + point after reasoning, so the breakpoint goes on the last part that can carry one.""" + messages: Final[list[AllMessageValues]] = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": copy.deepcopy(content)}, + {"role": "user", "content": "again"}, + ] + + _, processed, _ = AnthropicCacheControlHook().get_chat_completion_prompt( + model="anthropic/claude-sonnet-4-6", + messages=messages, + non_default_params={"cache_control_injection_points": [{"location": "message", "role": "assistant"}]}, + prompt_id=None, + prompt_variables=None, + dynamic_callback_params={}, + ) + + assert [index for index, part in enumerate(processed[1]["content"]) if "cache_control" in part] == marked_parts + + @pytest.mark.asyncio async def test_cache_control_hook_bedrock_payload_caps_cachepoints_at_four(monkeypatch: pytest.MonkeyPatch): """End-to-end: outgoing Bedrock payload must not exceed 4 cachePoint blocks. diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 046f448046b..fad1c6178dc 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -541,25 +541,66 @@ def test_bedrock_converse_assistant_with_empty_thinking_block_and_tool_calls(): _NATIVE_REDACTED_DATA: Final = "EmwKAhgBEgy3va3pzix/LafPsn4aDFIT2Xlxh0L5L8rLVyIwxtE3rAFBa8cr3qpP" _CONVERSE_REDACTED_DATA: Final = base64.b64encode(_NATIVE_REDACTED_DATA.encode()).decode() +REASONING_PLACEMENTS: Final = pytest.mark.parametrize( + "placement", ["thinking_blocks", "content", "both", "first-block-in-content", "last-block-in-content"] +) + + +def _assistant_with_reasoning( + placement: str, reasoning: Sequence[Mapping[str, object]], text: str +) -> Mapping[str, object]: + in_content: Final = {"role": "assistant", "content": [*reasoning, {"type": "text", "text": text}]} + return { + "thinking_blocks": {"role": "assistant", "content": text, "thinking_blocks": list(reasoning)}, + "content": in_content, + "both": {**in_content, "thinking_blocks": list(reasoning)}, + "first-block-in-content": { + "role": "assistant", + "content": [*reasoning[:1], {"type": "text", "text": text}], + "thinking_blocks": list(reasoning), + }, + "last-block-in-content": { + "role": "assistant", + "content": [*reasoning[-1:], {"type": "text", "text": text}], + "thinking_blocks": list(reasoning), + }, + }[placement] + + +SPLIT_REASONING: Final = pytest.mark.parametrize( + ("in_content", "in_thinking_blocks"), + [((0, 2), (0, 1, 2)), ((0, 1, 2), (1,))], + ids=["middle-block-only-in-thinking-blocks", "outer-blocks-only-in-content"], +) + + +COPY_TEXT: Final = pytest.mark.parametrize("copy_text", ["", "First"], ids=["emptied", "trimmed"]) + + +def _assistant_with_split_reasoning( + reasoning: Sequence[Mapping[str, object]], in_content: Sequence[int], in_thinking_blocks: Sequence[int], text: str +) -> Mapping[str, object]: + return { + "role": "assistant", + "content": [*(reasoning[index] for index in in_content), {"type": "text", "text": text}], + "thinking_blocks": [reasoning[index] for index in in_thinking_blocks], + } @pytest.mark.asyncio @pytest.mark.parametrize("use_async", [False, True]) -@pytest.mark.parametrize("placement", ["thinking_blocks", "content"]) +@REASONING_PLACEMENTS async def test_bedrock_converse_replays_redacted_thinking_as_redacted_content(use_async: bool, placement: str): """A redacted_thinking block that Bedrock returned must go back as reasoningContent.redactedContent, - in order, ahead of text and toolUse. It used to collapse into an empty unsigned reasoningText and get dropped.""" + in order, ahead of text and toolUse. It used to collapse into an empty unsigned reasoningText and get dropped. + Sent twice, when the history carries it in content and thinking_blocks, Bedrock 400s that it was modified.""" bedrock_reasoning_blocks: Final[list[BedrockConverseReasoningContentBlock]] = [ {"reasoningText": {"text": "Let me look this up.", "signature": "sig-abc"}}, {"redactedContent": _CONVERSE_REDACTED_DATA}, ] thinking_blocks: Final = litellm.AmazonConverseConfig()._transform_thinking_blocks(bedrock_reasoning_blocks) assistant_text: Final = "Checking the weather." - assistant_message: Final = ( - {"role": "assistant", "content": assistant_text, "thinking_blocks": thinking_blocks} - if placement == "thinking_blocks" - else {"role": "assistant", "content": [*thinking_blocks, {"type": "text", "text": assistant_text}]} - ) + assistant_message: Final = _assistant_with_reasoning(placement, thinking_blocks, assistant_text) messages: Final = [ {"role": "user", "content": "What's the weather in Paris?"}, { @@ -593,6 +634,231 @@ async def test_bedrock_converse_replays_redacted_thinking_as_redacted_content(us ] +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@COPY_TEXT +async def test_bedrock_converse_sends_a_signed_block_once_when_thinking_blocks_holds_other_text( + use_async: bool, copy_text: str +): + bedrock_reasoning_blocks: Final[list[BedrockConverseReasoningContentBlock]] = [ + {"reasoningText": {"text": "First.", "signature": "sig-1"}}, + {"redactedContent": _CONVERSE_REDACTED_DATA}, + ] + reasoning: Final = litellm.AmazonConverseConfig()._transform_thinking_blocks(bedrock_reasoning_blocks) + messages: Final = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [*reasoning, {"type": "text", "text": "hello"}], + "thinking_blocks": [{**reasoning[0], "thinking": copy_text}, reasoning[1]], + }, + {"role": "user", "content": "again"}, + ] + model: Final = "us.anthropic.claude-sonnet-4-6" + + result: Final = ( + await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( + messages=messages, model=model, llm_provider="bedrock" + ) + if use_async + else _bedrock_converse_messages_pt(messages=messages, model=model, llm_provider="bedrock") + ) + + assert result[1] == { + "role": "assistant", + "content": [*({"reasoningContent": block} for block in bedrock_reasoning_blocks), {"text": "hello"}], + } + + +_WEATHER_TOOL_CALL: Final = { + "id": "tooluse_weather_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, +} + + +def _block_split_across_assistant_messages(blocks: Sequence[Mapping[str, object]]) -> list[dict[str, object]]: + return [ + {"role": "user", "content": "What's the weather in Paris?"}, + {"role": "assistant", "content": list(blocks)}, + {"role": "assistant", "content": None, "thinking_blocks": list(blocks), "tool_calls": [_WEATHER_TOOL_CALL]}, + {"role": "tool", "tool_call_id": "tooluse_weather_1", "content": "18C, sunny"}, + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +async def test_bedrock_converse_sends_a_block_once_when_two_assistant_messages_carry_it(use_async: bool): + bedrock_reasoning_blocks: Final[list[BedrockConverseReasoningContentBlock]] = [ + {"reasoningText": {"text": "Let me look this up.", "signature": "sig-abc"}}, + {"redactedContent": _CONVERSE_REDACTED_DATA}, + ] + reasoning: Final = litellm.AmazonConverseConfig()._transform_thinking_blocks(bedrock_reasoning_blocks) + messages: Final = _block_split_across_assistant_messages(reasoning) + model: Final = "us.anthropic.claude-sonnet-4-6" + + result: Final = ( + await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( + messages=messages, model=model, llm_provider="bedrock" + ) + if use_async + else _bedrock_converse_messages_pt(messages=messages, model=model, llm_provider="bedrock") + ) + + assert result[1] == { + "role": "assistant", + "content": [ + *({"reasoningContent": block} for block in bedrock_reasoning_blocks), + {"toolUse": {"toolUseId": "tooluse_weather_1", "name": "get_weather", "input": {"city": "Paris"}}}, + ], + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@SPLIT_REASONING +async def test_bedrock_converse_keeps_reasoning_order_across_content_and_thinking_blocks( + use_async: bool, in_content: Sequence[int], in_thinking_blocks: Sequence[int] +): + """A reasoning block only one of content and thinking_blocks carries goes back in its original place.""" + bedrock_reasoning_blocks: Final[list[BedrockConverseReasoningContentBlock]] = [ + {"reasoningText": {"text": "First.", "signature": "sig-1"}}, + {"redactedContent": _CONVERSE_REDACTED_DATA}, + {"reasoningText": {"text": "Third.", "signature": "sig-3"}}, + ] + reasoning: Final = litellm.AmazonConverseConfig()._transform_thinking_blocks(bedrock_reasoning_blocks) + messages: Final = [ + {"role": "user", "content": "hi"}, + _assistant_with_split_reasoning(reasoning, in_content, in_thinking_blocks, "hello"), + {"role": "user", "content": "again"}, + ] + model: Final = "us.anthropic.claude-sonnet-4-6" + + result: Final = ( + await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( + messages=messages, model=model, llm_provider="bedrock" + ) + if use_async + else _bedrock_converse_messages_pt(messages=messages, model=model, llm_provider="bedrock") + ) + + assert result[1] == { + "role": "assistant", + "content": [*({"reasoningContent": block} for block in bedrock_reasoning_blocks), {"text": "hello"}], + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("entry", [None, "not a block"], ids=["none", "string"]) +async def test_bedrock_converse_skips_a_thinking_blocks_entry_that_is_not_a_block(use_async: bool, entry: object): + messages: Final = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello", "thinking_blocks": [entry]}, + {"role": "user", "content": "again"}, + ] + model: Final = "us.anthropic.claude-sonnet-4-6" + + result: Final = ( + await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( + messages=messages, model=model, llm_provider="bedrock" + ) + if use_async + else _bedrock_converse_messages_pt(messages=messages, model=model, llm_provider="bedrock") + ) + + assert result[1] == {"role": "assistant", "content": [{"text": "hello"}]} + + +MALFORMED_THINKING: Final = pytest.mark.parametrize( + "malformed", + [ + {"type": "thinking", "thinking": {"a": 1}}, + {"type": "thinking", "thinking": ["a"], "signature": ""}, + {"type": "thinking", "thinking": "draft", "signature": ["a"]}, + ], + ids=["dict-thinking", "list-thinking", "list-signature"], +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@REASONING_PLACEMENTS +@MALFORMED_THINKING +async def test_bedrock_converse_replays_redacted_thinking_once_beside_a_malformed_thinking_item( + use_async: bool, placement: str, malformed: dict[str, object] +): + messages: Final = [ + {"role": "user", "content": "hi"}, + _assistant_with_reasoning( + placement, [{"type": "redacted_thinking", "data": _NATIVE_REDACTED_DATA}, malformed], "hello" + ), + {"role": "user", "content": "again"}, + ] + model: Final = "us.anthropic.claude-sonnet-4-6" + + result: Final = ( + await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( + messages=messages, model=model, llm_provider="bedrock" + ) + if use_async + else _bedrock_converse_messages_pt(messages=messages, model=model, llm_provider="bedrock") + ) + + assert result[1]["content"].count({"reasoningContent": {"redactedContent": _CONVERSE_REDACTED_DATA}}) == 1 + + +@REASONING_PLACEMENTS +@MALFORMED_THINKING +def test_anthropic_messages_pt_replays_redacted_thinking_once_beside_a_malformed_thinking_item( + placement: str, malformed: dict[str, object] +): + redacted: Final = {"type": "redacted_thinking", "data": _NATIVE_REDACTED_DATA} + messages: Final = [ + {"role": "user", "content": "hi"}, + _assistant_with_reasoning(placement, [redacted, malformed], "hello"), + {"role": "user", "content": "again"}, + ] + + result: Final = anthropic_messages_pt(messages=messages, model="claude-sonnet-4-6", llm_provider="anthropic") + + assert result[1]["content"].count(redacted) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("drop_params", [False, True]) +async def test_bedrock_converse_drops_a_cache_point_after_reasoning_only_under_drop_params( + use_async: bool, drop_params: bool +): + """Bedrock 400s "Cache point cannot be inserted after reasoning block", so drop_params drops the + breakpoint and without it the request reaches Bedrock as sent.""" + messages: Final = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "redacted_thinking", "data": _NATIVE_REDACTED_DATA, "cache_control": {"type": "ephemeral"}}, + {"type": "text", "text": "hello"}, + ], + }, + {"role": "user", "content": "again"}, + ] + model: Final = "us.anthropic.claude-sonnet-4-6" + + result: Final = ( + await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( + messages=messages, model=model, llm_provider="bedrock", drop_params=drop_params + ) + if use_async + else _bedrock_converse_messages_pt( + messages=messages, model=model, llm_provider="bedrock", drop_params=drop_params + ) + ) + + assert ({"cachePoint": {"type": "default"}} in result[1]["content"]) is not drop_params + + _SENT_REDACTED: Final = [{"reasoningContent": {"redactedContent": _CONVERSE_REDACTED_DATA}}] @@ -813,17 +1079,17 @@ def test_anthropic_messages_pt_drops_unsignable_thinking_block(thinking_block): ) +@REASONING_PLACEMENTS @pytest.mark.parametrize("data", [_NATIVE_REDACTED_DATA, _CONVERSE_REDACTED_DATA], ids=["native", "legacy-wrapped"]) -def test_anthropic_messages_pt_replays_redacted_thinking_as_claude_native_data(data: str): +def test_anthropic_messages_pt_replays_redacted_thinking_as_claude_native_data(data: str, placement: str): """Histories saved from Converse before its replies were decoded hold the data base64-wrapped, and the - Anthropic API, Vertex and Bedrock Invoke 400 on it with "Invalid `data` in `redacted_thinking` block".""" + Anthropic API, Vertex and Bedrock Invoke 400 on it with "Invalid `data` in `redacted_thinking` block". + A redacted block in the content list, the shape the Anthropic SDK replays, used to be dropped.""" + signed: Final = {"type": "thinking", "thinking": "Let me look this up.", "signature": "sig-abc"} + redacted: Final = {"type": "redacted_thinking", "data": data} messages: Final = [ {"role": "user", "content": "hi"}, - { - "role": "assistant", - "content": "hello", - "thinking_blocks": [{"type": "redacted_thinking", "data": data, "cache_control": {"type": "ephemeral"}}], - }, + _assistant_with_reasoning(placement, [signed, redacted], "hello"), {"role": "user", "content": "again"}, ] @@ -832,12 +1098,125 @@ def test_anthropic_messages_pt_replays_redacted_thinking_as_claude_native_data(d assert result[1] == { "role": "assistant", "content": [ - {"type": "redacted_thinking", "data": _NATIVE_REDACTED_DATA, "cache_control": {"type": "ephemeral"}}, + signed, + {"type": "redacted_thinking", "data": _NATIVE_REDACTED_DATA}, {"type": "text", "text": "hello"}, ], } +@SPLIT_REASONING +def test_anthropic_messages_pt_keeps_reasoning_order_across_content_and_thinking_blocks( + in_content: Sequence[int], in_thinking_blocks: Sequence[int] +): + """A thinking block only one of content and thinking_blocks carries goes back in its original place.""" + reasoning: Final = ( + {"type": "thinking", "thinking": "First.", "signature": "sig-1"}, + {"type": "redacted_thinking", "data": _NATIVE_REDACTED_DATA}, + {"type": "thinking", "thinking": "Third.", "signature": "sig-3"}, + ) + messages: Final = [ + {"role": "user", "content": "hi"}, + _assistant_with_split_reasoning(reasoning, in_content, in_thinking_blocks, "hello"), + {"role": "user", "content": "again"}, + ] + + result: Final = anthropic_messages_pt(messages=messages, model="claude-sonnet-4-6", llm_provider="anthropic") + + assert result[1] == {"role": "assistant", "content": [*reasoning, {"type": "text", "text": "hello"}]} + + +@COPY_TEXT +def test_anthropic_messages_pt_sends_a_signed_block_once_when_thinking_blocks_holds_other_text(copy_text: str): + signed: Final = {"type": "thinking", "thinking": "First.", "signature": "sig-1"} + redacted: Final = {"type": "redacted_thinking", "data": _NATIVE_REDACTED_DATA} + messages: Final = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [signed, redacted, {"type": "text", "text": "hello"}], + "thinking_blocks": [{**signed, "thinking": copy_text}, redacted], + }, + {"role": "user", "content": "again"}, + ] + + result: Final = anthropic_messages_pt(messages=messages, model="claude-sonnet-4-6", llm_provider="anthropic") + + assert result[1] == {"role": "assistant", "content": [signed, redacted, {"type": "text", "text": "hello"}]} + + + +def test_anthropic_messages_pt_sends_a_block_once_when_two_assistant_messages_carry_it(): + signed: Final = {"type": "thinking", "thinking": "Let me look this up.", "signature": "sig-abc"} + redacted: Final = {"type": "redacted_thinking", "data": _NATIVE_REDACTED_DATA} + + result: Final = anthropic_messages_pt( + messages=_block_split_across_assistant_messages([signed, redacted]), + model="claude-sonnet-4-6", + llm_provider="anthropic", + ) + + assert result[1] == { + "role": "assistant", + "content": [ + signed, + redacted, + {"type": "tool_use", "id": "tooluse_weather_1", "name": "get_weather", "input": {"city": "Paris"}}, + ], + } + +@REASONING_PLACEMENTS +@pytest.mark.parametrize("drop_params", [False, True]) +def test_anthropic_messages_pt_drops_cache_control_from_thinking_blocks_only_under_drop_params( + placement: str, drop_params: bool +): + """Claude 400s "thinking.cache_control: Extra inputs are not permitted" on a thinking or redacted_thinking + block, so drop_params drops the breakpoint and without it the request reaches Claude as sent. The streamed + list-position `index` is LiteLLM's own field and never reaches Claude.""" + breakpoint_mark: Final = {"cache_control": {"type": "ephemeral"}} + signed: Final = {"type": "thinking", "thinking": "Let me look this up.", "signature": "sig-abc"} + redacted: Final = {"type": "redacted_thinking", "data": _NATIVE_REDACTED_DATA} + messages: Final = [ + {"role": "user", "content": "hi"}, + _assistant_with_reasoning( + placement, [{**signed, **breakpoint_mark, "index": 0}, {**redacted, **breakpoint_mark, "index": 1}], "hello" + ), + {"role": "user", "content": "again"}, + ] + + result: Final = anthropic_messages_pt( + messages=messages, model="claude-sonnet-4-6", llm_provider="anthropic", drop_params=drop_params + ) + + sent_reasoning: Final = ( + [signed, redacted] if drop_params else [{**signed, **breakpoint_mark}, {**redacted, **breakpoint_mark}] + ) + assert result[1] == {"role": "assistant", "content": [*sent_reasoning, {"type": "text", "text": "hello"}]} + + +@REASONING_PLACEMENTS +@pytest.mark.parametrize("drop_params", [False, True]) +@pytest.mark.parametrize("data", ["", " ", None], ids=["empty", "whitespace", "null"]) +def test_anthropic_messages_pt_drops_a_blank_redacted_thinking_block_only_under_drop_params( + data: str | None, placement: str, drop_params: bool +): + """Claude 400s "Invalid `data` in `redacted_thinking` block" on blank data, so drop_params drops the block + and without it the request reaches Claude as sent.""" + blank: Final = {"type": "redacted_thinking", "data": data} + messages: Final = [ + {"role": "user", "content": "hi"}, + _assistant_with_reasoning(placement, [blank], "hello"), + {"role": "user", "content": "again"}, + ] + + result: Final = anthropic_messages_pt( + messages=messages, model="claude-sonnet-4-6", llm_provider="anthropic", drop_params=drop_params + ) + + sent_reasoning: Final = [] if drop_params else [blank] + assert result[1] == {"role": "assistant", "content": [*sent_reasoning, {"type": "text", "text": "hello"}]} + + def test_anthropic_messages_pt_keeps_signed_thinking_block(): """A genuine Anthropic round-trip still holds its original signature, so that thinking block must be forwarded unchanged (we only drop unsignable blocks). diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py index 91d7acd331e..aab57b87a90 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py @@ -308,16 +308,8 @@ EMPTY_ASSISTANT = pytest.mark.parametrize( "content": None, "thinking_blocks": [{"type": "redacted_thinking", "data": encrypted_reasoning_signature("abc")}], }, - { - "role": "assistant", - "content": [{"type": "thinking", "thinking": "unsigned"}], - "thinking_blocks": [{"type": "thinking", "thinking": "hm", "signature": "s"}], - }, - { - "role": "assistant", - "content": [{"type": "redacted_thinking", "data": "x"}], - "thinking_blocks": [{"type": "redacted_thinking", "data": "x"}], - }, + {"role": "assistant", "content": [{"type": "redacted_thinking", "data": ""}]}, + {"role": "assistant", "content": None, "thinking_blocks": [{"type": "redacted_thinking", "data": " "}]}, litellm.Message(role="assistant", content=None), {"role": "assistant", "content": None, "tool_calls": [_TOOL_SEARCH_USE]}, ], @@ -327,8 +319,8 @@ EMPTY_ASSISTANT = pytest.mark.parametrize( "unsigned-thinking-part", "encrypted-thinking-part", "encrypted-redacted-thinking-block", - "unsigned-inline-part-hides-signed-block", - "inline-redacted-part-hides-redacted-block", + "blank-data-redacted-part", + "blank-data-redacted-block", "pydantic-none", "unanswered-tool-search-use", ], @@ -464,6 +456,16 @@ def test_flagged_placement_keeps_a_system_whose_empty_assistant_follower_ends_th "thinking_blocks": [{"type": "thinking", "thinking": "hm", "signature": "s"}], }, {"role": "assistant", "content": None, "thinking_blocks": [{"type": "redacted_thinking", "data": "x"}]}, + { + "role": "assistant", + "content": [{"type": "redacted_thinking", "data": "x"}], + "thinking_blocks": [{"type": "redacted_thinking", "data": "x"}], + }, + { + "role": "assistant", + "content": [{"type": "thinking", "thinking": "unsigned"}], + "thinking_blocks": [{"type": "thinking", "thinking": "hm", "signature": "s"}], + }, {"role": "assistant", "content": [{"type": "thinking", "thinking": "hm", "signature": "s"}]}, {"role": "assistant", "content": [{"type": "thinking", "thinking": "", "signature": "s"}]}, {"role": "assistant", "content": None, "function_call": {"name": "f", "arguments": "{}"}}, @@ -485,6 +487,8 @@ def test_flagged_placement_keeps_a_system_whose_empty_assistant_follower_ends_th "tool-calls", "signed-thinking-block", "redacted-thinking-block", + "redacted-part-beside-the-same-redacted-block", + "signed-block-beside-an-unsigned-part", "signed-thinking-part", "signed-empty-thinking-part", "function-call", From 422252326153b3006ed17abe0c0549bf795c3526 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 20 Oct 2026 08:02:23 -0700 Subject: [PATCH 13/33] fix(logging): report the pay-as-you-go entry billing used on Azure PTU spillover (#47567) get_standard_logging_object_payload decided custom_pricing from litellm_params alone, so on an Azure PTU spillover response model_map_information kept the deployment's zeroed PTU entry while response_cost already used the standard model rate. Reuse the billing side LitellmLogging._custom_pricing_for so both read the same decision Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 8 ++--- .../test_litellm_logging.py | 34 +++++++++++++++++++ 2 files changed, 38 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index e2eed3cf6d5..78ff073fdcb 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1990,7 +1990,7 @@ class Logging(LiteLLMLoggingBaseClass): router_model_id = self.get_router_model_id() ## RESPONSE COST ## - custom_pricing: Final = self._custom_pricing_for(priced_result) + custom_pricing: Final = self.custom_pricing_for(priced_result) prompt: Final = self._prompt_for_cost_calculation() @@ -2165,7 +2165,7 @@ class Logging(LiteLLMLoggingBaseClass): model=model, completion_response=result, custom_llm_provider=custom_llm_provider, - custom_pricing=self._custom_pricing_for(result), + custom_pricing=self.custom_pricing_for(result), base_model=get_base_model_from_metadata(model_call_details=self.model_call_details), router_model_id=router_model_id or self.get_router_model_id(), region_name=_resolve_mantle_region_for_cost( @@ -2190,7 +2190,7 @@ class Logging(LiteLLMLoggingBaseClass): usage=usage, ) - def _custom_pricing_for(self, result: object) -> bool: + def custom_pricing_for(self, result: object) -> bool: litellm_params: Final = getattr(self, "litellm_params", None) result_hidden_params: Final = getattr(result, HIDDEN_PARAMS_ATTR, None) or MappingProxyType({}) additional_headers: Final = ( @@ -7201,7 +7201,7 @@ def get_standard_logging_object_payload( # Fall back to the actual deployment model set by the router in metadata. if base_model is None: base_model = metadata.get("deployment") - custom_pricing: Final = use_custom_pricing_for_model(litellm_params=litellm_params) + custom_pricing: Final = logging_obj.custom_pricing_for(init_response_obj) raw_response_cost: Final = kwargs.get("response_cost") llm_response_cost: Final[float] = raw_response_cost or 0.0 guardrail_cost: Final = guardrail_information_cost(metadata.get("standard_logging_guardrail_information")) diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 9d651b4ae3e..7aac43298d4 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -10095,6 +10095,40 @@ class TestAzurePTUSpilloverCost: finally: self._unregister_models() + def test_spillover_payload_reports_the_pay_as_you_go_entry_billing_used(self, monkeypatch) -> None: + self._register_models() + try: + obj = self._logging_obj(dict(self.PTU_MODEL_INFO), flag="True", litellm_rate=0.0, monkeypatch=monkeypatch) + response = self._response() + response._hidden_params["additional_headers"] = {"llm_provider-x-ms-is-spilled-over": "true"} + response_cost: Final = obj.response_cost_calculator(result=response) + + payload: Final = get_standard_logging_object_payload( + kwargs={ + "model": self.SERVED_MODEL, + "custom_llm_provider": "azure", + "litellm_params": obj.litellm_params, + "response_cost": response_cost, + }, + init_response_obj=response, + start_time=datetime.datetime.now(), + end_time=datetime.datetime.now(), + logging_obj=obj, + status="success", + ) + + assert response_cost == pytest.approx(self.EXPECTED_SPILL_COST) + assert payload is not None + model_map_value: Final = payload["model_map_information"]["model_map_value"] + assert model_map_value is not None + assert model_map_value["input_cost_per_token"] == litellm.model_cost[self.SERVED_MODEL]["input_cost_per_token"] + assert model_map_value["output_cost_per_token"] == litellm.model_cost[self.SERVED_MODEL]["output_cost_per_token"] + assert response_cost == pytest.approx( + 100 * model_map_value["input_cost_per_token"] + 50 * model_map_value["output_cost_per_token"] + ) + finally: + self._unregister_models() + def test_spillover_header_does_not_touch_non_ptu_custom_pricing(self, monkeypatch) -> None: self._register_models() custom_model_id: Final = "non-ptu-custom-router-model-id" From c38b86cc93badcee75d822d4776721fce824e982 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 20 Oct 2026 08:03:31 -0700 Subject: [PATCH 14/33] fix(router): skip non-dict tool entries in the prompt caching affinity key instead of raising (#47565) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router_utils/prompt_caching_cache.py | 6 ++++-- .../router_utils/test_prompt_caching_cache.py | 17 +++++++++++++++++ 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/litellm/router_utils/prompt_caching_cache.py b/litellm/router_utils/prompt_caching_cache.py index 4ec0bc62e7f..f78b5982186 100644 --- a/litellm/router_utils/prompt_caching_cache.py +++ b/litellm/router_utils/prompt_caching_cache.py @@ -94,8 +94,10 @@ def _is_ephemeral_mark(mark: object) -> bool: return mark_map.get("type") == "ephemeral" -def _tool_carries_ephemeral_mark(tool: ChatCompletionToolParam) -> bool: - tool_map: Final = cast( # cast-ok: TypedDict.get rejects keys the type does not declare +def _tool_carries_ephemeral_mark(tool: object) -> bool: + if not isinstance(tool, dict): + return False + tool_map: Final = cast( # cast-ok: tools are plain JSON dicts Mapping[str, object], tool ) if _is_ephemeral_mark(tool_map.get("cache_control")): diff --git a/tests/unit/router_utils/test_prompt_caching_cache.py b/tests/unit/router_utils/test_prompt_caching_cache.py index 46af2454dea..ce8f04b8c54 100644 --- a/tests/unit/router_utils/test_prompt_caching_cache.py +++ b/tests/unit/router_utils/test_prompt_caching_cache.py @@ -608,6 +608,23 @@ def test_unmarked_tools_change_the_key_marked_tools_do_not_share(): assert key_a != key_b +def test_a_non_dict_tool_entry_still_derives_a_key_instead_of_raising(): + """ + A client can send `tools: ["not-a-tool"]`. The provider rejects that with its own 400, so the + routing check has to key on what it was given and let the request through, never raise a 500. + """ + messages = cast( + list[AllMessageValues], + [{"role": "user", "content": [{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}}]}], + ) + tools: Final = cast(list[ChatCompletionToolParam], ["not-a-tool", _marked_tool("read", {"type": "ephemeral"})]) + + positions = PromptCachingCache.prefix_positions(messages, tools) + + assert [position.position for position in positions] == [0, 1] + assert PromptCachingCache.get_prompt_caching_cache_key(messages, tools) == positions[-1].cache_key + + def test_empty_and_missing_tools_seed_the_same_key(): """`tools=[]` and no tools at all are the same request to the provider, so they must key alike.""" messages = cast( From 04e9ed4c52505b5a50840843bc8b367aa32d09eb Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 20 Oct 2026 08:04:55 -0700 Subject: [PATCH 15/33] fix(litellm_proxy): forward tool_config cache_control_injection_points to the downstream proxy (#47570) PR #47394 withheld cache_control_injection_points from the OpenAI SDK extra_body for every provider on that path, which also covers litellm_proxy. The downstream LiteLLM proxy applies the point itself, so litellm_proxy keeps forwarding it while OpenAI and the other OpenAI-compatible providers still never receive it Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/utils.py | 3 ++- tests/unit/test_utils.py | 15 +++++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/litellm/utils.py b/litellm/utils.py index 14ee35e571e..a6108a15bb3 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5552,8 +5552,9 @@ def add_provider_specific_params_to_optional_params( # for openai, azure we should pass the extra/passed params within `extra_body` https://github.com/openai/openai-python/blob/ac33853ba10d13ac149b1fa3ca6dba7d613065c9/src/openai/resources/models.py#L46 if _should_drop_param(k="extra_body", additional_drop_params=additional_drop_params) is False: extra_body: Final = dict(passed_params.pop("extra_body", None) or {}) + withheld_params: Final = frozenset[str]() if custom_llm_provider == "litellm_proxy" else LITELLM_ONLY_PARAMS for k in passed_params: - if k not in openai_params and k not in LITELLM_ONLY_PARAMS and passed_params[k] is not None: + if k not in openai_params and k not in withheld_params and passed_params[k] is not None: extra_body[k] = passed_params[k] if not isinstance(optional_params.get("extra_body"), dict): optional_params["extra_body"] = {} diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 44939ca7585..0be3ae05e97 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -10703,3 +10703,18 @@ def test_get_optional_params_keeps_a_tool_config_cache_point_out_of_the_openai_e assert optional_params["extra_body"] == {"kept": True} assert "cache_control_injection_points" not in optional_params + + +def test_get_optional_params_forwards_a_tool_config_cache_point_to_a_downstream_litellm_proxy(): + """A litellm_proxy deployment calls the downstream proxy through the OpenAI SDK, and that proxy + applies cache_control_injection_points itself, so the point must ride along in extra_body.""" + point: Final = {"location": "tool_config", "control": {"type": "ephemeral"}} + optional_params: Final[dict[str, object]] = litellm.utils.get_optional_params( + model="claude-sonnet-4-5", + custom_llm_provider="litellm_proxy", + tools=[{"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}}], + cache_control_injection_points=[point], + extra_body={"kept": True}, + ) + + assert optional_params["extra_body"] == {"kept": True, "cache_control_injection_points": [point]} From 91724bb0a7c88e54621f149b5cdb5f133f58dd12 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 15:13:59 +0000 Subject: [PATCH 16/33] fix(responses)!: count input tokens the way inference sends the request (#47363) * fix(count-tokens): translate chat tool use before Claude token counts Responses and chat tool use (function tools, function_call items, tool messages) is now converted the way Claude chat inference converts it before the Anthropic, Azure AI and Vertex partner counters call count_tokens. Translation errors answer 400 before any provider call. Anthropic-shaped requests and /v1/messages/count_tokens keep their body as sent. * test(count-tokens): compare Claude count bodies with inference for Responses tool use * fix(responses): forward the caller's Responses body to OpenAI token counts OpenAI deployments now get the caller's /v1/responses/input_tokens body as sent, minus LiteLLM-only fields, so OpenAI's own count and errors come back. previous_response_id is checked for ownership like inference and routes to the deployment that served it * fix(responses)!: count Claude Responses bodies the way inference sends them Non-native Responses counts now run the body through the inference bridge, so text.format, tool_choice, reasoning, built-in tools, and file_search count the way inference sends them on Anthropic, Azure AI, Vertex, and Bedrock. The OpenAI native count body uses the drop_params-resolved tool fields. * fix(responses): honor drop_params and replay previous turns in Claude counts Claude input_tokens counts now read the same drop_params setting inference merges from the deployment, router defaults, and request, so a forced tool_choice downgrades to auto under drop_params. previous_response_id replays the stored turns through inference's session loader, and chat messages carrying fields beyond role and content go through the provider's chat translation * fix(responses): keep proxy-only request keys out of OpenAI counts and map chat tools in Bedrock counts * fix(count-tokens): import on Python 3.11 and keep stored prompts in OpenAI counts under drop_params * fix(responses): leave tool_choice out of counts when drop_params drops every tool * fix(count-tokens): count Mantle options, drop cache marks, and reject input inference cannot send The Bedrock Mantle fallback now counts thinking, tool_choice and output_config like the runtime count. Claude Responses counts leave out cache_control marks again, so more than four marks count instead of failing. Counts answer 400 before any provider call for user blocks chat inference refuses and for Responses tools with no Chat Completions equivalent, instead of dropping the block or counting a placeholder tool. Adds the wire tests for these paths * fix(count-tokens): count Responses tools the way the inference bridge sends them A Claude count now hoists additional_tools input items into tools, strips a malformed tool flag under drop_params and rejects it without, and leaves tools out when the deployment lists tools in additional_drop_params, so the count matches the body Claude inference sends for the same request * fix(count-tokens): keep Bedrock runtime counts option-free and send OpenAI the provider file id * test(count-tokens): capture the Anthropic count request at the HTTP layer instead of patching the handler * fix(count-tokens): ignore a caller-sent file id mapping and keep non-string file ids as sent --- .../prompt_templates/common_utils.py | 6 +- .../anthropic/count_tokens/token_counter.py | 34 ++ .../anthropic/count_tokens/transformation.py | 225 ++++++++++- .../anthropic/count_tokens/handler.py | 3 + .../anthropic/count_tokens/token_counter.py | 34 ++ litellm/llms/base_llm/base_utils.py | 2 + litellm/llms/base_llm/files/batch_records.py | 2 + .../count_tokens/bedrock_token_counter.py | 62 ++- .../bedrock/count_tokens/mantle_handler.py | 7 +- .../bedrock/count_tokens/transformation.py | 9 +- litellm/llms/gemini/common_utils.py | 2 + .../openai/responses/count_tokens/handler.py | 16 +- .../responses/count_tokens/token_counter.py | 163 ++++---- .../responses/count_tokens/transformation.py | 36 ++ litellm/llms/vertex_ai/common_utils.py | 66 ++- .../vertex_ai_partner_models/main.py | 4 +- litellm/main.py | 1 + litellm/proxy/proxy_server.py | 141 ++++++- .../proxy/response_api_endpoints/endpoints.py | 93 ++++- .../responses/file_search/emulated_handler.py | 6 +- .../count_request.py | 207 ++++++++++ .../transformation.py | 2 +- litellm/responses/main.py | 4 +- .../providers/_count_tokens_system_lift.py | 81 ++++ ...anthropic_count_tokens_system_lift_wire.py | 176 +++++++- .../test_bedrock_mantle_count_tokens_wire.py | 107 +++++ .../providers/test_malformed_tools_wire.py | 11 +- .../test_openai_count_tokens_tools_wire.py | 74 +++- .../providers/test_responses_core_wire.py | 37 ++ .../test_vertex_gemini_count_tokens_wire.py | 54 +++ .../test_vertex_partner_count_tokens_wire.py | 86 +++- .../test_vertex_partner_count_tokens_sdk.py | 76 +++- ...ore_utils_prompt_templates_common_utils.py | 24 ++ ...t_anthropic_count_tokens_transformation.py | 306 ++++++++++++++ ...est_bedrock_count_tokens_transformation.py | 19 + .../test_bedrock_token_counter.py | 125 ++++++ .../count_tokens/test_token_counter.py | 121 ++++++ tests/unit/llms/test_file_search_responses.py | 16 +- .../vertex_ai/test_vertex_ai_common_utils.py | 149 ++++++- .../response_api_endpoints/test_endpoints.py | 149 ++++++- tests/unit/proxy/test_proxy_token_counter.py | 378 +++++++++++++++++- .../test_count_request.py | 127 ++++++ tests/unit/test_count_tokens_public_api.py | 46 +++ 43 files changed, 3127 insertions(+), 160 deletions(-) create mode 100644 litellm/responses/litellm_completion_transformation/count_request.py create mode 100644 tests/unit/responses/litellm_completion_transformation/test_count_request.py diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index b89e4d609ad..60be4d44706 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -624,7 +624,7 @@ def update_messages_with_model_file_ids( def update_responses_input_with_model_file_ids( input: object, model_id: str | None = None, - model_file_id_mapping: dict[str, dict[str, str]] | None = None, + model_file_id_mapping: Mapping[str, Mapping[str, str]] | None = None, ) -> object: """ Updates responses API input with provider-specific file IDs. @@ -665,7 +665,7 @@ def update_responses_input_with_model_file_ids( for content_item in content: if isinstance(content_item, dict) and content_item.get("type") == "input_file": file_id = content_item.get("file_id") - if file_id: + if isinstance(file_id, str) and file_id: provider_file_id = file_id # Default to original # Check if we have a mapping for this file ID @@ -761,7 +761,7 @@ def _decode_vector_store_ids_in_tools( def update_responses_tools_with_model_file_ids( tools: list[dict[str, object]] | None, model_id: str | None = None, - model_file_id_mapping: dict[str, dict[str, str]] | None = None, + model_file_id_mapping: Mapping[str, Mapping[str, str]] | None = None, ) -> list[dict[str, object]] | None: """ Updates responses API tools with provider-specific file IDs. diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index ecf050d185e..0764fb41941 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -2,6 +2,8 @@ Anthropic Token Counter implementation using the CountTokens API. """ +from collections.abc import Mapping +from types import MappingProxyType from typing import Any, Final from pydantic import JsonValue @@ -10,6 +12,11 @@ import litellm from litellm._logging import verbose_logger from litellm.exceptions import AuthenticationError from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler +from litellm.llms.anthropic.count_tokens.transformation import ( + AnthropicCountFields, + ChatCountTranslationError, + anthropic_count_fields, +) from litellm.llms.base_llm.base_utils import BaseTokenCounter from litellm.types.utils import CountContextManagement, LlmProviders, TokenCountMessagesFormat, TokenCountResponse @@ -38,6 +45,8 @@ class AnthropicTokenCounter(BaseTokenCounter): messages_format: TokenCountMessagesFormat = "anthropic", output_format: JsonValue = None, context_management: CountContextManagement | None = None, + chat_params: Mapping[str, JsonValue] = MappingProxyType({}), + count_options: Mapping[str, JsonValue] = MappingProxyType({}), ) -> TokenCountResponse | None: """ Count tokens using Anthropic's CountTokens API. @@ -57,6 +66,30 @@ class AnthropicTokenCounter(BaseTokenCounter): if not messages: return None + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + match translated := await anthropic_count_fields( + model_to_use, messages, system, tools, messages_format, AnthropicConfig(), chat_params + ): + case ChatCountTranslationError(): + return translated.response(request_model, model_to_use, "anthropic_api") + case AnthropicCountFields(): + return await self.count_tokens( + model_to_use=model_to_use, + messages=[dict(message) for message in translated.messages], + contents=contents, + deployment=deployment, + request_model=request_model, + tools=[dict(tool) for tool in translated.tools] if translated.tools is not None else None, + system=translated.system, + messages_format="anthropic", + output_format=output_format, + context_management=context_management, + count_options=translated.options, + ) + case None: + pass + deployment = deployment or {} litellm_params: Final = deployment.get("litellm_params", {}) api_base: Final = litellm_params.get("api_base") @@ -89,6 +122,7 @@ class AnthropicTokenCounter(BaseTokenCounter): api_base=api_base, tools=tools, system=system, + optional_params=count_options, ) if result is not None: diff --git a/litellm/llms/anthropic/count_tokens/transformation.py b/litellm/llms/anthropic/count_tokens/transformation.py index 8e2688d4360..592b5fb0f00 100644 --- a/litellm/llms/anthropic/count_tokens/transformation.py +++ b/litellm/llms/anthropic/count_tokens/transformation.py @@ -4,25 +4,37 @@ Anthropic CountTokens API transformation logic. This module handles the transformation of requests to Anthropic's CountTokens API format. """ +import copy from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from itertools import chain from types import MappingProxyType -from typing import Final +from typing import ( + Final, + Protocol, + TypeAlias, + cast, # noqa: TID251 # caller chat messages go to chat inference translation as-is +) -from pydantic import JsonValue, TypeAdapter +from pydantic import ConfigDict, JsonValue, TypeAdapter, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import ( anthropic_system_blocks, split_leading_system_run, ) from litellm.llms.anthropic.common_utils import merge_anthropic_beta_headers from litellm.llms.anthropic.wif import resolve_anthropic_base -from litellm.types.llms.openai import ChatCompletionImageObject +from litellm.types.llms.openai import AllMessageValues, ChatCompletionImageObject, ValidUserMessageContentTypes +from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse _COUNT_REQUEST: Final = TypeAdapter(dict[str, JsonValue]) _SYSTEM_BLOCKS: Final = TypeAdapter(list[JsonValue]) _IMAGE_BLOCK: Final = TypeAdapter(ChatCompletionImageObject) COUNT_TOKEN_OPTION_NAMES: Final = ("thinking", "tool_choice", "output_config") +_PLAIN_MESSAGE_KEYS: Final = frozenset({"role", "content"}) def _count_image(block: JsonValue) -> JsonValue: @@ -74,6 +86,213 @@ def _lift_leading_system( return conversation, lifted or system +@dataclass(frozen=True, slots=True) +class AnthropicCountFields: + messages: tuple[Mapping[str, JsonValue], ...] + system: JsonValue + tools: tuple[Mapping[str, JsonValue], ...] | None + options: Mapping[str, JsonValue] + + +@dataclass(frozen=True, slots=True) +class ChatCountTranslationError: + message: str + + def response(self, request_model: str, model_used: str, tokenizer_type: str) -> TokenCountResponse: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_used, + tokenizer_type=tokenizer_type, + error=True, + error_message=self.message, + status_code=400, + rejected_before_provider_call=True, + ) + + +@with_config(ConfigDict(extra="ignore")) +class _AnthropicChatRequest(TypedDict): + messages: ReadOnly[tuple[Mapping[str, JsonValue], ...]] + system: NotRequired[ReadOnly[JsonValue]] + tools: NotRequired[ReadOnly[tuple[Mapping[str, JsonValue], ...]]] + thinking: NotRequired[ReadOnly[JsonValue]] + tool_choice: NotRequired[ReadOnly[JsonValue]] + output_config: NotRequired[ReadOnly[JsonValue]] + output_format: NotRequired[ReadOnly[JsonValue]] + + +_ANTHROPIC_CHAT_REQUEST: Final = TypeAdapter(_AnthropicChatRequest) +_ChatParams: TypeAlias = dict[str, object] +_ChatMessages: TypeAlias = list[AllMessageValues] + + +class _ChatRequestTranslator(Protocol): + def map_openai_params( + self, non_default_params: _ChatParams, optional_params: _ChatParams, model: str, drop_params: bool + ) -> _ChatParams: ... + + def transform_request( + self, + model: str, + messages: _ChatMessages, + optional_params: _ChatParams, + litellm_params: _ChatParams, + headers: _ChatParams, + ) -> _ChatParams: ... + + +def _needs_chat_translation(message: Mapping[str, object]) -> bool: + return message.get("role") == "tool" or not _PLAIN_MESSAGE_KEYS.issuperset(message.keys()) + + +def _is_chat_tool(tool: Mapping[str, object]) -> bool: + from litellm.litellm_core_utils.prompt_templates.factory import is_flat_responses_tool + + return is_flat_responses_tool(tool) or (tool.get("type") == "function" and "function" in tool) + + +def _has_non_chat_user_block(message: Mapping[str, object]) -> bool: + content: Final = message.get("content") + if message.get("role") != "user" or not isinstance(content, list): + return False + blocks: Final = _SYSTEM_BLOCKS.validate_python(content) + return any(isinstance(block, dict) and block.get("type") not in ValidUserMessageContentTypes for block in blocks) + + +def _gets_placeholder_name(tool: Mapping[str, object]) -> bool: + from litellm.llms.anthropic.pass_through.adapters.transformation import is_provider_native_tool_dict + + name: Final = tool.get("name") + return not (_is_chat_tool(tool) or is_provider_native_tool_dict(tool) or (isinstance(name, str) and name.strip())) + + +def _untranslatable_input( + messages: Sequence[Mapping[str, object]], tools: Sequence[Mapping[str, object]] +) -> ChatCountTranslationError | None: + invalid_index: Final = next( + (index for index, message in enumerate(messages) if _has_non_chat_user_block(message)), None + ) + if invalid_index is not None: + return ChatCountTranslationError( + message=f"Invalid user message at index {invalid_index}. " + "Please ensure all user messages are valid OpenAI chat completion messages." + ) + unnamed_tool: Final = next(filter(_gets_placeholder_name, tools), None) + if unnamed_tool is not None: + return ChatCountTranslationError( + message=f"Invalid request for token counting: tool type {unnamed_tool.get('type')!r} " + "has no Chat Completions equivalent" + ) + return None + + +def _system_messages(system: object) -> tuple[Mapping[str, object], ...]: + message: Final[Mapping[str, object]] = {"role": "system", "content": system} + return (message,) if isinstance(system, (str, list)) and system else () + + +def _count_options(request: _AnthropicChatRequest) -> Mapping[str, JsonValue]: + """Structured output reaches chat inference as the beta ``output_format``, which count_tokens + refuses; its GA ``output_config.format`` form counts the same schema the same way.""" + output_config: Final = request.get("output_config") + output_format: Final = request.get("output_format") + counted_config: Final = { + **(output_config if isinstance(output_config, dict) else {}), + **({} if output_format is None else {"format": output_format}), + } + return MappingProxyType( + { + key: value + for key, value in ( + ("thinking", request.get("thinking")), + ("tool_choice", request.get("tool_choice")), + ("output_config", counted_config or None), + ) + if value is not None + } + ) + + +def _translated_chat_request( + model: str, + messages: Sequence[Mapping[str, object]], + system: object, + tools: Sequence[Mapping[str, object]], + chat_params: Mapping[str, object], + translator: _ChatRequestTranslator, +) -> _ChatParams: + from litellm.litellm_core_utils.prompt_templates.factory import chat_completion_tools_for + + chat_tools: Final = [dict(chat_tool) for chat_tool in chain.from_iterable(map(chat_completion_tools_for, tools))] + openai_params: Final = {name: value for name, value in chat_params.items() if name != "drop_params"} + mapped_params: Final = {**openai_params, **({"tools": chat_tools} if chat_tools else {})} + chat_messages: Final = cast( # cast-ok: chat-shaped request messages, read the way chat inference reads them + _ChatMessages, copy.deepcopy([*_system_messages(system), *messages]) + ) + return translator.transform_request( + model=model, + messages=chat_messages, + optional_params=translator.map_openai_params( + mapped_params, {}, model, drop_params=chat_params.get("drop_params") is True + ) + if mapped_params + else {}, + litellm_params={}, + headers={}, + ) + + +def anthropic_fields_for_chat_tool_use( + model: str, + messages: Sequence[Mapping[str, object]], + system: object, + tools: Sequence[Mapping[str, object]] | None, + translator: _ChatRequestTranslator, + chat_params: Mapping[str, object] = MappingProxyType({}), +) -> AnthropicCountFields | ChatCountTranslationError | None: + counted_tools: Final = tools or () + if ( + not chat_params + and not any(map(_is_chat_tool, counted_tools)) + and not any(map(_needs_chat_translation, messages)) + ): + return None + untranslatable: Final = _untranslatable_input(messages, counted_tools) + if untranslatable is not None: + return untranslatable + from litellm.exceptions import BadRequestError + from litellm.llms.anthropic.common_utils import AnthropicError + + try: + request: Final = _translated_chat_request(model, messages, system, counted_tools, chat_params, translator) + except (AnthropicError, BadRequestError) as e: + return ChatCountTranslationError(message=str(e.message)) + except (ValueError, KeyError) as e: + return ChatCountTranslationError(message=f"Invalid request for token counting: {e}") + translated: Final = _ANTHROPIC_CHAT_REQUEST.validate_python(request) + return AnthropicCountFields( + messages=translated["messages"], + system=translated.get("system"), + tools=translated.get("tools"), + options=_count_options(translated), + ) + + +async def anthropic_count_fields( + model: str, + messages: Sequence[Mapping[str, object]], + system: object, + tools: Sequence[Mapping[str, object]] | None, + messages_format: TokenCountMessagesFormat, + translator: _ChatRequestTranslator, + chat_params: Mapping[str, object] = MappingProxyType({}), +) -> AnthropicCountFields | ChatCountTranslationError | None: + if messages_format != "openai": + return None + return await asyncify(anthropic_fields_for_chat_tool_use)(model, messages, system, tools, translator, chat_params) + + class AnthropicCountTokensConfig: """ Configuration and transformation logic for Anthropic CountTokens API. diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index 53734a77f3c..0f8afaae681 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -4,6 +4,7 @@ Azure AI Anthropic CountTokens API handler. Uses httpx for HTTP requests with Azure authentication. """ +from collections.abc import Mapping from typing import Final import httpx @@ -36,6 +37,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): timeout: float | httpx.Timeout | None = None, tools: list[dict[str, JsonValue]] | None = None, system: object = None, + optional_params: Mapping[str, JsonValue] | None = None, ) -> dict[str, object]: """ Handle a CountTokens request using httpx with Azure authentication. @@ -66,6 +68,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): messages=messages, tools=tools, system=system, + optional_params=optional_params, ) verbose_logger.debug("Transformed request: %s", request_body) diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py index d1f45751d3e..af7cf9d81ce 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -3,12 +3,19 @@ Azure AI Anthropic Token Counter implementation using the CountTokens API. """ import os +from collections.abc import Mapping +from types import MappingProxyType from typing import Any, Final from pydantic import JsonValue import litellm from litellm._logging import verbose_logger +from litellm.llms.anthropic.count_tokens.transformation import ( + AnthropicCountFields, + ChatCountTranslationError, + anthropic_count_fields, +) from litellm.llms.azure_ai.anthropic.count_tokens.handler import ( AzureAIAnthropicCountTokensHandler, ) @@ -55,6 +62,8 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): messages_format: TokenCountMessagesFormat = "anthropic", output_format: JsonValue = None, context_management: CountContextManagement | None = None, + chat_params: Mapping[str, JsonValue] = MappingProxyType({}), + count_options: Mapping[str, JsonValue] = MappingProxyType({}), ) -> TokenCountResponse | None: """ Count tokens using Azure AI Anthropic's CountTokens API. @@ -74,6 +83,30 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): if not messages: return None + from litellm.llms.azure_ai.anthropic.transformation import AzureAnthropicConfig + + match translated := await anthropic_count_fields( + model_to_use, messages, system, tools, messages_format, AzureAnthropicConfig(), chat_params + ): + case ChatCountTranslationError(): + return translated.response(request_model, model_to_use, "azure_ai_anthropic_api") + case AnthropicCountFields(): + return await self.count_tokens( + model_to_use=model_to_use, + messages=[dict(message) for message in translated.messages], + contents=contents, + deployment=deployment, + request_model=request_model, + tools=[dict(tool) for tool in translated.tools] if translated.tools is not None else None, + system=translated.system, + messages_format="anthropic", + output_format=output_format, + context_management=context_management, + count_options=translated.options, + ) + case None: + pass + deployment = deployment or {} litellm_params: Final = deployment.get("litellm_params", {}) @@ -101,6 +134,7 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): litellm_params=litellm_params, tools=tools, system=system, + optional_params=count_options, ) if result is not None: diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py index 640930a780d..a9ec1f97d20 100644 --- a/litellm/llms/base_llm/base_utils.py +++ b/litellm/llms/base_llm/base_utils.py @@ -9,6 +9,7 @@ import re from abc import ABC, abstractmethod from collections.abc import Iterator, Mapping, Sequence from dataclasses import dataclass +from types import MappingProxyType from typing import Any, Final from openai.lib import _parsing, _pydantic # pyright: ignore[reportPrivateUsage] # SDK parser internals @@ -40,6 +41,7 @@ class BaseTokenCounter(ABC): messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, context_management: CountContextManagement | None = None, + chat_params: Mapping[str, JsonValue] = MappingProxyType({}), ) -> TokenCountResponse | None: pass diff --git a/litellm/llms/base_llm/files/batch_records.py b/litellm/llms/base_llm/files/batch_records.py index 6bb98456e69..070cf73cd57 100644 --- a/litellm/llms/base_llm/files/batch_records.py +++ b/litellm/llms/base_llm/files/batch_records.py @@ -18,6 +18,7 @@ def _responses_request_keys() -> frozenset[str]: def responses_batch_body_to_chat_body( openai_request_body: Mapping[str, object], custom_llm_provider: str | None = None, + drop_params: bool | None = None, ) -> dict[str, object]: # mutable-ok: provider transforms take the bridged chat body as a plain dict """ Rewrite the body of an OpenAI `/v1/responses` batch record as a Chat Completions body. @@ -50,5 +51,6 @@ def responses_batch_body_to_chat_body( input=chat_input, responses_api_request=responses_request, custom_llm_provider=custom_llm_provider, + drop_params=drop_params, metadata=openai_request_body.get("metadata"), ) diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 068cf4f2e67..21fbdb04896 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -4,12 +4,15 @@ Bedrock Token Counter implementation using the CountTokens API. from collections.abc import Mapping, Sequence from dataclasses import dataclass +from types import MappingProxyType from typing import Any, Final from pydantic import JsonValue from typing_extensions import assert_never from litellm._logging import verbose_logger +from litellm.llms.anthropic.chat.transformation import AnthropicConfig +from litellm.llms.anthropic.count_tokens.transformation import ChatCountTranslationError, anthropic_count_fields from litellm.llms.base_llm.base_utils import BaseTokenCounter from litellm.llms.bedrock.common_utils import BedrockError, get_bedrock_base_model from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler @@ -58,6 +61,17 @@ def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model: ) +def _request_data( + model: str, messages: Sequence[Mapping[str, object]], tools: Sequence[Mapping[str, object]] | None, system: object +) -> Mapping[str, object]: + return { + "model": model, + "messages": [dict(message) for message in messages], + **({"tools": [dict(tool) for tool in tools]} if tools else {}), + **({"system": system} if system else {}), + } + + class BedrockTokenCounter(BaseTokenCounter): """Counts on the configured Bedrock endpoint, with Mantle retry when Runtime rejects a Claude model""" @@ -81,13 +95,13 @@ class BedrockTokenCounter(BaseTokenCounter): self, handler: BedrockCountTokensHandler | BedrockMantleCountTokensHandler, tokenizer_type: str, - request_data: dict[str, object], + request_data: Mapping[str, object], litellm_params: dict[str, object], resolved_model: str, ) -> CountTokensOutcome: try: result: Final = await handler.handle_count_tokens_request( - request_data=request_data, + request_data=dict(request_data), litellm_params=litellm_params, resolved_model=resolved_model, client=self._client, @@ -115,30 +129,44 @@ class BedrockTokenCounter(BaseTokenCounter): messages_format: TokenCountMessagesFormat = "anthropic", output_format: JsonValue = None, context_management: CountContextManagement | None = None, + chat_params: Mapping[str, JsonValue] = MappingProxyType({}), ) -> TokenCountResponse | None: if not messages: return None - litellm_params: Final = (deployment or {}).get("litellm_params", {}) - request_data: Final[dict[str, object]] = { - "model": model_to_use, - "messages": messages, - **({"tools": tools} if tools else {}), - **({"system": system} if system else {}), - } use_mantle: Final = model_to_use.startswith("bedrock_mantle/") - resolved_model: Final = get_bedrock_base_model(model_to_use.removeprefix("bedrock_mantle/")) + counted_model: Final = model_to_use.removeprefix("bedrock_mantle/") + resolved_model: Final = get_bedrock_base_model(counted_model) + primary_tokenizer: Final = MANTLE_TOKENIZER_TYPE if use_mantle else RUNTIME_TOKENIZER_TYPE + translated: Final = ( + await anthropic_count_fields( + counted_model, messages, system, tools, messages_format, AnthropicConfig(), chat_params + ) + if is_mantle_claude_model(resolved_model) + else None + ) + if isinstance(translated, ChatCountTranslationError): + return translated.response(request_model, model_to_use, primary_tokenizer) + litellm_params: Final = (deployment or {}).get("litellm_params", {}) + runtime_request: Final = ( + _request_data(model_to_use, messages, tools, system) + if translated is None + else _request_data(model_to_use, translated.messages, translated.tools, translated.system) + ) + mantle_request: Final = runtime_request if translated is None else {**runtime_request, **translated.options} - primary: Final = await self._count_with( - self._mantle_handler if use_mantle else self._runtime_handler, - MANTLE_TOKENIZER_TYPE if use_mantle else RUNTIME_TOKENIZER_TYPE, - request_data, - litellm_params, - resolved_model, + primary: Final = ( + await self._count_with( + self._mantle_handler, primary_tokenizer, mantle_request, litellm_params, resolved_model + ) + if use_mantle + else await self._count_with( + self._runtime_handler, primary_tokenizer, runtime_request, litellm_params, resolved_model + ) ) outcome: Final = ( await self._count_with( - self._mantle_handler, MANTLE_TOKENIZER_TYPE, request_data, litellm_params, resolved_model + self._mantle_handler, MANTLE_TOKENIZER_TYPE, mantle_request, litellm_params, resolved_model ) if not use_mantle and _runtime_rejected_claude_model(primary, resolved_model) else primary diff --git a/litellm/llms/bedrock/count_tokens/mantle_handler.py b/litellm/llms/bedrock/count_tokens/mantle_handler.py index 4ea4440a3c4..4b1a6a75a7a 100644 --- a/litellm/llms/bedrock/count_tokens/mantle_handler.py +++ b/litellm/llms/bedrock/count_tokens/mantle_handler.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Final import httpx @@ -6,7 +7,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger -from litellm.llms.anthropic.count_tokens.transformation import AnthropicCountTokensConfig +from litellm.llms.anthropic.count_tokens.transformation import COUNT_TOKEN_OPTION_NAMES, AnthropicCountTokensConfig from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing from litellm.llms.bedrock.common_utils import BedrockError, build_mantle_messages_url from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client @@ -23,6 +24,7 @@ class MantleCountTokensRequest(TypedDict): _COUNT_REQUEST: Final = TypeAdapter(MantleCountTokensRequest) _COUNT_RESPONSE: Final = TypeAdapter(dict[str, JsonValue]) +_COUNT_OPTIONS: Final[TypeAdapter[Mapping[str, JsonValue]]] = TypeAdapter(Mapping[str, JsonValue]) class BedrockMantleCountTokensHandler(BaseAWSLLM): @@ -61,6 +63,9 @@ class BedrockMantleCountTokensHandler(BaseAWSLLM): messages=request["messages"], tools=request.get("tools"), system=request.get("system"), + optional_params=_COUNT_OPTIONS.validate_python( + {name: value for name, value in request_data.items() if name in COUNT_TOKEN_OPTION_NAMES} + ), ) verbose_logger.debug("Making bedrock-mantle count_tokens request to: %s", endpoint_url) api_key: Final = litellm_params.get("api_key") diff --git a/litellm/llms/bedrock/count_tokens/transformation.py b/litellm/llms/bedrock/count_tokens/transformation.py index 48fc41ed12b..7cca4f69acc 100644 --- a/litellm/llms/bedrock/count_tokens/transformation.py +++ b/litellm/llms/bedrock/count_tokens/transformation.py @@ -52,12 +52,15 @@ def _sanitized_bedrock_tool_name(raw_name: JsonValue) -> str: def _to_bedrock_tool_spec(tool: JsonValue) -> dict[str, JsonValue]: fields: Final = _json_dict(tool) - name: Final = _sanitized_bedrock_tool_name(fields.get("name", "")) + chat_function: Final = _json_dict(fields.get("function")) + spec: Final = chat_function or fields + name: Final = _sanitized_bedrock_tool_name(spec.get("name", "")) + schema_field: Final = "parameters" if chat_function else "input_schema" return { "toolSpec": { "name": name, - "description": fields.get("description") or name, - "inputSchema": {"json": fields.get("input_schema", {"type": "object", "properties": {}})}, + "description": spec.get("description") or name, + "inputSchema": {"json": spec.get(schema_field, {"type": "object", "properties": {}})}, } } diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index c5ebd84e00f..b1a726be876 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -3,6 +3,7 @@ import datetime import json import math from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import Any, Final import httpx @@ -498,6 +499,7 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, context_management: CountContextManagement | None = None, + chat_params: Mapping[str, JsonValue] = MappingProxyType({}), ) -> TokenCountResponse | None: import copy diff --git a/litellm/llms/openai/responses/count_tokens/handler.py b/litellm/llms/openai/responses/count_tokens/handler.py index 6f911e9c40f..739966c95d1 100644 --- a/litellm/llms/openai/responses/count_tokens/handler.py +++ b/litellm/llms/openai/responses/count_tokens/handler.py @@ -34,7 +34,7 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): tools: Sequence[Mapping[str, object]] | None = None, instructions: str | None = None, text: Mapping[str, object] | None = None, - ) -> dict[str, object]: + ) -> Mapping[str, object]: """ Handle a token counting request to OpenAI's Responses API. @@ -56,7 +56,19 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): instructions=instructions, text=text, ) + except ValueError as e: + verbose_logger.error("Error in CountTokens handler: %s", e) + raise OpenAIError(status_code=500, message=f"CountTokens processing error: {e}") + return await self.post_count_request(request_body, api_key, api_base, timeout) + async def post_count_request( + self, + request_body: Mapping[str, object], + api_key: str, + api_base: str | None = None, + timeout: float | httpx.Timeout | None = None, + ) -> Mapping[str, object]: + try: endpoint_url: Final = self.get_openai_count_tokens_endpoint(api_base) verbose_logger.debug("Making request to: %s", endpoint_url) @@ -70,7 +82,7 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): response: Final = await async_client.post( endpoint_url, headers=headers, - json=request_body, + json=dict(request_body), timeout=request_timeout, ) diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index b12fdb95765..437f19a4dc6 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -2,7 +2,9 @@ OpenAI Token Counter implementation using the Responses API /input_tokens endpoint. """ -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType from typing import Any, Final import openai @@ -11,6 +13,8 @@ from typing_extensions import assert_never from litellm._logging import verbose_logger from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.tool_fields import DropPolicy from litellm.llms.anthropic.pass_through.convertible_blocks import deployment_drops_params, unconvertible_blocks_error from litellm.llms.anthropic.pass_through.responses_adapters.transformation import ( LiteLLMAnthropicToResponsesAPIAdapter, @@ -23,6 +27,7 @@ from litellm.llms.openai.responses.count_tokens.handler import ( from litellm.llms.openai.responses.count_tokens.transformation import ( OpenAICountTokensConfig, ResponsesCountRequestParts, + native_count_request, ) from litellm.llms.openai.workload_identity import ( OpenAIWorkloadIdentityConfig, @@ -50,6 +55,46 @@ def _workload_identity_config( return e +@dataclass(frozen=True, slots=True) +class _CountCredential: + api_base: str | None + credential: str | OpenAIWorkloadIdentityConfig + + +def _count_credential(deployment: Mapping[str, object] | None, model_label: str) -> _CountCredential | OpenAIError: + litellm_params: Final = as_str_mapping(None if deployment is None else deployment.get("litellm_params")) or {} + api_base: Final = litellm_params.get("api_base") + api_key: Final = litellm_params.get("api_key") + credentials: Final = get_openai_credentials( + api_base=api_base if isinstance(api_base, str) else None, + api_key=api_key if isinstance(api_key, str) else None, + ) + credential: Final = openai_static_api_key(credentials.api_key) or _workload_identity_config(litellm_params) + if isinstance(credential, OpenAIError): + return credential + if credential is None: + return OpenAIError( + status_code=401, + message=( + f"No OpenAI API key is set for model {model_label}, so tokens can't be counted. Set api_key or " + "workload identity on the deployment, or OPENAI_API_KEY in the proxy environment." + ), + ) + return _CountCredential(api_base=credentials.api_base, credential=credential) + + +def _count_error(model_to_use: str, request_model: str, message: str, status_code: int) -> TokenCountResponse: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="openai_api", + error=True, + error_message=message, + status_code=status_code, + ) + + class OpenAITokenCounter(BaseTokenCounter): """Token counter implementation for OpenAI provider using the Responses API.""" @@ -102,6 +147,7 @@ class OpenAITokenCounter(BaseTokenCounter): messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, context_management: CountContextManagement | None = None, + chat_params: Mapping[str, JsonValue] = MappingProxyType({}), ) -> TokenCountResponse | None: """ Count tokens using OpenAI's Responses API /input_tokens endpoint. @@ -129,41 +175,9 @@ class OpenAITokenCounter(BaseTokenCounter): rejected_before_provider_call=True, ) - deployment = deployment or {} - litellm_params: Final = deployment.get("litellm_params", {}) - - credentials: Final = get_openai_credentials( - api_base=litellm_params.get("api_base"), api_key=litellm_params.get("api_key") - ) - credential: Final = openai_static_api_key(credentials.api_key) or _workload_identity_config(litellm_params) - - if isinstance(credential, OpenAIError): - return TokenCountResponse( - total_tokens=0, - request_model=request_model, - model_used=model_to_use, - tokenizer_type="openai_api", - error=True, - error_message=credential.message, - status_code=credential.status_code, - ) - - if credential is None: - return TokenCountResponse( - total_tokens=0, - request_model=request_model, - model_used=model_to_use, - tokenizer_type="openai_api", - error=True, - error_message=( - f"No OpenAI API key is set for model {request_model or model_to_use}, so tokens can't be " - "counted. Set api_key or workload identity on the deployment, or OPENAI_API_KEY in the proxy " - "environment." - ), - status_code=401, - ) - - api_base: Final = credentials.api_base + resolved: Final = _count_credential(deployment, request_model or model_to_use) + if isinstance(resolved, OpenAIError): + return _count_error(model_to_use, request_model, resolved.message, resolved.status_code) input_items, instructions, counted_tools, text = self._responses_request_parts( model_to_use=model_to_use, @@ -178,13 +192,9 @@ class OpenAITokenCounter(BaseTokenCounter): if not input_items: return None - try: - api_key: Final = ( - credential - if isinstance(credential, str) - else await asyncify(get_workload_identity_bearer_token)(credential) - ) - result: Final = await openai_count_tokens_handler.handle_count_tokens_request( + return await self._counted( + resolved, + lambda api_key, api_base: openai_count_tokens_handler.handle_count_tokens_request( model=model_to_use, input=input_items, api_key=api_key, @@ -192,7 +202,44 @@ class OpenAITokenCounter(BaseTokenCounter): tools=counted_tools, instructions=instructions, text=text, + ), + model_to_use, + request_model, + ) + + async def count_responses_request( + self, + model_to_use: str, + request: Mapping[str, object], + policy: DropPolicy, + deployment: Mapping[str, object] | None, + request_model: str, + ) -> TokenCountResponse | None: + resolved: Final = _count_credential(deployment, request_model or model_to_use) + if isinstance(resolved, OpenAIError): + return _count_error(model_to_use, request_model, resolved.message, resolved.status_code) + body: Final = native_count_request(request, model_to_use, policy) + return await self._counted( + resolved, + lambda api_key, api_base: openai_count_tokens_handler.post_count_request(body, api_key, api_base), + model_to_use, + request_model, + ) + + @staticmethod + async def _counted( + resolved: _CountCredential, + send: Callable[[str, str | None], Awaitable[Mapping[str, object]]], + model_to_use: str, + request_model: str, + ) -> TokenCountResponse | None: + try: + api_key: Final = ( + resolved.credential + if isinstance(resolved.credential, str) + else await asyncify(get_workload_identity_bearer_token)(resolved.credential) ) + result: Final = await send(api_key, resolved.api_base) if result is not None: return TokenCountResponse.model_validate( @@ -206,36 +253,12 @@ class OpenAITokenCounter(BaseTokenCounter): ) except OpenAIError as e: verbose_logger.warning("OpenAI CountTokens API error: status=%s, message=%s", e.status_code, e.message) - return TokenCountResponse( - total_tokens=0, - request_model=request_model, - model_used=model_to_use, - tokenizer_type="openai_api", - error=True, - error_message=e.message, - status_code=e.status_code, - ) + return _count_error(model_to_use, request_model, e.message, e.status_code) except openai.AuthenticationError as e: verbose_logger.warning("OpenAI workload identity token exchange failed: %s", e.message) - return TokenCountResponse( - total_tokens=0, - request_model=request_model, - model_used=model_to_use, - tokenizer_type="openai_api", - error=True, - error_message=e.message, - status_code=401, - ) + return _count_error(model_to_use, request_model, e.message, 401) except Exception as e: verbose_logger.warning("Error calling OpenAI CountTokens API: %s", e) - return TokenCountResponse( - total_tokens=0, - request_model=request_model, - model_used=model_to_use, - tokenizer_type="openai_api", - error=True, - error_message=str(e), - status_code=500, - ) + return _count_error(model_to_use, request_model, str(e), 500) return None diff --git a/litellm/llms/openai/responses/count_tokens/transformation.py b/litellm/llms/openai/responses/count_tokens/transformation.py index 438aa05a4e9..6d882e0eeb1 100644 --- a/litellm/llms/openai/responses/count_tokens/transformation.py +++ b/litellm/llms/openai/responses/count_tokens/transformation.py @@ -4,12 +4,15 @@ OpenAI Responses API token counting transformation logic. This module handles the transformation of requests to OpenAI's /v1/responses/input_tokens endpoint. """ +import copy from collections.abc import Mapping, Sequence from typing import Any, Final, Literal, TypeAlias +from openai.types.responses.input_token_count_params import InputTokenCountParams from pydantic import JsonValue, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict +from litellm.litellm_core_utils.tool_fields import DropPolicy from litellm.llms.anthropic.pass_through.convertible_blocks import malformed_tool_result_text_error from litellm.llms.anthropic.pass_through.messages.utils import normalize_native_messages from litellm.llms.anthropic.pass_through.responses_adapters.transformation import ( @@ -17,6 +20,8 @@ from litellm.llms.anthropic.pass_through.responses_adapters.transformation impor ) from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.anthropic import AnthropicMessagesRequest, AnthropicOutputSchema +from litellm.types.llms.openai import ResponsesAPIRequestParams +from litellm.types.utils import all_litellm_params class ResponsesInputTextPart(TypedDict): @@ -340,3 +345,34 @@ class OpenAICountTokensConfig: input_items: object, instructions: object, tools: object, text: object = None ) -> ResponsesCountRequestParts: return _RESPONSES_COUNT_REQUEST_PARTS.validate_python((input_items, instructions, tools, text)) + + +_COUNT_FIELDS_THE_SDK_TYPE_OMITS: Final = frozenset({"prompt"}) +_COUNT_REQUEST_FIELDS: Final = ( + InputTokenCountParams.__required_keys__ | InputTokenCountParams.__optional_keys__ | _COUNT_FIELDS_THE_SDK_TYPE_OMITS +) +_INFERENCE_REQUEST_FIELDS: Final = ( + ResponsesAPIRequestParams.__required_keys__ | ResponsesAPIRequestParams.__optional_keys__ +) +_LITELLM_REQUEST_FIELDS: Final = frozenset((*all_litellm_params, "drop_params", "additional_drop_params")) + + +def _forwards(field: str, policy: DropPolicy) -> bool: + if field in _LITELLM_REQUEST_FIELDS or field in policy.additional_drop_params: + return False + return field in _COUNT_REQUEST_FIELDS or (not policy.drop_params and field in _INFERENCE_REQUEST_FIELDS) + + +def _restored(field: str, value: object) -> object: + match field, value: + case "input", _: + return ResponsesAPIRequestUtils.restore_encrypted_content_item_ids_in_input(copy.deepcopy(value)) + case "previous_response_id", str(): + return ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(value) + case _: + return value + + +def native_count_request(request: Mapping[str, object], model: str, policy: DropPolicy) -> Mapping[str, object]: + forwarded: Final = {field: _restored(field, value) for field, value in request.items() if _forwards(field, policy)} + return {**forwarded, "model": model} diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 8298edd6744..f85a5a8cda9 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1,5 +1,5 @@ import re -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from copy import deepcopy from datetime import datetime, timezone from enum import Enum @@ -1263,6 +1263,52 @@ class VertexAITokenCounter(BaseTokenCounter): return custom_llm_provider == LlmProviders.VERTEX_AI.value + async def _count_translated_claude_request( + self, + model_to_use: str, + messages: Sequence[Mapping[str, object]] | None, + deployment: Mapping[str, object], + request_model: str, + tools: Sequence[Mapping[str, object]] | None, + system: object | None, + messages_format: TokenCountMessagesFormat, + output_format: JsonValue, + context_management: CountContextManagement | None, + chat_params: Mapping[str, JsonValue], + ) -> TokenCountResponse | None: + from litellm.llms.anthropic.count_tokens.transformation import ( + AnthropicCountFields, + ChatCountTranslationError, + anthropic_count_fields, + ) + from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import count_tokens_publisher + + if count_tokens_publisher(model_to_use) != "anthropic": + return None + from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import VertexAIAnthropicConfig + + match translated := await anthropic_count_fields( + model_to_use, messages or [], system, tools, messages_format, VertexAIAnthropicConfig(), chat_params + ): + case ChatCountTranslationError(): + return translated.response(request_model, model_to_use, "vertex_ai_partner_models") + case AnthropicCountFields(): + return await self.count_tokens( + model_to_use=model_to_use, + messages=[dict(message) for message in translated.messages], + contents=None, + deployment=dict(deployment), + request_model=request_model, + tools=[dict(tool) for tool in translated.tools] if translated.tools is not None else None, + system=translated.system, + messages_format="anthropic", + output_format=output_format, + context_management=context_management, + count_options=translated.options, + ) + case None: + return None + async def count_tokens( self, model_to_use: str, @@ -1275,6 +1321,8 @@ class VertexAITokenCounter(BaseTokenCounter): messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, context_management: CountContextManagement | None = None, + chat_params: Mapping[str, JsonValue] = MappingProxyType({}), + count_options: Mapping[str, JsonValue] = MappingProxyType({}), ) -> TokenCountResponse | None: import copy @@ -1313,6 +1361,21 @@ class VertexAITokenCounter(BaseTokenCounter): # Check if this is a partner model (Claude, Mistral, etc.) if VertexAIPartnerModels.is_vertex_partner_model(model_to_use): + translated_count: Final = await self._count_translated_claude_request( + model_to_use=model_to_use, + messages=messages, + deployment=deployment, + request_model=request_model, + tools=tools, + system=system, + messages_format=messages_format, + output_format=output_format, + context_management=context_management, + chat_params=chat_params, + ) + if translated_count is not None: + return translated_count + # Use partner models token counter partner_models_handler: Final = VertexAIPartnerModels() @@ -1343,6 +1406,7 @@ class VertexAITokenCounter(BaseTokenCounter): vertex_credentials=vertex_credentials, system=system, tools=tools, + count_options=count_options, ) except (PartnerVertexAIError, httpx.HTTPStatusError) as e: status_code: Final = e.response.status_code diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py index c855a073648..bb200fd542d 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py @@ -1,6 +1,6 @@ # What is this? ## API Handler for calling Vertex AI Partner Models -from collections.abc import Callable +from collections.abc import Callable, Mapping from enum import Enum from types import MappingProxyType from typing import Final @@ -266,6 +266,7 @@ class VertexAIPartnerModels(VertexBase): vertex_credentials=None, system: object | None = None, tools: list[dict[str, object]] | None = None, + count_options: Mapping[str, object] = MappingProxyType({}), ): """ Count tokens for Vertex AI partner models (Anthropic Claude, Mistral, etc.) @@ -302,6 +303,7 @@ class VertexAIPartnerModels(VertexBase): **MappingProxyType( {key: value for key, value in (("system", system), ("tools", tools)) if value is not None} ), + **count_options, } # Prepare litellm_params with credentials diff --git a/litellm/main.py b/litellm/main.py index 71cb7af1fbb..f9bdf64c667 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -9247,6 +9247,7 @@ async def acount_tokens( request_model=model, tools=tools, system=system, + messages_format="openai", ) if result is not None and not result.error: return result diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c5128ddab87..f87b1edc2e5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -14701,7 +14701,9 @@ async def run_thread( # dependencies=[Depends(user_api_key_auth)], # ) # async def get_available_routes(user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth)): +from litellm.litellm_core_utils.tool_fields import DropPolicy, drop_policy from litellm.llms.base_llm.base_utils import BaseTokenCounter +from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient from litellm.repositories.config_repository import ConfigRepository from litellm.repositories.model_repository import ModelRepository @@ -14714,6 +14716,11 @@ from litellm.repositories.table_repositories import ( ) from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository +from litellm.responses.litellm_completion_transformation.count_request import ( + ChatCountBridgeError, + chat_count_request, + with_provider_file_ids, +) def _get_provider_token_counter( @@ -14771,6 +14778,18 @@ def _get_provider_token_counter( return None, None, None +def _provider_count_outcome( + result: Union["TokenCountResponse", None], +) -> Union["TokenCountResponse", ProviderCountFailure, None]: + if result is not None and result.error is True: + return ProviderCountFailure( + status_code=result.status_code or 500, + message=result.error_message or "Token counting failed", + rejected_before_provider_call=result.rejected_before_provider_call, + ) + return result + + def _with_resolved_api_key(litellm_params: object) -> Mapping[str, object]: call_params: Final = _LITELLM_PARAMS_COPY.validate_python(litellm_params) try: @@ -14785,6 +14804,12 @@ def _with_resolved_api_key(litellm_params: object) -> Mapping[str, object]: return call_params +def _with_resolved_credentials(deployment: Mapping[str, object] | None) -> dict[str, object] | None: + if deployment is None: + return None + return {**deployment, "litellm_params": _with_resolved_api_key(deployment.get("litellm_params", {}))} + + _DEPLOYMENT_LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object]) @@ -14825,6 +14850,7 @@ async def _try_provider_token_count( messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, context_management: CountContextManagement | None = None, + chat_params: Mapping[str, JsonValue] = MappingProxyType({}), ) -> Union["TokenCountResponse", ProviderCountFailure, None]: """Attempt provider-specific token counting. Returns None to fall through to local counting.""" if not provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider): @@ -14844,31 +14870,44 @@ async def _try_provider_token_count( model_to_use=model_to_use or "", messages=messages, contents=contents, - deployment=None - if deployment_in_region is None - else { - **deployment_in_region, - "litellm_params": _with_resolved_api_key(deployment_in_region.get("litellm_params", {})), - }, + deployment=_with_resolved_credentials(deployment_in_region), request_model=request_model, tools=tools, system=system, messages_format=messages_format, output_format=output_format, context_management=context_management, + chat_params=chat_params, ) except httpx.HTTPStatusError as e: return ProviderCountFailure( status_code=getattr(e, "status_code", None) or e.response.status_code, message=getattr(e, "message", None) or str(e), ) - if result is not None and result.error is True: - return ProviderCountFailure( - status_code=result.status_code or 500, - message=result.error_message or "Token counting failed", - rejected_before_provider_call=result.rejected_before_provider_call, + return _provider_count_outcome(result) + + +async def _try_native_responses_count( + provider_counter: OpenAITokenCounter, + custom_llm_provider: str | None, + model_to_use: str, + _region: str | None, + responses_request: Mapping[str, object], + policy: DropPolicy, + deployment: Mapping[str, object] | None, + request_model: str, +) -> Union["TokenCountResponse", ProviderCountFailure, None]: + if not provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider): + return None + return _provider_count_outcome( + await provider_counter.count_responses_request( + model_to_use=model_to_use, + request=responses_request, + policy=policy, + deployment=_with_resolved_credentials(deployment), + request_model=request_model, ) - return result + ) async def _provider_token_count( @@ -15150,6 +15189,7 @@ class _TokenCountDeployment(TypedDict, total=False): _TOKEN_COUNT_DEPLOYMENT: Final = TypeAdapter(_TokenCountDeployment) +_ROUTED_DEPLOYMENT: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) def _model_group_deployment(router: Router | None, model_group: str) -> Mapping[str, object] | None: @@ -15228,6 +15268,43 @@ def _count_context_management( ) +def _responses_routing_kwargs(responses_request: Mapping[str, object] | None) -> Mapping[str, object]: + previous_response_id: Final = None if responses_request is None else responses_request.get("previous_response_id") + return {"previous_response_id": previous_response_id} if isinstance(previous_response_id, str) else {} + + +def _count_drop_policy(deployment: Mapping[str, object] | None, responses_request: Mapping[str, object]) -> DropPolicy: + litellm_params: Final = as_str_mapping(None if deployment is None else deployment.get("litellm_params")) or {} + router_defaults: Final[Mapping[str, object]] = ( + {} if llm_router is None else _router_default_litellm_params(llm_router) + ) + return drop_policy({**litellm_params, **router_defaults, **responses_request}) + + +def _native_responses_count( + provider_counter: BaseTokenCounter, + custom_llm_provider: str | None, + responses_request: Mapping[str, object] | None, + deployment: Mapping[str, object] | None, + request_model: str, +) -> Callable[[str, str | None], Awaitable[Union["TokenCountResponse", ProviderCountFailure, None]]] | None: + if responses_request is None or not isinstance(provider_counter, OpenAITokenCounter): + return None + return partial( + _try_native_responses_count, + provider_counter, + custom_llm_provider, + responses_request=responses_request, + policy=_count_drop_policy(deployment, responses_request), + deployment=deployment, + request_model=request_model, + ) + + +def _deployment_api_base(deployment: _TokenCountDeployment) -> str | None: + return _non_empty_str(deployment.get("litellm_params", {}).get("api_base")) + + def _system_message(system: object) -> ChatCompletionSystemMessage | None: if not isinstance(system, (str, list)) or not system: return None @@ -15260,6 +15337,7 @@ async def count_request_tokens( output_format: JsonValue = None, context_management: JsonValue = None, request_litellm_params: Mapping[str, JsonValue] = MappingProxyType({}), + responses_request: Mapping[str, object] | None = None, ) -> TokenCountResponse: prompt: Final = request.prompt messages: Final = request.messages @@ -15281,7 +15359,7 @@ async def count_request_tokens( try: deployment = await llm_router.async_get_available_deployment( model=request.model, - request_kwargs={}, + request_kwargs=dict(_responses_routing_kwargs(responses_request)), ) except Exception: verbose_proxy_logger.exception( @@ -15335,23 +15413,54 @@ async def count_request_tokens( request_messages: Final = _counted_request_messages(messages, messages_format, count_deployment) contents_estimable: Final = not contents or (bool(messages) and _estimable_contents(contents)) prompt_as_message: Final = prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES + routed_deployment: Final = _ROUTED_DEPLOYMENT.validate_python(deployment) + routed_request: Final = ( + None if responses_request is None else with_provider_file_ids(responses_request, routed_deployment) + ) + chat_count: Final = ( + None + if routed_request is None or provider_counter is None or isinstance(provider_counter, OpenAITokenCounter) + else await chat_count_request( + model_to_use, + custom_llm_provider, + _deployment_api_base(deployment_view), + routed_request, + _count_drop_policy(routed_deployment, routed_request), + ) + ) + if isinstance(chat_count, ChatCountBridgeError): + raise ProxyException(message=chat_count.message, type="invalid_request_error", param=None, code=400) ladder: Final = ( None if provider_counter is None else await _provider_token_count( - count_with=partial( + count_with=_native_responses_count( + provider_counter, + custom_llm_provider, + routed_request, + routed_deployment, + request.model, + ) + or partial( _try_provider_token_count, provider_counter, custom_llm_provider, - messages=[{"role": "user", "content": prompt}] if prompt_as_message else request_messages, + messages=( + [dict(message) for message in chat_count.messages] + if chat_count is not None + else [{"role": "user", "content": prompt}] + if prompt_as_message + else request_messages + ), contents=None if prompt_as_message else contents, deployment=count_deployment, request_model=request.model, - tools=tools, + tools=tools if chat_count is None else [dict(tool) for tool in chat_count.tools or ()] or None, system=system, messages_format=messages_format, output_format=output_format, context_management=_count_context_management(context_management, request.model, request_litellm_params), + chat_params=MappingProxyType({}) if chat_count is None else chat_count.params, ), model_to_use=model_to_use, settings=settings, diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index a0fb2b190e0..c918ea3f1d1 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -32,12 +32,13 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger from litellm.anthropic_interface.exceptions.exception_mapping_utils import AnthropicExceptionMapping +from litellm.caching.dual_cache import DualCache from litellm.constants import EMPTY_MAPPING, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, TERMINAL_RESPONSE_STATUSES from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_metadata_variable_name_from_kwargs from litellm.litellm_core_utils.internal_call_metadata import is_background_response, requests_background from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup -from litellm.litellm_core_utils.tool_fields import first_validation_error +from litellm.litellm_core_utils.tool_fields import TOOL_FIELDS, first_validation_error from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, build_blocked_response, @@ -79,6 +80,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio ) from litellm.proxy.route_llm_request import raise_if_required_body_param_missing from litellm.proxy.utils import handle_exception_on_proxy +from litellm.responses.litellm_completion_transformation.count_request import MODEL_FILE_ID_MAPPING_FIELD from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ResponsesWsFrameGuard, spawn_logging_task from litellm.router import Router from litellm.types.llms.base import LiteLLMBaseModel @@ -89,7 +91,7 @@ from litellm.types.llms.openai import ( ) from litellm.types.responses.main import DeleteResponseResult from litellm.types.token_count import TokenCountSourceFields -from litellm.types.utils import TokenCountResponse +from litellm.types.utils import ADDRESSED_RESPONSE_ID_FIELD, TokenCountResponse if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -1388,13 +1390,15 @@ class _TokenCountPayload(TypedDict): class _TokenCounter(Protocol): - def __call__(self, request: TokenCountRequest, call_endpoint: bool) -> Awaitable[TokenCountResponse]: ... + def __call__( + self, *, request: TokenCountRequest, call_endpoint: bool, responses_request: Mapping[str, object] | None + ) -> Awaitable[TokenCountResponse]: ... def _proxy_token_counter() -> _TokenCounter: - from litellm.proxy.proxy_server import token_counter + from litellm.proxy.proxy_server import count_request_tokens - return token_counter + return count_request_tokens _token_counter_dep: Final = Depends(_proxy_token_counter) @@ -1409,6 +1413,67 @@ def _proxy_llm_router() -> Router | None: _llm_router_dep: Final = Depends(_proxy_llm_router) +def _proxy_response_id_security() -> ResponsesIDSecurity | None: + from litellm.proxy.proxy_server import proxy_logging_obj + + security_hook: Final = proxy_logging_obj.get_proxy_hook("responses_id_security") + return security_hook if isinstance(security_hook, ResponsesIDSecurity) else None + + +_response_id_security_dep: Final = Depends(_proxy_response_id_security) + + +class _ManagedFileAccess(Protocol): + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: Literal["aresponses"], + ) -> object: ... + + +def _proxy_managed_files_hook() -> _ManagedFileAccess | None: + from litellm.proxy.proxy_server import proxy_logging_obj + + return proxy_logging_obj.get_proxy_hook("managed_files") + + +_managed_files_hook_dep: Final = Depends(_proxy_managed_files_hook) + + +async def _with_managed_file_id_mapping( + data: Mapping[str, object], managed_files: _ManagedFileAccess | None, user_api_key_dict: UserAPIKeyAuth +) -> Mapping[str, object]: + caller_fields: Final = {name: value for name, value in data.items() if name != MODEL_FILE_ID_MAPPING_FIELD} + if managed_files is None: + return caller_fields + from litellm.proxy.proxy_server import user_api_key_cache + + hooked: Final = dict(caller_fields) + try: + await managed_files.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, cache=user_api_key_cache, data=hooked, call_type="aresponses" + ) + except HTTPException as e: + raise ProxyException(message=str(e.detail), type="permission_error", param=None, code=e.status_code) from e + mapping: Final = hooked.get(MODEL_FILE_ID_MAPPING_FIELD) + return caller_fields if mapping is None else {**caller_fields, MODEL_FILE_ID_MAPPING_FIELD: mapping} + + +def _with_authorized_previous_response( + data: Mapping[str, object], security: ResponsesIDSecurity | None, user_api_key_dict: UserAPIKeyAuth +) -> Mapping[str, object]: + previous_response_id: Final = data.get("previous_response_id") + if security is None or not isinstance(previous_response_id, str) or not previous_response_id: + return data + try: + authorized_id: Final = security.authorize_response_id(previous_response_id, user_api_key_dict) + except HTTPException as e: + raise ProxyException(message=str(e.detail), type="permission_error", param=None, code=e.status_code) from e + return {**data, "previous_response_id": authorized_id, ADDRESSED_RESPONSE_ID_FIELD: previous_response_id} + + def _responses_invalid_request_response(message: str, param: str | None, code: str | None) -> JSONResponse: body: Final[_ResponsesApiErrorBody] = { "error": { @@ -1546,6 +1611,8 @@ async def responses_input_tokens( token_counter: _TokenCounter = _token_counter_dep, llm_router: Router | None = _llm_router_dep, user_api_key_dict: UserAPIKeyAuth = _user_api_key_auth_dep, + response_id_security: ResponsesIDSecurity | None = _response_id_security_dep, + managed_files: _ManagedFileAccess | None = _managed_files_hook_dep, ): """ Count the input tokens of a Responses API request without calling the model. @@ -1579,7 +1646,7 @@ async def responses_input_tokens( ) counted_fields: Final = { - name: data[name] for name in ("model", "tools", "drop_params", "additional_drop_params") if name in data + name: data[name] for name in ("model", "drop_params", "additional_drop_params", *TOOL_FIELDS) if name in data } token_fields: Final = resolve_counted_tool_fields(counted_fields, llm_router, user_api_key_dict.team_id) if isinstance(token_fields, ProxyException): @@ -1604,8 +1671,20 @@ async def responses_input_tokens( message=f"Invalid request for token counting: {e}", param=None, code=None ) + request_fields: Final = _PROXY_SERVER_REQUEST_ADAPTER.validate_python(data) + responses_request: Final = await _with_managed_file_id_mapping( + _with_authorized_previous_response( + {**{name: value for name, value in request_fields.items() if name not in counted_fields}, **token_fields}, + response_id_security, + user_api_key_dict, + ), + managed_files, + user_api_key_dict, + ) try: - token_response: Final = await token_counter(request=token_request, call_endpoint=True) + token_response: Final = await token_counter( + request=token_request, call_endpoint=True, responses_request=responses_request + ) except ProxyException as e: if e.type != _PROVIDER_COUNT_ERROR_TYPE: raise diff --git a/litellm/responses/file_search/emulated_handler.py b/litellm/responses/file_search/emulated_handler.py index e2bb2ae40ea..7121908517b 100644 --- a/litellm/responses/file_search/emulated_handler.py +++ b/litellm/responses/file_search/emulated_handler.py @@ -45,7 +45,7 @@ class FileSearchToolCallArgs(TypedDict): def should_use_emulated_file_search( - tools: Iterable[ToolParam] | None, + tools: Iterable[object] | None, provider_config: "BaseResponsesAPIConfig | None", ) -> bool: """Return True when there is a file_search tool and the provider can't handle it natively.""" @@ -105,7 +105,7 @@ def _build_function_tool(vector_store_ids: list[str]) -> dict[str, object]: } -def _replace_file_search_tools( +def replace_file_search_tools( tools: Iterable[ToolParam] | None, ) -> tuple[list[object], list[str]]: """ @@ -529,7 +529,7 @@ async def aresponses_with_emulated_file_search( _include_search_results, call_kwargs = _prepare_emulated_file_search_call(kwargs=kwargs) # 1. Replace file_search tools with function tool - transformed_tools, all_vs_ids = _replace_file_search_tools(tools) + transformed_tools, all_vs_ids = replace_file_search_tools(tools) # 2. First provider call — provider will call the file_search function. # Mark as an internal sub-call so wrapper_async skips billing callbacks; diff --git a/litellm/responses/litellm_completion_transformation/count_request.py b/litellm/responses/litellm_completion_transformation/count_request.py new file mode 100644 index 00000000000..56a0b9024c5 --- /dev/null +++ b/litellm/responses/litellm_completion_transformation/count_request.py @@ -0,0 +1,207 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final, cast # noqa: TID251 # input items are forwarded as sent, as inference does + +from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError + +from litellm.exceptions import BadRequestError +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + update_responses_input_with_model_file_ids, + update_responses_tools_with_model_file_ids, +) +from litellm.litellm_core_utils.tool_fields import DropPolicy +from litellm.llms.base_llm.files.batch_records import responses_batch_body_to_chat_body +from litellm.responses.additional_tools import hoist_additional_tools +from litellm.responses.file_search.emulated_handler import replace_file_search_tools +from litellm.responses.litellm_completion_transformation.transformation import LiteLLMCompletionResponsesConfig +from litellm.responses.main import previous_response_client_ids +from litellm.types.llms.openai import ResponseInputParam +from litellm.utils import ProviderConfigManager + +_COUNTED_CHAT_PARAMS: Final = ( + "tool_choice", + "parallel_tool_calls", + "response_format", + "web_search_options", + "reasoning_effort", +) +_ITEMS: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[object, ...]) +_MAPPINGS: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter(tuple[Mapping[str, object], ...]) +_JSON_PARAMS: Final[TypeAdapter[Mapping[str, JsonValue]]] = TypeAdapter(Mapping[str, JsonValue]) +_ENTRY: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) +_MARKED_ENTRY_LISTS: Final = frozenset({"content", "tool_calls"}) +_FILE_ID_MAPPING: Final[TypeAdapter[Mapping[str, Mapping[str, str]] | None]] = TypeAdapter( + Mapping[str, Mapping[str, str]] | None +) +MODEL_FILE_ID_MAPPING_FIELD: Final = "model_file_id_mapping" + + +@dataclass(frozen=True, slots=True) +class ChatCountRequest: + messages: tuple[Mapping[str, object], ...] + tools: tuple[Mapping[str, object], ...] | None + params: Mapping[str, JsonValue] + + +@dataclass(frozen=True, slots=True) +class ChatCountBridgeError: + message: str + + +def _plain(item: object) -> object: + return item.model_dump(exclude_none=True) if isinstance(item, BaseModel) else item + + +def _mappings(value: object) -> tuple[Mapping[str, object], ...]: + return _MAPPINGS.validate_python(tuple(_plain(item) for item in _ITEMS.validate_python(value or ()))) + + +def _unmarked(entry: Mapping[str, object]) -> Mapping[str, object]: + return { + key: _unmarked_entries(value) if key in _MARKED_ENTRY_LISTS else value + for key, value in entry.items() + if key != "cache_control" + } + + +def _unmarked_entries(value: object) -> object: + return ( + [ + _unmarked(_ENTRY.validate_python(item)) if isinstance(item, Mapping) else item + for item in _ITEMS.validate_python(value) + ] + if isinstance(value, list) + else value + ) + + +def _routed_model_id(deployment: Mapping[str, object] | None) -> str | None: + model_info: Final = as_str_mapping(None if deployment is None else deployment.get("model_info")) or {} + model_id: Final = model_info.get("id") + return model_id if isinstance(model_id, str) else None + + +def with_provider_file_ids( + responses_request: Mapping[str, object], deployment: Mapping[str, object] | None +) -> Mapping[str, object]: + mapping: Final = _FILE_ID_MAPPING.validate_python(responses_request.get(MODEL_FILE_ID_MAPPING_FIELD)) + model_id: Final = _routed_model_id(deployment) + fields: Final = {name: value for name, value in responses_request.items() if name != MODEL_FILE_ID_MAPPING_FIELD} + tools: Final = fields.get("tools") + return { + **fields, + **( + {"input": update_responses_input_with_model_file_ids(fields["input"], model_id, mapping)} + if "input" in fields + else {} + ), + **( + { + "tools": update_responses_tools_with_model_file_ids( + cast(list[dict[str, object]], tools), # cast-ok: entries that are not tools pass through as sent + model_id, + mapping, + ) + } + if isinstance(tools, list) + else {} + ), + } + + +def _with_hoisted_additional_tools(responses_request: Mapping[str, object]) -> Mapping[str, object]: + responses_input: Final = responses_request.get("input") + if not isinstance(responses_input, list): + return responses_request + input_items: Final = cast(ResponseInputParam, responses_input) # cast-ok: forwarded as sent, as inference does + hoisted: Final = hoist_additional_tools(input_items, None) + if not hoisted.hoisted: + return responses_request + return { + **responses_request, + "input": hoisted.input, + "tools": [*_ITEMS.validate_python(responses_request.get("tools") or ()), *hoisted.hoisted], + } + + +def _with_emulated_file_search(responses_request: Mapping[str, object]) -> Mapping[str, object]: + tools: Final = responses_request.get("tools") + if tools is None: + return responses_request + return {**responses_request, "tools": replace_file_search_tools(_ITEMS.validate_python(tools))[0]} + + +async def _with_previous_turns( + chat_body: Mapping[str, object], + model: str, + custom_llm_provider: str, + responses_request: Mapping[str, object], + policy: DropPolicy, +) -> Mapping[str, object]: + previous_response_id: Final = responses_request.get("previous_response_id") + if not isinstance(previous_response_id, str) or not previous_response_id: + return chat_body + instructions: Final = responses_request.get("instructions") + return await LiteLLMCompletionResponsesConfig.async_responses_api_session_handler( # pyright: ignore[reportUnknownMemberType] # the session handler takes a bare dict request + previous_response_id=previous_response_id, + litellm_completion_request={**chat_body, "model": model, "custom_llm_provider": custom_llm_provider}, + instructions=instructions if isinstance(instructions, str) else None, + request_kwargs={"drop_params": policy.drop_params}, + previous_response_client_ids=previous_response_client_ids(previous_response_id, responses_request), + ) + + +def _counted_params( + chat_body: Mapping[str, object], tools: tuple[Mapping[str, object], ...], policy: DropPolicy +) -> Mapping[str, JsonValue]: + sent: Final = { + name: chat_body[name] + for name in _COUNTED_CHAT_PARAMS + if chat_body.get(name) is not None and name not in policy.additional_drop_params + } + return _JSON_PARAMS.validate_python( + {**sent, **({"drop_params": True} if policy.drop_params and (sent or tools) else {})} + ) + + +async def chat_count_request( + model: str, + custom_llm_provider: str | None, + api_base: str | None, + responses_request: Mapping[str, object], + policy: DropPolicy, +) -> ChatCountRequest | ChatCountBridgeError | None: + """The chat request Responses inference sends a provider with no native Responses API, or None + when inference sends this provider the Responses body itself.""" + if custom_llm_provider is None or ( + ProviderConfigManager.get_provider_responses_api_config( + provider=custom_llm_provider, model=model, api_base=api_base + ) + is not None + ): + return None + try: + chat_body: Final = await _with_previous_turns( + responses_batch_body_to_chat_body( + _with_emulated_file_search(_with_hoisted_additional_tools(responses_request)), + custom_llm_provider, + policy.drop_params, + ), + model, + custom_llm_provider, + responses_request, + policy, + ) + tools: Final = ( + () if "tools" in policy.additional_drop_params else tuple(map(_unmarked, _mappings(chat_body.get("tools")))) + ) + return ChatCountRequest( + messages=tuple(map(_unmarked, _mappings(chat_body.get("messages")))), + tools=tools or None, + params=_counted_params(chat_body, tools, policy), + ) + except BadRequestError as e: + return ChatCountBridgeError(message=str(e.message)) + except (ValidationError, ValueError, KeyError) as e: + return ChatCountBridgeError(message=f"Invalid request for token counting: {e}") diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index b613b817225..d9ebf50cc1a 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -1005,7 +1005,7 @@ class LiteLLMCompletionResponsesConfig: session_lookup: Callable[[str], Awaitable[ChatCompletionSession]] | None = None, request_kwargs: Mapping[str, object] = MappingProxyType({}), previous_response_client_ids: tuple[str, ...] = (), - ) -> dict: + ) -> dict[str, object]: """ Async hook to get the chain of previous input and output pairs and return a list of Chat Completion messages """ diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 51cd3be1b63..995921beed1 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -548,7 +548,7 @@ def _bridges_to_chat_completions( _RESPONSES_ONLY_REQUEST_FIELDS_NEVER_BRIDGED: Final = frozenset({"client_metadata"}) -def _previous_response_client_ids(previous_response_id: str | None, kwargs: Mapping[str, object]) -> tuple[str, ...]: +def previous_response_client_ids(previous_response_id: str | None, kwargs: Mapping[str, object]) -> tuple[str, ...]: """The forms of the previous response's id the client held, which its spend log row can be stored under.""" if not previous_response_id: return () @@ -1622,7 +1622,7 @@ def responses( if _bridges_to_chat_completions(responses_api_provider_config, use_chat_completions_api): bridge_kwargs: Final = _bridge_kwargs(kwargs, responses_api_provider_config, allowed_openai_params) return litellm_completion_transformation_handler.response_api_handler( - _previous_response_client_ids(previous_response_id, kwargs), + previous_response_client_ids(previous_response_id, kwargs), model=model, input=input, responses_api_request=response_api_optional_params, diff --git a/tests/integration/providers/_count_tokens_system_lift.py b/tests/integration/providers/_count_tokens_system_lift.py index 56619d1c963..dedf6b5006d 100644 --- a/tests/integration/providers/_count_tokens_system_lift.py +++ b/tests/integration/providers/_count_tokens_system_lift.py @@ -49,6 +49,87 @@ TOOLS: Final[list[JsonValue]] = [ } ] +WEATHER_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"city": {"type": "string", "description": "City to look up"}}, + "required": ["city"], +} +RESPONSES_TOOLS: Final[list[JsonValue]] = [ + { + "type": "function", + "name": "get_weather", + "description": "Look up the current weather for a city", + "parameters": WEATHER_SCHEMA, + } +] +RESPONSES_TOOL_USE: Final[Mapping[str, dict[str, JsonValue]]] = MappingProxyType( + { + "function-tool": {"input": "Count this message", "instructions": "Be terse", "tools": RESPONSES_TOOLS}, + "tool-round-trip": { + "input": [ + {"role": "user", "content": "What is the weather in Paris?"}, + {"type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": '{"city": "Paris"}'}, + {"type": "function_call_output", "call_id": "call_1", "output": "Sunny, 21C"}, + ], + "tools": RESPONSES_TOOLS, + }, + } +) + + +def counted_fields(body: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return {field: body[field] for field in ("messages", "system", "tools") if field in body} + + +REPORT_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"city": {"type": "string"}, "summary": {"type": "string"}}, + "required": ["city", "summary"], + "additionalProperties": False, +} +WEATHER_QUESTION: Final = "What is the weather in Paris?" +RESPONSES_REQUEST_OPTIONS: Final[Mapping[str, dict[str, JsonValue]]] = MappingProxyType( + { + "text-format": { + "input": WEATHER_QUESTION, + "text": {"format": {"type": "json_schema", "name": "weather_report", "schema": REPORT_SCHEMA}}, + }, + "forced-tool-choice-dropped": { + "input": WEATHER_QUESTION, + "tools": RESPONSES_TOOLS, + "tool_choice": {"type": "function", "name": "get_weather"}, + "parallel_tool_calls": False, + "drop_params": True, + }, + "reasoning": { + "input": WEATHER_QUESTION, + "tools": RESPONSES_TOOLS, + "tool_choice": "auto", + "reasoning": {"effort": "low"}, + }, + } +) +_REQUEST_OPTIONS: Final = ("tool_choice", "thinking") + + +def count_request_fields(count_body: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + options: Final = (*_REQUEST_OPTIONS, "output_config", "output_format") + return {**counted_fields(count_body), **{field: count_body[field] for field in options if field in count_body}} + + +def inference_count_fields(inference_body: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + output_config: Final = inference_body.get("output_config") + output_format: Final = inference_body.get("output_format") + counted_config: Final = { + **(output_config if isinstance(output_config, dict) else {}), + **({} if output_format is None else {"format": output_format}), + } + return { + **counted_fields(inference_body), + **{field: inference_body[field] for field in _REQUEST_OPTIONS if field in inference_body}, + **({"output_config": counted_config} if counted_config else {}), + } + @dataclass(frozen=True, slots=True) class LiftCase: diff --git a/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py b/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py index afd25fcfa04..27cbbe81a13 100644 --- a/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py +++ b/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py @@ -6,6 +6,7 @@ from collections.abc import Callable, Iterator, Mapping from concurrent.futures import ThreadPoolExecutor from contextlib import AbstractContextManager, ExitStack from dataclasses import dataclass +from pathlib import Path from queue import SimpleQueue from types import MappingProxyType from typing import Final, cast @@ -14,8 +15,10 @@ import httpcore import httpx import psutil import pytest +import yaml from anthropic.types import MessageParam -from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment +from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment, object_value +from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.wire import Reply, Request, Wire, wire_server from integration._support.ports import allocate_port from integration.providers._count_tokens_system_lift import ( @@ -32,18 +35,24 @@ from integration.providers._count_tokens_system_lift import ( MID_SYSTEM, REJECTION, REMINDER, + RESPONSES_REQUEST_OPTIONS, + RESPONSES_TOOL_USE, STRING_CASE, TOKEN_COUNTING_BETA, TOOLS, USER, USER_TEXT, + WEATHER_QUESTION, LiftCase, accepts_count_body, anthropic_client, async_anthropic_client, async_openai_client, count_request, + count_request_fields, + counted_fields, expected_count_body, + inference_count_fields, openai_client, ) from pydantic import JsonValue, TypeAdapter @@ -480,9 +489,7 @@ def test_messages_count_tokens_unauthenticated_request_never_reaches_the_peer( assert peer.drain() == () -def test_peer_outage_between_concurrent_waves_errors_then_recovers( - deployment: _Deployment, gateway: Gateway -) -> None: +def test_peer_outage_between_concurrent_waves_errors_then_recovers(deployment: _Deployment, gateway: Gateway) -> None: with ExitStack() as stack: clients: Final = _clients(stack, str(gateway.client.base_url), 8) pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) @@ -573,3 +580,164 @@ def test_chat_completions_on_the_same_deployment_keeps_a_mid_conversation_system {"role": "assistant", "content": [{"type": "text", "text": "One."}]}, {"role": "user", "content": [{"type": "text", "text": "Again"}]}, ], body + + +_GENERATED: Final[dict[str, JsonValue]] = { + "id": "msg_generated", + "type": "message", + "role": "assistant", + "model": _MODEL, + "content": [{"type": "text", "text": "done"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 3}, +} + + +def _generating(provider: _Provider) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.target == provider.target.removesuffix("/count_tokens"): + return _json_reply(200, _GENERATED) + return _peer(provider)(request) + + return respond + + +@pytest.mark.parametrize("request_body", RESPONSES_TOOL_USE.values(), ids=RESPONSES_TOOL_USE.keys()) +def test_responses_input_tokens_counts_what_responses_inference_sends( + deployment: _Deployment, gateway: Gateway, request_body: dict[str, JsonValue] +) -> None: + with wire_server(_generating(deployment.provider), port=deployment.port) as peer: + generated: Final = gateway.request( + "POST", "/v1/responses", {**request_body, "model": deployment.model, "max_output_tokens": 16} + ) + assert generated.status_code == 200, generated.text + response: Final = gateway.request( + "POST", "/v1/responses/input_tokens", {**request_body, "model": deployment.model} + ) + assert _payload(response)["input_tokens"] == _COUNT, response.text + inference, counted = (counted_fields(JSON_OBJECT.validate_json(request.body)) for request in peer.drain()) + assert counted == inference, (counted, inference) + + +def _inference_and_count(peer: Wire) -> tuple[dict[str, JsonValue], dict[str, JsonValue]]: + inference, counted = (JSON_OBJECT.validate_json(request.body) for request in peer.drain()) + return inference, counted + + +@pytest.mark.parametrize("request_body", RESPONSES_REQUEST_OPTIONS.values(), ids=RESPONSES_REQUEST_OPTIONS.keys()) +def test_responses_input_tokens_counts_the_request_options_responses_inference_sends( + deployment: _Deployment, gateway: Gateway, request_body: dict[str, JsonValue] +) -> None: + with wire_server(_generating(deployment.provider), port=deployment.port) as peer: + generated: Final = gateway.request( + "POST", "/v1/responses", {**request_body, "model": deployment.model, "max_output_tokens": 4096} + ) + assert generated.status_code == 200, generated.text + response: Final = gateway.request( + "POST", "/v1/responses/input_tokens", {**request_body, "model": deployment.model} + ) + assert _payload(response)["input_tokens"] == _COUNT, response.text + inference, counted = _inference_and_count(peer) + assert count_request_fields(counted) == inference_count_fields(inference), (counted, inference) + + +def _text_blocks(message: JsonValue) -> JsonValue: + if not isinstance(message, dict) or not isinstance(message.get("content"), str): + return message + return {**message, "content": [{"type": "text", "text": message["content"]}]} + + +def _with_text_blocks(fields: dict[str, JsonValue]) -> dict[str, JsonValue]: + messages: Final = fields["messages"] + assert isinstance(messages, list), fields + return {**fields, "messages": [_text_blocks(message) for message in messages]} + + +def test_responses_input_tokens_leaves_tool_choice_out_when_drop_params_drops_every_tool( + deployment: _Deployment, gateway: Gateway +) -> None: + request_body: Final[dict[str, JsonValue]] = { + "input": WEATHER_QUESTION, + "tools": [5], + "tool_choice": "required", + "drop_params": True, + } + with wire_server(_generating(deployment.provider), port=deployment.port) as peer: + generated: Final = gateway.request( + "POST", "/v1/responses", {**request_body, "model": deployment.model, "max_output_tokens": 16} + ) + assert generated.status_code == 200, generated.text + response: Final = gateway.request( + "POST", "/v1/responses/input_tokens", {**request_body, "model": deployment.model} + ) + assert _payload(response)["input_tokens"] == _COUNT, response.text + inference, counted = _inference_and_count(peer) + assert "tools" not in counted and "tool_choice" not in counted, counted + assert _with_text_blocks(count_request_fields(counted)) == inference_count_fields(inference), ( + counted, + inference, + ) + + +def _config_storing_prompts(directory: Path) -> Path: + config: Final = object_value(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + general_settings: Final = object_value(config["general_settings"]) + path: Final = directory / "store-prompts.yaml" + path.write_text( + yaml.safe_dump({**config, "general_settings": {**general_settings, "store_prompts_in_spend_logs": True}}) + ) + return path + + +def _replying(answer: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.target == _ANTHROPIC.target: + return _counted(request) + if request.target == "/v1/messages": + return _json_reply( + 200, + {**_GENERATED, "id": f"msg_{uuid.uuid4().hex}", "content": [{"type": "text", "text": answer}]}, + ) + return _json_reply(404, {"error": f"unscripted target {request.target}"}) + + return respond + + +def _spend_rows_with(prompt_gateway: Gateway, marker: str) -> list[dict[str, JsonValue]]: + return prompt_gateway.rows( + 'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE proxy_server_request::text LIKE %s', (f"%{marker}%",) + ) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_responses_input_tokens_counts_the_turns_previous_response_id_replays(gateway: Gateway, tmp_path: Path) -> None: + first_prompt: Final = f"Remember the number {uuid.uuid4().hex}" + answer: Final = f"Remembered {uuid.uuid4().hex}" + follow_up: Final[dict[str, JsonValue]] = {"input": "What number did I ask you to remember?"} + port: Final = allocate_port() + with ( + wire_server(_replying(answer), port=port) as peer, + owned_proxy(gateway, tmp_path, {}, config=_config_storing_prompts(tmp_path), workers=1) as prompt_gateway, + prompt_gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model_info=None, model=f"anthropic/{_MODEL}", api_base=f"http://127.0.0.1:{port}", api_key=_API_KEY + ) + first: Final = prompt_gateway.request( + "POST", "/v1/responses", {"model": model, "input": first_prompt, "max_output_tokens": 16} + ) + assert first.status_code == 200, first.text + previous: Final = { + **follow_up, + "model": model, + "previous_response_id": JSON_OBJECT.validate_json(first.content)["id"], + } + eventually(lambda: _spend_rows_with(prompt_gateway, first_prompt), lambda rows: len(rows) == 1, seconds=70) + response: Final = prompt_gateway.request("POST", "/v1/responses/input_tokens", previous) + assert _payload(response)["input_tokens"] == _COUNT, response.text + generated: Final = prompt_gateway.request("POST", "/v1/responses", {**previous, "max_output_tokens": 16}) + assert generated.status_code == 200, generated.text + _, counted, inference = (JSON_OBJECT.validate_json(request.body) for request in peer.drain()) + assert first_prompt in json.dumps(inference["messages"]) and answer in json.dumps(inference["messages"]) + assert counted_fields(counted) == counted_fields(inference), (counted, inference) diff --git a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py index 4fe29fa49d8..063b0e46479 100644 --- a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py @@ -32,9 +32,12 @@ from integration.providers._count_tokens_system_lift import ( LIFTED, MID_SYSTEM, REMINDER, + RESPONSES_TOOLS, STRING_CASE, USER, USER_TEXT, + WEATHER_QUESTION, + WEATHER_SCHEMA, LiftCase, accepts_count_body, async_openai_client, @@ -1636,3 +1639,107 @@ def test_utils_token_counter_local_mode_never_calls_either_peer(counting_proxy: model: Final = _fallback_deployment(scenario, runtime.url) assert _local_count(gateway, _full(model)) > 0 assert runtime.drain() == () and mantle.drain() == () + + +_FORCED_TOOL_CHOICE: Final[dict[str, JsonValue]] = { + "input": WEATHER_QUESTION, + "tools": RESPONSES_TOOLS, + "tool_choice": {"type": "function", "name": "get_weather"}, +} +_CHAT_TOOLS: Final[list[JsonValue]] = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Look up the current weather for a city", + "parameters": WEATHER_SCHEMA, + }, + } +] +_WEATHER_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": WEATHER_QUESTION}] +_TOOL_USE_FIELDS: Final = ("tools", "tool_choice") + + +def _tool_use_fields(body: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return {field: body[field] for field in _TOOL_USE_FIELDS if field in body} + + +def _converse_count_tools(request: Request) -> JsonValue: + count_input: Final = _JSON_OBJECT.validate_python(_JSON_OBJECT.validate_json(request.body)["input"]) + converse: Final = _JSON_OBJECT.validate_python(count_input["converse"]) + return _JSON_OBJECT.validate_python(converse["toolConfig"])["tools"] + + +def _requests_to(requests: Sequence[Request], pattern: re.Pattern[str]) -> tuple[Request, ...]: + return tuple(request for request in requests if pattern.match(target_of(request))) + + +def test_responses_input_tokens_sends_mantle_the_tool_use_invoke_inference_sends( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + counted_model: Final = _deployment(scenario, runtime.url) + invoked_model: Final = _deployment(scenario, runtime.url, f"invoke/{_OPUS}") + generated: Final = gateway.request( + "POST", "/v1/responses", {**_FORCED_TOOL_CHOICE, "model": invoked_model, "max_output_tokens": 16} + ) + assert generated.status_code == 200, generated.text + response: Final = gateway.request( + "POST", "/v1/responses/input_tokens", {**_FORCED_TOOL_CHOICE, "model": counted_model} + ) + assert _payload(response)["input_tokens"] == _MANTLE_COUNT, response.text + requests: Final = runtime.drain() + (invoked,) = _requests_to(requests, _INVOKE_TARGET) + (mantle_count,) = _mantle_bodies(mantle) + expected: Final = _tool_use_fields(_JSON_OBJECT.validate_json(invoked.body)) + assert set(expected) == set(_TOOL_USE_FIELDS), expected + assert _requests_to(requests, _COUNT_TARGET) == () + assert _tool_use_fields(mantle_count) == expected, mantle_count + + +def _counting_every_model(request: Request) -> Reply: + if _COUNT_TARGET.match(target_of(request)): + return _json_reply(200, {"inputTokens": _RUNTIME_COUNT}) + return runtime_generation(request) + + +_NOVA_COUNTS: Final[Mapping[str, tuple[str, dict[str, JsonValue], str]]] = MappingProxyType( + { + "utils-token-counter": ( + "/utils/token_counter?call_endpoint=true", + {"messages": _WEATHER_MESSAGES, "tools": _CHAT_TOOLS}, + "total_tokens", + ), + "responses-input-tokens": ( + "/v1/responses/input_tokens", + {"input": WEATHER_QUESTION, "tools": RESPONSES_TOOLS}, + "input_tokens", + ), + } +) + + +@pytest.mark.parametrize("route", _NOVA_COUNTS.keys()) +def test_nova_count_sends_the_tool_specs_converse_inference_sends(counting_proxy: OwnedProxy, route: str) -> None: + path, body, count_field = _NOVA_COUNTS[route] + gateway: Final = counting_proxy.gateway + with wire_server(_counting_every_model) as runtime, gateway.scenario() as scenario: + model: Final = _deployment(scenario, runtime.url, _NOVA) + generated: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": _WEATHER_MESSAGES, "tools": _CHAT_TOOLS, "max_tokens": 16}, + ) + assert generated.status_code == 200, generated.text + response: Final = gateway.request("POST", path, {**body, "model": model}) + assert _payload(response)[count_field] == _RUNTIME_COUNT, response.text + requests: Final = runtime.drain() + (converse,) = (request for request in requests if target_of(request) == f"/model/{_NOVA}/converse") + (counted,) = _requests_to(requests, _COUNT_TARGET) + inference_tools: Final = _JSON_OBJECT.validate_python(_JSON_OBJECT.validate_json(converse.body)["toolConfig"]) + assert _converse_count_tools(counted) == inference_tools["tools"], counted.body diff --git a/tests/integration/providers/test_malformed_tools_wire.py b/tests/integration/providers/test_malformed_tools_wire.py index 706babac398..910f3faf312 100644 --- a/tests/integration/providers/test_malformed_tools_wire.py +++ b/tests/integration/providers/test_malformed_tools_wire.py @@ -1418,7 +1418,7 @@ def test_count_tokens_rejects_malformed_tools_when_a_request_list_replaces_the_d assert _delivered(rig, marker) == () -def test_responses_input_tokens_drops_the_bad_entry_a_deployment_list_names( +def test_responses_input_tokens_drops_the_tools_a_deployment_list_names( rig: _Rig, field_groups: _FieldGroups ) -> None: marker: Final = _marker() @@ -1428,8 +1428,7 @@ def test_responses_input_tokens_drops_the_bad_entry_a_deployment_list_names( {"model": field_groups.anthropic, "input": marker, "tools": [RESPONSES.tool, 5]}, ) assert response.status_code == 200, response.text - tools: Final = _one_delivery(rig, marker, "/v1/messages/count_tokens")["tools"] - assert isinstance(tools, list) and len(tools) == 1 and "get_weather" in json.dumps(tools), tools + assert "tools" not in _one_delivery(rig, marker, "/v1/messages/count_tokens") def test_responses_input_tokens_rejects_malformed_tools_when_a_request_list_replaces_the_deployment_list( @@ -1669,7 +1668,7 @@ _CHECKED: Final[Mapping[str, _Checked]] = MappingProxyType( send=_input_tokens, target="/v1/messages/count_tokens", answered=_counted({"object": "response.input_tokens", "input_tokens": _COUNT, **_PROVIDER_SOURCE}), - kept_by_a_tools_list=("get_weather",), + kept_by_a_tools_list=(), bad_entry=_RESPONSES_BAD_ENTRY, ), } @@ -2336,14 +2335,14 @@ def test_a_deployment_list_with_a_non_string_entry_drops_tools_it_names_on_messa @_TOKEN_COUNTS -def test_token_counting_with_a_request_list_with_a_non_string_entry_drops_the_bad_entry( +def test_token_counting_with_a_request_list_with_a_non_string_entry_counts_the_tools_its_route_keeps( rig: _Rig, checked: _Checked ) -> None: marker: Final = _marker() response: Final = checked.send( rig, rig.anthropic, marker, {"tools": [checked.tool, 5], "additional_drop_params": ["tools", 5]} ) - _assert_forwarded(rig, checked, response, marker, ("get_weather",)) + _assert_forwarded(rig, checked, response, marker, checked.kept_by_a_tools_list) def _alias(suffix: str, variant: str) -> str: diff --git a/tests/integration/providers/test_openai_count_tokens_tools_wire.py b/tests/integration/providers/test_openai_count_tokens_tools_wire.py index 16a5b8685ad..96410b03a92 100644 --- a/tests/integration/providers/test_openai_count_tokens_tools_wire.py +++ b/tests/integration/providers/test_openai_count_tokens_tools_wire.py @@ -415,4 +415,76 @@ def test_responses_input_tokens_passes_an_openai_count_error_through( response: Final = gateway.request("POST", "/v1/responses/input_tokens", {"model": model, "input": "hello"}) assert response.status_code == 400, response.text assert response.json() == {"error": {"param": None, "code": None, **error}}, response.text - assert _count_body(wire)["input"] == [{"role": "user", "content": "hello"}] + assert _count_body(wire)["input"] == "hello" + + +_NATIVE_TEXT: Final[dict[str, JsonValue]] = { + "format": {"type": "json_schema", "name": "answer", "schema": _ANSWER_SCHEMA}, +} +_NATIVE_BODY: Final[Mapping[str, JsonValue]] = { + "input": [ + {"role": "user", "content": "hi"}, + {"type": "reasoning", "summary": [], "encrypted_content": _WRAPPED_REASONING}, + {"role": "user", "content": "again"}, + ], + "text": _NATIVE_TEXT, + "max_output_tokens": "lots", + "metadata": {"run": "nightly"}, +} + + +@pytest.mark.parametrize( + ("drop_params", "expected_extra"), + [pytest.param(False, {"max_output_tokens": "lots"}, id="as-sent"), pytest.param(True, {}, id="drop-params")], +) +def test_responses_input_tokens_sends_the_callers_body_to_openai( + counting_proxy: OwnedProxy, drop_params: bool, expected_extra: dict[str, JsonValue] +) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_count_response) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/gpt-4o-mini", + api_base=f"{wire.url}/v1", + api_key="sk-integration-openai", + drop_params=drop_params, + ) + response: Final = gateway.request("POST", "/v1/responses/input_tokens", {"model": model, **_NATIVE_BODY}) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["input_tokens"] == _COUNT + assert _count_body(wire) == { + "model": "gpt-4o-mini", + "input": [ + {"role": "user", "content": "hi"}, + {"type": "reasoning", "summary": [], "encrypted_content": "gAAAAoriginal"}, + {"role": "user", "content": "again"}, + ], + "text": _NATIVE_TEXT, + **expected_extra, + } + + +def test_responses_input_tokens_sends_openai_the_owners_previous_response_id(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_responses_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + team: Final = scenario.team(models=[model]) + owner_key: Final = scenario.key(team_id=team, user_id=scenario.member(team), models=[model]) + other_key: Final = scenario.key(team_id=team, user_id=scenario.member(team), models=[model]) + generated: Final = gateway.request("POST", "/v1/responses", {"model": model, "input": "hi"}, key=owner_key) + assert generated.status_code == 200, generated.text + follow_up: Final = { + "model": model, + "input": "again", + "previous_response_id": _JSON_OBJECT.validate_json(generated.content)["id"], + } + _body_for(wire.drain(), "/v1/responses") + + refused: Final = gateway.request("POST", "/v1/responses/input_tokens", follow_up, key=other_key) + assert refused.status_code == 403, refused.text + assert _JSON_OBJECT.validate_json(refused.content)["error"]["type"] == "permission_error", refused.text + assert wire.drain() == () + + counted: Final = gateway.request("POST", "/v1/responses/input_tokens", follow_up, key=owner_key) + assert counted.status_code == 200, counted.text + count_body: Final = _body_for(wire.drain(), "/v1/responses/input_tokens") + assert count_body["previous_response_id"] == _RESPONSES_REPLY["id"] diff --git a/tests/integration/providers/test_responses_core_wire.py b/tests/integration/providers/test_responses_core_wire.py index 791b2b8783f..9a24903a20d 100644 --- a/tests/integration/providers/test_responses_core_wire.py +++ b/tests/integration/providers/test_responses_core_wire.py @@ -3,6 +3,7 @@ from __future__ import annotations import json import os import uuid +from collections.abc import Callable from typing import Final import openai @@ -192,6 +193,42 @@ def test_previous_response_id_uses_the_raw_id_and_affinity( ], [request.body for request in requests] +def _counting_peer(raw_id: str, input_tokens: int) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return model_discovery_reply(_MODEL) + if request.target == "/v1/responses/input_tokens": + return Reply(body=json.dumps({"object": "response.input_tokens", "input_tokens": input_tokens}).encode()) + return Reply(body=_response(raw_id)) + + return respond + + +def test_input_tokens_count_goes_to_the_deployment_that_served_the_previous_response(gateway: Gateway) -> None: + raw_id: Final = f"resp_count_affinity_{uuid.uuid4().hex}" + with ( + wire_server(_counting_peer(raw_id, 37)) as first_wire, + wire_server(_counting_peer(raw_id, 99)) as second_wire, + gateway.scenario() as scenario, + ): + model: Final = f"integration-{uuid.uuid4().hex}" + _enable_responses_affinity(scenario) + assert _register(scenario, first_wire, model) == model + client: Final = openai.OpenAI(base_url=f"{gateway.client.base_url}/v1", api_key=gateway.key, max_retries=0) + first: Final = client.responses.create(model=model, input="turn one") + _register_in_group(scenario, second_wire, model, order=1) + counted: Final = client.responses.input_tokens.count( + model=model, input="turn two", previous_response_id=first.id + ) + assert counted.input_tokens == 37, counted.model_dump_json() + assert drain_contract_requests(second_wire) == (), "the count left the previous response's deployment" + requests: Final = drain_contract_requests(first_wire) + assert [(request.target, _JSON_OBJECT.validate_json(request.body)) for request in requests] == [ + ("/v1/responses", {"model": _MODEL, "input": "turn one"}), + ("/v1/responses/input_tokens", {"model": _MODEL, "input": "turn two", "previous_response_id": raw_id}), + ], [request.body for request in requests] + + @pytest.mark.parametrize("prefix", _ALIASES) def test_codex_responses_request_preserves_all_supported_fields(gateway: Gateway, prefix: str) -> None: response_id: Final = f"resp_codex_{uuid.uuid4().hex}" diff --git a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py index ec3fb941037..1632e424c55 100644 --- a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py @@ -1096,3 +1096,57 @@ def test_utils_token_counter_counts_a_local_prompt_with_system_or_tools_like_its assert as_message.status_code == 200, as_message.text assert _total_tokens(prompted) == _total_tokens(as_message) assert requests == () + + +_RESPONSES_TOOL_ROUND_TRIP: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "instructions": "Follow the system instruction.", + "input": [ + {"role": "user", "content": "Look up the value."}, + {"type": "function_call", "call_id": "call_lookup", "name": "lookup", "arguments": '{"query":"value"}'}, + {"type": "function_call_output", "call_id": "call_lookup", "output": "result"}, + ], + "tools": [ + { + "type": "function", + "name": "lookup", + "description": "Look up a value", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}, + } + ], + } +) + + +@pytest.mark.parametrize("provider", ("vertex_ai", "gemini")) +def test_responses_input_tokens_counts_the_tool_round_trip_contents_responses_inference_sends( + gateway: Gateway, provider: str +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = ( + _deployment(gateway, scenario, wire.url) + if provider == "vertex_ai" + else _ai_studio_deployment(scenario, wire.url) + ) + body: Final = {"model": model, **_RESPONSES_TOOL_ROUND_TRIP} + generated: Final = gateway.request("POST", "/v1/responses", {**body, "max_output_tokens": 16}) + counted: Final = gateway.request("POST", "/v1/responses/input_tokens", body) + assert generated.status_code == 200, generated.text + assert counted.status_code == 200, counted.text + assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == _COUNT, counted.text + requests: Final = peer.drain() + (generated_request,) = (request for request in requests if request.target.endswith(":generateContent")) + (count_request,) = (request for request in requests if request.target.endswith(":countTokens")) + instruction: Final[JsonValue] = {"text": _RESPONSES_TOOL_ROUND_TRIP["instructions"]} + first_turn, *later_turns = _JSON_LIST.validate_python(_counted_contents(_json_body(count_request))) + first_parts: Final = _JSON_LIST.validate_python(_JSON_OBJECT.validate_python(first_turn)["parts"]) + assert instruction in first_parts, count_request.body + counted_turns: Final = [ + { + **_JSON_OBJECT.validate_python(first_turn), + "parts": [part for part in first_parts if part != instruction], + }, + *later_turns, + ] + assert counted_turns == _json_body(generated_request)["contents"], count_request.body diff --git a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py index 1956301f1be..c679311e4cc 100644 --- a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py @@ -19,6 +19,13 @@ from integration._support.client import Gateway, Scenario, eventually from integration._support.process import graceful_stop_seconds, owned_proxy_process from integration._support.vertex import service_account_json from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._count_tokens_system_lift import ( + RESPONSES_REQUEST_OPTIONS, + RESPONSES_TOOL_USE, + count_request_fields, + counted_fields, + inference_count_fields, +) from integration._support.ports import allocate_port from pydantic import JsonValue, TypeAdapter from tokenizers import Regex, Tokenizer, models, pre_tokenizers @@ -427,14 +434,22 @@ def test_utils_token_counter_local_mode_never_calls_the_peer(gateway: Gateway) - assert wire.drain() == () -def test_utils_token_counter_surfaces_the_peer_error_for_openai_tools(gateway: Gateway) -> None: +def _inference_and_count_fields(wire: Wire) -> tuple[dict[str, JsonValue], ...]: + requests: Final = wire.drain() + (inference,) = (request for request in requests if request.target == _MESSAGE_TARGET) + return (counted_fields(_JSON_OBJECT.validate_json(inference.body)), *map(counted_fields, _count_bodies(requests))) + + +def test_utils_token_counter_counts_openai_tools_the_way_chat_inference_sends_them(gateway: Gateway) -> None: with wire_server(_peer(_strict)) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) body: Final = {**_bare(model), "tools": _OPENAI_TOOLS} + generated: Final = gateway.request("POST", "/v1/chat/completions", {**body, "max_tokens": 16}) + assert generated.status_code == 200, generated.text response: Final = gateway.request("POST", "/utils/token_counter", body, params={"call_endpoint": "true"}) - assert response.status_code == 400, response.text - assert "total_tokens" not in response.text, response.text - assert _counted_bodies(wire) == ({**_PEER_BARE, "tools": _OPENAI_TOOLS},) + assert _payload(response)["total_tokens"] == _PEER_COUNT, response.text + inference, counted = _inference_and_count_fields(wire) + assert counted == inference, (counted, inference) def test_responses_input_tokens_counts_through_the_partner_peer(gateway: Gateway) -> None: @@ -451,8 +466,24 @@ def test_responses_input_tokens_counts_through_the_partner_peer(gateway: Gateway assert _counted_bodies(wire) == (_PEER_BARE,) -def test_responses_input_tokens_surfaces_the_peer_error(gateway: Gateway) -> None: +@pytest.mark.parametrize("request_body", RESPONSES_TOOL_USE.values(), ids=RESPONSES_TOOL_USE.keys()) +def test_responses_input_tokens_counts_what_responses_inference_sends( + gateway: Gateway, request_body: dict[str, JsonValue] +) -> None: with wire_server(_peer(_strict)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + generated: Final = gateway.request( + "POST", "/v1/responses", {**request_body, "model": model, "max_output_tokens": 16} + ) + assert generated.status_code == 200, generated.text + response: Final = gateway.request("POST", "/v1/responses/input_tokens", {**request_body, "model": model}) + assert _payload(response)["input_tokens"] == _PEER_COUNT, response.text + inference, counted = _inference_and_count_fields(wire) + assert counted == inference, (counted, inference) + + +def test_responses_input_tokens_surfaces_the_peer_error(gateway: Gateway) -> None: + with wire_server(_peer(_rejecting(400))) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) response: Final = gateway.request( "POST", @@ -462,7 +493,33 @@ def test_responses_input_tokens_surfaces_the_peer_error(gateway: Gateway) -> Non assert response.status_code == 400, response.text assert response.json() == _RESPONSES_ROUTE.rejection(_REJECTION), response.text (sent,) = _counted_bodies(wire) - assert sent["tools"] == _RESPONSES_TOOLS, sent + assert sent["system"] == [{"type": "text", "text": "Be terse"}], sent + assert sent["tools"] == [ + { + "name": "get_weather", + "input_schema": _WEATHER_SCHEMA, + "type": "custom", + "description": "Look up the current weather for a city", + } + ], sent + + +def test_responses_input_tokens_rejects_an_unpaired_tool_output_without_calling_vertex(gateway: Gateway) -> None: + with wire_server(_peer(_strict)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + body: Final[dict[str, JsonValue]] = { + "model": model, + "input": [ + {"role": "user", "content": "What is the weather in Paris?"}, + {"type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": "Paris?"}, + {"type": "function_call_output", "call_id": "call_1", "output": "Sunny"}, + ], + "tools": _RESPONSES_TOOLS, + } + response: Final = gateway.request("POST", "/v1/responses/input_tokens", body) + assert response.status_code == 400, response.text + assert "get_weather" in _error_message(_JSON_OBJECT.validate_json(response.content)), response.text + assert wire.drain() == () def test_responses_input_tokens_types_a_peer_error_by_its_status(gateway: Gateway) -> None: @@ -890,3 +947,20 @@ def test_responses_endpoint_on_the_same_deployment_still_generates(gateway: Gate (sent,) = wire.drain() assert sent.target == _MESSAGE_TARGET, sent.target assert marker in sent.body.decode(), sent.body + + +def test_responses_input_tokens_counts_the_structured_output_responses_inference_sends(gateway: Gateway) -> None: + request_body: Final = RESPONSES_REQUEST_OPTIONS["text-format"] + with wire_server(_peer(_strict)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + generated: Final = gateway.request( + "POST", "/v1/responses", {**request_body, "model": model, "max_output_tokens": 4096} + ) + assert generated.status_code == 200, generated.text + response: Final = gateway.request("POST", "/v1/responses/input_tokens", {**request_body, "model": model}) + assert _payload(response)["input_tokens"] == _PEER_COUNT, response.text + requests: Final = wire.drain() + (inference,) = (request for request in requests if request.target == _MESSAGE_TARGET) + (counted,) = _count_bodies(requests) + expected: Final = inference_count_fields(_JSON_OBJECT.validate_json(inference.body)) + assert count_request_fields(counted) == expected, (counted, expected) diff --git a/tests/integration/sdk/test_vertex_partner_count_tokens_sdk.py b/tests/integration/sdk/test_vertex_partner_count_tokens_sdk.py index e59d3f9264a..2e1205be688 100644 --- a/tests/integration/sdk/test_vertex_partner_count_tokens_sdk.py +++ b/tests/integration/sdk/test_vertex_partner_count_tokens_sdk.py @@ -1,5 +1,5 @@ import json -from collections.abc import Callable +from collections.abc import Callable, Mapping from typing import Final import litellm @@ -88,3 +88,77 @@ async def test_acount_tokens_falls_back_to_the_local_tokenizer_when_the_peer_rej assert len(_count_requests(wire.drain())) == 1 assert counted.tokenizer_type == "local_tokenizer", counted assert counted.total_tokens > 0 and counted.total_tokens != _PEER_COUNT, counted + + +_WEATHER_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], +} +_CHAT_TOOLS: Final[list[dict[str, JsonValue]]] = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Look up the current weather for a city", + "parameters": _WEATHER_SCHEMA, + }, + } +] +_CHAT_TOOL_USE: Final[list[dict[str, JsonValue]]] = [ + {"role": "user", "content": "What is the weather in Paris?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "Sunny, 21C"}, +] +_GENERATED: Final[dict[str, JsonValue]] = { + "id": "msg_sdk_count", + "type": "message", + "role": "assistant", + "model": _BACKEND, + "content": [{"type": "text", "text": "It is sunny."}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 3}, +} +_COUNTED_FIELDS: Final = ("messages", "system", "tools") + + +def _generating_peer(request: Request) -> Reply: + if request.target == "/_oauth/token": + token: Final = {"access_token": "scripted-token", "token_type": "Bearer", "expires_in": 3600} + return Reply(body=json.dumps(token).encode()) + if "count-tokens" in request.target: + return Reply(body=json.dumps({"input_tokens": _PEER_COUNT}).encode()) + if request.target.endswith(f"/{_BACKEND}:rawPredict"): + return Reply(body=json.dumps(_GENERATED).encode()) + return Reply(status=404, body=json.dumps({"error": request.target}).encode()) + + +def _counted_fields(body: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return {field: body[field] for field in _COUNTED_FIELDS if field in body} + + +async def test_acount_tokens_sends_chat_tool_use_the_way_acompletion_sends_it( + vertex_environment: Callable[[str], None], +) -> None: + with wire_server(_generating_peer) as wire: + vertex_environment(wire.url) + generated: Final = await litellm.acompletion( + model=f"vertex_ai/{_BACKEND}", messages=_CHAT_TOOL_USE, tools=_CHAT_TOOLS, max_tokens=16, api_base=wire.url + ) + assert generated.choices[0].message.content == "It is sunny.", generated + counted: Final = await litellm.acount_tokens( + model=f"vertex_ai/{_BACKEND}", messages=_CHAT_TOOL_USE, tools=_CHAT_TOOLS, api_base=wire.url + ) + assert (counted.total_tokens, counted.tokenizer_type) == (_PEER_COUNT, "vertex_ai_partner_models"), counted + requests: Final = wire.drain() + (inference,) = (request for request in requests if request.target.endswith(f"/{_BACKEND}:rawPredict")) + (count,) = _count_requests(requests) + expected: Final = _counted_fields(_JSON_OBJECT.validate_json(inference.body)) + assert _counted_fields(_JSON_OBJECT.validate_json(count.body)) == expected, (count.body, expected) diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index c08eeea4bb9..cc0ef062e4b 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -1085,6 +1085,30 @@ def test_update_responses_input_with_model_file_ids_decodes_xlitellm_encoded_id( assert updated[0]["content"][1]["file_id"] == raw_id + +@pytest.mark.parametrize("unusable_file_id", [[1], {"a": 1}]) +def test_update_responses_input_with_model_file_ids_keeps_a_non_string_file_id_as_sent(unusable_file_id): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + update_responses_input_with_model_file_ids, + ) + + encoded_id = _xlitellm_encoded("file-ExTuCawUqxEMjVFK6xwR9B", "gpt-5.1") + input_items = [ + { + "role": "user", + "content": [ + {"type": "input_file", "file_id": encoded_id}, + {"type": "input_file", "file_id": unusable_file_id}, + ], + } + ] + + updated = update_responses_input_with_model_file_ids( + input_items, "model-A", {encoded_id: {"model-A": "file-on-a"}} + ) + + assert [item["file_id"] for item in updated[0]["content"]] == ["file-on-a", unusable_file_id] + def test_update_messages_xlitellm_decode_does_not_override_mapping(): """If the call-site already resolved a provider id via the mapping, that wins. The new decode fallback runs only when no mapping match.""" diff --git a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py index 9af7497bc82..63f149fabe4 100644 --- a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py +++ b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py @@ -1,7 +1,10 @@ import asyncio from base64 import b64encode +from collections.abc import Mapping from copy import deepcopy +from queue import SimpleQueue from threading import get_ident +from types import MappingProxyType from typing import Final import httpx @@ -10,16 +13,20 @@ import respx from pydantic import JsonValue, TypeAdapter import litellm +from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler +from litellm.llms.anthropic.count_tokens.token_counter import AnthropicTokenCounter from litellm.llms.anthropic.count_tokens.transformation import ( AnthropicCountTokensConfig, ) from litellm.llms.azure_ai.anthropic.count_tokens.handler import ( AzureAIAnthropicCountTokensHandler, ) +from litellm.llms.azure_ai.anthropic.count_tokens.token_counter import AzureAIAnthropicTokenCounter from litellm.llms.azure_ai.anthropic.count_tokens.transformation import ( AzureAIAnthropicCountTokensConfig, ) +from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse @pytest.mark.parametrize( @@ -441,3 +448,302 @@ async def test_handler_sends_the_leading_system_run_as_system_not_as_a_message(h "system": [{"type": "text", "text": "Be terse"}], "messages": [{"role": "user", "content": "Hello, how are you?"}], } + + +_WEATHER_SCHEMA: Final[dict[str, JsonValue]] = {"type": "object", "properties": {"city": {"type": "string"}}} + +_CHAT_TOOL_TURNS: Final[list[dict[str, JsonValue]]] = [ + {"role": "system", "content": "Be terse."}, + {"role": "user", "content": "Weather in SF?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": '{"city": "SF"}'}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "sunny"}, +] + + +_COUNTERS: Final = pytest.mark.parametrize( + ("counter_type", "count_endpoint"), + ( + (AnthropicTokenCounter, "https://gateway.example/v1/messages/count_tokens"), + (AzureAIAnthropicTokenCounter, "https://gateway.example/anthropic/v1/messages/count_tokens"), + ), +) + + +async def _counted_body( + counter_type: type[AnthropicTokenCounter] | type[AzureAIAnthropicTokenCounter], + count_endpoint: str, + messages: list[dict[str, JsonValue]], + tools: list[dict[str, JsonValue]] | None, + messages_format: TokenCountMessagesFormat, + chat_params: Mapping[str, JsonValue] = MappingProxyType({}), +) -> tuple[TokenCountResponse | None, tuple[dict[str, JsonValue], ...]]: + sent: Final[SimpleQueue[bytes]] = SimpleQueue() + + def count(request: httpx.Request) -> httpx.Response: + sent.put(request.content) + return httpx.Response(200, json={"input_tokens": 31}) + + with respx.mock: + respx.post(count_endpoint).mock(side_effect=count) + result: Final = await counter_type().count_tokens( + model_to_use="claude-opus-5-5", + messages=messages, + contents=None, + deployment={"litellm_params": {"api_key": "sk-ant-api03-test-key", "api_base": "https://gateway.example"}}, + request_model="claude-opus-5-5", + tools=tools, + messages_format=messages_format, + chat_params=chat_params, + ) + return result, tuple(TypeAdapter(dict[str, JsonValue]).validate_json(sent.get()) for _ in range(sent.qsize())) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport_clients") +@_COUNTERS +@pytest.mark.parametrize( + "tool", + ( + {"type": "function", "name": "get_weather", "description": "Look up the weather", "parameters": _WEATHER_SCHEMA}, + { + "type": "function", + "function": {"name": "get_weather", "description": "Look up the weather", "parameters": _WEATHER_SCHEMA}, + }, + ), +) +async def test_counter_sends_chat_tool_use_in_the_shape_chat_inference_sends( + counter_type: type[AnthropicTokenCounter] | type[AzureAIAnthropicTokenCounter], + count_endpoint: str, + tool: dict[str, JsonValue], +) -> None: + """Responses and chat bodies reach the counter as chat messages and tools, which Anthropic rejects + with a 400 unless they are translated the way the chat inference request translates them.""" + original: Final = deepcopy(_CHAT_TOOL_TURNS) + result, bodies = await _counted_body(counter_type, count_endpoint, _CHAT_TOOL_TURNS, [tool], "openai") + + assert result is not None + assert result.error is not True + assert result.total_tokens == 31 + assert bodies == ( + { + "model": "claude-opus-5-5", + "system": [{"type": "text", "text": "Be terse."}], + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "Weather in SF?"}]}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "call_1", "name": "get_weather", "input": {"city": "SF"}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_1", "content": "sunny"}]}, + ], + "tools": [ + { + "name": "get_weather", + "input_schema": _WEATHER_SCHEMA, + "type": "custom", + "description": "Look up the weather", + } + ], + }, + ) + assert _CHAT_TOOL_TURNS == original + + +_ANTHROPIC_TOOL_USE: Final[list[dict[str, JsonValue]]] = [ + {"role": "user", "content": "Weather in SF?"}, + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": "sunny"}]}, +] + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport_clients") +@_COUNTERS +@pytest.mark.parametrize( + ("messages", "tool", "messages_format"), + ( + (_ANTHROPIC_TOOL_USE, {"name": "get_weather", "input_schema": _WEATHER_SCHEMA}, "openai"), + ( + [{"role": "user", "content": "Weather in SF?"}], + {"type": "function", "name": "get_weather", "parameters": _WEATHER_SCHEMA}, + "anthropic", + ), + ), +) +async def test_counter_sends_anthropic_shaped_and_messages_api_requests_as_sent( + counter_type: type[AnthropicTokenCounter] | type[AzureAIAnthropicTokenCounter], + count_endpoint: str, + messages: list[dict[str, JsonValue]], + tool: dict[str, JsonValue], + messages_format: TokenCountMessagesFormat, +) -> None: + _, bodies = await _counted_body(counter_type, count_endpoint, messages, [tool], messages_format) + + assert bodies == ({"model": "claude-opus-5-5", "messages": messages, "tools": [tool]},) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport_clients") +@_COUNTERS +async def test_counter_rejects_a_tool_call_chat_inference_rejects_without_calling_anthropic( + counter_type: type[AnthropicTokenCounter] | type[AzureAIAnthropicTokenCounter], count_endpoint: str +) -> None: + messages: Final[list[dict[str, JsonValue]]] = [ + {"role": "user", "content": "Weather in SF?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": "sunny?"}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "sunny"}, + ] + result, bodies = await _counted_body(counter_type, count_endpoint, messages, None, "openai") + + assert bodies == () + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert result.rejected_before_provider_call is True + assert "Failed to parse tool call arguments for tool 'get_weather'" in (result.error_message or "") + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport_clients") +@_COUNTERS +async def test_counter_rejects_a_chat_tool_without_a_name_without_calling_anthropic( + counter_type: type[AnthropicTokenCounter] | type[AzureAIAnthropicTokenCounter], count_endpoint: str +) -> None: + result, bodies = await _counted_body( + counter_type, + count_endpoint, + [{"role": "user", "content": "Weather in SF?"}], + [{"type": "function", "function": {"parameters": _WEATHER_SCHEMA}}], + "openai", + ) + + assert bodies == () + assert result is not None + assert (result.error, result.status_code, result.rejected_before_provider_call) == (True, 400, True) + + +_CHAT_WEATHER_TOOL: Final[dict[str, JsonValue]] = { + "type": "function", + "function": {"name": "get_weather", "parameters": _WEATHER_SCHEMA}, +} + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport_clients") +@_COUNTERS +@pytest.mark.parametrize( + ("messages", "tools"), + ( + pytest.param( + [ + { + "role": "user", + "name": "ana", + "content": [ + {"type": "text", "text": "Describe this"}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "iVBORw0K"}}, + ], + } + ], + None, + id="native-image-block", + ), + pytest.param( + [ + {"role": "user", "content": "Weather in SF?"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {"city": "SF"}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": "sunny"}]}, + ], + [_CHAT_WEATHER_TOOL], + id="native-tool-result-block", + ), + pytest.param( + [{"role": "user", "content": "Open the docs"}], + [ + {"type": "computer_use_preview", "display_width": 1024, "display_height": 768, "environment": "browser"}, + _CHAT_WEATHER_TOOL, + ], + id="computer-use-preview-tool", + ), + pytest.param( + [{"role": "user", "content": "Ask the wiki"}], + [{"type": "mcp", "server_label": "wiki", "server_url": "https://mcp.example/mcp"}, _CHAT_WEATHER_TOOL], + id="mcp-tool", + ), + ), +) +async def test_counter_rejects_input_chat_translation_would_leave_out_without_calling_anthropic( + counter_type: type[AnthropicTokenCounter] | type[AzureAIAnthropicTokenCounter], + count_endpoint: str, + messages: list[dict[str, JsonValue]], + tools: list[dict[str, JsonValue]] | None, +) -> None: + result, bodies = await _counted_body(counter_type, count_endpoint, messages, tools, "openai") + + assert bodies == () + assert result is not None + assert (result.error, result.status_code, result.rejected_before_provider_call) == (True, 400, True) + + +_ANSWER_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"answer": {"type": "string"}}, + "required": ["answer"], + "additionalProperties": False, +} + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport_clients") +@_COUNTERS +async def test_counter_sends_the_structured_output_tool_choice_and_thinking_chat_inference_sends( + counter_type: type[AnthropicTokenCounter] | type[AzureAIAnthropicTokenCounter], count_endpoint: str +) -> None: + """count_tokens refuses the beta output_format that chat inference sends, so the count carries the same + schema as output_config.format, next to the inference request's thinking and tool_choice.""" + messages: Final[list[dict[str, JsonValue]]] = [{"role": "user", "content": "Weather in SF?"}] + tool: Final[dict[str, JsonValue]] = { + "type": "function", + "function": {"name": "get_weather", "parameters": _WEATHER_SCHEMA}, + } + chat_params: Final[dict[str, JsonValue]] = { + "response_format": {"type": "json_schema", "json_schema": {"name": "answer", "schema": _ANSWER_SCHEMA}}, + "reasoning_effort": "low", + "tool_choice": "auto", + } + config: Final = AnthropicConfig() + inference: Final = config.transform_request( + model="claude-opus-5-5", + messages=deepcopy(messages), + optional_params=config.map_openai_params( + {**chat_params, "tools": [tool]}, {}, "claude-opus-5-5", drop_params=False + ), + litellm_params={}, + headers={}, + ) + + _, bodies = await _counted_body(counter_type, count_endpoint, messages, [tool], "openai", chat_params) + + (body,) = bodies + assert "output_format" not in body + assert body["output_config"] == {**inference.get("output_config", {}), "format": inference["output_format"]} + assert (body["thinking"], body["tool_choice"], body["tools"]) == ( + inference["thinking"], + inference["tool_choice"], + inference["tools"], + ) diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py index b357c5ac126..052baf601f9 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py @@ -168,6 +168,25 @@ def test_transform_includes_tools(): assert tools[0]["toolSpec"]["inputSchema"]["json"]["type"] == "object" +def test_transform_maps_chat_function_tools(): + parameters = {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]} + request = { + "model": "amazon.nova-pro-v1:0", + "messages": [{"role": "user", "content": "Weather in Paris?"}], + "tools": [ + { + "type": "function", + "function": {"name": "get_weather", "description": "Look up the weather", "parameters": parameters}, + } + ], + } + + result = BedrockCountTokensConfig().transform_anthropic_to_bedrock_count_tokens(request) + + assert result["input"]["converse"]["toolConfig"]["tools"] == [ + {"toolSpec": {"name": "get_weather", "description": "Look up the weather", "inputSchema": {"json": parameters}}} + ] + def test_transform_includes_system_and_tools_together(): """Test that both system and tools are included together.""" config = BedrockCountTokensConfig() diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py index dd0d6429f8a..6fb47491323 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py @@ -1,3 +1,4 @@ +import base64 from collections.abc import Mapping from typing import Final @@ -176,6 +177,130 @@ async def test_claude_runtime_404_retries_with_system_and_tools_on_mantle() -> N } +def _invoke_body(request: httpx.Request) -> dict[str, JsonValue]: + invoke_model: Final = _JSON_BODY.validate_json(request.content)["input"] + assert isinstance(invoke_model, dict) + encoded: Final = invoke_model["invokeModel"] + assert isinstance(encoded, dict) and isinstance(encoded["body"], str) + return _JSON_BODY.validate_json(base64.b64decode(encoded["body"])) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "chat_params", + [ + {"tool_choice": "required"}, + { + "response_format": { + "type": "json_schema", + "json_schema": {"name": "city", "schema": {"type": "object", "properties": {"city": {"type": "string"}}}}, + } + }, + {"reasoning_effort": "low"}, + ], +) +async def test_runtime_count_leaves_out_the_options_only_mantle_counts(chat_params: dict[str, JsonValue]) -> None: + runtime_counts: Final = _Bedrock(runtime=(200, {"inputTokens": 412}), mantle=(200, {"input_tokens": 2177})) + mantle_counts: Final = _Bedrock(runtime=(404, UNSUPPORTED), mantle=(200, {"input_tokens": 2177})) + request: Final = { + "model_to_use": "global.anthropic.claude-opus-4-8", + "messages": [{"role": "user", "content": "Weather in SF?"}], + "contents": None, + "deployment": DEPLOYMENT, + "request_model": "claude-opus-4-8", + "tools": [{"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}}], + "messages_format": "openai", + "chat_params": chat_params, + } + + await _counter(runtime_counts).count_tokens(**request) + await _counter(mantle_counts).count_tokens(**request) + + assert runtime_counts.posted_hosts() == (RUNTIME_HOST,) + assert mantle_counts.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST) + options: Final = ("thinking", "tool_choice", "output_config") + invoke_body: Final = _invoke_body(runtime_counts.requests[0]) + mantle_body: Final = _JSON_BODY.validate_json(mantle_counts.requests[1].content) + assert [name for name in options if name in invoke_body] == [] + assert [name for name in options if name in mantle_body] != [] + assert {name: mantle_body[name] for name in ("messages", "tools")} == { + name: invoke_body[name] for name in ("messages", "tools") + } + + +@pytest.mark.asyncio +async def test_claude_chat_tool_use_is_counted_in_the_shape_claude_inference_sends() -> None: + bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 412}), mantle=(200, {"input_tokens": 2177})) + schema: Final = {"type": "object", "properties": {"city": {"type": "string"}}} + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="global.anthropic.claude-opus-4-8", + messages=[ + {"role": "user", "content": "Weather in SF?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": "{}"}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "sunny"}, + ], + contents=None, + deployment=DEPLOYMENT, + request_model="claude-opus-4-8", + tools=[{"type": "function", "function": {"name": "get_weather", "parameters": schema}}], + messages_format="openai", + ) + + assert result is not None + assert (result.error, result.total_tokens) == (False, 412) + assert bedrock.posted_hosts() == (RUNTIME_HOST,) + body: Final = _invoke_body(bedrock.requests[0]) + assert body["tools"] == [{"name": "get_weather", "input_schema": schema, "type": "custom"}] + assert body["messages"] == [ + {"role": "user", "content": [{"type": "text", "text": "Weather in SF?"}]}, + {"role": "assistant", "content": [{"type": "tool_use", "id": "call_1", "name": "get_weather", "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_1", "content": "sunny"}]}, + ] + +@pytest.mark.asyncio +async def test_claude_chat_tool_use_counted_on_mantle_metadata_route_is_translated_too() -> None: + bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 412}), mantle=(200, {"input_tokens": 2177})) + schema: Final = {"type": "object", "properties": {"city": {"type": "string"}}} + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="bedrock_mantle/global.anthropic.claude-opus-4-8", + messages=[ + {"role": "user", "content": "Weather in SF?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": "{}"}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "sunny"}, + ], + contents=None, + deployment=DEPLOYMENT, + request_model="claude-opus-4-8", + tools=[{"type": "function", "function": {"name": "get_weather", "parameters": schema}}], + messages_format="openai", + ) + + assert result is not None + assert (result.error, result.total_tokens) == (False, 2177) + assert bedrock.posted_hosts() == (MANTLE_HOST,) + body: Final = _JSON_BODY.validate_json(bedrock.requests[0].content) + assert body["tools"] == [{"name": "get_weather", "input_schema": schema, "type": "custom"}] + assert body["messages"] == [ + {"role": "user", "content": [{"type": "text", "text": "Weather in SF?"}]}, + {"role": "assistant", "content": [{"type": "tool_use", "id": "call_1", "name": "get_weather", "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_1", "content": "sunny"}]}, + ] + + @pytest.mark.asyncio @pytest.mark.parametrize( ("model", "host", "tokenizer", "expected_tokens"), diff --git a/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py index 7495d98b42a..17c1b35e211 100644 --- a/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py +++ b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py @@ -1,3 +1,4 @@ +import copy import json import operator from collections.abc import Mapping @@ -9,6 +10,7 @@ import respx from pydantic import JsonValue, ValidationError import litellm +from litellm.litellm_core_utils.tool_fields import DropPolicy from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse @@ -521,3 +523,122 @@ async def test_count_sends_the_global_key_only_to_the_global_api_base(monkeypatc assert result is not None assert result.total_tokens == 4 assert route.calls.last.request.headers["authorization"] == "Bearer sk-private-gateway" + + +_KEEP_ALL: Final = DropPolicy(drop_params=False, additional_drop_params=frozenset()) +_ANSWER_FORMAT: Final[dict[str, object]] = { + "format": {"type": "json_schema", "name": "answer", "schema": {"type": "object", "properties": {}}} +} +_STORED_PROMPT: Final[dict[str, object]] = {"id": "pmpt_abc", "variables": {"city": "Paris"}} +_CALLER_BODY: Final[dict[str, object]] = { + "model": "my-gpt", + "input": "hi", + "text": _ANSWER_FORMAT, + "prompt": _STORED_PROMPT, + "max_output_tokens": "lots", + "metadata": {"run": "nightly"}, + "drop_params": False, + "additional_drop_params": [], + "timeout": 30, + "extra_headers": {"x-trace": "1"}, + "disable_fallbacks": True, + "spend_logs_metadata": {"team": "search"}, + "mock_testing_fallbacks": True, +} + + +async def _native_count( + monkeypatch: pytest.MonkeyPatch, request: Mapping[str, object], policy: DropPolicy, response: httpx.Response +) -> tuple[TokenCountResponse | None, object]: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=response) + result: Final = await OpenAITokenCounter().count_responses_request( + model_to_use="gpt-4o", + request=request, + policy=policy, + deployment={"litellm_params": {"api_key": "sk-test"}}, + request_model="my-gpt", + ) + return result, json.loads(route.calls.last.request.content) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("policy", "expected_sent"), + [ + pytest.param( + _KEEP_ALL, + { + "model": "gpt-4o", + "input": "hi", + "text": _ANSWER_FORMAT, + "prompt": _STORED_PROMPT, + "max_output_tokens": "lots", + }, + id="keeps-responses-fields", + ), + pytest.param( + DropPolicy(drop_params=True, additional_drop_params=frozenset()), + {"model": "gpt-4o", "input": "hi", "text": _ANSWER_FORMAT, "prompt": _STORED_PROMPT}, + id="drop-params-keeps-count-fields", + ), + pytest.param( + DropPolicy(drop_params=False, additional_drop_params=frozenset({"text"})), + {"model": "gpt-4o", "input": "hi", "prompt": _STORED_PROMPT, "max_output_tokens": "lots"}, + id="additional-drop-params", + ), + ], +) +async def test_native_count_sends_the_callers_body_without_litellm_fields( + monkeypatch: pytest.MonkeyPatch, policy: DropPolicy, expected_sent: dict[str, object] +) -> None: + result, sent = await _native_count( + monkeypatch, _CALLER_BODY, policy, httpx.Response(200, json={"input_tokens": 76}) + ) + + assert sent == expected_sent + assert result is not None + assert (result.total_tokens, result.model_used, result.error) == (76, "gpt-4o", False) + + +@pytest.mark.asyncio +async def test_native_count_restores_provider_ids_on_a_copy_of_the_callers_body( + monkeypatch: pytest.MonkeyPatch, +) -> None: + wrapped: Final = ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id("gAAAAoriginal", "deployment-1") + request: Final[dict[str, object]] = { + "model": "my-gpt", + "previous_response_id": ResponsesAPIRequestUtils._build_responses_api_response_id( + "openai", "deployment-1", "resp_original" + ), + "input": [ + {"type": "reasoning", "summary": [], "encrypted_content": wrapped}, + {"role": "user", "content": "again"}, + ], + } + before: Final = copy.deepcopy(request) + + _, sent = await _native_count(monkeypatch, request, _KEEP_ALL, httpx.Response(200, json={"input_tokens": 1})) + + assert sent == { + "model": "gpt-4o", + "previous_response_id": "resp_original", + "input": [ + {"type": "reasoning", "summary": [], "encrypted_content": "gAAAAoriginal"}, + {"role": "user", "content": "again"}, + ], + } + assert request == before + + +@pytest.mark.asyncio +async def test_native_count_answers_openais_rejection_with_its_status_and_text( + monkeypatch: pytest.MonkeyPatch, +) -> None: + rejection: Final = json.dumps({"error": {"message": "Unknown parameter: 'max_output_tokens'."}}) + + result, _ = await _native_count(monkeypatch, _CALLER_BODY, _KEEP_ALL, httpx.Response(400, text=rejection)) + + assert result is not None + assert (result.error, result.status_code, result.error_message) == (True, 400, rejection) diff --git a/tests/unit/llms/test_file_search_responses.py b/tests/unit/llms/test_file_search_responses.py index 9d2a5dbaf59..b2cbcd4a6dd 100644 --- a/tests/unit/llms/test_file_search_responses.py +++ b/tests/unit/llms/test_file_search_responses.py @@ -573,11 +573,11 @@ class TestEmulatedFileSearchHandler: def test_H1_file_search_replaced_with_function_tool(self): from litellm.responses.file_search.emulated_handler import ( - _replace_file_search_tools, + replace_file_search_tools, ) tools = [{"type": "file_search", "vector_store_ids": ["vs_abc", "vs_def"]}] - new_tools, vs_ids = _replace_file_search_tools(tools) + new_tools, vs_ids = replace_file_search_tools(tools) assert vs_ids == ["vs_abc", "vs_def"] assert len(new_tools) == 1 @@ -590,14 +590,14 @@ class TestEmulatedFileSearchHandler: def test_H2_non_file_search_tools_preserved(self): from litellm.responses.file_search.emulated_handler import ( - _replace_file_search_tools, + replace_file_search_tools, ) tools = [ {"type": "web_search"}, {"type": "file_search", "vector_store_ids": ["vs_abc"]}, ] - new_tools, vs_ids = _replace_file_search_tools(tools) + new_tools, vs_ids = replace_file_search_tools(tools) assert len(new_tools) == 2 # web_search + generated function tool assert new_tools[0]["type"] == "web_search" @@ -605,22 +605,22 @@ class TestEmulatedFileSearchHandler: def test_H3_no_file_search_tools_returns_unchanged(self): from litellm.responses.file_search.emulated_handler import ( - _replace_file_search_tools, + replace_file_search_tools, ) tools = [{"type": "web_search"}] - new_tools, vs_ids = _replace_file_search_tools(tools) + new_tools, vs_ids = replace_file_search_tools(tools) assert vs_ids == [] assert new_tools == [{"type": "web_search"}] def test_H4_empty_vector_store_ids_no_function_tool(self): from litellm.responses.file_search.emulated_handler import ( - _replace_file_search_tools, + replace_file_search_tools, ) tools = [{"type": "file_search", "vector_store_ids": []}] - new_tools, vs_ids = _replace_file_search_tools(tools) + new_tools, vs_ids = replace_file_search_tools(tools) assert vs_ids == [] assert new_tools == [] # no function tool added without store IDs diff --git a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py index 7d380205ca1..38737df7545 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -28,6 +28,7 @@ from litellm.llms.vertex_ai.common_utils import ( vertex_request_labels_from_litellm_params, ) from litellm.llms.vertex_ai.count_tokens import handler as vertex_count_tokens_handler +from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import VertexAIAnthropicConfig from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat @@ -64,7 +65,9 @@ async def test_vertex_gemini_rejects_non_text_system_blocks_before_request( "google/gemma-4-26b-a4b-it-maas", ), ) -async def test_vertex_partner_model_without_an_inference_matching_count_answers_400_naming_the_model(model: str) -> None: +async def test_vertex_partner_model_without_an_inference_matching_count_answers_400_naming_the_model( + model: str, +) -> None: result: Final = await VertexAITokenCounter().count_tokens( model_to_use=model, messages=[{"role": "user", "content": "Count this"}], @@ -488,7 +491,6 @@ def test_anyof_conversion_forwards_a_non_object_branch_for_the_provider_to_judge assert schema == {"type": "object", "properties": {"example": {"anyOf": expected}}} - @pytest.mark.parametrize( ("field", "expected"), [ @@ -534,6 +536,7 @@ def test_vertex_schema_folds_every_null_branch_into_nullable(field: JsonValue, e assert schema["properties"]["value"] == expected + def test_anyof_with_excessive_nesting(): """Test conversion with excessive nesting > max levels +1 deep.""" # generate a schema with excessive nesting @@ -1700,6 +1703,148 @@ async def test_vertex_ai_token_counter_forwards_system_and_tools_to_partner_requ assert without_optional_fields.total_tokens == 37 +_PARTNER_CHAT_TOOL_TURNS: Final[list[dict[str, object]]] = [ + {"role": "user", "content": "Weather in SF?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "sunny"}, +] +_PARTNER_RESPONSES_TOOL: Final[dict[str, object]] = { + "type": "function", + "name": "get_weather", + "parameters": {"type": "object", "properties": {}}, +} + + +async def _partner_counted_bodies( + model: str, chat_params: dict[str, JsonValue] | None = None +) -> tuple[dict[str, object], ...]: + from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens import handler + + class FakeResponse: + status_code = 200 + + def json(self) -> dict[str, int]: + return {"input_tokens": 37} + + class FakeHttpClient: + posted_bodies: tuple[dict[str, object], ...] = () + + async def post( + self, url: str, headers: dict[str, str], json: dict[str, object], timeout: float + ) -> FakeResponse: + self.posted_bodies = (*self.posted_bodies, json) + return FakeResponse() + + fake_http_client: Final = FakeHttpClient() + with ( + patch.object(handler, "get_async_httpx_client", return_value=fake_http_client), + patch.object( + handler.VertexAIPartnerModelsTokenCounter, + "_ensure_access_token_async", + new=AsyncMock(return_value=("fake-token", "test-project")), + ), + ): + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use=model, + messages=copy.deepcopy(_PARTNER_CHAT_TOOL_TURNS), + contents=None, + deployment={"litellm_params": {"vertex_project": "test-project", "vertex_location": "us-east5"}}, + tools=[_PARTNER_RESPONSES_TOOL], + chat_params=chat_params or {}, + ) + + assert result is not None + assert result.total_tokens == 37 + return fake_http_client.posted_bodies + + +@pytest.mark.asyncio +async def test_vertex_claude_count_sends_chat_tool_use_in_the_shape_chat_inference_sends() -> None: + assert await _partner_counted_bodies("claude-opus-5-5") == ( + { + "model": "claude-opus-5-5", + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "Weather in SF?"}]}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "call_1", "name": "get_weather", "input": {}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_1", "content": "sunny"}]}, + ], + "tools": [ + { + "name": "get_weather", + "input_schema": {"type": "object", "properties": {}}, + "type": "custom", + "description": "", + } + ], + }, + ) + + +@pytest.mark.asyncio +async def test_vertex_claude_count_sends_the_structured_output_vertex_inference_sends() -> None: + model: Final = "claude-haiku-4-5" + chat_params: Final[dict[str, JsonValue]] = { + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "answer", + "schema": {"type": "object", "properties": {"answer": {"type": "string"}}, "required": ["answer"]}, + }, + } + } + config: Final = VertexAIAnthropicConfig() + inference: Final = config.transform_request( + model=model, + messages=copy.deepcopy(_PARTNER_CHAT_TOOL_TURNS), + optional_params=config.map_openai_params( + { + **chat_params, + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "", + "parameters": _PARTNER_RESPONSES_TOOL["parameters"], + }, + } + ], + }, + {}, + model, + drop_params=False, + ), + litellm_params={}, + headers={}, + ) + + (body,) = await _partner_counted_bodies(model, chat_params) + + assert (body["tools"], body.get("tool_choice"), body.get("output_config")) == ( + inference["tools"], + inference.get("tool_choice"), + inference.get("output_config"), + ) + + +@pytest.mark.asyncio +async def test_vertex_llama_count_sends_chat_tool_use_as_sent() -> None: + assert await _partner_counted_bodies("meta/llama-4-maverick-17b-128e-instruct-maas") == ( + { + "model": "meta/llama-4-maverick-17b-128e-instruct-maas", + "messages": _PARTNER_CHAT_TOOL_TURNS, + "tools": [_PARTNER_RESPONSES_TOOL], + }, + ) + + @pytest.mark.asyncio @pytest.mark.parametrize( ("provider_failure", "expected_status", "expected_message"), diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index e3ce513dbcb..32b77c3b564 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -4171,10 +4171,17 @@ class TestResponsesInputTokens: path: str = "/v1/responses/input_tokens", counter: AsyncMock | None = None, llm_router: litellm.Router | None = None, + response_id_security: object | None = None, + managed_files: object | None = None, ) -> tuple[Response, AsyncMock]: from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth - from litellm.proxy.response_api_endpoints.endpoints import _proxy_llm_router, _proxy_token_counter + from litellm.proxy.response_api_endpoints.endpoints import ( + _proxy_llm_router, + _proxy_managed_files_hook, + _proxy_response_id_security, + _proxy_token_counter, + ) from litellm.types.utils import TokenCountResponse token_counter_mock = ( @@ -4192,6 +4199,10 @@ class TestResponsesInputTokens: app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(api_key="sk-test", request_route=path) app.dependency_overrides[_proxy_token_counter] = lambda: token_counter_mock app.dependency_overrides[_proxy_llm_router] = lambda: llm_router + if response_id_security is not None: + app.dependency_overrides[_proxy_response_id_security] = lambda: response_id_security + if managed_files is not None: + app.dependency_overrides[_proxy_managed_files_hook] = lambda: managed_files try: client = TestClient(app) response = client.post(path, json=body, headers={"Authorization": "Bearer sk-9876"}) @@ -4200,6 +4211,8 @@ class TestResponsesInputTokens: app.dependency_overrides.pop(user_api_key_auth, None) app.dependency_overrides.pop(_proxy_token_counter, None) app.dependency_overrides.pop(_proxy_llm_router, None) + app.dependency_overrides.pop(_proxy_response_id_security, None) + app.dependency_overrides.pop(_proxy_managed_files_hook, None) def test_string_input_returns_openai_input_tokens_shape(self): response, counter = self._post_input_tokens({"model": "gpt-4o", "input": "Hello, how are you?"}) @@ -4324,6 +4337,139 @@ class TestResponsesInputTokens: assert response.status_code == 200, response.text assert counter.await_args.kwargs["request"].tools == [{"type": "web_search"}] + assert counter.await_args.kwargs["responses_request"]["tools"] == [{"type": "web_search"}] + + @pytest.mark.parametrize( + ("tools", "expected"), + [ + ([5], {}), + ( + [5, {"type": "web_search"}], + {"tools": [{"type": "web_search"}], "tool_choice": "required", "parallel_tool_calls": False}, + ), + ], + ) + def test_tool_choice_leaves_the_count_with_the_tools_drop_params_drops( + self, tools: list[object], expected: dict[str, object] + ) -> None: + response, counter = self._post_input_tokens( + { + "model": "gpt-4o", + "input": "hi", + "tools": tools, + "tool_choice": "required", + "parallel_tool_calls": False, + "drop_params": True, + } + ) + + assert response.status_code == 200, response.text + sent = counter.await_args.kwargs["responses_request"] + assert {name: sent[name] for name in ("tools", "tool_choice", "parallel_tool_calls") if name in sent} == expected + + def test_previous_response_id_reaches_the_counter_authorized_next_to_the_id_the_client_held(self) -> None: + from litellm.types.utils import ADDRESSED_RESPONSE_ID_FIELD + + class DecodingSecurity: + def authorize_response_id(self, response_id: str, user_api_key_dict: object) -> str: + return {"resp_client_form": "resp_raw"}[response_id] + + response, counter = self._post_input_tokens( + {"model": "gpt-4o", "input": "hi", "previous_response_id": "resp_client_form"}, + response_id_security=DecodingSecurity(), + ) + + assert response.status_code == 200, response.text + responses_request: Final = counter.await_args.kwargs["responses_request"] + assert (responses_request["previous_response_id"], responses_request[ADDRESSED_RESPONSE_ID_FIELD]) == ( + "resp_raw", + "resp_client_form", + ) + + def test_managed_file_ids_reach_the_counter_with_the_mapping_inference_uses(self) -> None: + from litellm.integrations.custom_logger import CustomLogger + + mapping: Final = {"litellm-managed-file": {"deployment-a": "file-on-a"}} + + class ManagedFiles(CustomLogger): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + if call_type == "aresponses": + data["model_file_id_mapping"] = mapping + return None + + response, counter = self._post_input_tokens( + { + "model": "gpt-4o", + "input": [{"role": "user", "content": [{"type": "input_file", "file_id": "litellm-managed-file"}]}], + }, + managed_files=ManagedFiles(), + ) + + assert response.status_code == 200, response.text + assert counter.await_args.kwargs["responses_request"]["model_file_id_mapping"] == mapping + + @pytest.mark.parametrize("caller_mapping", [5, "lots", [1], {"file-x": {"deployment-a": 7}}]) + def test_a_caller_sent_model_file_id_mapping_does_not_break_the_count( + self, monkeypatch: pytest.MonkeyPatch, caller_mapping: object + ) -> None: + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + from litellm.llms.openai.responses.count_tokens import handler as openai_count_tokens_http + from litellm.proxy.proxy_server import count_request_tokens + + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-files", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "test-key", + "api_base": "https://count.test/v1", + }, + "model_info": {"id": "deployment-a"}, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + sent: Final[list[object]] = [] # mutable-ok: collects what the mock peer receives + + def count_peer(request: httpx.Request) -> httpx.Response: + sent.append(json.loads(request.content)) + return httpx.Response(200, json={"object": "response.input_tokens", "input_tokens": 7}, request=request) + + with patch.object( + openai_count_tokens_http, + "get_async_httpx_client", + return_value=AsyncHTTPHandler(transport=httpx.MockTransport(count_peer)), + ): + response, _ = self._post_input_tokens( + {"model": "gpt-files", "input": "hi", "model_file_id_mapping": caller_mapping}, + counter=AsyncMock(side_effect=count_request_tokens), + llm_router=router, + ) + + assert (response.status_code, response.json().get("input_tokens")) == (200, 7), response.text + assert sent == [{"model": "gpt-4o-mini", "input": "hi"}] + + def test_a_managed_file_the_key_cannot_read_answers_403_without_a_count(self) -> None: + from litellm.integrations.custom_logger import CustomLogger + + denied: Final = "User u1 does not have access to the file litellm-managed-file" + + class ManagedFiles(CustomLogger): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + raise HTTPException(status_code=403, detail=denied) + + response, counter = self._post_input_tokens( + { + "model": "gpt-4o", + "input": [{"role": "user", "content": [{"type": "input_file", "file_id": "litellm-managed-file"}]}], + }, + managed_files=ManagedFiles(), + ) + + assert response.status_code == 403, response.text + assert response.json()["error"]["message"] == denied + counter.assert_not_awaited() def test_malformed_tools_follow_the_counted_deployment_even_when_fallbacks_are_configured(self) -> None: router: Final = litellm.Router( @@ -4345,6 +4491,7 @@ class TestResponsesInputTokens: assert response.status_code == 200, response.text assert counter.await_args.kwargs["request"].tools is None + assert "tools" not in counter.await_args.kwargs["responses_request"] def test_malformed_tools_for_a_model_outside_the_router_get_a_clean_400(self) -> None: router: Final = litellm.Router( diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 714d8bb77d7..46e88614bb6 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -57,8 +57,10 @@ from litellm.llms.bedrock.count_tokens.kimi_k3 import ( ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler +from litellm.llms.openai.responses.count_tokens import handler as openai_count_tokens_http from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import VertexAIPartnerModelsTokenCounter from litellm.proxy._types import ProxyException, TokenCountRequest, UserAPIKeyAuth +from litellm.responses.litellm_completion_transformation.session_handler import ResponsesSessionHandler from litellm.proxy.anthropic_endpoints.endpoints import ( count_tokens as anthropic_count_tokens, ) @@ -77,9 +79,16 @@ from litellm.proxy.proxy_server import ( count_request_tokens, token_counter, ) +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem -from litellm.types.utils import SelectTokenizerResponse, TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import ( + ADDRESSED_RESPONSE_ID_FIELD, + SelectTokenizerResponse, + SpecialEnums, + TokenCountMessagesFormat, + TokenCountResponse, +) def _hugging_face_offline() -> NoReturn: @@ -270,6 +279,50 @@ async def test_token_counter_reads_an_oidc_api_key_only_when_it_calls_the_provid assert str(tmp_path) not in excinfo.value.message +@pytest.mark.asyncio +async def test_native_openai_responses_count_sends_the_token_an_oidc_api_key_points_to( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path)) + token_file: Final = tmp_path / "token" + token_file.write_text("sa-token-from-file") + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "gpt-oidc", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": f"oidc/file/{token_file}", + "api_base": "https://count.test/v1", + }, + } + ] + ), + ) + authorizations: Final[list[str | None]] = [] # mutable-ok: collects what the mock peer receives + + def count_peer(request: httpx.Request) -> httpx.Response: + authorizations.append(request.headers.get("authorization")) + return httpx.Response(200, json={"object": "response.input_tokens", "input_tokens": 9}, request=request) + + with patch.object( + openai_count_tokens_http, + "get_async_httpx_client", + return_value=AsyncHTTPHandler(transport=httpx.MockTransport(count_peer)), + ): + result: Final = await count_request_tokens( + request=TokenCountRequest(model="gpt-oidc", messages=[{"role": "user", "content": "hello"}]), + call_endpoint=True, + responses_request={"model": "gpt-oidc", "input": "hello"}, + ) + + assert (result.error, result.total_tokens) == (False, 9) + assert authorizations == ["Bearer sa-token-from-file"] + + @pytest.mark.asyncio async def test_anthropic_messages_count_tokens_endpoint(): """ @@ -1567,6 +1620,7 @@ class _ScriptedCounter(BaseTokenCounter): messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, context_management: object = None, + chat_params: object = None, ) -> TokenCountResponse | None: self.calls.append((model_to_use, messages, tools, system)) self.regions.append(None if deployment is None else deployment["litellm_params"].get("aws_region_name")) @@ -3505,3 +3559,325 @@ async def test_call_endpoint_prompt_keeps_the_local_count_on_a_counter_that_need local.total_tokens, local.tokenizer_type, ) + + +_PINNED_COUNTS: Final = MappingProxyType({"deployment-a.test": 11, "deployment-b.test": 22}) + + +def _pinned_count_peer(request: httpx.Request) -> httpx.Response: + if json.loads(request.content).get("previous_response_id") != "resp_original": + return httpx.Response(400, json={"error": {"message": "previous_response_id was not decoded"}}, request=request) + return httpx.Response(200, json={"input_tokens": _PINNED_COUNTS[request.url.host]}, request=request) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("pinned", ["deployment-a", "deployment-b"]) +async def test_responses_count_goes_to_the_deployment_that_served_the_previous_response( + monkeypatch: pytest.MonkeyPatch, pinned: str +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "gpt-group", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "test-key", + "api_base": f"https://{identity}.test/v1", + }, + "model_info": {"id": identity}, + } + for identity in ("deployment-a", "deployment-b") + ], + optional_pre_call_checks=["responses_api_deployment_check"], + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_pinned_count_peer)) + previous_response_id: Final = ResponsesAPIRequestUtils._build_responses_api_response_id( + "openai", pinned, "resp_original" + ) + + with patch.object(openai_count_tokens_http, "get_async_httpx_client", return_value=client): + counts: Final = [ + ( + await count_request_tokens( + request=TokenCountRequest(model="gpt-group", messages=[{"role": "user", "content": "again"}]), + call_endpoint=True, + responses_request={ + "model": "gpt-group", + "input": "again", + "previous_response_id": previous_response_id, + }, + ) + ).total_tokens + for _ in range(8) + ] + + assert counts == [_PINNED_COUNTS[f"{pinned}.test"]] * 8 + + +def _managed_file_id(provider_file_id: str, model_id: str) -> str: + unified: Final = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/pdf", "3f0c9a", "gpt-files", provider_file_id, model_id + ) + return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("with_mapping", "expected_file_id"), [(False, "file-uploaded"), (True, "file-on-a")]) +async def test_responses_count_sends_openai_the_provider_file_id_inference_sends( + monkeypatch: pytest.MonkeyPatch, with_mapping: bool, expected_file_id: str +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "gpt-files", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key", "api_base": "https://count.test/v1"}, + "model_info": {"id": "deployment-a"}, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + managed_id: Final = _managed_file_id("file-uploaded", "deployment-a") + bodies: Final[list[dict[str, JsonValue]]] = [] # mutable-ok: collects what the mock peer receives + + def count_peer(request: httpx.Request) -> httpx.Response: + bodies.append(json.loads(request.content)) + return httpx.Response(200, json={"object": "response.input_tokens", "input_tokens": 49}, request=request) + + with patch.object( + openai_count_tokens_http, + "get_async_httpx_client", + return_value=AsyncHTTPHandler(transport=httpx.MockTransport(count_peer)), + ): + result: Final = await count_request_tokens( + request=TokenCountRequest(model="gpt-files", messages=[{"role": "user", "content": "Summarize"}]), + call_endpoint=True, + responses_request={ + "model": "gpt-files", + "input": [{"role": "user", "content": [{"type": "input_file", "file_id": managed_id}]}], + **({"model_file_id_mapping": {managed_id: {"deployment-a": "file-on-a"}}} if with_mapping else {}), + }, + ) + + assert (result.error, result.total_tokens) == (False, 49) + (body,) = bodies + assert body["input"] == [{"role": "user", "content": [{"type": "input_file", "file_id": expected_file_id}]}] + assert "model_file_id_mapping" not in body + + +_ANSWER_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"answer": {"type": "string"}}, + "required": ["answer"], + "additionalProperties": False, +} + + +@pytest.mark.asyncio +async def test_responses_count_on_claude_counts_the_chat_request_responses_inference_sends( + monkeypatch: pytest.MonkeyPatch, +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "claude-counted", + "litellm_params": { + "model": "anthropic/claude-opus-5-5", + "api_key": "test-key", + "api_base": "https://count-tokens.test", + }, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + bodies: Final[list[dict[str, JsonValue]]] = [] # mutable-ok: collects what the mock peer receives + + def count_peer(request: httpx.Request) -> httpx.Response: + bodies.append(json.loads(request.content)) + return httpx.Response(200, json={"input_tokens": 77}, request=request) + + responses_tools: Final[list[dict[str, JsonValue]]] = [ + {"type": "function", "name": "get_weather", "parameters": {"type": "object", "properties": {}}}, + {"type": "web_search"}, + {"type": "file_search", "vector_store_ids": ["vs_docs"]}, + ] + with patch.object( + anthropic_count_tokens_handler, + "get_async_httpx_client", + return_value=AsyncHTTPHandler(transport=httpx.MockTransport(count_peer)), + ): + result: Final = await count_request_tokens( + request=TokenCountRequest( + model="claude-counted", messages=[{"role": "user", "content": "Weather?"}], tools=responses_tools + ), + call_endpoint=True, + responses_request={ + "model": "claude-counted", + "input": "Weather?", + "tools": responses_tools, + "tool_choice": "auto", + "text": {"format": {"type": "json_schema", "name": "answer", "schema": _ANSWER_SCHEMA}}, + "reasoning": {"effort": "low"}, + }, + ) + + assert (result.error, result.total_tokens) == (False, 77) + (body,) = bodies + tools: Final = body["tools"] + assert isinstance(tools, list) + tool_names: Final = [tool["name"] for tool in tools if isinstance(tool, dict)] + assert sorted(tool_names) == ["get_weather", "litellm_file_search", "web_search"] + assert any(isinstance(tool, dict) and str(tool.get("type")).startswith("web_search") for tool in tools) + assert body["tool_choice"] == {"type": "auto"} + output_config: Final = body["output_config"] + assert isinstance(output_config, dict) + assert output_config["format"] == {"type": "json_schema", "schema": _ANSWER_SCHEMA} + assert "thinking" in body + + +_WEATHER_TOOL: Final[dict[str, JsonValue]] = { + "type": "function", + "name": "get_weather", + "parameters": {"type": "object", "properties": {}}, +} + + +async def _claude_count( + monkeypatch: pytest.MonkeyPatch, + deployment_params: Mapping[str, JsonValue], + responses_request: Mapping[str, object], +) -> tuple[TokenCountResponse, list[dict[str, JsonValue]]]: + router: Final = Router( + model_list=[ + { + "model_name": "claude-counted", + "litellm_params": { + "model": "anthropic/claude-zeta-9", + "api_key": "test-key", + "api_base": "https://count-tokens.test", + **deployment_params, + }, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + bodies: Final[list[dict[str, JsonValue]]] = [] # mutable-ok: collects what the mock peer receives + + def count_peer(request: httpx.Request) -> httpx.Response: + bodies.append(json.loads(request.content)) + return httpx.Response(200, json={"input_tokens": 77}, request=request) + + with patch.object( + anthropic_count_tokens_handler, + "get_async_httpx_client", + return_value=AsyncHTTPHandler(transport=httpx.MockTransport(count_peer)), + ): + result: Final = await count_request_tokens( + request=TokenCountRequest(model="claude-counted", messages=[{"role": "user", "content": "Weather?"}]), + call_endpoint=True, + responses_request={"model": "claude-counted", **responses_request}, + ) + return result, bodies + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("deployment_params", "request_params"), + [({"drop_params": True}, {}), ({}, {"drop_params": True})], + ids=["deployment", "request"], +) +async def test_responses_count_on_claude_downgrades_forced_tool_choice_under_drop_params( + monkeypatch: pytest.MonkeyPatch, + deployment_params: Mapping[str, JsonValue], + request_params: Mapping[str, JsonValue], +) -> None: + monkeypatch.setitem(litellm.model_cost, "claude-zeta-9", {"supports_forced_tool_use": False}) + + result, bodies = await _claude_count( + monkeypatch, + deployment_params, + {"input": "Weather?", "tools": [_WEATHER_TOOL], "tool_choice": "required", **request_params}, + ) + + assert (result.error, result.total_tokens) == (False, 77) + (body,) = bodies + assert body["tool_choice"] == {"type": "auto"} + + +@pytest.mark.asyncio +async def test_responses_count_on_claude_rejects_forced_tool_choice_without_drop_params( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setitem(litellm.model_cost, "claude-zeta-9", {"supports_forced_tool_use": False}) + + with pytest.raises(ProxyException) as rejected: + await _claude_count( + monkeypatch, {}, {"input": "Weather?", "tools": [_WEATHER_TOOL], "tool_choice": "required"} + ) + + assert (rejected.value.code, "forced tool use" in rejected.value.message) == ("400", True) + + +def _texts(message: JsonValue) -> list[str]: + assert isinstance(message, dict) + content: Final = message["content"] + if isinstance(content, str): + return [content] + assert isinstance(content, list) + return [str(block["text"]) for block in content if isinstance(block, dict) and "text" in block] + + +@pytest.mark.asyncio +async def test_responses_count_on_claude_counts_the_turns_previous_response_id_replays( + monkeypatch: pytest.MonkeyPatch, +) -> None: + stored_turn: Final = { + "request_id": "chatcmpl-stored", + "call_type": "aresponses", + "session_id": "session-1", + "proxy_server_request": {"input": "Who won in 1998?", "model": "claude-counted"}, + "response": { + "id": "chatcmpl-stored", + "object": "chat.completion", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "The Bulls.", + "reasoning_content": "Chicago beat Utah in six games.", + }, + "finish_reason": "stop", + } + ], + }, + "status": "success", + } + + with patch.object( + ResponsesSessionHandler, + "get_all_spend_logs_for_previous_response_id", + new_callable=AsyncMock, + return_value=[stored_turn], + ) as spend_logs: + result, bodies = await _claude_count( + monkeypatch, + {}, + { + "input": "And in 1999?", + "previous_response_id": "chatcmpl-stored", + ADDRESSED_RESPONSE_ID_FIELD: "resp_client_form", + }, + ) + + assert (result.error, result.total_tokens) == (False, 77) + (body,) = bodies + messages: Final = body["messages"] + assert isinstance(messages, list) + assert list(itertools.chain.from_iterable(map(_texts, messages))) == [ + "Who won in 1998?", + "The Bulls.", + "And in 1999?", + ] + assert [sorted(message) for message in messages if isinstance(message, dict)] == [["content", "role"]] * 3 + spend_logs.assert_awaited_once_with("chatcmpl-stored", ("chatcmpl-stored", "resp_client_form")) diff --git a/tests/unit/responses/litellm_completion_transformation/test_count_request.py b/tests/unit/responses/litellm_completion_transformation/test_count_request.py new file mode 100644 index 00000000000..464bdb1a367 --- /dev/null +++ b/tests/unit/responses/litellm_completion_transformation/test_count_request.py @@ -0,0 +1,127 @@ +import json +from typing import Final + +import pytest +from pydantic import JsonValue + +from litellm.litellm_core_utils.tool_fields import DropPolicy +from litellm.responses.litellm_completion_transformation.count_request import ( + ChatCountBridgeError, + ChatCountRequest, + chat_count_request, +) + +_MARK: Final[dict[str, JsonValue]] = {"type": "ephemeral"} +_MARKED_RESPONSES_REQUEST: Final[dict[str, JsonValue]] = { + "model": "claude-opus-5-5", + "instructions": "Be terse.", + "input": [ + {"type": "message", "role": "user", "content": "Weather in SF?", "cache_control": _MARK}, + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "And in LA?", "cache_control": _MARK}], + }, + { + "type": "function_call", + "call_id": "call_1", + "name": "get_weather", + "arguments": '{"city": "SF"}', + "cache_control": _MARK, + }, + {"type": "function_call_output", "call_id": "call_1", "output": "sunny", "cache_control": _MARK}, + ], + "tools": [ + { + "type": "function", + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + "cache_control": _MARK, + } + ], +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("drop_params", (False, True)) +async def test_claude_count_request_leaves_out_every_cache_control_mark(drop_params: bool) -> None: + request: Final = await chat_count_request( + model="claude-opus-5-5", + custom_llm_provider="anthropic", + api_base=None, + responses_request=_MARKED_RESPONSES_REQUEST, + policy=DropPolicy(drop_params=drop_params, additional_drop_params=frozenset()), + ) + + assert isinstance(request, ChatCountRequest) + counted: Final = json.dumps({"messages": request.messages, "tools": request.tools}, default=str) + assert "cache_control" not in counted + for kept in ("Weather in SF?", "And in LA?", "call_1", "sunny", "get_weather"): + assert kept in counted + + +_WEATHER_TOOL: Final[dict[str, JsonValue]] = { + "type": "function", + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, +} +_QUESTION: Final[dict[str, JsonValue]] = {"role": "user", "content": "Weather in SF?"} + + +async def _claude_count( + responses_request: dict[str, JsonValue], policy: DropPolicy +) -> ChatCountRequest | ChatCountBridgeError | None: + return await chat_count_request( + model="claude-opus-5-5", + custom_llm_provider="anthropic", + api_base=None, + responses_request={"model": "claude-opus-5-5", **responses_request}, + policy=policy, + ) + + +@pytest.mark.asyncio +async def test_claude_count_request_counts_additional_tools_input_items_as_tools_like_inference() -> None: + policy: Final = DropPolicy(drop_params=False, additional_drop_params=frozenset()) + + hoisted: Final = await _claude_count( + {"input": [{"type": "additional_tools", "tools": [_WEATHER_TOOL]}, _QUESTION]}, policy + ) + + assert isinstance(hoisted, ChatCountRequest) + assert hoisted.tools is not None + assert hoisted == await _claude_count({"input": [_QUESTION], "tools": [_WEATHER_TOOL]}, policy) + + +@pytest.mark.asyncio +async def test_claude_count_request_drops_a_malformed_tool_flag_under_drop_params_like_inference() -> None: + policy: Final = DropPolicy(drop_params=True, additional_drop_params=frozenset()) + + flagged: Final = await _claude_count( + {"input": [_QUESTION], "tools": [{**_WEATHER_TOOL, "defer_loading": "yes"}]}, policy + ) + + assert isinstance(flagged, ChatCountRequest) + assert flagged == await _claude_count({"input": [_QUESTION], "tools": [_WEATHER_TOOL]}, policy) + + +@pytest.mark.asyncio +async def test_claude_count_request_rejects_a_malformed_tool_flag_without_drop_params() -> None: + flagged: Final = await _claude_count( + {"input": [_QUESTION], "tools": [{**_WEATHER_TOOL, "defer_loading": "yes"}]}, + DropPolicy(drop_params=False, additional_drop_params=frozenset()), + ) + + assert isinstance(flagged, ChatCountBridgeError) + assert "defer_loading" in flagged.message + + +@pytest.mark.asyncio +async def test_claude_count_request_leaves_out_tools_the_deployment_drops() -> None: + policy: Final = DropPolicy(drop_params=False, additional_drop_params=frozenset({"tools"})) + + dropped: Final = await _claude_count({"input": [_QUESTION], "tools": [_WEATHER_TOOL]}, policy) + + assert isinstance(dropped, ChatCountRequest) + assert dropped.tools is None + assert dropped == await _claude_count({"input": [_QUESTION]}, policy) diff --git a/tests/unit/test_count_tokens_public_api.py b/tests/unit/test_count_tokens_public_api.py index e33bc1e5907..4eb8bea3427 100644 --- a/tests/unit/test_count_tokens_public_api.py +++ b/tests/unit/test_count_tokens_public_api.py @@ -3,9 +3,13 @@ Tests for litellm.acount_tokens() public API. """ import asyncio +import json import os from unittest.mock import AsyncMock, patch +import httpx +import pytest +import respx import litellm from litellm.types.utils import TokenCountResponse @@ -157,3 +161,45 @@ def test_acount_tokens_no_api_key_falls_back(monkeypatch): assert result.tokenizer_type == "local_tokenizer" + + +@pytest.fixture +def httpx_transport_clients(monkeypatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + for name in ("ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"): + monkeypatch.delenv(name, raising=False) + litellm.in_memory_llm_clients_cache.flush_cache() + yield + litellm.in_memory_llm_clients_cache.flush_cache() + + +def test_acount_tokens_sends_chat_tool_use_to_anthropic_in_the_shape_claude_inference_sends(httpx_transport_clients): + schema = {"type": "object", "properties": {"city": {"type": "string"}}} + with respx.mock(assert_all_called=True) as router: + route = router.post("https://api.anthropic.com/v1/messages/count_tokens").mock( + return_value=httpx.Response(200, json={"input_tokens": 52}) + ) + result = asyncio.run( + litellm.acount_tokens( + model="anthropic/claude-opus-5-5", + messages=[ + {"role": "user", "content": "Weather in SF?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": "{}"}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "sunny"}, + ], + tools=[{"type": "function", "function": {"name": "get_weather", "parameters": schema}}], + api_key="sk-ant-test-key", + ) + ) + + assert (result.error, result.total_tokens) == (False, 52) + sent = json.loads(route.calls.last.request.content) + assert sent["tools"] == [{"name": "get_weather", "input_schema": schema, "type": "custom"}] + assert [message["role"] for message in sent["messages"]] == ["user", "assistant", "user"] + assert sent["messages"][2]["content"] == [{"type": "tool_result", "tool_use_id": "call_1", "content": "sunny"}] From 02199b670bbbda106e6a35afe9263ca6bfb91f6a Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Tue, 20 Oct 2026 08:21:49 -0700 Subject: [PATCH 17/33] fix(proxy): share the billing logger across callback registries (#47510) * fix(logging): deduplicate equivalent completed success callbacks * fix(proxy): share the billing logger across callback registries * test(proxy): use fixed timestamps for billing regressions --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 5 +- .../test_litellm_logging.py | 72 +++++++++++++++---- .../unit/proxy/proxy_server/test_lifecycle.py | 72 +++++++++++++------ 3 files changed, 112 insertions(+), 37 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f87b1edc2e5..2c6d2e2736a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3118,8 +3118,9 @@ def cost_tracking(): from litellm.integrations.shadow_eval_logger import ShadowEvalLogger spend_event_producer = build_spend_event_producer(CollectorSettings(), fallback=run_spend_event) - litellm.logging_callback_manager.add_litellm_callback(ProxyDBLogger(spend_event_producer)) - litellm.logging_callback_manager.add_litellm_async_success_callback(ProxyDBLogger(spend_event_producer)) + db_logger: Final = ProxyDBLogger(spend_event_producer) + litellm.logging_callback_manager.add_litellm_callback(db_logger) + litellm.logging_callback_manager.add_litellm_async_success_callback(db_logger) litellm.logging_callback_manager.add_litellm_callback(ShadowEvalLogger()) diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 7aac43298d4..981c43a9792 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -2462,7 +2462,41 @@ def test_success_handler_runs_sync_callbacks_for_sync_requests( @pytest.mark.asyncio -@pytest.mark.parametrize("registration", ["sync", "async", "overlapping"]) +@pytest.mark.parametrize("different_headers", [False, True]) +async def test_completed_success_preserves_generic_logger_configuration( + monkeypatch: pytest.MonkeyPatch, different_headers: bool +) -> None: + from litellm.integrations.generic_api.generic_api_callback import GenericAPILogger + + first: Final = GenericAPILogger( + endpoint="https://logs.example.test/events", + headers={"Authorization": "Bearer first"}, + event_types=["llm_api_success"], + ) + second: Final = GenericAPILogger( + endpoint="https://logs.example.test/events", + headers={"Authorization": "Bearer second" if different_headers else "Bearer first"}, + event_types=["llm_api_success"] if different_headers else ["llm_api_failure"], + ) + monkeypatch.setattr(litellm, "success_callback", [first]) + monkeypatch.setattr(litellm, "_async_success_callback", [second]) + monkeypatch.setattr(litellm, "generic_api_use_v1", False) + logging_obj: Final = _async_request_logging_obj("distinct-generic-loggers", []) + response: Final = ModelResponse(id="chatcmpl-distinct-generic-loggers") + now: Final = datetime.datetime(2026, 1, 1) + + await logging_obj.async_success_handler(response, now, now) + + assert len(first.log_queue) == 1 + assert first.log_queue[0]["id"] == response.id + if different_headers: + assert second.log_queue == first.log_queue + else: + assert second.log_queue == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("registration", ["sync", "async", "overlapping", "equivalent", "distinct"]) @pytest.mark.parametrize("kind", ["sync", "async", "dual"]) @pytest.mark.parametrize("compatibility_first", [False, True]) async def test_completed_success_selects_one_implemented_hook( @@ -2517,9 +2551,15 @@ async def test_completed_success_selects_one_implemented_hook( ) -> None: record("async", kwargs, response_obj) - callback: Final = {"sync": SyncRecorder, "async": AsyncRecorder, "dual": DualRecorder}[kind]() + recorder: Final = {"sync": SyncRecorder, "async": AsyncRecorder, "dual": DualRecorder}[kind] + callback: Final = recorder() + async_callback: Final = ( + recorder(turn_off_message_logging=registration == "distinct") + if registration in ("equivalent", "distinct") + else callback + ) monkeypatch.setattr(litellm, "success_callback", [callback] if registration != "async" else []) - monkeypatch.setattr(litellm, "_async_success_callback", [callback] if registration != "sync" else []) + monkeypatch.setattr(litellm, "_async_success_callback", [async_callback] if registration != "sync" else []) logging_obj: Final = _async_request_logging_obj("completed-hook", []) response: Final = ModelResponse(id="chatcmpl-completed-hook") now: Final = datetime.datetime.now() @@ -2531,20 +2571,26 @@ async def test_completed_success_selects_one_implemented_hook( await logging_obj.async_success_handler(response, now, now) logging_obj.handle_sync_success_callbacks_for_async_calls(response, now, now) - assert events.qsize() == 1 - hook, response_id, cost, thread_id = events.get_nowait() - assert hook == ("sync" if kind == "sync" else "async") - assert response_id == response.id - assert cost == logging_obj.model_call_details["response_cost"] - assert (thread_id == loop_thread) is (kind != "sync") + expected_count: Final = 2 if registration in ("equivalent", "distinct") else 1 + assert events.qsize() == expected_count + for _ in range(expected_count): + hook, response_id, cost, thread_id = events.get_nowait() + assert hook == ("sync" if kind == "sync" else "async") + assert response_id == response.id + assert cost == logging_obj.model_call_details["response_cost"] + assert (thread_id == loop_thread) is (kind != "sync") -@pytest.mark.parametrize("for_async_request", [False, True]) +@pytest.mark.parametrize("for_async_request,split_registries", [(False, False), (True, False), (True, True)]) @pytest.mark.parametrize("private_first", [False, True]) @pytest.mark.parametrize("response_kind", ["responses", "completion"]) def test_sync_success_redacts_only_opted_out_loggers( - monkeypatch: pytest.MonkeyPatch, for_async_request: bool, private_first: bool, response_kind: str + monkeypatch: pytest.MonkeyPatch, + for_async_request: bool, + split_registries: bool, + private_first: bool, + response_kind: str, ) -> None: from queue import SimpleQueue @@ -2563,8 +2609,8 @@ def test_sync_success_redacts_only_opted_out_loggers( private_logger: Final = PayloadRecorder(turn_off_message_logging=True) public_logger: Final = PayloadRecorder() callbacks: Final = [private_logger, public_logger] if private_first else [public_logger, private_logger] - monkeypatch.setattr(litellm, "success_callback", callbacks) - monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "success_callback", callbacks[:1] if split_registries else callbacks) + monkeypatch.setattr(litellm, "_async_success_callback", callbacks[1:] if split_registries else []) monkeypatch.setattr(litellm, "turn_off_message_logging", False) logging_obj: Final = LitellmLogging( model="openai/gpt-6.1-sol", diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index 3914c337f59..11a44a0ec16 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/tests/unit/proxy/proxy_server/test_lifecycle.py @@ -697,33 +697,61 @@ def test_load_from_azure_key_vault_missing_uri_failure_is_swallowed(monkeypatch) # --------------------------------------------------------------------------- -def test_cost_tracking_adds_db_and_shadow_eval_callbacks_when_prisma_set(monkeypatch): +@pytest.mark.asyncio +@pytest.mark.parametrize("setup_count", [1, 2]) +async def test_cost_tracking_adds_db_and_shadow_eval_callbacks_when_prisma_set( + monkeypatch: pytest.MonkeyPatch, setup_count: int +) -> None: + from datetime import datetime + import litellm from litellm.integrations.shadow_eval_logger import ShadowEvalLogger + from litellm.proxy.spend_tracking.spend_event import SpendEventDecodeError, decode_spend_event + from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer + from litellm.types.utils import ModelResponse, Usage + from litellm.utils import Rules, function_setup - fake_prisma = MagicMock() - monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False) - monkeypatch.setattr(litellm, "callbacks", [], raising=False) - monkeypatch.setattr(litellm, "_async_success_callback", [], raising=False) + producer: Final = AsyncMock(spec=SpendEventProducer) + monkeypatch.setattr(ps, "prisma_client", MagicMock()) + monkeypatch.setattr(ps, "spend_event_producer", None) + monkeypatch.setattr(ps, "build_spend_event_producer", lambda settings, fallback: producer) + for name in ( + "callbacks", + "input_callback", + "success_callback", + "failure_callback", + "_async_success_callback", + "_async_failure_callback", + ): + monkeypatch.setattr(litellm, name, []) - before_callbacks = len(litellm.callbacks) - before_async = len(litellm._async_success_callback) + for _ in range(setup_count): + cost_tracking() - cost_tracking() - cost_tracking() - - observed = { - "added_to_callbacks": len(litellm.callbacks) - before_callbacks, - "added_to_async_success": len(litellm._async_success_callback) - before_async, - "shadow_eval_loggers": sum(isinstance(cb, ShadowEvalLogger) for cb in litellm.callbacks), - "prisma_was_set": True, - } - assert normalize(observed) == { - "added_to_callbacks": 2, - "added_to_async_success": 1, - "shadow_eval_loggers": 1, - "prisma_was_set": True, - } + assert len(litellm.callbacks) == 2 + assert sum(isinstance(cb, ShadowEvalLogger) for cb in litellm.callbacks) == 1 + assert len(litellm._async_success_callback) == 1 + now: Final = datetime(2026, 1, 1) + logging_obj, _ = function_setup( + "acompletion", + Rules(), + now, + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + litellm_call_id="single-accounting-event", + ) + logging_obj.update_environment_variables(litellm_params={"acompletion": True, "metadata": {}}, optional_params={}) + response: Final = ModelResponse( + id="chatcmpl-single-accounting-event", + model="gpt-4o-mini", + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) + await logging_obj.async_success_handler(response, now, now) + producer.publish.assert_awaited_once() + event: Final = decode_spend_event(producer.publish.call_args.args[0]) + assert not isinstance(event, SpendEventDecodeError) + assert event.response_cost == logging_obj.model_call_details["response_cost"] + assert event.response_cost > 0 def test_cost_tracking_no_op_when_prisma_missing(monkeypatch): From ce63f48ed2ee3a18b7860b72f292e90ebeb1a986 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 20 Oct 2026 08:22:06 -0700 Subject: [PATCH 18/33] fix(decisions): log non-UTF-8 bytes as base64 so the request reaches validation and gets a 400 (#47564) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/decisions/call.py | 2 +- tests/unit/decisions/test_call.py | 20 ++++++++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/litellm/decisions/call.py b/litellm/decisions/call.py index b31cb1f69db..b94b3380ad5 100644 --- a/litellm/decisions/call.py +++ b/litellm/decisions/call.py @@ -427,7 +427,7 @@ def _jsonable_fallback(value: object) -> object: def _logged_request_message(caller_request: object) -> ChatCompletionUserMessage: caller_fields: Final = _LOGGED_FIELDS_ADAPTER.validate_python( - to_jsonable_python(caller_request, fallback=_jsonable_fallback) + to_jsonable_python(caller_request, fallback=_jsonable_fallback, bytes_mode="base64") ) sent_fields: Final = json.dumps( dict((name, value) for name, value in caller_fields.items() if value not in (None, [], {})) diff --git a/tests/unit/decisions/test_call.py b/tests/unit/decisions/test_call.py index 03c135369b8..36a16dd5ad1 100644 --- a/tests/unit/decisions/test_call.py +++ b/tests/unit/decisions/test_call.py @@ -1,3 +1,4 @@ +import base64 import json from collections.abc import Mapping from datetime import datetime, timezone @@ -112,3 +113,22 @@ def test_decisions_request_is_logged_even_when_its_questions_fail_validation(): logged: Final = _logged_request(logging_obj) assert logged["questions"][0]["instructions"] == QUESTION_TEXT + + +def test_non_utf8_bytes_in_the_request_are_logged_and_rejected_with_a_400(): + logging_obj: Final = _logging_obj("adecisions") + raw_input: Final = b"\xff\xfe not utf-8" + + with pytest.raises(litellm.BadRequestError): + _prepare_decisions_call( + **CONNECTION, + decision_input=raw_input, + questions=[{"type": "predicate", "instructions": QUESTION_TEXT}], + audio=None, + videos=None, + kwargs={"litellm_logging_obj": logging_obj}, + ) + + logged: Final = _logged_request(logging_obj) + assert base64.urlsafe_b64decode(logged["input"]) == raw_input + assert logged["questions"][0]["instructions"] == QUESTION_TEXT From 0ede0e10e7ee694c46f14afd1606885d8c369342 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Tue, 20 Oct 2026 08:22:19 -0700 Subject: [PATCH 19/33] fix(batch): preserve recoverable costs and prevent poll starvation (#47190) * fix(batch): preserve recoverable costs and rotate polling pages * test(batch): type retirement failure regression * fix(batch): alternate new arrivals with retry polling * test(batch): defer polling row annotation on Python 3.12 * test(batch): reconcile fixtures with atomic metadata claims * refactor(batch): simplify reconciliation to one bounded scan * fix(batch): restore arrival priority during retry scans * fix(batch): preserve recoverable jobs on unrelated resource errors * fix: revisit recent batches before old retry backlog --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> --- .../proxy/common_utils/check_batch_cost.py | 128 +++++--- .../managed_object_attribution.py | 7 + .../spend/test_batch_poll_starvation.py | 230 ++++++++++++- .../common_utils/test_check_batch_cost.py | 304 +++++++++++++++--- 4 files changed, 557 insertions(+), 112 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 7449ec9882f..05e3e0d839d 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -4,6 +4,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t import asyncio import json +import re from collections.abc import Mapping, Sequence from dataclasses import replace as dataclasses_replace from datetime import datetime, timedelta, timezone @@ -39,7 +40,12 @@ if TYPE_CHECKING: from litellm.types.utils import LiteLLMBatch +def _error_names_resource(error: Exception, resource_id: str) -> bool: + return resource_id in re.findall(r"""[^\s"'`]+""", str(error)) + + CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost" +RECENT_BATCH_POLL_WINDOW: Final = timedelta(hours=48) PROVIDER_TERMINAL_BATCH_STATUSES: Final[Tuple[str, ...]] = ( "completed", @@ -109,6 +115,11 @@ class CheckBatchCost: # the guaranteed-failing primary query on every subsequent cycle. self._has_batch_processed_column: bool = True self.batch_processed_support_confirmed: bool = False + self._poll_cursor: tuple[datetime, str] | None = None + self._poll_before: datetime = datetime.now(timezone.utc) + self._arrival_cursor: tuple[datetime, str] | None = None + self._arrival_before: datetime = self._poll_before + self._poll_arrivals: bool = False self._line_item_stores: Final[set[asyncio.Task[None]]] = set() self._line_item_store_turn: Final = asyncio.Lock() self._poll_idle: Final = asyncio.Event() @@ -161,29 +172,9 @@ class CheckBatchCost: f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired" ) - if not self._has_batch_processed_column: - return - - # A row already in a terminal status is never rewritten by the sweep above, so - # without this it keeps a poll-page slot forever and starves newer batches. - retired: Final = await _managed_object_table(self.prisma_client).update_many( - where={ - "file_purpose": "batch", - "batch_processed": False, - "status": {"in": ["complete", "completed"]}, - "created_at": {"lt": cutoff}, - }, - data={"batch_processed": True}, - ) - if retired > 0: - verbose_proxy_logger.warning( - f"CheckBatchCost: gave up on {retired} completed managed objects older than " - f"{MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days that were never costed" - ) - async def _fallback_find_jobs(self) -> "Sequence[ManagedObjectRow]": """Query batch jobs without the batch_processed filter (for older schemas).""" - return await _managed_object_table(self.prisma_client).find_many( + return await self._find_jobs( where={ "file_purpose": "batch", "status": { @@ -197,10 +188,59 @@ class CheckBatchCost: ] }, }, - take=MAX_OBJECTS_PER_POLL_CYCLE, - order={"created_at": "asc"}, ) + async def _find_jobs(self, where: Mapping[str, object]) -> Sequence[ManagedObjectRow]: + arrivals: Final = self._poll_arrivals + page: Final = await self._find_job_page(where, arrivals=arrivals) + self._poll_arrivals = not arrivals if page else arrivals + if page: + return page + return await self._find_job_page(where, arrivals=not arrivals) + + async def _find_job_page(self, where: Mapping[str, object], *, arrivals: bool) -> Sequence[ManagedObjectRow]: + cursor: Final = self._arrival_cursor if arrivals else self._poll_cursor + if cursor is None: + if arrivals: + self._arrival_before = datetime.now(timezone.utc) + else: + self._poll_before = datetime.now(timezone.utc) - RECENT_BATCH_POLL_WINDOW + before: Final = self._arrival_before if arrivals else self._poll_before + page: Final = await _managed_object_table(self.prisma_client).find_many( + where={ + **where, + "created_at": { + "lte": before, + **({"gt": before - RECENT_BATCH_POLL_WINDOW} if arrivals else {}), + }, + **( + { + "OR": [ + {"created_at": {"gt": cursor[0]}}, + {"created_at": cursor[0], "id": {"gt": cursor[1]}}, + ] + } + if cursor is not None + else {} + ), + }, + take=MAX_OBJECTS_PER_POLL_CYCLE, + order=[{"created_at": "asc"}, {"id": "asc"}], + ) + next_cursor: Final = (page[-1].created_at, page[-1].id) if len(page) == MAX_OBJECTS_PER_POLL_CYCLE else None + if arrivals: + self._arrival_cursor = next_cursor + else: + self._poll_cursor = next_cursor + return page + + def _unprocessed_snapshot(self, job: ManagedObjectRow) -> Mapping[str, object]: + return { + "id": job.id, + "updated_at": job.updated_at, + **({"batch_processed": False} if self._has_batch_processed_column else {}), + } + async def _retire_job(self, job: "ManagedObjectRow", reason: str) -> None: """ Take a row that can never be costed out of the poll page. Leaving it selectable @@ -208,24 +248,19 @@ class CheckBatchCost: once enough such rows accumulate no newer batch is ever reached. Older schemas without batch_processed can only be excluded through the status filter. """ - data: Final = ( - {"batch_processed": True} - if self._has_batch_processed_column - else {"status": "stale_expired"} - ) + data: Final = {"batch_processed": True} if self._has_batch_processed_column else {"status": "stale_expired"} try: - await _managed_object_table(self.prisma_client).update( - where={"id": job.id}, + retired: Final = await _managed_object_table(self.prisma_client).update_many( + where=self._unprocessed_snapshot(job), data=data, ) except Exception as db_err: - verbose_proxy_logger.error( - f"CheckBatchCost: failed to retire uncostable job {job.id} ({reason}): {db_err}" - ) + verbose_proxy_logger.error(f"CheckBatchCost: failed to retire uncostable job {job.id} ({reason}): {db_err}") + return + if retired == 0: return verbose_proxy_logger.warning( - f"CheckBatchCost: job {job.id} can never be costed ({reason}), " - "so it will no longer be polled" + f"CheckBatchCost: job {job.id} can never be costed ({reason}), so it will no longer be polled" ) async def _claim_job_for_costing(self, job: "ManagedObjectRow", response: "LiteLLMBatch") -> bool: @@ -312,13 +347,13 @@ class CheckBatchCost: from litellm.exceptions import NotFoundError - return isinstance(error, (NotFoundError, openai.NotFoundError)) and batch_id in str(error) + return isinstance(error, (NotFoundError, openai.NotFoundError)) and _error_names_resource(error, batch_id) def _batch_deployment_exists(self, model_id: str) -> bool: """A 404 only proves the batch is gone when it came from the batch's own deployment. Once that deployment leaves the router, default fallbacks can silently send the retrieve to a provider that never saw the batch, so its - 404 must not retire the row; the staleness sweep bounds it instead.""" + 404 must not retire a potentially recoverable row.""" return self.llm_router.get_deployment(model_id=model_id) is not None @staticmethod @@ -333,7 +368,7 @@ class CheckBatchCost: if not output_file_id: return False - return isinstance(error, (NotFoundError, openai.NotFoundError)) and output_file_id in str(error) + return isinstance(error, (NotFoundError, openai.NotFoundError)) and _error_names_resource(error, output_file_id) async def _finalize_unbilled_terminal_job( self, job: "ManagedObjectRow", response: "LiteLLMBatch" @@ -355,18 +390,18 @@ class CheckBatchCost: db_batch_object=job, unified_batch_id=_is_base64_encoded_unified_file_id(job.unified_object_id), ) - update_data: Final[dict] = { + update_data: Final[dict[str, object]] = { "status": response.status, "file_object": response.model_dump_json(), **({"batch_processed": True} if self._has_batch_processed_column else {}), } - await _managed_object_table(self.prisma_client).update( - where={"id": job.id}, + finalized: Final = await _managed_object_table(self.prisma_client).update_many( + where=self._unprocessed_snapshot(job), data=update_data, ) - verbose_proxy_logger.info( - f"CheckBatchCost: marked job {job.id} as {response.status} in DB" - ) + if finalized == 0: + return False + verbose_proxy_logger.info(f"CheckBatchCost: marked job {job.id} as {response.status} in DB") return True except Exception as db_err: verbose_proxy_logger.error( @@ -1047,6 +1082,7 @@ class CheckBatchCost: """ try: from litellm.integrations.prometheus import PrometheusLogger + prom_logger = PrometheusLogger.get_instance() except Exception as e: verbose_proxy_logger.error(f"CheckBatchCost: could not get Prometheus logger: {e}") @@ -1071,7 +1107,7 @@ class CheckBatchCost: # endpoint may transition a batch to "complete" before # CheckBatchCost runs. The batch_processed=False filter # already prevents reprocessing finished batches. - jobs = await _managed_object_table(self.prisma_client).find_many( + jobs = await self._find_jobs( where={ "file_purpose": "batch", "batch_processed": False, @@ -1084,8 +1120,6 @@ class CheckBatchCost: ] }, }, - take=MAX_OBJECTS_PER_POLL_CYCLE, - order={"created_at": "asc"}, ) self.batch_processed_support_confirmed = True except Exception as query_err: diff --git a/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py b/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py index 16b1b0e973c..2e2c07fd544 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py @@ -9,6 +9,7 @@ from __future__ import annotations from collections.abc import Mapping from dataclasses import dataclass +from datetime import datetime from typing import TYPE_CHECKING, Final, Protocol, cast from pydantic import TypeAdapter @@ -27,6 +28,12 @@ if TYPE_CHECKING: class ManagedObjectRow(Protocol): + @property + def created_at(self) -> datetime: ... + + @property + def updated_at(self) -> datetime: ... + @property def id(self) -> str: ... diff --git a/tests/integration/spend/test_batch_poll_starvation.py b/tests/integration/spend/test_batch_poll_starvation.py index f68f10e3d8a..51f5f0871f5 100644 --- a/tests/integration/spend/test_batch_poll_starvation.py +++ b/tests/integration/spend/test_batch_poll_starvation.py @@ -1,14 +1,20 @@ import json import os +from datetime import datetime, timezone +from functools import partial from hashlib import sha256 +from itertools import chain from typing import Final +from uuid import uuid4 +import httpx import pytest from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value -from integration._support.database import read_rows -from integration._support.upstream import delete_scenario, register_scenario +from integration._support.database import read_rows, write_rows +from integration._support.prisma_clients import connected_proxy_prisma +from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, TextResponse -from pydantic import JsonValue +from pydantic import JsonValue, TypeAdapter from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE @@ -59,17 +65,17 @@ def _accepting_routes() -> dict[str, JsonResponse | TextResponse]: } -def _gone_at_provider_routes() -> RoutedResponse: +def _gone_at_provider_routes(*, retryable: bool = False) -> RoutedResponse: return RoutedResponse( content_type="application/x-routed", routes={ **_accepting_routes(), "GET /batches/batch-$REQUEST_ID": JsonResponse( content_type="application/json", - status=404, + status=503 if retryable else 404, body={ "error": { - "message": "No batch found with id 'batch-$REQUEST_ID'.", + "message": "Temporary failure" if retryable else "No batch found with id 'batch-$REQUEST_ID'.", "type": "invalid_request_error", "param": "id", "code": "batch_not_found", @@ -80,7 +86,7 @@ def _gone_at_provider_routes() -> RoutedResponse: ) -def _completed_routes() -> RoutedResponse: +def _completed_routes(*, ready: bool = True) -> RoutedResponse: output_line: Final = { "id": "batch_req_1", "custom_id": "r1", @@ -106,7 +112,8 @@ def _completed_routes() -> RoutedResponse: routes={ **_accepting_routes(), "GET /batches/batch-$REQUEST_ID": JsonResponse( - content_type="application/json", body=_batch("completed", "file-out-$REQUEST_ID") + content_type="application/json", + body=_batch("completed", "file-out-$REQUEST_ID") if ready else _batch("in_progress", None), ), "GET /files/file-out-$REQUEST_ID/content": TextResponse( content_type="application/jsonl", body=json.dumps(output_line, separators=(",", ":")) + "\n" @@ -115,7 +122,37 @@ def _completed_routes() -> RoutedResponse: ) -def _scripted_deployment(scenario: Scenario, marker: str, routes: RoutedResponse) -> str: +def _batch_was_polled(upstream: ScenarioHandle) -> bool: + response: Final = httpx.get(f"{upstream.control_url}/__observations", timeout=15, trust_env=False) + response.raise_for_status() + observations: Final = TypeAdapter(dict[str, list[dict[str, JsonValue]]]).validate_json(response.content) + return any( + item.get("method") == "GET" and string_value(item["path"]).startswith(f"/{upstream.scenario_id}/batches/") + for item in observations["requests"] + ) + + +def _missing_output_routes() -> RoutedResponse: + return RoutedResponse( + content_type="application/x-routed", + routes={ + **_completed_routes().routes, + "GET /files/file-out-$REQUEST_ID/content": JsonResponse( + content_type="application/json", + status=404, + body={ + "error": { + "message": "No file found with id 'file-out-$REQUEST_ID'.", + "type": "invalid_request_error", + "code": "file_not_found", + } + }, + ), + }, + ) + + +def _scripted_deployment(scenario: Scenario, marker: str, routes: RoutedResponse) -> tuple[str, ScenarioHandle]: scenario_id: Final = f"poll-{marker}-{sha256(os.urandom(16)).hexdigest()[:12]}" handle: Final = register_scenario(scenario_id, routes) scenario.cleanups.callback(delete_scenario, handle) @@ -133,7 +170,7 @@ def _scripted_deployment(scenario: Scenario, marker: str, routes: RoutedResponse }, ) scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"])) - return string_value(created["model_name"]) + return string_value(created["model_name"]), handle def _submitted_batch_id(gateway: Gateway, key: str, model_name: str) -> str: @@ -173,20 +210,51 @@ def _managed_rows(batch_ids: tuple[str, ...]) -> list[dict[str, JsonValue]]: ) +@pytest.mark.parametrize( + ("retryable", "old_completed", "backlog"), + [(False, False, 0), (False, True, 0), (True, False, 0), (True, True, 0), (True, False, 1000)], +) @pytest.mark.timeout(180) @pytest.mark.covers("quota_management.spend_tracking.batch_costs.uncostable_rows_retire_so_newer_batches_are_costed") -def test_batches_gone_at_provider_do_not_starve_a_newer_batch_out_of_cost_polling(gateway: Gateway) -> None: +def test_uncostable_batches_do_not_starve_billable_batches( + gateway: Gateway, retryable: bool, old_completed: bool, backlog: int +) -> None: with gateway.scenario() as scenario: - key: Final = scenario.key() + owner: Final = f"batch-backlog-{uuid4().hex}" + if backlog: + _insert_scan_rows(owner, "old", backlog, datetime(2000, 1, 1, tzinfo=timezone.utc)) + scenario.cleanups.callback( + write_rows, 'DELETE FROM "LiteLLM_ManagedObjectTable" WHERE created_by = %s', (owner,) + ) + team_id: Final = scenario.team() + key: Final = scenario.key(team_id=team_id) gone_batch_ids: Final = tuple( _submitted_batch_id( - gateway, key, _scripted_deployment(scenario, f"gone{index}", _gone_at_provider_routes()) + gateway, + key, + _scripted_deployment(scenario, f"gone{index}", _gone_at_provider_routes(retryable=retryable))[0], ) for index in range(MAX_OBJECTS_PER_POLL_CYCLE) ) - costable_batch_id: Final = _submitted_batch_id( - gateway, key, _scripted_deployment(scenario, "costable", _completed_routes()) + missing_output_id: Final = _submitted_batch_id( + gateway, key, _scripted_deployment(scenario, "missing-output", _missing_output_routes())[0] ) + model, upstream = _scripted_deployment(scenario, "costable", _completed_routes(ready=False)) + costable_batch_id: Final = _submitted_batch_id(gateway, key, model) + for batch_id in (*gone_batch_ids, missing_output_id, costable_batch_id): + scenario.cleanups.callback( + write_rows, + 'DELETE FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id = %s', + (batch_id,), + ) + if old_completed: + write_rows( + 'UPDATE "LiteLLM_ManagedObjectTable" SET status = %s, ' + "created_at = NOW() - INTERVAL '30 days' WHERE unified_object_id = %s", + ("completed", costable_batch_id), + ) + assert eventually(lambda: _batch_was_polled(upstream), bool, seconds=120) + register_scenario(upstream.scenario_id, _completed_routes(), control_url=upstream.control_url) spend_rows: Final = eventually( lambda: read_rows( 'SELECT call_type, status, prompt_tokens, completion_tokens, spend FROM "LiteLLM_SpendLogs" ' @@ -209,4 +277,134 @@ def test_batches_gone_at_provider_do_not_starve_a_newer_batch_out_of_cost_pollin } ] assert _managed_rows((costable_batch_id,)) == [{"batch_processed": True}] - assert _managed_rows(gone_batch_ids) == [{"batch_processed": True}] * MAX_OBJECTS_PER_POLL_CYCLE + assert eventually( + lambda: _managed_rows((missing_output_id,)), + lambda rows: rows == [{"batch_processed": True}], + seconds=120, + ) == [{"batch_processed": True}] + expected_blockers: Final = [{"batch_processed": not retryable}] * MAX_OBJECTS_PER_POLL_CYCLE + blockers: Final = eventually( + lambda: _managed_rows(gone_batch_ids), lambda rows: rows == expected_blockers, seconds=120 + ) + assert blockers == expected_blockers + + expected_spend: Final = BATCH_COST_SHARE * ( + PROMPT_TOKENS * INPUT_COST_PER_TOKEN + COMPLETION_TOKENS * OUTPUT_COST_PER_TOKEN + ) + for table, column, identity in ( + ("LiteLLM_VerificationToken", "token", sha256(key.encode()).hexdigest()), + ("LiteLLM_TeamTable", "team_id", team_id), + ): + balance: Final = eventually( + partial(read_rows, f'SELECT spend FROM "{table}" WHERE {column} = %s', (identity,)), + lambda rows: len(rows) == 1 and rows[0]["spend"] == pytest.approx(expected_spend), + seconds=70, + ) + assert balance == [{"spend": pytest.approx(expected_spend)}] + + +def _insert_scan_rows(owner: str, group: str, size: int, created_at: datetime) -> tuple[str, ...]: + prefix: Final = f"{owner}-{group}-" + write_rows( + 'INSERT INTO "LiteLLM_ManagedObjectTable" ' + "(id, unified_object_id, model_object_id, file_object, file_purpose, status, created_by, created_at, updated_at) " + "SELECT %s || lpad(n::text, 8, '0'), %s || n, %s || n, '{}'::jsonb, 'batch', 'completed', %s, " + "%s::timestamptz, NOW() FROM generate_series(0, %s::int - 1) n", + (prefix, prefix, prefix, owner, created_at.isoformat(), str(size)), + ) + return tuple(f"{prefix}{index:08}" for index in range(size)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("count", [2 * MAX_OBJECTS_PER_POLL_CYCLE, 2 * MAX_OBJECTS_PER_POLL_CYCLE + 1, 1000]) +async def test_arrivals_progress_while_retry_scan_advances_and_wraps(count: int) -> None: + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.utils import ProxyLogging + from litellm.router import Router + from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + + owner: Final = f"batch-scan-{uuid4().hex}" + + try: + old_ids: Final = _insert_scan_rows(owner, "old", count, datetime(2000, 1, 1, tzinfo=timezone.utc)) + async with connected_proxy_prisma() as prisma: + logging: Final = ProxyLogging(UserApiKeyCache()) + router: Final = Router() + poller: Final = CheckBatchCost(logging, prisma, router) + where: Final = {"created_by": owner, "batch_processed": False} + first: Final = await poller._find_jobs(where) + assert tuple(row.id for row in first) == old_ids[:MAX_OBJECTS_PER_POLL_CYCLE] + + recent_ids: Final = _insert_scan_rows( + owner, "recent", MAX_OBJECTS_PER_POLL_CYCLE + 1, datetime.now(timezone.utc) + ) + + async def poll_with_arrival(turn: int) -> tuple[str, ...]: + if turn % 4 == 0: + _insert_scan_rows(owner, f"arrival-{turn}", 1, datetime.now(timezone.utc)) + page: Final = await poller._find_jobs(where) + assert len(page) <= MAX_OBJECTS_PER_POLL_CYCLE + write_rows( + 'UPDATE "LiteLLM_ManagedObjectTable" SET batch_processed = true ' + "WHERE created_by = %s AND id = ANY(%s) AND id LIKE %s", + (owner, [row.id for row in page], f"{owner}-arrival-%"), + ) + return tuple(row.id for row in page) + + pages: Final = tuple( + [await poll_with_arrival(turn) for turn in range(4 * (count // MAX_OBJECTS_PER_POLL_CYCLE + 2))] + ) + selected: Final = tuple(chain.from_iterable(pages)) + assert all(selected.count(identity) >= 2 for identity in recent_ids) + selected_old: Final = tuple(identity for identity in selected if identity in old_ids) + assert selected_old[: count - MAX_OBJECTS_PER_POLL_CYCLE] == old_ids[MAX_OBJECTS_PER_POLL_CYCLE:] + assert set(old_ids).issubset(selected_old) + + latest_ids: Final = _insert_scan_rows(owner, "latest", 1, datetime.now(timezone.utc)) + followups: Final = tuple([await poller._find_jobs(where) for _ in range(12)]) + assert any(latest_ids[0] in tuple(row.id for row in page) for page in followups) + + restarted: Final = CheckBatchCost(logging, prisma, router) + restart_page: Final = await restarted._find_jobs(where) + assert tuple(row.id for row in restart_page) == old_ids[:MAX_OBJECTS_PER_POLL_CYCLE] + finally: + write_rows('DELETE FROM "LiteLLM_ManagedObjectTable" WHERE created_by = %s', (owner,)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("steady_arrivals", [False, True]) +@pytest.mark.parametrize("backlog", [0, 1000]) +async def test_recent_batches_are_revisited_without_scanning_old_backlog(steady_arrivals: bool, backlog: int) -> None: + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.utils import ProxyLogging + from litellm.router import Router + from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + + owner: Final = f"batch-repeat-{uuid4().hex}" + try: + _insert_scan_rows(owner, "old", backlog, datetime(2000, 1, 1, tzinfo=timezone.utc)) + async with connected_proxy_prisma() as prisma: + poller: Final = CheckBatchCost(ProxyLogging(UserApiKeyCache()), prisma, Router()) + where: Final = {"created_by": owner, "batch_processed": False} + await poller._find_jobs(where) + target: Final = _insert_scan_rows(owner, "target", 1, datetime.now(timezone.utc))[0] + assert target in tuple(row.id for row in await poller._find_jobs(where)) + for cycle in range(3): + if steady_arrivals: + _insert_scan_rows(owner, f"arrival-{cycle}", 1, datetime.now(timezone.utc)) + pages: Final = tuple([await poller._find_jobs(where) for _ in range(4)]) + assert any(target in tuple(row.id for row in page) for page in pages) + assert all(len(page) <= MAX_OBJECTS_PER_POLL_CYCLE for page in pages) + write_rows( + 'UPDATE "LiteLLM_ManagedObjectTable" SET batch_processed = true WHERE created_by = %s AND id != %s', + (owner, target), + ) + write_rows( + 'UPDATE "LiteLLM_ManagedObjectTable" SET created_at = %s WHERE id = %s', + (datetime(2000, 1, 1, tzinfo=timezone.utc), target), + ) + recovery: Final = tuple([await poller._find_jobs(where) for _ in range(4)]) + assert any(target in tuple(row.id for row in page) for page in recovery) + + finally: + write_rows('DELETE FROM "LiteLLM_ManagedObjectTable" WHERE created_by = %s', (owner,)) diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index e3012dfcd13..241cf90eb5d 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -11,7 +11,8 @@ import json import logging from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager -from datetime import datetime +from copy import copy +from datetime import datetime, timezone from hashlib import sha256 from queue import SimpleQueue from types import SimpleNamespace @@ -217,7 +218,7 @@ class TestCheckBatchCost: find_call = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args assert find_call[1]["take"] == MAX_OBJECTS_PER_POLL_CYCLE - assert find_call[1]["order"] == {"created_at": "asc"} + assert find_call[1]["order"] == [{"created_at": "asc"}, {"id": "asc"}] not_in = find_call[1]["where"]["status"]["not_in"] assert "stale_expired" in not_in # "complete"/"completed" are intentionally NOT excluded from the @@ -240,13 +241,13 @@ class TestCheckBatchCost: mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) # First find_many (primary query) raises with a schema error; second (fallback) returns empty mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - side_effect=[Exception("column batch_processed does not exist"), []] + side_effect=[Exception("column batch_processed does not exist"), [], []] ) await check_batch_cost_instance.check_batch_cost() calls = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list - assert len(calls) == 2 + assert len(calls) == 3 fallback_where = calls[1][1]["where"] assert "batch_processed" not in fallback_where assert "stale_expired" in fallback_where["status"]["not_in"] @@ -267,8 +268,7 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - # Only one find_many call — the fallback directly, no primary query attempt - assert mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 1 + assert mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 2 fallback_where = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args[1]["where"] assert "batch_processed" not in fallback_where @@ -744,7 +744,7 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count >= 3, ( + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, ( "Expected update() to be called exactly once for the completed job" ) update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"] @@ -877,7 +877,7 @@ class TestCheckBatchCost: mock_update_database.assert_awaited_once() assert mock_update_database.call_args.kwargs["response_cost"] == 0.01 - assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count >= 3, ( + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, ( "the job must still be marked processed once cost tracking succeeds" ) @@ -979,10 +979,10 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, ( f"Expected update() to be called exactly once for a {terminal_status} job" ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["status"] == terminal_status assert update_data["batch_processed"] is True, ( "terminal-status update() must set batch_processed=True so polling stops" @@ -1085,9 +1085,13 @@ class TestCheckBatchCost: assert owner.user_id == "user-1" assert owner.team_id == "team-1" - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - update_call = mock_prisma_client.db.litellm_managedobjecttable.update.call_args - assert update_call.kwargs["where"] == {"id": "job-terminal-mint-1"} + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2 + update_call = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args + assert update_call.kwargs["where"] == { + "id": "job-terminal-mint-1", + "batch_processed": False, + "updated_at": mock_job.updated_at, + } update_data = update_call.kwargs["data"] assert update_data["status"] == terminal_status assert update_data["batch_processed"] is True @@ -1152,10 +1156,10 @@ class TestCheckBatchCost: ) as mock_afile_content: await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, ( "a completed batch with no output file must be marked processed exactly once" ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["status"] == completed_status assert update_data["batch_processed"] is True, ( "completed-without-output update() must set batch_processed=True so polling stops" @@ -1371,7 +1375,7 @@ class TestCheckBatchCost: f"{terminal_status} batch with an output file must fetch results and be billed" ) mock_logging_obj.async_success_handler.assert_awaited_once() - assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count >= 3 + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2 update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["batch_processed"] is True assert update_data["status"] == terminal_status, ( @@ -1544,10 +1548,10 @@ class TestCheckBatchCost: assert mock_afile_content.await_count == 1 mock_calculate.assert_not_awaited() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, ( "a terminal batch with a 404ing output file must be retired, not retried forever" ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["status"] == "failed" assert update_data["batch_processed"] is True @@ -1942,7 +1946,7 @@ class TestUnmanagedVertexRouting: mock_logging_obj.async_success_handler.assert_awaited_once() assert mock_logging_obj.async_success_handler.call_args[1]["batch_cost"] == 0.01 - assert prisma.db.litellm_managedobjecttable.update_many.call_count >= 3 + assert prisma.db.litellm_managedobjecttable.update_many.call_count == 2 update_data = prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["batch_processed"] is True assert update_data["status"] == "complete" @@ -2159,7 +2163,7 @@ class TestUnmanagedBedrockRouting: mock_logging_obj.async_success_handler.assert_awaited_once() assert mock_logging_obj.async_success_handler.call_args[1]["batch_cost"] == 0.02 - assert prisma.db.litellm_managedobjecttable.update_many.call_count >= 3 + assert prisma.db.litellm_managedobjecttable.update_many.call_count == 2 update_data = prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["batch_processed"] is True assert update_data["status"] == "complete" @@ -2740,9 +2744,10 @@ class TestPollPageStarvation: await self._instance(prisma, llm_router).check_batch_cost() llm_router.aretrieve_batch.assert_not_awaited() - prisma.db.litellm_managedobjecttable.update.assert_awaited_once() - call = prisma.db.litellm_managedobjecttable.update.call_args[1] - assert call["where"] == {"id": "job-no-model"} + prisma.db.litellm_managedobjecttable.update_many.assert_awaited() + call = prisma.db.litellm_managedobjecttable.update_many.call_args[1] + assert call["where"]["id"] == "job-no-model" + assert call["where"]["batch_processed"] is False assert call["data"] == {"batch_processed": True} @pytest.mark.asyncio @@ -2770,8 +2775,8 @@ class TestPollPageStarvation: await self._instance(prisma, llm_router).check_batch_cost() - prisma.db.litellm_managedobjecttable.update.assert_awaited_once() - assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"batch_processed": True} + prisma.db.litellm_managedobjecttable.update_many.assert_awaited() + assert prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"] == {"batch_processed": True} @pytest.mark.asyncio async def test_provider_404_with_deployment_gone_keeps_job(self): @@ -2830,24 +2835,7 @@ class TestPollPageStarvation: await instance.check_batch_cost() - assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"status": "stale_expired"} - - @pytest.mark.asyncio - async def test_stale_cleanup_gives_up_on_never_costed_completed_rows(self): - """A row already in a terminal status is never rewritten by the staleness sweep, so - it needs its own bound or it starves newer batches indefinitely.""" - prisma = self._prisma([]) - - await self._instance(prisma, MagicMock()).check_batch_cost() - - calls = prisma.db.litellm_managedobjecttable.update_many.call_args_list - assert len(calls) == 2, "expected the staleness sweep plus the never-costed sweep" - where = calls[1][1]["where"] - assert where["file_purpose"] == "batch" - assert where["batch_processed"] is False - assert where["status"] == {"in": ["complete", "completed"]} - assert "created_at" in where - assert calls[1][1]["data"] == {"batch_processed": True} + assert prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"] == {"status": "stale_expired"} @pytest.mark.asyncio async def test_newer_batch_is_polled_once_dead_rows_are_retired(self): @@ -2885,7 +2873,11 @@ class TestPollPageStarvation: await self._instance(prisma, llm_router).check_batch_cost() - retired = [call[1]["where"]["id"] for call in prisma.db.litellm_managedobjecttable.update.call_args_list] + retired = [ + call[1]["where"]["id"] + for call in prisma.db.litellm_managedobjecttable.update_many.call_args_list + if "id" in call[1]["where"] + ] assert retired == ["job-no-model", "job-gone"] assert llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live", ( "the newer healthy batch must still be polled in the same cycle" @@ -2933,7 +2925,8 @@ class _FakeManagedObjectRow: self.team_id = None self.api_key = None self.request_tags = None - self.created_at = 1700000000 + self.created_at = datetime(9999, 1, 1, tzinfo=timezone.utc) + self.updated_at = datetime(2000, 1, 1, tzinfo=timezone.utc) self.file_object = json.dumps( { "id": "batch-456", @@ -2963,7 +2956,9 @@ class _FakeManagedObjectTable: def _matches(self, where: dict) -> bool: for key, value in where.items(): if key == "created_at": - return False + if "lt" in value and not self.row.created_at < value["lt"]: + return False + continue if key == "status": if self.row.status in value.get("not_in", []): return False @@ -2976,7 +2971,9 @@ class _FakeManagedObjectTable: async def _update_many(self, *, where: dict, data: dict) -> int: if not self._matches(where): return 0 - if "batch_processed" in where: + if "updated_at" in where: + self.journal.append("finalize") + elif "batch_processed" in where: self.journal.append("claim" if data.get("batch_processed") else "release") for key, value in data.items(): setattr(self.row, key, value) @@ -3224,6 +3221,89 @@ class TestMultiPodBatchCostClaim: ["batch-test-model"] if deployment_model == "openai/batch-test-model" else ["batch-test-model"] * successes ) + @pytest.mark.asyncio + @pytest.mark.parametrize("error_source", ("batch", "output")) + async def test_a_404_for_another_identifier_preserves_recovery(self, error_source: str) -> None: + import base64 + + import httpx + import litellm + import respx + + from litellm.integrations.custom_logger import CustomLogger + from litellm.types.utils import SpecialEnums + + class SpendRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.costs: tuple[float, ...] = () + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + if kwargs["call_type"] == "aretrieve_batch": + cost: Final = kwargs["response_cost"] + assert isinstance(cost, float) + self.costs += (cost,) + + recorder: Final = SpendRecorder() + litellm.logging_callback_manager.add_litellm_async_success_callback(recorder) + litellm.register_model( + { + "batch-recovery-model": { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token_batches": 0.01, + "output_cost_per_token_batches": 0.02, + } + } + ) + row: Final = _FakeManagedObjectRow() + row.unified_object_id = ( + base64.urlsafe_b64encode( + SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format("model-123", "batch-456").encode() + ) + .decode() + .rstrip("=") + ) + router: Final = self._router() + response: Final = router.aretrieve_batch.return_value + error: Final = litellm.NotFoundError( + message=f"No object found with id '{'batch-456' if error_source == 'batch' else _CLAIM_OUTPUT_FILE_ID}0'.", + model="gpt-4", + llm_provider="openai", + ) + router.aretrieve_batch.side_effect = (error, response) if error_source == "batch" else (response, response) + prisma: Final = self._prisma(row, []) + output: Final = json.dumps( + { + "response": { + "status_code": 200, + "body": { + "model": "batch-recovery-model", + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + }, + "error": None, + } + ).encode() + with respx.mock: + route: Final = respx.get(f"https://api.openai.com/v1/files/{_CLAIM_OUTPUT_FILE_ID}/content") + route.mock( + side_effect=error if error_source == "output" else None, + return_value=None if error_source == "output" else httpx.Response(200, content=output), + ) + await self._instance(prisma, router).check_batch_cost() + assert row.batch_processed is False + assert recorder.costs == () + route.mock(return_value=httpx.Response(200, content=output)) + await self._instance(prisma, router).check_batch_cost() + await self._instance(prisma, router).check_batch_cost() + + assert row.batch_processed is True + assert row.status == "complete" + assert recorder.costs == (pytest.approx(10 * 0.01 + 5 * 0.02),) + @staticmethod def _claim_calls(prisma) -> list: return [ @@ -3696,12 +3776,17 @@ class TestBatchLineItemsFromThePoller: assert logger.payloads == [] @pytest.mark.asyncio - async def test_a_batch_that_could_not_be_retired_sends_no_lines_until_it_is(self, monkeypatch): + @pytest.mark.parametrize("database_error", [False, True]) + async def test_a_batch_that_could_not_be_retired_sends_no_lines_until_it_is( + self, monkeypatch: pytest.MonkeyPatch, database_error: bool + ) -> None: import litellm journal: list = [] row, instance = self._setup(journal, status="expired", output_file_id=None, error_file_id="file-err-1") - instance.prisma_client.db.litellm_managedobjecttable.update = AsyncMock(side_effect=RuntimeError("db down")) + instance.prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + side_effect=RuntimeError("db down") if database_error else None, return_value=0 + ) logger = self._recording_logger(journal) monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,)) @@ -4111,6 +4196,127 @@ class TestCreatorTurnedOffMessageLogging: instance.prisma_client.db.litellm_verificationtoken.find_unique.assert_not_called() +@pytest.mark.asyncio +async def test_old_completed_batch_remains_available_for_costing() -> None: + row: Final = _FakeManagedObjectRow() + row.status = "completed" + row.created_at = datetime(2000, 1, 1, tzinfo=timezone.utc) + prisma: Final = TestMultiPodBatchCostClaim._prisma(row, []) + poller: Final = TestMultiPodBatchCostClaim._instance(prisma, TestMultiPodBatchCostClaim._router()) + + await poller._cleanup_stale_managed_objects() + + assert row.batch_processed is False + assert row.status == "completed" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("progress", ["claimed", "refreshed"]) +async def test_stale_retirement_preserves_concurrent_progress(progress: str) -> None: + row: Final = _FakeManagedObjectRow() + selected: Final = copy(row) + if progress == "claimed": + row.batch_processed = True + row.updated_at = datetime(2000, 1, 2, tzinfo=timezone.utc) + prisma: Final = TestMultiPodBatchCostClaim._prisma(row, []) + poller: Final = TestMultiPodBatchCostClaim._instance(prisma, TestMultiPodBatchCostClaim._router()) + + await poller._retire_job(selected, "provider batch disappeared") + + assert row.batch_processed is (progress == "claimed") + assert row.updated_at > selected.updated_at + + +@pytest.mark.asyncio +@pytest.mark.parametrize("progress", ["claimed", "refreshed"]) +async def test_stale_terminal_response_preserves_concurrent_progress(progress: str) -> None: + from litellm.types.utils import LiteLLMBatch + + row: Final = _FakeManagedObjectRow() + selected: Final = copy(row) + if progress == "claimed": + row.batch_processed = True + row.updated_at = datetime(2000, 1, 2, tzinfo=timezone.utc) + row.file_object = '{"revision":"fresh"}' + prisma: Final = TestMultiPodBatchCostClaim._prisma(row, []) + poller: Final = TestMultiPodBatchCostClaim._instance(prisma, TestMultiPodBatchCostClaim._router()) + response: Final = LiteLLMBatch( + id="batch-456", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-input-1", + object="batch", + status="failed", + ) + + assert await poller._finalize_unbilled_terminal_job(selected, response) is False + + assert row.file_object == '{"revision":"fresh"}' + assert row.status == "in_progress" + assert row.batch_processed is (progress == "claimed") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("legacy", [False, True]) +async def test_permanently_lost_batch_retires_unchanged_row(legacy: bool) -> None: + row: Final = _FakeManagedObjectRow() + prisma: Final = TestMultiPodBatchCostClaim._prisma(row, []) + poller: Final = TestMultiPodBatchCostClaim._instance(prisma, TestMultiPodBatchCostClaim._router()) + poller._has_batch_processed_column = not legacy + + await poller._retire_job(row, "provider batch disappeared") + + assert row.batch_processed is (not legacy) + assert row.status == ("stale_expired" if legacy else "in_progress") + + +@pytest.mark.asyncio +async def test_failed_retirement_keeps_batch_available() -> None: + row: Final = _FakeManagedObjectRow() + prisma: Final = TestMultiPodBatchCostClaim._prisma(row, []) + prisma.db.litellm_managedobjecttable.update_many = AsyncMock(side_effect=RuntimeError("database unavailable")) + poller: Final = TestMultiPodBatchCostClaim._instance(prisma, TestMultiPodBatchCostClaim._router()) + + await poller._retire_job(row, "provider batch disappeared") + + assert row.batch_processed is False + assert row.status == "in_progress" + + +@pytest.mark.parametrize("resource_id", ("batch-123", "file-123", "gs://bucket/job/predictions.jsonl")) +@pytest.mark.parametrize("other_id", ("{id}4", "prefix-{id}", "{id}/child", "{id}.backup", "{id}?generation=2")) +def test_another_resource_in_a_404_is_not_permanent_loss(resource_id: str, other_id: str) -> None: + from litellm.exceptions import NotFoundError + from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + + error: Final = NotFoundError( + message=f"No object found with id '{other_id.format(id=resource_id)}'.", model="test", llm_provider="openai" + ) + assert CheckBatchCost._is_batch_gone_at_provider(error, resource_id) is False + assert CheckBatchCost._is_output_file_gone_at_provider(error, resource_id) is False + + +@pytest.mark.parametrize("resource_id", ("batch-123", "file-123", "gs://bucket/job/predictions.jsonl")) +@pytest.mark.parametrize("message", ("No object found with id '{id}'.", 'File "{id}" does not exist', "{id} not found")) +def test_a_404_naming_the_exact_resource_is_permanent_loss(resource_id: str, message: str) -> None: + from litellm.exceptions import NotFoundError + from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + + error: Final = NotFoundError(message=message.format(id=resource_id), model="test", llm_provider="openai") + assert CheckBatchCost._is_batch_gone_at_provider(error, resource_id) is True + assert CheckBatchCost._is_output_file_gone_at_provider(error, resource_id) is True + + +@pytest.mark.parametrize("output_file_id", (None, "")) +def test_a_404_without_an_output_identity_remains_retryable(output_file_id: str | None) -> None: + from litellm.exceptions import NotFoundError + from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + + error: Final = NotFoundError(message="Output file not found", model="test", llm_provider="openai") + assert CheckBatchCost._is_output_file_gone_at_provider(error, output_file_id) is False + + @pytest.mark.asyncio @pytest.mark.parametrize( "kind", ["sync", "async", "custom", "overlapping", "general", "mixed", "sync_only", "executor_shutdown"] From 047cb4a8e2b2c0853baaf9b6a3f082d851d58a9f Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 15:29:30 +0000 Subject: [PATCH 20/33] test(integration): expect the Exa search envelope's provider_metadata after #47423 (#47572) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../integration/authorization/test_search_tool_allowlist.py | 5 ++++- tests/integration/management/test_search_tool_crud.py | 1 + 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/integration/authorization/test_search_tool_allowlist.py b/tests/integration/authorization/test_search_tool_allowlist.py index 84f2b7ebaf2..3782f31991b 100644 --- a/tests/integration/authorization/test_search_tool_allowlist.py +++ b/tests/integration/authorization/test_search_tool_allowlist.py @@ -92,7 +92,10 @@ def test_search_tool_allowlist_refuses_other_tools_before_any_upstream_call( assert (refused.status_code, refused.json()) == (403, refusal), refused.text assert wire.drain() == () served: Final = gateway.request("POST", f"/v1/search/{allowed}", _SEARCH_BODY, key=key) - assert (served.status_code, served.json()) == (200, {"object": "search", "results": []}), served.text + assert (served.status_code, served.json()) == ( + 200, + {"object": "search", "results": [], "provider_metadata": {}}, + ), served.text assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/allowed/search")] diff --git a/tests/integration/management/test_search_tool_crud.py b/tests/integration/management/test_search_tool_crud.py index f46da964e14..e97f2727595 100644 --- a/tests/integration/management/test_search_tool_crud.py +++ b/tests/integration/management/test_search_tool_crud.py @@ -113,6 +113,7 @@ def test_search_tool_create_update_delete_is_read_back_and_applied_without_resta "results": [ {"title": target, "url": "https://crud.invalid", "snippet": key, "date": None, "last_updated": None} ], + "provider_metadata": {}, }, response.text read_back(original) From 7a4a255d1d33cb6145546c65b69b444a5f136da2 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 20 Oct 2026 09:08:11 -0700 Subject: [PATCH 21/33] feat(proxy): lift every earlier billed timed-out retry attempt onto the failed request (#47490) * feat(proxy): lift every earlier billed timed-out retry attempt onto the failed request _billed_earlier_attempts reads the router's previous_models records, keeps the Timeout ones other than the final attempt (matched on deployment_id plus attempted_retries, since the count restarts after a fallback hop and the fallback handler re-logs the exhausting attempt) and prices each at the deployment that served it. post_call_failure_hook lifts the tuple as billed_failed_attempts; nothing writes it yet. previous_models joins the router-reserved metadata fields so a client-planted record never reaches it * fix(proxy): price earlier timed-out attempts from the prompt estimate only The final row may lift a mid-stream recovered usage (prompt plus delivered chunks) or a failed_attempts_usage sum; earlier attempts never produced output, so they are priced from _estimated_prompt_usage, and when failed_attempts_usage already bills them into the final row none is lifted again --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/litellm_pre_call_utils.py | 1 + litellm/proxy/utils.py | 153 ++++++++++++++++-- .../unit/proxy/test_litellm_pre_call_utils.py | 6 +- ..._utils_model_creation_and_error_logging.py | 137 ++++++++++++++++ 4 files changed, 282 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index c0222dbcd0b..efc8575feb3 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -467,6 +467,7 @@ _ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset( { "attempted_fallbacks", "original_model_group", + "previous_models", "request_retry_count", CLIENT_OUTPUT_CEILING_METADATA_KEY, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index a8bd35644c9..735df38efe6 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1031,9 +1031,7 @@ def _failure_usage_to_lift( at the rate the router registered under ``deployment_id`` (``Logging.get_router_model_id``) when the deployment declares one, else the cost map's rate for ``base_model`` (an Azure deployment's ``model_info.base_model``) or the deployment model; any other post-dispatch - failure stays at zero cost. The raw request body backfills - the system prompt when the SDK bridges an endpoint (e.g. /v1/messages on a - chat-completions provider) without filling optional_params. Returns the + failure stays at zero cost. Returns the (combined_usage_object, response_cost) pair to lift, or None.""" recovered_usage: Final = model_call_details.get("combined_usage_object") failed_attempts: Final = _FAILED_ATTEMPTS_USAGE.validate_python(model_call_details.get("failed_attempts_usage", ())) @@ -1047,6 +1045,22 @@ def _failure_usage_to_lift( # The proxy-only error handler stamps this flag on a client disconnect too, even though the # request had already reached the provider, so a billed failure keeps its estimate. return None + estimated_usage: Final = _estimated_prompt_usage(model_call_details, request_body) + if estimated_usage is None: + return None + if not provider_billed: + return estimated_usage, 0.0 + return estimated_usage, _estimated_input_cost( + model_call_details, estimated_usage.prompt_tokens, deployment_id=deployment_id, base_model=base_model + ) + + +def _estimated_prompt_usage( + model_call_details: Mapping[str, object], request_body: Mapping[str, object] +) -> Usage | None: + """The input-only Usage estimate for the prompt this request sent, or None when the call type has no + estimable input. The raw request body backfills the system prompt when the SDK bridges an endpoint + (e.g. /v1/messages on a chat-completions provider) without filling optional_params.""" if str(model_call_details.get("call_type")) not in _INPUT_ESTIMABLE_CALL_TYPES: return None optional_params: Final = model_call_details.get("optional_params") @@ -1056,17 +1070,118 @@ def _failure_usage_to_lift( else None ) system_input: Final = dispatched_system or request_body.get("system") or request_body.get("instructions") - estimated_usage: Final = _estimate_dispatched_failure_usage( + return _estimate_dispatched_failure_usage( model=str(model_call_details.get("model") or ""), request_input=model_call_details.get("messages"), system_input=system_input, ) - if estimated_usage is None: - return None - if not provider_billed: - return estimated_usage, 0.0 - return estimated_usage, _estimated_input_cost( - model_call_details, estimated_usage.prompt_tokens, deployment_id=deployment_id, base_model=base_model + + +@dataclass(frozen=True) +class BilledFailedAttempt: + """An earlier retry attempt the provider also billed, written as its own spend row.""" + + attempted_retries: int + deployment_id: str | None + usage: Usage + response_cost: float + + +_NO_RETRY_METADATA: Final[Mapping[str, object]] = MappingProxyType({}) + + +def _retry_metadata_bucket(litellm_params: Mapping[str, object] | None) -> Mapping[str, object]: + """The request metadata bucket the router writes retry state into, ``litellm_metadata`` when the + route carries one (/messages, /responses) and ``metadata`` otherwise, same order as the deployment id.""" + if litellm_params is None: + return _NO_RETRY_METADATA + for metadata_key in ("litellm_metadata", "metadata"): + metadata = litellm_params.get(metadata_key) + if isinstance(metadata, dict): + return cast("Mapping[str, object]", metadata) # cast-ok: untyped router dict, keys read with .get + return _NO_RETRY_METADATA + + +def _attempt_number(record: Mapping[str, object], index: int) -> int: + """The attempt a retry record describes: its own ``attempted_retries`` count, else its position.""" + recorded: Final = record.get("attempted_retries") + return recorded if isinstance(recorded, int) else index + + +def _retry_attempt_records(litellm_params: Mapping[str, object] | None) -> tuple[Mapping[str, object], ...]: + """The router's ``RetryAttemptRecord`` breadcrumbs, one per failed attempt, read from the same bucket.""" + if litellm_params is None: + return () + for metadata_key in ("litellm_metadata", "metadata"): + metadata = litellm_params.get(metadata_key) + if not isinstance(metadata, dict): + continue + records = cast("Mapping[str, object]", metadata).get("previous_models") # cast-ok: untyped router dict + if isinstance(records, (list, tuple)): + return tuple( + cast("Mapping[str, object]", record) # cast-ok: RetryAttemptRecord TypedDict, read with .get + for record in cast("Sequence[object]", records) # cast-ok: untyped router tuple + if isinstance(record, dict) + ) + return () + + +def _record_deployment_id(record: Mapping[str, object]) -> str | None: + deployment_id: Final = record.get("deployment_id") + return deployment_id if isinstance(deployment_id, str) else None + + +def _attempt_identity(record: Mapping[str, object]) -> tuple[str | None, int | None]: + """What tells one provider call from another across the request: the deployment that served it + and the retry count within its model group. Two records with the same identity describe the + same call (the fallback handler re-logs the attempt that exhausted the previous group).""" + attempted_retries: Final = record.get("attempted_retries") + return _record_deployment_id(record), attempted_retries if isinstance(attempted_retries, int) else None + + +def _timed_out_earlier_records( + litellm_params: Mapping[str, object] | None, final_identity: tuple[str | None, int | None] +) -> tuple[Mapping[str, object], ...]: + """The ``Timeout`` breadcrumbs for calls other than the final attempt, one per provider call.""" + by_identity: Final = { + _attempt_identity(record): record + for record in _retry_attempt_records(litellm_params) + if record.get("exception_type") == litellm.Timeout.__name__ + } + return tuple(record for identity, record in by_identity.items() if identity != final_identity) + + +def _billed_earlier_attempts( + model_call_details: Mapping[str, object], + litellm_params: Mapping[str, object] | None, + estimated_usage: Usage, + base_model: str | None, + deployment_id: str | None, +) -> tuple[BilledFailedAttempt, ...]: + """Every earlier retry attempt the provider also billed. The router records one + ``RetryAttemptRecord`` per failed attempt in ``metadata.previous_models``, the final attempt + included, which the failure row itself bills. The final attempt is the record with the + request's own deployment and retry count (``attempted_retries`` restarts at 0 after a fallback + hop, so the count alone cannot tell the groups apart). An earlier attempt that ended in a + ``Timeout`` had the same prompt accepted by its deployment, so it is billed like the final + attempt, at that deployment's input rate (else the cost map's rate for ``base_model`` or the + model). Only the exception class survives on the record, so an earlier connect timeout is + billed too; any other failure class is not.""" + final_attempt: Final = _retry_metadata_bucket(litellm_params).get("attempted_retries") + final_identity: Final = (deployment_id, final_attempt if isinstance(final_attempt, int) else None) + return tuple( + BilledFailedAttempt( + attempted_retries=_attempt_number(record, index), + deployment_id=_record_deployment_id(record), + usage=estimated_usage, + response_cost=_estimated_input_cost( + model_call_details, + estimated_usage.prompt_tokens, + deployment_id=_record_deployment_id(record), + base_model=base_model, + ), + ) + for index, record in enumerate(_timed_out_earlier_records(litellm_params, final_identity)) ) @@ -1265,13 +1380,26 @@ def _failure_fields_to_lift(request_data: Mapping[str, object], provider_billed: return _EMPTY_LIFT _model_call_details: Final = getattr(_logging_obj, "model_call_details", {}) _first_handoff: Final = _model_call_details.get("first_api_call_start_time") + _litellm_params: Final = getattr(_logging_obj, "litellm_params", None) + _base_model: Final = get_base_model_from_metadata(model_call_details=_model_call_details) + _deployment_id: Final = _logging_obj.get_router_model_id() if isinstance(_logging_obj, Logging) else None _usage_to_lift: Final = _failure_usage_to_lift( model_call_details=_model_call_details, request_body=request_data, dispatched=_first_handoff is not None, provider_billed=provider_billed, - deployment_id=_logging_obj.get_router_model_id() if isinstance(_logging_obj, Logging) else None, - base_model=get_base_model_from_metadata(model_call_details=_model_call_details), + deployment_id=_deployment_id, + base_model=_base_model, + ) + _prompt_estimate: Final = ( + _estimated_prompt_usage(_model_call_details, request_data) + if provider_billed and _first_handoff is not None and not _model_call_details.get("failed_attempts_usage") + else None + ) + _earlier_attempts: Final = ( + _billed_earlier_attempts(_model_call_details, _litellm_params, _prompt_estimate, _base_model, _deployment_id) + if _prompt_estimate is not None + else () ) _entries: Final = ( ("start_time", _model_call_details.get("start_time")), @@ -1280,6 +1408,7 @@ def _failure_fields_to_lift(request_data: Mapping[str, object], provider_billed: ("combined_usage_object", None if _usage_to_lift is None else _usage_to_lift[0]), ("response_cost", None if _usage_to_lift is None else (_usage_to_lift[1] or 0.0)), ("standard_logging_object", _model_call_details.get("standard_logging_object")), + ("billed_failed_attempts", _earlier_attempts or None), ) return MappingProxyType({key: value for key, value in _entries if value is not None}) diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 4a97106b3ac..246496ec48f 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -3330,9 +3330,7 @@ def test_x_litellm_trace_id_keeps_the_body_session_beside_it_whatever_the_log_se assert data["litellm_trace_id"] == "req-trace-1" -@pytest.mark.parametrize( - "session_header", ["x-litellm-session-id", "x-claude-code-session-id", "x-session-id"] -) +@pytest.mark.parametrize("session_header", ["x-litellm-session-id", "x-claude-code-session-id", "x-session-id"]) @pytest.mark.parametrize("with_trace_header", [False, True], ids=["session-header-only", "beside-trace-header"]) def test_session_headers_still_beat_the_body_session_with_correlation_logging( monkeypatch: pytest.MonkeyPatch, session_header: str, with_trace_header: bool @@ -8282,6 +8280,7 @@ _PLANTED_STAMPS = { "attempted_fallbacks": 99, "original_model_group": "spoofed-group", "request_retry_count": -100, + "previous_models": [{"exception_type": "Timeout", "deployment_id": "spoofed", "attempted_retries": 101}], "_client_output_ceiling": {"api_base": "https://attacker.example"}, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: 10**9, "client_key": "client_value", @@ -8316,6 +8315,7 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_bo assert "original_model_group" not in updated["metadata"] assert "_client_output_ceiling" not in updated["metadata"] assert "request_retry_count" not in updated["metadata"] + assert "previous_models" not in updated["metadata"] assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in updated["metadata"] assert updated["metadata"]["client_key"] == "client_value" diff --git a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index 77d17741a18..5ab04724bc4 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -2135,6 +2135,143 @@ def test_a_timed_out_request_is_priced_at_the_rate_the_router_registered_for_the assert response_cost == pytest.approx(usage.prompt_tokens * 0.25) +def test_every_earlier_timed_out_attempt_is_billed_at_its_own_deployment_rate(monkeypatch): + """The router leaves one RetryAttemptRecord per failed attempt in metadata.previous_models, the + final attempt included; the earlier Timeout ones are billed (each at the deployment that served + it), a 429 attempt is not, and the final attempt is left to the failure row itself.""" + from litellm.proxy.utils import _billed_earlier_attempts + + monkeypatch.setitem( + litellm.model_cost, + "dep-lit145-b", + {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 0.25, "output_cost_per_token": 1.0}, + ) + usage = litellm.Usage(prompt_tokens=24, completion_tokens=0, total_tokens=24) + public_rate, _ = litellm.cost_per_token(model=FAILURE_USAGE_MODEL, prompt_tokens=24) + + attempts = _billed_earlier_attempts( + model_call_details={"model": FAILURE_USAGE_MODEL, "custom_llm_provider": "openai"}, + litellm_params={ + "metadata": { + "attempted_retries": 3, + "previous_models": ( + {"deployment_id": "dep-lit145-a", "exception_type": "Timeout", "attempted_retries": 0}, + {"deployment_id": "dep-lit145-b", "exception_type": "RateLimitError", "attempted_retries": 1}, + {"deployment_id": "dep-lit145-b", "exception_type": "Timeout", "attempted_retries": 2}, + {"deployment_id": "dep-lit145-a", "exception_type": "Timeout", "attempted_retries": 3}, + ), + } + }, + estimated_usage=usage, + base_model=None, + deployment_id="dep-lit145-a", + ) + + assert [(a.attempted_retries, a.deployment_id) for a in attempts] == [(0, "dep-lit145-a"), (2, "dep-lit145-b")] + assert attempts[0].response_cost == pytest.approx(public_rate) + assert attempts[0].response_cost > 0 + assert attempts[1].response_cost == pytest.approx(24 * 0.25) + assert all(a.usage is usage for a in attempts) + + +def test_earlier_attempts_are_priced_from_the_prompt_estimate_not_the_final_streams_recovered_usage(): + """A final attempt whose stream broke mid-flight lifts its recovered usage (prompt plus the chunks + delivered); the earlier timed-out attempts never produced output, so they are priced from the + input-only estimate and not from that recovered usage.""" + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.proxy.utils import _estimated_prompt_usage, _failure_fields_to_lift + + logging_obj = _dispatched_failure_logging_obj( + FAILURE_USAGE_MODEL, + get_litellm_params( + metadata={ + "attempted_retries": 1, + "model_info": {"id": "dep-lit145-final"}, + "previous_models": ( + {"deployment_id": "dep-lit145-final", "exception_type": "Timeout", "attempted_retries": 0}, + ), + } + ), + ) + recovered = litellm.Usage(prompt_tokens=500, completion_tokens=40, total_tokens=540) + logging_obj.model_call_details["combined_usage_object"] = recovered + logging_obj.model_call_details["response_cost"] = 0.01 + + lifted = _failure_fields_to_lift({"litellm_logging_obj": logging_obj}, provider_billed=True) + + estimate = _estimated_prompt_usage(logging_obj.model_call_details, {}) + assert lifted["combined_usage_object"] is recovered + assert [(a.usage.prompt_tokens, a.usage.completion_tokens) for a in lifted["billed_failed_attempts"]] == [ + (estimate.prompt_tokens, 0) + ] + assert estimate.prompt_tokens < recovered.prompt_tokens + + +def test_attempts_already_summed_into_the_final_row_are_not_lifted_again(): + """``Logging.carry_failed_attempt_usage`` moves a cut stream's usage into ``failed_attempts_usage`` + and the final row already bills every entry of it, so no earlier attempt gets a row of its own.""" + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.proxy.utils import _failure_fields_to_lift + + logging_obj = _dispatched_failure_logging_obj( + FAILURE_USAGE_MODEL, + get_litellm_params( + metadata={ + "attempted_retries": 1, + "model_info": {"id": "dep-lit145-final"}, + "previous_models": ( + {"deployment_id": "dep-lit145-final", "exception_type": "Timeout", "attempted_retries": 0}, + ), + } + ), + ) + logging_obj.model_call_details["failed_attempts_usage"] = ( + (litellm.Usage(prompt_tokens=24, completion_tokens=3, total_tokens=27), 0.001), + ) + + lifted = _failure_fields_to_lift({"litellm_logging_obj": logging_obj}, provider_billed=True) + + assert lifted["combined_usage_object"].prompt_tokens == 24 + assert "billed_failed_attempts" not in lifted + + +def test_a_timed_out_attempt_from_the_group_a_fallback_left_is_still_billed(): + """After a fallback hop the router restarts ``attempted_retries`` at 0 for the new model group and + re-logs the attempt that exhausted the previous group, so the final attempt (dep-b, retry 1) must be + told from the previous group's attempt with the same count (dep-a, retry 1), and the re-logged + record must bill once.""" + from litellm.proxy.utils import _billed_earlier_attempts + + usage = litellm.Usage(prompt_tokens=24, completion_tokens=0, total_tokens=24) + + attempts = _billed_earlier_attempts( + model_call_details={"model": FAILURE_USAGE_MODEL, "custom_llm_provider": "openai"}, + litellm_params={ + "metadata": { + "attempted_retries": 1, + "model_info": {"id": "dep-lit145-b"}, + "previous_models": ( + {"deployment_id": "dep-lit145-a", "exception_type": "Timeout", "attempted_retries": 0}, + {"deployment_id": "dep-lit145-a", "exception_type": "Timeout", "attempted_retries": 1}, + {"deployment_id": "dep-lit145-a", "exception_type": "Timeout", "attempted_retries": 1}, + {"deployment_id": "dep-lit145-b", "exception_type": "Timeout", "attempted_retries": 0}, + {"deployment_id": "dep-lit145-b", "exception_type": "Timeout", "attempted_retries": 1}, + ), + } + }, + estimated_usage=usage, + base_model=None, + deployment_id="dep-lit145-b", + ) + + assert [(a.attempted_retries, a.deployment_id) for a in attempts] == [ + (0, "dep-lit145-a"), + (1, "dep-lit145-a"), + (0, "dep-lit145-b"), + ] + assert all(a.response_cost > 0 for a in attempts) + + def _dispatched_failure_logging_obj(model: str, litellm_params: dict) -> object: """A Logging object as the failure hook sees it: litellm_params in the real get_litellm_params() shape (every key present, ``litellm_metadata`` None on chat completions) and a first-handoff From ccd0d4ea493d1c5f9c2f6574261002407687663d Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 20 Oct 2026 09:08:18 -0700 Subject: [PATCH 22/33] fix(proxy): requeue spend logs when a Postgres restart tears down the prisma engine mid-write (#47571) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/utils.py | 6 ++- .../test_proxy_update_spend.py | 50 +++++++++++++++++++ 2 files changed, 55 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 735df38efe6..d6dfd3ea359 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -8504,7 +8504,11 @@ MAX_SPEND_LOG_ISOLATION_FAILURES_PER_BATCH: Final = 256 def _is_transient_spend_log_write_error(e: Exception) -> bool: - return PrismaDBExceptionHandler.is_database_transport_error(e) or PrismaDBExceptionHandler.is_deadlock_error(e) + return ( + PrismaDBExceptionHandler.is_database_transport_error(e) + or PrismaDBExceptionHandler.is_prisma_engine_internal_error(e) + or PrismaDBExceptionHandler.is_deadlock_error(e) + ) async def _create_spend_logs_with_poison_isolation( diff --git a/tests/unit/proxy/utils/prisma_and_spend/test_proxy_update_spend.py b/tests/unit/proxy/utils/prisma_and_spend/test_proxy_update_spend.py index fc3deb0de18..e3b7bb5b49f 100644 --- a/tests/unit/proxy/utils/prisma_and_spend/test_proxy_update_spend.py +++ b/tests/unit/proxy/utils/prisma_and_spend/test_proxy_update_spend.py @@ -952,3 +952,53 @@ async def test_update_spend_logs_requeues_batch_when_postgres_is_out_of_connecti assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 1 sleep.assert_not_awaited() assert [row["request_id"] for row in mock_prisma_client.spend_log_transactions] == ["r0", "r1", "r2", "r3", "e"] + + +async def _create_many_while_postgres_is_restarting(**_: object) -> int: + from prisma.engine.utils import handle_response_errors + + handle_response_errors( + MagicMock(), + [ + { + "error": "Error in connector: Can't reach database server at `127.0.0.1:5432`", + "user_facing_error": {"error_code": "P1001", "message": "Can't reach database server", "meta": None}, + } + ], + ) + return 0 + + +@pytest.mark.asyncio +async def test_update_spend_logs_retries_and_requeues_batch_when_postgres_restart_tears_down_prisma_engine( + mock_prisma_client: Any, make_spend_log_row: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + """While Postgres restarts the prisma engine answers the write with a P1001 + payload whose ``meta`` is ``None``, which prisma's own response parser trips + over before it can build a typed error. That outage must be retried like + any other, and once retries exhaust the dequeued batch goes back to the head + of the queue so the rows land when Postgres is back. + """ + + async def _fake_sleep(_: float) -> None: + return None + + monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep) + mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock( + side_effect=_create_many_while_postgres_is_restarting + ) + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="c")] + + with pytest.raises(AttributeError): + await ProxyUpdateSpend.update_spend_logs( + n_retry_times=2, + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging, + logs_to_process=[make_spend_log_row(request_id="a"), make_spend_log_row(request_id="b")], + ) + + assert [row["request_id"] for row in mock_prisma_client.spend_log_transactions] == ["a", "b", "c"] + assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 3 From 97e5438ffba9a8c4e752686fd4fb040180635572 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 20 Oct 2026 09:08:29 -0700 Subject: [PATCH 23/33] fix(proxy): fall back to the deployment model when a billed failure row's base_model is not in the cost map (#47568) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/utils.py | 49 ++++++++++++------- ..._utils_model_creation_and_error_logging.py | 30 +++++++++++- 2 files changed, 61 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d6dfd3ea359..b01f53c8c6b 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1200,6 +1200,25 @@ def _deployment_cost_per_token(deployment_id: str | None, model: str) -> CostPer ) +def _priced_input_cost( + pricing_model: str, + custom_llm_provider: str | None, + prompt_tokens: int, + custom_cost_per_token: CostPerToken | None, +) -> float | None: + try: + estimated_input_cost, _ = litellm.cost_per_token( + model=pricing_model, + custom_llm_provider=custom_llm_provider, + prompt_tokens=prompt_tokens, + custom_cost_per_token=custom_cost_per_token, + ) + except Exception as exc: # noqa: BLE001 # cost_per_token raises BadRequestError for an unmapped model and the row must still be written + verbose_proxy_logger.debug("failure spend estimate: could not price %s: %s", pricing_model, exc) + return None + return estimated_input_cost + + def _estimated_input_cost( model_call_details: Mapping[str, object], prompt_tokens: int, @@ -1207,23 +1226,19 @@ def _estimated_input_cost( base_model: str | None, ) -> float: """Input price for the estimated prompt: the deployment's registered rate when it has one, else - the cost map's rate for ``base_model`` or the deployment model (an Azure deployment name is not - in the cost map, its ``base_model`` is). An unpriceable model keeps the row at $0 rather than - dropping the failure row.""" - try: - model: Final = str(model_call_details.get("model") or "") - estimated_input_cost, _ = litellm.cost_per_token( - model=base_model or model, - custom_llm_provider=str(model_call_details.get("custom_llm_provider") or "") or None, - prompt_tokens=prompt_tokens, - custom_cost_per_token=_deployment_cost_per_token(deployment_id, model), - ) - except Exception as exc: # noqa: BLE001 # cost_per_token raises BadRequestError for an unmapped model and the row must still be written - verbose_proxy_logger.debug( - "failure spend estimate: could not price %s: %s", model_call_details.get("model"), exc - ) - return 0.0 - return estimated_input_cost + the cost map's rate for ``base_model``, else the deployment model, the same candidate order + ``completion_cost`` walks (an Azure deployment name is not in the cost map, its ``base_model`` + is; a ``base_model`` the cost map does not know falls through to the deployment model). A model + no candidate can price keeps the row at $0 rather than dropping the failure row.""" + model: Final = str(model_call_details.get("model") or "") + custom_llm_provider: Final = str(model_call_details.get("custom_llm_provider") or "") or None + custom_cost_per_token: Final = _deployment_cost_per_token(deployment_id, model) + pricing_models: Final = (base_model, model) if base_model is not None else (model,) + candidate_costs: Final = ( + _priced_input_cost(pricing_model, custom_llm_provider, prompt_tokens, custom_cost_per_token) + for pricing_model in pricing_models + ) + return next((cost for cost in candidate_costs if cost is not None), 0.0) _EMPTY_LIFT: Final = MappingProxyType({}) diff --git a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index 5ab04724bc4..99cd569eb63 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -1174,7 +1174,7 @@ from typing import cast import litellm from litellm.proxy.utils import create_model_info_response from litellm.types.router import DeploymentModelListingInfo -from litellm.types.utils import ModelInfo +from litellm.types.utils import ModelInfo, ModelResponse, Usage def _fake_model_info(**fields: object) -> ModelInfo: @@ -2337,6 +2337,34 @@ def test_an_azure_deployment_is_priced_at_its_base_model_like_the_success_path() assert lifted["response_cost"] == pytest.approx(expected_cost) +def test_an_unmapped_base_model_falls_back_to_the_deployment_model_like_the_success_path(): + """``completion_cost`` tries ``base_model`` first and then the deployment model, so a + ``model_info.base_model`` the cost map does not know still prices a success row at the + deployment model's rate. The failure row must land on the same number instead of $0.""" + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.proxy.utils import _failure_fields_to_lift + + logging_obj = _dispatched_failure_logging_obj( + FAILURE_USAGE_MODEL, + get_litellm_params(metadata={"model_info": {"base_model": "not-in-the-cost-map"}}), + ) + logging_obj.model_call_details["custom_llm_provider"] = "openai" + + lifted = _failure_fields_to_lift({"litellm_logging_obj": logging_obj}, provider_billed=True) + + prompt_tokens = lifted["combined_usage_object"].prompt_tokens + success_cost = litellm.completion_cost( + completion_response=ModelResponse( + model=FAILURE_USAGE_MODEL, usage=Usage(prompt_tokens=prompt_tokens, completion_tokens=0) + ), + model=FAILURE_USAGE_MODEL, + custom_llm_provider="openai", + base_model="not-in-the-cost-map", + ) + assert success_cost > 0 + assert lifted["response_cost"] == pytest.approx(success_cost) + + def test_a_litellm_params_key_in_the_request_body_cannot_zero_the_timeout_price(): """``model_call_details`` is flattened with the request's optional_params, so a body key named ``litellm_params`` can land there; the deployment rate must come from the logging From 77e08731eba09b57182283e8bc015c74a449c55b Mon Sep 17 00:00:00 2001 From: Muspi Merol Date: Wed, 21 Oct 2026 00:18:13 +0800 Subject: [PATCH 24/33] test: isolate model cost maps in runtime registration tests (#47559) * test: isolate model cost maps in runtime registration tests * test: restore cost map entries and provider model lists after unit tests * test: restore provider model sets after unit tests and cover the catalog restore * test: snapshot litellm globals before scanning provider model sets --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- tests/unit/conftest.py | 24 ++++++++++++++++++++++++ tests/unit/test_conftest_isolation.py | 20 ++++++++++++++++++++ tests/unit/test_utils.py | 7 ++----- 3 files changed, 46 insertions(+), 5 deletions(-) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 9028b8461b0..c96e3e2c031 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -307,6 +307,30 @@ def isolate_router_model_cost_state() -> Iterator[None]: litellm.get_model_info.cache_clear() +def _provider_model_sets() -> tuple[set[str], ...]: + module_globals: Final = tuple(vars(litellm).items()) + return tuple(models for name, models in module_globals if name.endswith("_models") and isinstance(models, set)) + + +@pytest.fixture(autouse=True) +def isolate_model_catalog() -> Iterator[None]: + model_cost: Final = litellm.model_cost + model_cost_entries: Final = dict(model_cost) + models_by_provider: Final = litellm.models_by_provider + models_by_provider_entries: Final = dict(models_by_provider) + provider_model_sets: Final = tuple((models, frozenset(models)) for models in _provider_model_sets()) + yield + model_cost.clear() + model_cost.update(model_cost_entries) + models_by_provider.clear() + models_by_provider.update(models_by_provider_entries) + for models, entries in provider_model_sets: + models.clear() + models.update(entries) + litellm_utils_module._invalidate_model_cost_lowercase_map() + litellm.get_model_info.cache_clear() + + @pytest.fixture def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") diff --git a/tests/unit/test_conftest_isolation.py b/tests/unit/test_conftest_isolation.py index 63db0eb5bba..92457119a82 100644 --- a/tests/unit/test_conftest_isolation.py +++ b/tests/unit/test_conftest_isolation.py @@ -88,6 +88,26 @@ def test_register_model_ledger_entry_was_rolled_back(): assert CANARY_MODEL not in litellm_utils_module._runtime_registered_model_cost +def test_register_model_catalog_writes_are_scoped_to_this_test(): + litellm.register_model( + { + CANARY_MODEL: {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 0.001}, + "collection-isolation-ledger": {"input_cost_per_token": 0.5}, + } + ) + assert CANARY_MODEL in litellm.model_cost + assert CANARY_MODEL in litellm.open_ai_chat_completion_models + assert CANARY_MODEL in litellm.models_by_provider["openai"] + assert litellm.model_cost["collection-isolation-ledger"]["input_cost_per_token"] == 0.5 + + +def test_register_model_catalog_writes_were_rolled_back(): + assert CANARY_MODEL not in litellm.model_cost + assert CANARY_MODEL not in litellm.open_ai_chat_completion_models + assert CANARY_MODEL not in litellm.models_by_provider["openai"] + assert litellm.model_cost["collection-isolation-ledger"]["input_cost_per_token"] == 0.001 + + def test_live_router_membership_is_scoped_to_this_test(): _CanaryRouterHolder.router = Router( model_list=[ diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 0be3ae05e97..40eb5c31fda 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -2358,7 +2358,6 @@ def test_register_model_with_scientific_notation(): _invalidate_model_cost_lowercase_map() -@pytest.mark.usefixtures("isolated_openai_model_sets") @pytest.mark.usefixtures("isolated_openai_model_sets") @respx.mock def test_register_model_url_fetch_uses_single_attempt(monkeypatch): @@ -4803,8 +4802,7 @@ async def test_builtin_string_callback_registers_when_subclass_already_active( assert any(type(cb) is S3Logger for cb in litellm._async_success_callback) -@pytest.mark.usefixtures("isolated_openai_model_sets") -@pytest.mark.usefixtures("isolated_openai_model_sets") +@pytest.mark.usefixtures("isolated_openai_model_sets", "local_model_cost_map") @pytest.mark.parametrize("rate", (0.0, 0.000123)) @pytest.mark.parametrize("updated_rate", (None, 0.0, 0.000456)) def test_reapply_runtime_registrations_replays_register_model_overrides( @@ -4871,8 +4869,7 @@ def test_reapply_runtime_registrations_replays_register_model_overrides( _invalidate_model_cost_lowercase_map() -@pytest.mark.usefixtures("isolated_openai_model_sets") -@pytest.mark.usefixtures("isolated_openai_model_sets") +@pytest.mark.usefixtures("isolated_openai_model_sets", "local_model_cost_map") def test_reapply_runtime_registrations_drops_request_scoped_registrations(monkeypatch): """ Per-request custom pricing describes one call, so it must not be re-asserted From dd22dc9132ba2d32cb809fbcb482c2b76b74d352 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 20 Oct 2026 09:18:44 -0700 Subject: [PATCH 25/33] test(perf): measure Admin UI journey latency in a browser on the 1TB SpendLogs restore (#47486) * test(perf): measure Admin UI journey latency in a browser on the 1TB restore * test(perf): settle on rendered rows, chart bars and the usage view trigger * test(perf): cover the UI journey summaries in the e2e harness tests * test(perf): fail UI journeys on failed samples and p50/p90 ceilings * test(perf): add Usage, Models and Keys journeys and seed realistic keys, teams, users and models on the disposable restore --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/perf/UI_QUERIES.md | 111 ++--- tests/e2e/perf/setup/ui_journeys.py | 397 ++++++++++++++++++ tests/e2e/perf/setup/ui_queries.py | 227 ---------- tests/e2e/perf/setup/ui_query_cases.py | 206 --------- tests/e2e/perf/setup/ui_query_fixture.py | 72 +++- tests/e2e/perf/setup/ui_query_isolation.py | 4 +- tests/e2e/perf/test_ui_journeys.py | 108 +++++ tests/e2e/perf/test_ui_queries.py | 159 ------- .../perf/setup/test_ui_journeys.py | 144 +++++++ .../e2e_harness/perf/setup/test_ui_queries.py | 257 ------------ tests/e2e_harness/perf/test_upgrade.py | 4 +- 11 files changed, 761 insertions(+), 928 deletions(-) create mode 100644 tests/e2e/perf/setup/ui_journeys.py delete mode 100644 tests/e2e/perf/setup/ui_queries.py delete mode 100644 tests/e2e/perf/setup/ui_query_cases.py create mode 100644 tests/e2e/perf/test_ui_journeys.py delete mode 100644 tests/e2e/perf/test_ui_queries.py create mode 100644 tests/e2e_harness/perf/setup/test_ui_journeys.py delete mode 100644 tests/e2e_harness/perf/setup/test_ui_queries.py diff --git a/tests/e2e/perf/UI_QUERIES.md b/tests/e2e/perf/UI_QUERIES.md index 552f1b953fd..27780537042 100644 --- a/tests/e2e/perf/UI_QUERIES.md +++ b/tests/e2e/perf/UI_QUERIES.md @@ -1,92 +1,55 @@ -# SpendLog UI query performance +# Admin UI journey latency on the 1TB SpendLogs restore -Run `test_ui_queries.py` against a dedicated proxy and a disposable database restored from the metadata-only performance seed. These tests do not call an LLM or measure browser rendering +Run `test_ui_journeys.py` against a dedicated proxy, its Admin UI and a disposable database restored from the metadata-only performance seed. A headless Chromium logs into the dashboard and performs the journeys below, timing each one from the user's click or navigation until the data is on screen. No LLM is called. The test does not look at which routes, SQL or tables served the page, so it stays valid whether the backend gets faster queries or new aggregate tables -The suite measures 42 request shapes across 20 routes used by current dashboard screens. It replays direct HTTP requests against the existing bulk seed, with the seed anchor replacing the wall clock and a UTC browser timezone. It does not run a browser or claim to cover every dashboard request +## Journeys -## UI callers +| name | what the user does | timer stops when | +| --- | --- | --- | +| `logs_open_24h` | opens Logs (default Last 24 Hours) | request rows are rendered and the refresh control is idle | +| `logs_last_7_days` | switches the time range to Last 7 Days | same | +| `logs_next_page` | clicks the next page control | same | +| `logs_filter_key_hash` | filters by the seeded key hash and applies | same | +| `logs_filter_key_alias` | picks the first key alias in the filter drawer and applies | same | +| `logs_filter_team` | picks the first team in the filter drawer and applies | same | +| `logs_filter_model` | picks the first model in the filter drawer and applies | same | +| `usage_open` | opens Usage | summary tiles show dollar values with no skeletons left | +| `usage_last_30_days` | picks Last 30 days in the Usage date picker | same | +| `usage_team_view` | switches the Usage view to Team Usage | same | +| `usage_customer_view` | switches the Usage view to Customer Usage | same | +| `usage_model_activity_30_days` | on the Model Activity tab, picks Last 30 days | same | +| `cache_open` | opens Caching | the Cache Hits tile is visible and the Cache Hits vs API Requests chart has bars | +| `models_open` | opens Models + Endpoints | model rows are rendered and the loading state is gone | +| `keys_open` | opens Virtual Keys | key rows are rendered and the loading state is gone | -All paths below are relative to `ui/litellm-dashboard/src/`. Each case requires an authenticated admin session with access to that screen, and selected resources must exist in the fixture +The mix follows what customers report as slow: the Logs page on `LiteLLM_SpendLogs`, the Usage page on the `LiteLLM_Daily*Spend` rollups (30 day ranges, Customer and Model views), and the Models and Keys pages. -| Cases | UI action and source | -| --- | --- | -| `logs_*` | Logs > Request Logs, change date range, filter, sort or next page. `components/logs/request/useLogFilterLogic.ts` and `RequestLogsFilters.tsx` | -| `usage_*` | Usage overview or entity tab, open key activity, search keys, expand a model group. `app/(dashboard)/usage/_components/components/UsagePageView.tsx`, `EntityUsage/EntityUsage.tsx`, `EntityUsage/entityFetchFns.ts`, `components/UsagePage/components/KeyActivityPanel.tsx` | -| `team_users_*` | Team Usage. `EntityUsage/EntityUsage.tsx` passes selected teams or all available non-dashboard teams to `TeamUserSpendCard.tsx` | -| `cache_activity_*` | Caching > Cache Analytics, optionally filter aliases/models. `app/(dashboard)/hooks/caching/useCacheActivity.ts` | -| `cache_leakage_keys` | Cost Optimization, switch cache leakage to Keys. `app/(dashboard)/cost-optimization/_components/useCacheLeakageKeys.ts` | +## Seeded rows -Request serialization lives in `components/networking.tsx`. Logs use page size 25 and session grouping, including the next-session cursor returned by page one. Individual Logs filters use a 30-day window to exercise larger matching datasets. Their case names end in `_30d` so the measurements do not share history with the old 24-hour filters. The synthetic cursorless page-three request is excluded +The snapshot carries the SpendLogs rows and their Daily* rollups. `setup/ui_query_fixture.py` checks that `LiteLLM_DailyUserSpend` and `LiteLLM_DailyTeamSpend` have rows in the week before the anchor, then fills the tables the snapshot does not have so the Keys, Models and filter journeys see realistic counts: 100 teams (`perf-team-N`), 2,000 users (`perf-user-N`), 5,000 keys whose token hashes and `perf-key-N` aliases match the seeded log rows, and 200 DB models pointed at the perf mock provider. Every insert is `ON CONFLICT DO NOTHING`, so reruns on the same restore are no-ops. Nothing here runs outside the disposable restore -Usage list, search and model drilldown requests do not inject an `api_key` filter. Team activity excludes `litellm-dashboard`. Cache leakage supplies `include_current_utc_day=true`. Dates and filter values are deterministic fixture values, not captured production traffic +Each journey is measured once cold, then `E2E_UI_JOURNEY_REPEATS` more times (default 9) from the same starting state, inside one logged-in tab whose clock is pinned to the seed anchor so relative windows land on seeded rows. The browser timezone is UTC to match the seed -Old Usage is excluded because its expensive-query gate normally hides those calls at this database size. LiteAdmin's agent-selected reports are excluded from this deterministic dashboard set. Missing-token alias recovery is backend functional coverage, not a separate UI action, and its artificial requests and orphan fixtures are removed +## Safety -This mapping is source-verified, not browser-recorded. Recheck the named callers when changing cases. Neither browser capture nor an automatic UI-drift checker is installed by this suite +The seeded source is never the target. `E2E_UI_PERF_DISPOSABLE=1`, `E2E_RESET_SPEND_LOGS=0` and a hostname matching `db-perf-perf-.litellm.svc` are all required, and in-flight SQL is cancelled only on that disposable restore between journeys (see `setup/ui_query_isolation.py`) -## Dataset scope +## Running -Logs, team-user spend and cache activity query the bulk SpendLogs data. User, team, tag and customer usage, plus cache leakage, query daily tables aggregated from that dataset. These are sequential large-dataset queries, not concurrent traffic or saturation tests - -A read-only restore of `db-perf-v1-105-1tb` gave exact request totals from daily aggregates: 9,442,606 across 24 hours, 66,077,435 across seven days and 283,189,684 across 30 days. The retained 30-day key filter matches 1,134,559 requests, user filter 11,316,454, team filter 113,298,119 and model filter 94,395,473. These describe matching requests, not rows scanned or returned by the API, and are documentation rather than performance assertions - -PostgreSQL table-wide estimates (`pg_class.reltuples`) are approximately 2.70 million daily user rows, 2.70 million daily team rows, 5.40 million daily tag rows and 78.27 million daily end-user rows. These are not exact counts within each Usage case's seven-day window - -The individual end-user filter is excluded because it matched only 104 requests in 24 hours and 8,616 across all dates. User/end-user dropdown requests are also excluded because their queries cap the input at 10,000 matching SpendLogs before extracting distinct values. Neither is a broad matching-data workload in this seed - -Excluded because they only had tiny matching fixtures or individual-record lookups: organization/agent usage, guardrail/tool logs, request-ID lookup, payload/session drawers, session/search filters for `perf-ui-*`, the OAuth-token filter, and all prompt-caching request cases. The seed has response-cache hits but no prompt-cache usage or injection metadata, so the prompt-caching cases only matched the 320 added rows. No UI-specific rows are inserted anymore - -| Remaining cases | Count | Variations | -| --- | ---: | --- | -| Logs | 21 | 24h/7d/30d, page 2 with session cursor, 12 individual filters over 30 days, five additional descending sorts | -| User/team/tag/customer usage | 16 | Aggregated totals, keys, key search, model top keys for each | -| Team-user spend | 2 | All 100 seeded teams or `perf-team-0` | -| Cache activity | 2 | All or selected aliases/model | -| Cache leakage keys | 1 | 30-day window | - -## Deployment prerequisites - -Use a fresh proxy against the disposable restore, never the active seed source or a shared test database. Apply the target migrations and wait for readiness first. The source seed's `litellm_perf_seed.checkpoint` must exist and contain rows - -Configure `litellm_settings.service_callback: [prometheus_system]` and expose metrics for every proxy worker/pod. Set `E2E_DATABASE_METRICS_URLS` to their comma-separated URLs, not a load-balanced endpoint that can miss a pod. Absent failure counters count as zero only when the endpoint exposes a positive PostgreSQL request counter. Missing instrumentation stops the run - -No expensive-query override or legacy view creation is needed for these cases. Keep the normal large-database UI restrictions - -The cases use explicit dates from the seed anchor where the UI supplies dates. A missing nonempty fixture result records a contract failure rather than reporting a fast empty query as success - -## Run - -Set `DATABASE_URL` and `LITELLM_MASTER_KEY` securely for this disposable deployment. The URL must name the same database used by the proxy, with hostname `db-perf-perf-.litellm.svc`. The explicit disposable confirmation permits cancellation of leftover queries in this database. Setup reads the existing seed checkpoint without adding UI-specific rows. Destroy the disposable database after the run +The componentized perf stack serves the dashboard behind the same router as the proxy, so `LITELLM_PROXY_URL` covers `/ui`. Set `E2E_UI_BASE_URL` only if the UI lives elsewhere. The admin login uses `E2E_UI_USERNAME` (default `admin`) and `E2E_UI_PASSWORD` (default the master key) ```bash -export LITELLM_PROXY_URL=http://localhost:4000 -export E2E_DATABASE_METRICS_URLS=http://localhost:4000/metrics -export E2E_PERF_TESTS=1 -export E2E_UI_PERF_DISPOSABLE=1 -export E2E_RESET_SPEND_LOGS=0 -export E2E_ARTIFACT_DIR="$PWD/ui-query-results" -mkdir -p "$E2E_ARTIFACT_DIR" -uv run pytest tests/e2e/perf/test_ui_queries.py --reruns 0 --junitxml="$E2E_ARTIFACT_DIR/junit.xml" -q +export E2E_PERF_TESTS=1 E2E_UI_PERF_DISPOSABLE=1 E2E_RESET_SPEND_LOGS=0 +export DATABASE_URL=postgresql://...@db-perf-perf-.litellm.svc:5432/litellm +export E2E_DATABASE_METRICS_URLS=http://:4000/metrics +export E2E_ARTIFACT_DIR=/tmp/e2e-artifacts +uv run pytest tests/e2e/perf/test_ui_journeys.py --reruns 0 --junitxml="$E2E_ARTIFACT_DIR/junit.xml" -q ``` -Do not enable SpendLog retention or run other traffic on this stack. The `quiet_stack` marker coordinates local pytest callers but cannot stop external traffic. Use one pytest worker for comparable timings +`E2E_UI_JOURNEY_TIMEOUT_SECONDS` (default 120) bounds one journey. A journey that never settles is recorded as a `timeout` sample and a journey whose page rendered empty or threw in the browser as a `ui` sample. The test fails when any sample failed, or when the p50 or p90 across all samples is over the journey's `Ceiling` in `UI_JOURNEY_LIMITS` (`setup/ui_journeys.py`, 10s/20s for Logs journeys, 20s/40s for Usage and Caching, 5s/10s for the Models and Keys tables, generous until the first 1TB run sets real numbers). The artifacts are written before the assertion so a failed run still ships its samples -`E2E_UI_QUERY_REPEATS` defaults to 9, giving ten measured requests including the first, and `E2E_UI_QUERY_TIMEOUT_SECONDS` defaults to 120. Every request is issued once without transport retries. A failed first request stops that case's repeated samples. Cursor prerequisites are recorded separately, and failed prerequisites stop that case's measurements +## Artifacts -## Reading results +Each JUnit testcase `test_ui_journey_latency[]` carries `perf_ui_journeys`, `perf_ui_failures`, `perf_ui_seconds_p50`, `perf_ui_seconds_p90`, `perf_ui_seconds_max`, `perf_ui_first_seconds`, `perf_ui_repeated_seconds_p50`, `perf_ui_repeated_seconds_max` and `perf_ui_database_failures`. `ui-journey-.json` holds every sample with its on-screen row count and `ui-isolation-.jsonl` the SQL isolation evidence -Each case writes `ui-query-.json` with its exact method, path, parameters and body, response status counts, elapsed times, row counts, errors and before/after PostgreSQL failure counters. Response bodies and authentication headers are not saved - -The dashboard's `perf_ui_seconds_p50`, `_p90` and `_max` include all successful measured requests, including the first. `perf_ui_samples` is their sample count. p50 is the median and p90 uses nearest rank, the ninth sorted value with ten successes. Failures and cursor prerequisites contribute no latency. Pass percentage is `(perf_ui_requests - perf_ui_failures) / perf_ui_requests * 100`, including cursor prerequisites in the request counts - -The first-request and repeated-request metrics are retained for existing consumers. This is not a cold-cache claim: previous cases and previous runs can warm database buffers or application caches. Restart the disposable proxy between independent runs when comparing independent runs - -Repeated p50 is reported for successful samples, p95 only with at least 20 and p99 only with at least 100. First-request latency is omitted when that request fails. HTTP, JSON, response-shape, nonempty-result, cursor failures are recorded, never asserted. Endpoint failures, database counter increases and latency do not fail CI - -Each JUnit testcase `test_ui_query_performance[]` carries `perf_ui_requests`, `perf_ui_failures`, `perf_ui_repeated_samples`, and nonzero `perf_ui_error_` counts. Codes distinguish `http_400`, `http_500`, `http_503`, other HTTP statuses, `timeout`, `network`, `invalid_json` and `contract`. Successful latency properties are `perf_ui_first_seconds` and `perf_ui_repeated_seconds_p50`, `_p95`, `_p99`, `_max`. A failed cursor prerequisite is counted once and contributes no latency - -Before every case the runner requires three idle database observations and stable failure counters. After each case it cancels leftover SQL, terminates idle transactions and waits for idle and stable counters again. `ui-isolation-.jsonl` retains active SQL, waits, elapsed times, raw samples and counter observations. `perf_ui_database_failures` measures the increase before cleanup, while `perf_ui_cleanup_database_failures` measures the increase observed afterward, including cancellations and potentially delayed request failures. Isolation, instrumentation or setup failures still stop the run rather than publish contaminated measurements. Disk and process caches are not reset between cases - -The Buildkite `ui_queries` lane restores the pinned snapshot independently, waits for migrations and indexes to complete, runs these cases with no retries, uploads JUnit and JSON artifacts, and removes the temporary resources. It follows the nightly regular perf and database-upgrade builds using the same LiteLLM SHA. Shinkiro ingestion is a separate change - -Validate a small restore first. Small-database numbers validate the runner only, not full-size performance. Run the same cases against the completed full-size snapshot before publishing large-database results +The Buildkite `ui_queries` lane restores the pinned snapshot, waits for migrations and indexes, runs these journeys with no retries and uploads the JUnit and JSON artifacts. Validate selectors on a small restore first, then take numbers from the completed 1TB snapshot diff --git a/tests/e2e/perf/setup/ui_journeys.py b/tests/e2e/perf/setup/ui_journeys.py new file mode 100644 index 00000000000..e9d4360d1a4 --- /dev/null +++ b/tests/e2e/perf/setup/ui_journeys.py @@ -0,0 +1,397 @@ +from __future__ import annotations + +import math +import re +import statistics +import time +from collections.abc import Callable, Generator +from contextlib import contextmanager +from dataclasses import dataclass +from datetime import datetime +from itertools import chain +from types import MappingProxyType +from typing import Final + +from playwright.sync_api import Error as PlaywrightError +from playwright.sync_api import Locator, Page, Response, expect, sync_playwright +from playwright.sync_api import TimeoutError as PlaywrightTimeoutError +from prometheus_client.parser import text_string_to_metric_families +from pydantic import BaseModel, ConfigDict + +MONEY: Final = re.compile(r"^\$[\d,]+\.\d{2}$") +LOADING_LOGS: Final = "Loading request logs…" +LOADING_MODELS: Final = "Loading models…" +LOADING_KEYS: Final = "Loading keys..." + + +class Sample(BaseModel): + model_config = ConfigDict(frozen=True) + + seconds: float + rows: int = 0 + error: str = "" + + +class Timings(BaseModel): + model_config = ConfigDict(frozen=True) + + first_seconds: float | None + repeated_samples: int + repeated_p50_seconds: float | None + repeated_p95_seconds: float | None + repeated_p99_seconds: float | None + repeated_max_seconds: float | None + + +@dataclass(frozen=True, slots=True) +class Journey: + """One dashboard interaction. ``start`` brings the page to the state the user acts from (untimed), + ``act`` is the click or navigation being measured and ``settled`` blocks until the data the user + was waiting for is on screen, returning the number of rows or tiles that rendered.""" + + name: str + surface: str + start: Callable[[Page, float], None] + act: Callable[[Page], None] + settled: Callable[[Page, float], int] + + +@dataclass(frozen=True, slots=True) +class Dashboard: + base_url: str + username: str + password: str + anchor: datetime + + +def summarize(samples: tuple[Sample, ...]) -> Timings: + assert samples, "No UI journeys were measured" + repeated: Final = tuple(sorted(sample.seconds for sample in samples[1:] if not sample.error)) + return Timings( + first_seconds=samples[0].seconds if not samples[0].error else None, + repeated_samples=len(repeated), + repeated_p50_seconds=statistics.median(repeated) if repeated else None, + repeated_p95_seconds=repeated[math.ceil(len(repeated) * 0.95) - 1] if len(repeated) >= 20 else None, + repeated_p99_seconds=repeated[math.ceil(len(repeated) * 0.99) - 1] if len(repeated) >= 100 else None, + repeated_max_seconds=max(repeated) if repeated else None, + ) + + +def error_code(sample: Sample) -> str: + if not sample.error: + return "" + return "timeout" if sample.error == "timeout" else "ui" + + +def ui_junit_properties(samples: tuple[Sample, ...]) -> tuple[tuple[str, float], ...]: + codes: Final = tuple(error_code(sample) for sample in samples if sample.error) + timings: Final = summarize(samples) + successful: Final = tuple(sorted(sample.seconds for sample in samples if not sample.error)) + latency: Final = ( + ("seconds_p50", statistics.median(successful) if successful else None), + ("seconds_p90", successful[math.ceil(len(successful) * 0.90) - 1] if successful else None), + ("seconds_max", max(successful) if successful else None), + ("first_seconds", timings.first_seconds), + ("repeated_seconds_p50", timings.repeated_p50_seconds), + ("repeated_seconds_p95", timings.repeated_p95_seconds), + ("repeated_seconds_p99", timings.repeated_p99_seconds), + ("repeated_seconds_max", timings.repeated_max_seconds), + ) + return ( + ("perf_ui_journeys", len(samples)), + ("perf_ui_failures", len(codes)), + ("perf_ui_samples", len(successful)), + ("perf_ui_repeated_samples", timings.repeated_samples), + *((f"perf_ui_error_{code}", codes.count(code)) for code in sorted(frozenset(codes))), + *((f"perf_ui_{name}", value) for name, value in latency if value is not None), + ) + + +@dataclass(frozen=True, slots=True) +class Ceiling: + p50_seconds: float + p90_seconds: float + + +LOGS_CEILING: Final = Ceiling(p50_seconds=10, p90_seconds=20) +CHART_CEILING: Final = Ceiling(p50_seconds=20, p90_seconds=40) +TABLE_CEILING: Final = Ceiling(p50_seconds=5, p90_seconds=10) +UI_JOURNEY_LIMITS: Final = MappingProxyType( + { + "logs_open_24h": LOGS_CEILING, + "logs_last_7_days": LOGS_CEILING, + "logs_next_page": LOGS_CEILING, + "logs_filter_key_hash": LOGS_CEILING, + "logs_filter_key_alias": LOGS_CEILING, + "logs_filter_team": LOGS_CEILING, + "logs_filter_model": LOGS_CEILING, + "usage_open": CHART_CEILING, + "usage_last_30_days": CHART_CEILING, + "usage_team_view": CHART_CEILING, + "usage_customer_view": CHART_CEILING, + "usage_model_activity_30_days": CHART_CEILING, + "cache_open": CHART_CEILING, + "models_open": TABLE_CEILING, + "keys_open": TABLE_CEILING, + } +) + + +def assert_within_ceiling(journey_name: str, samples: tuple[Sample, ...]) -> None: + failed: Final = tuple(sample.error for sample in samples if sample.error) + assert not failed, f"{len(failed)} of {len(samples)} {journey_name} samples failed: {failed[0]}" + metrics: Final = dict(ui_junit_properties(samples)) + ceiling: Final = UI_JOURNEY_LIMITS[journey_name] + assert metrics["perf_ui_seconds_p50"] <= ceiling.p50_seconds, ( + f"{journey_name} p50 {metrics['perf_ui_seconds_p50']:.1f}s over {ceiling.p50_seconds}s" + ) + assert metrics["perf_ui_seconds_p90"] <= ceiling.p90_seconds, ( + f"{journey_name} p90 {metrics['perf_ui_seconds_p90']:.1f}s over {ceiling.p90_seconds}s" + ) + + +@contextmanager +def dashboard_page(dashboard: Dashboard, timeout: float) -> Generator[Page]: + """A logged-in admin tab whose clock reads the seed anchor, so relative windows such as + Last 24 Hours land on seeded rows. Live Tail is off so its 15 second refetch cannot land + inside a sample. Playwright keeps a UTC timezone to match the seed.""" + with sync_playwright() as playwright: + browser: Final = playwright.chromium.launch() + context: Final = browser.new_context( + base_url=dashboard.base_url, timezone_id="UTC", viewport={"width": 1600, "height": 1000} + ) + context.set_default_timeout(timeout * 1000) + context.add_init_script('sessionStorage.setItem("isLiveTail", "false")') + page: Final = context.new_page() + page.clock.set_fixed_time(dashboard.anchor) + _login(page, dashboard) + try: + yield page + finally: + context.close() + browser.close() + + +def _login(page: Page, dashboard: Dashboard) -> None: + _ = page.goto("/ui/login") + page.get_by_placeholder("Enter your username").fill(dashboard.username) + page.get_by_placeholder("Enter your password").fill(dashboard.password) + page.get_by_role("button", name="Login", exact=True).click() + page.wait_for_url(lambda url: "/ui" in url and "/login" not in url) + expect(page.get_by_label("LiteLLM home")).to_be_visible() + + +def _is_data_fetch(response: Response) -> bool: + return response.request.resource_type in ("fetch", "xhr") and "/_next/" not in response.url + + +def measure(page: Page, journey: Journey, timeout: float) -> Sample: + """Wall time from the user's action to the data being on screen. Every sample starts from the + same pre-action state so repeats measure the same journey, not a warmed client cache.""" + try: + journey.start(page, timeout) + except (PlaywrightError, AssertionError) as error: + return Sample(seconds=0.0, error=f"start {type(error).__name__}: {error}"[:300]) + started: Final = time.monotonic() + try: + with page.expect_response(_is_data_fetch, timeout=timeout * 1000): + journey.act(page) + rows: Final = journey.settled(page, timeout) + except PlaywrightTimeoutError: + return Sample(seconds=time.monotonic() - started, error="timeout") + except (PlaywrightError, AssertionError) as error: + return Sample(seconds=time.monotonic() - started, error=f"{type(error).__name__}: {error}"[:300]) + return Sample(seconds=time.monotonic() - started, rows=rows) + + +def _visible(page: Page, test_id: str) -> Locator: + return page.get_by_test_id(test_id).filter(visible=True).first + + +def _logs_rows(page: Page) -> Locator: + return page.locator("table").filter(visible=True).first.locator("tbody tr") + + +def _logs_settled(page: Page, timeout: float) -> int: + millis: Final = timeout * 1000 + rows: Final = _logs_rows(page).filter(has_not_text="No results") + expect(rows.first).to_be_visible(timeout=millis) + expect(_visible(page, "datatable-refresh")).to_be_enabled(timeout=millis) + expect(page.get_by_text(LOADING_LOGS)).to_have_count(0, timeout=millis) + return rows.count() + + +def _no_skeletons(page: Page, timeout: float) -> None: + expect(page.locator('[data-slot="skeleton"]')).to_have_count(0, timeout=timeout * 1000) + + +def _usage_settled(page: Page, timeout: float) -> int: + millis: Final = timeout * 1000 + expect(page.get_by_text("Total Spend", exact=True).first).to_be_visible(timeout=millis) + _no_skeletons(page, timeout) + money: Final = page.get_by_text(MONEY) + expect(money.first).to_be_visible(timeout=millis) + return money.count() + + +def _cache_settled(page: Page, timeout: float) -> int: + millis: Final = timeout * 1000 + expect(page.get_by_text("Cache Hits", exact=True)).to_be_visible(timeout=millis) + bars: Final = page.locator(".recharts-bar-rectangle") + expect(bars.first).to_be_visible(timeout=millis) + return bars.count() + + +def _table_settled(loading: str) -> Callable[[Page, float], int]: + def settled(page: Page, timeout: float) -> int: + millis: Final = timeout * 1000 + rows: Final = page.locator("[data-row-id]") + expect(rows.first).to_be_visible(timeout=millis) + expect(page.get_by_text(loading)).to_have_count(0, timeout=millis) + return rows.count() + + return settled + + +def _open_logs(page: Page, timeout: float) -> None: + _ = page.goto("/ui/logs") + _ = _logs_settled(page, timeout) + + +def _open_usage(page: Page, timeout: float) -> None: + _ = page.goto("/ui/usage") + _ = _usage_settled(page, timeout) + + +def _filters_drawer(page: Page) -> Locator: + _visible(page, "datatable-filters-trigger").click() + drawer: Final = page.get_by_role("dialog", name="Filters") + expect(drawer).to_be_visible() + return drawer + + +def _apply_filters(page: Page) -> None: + page.get_by_role("dialog", name="Filters").get_by_role("button", name="Apply Filters").click() + + +def _pick_first_option(page: Page, combobox: str) -> None: + _filters_drawer(page).get_by_role("combobox", name=combobox).click() + page.get_by_role("option").first.click() + + +def _open_time_range(page: Page, timeout: float) -> None: + _open_logs(page, timeout) + page.get_by_role("button", name="Last 24 Hours").click() + + +def _open_usage_range(page: Page, timeout: float) -> None: + _open_usage(page, timeout) + page.locator('[data-slot="advanced-date-picker-trigger"]').click() + + +def _model_activity_settled(page: Page, timeout: float) -> int: + millis: Final = timeout * 1000 + expect(page.locator('[role="tab"][data-active]', has_text="Model Activity")).to_be_visible(timeout=millis) + _no_skeletons(page, timeout) + money: Final = page.get_by_text(MONEY).filter(visible=True) + expect(money.first).to_be_visible(timeout=millis) + return money.count() + + +def _open_model_activity_range(page: Page, timeout: float) -> None: + _open_usage(page, timeout) + page.get_by_role("tab", name="Model Activity").click() + _ = _model_activity_settled(page, timeout) + page.locator('[data-slot="advanced-date-picker-trigger"]').click() + + +def _pick_usage_preset(preset: str) -> Callable[[Page], None]: + def act(page: Page) -> None: + page.locator('[data-slot="advanced-date-picker-preset"]', has_text=preset).click() + + return act + + +def _select_usage_view(page: Page, view: str) -> None: + page.locator('[data-slot="select-trigger"]').filter(has_text="Usage").first.click() + page.get_by_role("option", name=re.compile(rf"^{re.escape(view)}")).click() + + +def _nothing(page: Page, timeout: float) -> None: + return None + + +def _goto(path: str) -> Callable[[Page], None]: + def act(page: Page) -> None: + _ = page.goto(path) + + return act + + +def _logs_filter(name: str, prepare: Callable[[Page], None]) -> Journey: + def start(page: Page, timeout: float) -> None: + _open_logs(page, timeout) + prepare(page) + + return Journey(name, "logs", start, _apply_filters, _logs_settled) + + +def journeys(api_key: str) -> tuple[Journey, ...]: + return ( + Journey("logs_open_24h", "logs", _nothing, _goto("/ui/logs"), _logs_settled), + Journey( + "logs_last_7_days", + "logs", + _open_time_range, + lambda page: page.get_by_role("button", name="Last 7 Days", exact=True).click(), + _logs_settled, + ), + Journey( + "logs_next_page", "logs", _open_logs, lambda page: _visible(page, "pagination-next").click(), _logs_settled + ), + _logs_filter( + "logs_filter_key_hash", + lambda page: _filters_drawer(page).get_by_placeholder("Enter key hash…").fill(api_key), + ), + _logs_filter("logs_filter_key_alias", lambda page: _pick_first_option(page, "Search a key alias")), + _logs_filter("logs_filter_team", lambda page: _pick_first_option(page, "Search or select a team")), + _logs_filter("logs_filter_model", lambda page: _pick_first_option(page, "Search a model")), + Journey("usage_open", "usage", _nothing, _goto("/ui/usage"), _usage_settled), + Journey("usage_last_30_days", "usage", _open_usage_range, _pick_usage_preset("Last 30 days"), _usage_settled), + Journey( + "usage_team_view", "usage", _open_usage, lambda page: _select_usage_view(page, "Team Usage"), _usage_settled + ), + Journey( + "usage_customer_view", + "usage", + _open_usage, + lambda page: _select_usage_view(page, "Customer Usage"), + _usage_settled, + ), + Journey( + "usage_model_activity_30_days", + "usage", + _open_model_activity_range, + _pick_usage_preset("Last 30 days"), + _model_activity_settled, + ), + Journey("cache_open", "caching", _nothing, _goto("/ui/caching"), _cache_settled), + Journey("models_open", "models", _nothing, _goto("/ui/models-and-endpoints"), _table_settled(LOADING_MODELS)), + Journey("keys_open", "keys", _nothing, _goto("/ui/api-keys"), _table_settled(LOADING_KEYS)), + ) + + +def database_failures(text: str) -> float: + families: Final = tuple(text_string_to_metric_families(text)) + required: Final = frozenset(("litellm_postgres_failed_requests", "litellm_batch_write_to_db_failed_requests")) + missing: Final = required - frozenset(family.name for family in families) + all_samples: Final = tuple(chain.from_iterable(family.samples for family in families)) + has_activity: Final = any( + sample.name == "litellm_postgres_total_requests_total" and math.isfinite(sample.value) and sample.value > 0 + for sample in all_samples + ) + assert not missing or has_activity, f"Missing PostgreSQL failure counters {missing}" + names: Final = frozenset(f"{name}_total" for name in required) + values: Final = tuple(sample.value for sample in all_samples if sample.name in names) + assert all(math.isfinite(value) and value >= 0 for value in values), "Invalid PostgreSQL failure counter" + return sum(values) diff --git a/tests/e2e/perf/setup/ui_queries.py b/tests/e2e/perf/setup/ui_queries.py deleted file mode 100644 index b6945d79896..00000000000 --- a/tests/e2e/perf/setup/ui_queries.py +++ /dev/null @@ -1,227 +0,0 @@ -from __future__ import annotations - -import math -import statistics -import time -from collections.abc import Callable -from dataclasses import dataclass -from itertools import chain -from typing import Final, Literal -from urllib.parse import urlencode - -from e2e_http import ( - NetworkError, - RateLimitedError, - Result, - Success, - UnauthorizedError, - UnknownApiError, - ValidationError, - without_retries, -) -from e2e_metadata import step -from prometheus_client.parser import text_string_to_metric_families -from pydantic import BaseModel, ConfigDict, JsonValue, RootModel -from transport import HttpTransport - - -class Query(BaseModel): - model_config = ConfigDict(frozen=True) - - -class Document(RootModel[JsonValue]): - model_config = ConfigDict(frozen=True) - - -@dataclass(frozen=True, slots=True) -class Contract: - collection: tuple[str, ...] = () - row_fields: tuple[str, ...] = () - positive: tuple[str, ...] = () - object_fields: tuple[str, ...] = () - excluded_ids: frozenset[str] = frozenset() - - def count(self, document: Document) -> int: - value: Final = descend(document.root, self.collection) - if self.object_fields: - assert isinstance(value, dict), "Expected a JSON object" - assert set(self.object_fields) <= value.keys(), f"Missing fields {self.object_fields}" - return 1 - assert isinstance(value, list) and value, f"Empty/missing collection {self.collection}, check the fixture" - for row in value: - if self.excluded_ids: - assert isinstance(row, dict) and isinstance(row_id := row.get("request_id"), str), "Missing request_id" - assert row_id not in self.excluded_ids, "Repeated page-one row" - if self.row_fields or self.positive: - assert isinstance(row, dict), "Expected object rows" - assert set(self.row_fields) <= row.keys(), f"Missing row fields {self.row_fields}" - for name in self.positive: - assert any(tuple(_positive(row, name) for row in value)), f"Expected positive {name}" - return len(value) - - -def _positive(row: JsonValue, name: str) -> bool: - assert isinstance(row, dict), "Expected object row" - number: Final = row.get(name) - assert isinstance(number, (int, float)) and not isinstance(number, bool), f"Expected numeric {name}" - assert math.isfinite(number), f"Expected finite {name}" - return number > 0 - - -def descend(value: JsonValue, path: tuple[str, ...]) -> JsonValue: - if not path: - return value - assert isinstance(value, dict) and path[0] in value, f"Missing response path {path}" - return descend(value[path[0]], path[1:]) - - -@dataclass(frozen=True, slots=True) -class UiQueryCase: - name: str - path: str - params: BaseModel - contract: Contract - surface: Literal["current"] = "current" - body: BaseModel | None = None - cursor: Literal["none", "session"] = "none" - - -class Sample(BaseModel): - model_config = ConfigDict(frozen=True) - - seconds: float - status: int - rows: int = 0 - error: str = "" - - -class Timings(BaseModel): - model_config = ConfigDict(frozen=True) - - first_seconds: float | None - repeated_samples: int - repeated_p50_seconds: float | None - repeated_p95_seconds: float | None - repeated_p99_seconds: float | None - repeated_max_seconds: float | None - - -def summarize(samples: tuple[Sample, ...]) -> Timings: - assert samples, "No UI requests were measured" - repeated: Final = tuple(sorted(sample.seconds for sample in samples[1:] if not sample.error)) - return Timings( - first_seconds=samples[0].seconds if not samples[0].error else None, - repeated_samples=len(repeated), - repeated_p50_seconds=statistics.median(repeated) if repeated else None, - repeated_p95_seconds=repeated[math.ceil(len(repeated) * 0.95) - 1] if len(repeated) >= 20 else None, - repeated_p99_seconds=repeated[math.ceil(len(repeated) * 0.99) - 1] if len(repeated) >= 100 else None, - repeated_max_seconds=max(repeated) if repeated else None, - ) - - -def error_code(sample: Sample) -> str: - if not sample.error: - return "" - if sample.status > 0 and sample.status != 200: - return f"http_{sample.status}" - if sample.error in ("timeout", "network"): - return sample.error - if sample.error == "Response is not valid JSON": - return "invalid_json" - return "contract" - - -def ui_junit_properties( - samples: tuple[Sample, ...], prerequisite: Sample | None = None -) -> tuple[tuple[str, float], ...]: - requests: Final = (prerequisite, *samples) if prerequisite is not None and not prerequisite.error else samples - codes: Final = tuple(error_code(sample) for sample in requests if sample.error) - timings: Final = summarize(samples) - successful: Final = tuple(sorted(sample.seconds for sample in samples if not sample.error)) - latency: Final = ( - ("seconds_p50", statistics.median(successful) if successful else None), - ("seconds_p90", successful[math.ceil(len(successful) * 0.90) - 1] if successful else None), - ("seconds_max", max(successful) if successful else None), - ("first_seconds", timings.first_seconds), - ("repeated_seconds_p50", timings.repeated_p50_seconds), - ("repeated_seconds_p95", timings.repeated_p95_seconds), - ("repeated_seconds_p99", timings.repeated_p99_seconds), - ("repeated_seconds_max", timings.repeated_max_seconds), - ) - return ( - ("perf_ui_requests", len(requests)), - ("perf_ui_failures", len(codes)), - ("perf_ui_samples", len(successful)), - ("perf_ui_repeated_samples", timings.repeated_samples), - *((f"perf_ui_error_{code}", codes.count(code)) for code in sorted(frozenset(codes))), - *((f"perf_ui_{name}", value) for name, value in latency if value is not None), - ) - - -def _failure(result: Result[Document], seconds: float) -> Sample: - match result: - case Success(status_code=status): - return Sample(seconds=seconds, status=status, error=f"Unexpected HTTP {status}") - case UnknownApiError(status_code=status): - return Sample(seconds=seconds, status=status, error=f"HTTP {status}") - case UnauthorizedError(): - return Sample(seconds=seconds, status=401, error="Unauthorized") - case RateLimitedError(): - return Sample(seconds=seconds, status=429, error="Rate limited") - case NetworkError(message=message): - return Sample(seconds=seconds, status=-1, error="timeout" if "timed out" in message.lower() else "network") - case ValidationError(): - return Sample(seconds=seconds, status=-1, error="Response is not valid JSON") - - -def measure( - issue: Callable[[], Result[Document]], - contract: Contract, - *, - clock: Callable[[], float] = time.monotonic, -) -> tuple[Sample, Document | None]: - started: Final = clock() - with without_retries(): - result: Final = issue() - seconds: Final = clock() - started - match result: - case Success(status_code=200, data=document): - try: - count: Final = contract.count(document) - except AssertionError as error: - return Sample(seconds=seconds, status=200, error=str(error)), document - return Sample(seconds=seconds, status=200, rows=count), document - case _: - return _failure(result, seconds), None - - -def query_path(case: UiQueryCase) -> str: - query: Final = urlencode(case.params.model_dump(by_alias=True, exclude_none=True), doseq=True) - return f"{case.path}?{query}" if query else case.path - - -@step("Measure UI query {case.name} at {case.path}") -def issue_query(transport: HttpTransport, case: UiQueryCase, *, timeout: float) -> Result[Document]: - if case.body is not None: - return transport.post( - case.path, headers=transport.master, json=case.body, response_type=Document, timeout=timeout - ) - return transport.get( - query_path(case), headers=transport.master, params=Query(), response_type=Document, timeout=timeout - ) - - -def database_failures(text: str) -> float: - families: Final = tuple(text_string_to_metric_families(text)) - required: Final = frozenset(("litellm_postgres_failed_requests", "litellm_batch_write_to_db_failed_requests")) - missing: Final = required - frozenset(family.name for family in families) - all_samples: Final = tuple(chain.from_iterable(family.samples for family in families)) - has_activity: Final = any( - sample.name == "litellm_postgres_total_requests_total" and math.isfinite(sample.value) and sample.value > 0 - for sample in all_samples - ) - assert not missing or has_activity, f"Missing PostgreSQL failure counters {missing}" - names: Final = frozenset(f"{name}_total" for name in required) - values: Final = tuple(sample.value for sample in all_samples if sample.name in names) - assert all(math.isfinite(value) and value >= 0 for value in values), "Invalid PostgreSQL failure counter" - return sum(values) diff --git a/tests/e2e/perf/setup/ui_query_cases.py b/tests/e2e/perf/setup/ui_query_cases.py deleted file mode 100644 index e9b70aad713..00000000000 --- a/tests/e2e/perf/setup/ui_query_cases.py +++ /dev/null @@ -1,206 +0,0 @@ -from __future__ import annotations - -from dataclasses import replace -from datetime import datetime, timedelta -from itertools import chain -from typing import Final, Literal - -from pydantic import ConfigDict - -from setup.ui_queries import Contract, Document, Query, UiQueryCase, descend - - -class Params(Query): - model_config = ConfigDict(frozen=True, extra="forbid", populate_by_name=True) - - start_date: str | None = None - timezone: str | None = None - offset: int | None = None - model_group: str | None = None - by_model_group: str | None = None - end_date: str | None = None - page: int | None = None - page_size: int | None = None - sort_by: str | None = None - sort_order: str | None = None - group_by_session: Literal["true"] | None = None - session_cursor: str | None = None - api_key: str | None = None - team_id: str | None = None - team_ids: str | None = None - user_id: str | None = None - model: str | None = None - model_id: str | None = None - key_alias: str | None = None - search: str | None = None - status_filter: str | None = None - cache_hit_filter: str | None = None - error_code: str | None = None - error_message: str | None = None - span_type: str | None = None - exclude_internal_health_checks: Literal["true"] | None = None - key_aliases: tuple[str, ...] | None = None - models: tuple[str, ...] | None = None - limit: int | None = None - tags: str | None = None - exclude_team_ids: str | None = None - include_current_utc_day: Literal["true"] | None = None - - -def _timestamp(value: datetime) -> str: - return value.strftime("%Y-%m-%d %H:%M:%S") - - -def next_page(case: UiQueryCase, document: Document) -> UiQueryCase: - params: Final = Params.model_validate(case.params.model_dump()) - rows: Final = descend(document.root, case.contract.collection) - assert isinstance(rows, list), "Missing page-one rows" - ids: Final = frozenset(str(row["request_id"]) for row in rows if isinstance(row, dict)) - contract: Final = replace(case.contract, excluded_ids=ids) - cursor: Final = descend(document.root, ("next_session_cursor",)) - assert isinstance(cursor, str) and cursor, "Missing session cursor" - return replace(case, params=params.model_copy(update={"page": 2, "session_cursor": cursor}), contract=contract) - - -_LOGS: Final = Contract(("data",), ("request_id", "startTime", "metadata")) - - -def log_cases(anchor: datetime, api_key: str) -> tuple[UiQueryCase, ...]: - params: Final = Params( - start_date=_timestamp(anchor - timedelta(days=1)), - end_date=_timestamp(anchor), - page=1, - page_size=25, - sort_by="startTime", - sort_order="desc", - group_by_session="true", - ) - base: Final = UiQueryCase("logs_default_24h", "/spend/logs/ui", params, _LOGS) - filtered: Final = params.model_copy(update={"start_date": _timestamp(anchor - timedelta(days=30))}) - filters: Final = ( - ("api_key", api_key), - ("team_id", "perf-team-0"), - ("user_id", "perf-user-0"), - ("model", "perf-chat"), - ("model_id", "perf-model-0"), - ("key_alias", "perf-key-0"), - ("status_filter", "failure"), - ("cache_hit_filter", "hit"), - ("error_code", "429"), - ("error_message", "Synthetic"), - ("span_type", "llm"), - ("exclude_internal_health_checks", "true"), - ) - return ( - base, - replace(base, name="logs_cursor_page_2", params=params.model_copy(update={"page": 2}), cursor="session"), - replace( - base, - name="logs_7d", - params=params.model_copy(update={"start_date": _timestamp(anchor - timedelta(days=7))}), - ), - replace( - base, - name="logs_30d", - params=params.model_copy(update={"start_date": _timestamp(anchor - timedelta(days=30))}), - ), - *( - replace(base, name=f"logs_{name}_30d", params=filtered.model_copy(update={name: value})) - for name, value in filters - ), - *( - replace(base, name=f"logs_sort_{name}", params=params.model_copy(update={"sort_by": name})) - for name in ("spend", "total_tokens", "request_duration_ms", "model", "ttft_ms") - ), - ) - - -def analytics_cases(anchor: datetime) -> tuple[UiQueryCase, ...]: - dates: Final = Params( - start_date=(anchor - timedelta(days=7)).date().isoformat(), end_date=anchor.date().isoformat() - ) - cache: Final = UiQueryCase( - "cache_activity_7d", "/global/activity/cache_hits", dates, Contract(("groups",), ("call_type",)) - ) - return ( - UiQueryCase( - "team_users_all_teams", - "/team/spend/by_user", - dates.model_copy(update={"team_ids": ",".join(f"perf-team-{i}" for i in range(100))}), - Contract(("results",), ("team_id", "user_id", "spend"), ("api_requests",)), - ), - UiQueryCase( - "team_users_filtered", - "/team/spend/by_user", - dates.model_copy(update={"team_ids": "perf-team-0"}), - Contract(("results",), ("team_id", "user_id", "spend"), ("api_requests",)), - ), - cache, - replace( - cache, - name="cache_activity_filtered", - params=dates.model_copy(update={"key_aliases": ("perf-key-0", "perf-key-1"), "models": ("perf-chat",)}), - ), - ) - - -def _entity_activity_cases(anchor: datetime, entity: str) -> tuple[UiQueryCase, ...]: - dates: Final = Params( - start_date=(anchor - timedelta(days=7)).date().isoformat(), - end_date=anchor.date().isoformat(), - timezone="0", - exclude_team_ids="litellm-dashboard" if entity == "team" else None, - ) - routes: Final = ( - ("aggregated", dates, Contract(("results",), ("date", "breakdown"))), - ( - "aggregated/keys", - dates.model_copy(update={"offset": 0, "limit": 50}), - Contract(("api_keys",), ("api_key", "metrics", "metadata")), - ), - ( - "aggregated/search", - dates.model_copy(update={"search": "perf-key"}), - Contract(("api_keys",), ("api_key", "metrics", "metadata")), - ), - ( - "aggregated/model_top_keys", - dates.model_copy(update={"model_group": "perf-chat", "by_model_group": "true"}), - Contract(("api_keys",), ("api_key", "metrics", "metadata")), - ), - ) - return tuple( - UiQueryCase( - f"usage_{entity}_{route.replace('/', '_')}", - f"/{entity}/daily/activity/{route}", - params, - contract, - ) - for route, params, contract in routes - ) - - -def activity_cases(anchor: datetime) -> tuple[UiQueryCase, ...]: - groups: Final = tuple(_entity_activity_cases(anchor, entity) for entity in ("user", "team", "tag", "customer")) - return ( - *chain.from_iterable(groups), - UiQueryCase( - "cache_leakage_keys", - "/user/daily/activity/aggregated/cache_leakage_keys", - Params( - start_date=(anchor - timedelta(days=30)).date().isoformat(), - end_date=anchor.date().isoformat(), - timezone="0", - include_current_utc_day="true", - ), - Contract(("api_keys",), ("api_key", "metrics", "metadata")), - ), - ) - - -def cases(anchor: datetime, api_key: str) -> tuple[UiQueryCase, ...]: - return ( - *log_cases(anchor, api_key), - *analytics_cases(anchor), - *activity_cases(anchor), - ) diff --git a/tests/e2e/perf/setup/ui_query_fixture.py b/tests/e2e/perf/setup/ui_query_fixture.py index e3f5236ba47..1cbb5c41ea3 100644 --- a/tests/e2e/perf/setup/ui_query_fixture.py +++ b/tests/e2e/perf/setup/ui_query_fixture.py @@ -3,7 +3,7 @@ from __future__ import annotations import hashlib import os import time -from datetime import datetime +from datetime import datetime, timedelta from typing import Final import psycopg @@ -20,6 +20,45 @@ class UiFixture(BaseModel): api_key: str +SEEDED_TEAMS: Final = 100 +SEEDED_USERS: Final = 2000 +SEEDED_KEYS: Final = 5000 +SEEDED_MODELS: Final = 200 +DAILY_TABLES: Final = ("LiteLLM_DailyUserSpend", "LiteLLM_DailyTeamSpend") +_SEED_SQL: Final = ( + """ + INSERT INTO "LiteLLM_TeamTable" (team_id, team_alias, members_with_roles, metadata, model_spend, model_max_budget) + SELECT 'perf-team-' || n, 'perf-team-' || n, '[]', '{}', '{}', '{}' FROM generate_series(0, %(teams)s - 1) AS n + ON CONFLICT (team_id) DO NOTHING + """, + """ + INSERT INTO "LiteLLM_UserTable" (user_id, user_alias, user_email, user_role, team_id, metadata, model_spend, model_max_budget) + SELECT 'perf-user-' || n, 'perf-user-' || n, 'perf-user-' || n || '@example.com', 'internal_user', + 'perf-team-' || (n %% %(teams)s), '{}', '{}', '{}' + FROM generate_series(0, %(users)s - 1) AS n + ON CONFLICT (user_id) DO NOTHING + """, + """ + INSERT INTO "LiteLLM_VerificationToken" (token, key_name, key_alias, team_id, user_id, aliases, config, permissions, + metadata, model_spend, model_max_budget, created_at, updated_at) + SELECT md5('perf-key:' || n) || md5('perf-key-tail:' || n), 'sk-...perf' || n, 'perf-key-' || n, + 'perf-team-' || (n %% %(teams)s), 'perf-user-' || (n %% %(users)s), '{}', '{}', '{}', '{}', '{}', '{}', + %(anchor)s - make_interval(days => n %% 365), %(anchor)s + FROM generate_series(0, %(keys)s - 1) AS n + ON CONFLICT (token) DO NOTHING + """, + """ + INSERT INTO "LiteLLM_ProxyModelTable" (model_id, model_name, litellm_params, model_info, created_by, updated_by) + SELECT 'perf-model-' || n, 'perf-chat-' || (n %% 20), + jsonb_build_object('model', 'openai/perf-chat', 'api_key', 'perf-seed', + 'api_base', 'os.environ/E2E_PERF_MOCK_OPENAI_URL'), + jsonb_build_object('id', 'perf-model-' || n, 'db_model', true), 'perf-seed', 'perf-seed' + FROM generate_series(0, %(models)s - 1) AS n + ON CONFLICT (model_id) DO NOTHING + """, +) + + def prepare_fixture() -> UiFixture: validate_target(os.environ["DATABASE_URL"]) _finish_upgrade() @@ -27,6 +66,7 @@ def prepare_fixture() -> UiFixture: _ = connection.execute("SET statement_timeout='30s'") row: Final = connection.execute("SELECT spec->>'anchor', rows FROM litellm_perf_seed.checkpoint").fetchone() assert row is not None and TypeAdapter(int).validate_python(row[1]) > 0, "Missing completed seed rows" + seed_ui_tables(connection, anchor=TypeAdapter(datetime).validate_python(row[0])) fixture: Final = UiFixture( anchor=TypeAdapter(datetime).validate_python(row[0]), api_key=hashlib.md5(b"perf-key:0", usedforsecurity=False).hexdigest() @@ -35,6 +75,36 @@ def prepare_fixture() -> UiFixture: return fixture +def seed_ui_tables(connection: psycopg.Connection, anchor: datetime) -> None: + """Give the Keys, Teams, Users and Models pages realistic row counts on the disposable restore. + + Names follow the SpendLogs seed (perf-key-N aliases, perf-team-N, perf-user-N) so the Logs filters and + Team Usage line up with real spend. The Daily* rollups come from the snapshot itself and are only checked. + """ + window: Final = { + "since": (anchor - timedelta(days=7)).strftime("%Y-%m-%d"), + "until": anchor.strftime("%Y-%m-%d"), + } + for table in DAILY_TABLES: + row = connection.execute(f'SELECT count(*) FROM "{table}" WHERE date BETWEEN %(since)s AND %(until)s', window) + assert _count(row.fetchone()) > 0, f"{table} has no rows in the week before the anchor" + params: Final = { + "teams": SEEDED_TEAMS, + "users": SEEDED_USERS, + "keys": SEEDED_KEYS, + "models": SEEDED_MODELS, + "anchor": anchor, + } + for statement in _SEED_SQL: + _ = connection.execute(statement, params) + connection.commit() + + +def _count(row: tuple[object, ...] | None) -> int: + assert row is not None + return TypeAdapter(int).validate_python(row[0]) + + def _finish_upgrade() -> None: if os.environ.get("E2E_DATABASE_UPGRADE") != "1": return diff --git a/tests/e2e/perf/setup/ui_query_isolation.py b/tests/e2e/perf/setup/ui_query_isolation.py index d7f6d725edc..d892c268fa2 100644 --- a/tests/e2e/perf/setup/ui_query_isolation.py +++ b/tests/e2e/perf/setup/ui_query_isolation.py @@ -62,7 +62,9 @@ def wait_idle(connection: psycopg.Connection[dict[str, object]], evidence: Path, for query in queries: if cancel and query.leader_pid is None: _ = connection.execute( - "SELECT pg_cancel_backend(%s)" if query.state == "active" else "SELECT pg_terminate_backend(%s)", + "SELECT pg_cancel_backend(%s)" + if query.state == "active" + else "SELECT pg_terminate_backend(%s)", (query.pid,), ) else: diff --git a/tests/e2e/perf/test_ui_journeys.py b/tests/e2e/perf/test_ui_journeys.py new file mode 100644 index 00000000000..888a44d2e27 --- /dev/null +++ b/tests/e2e/perf/test_ui_journeys.py @@ -0,0 +1,108 @@ +from __future__ import annotations + +import json +import os +from datetime import UTC, datetime +from pathlib import Path +from typing import Final + +import pytest +from e2e_config import MASTER_KEY, UI_BASE_URL, UI_PASSWORD, UI_USERNAME +from e2e_http import send_text_external, without_retries +from setup.ui_journeys import ( + assert_within_ceiling, + Dashboard, + Journey, + Sample, + dashboard_page, + database_failures, + error_code, + journeys, + measure, + summarize, + ui_junit_properties, +) +from setup.ui_query_fixture import UiFixture, prepare_fixture +from setup.ui_query_isolation import isolated_database, settled_counter +from transport import HttpTransport + +pytestmark = [ + pytest.mark.perf, + pytest.mark.quiet_stack, + pytest.mark.skipif(os.environ.get("E2E_UI_PERF_DISPOSABLE") != "1", reason="Dedicated disposable database only"), +] + + +def pytest_generate_tests(metafunc: pytest.Metafunc) -> None: + if "journey_name" in metafunc.fixturenames: + metafunc.parametrize("journey_name", [journey.name for journey in journeys("")]) + + +@pytest.fixture(scope="module") +def ui_fixture() -> UiFixture: + return prepare_fixture() + + +def _metrics(transport: HttpTransport) -> float: + urls: Final = tuple(url.strip() for url in os.environ["E2E_DATABASE_METRICS_URLS"].split(",") if url.strip()) + try: + assert urls, "Provide a metrics URL for every proxy worker/pod" + readings: Final = tuple(send_text_external("GET", url, headers=transport.master, timeout=10) for url in urls) + assert all(reading.status_code == 200 for reading in readings), "Cannot read database failure metrics" + return sum(database_failures(reading.body) for reading in readings) + except Exception as error: + pytest.exit(f"Database metrics unavailable ({type(error).__name__}). Stopping UI measurements", returncode=2) + + +def _samples(dashboard: Dashboard, journey: Journey, timeout: float, repeats: int) -> tuple[Sample, ...]: + with dashboard_page(dashboard, timeout) as page: + return tuple(measure(page, journey, timeout) for _ in range(repeats + 1)) + + +def test_ui_journey_latency(journey_name: str, ui_fixture: UiFixture, request: pytest.FixtureRequest) -> None: + journey: Final = next(journey for journey in journeys(ui_fixture.api_key) if journey.name == journey_name) + transport: Final = HttpTransport(base_url=UI_BASE_URL, master_key=MASTER_KEY) + timeout: Final = float(os.environ.get("E2E_UI_JOURNEY_TIMEOUT_SECONDS", "120")) + repeats: Final = int(os.environ.get("E2E_UI_JOURNEY_REPEATS", "9")) + assert 0 < timeout <= 600 and 1 <= repeats <= 1000, "Invalid UI performance sample settings" + target: Final = Path(os.environ["E2E_ARTIFACT_DIR"]) + assert target.is_dir(), "E2E_ARTIFACT_DIR must exist" + evidence: Final = target / f"ui-isolation-{journey.name}.jsonl" + dashboard: Final = Dashboard( + base_url=UI_BASE_URL, username=UI_USERNAME, password=UI_PASSWORD, anchor=ui_fixture.anchor + ) + with without_retries(): + with isolated_database(journey.name): + before: Final = settled_counter(lambda: _metrics(transport), evidence) + samples: Final = _samples(dashboard, journey, timeout, repeats) + with evidence.open("a", encoding="utf-8") as sample_output: + _ = sample_output.write( + json.dumps({"samples": tuple(sample.model_dump() for sample in samples)}) + "\n" + ) + after_requests: Final = _metrics(transport) + item: Final = request.node # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # pytest: untyped + assert isinstance(item, pytest.Item) + item.user_properties.extend(ui_junit_properties(samples)) + after: Final = settled_counter(lambda: _metrics(transport), evidence) + item.user_properties.extend( + ( + ("perf_ui_database_failures", after_requests - before), + ("perf_ui_cleanup_database_failures", after - after_requests), + ) + ) + artifact: Final = { + "name": journey.name, + "surface": journey.surface, + "measured_at": datetime.now(UTC).isoformat(), + "samples": [sample.model_dump() for sample in samples], + "error_codes": [error_code(sample) for sample in samples], + "timings": summarize(samples).model_dump(), + "timeouts": sum(sample.error == "timeout" for sample in samples), + "failures": sum(bool(sample.error) for sample in samples), + "database_failures_before": before, + "database_failures_after_requests": after_requests, + "database_failures_after_cleanup": after, + } + output: Final = target / f"ui-journey-{journey.name}.json" + _ = output.write_text(json.dumps(artifact, indent=2), encoding="utf-8") + assert_within_ceiling(journey.name, samples) diff --git a/tests/e2e/perf/test_ui_queries.py b/tests/e2e/perf/test_ui_queries.py deleted file mode 100644 index 75ac4d3055f..00000000000 --- a/tests/e2e/perf/test_ui_queries.py +++ /dev/null @@ -1,159 +0,0 @@ -from __future__ import annotations - -import json -import os -from dataclasses import dataclass, replace -from datetime import UTC, datetime -from pathlib import Path -from typing import Final - -import pytest -from e2e_config import CONTROL_PLANE_BASE_URL, MASTER_KEY -from e2e_http import send_text_external, without_retries -from setup.ui_queries import ( - Sample, - UiQueryCase, - database_failures, - error_code, - issue_query, - measure, - summarize, - ui_junit_properties, -) -from setup.ui_query_cases import Params, cases, next_page -from setup.ui_query_fixture import UiFixture, prepare_fixture -from setup.ui_query_isolation import isolated_database, settled_counter -from transport import HttpTransport - -pytestmark = [ - pytest.mark.perf, - pytest.mark.quiet_stack, - pytest.mark.skipif(os.environ.get("E2E_UI_PERF_DISPOSABLE") != "1", reason="Dedicated disposable database only"), -] - - -def pytest_generate_tests(metafunc: pytest.Metafunc) -> None: - if "ui_case_name" not in metafunc.fixturenames: - return - if os.environ.get("E2E_UI_PERF_DISPOSABLE") == "1" and os.environ.get("E2E_RESET_SPEND_LOGS") == "1": - raise pytest.UsageError("UI performance tests cannot run with E2E_RESET_SPEND_LOGS=1") - names: Final = tuple(case.name for case in cases(datetime(2000, 1, 1, tzinfo=UTC), "collection-only")) - metafunc.parametrize("ui_case_name", names, ids=names) - - -@pytest.fixture(scope="module") -def ui_fixture() -> UiFixture: - return prepare_fixture() - - -def _metrics(transport: HttpTransport) -> float: - urls: Final = tuple(url.strip() for url in os.environ["E2E_DATABASE_METRICS_URLS"].split(",") if url.strip()) - try: - assert urls, "Provide a metrics URL for every proxy worker/pod" - readings: Final = tuple(send_text_external("GET", url, headers=transport.master, timeout=10) for url in urls) - assert all(reading.status_code == 200 for reading in readings), "Cannot read database failure metrics" - return sum(database_failures(reading.body) for reading in readings) - except Exception as error: - pytest.exit(f"Database metrics unavailable ({type(error).__name__}). Stopping UI measurements", returncode=2) - - -@dataclass(frozen=True, slots=True) -class Prepared: - case: UiQueryCase - prerequisite: Sample | None = None - - -def _page_two(transport: HttpTransport, case: UiQueryCase, timeout: float) -> Prepared: - if case.cursor == "none": - return Prepared(case) - params: Final = Params.model_validate(case.params.model_dump()).model_copy(update={"page": 1}) - first: Final = replace(case, params=params) - sample, document = measure(lambda: issue_query(transport, first, timeout=timeout), case.contract) - if sample.error or document is None: - return Prepared(case, sample) - try: - return Prepared(next_page(case, document), sample) - except AssertionError as error: - return Prepared(case, sample.model_copy(update={"error": str(error)})) - - -def _samples(transport: HttpTransport, case: UiQueryCase, timeout: float, repeats: int) -> tuple[Sample, ...]: - measured, _ = measure(lambda: issue_query(transport, case, timeout=timeout), case.contract) - if measured.error: - return (measured,) - repeated: Final = tuple( - measure(lambda: issue_query(transport, case, timeout=timeout), case.contract)[0] for _ in range(repeats) - ) - return (measured, *repeated) - - -def test_ui_query_performance(ui_case_name: str, ui_fixture: UiFixture, request: pytest.FixtureRequest) -> None: - case: Final = next(case for case in cases(ui_fixture.anchor, ui_fixture.api_key) if case.name == ui_case_name) - transport: Final = HttpTransport(base_url=CONTROL_PLANE_BASE_URL, master_key=MASTER_KEY) - timeout: Final = float(os.environ.get("E2E_UI_QUERY_TIMEOUT_SECONDS", "120")) - repeats: Final = int(os.environ.get("E2E_UI_QUERY_REPEATS", "9")) - assert 0 < timeout <= 600 and 1 <= repeats <= 1000, "Invalid UI performance sample settings" - target: Final = Path(os.environ["E2E_ARTIFACT_DIR"]) - assert target.is_dir(), "E2E_ARTIFACT_DIR must exist" - evidence: Final = target / f"ui-isolation-{case.name}.jsonl" - with without_retries(): - with isolated_database(case.name): - before: Final = settled_counter(lambda: _metrics(transport), evidence) - paginated: Final = _page_two(transport, case, timeout) - samples: Final = ( - (paginated.prerequisite,) - if paginated.prerequisite is not None and paginated.prerequisite.error - else _samples(transport, paginated.case, timeout, repeats) - ) - with evidence.open("a", encoding="utf-8") as sample_output: - _ = sample_output.write( - json.dumps( - { - "samples": tuple(sample.model_dump() for sample in samples), - "prerequisite": paginated.prerequisite.model_dump() - if paginated.prerequisite is not None - else None, - } - ) - + "\n" - ) - after_requests: Final = _metrics(transport) - item: Final = request.node # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # pytest: untyped - assert isinstance(item, pytest.Item) - item.user_properties.extend(ui_junit_properties(samples, paginated.prerequisite)) - after: Final = settled_counter(lambda: _metrics(transport), evidence) - item.user_properties.extend( - ( - ("perf_ui_database_failures", after_requests - before), - ("perf_ui_cleanup_database_failures", after - after_requests), - ) - ) - timings: Final = summarize(samples) - artifact: Final = { - "name": case.name, - "surface": case.surface, - "path": case.path, - "method": "POST" if case.body is not None else "GET", - "params": paginated.case.params.model_dump(mode="json", by_alias=True, exclude_none=True), - "prerequisite": paginated.prerequisite.model_dump() if paginated.prerequisite is not None else None, - "prerequisite_failed": paginated.prerequisite is not None and bool(paginated.prerequisite.error), - "body": case.body.model_dump(mode="json", exclude_none=True) if case.body is not None else None, - "samples": [sample.model_dump() for sample in samples], - "error_codes": [error_code(sample) for sample in samples], - "timings": timings.model_dump(), - "timeouts": sum(sample.error == "timeout" for sample in samples), - "failures": sum(bool(sample.error) for sample in samples), - "status_counts": { - str(status): sum(sample.status == status for sample in samples) - for status in frozenset(sample.status for sample in samples) - }, - "validation_failures": sum( - bool(sample.error) and (sample.status == 200 or sample.error == "Response is not valid JSON") - for sample in samples - ), - "database_failures_before": before, - "database_failures_after_requests": after_requests, - "database_failures_after_cleanup": after, - } - output: Final = target / f"ui-query-{case.name}.json" - _ = output.write_text(json.dumps(artifact, indent=2), encoding="utf-8") diff --git a/tests/e2e_harness/perf/setup/test_ui_journeys.py b/tests/e2e_harness/perf/setup/test_ui_journeys.py new file mode 100644 index 00000000000..7d5b0d531ec --- /dev/null +++ b/tests/e2e_harness/perf/setup/test_ui_journeys.py @@ -0,0 +1,144 @@ +from __future__ import annotations + +from typing import Final + +import pytest +from setup.ui_journeys import ( + UI_JOURNEY_LIMITS, + Sample, + assert_within_ceiling, + database_failures, + journeys, + summarize, + ui_junit_properties, +) + + +def test_first_journey_is_separate_from_repeated_latency() -> None: + samples: Final = (Sample(seconds=80), *(Sample(seconds=i) for i in range(1, 101))) + result: Final = summarize(samples) + assert result.first_seconds == 80 + assert result.repeated_samples == 100 + assert result.repeated_p50_seconds == 50.5 + assert result.repeated_p95_seconds == 95 + assert result.repeated_p99_seconds == 99 + assert result.repeated_max_seconds == 100 + + +def test_failed_repeats_do_not_look_like_fast_successes() -> None: + result: Final = summarize((Sample(seconds=10), Sample(seconds=0.01, error="AssertionError: no rows"))) + assert result.repeated_samples == 0 + assert result.repeated_p50_seconds is None + assert result.repeated_p95_seconds is None + assert result.repeated_p99_seconds is None + assert result.repeated_max_seconds is None + + +def test_small_sample_does_not_report_tail_percentiles() -> None: + result: Final = summarize(tuple(Sample(seconds=1) for _ in range(6))) + assert result.repeated_samples == 5 + assert result.repeated_p50_seconds == 1 + assert result.repeated_p95_seconds is None + assert result.repeated_p99_seconds is None + + +def test_failed_first_journey_emits_error_code_without_latency() -> None: + sample: Final = Sample(seconds=120, error="timeout") + assert summarize((sample,)).first_seconds is None + assert dict(ui_junit_properties((sample,))) == { + "perf_ui_journeys": 1, + "perf_ui_failures": 1, + "perf_ui_samples": 0, + "perf_ui_repeated_samples": 0, + "perf_ui_error_timeout": 1, + } + + +def test_junit_counts_failures_but_only_measures_successful_latency() -> None: + samples: Final = ( + Sample(seconds=2, rows=25), + Sample(seconds=3, rows=25), + Sample(seconds=120, error="timeout"), + Sample(seconds=1, error="AssertionError: Logs table rendered empty"), + Sample(seconds=4, rows=25), + ) + assert dict(ui_junit_properties(samples)) == { + "perf_ui_journeys": 5, + "perf_ui_failures": 2, + "perf_ui_samples": 3, + "perf_ui_repeated_samples": 2, + "perf_ui_error_timeout": 1, + "perf_ui_error_ui": 1, + "perf_ui_seconds_p50": 3, + "perf_ui_seconds_p90": 4, + "perf_ui_seconds_max": 4, + "perf_ui_first_seconds": 2, + "perf_ui_repeated_seconds_p50": 3.5, + "perf_ui_repeated_seconds_max": 4, + } + + +def test_all_journey_percentiles_include_first_and_use_nearest_rank_p90() -> None: + samples: Final = tuple(Sample(seconds=seconds) for seconds in (10, 1, 9, 2, 8, 3, 7, 4, 6, 5)) + metrics: Final = dict(ui_junit_properties(samples)) + assert {key: value for key, value in metrics.items() if key.startswith("perf_ui_seconds_")} == { + "perf_ui_seconds_p50": 5.5, + "perf_ui_seconds_p90": 9, + "perf_ui_seconds_max": 10, + } + assert metrics["perf_ui_samples"] == 10 + + +def test_journeys_have_unique_names_and_cover_every_dashboard_surface() -> None: + names: Final = tuple(journey.name for journey in journeys("key")) + assert len(names) == len(frozenset(names)) + assert frozenset(journey.surface for journey in journeys("key")) == {"logs", "usage", "caching", "models", "keys"} + + +def _metrics(postgres: str, batch: str) -> str: + return ( + "# TYPE litellm_postgres_failed_requests_total counter\n" + f'litellm_postgres_failed_requests_total{{call_type="query_raw"}} {postgres}\n' + "# TYPE litellm_batch_write_to_db_failed_requests_total counter\n" + f'litellm_batch_write_to_db_failed_requests_total{{call_type="batch"}} {batch}\n' + ) + + +def test_postgres_failures_include_background_batch_writes() -> None: + assert database_failures(_metrics("0", "0")) == 0 + assert database_failures(_metrics("2", "3")) == 5 + + +def test_lazy_failure_counters_require_evidence_of_database_activity() -> None: + assert ( + database_failures( + "# TYPE litellm_postgres_total_requests_total counter\nlitellm_postgres_total_requests_total 3\n" + ) + == 0 + ) + + +@pytest.mark.parametrize( + "text", ("", "# TYPE unrelated counter\nunrelated_total 0\n", _metrics("NaN", "0"), _metrics("-1", "0")) +) +def test_missing_or_invalid_metrics_cannot_pass_as_zero(text: str) -> None: + with pytest.raises(AssertionError): + _ = database_failures(text) + + +def test_every_journey_has_a_latency_ceiling() -> None: + assert frozenset(journey.name for journey in journeys("key")) == frozenset(UI_JOURNEY_LIMITS) + + +def test_failed_sample_fails_the_journey_before_latency_is_judged() -> None: + with pytest.raises(AssertionError, match="1 of 2 logs_open_24h samples failed: timeout"): + assert_within_ceiling("logs_open_24h", (Sample(seconds=0.2, rows=25), Sample(seconds=120, error="timeout"))) + + +def test_slow_journey_fails_its_ceiling() -> None: + fast: Final = tuple(Sample(seconds=1, rows=25) for _ in range(9)) + assert_within_ceiling("logs_open_24h", fast) + with pytest.raises(AssertionError, match=r"logs_open_24h p90 25\.0s over 20s"): + assert_within_ceiling("logs_open_24h", (*fast[:8], Sample(seconds=25, rows=25), Sample(seconds=25, rows=25))) + with pytest.raises(AssertionError, match=r"logs_open_24h p50 11\.0s over 10s"): + assert_within_ceiling("logs_open_24h", tuple(Sample(seconds=11, rows=25) for _ in range(3))) diff --git a/tests/e2e_harness/perf/setup/test_ui_queries.py b/tests/e2e_harness/perf/setup/test_ui_queries.py deleted file mode 100644 index f24d03c6bba..00000000000 --- a/tests/e2e_harness/perf/setup/test_ui_queries.py +++ /dev/null @@ -1,257 +0,0 @@ -from __future__ import annotations - -from datetime import UTC, datetime -from typing import Final -from urllib.parse import parse_qs, urlsplit - -import pytest -from e2e_http import NetworkError, Result, Success, UnknownApiError -from setup.ui_queries import ( - Contract, - Document, - Sample, - database_failures, - measure, - query_path, - summarize, - ui_junit_properties, -) -from setup.ui_query_cases import Params, cases, next_page - - -def test_query_string_preserves_ui_arrays_and_boolean_strings() -> None: - requests: Final = cases(datetime(2026, 1, 1, tzinfo=UTC), "key") - cache: Final = next(case for case in requests if case.name == "cache_activity_filtered") - assert parse_qs(urlsplit(query_path(cache)).query) == { - "start_date": ["2025-12-25"], - "end_date": ["2026-01-01"], - "key_aliases": ["perf-key-0", "perf-key-1"], - "models": ["perf-chat"], - } - health: Final = next(case for case in requests if case.name == "logs_exclude_internal_health_checks_30d") - query: Final = parse_qs(urlsplit(query_path(health)).query) - assert query["exclude_internal_health_checks"] == ["true"] - assert query["group_by_session"] == ["true"] - - -@pytest.mark.parametrize("entity", ("user", "team", "tag", "customer")) -@pytest.mark.parametrize( - ("route", "extra"), - ( - ("aggregated", {}), - ("aggregated/keys", {"offset": 0, "limit": 50}), - ("aggregated/search", {"search": "perf-key"}), - ("aggregated/model_top_keys", {"model_group": "perf-chat", "by_model_group": "true"}), - ), -) -def test_usage_requests_list_keys_without_a_single_key_filter( - entity: str, route: str, extra: dict[str, str | int] -) -> None: - case: Final = next( - case - for case in cases(datetime(2026, 1, 1, tzinfo=UTC), "key") - if case.path == f"/{entity}/daily/activity/{route}" - ) - assert case.params.model_dump(exclude_none=True) == { - "start_date": "2025-12-25", - "end_date": "2026-01-01", - "timezone": "0", - **({"exclude_team_ids": "litellm-dashboard"} if entity == "team" else {}), - **extra, - } - - -@pytest.mark.parametrize(("name", "value"), (("api_key", "key"), ("key_alias", "perf-key-0"))) -def test_key_filters_request_30_days_of_logs(name: str, value: str) -> None: - case: Final = next( - case for case in cases(datetime(2026, 1, 1, tzinfo=UTC), "key") if case.name == f"logs_{name}_30d" - ) - assert parse_qs(urlsplit(query_path(case)).query) == { - "start_date": ["2025-12-02 00:00:00"], - "end_date": ["2026-01-01 00:00:00"], - "page_size": ["25"], - "page": ["1"], - "sort_by": ["startTime"], - "sort_order": ["desc"], - "group_by_session": ["true"], - name: [value], - } - - -@pytest.mark.parametrize("invalid", ("not-a-number", True, float("nan"), float("inf"))) -def test_positive_row_does_not_hide_a_malformed_later_row(invalid: str | bool | float) -> None: - with pytest.raises(AssertionError, match=r"Expected (numeric|finite) spend"): - Contract(positive=("spend",)).count(Document([{"spend": 1}, {"spend": invalid}])) - - -def test_session_cursor_is_reused_and_repeated_rows_rejected() -> None: - case: Final = next(case for case in cases(datetime(2026, 1, 1, tzinfo=UTC), "key") if case.cursor == "session") - first: Final = Document.model_validate_json('{"data":[{"request_id":"one"}],"next_session_cursor":"opaque"}') - second: Final = next_page(case, first) - assert Params.model_validate(second.params).session_cursor == "opaque" - assert Params.model_validate(second.params).page == 2 - with pytest.raises(AssertionError, match="Repeated page-one row"): - _ = second.contract.count(first) - - -def test_first_query_is_separate_from_repeated_latency() -> None: - samples: Final = (Sample(seconds=80, status=200), *(Sample(seconds=i, status=200) for i in range(1, 101))) - result: Final = summarize(samples) - assert result.first_seconds == 80 - assert result.repeated_samples == 100 - assert result.repeated_p50_seconds == 50.5 - assert result.repeated_p95_seconds == 95 - assert result.repeated_p99_seconds == 99 - assert result.repeated_max_seconds == 100 - - -def test_failed_repeats_do_not_look_like_fast_successes() -> None: - result: Final = summarize((Sample(seconds=10, status=200), Sample(seconds=0.01, status=503, error="HTTP 503"))) - assert result.repeated_samples == 0 - assert result.repeated_p50_seconds is None - assert result.repeated_p95_seconds is None - assert result.repeated_p99_seconds is None - assert result.repeated_max_seconds is None - - -def test_small_sample_does_not_report_tail_percentiles() -> None: - result: Final = summarize(tuple(Sample(seconds=1, status=200) for _ in range(6))) - assert result.repeated_samples == 5 - assert result.repeated_p50_seconds == 1 - assert result.repeated_p95_seconds is None - assert result.repeated_p99_seconds is None - - -def test_failed_first_request_emits_error_code_without_latency() -> None: - sample: Final = Sample(seconds=30, status=500, error="HTTP 500") - assert summarize((sample,)).first_seconds is None - assert dict(ui_junit_properties((sample,))) == { - "perf_ui_requests": 1, - "perf_ui_failures": 1, - "perf_ui_samples": 0, - "perf_ui_repeated_samples": 0, - "perf_ui_error_http_500": 1, - } - - -def test_junit_counts_failures_but_only_measures_successful_latency() -> None: - samples: Final = ( - Sample(seconds=2, status=200), - Sample(seconds=3, status=200), - Sample(seconds=30, status=-1, error="timeout"), - Sample(seconds=30, status=503, error="HTTP 503"), - Sample(seconds=1, status=200, error="Empty collection"), - ) - assert dict(ui_junit_properties(samples, Sample(seconds=4, status=200))) == { - "perf_ui_requests": 6, - "perf_ui_failures": 3, - "perf_ui_samples": 2, - "perf_ui_seconds_p50": 2.5, - "perf_ui_seconds_p90": 3, - "perf_ui_seconds_max": 3, - "perf_ui_repeated_samples": 1, - "perf_ui_error_timeout": 1, - "perf_ui_error_http_503": 1, - "perf_ui_error_contract": 1, - "perf_ui_first_seconds": 2, - "perf_ui_repeated_seconds_p50": 3, - "perf_ui_repeated_seconds_max": 3, - } - - -def test_failed_cursor_prerequisite_is_counted_once() -> None: - sample: Final = Sample(seconds=30, status=503, error="HTTP 503") - assert dict(ui_junit_properties((sample,), sample)) == dict(ui_junit_properties((sample,))) - - -def test_all_request_percentiles_include_first_and_use_nearest_rank_p90() -> None: - samples: Final = tuple(Sample(seconds=seconds, status=200) for seconds in (10, 1, 9, 2, 8, 3, 7, 4, 6, 5)) - metrics: Final = dict(ui_junit_properties(samples)) - assert {key: value for key, value in metrics.items() if key.startswith("perf_ui_seconds_")} == { - "perf_ui_seconds_p50": 5.5, - "perf_ui_seconds_p90": 9, - "perf_ui_seconds_max": 10, - } - assert metrics["perf_ui_samples"] == 10 - - -@pytest.mark.parametrize("body", ("{}", '{"data": []}', '{"data": [{"spend": 1}]}', '{"data": [null]}')) -def test_empty_or_malformed_rows_fail_even_with_http_200(body: str) -> None: - ticks: Final = iter((10.0, 12.5)) - result, _ = measure( - lambda: Success(status_code=200, data=Document.model_validate_json(body)), - Contract(("data",), ("request_id",)), - clock=lambda: next(ticks), - ) - assert result.status == 200 - assert result.seconds == 2.5 - assert result.rows == 0 - assert result.error - - -def test_valid_response_records_elapsed_seconds_and_rows() -> None: - ticks: Final = iter((5.0, 5.125)) - result, document = measure( - lambda: Success(status_code=200, data=Document.model_validate_json('{"data":[{"request_id":"one"}]}')), - Contract(("data",), ("request_id",)), - clock=lambda: next(ticks), - ) - assert result == Sample(seconds=0.125, status=200, rows=1) - assert document is not None - - -def test_zero_spend_rows_are_allowed_alongside_nonzero_fixture_data() -> None: - contract: Final = Contract(positive=("spend",)) - assert contract.count(Document.model_validate_json('[{"spend":0},{"spend":2}]')) == 2 - with pytest.raises(AssertionError, match="positive spend"): - _ = contract.count(Document.model_validate_json('[{"spend":0}]')) - - -@pytest.mark.parametrize( - "failure,error,status", - ( - (UnknownApiError(status_code=500, body="sensitive error detail"), "HTTP 500", 500), - (NetworkError(message="Read timed out"), "timeout", -1), - (NetworkError(message="connection refused"), "network", -1), - ), -) -def test_failed_requests_are_recorded_without_response_bodies( - failure: Result[Document], - error: str, - status: int, -) -> None: - ticks: Final = iter((1.0, 2.0)) - result, document = measure(lambda: failure, Contract(), clock=lambda: next(ticks)) - assert result == Sample(seconds=1, status=status, error=error) - assert document is None - - -def _metrics(postgres: str, batch: str) -> str: - return ( - "# TYPE litellm_postgres_failed_requests_total counter\n" - f'litellm_postgres_failed_requests_total{{call_type="query_raw"}} {postgres}\n' - "# TYPE litellm_batch_write_to_db_failed_requests_total counter\n" - f'litellm_batch_write_to_db_failed_requests_total{{call_type="batch"}} {batch}\n' - ) - - -def test_postgres_failures_include_background_batch_writes() -> None: - assert database_failures(_metrics("0", "0")) == 0 - assert database_failures(_metrics("2", "3")) == 5 - - -def test_lazy_failure_counters_require_evidence_of_database_activity() -> None: - assert ( - database_failures( - "# TYPE litellm_postgres_total_requests_total counter\nlitellm_postgres_total_requests_total 3\n" - ) - == 0 - ) - - -@pytest.mark.parametrize( - "text", ("", "# TYPE unrelated counter\nunrelated_total 0\n", _metrics("NaN", "0"), _metrics("-1", "0")) -) -def test_missing_or_invalid_metrics_cannot_pass_as_zero(text: str) -> None: - with pytest.raises(AssertionError): - _ = database_failures(text) diff --git a/tests/e2e_harness/perf/test_upgrade.py b/tests/e2e_harness/perf/test_upgrade.py index 0904b7a7b5a..89e3a620f48 100644 --- a/tests/e2e_harness/perf/test_upgrade.py +++ b/tests/e2e_harness/perf/test_upgrade.py @@ -308,9 +308,7 @@ def test_invalid_200_is_not_an_allowed_non_200(failed: WireRecord) -> None: def test_upgrade_verdict_uses_regular_limits_without_cpu_or_datastore_samples( failures: int, overhead: float, ttft: float, growth: float, expected: tuple[str, ...] ) -> None: - nonstream: Final = tuple( - record(i).model_copy(update={"total_s": overhead, "mock_delay_s": 0}) for i in range(100) - ) + nonstream: Final = tuple(record(i).model_copy(update={"total_s": overhead, "mock_delay_s": 0}) for i in range(100)) stream: Final = record(100).model_copy(update={"stream": True, "first_token_s": ttft}) observed: Final = Observation( records=(*nonstream, stream, *(failed_record(i) for i in range(101, 101 + failures))), From 8d2bf426f59ffbda9aaa49608706aee4acb93395 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 20 Oct 2026 09:26:38 -0700 Subject: [PATCH 26/33] fix(cost): price base_model deployments past rule-only cost map entries (#47497) * fix(cost): price base_model deployments past rule-only cost map entries and warn when base_model has no rates * fix(router): warn on unpriced base_model from add_deployment too and honour model_info prices * refactor(cost): drop the unpriced base_model startup warning, keep the one-line ladder fix --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/cost_calculator.py | 2 +- tests/unit/test_cost_calculator.py | 32 ++++++++++++++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 20914df8fe2..e8ca60e00af 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1642,7 +1642,7 @@ def completion_cost( *(((model, request_pricing_provider),) if model is not None else ()), ) potential_model_names: Final = ( - raw_candidates if explicit_pricing else _priced_candidates_before_unpriced_rules(raw_candidates) + raw_candidates if custom_pricing is True else _priced_candidates_before_unpriced_rules(raw_candidates) ) for idx, (model, pricing_provider) in enumerate(potential_model_names): diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index fee4db586cd..2854f9010b2 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -391,6 +391,38 @@ def test_completion_cost_bills_undated_azure_data_zone_base_model_at_its_dated_r ) +def test_completion_cost_prices_rule_only_base_model_from_the_response_model( + _local_model_cost_map: None, + _price_free_gpt_family_generalization: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """LIT-9812: a deployment whose base_model is an undated name only a capability rule matches must not + stop the ladder at $0 when the response names a priced model""" + monkeypatch.setattr( + litellm, + "model_cost", + { + "azure/gpt-9.9-2026-01-01": { + "litellm_provider": "azure", + "mode": "chat", + "input_cost_per_token": 0.0000111, + "output_cost_per_token": 0.0000222, + } + }, + ) + usage: Final = Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26) + response: Final = ModelResponse(model="gpt-9.9-2026-01-01", choices=[], usage=usage) + + cost: Final = completion_cost( + completion_response=response, + model="azure/my-eu-deployment", + custom_llm_provider="azure", + base_model="azure/eu/gpt-9.9", + ) + + assert cost == pytest.approx(usage.prompt_tokens * 0.0000111 + usage.completion_tokens * 0.0000222) + + def test_completion_cost_orders_candidates_without_extra_huggingface_config_requests( _local_model_cost_map: None, ) -> None: From 9bbf70edb6372dd1f7436562f4bb1dc248af793d Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 20 Oct 2026 09:41:22 -0700 Subject: [PATCH 27/33] fix(completion): send prompt_cache_retention to OpenAI instead of silently dropping it (#47480) * fix(completion): send prompt_cache_retention to OpenAI instead of silently dropping it litellm.completion(), Router.completion() and the proxy's /v1/chat/completions accepted prompt_cache_retention and listed it as a supported OpenAI param, but completion() never handed it to get_optional_params, and the leftover-kwargs paths skip known OpenAI params, so it never reached the request body. Make it a named argument of completion() and acompletion() like prompt_cache_key and pass it on. A unit guard sends every OpenAI body param through completion() and fails when a param added to OPENAI_CHAT_COMPLETION_PARAMS has no case proving it reaches the wire, the same hole store and prompt_cache_key fell into in #33184. Live e2e tests capture the body OpenAI receives for the proxy, the SDK and the Router on both openai/gpt-5.5 and openai/responses/gpt-5.5. * chore(completion): type the param-forwarding unit test and keep the basedpyright count flat Every new completion() kwarg adds one unknown-argument error at each call site that splats an untyped dict into it; typing print_verbose's argument and Chat's params removes more than prompt_cache_retention added. * test: freeze the outbound-body model and read captured requests from the mock Addresses review: the e2e body model is frozen, the backend params are a tuple, and the unit test reads the request from the respx route instead of appending to a list. --- litellm/main.py | 13 +- .../coverage_registry/llm_conversational.yaml | 2 + tests/e2e/coverage_registry/schema.py | 1 + .../test_provider_features_e2e.py | 168 +++++++++++++++++- tests/e2e/models.py | 1 + tests/unit/test_main.py | 95 ++++++++++ 6 files changed, 275 insertions(+), 5 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index f9bdf64c667..63375292077 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -400,7 +400,11 @@ class LiteLLM: class Chat: - def __init__(self, params, router_obj: "Router | None"): + def __init__( + self, + params: dict[str, object], # mutable-ok: pops the acompletion flag in place + router_obj: "Router | None", + ): self.params = params if self.params.get("acompletion", False) is True: self.params.pop("acompletion") @@ -480,6 +484,7 @@ async def acompletion( service_tier: str | None = None, store: bool | None = None, prompt_cache_key: str | None = None, + prompt_cache_retention: str | None = None, # set api_base, api_version, api_key base_url: str | None = None, api_version: str | None = None, @@ -657,6 +662,7 @@ async def acompletion( "service_tier": service_tier, "store": store, "prompt_cache_key": prompt_cache_key, + "prompt_cache_retention": prompt_cache_retention, "extra_headers": extra_headers, "acompletion": True, # assuming this is a required parameter "thinking": thinking, @@ -5356,6 +5362,7 @@ def completion( service_tier: str | None = None, store: bool | None = None, prompt_cache_key: str | None = None, + prompt_cache_retention: str | None = None, # soon to be deprecated params by OpenAI functions: list | None = None, function_call: str | None = None, @@ -5478,6 +5485,7 @@ def completion( service_tier=service_tier, store=store, prompt_cache_key=prompt_cache_key, + prompt_cache_retention=prompt_cache_retention, base_url=base_url, api_version=api_version, api_key=api_key, @@ -5843,6 +5851,7 @@ def completion( "service_tier": service_tier, "store": store, "prompt_cache_key": prompt_cache_key, + "prompt_cache_retention": prompt_cache_retention, "allowed_openai_params": allowed_openai_params, "base_model": base_model, } @@ -9143,7 +9152,7 @@ async def ahealth_check( ####### HELPER FUNCTIONS ################ ## Set verbose to true -> ```litellm.set_verbose = True``` -def print_verbose(print_statement): +def print_verbose(print_statement: object): try: verbose_logger.debug(print_statement) if litellm.set_verbose: diff --git a/tests/e2e/coverage_registry/llm_conversational.yaml b/tests/e2e/coverage_registry/llm_conversational.yaml index c40aeade1c1..e3994e598c6 100644 --- a/tests/e2e/coverage_registry/llm_conversational.yaml +++ b/tests/e2e/coverage_registry/llm_conversational.yaml @@ -13,6 +13,8 @@ - {id: llm.chat_completions.openai.vision.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: openai, capability: vision, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "gpt-4o vision; high usage"} - {id: llm.chat_completions.openai.prompt_cache_5m.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: openai, capability: prompt_cache_5m, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Prompt caching cost optimization"} - {id: llm.chat_completions.openai.service_tier.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: openai, capability: service_tier, streaming: nonstream, assertions: [works], source: "OpenAI service_tier param", rationale: "OpenAI scale-tier request option is forwarded and echoed"} +- {id: llm.chat_completions.openai.prompt_cache_24h.nonstream.forwards_retention, module: llm, tier: P1, subject_endpoint: chat_completions, route: openai, capability: prompt_cache_24h, streaming: nonstream, assertions: [forwards_retention], source: "main.py completion()", fail_before_fix: proven, rationale: "prompt_cache_retention sent through the proxy, litellm.completion or Router.completion reaches the body sent to OpenAI /v1/chat/completions, so a caller asking for 24h prompt caching gets it instead of a silent drop"} +- {id: llm.chat_completions.openai.prompt_cache_24h.nonstream.bridge_forwards_retention, module: llm, tier: P1, subject_endpoint: chat_completions, route: openai, capability: prompt_cache_24h, streaming: nonstream, assertions: [forwards_retention], source: "main.py completion()", fail_before_fix: proven, rationale: "prompt_cache_retention on an openai/responses/ model (chat completions in, Responses out) reaches the body sent to OpenAI /v1/responses"} - {id: llm.chat_completions.openai.service_tier.stream.echoes_served_tier, module: llm, tier: P1, subject_endpoint: chat_completions, route: openai, capability: service_tier, streaming: stream, assertions: [works], source: "litellm_core_utils/streaming_handler.py", fail_before_fix: proven, rationale: "Every relayed stream chunk carries the service_tier OpenAI stamped on it, so a streaming caller can see which tier served the request"} - {id: llm.chat_completions.openai.thinking.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: openai, capability: thinking, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "o-series reasoning; emerging"} - {id: llm.chat_completions.openai.structured_output.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: openai, capability: structured_output, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "response_schema extraction"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index abef3266aa8..d2ba2c325df 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -81,6 +81,7 @@ LlmCapability = Literal[ "pdf_input", "logprobs", "prompt_cache_1h", + "prompt_cache_24h", "prompt_cache_5m", "response_headers", "service_tier", diff --git a/tests/e2e/llm_translation/test_provider_features_e2e.py b/tests/e2e/llm_translation/test_provider_features_e2e.py index 9a3de0aba7d..60dfd49638f 100644 --- a/tests/e2e/llm_translation/test_provider_features_e2e.py +++ b/tests/e2e/llm_translation/test_provider_features_e2e.py @@ -1,4 +1,4 @@ -"""Live e2e for model-specific request features: service_tier. +"""Live e2e for model-specific request features: service_tier and prompt_cache_retention. Each case asserts the feature took effect, not just a 200. @@ -11,26 +11,38 @@ avoided here because it is capacity-constrained and returns a transient 429 when flex resources are unavailable. Bedrock and Vertex do not accept service_tier, so that cell is OpenAI-only by design. -Prompt caching lives in test_cache_control.py. +prompt_cache_retention is judged on the body OpenAI receives, captured by a live +provider edge, not on the reply: OpenAI echoes "24h" on gpt-5.5 even when the +request never carried the param, so an echo would pass with the param dropped. + +Prompt caching hits and misses live in test_cache_control.py. """ from __future__ import annotations +import queue +from collections.abc import Mapping +from dataclasses import dataclass, field from typing import Final import pytest +from pydantic import BaseModel -from e2e_config import unique_marker +from e2e_config import PROVIDER_EDGE_ADVERTISE_HOST, PROVIDER_EDGE_BIND_HOST, unique_marker from e2e_http import unwrap from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta from lifecycle import ResourceManager from models import ChatBody, ChatMessage, LiteLLMParamsBody from passthrough_client import PassthroughClient +from provider_edge import LiveEdge, start_provider_edge +from proxy_client import ProxyClient pytestmark = pytest.mark.e2e SERVICE_TIER = "priority" OPENAI_BACKEND: Final = "openai/gpt-5.5" +OPENAI_RESPONSES_BRIDGE_BACKEND: Final = "openai/responses/gpt-5.5" +PROMPT_CACHE_RETENTION: Final = "24h" class TestServiceTier: @@ -74,3 +86,153 @@ class TestServiceTier: f"service_tier not honored: sent {SERVICE_TIER!r}, response reported " f"{response.service_tier!r} ({response})" ) + + +class _OutboundRetention(BaseModel, frozen=True): + prompt_cache_retention: str | None = None + + +@dataclass(frozen=True, slots=True) +class _OpenAIBodyCapture: + """The raw bodies sent to one OpenAI path, seen by a live edge in front of + api.openai.com. Pass ``observe`` as the edge's ``observe_request``.""" + + path: str + _bodies: queue.SimpleQueue[bytes] = field(default_factory=queue.SimpleQueue) + + def observe(self, url: str, headers: Mapping[str, str], body: bytes | None) -> None: + if body is not None and url.split("?")[0].endswith(self.path): + self._bodies.put(body) + + def retentions(self) -> tuple[str | None, ...]: + bodies: Final = tuple(self._bodies.get_nowait() for _ in range(self._bodies.qsize())) + return tuple(_OutboundRetention.model_validate_json(body).prompt_cache_retention for body in bodies) + + +def _openai_edge_api_base(resources: ResourceManager, capture: _OpenAIBodyCapture, *, reachable_from_proxy: bool) -> str: + """The proxy may run in a container, so its edge binds where the proxy can + reach it; an in-process SDK call reaches the edge on loopback.""" + edge: Final = ( + start_provider_edge( + LiveEdge(observe_request=capture.observe), + bind_host=PROVIDER_EDGE_BIND_HOST, + advertise_host=PROVIDER_EDGE_ADVERTISE_HOST, + ) + if reachable_from_proxy + else start_provider_edge(LiveEdge(observe_request=capture.observe)) + ) + resources.defer(edge.shutdown) + return f"{edge.edge.api_base('openai')}/v1" + + +def _retention_backends() -> tuple[object, ...]: + """gpt-5.5 straight to /v1/chat/completions, and the openai/responses/ bridge + (chat completions in, Responses out) to /v1/responses.""" + return ( + pytest.param( + OPENAI_BACKEND, + "/chat/completions", + marks=( + pytest.mark.covers("llm.chat_completions.openai.prompt_cache_24h.nonstream.forwards_retention"), + meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.CHAT_COMPLETIONS, + providers=(Provider.OPENAI,), + models=(OPENAI_BACKEND,), + mode=Mode.NONSTREAM, + ) + ), + ), + id="chat_completions", + ), + pytest.param( + OPENAI_RESPONSES_BRIDGE_BACKEND, + "/responses", + marks=( + pytest.mark.covers("llm.chat_completions.openai.prompt_cache_24h.nonstream.bridge_forwards_retention"), + meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.CHAT_COMPLETIONS, + providers=(Provider.OPENAI,), + models=(OPENAI_RESPONSES_BRIDGE_BACKEND,), + mode=Mode.NONSTREAM, + ) + ), + ), + id="responses_bridge", + ), + ) + + +class TestPromptCacheRetention: + @pytest.mark.provider_edge_host + @pytest.mark.parametrize(("backend", "openai_path"), _retention_backends()) + def test_proxy_chat_completions_forwards_prompt_cache_retention_to_openai( + self, proxy: ProxyClient, resources: ResourceManager, backend: str, openai_path: str + ) -> None: + capture: Final = _OpenAIBodyCapture(openai_path) + api_base: Final = _openai_edge_api_base(resources, capture, reachable_from_proxy=True) + model: Final = f"e2e-cache-retention-{unique_marker()}" + model_id: Final = proxy.create_model( + model, LiteLLMParamsBody(model=backend, api_key="os.environ/OPENAI_API_KEY", api_base=api_base) + ) + resources.defer(lambda: proxy.delete_model(model_id)) + + unwrap( + proxy.chat( + resources.key(), + ChatBody( + model=model, + messages=[ChatMessage(role="user", content="reply with one word")], + prompt_cache_retention=PROMPT_CACHE_RETENTION, + ), + ) + ) + + assert capture.retentions() == (PROMPT_CACHE_RETENTION,), ( + f"one {openai_path} call to OpenAI must carry prompt_cache_retention={PROMPT_CACHE_RETENTION!r}" + ) + + @pytest.mark.parametrize(("backend", "openai_path"), _retention_backends()) + def test_sdk_completion_forwards_prompt_cache_retention_to_openai( + self, resources: ResourceManager, backend: str, openai_path: str + ) -> None: + import litellm + + capture: Final = _OpenAIBodyCapture(openai_path) + api_base: Final = _openai_edge_api_base(resources, capture, reachable_from_proxy=False) + + litellm.completion( + model=backend, + messages=[{"role": "user", "content": "reply with one word"}], + api_base=api_base, + prompt_cache_retention=PROMPT_CACHE_RETENTION, + ) + + assert capture.retentions() == (PROMPT_CACHE_RETENTION,), ( + f"one {openai_path} call to OpenAI must carry prompt_cache_retention={PROMPT_CACHE_RETENTION!r}" + ) + + @pytest.mark.parametrize(("backend", "openai_path"), _retention_backends()) + def test_router_completion_forwards_prompt_cache_retention_to_openai( + self, resources: ResourceManager, backend: str, openai_path: str + ) -> None: + import litellm + + capture: Final = _OpenAIBodyCapture(openai_path) + api_base: Final = _openai_edge_api_base(resources, capture, reachable_from_proxy=False) + router: Final = litellm.Router( + model_list=[{"model_name": "cache-retention", "litellm_params": {"model": backend, "api_base": api_base}}] + ) + + router.completion( # pyright: ignore[reportUnknownMemberType] # Router.completion forwards an untyped **kwargs to litellm.completion + model="cache-retention", + messages=[{"role": "user", "content": "reply with one word"}], + prompt_cache_retention=PROMPT_CACHE_RETENTION, + ) + + assert capture.retentions() == (PROMPT_CACHE_RETENTION,), ( + f"one {openai_path} call to OpenAI must carry prompt_cache_retention={PROMPT_CACHE_RETENTION!r}" + ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index f0a8d67b444..afddc7e25d6 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -364,6 +364,7 @@ class ChatBody(BaseModel): thinking: ThinkingParam | None = None service_tier: str | None = None prompt_cache_key: str | None = None + prompt_cache_retention: str | None = None tools: Sequence[ChatTool | McpChatTool | GeminiUrlContextTool] | None = None tool_choice: str | None = None guardrails: list[str] | None = None diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 3e5cfeb455e..f5d13aecb6f 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -23,6 +23,7 @@ import respx from openai import APITimeoutError, AsyncOpenAI from openai._legacy_response import LegacyAPIResponse from openai.types.chat.chat_completion import ChatCompletion +from pydantic import TypeAdapter import litellm from litellm import ( @@ -3937,6 +3938,100 @@ async def test_acompletion_forwards_store_and_prompt_cache_key_to_openai(): assert request_body["prompt_cache_key"] == "test-cache-key" +_TOOL: Final[Mapping[str, object]] = {"type": "function", "function": {"name": "f", "parameters": {"type": "object", "properties": {}}}} + +# Client settings ride in OPENAI_CHAT_COMPLETION_PARAMS but configure the HTTP client, never the body. +_OPENAI_CLIENT_SETTINGS: Final = frozenset( + { + "api_base", + "api_key", + "api_version", + "base_url", + "default_headers", + "deployment_id", + "extra_headers", + "max_retries", + "organization", + "request_timeout", + "timeout", + } +) + +# Accepted by completion() for other providers, never sent to OpenAI's /chat/completions body. +_NOT_OPENAI_BODY_PARAMS: Final = frozenset({"include_server_side_tool_invocations", "thinking"}) + +# param -> (model that accepts it, the call kwargs that carry it) +_OPENAI_BODY_PARAM_CASES: Final[Mapping[str, tuple[str, Mapping[str, object]]]] = { + "audio": ("openai/gpt-4o", {"audio": {"voice": "alloy", "format": "wav"}}), + "frequency_penalty": ("openai/gpt-4o", {"frequency_penalty": 0.5}), + "function_call": ("openai/gpt-4o", {"function_call": "auto"}), + "functions": ("openai/gpt-4o", {"functions": [_TOOL["function"]]}), + "logit_bias": ("openai/gpt-4o", {"logit_bias": {"1": 1}}), + "logprobs": ("openai/gpt-4o", {"logprobs": True}), + "max_completion_tokens": ("openai/gpt-4o", {"max_completion_tokens": 7}), + "max_tokens": ("openai/gpt-4o", {"max_tokens": 7}), + "modalities": ("openai/gpt-4o", {"modalities": ["text"]}), + "n": ("openai/gpt-4o", {"n": 2}), + "parallel_tool_calls": ("openai/gpt-4o", {"tools": [_TOOL], "parallel_tool_calls": False}), + "prediction": ("openai/gpt-4o", {"prediction": {"type": "content", "content": "x"}}), + "presence_penalty": ("openai/gpt-4o", {"presence_penalty": 0.5}), + "prompt_cache_key": ("openai/gpt-4o", {"prompt_cache_key": "k"}), + "prompt_cache_retention": ("openai/gpt-5.5", {"prompt_cache_retention": "24h"}), + "reasoning_effort": ("openai/gpt-5.5", {"reasoning_effort": "low"}), + "response_format": ("openai/gpt-4o", {"response_format": {"type": "json_object"}}), + "safety_identifier": ("openai/gpt-4o", {"safety_identifier": "s"}), + "seed": ("openai/gpt-4o", {"seed": 3}), + "service_tier": ("openai/gpt-4o", {"service_tier": "priority"}), + "stop": ("openai/gpt-4o", {"stop": ["x"]}), + "store": ("openai/gpt-4o", {"store": False}), + "stream": ("openai/gpt-4o", {"stream": True}), + "stream_options": ("openai/gpt-4o", {"stream": True, "stream_options": {"include_usage": True}}), + "temperature": ("openai/gpt-4o", {"temperature": 0.5}), + "tool_choice": ("openai/gpt-4o", {"tools": [_TOOL], "tool_choice": "auto"}), + "tools": ("openai/gpt-4o", {"tools": [_TOOL]}), + "top_logprobs": ("openai/gpt-4o", {"logprobs": True, "top_logprobs": 2}), + "top_p": ("openai/gpt-4o", {"top_p": 0.5}), + "user": ("openai/gpt-4o", {"user": "u1"}), + "verbosity": ("openai/gpt-5.5", {"verbosity": "low"}), + "web_search_options": ("openai/gpt-4o", {"web_search_options": {"search_context_size": "low"}}), +} + + +def test_every_openai_body_param_has_a_forwarding_case(): + body_params: Final = frozenset(litellm.OPENAI_CHAT_COMPLETION_PARAMS) - _OPENAI_CLIENT_SETTINGS - _NOT_OPENAI_BODY_PARAMS + assert frozenset(_OPENAI_BODY_PARAM_CASES) == body_params, ( + "a param added to OPENAI_CHAT_COMPLETION_PARAMS needs a case in _OPENAI_BODY_PARAM_CASES proving " + "completion() sends it to OpenAI; prompt_cache_retention (LIT-9758) and store/prompt_cache_key " + "(#33184) were listed as supported but silently dropped" + ) + + +_REQUEST_BODY: Final = TypeAdapter(dict[str, object]) + + +@pytest.mark.parametrize("param", sorted(_OPENAI_BODY_PARAM_CASES)) +def test_completion_sends_openai_body_param_to_openai( + param: str, respx_mock: respx.MockRouter, openai_api_response: Mapping[str, object] +): + api_base: Final = "http://localhost:12346/v1" + route: Final = respx_mock.post(f"{api_base}/chat/completions").mock( + return_value=httpx.Response(status_code=200, json=openai_api_response) + ) + model, call_kwargs = _OPENAI_BODY_PARAM_CASES[param] + + litellm.completion( + model=model, + messages=[{"role": "user", "content": "hi"}], + api_base=api_base, + api_key="fake_openai_api_key", + **call_kwargs, + ) + + assert route.call_count == 1 + sent: Final = _REQUEST_BODY.validate_json(route.calls.last.request.content) + assert sent.get(param) == call_kwargs[param], f"{param} did not reach the OpenAI request body: {sent}" + + def test_completion_omits_store_and_prompt_cache_key_when_not_passed(): """ When store and prompt_cache_key are not passed, they must not appear in the From 5ab0130abcdf8725dc67afd8e5f30c023589a1f2 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 20 Oct 2026 09:50:17 -0700 Subject: [PATCH 28/33] test(integration): call _fallback_deployment in the count_tokens wire tests #47363 added (#47593) #47525 renamed _deployment to _fallback_deployment; #47363 was written against the old name and merged after it, leaving three F821 undefined-name errors that fail the python lint job on main Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../providers/test_bedrock_mantle_count_tokens_wire.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py index 063b0e46479..25f46be4c5d 100644 --- a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py @@ -1683,8 +1683,8 @@ def test_responses_input_tokens_sends_mantle_the_tool_use_invoke_inference_sends wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - counted_model: Final = _deployment(scenario, runtime.url) - invoked_model: Final = _deployment(scenario, runtime.url, f"invoke/{_OPUS}") + counted_model: Final = _fallback_deployment(scenario, runtime.url) + invoked_model: Final = _fallback_deployment(scenario, runtime.url, f"invoke/{_OPUS}") generated: Final = gateway.request( "POST", "/v1/responses", {**_FORCED_TOOL_CHOICE, "model": invoked_model, "max_output_tokens": 16} ) @@ -1729,7 +1729,7 @@ def test_nova_count_sends_the_tool_specs_converse_inference_sends(counting_proxy path, body, count_field = _NOVA_COUNTS[route] gateway: Final = counting_proxy.gateway with wire_server(_counting_every_model) as runtime, gateway.scenario() as scenario: - model: Final = _deployment(scenario, runtime.url, _NOVA) + model: Final = _fallback_deployment(scenario, runtime.url, _NOVA) generated: Final = gateway.request( "POST", "/v1/chat/completions", From 2c6b280755ba617cd6647c28706abdd581a7dd19 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 20 Oct 2026 09:51:55 -0700 Subject: [PATCH 29/33] chore: lead the PR template with a short plain description and collapse the rest (#47594) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/PULL_REQUEST_TEMPLATE/general.md | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/.github/PULL_REQUEST_TEMPLATE/general.md b/.github/PULL_REQUEST_TEMPLATE/general.md index 644d1c4a450..7d920b715e8 100644 --- a/.github/PULL_REQUEST_TEMPLATE/general.md +++ b/.github/PULL_REQUEST_TEMPLATE/general.md @@ -4,6 +4,15 @@ Drop every section you have nothing to put in, heading included: a bare "## Relevant issues" or "## Affected release" with nothing under it must not appear in the final description --> +## Description + + + +
+Details + ## TLDR