From f1795e6d447a4e5a3660de0d2fde5e1f44f52929 Mon Sep 17 00:00:00 2001 From: Shivam Rawat Date: Wed, 28 Oct 2026 00:51:10 -0700 Subject: [PATCH 001/128] fix(responses): keep every MCP tool round in streaming follow-up input (#47698) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: ryan-crabbe-berri --- .../responses/mcp/mcp_streaming_iterator.py | 7 +- ...t_responses_mcp_streaming_round_history.py | 174 ++++++++++++++++++ .../mcp/test_mcp_streaming_iterator.py | 53 +++++- 3 files changed, 232 insertions(+), 2 deletions(-) create mode 100644 tests/integration/mcp/test_responses_mcp_streaming_round_history.py diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index cc79d1c88bb..189366feab4 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -20,6 +20,7 @@ from litellm.types.llms.openai import ( MCPListToolsCompletedEvent, MCPListToolsFailedEvent, MCPListToolsInProgressEvent, + ResponseInputParam, ResponsesAPIResponse, ResponsesAPIStreamEvents, ResponsesAPIStreamingResponse, @@ -290,6 +291,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self.mcp_tools_with_litellm_proxy = mcp_tools_with_litellm_proxy or [] self.user_api_key_auth = user_api_key_auth self.original_request_params = original_request_params or {} + self._follow_up_history: str | ResponseInputParam | None = self.original_request_params.get("input") configured_rounds: Final = validated_max_agentic_loops( self.original_request_params.get("max_agentic_loops"), field="litellm_params.max_agentic_loops" ) @@ -1003,9 +1005,12 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): follow_up_input: Final = LiteLLM_Proxy_MCP_Handler.create_follow_up_input( response=self.collected_response, tool_results=self.tool_results, - original_input=self.original_request_params.get("input"), + original_input=self._follow_up_history, preserve_reasoning=persistence_disabled, ) + self._follow_up_history = cast( # cast-ok: helper returns valid input items + "ResponseInputParam", follow_up_input + ) # Make follow-up call with streaming follow_up_params: Final = self.original_request_params.copy() diff --git a/tests/integration/mcp/test_responses_mcp_streaming_round_history.py b/tests/integration/mcp/test_responses_mcp_streaming_round_history.py new file mode 100644 index 00000000000..8313d96e23c --- /dev/null +++ b/tests/integration/mcp/test_responses_mcp_streaming_round_history.py @@ -0,0 +1,174 @@ +from __future__ import annotations + +import json +import uuid +from collections.abc import Mapping +from typing import Final + +from integration._support.client import Gateway +from integration._support.mcp import mcp_peer, register_mcp, tool_calls +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +TOOL_ARGUMENTS: Final = '{"a":2,"b":3}' +ANSWER: Final = "All three tool rounds are done." + + +def _event_frame(event: Mapping[str, JsonValue]) -> bytes: + return f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() + + +def _streamed_response(round_number: int, model: str, tool_name: str) -> Reply: + response_id: Final = f"resp_mcp_round_{round_number}" + call_id: Final = f"call_{round_number}" + is_tool_call: Final = round_number <= 3 + output_item: Final[dict[str, JsonValue]] = ( + { + "type": "function_call", + "id": f"fc_{round_number}", + "call_id": call_id, + "name": tool_name, + "arguments": TOOL_ARGUMENTS, + "status": "completed", + } + if is_tool_call + else { + "type": "message", + "id": f"msg_{round_number}", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": ANSWER, "annotations": []}], + } + ) + added_item: Final[dict[str, JsonValue]] = ( + {**output_item, "arguments": "", "status": "in_progress"} + if is_tool_call + else {**output_item, "content": [], "status": "in_progress"} + ) + item_events: Final[tuple[dict[str, JsonValue], ...]] = ( + ( + {"type": "response.output_item.added", "output_index": 0, "item": added_item}, + { + "type": "response.function_call_arguments.delta", + "item_id": f"fc_{round_number}", + "output_index": 0, + "delta": TOOL_ARGUMENTS, + }, + { + "type": "response.function_call_arguments.done", + "item_id": f"fc_{round_number}", + "output_index": 0, + "arguments": TOOL_ARGUMENTS, + }, + {"type": "response.output_item.done", "output_index": 0, "item": output_item}, + ) + if is_tool_call + else ( + {"type": "response.output_item.added", "output_index": 0, "item": added_item}, + { + "type": "response.output_text.delta", + "item_id": f"msg_{round_number}", + "output_index": 0, + "content_index": 0, + "delta": ANSWER, + }, + {"type": "response.output_item.done", "output_index": 0, "item": output_item}, + ) + ) + response: Final[dict[str, JsonValue]] = { + "id": response_id, + "object": "response", + "created_at": 1, + "status": "completed", + "model": model, + "output": [output_item], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + events: Final[tuple[dict[str, JsonValue], ...]] = ( + { + "type": "response.created", + "response": {**response, "status": "in_progress", "output": []}, + }, + { + "type": "response.in_progress", + "response": {**response, "status": "in_progress", "output": []}, + }, + *item_events, + {"type": "response.completed", "response": response}, + ) + chunks: Final = tuple( + _event_frame({**event, "sequence_number": sequence_number}) for sequence_number, event in enumerate(events) + ) + return Reply(content_type="text/event-stream", chunks=chunks) + + +def _tool_call_history(body: Mapping[str, JsonValue]) -> tuple[tuple[str, str], ...]: + input_items: Final = body["input"] + assert isinstance(input_items, list) + return tuple( + (str(item["type"]), str(item["call_id"])) + for item in input_items + if isinstance(item, dict) and item.get("type") in ("function_call", "function_call_output") + ) + + +def test_streaming_responses_mcp_follow_up_keeps_every_tool_round(gateway: Gateway) -> None: + alias: Final = "hist" + uuid.uuid4().hex[:8] + upstream_requests: Final[list[dict[str, JsonValue]]] = [] + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target.endswith("/models"): + return Reply(body=b'{"object":"list","data":[]}') + assert request.target.endswith("/responses"), request.target + request_body: Final = JSON_OBJECT.validate_json(request.body) + upstream_requests.append(request_body) + return _streamed_response(len(upstream_requests), str(request_body["model"]), f"{alias}-add") + + with mcp_peer() as peer, wire_server(respond) as wire, gateway.scenario() as scenario: + server_id: Final = register_mcp(scenario, peer, alias) + model: Final = scenario.model(model="openai/responses/gpt-4o-mini", api_base=wire.url + "/v1") + key: Final = scenario.key(object_permission={"mcp_servers": [server_id]}) + peer.drain() + response: Final = gateway.client.post( + "/v1/responses", + headers={"Authorization": f"Bearer {key}"}, + json={ + "model": model, + "input": "run each tool round in order", + "stream": True, + "tools": [ + { + "type": "mcp", + "server_url": "litellm_proxy", + "server_label": "litellm", + "require_approval": "never", + } + ], + }, + timeout=90, + ) + assert response.status_code == 200, response.text + assert "response.completed" in response.text and ANSWER in response.text, response.text + + assert len(upstream_requests) == 4, upstream_requests + assert _tool_call_history(upstream_requests[2]) == ( + ("function_call", "call_1"), + ("function_call_output", "call_1"), + ("function_call", "call_2"), + ("function_call_output", "call_2"), + ) + assert _tool_call_history(upstream_requests[3]) == ( + ("function_call", "call_1"), + ("function_call_output", "call_1"), + ("function_call", "call_2"), + ("function_call_output", "call_2"), + ("function_call", "call_3"), + ("function_call_output", "call_3"), + ) + third_input: Final = upstream_requests[2]["input"] + fourth_input: Final = upstream_requests[3]["input"] + assert isinstance(third_input, list) and isinstance(fourth_input, list) + assert sum(1 for item in third_input if isinstance(item, dict) and item.get("role") == "user") == 1 + assert sum(1 for item in fourth_input if isinstance(item, dict) and item.get("role") == "user") == 1 + assert len(tool_calls(peer.drain())) == 3 diff --git a/tests/unit/responses/mcp/test_mcp_streaming_iterator.py b/tests/unit/responses/mcp/test_mcp_streaming_iterator.py index 78b0e3d4e5b..f745b5bfadf 100644 --- a/tests/unit/responses/mcp/test_mcp_streaming_iterator.py +++ b/tests/unit/responses/mcp/test_mcp_streaming_iterator.py @@ -2,7 +2,7 @@ import sys import types from contextlib import nullcontext from types import SimpleNamespace -from typing import Final +from typing import Final, cast from unittest.mock import AsyncMock, MagicMock import pytest @@ -66,6 +66,14 @@ def _item_type(item: dict[str, object] | BaseLiteLLMOpenAIResponseObject) -> str return str(item["type"]) if isinstance(item, dict) else str(item.type) +def _follow_up_tool_history(input_items: list[object]) -> tuple[tuple[str, str], ...]: + return tuple( + (str(item.get("type")), str(item.get("call_id"))) + for item in input_items + if isinstance(item, dict) and item.get("type") in ("function_call", "function_call_output") + ) + + def _tool_call_stream(call_id: str, tool_name: str, response_id: str = "resp-1") -> _FakeAsyncStream: return _FakeAsyncStream([_completed_chunk([_function_call(call_id, tool_name)], response_id=response_id)]) @@ -157,6 +165,49 @@ async def test_second_round_tool_call_is_executed_and_reaches_final_text(monkeyp assert final_output[-1]["content"][0]["text"] == "Here's what I found after retrying." +@pytest.mark.asyncio +async def test_follow_up_input_keeps_every_previous_tool_round(monkeypatch: pytest.MonkeyPatch) -> None: + _mock_mcp_environment(monkeypatch) + + aresponses_mock: Final = AsyncMock( + side_effect=[ + _tool_call_stream("call_2", "read_wiki_contents"), + _tool_call_stream("call_3", "read_wiki_contents"), + _text_only_stream("All three tool rounds are done."), + ] + ) + monkeypatch.setattr(responses_main_module, "aresponses", aresponses_mock) + iterator: Final = _make_iterator( + [ + _output_item_added_chunk(), + _completed_chunk([_function_call("call_1", "read_wiki_contents")]), + ] + ) + + _ = [chunk async for chunk in iterator] + + second_follow_up_input: Final = cast("list[object]", aresponses_mock.call_args_list[1].kwargs["input"]) + third_follow_up_input: Final = cast("list[object]", aresponses_mock.call_args_list[2].kwargs["input"]) + assert _follow_up_tool_history(second_follow_up_input) == ( + ("function_call", "call_1"), + ("function_call_output", "call_1"), + ("function_call", "call_2"), + ("function_call_output", "call_2"), + ) + assert _follow_up_tool_history(third_follow_up_input) == ( + ("function_call", "call_1"), + ("function_call_output", "call_1"), + ("function_call", "call_2"), + ("function_call_output", "call_2"), + ("function_call", "call_3"), + ("function_call_output", "call_3"), + ) + user_messages: Final = [ + item for item in third_follow_up_input if isinstance(item, dict) and item.get("role") == "user" + ] + assert user_messages == [{"type": "message", "role": "user", "content": "what is berriai/litellm?"}] + + @pytest.mark.asyncio @pytest.mark.parametrize( ("max_agentic_loops", "expected_cap"), From 220a70653e5b57d2ca75815e1ad71632edea726e Mon Sep 17 00:00:00 2001 From: zoveki <349028329+zoveki@users.noreply.github.com> Date: Wed, 28 Oct 2026 01:38:07 -0700 Subject: [PATCH 002/128] docs(readme): link the LiteLLM OSS Champions program from Contributing (#48734) * docs(readme): link the LiteLLM OSS Champions program from Contributing * docs(readme): frame the champions perks the way the program page does * docs(readme): drop the champions swag claim the page gates on merged PRs * docs(readme): word the champions line as the invitation the page makes * docs(readme): invite readers to the champions page without promising its perks --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/README.md b/README.md index 5dfa8d01fc4..66b6acb7a71 100644 --- a/README.md +++ b/README.md @@ -660,6 +660,8 @@ This covers: We welcome contributions to LiteLLM! Whether you're fixing bugs, adding features, or improving documentation, we appreciate your help. +If you want to work closely with our team, check out [LiteLLM OSS Champions](https://www.litellm.ai/champions), our program for trusted OSS contributors with their own Slack channel and a voice in the roadmap + ## Quick Start for Contributors This requires uv to be installed. From 1176e608b199c4a6572c22490839147a9f134ece Mon Sep 17 00:00:00 2001 From: norkelbyte <348963847+norkelbyte@users.noreply.github.com> Date: Wed, 28 Oct 2026 01:57:41 -0700 Subject: [PATCH 003/128] test(harness): audit #47530 (#48751) Check that cancelling a LocalSandbox run stops the grandchild process it started, using a real child process Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../integration/sdk/test_local_sandbox_sdk.py | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/tests/integration/sdk/test_local_sandbox_sdk.py b/tests/integration/sdk/test_local_sandbox_sdk.py index b30010ac3f8..b1b9c89adfd 100644 --- a/tests/integration/sdk/test_local_sandbox_sdk.py +++ b/tests/integration/sdk/test_local_sandbox_sdk.py @@ -1,16 +1,25 @@ +import asyncio import os import sys from collections.abc import AsyncIterator from pathlib import Path from typing import Final +import psutil import pytest import pytest_asyncio +from integration._support.client import eventually +from integration._support.process import kill_stray_process from litellm import sandbox from litellm.harness.sandbox import LocalSandbox WRITE_AT_PWD: Final = "import os;open(os.path.join(os.environ['PWD'],'hello.txt'),'w').write('hi')" +SLEEPER_SPAWNED_ONCE_STDIN_CLOSES: Final = ( + "import subprocess,sys;sys.stdin.read();" + "p=subprocess.Popen([sys.executable,'-c','import time;time.sleep(60)']);" + "open('grandchild.pid','w').write(str(p.pid));p.wait()" +) @pytest_asyncio.fixture(loop_scope="function") @@ -37,3 +46,27 @@ async def test_child_trusting_pwd_writes_inside_its_cwd( assert os.listdir(elsewhere) == [] assert await box.read("hello.txt") == b"hi" assert await box.read("sub/hello.txt") == b"hi" + + +def running(pid: int) -> bool: + try: + return psutil.Process(pid).status() != psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: + return False + + +@pytest.mark.asyncio(loop_scope="function") +async def test_cancelled_run_stops_the_whole_command(box: LocalSandbox) -> None: + run: Final = asyncio.ensure_future(box.run([sys.executable, "-I", "-c", SLEEPER_SPAWNED_ONCE_STDIN_CLOSES])) + pid_file: Final = Path(box.workdir) / "grandchild.pid" + grandchild: Final = int( + await asyncio.to_thread(eventually, lambda: pid_file.read_text() if pid_file.exists() else "", bool) + ) + try: + run.cancel() + with pytest.raises(asyncio.CancelledError): + await run + + assert not eventually(lambda: running(grandchild), lambda alive: not alive, return_last_on_timeout=True) + finally: + kill_stray_process(grandchild) From 000d4f6bb3c8d3cf0e2bb9540a61ccd77ffcd5b9 Mon Sep 17 00:00:00 2001 From: rj-cairn <349020452+rj-cairn@users.noreply.github.com> Date: Wed, 28 Oct 2026 02:40:45 -0700 Subject: [PATCH 004/128] test(sagemaker): audit #47495 (#48767) Checks that a SageMaker stream failing after its first chunks reaches /v1/chat/completions, /v1/messages and /v1/responses callers as an error with no success end event, and that each writes a failed spend log row Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../providers/test_sagemaker_chat_wire.py | 122 ++++++++++++++++++ 1 file changed, 122 insertions(+) diff --git a/tests/integration/providers/test_sagemaker_chat_wire.py b/tests/integration/providers/test_sagemaker_chat_wire.py index 261787fec38..fdf0d94f7c7 100644 --- a/tests/integration/providers/test_sagemaker_chat_wire.py +++ b/tests/integration/providers/test_sagemaker_chat_wire.py @@ -7,6 +7,7 @@ from dataclasses import dataclass from types import MappingProxyType from typing import Final, Literal, TypeAlias +import anthropic import httpx import openai import pytest @@ -703,3 +704,124 @@ async def test_sagemaker_chat_container_outage_mid_burst_answers_every_call_once assert statuses == { served.call_id: "failure" if served.call.marker in down else "success" for served in answered }, rows + + +_MID_STREAM_ERROR: Final = "synthetic mid-stream container failure" + + +@dataclass(frozen=True, slots=True) +class _CutStream: + call_id: str + text: str + errors: tuple[str, ...] + finished: bool + + +async def _chat_chunks_until_error( + stream: openai.AsyncStream[openai.types.chat.ChatCompletionChunk], +) -> AsyncIterator[openai.types.chat.ChatCompletionChunk | openai.APIError]: + try: + async for chunk in stream: + yield chunk + except openai.APIError as error: + yield error + + +async def _chat_cut(gateway: Gateway, model: str, prompt: str) -> _CutStream: + async with async_openai_client(gateway) as client: + stream: Final = await client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": prompt}], + max_tokens=16, + stream=True, + extra_body=dict(NO_CACHE), + ) + items: Final = tuple([item async for item in _chat_chunks_until_error(stream)]) + chunks: Final = tuple(item for item in items if not isinstance(item, openai.APIError)) + choices: Final = tuple(itertools.chain.from_iterable(chunk.choices for chunk in chunks)) + return _CutStream( + call_id=stream.response.headers["x-litellm-call-id"], + text="".join(choice.delta.content or "" for choice in choices), + errors=tuple(item.message for item in items if isinstance(item, openai.APIError)), + finished=any(choice.finish_reason for choice in choices), + ) + + +async def _anthropic_events_until_error( + stream: anthropic.AsyncStream[anthropic.types.RawMessageStreamEvent], +) -> AsyncIterator[anthropic.types.RawMessageStreamEvent | anthropic.APIError]: + try: + async for event in stream: + yield event + except anthropic.APIError as error: + yield error + + +async def _messages_cut(gateway: Gateway, model: str, prompt: str) -> _CutStream: + async with async_anthropic_client(gateway) as client: + stream: Final = await client.messages.create( + model=model, + max_tokens=16, + messages=[{"role": "user", "content": prompt}], + stream=True, + extra_body=dict(NO_CACHE), + ) + items: Final = tuple([item async for item in _anthropic_events_until_error(stream)]) + events: Final = tuple(item for item in items if not isinstance(item, anthropic.APIError)) + deltas: Final = tuple(event.delta for event in events if event.type == "content_block_delta") + return _CutStream( + call_id=stream.response.headers["x-litellm-call-id"], + text="".join(delta.text for delta in deltas if delta.type == "text_delta"), + errors=tuple(str(item) for item in items if isinstance(item, anthropic.APIError)), + finished=any(event.type == "message_stop" for event in events), + ) + + +def _is_error_event(event: Mapping[str, JsonValue]) -> bool: + return "error" in event or event.get("type") in ("error", "response.failed") + + +def _responses_cut(gateway: Gateway, model: str, prompt: str) -> _CutStream: + with gateway.client.stream( + "POST", + "/v1/responses", + json={"model": model, "input": prompt, "max_output_tokens": 16, "stream": True, **NO_CACHE}, + headers={"Authorization": f"Bearer {gateway.key}"}, + ) as response: + lines: Final = tuple(response.iter_lines()) + events: Final = tuple( + JSON_OBJECT.validate_json(line.removeprefix("data: ")) for line in lines if line.startswith("data: {") + ) + return _CutStream( + call_id=response.headers["x-litellm-call-id"], + text="".join( + string_value(event["delta"]) for event in events if event.get("type") == "response.output_text.delta" + ), + errors=tuple(json.dumps(event) for event in events if _is_error_event(event)), + finished=any(event.get("type") == "response.completed" for event in events), + ) + + +_CUT_DRIVERS: Final[Mapping[Endpoint, Callable[[Gateway, str, str], _CutStream]]] = MappingProxyType( + { + "chat": lambda gateway, model, prompt: asyncio.run(_chat_cut(gateway, model, prompt)), + "messages": lambda gateway, model, prompt: asyncio.run(_messages_cut(gateway, model, prompt)), + "responses": _responses_cut, + } +) + + +@pytest.mark.parametrize("endpoint", _ENDPOINTS) +def test_sagemaker_chat_stream_failure_after_the_first_chunks_reaches_every_endpoint_as_an_error_and_a_failed_spend_row( + gateway: Gateway, endpoint: Endpoint +) -> None: + events: Final = vllm_stream_events(f"sagemaker-chat-{uuid.uuid4().hex}") + reply: Final = event_stream_reply((*events[:2], _sse({"error": {"message": _MID_STREAM_ERROR, "code": 500}}))) + + with wire_server(lambda _request: reply) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, {}) + cut: Final = _CUT_DRIVERS[endpoint](gateway, model, marked_prompt()) + only_request(wire) + assert (cut.text, cut.finished) == ("".join(CONTENT), False), cut + assert any(_MID_STREAM_ERROR in error for error in cut.errors), cut + assert _spend_row(cut.call_id)["status"] == "failure", cut From d8a43a7cd31e6bf7ee6bfd192249513497c39ce0 Mon Sep 17 00:00:00 2001 From: zoveki <349028329+zoveki@users.noreply.github.com> Date: Wed, 28 Oct 2026 02:42:45 -0700 Subject: [PATCH 005/128] test(gemini): audit #47542 (#48758) * test(gemini): audit #47542 Proxy integration cells for the Gemini and Vertex default output budget: the env budget reaches Google on chat and responses, turns on long decode, yields to the caller's own limit, caps per deployment, skips resumes, and a camelCase two-candidate request is never continued * test(gemini): annotate the budget wire module's pytestmark as Final --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../test_gemini_default_output_budget_wire.py | 207 ++++++++++++++++++ .../providers/test_gemini_long_decode_wire.py | 17 ++ 2 files changed, 224 insertions(+) create mode 100644 tests/integration/providers/test_gemini_default_output_budget_wire.py diff --git a/tests/integration/providers/test_gemini_default_output_budget_wire.py b/tests/integration/providers/test_gemini_default_output_budget_wire.py new file mode 100644 index 00000000000..3eab626690d --- /dev/null +++ b/tests/integration/providers/test_gemini_default_output_budget_wire.py @@ -0,0 +1,207 @@ +from __future__ import annotations + +import itertools +import uuid +from collections.abc import Iterator, Mapping +from typing import Final + +import pytest +from _pytest.mark.structures import ParameterSet +from integration._support.client import JSON_OBJECT, Gateway, Scenario +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.wire import Wire, wire_server +from integration.providers._gemini_long_decode_support import ( + API_KEY, + BACKEND, + CAP, + CUT_TOKEN, + FIRST_HALF, + HOP_TOKEN, + INPUT_COST, + NO_CACHE, + OUTPUT_COST, + PLAIN, + PROVIDERS, + RESUMED, + SECOND_HALF, + Endpoint, + Provider, + answered, + answers, + at, + complete, + post, + received, + register, + request_body, + upstream, +) +from pydantic import JsonValue + +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_BUDGET_ENV: Final = "LITELLM_GEMINI_DEFAULT_MAX_OUTPUT_TOKENS" +_SMALL_BACKEND: Final = "gemini-2.5-flash" +_LIMIT_KEYS: Final = ("max_output_tokens", "maxOutputTokens") +_JOINED: Final = FIRST_HALF.text + SECOND_HALF.text + + +def _row_limit(key: str) -> int: + limit: Final = JSON_OBJECT.validate_python(GetModelCostMap.load_local_model_cost_map()[key]).get( + "max_output_tokens" + ) + assert isinstance(limit, int), (key, limit) + return limit + + +_SMALL_LIMIT: Final = _row_limit(f"gemini/{_SMALL_BACKEND}") +_LARGE_LIMIT: Final = min(_row_limit(f"{provider}/{BACKEND}") for provider in PROVIDERS) +_BUDGET: Final = (_SMALL_LIMIT + _LARGE_LIMIT) // 2 + + +@pytest.fixture(scope="module") +def budgeted(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + assert _SMALL_LIMIT < _BUDGET < _LARGE_LIMIT, (_SMALL_LIMIT, _BUDGET, _LARGE_LIMIT) + directory: Final = tmp_path_factory.mktemp("gemini-default-output-budget") + with owned_proxy(session_gateway, directory, {_BUDGET_ENV: str(_BUDGET)}, workers=2) as owned: + yield owned + + +def _limits_sent(body: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + config: Final = body.get("generationConfig") + assert isinstance(config, dict), body + return {key: config[key] for key in _LIMIT_KEYS if key in config} + + +def _mode(stream: bool) -> str: + return "stream" if stream else "sync" + + +_REACHED: Final[tuple[ParameterSet, ...]] = tuple( + pytest.param(provider, endpoint, stream, id=f"{provider}-{endpoint}-{_mode(stream)}") + for provider, endpoint, stream in ( + ("gemini", "chat", False), + ("gemini", "chat", True), + ("vertex_ai", "chat", False), + ("vertex_ai", "chat", True), + ("gemini", "responses", False), + ("vertex_ai", "responses", True), + ) +) + + +@pytest.mark.parametrize(("provider", "endpoint", "stream"), _REACHED) +def test_request_with_no_output_limit_reaches_google_with_the_env_budget( + budgeted: Gateway, provider: Provider, endpoint: Endpoint, stream: bool +) -> None: + with wire_server(upstream(answers(f"plain-{uuid.uuid4().hex}", PLAIN))) as wire, budgeted.scenario() as scenario: + model: Final = register(budgeted, scenario, provider, wire) + response: Final = post(budgeted, endpoint, request_body(endpoint, model, stream, cap=None)) + observed: Final = answered(endpoint, stream, response) + assert (observed.answer, observed.finishes) == (PLAIN.text, (complete(endpoint),)), observed + (only,) = received(wire, provider, stream, (False,)) + assert _limits_sent(only) == {"max_output_tokens": _BUDGET}, only + + +@pytest.mark.parametrize( + ("provider", "stream"), + [ + pytest.param(provider, stream, id=f"{provider}-{_mode(stream)}") + for provider, stream in itertools.product(PROVIDERS, (False, True)) + ], +) +def test_request_with_no_output_limit_is_continued_into_one_joined_answer( + budgeted: Gateway, provider: Provider, stream: bool +) -> None: + respond: Final = upstream( + answers(f"first-{uuid.uuid4().hex}", FIRST_HALF), answers(f"second-{uuid.uuid4().hex}", SECOND_HALF) + ) + with wire_server(respond) as wire, budgeted.scenario() as scenario: + model: Final = register(budgeted, scenario, provider, wire) + response: Final = post(budgeted, "chat", request_body("chat", model, stream, cap=None)) + observed: Final = answered("chat", stream, response) + assert (observed.answer, observed.streamed, observed.finishes) == (_JOINED, _JOINED, ("stop",)), observed + assert observed.tokens == () and HOP_TOKEN not in observed.raw, observed.raw + opening, continuing = received(wire, provider, stream, (False, True)) + assert _limits_sent(opening) == {"max_output_tokens": _BUDGET}, opening + assert continuing == {**opening, "continuationToken": HOP_TOKEN}, (opening, continuing) + + +@pytest.mark.parametrize( + ("endpoint", "limit", "sent"), + [ + pytest.param("chat", {"max_tokens": CAP}, {"max_output_tokens": CAP}, id="chat-max_tokens"), + pytest.param( + "chat", {"max_completion_tokens": CAP}, {"max_output_tokens": CAP}, id="chat-max_completion_tokens" + ), + pytest.param( + "chat", + {"extra_body": {"generationConfig": {"maxOutputTokens": CAP}}}, + {"maxOutputTokens": CAP}, + id="chat-extra_body-generationConfig", + ), + pytest.param( + "responses", {"max_output_tokens": CAP}, {"max_output_tokens": CAP}, id="responses-max_output_tokens" + ), + ], +) +def test_callers_own_output_limit_wins_over_the_env_budget( + budgeted: Gateway, endpoint: Endpoint, limit: Mapping[str, JsonValue], sent: Mapping[str, JsonValue] +) -> None: + with wire_server(upstream(answers(f"plain-{uuid.uuid4().hex}", PLAIN))) as wire, budgeted.scenario() as scenario: + model: Final = register(budgeted, scenario, "gemini", wire) + body: Final = request_body(endpoint, model, False, {**NO_CACHE, **limit}, cap=None) + observed: Final = answered(endpoint, False, post(budgeted, endpoint, body)) + assert observed.answer == PLAIN.text, observed + (only,) = received(wire, "gemini", False, (False,)) + assert _limits_sent(only) == sent, only + + +def _deploy(scenario: Scenario, wire: Wire, backend: str, model_info: Mapping[str, JsonValue] | None) -> str: + return scenario.model( + model=f"gemini/{backend}", + api_base=wire.url, + api_key=API_KEY, + input_cost_per_token=INPUT_COST, + output_cost_per_token=OUTPUT_COST, + model_info=model_info, + ) + + +def test_env_budget_is_capped_per_deployment_and_never_at_a_siblings_limit(budgeted: Gateway) -> None: + lowered: Final = _BUDGET // 4 + with wire_server(upstream(answers(f"plain-{uuid.uuid4().hex}", PLAIN))) as wire, budgeted.scenario() as scenario: + deployments: Final = { + "row-below-budget": _deploy(scenario, wire, _SMALL_BACKEND, None), + "model-info-above-row": _deploy(scenario, wire, _SMALL_BACKEND, {"max_output_tokens": 2 * _SMALL_LIMIT}), + "model-info-below-budget": _deploy(scenario, wire, BACKEND, {"max_output_tokens": lowered}), + "sibling-without-model-info": _deploy(scenario, wire, BACKEND, None), + } + + def limit_sent(model: str) -> dict[str, JsonValue]: + _ = answered("chat", False, post(budgeted, "chat", request_body("chat", model, False, cap=None))) + (only,) = wire.drain() + return _limits_sent(JSON_OBJECT.validate_json(only.body)) + + assert {name: limit_sent(model) for name, model in deployments.items()} == { + "row-below-budget": {"max_output_tokens": _SMALL_LIMIT}, + "model-info-above-row": {"max_output_tokens": _SMALL_LIMIT}, + "model-info-below-budget": {"max_output_tokens": lowered}, + "sibling-without-model-info": {"max_output_tokens": _BUDGET}, + } + + +def test_resume_with_a_continuation_token_gets_no_env_budget(budgeted: Gateway) -> None: + respond: Final = upstream( + answers(f"restarted-{uuid.uuid4().hex}", PLAIN), answers(f"resumed-{uuid.uuid4().hex}", RESUMED) + ) + with wire_server(respond) as wire, budgeted.scenario() as scenario: + model: Final = register(budgeted, scenario, "gemini", wire) + body: Final = request_body("chat", model, False, {**NO_CACHE, "continuation_token": CUT_TOKEN}, cap=None) + observed: Final = answered("chat", False, post(budgeted, "chat", body)) + assert (observed.answer, observed.finishes) == (RESUMED.text, ("stop",)), observed + (only,) = received(wire, "gemini", False, (True,)) + assert only["continuationToken"] == CUT_TOKEN, only + assert [at(only, "generationConfig", key) for key in _LIMIT_KEYS] == [None, None], only diff --git a/tests/integration/providers/test_gemini_long_decode_wire.py b/tests/integration/providers/test_gemini_long_decode_wire.py index 1270b864cb7..5ecb5e2a942 100644 --- a/tests/integration/providers/test_gemini_long_decode_wire.py +++ b/tests/integration/providers/test_gemini_long_decode_wire.py @@ -1163,6 +1163,23 @@ def test_request_for_two_candidates_is_never_continued_and_each_choice_keeps_its assert at(only, "generationConfig", "candidate_count") == 2, only +@pytest.mark.parametrize("provider", PROVIDERS) +def test_request_for_two_candidates_in_camel_case_is_never_continued(gateway: Gateway, provider: Provider) -> None: + with wire_server(_two_drafts) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, provider, wire) + two_candidates: Final = {"extra_body": {"generationConfig": {"candidateCount": 2}}} + body: Final = request_body("chat", model, False, {**NO_CACHE, **two_candidates}) + response: Final = post(gateway, "chat", body) + assert response.status_code == 200, response.text + choices: Final = items_of(JSON_OBJECT.validate_json(response.content).get("choices")) + assert [ + (at(choice, "message", "content"), at(choice, "message", "provider_specific_fields", "continuation_token")) + for choice in choices + ] == list(_DRAFTS), response.text + (only,) = received(wire, provider, False, (False,)) + assert at(only, "generationConfig", "candidateCount") == 2, only + + def _uncapped_cells() -> tuple[ParameterSet, ...]: reached: Final = (("gemini", "chat"), ("vertex_ai", "chat"), ("gemini", "responses")) return tuple( From 94e4ebd63d9384efd933812d4e8a786d1c2f9652 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 02:47:27 -0700 Subject: [PATCH 006/128] test(proxy): pin the max_parallel_requests slot release when a streamed /v1/messages MCP auto-execute ends (#47788) * test(proxy): cover the max_parallel_requests slot a streamed /v1/messages MCP auto-execute keeps A streamed /v1/messages request that auto-executes MCP tools replays its answer as a stream after its model rounds already logged success. The stream path frees the request's max_parallel_requests slot only on a client disconnect, so after message_stop the slot stays taken and the key's next request answers 429 The test streams such a replay to the end and expects the slot to be free. It fails until the fix lands * fix(proxy): free the max_parallel_requests slot when a replayed /v1/messages MCP stream completes A streamed /v1/messages request that auto-executes MCP tools runs its model rounds without streaming and replays the final answer as a stream. The first round claims the request's only success log, and its limiter callback keeps the slot because the tool call reused it afterwards. The stream path released the slot only on a client disconnect, so after message_stop the key stayed at its max_parallel_requests limit and its next request answered 429 Stream cleanup now releases the stashed slot when the stream completed and its logging object already claimed the success log, which only happens when the stream replays a response that was logged before it started. A real stream still releases through its own success callback * test(mcp): free the key's only parallel slot after a streamed auto-execute ends Run a streamed MCP auto-execute on a key with max_parallel_requests 1 on every LLM surface, then send the next request on that key and expect 200. On main both /v1/messages surfaces answer 429 here, while /v1/responses and /v1/chat/completions already pass * test(proxy): keep only the MCP stream regression tests, since #48628 already frees the slot at stream end * test(proxy): require the MCP stream regression test's response to stream and reach its end event --- .../integration/mcp/test_mcp_llm_endpoints.py | 30 ++++++ .../hooks/test_parallel_request_limiter_v3.py | 91 +++++++++++++++++++ 2 files changed, 121 insertions(+) diff --git a/tests/integration/mcp/test_mcp_llm_endpoints.py b/tests/integration/mcp/test_mcp_llm_endpoints.py index c5ffb88506a..490a8f831bc 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/tests/integration/mcp/test_mcp_llm_endpoints.py @@ -702,6 +702,36 @@ def test_streaming_responses_opens_the_gateway_tool_call_empty_and_streams_its_a assert len(rig.upstream_tools()) == 2 +STREAM_END_EVENT: Final[Mapping[Surface, str]] = { + "chat": "data: [DONE]", + "responses": "response.completed", + "messages": "message_stop", + "messages_bridge": "message_stop", +} + + +@pytest.mark.parametrize("surface", SURFACES) +def test_streamed_auto_execute_gives_back_the_keys_only_parallel_slot_once_the_stream_ends( + gateway: Gateway, surface: Surface +) -> None: + with _rig(gateway, surface) as rig: + key: Final = rig.scenario.key(object_permission={"mcp_servers": [rig.server_id]}, max_parallel_requests=1) + streamed: Final = rig.send(key, [AUTO], stream=True) + assert streamed.status_code == 200, streamed.text + assert streamed.headers["content-type"].startswith("text/event-stream"), streamed.headers + assert STREAM_END_EVENT[surface] in streamed.text, streamed.text + assert ANSWER in streamed.text, streamed.text + assert [call["body"]["params"]["name"] for call in _peer_add_calls(rig.peer)] == ["add"] + follow_up: Final = eventually( + lambda: rig.send(key, [], prompt=f"{DIRECT} {rig.alias}"), + lambda response: response.status_code == 200, + seconds=20, + return_last_on_timeout=True, + ) + assert follow_up.status_code == 200, follow_up.text + assert rig.final_text(follow_up.json()) == ANSWER + + def test_streaming_chat_through_a_toolset_gateway_url_serves_a_team_key_without_its_own_grant( gateway: Gateway, ) -> None: diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index cd0d786f841..e056e723a04 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -4332,6 +4332,97 @@ async def test_first_round_logged_after_tool_call_reuse_releases_only_at_request assert await in_flight() == 0 +@pytest.mark.asyncio +async def test_completed_mcp_replay_stream_frees_slot_its_logged_first_round_kept() -> None: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.anthropic.messages.native.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, + ) + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + reused_at: Final = datetime(2026, 1, 5, 12, 0, 0) + api_key: Final = hash_token("sk-mcp-replayed-stream") + local_cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache), + time_provider=lambda: reused_at, + ) + proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = handler + user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) + parallel_key: Final = f"{{api_key:{api_key}}}:max_parallel_requests" + owner_call_id: Final = "mcp-replayed-stream-owner" + messages: Final = [{"role": "user", "content": "add 2 and 3"}] + + async def in_flight() -> int: + return handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) + + async def pre_call(data: dict[str, object], call_type: str) -> None: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=data, + call_type=call_type, + ) + + await pre_call( + {"model": "claude-haiku-5-5", "messages": messages, "stream": True, "litellm_call_id": owner_call_id}, + CallTypes.anthropic_messages.value, + ) + await pre_call( + {"model": "MCP: add", "messages": [{"role": "user", "content": "Tool: add"}], "mcp_server_name": "math"}, + CallTypes.call_mcp_tool.value, + ) + logging_obj: Final = LiteLLMLoggingObj( + model="claude-haiku-5-5", + messages=messages, + stream=True, + call_type=CallTypes.anthropic_messages.value, + start_time=None, + litellm_call_id=owner_call_id, + function_id="mcp-replayed-stream", + ) + assert logging_obj.claim_async_success_log() + await handler.async_log_success_event( + kwargs={"litellm_call_id": owner_call_id, "stream": False}, + response_obj=None, + start_time=None, + end_time=reused_at - timedelta(seconds=1), + ) + assert await in_flight() == 1 + + replay: Final = FakeAnthropicMessagesStreamIterator( + response={ + "id": "msg_replayed", + "type": "message", + "role": "assistant", + "model": "claude-haiku-5-5", + "content": [{"type": "text", "text": "2 + 3 = 5"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 20, "output_tokens": 8}, + } + ) + with _override_litellm_callbacks([]): + streamed: Final = [ + chunk + async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=replay, + user_api_key_dict=user_api_key_dict, + request_data={ + "model": "claude-haiku-5-5", + "litellm_call_id": owner_call_id, + "litellm_logging_obj": logging_obj, + }, + proxy_logging_obj=proxy_logging_obj, + ) + ] + + assert b"event: message_stop" in streamed[-1] + assert await in_flight() == 0 + + @pytest.mark.asyncio async def test_mcp_tool_call_reusing_request_slot_keeps_request_max_parallel_headers() -> None: api_key: Final = hash_token("sk-mcp-reuse-headers") From 4b0aa0f44679e3aa1dd58f8208ff5ccfcbdc1c27 Mon Sep 17 00:00:00 2001 From: tessellark <349031000+tessellark@users.noreply.github.com> Date: Wed, 28 Oct 2026 03:01:40 -0700 Subject: [PATCH 007/128] fix(cost-map): drop the gpt-5-nano priority input price OpenAI does not sell (#48763) gpt-5-nano and gpt-5-nano-2025-08-07 carried gpt-5's priority input price, 50x the nano standard rate. OpenAI sells no priority tier for nano and answers a nano priority request with service_tier default, so proxy spend on real calls was already billed at the standard rate. The wrong price showed in /model/info, in the tier pricing a custom-priced nano deployment inherited, and in SDK cost estimates like litellm.cost_per_token with service_tier priority. The key is removed, and a new test keeps every priority rate between 1x and 5x of its standard rate in both cost map files --- ...odel_prices_and_context_window_backup.json | 2 - model_prices_and_context_window.json | 2 - .../test_model_prices_and_context_window.py | 46 ++++++++++++++++++- 3 files changed, 45 insertions(+), 5 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 5774e699ad7..d053300e299 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -37504,7 +37504,6 @@ "input_cost_per_token": 5e-08, "input_cost_per_token_batches": 2.5e-08, "input_cost_per_token_flex": 2.5e-08, - "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", "max_input_tokens": 272000, "max_output_tokens": 128000, @@ -37552,7 +37551,6 @@ "deprecation_date": "2026-12-11", "input_cost_per_token": 5e-08, "input_cost_per_token_batches": 2.5e-08, - "input_cost_per_token_priority": 2.5e-06, "input_cost_per_token_flex": 2.5e-08, "litellm_provider": "openai", "max_input_tokens": 272000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 5774e699ad7..d053300e299 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -37504,7 +37504,6 @@ "input_cost_per_token": 5e-08, "input_cost_per_token_batches": 2.5e-08, "input_cost_per_token_flex": 2.5e-08, - "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", "max_input_tokens": 272000, "max_output_tokens": 128000, @@ -37552,7 +37551,6 @@ "deprecation_date": "2026-12-11", "input_cost_per_token": 5e-08, "input_cost_per_token_batches": 2.5e-08, - "input_cost_per_token_priority": 2.5e-06, "input_cost_per_token_flex": 2.5e-08, "litellm_provider": "openai", "max_input_tokens": 272000, diff --git a/tests/unit/test_model_prices_and_context_window.py b/tests/unit/test_model_prices_and_context_window.py index 1f84be7710f..08a39b94c12 100644 --- a/tests/unit/test_model_prices_and_context_window.py +++ b/tests/unit/test_model_prices_and_context_window.py @@ -2,7 +2,7 @@ from __future__ import annotations import json import re -from collections.abc import Mapping +from collections.abc import Iterator, Mapping from pathlib import Path from types import MappingProxyType from typing import Final @@ -224,6 +224,50 @@ def test_dated_variants_carry_base_alias_service_tier_pricing(prices: dict): ) +PRIORITY_SUFFIX: Final = "_priority" +MIN_PRIORITY_MARKUP: Final = 1.0 +MAX_PRIORITY_MARKUP: Final = 5.0 + + +def _rate(value: JsonValue | None) -> float | None: + return float(value) if isinstance(value, (int, float)) and not isinstance(value, bool) else None + + +def _priority_markup(entry: Mapping[str, JsonValue], key: str) -> float | None: + if not key.endswith(PRIORITY_SUFFIX): + return None + priority: Final = _rate(entry[key]) + standard: Final = _rate(entry.get(key.removesuffix(PRIORITY_SUFFIX))) + if priority is None or standard is None or standard <= 0: + return None + return priority / standard + + +def priority_markups(prices: Mapping[str, JsonValue]) -> Iterator[tuple[str, str, float]]: + for name, entry in prices.items(): + if not isinstance(entry, dict): + continue + for key in entry: + if (markup := _priority_markup(entry, key)) is not None: + yield name, key, markup + + +@pytest.mark.parametrize("path", (PRICES_PATH, BACKUP_PRICES_PATH), ids=("main", "backup")) +def test_priority_rates_stay_within_a_premium_of_their_standard_rate(path: Path) -> None: + prices: Final = TypeAdapter(dict[str, JsonValue]).validate_json(path.read_bytes()) + markups: Final = tuple(priority_markups(prices)) + assert markups + overpriced: Final = tuple( + f"{name}.{key} is {markup:.1f}x its standard rate" + for name, key, markup in markups + if not MIN_PRIORITY_MARKUP <= markup <= MAX_PRIORITY_MARKUP + ) + assert overpriced == (), ( + "a priority rate is a premium on the standard rate, so one below it or far above it was likely copied from " + "another model; drop the key when the provider sells no priority tier for the model:\n" + "\n".join(overpriced) + ) + + # Azure Foundry GPT-6.1 Sol announcement, 2026-09-29: https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/introducing-gpt-6-1-sol-in-microsoft-foundry-advanced-intelligence-optimized-for/4560811 AZURE_GPT_6_1_SOL_REGIONAL_PREMIUM: Final = MappingProxyType({"us": 1.1, "eu": 1.2, "apac": 1.2}) From d84a2d4452d37420876573234e810bf4588882b0 Mon Sep 17 00:00:00 2001 From: norkelbyte <348963847+norkelbyte@users.noreply.github.com> Date: Wed, 28 Oct 2026 03:04:11 -0700 Subject: [PATCH 008/128] test(spend): audit #47815 (#48774) Adds an integration cell proving a failed /v1/responses request counts its input from the newest compaction item on its failure spend row Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../spend/test_failed_dispatch_tokens.py | 48 +++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/tests/integration/spend/test_failed_dispatch_tokens.py b/tests/integration/spend/test_failed_dispatch_tokens.py index b92606ebc88..c9875be562c 100644 --- a/tests/integration/spend/test_failed_dispatch_tokens.py +++ b/tests/integration/spend/test_failed_dispatch_tokens.py @@ -167,3 +167,51 @@ def test_provider_500_failure_rows_count_deepseek_v4_prompt_tokens_with_deepseek short: Final = _failure_row_prompt_tokens(gateway, model, key, SHORT_TEXT) assert wire.drain() assert_counted_with_deepseek(chinese - short, chinese_minus_short) + + +COMPACTION_INSTRUCTIONS: Final[JsonValue] = {"role": "system", "content": "answer briefly"} +COMPACTED_HISTORY: Final[tuple[JsonValue, ...]] = ( + {"role": "user", "content": "a long turn " * 200}, + {"role": "assistant", "content": "a long reply " * 200}, +) +COMPACTION_ITEM: Final[JsonValue] = { + "type": "compaction", + "encrypted_content": json.dumps({"type": "compaction", "content": "the cat is named Biscuit"}), +} +TURN_AFTER_COMPACTION: Final[JsonValue] = {"role": "user", "content": "what is the cat's name?"} + + +def _failed_responses_prompt_tokens(gateway: Gateway, model: str, key: str, items: list[JsonValue]) -> int: + failed: Final = gateway.request("POST", "/v1/responses", {"model": model, "input": items}, key=key) + assert failed.status_code == 500 and "synthetic provider outage" in failed.text, failed.text + rows: Final = eventually( + lambda: read_rows( + 'SELECT status, prompt_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (failed.headers["x-litellm-call-id"],), + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert rows[0]["status"] == "failure", rows + prompt_tokens: Final = rows[0]["prompt_tokens"] + assert isinstance(prompt_tokens, int) and prompt_tokens > 0, rows + return prompt_tokens + + +def test_a_failed_responses_request_counts_its_input_from_the_newest_compaction_item(gateway: Gateway) -> None: + with wire_server(_provider_500) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(api_base=wire.url + "/v1", num_retries=0) + key: Final = scenario.key(models=[model]) + resent: Final = _failed_responses_prompt_tokens( + gateway, model, key, [COMPACTION_INSTRUCTIONS, *COMPACTED_HISTORY, COMPACTION_ITEM, TURN_AFTER_COMPACTION] + ) + compacted: Final = _failed_responses_prompt_tokens( + gateway, model, key, [COMPACTION_INSTRUCTIONS, COMPACTION_ITEM, TURN_AFTER_COMPACTION] + ) + uncompacted: Final = _failed_responses_prompt_tokens( + gateway, model, key, [COMPACTION_INSTRUCTIONS, *COMPACTED_HISTORY, TURN_AFTER_COMPACTION] + ) + assert len(wire.drain()) == 3 + + assert resent == compacted + assert resent < uncompacted From f6f49d16cc2422d2a83c93c074e0ac4d6a889cc5 Mon Sep 17 00:00:00 2001 From: islomjon-aiand Date: Wed, 28 Oct 2026 19:15:56 +0900 Subject: [PATCH 009/128] feat: add ai& (aiand) provider (#38958) * feat: add aiand provider * feat: expose aiand provider in admin UI * feat(aiand): add ai& model pricing to the cost map * feat(aiand): list ai& models for aiand/* wildcards and mark /v1/messages supported * feat(aiand): send Anthropic /v1/messages to ai& natively * fix(aiand): drop the retired max_tokens field and regenerate endpoint support --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/__init__.py | 5 + litellm/constants.py | 2 + litellm/llms/openai_like/providers.json | 6 + ...odel_prices_and_context_window_backup.json | 195 ++++++++++++++++++ .../provider_endpoints_support_backup.json | 26 +++ .../provider_create_fields.json | 28 +++ litellm/types/utils.py | 1 + model_prices_and_context_window.json | 195 ++++++++++++++++++ provider_endpoints_support.json | 26 +++ .../llms/openai_like/test_aiand_provider.py | 145 +++++++++++++ 10 files changed, 629 insertions(+) create mode 100644 tests/unit/llms/openai_like/test_aiand_provider.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 41e3fa34323..2cff21bf29a 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -710,6 +710,7 @@ transcribe_models: Set = set() gigachat_models: Set = set() llamagate_models: Set = set() coralbricks_models: Set[str] = set() +aiand_models: Set[str] = set() reducto_models: Set = set() bedrock_mantle_models: Set = set() @@ -1022,6 +1023,8 @@ def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, obje llamagate_models.add(key) elif value.get("litellm_provider") == "coralbricks": coralbricks_models.add(key) + elif value.get("litellm_provider") == "aiand": + aiand_models.add(key) elif value.get("litellm_provider") == "reducto": reducto_models.add(key) elif value.get("litellm_provider") == "bedrock_mantle": @@ -1159,6 +1162,7 @@ model_list = list( | reducto_models | bedrock_mantle_models | coralbricks_models + | aiand_models | set(clarifai_models) ) @@ -1283,6 +1287,7 @@ def _build_models_by_provider() -> dict: "gigachat": gigachat_models, "llamagate": llamagate_models, "coralbricks": coralbricks_models, + "aiand": aiand_models, "reducto": reducto_models, "bedrock_mantle": bedrock_mantle_models, } diff --git a/litellm/constants.py b/litellm/constants.py index 8cc51d0ab5e..0d6dfb5c0d4 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1035,6 +1035,7 @@ openai_compatible_endpoints: Final[list[str]] = [ "https://api-inference.modelscope.cn/v1", "https://api.moonshot.ai/v1", "https://api.publicai.co/v1", + "https://api.aiand.com/v1", "https://api.synthetic.new/openai/v1", "https://serverless.tensormesh.ai/v1", "https://api.stima.tech/v1", @@ -1100,6 +1101,7 @@ openai_compatible_providers: Final[list[str]] = [ "novita", "meta_llama", "publicai", # PublicAI - JSON-configured provider + "aiand", # ai& - JSON-configured provider "synthetic", # Synthetic - JSON-configured provider "tensormesh", # Tensormesh - JSON-configured provider "apertis", # Apertis - JSON-configured provider diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index 115d914d3ac..0034ed2b9d4 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -107,6 +107,12 @@ "api_key_env": "AIHUBMIX_API_KEY", "api_base_env": "AIHUBMIX_API_BASE" }, + "aiand": { + "base_url": "https://api.aiand.com/v1", + "api_key_env": "AIAND_API_KEY", + "api_base_env": "AIAND_API_BASE", + "supported_endpoints": ["/v1/chat/completions", "/v1/responses", "/v1/messages"] + }, "crusoe": { "base_url": "https://api.inference.crusoecloud.com/v1", "api_key_env": "CRUSOE_API_KEY", diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d053300e299..8207c3122f9 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -69825,6 +69825,201 @@ "supports_reasoning": true, "supports_vision": true }, + "aiand/deepseek-ai/deepseek-v4-flash": { + "cache_read_input_token_cost": 8e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "aiand", + "max_input_tokens": 1048576, + "max_output_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 2.5e-07, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "aiand/deepseek-ai/deepseek-v4-pro": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "aiand", + "max_input_tokens": 1048576, + "max_output_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "aiand/deepseek-ai/deepseek-v4.1-flash": { + "cache_read_input_token_cost": 2e-08, + "input_cost_per_token": 3e-07, + "litellm_provider": "aiand", + "max_input_tokens": 1048576, + "max_output_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "aiand/google/gemma-4-31b-it": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "aiand", + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "aiand/moonshotai/kimi-k2.7-code": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 7.5e-07, + "litellm_provider": "aiand", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.5e-06, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "aiand/moonshotai/kimi-k3": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "aiand", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.25e-05, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "aiand/motif-technologies/motif-3": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 5e-07, + "litellm_provider": "aiand", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2e-06, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "aiand/openai/gpt-oss-120b": { + "cache_read_input_token_cost": 8e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "aiand", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "aiand/qwen/qwen3.6-27b": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 3.2e-07, + "litellm_provider": "aiand", + "max_input_tokens": 262144, + "max_output_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 3.2e-06, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "aiand/qwen/qwen3.8-27b": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 4e-07, + "litellm_provider": "aiand", + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "aiand/zai-org/glm-5.2": { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "aiand", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "aiand/zai-org/glm-5.3": { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "aiand", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "aiand/zai-org/glm-5.3-flash": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "aiand", + "max_input_tokens": 1048550, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "bespoke/nimble-latest": { "input_cost_per_token": 0.0, "litellm_provider": "bespoke", diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index 9ba6ed342ea..1ab03fb3fdc 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -189,6 +189,32 @@ "generateContent": false } }, + "aiand": { + "display_name": "ai& (`aiand`)", + "url": "https://docs.aiand.com/api/chat-completions/", + "endpoints": { + "chat_completions": true, + "responses": true, + "text_completion": true, + "embeddings": false, + "image_generations": false, + "image_edits": false, + "audio_transcriptions": false, + "audio_speech": false, + "rerank": false, + "ocr": false, + "search": false, + "moderations": false, + "realtime": false, + "video_generations": false, + "batches": false, + "messages": true, + "interactions": true, + "systemone": false, + "a2a": true, + "generateContent": true + } + }, "amazon_nova": { "display_name": "Amazon Nova (`amazon_nova`)", "url": "https://docs.litellm.ai/docs/providers/amazon_nova", diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index d96f2976508..0dd96c856f8 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -1,4 +1,32 @@ [ + { + "provider": "AIAND", + "provider_display_name": "ai&", + "litellm_provider": "aiand", + "credential_fields": [ + { + "key": "api_base", + "label": "API Base", + "placeholder": "https://api.aiand.com/v1", + "tooltip": null, + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "api_key", + "label": "API Key", + "placeholder": null, + "tooltip": null, + "required": true, + "field_type": "password", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "aiand/openai/gpt-oss-120b" + }, { "provider": "SCALEDOWN", "provider_display_name": "ScaleDown", diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 2869f9b2577..055c18f2e36 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4573,6 +4573,7 @@ class LlmProviders(str, Enum): POE = "poe" CHUTES = "chutes" NEOSANTARA = "neosantara" + AIAND = "aiand" PARASAIL = "parasail" XIAOMI_MIMO = "xiaomi_mimo" TENSORMESH = "tensormesh" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d053300e299..8207c3122f9 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -69825,6 +69825,201 @@ "supports_reasoning": true, "supports_vision": true }, + "aiand/deepseek-ai/deepseek-v4-flash": { + "cache_read_input_token_cost": 8e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "aiand", + "max_input_tokens": 1048576, + "max_output_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 2.5e-07, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "aiand/deepseek-ai/deepseek-v4-pro": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "aiand", + "max_input_tokens": 1048576, + "max_output_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "aiand/deepseek-ai/deepseek-v4.1-flash": { + "cache_read_input_token_cost": 2e-08, + "input_cost_per_token": 3e-07, + "litellm_provider": "aiand", + "max_input_tokens": 1048576, + "max_output_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "aiand/google/gemma-4-31b-it": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "aiand", + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "aiand/moonshotai/kimi-k2.7-code": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 7.5e-07, + "litellm_provider": "aiand", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.5e-06, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "aiand/moonshotai/kimi-k3": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "aiand", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.25e-05, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "aiand/motif-technologies/motif-3": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 5e-07, + "litellm_provider": "aiand", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2e-06, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "aiand/openai/gpt-oss-120b": { + "cache_read_input_token_cost": 8e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "aiand", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "aiand/qwen/qwen3.6-27b": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 3.2e-07, + "litellm_provider": "aiand", + "max_input_tokens": 262144, + "max_output_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 3.2e-06, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "aiand/qwen/qwen3.8-27b": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 4e-07, + "litellm_provider": "aiand", + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "aiand/zai-org/glm-5.2": { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "aiand", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "aiand/zai-org/glm-5.3": { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "aiand", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "aiand/zai-org/glm-5.3-flash": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "aiand", + "max_input_tokens": 1048550, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "bespoke/nimble-latest": { "input_cost_per_token": 0.0, "litellm_provider": "bespoke", diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 9ba6ed342ea..1ab03fb3fdc 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -189,6 +189,32 @@ "generateContent": false } }, + "aiand": { + "display_name": "ai& (`aiand`)", + "url": "https://docs.aiand.com/api/chat-completions/", + "endpoints": { + "chat_completions": true, + "responses": true, + "text_completion": true, + "embeddings": false, + "image_generations": false, + "image_edits": false, + "audio_transcriptions": false, + "audio_speech": false, + "rerank": false, + "ocr": false, + "search": false, + "moderations": false, + "realtime": false, + "video_generations": false, + "batches": false, + "messages": true, + "interactions": true, + "systemone": false, + "a2a": true, + "generateContent": true + } + }, "amazon_nova": { "display_name": "Amazon Nova (`amazon_nova`)", "url": "https://docs.litellm.ai/docs/providers/amazon_nova", diff --git a/tests/unit/llms/openai_like/test_aiand_provider.py b/tests/unit/llms/openai_like/test_aiand_provider.py new file mode 100644 index 00000000000..804f3a3b6eb --- /dev/null +++ b/tests/unit/llms/openai_like/test_aiand_provider.py @@ -0,0 +1,145 @@ +import json +from typing import Final + +import httpx +import pytest +import respx + +import litellm +from litellm.caching.llm_caching_handler import LLMClientCache +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.llms.openai_like.json_loader import JSONProviderRegistry +from litellm.proxy.auth.model_checks import get_known_models_from_wildcard +from litellm.utils import ProviderConfigManager + +AIAND_API_BASE = "https://api.aiand.com/v1" + + +def test_aiand_chat_provider(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("AIAND_API_KEY", "test-key") + + provider_config = JSONProviderRegistry.get("aiand") + assert provider_config is not None + assert litellm.LlmProviders.AIAND.value == "aiand" + assert "aiand" in litellm.openai_compatible_providers + assert provider_config.base_url == AIAND_API_BASE + assert provider_config.api_key_env == "AIAND_API_KEY" + assert provider_config.api_base_env == "AIAND_API_BASE" + + model, provider, api_key, api_base = get_llm_provider("aiand/openai/gpt-oss-120b") + assert model == "openai/gpt-oss-120b" + assert provider == "aiand" + assert api_key == "test-key" + assert api_base == AIAND_API_BASE + + +def test_aiand_completion_request(respx_mock: respx.MockRouter): + messages = [{"role": "user", "content": "Hello"}] + route = respx_mock.post(f"{AIAND_API_BASE}/chat/completions").mock( + return_value=httpx.Response( + 200, + json={ + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1, + "model": "openai/gpt-oss-120b", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hello"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + }, + ) + ) + + litellm.completion( + model="aiand/openai/gpt-oss-120b", + api_key="test-key", + messages=messages, + ) + + assert route.called + request = route.calls.last.request + assert str(request.url) == f"{AIAND_API_BASE}/chat/completions" + assert request.headers["Authorization"] == "Bearer test-key" + request_body = json.loads(request.read()) + assert request_body["model"] == "openai/gpt-oss-120b" + assert request_body["messages"] == messages + + +def test_aiand_responses_provider(): + config = ProviderConfigManager.get_provider_responses_api_config( + provider="aiand", + model="aiand/openai/gpt-oss-120b", + ) + + assert config is not None + assert config.custom_llm_provider == "aiand" + assert config.get_complete_url(api_base=None, litellm_params={}) == f"{AIAND_API_BASE}/responses" + + +def _aiand_catalog() -> frozenset[str]: + return frozenset(name for name, info in litellm.model_cost.items() if info.get("litellm_provider") == "aiand") + + +def test_aiand_wildcard_lists_its_priced_models(): + catalog: Final = _aiand_catalog() + + assert catalog + assert frozenset(get_known_models_from_wildcard("aiand/*")) == catalog + + +def test_aiand_calls_are_billed_from_their_cost_rows(): + catalog: Final = _aiand_catalog() + assert catalog + + for model in sorted(catalog): + row: Final = litellm.model_cost[model] + prompt_cost, completion_cost = litellm.cost_per_token(model=model, prompt_tokens=1000, completion_tokens=500) + + assert prompt_cost == pytest.approx(1000 * row["input_cost_per_token"]) + assert completion_cost == pytest.approx(500 * row["output_cost_per_token"]) + assert prompt_cost > 0 and completion_cost > 0 + + +@pytest.mark.asyncio +async def test_aiand_anthropic_messages_reach_ai_and_with_the_thinking_budget(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + thinking: Final = {"type": "enabled", "budget_tokens": 1024} + with respx.mock() as upstream: + route: Final = upstream.post(f"{AIAND_API_BASE}/messages").respond( + 200, + json={ + "id": "msg_aiand", + "type": "message", + "role": "assistant", + "model": "zai-org/glm-5.3", + "content": [{"type": "thinking", "thinking": "A greeting", "signature": ""}, {"type": "text", "text": "Hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 4, "output_tokens": 3}, + }, + ) + response: Final = await litellm.anthropic.messages.acreate( + model="aiand/zai-org/glm-5.3", + messages=[{"role": "user", "content": "Say hello"}], + max_tokens=2048, + thinking=thinking, + api_key="aiand-test-key", + ) + + request: Final = route.calls.last.request + body: Final = json.loads(request.content) + assert route.call_count == 1 + assert body["model"] == "zai-org/glm-5.3" + assert body["thinking"] == thinking + assert "reasoning_effort" not in body + assert [block["type"] for block in response["content"]] == ["thinking", "text"] From ff65785717e7b682d5b174bfca3448c8511c5057 Mon Sep 17 00:00:00 2001 From: Max Azatian <66359507+HardMax71@users.noreply.github.com> Date: Wed, 28 Oct 2026 11:19:26 +0100 Subject: [PATCH 010/128] fix(router)!: count an explicit allowed_fails of 3 and keep fallback deployments out of cooldown (#46183) * fix(router): count an explicit allowed_fails equal to the default as set Router(allowed_fails=3) stored the same value as an unset Router, so _is_allowed_fails_set_on_router treated it as "not set" and kept the rate-based cooldown logic. Keep the passed value in explicit_allowed_fails, like _explicit_timeout, and count it as set. * chore(router): drop the comment on explicit_allowed_fails * fix(router): keep the fallback deployment out of cooldown when no deployment can be picked When every deployment of the requested model is in cooldown, the router raises "No deployments available" and schedules the failure callbacks on the request's logging object. By the time they run, the fallback has already written its own deployment into that object, so the 429 counted as a failure of the healthy fallback deployment and cooled it down too, leaving the model group with nothing to serve. Tag the selection error and skip cooldown accounting for it, the same way advisor orchestration failures are skipped. * test(router): annotate the new cooldown test variables with Final * test(router): pin that a saved allowed_fails of 3 stays unset and type the fallback helper * test(router): use a fixed start time in the fallback helper * test(router): cover a fallback group staying up after the primary's selection fails --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/router.py | 10 +- litellm/router_utils/cooldown_handlers.py | 20 ++- .../router_utils/test_cooldown_handlers.py | 141 +++++++++++++++++- 3 files changed, 164 insertions(+), 7 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index c8153b596d6..18c5c67abce 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -238,11 +238,12 @@ from litellm.router_utils.cooldown_handlers import ( async_get_cooldown_deployments_with_debug_info, get_cooldown_deployment_info, get_cooldown_deployments, - is_advisor_orchestration_failure, is_background_response_poll_not_found, is_caller_scoped_auth_failure, is_caller_timeout_408, + is_failure_outside_selected_deployment, is_video_lookup_not_found, + mark_deployment_selection_failure, set_cooldown_deployments, ) from litellm.router_utils.deployment_pricing import deployment_pricing_models @@ -1519,6 +1520,7 @@ class Router: self.allowed_fails = allowed_fails else: self.allowed_fails = litellm.allowed_fails + self.explicit_allowed_fails = allowed_fails self.cooldown_time = cooldown_time or DEFAULT_COOLDOWN_TIME_SECONDS self.cooldown_cache = CooldownCache(cache=self.cache, default_cooldown_time=self.cooldown_time) self.disable_cooldowns = disable_cooldowns @@ -9579,10 +9581,10 @@ class Router: try: exception: Final = kwargs.get("exception", None) - if is_advisor_orchestration_failure(exception): + if is_failure_outside_selected_deployment(exception): verbose_router_logger.debug( "Router: Exiting 'deployment_callback_on_failure' without cooldown. " - "Failure originated from advisor orchestration, not the selected deployment." + "Failure came from advisor orchestration or deployment selection, not the selected deployment." ) return False @@ -15473,6 +15475,7 @@ class Router: phase_event(_DEPLOYMENT_SELECTED_EVENT, pick_attributes) return deployment except Exception as e: + mark_deployment_selection_failure(e) traceback_exception: Final = traceback.format_exc() # if router rejects call -> log to langfuse/otel/etc. if request_kwargs is not None: @@ -15641,6 +15644,7 @@ class Router: return deployment except Exception as e: + mark_deployment_selection_failure(e) traceback_exception: Final = traceback.format_exc() if request_kwargs is not None: logging_obj: Final = request_kwargs.get("litellm_logging_obj", None) diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index 330500789a4..e1baf66b201 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -57,6 +57,7 @@ else: Span = Any _ADVISOR_ORCHESTRATION_FAILURE_ATTR: Final = "_litellm_advisor_orchestration_failure" +_DEPLOYMENT_SELECTION_FAILURE_ATTR: Final = "_litellm_deployment_selection_failure" _CREDENTIAL_VALUES_ADAPTER: Final = TypeAdapter(Mapping[str, object]) @@ -79,6 +80,16 @@ def is_advisor_orchestration_failure(exception: BaseException | None) -> bool: return bool(getattr(exception, _ADVISOR_ORCHESTRATION_FAILURE_ATTR, False)) +def mark_deployment_selection_failure(exception: BaseException) -> None: + setattr(exception, _DEPLOYMENT_SELECTION_FAILURE_ATTR, True) + + +def is_failure_outside_selected_deployment(exception: BaseException | None) -> bool: + return is_advisor_orchestration_failure(exception) or bool( + getattr(exception, _DEPLOYMENT_SELECTION_FAILURE_ATTR, False) + ) + + _BACKGROUND_RESPONSE_POLL_CALL_ORIGINS: Final[frozenset[InternalCallOrigin]] = frozenset( {BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN} ) @@ -781,14 +792,17 @@ def _is_allowed_fails_set_on_router( litellm_router_instance: LitellmRouter, ) -> bool: """ - Check if Router.allowed_fails is set or is Non-default Value + Check if Router.allowed_fails was passed to the Router or is a Non-default Value Returns: - - True if Router.allowed_fails is set or is Non-default Value - - False if Router.allowed_fails is None or is Default Value + - True if allowed_fails was passed to the Router, even when it equals the default + - True if Router.allowed_fails is a Non-default Value + - False if Router.allowed_fails is None, or is the Default Value and was not passed """ if litellm_router_instance.allowed_fails is None: return False + if litellm_router_instance.explicit_allowed_fails is not None: + return True if litellm_router_instance.allowed_fails != litellm.allowed_fails: return True return False diff --git a/tests/unit/router_utils/test_cooldown_handlers.py b/tests/unit/router_utils/test_cooldown_handlers.py index bab6d257579..4299b61c703 100644 --- a/tests/unit/router_utils/test_cooldown_handlers.py +++ b/tests/unit/router_utils/test_cooldown_handlers.py @@ -1,10 +1,12 @@ import asyncio import copy +import datetime import importlib import itertools import random import time -from collections.abc import Mapping +import uuid +from collections.abc import Mapping, Sequence from typing import Final from unittest.mock import MagicMock, patch @@ -682,6 +684,45 @@ class TestDeploymentLevelAllowedFails: ) +def _gpt_4_group_with_a_sibling() -> Sequence[Mapping[str, object]]: + return ( + {"model_name": "gpt-4", "litellm_params": {"model": "openai/gpt-4"}, "model_info": {"id": "primary"}}, + {"model_name": "gpt-4", "litellm_params": {"model": "openai/gpt-4"}, "model_info": {"id": "secondary"}}, + ) + + +def _primary_cooldown_decisions_on_repeated_503s(router: Router) -> tuple[bool, ...]: + exception: Final = litellm.ServiceUnavailableError("down", "openai", "gpt-4") + return tuple( + _should_cooldown_deployment( + litellm_router_instance=router, + deployment="primary", + exception_status=503, + original_exception=exception, + ) + for _ in range(litellm.allowed_fails + 1) + ) + + +@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") +class TestRouterLevelAllowedFailsEqualToDefault: + def test_allowed_fails_equal_to_the_default_cools_down_after_that_many_fails(self): + router: Final = _make_router(model_list=list(_gpt_4_group_with_a_sibling()), allowed_fails=litellm.allowed_fails) + + assert _primary_cooldown_decisions_on_repeated_503s(router) == (False,) * litellm.allowed_fails + (True,) + + def test_unset_allowed_fails_keeps_the_rate_based_default(self): + router: Final = _make_router(model_list=list(_gpt_4_group_with_a_sibling())) + + assert _primary_cooldown_decisions_on_repeated_503s(router) == (False,) * (litellm.allowed_fails + 1) + + def test_default_allowed_fails_saved_through_update_settings_keeps_the_rate_based_default(self): + router: Final = _make_router(model_list=list(_gpt_4_group_with_a_sibling())) + router.update_settings(allowed_fails=litellm.allowed_fails) + + assert _primary_cooldown_decisions_on_repeated_503s(router) == (False,) * (litellm.allowed_fails + 1) + + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestDeploymentLevelAllowedFailsPolicyByExceptionType: def test_rate_limit_error_triggers_cooldown_with_zero_threshold(self): @@ -2931,6 +2972,104 @@ async def test_bad_request_error_does_not_cool_down_a_deployment( assert tuple(cooldown_deployments) == () +async def _complete_like_the_proxy(router: Router, model: str) -> str: + logging_obj, request = litellm.utils.function_setup( + original_function="acompletion", + rules_obj=litellm.utils.Rules(), + start_time=datetime.datetime(2026, 1, 1), + model=model, + messages=[{"role": "user", "content": "hi"}], + litellm_call_id=str(uuid.uuid4()), + ) + response: Final = await router.acompletion(**request, litellm_logging_obj=logging_obj) + model_id: Final = response._hidden_params["model_id"] + assert isinstance(model_id, str) + return model_id + + +_CHAT_COMPLETION_OK: Final = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} + + +def _openai_deployment(model_name: str, host: str, deployment_id: str) -> Mapping[str, object]: + return { + "model_name": model_name, + "litellm_params": {"model": "openai/gpt-4o", "api_key": "k", "api_base": f"https://{host}/v1"}, + "model_info": {"id": deployment_id}, + } + + +@pytest.mark.asyncio +async def test_fallback_deployment_stays_up_while_the_primary_is_in_cooldown( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://primary.test/v1/chat/completions").mock( + return_value=httpx.Response(503, json={"error": {"message": "down", "type": "server_error"}}) + ) + backup_route: Final = respx_mock.post("https://backup.test/v1/chat/completions").mock( + return_value=httpx.Response(200, json=_CHAT_COMPLETION_OK) + ) + router: Final = litellm.Router( + model_list=[ + _openai_deployment("primary", "primary.test", "primary-dep"), + _openai_deployment("backup", "backup.test", "backup-dep"), + ], + fallbacks=[{"primary": ["backup"]}], + allowed_fails=litellm.allowed_fails, + num_retries=0, + ) + + served_by: Final = [await _complete_like_the_proxy(router, "primary") for _ in range(20)] + cooldown_deployments: Final = await async_get_cooldown_deployments( + litellm_router_instance=router, + parent_otel_span=None, + ) + + assert served_by == ["backup-dep"] * 20 + assert len(backup_route.calls) == 20 + assert tuple(cooldown_deployments) == ("primary-dep",) + + +@pytest.mark.asyncio +async def test_fallback_group_stays_up_after_a_bad_key_cools_the_primary_under_the_rate_based_default( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://primary.test/v1/chat/completions").mock( + return_value=httpx.Response(401, json={"error": {"message": "bad key", "type": "invalid_request_error"}}) + ) + respx_mock.post("https://backup.test/v1/chat/completions").mock( + return_value=httpx.Response(200, json=_CHAT_COMPLETION_OK) + ) + router: Final = litellm.Router( + model_list=[ + _openai_deployment("primary", "primary.test", "primary-dep"), + _openai_deployment("backup", "backup.test", "backup-1"), + _openai_deployment("backup", "backup.test", "backup-2"), + ], + fallbacks=[{"primary": ["backup"]}], + num_retries=0, + ) + + served_by: Final = [await _complete_like_the_proxy(router, "primary") for _ in range(12)] + cooldown_deployments: Final = await async_get_cooldown_deployments( + litellm_router_instance=router, + parent_otel_span=None, + ) + + assert set(served_by) <= {"backup-1", "backup-2"} + assert tuple(cooldown_deployments) == ("primary-dep",) + + TIMED_OUT_DEPLOYMENT: Final = { "model_name": "test-model", "litellm_params": {"model": "openai/test-model", "api_key": "test-key"}, From 5dc18fe909a840c527b80592ce1b93a7354de929 Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Wed, 28 Oct 2026 03:23:34 -0700 Subject: [PATCH 011/128] test(proxy): audit #47418 (#48766) * test(proxy): audit #47418 * test(proxy): hold the stream open past its first chunk and require cost headers in the audit cells --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../test_compressed_request_bodies.py | 392 ++++++++++++++++++ 1 file changed, 392 insertions(+) create mode 100644 tests/integration/compatibility/test_compressed_request_bodies.py diff --git a/tests/integration/compatibility/test_compressed_request_bodies.py b/tests/integration/compatibility/test_compressed_request_bodies.py new file mode 100644 index 00000000000..ef2490add40 --- /dev/null +++ b/tests/integration/compatibility/test_compressed_request_bodies.py @@ -0,0 +1,392 @@ +from __future__ import annotations + +import gzip +import json +import os +import secrets +import threading +import uuid +from collections.abc import Callable, Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from functools import partial +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.budget import case_timeout_seconds +from integration._support.client import ( + JSON_OBJECT, + Gateway, + eventually, + gateway_from_environment, + list_value, + object_value, + string_value, +) +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +MIB: Final = 1024 * 1024 +CHAT: Final = "/v1/chat/completions" +CHAT_MODEL: Final = "compressed-chat" +OTHER_MODEL: Final = "compressed-chat-other" +ANTHROPIC_MODEL: Final = "claude-haiku-5-5" +GATED_PREFIX: Final = "gated-" +ORIGIN: Final = "http://browser.example" +MIDDLEWARE_HEADERS: Final = ( + "access-control-allow-origin", + "x-frame-options", + "content-security-policy", + "x-content-type-options", +) +ZSTD_MAGIC: Final = b"\x28\xb5\x2f\xfd" +ZSTD_SINGLE_SEGMENT_FOUR_BYTE_SIZE: Final = b"\xa0" +ZSTD_MAX_BLOCK: Final = 128 * 1024 +NO_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) + +pytestmark: Final = [ + pytest.mark.timeout(2 * graceful_stop_seconds() + 120), + pytest.mark.xdist_group("compressed-request-bodies"), +] + + +def zstd_frame(data: bytes) -> bytes: + blocks: Final = tuple(data[start : start + ZSTD_MAX_BLOCK] for start in range(0, len(data), ZSTD_MAX_BLOCK)) + framed: Final = blocks or (b"",) + return ( + ZSTD_MAGIC + + ZSTD_SINGLE_SEGMENT_FOUR_BYTE_SIZE + + len(data).to_bytes(4, "little") + + b"".join( + (len(block) << 3 | (index == len(framed) - 1)).to_bytes(3, "little") + block + for index, block in enumerate(framed) + ) + ) + + +def gzip_body(data: bytes) -> bytes: + return gzip.compress(data, mtime=0) + + +def as_sent(data: bytes) -> bytes: + return data + + +@dataclass(frozen=True, slots=True) +class Rig: + gateway: Gateway + upstream: Wire + gate: threading.Event + limited_key: str + + def calls(self) -> tuple[Request, ...]: + return tuple(request for request in self.upstream.drain() if request.method == "POST") + + +def _chat_completion() -> bytes: + return json.dumps( + { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "scripted"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + } + ).encode() + + +def _stream_frame(identity: str, delta: Mapping[str, JsonValue], finish_reason: str | None) -> bytes: + chunk: Final = { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "delta": dict(delta), "finish_reason": finish_reason}], + } + return b"data: " + json.dumps(chunk).encode() + b"\n\n" + + +def _chat_stream(gate: threading.Event | None) -> Reply: + identity: Final = f"chatcmpl-{uuid.uuid4().hex}" + return Reply( + content_type="text/event-stream", + chunks=( + _stream_frame(identity, {"role": "assistant", "content": "streamed "}, None), + _stream_frame(identity, {"content": "answer"}, "stop"), + b"data: [DONE]\n\n", + ), + gate_after_first=gate, + ) + + +def _anthropic_message() -> bytes: + return json.dumps( + { + "id": f"msg_{uuid.uuid4().hex}", + "type": "message", + "role": "assistant", + "model": ANTHROPIC_MODEL, + "content": [{"type": "text", "text": "scripted"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 3, "output_tokens": 2}, + } + ).encode() + + +def _respond(gate: threading.Event, request: Request) -> Reply: + if request.method == "GET": + return Reply(body=b'{"object":"list","data":[]}') + if request.target == "/v1/messages": + return Reply(body=_anthropic_message()) + body: Final = JSON_OBJECT.validate_json(request.body) + if body.get("stream") is not True: + return Reply(body=_chat_completion()) + text: Final = string_value(object_value(list_value(body["messages"])[0])["content"]) + return _chat_stream(gate if text.startswith(GATED_PREFIX) else None) + + +def _deployment(name: str, upstream: Wire) -> Mapping[str, JsonValue]: + return { + "model_name": name, + "litellm_params": {"model": "openai/gpt-4o-mini", "api_base": f"{upstream.url}/v1", "api_key": "sk-fixture"}, + } + + +def _config(directory: Path, upstream: Wire) -> Path: + base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + configuration: Final = { + **base, + "model_list": [_deployment(CHAT_MODEL, upstream), _deployment(OTHER_MODEL, upstream)], + "general_settings": { + **object_value(base.get("general_settings", {})), + "max_request_size_mb": 1, + "enable_claude_code_gateway": True, + }, + "router_settings": {**object_value(base.get("router_settings", {})), "num_retries": 0}, + } + path: Final = directory / "compressed-request-bodies.yaml" + path.write_text(yaml.safe_dump(configuration)) + return path + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + assert os.environ.get("LITELLM_LICENSE"), ( + "LITELLM_LICENSE must be set: max_request_size_mb is an enterprise feature" + ) + gate: Final = threading.Event() + directory: Final = tmp_path_factory.mktemp("compressed-request-bodies") + with gateway_from_environment() as root, wire_server(partial(_respond, gate)) as upstream: + overrides: Final = { + "ANTHROPIC_API_BASE": upstream.url, + "ANTHROPIC_API_KEY": "sk-ant-synthetic", + "MAX_DECOMPRESSED_REQUEST_BODIES_IN_FLIGHT_MB": "1", + } + with ( + owned_proxy_process(root, directory, overrides, config=_config(directory, upstream), workers=1) as owned, + owned.gateway.scenario() as scenario, + ): + yield Rig(owned.gateway, upstream, gate, scenario.key(models=[CHAT_MODEL])) + + +@pytest.fixture(autouse=True) +def empty_upstream_log(rig: Rig) -> None: + rig.upstream.drain() + + +def _headers(key: str, coding: str | None, extra: Mapping[str, str] = NO_HEADERS) -> Mapping[str, str]: + return { + "Authorization": f"Bearer {key}", + "content-type": "application/json", + **({"content-encoding": coding} if coding is not None else {}), + **extra, + } + + +def _send( + rig: Rig, + data: bytes, + coding: str | None, + *, + key: str | None = None, + extra: Mapping[str, str] = NO_HEADERS, +) -> httpx.Response: + return rig.gateway.client.post( + CHAT, + content=data, + headers=_headers(rig.gateway.key if key is None else key, coding, extra), + timeout=case_timeout_seconds(), + ) + + +def _chat_body(model: str, text: str, *, stream: bool = False) -> bytes: + return json.dumps( + { + "model": model, + "messages": [{"role": "user", "content": text}], + "stream": stream, + "cache": {"no-cache": True}, + } + ).encode() + + +def _error_type(response: httpx.Response) -> JsonValue: + return object_value(JSON_OBJECT.validate_json(response.content)["error"])["type"] + + +def _answer(response: httpx.Response) -> JsonValue: + choice: Final = object_value(list_value(JSON_OBJECT.validate_json(response.content)["choices"])[0]) + return object_value(choice["message"])["content"] + + +def _stream(gateway: Gateway, data: bytes, first_text: threading.Event) -> tuple[int, str]: + with ( + httpx.Client(base_url=gateway.client.base_url, timeout=case_timeout_seconds(), trust_env=False) as client, + client.stream("POST", CHAT, content=data, headers=_headers(gateway.key, "gzip")) as response, + ): + pieces: Final = response.iter_text() + first: Final = next(pieces, "") + first_text.set() + return response.status_code, first + "".join(pieces) + + +CODINGS: Final = ( + pytest.param("zstd", zstd_frame, id="zstd"), + pytest.param("gzip", gzip_body, id="gzip"), + pytest.param("gzip", as_sent, id="plain-labeled-gzip"), +) + + +@pytest.mark.parametrize(("coding", "encode"), CODINGS) +def test_a_compressed_chat_body_reaches_the_provider_as_its_plain_twin( + rig: Rig, coding: str, encode: Callable[[bytes], bytes] +) -> None: + body: Final = _chat_body(CHAT_MODEL, f"twin {uuid.uuid4().hex}") + plain: Final = _send(rig, body, None) + compressed: Final = _send(rig, encode(body), coding) + assert plain.status_code == 200, plain.text + assert compressed.status_code == 200, compressed.text + assert _answer(compressed) == _answer(plain) + assert compressed.headers["x-litellm-response-cost"] == plain.headers["x-litellm-response-cost"] + calls: Final = rig.calls() + assert len(calls) == 2, [call.target for call in calls] + assert JSON_OBJECT.validate_json(calls[1].body) == JSON_OBJECT.validate_json(calls[0].body) + assert "content-encoding" not in calls[1].headers, calls[1].headers + + +def test_auth_reads_the_model_from_a_compressed_body(rig: Rig) -> None: + refused: Final = _send(rig, zstd_frame(_chat_body(OTHER_MODEL, "outside the key")), "zstd", key=rig.limited_key) + assert refused.status_code == 403, refused.text + assert _error_type(refused) == "key_model_access_denied", refused.text + assert rig.calls() == () + served: Final = _send(rig, zstd_frame(_chat_body(CHAT_MODEL, "inside the key")), "zstd", key=rig.limited_key) + assert served.status_code == 200, served.text + assert len(rig.calls()) == 1 + + +@pytest.mark.parametrize( + ("coding", "encode"), + (pytest.param("zstd", zstd_frame, id="zstd"), pytest.param("gzip", as_sent, id="plain-labeled-gzip")), +) +def test_a_compressed_pass_through_body_reaches_the_vendor_decoded( + rig: Rig, coding: str, encode: Callable[[bytes], bytes] +) -> None: + body: Final = json.dumps( + { + "model": ANTHROPIC_MODEL, + "max_tokens": 8, + "messages": [{"role": "user", "content": f"pass through {uuid.uuid4().hex}"}], + } + ).encode() + response: Final = rig.gateway.client.post( + "/anthropic/v1/messages", + content=encode(body), + headers={ + "x-litellm-api-key": rig.gateway.key, + "anthropic-version": "2023-06-01", + "content-type": "application/json", + "content-encoding": coding, + }, + timeout=case_timeout_seconds(), + ) + assert response.status_code == 200, response.text + calls: Final = rig.calls() + assert [call.target for call in calls] == ["/v1/messages"] + assert JSON_OBJECT.validate_json(calls[0].body) == JSON_OBJECT.validate_json(body) + assert "content-encoding" not in calls[0].headers, calls[0].headers + + +def test_a_gzip_bomb_gets_a_413_with_the_headers_other_errors_carry(rig: Rig) -> None: + origin: Final = {"origin": ORIGIN} + rejected: Final = _send(rig, gzip_body(b"\0" * 800_000), "gzip", extra=origin) + unauthorized: Final = _send( + rig, _chat_body(CHAT_MODEL, "unknown key"), None, key=f"sk-unknown-{uuid.uuid4().hex}", extra=origin + ) + assert rejected.status_code == 413, rejected.text + assert unauthorized.status_code == 401, unauthorized.text + assert unauthorized.headers.get("access-control-allow-origin") is not None, dict(unauthorized.headers) + assert {name: rejected.headers.get(name) for name in MIDDLEWARE_HEADERS} == { + name: unauthorized.headers.get(name) for name in MIDDLEWARE_HEADERS + } + assert rig.calls() == () + + +def test_the_operator_size_limit_holds_a_compressed_body_to_its_plain_twins_limit(rig: Rig) -> None: + over: Final = _chat_body(CHAT_MODEL, secrets.token_hex(600_000)) + compressed_over: Final = gzip_body(over) + assert len(compressed_over) < MIB < len(over) + plain: Final = _send(rig, over, None) + rejected: Final = _send(rig, compressed_over, "gzip") + served: Final = _send(rig, gzip_body(_chat_body(CHAT_MODEL, secrets.token_hex(300_000))), "gzip") + assert plain.status_code == 413, plain.text + assert rejected.status_code == plain.status_code, rejected.text + assert served.status_code == 200, served.text + assert len(rig.calls()) == 1 + + +def test_the_decoded_body_budget_turns_away_a_second_body_until_the_stream_ends(rig: Rig) -> None: + held: Final = gzip_body(_chat_body(CHAT_MODEL, GATED_PREFIX + secrets.token_hex(300_000), stream=True)) + second: Final = gzip_body(_chat_body(CHAT_MODEL, secrets.token_hex(300_000))) + keyless: Final = _send(rig, second, "gzip", key=f"sk-unknown-{uuid.uuid4().hex}") + assert keyless.status_code == 401, keyless.text + first_text: Final = threading.Event() + with ThreadPoolExecutor(1) as pool: + streaming: Final = pool.submit(_stream, rig.gateway, held, first_text) + try: + eventually(lambda: first_text.is_set() or streaming.done(), bool, seconds=case_timeout_seconds()) + reached: Final = rig.calls() + turned_away: Final = _send(rig, second, "gzip") + finally: + rig.gate.set() + status, streamed = streaming.result(timeout=case_timeout_seconds()) + assert len(reached) == 1, streamed + assert turned_away.status_code == 503, turned_away.text + assert _error_type(turned_away) == "overloaded_error", turned_away.text + assert turned_away.headers.get("retry-after") == "1", dict(turned_away.headers) + assert status == 200, streamed + assert "streamed " in streamed and "answer" in streamed and "data: [DONE]" in streamed, streamed + assert rig.calls() == () + after_the_stream: Final = _send(rig, second, "gzip") + assert after_the_stream.status_code == 200, after_the_stream.text + assert len(rig.calls()) == 1 + + +def test_a_claude_code_otlp_export_keeps_a_two_member_gzip_body(rig: Rig) -> None: + response: Final = rig.gateway.client.post( + "/claude_code_gateway/v1/metrics", + content=gzip_body(b"first member") + gzip_body(b"second member"), + headers={ + "Authorization": f"Bearer {rig.gateway.key}", + "content-type": "application/x-protobuf", + "content-encoding": "gzip", + }, + timeout=case_timeout_seconds(), + ) + assert response.status_code == 200, response.text From b97fcf1af0c68767e6c42990f2a6815df972cfb1 Mon Sep 17 00:00:00 2001 From: Aidan Singh <95709291+aidanasingh@users.noreply.github.com> Date: Wed, 28 Oct 2026 06:24:26 -0400 Subject: [PATCH 012/128] fix(router): let affinity pins yield to deployments that already failed this request (#32309) * fix(router): apply weighted-failover exclusion before deployment-affinity callback * fix(router): let affinity pins yield to deployments that already failed this request Moving the weighted failover exclusion ahead of the affinity check also changed what tag, order, plugin and managed-file filters saw. Restore the original filter order and have the affinity check ignore session and API key pins that point at a deployment already excluded or retry-skipped in this request. previous_response_id pins still win. * fix(router): release affinity pins on deployments that failed this request Delete a session or API key pin whose deployment already failed the request, so the deployment that serves it claims the pin. Prompt caching pins skip failed deployments too, and the proxy strips the internal failed-id fields from client bodies. * fix(router): keep affinity pins in place and skip failed origins in encrypted content affinity A pin now only yields for the request whose pinned deployment failed. Deleting it moved pins after one 429 or a request-caused 400 and added a Redis call on the routing path. Encrypted content affinity also ignores deployments that failed this request, so a failed origin takes the strip-and-dispatch path. * fix(router): keep encrypted reasoning when a retry can only go back to its origin The encrypted content check now drops failed deployments the way the router does: a weighted failover exclusion always holds, and a retry skip that would leave nothing falls back to the whole group. The proxy also drops client-sent metadata._failover_excluded_ids, and the encrypted content tests use a real Router * fix(router): keep encrypted reasoning when previous_response_id binds the failed origin --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/proxy/litellm_pre_call_utils.py | 3 + litellm/router.py | 11 +- litellm/router_utils/common_utils.py | 28 ++++- .../deployment_affinity_check.py | 8 +- .../encrypted_content_affinity_check.py | 39 ++++++- .../prompt_caching_deployment_check.py | 5 +- .../unit/proxy/test_litellm_pre_call_utils.py | 4 + .../test_deployment_affinity_check.py | 90 +++++++++++++++ .../test_encrypted_content_affinity_check.py | 109 ++++++++++++++++++ .../test_prompt_caching_deployment_check.py | 21 ++++ tests/unit/test_router/test_router.py | 4 +- tests/unit/test_router_weighted_failover.py | 77 ++++++++++++- 12 files changed, 382 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 1e7cdc44c70..0f8ac58861b 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -363,6 +363,8 @@ _UNTRUSTED_ROOT_CONTROL_FIELDS: Final = ( "fallback_depth", "_target_order", "attempted_targets", + "_excluded_deployment_ids", + "_retry_skipped_deployment_ids", "proxy_server_request", "standard_logging_object", "secret_fields", @@ -444,6 +446,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = ( "_pipeline_managed_guardrails", "client_disconnected", "error_information", + "_failover_excluded_ids", PRE_CALL_EXECUTED_GUARDRAILS_KEY, ) diff --git a/litellm/router.py b/litellm/router.py index 18c5c67abce..e53bab4659e 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -212,6 +212,7 @@ from litellm.router_utils.clientside_credential_handler import ( ) from litellm.router_utils.common_utils import ( MID_STREAM_CONTINUATION_KWARG, + as_retry_skipped_deployment_ids, deployments_the_caller_may_use, filter_team_based_models, filter_web_search_deployments, @@ -685,10 +686,6 @@ def _configured_model_info(deployment: Mapping[str, object]) -> Mapping[str, obj return _MODEL_INFO_ADAPTER.validate_python(deployment.get("model_info") or {}) -def _as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]: - return tuple(item for item in value if isinstance(item, str)) if isinstance(value, tuple) else () - - def _silent_experiment_targets(silent_model: object) -> tuple[str, ...]: if silent_model is None: return () @@ -8942,7 +8939,7 @@ class Router: return () if litellm.should_retry(status_code): return () - already_skipped_ids: Final = _as_retry_skipped_deployment_ids(already_skipped) + already_skipped_ids: Final = as_retry_skipped_deployment_ids(already_skipped) skipped: Final = tuple(sorted(frozenset((*already_skipped_ids, failed_deployment_id)))) verbose_router_logger.debug( "Retry skips deployments that already answered %s to this request: %s", status_code, skipped @@ -15186,7 +15183,7 @@ class Router: ## RETRY SKIP ## -> drop deployments that already refused this request with a ## non-retryable status, unless that leaves nothing, so the caller still gets ## the provider's own error instead of a no-deployments error. - _retry_skipped_deployment_ids: Final = _as_retry_skipped_deployment_ids( + _retry_skipped_deployment_ids: Final = as_retry_skipped_deployment_ids( request_kwargs.pop("_retry_skipped_deployment_ids", None) if request_kwargs else None ) healthy_deployments = ( @@ -16418,7 +16415,7 @@ class Router: ) ## RETRY SKIP ## -> see async counterpart in async_get_healthy_deployments. - _retry_skipped_deployment_ids: Final = _as_retry_skipped_deployment_ids( + _retry_skipped_deployment_ids: Final = as_retry_skipped_deployment_ids( request_kwargs.pop("_retry_skipped_deployment_ids", None) if request_kwargs else None ) healthy_deployments = ( diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index a00d84babf3..91b9e5cb114 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -4,7 +4,7 @@ from collections.abc import Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Final, TypeVar -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError if TYPE_CHECKING: from litellm.types.llms.openai import OpenAIFileObject @@ -38,6 +38,7 @@ _OBJECT_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[object, object]]] = TypeAdapt _DEPLOYMENT_ROWS_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter( tuple[Mapping[str, object], ...] ) +_DEPLOYMENT_IDS_ADAPTER: Final = TypeAdapter(tuple[str, ...]) def with_admin_api_key_secret_reference(deployment: Deployment) -> Deployment: @@ -57,6 +58,31 @@ def without_router_only_kwargs(kwargs: Mapping[str, _V]) -> dict[str, _V]: return {key: value for key, value in kwargs.items() if key not in ROUTER_ONLY_CALL_KWARGS} +def as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]: + if not isinstance(value, tuple): + return () + items: Final[tuple[object, ...]] = value + return tuple(item for item in items if isinstance(item, str)) + + +def as_excluded_deployment_ids(value: object) -> tuple[str, ...]: + if value is None: + return () + try: + return _DEPLOYMENT_IDS_ADAPTER.validate_python(value) + except ValidationError: + return () + + +def deployment_ids_failed_in_request(request_kwargs: Mapping[str, object]) -> frozenset[str]: + return frozenset( + ( + *as_excluded_deployment_ids(request_kwargs.get("_excluded_deployment_ids")), + *as_retry_skipped_deployment_ids(request_kwargs.get("_retry_skipped_deployment_ids")), + ) + ) + + def request_user_api_key_auth(request_kwargs: Mapping[str, object]) -> object: """The caller's auth object, from whichever metadata bucket this surface writes to.""" metadata_value: Final = request_kwargs.get("metadata") diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index 8693eaa5151..daf7d6bb345 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -31,6 +31,7 @@ from litellm.caching.dual_cache import DualCache from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger, Span from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.router_utils.common_utils import deployment_ids_failed_in_request from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CallTypes @@ -429,6 +430,9 @@ class DeploymentAffinityCheck(CustomLogger): session_affinity_active: Final = ( enable_session_id or self._get_marker_session_affinity_ttl(request_kwargs=request_kwargs) is not None ) + failed_deployment_ids: Final = deployment_ids_failed_in_request( + cast(Mapping[str, object], request_kwargs) # cast-ok: request kwargs keys are param names + ) user_key: Final = ( self.get_user_key_from_request_kwargs(request_kwargs=request_kwargs) if (session_affinity_active or enable_user_key) @@ -458,7 +462,7 @@ class DeploymentAffinityCheck(CustomLogger): elif isinstance(session_cache_result, str): session_model_id = session_cache_result - if session_model_id: + if session_model_id and session_model_id not in failed_deployment_ids: session_deployment: Final = self._find_deployment_by_model_id( healthy_deployments=typed_healthy_deployments, model_id=session_model_id, @@ -493,7 +497,7 @@ class DeploymentAffinityCheck(CustomLogger): # Backwards / safety: allow raw string values. model_id = cache_result - if not model_id: + if not model_id or model_id in failed_deployment_ids: return typed_healthy_deployments deployment = self._find_deployment_by_model_id( diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index 5d47bf63a38..806c3db7601 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -51,6 +51,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( strip_encrypted_reasoning_from_messages, ) from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.router_utils.common_utils import as_excluded_deployment_ids, as_retry_skipped_deployment_ids from litellm.types.llms.openai import AllMessageValues from litellm.types.router import Deployment from litellm.utils import get_order_filtered_deployments @@ -182,6 +183,14 @@ class EncryptedContentAffinityCheck(CustomLogger): None, ) + @staticmethod + def _deployment_model_id(deployment: Mapping[str, object]) -> str | None: + model_info: Final = deployment.get("model_info") + if not isinstance(model_info, Mapping): + return None + deployment_id: Final = cast(Mapping[str, object], model_info).get("id") # cast-ok: narrowed by isinstance + return None if deployment_id is None else str(deployment_id) + @staticmethod def _find_deployment_by_model_id(healthy_deployments: list[dict], model_id: str) -> dict | None: for deployment in healthy_deployments: @@ -345,7 +354,8 @@ class EncryptedContentAffinityCheck(CustomLogger): typed_healthy_deployments: Final = cast(list[dict[str, object]], healthy_deployments) if not self._is_enabled_for_model_group(model): return typed_healthy_deployments - target_order: Final = _REQUEST_KWARGS_ADAPTER.validate_python(request_kwargs).get("_target_order") + typed_request_kwargs: Final = _REQUEST_KWARGS_ADAPTER.validate_python(request_kwargs) + target_order: Final = typed_request_kwargs.get("_target_order") candidates: Final = ( get_order_filtered_deployments(typed_healthy_deployments, target_order=target_order) if isinstance(target_order, int) @@ -375,8 +385,31 @@ class EncryptedContentAffinityCheck(CustomLogger): model_id, ) + previous_response_id: Final = typed_request_kwargs.get("previous_response_id") + previous_response_origin: Final = ResponsesAPIRequestUtils.get_model_id_from_response_id( + None if previous_response_id is None else str(previous_response_id) + ) + previous_response_binds_origin: Final = previous_response_origin == model_id + excluded_ids: Final[frozenset[str]] = ( + frozenset() + if previous_response_binds_origin + else frozenset(as_excluded_deployment_ids(typed_request_kwargs.get("_excluded_deployment_ids"))) + ) + retry_skipped_ids: Final[frozenset[str]] = ( + frozenset() + if previous_response_binds_origin + else frozenset(as_retry_skipped_deployment_ids(typed_request_kwargs.get("_retry_skipped_deployment_ids"))) + ) + not_excluded_candidates: Final = [ + candidate for candidate in candidates if self._deployment_model_id(candidate) not in excluded_ids + ] + untried_candidates: Final = [ + candidate + for candidate in not_excluded_candidates + if self._deployment_model_id(candidate) not in retry_skipped_ids + ] or not_excluded_candidates deployment: Final = self._find_deployment_by_model_id( - healthy_deployments=candidates, + healthy_deployments=untried_candidates, model_id=model_id, ) if deployment is not None: @@ -392,7 +425,7 @@ class EncryptedContentAffinityCheck(CustomLogger): # Follow-up switched model_name (LIT-2531): pin by Azure resource instead. boundary_matches, _originating = self._find_deployments_on_same_encryption_boundary( - healthy_deployments=candidates, + healthy_deployments=untried_candidates, model_id=model_id, ) if boundary_matches: diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index bb3475cdcbe..c04bab4d360 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -19,6 +19,7 @@ from litellm.integrations.anthropic_cache_control_hook import ( from litellm.integrations.custom_logger import CustomLogger, Span from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.llms.bedrock.common_utils import bedrock_cache_point_ttl_supported, bedrock_capability_model +from litellm.router_utils.common_utils import deployment_ids_failed_in_request from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam from litellm.types.utils import CallTypes, StandardLoggingPayload from litellm.utils import ( @@ -152,7 +153,9 @@ class PromptCachingDeploymentCheck(CustomLogger): request_tools, ), ) - if model_id_dict is not None: + if model_id_dict is not None and model_id_dict["model_id"] not in deployment_ids_failed_in_request( + request_data or {} + ): model_id: Final = model_id_dict["model_id"] for deployment in healthy_deployments: if deployment["model_info"]["id"] == model_id: diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index a44373b87c8..f743a14c7a2 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -1393,6 +1393,8 @@ async def test_add_litellm_data_to_request_strips_string_encoded_admin_injection ("fallback_depth", True), ("_target_order", 2), ("attempted_targets", ["forged"]), + ("_excluded_deployment_ids", ["pinned-deployment"]), + ("_retry_skipped_deployment_ids", ["pinned-deployment"]), ], ) async def test_add_litellm_data_to_request_strips_forged_fallback_hop_state( @@ -1477,6 +1479,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): "_managed_background_response": True, "_guardrail_pipelines": [{"name": "spoofed"}], "_pipeline_managed_guardrails": ["evaded"], + "_failover_excluded_ids": ["pinned-deployment"], "safe_user_metadata": "kept", } data = { @@ -1537,6 +1540,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): "_managed_background_response", "_guardrail_pipelines", "_pipeline_managed_guardrails", + "_failover_excluded_ids", } assert "litellm_metadata" not in updated for stripped_key in stripped_keys: diff --git a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py index d8bc4c45ab8..fa75f980cc0 100644 --- a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -9,6 +9,7 @@ import pytest import litellm from litellm.caching.dual_cache import DualCache +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, ) @@ -1209,3 +1210,92 @@ async def test_virtual_key_hash_wins_over_user_id_for_affinity(): parent_otel_span=None, ) assert other_key_same_user == healthy_deployments + + +class _ForbiddenError(Exception): + status_code: Final = 403 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("pre_call_check", "metadata"), + [ + ("session_affinity", {"session_id": "sess-retry"}), + ("deployment_affinity", {"user_api_key_hash": "hashed-key-retry"}), + ], +) +async def test_retry_reaches_healthy_sibling_when_affinity_pins_deployment_that_refused( + pre_call_check: str, metadata: dict[str, str] +): + router: Final = litellm.Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "bad", + "mock_response": _ForbiddenError("A refuses"), + "weight": 1, + }, + "model_info": {"id": "A"}, + }, + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "good", + "mock_response": "ok from B", + "weight": 0, + }, + "model_info": {"id": "B"}, + }, + ], + routing_strategy="simple-shuffle", + num_retries=1, + retry_after=0, + disable_cooldowns=True, + optional_pre_call_checks=[pre_call_check], + ) + + responses: Final = [ + await router.acompletion( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + metadata=dict(metadata), + ) + for _ in range(2) + ] + + assert [response._hidden_params["model_id"] for response in responses] == ["B", "B"] + assert [response._hidden_params["additional_headers"]["x-litellm-attempted-retries"] for response in responses] == [ + 1, + 1, + ] + + +@pytest.mark.asyncio +async def test_previous_response_id_pin_holds_even_after_its_deployment_failed(): + callback: Final = DeploymentAffinityCheck( + cache=DualCache(), + ttl_seconds=60, + enable_user_key_affinity=True, + enable_responses_api_affinity=True, + ) + healthy_deployments: Final = [ + {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}, "model_info": {"id": "A"}}, + {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}, "model_info": {"id": "B"}}, + ] + + assert await callback.async_filter_deployments( + model="gpt-4o", + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs={ + "previous_response_id": ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="openai", model_id="A", response_id="resp_upstream" + ), + "_excluded_deployment_ids": ["A"], + "metadata": {"user_api_key_hash": "hashed-key"}, + }, + parent_otel_span=None, + ) == [healthy_deployments[0]] diff --git a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index 87fbf4433a1..a3745c7b23f 100644 --- a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -16,6 +16,7 @@ The mechanism works without any cache and supports two encoding strategies: """ import time +from typing import Final from unittest.mock import AsyncMock, patch import pytest @@ -2606,3 +2607,111 @@ async def test_affinity_strips_for_team_and_pattern_routes(): assert result is sibling_pool assert not request_kwargs["input"][1].get("encrypted_content") + + +def _region_deployment(region: str) -> dict[str, object]: + return { + "model_name": "gpt-5.4", + "litellm_params": { + "model": "openai/gpt-5.4", + "api_base": f"https://{region}.example.com/v1", + "api_key": f"key-{region}", + }, + "model_info": {"id": region}, + } + + +def _reasoning_input_minted_by(origin: str) -> list[dict[str, object]]: + return [ + { + "type": "reasoning", + "id": ResponsesAPIRequestUtils._build_encrypted_item_id(origin, "rs_test"), + "encrypted_content": ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id("gAAAAA-blob", origin), + "summary": [{"type": "summary_text", "text": "scattering"}], + }, + {"role": "user", "content": "and sunsets?"}, + ] + + +def _previous_response_id_from(origin: str) -> str: + return ResponsesAPIRequestUtils._build_responses_api_response_id("openai", origin, "resp_test") + + +_ORIGIN_FAILED_THIS_REQUEST: Final = pytest.mark.parametrize( + "failed_in_request", + [{"_excluded_deployment_ids": ["region-a"]}, {"_retry_skipped_deployment_ids": ("region-a",)}], +) + + +@pytest.mark.asyncio +@_ORIGIN_FAILED_THIS_REQUEST +@pytest.mark.parametrize("continuing", [{}, {"previous_response_id": _previous_response_id_from("region-b")}]) +async def test_affinity_strips_and_dispatches_when_origin_already_failed_this_request( + failed_in_request: dict[str, object], continuing: dict[str, object] +): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + router: Final = litellm.Router(model_list=[_region_deployment("region-a"), _region_deployment("region-b")]) + group: Final = router.model_list + reasoning_input: Final = _reasoning_input_minted_by("region-a") + + result: Final = await EncryptedContentAffinityCheck(router=router).async_filter_deployments( + model="gpt-5.4", + healthy_deployments=group, + messages=None, + request_kwargs={"litellm_metadata": {}, "input": reasoning_input, **failed_in_request, **continuing}, + ) + + assert result is group + assert reasoning_input[0] == {"type": "reasoning", "summary": [{"type": "summary_text", "text": "scattering"}]} + + +@pytest.mark.asyncio +@_ORIGIN_FAILED_THIS_REQUEST +async def test_affinity_keeps_origin_and_reasoning_when_previous_response_id_binds_the_failed_origin( + failed_in_request: dict[str, object], +): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + pool: Final = (_region_deployment("region-a"), _region_deployment("region-b")) + reasoning_input: Final = _reasoning_input_minted_by("region-a") + original_reasoning: Final = dict(reasoning_input[0]) + + assert await EncryptedContentAffinityCheck(router=litellm.Router(model_list=list(pool))).async_filter_deployments( + model="gpt-5.4", + healthy_deployments=list(pool), + messages=None, + request_kwargs={ + "litellm_metadata": {}, + "input": reasoning_input, + "previous_response_id": _previous_response_id_from("region-a"), + **failed_in_request, + }, + ) == [pool[0]] + assert reasoning_input[0] == original_reasoning + + +@pytest.mark.asyncio +@pytest.mark.parametrize("regions", [("region-a",), ("region-a", "region-b")]) +async def test_affinity_keeps_reasoning_when_every_deployment_already_refused_the_request(regions: tuple[str, ...]): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + router: Final = litellm.Router(model_list=[_region_deployment(region) for region in regions]) + reasoning_input: Final = _reasoning_input_minted_by("region-a") + original_reasoning: Final = dict(reasoning_input[0]) + + result: Final = await EncryptedContentAffinityCheck(router=router).async_filter_deployments( + model="gpt-5.4", + healthy_deployments=router.model_list, + messages=None, + request_kwargs={"litellm_metadata": {}, "input": reasoning_input, "_retry_skipped_deployment_ids": regions}, + ) + + assert [deployment["model_info"]["id"] for deployment in result] == ["region-a"] + assert reasoning_input[0] == original_reasoning diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 8d893160780..9487e4d6321 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -313,6 +313,27 @@ async def test_async_filter_deployments_does_not_pin_when_target_order_is_set(): assert filtered == deployments +@pytest.mark.asyncio +@pytest.mark.parametrize( + "failed_in_request", + [{"_excluded_deployment_ids": ["dep-2"]}, {"_retry_skipped_deployment_ids": ("dep-2",)}], +) +async def test_async_filter_deployments_does_not_pin_to_a_deployment_that_failed_this_request( + failed_in_request: dict[str, object], +): + cache: Final = DualCache() + check: Final = PromptCachingDeploymentCheck(cache=cache) + messages: Final = _messages(word_count=5000) + await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=messages, tools=None) + + assert await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=_deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6"), + messages=messages, + request_kwargs=failed_in_request, + ) == _deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6") + + @pytest.mark.asyncio async def test_async_filter_deployments_narrows_for_group_whose_model_minimum_is_lower(): """ diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index d9b655a0aff..30d675e130d 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -21239,9 +21239,9 @@ def test_router_retry_skip_stamp_feeds_deployment_ids_to_skip_on_retry( ], ) def test_router_as_retry_skipped_deployment_ids_keeps_only_a_tuple_of_strings(value, expected): - from litellm.router import _as_retry_skipped_deployment_ids + from litellm.router_utils.common_utils import as_retry_skipped_deployment_ids - assert _as_retry_skipped_deployment_ids(value) == expected + assert as_retry_skipped_deployment_ids(value) == expected @pytest.mark.parametrize( diff --git a/tests/unit/test_router_weighted_failover.py b/tests/unit/test_router_weighted_failover.py index 9cb38b00e1e..8b7fc81dbde 100644 --- a/tests/unit/test_router_weighted_failover.py +++ b/tests/unit/test_router_weighted_failover.py @@ -8,13 +8,14 @@ cross-group fallback runs. """ from collections import Counter -from typing import Optional +from typing import Final, Optional from unittest.mock import AsyncMock, patch import pytest import litellm from litellm import Router +from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo from litellm.utils import get_excluded_filtered_deployments @@ -944,3 +945,77 @@ async def test_failover_falls_through_to_external_fallback_when_remaining_in_coo ) assert response._hidden_params["model_id"] == "fallback" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("pre_call_check", "metadata"), + [ + ("session_affinity", {"session_id": "sess-1"}), + ("deployment_affinity", {"user_api_key_hash": "hashed-key-1"}), + ], +) +async def test_failover_reaches_healthy_sibling_when_affinity_pins_failed_deployment( + pre_call_check: str, metadata: dict[str, str] +): + router: Final = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "bad", + "mock_response": Exception("A down"), + "weight": 1, + }, + "model_info": {"id": "A"}, + }, + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "good", + "mock_response": "ok from B", + "weight": 0, + }, + "model_info": {"id": "B"}, + }, + { + "model_name": "fallback-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "good", + "mock_response": "ok from fallback", + }, + "model_info": {"id": "fallback"}, + }, + ], + routing_strategy="simple-shuffle", + num_retries=0, + disable_cooldowns=True, + enable_weighted_failover=True, + optional_pre_call_checks=[pre_call_check], + fallbacks=[{"test-model": ["fallback-model"]}], + ) + + while_a_is_down: Final = await router.acompletion( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + metadata=dict(metadata), + ) + router.upsert_deployment( + Deployment( + model_name="test-model", + litellm_params=LiteLLM_Params(model="gpt-4o", api_key="bad", mock_response="ok from A", weight=1), + model_info=ModelInfo(id="A"), + ) + ) + after_a_recovers: Final = await router.acompletion( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + metadata=dict(metadata), + ) + + assert while_a_is_down._hidden_params["model_id"] == "B" + assert while_a_is_down._hidden_params["additional_headers"]["x-litellm-attempted-fallbacks"] == 1 + assert after_a_recovers._hidden_params["model_id"] == "A" From de40c0aa39e26d318eaa3fc0337a788165862972 Mon Sep 17 00:00:00 2001 From: tessellark <349031000+tessellark@users.noreply.github.com> Date: Wed, 28 Oct 2026 03:38:21 -0700 Subject: [PATCH 013/128] test(llms): audit #47535 (#48786) Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../providers/_api_base_url_join_support.py | 9 ++ ...test_api_base_url_join_host_owned_proxy.py | 37 ++++++- .../providers/test_api_base_url_join_wire.py | 99 +++++++++++++++++-- 3 files changed, 134 insertions(+), 11 deletions(-) diff --git a/tests/integration/providers/_api_base_url_join_support.py b/tests/integration/providers/_api_base_url_join_support.py index 44bd6ec4dca..5711cc77309 100644 --- a/tests/integration/providers/_api_base_url_join_support.py +++ b/tests/integration/providers/_api_base_url_join_support.py @@ -13,6 +13,7 @@ from pydantic import JsonValue COUNTED_TOKENS: Final = 7 RELEVANCE: Final = 0.75 +AZURE_DEPLOYMENT: Final = "gpt-url-join" _USAGE: Final = {"input_tokens": 3, "output_tokens": 2} @@ -146,12 +147,20 @@ def _cohere_rerank(request: Request) -> Reply: ) +def _titan_embedding(_: Request) -> Reply: + return _json({"embedding": list(VECTOR), "inputTextTokenCount": 2}) + + _BY_PREFIX: Final[Mapping[str, Callable[[Request], Reply]]] = MappingProxyType( { "an": _anthropic, "bi": lambda request: _json(anthropic_message(marker_of(request))), "bc": _bedrock_converse, + "be": _titan_embedding, "gm": _gemini, + "vx": _gemini, + "fd": _anthropic, + "az": _openai_like(AZURE_DEPLOYMENT), "mi": _openai_like("mistral-small-latest"), "db": _openai_like("databricks-meta-llama-3-3-70b-instruct"), "co": _cohere_rerank, diff --git a/tests/integration/providers/test_api_base_url_join_host_owned_proxy.py b/tests/integration/providers/test_api_base_url_join_host_owned_proxy.py index 4a075138305..687c15283f4 100644 --- a/tests/integration/providers/test_api_base_url_join_host_owned_proxy.py +++ b/tests/integration/providers/test_api_base_url_join_host_owned_proxy.py @@ -1,6 +1,7 @@ import os from collections.abc import Callable, Iterator from pathlib import Path +from types import MappingProxyType from typing import Final import pytest @@ -8,6 +9,7 @@ import yaml from integration._support.client import JSON_OBJECT, Gateway, gateway_from_environment from integration._support.forward_proxy import ForwardProxy, refusing_forward_proxy from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.vertex import service_account_json from integration.providers._api_base_trailing_slash_support import KEY, new_marker from pydantic import JsonValue @@ -17,6 +19,12 @@ _RERANK: Final = "host-rerank" _AZURE_CLAUDE: Final = "host-azure-claude" _AZURE_COUNT: Final = "host-azure-count" _EMPTY_BASE: Final = "host-empty-base" +_VERTEX_HOST: Final = "vertex-gateway" +_VERTEX_PROJECT: Final = "url-join-project" +_VERTEX_LOCATION: Final = "us-east5" +_VERTEX_CLAUDE_BASES: Final = MappingProxyType( + {"host-vertex-claude": f"http://{_VERTEX_HOST}", "host-vertex-claude-slash": f"http://{_VERTEX_HOST}/"} +) _SCRUBBED_PREFIXES: Final = ("ANTHROPIC_", "AZURE_AI_", "INFINITY_") _PROXY_VARIABLES: Final = frozenset({"HTTPS_PROXY", "HTTP_PROXY", "ALL_PROXY", "NO_PROXY"}) _load_yaml: Final[Callable[[str], object]] = yaml.safe_load @@ -28,7 +36,20 @@ def _deployment(name: str, model: str, api_base: str) -> dict[str, JsonValue]: return {"model_name": name, "litellm_params": {"model": model, "api_base": api_base, "api_key": KEY}} -def _config(directory: Path) -> Path: +def _vertex_claude_deployment(name: str, api_base: str, token_url: str) -> dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": { + "model": "vertex_ai/claude-haiku-5-5", + "api_base": api_base, + "vertex_project": _VERTEX_PROJECT, + "vertex_location": _VERTEX_LOCATION, + "vertex_credentials": service_account_json(_VERTEX_PROJECT, token_url), + }, + } + + +def _config(directory: Path, token_url: str) -> Path: base: Final = JSON_OBJECT.validate_python(_load_yaml(Path("tests/integration/proxy_config.yaml").read_text())) configuration: Final = { **base, @@ -38,6 +59,7 @@ def _config(directory: Path) -> Path: _deployment(_AZURE_CLAUDE, "azure_ai/claude-haiku-5-5", f"http://{_AZURE_HOST}/anthropic"), _deployment(_AZURE_COUNT, "azure_ai/claude-haiku-5-5", f"http://{_AZURE_HOST}"), _deployment(_EMPTY_BASE, "anthropic/claude-haiku-5-5", ""), + *(_vertex_claude_deployment(name, base, token_url) for name, base in _VERTEX_CLAUDE_BASES.items()), ], "router_settings": {"disable_cooldowns": True, "num_retries": 0}, } @@ -72,7 +94,7 @@ def isolated(forward: ForwardProxy, tmp_path_factory: pytest.TempPathFactory) -> "NO_PROXY": "127.0.0.1,localhost", "PROXY_BATCH_POLLING_ENABLED": "false", }, - config=_config(directory), + config=_config(directory, shared.upstream_url.rstrip("/")), remove_environment=_scrubbed(), ) as candidate: yield candidate @@ -138,3 +160,14 @@ def test_azure_claude_count_tokens_on_a_bare_host_starting_with_anthropic_adds_t def test_empty_api_base_falls_back_to_the_provider_default_host(isolated: Gateway, forward: ForwardProxy) -> None: targets: Final = _targets_after(forward, lambda: _chat(isolated, _EMPTY_BASE)) assert targets == {"api.anthropic.com:443"} + + +@pytest.mark.parametrize("model", tuple(_VERTEX_CLAUDE_BASES), ids=("no-slash", "slash")) +def test_vertex_partner_claude_on_a_bare_host_reaches_one_single_slash_url( + isolated: Gateway, forward: ForwardProxy, model: str +) -> None: + targets: Final = _targets_after(forward, lambda: _chat(isolated, model)) + assert targets == { + f"http://{_VERTEX_HOST}/v1/projects/{_VERTEX_PROJECT}/locations/{_VERTEX_LOCATION}" + "/publishers/anthropic/models/claude-haiku-5-5:rawPredict" + } diff --git a/tests/integration/providers/test_api_base_url_join_wire.py b/tests/integration/providers/test_api_base_url_join_wire.py index 42e53ff9f95..74f8dde3cf7 100644 --- a/tests/integration/providers/test_api_base_url_join_wire.py +++ b/tests/integration/providers/test_api_base_url_join_wire.py @@ -7,6 +7,7 @@ import openai import pytest from integration._support.client import Gateway, Scenario, eventually, list_value, object_value, string_value from integration._support.database import read_rows +from integration._support.vertex import service_account_json from integration._support.wire import Reply, Request, wire_server from integration.providers._api_base_trailing_slash_support import ( KEY, @@ -22,7 +23,7 @@ from integration.providers._api_base_trailing_slash_support import ( reply_text, request_json, ) -from integration.providers._api_base_url_join_support import COUNTED_TOKENS, RELEVANCE, join_peer +from integration.providers._api_base_url_join_support import AZURE_DEPLOYMENT, COUNTED_TOKENS, RELEVANCE, join_peer from pydantic import JsonValue _HAIKU: Final = "anthropic.claude-haiku-5-5" @@ -49,10 +50,20 @@ _RESPONSES_ROWS: Final = (_ANTHROPIC, _GEMINI) _EMBEDDING_ROWS: Final = ( Row("mistral/mistral-embed", "/mi/v1", "/embeddings"), Row("gemini/gemini-embedding-2", "/gm/v1beta", "/models/gemini-embedding-2:batchEmbedContents"), + Row("bedrock/amazon.titan-embed-text-v2:0", "/be", "/model/amazon.titan-embed-text-v2%3A0/invoke"), ) _TRANSCRIPTION: Final = Row("mistral/voxtral-mini-latest", "/mi/v1", "/audio/transcriptions") -_FULL_ROUTE_SUFFIXES: Final = ("", "/", "?beta=true") -_FULL_ROUTE_IDS: Final = ("exact", "slash", "query") +_FULL_ROUTE_SUFFIXES: Final = (("", ""), ("/", ""), ("?beta=true", "?beta=true"), ("/?beta=true", "?beta=true")) +_FULL_ROUTE_IDS: Final = ("exact", "slash", "query", "slash-query") +_VERTEX_PROJECT: Final = "url-join-project" +_VERTEX_LOCATION: Final = "us-east5" +_VERTEX_GEMINI_PATH: Final = ( + f"/vx/v1/projects/{_VERTEX_PROJECT}/locations/{_VERTEX_LOCATION}/publishers/google/models/gemini-3.8-flash" +) +_AZURE_V1_BASES: Final = ("/az", "/az/", "/az/openai/v1", "/az/openai/v1/") +_AZURE_V1_IDS: Final = ("host", "host-slash", "route", "route-slash") +_AZURE_ANTHROPIC_BASES: Final = ("/fd/anthropic", "/fd/anthropic/v1/messages") +_AZURE_ANTHROPIC_IDS: Final = ("anthropic-root", "messages-route") def _ids(rows: tuple[Row, ...]) -> tuple[str, ...]: @@ -65,6 +76,17 @@ def _deploy(scenario: Scenario, row: Row, api_base: str) -> str: return scenario.model(model=row.model, api_base=api_base, api_key=KEY) +def _deploy_vertex(gateway: Gateway, scenario: Scenario, model: str, api_base: str) -> str: + return scenario.model( + model=model, + api_base=api_base, + api_key=None, + vertex_project=_VERTEX_PROJECT, + vertex_location=_VERTEX_LOCATION, + vertex_credentials=service_account_json(_VERTEX_PROJECT, gateway.upstream_url.rstrip("/")), + ) + + def _openai(gateway: Gateway) -> openai.OpenAI: return openai.OpenAI(base_url=str(gateway.client.base_url) + "/v1", api_key=gateway.key, max_retries=0) @@ -80,6 +102,13 @@ def _chat(gateway: Gateway, model: str, marker: str) -> str | None: return completion.choices[0].message.content +def _streamed_chat(gateway: Gateway, model: str, marker: str) -> str: + stream: Final = _openai(gateway).chat.completions.create( + model=model, messages=[{"role": "user", "content": marker}], stream=True + ) + return "".join(chunk.choices[0].delta.content or "" for chunk in stream if chunk.choices) + + def _message_text(gateway: Gateway, model: str, marker: str) -> list[str]: message: Final = _anthropic(gateway).messages.create( model=model, max_tokens=16, messages=[{"role": "user", "content": marker}] @@ -270,22 +299,74 @@ def test_cache_hit_on_a_slash_base_skips_the_upstream_and_logs_its_own_row(gatew assert "//" not in urlsplit(string_value(hit["api_base"])).path, hit -@pytest.mark.parametrize("suffix", _FULL_ROUTE_SUFFIXES, ids=_FULL_ROUTE_IDS) -def test_chat_base_that_already_names_the_route_keeps_its_path_and_query(gateway: Gateway, suffix: str) -> None: +@pytest.mark.parametrize(("suffix", "sent"), _FULL_ROUTE_SUFFIXES, ids=_FULL_ROUTE_IDS) +def test_chat_base_that_already_names_the_route_keeps_its_path_and_query( + gateway: Gateway, suffix: str, sent: str +) -> None: marker: Final = new_marker() with wire_server(join_peer) as wire, gateway.scenario() as scenario: model: Final = _deploy(scenario, _ANTHROPIC, f"{wire.url}/an/v1/messages{suffix}") assert _chat(gateway, model, marker) == reply_text(marker) - assert_one_request(wire.drain(), f"/an/v1/messages{suffix.rstrip('/')}", marker) + assert_one_request(wire.drain(), f"/an/v1/messages{sent}", marker) -@pytest.mark.parametrize("suffix", _FULL_ROUTE_SUFFIXES, ids=_FULL_ROUTE_IDS) -def test_messages_base_that_already_names_the_route_keeps_its_path_and_query(gateway: Gateway, suffix: str) -> None: +@pytest.mark.parametrize(("suffix", "sent"), _FULL_ROUTE_SUFFIXES, ids=_FULL_ROUTE_IDS) +def test_messages_base_that_already_names_the_route_keeps_its_path_and_query( + gateway: Gateway, suffix: str, sent: str +) -> None: marker: Final = new_marker() with wire_server(join_peer) as wire, gateway.scenario() as scenario: model: Final = _deploy(scenario, _ANTHROPIC, f"{wire.url}/an/v1/messages{suffix}") assert _message_text(gateway, model, marker) == [reply_text(marker)] - assert_one_request(wire.drain(), f"/an/v1/messages{suffix.rstrip('/')}", marker) + assert_one_request(wire.drain(), f"/an/v1/messages{sent}", marker) + + +@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS) +def test_vertex_model_url_base_calls_its_method_with_no_slash_before_the_colon(gateway: Gateway, trailing: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy_vertex( + gateway, scenario, "vertex_ai/gemini-3.8-flash", f"{wire.url}{_VERTEX_GEMINI_PATH}{trailing}" + ) + assert _chat(gateway, model, marker) == reply_text(marker) + assert_one_request(wire.drain(), f"{_VERTEX_GEMINI_PATH}:generateContent", marker, keyed=False) + + +@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS) +def test_vertex_model_url_base_streams_with_no_slash_before_the_colon(gateway: Gateway, trailing: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy_vertex( + gateway, scenario, "vertex_ai/gemini-3.8-flash", f"{wire.url}{_VERTEX_GEMINI_PATH}{trailing}" + ) + assert _streamed_chat(gateway, model, marker) == reply_text(marker) + assert_one_request(wire.drain(), f"{_VERTEX_GEMINI_PATH}:streamGenerateContent?alt=sse", marker, keyed=False) + + +@pytest.mark.parametrize("base", _AZURE_V1_BASES, ids=_AZURE_V1_IDS) +def test_azure_v1_chat_reaches_one_openai_v1_route(gateway: Gateway, base: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"azure/{AZURE_DEPLOYMENT}", api_base=f"{wire.url}{base}", api_key=KEY, api_version="v1" + ) + assert _chat(gateway, model, marker) == reply_text(marker) + assert_one_request(wire.drain(), "/az/openai/v1/chat/completions", marker) + + +@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS) +@pytest.mark.parametrize("base", _AZURE_ANTHROPIC_BASES, ids=_AZURE_ANTHROPIC_IDS) +def test_azure_ai_count_tokens_reaches_one_count_tokens_route(gateway: Gateway, base: str, trailing: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="azure_ai/claude-haiku-5-5", api_base=f"{wire.url}{base}{trailing}", api_key=KEY + ) + response: Final = gateway.request( + "POST", "/v1/messages/count_tokens", {"model": model, "messages": [{"role": "user", "content": marker}]} + ) + assert body_of(response)["input_tokens"] == COUNTED_TOKENS + assert_one_request(wire.drain(), "/fd/anthropic/v1/messages/count_tokens", marker) @pytest.mark.parametrize("suffix", ("", "/", "?tier=a"), ids=("exact", "slash", "query")) From 91db681532d35f318726df20c605e14c82dd2219 Mon Sep 17 00:00:00 2001 From: norkelbyte <348963847+norkelbyte@users.noreply.github.com> Date: Wed, 28 Oct 2026 03:50:24 -0700 Subject: [PATCH 014/128] test(responses): audit #47504 (#48787) Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../test_background_response_billed_once.py | 31 ++++++++++++++++--- 1 file changed, 26 insertions(+), 5 deletions(-) diff --git a/tests/integration/spend/test_background_response_billed_once.py b/tests/integration/spend/test_background_response_billed_once.py index 1b7fc975217..b207e46204a 100644 --- a/tests/integration/spend/test_background_response_billed_once.py +++ b/tests/integration/spend/test_background_response_billed_once.py @@ -642,6 +642,15 @@ def _key_budget(scenario: Scenario, budget_owner: str, model_budget: JsonValue) return {} +def _assert_the_next_create_exceeds_the_model_budget(scenario: Scenario, key: str, model: str) -> None: + refused: Final = eventually( + lambda: _next_create(scenario, key, model), + lambda response: response.status_code != 200, + seconds=SPEND_LANDING_SECONDS, + ) + assert f"exceeded budget for model={model}" in refused.text, refused.text + + @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( @@ -663,12 +672,24 @@ def test_a_billed_background_response_counts_toward_its_creators_project_and_mod 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_the_next_create_exceeds_the_model_budget(scenario, key, upstream.model) + + +@pytest.mark.timeout(240) +def test_a_billed_background_response_on_a_team_model_counts_toward_the_budget_of_the_model_name_it_was_created_with( + gateway: Gateway, +) -> None: + with gateway.scenario() as scenario: + team_id: Final = scenario.team() + upstream: Final = scripted_upstream(scenario, background_routes("completed"), model_info={"team_id": team_id}) + key: Final = scenario.key( + team_id=team_id, model_max_budget={upstream.model: {"max_budget": PRICE / 2, "budget_duration": "1d"}} ) - assert f"exceeded budget for model={upstream.model}" in refused.text, refused.text + 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), lambda spend: spend > 0, seconds=SPEND_LANDING_SECONDS) + _assert_the_next_create_exceeds_the_model_budget(scenario, key, upstream.model) @pytest.mark.timeout(240) From 6ba77b7013666c5b16ddc369c88f926b6318f486 Mon Sep 17 00:00:00 2001 From: rj-cairn <349020452+rj-cairn@users.noreply.github.com> Date: Wed, 28 Oct 2026 03:50:50 -0700 Subject: [PATCH 015/128] test(spend): audit #47353 (#48785) A DeepSeek V4 stream that ends without a usage block now has a spend-row cell proving its prompt_tokens are counted with DeepSeek's tokenizer, the one path #47353's own cells did not reach Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../spend/test_streamed_usage_estimate.py | 73 +++++++++++++++++++ 1 file changed, 73 insertions(+) create mode 100644 tests/integration/spend/test_streamed_usage_estimate.py diff --git a/tests/integration/spend/test_streamed_usage_estimate.py b/tests/integration/spend/test_streamed_usage_estimate.py new file mode 100644 index 00000000000..c7bcdf4401a --- /dev/null +++ b/tests/integration/spend/test_streamed_usage_estimate.py @@ -0,0 +1,73 @@ +import json +import uuid +from hashlib import sha256 +from typing import Final + +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.deepseek_tokens import ( + CHINESE_TEXT, + SHORT_TEXT, + assert_counted_with_deepseek, + chinese_minus_short, +) +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +_DEEPSEEK_V4: Final = "DeepSeek-V4-Flash-0731" + + +def _data_frame(frame: dict[str, JsonValue]) -> bytes: + return b"data: " + json.dumps(frame).encode() + b"\n\n" + + +def _stream_without_usage(request: Request) -> Reply: + assert request.method == "POST" + chunk: Final[dict[str, JsonValue]] = { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion.chunk", + "created": 1, + "model": _DEEPSEEK_V4, + } + return Reply( + content_type="text/event-stream", + chunks=( + _data_frame({**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "OK"}}]}), + _data_frame({**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}), + b"data: [DONE]\n\n", + ), + ) + + +def _streamed_prompt_tokens(gateway: Gateway, key: str, model: str, text: str) -> int: + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": text}], "stream": True}, + key=key, + ) + assert response.status_code == 200, response.text + assert response.text.rstrip().endswith("data: [DONE]"), response.text + rows: Final = eventually( + lambda: read_rows( + 'SELECT status, prompt_tokens FROM "LiteLLM_SpendLogs" WHERE api_key=%s', + (sha256(key.encode()).hexdigest(),), + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert rows[0]["status"] == "success", rows + prompt_tokens: Final = rows[0]["prompt_tokens"] + assert isinstance(prompt_tokens, int), rows + return prompt_tokens + + +def test_a_deepseek_v4_stream_without_usage_bills_prompt_tokens_counted_with_deepseeks_tokenizer( + gateway: Gateway, +) -> None: + with wire_server(_stream_without_usage) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"azure_ai/{_DEEPSEEK_V4}", api_base=wire.url, num_retries=0) + chinese: Final = _streamed_prompt_tokens(gateway, scenario.key(models=[model]), model, CHINESE_TEXT) + short: Final = _streamed_prompt_tokens(gateway, scenario.key(models=[model]), model, SHORT_TEXT) + assert len(wire.drain()) == 2 + assert_counted_with_deepseek(chinese - short, chinese_minus_short) From e2d78b69cd8b8e51569c9334eda6dd2aa491b002 Mon Sep 17 00:00:00 2001 From: m6atlas <349022726+m6atlas@users.noreply.github.com> Date: Wed, 28 Oct 2026 03:55:25 -0700 Subject: [PATCH 016/128] test(proxy): audit #46815 (#48768) * test(proxy): audit #46815 Add integration tests for every count route #46815 changed, and count a Claude subscription caller's tokens with their forwarded token when the deployment has no configured key, which #46815 turned into a 401 * test(proxy): accept the forwarded subscription token in the count endpoint mocks --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../anthropic/count_tokens/token_counter.py | 10 +- .../proxy/anthropic_endpoints/endpoints.py | 6 +- litellm/proxy/litellm_pre_call_utils.py | 39 +++-- litellm/proxy/proxy_server.py | 13 +- .../test_count_tokens_error_shapes_wire.py | 143 +++++++++++++++++- .../llms/anthropic/test_count_tokens_oauth.py | 39 +++++ .../unit/proxy/test_litellm_pre_call_utils.py | 20 +++ tests/unit/proxy/test_proxy_token_counter.py | 5 + 8 files changed, 250 insertions(+), 25 deletions(-) diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index 0764fb41941..368bb929eeb 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -27,6 +27,9 @@ anthropic_count_tokens_handler: Final = AnthropicCountTokensHandler() class AnthropicTokenCounter(BaseTokenCounter): """Token counter implementation for Anthropic provider using the CountTokens API.""" + def __init__(self, client_oauth_token: str | None = None) -> None: + self.client_oauth_token: Final = client_oauth_token + def should_use_token_counting_api( self, custom_llm_provider: str | None = None, @@ -95,12 +98,17 @@ class AnthropicTokenCounter(BaseTokenCounter): api_base: Final = litellm_params.get("api_base") try: - auth_header: Final = await AnthropicModelInfo.aget_auth_header( + configured_auth_header: Final = await AnthropicModelInfo.aget_auth_header( api_key=litellm_params.get("api_key") or litellm.anthropic_key or litellm.api_key, api_base=api_base, litellm_params=litellm_params, allow_workload_identity=True, ) + auth_header: Final = ( + configured_auth_header + if configured_auth_header is not None or self.client_oauth_token is None + else AnthropicModelInfo.get_auth_header(api_key=self.client_oauth_token, api_base=api_base) + ) if auth_header is None: return TokenCountResponse( total_tokens=0, diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 13a1844a6de..1d2045eb743 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -50,7 +50,10 @@ from litellm.proxy.common_utils.openai_error_payload import ( with_litellm_call_id, ) from litellm.proxy.common_utils.token_count_settings import count_source_fields -from litellm.proxy.litellm_pre_call_utils import should_auto_drop_params_for_agentic_cli +from litellm.proxy.litellm_pre_call_utils import ( + forwarded_anthropic_oauth_token, + should_auto_drop_params_for_agentic_cli, +) from litellm.router import Router router: Final = APIRouter() @@ -412,6 +415,7 @@ async def count_tokens( request=token_request, call_endpoint=True, messages_format="anthropic", + client_oauth_token=forwarded_anthropic_oauth_token(request.headers), **_count_tokens_body_options(data, auto_drop_params), ) except UnsupportedContentTypeError as e: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 0f8ac58861b..314f14b6279 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1415,12 +1415,34 @@ def _tenant_service_name(user_api_key_dict: UserAPIKeyAuth) -> str | None: ) +def litellm_auth_header_name(headers: Headers) -> str: + if "x-litellm-api-key" in headers: + return "x-litellm-api-key" + if "authorization" in headers: + return "authorization" + return "x-api-key" + + +def forwarded_anthropic_oauth_token(headers: Headers) -> str | None: + from litellm.llms.anthropic.common_utils import is_anthropic_oauth_key + + forwarded: Final = clean_headers(headers, authenticated_with_header=litellm_auth_header_name(headers)) + return next( + ( + value.removeprefix("Bearer ") + for header, value in forwarded.items() + if header.lower() == "authorization" and is_anthropic_oauth_key(value) + ), + None, + ) + + def clean_headers( headers: Headers, litellm_key_header_name: str | None = None, forward_llm_provider_auth_headers: bool = False, authenticated_with_header: str | None = None, -) -> dict: +) -> dict[str, str]: """ Removes litellm api key from headers @@ -2513,26 +2535,13 @@ async def add_litellm_data_to_request( forward_llm_auth = general_settings.get("forward_llm_provider_auth_headers", False) if not forward_llm_auth: forward_llm_auth = getattr(litellm, "forward_llm_provider_auth_headers", False) - # Determine which header was used for authentication - # This enables forwarding provider keys (e.g., x-api-key) when they weren't used for LiteLLM auth - authenticated_with_header = None - if "x-litellm-api-key" in request.headers: - # If x-litellm-api-key is present, it was used for auth - authenticated_with_header = "x-litellm-api-key" - elif "authorization" in request.headers: - # Authorization header was used for auth - authenticated_with_header = "authorization" - else: - # x-api-key or another header was used for auth - authenticated_with_header = "x-api-key" - _headers: Final[dict[str, str]] = clean_headers( request.headers, litellm_key_header_name=( general_settings.get("litellm_key_header_name") if general_settings is not None else None ), forward_llm_provider_auth_headers=forward_llm_auth, - authenticated_with_header=authenticated_with_header, + authenticated_with_header=litellm_auth_header_name(request.headers), ) _strip_untrusted_request_header_controls( _headers, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3b36ea5afce..04f84f1f7b5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -14628,6 +14628,7 @@ async def run_thread( # ) # 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.anthropic.count_tokens.token_counter import AnthropicTokenCounter 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 @@ -14649,6 +14650,14 @@ from litellm.responses.litellm_completion_transformation.count_request import ( ) +def _with_client_oauth_token( + counter: BaseTokenCounter | None, client_oauth_token: str | None +) -> BaseTokenCounter | None: + if client_oauth_token is None or not isinstance(counter, AnthropicTokenCounter): + return counter + return AnthropicTokenCounter(client_oauth_token=client_oauth_token) + + def _get_provider_token_counter( deployment: dict, model_to_use: str ) -> tuple[BaseTokenCounter | None, str | None, str | None]: @@ -15264,6 +15273,7 @@ async def count_request_tokens( context_management: JsonValue = None, request_litellm_params: Mapping[str, JsonValue] = MappingProxyType({}), responses_request: Mapping[str, object] | None = None, + client_oauth_token: str | None = None, ) -> TokenCountResponse: prompt: Final = request.prompt messages: Final = request.messages @@ -15321,7 +15331,8 @@ async def count_request_tokens( custom_llm_provider: str | None = None if call_endpoint is True and deployment is not None: # Auto-route to the correct provider based on model - provider_counter, _model, custom_llm_provider = _get_provider_token_counter(deployment, model_to_use) + routed_counter, _model, custom_llm_provider = _get_provider_token_counter(deployment, model_to_use) + provider_counter = _with_client_oauth_token(routed_counter, client_oauth_token) if _model is not None: model_to_use = _model diff --git a/tests/integration/providers/test_count_tokens_error_shapes_wire.py b/tests/integration/providers/test_count_tokens_error_shapes_wire.py index a6c3a6dc62f..59114e4aced 100644 --- a/tests/integration/providers/test_count_tokens_error_shapes_wire.py +++ b/tests/integration/providers/test_count_tokens_error_shapes_wire.py @@ -1,5 +1,6 @@ import json -from collections.abc import Iterator +from collections.abc import Callable, Iterator +from types import MappingProxyType from typing import Final import anthropic @@ -8,6 +9,7 @@ import pytest from integration._support.client import Gateway, gateway_from_environment from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process from integration._support.wire import Reply, Request, wire_server +from litellm.types.llms.anthropic import ANTHROPIC_OAUTH_BETA_HEADER from pydantic import JsonValue, TypeAdapter pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) @@ -26,6 +28,14 @@ _PROVIDER_KEY_VARIABLES: Final = ( "AZURE_TENANT_ID", "OPENAI_API_KEY", ) +_ENVIRONMENT_KEYS: Final = MappingProxyType( + { + "ANTHROPIC_API_KEY": "synthetic-anthropic-env-key", + "AZURE_API_KEY": "synthetic-azure-env-key", + "OPENAI_API_KEY": "synthetic-openai-env-key", + } +) +_SUBSCRIPTION_TOKEN: Final = "sk-ant-oat01-synthetic-subscription-token" @pytest.fixture(scope="module") @@ -44,14 +54,14 @@ def keyless_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedPro @pytest.fixture(scope="module") -def azure_environment_key_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]: +def environment_key_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]: with ( gateway_from_environment() as gateway, owned_proxy_process( gateway, - tmp_path_factory.mktemp("count-tokens-azure-environment-key"), - {"AZURE_API_KEY": "synthetic-azure-env-key"}, - remove_environment=tuple(name for name in _PROVIDER_KEY_VARIABLES if name != "AZURE_API_KEY"), + tmp_path_factory.mktemp("count-tokens-environment-keys"), + _ENVIRONMENT_KEYS, + remove_environment=tuple(name for name in _PROVIDER_KEY_VARIABLES if name not in _ENVIRONMENT_KEYS), workers=1, ) as owned, ): @@ -147,9 +157,9 @@ def test_anthropic_count_without_a_provider_credential_is_an_anthropic_401( def test_azure_ai_anthropic_count_sends_the_azure_environment_key_as_x_api_key( - azure_environment_key_proxy: OwnedProxy, + environment_key_proxy: OwnedProxy, ) -> None: - gateway: Final = azure_environment_key_proxy.gateway + gateway: Final = environment_key_proxy.gateway def count(request: Request) -> Reply: assert request.headers.get("x-api-key") == "synthetic-azure-env-key", request.headers @@ -257,3 +267,122 @@ def test_google_count_tokens_errors_use_the_google_envelope(keyless_proxy: Owned assert _JSON_OBJECT.validate_json(malformed.content) == { "error": {"code": 400, "message": "Invalid Google countTokens request body", "status": "INVALID_ARGUMENT"} } + + +def _counted(tokens: int) -> Reply: + return Reply(body=json.dumps({"object": "response.input_tokens", "input_tokens": tokens}).encode()) + + +def _anthropic_count(model: str) -> tuple[str, dict[str, JsonValue]]: + return "/v1/messages/count_tokens", {"model": model, "messages": _MESSAGES} + + +def _responses_count(model: str) -> tuple[str, dict[str, JsonValue]]: + return "/v1/responses/input_tokens", {"model": model, "input": "Count this synthetic request."} + + +@pytest.mark.parametrize( + ("provider_model", "base_path", "count_request", "header", "credential"), + ( + ("anthropic/claude-sonnet-5-5", "", _anthropic_count, "x-api-key", "synthetic-anthropic-env-key"), + ("openai/gpt-4o-mini", "/v1", _responses_count, "authorization", "Bearer synthetic-openai-env-key"), + ), + ids=("anthropic", "openai"), +) +def test_count_uses_a_provider_key_set_only_in_the_proxy_environment( + environment_key_proxy: OwnedProxy, + provider_model: str, + base_path: str, + count_request: Callable[[str], tuple[str, dict[str, JsonValue]]], + header: str, + credential: str, +) -> None: + gateway: Final = environment_key_proxy.gateway + with wire_server(lambda _request: _counted(47)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=provider_model, api_base=f"{wire.url}{base_path}", api_key=None, num_retries=0 + ) + path, body = count_request(model) + counted: Final = gateway.request("POST", path, body) + assert counted.status_code == 200, counted.text + assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == 47, counted.text + received: Final = wire.drain() + assert [request.headers.get(header) for request in received] == [credential], received + + +def _google_count(model: str) -> tuple[str, dict[str, JsonValue]]: + return f"/v1beta/models/{model}:countTokens", { + "contents": [{"role": "user", "parts": [{"text": "Count this synthetic request."}]}] + } + + +def _utils_count(model: str) -> tuple[str, dict[str, JsonValue]]: + return "/utils/token_counter?call_endpoint=true", {"model": model, "messages": _MESSAGES} + + +@pytest.mark.parametrize( + ("count_request", "google_status"), + ((_google_count, "UNAUTHENTICATED"), (_utils_count, None)), + ids=("google", "utils"), +) +def test_count_entrypoints_report_a_missing_provider_key_in_their_own_shape( + keyless_proxy: OwnedProxy, + count_request: Callable[[str], tuple[str, dict[str, JsonValue]]], + google_status: str | None, +) -> None: + gateway: Final = keyless_proxy.gateway + with wire_server(_never_called) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="anthropic/claude-sonnet-5-5", api_base=wire.url, api_key=None, num_retries=0 + ) + path, body = count_request(model) + rejected: Final = gateway.request("POST", path, body) + assert rejected.status_code == 401, rejected.text + error: Final = _JSON_OBJECT.validate_json(rejected.content)["error"] + assert isinstance(error, dict), rejected.text + assert f"No Anthropic API key is set for model {model}" in str(error["message"]), rejected.text + if google_status is not None: + assert (error["code"], error["status"]) == (401, google_status), rejected.text + assert wire.drain() == () + + +def test_anthropic_count_uses_the_callers_subscription_token_when_the_deployment_has_no_key( + keyless_proxy: OwnedProxy, +) -> None: + gateway: Final = keyless_proxy.gateway + with wire_server(lambda _request: _counted(41)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="anthropic/claude-sonnet-5-5", api_base=wire.url, api_key=None, num_retries=0 + ) + counted: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + {"model": model, "messages": _MESSAGES}, + headers={"Authorization": f"Bearer {_SUBSCRIPTION_TOKEN}", "x-litellm-api-key": f"Bearer {gateway.key}"}, + ) + assert counted.status_code == 200, counted.text + assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == 41, counted.text + (received,) = wire.drain() + assert received.headers.get("authorization") == f"Bearer {_SUBSCRIPTION_TOKEN}", received.headers + assert "x-api-key" not in received.headers, received.headers + assert ANTHROPIC_OAUTH_BETA_HEADER in received.headers.get("anthropic-beta", ""), received.headers + + +def test_anthropic_count_keeps_the_deployments_key_over_the_callers_subscription_token( + keyless_proxy: OwnedProxy, +) -> None: + gateway: Final = keyless_proxy.gateway + with wire_server(lambda _request: _counted(43)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="anthropic/claude-sonnet-5-5", api_base=wire.url, api_key="synthetic-deployment-key", num_retries=0 + ) + counted: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + {"model": model, "messages": _MESSAGES}, + headers={"Authorization": f"Bearer {_SUBSCRIPTION_TOKEN}", "x-litellm-api-key": f"Bearer {gateway.key}"}, + ) + assert counted.status_code == 200, counted.text + (received,) = wire.drain() + assert received.headers.get("x-api-key") == "synthetic-deployment-key", received.headers + assert _SUBSCRIPTION_TOKEN not in json.dumps(dict(received.headers)), received.headers diff --git a/tests/unit/llms/anthropic/test_count_tokens_oauth.py b/tests/unit/llms/anthropic/test_count_tokens_oauth.py index 16ec513916c..b8a7c001b37 100644 --- a/tests/unit/llms/anthropic/test_count_tokens_oauth.py +++ b/tests/unit/llms/anthropic/test_count_tokens_oauth.py @@ -221,6 +221,45 @@ class TestCountTokensUsesWorkloadIdentity: assert result.total_tokens == 13 assert route.calls.last.request.headers["x-api-key"] == FAKE_REGULAR_KEY + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("deployment_key", "expected_credentials"), + [ + (None, {"authorization": f"Bearer {FAKE_OAUTH_TOKEN}", "x-api-key": None}), + (FAKE_REGULAR_KEY, {"authorization": None, "x-api-key": FAKE_REGULAR_KEY}), + ], + ids=["keyless-deployment-uses-the-callers-token", "deployment-key-wins-over-the-callers-token"], + ) + async def test_a_callers_subscription_token_counts_only_when_no_key_is_configured( + self, + monkeypatch: pytest.MonkeyPatch, + httpx_transport_clients: None, + deployment_key: str | None, + expected_credentials: dict[str, str | None], + ) -> None: + for name in ("ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(litellm, "anthropic_key", None) + monkeypatch.setattr(litellm, "api_key", None) + + 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": 17}) + ) + result = await AnthropicTokenCounter(client_oauth_token=FAKE_OAUTH_TOKEN).count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": deployment_key}}, + request_model="claude-sonnet-4-5", + ) + + assert result is not None + assert result.total_tokens == 17 + sent = route.calls.last.request.headers + assert {name: sent.get(name) for name in expected_credentials} == expected_credentials + assert (ANTHROPIC_OAUTH_BETA_HEADER in sent["anthropic-beta"].split(",")) is (deployment_key is None) + @pytest.mark.asyncio async def test_a_failed_mint_degrades_like_an_anthropic_error(self, monkeypatch): from litellm.llms.anthropic.count_tokens import token_counter as token_counter_module diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index f743a14c7a2..a32c40e5a43 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -53,6 +53,7 @@ from litellm.proxy.litellm_pre_call_utils import ( add_provider_specific_headers_to_request, check_if_token_is_service_account, clean_headers, + forwarded_anthropic_oauth_token, get_dynamic_logging_metadata, get_metadata_variable_name, move_guardrails_to_metadata, @@ -5966,6 +5967,25 @@ def test_clean_headers_strips_x_api_key_when_byok_disabled(): assert "x-api-key" not in result +@pytest.mark.parametrize( + ("headers", "expected_token"), + [ + ( + {"authorization": "Bearer sk-ant-oat01-caller", "x-litellm-api-key": "Bearer sk-proxy"}, + "sk-ant-oat01-caller", + ), + ({"authorization": "Bearer sk-ant-oat01-caller"}, None), + ({"authorization": "Bearer sk-ant-api03-caller", "x-litellm-api-key": "Bearer sk-proxy"}, None), + ({"x-litellm-api-key": "Bearer sk-proxy"}, None), + ], + ids=["oauth-beside-proxy-key", "oauth-used-as-proxy-key", "non-oauth-authorization", "no-authorization"], +) +def test_forwarded_anthropic_oauth_token_follows_chats_forwarding_rule( + headers: dict[str, str], expected_token: str | None +) -> None: + assert forwarded_anthropic_oauth_token(Headers(headers)) == expected_token + + def test_clean_headers_strips_x_api_key_when_byok_enabled_but_x_api_key_was_auth_header(): """ Anti-replay regression: even when forward_llm_provider_auth_headers=True, diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index fa0c9eda02d..ebe3c7407d8 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -364,6 +364,7 @@ async def test_anthropic_messages_count_tokens_endpoint(): output_format=None, context_management=None, request_litellm_params=None, + client_oauth_token=None, ): assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" assert request.model == "claude-3-sonnet-20240229" @@ -444,6 +445,7 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): output_format=None, context_management=None, request_litellm_params=None, + client_oauth_token=None, ): assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" assert request.model == "gpt-4" @@ -1262,6 +1264,7 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): output_format=None, context_management=None, request_litellm_params=None, + client_oauth_token=None, ): raise ProxyException( message='{"detail":{"message":"Input is too long for requested model."}}', @@ -1319,6 +1322,7 @@ async def test_anthropic_endpoint_403_permission_error_format(): output_format=None, context_management=None, request_litellm_params=None, + client_oauth_token=None, ): raise ProxyException( message='{"Message":"Bearer Token has expired"}', @@ -1376,6 +1380,7 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): output_format=None, context_management=None, request_litellm_params=None, + client_oauth_token=None, ): raise ProxyException( message="Rate limit exceeded", From a0c490ea5471c8ef2c09b26e0ea744fe1de9961a Mon Sep 17 00:00:00 2001 From: m6atlas <349022726+m6atlas@users.noreply.github.com> Date: Wed, 28 Oct 2026 04:04:39 -0700 Subject: [PATCH 017/128] fix(bedrock): retry without strict when Bedrock rejects strict tool schemas under drop_params (#48733) * fix(bedrock): retry without strict when Bedrock rejects strict tool schemas under drop_params * test(bedrock): cover the strict tools 400 retry under drop_params * fix(bedrock): anchor the strict tools rejection match and strip toolSpec strict * test(bedrock): return captured requests from the strict tools test helper * fix(bedrock): retry without strict only when a tool sent strict to Bedrock A strict-limit 400 on a request whose tools carry strict: false, or on a model litellm never sends strict to, was sent again with an identical body. The fallback now takes the model and returns nothing unless a tool puts strict on the wire. Also types the fallback test parameter and records the mock transport's requests on the recorder instead of a mutated list * fix(bedrock): decide the strict retry on the profile's base model and retry grammar timeouts An application inference profile with base_model sends strict on the wire, so the retry check now resolves the capability model the way the request transformation does. Bedrock's "Grammar compilation timed out" 400 is a strict-grammar rejection too, so it takes the same single retry under drop_params. * test(bedrock): annotate the strict-tools retry test helper's values as Final --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../bedrock/chat/native_output_fallback.py | 48 +++++-- litellm/llms/bedrock/common_utils.py | 54 ++++++++ litellm/main.py | 13 +- .../chat/test_native_output_fallback.py | 102 ++++++++++++++ .../llms/bedrock/test_bedrock_common_utils.py | 79 +++++++++++ tests/unit/test_main.py | 128 +++++++++++++++++- 6 files changed, 412 insertions(+), 12 deletions(-) diff --git a/litellm/llms/bedrock/chat/native_output_fallback.py b/litellm/llms/bedrock/chat/native_output_fallback.py index 1b8ce5fcda4..6892cd37beb 100644 --- a/litellm/llms/bedrock/chat/native_output_fallback.py +++ b/litellm/llms/bedrock/chat/native_output_fallback.py @@ -1,11 +1,14 @@ from collections.abc import Callable, Coroutine +from functools import partial from typing import Final, TypeAlias from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.bedrock.common_utils import ( NativeOutputFallback, is_bedrock_native_output_rejection, + is_bedrock_strict_tools_rejection, warn_bedrock_native_output_rejected, + warn_bedrock_strict_tools_dropped, ) from litellm.types.utils import ModelResponse from litellm.utils import CustomStreamWrapper @@ -14,12 +17,13 @@ NativeOutputResponse: TypeAlias = ModelResponse | CustomStreamWrapper NativeOutputAttemptResult: TypeAlias = NativeOutputResponse | Coroutine[object, object, NativeOutputResponse] NativeOutputAttempt: TypeAlias = Callable[[NativeOutputFallback | None], NativeOutputAttemptResult] NativeOutputFallbackFactory: TypeAlias = Callable[[], NativeOutputFallback | None] +_RetryFor: TypeAlias = Callable[[BaseLLMException], NativeOutputFallback | None] -def _fallback_for( - error: BaseLLMException, fallback: NativeOutputFallbackFactory, model: str +def _native_output_retry( + error: BaseLLMException, fallback: NativeOutputFallbackFactory | None, model: str ) -> NativeOutputFallback | None: - if not is_bedrock_native_output_rejection(error.status_code, error.message): + if fallback is None or not is_bedrock_native_output_rejection(error.status_code, error.message): return None retry: Final = fallback() if retry is not None: @@ -27,16 +31,38 @@ def _fallback_for( return retry +def _strict_tools_retry( + error: BaseLLMException, fallback: NativeOutputFallbackFactory | None, model: str +) -> NativeOutputFallback | None: + if fallback is None or not is_bedrock_strict_tools_rejection(error.status_code, error.message): + return None + retry: Final = fallback() + if retry is not None: + warn_bedrock_strict_tools_dropped(model) + return retry + + +def _fallback_for( + error: BaseLLMException, + fallback: NativeOutputFallbackFactory | None, + strict_tools_fallback: NativeOutputFallbackFactory | None, + model: str, +) -> NativeOutputFallback | None: + native_output_retry: Final = _native_output_retry(error, fallback, model) + if native_output_retry is not None: + return native_output_retry + return _strict_tools_retry(error, strict_tools_fallback, model) + + async def _await_with_native_output_fallback( first_attempt: Coroutine[object, object, NativeOutputResponse], attempt: NativeOutputAttempt, - fallback: NativeOutputFallbackFactory, - model: str, + retry_for: _RetryFor, ) -> NativeOutputResponse: try: return await first_attempt except BaseLLMException as error: - retry: Final = _fallback_for(error, fallback, model) + retry: Final = retry_for(error) if retry is None: raise second_attempt: Final = attempt(retry) @@ -47,16 +73,20 @@ def call_with_native_output_fallback( attempt: NativeOutputAttempt, fallback: NativeOutputFallbackFactory | None, model: str, + strict_tools_fallback: NativeOutputFallbackFactory | None = None, ) -> NativeOutputAttemptResult: - if fallback is None: + if fallback is None and strict_tools_fallback is None: return attempt(None) + retry_for: Final = partial( + _fallback_for, fallback=fallback, strict_tools_fallback=strict_tools_fallback, model=model + ) try: first_attempt: Final = attempt(None) except BaseLLMException as error: - retry: Final = _fallback_for(error, fallback, model) + retry: Final = retry_for(error) if retry is None: raise return attempt(retry) if isinstance(first_attempt, Coroutine): - return _await_with_native_output_fallback(first_attempt, attempt, fallback, model) + return _await_with_native_output_fallback(first_attempt, attempt, retry_for) return first_attempt diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 28888f1a23a..fde995c7907 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -362,6 +362,60 @@ class NativeOutputFallback: litellm_params: Mapping[str, object] +_BEDROCK_STRICT_TOOLS_REJECTION: Final = re.compile( + r"Too many strict tools \(\d+\)" + r"|Schemas contains too many (?:optional parameters|parameters with union types)" + r"|[Cc]ompiled grammar (?:is too large|size)" + r"|tools\.\d+\.custom(?:\.strict)?: " + r"|Grammar compilation timed out" +) +_TOOLS_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, object], ...]) +_TOOL_SPEC_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + + +def is_bedrock_strict_tools_rejection(status_code: int, error_text: str) -> bool: + return status_code == 400 and _BEDROCK_STRICT_TOOLS_REJECTION.search(error_text) is not None + + +def warn_bedrock_strict_tools_dropped(model: str) -> None: + verbose_logger.warning( + "Bedrock rejected the strict tool schemas for model=%s, so litellm sent the request again without `strict` " + "because drop_params is set.", + model, + ) + + +def _without_strict(tool: Mapping[str, object]) -> Mapping[str, object]: + spec_key: Final = next((key for key in ("function", "toolSpec") if key in tool), None) + if spec_key is None: + return tool + spec: Final = _TOOL_SPEC_ADAPTER.validate_python(tool[spec_key]) + return {**tool, spec_key: {key: value for key, value in spec.items() if key != "strict"}} + + +def _sends_strict(tool: Mapping[str, object], supports_strict_tools: bool) -> bool: + if "toolSpec" in tool: + return "strict" in _TOOL_SPEC_ADAPTER.validate_python(tool["toolSpec"]) + if "function" not in tool: + return False + return supports_strict_tools and bool(_TOOL_SPEC_ADAPTER.validate_python(tool["function"]).get("strict")) + + +def bedrock_strict_tools_fallback( + model: str, optional_params: Mapping[str, object], litellm_params: Mapping[str, object] +) -> NativeOutputFallback | None: + tools: Final = _TOOLS_ADAPTER.validate_python(optional_params.get("tools", ())) + supports_strict_tools: Final = bedrock_converse_supports_strict_tools( + bedrock_capability_model(model, litellm_params.get("base_model")) + ) + if not any(_sends_strict(tool, supports_strict_tools) for tool in tools): + return None + return NativeOutputFallback( + optional_params={**optional_params, "tools": [_without_strict(tool) for tool in tools]}, + litellm_params=litellm_params, + ) + + def requests_output_config_format(optional_params: Mapping[str, object]) -> bool: match optional_params: case {"output_config": {"format": object()}}: diff --git a/litellm/main.py b/litellm/main.py index 9f412e1f24b..39a0a0cd3c1 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -123,6 +123,7 @@ from litellm.llms.bedrock.common_utils import ( BedrockModelInfo, NativeOutputFallback, bedrock_route_for_request, + bedrock_strict_tools_fallback, deployment_model_info, without_bedrock_native_output_opt_in, without_bedrock_route_prefix, @@ -206,7 +207,12 @@ from litellm.utils import ( from ._logging import verbose_logger from .caching.caching import disable_cache, enable_cache, update_cache -from .litellm_core_utils.core_helpers import as_str_mapping, normalize_drop_params, safe_deep_copy +from .litellm_core_utils.core_helpers import ( + as_str_mapping, + drops_unsupported_params, + normalize_drop_params, + safe_deep_copy, +) from .litellm_core_utils.fallback_utils import ( async_completion_with_fallbacks, completion_with_fallbacks, @@ -4378,6 +4384,11 @@ def _complete_bedrock(ctx: CompletionDispatchContext) -> _CompletionDispatchResu else None ), converse_model, + strict_tools_fallback=( + partial(bedrock_strict_tools_fallback, converse_model, converse_params, dict(litellm_params)) + if drops_unsupported_params(litellm_params.get("drop_params")) + else None + ), ) elif bedrock_route == "converse_like": model = model.replace("converse_like/", "") diff --git a/tests/unit/llms/bedrock/chat/test_native_output_fallback.py b/tests/unit/llms/bedrock/chat/test_native_output_fallback.py index 301fdaeb245..9ee8c14f5d6 100644 --- a/tests/unit/llms/bedrock/chat/test_native_output_fallback.py +++ b/tests/unit/llms/bedrock/chat/test_native_output_fallback.py @@ -121,3 +121,105 @@ def test_a_rejection_logs_the_retry_and_advises_opting_out_only_when_native_outp assert _MODEL in caplog.text assert ("supports_bedrock_native_structured_output: false" in caplog.text) is advises_opting_out + + +_STRICT_FALLBACK: Final = NativeOutputFallback(optional_params={"tools": "without_strict"}, litellm_params={}) +# Bedrock's 400s for strict tools, seen live on us.anthropic.claude-sonnet-4-6 from 2026-10-20 to 2026-10-28 and quoted +# in issue #34388, each answered 200 once the request dropped strict +_STRICT_TOOLS_REJECTIONS: Final = tuple( + BedrockError(status_code=400, message=f'{{"message":"The model returned the following errors: {text}"}}') + for text in ( + "Too many strict tools (21). The maximum number of strict tools supported is 20.", + "Schemas contains too many optional parameters (27), which would make grammar compilation inefficient.", + "Schemas contains too many parameters with union types (18 parameters with type arrays or anyOf). " + "This causes exponential compilation cost.", + "The compiled grammar is too large, which would cause performance issues.", + "tools.0.custom: Invalid schema: Schema type is missing for schema: {'description': 'tag'}", + "tools.3.custom: For 'integer' type, property 'minimum' is not supported", + "Compiled grammar size (410.4MB) exceeds maximum allowed size (300MB)", + "tools.1.custom.strict: Extra inputs are not permitted", + "Grammar compilation timed out.", + ) +) +_TOO_MANY_STRICT_TOOLS: Final = _STRICT_TOOLS_REJECTIONS[0] +# Bedrock's 400s for requests with a strict tool that fail for other reasons, seen live on +# us.anthropic.claude-sonnet-4-6 on 2026-10-21, which echo tool names and tool call ids +_OTHER_400S_WITH_A_STRICT_TOOL: Final = tuple( + BedrockError(status_code=400, message=f'{{"message":"{text}"}}') + for text in ( + "The tool check_grammar is already defined at toolConfig.tools.1.", + "Expected toolResult blocks at messages.0.content for the following Ids: call_grammar_9", + "tool_use ids were found without tool_result blocks immediately after: call_grammar_1", + ) +) + + +def _call_with_strict_tools_fallback( + bedrock: ScriptedBedrock, fallback: NativeOutputFallback | None, strict_tools_fallback: NativeOutputFallback | None +) -> object: + return _resolve( + call_with_native_output_fallback( + bedrock, + None if fallback is None else lambda: fallback, + _MODEL, + strict_tools_fallback=None if strict_tools_fallback is None else lambda: strict_tools_fallback, + ) + ) + + +@_MODES +@pytest.mark.parametrize("rejection", _STRICT_TOOLS_REJECTIONS, ids=lambda error: error.message[60:100]) +def test_a_strict_tools_rejection_is_sent_again_once_with_the_strict_tools_fallback( + asynchronous: bool, rejection: BedrockError +) -> None: + bedrock: Final = ScriptedBedrock(errors=(rejection, None), asynchronous=asynchronous) + + response: Final = _call_with_strict_tools_fallback(bedrock, None, _STRICT_FALLBACK) + + assert isinstance(response, ModelResponse) and response.id == "answer-2" + assert bedrock.sent == (None, _STRICT_FALLBACK) + + +@_MODES +@pytest.mark.parametrize( + ("error", "strict_tools_fallback"), + [ + (_OTHER_400, _STRICT_FALLBACK), + *((error, _STRICT_FALLBACK) for error in _OTHER_400S_WITH_A_STRICT_TOOL), + (BedrockError(status_code=500, message=_TOO_MANY_STRICT_TOOLS.message), _STRICT_FALLBACK), + (_TOO_MANY_STRICT_TOOLS, None), + ], + ids=[ + "other_400", + "duplicate_tool_name", + "missing_tool_result", + "orphan_tool_use", + "server_error", + "no_strict_tools_fallback", + ], +) +def test_no_second_request_without_both_a_strict_tools_rejection_and_a_strict_tools_fallback( + asynchronous: bool, error: BedrockError, strict_tools_fallback: NativeOutputFallback | None +) -> None: + bedrock: Final = ScriptedBedrock(errors=(error,), asynchronous=asynchronous) + + with pytest.raises(BedrockError) as raised: + _call_with_strict_tools_fallback(bedrock, None, strict_tools_fallback) + + assert raised.value is error + assert bedrock.sent == (None,) + + +@pytest.mark.parametrize( + ("error", "retried_with"), + [(_REJECTION, _FALLBACK), (_TOO_MANY_STRICT_TOOLS, _STRICT_FALLBACK)], + ids=["native_output_rejection", "strict_tools_rejection"], +) +def test_each_rejection_is_sent_again_with_its_own_fallback_when_both_are_set( + error: BedrockError, retried_with: NativeOutputFallback +) -> None: + bedrock: Final = ScriptedBedrock(errors=(error, None), asynchronous=False) + + _call_with_strict_tools_fallback(bedrock, _FALLBACK, _STRICT_FALLBACK) + + assert bedrock.sent == (None, retried_with) diff --git a/tests/unit/llms/bedrock/test_bedrock_common_utils.py b/tests/unit/llms/bedrock/test_bedrock_common_utils.py index 7bde1805e18..dd89d9c46b6 100644 --- a/tests/unit/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/unit/llms/bedrock/test_bedrock_common_utils.py @@ -6,7 +6,9 @@ from typing import Final from litellm.llms.bedrock.common_utils import( AmazonBedrockGlobalConfig, BedrockModelInfo, + NativeOutputFallback, bedrock_native_output_is_opt_in, + bedrock_strict_tools_fallback, deployment_model_info, extract_model_name_from_bedrock_arn, get_bedrock_base_model, @@ -2150,3 +2152,80 @@ def test_bedrock_response_stream_shape_is_unavailable_when_botocore_describes_th assert shape is None (logged,) = warning.call_args_list assert "only-in-the-description" not in logged.args[0] % logged.args[1:] + + +def test_bedrock_strict_tools_fallback_drops_only_strict_and_keeps_everything_else(): + strict_tool: Final = { + "type": "function", + "cache_control": {"type": "ephemeral"}, + "function": {"name": "lookup", "description": "Look up", "strict": True, "parameters": {"type": "object"}}, + } + plain_tool: Final = {"type": "function", "function": {"name": "plain", "parameters": {"type": "object"}}} + anthropic_tool: Final = {"name": "native", "input_schema": {"type": "object"}} + bedrock_tool: Final = {"toolSpec": {"name": "spec", "strict": True, "inputSchema": {"json": {"type": "object"}}}} + litellm_params: Final = {"drop_params": True} + + fallback: Final = bedrock_strict_tools_fallback( + "us.anthropic.claude-sonnet-4-6", + {"tools": [strict_tool, plain_tool, anthropic_tool, bedrock_tool], "maxTokens": 16, "stream": True}, + litellm_params, + ) + + assert fallback == NativeOutputFallback( + optional_params={ + "tools": [ + {**strict_tool, "function": {"name": "lookup", "description": "Look up", "parameters": {"type": "object"}}}, + plain_tool, + anthropic_tool, + {"toolSpec": {"name": "spec", "inputSchema": {"json": {"type": "object"}}}}, + ], + "maxTokens": 16, + "stream": True, + }, + litellm_params=litellm_params, + ) + assert strict_tool["function"]["strict"] is True + assert bedrock_tool["toolSpec"]["strict"] is True + + +@pytest.mark.parametrize( + "optional_params", + [ + {"maxTokens": 16}, + {"tools": [{"type": "function", "function": {"name": "plain", "parameters": {}}}]}, + {"tools": [{"toolSpec": {"name": "spec", "inputSchema": {"json": {}}}}]}, + ], + ids=["no_tools", "no_strict_tool", "no_strict_tool_spec"], +) +def test_bedrock_strict_tools_fallback_is_none_without_a_strict_tool(optional_params: dict[str, object]) -> None: + assert bedrock_strict_tools_fallback("us.anthropic.claude-sonnet-4-6", optional_params, {}) is None + + +_CLAUDE_PROFILE_ARN: Final = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/claude" + + +@pytest.mark.parametrize( + ("model", "litellm_params"), + [ + ("us.anthropic.claude-sonnet-4-6", {}), + ("us.amazon.nova-pro-v1:0", {}), + (_CLAUDE_PROFILE_ARN, {"base_model": "bedrock/us.anthropic.claude-sonnet-4-6"}), + (_CLAUDE_PROFILE_ARN, {}), + ], + ids=["claude", "nova", "claude_profile_with_base_model", "profile_without_base_model"], +) +@pytest.mark.parametrize("strict", [True, False]) +def test_bedrock_strict_tools_fallback_retries_only_when_the_request_sent_strict( + model: str, litellm_params: dict[str, object], strict: bool +) -> None: + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + + tools: Final = [{"type": "function", "function": {"name": "lookup", "strict": strict, "parameters": {}}}] + request: Final = AmazonConverseConfig().transform_request( + model, [{"role": "user", "content": "hi"}], {"tools": tools}, {**litellm_params}, {} + ) + sent_strict: Final = any("strict" in tool["toolSpec"] for tool in request["toolConfig"]["tools"]) + + fallback: Final = bedrock_strict_tools_fallback(model, {"tools": tools}, litellm_params) + + assert (fallback is not None) == sent_strict diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 70b268bf2a2..f12129af1fc 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -14,7 +14,7 @@ from dataclasses import dataclass from datetime import datetime from importlib import import_module from pathlib import Path -from typing import Final +from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -44,12 +44,13 @@ from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kw from litellm.litellm_core_utils.get_litellm_params import stored_control_options from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt -from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.types.litellm_params import ControlOptions from litellm.types.llms.openai import AllMessageValues, HttpxBinaryResponseContent from litellm.types.prompts.init_prompts import PromptSpec from litellm.types.utils import Delta, ModelResponseStream, StandardCallbackDynamicParams, StreamingChoices, Usage from tests.unit.llms.bedrock.slow_upstream import ( + CONVERSE_ANSWER, STREAM_TIMEOUT_SECONDS, UPSTREAM_FIRST_BYTE_SECONDS, slow_upstream_sync_client, @@ -9950,3 +9951,126 @@ async def test_converse_fallback_prices_the_current_attempt( else: client.close() configuration.reset_rust_configuration() + + +_LOOKUP_FUNCTION: Final = { + "name": "lookup_record", + "description": "Look up a record", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + "additionalProperties": False, + }, +} +_STRICT_LOOKUP_TOOL: Final = {"type": "function", "function": {**_LOOKUP_FUNCTION, "strict": True}} +# Bedrock's 400 for a request past Claude's strict tool limits, seen live on us.anthropic.claude-sonnet-4-6 on 2026-10-20 +_TOO_MANY_STRICT_TOOLS: Final = ( + "The model returned the following errors: Too many strict tools (21). The maximum number of strict tools " + "supported is 20. Try reducing the number of tools marked as strict." +) + + +@dataclass(slots=True) +class _StrictToolsLimitRejectionFirst: + stream: Literal["off", "fake", "real"] + sent: tuple[tuple[str, dict[str, object]], ...] = () + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.sent = (*self.sent, (request.url.path.rsplit("/", 1)[-1], json.loads(request.content))) + if len(self.sent) == 1: + return httpx.Response(400, json={"message": _TOO_MANY_STRICT_TOOLS}) + return httpx.Response(200, content=b"" if self.stream == "real" else CONVERSE_ANSWER) + + +async def _converse_with_a_strict_tools_limit_rejection_first( + asynchronous: bool, + stream: Literal["off", "fake", "real"], + drop_params: bool | None, + tool: Mapping[str, object] = _STRICT_LOOKUP_TOOL, +) -> tuple[tuple[tuple[str, dict[str, object]], ...], litellm.BadRequestError | None]: + respond: Final = _StrictToolsLimitRejectionFirst(stream) + kwargs: Final = { + "model": "bedrock/us.anthropic.claude-sonnet-4-6", + "messages": [{"role": "user", "content": "Find invoice 42"}], + "tools": [tool], + "aws_access_key_id": "AKIASTRICTTOOLS", + "aws_secret_access_key": "strict-tools-secret", + "aws_region_name": "us-east-1", + "num_retries": 0, + "drop_params": drop_params, + "stream": stream != "off", + "fake_stream": stream == "fake", + } + try: + chunks: Final = ( + await _acompletion_chunks(AsyncHTTPHandler(transport=httpx.MockTransport(respond)), kwargs, stream) + if asynchronous + else _completion_chunks( + HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))), kwargs, stream + ) + ) + except litellm.BadRequestError as error: + return respond.sent, error + assert chunks + return respond.sent, None + + +async def _acompletion_chunks( + client: AsyncHTTPHandler, kwargs: Mapping[str, object], stream: Literal["off", "fake", "real"] +) -> tuple[object, ...]: + response: Final = await litellm.acompletion(client=client, **kwargs) + return tuple([chunk async for chunk in response]) if stream == "fake" else (response,) + + +def _completion_chunks( + client: HTTPHandler, kwargs: Mapping[str, object], stream: Literal["off", "fake", "real"] +) -> tuple[object, ...]: + response: Final = litellm.completion(client=client, **kwargs) + return tuple(response) if stream == "fake" else (response,) + + +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("stream", ["off", "fake", "real"], ids=["completion", "fake_stream", "stream"]) +async def test_bedrock_strict_tools_limit_rejection_is_sent_again_without_strict_under_drop_params( + asynchronous: bool, stream: Literal["off", "fake", "real"], local_model_cost_map: None +) -> None: + sent, error = await _converse_with_a_strict_tools_limit_rejection_first(asynchronous, stream, drop_params=True) + + assert error is None + (first_path, first_body), (second_path, second_body) = sent + first_tool_spec: Final = first_body["toolConfig"]["tools"][0]["toolSpec"] + assert first_tool_spec["strict"] is True + assert first_path == second_path == ("converse-stream" if stream == "real" else "converse") + assert second_body == { + **first_body, + "toolConfig": { + **first_body["toolConfig"], + "tools": [{"toolSpec": {key: value for key, value in first_tool_spec.items() if key != "strict"}}], + }, + } + + +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_bedrock_strict_tools_limit_rejection_reaches_the_caller_without_drop_params( + asynchronous: bool, local_model_cost_map: None +) -> None: + sent, error = await _converse_with_a_strict_tools_limit_rejection_first(asynchronous, stream="off", drop_params=None) + + assert error is not None and "Too many strict tools" in str(error) + assert len(sent) == 1 + + +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_bedrock_strict_tools_limit_rejection_is_not_sent_again_when_no_tool_sent_strict( + asynchronous: bool, local_model_cost_map: None +) -> None: + sent, error = await _converse_with_a_strict_tools_limit_rejection_first( + asynchronous, + stream="off", + drop_params=True, + tool={"type": "function", "function": {**_LOOKUP_FUNCTION, "strict": False}}, + ) + + assert error is not None and "Too many strict tools" in str(error) + assert len(sent) == 1 From 5a18e6b6da67f273c14fcbd47000145294ef6a40 Mon Sep 17 00:00:00 2001 From: zoveki <349028329+zoveki@users.noreply.github.com> Date: Wed, 28 Oct 2026 04:13:03 -0700 Subject: [PATCH 018/128] test(files): audit #40621 (#48788) * test(files): audit #40621 Pin the raw cloud storage file id gate on /v1/files content and delete against a running proxy: twice-encoded gs:// ids stay blocked for non-admin keys, proxy admin keys read s3:// objects, and ordinary provider ids keep working * test(files): check the full SigV4 signature on the admin raw s3 read --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../security/test_files_batches_access.py | 126 +++++++++++++++--- 1 file changed, 104 insertions(+), 22 deletions(-) diff --git a/tests/integration/security/test_files_batches_access.py b/tests/integration/security/test_files_batches_access.py index 7b6fd4652cf..119055c64c7 100644 --- a/tests/integration/security/test_files_batches_access.py +++ b/tests/integration/security/test_files_batches_access.py @@ -48,6 +48,7 @@ VERTEX_ACCESS_TOKEN: Final = "ya29.scripted-delete-token" AWS_ACCESS_KEY_ID: Final = "AKIAIOSFODNN7EXAMPLE" AWS_SECRET_ACCESS_KEY: Final = "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" AWS_REGION: Final = "us-east-1" +S3_OBJECT_BODY: Final = b'{"recordId":"r1","modelOutput":{"stop_reason":"end_turn"}}\n' _ForbiddenSurface = Literal[ "upload_sdk_extra_body", "upload_form", @@ -708,6 +709,86 @@ def test_non_admin_cannot_delete_a_raw_gcs_uri_before_vertex_auth_or_storage(gat assert token_requests == (), "non-admin raw GCS delete reached the Vertex token endpoint" +def test_non_admin_still_reads_an_ordinary_provider_file_id(gateway: Gateway) -> None: + with wire_server(_openai_backend()) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="provider-key-" + uuid.uuid4().hex[:8] + ) + _wait_until_every_worker_serves(gateway, model) + member: Final = _member(scenario, model) + file_id: Final = "file-" + uuid.uuid4().hex + + read: Final = gateway.request("GET", f"/v1/files/{file_id}/content", key=member.key, params={"model": model}) + assert read.status_code == 200, read.text + assert read.content == b"file-content\n", read.content + assert [ + (request.method, urlsplit(request.target).path) for request in _drained_other_than_model_list_probes(wire) + ] == [("GET", f"/v1/files/{file_id}/content")] + + +@pytest.mark.parametrize( + ("method", "route_suffix", "verb"), + [("GET", "/content", "retrieved"), ("DELETE", "", "deleted")], + ids=["content", "delete"], +) +def test_non_admin_cannot_reach_vertex_with_a_twice_encoded_raw_gcs_uri( + gateway: Gateway, method: Literal["GET", "DELETE"], route_suffix: str, verb: str +) -> None: + with ( + wire_server(_vertex_token_backend) as token_wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model=VERTEX_MODEL, + api_key=None, + vertex_project=VERTEX_PROJECT, + vertex_location=VERTEX_LOCATION, + vertex_credentials=_vertex_credentials(token_wire.url), + gcs_bucket_name=BUCKET, + ) + _wait_until_every_worker_serves(gateway, model) + caller_user: Final = scenario.user(user_role="internal_user") + caller_key: Final = scenario.key(user_id=caller_user, models=[model]) + object_uri: Final = f"gs://{BUCKET}/litellm-vertex-files/litellm-vertex-files-obj.jsonl" + twice_encoded: Final = quote(quote(object_uri, safe=""), safe="") + token_wire.drain() + + denied: Final = gateway.request( + method, f"/v1/files/{twice_encoded}{route_suffix}", key=caller_key, params={"model": model} + ) + token_requests: Final = token_wire.drain() + assert denied.status_code == 403, f"{denied.text}; token requests={token_requests}" + assert _error_message(denied, 403) == ( + f"Raw cloud storage file ids can only be {verb} by a proxy admin key. " + "Use the LiteLLM managed file id returned when the file was created." + ), denied.text + assert token_requests == (), f"non-admin twice-encoded raw GCS {method} reached the Vertex token endpoint" + + +def _assert_signed_empty_body_s3_request(request: Request, method: Literal["GET", "DELETE"], wire_url: str) -> None: + assert request.method == method, request.method + request_date: Final = request.headers["x-amz-date"] + content_hash: Final = hashlib.sha256(b"").hexdigest() + assert request.headers["x-amz-content-sha256"] == content_hash, request.headers + aws_request: Final = AWSRequest( + method=method, + url=wire_url + request.target, + headers={ + "host": request.headers["host"], + "x-amz-date": request_date, + "x-amz-content-sha256": content_hash, + }, + ) + aws_request.context["timestamp"] = request_date + auth: Final = S3SigV4Auth(Credentials(AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY), "s3", AWS_REGION) + signature: Final = auth.signature(auth.string_to_sign(aws_request, auth.canonical_request(aws_request)), aws_request) + expected_authorization: Final = ( + f"AWS4-HMAC-SHA256 Credential={AWS_ACCESS_KEY_ID}/{request_date[:8]}/{AWS_REGION}/s3/aws4_request, " + f"SignedHeaders={auth.signed_headers(auth.headers_to_sign(aws_request))}, Signature={signature}" + ) + assert request.headers["authorization"] == expected_authorization, request.headers["authorization"] + + def test_proxy_admin_deletes_a_raw_s3_uri_with_one_signed_delete(gateway: Gateway) -> None: with wire_server(_s3_backend()) as wire, gateway.scenario() as scenario: model: Final = _bedrock_model(scenario, wire) @@ -721,29 +802,30 @@ def test_proxy_admin_deletes_a_raw_s3_uri_with_one_signed_delete(gateway: Gatewa assert parsed.deleted is True and parsed.id == object_uri, deleted.text (request,) = _drained_other_than_model_list_probes(wire) assert request.target == f"/{BUCKET}/litellm-bedrock-files-obj.jsonl", request.target - request_date: Final = request.headers["x-amz-date"] - content_hash: Final = hashlib.sha256(b"").hexdigest() - assert request.headers["x-amz-content-sha256"] == content_hash, request.headers - aws_request: Final = AWSRequest( - method="DELETE", - url=wire.url + request.target, - headers={ - "host": request.headers["host"], - "x-amz-date": request_date, - "x-amz-content-sha256": content_hash, - }, + _assert_signed_empty_body_s3_request(request, "DELETE", wire.url) + + +def _s3_object_backend(request: Request) -> Reply: + if request.method == "GET": + return Reply(body=S3_OBJECT_BODY, content_type="application/octet-stream") + return _json_reply({"error": {"message": f"unscripted {request.method} {request.target}"}}, 404) + + +def test_proxy_admin_reads_a_raw_s3_uri_with_one_signed_get(gateway: Gateway) -> None: + with wire_server(_s3_object_backend) as wire, gateway.scenario() as scenario: + model: Final = _bedrock_model(scenario, wire) + _wait_until_every_worker_serves(gateway, model) + object_key: Final = f"litellm-batch-outputs/{uuid.uuid4().hex}.jsonl.out" + object_uri: Final = f"s3://{BUCKET}/{object_key}" + + read: Final = gateway.request( + "GET", "/v1/files/" + quote(object_uri, safe="") + "/content", params={"model": model} ) - aws_request.context["timestamp"] = request_date - auth: Final = S3SigV4Auth(Credentials(AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY), "s3", AWS_REGION) - signature: Final = auth.signature( - auth.string_to_sign(aws_request, auth.canonical_request(aws_request)), aws_request - ) - expected_authorization: Final = ( - f"AWS4-HMAC-SHA256 Credential={AWS_ACCESS_KEY_ID}/{request_date[:8]}/{AWS_REGION}/s3/aws4_request, " - f"SignedHeaders={auth.signed_headers(auth.headers_to_sign(aws_request))}, Signature={signature}" - ) - assert request.headers["authorization"] == expected_authorization, request.headers["authorization"] - assert request.method == "DELETE", request.method + storage_requests: Final = _drained_other_than_model_list_probes(wire) + assert read.status_code == 200, f"{read.text}; storage requests={storage_requests}" + assert read.content == S3_OBJECT_BODY, read.content + assert [request.target for request in storage_requests] == [f"/{BUCKET}/{object_key}"], storage_requests + _assert_signed_empty_body_s3_request(storage_requests[0], "GET", wire.url) def _granting_vertex_token_backend(request: Request) -> Reply: From 53c6939ad406b501d75d2f3cb13f0e1fc2ca4c5d Mon Sep 17 00:00:00 2001 From: zoveki <349028329+zoveki@users.noreply.github.com> Date: Wed, 28 Oct 2026 04:24:50 -0700 Subject: [PATCH 019/128] test(anthropic, bedrock): audit #47489 (#48782) * test(anthropic, bedrock): audit #47489 thinking replay across Claude routes and Bedrock Converse * test(bedrock): compare the whole replayed Converse turn in the merged-message and cache point cells --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- ...bedrock_converse_redacted_thinking_wire.py | 91 +++++++- ...est_claude_thinking_content_replay_wire.py | 195 ++++++++++++++++++ ...k_converse_redacted_thinking_router_sdk.py | 39 ++++ 3 files changed, 324 insertions(+), 1 deletion(-) create mode 100644 tests/integration/providers/test_claude_thinking_content_replay_wire.py diff --git a/tests/integration/providers/test_bedrock_converse_redacted_thinking_wire.py b/tests/integration/providers/test_bedrock_converse_redacted_thinking_wire.py index 79ec0545c44..efb0609d560 100644 --- a/tests/integration/providers/test_bedrock_converse_redacted_thinking_wire.py +++ b/tests/integration/providers/test_bedrock_converse_redacted_thinking_wire.py @@ -6,7 +6,7 @@ from typing import Final from urllib.parse import unquote import pytest -from integration._support.client import Gateway +from integration._support.client import Gateway, Scenario from integration._support.upstream import _aws_event_frame from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import BaseModel, ConfigDict, JsonValue @@ -414,3 +414,92 @@ def test_converse_redacted_thinking_replays_to_the_anthropic_api_as_claude_nativ {"type": "redacted_thinking", "data": NATIVE_REDACTED}, {"type": "text", "text": FIRST_ANSWER}, ) + + +_SIGNED_THINKING: Final[dict[str, JsonValue]] = { + "type": "thinking", + "thinking": "Checking the forecast.", + "signature": "scripted-signature-" + "s" * 32, +} +_CONVERSE_REDACTED_PART: Final[dict[str, JsonValue]] = {"reasoningContent": {"redactedContent": CONVERSE_REDACTED}} +_CONVERSE_SIGNED_PART: Final[dict[str, JsonValue]] = { + "reasoningContent": { + "reasoningText": {"text": _SIGNED_THINKING["thinking"], "signature": _SIGNED_THINKING["signature"]} + } +} + + +def _replayed_to_converse( + gateway: Gateway, wire: Wire, model: str, assistants: list[JsonValue], **fields: JsonValue +) -> tuple[dict[str, JsonValue], ...]: + reply: Final = _chat( + gateway, + { + "model": model, + "thinking": THINKING, + "max_tokens": 2048, + "messages": [{"role": "user", "content": QUESTION}, *assistants, {"role": "user", "content": FOLLOW_UP}], + **fields, + }, + ) + assert reply.content == FINAL_ANSWER + return _replayed_assistant_turn(wire).content + + +def _claude_on_converse(scenario: Scenario, wire: Wire) -> str: + return scenario.model( + model=f"bedrock/converse/{CLAUDE_MODEL_ID}", + api_key=TOKEN, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire.url, + ) + + +@pytest.mark.parametrize( + "thinking_blocks", [None, [_redacted(NATIVE_REDACTED)]], ids=["content", "content-and-thinking-blocks"] +) +def test_claude_redacted_thinking_replayed_in_content_reaches_converse_once( + gateway: Gateway, thinking_blocks: list[JsonValue] | None +) -> None: + assistant: Final[dict[str, JsonValue]] = { + "role": "assistant", + "content": [_redacted(NATIVE_REDACTED), {"type": "text", "text": FIRST_ANSWER}], + **({} if thinking_blocks is None else {"thinking_blocks": thinking_blocks}), + } + with wire_server(_bedrock_peer()) as wire, gateway.scenario() as scenario: + sent: Final = _replayed_to_converse(gateway, wire, _claude_on_converse(scenario, wire), [assistant]) + assert sent == (_CONVERSE_REDACTED_PART, {"text": FIRST_ANSWER}) + + +@pytest.mark.parametrize( + "block, converse_part", + [(_redacted(NATIVE_REDACTED), _CONVERSE_REDACTED_PART), (_SIGNED_THINKING, _CONVERSE_SIGNED_PART)], + ids=["redacted", "signed"], +) +def test_block_split_across_merged_assistant_messages_reaches_converse_once( + gateway: Gateway, block: dict[str, JsonValue], converse_part: dict[str, JsonValue] +) -> None: + assistants: Final[list[JsonValue]] = [ + {"role": "assistant", "content": [block, {"type": "text", "text": FIRST_ANSWER}]}, + {"role": "assistant", "content": "Checking again.", "thinking_blocks": [block]}, + ] + with wire_server(_bedrock_peer()) as wire, gateway.scenario() as scenario: + sent: Final = _replayed_to_converse(gateway, wire, _claude_on_converse(scenario, wire), assistants) + assert sent == (converse_part, {"text": FIRST_ANSWER}, {"text": "Checking again."}) + + +@pytest.mark.parametrize("drop", [True, False], ids=["drop-params", "no-drop-params"]) +def test_converse_cache_point_after_replayed_thinking_follows_drop_params(gateway: Gateway, drop: bool) -> None: + assistant: Final[dict[str, JsonValue]] = { + "role": "assistant", + "content": [ + {**_redacted(NATIVE_REDACTED), "cache_control": {"type": "ephemeral"}}, + {"type": "text", "text": FIRST_ANSWER}, + ], + } + with wire_server(_bedrock_peer()) as wire, gateway.scenario() as scenario: + sent: Final = _replayed_to_converse( + gateway, wire, _claude_on_converse(scenario, wire), [assistant], drop_params=drop + ) + replayed: Final = (_CONVERSE_REDACTED_PART, {"text": FIRST_ANSWER}) + assert sent == (replayed if drop else (*replayed, {"cachePoint": {"type": "default"}})) diff --git a/tests/integration/providers/test_claude_thinking_content_replay_wire.py b/tests/integration/providers/test_claude_thinking_content_replay_wire.py new file mode 100644 index 00000000000..4f2f3cc66c4 --- /dev/null +++ b/tests/integration/providers/test_claude_thinking_content_replay_wire.py @@ -0,0 +1,195 @@ +import uuid +from collections.abc import Mapping +from typing import Final, Literal + +import pytest +from integration._support.anthropic_thinking import ( + BEDROCK_MODEL, + MODEL, + NO_CACHE, + SIGNATURE, + THINKING, + answer, + marker_of, + prompt, + standard_peer, + thinking_block, +) +from integration._support.client import Gateway, Scenario +from integration._support.vertex import service_account_json +from integration._support.wire import Wire, wire_server +from pydantic import BaseModel, ConfigDict, JsonValue + +Route = Literal["anthropic", "bedrock_invoke", "claude_platform", "vertex_ai", "azure_ai"] +Endpoint = Literal["chat", "messages"] + +_VERTEX_PROJECT: Final = "scripted-project" +_VERTEX_LOCATION: Final = "us-east5" +_ENABLED: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": 1024} +_EPHEMERAL: Final[dict[str, JsonValue]] = {"type": "ephemeral"} +_SIGNED: Final = thinking_block(THINKING, SIGNATURE) +_REDACTED: Final[dict[str, JsonValue]] = {"type": "redacted_thinking", "data": "scripted-redacted-" + "r" * 32} +_BLANK_REDACTED: Final[dict[str, JsonValue]] = {"type": "redacted_thinking", "data": " "} +_ROUTES: Final[tuple[Route, ...]] = ("anthropic", "bedrock_invoke", "claude_platform", "vertex_ai", "azure_ai") + + +class _Frozen(BaseModel): + model_config = ConfigDict(frozen=True) + + +class _SentMessage(_Frozen): + role: str + content: str | tuple[dict[str, JsonValue], ...] + + +class _SentRequest(_Frozen): + messages: tuple[_SentMessage, ...] + + +def _deployment(scenario: Scenario, route: Route, wire_url: str, upstream_url: str) -> str: + match route: + case "anthropic": + return scenario.model(model=f"anthropic/{MODEL}", api_base=wire_url, api_key="scripted-key") + case "bedrock_invoke": + return scenario.model( + model=f"bedrock/invoke/{BEDROCK_MODEL}", + api_base=wire_url, + aws_access_key_id="AKIASCRIPTEDPROVIDER", + aws_secret_access_key="scripted-secret", + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire_url, + ) + case "claude_platform": + return scenario.model( + model=f"bedrock/claude_platform/{MODEL}", + api_base=wire_url, + api_key="scripted-platform-key", + aws_region_name="us-east-1", + workspace_id="scripted-workspace", + ) + case "vertex_ai": + return scenario.model( + model=f"vertex_ai/{MODEL}", + api_base=( + f"{wire_url}/v1/projects/{_VERTEX_PROJECT}/locations/{_VERTEX_LOCATION}" + f"/publishers/anthropic/models/{MODEL}" + ), + api_key=None, + vertex_project=_VERTEX_PROJECT, + vertex_location=_VERTEX_LOCATION, + vertex_credentials=service_account_json(_VERTEX_PROJECT, upstream_url.rstrip("/")), + ) + case "azure_ai": + return scenario.model(model=f"azure_ai/{MODEL}", api_base=wire_url, api_key="scripted-azure-key") + + +def _text(text: str) -> dict[str, JsonValue]: + return {"type": "text", "text": text} + + +def _replayed_turn( + gateway: Gateway, + wire: Wire, + *, + endpoint: Endpoint = "chat", + model: str, + marker: str, + assistant: Mapping[str, JsonValue], + **fields: JsonValue, +) -> tuple[dict[str, JsonValue], ...]: + follow_up: Final = uuid.uuid4().hex + body: Final[dict[str, JsonValue]] = { + "model": model, + "max_tokens": 2048, + "thinking": _ENABLED, + "messages": [ + {"role": "user", "content": prompt(marker)}, + {"role": "assistant", **assistant}, + {"role": "user", "content": prompt(follow_up)}, + ], + **NO_CACHE, + **fields, + } + path: Final = "/v1/chat/completions" if endpoint == "chat" else "/v1/messages" + response: Final = gateway.request("POST", path, body) + assert response.status_code == 200, response.text + assert answer(follow_up) in response.text, response.text + received: Final = wire.drain() + assert len(received) == 1, [request.body for request in received] + assert marker_of(received[0]) == follow_up, received[0].body + sent: Final = _SentRequest.model_validate_json(received[0].body).messages[1] + assert sent.role == "assistant", received[0].body + assert not isinstance(sent.content, str), received[0].body + return sent.content + + +@pytest.mark.parametrize("route", _ROUTES) +def test_redacted_thinking_replayed_in_assistant_content_reaches_claude_once(gateway: Gateway, route: Route) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(standard_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, route, wire.url, gateway.upstream_url) + sent: Final = _replayed_turn( + gateway, wire, model=model, marker=marker, assistant={"content": [_REDACTED, _text(answer(marker))]} + ) + assert sent == (_REDACTED, _text(answer(marker))) + + +def test_block_in_content_and_thinking_blocks_reaches_claude_once_in_order(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(standard_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", wire.url, gateway.upstream_url) + sent: Final = _replayed_turn( + gateway, + wire, + model=model, + marker=marker, + assistant={"content": [_REDACTED, _text(answer(marker))], "thinking_blocks": [_SIGNED, _REDACTED]}, + ) + assert sent == (_SIGNED, _REDACTED, _text(answer(marker))) + + +def test_streamed_thinking_replayed_in_content_reaches_claude_without_its_index(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(standard_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", wire.url, gateway.upstream_url) + sent: Final = _replayed_turn( + gateway, + wire, + model=model, + marker=marker, + assistant={"content": [{**_SIGNED, "index": 0}, _text(answer(marker))]}, + ) + assert sent == (_SIGNED, _text(answer(marker))) + + +@pytest.mark.parametrize("drop", [True, False], ids=["drop-params", "no-drop-params"]) +def test_replayed_thinking_cache_control_and_blank_data_follow_drop_params(gateway: Gateway, drop: bool) -> None: + marker: Final = uuid.uuid4().hex + content: Final = ({**_SIGNED, "cache_control": _EPHEMERAL}, _BLANK_REDACTED, _REDACTED, _text(answer(marker))) + with wire_server(standard_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", wire.url, gateway.upstream_url) + sent: Final = _replayed_turn( + gateway, wire, model=model, marker=marker, assistant={"content": list(content)}, drop_params=drop + ) + assert sent == ((_SIGNED, _REDACTED, _text(answer(marker))) if drop else content) + + +@pytest.mark.parametrize("endpoint", ["chat", "messages"]) +def test_assistant_cache_injection_point_never_marks_a_thinking_part(gateway: Gateway, endpoint: Endpoint) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(standard_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{MODEL}", + api_base=wire.url, + api_key="scripted-key", + cache_control_injection_points=[{"location": "message", "role": "assistant"}], + ) + sent: Final = _replayed_turn( + gateway, + wire, + endpoint=endpoint, + model=model, + marker=marker, + assistant={"content": [_text(answer(marker)), _SIGNED]}, + ) + assert sent == ({**_text(answer(marker)), "cache_control": _EPHEMERAL}, _SIGNED) diff --git a/tests/integration/sdk/test_bedrock_converse_redacted_thinking_router_sdk.py b/tests/integration/sdk/test_bedrock_converse_redacted_thinking_router_sdk.py index 6b2f4cb8eea..2988f9b5f53 100644 --- a/tests/integration/sdk/test_bedrock_converse_redacted_thinking_router_sdk.py +++ b/tests/integration/sdk/test_bedrock_converse_redacted_thinking_router_sdk.py @@ -88,3 +88,42 @@ def test_sync_router_reads_base_model_to_replay_redacted_thinking_to_an_applicat assistant: Final = next(message for message in received[-1].messages if message.role == "assistant") redacted: Final = ({"reasoningContent": {"redactedContent": REDACTED}},) if replays_redacted else () assert assistant.content == (*redacted, {"text": FIRST_ANSWER}) + + +def test_sync_router_replays_a_block_in_content_and_thinking_blocks_once() -> None: + redacted: Final[dict[str, JsonValue]] = {"type": "redacted_thinking", "data": REDACTED} + with wire_server(_respond) as wire: + router: Final = Router( + model_list=[ + { + "model_name": "claude-profile", + "litellm_params": { + "model": f"bedrock/converse/{APPLICATION_PROFILE_ARN}", + "api_base": wire.url, + "api_key": "synthetic-bedrock-bearer", + "aws_region_name": "us-east-1", + }, + "model_info": {"base_model": "bedrock/us.anthropic.claude-sonnet-4-6"}, + } + ], + num_retries=0, + ) + + response: Final = router.completion( + model="claude-profile", + max_tokens=256, + messages=[ + {"role": "user", "content": "What's the weather in Paris?"}, + { + "role": "assistant", + "content": [redacted, {"type": "text", "text": FIRST_ANSWER}], + "thinking_blocks": [redacted], + }, + {"role": "user", "content": "Reply with the word ok."}, + ], + ) + + assert response.choices[0].message.content == FINAL_ANSWER + received: Final = tuple(_ConverseRequest.model_validate_json(request.body) for request in wire.drain()) + assistant: Final = next(message for message in received[-1].messages if message.role == "assistant") + assert assistant.content == ({"reasoningContent": {"redactedContent": REDACTED}}, {"text": FIRST_ANSWER}) From 8033f311f27e9ad5f0776d9eed39a0f484fd9e92 Mon Sep 17 00:00:00 2001 From: norkelbyte <348963847+norkelbyte@users.noreply.github.com> Date: Wed, 28 Oct 2026 04:27:58 -0700 Subject: [PATCH 020/128] fix(anthropic): let drop_params drop speed on claude-opus-4-7 after Anthropic removed its fast mode (#48764) * fix(anthropic): let drop_params drop speed on claude-opus-4-7 after Anthropic removed its fast mode * test(anthropic): cite the Claude Code fast-mode page for the Opus 4.7 removal date * test(anthropic): cover the drop_params speed drop for claude-opus-4-7 * test(anthropic): type the model parameter of the speed drop test --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 2 -- model_prices_and_context_window.json | 2 -- .../chat/test_anthropic_chat_transformation.py | 13 ++++++++----- 3 files changed, 8 insertions(+), 9 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 8207c3122f9..9c3bd9dacbb 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -15764,7 +15764,6 @@ "us": 1.1 }, "supports_output_config": true, - "supports_speed": true, "prompt_cache_min_tokens": 2048, "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, @@ -15807,7 +15806,6 @@ "us": 1.1 }, "supports_output_config": true, - "supports_speed": true, "prompt_cache_min_tokens": 2048 }, "claude-fable-5": { diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8207c3122f9..9c3bd9dacbb 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -15764,7 +15764,6 @@ "us": 1.1 }, "supports_output_config": true, - "supports_speed": true, "prompt_cache_min_tokens": 2048, "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, @@ -15807,7 +15806,6 @@ "us": 1.1 }, "supports_output_config": true, - "supports_speed": true, "prompt_cache_min_tokens": 2048 }, "claude-fable-5": { diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py index ef002a73e52..4d5d6cc17a9 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -2378,10 +2378,8 @@ def test_anthropic_model_supports_effort_param_rejects_non_supporting_models(mod "model", [ "claude-opus-4-6", - "claude-opus-4-7", "claude-opus-4-8", "claude-opus-4-6-20260205", - "claude-opus-4-7-20260416", ], ) def test_anthropic_model_supports_speed_param_recognizes_supporting_models(model): @@ -2397,6 +2395,10 @@ def test_anthropic_model_supports_speed_param_recognizes_supporting_models(model "vertex_ai/claude-opus-4-8", "azure_ai/claude-opus-4-8", "anthropic.claude-opus-4-8", + # Anthropic removed Opus 4.7 fast mode on 2026-07-24 and answers 400 for `speed` on it + # (code.claude.com/docs/en/fast-mode, platform.claude.com/docs/en/build-with-claude/fast-mode, read 2026-10-28) + "claude-opus-4-7", + "claude-opus-4-7-20260416", ], ) def test_anthropic_model_supports_speed_param_rejects_non_supporting_models(model): @@ -4677,8 +4679,9 @@ def test_fast_mode_parameter_mapping(): assert result["speed"] == "fast" -def test_anthropic_drop_params_strips_speed_for_unsupported_models(): - """``drop_params=True`` strips unsupported ``speed`` for non-Opus models.""" +@pytest.mark.parametrize("model", ["claude-sonnet-4-6", "claude-opus-4-7"]) +def test_anthropic_drop_params_strips_speed_for_unsupported_models(model: str): + """``drop_params=True`` strips ``speed`` for models without fast mode.""" config = AnthropicConfig() messages = [{"role": "user", "content": "Hello"}] @@ -4686,7 +4689,7 @@ def test_anthropic_drop_params_strips_speed_for_unsupported_models(): litellm.drop_params = True try: result = config.transform_request( - model="claude-sonnet-4-6", + model=model, messages=messages, optional_params={"speed": "fast", "max_tokens": 1024}, litellm_params={}, From e838957b1d194b8ff01587f341f07052f8bdbf24 Mon Sep 17 00:00:00 2001 From: simonp9 <349053457+simonp9@users.noreply.github.com> Date: Wed, 28 Oct 2026 04:30:19 -0700 Subject: [PATCH 021/128] test(token_counter): audit #47586 (#48800) * test(token_counter): audit #47586 * test(token_counter): check streamed usage equals the local token count --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- ...st_huggingface_tokenizer_downloads_wire.py | 201 ++++++++++++++++++ 1 file changed, 201 insertions(+) create mode 100644 tests/integration/providers/test_huggingface_tokenizer_downloads_wire.py diff --git a/tests/integration/providers/test_huggingface_tokenizer_downloads_wire.py b/tests/integration/providers/test_huggingface_tokenizer_downloads_wire.py new file mode 100644 index 00000000000..ca1bfb5f78a --- /dev/null +++ b/tests/integration/providers/test_huggingface_tokenizer_downloads_wire.py @@ -0,0 +1,201 @@ +import json +import threading +import uuid +from collections.abc import Callable, Iterator, Mapping +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import pytest +from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter +from tokenizers import Regex, Tokenizer, models, pre_tokenizers + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_LLAMA_3_BACKEND: Final = "databricks-meta-llama-3-3-70b-instruct" +_DEEPSEEK_V4: Final = "azure_ai/DeepSeek-V4-Flash-0731" +_API_KEY: Final = "synthetic-databricks-key" +_PROMPT: Final = "Count this message before the tokenizer comes back" +_COUNTS_AFTER_FAILURE: Final = 4 +_TOKENIZER_FILE: Final = "/resolve/main/tokenizer.json" +_TOKENIZER_HEADERS: Final = MappingProxyType({"X-Repo-Commit": "0" * 40, "ETag": '"scripted-char-tokenizer"'}) +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +class _Usage(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + prompt_tokens: int + + +class _Chunk(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + id: str + usage: _Usage | None = None + + +def _hf_environment(directory: Path) -> dict[str, str]: + return { + "HF_HOME": str(directory / "hf-home"), + "HF_HUB_CACHE": str(directory / "hf-cache"), + "HF_HUB_OFFLINE": "0", + } + + +def _missing(request: Request) -> Reply: + return Reply(status=404, body=json.dumps({"error": f"unscripted target {request.target}"}).encode()) + + +@pytest.fixture(scope="module") +def unreachable_hub_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[OwnedProxy, Wire]]: + directory: Final = tmp_path_factory.mktemp("hf-tokenizer-unreachable") + with ( + gateway_from_environment() as gateway, + wire_server(_missing) as hub, + owned_proxy_process(gateway, directory, {**_hf_environment(directory), "HF_ENDPOINT": hub.url}) as owned, + ): + yield owned, hub + + +def _llama_3(scenario: Scenario, api_base: str) -> str: + return scenario.model(model=f"databricks/{_LLAMA_3_BACKEND}", api_base=api_base, api_key=_API_KEY) + + +def _local_count(gateway: Gateway, model: str) -> tuple[JsonValue, JsonValue]: + response: Final = gateway.request( + "POST", + "/utils/token_counter", + {"model": model, "messages": [{"role": "user", "content": _PROMPT}]}, + params={"call_endpoint": "false"}, + ) + assert response.status_code == 200, response.text + payload: Final = _JSON_OBJECT.validate_json(response.content) + return payload["tokenizer_type"], payload["total_tokens"] + + +def _frame(identity: str, delta: Mapping[str, JsonValue], finish_reason: str | None) -> bytes: + chunk: Final = { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": _LLAMA_3_BACKEND, + "choices": [{"index": 0, "delta": dict(delta), "finish_reason": finish_reason}], + } + return b"data: " + json.dumps(chunk).encode() + b"\n\n" + + +def _stream_without_usage(identity: str) -> Callable[[Request], Reply]: + frames: Final = ( + _frame(identity, {"role": "assistant", "content": "Counted "}, None), + _frame(identity, {"content": "locally."}, None), + _frame(identity, {}, "stop"), + b"data: [DONE]\n\n", + ) + + def respond(_request: Request) -> Reply: + return Reply(content_type="text/event-stream", chunks=frames) + + return respond + + +def test_counts_after_a_failed_llama_3_tokenizer_download_never_reach_the_hub_again( + unreachable_hub_proxy: tuple[OwnedProxy, Wire], +) -> None: + owned, hub = unreachable_hub_proxy + with owned.gateway.scenario() as scenario: + model: Final = _llama_3(scenario, owned.gateway.upstream_url) + first: Final = _local_count(owned.gateway, model) + attempted: Final = hub.received.qsize() + later: Final = tuple(_local_count(owned.gateway, model) for _ in range(_COUNTS_AFTER_FAILURE)) + retried: Final = hub.received.qsize() - attempted + assert attempted > 0, "the proxy never asked the scripted hub for the llama-3 tokenizer" + assert first[0] == "openai_tokenizer", first + assert later == (first,) * _COUNTS_AFTER_FAILURE, later + assert retried == 0, f"{retried} hub requests after the download already failed" + + +def test_streamed_llama_3_chat_after_a_failed_download_counts_usage_without_reaching_the_hub( + unreachable_hub_proxy: tuple[OwnedProxy, Wire], +) -> None: + owned, hub = unreachable_hub_proxy + identity: Final = f"chatcmpl-{uuid.uuid4().hex}" + with wire_server(_stream_without_usage(identity)) as upstream, owned.gateway.scenario() as scenario: + model: Final = _llama_3(scenario, upstream.url) + counted: Final = _local_count(owned.gateway, model) + attempted: Final = hub.received.qsize() + with owned.gateway.client.stream( + "POST", + "/v1/chat/completions", + json={ + "model": model, + "messages": [{"role": "user", "content": _PROMPT}], + "stream": True, + "stream_options": {"include_usage": True}, + }, + headers={"Authorization": f"Bearer {owned.gateway.key}"}, + ) as response: + assert response.status_code == 200, response.read() + lines: Final = tuple(line for line in response.iter_lines() if line.startswith("data: ")) + retried: Final = hub.received.qsize() - attempted + sent: Final = tuple((request.method, request.target) for request in upstream.drain()) + assert counted[0] == "openai_tokenizer", counted + assert lines[-1] == "data: [DONE]", lines + chunks: Final = tuple(_Chunk.model_validate_json(line.removeprefix("data: ")) for line in lines[:-1]) + assert {chunk.id for chunk in chunks} == {identity}, lines + usages: Final = tuple(chunk.usage for chunk in chunks if chunk.usage is not None) + assert tuple(usage.prompt_tokens for usage in usages) == (counted[1],), (counted, lines) + assert sent == (("POST", "/chat/completions"),), sent + assert retried == 0, f"{retried} hub requests while streaming after the download already failed" + + +def test_a_failed_llama_3_download_leaves_the_deepseek_tokenizer_counting( + unreachable_hub_proxy: tuple[OwnedProxy, Wire], +) -> None: + owned, _hub = unreachable_hub_proxy + with owned.gateway.scenario() as scenario: + llama: Final = _llama_3(scenario, owned.gateway.upstream_url) + deepseek: Final = scenario.model(model=_DEEPSEEK_V4, api_base=owned.gateway.upstream_url, api_key=_API_KEY) + assert _local_count(owned.gateway, llama)[0] == "openai_tokenizer" + counted: Final = _local_count(owned.gateway, deepseek) + assert counted[0] == "huggingface_tokenizer", counted + + +def _char_tokenizer() -> bytes: + tokenizer: Final = Tokenizer(models.WordLevel({"[UNK]": 0}, unk_token="[UNK]")) + tokenizer.pre_tokenizer = pre_tokenizers.Split(Regex("."), behavior="isolated") + return tokenizer.to_str().encode() + + +def _recovering_hub(recovered: threading.Event, tokenizer: bytes) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if recovered.is_set() and request.target.endswith(_TOKENIZER_FILE): + return Reply(body=tokenizer, headers=_TOKENIZER_HEADERS) + return _missing(request) + + return respond + + +def test_a_failed_llama_3_download_is_retried_once_the_retry_window_passes(gateway: Gateway, tmp_path: Path) -> None: + recovered: Final = threading.Event() + with ( + wire_server(_recovering_hub(recovered, _char_tokenizer())) as hub, + owned_proxy_process( + gateway, + tmp_path, + {**_hf_environment(tmp_path), "HF_ENDPOINT": hub.url, "HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS": "1"}, + ) as owned, + owned.gateway.scenario() as scenario, + ): + model: Final = _llama_3(scenario, owned.gateway.upstream_url) + failed: Final = _local_count(owned.gateway, model) + hub.drain() + recovered.set() + counted: Final = eventually( + lambda: _local_count(owned.gateway, model), lambda count: count[0] == "huggingface_tokenizer", seconds=30 + ) + downloads: Final = tuple(request.target for request in hub.drain() if request.target.endswith(_TOKENIZER_FILE)) + assert failed[0] == "openai_tokenizer", failed + assert counted[1] != failed[1], (failed, counted) + assert downloads, "the tokenizer came back without a new download" From 7f44cff2745d65ac3b3f11e58f08c721579b6474 Mon Sep 17 00:00:00 2001 From: paveln83 <349044532+paveln83@users.noreply.github.com> Date: Wed, 28 Oct 2026 04:34:21 -0700 Subject: [PATCH 022/128] test(anthropic): audit #47478 (#48803) Covers the json_schema envelope edge cases #47478 fixed on the Anthropic json tool path: a key listed only in `required`, a tool-named wrapper on a stream, non-ASCII and lone-surrogate answers, a stream cut before the tool block closes (proxy and sync SDK, Anthropic and Bedrock Invoke), and a string `output_format` with and without `drop_params` on Anthropic and Vertex chat and /v1/messages. The scripted Anthropic peer now refuses a non-object `output_format` the way Anthropic does Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- ...vertex_anthropic_structured_output_wire.py | 188 +++++++++++++++++- ...anthropic_json_tool_cut_stream_sync_sdk.py | 146 ++++++++++++++ 2 files changed, 323 insertions(+), 11 deletions(-) create mode 100644 tests/integration/sdk/test_anthropic_json_tool_cut_stream_sync_sdk.py diff --git a/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py index 2918f364c56..4e31122440d 100644 --- a/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py +++ b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py @@ -60,6 +60,8 @@ _OWNED_MODEL: Final = "partner-claude" _DEPRECATED_OUTPUT_FORMAT: Final = "output_format: this field is deprecated, use output_config.format" _BAD_OUTPUT_CONFIG: Final = "output_config: Input should be a valid dictionary" _BAD_FORMAT: Final = "output_config.format: Input should be a json_schema object" +_BAD_OUTPUT_FORMAT: Final = "output_format: Input should be a valid dictionary" +_MISTYPED_OUTPUT_FORMAT: Final = "json_schema" _SCHEMA: Final[dict[str, JsonValue]] = { "type": "object", "properties": {"tier": {"type": "integer"}, "summary": {"type": "string"}}, @@ -94,6 +96,11 @@ _VALUES_RESPONSE_FORMAT: Final[dict[str, JsonValue]] = { "type": "json_schema", "json_schema": {"name": "valued", "schema": _VALUES_SCHEMA}, } +_REQUIRED_ONLY_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"tier": {"type": "integer"}}, + "required": ["payload"], +} _ADDITIONAL_OBJECT_SCHEMA: Final[dict[str, JsonValue]] = { **_OPTIONAL_SCHEMA, "additionalProperties": {"type": "object"}, @@ -146,6 +153,7 @@ _STREAM_SHAPES: Final = ( "number_first", "bool_first", "null_first", + "named_wrapper", ) _HELD_STREAM_SHAPES: Final = frozenset( { @@ -166,12 +174,15 @@ _HELD_STREAM_SHAPES: Final = frozenset( "named_encoded", "undeclared_encoded", "encoded_deep", + "named_wrapper", } ) _AS_WRITTEN_SHAPES: Final = ("label_non_ascii", "label_two_objects_non_ascii") _TEXT_ANSWER_SHAPES: Final = frozenset({"undeclared_wrapper"}) _LABEL: Final = "RiskAssessment" _NON_ASCII_LABEL: Final = "Évaluation des risques 风险" +_NON_ASCII_SUMMARY: Final = "Hôtel 東京" +_LONE_SURROGATE: Final = "\ud800" _ENCODED_DEPTH: Final = 200_000 _FAMILY: Final = ( "anthropic_unmapped", @@ -181,6 +192,9 @@ _FAMILY: Final = ( "claude_platform_unmapped", "vertex_haiku", ) +_CUT_FAMILY: Final = ("anthropic_unmapped", "bedrock_invoke") +_CUT_EVENTS: Final = 3 +_OUTPUT_FORMAT_ROUTES: Final = ("anthropic_chat", "anthropic_messages", "vertex_chat", "vertex_messages") _FRAGMENT_SIZE: Final = 20 _TIMING_BYTES: Final = 512 * 1024 _TIMING_FRAGMENT: Final = 4096 @@ -228,9 +242,9 @@ def _nest(payload: JsonValue, depth: int) -> JsonValue: def _tool_input(shape: str, marker_hex: str) -> JsonValue: payload: Final = _payload(marker_hex) match shape: - case "plain": + case "plain" | "cut_plain": return payload - case "wrapped": + case "wrapped" | "cut_wrapped": return {"tool_name": _JSON_TOOL, "parameters": payload} case "named": return {"name": _JSON_TOOL, "parameters": payload} @@ -256,6 +270,14 @@ def _tool_input(shape: str, marker_hex: str) -> JsonValue: return {"tool_name": "other_tool", "parameters": payload} case "named_input": return {"name": _JSON_TOOL, "input": payload} + case "named_wrapper": + return {"name": _JSON_TOOL, "arguments": {"input": payload}} + case "required_only": + return {"payload": payload} + case "wrapped_non_ascii": + return {"tool_name": _JSON_TOOL, "parameters": _payload(f"{_NON_ASCII_SUMMARY} {marker_hex}")} + case "wrapped_lone_surrogate": + return {"tool_name": _JSON_TOOL, "parameters": _payload(marker_hex + _LONE_SURROGATE)} case "named_two_objects": return {"tool_name": _JSON_TOOL, "parameters": payload, "extra": {"note": marker_hex}} case "named_string": @@ -348,8 +370,13 @@ def _expected(shape: str, marker_hex: str) -> JsonValue: | "parameter_label" | "function_label" | "named_encoded" + | "named_wrapper" + | "cut_plain" + | "cut_wrapped" ): return _payload(marker_hex) + case "wrapped_non_ascii" | "wrapped_lone_surrogate": + return object_value(_tool_input(shape, marker_hex))["parameters"] case "large": return {**_payload(marker_hex), "summary": marker_hex + "x" * 5000} case "named_partial" | "inner_incomplete": @@ -423,9 +450,14 @@ def _aws_chunk(data: Mapping[str, JsonValue]) -> bytes: def _stream_reply( - message: dict[str, JsonValue], block: dict[str, JsonValue], deltas: tuple[dict[str, JsonValue], ...], framing: str + message: dict[str, JsonValue], + block: dict[str, JsonValue], + deltas: tuple[dict[str, JsonValue], ...], + framing: str, + cut: bool = False, ) -> Reply: - events: Final = _events(message, block, deltas) + complete: Final = _events(message, block, deltas) + events: Final = complete[:-_CUT_EVENTS] if cut else complete if framing == "aws": return Reply(content_type=_EVENT_STREAM, chunks=tuple(_aws_chunk(data) for _, data in events)) return Reply( @@ -445,6 +477,8 @@ def _error(status: int, message: str) -> Reply: def _format_fault(body: Mapping[str, JsonValue], vertex: bool) -> str | None: if vertex and "output_format" in body: return _DEPRECATED_OUTPUT_FORMAT + if body.get("output_format") is not None and not isinstance(body["output_format"], dict): + return _BAD_OUTPUT_FORMAT output_config: Final = body.get("output_config") if output_config is None: return None @@ -486,7 +520,7 @@ def _tool_reply(message_id: str, shape: str, marker_hex: str, user_tool: bool, f deltas: Final = tuple( {"type": "input_json_delta", "partial_json": piece} for piece in _fragments(partial, _fragment_size(shape)) ) - return _stream_reply(message, {**json_block, "input": {}}, deltas, framing) + return _stream_reply(message, {**json_block, "input": {}}, deltas, framing, cut=shape.startswith("cut_")) def _text_reply(message_id: str, text: str, framing: str | None) -> Reply: @@ -748,7 +782,9 @@ def _messages_stream_text(response: httpx.Response) -> JsonValue: return json.loads("".join(string_value(delta["text"]) for delta in deltas if delta["type"] == "text_delta")) -def _vertex_deployment(gateway: Gateway, scenario: Scenario, api_base: str, backend: str) -> str: +def _vertex_deployment( + gateway: Gateway, scenario: Scenario, api_base: str, backend: str, **litellm_params: JsonValue +) -> str: return scenario.model( model=f"vertex_ai/{backend}", api_base=api_base, @@ -756,6 +792,7 @@ def _vertex_deployment(gateway: Gateway, scenario: Scenario, api_base: str, back vertex_project=_PROJECT, vertex_location=_LOCATION, vertex_credentials=service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")), + **litellm_params, ) @@ -814,6 +851,33 @@ def _family_deployment(gateway: Gateway, scenario: Scenario, api_base: str, memb raise AssertionError(member) +def _output_format_deployment( + gateway: Gateway, scenario: Scenario, api_base: str, route: str, **litellm_params: JsonValue +) -> str: + if route.startswith("vertex_"): + return _vertex_deployment(gateway, scenario, api_base, _NATIVE_MODEL, **litellm_params) + return scenario.model( + model=f"anthropic/{_NATIVE_MODEL}", api_base=api_base, api_key=_ANTHROPIC_KEY, **litellm_params + ) + + +def _output_format_call( + gateway: Gateway, model: str, route: str, marker: str, output_format: JsonValue +) -> httpx.Response: + if route.endswith("_chat"): + return gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": marker}], "output_format": output_format}, + ) + return gateway.request( + "POST", + "/v1/messages", + _messages_body(model, marker, output_format=output_format), + headers={"anthropic-beta": _STRUCTURED_OUTPUTS_BETA}, + ) + + def _owned_config(path: Path, gateway: Gateway, api_base: str) -> Path: config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) path.write_text( @@ -1188,13 +1252,18 @@ def test_vertex_chat_json_schema_without_required_keys_unwraps_only_the_tool_nam @pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +@pytest.mark.parametrize( + ("shape", "schema"), + [("values", _VALUES_SCHEMA), ("required_only", _REQUIRED_ONLY_SCHEMA)], + ids=["declared_property", "required_only"], +) def test_vertex_chat_json_schema_that_declares_the_wrapper_key_keeps_the_answer_whole( - gateway: Gateway, stream: bool + gateway: Gateway, shape: str, schema: dict[str, JsonValue], stream: bool ) -> None: - marker, marker_hex = _marker("values") + marker, marker_hex = _marker(shape) response_format: Final[dict[str, JsonValue]] = { "type": "json_schema", - "json_schema": {"name": "valued", "schema": _VALUES_SCHEMA}, + "json_schema": {"name": "valued", "schema": schema}, } with wire_server(_peer) as wire, gateway.scenario() as scenario: model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) @@ -1202,9 +1271,9 @@ def test_vertex_chat_json_schema_that_declares_the_wrapper_key_keeps_the_answer_ "POST", "/v1/chat/completions", _chat_body(model, marker, stream=stream, response_format=response_format) ) content: Final = _stream_content(response) if stream else _chat_content(response) - assert content == _tool_input("values", marker_hex), response.text + assert content == _tool_input(shape, marker_hex), response.text if stream: - assert len(_content_frames(response)) >= _fragment_count("values", marker_hex), response.text + assert len(_content_frames(response)) >= _fragment_count(shape, marker_hex), response.text _one_delivery(wire, marker_hex) @@ -1265,6 +1334,22 @@ def test_vertex_chat_json_schema_unwraps_a_tool_named_envelope_only_around_one_o _one_delivery(wire, marker_hex) +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +@pytest.mark.parametrize("shape", ["wrapped_non_ascii", "wrapped_lone_surrogate"]) +def test_vertex_chat_json_schema_answer_keeps_non_ascii_text_and_escapes_a_lone_surrogate( + gateway: Gateway, shape: str, stream: bool +) -> None: + marker, marker_hex = _marker(shape) + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker, stream=stream)) + text: Final = _stream_text(response) if stream else string_value(_chat_message(response)["content"]) + assert json.loads(text) == _expected(shape, marker_hex), response.text + if shape == "wrapped_non_ascii": + assert _NON_ASCII_SUMMARY in text, text + _one_delivery(wire, marker_hex) + + @pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) @pytest.mark.parametrize("shape", ["name_property", "name_value_plain"]) def test_vertex_chat_json_schema_with_a_name_property_unwraps_only_the_tool_named_wrapper( @@ -1350,6 +1435,19 @@ def test_tool_path_family_stream_unwraps_the_envelope(gateway: Gateway, member: assert _json_tool_schema(sent)["properties"] == _SCHEMA["properties"], sent.body +@pytest.mark.parametrize("shape", ["plain", "cut_plain", "cut_wrapped"]) +@pytest.mark.parametrize("member", _CUT_FAMILY) +def test_tool_path_stream_cut_before_the_tool_block_closes_still_delivers_the_answer( + gateway: Gateway, member: str, shape: str +) -> None: + marker, marker_hex = _marker(shape) + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _family_deployment(gateway, scenario, wire.url, member) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker, stream=True)) + assert _stream_content(response) == _payload(marker_hex), response.text + _one_delivery(wire, marker_hex) + + @pytest.mark.parametrize("accepts_forced_tools", [True, False], ids=["accepts_forced_tools", "rejects_forced_tools"]) def test_vertex_chat_json_schema_forces_the_json_tool_only_when_the_model_accepts_forced_tools( gateway: Gateway, accepts_forced_tools: bool @@ -1865,6 +1963,74 @@ def test_vertex_messages_duplicate_output_format_keys_fold_the_last_value(gatewa assert _folded_format(_one_delivery(wire, marker_hex)) == _OUTPUT_FORMAT +@pytest.mark.parametrize("route", _OUTPUT_FORMAT_ROUTES) +def test_string_output_format_under_deployment_drop_params_is_dropped_before_the_peer( + gateway: Gateway, route: str +) -> None: + marker, marker_hex = _marker("plain_text") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _output_format_deployment(gateway, scenario, wire.url, route, drop_params=True) + response: Final = _output_format_call(gateway, model, route, marker, _MISTYPED_OUTPUT_FORMAT) + assert response.status_code == 200, response.text + assert f"plain text {marker_hex}" in response.text, response.text + sent: Final = _one_delivery(wire, marker_hex) + assert not {"output_format", "output_config"} & sent.body.keys(), sent.body + + +@pytest.mark.parametrize("route", ["anthropic_chat", "anthropic_messages"]) +def test_object_output_format_under_deployment_drop_params_still_reaches_anthropic( + gateway: Gateway, route: str +) -> None: + marker, marker_hex = _marker("native") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _output_format_deployment(gateway, scenario, wire.url, route, drop_params=True) + response: Final = _output_format_call(gateway, model, route, marker, _OUTPUT_FORMAT) + assert response.status_code == 200, response.text + assert _one_delivery(wire, marker_hex).body["output_format"] == _OUTPUT_FORMAT + + +def test_vertex_messages_string_output_format_under_drop_params_keeps_the_callers_output_config_format( + gateway: Gateway, +) -> None: + marker, marker_hex = _marker("native") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL, drop_params=True) + response: Final = gateway.request( + "POST", + "/v1/messages", + _messages_body( + model, marker, output_format=_MISTYPED_OUTPUT_FORMAT, output_config={"format": _OUTPUT_FORMAT} + ), + ) + assert _messages_text(response) == _payload(marker_hex), response.text + assert _folded_format(_one_delivery(wire, marker_hex)) == _OUTPUT_FORMAT + + +@pytest.mark.parametrize("route", ["anthropic_chat", "anthropic_messages"]) +def test_string_output_format_without_drop_params_reaches_anthropic_and_returns_its_400( + gateway: Gateway, route: str +) -> None: + marker, marker_hex = _marker("plain_text") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _output_format_deployment(gateway, scenario, wire.url, route) + response: Final = _output_format_call(gateway, model, route, marker, _MISTYPED_OUTPUT_FORMAT) + assert response.status_code == 400, response.text + assert _BAD_OUTPUT_FORMAT in response.text, response.text + assert _one_delivery(wire, marker_hex).body["output_format"] == _MISTYPED_OUTPUT_FORMAT + + +def test_vertex_chat_string_output_format_without_drop_params_is_refused_without_reaching_the_peer( + gateway: Gateway, +) -> None: + marker, _ = _marker("plain_text") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _output_format_deployment(gateway, scenario, wire.url, "vertex_chat") + response: Final = _output_format_call(gateway, model, "vertex_chat", marker, _MISTYPED_OUTPUT_FORMAT) + assert response.status_code == 400, response.text + assert "needs `output_format` to be a JSON object" in response.text, response.text + assert _delivered(wire) == () + + @pytest.mark.parametrize( "schema", [ diff --git a/tests/integration/sdk/test_anthropic_json_tool_cut_stream_sync_sdk.py b/tests/integration/sdk/test_anthropic_json_tool_cut_stream_sync_sdk.py new file mode 100644 index 00000000000..da091b1c12e --- /dev/null +++ b/tests/integration/sdk/test_anthropic_json_tool_cut_stream_sync_sdk.py @@ -0,0 +1,146 @@ +import base64 +import json +import uuid +from collections.abc import Callable, Mapping +from typing import Final + +import pytest +from integration._support.upstream import aws_event_stream_frame +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +import litellm +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + +_JSON_TOOL: Final = "json_tool_call" +_FRAGMENT_SIZE: Final = 20 +_CUT_EVENTS: Final = 3 +_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"tier": {"type": "integer"}, "summary": {"type": "string"}}, + "required": ["tier", "summary"], + "additionalProperties": False, +} +_RESPONSE_FORMAT: Final[dict[str, JsonValue]] = { + "type": "json_schema", + "json_schema": {"name": "risk", "schema": _SCHEMA}, +} +_ANTHROPIC_MODEL: Final = "anthropic/claude-scripted-9" +_BEDROCK_INVOKE_MODEL: Final = "bedrock/invoke/anthropic.claude-sonnet-5" + + +def _payload(marker_hex: str) -> dict[str, JsonValue]: + return {"tier": 2, "summary": marker_hex} + + +def _tool_input(shape: str, marker_hex: str) -> dict[str, JsonValue]: + if shape == "wrapped": + return {"tool_name": _JSON_TOOL, "parameters": _payload(marker_hex)} + return _payload(marker_hex) + + +def _events(marker_hex: str, shape: str, cut: bool) -> tuple[tuple[str, dict[str, JsonValue]], ...]: + text: Final = json.dumps(_tool_input(shape, marker_hex)) + deltas: Final = tuple( + ( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": text[start : start + _FRAGMENT_SIZE]}, + }, + ) + for start in range(0, len(text), _FRAGMENT_SIZE) + ) + message: Final[dict[str, JsonValue]] = { + "id": f"msg_{marker_hex}", + "type": "message", + "role": "assistant", + "model": "claude-scripted-9", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 1}, + } + complete: Final = ( + ("message_start", {"type": "message_start", "message": message}), + ( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "tool_use", "id": f"toolu_{marker_hex}", "name": _JSON_TOOL, "input": {}}, + }, + ), + *deltas, + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "tool_use", "stop_sequence": None}, + "usage": {"output_tokens": 3}, + }, + ), + ("message_stop", {"type": "message_stop"}), + ) + return complete[:-_CUT_EVENTS] if cut else complete + + +def _aws_chunk(data: Mapping[str, JsonValue]) -> bytes: + payload: Final = json.dumps({"bytes": base64.b64encode(json.dumps(data).encode()).decode()}).encode() + return aws_event_stream_frame( + {":event-type": "chunk", ":content-type": "application/json", ":message-type": "event"}, payload + ) + + +def _peer(marker_hex: str, shape: str, cut: bool) -> Callable[[Request], Reply]: + events: Final = _events(marker_hex, shape, cut) + + def respond(request: Request) -> Reply: + assert marker_hex in request.body.decode(), request.body + if request.target.endswith("/invoke-with-response-stream"): + return Reply( + content_type="application/vnd.amazon.eventstream", + chunks=tuple(_aws_chunk(data) for _, data in events), + ) + assert request.target.endswith("/v1/messages"), request.target + return Reply( + content_type="text/event-stream", + chunks=tuple(f"event: {name}\ndata: {json.dumps(data)}\n\n".encode() for name, data in events), + ) + + return respond + + +def _credentials(model: str) -> dict[str, str]: + if model == _ANTHROPIC_MODEL: + return {"api_key": "synthetic-anthropic-key"} + return {"api_key": "synthetic-bedrock-bearer", "aws_region_name": "us-east-1"} + + +@pytest.mark.parametrize("model", [_ANTHROPIC_MODEL, _BEDROCK_INVOKE_MODEL], ids=["anthropic", "bedrock_invoke"]) +@pytest.mark.parametrize( + ("shape", "cut"), + [("plain", False), ("plain", True), ("wrapped", True)], + ids=["plain_complete", "plain_cut", "wrapped_cut"], +) +def test_sync_stream_cut_before_the_json_tool_block_closes_still_yields_the_whole_answer( + model: str, shape: str, cut: bool +) -> None: + marker_hex: Final = uuid.uuid4().hex + with wire_server(_peer(marker_hex, shape, cut)) as wire: + response: Final = litellm.completion( + model=model, + messages=[{"role": "user", "content": f"Rate the risk {marker_hex}"}], + response_format=_RESPONSE_FORMAT, + max_tokens=64, + stream=True, + api_base=wire.url, + num_retries=0, + **_credentials(model), + ) + assert isinstance(response, CustomStreamWrapper), type(response) + text: Final = "".join(str(chunk.choices[0].delta.content or "") for chunk in tuple(response)) + assert len(wire.drain()) == 1 + assert json.loads(text) == _payload(marker_hex), text From 8e2e99670d2a8af603b63e975a42347559e74858 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 28 Oct 2026 07:34:40 -0400 Subject: [PATCH 023/128] fix(router): continue Responses mid-stream fallbacks on models without assistant prefill (#47672) * fix(router): avoid assistant prefill in Responses fallbacks * fix(router): keep prefill for Claude fallbacks that take it and skip constrained continuations A Responses fallback to a Claude model that takes assistant prefill now gets the partial response as prefill again, so it continues mid-word without the prompt-shape drift. Every other fallback keeps the request to continue. A stream cut after text no longer continues when the request constrains its output (a json_schema or json_object text format, or a forced tool_choice), since a continued JSON document or tool call would not parse. The client gets the primary's error, as on the chat path. The fallback's re-sent leading space at the seam is dropped, and the prefill shaping runs per deployment after the encrypted reasoning strip, so stripped reasoning is never replayed. * fix(router): guard SDK text_format, narrow forced tool_choice, read Claude base_model, and cap Responses continuation output Cut Responses streams that asked for structured output through the SDK's text_format now surface the primary's error instead of stitching two models' JSON. Only a required or named tool_choice counts as forced, so allowed_tools in auto mode and hosted tool choices still continue. A Claude deployment named only through base_model keeps its prefill continuation, and the continuation's max_output_tokens is lowered by what the cut stream already emitted. * fix(router): quote the partial's ending in the continue request and surface the primary's error after a cut tool call * fix(router): stop Responses fallbacks from continuing content-filter cuts and client-run tool calls A content filter that cuts a /v1/responses stream mid-answer now returns its error, as on chat, instead of letting a content_policy_fallbacks backup finish the filtered text. Shell, apply_patch, computer, local_shell, and MCP approval items count as a started tool call. Chat continuation goes back to treating any tool_choice object as constraining, as on main * fix(responses): count a stream of tool arguments as a started tool call * fix(router): let Responses fallbacks continue under allowed_tools auto and after server-run tool calls A Responses tool_choice of allowed_tools in auto mode no longer blocks the continuation, since the model may still answer in text. Hosted MCP calls and the bridge's server-run tool search no longer count as a started client tool call. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../streaming_iterator.py | 1 + litellm/responses/streaming_iterator.py | 27 +- litellm/router.py | 141 ++-- .../continuation_prefill_check.py | 149 ++++- .../router_code_coverage.py | 2 + .../base_llm/responses/test_transformation.py | 605 +++++++++++++++++- .../test_continuation_prefill_check.py | 250 +++++++- tests/unit/test_router/test_router.py | 13 +- 8 files changed, 1049 insertions(+), 139 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index b5ef9e95741..219944de521 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -423,6 +423,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if self._is_tool_search_call(call_id, fn_name): self._queue_tool_search_call_added(call_id, output_index, fn_args_delta) continue + self._started_tool_call = True if call_id not in self._tool_args_by_call_id: self._tool_args_by_call_id[call_id] = "" diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 50ab3e4f75a..c1f422c7d7b 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -365,6 +365,7 @@ class BaseResponsesAPIStreamingIterator: """ _pending_logging_tasks: tuple[asyncio.Task[object], ...] = () + _started_tool_call: bool = False _earlier_rounds_spend: BilledRoundsSpend | None = None def __init__( @@ -488,12 +489,16 @@ class BaseResponsesAPIStreamingIterator: _args_delta: Final = getattr(openai_responses_api_chunk, "delta", None) if isinstance(_args_delta, str): self._generated_tool_arguments += _args_delta + if _event_type in _CLIENT_RUN_TOOL_ARGUMENTS_DELTA_EVENTS: + self._started_tool_call = True _stream_model_id: Final = _model_id_from_metadata(self.litellm_metadata) if _event_type in ( ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, ): _item: Final[object] = getattr(openai_responses_api_chunk, "item", None) + if getattr(_item, "type", None) in _CLIENT_RUN_TOOL_CALL_ITEM_TYPES: + self._started_tool_call = True if _item is not None: ResponsesAPIRequestUtils.encode_container_id_on_output_item( item=_item, @@ -1044,6 +1049,11 @@ class BaseResponsesAPIStreamingIterator: if pending: await asyncio.wait(pending) + @property + def started_tool_call(self) -> bool: + """Whether the stream began a tool call, which a text continuation could never finish.""" + return self._started_tool_call + def partial_response_on_disconnect(self) -> ResponseIncompleteEvent | None: """ The response.incomplete event to bill when the client went away mid-stream, @@ -1690,13 +1700,26 @@ def _usage_as_model(usage: object) -> ResponseAPIUsage | None: return None -_TOOL_ARGUMENTS_DELTA_EVENTS: Final = frozenset( +_CLIENT_RUN_TOOL_CALL_ITEM_TYPES: Final = frozenset( + { + "function_call", + "custom_tool_call", + "mcp_approval_request", + "computer_call", + "local_shell_call", + "shell_call", + "apply_patch_call", + } +) +_CLIENT_RUN_TOOL_ARGUMENTS_DELTA_EVENTS: Final = frozenset( { ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA, ResponsesAPIStreamEvents.CUSTOM_TOOL_CALL_INPUT_DELTA, - ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DELTA, } ) +_TOOL_ARGUMENTS_DELTA_EVENTS: Final = _CLIENT_RUN_TOOL_ARGUMENTS_DELTA_EVENTS | { + ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DELTA +} def _estimate_usage_from_text( diff --git a/litellm/router.py b/litellm/router.py index e53bab4659e..f0bb6c3a7a9 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -305,14 +305,20 @@ from litellm.router_utils.kubernetes_pod_discovery import ( ) from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( MID_STREAM_CONTINUATION_MARKER, + RESPONSES_CONTINUATION_KWARG, ContinuationOutput, ContinuationPrefillDeploymentCheck, ContinuedPartial, + ResponsesContinuation, + ResponsesContinuationSeam, constrains_continuation, + constrains_responses_continuation, continuation_deployments, continuation_output_ceilings, continued_text, deployment_can_continue, + responses_continuation_input, + responses_continuation_input_for, ) from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, @@ -457,6 +463,7 @@ from litellm.utils import ( provider_rejectable_params, resolve_api_key_secret_reference, set_live_deployment_replay, + token_counter, ) from .router_utils.pattern_match_deployments import PatternMatchRouter @@ -488,7 +495,6 @@ if TYPE_CHECKING: from litellm.types.llms.openai import ( ResponseAPIUsage, ResponseIncompleteEvent, - ResponseInputParam, ResponsesAPIResponse, ) @@ -3865,7 +3871,8 @@ class Router: ) -> Mapping[str, object]: """The request as the deployment received it, merged in the order _update_kwargs_with_deployment applies.""" deployment_params: Final = deployment.litellm_params.model_dump(exclude_none=True) if deployment else {} - defaults: Final = {key: value for key, value in self.default_litellm_params.items() if value is not None} + default_params: Final = _REQUEST_KWARGS_ADAPTER.validate_python(self.default_litellm_params) + defaults: Final = {key: value for key, value in default_params.items() if value is not None} deployment_tool_choice: Final = ( {"tool_choice": deployment_params["tool_choice"]} if deployment_params.get("tool_choice") is not None @@ -3944,67 +3951,6 @@ class Router: extended: Final = _extend_trailing_text_block(blocks, new_text) return [*messages[:-1], {**last, "content": extended, "prefix": True}] - @staticmethod - def _build_responses_continuation_input( - input_val: Union[str, "ResponseInputParam"] | None, - generated_content: str, - ) -> "ResponseInputParam": - """ - Convert Responses-API input + partial assistant output into a - continuation input that asks the fallback model to pick up where the - prior assistant message stopped. - - Best effort across providers. The chat-completions path uses - Anthropic's `prefix: True` prefill trick on the assistant message; - the Responses-API input schema has no direct equivalent, so we - append an instruction (developer role) plus a prior assistant - message containing the partial output. Providers without prefill - semantics (OpenAI, Vertex) treat this as conversational context - and may regenerate — same trade-off as the chat-completions path - for non-Anthropic fallbacks. - """ - # base/continuation are List[Any] because ResponseInputParam items - # are a wide Union of TypedDicts (EasyInputMessageParam, Message, - # ResponseOutputMessageParam, ...) — annotating as List[Dict[str, Any]] - # rejects the list() spread of input_val. We cast the combined list to - # ResponseInputParam at the return. - base: list[object] - if isinstance(input_val, str): - base = [ - { - "type": "message", - "role": "user", - "content": [{"type": "input_text", "text": input_val}], - } - ] - elif isinstance(input_val, list): - base = list(input_val) - else: - base = [] - continuation: Final[list[object]] = [ - { - "type": "message", - "role": "developer", - "content": [ - { - "type": "input_text", - "text": ( - "The previous assistant response was interrupted " - "mid-stream. Continue exactly where it stopped — " - "do not repeat any of its content. Your response " - "must read as a seamless continuation." - ), - } - ], - }, - { - "type": "message", - "role": "assistant", - "content": [{"type": "output_text", "text": generated_content}], - }, - ] - return cast("ResponseInputParam", base + continuation) - async def _aresponses_streaming_iterator( self, response: "BaseResponsesAPIStreamingIterator", @@ -4026,9 +3972,12 @@ class Router: Full parity with the chat-completions path: - Pre-first-chunk: retry with the original input unchanged. - - Partial content: inject a developer instruction + prior - assistant message carrying the generated text so the fallback - model continues rather than restarts. + - Partial content: continue from the generated text, as an + assistant prefill on a fallback that takes one and as assistant + history plus a user request to continue on any other. A request + whose output is constrained (a JSON schema, a forced tool), or a + stream that began a tool call, is not continued, and the primary's + error surfaces instead. - Usage combining: merge partial-stream usage onto the fallback's response.completed event so accounting reflects both attempts. - Stream cleanup: shielded aclose() on both source and fallback @@ -4217,6 +4166,42 @@ class Router: wrapper: Final = FallbackResponsesStreamWrapper(stream_with_fallbacks()) return wrapper + def _responses_continuation( + self, + e: "MidStreamFallbackError", + source_iterator: "BaseResponsesAPIStreamingIterator", + request_kwargs: Mapping[str, object], + ) -> ResponsesContinuation | None: + """What a fallback continues from once a Responses stream broke after text, raising the primary's error + when the request constrains its output or the stream began a tool call, since a continued JSON document or + tool call would not parse, and when a content filter cut it, so no other model finishes filtered text.""" + if not e.generated_content or e.is_pre_first_chunk: + return None + if source_iterator.started_tool_call or isinstance(e.original_exception, litellm.ContentPolicyViolationError): + Router._raise_original_mid_stream_error(e) + cut_deployment_id: Final = (get_hidden_params(source_iterator) or {}).get("model_id") + cut_deployment: Final = ( + self.get_deployment(model_id=cut_deployment_id) if isinstance(cut_deployment_id, str) else None + ) + if constrains_responses_continuation(self._params_sent_to_deployment(request_kwargs, cut_deployment)): + Router._raise_original_mid_stream_error(e) + earlier: Final = request_kwargs.get(RESPONSES_CONTINUATION_KWARG) + if isinstance(earlier, ResponsesContinuation): + return earlier.extended(e.generated_content) + return ResponsesContinuation(request_kwargs.get("input"), e.generated_content) + + @staticmethod + def _responses_continuation_ceilings( + e: "MidStreamFallbackError", request_kwargs: Mapping[str, object] + ) -> Mapping[str, object]: + """The caller's output ceilings less what the cut stream already emitted, raising the primary's error once + nothing is left.""" + emitted_tokens: Final = token_counter(model=e.model or "", text=e.generated_content) + reduced: Final = continuation_output_ceilings(request_kwargs, emitted_tokens) + if reduced is None: + Router._raise_original_mid_stream_error(e) + return reduced + async def _aresponses_fallback_attempt( self, e: "MidStreamFallbackError", @@ -4235,6 +4220,11 @@ class Router: """ from litellm.exceptions import MidStreamFallbackError + continuation: Final = self._responses_continuation(e, source_iterator, initial_kwargs) + reduced_ceilings: Final[Mapping[str, object]] = ( + Router._responses_continuation_ceilings(e, initial_kwargs) if continuation is not None else {} + ) + seam: Final = ResponsesContinuationSeam(continuation.emitted_text if continuation is not None else "") partial_usage: Final = Router._extract_partial_responses_usage(source_iterator) fallback_response = None # rebind-ok: pre-init so finally can close it if a fallback was actually attempted fallback_yielded = False # rebind-ok: flipped on the first fallback item so a fallback that dies before its first event still replays the primary's held announcement @@ -4252,11 +4242,10 @@ class Router: initial_kwargs["original_function"] = ( # rebind-ok: the fallback chain re-enters on the same kwargs self._ageneric_api_call_with_fallbacks_responses_attempt ) - if e.generated_content and not e.is_pre_first_chunk: - initial_kwargs["input"] = Router._build_responses_continuation_input( # rebind-ok: fallback hop input - initial_kwargs.get("input"), - e.generated_content, - ) + if continuation is not None: + initial_kwargs.update(reduced_ceilings) + initial_kwargs[RESPONSES_CONTINUATION_KWARG] = continuation # rebind-ok: shaped per fallback deployment + initial_kwargs["input"] = responses_continuation_input(continuation) # rebind-ok: fallback hop input # The Responses-API path stores observability metadata # under "litellm_metadata" (not the default "metadata") — # see _ageneric_api_call_with_fallbacks. Mirroring that @@ -4285,7 +4274,7 @@ class Router: ) prepared_fallback_hidden_params: Final = adopt_headers(fallback_response) if hasattr(fallback_response, "__aiter__"): - async for fallback_item in fallback_response: + async for fallback_item in seam.splice(fallback_response): Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params) if partial_usage is not None: Router._combine_responses_fallback_usage(fallback_item, partial_usage) @@ -4808,11 +4797,12 @@ class Router: splits_stream_timeout: bool = False, ) -> None: """ - 4 jobs: + 5 jobs: - Adds selected deployment, model_info and api_base to kwargs["metadata"] (used for logging) - Adds default litellm params to kwargs, if set. - Merges tools from deployment with request (proxy-configured tools + request tools). - On a fallback hop, drops the encrypted reasoning this deployment cannot decrypt, keeping its summary. + - On a Responses continuation, sends the partial response as prefill where this deployment takes one. """ for key in self._forwarded_alias_marker_keys_the_deployment_sets( deployment=deployment, forwarded_keys=kwargs.pop(_ALIAS_MARKER_FORWARDED_PARAMS_KWARG, ()) @@ -5016,6 +5006,11 @@ class Router: (_FALLBACK_HOP_ADAPTER.validate_python(deployment),), unmarked_origin=deployment_that_just_failed, ) + if RESPONSES_CONTINUATION_KWARG in hop_kwargs: + del kwargs[RESPONSES_CONTINUATION_KWARG] # rebind-ok: this method prepares the call's kwargs in place + kwargs["input"] = responses_continuation_input_for( # rebind-ok: the deployment's continuation input + hop_kwargs, _FALLBACK_HOP_ADAPTER.validate_python(deployment) + ) def _get_async_openai_model_client(self, deployment: dict, kwargs: dict): """ diff --git a/litellm/router_utils/pre_call_checks/continuation_prefill_check.py b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py index 8f211ef67db..d2a46fc316d 100644 --- a/litellm/router_utils/pre_call_checks/continuation_prefill_check.py +++ b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py @@ -1,12 +1,13 @@ """Keeps a mid-stream continuation on deployments whose model supports assistant prefill and whose own params leave -the output unconstrained.""" +the output unconstrained, and shapes a Responses API continuation for the deployment that serves it.""" -from collections.abc import AsyncIterable, AsyncIterator, Mapping, Sequence +import json +from collections.abc import AsyncIterable, AsyncIterator, Callable, Mapping, Sequence from dataclasses import dataclass from itertools import accumulate from os.path import commonprefix from types import MappingProxyType -from typing import Final +from typing import Final, TypeVar from pydantic import TypeAdapter, ValidationError @@ -14,7 +15,7 @@ import litellm from litellm.integrations.custom_logger import CustomLogger, Span from litellm.litellm_core_utils.rules import Rules from litellm.router_utils.common_utils import MID_STREAM_CONTINUATION_KWARG -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.openai import AllMessageValues, OutputTextDeltaEvent from litellm.types.utils import ModelResponse, ModelResponseStream, Usage from litellm.utils import supports_assistant_prefill @@ -28,7 +29,9 @@ MID_STREAM_CONTINUATION_MARKER: Final = _ContinuationMarker() _STR_KEYED_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) _NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) _UNCONSTRAINED_RESPONSE_FORMAT: Final[Mapping[str, str]] = MappingProxyType({"type": "text"}) -_OUTPUT_CEILING_KEYS: Final = ("max_tokens", "max_completion_tokens") +_OUTPUT_CEILING_KEYS: Final = ("max_tokens", "max_completion_tokens", "max_output_tokens") +_QUOTED_ENDING_CHARS: Final = 100 +RESPONSES_CONTINUATION_KWARG: Final = "_responses_continuation" def _str_keyed(value: object) -> Mapping[str, object]: @@ -38,32 +41,55 @@ def _str_keyed(value: object) -> Mapping[str, object]: return _NO_FIELDS -def deployment_supports_prefill(deployment: object) -> bool: +def _underlying_model(deployment: object) -> str | None: deployment_map: Final = _str_keyed(deployment) model_info: Final = _str_keyed(deployment_map.get("model_info")) - declared: Final = model_info.get("supports_assistant_prefill") - if isinstance(declared, bool): - return declared litellm_params: Final = _str_keyed(deployment_map.get("litellm_params")) candidates: Final = (model_info.get("base_model"), litellm_params.get("base_model"), litellm_params.get("model")) - model: Final = next((candidate for candidate in candidates if isinstance(candidate, str) and candidate), None) + return next((candidate for candidate in candidates if isinstance(candidate, str) and candidate), None) + + +def deployment_supports_prefill(deployment: object) -> bool: + declared: Final = _str_keyed(_str_keyed(deployment).get("model_info")).get("supports_assistant_prefill") + if isinstance(declared, bool): + return declared + model: Final = _underlying_model(deployment) return model is not None and supports_assistant_prefill(model=model) -def constrains_continuation(params: Mapping[str, object]) -> bool: +def _forces_tool_call(tool_choice: object) -> bool: + return tool_choice == "required" or isinstance(tool_choice, Mapping) + + +def _forces_responses_tool_call(tool_choice: object) -> bool: + choice: Final = _str_keyed(tool_choice) + leaves_text_open: Final = choice.get("type") == "allowed_tools" and choice.get("mode") == "auto" + return _forces_tool_call(tool_choice) and not leaves_text_open + + +def _constrains(params: Mapping[str, object], forces_tool_call: Callable[[object], bool]) -> bool: sent: Final = {**params, **_str_keyed(params.get("extra_body"))} choice_count: Final = sent.get("n") response_format: Final = sent.get("response_format") - tool_choice: Final = sent.get("tool_choice") + text_format: Final = _str_keyed(sent.get("text")).get("format") return ( (isinstance(choice_count, int) and choice_count > 1) or (response_format is not None and response_format != _UNCONSTRAINED_RESPONSE_FORMAT) - or tool_choice == "required" - or isinstance(tool_choice, Mapping) + or (text_format is not None and text_format != _UNCONSTRAINED_RESPONSE_FORMAT) + or sent.get("text_format") is not None + or forces_tool_call(sent.get("tool_choice")) or sent.get("merge_reasoning_content_in_choices") is True ) +def constrains_continuation(params: Mapping[str, object]) -> bool: + return _constrains(params, _forces_tool_call) + + +def constrains_responses_continuation(params: Mapping[str, object]) -> bool: + return _constrains(params, _forces_responses_tool_call) + + def deployment_can_continue(deployment: object) -> bool: litellm_params: Final = _str_keyed(_str_keyed(deployment).get("litellm_params")) return deployment_supports_prefill(deployment) and not constrains_continuation(litellm_params) @@ -99,6 +125,7 @@ def _delta_text(chunk: ModelResponseStream) -> str: _CONTENT_BLOCKS_ADAPTER: Final = TypeAdapter(list[dict[str, object]]) +_INPUT_ITEMS_ADAPTER: Final = TypeAdapter(list[object]) def _prefill_text(prefill_message: Mapping[str, object]) -> str: @@ -114,7 +141,7 @@ def _prefill_text(prefill_message: Mapping[str, object]) -> str: ) -def _resendable_whitespace(prefill: str) -> str: +def resendable_whitespace(prefill: str) -> str: return prefill[len(prefill.rstrip()) :] if prefill.strip() else "" @@ -124,7 +151,7 @@ def continued_text(prefill_message: Mapping[str, object], generated_content: str return None echo: Final = len(prefill) if prefill and generated_content.startswith(prefill) else 0 after_echo: Final = generated_content[echo:] - return after_echo[len(commonprefix((after_echo, _resendable_whitespace(prefill)))) :] + return after_echo[len(commonprefix((after_echo, resendable_whitespace(prefill)))) :] def _reduced_ceilings(params: Mapping[str, object], emitted_tokens: int) -> Mapping[str, int]: @@ -211,7 +238,7 @@ class ContinuationOutput: self._rules: Final = Rules() self._held: tuple[ModelResponseStream, ...] = () self._echo_settled = not self._prefill_text - self._seam = _resendable_whitespace(self._prefill_text) + self._seam = resendable_whitespace(self._prefill_text) async def splice(self, stream: AsyncIterable[object]) -> AsyncIterator[object]: async for item in stream: @@ -262,3 +289,91 @@ class ContinuationOutput: self._text += _delta_text(chunk) self._rules.post_call_rules(input=self._text, model=self._model) return chunks + + +@dataclass(frozen=True, slots=True) +class ResponsesContinuation: + prior_input: object + emitted_text: str + + def extended(self, generated_content: str) -> "ResponsesContinuation": + return ResponsesContinuation(self.prior_input, f"{self.emitted_text}{generated_content}") + + +def deployment_takes_responses_prefill(deployment: object) -> bool: + """Claude takes a trailing assistant turn as prefill with no `prefix` flag, which a Responses input item cannot + carry.""" + model: Final = _underlying_model(deployment) + return model is not None and "claude" in model.lower() and deployment_supports_prefill(deployment) + + +def _continuation_request(emitted_text: str) -> Mapping[str, object]: + seam: Final = ( + "It already ends with whitespace, so do not add another space before your first word." + if emitted_text[-1:].isspace() + else "It ends without trailing whitespace, so begin with a space if your first word starts a new word." + ) + ending: Final = json.dumps(emitted_text[-_QUOTED_ENDING_CHARS:], ensure_ascii=False) + instruction: Final = ( + f"The previous assistant response was interrupted mid-stream. It ended with {ending}. Continue exactly where " + "it stopped, so your first characters are the ones that would have come right after it, and do not repeat " + f"any of its content. {seam} Your response must read as a seamless continuation." + ) + return {"type": "message", "role": "user", "content": [{"type": "input_text", "text": instruction}]} + + +def _input_items(request_input: object) -> tuple[object, ...]: + if isinstance(request_input, str): + return ({"type": "message", "role": "user", "content": [{"type": "input_text", "text": request_input}]},) + try: + return tuple(_INPUT_ITEMS_ADAPTER.validate_python(request_input)) + except ValidationError: + return () + + +def responses_continuation_input(continuation: ResponsesContinuation) -> Sequence[object]: + """The prior input, the partial response as assistant history, and a user request to continue it.""" + partial: Final = { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": continuation.emitted_text}], + } + return [*_input_items(continuation.prior_input), partial, _continuation_request(continuation.emitted_text)] + + +def responses_continuation_input_for(request_kwargs: Mapping[str, object], deployment: Mapping[str, object]) -> object: + """The continuation input as the deployment receives it, without the request to continue where the deployment + takes the partial response as prefill. It reads the input the router already stripped of reasoning this + deployment cannot decrypt, never the continuation's prior input.""" + request_input: Final = request_kwargs.get("input") + continuation: Final = request_kwargs.get(RESPONSES_CONTINUATION_KWARG) + if not isinstance(continuation, ResponsesContinuation): + return request_input + items: Final = _input_items(request_input) + if items[-1:] != (_continuation_request(continuation.emitted_text),) or not deployment_takes_responses_prefill( + deployment + ): + return request_input + return list(items[:-1]) + + +_StreamItem = TypeVar("_StreamItem") + + +class ResponsesContinuationSeam: + """Drops whitespace a fallback sends again at the start of its text when the client already has it, since + Anthropic strips a prefill's trailing whitespace and the model then writes it back.""" + + def __init__(self, emitted_text: str) -> None: + self._seam = resendable_whitespace(emitted_text) + + async def splice(self, stream: AsyncIterable[_StreamItem]) -> AsyncIterator[_StreamItem]: + async for item in stream: + yield self._past(item) + + def _past(self, item: _StreamItem) -> _StreamItem: + if not self._seam or not isinstance(item, OutputTextDeltaEvent): + return item + matched: Final = len(commonprefix((item.delta, self._seam))) + self._seam = self._seam[matched:] if matched == len(item.delta) else "" + return item.model_copy(update={"delta": item.delta[matched:]}) if matched else item diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 853b10ad97b..f4a295e697d 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -171,6 +171,8 @@ ignored_function_names = [ "_deployment_field", # Read by is_batch_capable, tested through the mixed-group tests in test_router.py "_deployment_provider", # Read by without_deployments_that_cannot_batch, tested through the mixed-group tests in test_router.py "_cannot_run_batches", # Tested through the Azure and OpenAI mixed-group acreate_batch tests in test_router.py + "_responses_continuation", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py + "_responses_continuation_ceilings", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py "_team_deployments_named", # Tested through the team lookups in the concurrent delete tests in test_router.py ] diff --git a/tests/unit/llms/base_llm/responses/test_transformation.py b/tests/unit/llms/base_llm/responses/test_transformation.py index 8d21af83668..50d918fa5be 100644 --- a/tests/unit/llms/base_llm/responses/test_transformation.py +++ b/tests/unit/llms/base_llm/responses/test_transformation.py @@ -1,14 +1,17 @@ """The shared Responses API config contract.""" import asyncio, importlib, json +from collections.abc import Mapping import httpx import pytest +import respx import litellm from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.types.router import GenericLiteLLMParams from litellm import Router +from pydantic import BaseModel, ConfigDict from litellm.constants import STREAM_SSE_DONE_STRING from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig @@ -879,38 +882,6 @@ def test_combine_responses_fallback_usage_passthrough_for_unknown_event(): partial = ResponseAPIUsage(input_tokens=1, output_tokens=1, total_tokens=2) Router._combine_responses_fallback_usage(other, partial) -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -def test_build_responses_continuation_input_from_string(): - out = Router._build_responses_continuation_input("Hello world", "partial assistant text") - assert len(out) == 3 - assert out[0]["role"] == "user" - assert out[0]["content"][0]["text"] == "Hello world" - assert out[1]["role"] == "developer" - assert out[2]["role"] == "assistant" - assert out[2]["content"][0]["text"] == "partial assistant text" - -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -def test_build_responses_continuation_input_from_list_preserves_items(): - existing: List[Any] = [ - { - "type": "message", - "role": "user", - "content": [{"type": "input_text", "text": "msg1"}], - } - ] - out = Router._build_responses_continuation_input(existing, "partial") - assert len(out) == 3 - assert out[0]["content"][0]["text"] == "msg1" - assert out[1]["role"] == "developer" - assert out[2]["role"] == "assistant" - -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -def test_build_responses_continuation_input_from_none(): - out = Router._build_responses_continuation_input(None, "partial") - assert len(out) == 2 - assert out[0]["role"] == "developer" - assert out[1]["role"] == "assistant" - @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") @pytest.mark.asyncio async def test_aresponses_streaming_iterator_passthrough(): @@ -1355,9 +1326,573 @@ async def test_aresponses_fallback_uses_continuation_input_after_partial_content continuation = mock_fallback.await_args.kwargs["kwargs"]["input"] assert isinstance(continuation, list) assert continuation[0]["content"][0]["text"] == "original question" - assert continuation[-2]["role"] == "developer" - assert continuation[-1]["role"] == "assistant" - assert continuation[-1]["content"][0]["text"] == "partial answer" + assert continuation[-2]["role"] == "assistant" + assert continuation[-2]["content"][0]["text"] == "partial answer" + assert continuation[-1]["role"] == "user" + + +def _sse(*events: Mapping[str, object]) -> httpx.Response: + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content="".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events), + ) + + +def _openai_stream_cut_after(partial: str) -> httpx.Response: + return _sse( + *( + ({ + "type": "response.output_text.delta", "sequence_number": 0, + "item_id": "msg_primary", "output_index": 0, "content_index": 0, "delta": partial, + },) + if partial else () + ), + { + "type": "error", "sequence_number": 1, + "error": {"type": "server_error", "code": "internal_error", "message": "interrupted"}, + }, + ) + + +def _anthropic_stream(text: str) -> httpx.Response: + return _sse( + {"type": "message_start", "message": { + "id": "msg_fallback", "type": "message", "role": "assistant", + "model": "claude-fallback", "content": [], "usage": {"input_tokens": 10, "output_tokens": 0}, + }}, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}, + {"type": "content_block_stop", "index": 0}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 5}}, + {"type": "message_stop"}, + ) + + +def _primary_with_fallback(fallback_model: str) -> Router: + return Router( + model_list=[ + {"model_name": "primary", "litellm_params": {"model": "openai/gpt-6-astra", "api_key": "synthetic"}}, + {"model_name": "fallback", "litellm_params": {"model": fallback_model, "api_key": "synthetic"}}, + ], + fallbacks=[{"primary": ["fallback"]}], + num_retries=0, + ) + + +def _streamed_text(events: tuple[Any, ...]) -> str: + return "".join(event.delta for event in events if event.type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA) + + +@pytest.fixture +def claude_with_prefill() -> str: + model: Final = "anthropic/claude-prefill-capable-test-model" + litellm.register_model( + {model: {"litellm_provider": "anthropic", "mode": "chat", "supports_assistant_prefill": True}}, + persist_across_reloads=False, + ) + litellm.get_model_info.cache_clear() + return model + + +@pytest.mark.asyncio +@pytest.mark.parametrize("partial", ["", "1, 2, "]) +async def test_responses_fallback_can_continue_on_a_provider_without_prefill( + partial: str, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + prompt: Final = "Print the integers 1 through 5, separated by commas." + primary: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=_openai_stream_cut_after(partial) + ) + + def no_prefill_provider(request: httpx.Request) -> httpx.Response: + body: Final = json.loads(request.content) + if body["messages"][-1]["role"] != "user": + return httpx.Response( + 400, + json={"type": "error", "error": { + "type": "invalid_request_error", "message": "The conversation must end with a user message", + }}, + ) + return _anthropic_stream("3, 4, 5" if partial else "1, 2, 3, 4, 5") + + fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(side_effect=no_prefill_provider) + router: Final = _primary_with_fallback("anthropic/claude-opus-5-5") + try: + stream: Final = await router.aresponses(model="primary", input=prompt, stream=True, max_output_tokens=100) + events: Final = tuple([event async for event in stream]) + assert _streamed_text(events) == "1, 2, 3, 4, 5" + assert events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + assert primary.call_count == fallback.call_count == 1 + messages: Final = json.loads(fallback.calls.last.request.content)["messages"] + assert messages[0] == {"role": "user", "content": [{"type": "text", "text": prompt}]} + if partial: + assert [message["role"] for message in messages] == ["user", "assistant", "user"] + assert messages[-2] == {"role": "assistant", "content": [{"type": "text", "text": partial}]} + assert json.dumps(partial) in messages[-1]["content"][0]["text"] + else: + assert len(messages) == 1 + finally: + router.discard() + + +@pytest.mark.asyncio +async def test_responses_fallback_continues_a_claude_that_takes_prefill_from_its_own_partial( + claude_with_prefill: str, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://api.openai.com/v1/responses").mock(return_value=_openai_stream_cut_after("1, 2, ")) + fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=_anthropic_stream(" 3, 4, 5") + ) + router: Final = _primary_with_fallback(claude_with_prefill) + try: + stream: Final = await router.aresponses( + model="primary", input="Print the integers 1 through 5.", stream=True, max_output_tokens=100 + ) + events: Final = tuple([event async for event in stream]) + assert _streamed_text(events) == "1, 2, 3, 4, 5" + assert events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + messages: Final = json.loads(fallback.calls.last.request.content)["messages"] + assert [message["role"] for message in messages] == ["user", "assistant"] + assert messages[-1]["content"][0]["text"].startswith("1, 2,") + finally: + router.discard() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fallback_takes_prefill", [True, False], ids=["prefill", "request_to_continue"]) +async def test_responses_fallback_continuation_drops_reasoning_the_fallback_cannot_decrypt( + fallback_takes_prefill: bool, + claude_with_prefill: str, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://api.openai.com/v1/responses").mock(return_value=_openai_stream_cut_after("1, 2, ")) + fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=_anthropic_stream(" 3, 4, 5") + ) + router: Final = _primary_with_fallback(claude_with_prefill if fallback_takes_prefill else "anthropic/claude-opus-5-5") + try: + stream: Final = await router.aresponses( + model="primary", + input=[ + {"role": "user", "content": "Print the integers 1 through 5."}, + {"type": "reasoning", "id": "rs_1", "summary": [], "encrypted_content": "primary-only-ciphertext"}, + ], + stream=True, + max_output_tokens=100, + ) + _ = [event async for event in stream] + assert fallback.call_count == 1 + assert "primary-only-ciphertext" not in fallback.calls.last.request.content.decode() + finally: + router.discard() + + +class _Numbers(BaseModel): + model_config = ConfigDict(frozen=True) + + numbers: tuple[int, ...] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "constraint", + [ + {"text": {"format": {"type": "json_schema", "name": "numbers", "strict": True, "schema": {"type": "object"}}}}, + {"text": {"format": {"type": "json_object"}}}, + {"tool_choice": "required", "tools": [{"type": "function", "name": "lookup", "parameters": {"type": "object"}}]}, + {"text_format": _Numbers}, + ], + ids=["json_schema", "json_object", "required_tool", "sdk_text_format"], +) +async def test_responses_fallback_surfaces_the_primary_error_once_constrained_output_was_streamed( + constraint: dict[str, object], + claude_with_prefill: str, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=_openai_stream_cut_after('{"numbers": [1, 2, ') + ) + fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=_anthropic_stream("3]}") + ) + router: Final = _primary_with_fallback(claude_with_prefill) + try: + stream: Final = await router.aresponses( + model="primary", input="Numbers 1 to 3 as JSON.", stream=True, max_output_tokens=100, **constraint + ) + with pytest.raises(litellm.InternalServerError, match="interrupted"): + [event async for event in stream] + assert fallback.call_count == 0 + finally: + router.discard() + + +@pytest.mark.asyncio +async def test_responses_fallback_continues_within_the_callers_output_ceiling( + claude_with_prefill: str, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://api.openai.com/v1/responses").mock(return_value=_openai_stream_cut_after("1, 2, ")) + fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=_anthropic_stream("3, 4, 5") + ) + router: Final = _primary_with_fallback(claude_with_prefill) + try: + stream: Final = await router.aresponses(model="primary", input="Count to 5.", stream=True, max_output_tokens=100) + events: Final = tuple([event async for event in stream]) + assert _streamed_text(events) == "1, 2, 3, 4, 5" + assert json.loads(fallback.calls.last.request.content)["max_tokens"] == 100 - litellm.token_counter( + text="1, 2, " + ) + finally: + router.discard() + + +@pytest.mark.asyncio +async def test_responses_fallback_surfaces_the_primary_error_once_the_cut_stream_used_the_output_ceiling( + claude_with_prefill: str, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://api.openai.com/v1/responses").mock(return_value=_openai_stream_cut_after("1, 2, ")) + fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=_anthropic_stream("3, 4, 5") + ) + router: Final = _primary_with_fallback(claude_with_prefill) + try: + stream: Final = await router.aresponses( + model="primary", input="Count to 5.", stream=True, max_output_tokens=litellm.token_counter(text="1, 2, ") + ) + with pytest.raises(litellm.InternalServerError, match="interrupted"): + [event async for event in stream] + assert fallback.call_count == 0 + finally: + router.discard() + + +def _openai_stream_cut_mid_tool_call(preamble: str, *, announce_item: bool = True) -> httpx.Response: + text: Final = { + "type": "response.output_text.delta", "sequence_number": 0, + "item_id": "msg_primary", "output_index": 0, "content_index": 0, "delta": preamble, + } + item_added: Final = { + "type": "response.output_item.added", "sequence_number": 1, "output_index": 1, + "item": { + "type": "function_call", "id": "fc_primary", "call_id": "call_primary", + "name": "lookup", "arguments": "", "status": "in_progress", + }, + } + arguments: Final = { + "type": "response.function_call_arguments.delta", "sequence_number": 2, + "item_id": "fc_primary", "output_index": 1, "delta": '{"city": "Par', + } + error: Final = { + "type": "error", "sequence_number": 3, + "error": {"type": "server_error", "code": "internal_error", "message": "interrupted"}, + } + return _sse(*((text, item_added, arguments, error) if announce_item else (text, arguments, error))) + + +def _anthropic_stream_cut_mid_tool_call(preamble: str, tool_block: Mapping[str, object] | None = None) -> httpx.Response: + return _sse( + {"type": "message_start", "message": { + "id": "msg_primary", "type": "message", "role": "assistant", + "model": "claude-primary", "content": [], "usage": {"input_tokens": 10, "output_tokens": 0}, + }}, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": preamble}}, + {"type": "content_block_stop", "index": 0}, + {"type": "content_block_start", "index": 1, "content_block": tool_block or { + "type": "tool_use", "id": "toolu_primary", "name": "lookup", "input": {}, + }}, + {"type": "content_block_delta", "index": 1, "delta": {"type": "input_json_delta", "partial_json": '{"city": "Par'}}, + {"type": "error", "error": {"type": "overloaded_error", "message": "interrupted"}}, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("announce_item", [True, False], ids=["announced_item", "arguments_only"]) +async def test_responses_fallback_surfaces_the_primary_error_once_the_cut_stream_began_a_tool_call( + announce_item: bool, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=_openai_stream_cut_mid_tool_call("Let me look that up. ", announce_item=announce_item) + ) + fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=_anthropic_stream("It is sunny in Paris.") + ) + router: Final = _primary_with_fallback("anthropic/claude-opus-5-5") + try: + stream: Final = await router.aresponses( + model="primary", + input="What is the weather in Paris?", + stream=True, + tools=[{"type": "function", "name": "lookup", "parameters": {"type": "object"}}], + ) + with pytest.raises(litellm.InternalServerError, match="interrupted"): + [event async for event in stream] + assert fallback.call_count == 0 + finally: + router.discard() + + +def _openai_stream_cut_after_event(preamble: str, event: Mapping[str, object]) -> httpx.Response: + return _sse( + { + "type": "response.output_text.delta", "sequence_number": 0, + "item_id": "msg_primary", "output_index": 0, "content_index": 0, "delta": preamble, + }, + {**event, "sequence_number": 1}, + { + "type": "error", "sequence_number": 2, + "error": {"type": "server_error", "code": "internal_error", "message": "interrupted"}, + }, + ) + + +def _openai_stream_cut_after_output_item(preamble: str, item: Mapping[str, object]) -> httpx.Response: + return _openai_stream_cut_after_event( + preamble, {"type": "response.output_item.added", "output_index": 1, "item": item} + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "item", + [ + {"type": "shell_call", "id": "sh_1", "call_id": "call_1", "action": {"commands": ["ls"]}, "status": "in_progress"}, + { + "type": "apply_patch_call", "id": "apc_1", "call_id": "call_1", "status": "in_progress", + "operation": {"type": "update_file", "path": "a.py", "diff": "@@"}, + }, + { + "type": "local_shell_call", "id": "lsh_1", "call_id": "call_1", "status": "in_progress", + "action": {"type": "exec", "command": ["ls"], "env": {}}, + }, + { + "type": "computer_call", "id": "cu_1", "call_id": "call_1", "status": "in_progress", + "action": {"type": "screenshot"}, "pending_safety_checks": [], + }, + {"type": "mcp_approval_request", "id": "mcpr_1", "name": "lookup", "arguments": "{}", "server_label": "docs"}, + ], + ids=lambda item: item["type"], +) +async def test_responses_fallback_surfaces_the_primary_error_once_the_cut_stream_began_a_client_run_tool_call( + item: Mapping[str, object], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=_openai_stream_cut_after_output_item("Let me check the repo. ", item) + ) + fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=_anthropic_stream("The repo has two files.") + ) + router: Final = _primary_with_fallback("anthropic/claude-opus-5-5") + try: + stream: Final = await router.aresponses(model="primary", input="What is in the repo?", stream=True) + with pytest.raises(litellm.InternalServerError, match="interrupted"): + [event async for event in stream] + assert fallback.call_count == 0 + finally: + router.discard() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "server_run_event", + [ + { + "type": "response.output_item.added", "output_index": 1, + "item": { + "type": "mcp_call", "id": "mcp_1", "name": "ask_question", "arguments": "{}", + "server_label": "docs", "status": "in_progress", + }, + }, + {"type": "response.mcp_call_arguments.delta", "output_index": 1, "item_id": "mcp_1", "delta": '{"question": "set'}, + ], + ids=["mcp_call_item", "mcp_call_arguments"], +) +async def test_responses_fallback_continues_after_a_server_run_mcp_call( + server_run_event: Mapping[str, object], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=_openai_stream_cut_after_event("Let me check the docs. ", server_run_event) + ) + fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=_anthropic_stream("The docs cover setup.") + ) + router: Final = _primary_with_fallback("anthropic/claude-opus-5-5") + try: + stream: Final = await router.aresponses(model="primary", input="What do the docs cover?", stream=True) + events: Final = tuple([event async for event in stream]) + assert _streamed_text(events).endswith("The docs cover setup.") + assert fallback.call_count == 1 + finally: + router.discard() + + +@pytest.mark.asyncio +async def test_responses_fallback_continues_when_allowed_tools_in_auto_mode_leaves_the_reply_open( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + openai: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + side_effect=[ + _openai_stream_cut_after("1, 2, "), + _sse( + { + "type": "response.output_text.delta", "sequence_number": 0, + "item_id": "msg_fallback", "output_index": 0, "content_index": 0, "delta": "3, 4, 5", + }, + { + "type": "response.completed", "sequence_number": 1, + "response": { + "id": "resp_fallback", "object": "response", "created_at": 0, "status": "completed", + "model": "gpt-6.1-terra", "output": [], + "usage": {"input_tokens": 5, "output_tokens": 5, "total_tokens": 10}, + }, + }, + ), + ] + ) + router: Final = _primary_with_fallback("openai/gpt-6.1-terra") + try: + stream: Final = await router.aresponses( + model="primary", + input="Print the integers 1 through 5.", + stream=True, + tools=[{"type": "function", "name": "lookup", "parameters": {"type": "object"}}], + tool_choice={"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "lookup"}]}, + ) + events: Final = tuple([event async for event in stream]) + assert _streamed_text(events) == "1, 2, 3, 4, 5" + assert openai.call_count == 2 + continued_input: Final = json.loads(openai.calls.last.request.content)["input"] + assert [item["role"] for item in continued_input[-2:]] == ["assistant", "user"] + finally: + router.discard() + + +@pytest.mark.asyncio +async def test_responses_fallback_surfaces_a_content_filter_that_cut_the_stream_mid_answer( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=_sse( + { + "type": "response.output_text.delta", "sequence_number": 0, + "item_id": "msg_primary", "output_index": 0, "content_index": 0, "delta": "Step one is to obtain the", + }, + { + "type": "error", "sequence_number": 1, + "error": { + "type": "invalid_request_error", "code": "content_policy_violation", + "message": "The response was halted mid-stream.", "param": None, + }, + }, + ) + ) + fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=_anthropic_stream(" reagents.") + ) + router: Final = Router( + model_list=[ + {"model_name": "primary", "litellm_params": {"model": "openai/gpt-6-astra", "api_key": "synthetic"}}, + {"model_name": "fallback", "litellm_params": {"model": "anthropic/claude-opus-5-5", "api_key": "synthetic"}}, + ], + fallbacks=[{"primary": ["fallback"]}], + content_policy_fallbacks=[{"primary": ["fallback"]}], + num_retries=0, + ) + try: + stream: Final = await router.aresponses(model="primary", input="How do I make it?", stream=True) + with pytest.raises(litellm.ContentPolicyViolationError, match="halted mid-stream"): + [event async for event in stream] + assert fallback.call_count == 0 + finally: + router.discard() + + +@pytest.mark.asyncio +async def test_bridged_responses_fallback_surfaces_the_primary_error_once_the_cut_stream_began_a_tool_call( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + provider: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + side_effect=[ + _anthropic_stream_cut_mid_tool_call("Let me look that up. "), + _anthropic_stream("It is sunny in Paris."), + ] + ) + router: Final = Router( + model_list=[ + {"model_name": "primary", "litellm_params": {"model": "anthropic/claude-primary", "api_key": "synthetic"}}, + {"model_name": "fallback", "litellm_params": {"model": "anthropic/claude-opus-5-5", "api_key": "synthetic"}}, + ], + fallbacks=[{"primary": ["fallback"]}], + num_retries=0, + ) + try: + stream: Final = await router.aresponses( + model="primary", + input="What is the weather in Paris?", + stream=True, + tools=[{"type": "function", "name": "lookup", "parameters": {"type": "object"}}], + ) + with pytest.raises(litellm.InternalServerError, match="interrupted"): + [event async for event in stream] + assert provider.call_count == 1 + finally: + router.discard() + + +@pytest.mark.asyncio +async def test_bridged_responses_fallback_continues_after_a_server_run_tool_search( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + provider: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + side_effect=[ + _anthropic_stream_cut_mid_tool_call( + "Let me find the right tool. ", + {"type": "server_tool_use", "id": "srvtoolu_primary", "name": "tool_search_tool_regex", "input": {}}, + ), + _anthropic_stream("It is sunny in Paris."), + ] + ) + router: Final = Router( + model_list=[ + {"model_name": "primary", "litellm_params": {"model": "anthropic/claude-primary", "api_key": "synthetic"}}, + {"model_name": "fallback", "litellm_params": {"model": "anthropic/claude-opus-5-5", "api_key": "synthetic"}}, + ], + fallbacks=[{"primary": ["fallback"]}], + num_retries=0, + ) + try: + stream: Final = await router.aresponses( + model="primary", + input="What is the weather in Paris?", + stream=True, + tools=[ + {"type": "tool_search", "execution": "server"}, + {"type": "function", "name": "lookup", "parameters": {"type": "object"}, "defer_loading": True}, + ], + ) + events: Final = tuple([event async for event in stream]) + assert _streamed_text(events).endswith("It is sunny in Paris.") + assert provider.call_count == 2 + finally: + router.discard() + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") @pytest.mark.asyncio diff --git a/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py b/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py index 8fa1c799863..f79c8a6fadb 100644 --- a/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py @@ -4,13 +4,24 @@ import litellm from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( MID_STREAM_CONTINUATION_KWARG, MID_STREAM_CONTINUATION_MARKER, + RESPONSES_CONTINUATION_KWARG, ContinuationOutput, ContinuationPrefillDeploymentCheck, + ResponsesContinuation, + ResponsesContinuationSeam, + constrains_continuation, + constrains_responses_continuation, continuation_output_ceilings, deployment_supports_prefill, + deployment_takes_responses_prefill, + responses_continuation_input, + responses_continuation_input_for, ) +from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents PREFILL_MODEL = "anthropic/prefill-capable-test-model" +CLAUDE_PREFILL_MODEL = "anthropic/claude-prefill-capable-test-model" +CLAUDE_NON_PREFILL_MODEL = "anthropic/claude-prefill-rejecting-test-model" NON_PREFILL_MODEL = "openai/prefill-unknown-test-model" UNMAPPED_MODEL = "openai/unmapped-test-model" @@ -21,6 +32,12 @@ def _synthetic_cost_map_entries() -> None: { PREFILL_MODEL: {"litellm_provider": "anthropic", "mode": "chat", "supports_assistant_prefill": True}, NON_PREFILL_MODEL: {"litellm_provider": "openai", "mode": "chat"}, + CLAUDE_PREFILL_MODEL: {"litellm_provider": "anthropic", "mode": "chat", "supports_assistant_prefill": True}, + CLAUDE_NON_PREFILL_MODEL: { + "litellm_provider": "anthropic", + "mode": "chat", + "supports_assistant_prefill": False, + }, }, persist_across_reloads=False, ) @@ -191,7 +208,11 @@ def _text_chunk(text: str) -> litellm.ModelResponseStream: ("Hello", (" world",), (" world",)), ("Hello", ("Hel", "p me"), ("Hel", "p me")), ("Hello", ("Hel",), ("Hel",)), - ([{"type": "text", "text": "Sure: "}, {"type": "text", "text": "Hello"}], ("Sure: Hello", " world"), ("", " world")), + ( + [{"type": "text", "text": "Sure: "}, {"type": "text", "text": "Hello"}], + ("Sure: Hello", " world"), + ("", " world"), + ), ], ids=[ "echo_split_across_chunks_is_dropped", @@ -206,9 +227,7 @@ def _text_chunk(text: str) -> litellm.ModelResponseStream: async def test_continuation_output_drops_only_a_full_echo_of_the_prefill( prefill_content, fallback_texts, expected_texts ): - output = ContinuationOutput( - "Hello", {"role": "assistant", "content": prefill_content, "prefix": True}, "group", () - ) + output = ContinuationOutput("Hello", {"role": "assistant", "content": prefill_content, "prefix": True}, "group", ()) async def fallback_stream(): for text in fallback_texts: @@ -266,7 +285,14 @@ async def test_continuation_output_drops_only_whitespace_the_client_already_has( ({"extra_body": {"max_completion_tokens": 45}}, {"extra_body": {"max_completion_tokens": 15}}), ({"extra_body": {"top_k": 3}}, {}), ], - ids=["max_tokens", "both_ceilings", "no_ceiling", "extra_body_max_tokens", "extra_body_only", "extra_body_no_ceiling"], + ids=[ + "max_tokens", + "both_ceilings", + "no_ceiling", + "extra_body_max_tokens", + "extra_body_only", + "extra_body_no_ceiling", + ], ) def test_continuation_output_ceilings_take_the_emitted_tokens_off_every_ceiling_the_request_sends( request_kwargs, expected @@ -286,3 +312,217 @@ def test_continuation_output_ceilings_take_the_emitted_tokens_off_every_ceiling_ ) def test_continuation_output_ceilings_decline_once_a_ceiling_has_no_tokens_left(request_kwargs): assert continuation_output_ceilings(request_kwargs, 30) is None + + +@pytest.mark.parametrize( + "sent,expected", + [ + ({"text": {"format": {"type": "json_schema", "name": "n", "schema": {"type": "object"}}}}, True), + ({"text": {"format": {"type": "json_object"}}}, True), + ({"text": {"format": {"type": "text"}}}, False), + ({"text": {"verbosity": "low"}}, False), + ({"tool_choice": "required"}, True), + ({"tool_choice": "auto"}, False), + ({"tool_choice": {"type": "function", "name": "lookup"}}, True), + ({"tool_choice": {"type": "allowed_tools", "mode": "required", "tools": [{"type": "function", "name": "f"}]}}, True), + ({"tool_choice": {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "f"}]}}, True), + ({"tool_choice": {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": []}}}, True), + ({"tool_choice": {"type": "web_search_preview"}}, True), + ], + ids=[ + "json_schema", + "json_object", + "text_format", + "verbosity_only", + "required_tool", + "auto_tool", + "forced_function", + "allowed_tools_required", + "allowed_tools_auto", + "chat_allowed_tools_auto", + "hosted_tool", + ], +) +def test_constrains_continuation_reads_the_responses_text_format(sent, expected): + assert constrains_continuation(sent) is expected + + +@pytest.mark.parametrize( + "sent,expected", + [ + ({"tool_choice": {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "f"}]}}, False), + ({"tool_choice": {"type": "allowed_tools", "mode": "required", "tools": [{"type": "function", "name": "f"}]}}, True), + ({"tool_choice": {"type": "function", "name": "lookup"}}, True), + ({"tool_choice": {"type": "web_search_preview"}}, True), + ({"tool_choice": "required"}, True), + ({"tool_choice": "auto"}, False), + ( + { + "tool_choice": {"type": "allowed_tools", "mode": "auto", "tools": []}, + "text": {"format": {"type": "json_object"}}, + }, + True, + ), + ], + ids=[ + "allowed_tools_auto", + "allowed_tools_required", + "forced_function", + "hosted_tool", + "required_tool", + "auto_tool", + "allowed_tools_auto_with_json", + ], +) +def test_constrains_responses_continuation_lets_allowed_tools_in_auto_mode_answer_in_text(sent, expected): + assert constrains_responses_continuation(sent) is expected + + +@pytest.mark.parametrize( + "deployment,expected", + [ + (_deployment(CLAUDE_PREFILL_MODEL, "a"), True), + (_deployment(CLAUDE_NON_PREFILL_MODEL, "b"), False), + (_deployment(PREFILL_MODEL, "c"), False), + ( + {"litellm_params": {"model": "anthropic/support-model-v1"}, "model_info": {"id": "d", "base_model": CLAUDE_PREFILL_MODEL}}, + True, + ), + ], + ids=["claude_with_prefill", "claude_without_prefill", "prefill_capable_but_not_claude", "claude_named_by_base_model"], +) +def test_deployment_takes_responses_prefill_only_on_claude_with_prefill(deployment, expected): + assert deployment_takes_responses_prefill(deployment) is expected + + +def _continuation_kwargs() -> dict[str, object]: + continuation = ResponsesContinuation("Count to 5.", "1, 2, ") + return {RESPONSES_CONTINUATION_KWARG: continuation, "input": responses_continuation_input(continuation)} + + +@pytest.mark.parametrize( + "prior_input,expected_prior", + [ + ( + "Count to 5.", + [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Count to 5."}]}], + ), + ( + [ + {"role": "user", "content": "Count to 5."}, + {"type": "function_call_output", "call_id": "c", "output": "x"}, + ], + [ + {"role": "user", "content": "Count to 5."}, + {"type": "function_call_output", "call_id": "c", "output": "x"}, + ], + ), + (None, []), + ], + ids=["string_input", "item_list_input", "no_input"], +) +def test_responses_continuation_input_keeps_the_prior_input_then_the_partial_then_a_request_to_continue( + prior_input, expected_prior +): + built = responses_continuation_input(ResponsesContinuation(prior_input, "1, 2, ")) + + partial = {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "1, 2, "}]} + assert built[:-1] == [*expected_prior, partial] + assert built[-1]["role"] == "user" + assert "already ends with whitespace" in built[-1]["content"][0]["text"] + + +def test_responses_continuation_request_asks_for_a_leading_space_only_after_a_word(): + built = responses_continuation_input(ResponsesContinuation("Hi", "Transformers use")) + + assert "begin with a space" in built[-1]["content"][0]["text"] + + +def test_responses_continuation_accumulates_across_hops(): + continuation = ResponsesContinuation("Count to 5.", "1, ").extended("2, ") + + assert continuation == ResponsesContinuation("Count to 5.", "1, 2, ") + + +@pytest.mark.parametrize( + "deployment,expected_roles", + [ + (_deployment(CLAUDE_PREFILL_MODEL, "a"), ["user", "assistant"]), + (_deployment(CLAUDE_NON_PREFILL_MODEL, "b"), ["user", "assistant", "user"]), + (_deployment(NON_PREFILL_MODEL, "c"), ["user", "assistant", "user"]), + ], + ids=["claude_prefill", "claude_without_prefill", "openai"], +) +def test_responses_continuation_input_for_shapes_the_input_for_the_deployment_that_serves_it( + deployment, expected_roles +): + shaped = responses_continuation_input_for(_continuation_kwargs(), deployment) + + assert [item["role"] for item in shaped] == expected_roles + + +def test_responses_continuation_input_for_keeps_what_the_router_stripped_from_the_input(): + continuation = ResponsesContinuation( + [ + {"type": "reasoning", "id": "rs_1", "summary": [], "encrypted_content": "opaque"}, + {"role": "user", "content": "Count to 5."}, + ], + "1, 2, ", + ) + stripped_input = responses_continuation_input(continuation)[1:] + + shaped = responses_continuation_input_for( + {RESPONSES_CONTINUATION_KWARG: continuation, "input": stripped_input}, _deployment(CLAUDE_PREFILL_MODEL, "a") + ) + + assert shaped == stripped_input[:-1] + + +def test_responses_continuation_input_for_leaves_an_input_without_its_request_to_continue_alone(): + kwargs = {**_continuation_kwargs(), "input": [{"role": "user", "content": "Count to 5."}]} + + assert responses_continuation_input_for(kwargs, _deployment(CLAUDE_PREFILL_MODEL, "a")) == kwargs["input"] + + +def test_responses_continuation_input_for_leaves_a_request_without_a_continuation_alone(): + kwargs = {"input": "Count to 5."} + + assert responses_continuation_input_for(kwargs, _deployment(CLAUDE_PREFILL_MODEL, "a")) == "Count to 5." + + +def _delta(text: str) -> OutputTextDeltaEvent: + return OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, item_id="m", output_index=0, content_index=0, delta=text + ) + + +@pytest.mark.parametrize( + "emitted,fallback_texts,expected_texts", + [ + ("1, 2, ", (" 3, 4", ", 5"), ("3, 4", ", 5")), + ("1, 2, ", (" ", " 3"), ("", " 3")), + ("1, 2,", (" 3",), (" 3",)), + ("1, 2, ", ("3",), ("3",)), + ("A.\n\n", ("\n Hi",), (" Hi",)), + ("", (" Hi",), (" Hi",)), + ], + ids=[ + "re_sent_space_is_dropped", + "seam_split_across_events", + "no_trailing_whitespace_keeps_the_space", + "no_re_sent_space_passes_through", + "seam_ends_at_the_first_divergence", + "fresh_start_passes_through", + ], +) +@pytest.mark.asyncio +async def test_responses_continuation_seam_drops_only_whitespace_the_client_already_has( + emitted, fallback_texts, expected_texts +): + async def fallback_stream(): + for text in fallback_texts: + yield _delta(text) + + released = [event async for event in ResponsesContinuationSeam(emitted).splice(fallback_stream())] + + assert tuple(event.delta for event in released) == expected_texts diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 30d675e130d..c6d02dddae0 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -7039,8 +7039,7 @@ async def test_aresponses_fallback_attempt_does_not_replay_held_lifecycle_events @pytest.mark.asyncio async def test_aresponses_streaming_iterator_partial_content_injects_continuation(): - """Mid-stream error: input is rewritten to include user prompt + - developer instruction + prior assistant message with partial output.""" + """Mid-stream error: preserve partial assistant output and request continuation in a user turn.""" from litellm.exceptions import MidStreamFallbackError router = _make_router_with_fallback() @@ -7076,11 +7075,11 @@ async def test_aresponses_streaming_iterator_partial_content_injects_continuatio assert isinstance(new_input, list) assert new_input[0]["role"] == "user" assert new_input[0]["content"][0]["text"] == "What's the capital of France?" - assert new_input[1]["role"] == "developer" - assert "do not repeat" in new_input[1]["content"][0]["text"].lower() - assert new_input[2]["role"] == "assistant" - assert new_input[2]["content"][0]["type"] == "output_text" - assert new_input[2]["content"][0]["text"] == "The capital of France is" + assert new_input[1]["role"] == "assistant" + assert new_input[1]["content"][0]["type"] == "output_text" + assert new_input[1]["content"][0]["text"] == "The capital of France is" + assert new_input[2]["role"] == "user" + assert "do not repeat" in new_input[2]["content"][0]["text"].lower() @pytest.mark.asyncio From 9546f9a69ff470dfcc5a21eccab89d2fd86a634f Mon Sep 17 00:00:00 2001 From: mika-lt <349045450+mika-lt@users.noreply.github.com> Date: Wed, 28 Oct 2026 04:36:50 -0700 Subject: [PATCH 024/128] test(router): audit #47498 (#48799) Pin the queue route contracts #47498 changed: the priority header on streamed answers, the 400 for a request with no priority to queue at, a provider bad request keeping its own status, and a deployment-id request served ahead of the queue Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../test_priority_scheduler_queue_cleanup.py | 61 ++++++++++++++++++- 1 file changed, 58 insertions(+), 3 deletions(-) diff --git a/tests/integration/routing/test_priority_scheduler_queue_cleanup.py b/tests/integration/routing/test_priority_scheduler_queue_cleanup.py index 9fc21843b04..11798542088 100644 --- a/tests/integration/routing/test_priority_scheduler_queue_cleanup.py +++ b/tests/integration/routing/test_priority_scheduler_queue_cleanup.py @@ -512,12 +512,55 @@ def test_other_prioritized_endpoints_are_served_once_and_billed( marker: Final = new_marker() answer: Final = call(gateway, endpoint, model, marker) identity: Final = assert_served(answer, marker) - if endpoint == "queue": + if endpoint in ("queue", "queue-stream"): assert answer.headers.get("x-litellm-priority") == "1", answer.headers assert_once(received_markers(wire), (marker,)) spend_row_lands(identity) +@pytest.mark.parametrize("field", ({}, {"priority": "high", "drop_params": True}), ids=("missing", "dropped")) +def test_queue_route_refuses_a_request_it_cannot_queue_before_the_upstream( + gateway: Gateway, rig_model: tuple[str, Wire], field: Mapping[str, JsonValue] +) -> None: + model, wire = rig_model + marker: Final = new_marker() + answer: Final = post(gateway, "/queue/chat/completions", chat_body(model, marker, **field)) + assert answer.status == 400, (answer.status, answer.text) + error: Final = JSON_OBJECT.validate_json(answer.text)["error"] + assert isinstance(error, dict), answer.text + assert (error["type"], error["param"]) == ("invalid_request_error", "priority"), answer.text + assert_never(received_markers(wire), (marker,)) + + +def anthropic_error(status: int) -> Reply: + body: Final = {"type": "error", "error": {"type": "invalid_request_error", "message": f"scripted {status}"}} + return Reply(status=status, body=json.dumps(body).encode()) + + +@pytest.mark.parametrize( + ("status", "provider_model", "base_path", "respond"), + ( + (400, "openai/gpt-4o-mini", "/v1", answering_model_discovery(lambda _: openai_error(400))), + (402, "anthropic/claude-sonnet-5-5", "", lambda _: anthropic_error(402)), + ), + ids=("openai-400", "anthropic-402"), +) +def test_queue_route_answers_a_provider_bad_request_with_its_own_status( + gateway: Gateway, status: int, provider_model: str, base_path: str, respond: Callable[[Request], Reply] +) -> None: + marker: Final = new_marker() + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=provider_model, api_base=wire.url + base_path) + answer: Final = post(gateway, "/queue/chat/completions", chat_body(model, marker, priority=1)) + assert marker in received_markers(wire) + assert answer.status == status, (answer.status, answer.text) + error: Final = JSON_OBJECT.validate_json(answer.text)["error"] + assert isinstance(error, dict), answer.text + assert error["type"] == "invalid_request_error", answer.text + assert f"scripted {status}" in str(error["message"]), answer.text + assert "Authentication Error" not in str(error["message"]), answer.text + + def upstream_body_for(wire: Wire, marker: str) -> dict[str, JsonValue]: (upstream_request,) = tuple(request for request in wire.drain() if marker.encode() in request.body) return JSON_OBJECT.validate_json(upstream_request.body) @@ -751,6 +794,19 @@ def test_waiter_proceeds_once_the_entry_ahead_is_cleared(pair: Pair) -> None: assert_drained(pair.cache, group) +@pytest.mark.timeout(PAIR_CELL_TIMEOUT) +def test_request_pinned_to_a_deployment_id_is_served_ahead_of_an_earlier_queue_entry(pair: Pair) -> None: + deployment: Final = f"{PAIR_GROUPS[0]}-deployment" + write_queue(pair.cache, deployment, [GHOST_ENTRY]) + marker: Final = new_marker() + answer: Final = post(pair.second, "/v1/chat/completions", chat_body(deployment, marker, priority=2)) + assert_served(answer, marker) + assert_prioritized(answer) + assert answer.seconds < PROMPT_SECONDS, answer.seconds + assert queue_entries(pair.cache, deployment) == [GHOST_ENTRY] + assert_once(received_markers(pair.wire), (marker,)) + + def pinned_connection(gateway: Gateway) -> http.client.HTTPConnection: address: Final = urlsplit(str(gateway.client.base_url)) assert address.hostname is not None and address.port is not None, address @@ -1313,8 +1369,7 @@ def test_queue_route_without_priority_takes_the_default(defaulted: Defaulted, en marker: Final = new_marker() answer: Final = call(defaulted.pair.first, endpoint, group, marker, priority=None) assert_served(answer, marker) - if endpoint == "queue": - assert answer.headers.get("x-litellm-priority") == str(DEFAULT_PRIORITY), answer.headers + assert answer.headers.get("x-litellm-priority") == str(DEFAULT_PRIORITY), answer.headers assert "priority" not in upstream_body_for(defaulted.pair.wire, marker) From 889d8aec15b11250eb3ea18c825be34eac94b6b1 Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Wed, 28 Oct 2026 04:38:21 -0700 Subject: [PATCH 025/128] test(decisions): audit #47508 (#48777) * test(decisions): audit #47508 * test(decisions): price the Databricks decision cell from its cost-map row --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../cost_calculation/cases/decisions.py | 32 +++++++++++++++++++ .../integration/cost_calculation/cost_map.py | 6 ++++ .../providers/test_decisions_wire.py | 23 ++++++++++--- 3 files changed, 56 insertions(+), 5 deletions(-) diff --git a/tests/integration/cost_calculation/cases/decisions.py b/tests/integration/cost_calculation/cases/decisions.py index 1cc57afe8ed..6594510b232 100644 --- a/tests/integration/cost_calculation/cases/decisions.py +++ b/tests/integration/cost_calculation/cases/decisions.py @@ -162,6 +162,37 @@ TYPESAFE_JEV_LATEST_SYSTEMONE: Final = CostTrackingTestCase( completion_tokens=22, ), ) +TYPESAFE_JEV_PREVIEW_ANSWERED_AS_PRICED_VERSION_SYSTEMONE: Final = CostTrackingTestCase( + name="typesafe/jev-preview-answered-as-priced-version-systemone", + covers="quota_management.spend_tracking.systemone.logs_cost", + model="typesafe/jev-preview", + endpoint="/v1/systemone", + request={ + "model": "$MODEL", + "state": "The replacement part arrived but the wrong size.", + "questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}}, + }, + upstream_body={ + "model": "jev-preview", + "state": "The replacement part arrived but the wrong size.", + "questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}}, + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "jev-latest", + "answers": {"is_refund": {"type": "noul", "noul": 0.31}}, + "usage": {"input_tokens": 311, "output_tokens": 9}, + }, + ), + expected=ExactExpected( + spend=311 * 4.2e-08, + input_cost=311 * 4.2e-08, + output_cost=0.0, + prompt_tokens=311, + completion_tokens=9, + ), +) # perplexity/pplx-decider-v1-27b PERPLEXITY_PPLX_DECIDER_V1_27B_DECISIONS: Final = CostTrackingTestCase( @@ -514,6 +545,7 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( OPENAI_GPT_6_LUNA_SYSTEMONE, TYPESAFE_JEV_LATEST_DECISIONS, TYPESAFE_JEV_LATEST_SYSTEMONE, + TYPESAFE_JEV_PREVIEW_ANSWERED_AS_PRICED_VERSION_SYSTEMONE, PERPLEXITY_PPLX_DECIDER_V1_27B_DECISIONS, PERPLEXITY_PPLX_DECIDER_V1_27B_SYSTEMONE, CLOUDFLARE_CLEF_FLASH_DECISIONS, diff --git a/tests/integration/cost_calculation/cost_map.py b/tests/integration/cost_calculation/cost_map.py index 16107def2ad..32722b37c9b 100644 --- a/tests/integration/cost_calculation/cost_map.py +++ b/tests/integration/cost_calculation/cost_map.py @@ -121,6 +121,12 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "input_cost_per_token": 4.2e-08, "output_cost_per_token": 0.0, }, + "typesafe/jev-preview": { + "litellm_provider": "typesafe", + "mode": "evaluation", + "input_cost_per_token": 8.4e-08, + "output_cost_per_token": 0.0, + }, "fireworks_ai/accounts/fireworks/routers/rune": { "litellm_provider": "fireworks_ai", "mode": "evaluation", diff --git a/tests/integration/providers/test_decisions_wire.py b/tests/integration/providers/test_decisions_wire.py index 89ca674ef31..04b78985b20 100644 --- a/tests/integration/providers/test_decisions_wire.py +++ b/tests/integration/providers/test_decisions_wire.py @@ -137,7 +137,7 @@ _PROVIDERS: Final = ( "databricks-openjev-qwen35-4b", _API_KEY, False, - None, + "databricks/databricks-openjev-qwen35-4b", ), ) _PERPLEXITY: Final = _PROVIDERS[0] @@ -576,13 +576,26 @@ def test_upstream_errors_keep_their_status_and_log_an_unbilled_failure(gateway: ) -def test_upstream_success_without_answers_is_a_gateway_side_server_error(gateway: Gateway) -> None: +@pytest.mark.parametrize( + ("reply", "field"), + ( + pytest.param({"model": _PERPLEXITY.body_model, "usage": _USAGE}, "answers", id="missing_answers"), + pytest.param( + {"model": _PERPLEXITY.body_model, "answers": {"defect": {"type": "Request timed out"}}, "usage": _USAGE}, + "answers.defect", + id="answer_type_timeout_text", + ), + ), +) +def test_a_malformed_upstream_success_is_a_server_error_naming_the_field( + gateway: Gateway, reply: dict[str, JsonValue], field: str +) -> None: with gateway.scenario() as scenario: - handle: Final = _register(scenario, {"model": _PERPLEXITY.body_model, "usage": _USAGE}) + handle: Final = _register(scenario, reply) model: Final = _deployment(scenario, handle, _PERPLEXITY) response: Final = _decide(gateway, model, num_retries=0) - assert 500 <= response.status_code < 600, response.text - assert "answers" in response.text + assert response.status_code == 500, response.text + assert field in response.text assert _spend_row(response.headers["x-litellm-call-id"])["status"] == "failure" From 64d545ba7c832996a45ec1f36c6a27c8e46c0617 Mon Sep 17 00:00:00 2001 From: norkelbyte <348963847+norkelbyte@users.noreply.github.com> Date: Wed, 28 Oct 2026 04:48:30 -0700 Subject: [PATCH 026/128] test(router): audit #47376 (#48792) * test(router): audit #47376 batch create routing and supported_endpoints listing * test(router): make the batch mode cells deterministic and assert concurrent creates --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../test_batch_mode_deployments_wire.py | 426 ++++++++++++++++++ 1 file changed, 426 insertions(+) create mode 100644 tests/integration/routing/test_batch_mode_deployments_wire.py diff --git a/tests/integration/routing/test_batch_mode_deployments_wire.py b/tests/integration/routing/test_batch_mode_deployments_wire.py new file mode 100644 index 00000000000..7d702f499f5 --- /dev/null +++ b/tests/integration/routing/test_batch_mode_deployments_wire.py @@ -0,0 +1,426 @@ +import json +import uuid +from collections.abc import Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from typing import Final, Literal +from urllib.parse import urlsplit + +import pytest +from integration._support.client import ( + DEPLOYMENT_IDS, + Gateway, + Scenario, + eventually, + gateway_from_environment, + list_value, + object_value, + string_value, +) +from integration._support.openai_wire import chat_reply +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._cache_control_marks_support import owned_config +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.shares_worker_state + +_AZURE_API_VERSION: Final = "2024-10-21" +_BACKEND: Final = "gpt-5.4-mini" +_BATCH_ENDPOINT: Final = "/v1/batch" +_CHAT_ENDPOINT: Final = "/v1/chat/completions" +_RESPONSES_ENDPOINT: Final = "/v1/responses" +_TEAM_GROUP: Final = "batch-team-mix" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_JSON_LIST: Final = TypeAdapter(list[JsonValue]) +_ENDPOINTS: Final = TypeAdapter(list[str] | None) +_OWNED_PROXY_TIMEOUT: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +Provider = Literal["azure", "openai"] +Mode = Literal["chat", "batch"] + + +@dataclass(frozen=True, slots=True) +class _Sku: + provider: Provider + mode: Mode + key: str + weight: int | None + endpoints: JsonValue + team_id: str | None + + +@dataclass(frozen=True, slots=True) +class _Create: + status: int + body: str + + +@dataclass(frozen=True, slots=True) +class _TeamRig: + owned: OwnedProxy + wire: Wire + team_key: str + other_team_key: str + no_team_key: str + chat: _Sku + batch: _Sku + + +def _sku( + provider: Provider, + mode: Mode, + *, + weight: int | None = None, + endpoints: JsonValue = None, + team_id: str | None = None, +) -> _Sku: + return _Sku(provider, mode, f"{mode}-sku-{uuid.uuid4().hex}", weight, endpoints, team_id) + + +def _serves_batches(key: str) -> bool: + return key.startswith("batch-sku-") + + +def _deployment(name: str, wire: Wire, sku: _Sku) -> dict[str, JsonValue]: + endpoint: Final[dict[str, JsonValue]] = ( + {"model": f"azure/{_BACKEND}", "api_base": wire.url, "api_version": _AZURE_API_VERSION} + if sku.provider == "azure" + else {"model": f"openai/{_BACKEND}", "api_base": wire.url + "/v1"} + ) + weight: Final[dict[str, JsonValue]] = {} if sku.weight is None else {"weight": sku.weight} + endpoints: Final[dict[str, JsonValue]] = {} if sku.endpoints is None else {"supported_endpoints": sku.endpoints} + team: Final[dict[str, JsonValue]] = {} if sku.team_id is None else {"team_id": sku.team_id} + return { + "model_name": name, + "litellm_params": {**endpoint, "api_key": sku.key, **weight}, + "model_info": {"mode": sku.mode, **endpoints, **team}, + } + + +def _json_reply(body: Mapping[str, JsonValue], status: int = 200) -> Reply: + return Reply(status=status, body=json.dumps(body).encode()) + + +def _provider_key(request: Request) -> str: + return request.headers.get("api-key") or request.headers.get("authorization", "").removeprefix("Bearer ") + + +def _is_azure(request: Request) -> bool: + return urlsplit(request.target).path.startswith("/openai/") + + +def _posted_to(request: Request, suffix: str) -> bool: + return request.method == "POST" and urlsplit(request.target).path.endswith(suffix) + + +def _vendor(request: Request) -> Reply: + key: Final = _provider_key(request) + if _posted_to(request, "/files"): + return _json_reply( + { + "id": f"file-{key}", + "object": "file", + "bytes": len(request.body), + "created_at": 1700000000, + "filename": "in.jsonl", + "purpose": "batch", + "status": "processed", + } + ) + if _posted_to(request, "/batches") and _is_azure(request) and not _serves_batches(key): + return _json_reply( + {"error": {"code": "invalid_deployment_type", "message": "The deployment does not support batch jobs."}}, + 400, + ) + if _posted_to(request, "/batches"): + return _json_reply( + { + "id": f"batch-{uuid.uuid4().hex}", + "object": "batch", + "endpoint": _CHAT_ENDPOINT, + "input_file_id": f"file-{key}", + "status": "validating", + "created_at": 1700000000, + "completion_window": "24h", + } + ) + if _posted_to(request, "/chat/completions") and _is_azure(request) and _serves_batches(key): + return _json_reply( + {"error": {"code": "OperationNotSupported", "message": "Batch deployments do not serve chat."}}, 400 + ) + if _posted_to(request, "/chat/completions"): + return chat_reply(f"chatcmpl-{uuid.uuid4().hex}", _BACKEND, "served", stream=False) + return _json_reply({"data": [], "object": "list"}) + + +def _deployment_ids_on_a_worker(gateway: Gateway, name: str) -> tuple[str, frozenset[str]]: + response: Final = gateway.request( + "GET", + "/health/integration-deployment-ids", + params={"model_name": name}, + headers={"connection": "close"}, + ) + assert response.status_code == 200, response.text + return response.headers["x-integration-worker-pid"], frozenset( + DEPLOYMENT_IDS.validate_json(response.content)["model_ids"] + ) + + +def _wait_until_every_worker_serves(gateway: Gateway, name: str, identities: frozenset[str]) -> None: + from litellm.constants import PROXY_CONFIG_RELOAD_INTERVAL_SECONDS + + def answers() -> tuple[tuple[str, frozenset[str]], ...]: + with ThreadPoolExecutor(max_workers=8) as pool: + return tuple(pool.map(lambda _: _deployment_ids_on_a_worker(gateway, name), range(16))) + + eventually( + answers, + lambda seen: ( + len({pid for pid, _ in seen}) >= gateway.workers and all(identities <= served for _, served in seen) + ), + seconds=2 * PROXY_CONFIG_RELOAD_INTERVAL_SECONDS, + ) + + +def _add_deployment(scenario: Scenario, deployment: Mapping[str, JsonValue]) -> str: + created: Final = scenario.gateway.post("/model/new", deployment) + identity: Final = string_value(object_value(created["model_info"])["id"]) + scenario.cleanups.callback(scenario.delete_model, identity) + return identity + + +def _group(scenario: Scenario, wire: Wire, skus: Sequence[_Sku]) -> str: + name: Final = f"integration-batch-{uuid.uuid4().hex}" + identities: Final = frozenset(_add_deployment(scenario, _deployment(name, wire, sku)) for sku in skus) + _wait_until_every_worker_serves(scenario.gateway, name, identities) + return name + + +def _managed_upload(gateway: Gateway, group: str, key: str) -> str: + line: Final = {"custom_id": "r1", "method": "POST", "url": _CHAT_ENDPOINT, "body": {"model": group}} + uploaded: Final = gateway.request_multipart( + "/v1/files", + {"purpose": "batch", "target_model_names": group}, + {"file": ("in.jsonl", json.dumps(line).encode() + b"\n", "application/jsonl")}, + key=key, + ) + assert uploaded.status_code == 200, uploaded.text + return string_value(_JSON_OBJECT.validate_json(uploaded.content)["id"]) + + +def _create(gateway: Gateway, managed_file: str, key: str) -> _Create: + response: Final = gateway.request( + "POST", + "/v1/batches", + {"input_file_id": managed_file, "endpoint": _CHAT_ENDPOINT, "completion_window": "24h"}, + key=key, + ) + return _Create(response.status_code, response.text) + + +def _chat(gateway: Gateway, group: str) -> int: + return gateway.request( + "POST", + "/v1/chat/completions", + {"model": group, "messages": [{"role": "user", "content": "hello"}], "cache": {"no-cache": True}}, + ).status_code + + +def _keys_posted_to(requests: Sequence[Request], suffix: str) -> tuple[str, ...]: + return tuple(_provider_key(request) for request in requests if _posted_to(request, suffix)) + + +def _assert_creates_reach_only(wire: Wire, creates: Sequence[_Create], sku: _Sku) -> None: + assert [create.status for create in creates] == [200] * len(creates), creates + assert _keys_posted_to(wire.drain(), "/batches") == (sku.key,) * len(creates) + + +@pytest.mark.parametrize( + ("provider", "creates_reach"), + [("azure", "batch"), ("openai", "chat")], + ids=["azure-skips-chat", "openai-keeps-chat"], +) +def test_a_batch_create_on_a_mixed_group_reaches_the_deployment_its_provider_allows( + gateway: Gateway, provider: Provider, creates_reach: Mode +) -> None: + chat: Final = _sku(provider, "chat", weight=1) + batch: Final = _sku(provider, "batch", weight=0) + with wire_server(_vendor) as wire, gateway.scenario() as scenario: + group: Final = _group(scenario, wire, (chat, batch)) + managed_file: Final = _managed_upload(gateway, group, gateway.key) + assert set(_keys_posted_to(wire.drain(), "/files")) == {chat.key, batch.key} + + creates: Final = tuple(_create(gateway, managed_file, gateway.key) for _ in range(3)) + + _assert_creates_reach_only(wire, creates, batch if creates_reach == "batch" else chat) + + +def test_concurrent_batch_creates_and_chats_each_reach_their_own_deployment(gateway: Gateway) -> None: + chat: Final = _sku("azure", "chat", weight=1) + batch: Final = _sku("azure", "batch", weight=0) + with wire_server(_vendor) as wire, gateway.scenario() as scenario: + group: Final = _group(scenario, wire, (chat, batch)) + managed_file: Final = _managed_upload(gateway, group, gateway.key) + wire.drain() + + with ThreadPoolExecutor(max_workers=16) as pool: + burst: Final = tuple( + pool.map( + lambda turn: _create(gateway, managed_file, gateway.key) if turn % 2 else _chat(gateway, group), + range(24), + ) + ) + after: Final = tuple(_chat(gateway, group) for _ in range(6)) + received: Final = wire.drain() + + chat_statuses: Final = tuple(result for result in burst if isinstance(result, int)) + creates: Final = tuple(result for result in burst if isinstance(result, _Create)) + assert chat_statuses + after == (200,) * 18, (burst, after) + assert [create.status for create in creates] == [200] * 12, creates + assert _keys_posted_to(received, "/chat/completions") == (chat.key,) * 18 + assert _keys_posted_to(received, "/batches") == (batch.key,) * 12 + + +def _group_endpoints(gateway: Gateway, group: str, key: str | None = None) -> list[str] | None: + response: Final = gateway.request("GET", "/model_group/info", params={"model_group": group}, key=key) + assert response.status_code == 200, response.text + (entry,) = list_value(_JSON_OBJECT.validate_json(response.content)["data"]) + return _ENDPOINTS.validate_python(object_value(entry).get("supported_endpoints")) + + +@pytest.mark.parametrize( + ("shape", "expected"), + [ + pytest.param( + ( + ("azure", "chat", [_CHAT_ENDPOINT, _RESPONSES_ENDPOINT, _BATCH_ENDPOINT]), + ("azure", "batch", [_CHAT_ENDPOINT]), + ), + [_BATCH_ENDPOINT, _CHAT_ENDPOINT, _RESPONSES_ENDPOINT], + id="azure-chat-and-batch", + ), + pytest.param( + (("azure", "chat", [_CHAT_ENDPOINT, _BATCH_ENDPOINT]),), + [_CHAT_ENDPOINT], + id="azure-chat-only", + ), + pytest.param( + (("azure", "batch", [_CHAT_ENDPOINT, _RESPONSES_ENDPOINT]),), + [_BATCH_ENDPOINT], + id="azure-batch-only", + ), + pytest.param( + ( + ("openai", "chat", [_CHAT_ENDPOINT, _BATCH_ENDPOINT]), + ("openai", "batch", [_RESPONSES_ENDPOINT]), + ), + [_BATCH_ENDPOINT, _CHAT_ENDPOINT, _RESPONSES_ENDPOINT], + id="openai-chat-and-batch", + ), + ], +) +def test_model_group_info_lists_the_endpoints_its_deployments_serve( + gateway: Gateway, shape: tuple[tuple[Provider, Mode, list[str]], ...], expected: list[str] +) -> None: + skus: Final = tuple(_sku(provider, mode, endpoints=endpoints) for provider, mode, endpoints in shape) + with wire_server(_vendor) as wire, gateway.scenario() as scenario: + group: Final = _group(scenario, wire, skus) + + listed: Final = _group_endpoints(gateway, group) + + assert listed is not None and sorted(listed) == expected, listed + + +@pytest.mark.parametrize("malformed", ["/v1/batch", [7], {"batch": True}], ids=["string", "number-list", "object"]) +def test_a_malformed_endpoint_list_keeps_group_info_and_chat_answering(gateway: Gateway, malformed: JsonValue) -> None: + with wire_server(_vendor) as wire, gateway.scenario() as scenario: + group: Final = _group(scenario, wire, (_sku("openai", "chat", endpoints=malformed),)) + + _group_endpoints(gateway, group) + + assert _chat(gateway, group) == 200 + + +@pytest.fixture(scope="module") +def team_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_TeamRig]: + directory: Final = tmp_path_factory.mktemp("batch-team") + team_id: Final = f"integration-batch-team-{uuid.uuid4().hex}" + chat: Final = _sku("azure", "chat", endpoints=[_CHAT_ENDPOINT]) + batch: Final = _sku("azure", "batch", team_id=team_id) + with gateway_from_environment() as gateway, wire_server(_vendor) as wire: + config: Final = owned_config( + directory, + [_deployment(_TEAM_GROUP, wire, chat), _deployment(_TEAM_GROUP, wire, batch)], + litellm_settings={"public_model_groups": [_TEAM_GROUP]}, + router_settings={"routing_strategy": "usage-based-routing"}, + ) + with owned_proxy_process(gateway, directory, {}, config=config) as owned, owned.gateway.scenario() as scenario: + scenario.team(team_id=team_id) + yield _TeamRig( + owned=owned, + wire=wire, + team_key=scenario.key(team_id=team_id), + other_team_key=scenario.key(team_id=scenario.team()), + no_team_key=scenario.key(), + chat=chat, + batch=batch, + ) + + +def _caller_key(rig: _TeamRig, caller: str) -> str: + return {"team": rig.team_key, "other-team": rig.other_team_key, "no-team": rig.no_team_key}[caller] + + +@_OWNED_PROXY_TIMEOUT +@pytest.mark.parametrize( + ("caller", "lists_batch"), + [("team", True), ("other-team", False), ("no-team", False)], + ids=["team", "other-team", "no-team"], +) +def test_model_group_info_lists_a_team_owned_batch_endpoint_only_to_that_team( + team_rig: _TeamRig, caller: str, lists_batch: bool +) -> None: + listed: Final = _group_endpoints(team_rig.owned.gateway, _TEAM_GROUP, _caller_key(team_rig, caller)) + + assert listed is not None and _CHAT_ENDPOINT in listed, listed + assert (_BATCH_ENDPOINT in listed) is lists_batch, listed + + +@_OWNED_PROXY_TIMEOUT +def test_public_model_hub_never_lists_a_team_owned_batch_endpoint(team_rig: _TeamRig) -> None: + response: Final = team_rig.owned.gateway.client.get("/public/model_hub") + assert response.status_code == 200, response.text + (entry,) = [ + object_value(row) + for row in _JSON_LIST.validate_json(response.content) + if object_value(row)["model_group"] == _TEAM_GROUP + ] + + assert _BATCH_ENDPOINT not in (_ENDPOINTS.validate_python(entry.get("supported_endpoints")) or []), entry + + +@_OWNED_PROXY_TIMEOUT +def test_a_legacy_strategy_sends_a_team_batch_create_to_its_batch_deployment(team_rig: _TeamRig) -> None: + gateway: Final = team_rig.owned.gateway + team_rig.wire.drain() + managed_file: Final = _managed_upload(gateway, _TEAM_GROUP, team_rig.team_key) + assert team_rig.batch.key in _keys_posted_to(team_rig.wire.drain(), "/files") + + creates: Final = tuple(_create(gateway, managed_file, team_rig.team_key) for _ in range(3)) + + _assert_creates_reach_only(team_rig.wire, creates, team_rig.batch) + + +@_OWNED_PROXY_TIMEOUT +@pytest.mark.parametrize("caller", ["other-team", "no-team"]) +def test_a_legacy_strategy_keeps_other_callers_off_a_team_batch_deployment(team_rig: _TeamRig, caller: str) -> None: + gateway: Final = team_rig.owned.gateway + key: Final = _caller_key(team_rig, caller) + team_rig.wire.drain() + managed_file: Final = _managed_upload(gateway, _TEAM_GROUP, key) + team_rig.wire.drain() + + tuple(_create(gateway, managed_file, key) for _ in range(3)) + + assert _keys_posted_to(team_rig.wire.drain(), "/batches") == (team_rig.chat.key,) * 3 From 93056ef93929fa2e2b5aeb2890172cb950c261af Mon Sep 17 00:00:00 2001 From: norkelbyte <348963847+norkelbyte@users.noreply.github.com> Date: Wed, 28 Oct 2026 04:55:44 -0700 Subject: [PATCH 027/128] test(proxy): audit #47358 (#48784) * test(proxy): audit #47358 Add integration cells for pass-through bodies and settings-endpoint mappings across workers * test(proxy): wait for every worker to hold the saved mapping before probing * test(proxy): compare passthrough bodies to the client body and probe each worker on its own connection --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../spend/test_baggage_tag_mappings.py | 229 +++++++++++++++++- 1 file changed, 222 insertions(+), 7 deletions(-) diff --git a/tests/integration/spend/test_baggage_tag_mappings.py b/tests/integration/spend/test_baggage_tag_mappings.py index 1bf90269484..e08348a074c 100644 --- a/tests/integration/spend/test_baggage_tag_mappings.py +++ b/tests/integration/spend/test_baggage_tag_mappings.py @@ -2,8 +2,9 @@ import itertools import json import uuid from collections.abc import Generator, Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor from contextlib import ExitStack, contextmanager -from dataclasses import dataclass +from dataclasses import dataclass, replace from hashlib import sha256 from pathlib import Path from types import MappingProxyType @@ -15,14 +16,15 @@ import yaml from integration._support.client import ( JSON_OBJECT, Gateway, + Scenario, eventually, gateway_from_environment, list_value, object_value, ) -from integration._support.database import scratch_database +from integration._support.database import read_rows, scratch_database, write_rows from integration._support.database_relay import dropped_connection_relay, plaintext_database_url -from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter from websockets.exceptions import InvalidStatus @@ -33,6 +35,8 @@ pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" PRICED: Final[dict[str, JsonValue]] = {"input_cost_per_token": 0.01, "output_cost_per_token": 0.01} WORKER_PROBES: Final = 6 +VLLM_UPSTREAM_MODEL: Final = "baggage-model" +PINNED_CONNECTION: Final = httpx.Limits(max_connections=1) JSON_LIST: Final = TypeAdapter(list[JsonValue]) GUARDRAIL_VERDICT: Final = b'{"action":"GUARDRAIL_INTERVENED","texts":["masked baggage prompt"]}' GUARDRAIL_PATH: Final = "/beta/litellm_basic_guardrail_api" @@ -40,6 +44,7 @@ NO_FIELDS: Final[Mapping[str, JsonValue]] = MappingProxyType({}) Route = Literal["chat", "chat_litellm_metadata", "messages", "responses"] CallerTag = Literal["none", "header", "metadata_json_with_header", "litellm_metadata"] +PassThrough = Literal["custom", "anthropic", "vllm"] RejectedClientTag = Literal["none", "header", "body", "metadata_json", "litellm_metadata"] @@ -120,11 +125,31 @@ def _responses_reply() -> Reply: ) +def _anthropic_reply() -> Reply: + marker: Final = uuid.uuid4().hex + return Reply( + body=json.dumps( + { + "id": f"msg_{marker}", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": marker}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 3}, + } + ).encode() + ) + + def _provider(request: Request) -> Reply: if request.target == GUARDRAIL_PATH: return Reply(body=b'{"action":"NONE"}') if request.target == "/v1/responses": return _responses_reply() + if request.target == "/v1/messages": + return _anthropic_reply() return _chat_reply() @@ -183,7 +208,15 @@ def _mapped_models(marker: str, provider: Wire, upstream_url: str) -> list[JsonV } guarded: Final = _deployment(f"guarded-{marker}", provider, tags=[f"guarded-{marker}:on"]) pre_guarded: Final = _deployment(f"guarded-pre-{marker}", provider, tags=[f"guarded-{marker}:pre"]) - return [_deployment(f"chat-{marker}", provider), guarded, pre_guarded, *routed, realtime] + vllm: Final[dict[str, JsonValue]] = { + "model_name": f"vllm-{marker}", + "litellm_params": { + "model": f"hosted_vllm/{VLLM_UPSTREAM_MODEL}", + "api_base": provider.url, + "api_key": "synthetic-vllm", + }, + } + return [_deployment(f"chat-{marker}", provider), guarded, pre_guarded, *routed, realtime, vllm] def _mapped_rules(marker: str, *, tenant_required: bool) -> list[JsonValue]: @@ -239,9 +272,12 @@ def _rig( {"callbacks": ["generic_api"], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1}, [_deployment_tag_mode_guardrail(marker, provider)] if tag_mode_guardrail else (), ) - with owned_proxy( - base, directory, {"GENERIC_LOGGER_ENDPOINT": callbacks.url}, config=config, workers=workers - ) as proxy: + environment: Final = { + "GENERIC_LOGGER_ENDPOINT": callbacks.url, + "ANTHROPIC_API_BASE": provider.url, + "ANTHROPIC_API_KEY": "synthetic-anthropic-key", + } + with owned_proxy(base, directory, environment, config=config, workers=workers) as proxy: yield Rig(proxy, provider, callbacks, marker) @@ -676,6 +712,185 @@ def test_custom_passthrough_spend_row_carries_the_mapped_tags(mapped: Rig) -> No assert {mapped.exact_tag, mapped.tenant_tag(tenant)} <= _spend_tags(mapped.proxy, key) +def _passthrough_path(rig: Rig, scenario: Scenario, entrypoint: PassThrough) -> str: + match entrypoint: + case "custom": + path: Final = f"/baggage-passthrough-{uuid.uuid4().hex}" + created: Final = rig.proxy.post( + "/config/pass_through_endpoint", {"path": path, "target": f"{rig.provider.url}/echo", "auth": True} + ) + endpoint_id: Final = str(object_value(list_value(created["endpoints"])[0])["id"]) + scenario.cleanups.callback( + lambda: rig.proxy.request( + "DELETE", "/config/pass_through_endpoint", params={"endpoint_id": endpoint_id} + ) + ) + return path + case "anthropic": + return "/anthropic/v1/messages" + case "vllm": + return "/vllm/v1/chat/completions" + + +def _passthrough_body(rig: Rig, entrypoint: PassThrough, text: str) -> dict[str, JsonValue]: + model: Final = {"custom": rig.chat_model, "anthropic": "claude-sonnet-4-5", "vllm": f"vllm-{rig.marker}"} + body: Final[dict[str, JsonValue]] = { + "model": model[entrypoint], + "max_tokens": 16, + "messages": [{"role": "user", "content": text}], + } + match entrypoint: + case "vllm": + return {**body, "metadata": {"user_id": f"user-{text}"}} + case "custom" | "anthropic": + return body + + +def _upstream_body(entrypoint: PassThrough, body: Mapping[str, JsonValue]) -> JsonValue: + match entrypoint: + case "vllm": + return {**body, "model": VLLM_UPSTREAM_MODEL} + case "custom" | "anthropic": + return dict(body) + + +def _forwarded_bodies(rig: Rig) -> tuple[JsonValue, ...]: + return tuple( + JSON_OBJECT.validate_json(request.body) for request in rig.provider.drain() if request.target != GUARDRAIL_PATH + ) + + +@pytest.mark.parametrize("entrypoint", ["custom", "anthropic", "vllm"]) +def test_passthrough_forwards_the_client_body_unchanged_while_its_spend_row_carries_the_mapped_tags( + optional: Rig, entrypoint: PassThrough +) -> None: + tenant: Final = uuid.uuid4().hex + with optional.proxy.scenario() as scenario: + path: Final = _passthrough_path(optional, scenario, entrypoint) + key_fields: Final[dict[str, JsonValue]] = ( + {"allowed_passthrough_routes": [path]} if entrypoint == "custom" else {} + ) + plain_key: Final = scenario.key(**key_fields) + key: Final = scenario.key(**key_fields) + body: Final = _passthrough_body(optional, entrypoint, tenant) + optional.provider.drain() + plain: Final = optional.proxy.request("POST", path, body, key=plain_key) + plain_forwarded: Final = _forwarded_bodies(optional) + response: Final = optional.proxy.request( + "POST", path, body, key=key, headers={"baggage": f"entrypoint=a,tenant={tenant}"} + ) + assert (plain.status_code, response.status_code) == (200, 200), (plain.text, response.text) + assert plain_forwarded == (_upstream_body(entrypoint, body),) + assert _forwarded_bodies(optional) == plain_forwarded + mapped_tags: Final = {optional.exact_tag, optional.tenant_tag(tenant)} + assert _spend_tags(optional.proxy, key) == _spend_tags(optional.proxy, plain_key) | mapped_tags + + +def _general_settings_mappings(database_url: str) -> JsonValue: + rows: Final = read_rows( + 'SELECT param_value FROM "LiteLLM_Config" WHERE param_name = %s', + ("general_settings",), + database_url=database_url, + ) + return object_value(rows[0]["param_value"]).get("request_tag_mappings") if rows else None + + +def _serving_pid(proxy: Gateway) -> int: + response: Final = proxy.request("GET", "/health/integration-readiness") + return int(response.headers["x-integration-worker-pid"]) + + +def _status_on_one_worker(owned: OwnedProxy, rig: Rig, key: str) -> tuple[int, int] | None: + with httpx.Client( + base_url=str(owned.gateway.client.base_url), timeout=15, trust_env=False, limits=PINNED_CONNECTION + ) as client: + pinned: Final = replace(rig, proxy=replace(owned.gateway, client=client)) + before: Final = _serving_pid(pinned.proxy) + status: Final = _chat(pinned, key, None).status_code + return (before, status) if _serving_pid(pinned.proxy) == before else None + + +def _worker_statuses(owned: OwnedProxy, rig: Rig, key: str) -> frozenset[tuple[int, int]]: + def answers() -> frozenset[tuple[int, int]]: + with ThreadPoolExecutor(max_workers=8) as pool: + found: Final = tuple(pool.map(lambda _: _status_on_one_worker(owned, rig, key), range(8))) + return frozenset(answer for answer in found if answer is not None) + + return eventually(answers, lambda found: frozenset(pid for pid, _ in found) >= owned.worker_pids(), seconds=30) + + +def _mappings_on_a_worker(proxy: Gateway) -> tuple[int, JsonValue]: + response: Final = proxy.request( + "GET", + "/health/integration-general-setting", + params={"name": "request_tag_mappings"}, + headers={"connection": "close"}, + ) + assert response.status_code == 200, response.text + return int(response.headers["x-integration-worker-pid"]), JSON_OBJECT.validate_json(response.content)["value"] + + +def _mappings_settled_on_every_worker(owned: OwnedProxy, mappings: JsonValue) -> None: + def answers() -> tuple[tuple[int, JsonValue], ...]: + with ThreadPoolExecutor(max_workers=8) as pool: + return tuple(pool.map(lambda _: _mappings_on_a_worker(owned.gateway), range(16))) + + eventually( + answers, + lambda found: ( + frozenset(pid for pid, _ in found) >= owned.worker_pids() and all(value == mappings for _, value in found) + ), + seconds=70, + ) + + +def test_mappings_saved_through_the_settings_endpoint_apply_on_every_worker_and_an_invalid_stored_value_is_ignored( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = uuid.uuid4().hex[:12] + rule: Final[dict[str, JsonValue]] = {"key": "tenant", "tag_prefix": f"tenant-{marker}:", "required": True} + with ExitStack() as stack: + provider: Final = stack.enter_context(wire_server(_provider)) + config: Final = _write_config(tmp_path, [_deployment(f"chat-{marker}", provider)], {}, {}) + owned: Final = stack.enter_context(owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2)) + proxy: Final = owned.gateway + database_url: Final = proxy.database_url + assert database_url is not None + scenario: Final = stack.enter_context(proxy.scenario()) + rig: Final = Rig(proxy, provider, provider, marker) + key: Final = scenario.key() + + def save(rules: list[JsonValue]) -> httpx.Response: + return proxy.request( + "POST", + "/config/field/update", + {"field_name": "request_tag_mappings", "field_value": rules, "config_type": "general_settings"}, + ) + + duplicate: Final = save([rule, {**rule, "required": False}]) + assert duplicate.status_code == 400, duplicate.text + assert _general_settings_mappings(database_url) is None + saved: Final = save([rule]) + assert saved.status_code == 200, saved.text + _mappings_settled_on_every_worker(owned, [rule]) + assert _worker_statuses(owned, rig, key) == frozenset((pid, 400) for pid in owned.worker_pids()) + tagged_key: Final = scenario.key() + provider.drain() + admitted: Final = _chat(rig, tagged_key, f"tenant={marker}") + assert admitted.status_code == 200, admitted.text + assert len(provider.drain()) == 1 + assert _spend_tags(proxy, tagged_key) == {f"tenant-{marker}:{marker}"} + + write_rows( + 'UPDATE "LiteLLM_Config" SET param_value = jsonb_set(param_value::jsonb, %s, %s::jsonb) ' + "WHERE param_name = %s", + ("{request_tag_mappings}", json.dumps([{"key": "tenant", "required": True}]), "general_settings"), + database_url=database_url, + ) + _mappings_settled_on_every_worker(owned, None) + assert _worker_statuses(owned, rig, key) == frozenset((pid, 200) for pid in owned.worker_pids()) + + def _rejected_client_tag(client_tag: RejectedClientTag, tag: str) -> tuple[Mapping[str, JsonValue], Mapping[str, str]]: match client_tag: case "none": From 826859588ce1f045fb1b9800e868eb5bb28dd9db Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Wed, 28 Oct 2026 05:14:59 -0700 Subject: [PATCH 028/128] test(router): audit #41368 (#48794) * test(router): cover silent_model shadow fan-out and streaming end to end Proxy cells check that each silent_model target gets one request in the caller's stream mode, that the caller sees only the primary, and that every shadow lands its own spend row under its own model group. SDK cells check the sync Router.completion thread path streams each shadow and fires its success callback. * test(router): flatten the spend row lookup with itertools.chain --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../routing/test_silent_model_shadow_wire.py | 195 ++++++++++++++++++ .../test_router_silent_model_shadow_sdk.py | 129 ++++++++++++ 2 files changed, 324 insertions(+) create mode 100644 tests/integration/routing/test_silent_model_shadow_wire.py create mode 100644 tests/integration/sdk/test_router_silent_model_shadow_sdk.py diff --git a/tests/integration/routing/test_silent_model_shadow_wire.py b/tests/integration/routing/test_silent_model_shadow_wire.py new file mode 100644 index 00000000000..b00fbca645e --- /dev/null +++ b/tests/integration/routing/test_silent_model_shadow_wire.py @@ -0,0 +1,195 @@ +from __future__ import annotations + +import json +import re +import uuid +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from itertools import chain +from types import MappingProxyType +from typing import Final, Literal +from urllib.parse import unquote + +import httpx +import pytest +from integration._support.anthropic_sse import parse_sse +from integration._support.client import Gateway, eventually, gateway_from_environment +from integration._support.openai_wire import chat_reply +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +Shape = Literal["single", "pair"] + +_MODEL: Final = "gpt-5.6" +_API_KEY: Final = "synthetic-shadow-key" +_ROUTER_ONLY_KEYS: Final = ("silent_model", "include_fallback_errors") +_ANSWER: Final = "answered by" +_NONCE: Final = re.compile(r"nonce=([0-9a-f]{32})") +_JSON: Final = TypeAdapter(dict[str, JsonValue]) +_CHOICES: Final = TypeAdapter(list[dict[str, JsonValue]]) +_FRESH_CONNECTION: Final = MappingProxyType({"Connection": "close"}) +_SPEND_ROW_SECONDS: Final = 70 +_TARGETS: Final[Mapping[Shape, tuple[str, ...]]] = MappingProxyType( + {"single": ("shadow-a",), "pair": ("shadow-a", "shadow-b")} +) +_CASES: Final = ( + pytest.param("single", False, id="non-stream-single"), + pytest.param("single", True, id="stream-single"), + pytest.param("pair", True, id="stream-pair"), +) + + +def _prompt(nonce: str) -> str: + return f"which deployment answers this? nonce={nonce}" + + +def _nonce_of(raw: str) -> str: + found: Final = _NONCE.search(raw) + assert found is not None, raw + return found.group(1) + + +def _identity(segment: str, nonce: str) -> str: + return f"chatcmpl-{segment}-{nonce}" + + +def _request(model: str, nonce: str, *, stream: bool) -> dict[str, JsonValue]: + body: Final[dict[str, JsonValue]] = { + "model": model, + "stream": stream, + "messages": [{"role": "user", "content": _prompt(nonce)}], + } + return {**body, "stream_options": {"include_usage": True}} if stream else body + + +def _segment_of(request: Request) -> str: + return unquote(request.target).lstrip("/").partition("/")[0] + + +def _peer(request: Request) -> Reply: + if request.method == "GET": + return Reply(body=json.dumps({"object": "list", "data": []}).encode()) + segment: Final = _segment_of(request) + body: Final = _JSON.validate_json(request.body) + return chat_reply( + _identity(segment, _nonce_of(request.body.decode())), + _MODEL, + f"{_ANSWER} {segment}", + stream=body.get("stream") is True, + ) + + +@dataclass(frozen=True, slots=True) +class _Posted: + segment: str + streamed: bool + leaked: tuple[str, ...] + + +def _posted(wire: Wire, nonce: str) -> tuple[_Posted, ...]: + return tuple( + _Posted( + _segment_of(request), + _JSON.validate_json(request.body).get("stream") is True, + tuple(key for key in _ROUTER_ONLY_KEYS if key in request.body.decode()), + ) + for request in wire.drain() + if request.method == "POST" and nonce in request.body.decode() + ) + + +def _delta_text(data: Mapping[str, JsonValue]) -> str: + deltas: Final = ( + _JSON.validate_python(choice.get("delta") or {}) for choice in _CHOICES.validate_python(data["choices"]) + ) + return "".join(str(delta.get("content") or "") for delta in deltas) + + +def _served(response: httpx.Response, *, stream: bool) -> tuple[frozenset[str], str]: + if not stream: + body: Final = _JSON.validate_json(response.content) + message: Final = _JSON.validate_python(_CHOICES.validate_python(body["choices"])[0]["message"]) + return frozenset({str(body["id"])}), str(message["content"]) + events: Final = tuple(event.data for event in parse_sse(response.text)) + return frozenset(str(data["id"]) for data in events), "".join(_delta_text(data) for data in events) + + +@dataclass(frozen=True, slots=True) +class _SpendRow: + request_id: str + model_group: str + spend: float + + +def _spend_row(gateway: Gateway, request_id: str) -> tuple[_SpendRow, ...]: + rows: Final = gateway.rows( + 'SELECT model_group, spend FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (request_id,) + ) + return tuple(_SpendRow(request_id, str(row["model_group"]), float(str(row["spend"]))) for row in rows) + + +def _spend_rows(gateway: Gateway, request_ids: tuple[str, ...]) -> tuple[_SpendRow, ...]: + return tuple(chain.from_iterable(_spend_row(gateway, request_id) for request_id in request_ids)) + + +@dataclass(frozen=True, slots=True) +class _Mirrors: + gateway: Gateway + wire: Wire + groups: Mapping[str, str] + + +@pytest.fixture(scope="module") +def mirrors() -> Iterator[_Mirrors]: + with gateway_from_environment() as gateway, wire_server(_peer) as wire, gateway.scenario() as scenario: + + def register(segment: str, silent_model: JsonValue) -> str: + return scenario.model( + model=f"openai/{_MODEL}", api_key=_API_KEY, api_base=f"{wire.url}/{segment}", silent_model=silent_model + ) + + shadow_a: Final = register("shadow-a", None) + shadow_b: Final = register("shadow-b", None) + groups: Final = MappingProxyType( + { + "shadow-a": shadow_a, + "shadow-b": shadow_b, + "primary-single": register("primary-single", shadow_a), + "primary-pair": register("primary-pair", [shadow_a, shadow_b]), + } + ) + wire.drain() + yield _Mirrors(gateway, wire, groups) + + +@pytest.mark.parametrize(("shape", "stream"), _CASES) +def test_each_shadow_target_gets_one_request_in_the_callers_stream_mode( + mirrors: _Mirrors, shape: Shape, stream: bool +) -> None: + nonce: Final = uuid.uuid4().hex + primary: Final = f"primary-{shape}" + targets: Final = _TARGETS[shape] + response: Final = mirrors.gateway.request( + "POST", + "/v1/chat/completions", + _request(mirrors.groups[primary], nonce, stream=stream), + headers=_FRESH_CONNECTION, + ) + assert response.status_code == 200, response.text + assert _served(response, stream=stream) == (frozenset({_identity(primary, nonce)}), f"{_ANSWER} {primary}") + + expected_groups: Final = MappingProxyType( + {_identity(segment, nonce): mirrors.groups[segment] for segment in (primary, *targets)} + ) + rows: Final = eventually( + lambda: _spend_rows(mirrors.gateway, tuple(expected_groups)), + lambda found: len(found) == len(expected_groups), + seconds=_SPEND_ROW_SECONDS, + ) + assert {row.request_id: row.model_group for row in rows} == expected_groups, rows + assert all(row.spend > 0 for row in rows), rows + + posted: Final = _posted(mirrors.wire, nonce) + assert sorted(item.segment for item in posted) == sorted((primary, *targets)), posted + assert all(item.streamed is stream for item in posted), posted + assert all(item.leaked == () for item in posted), posted diff --git a/tests/integration/sdk/test_router_silent_model_shadow_sdk.py b/tests/integration/sdk/test_router_silent_model_shadow_sdk.py new file mode 100644 index 00000000000..821e4ddd0b7 --- /dev/null +++ b/tests/integration/sdk/test_router_silent_model_shadow_sdk.py @@ -0,0 +1,129 @@ +from __future__ import annotations + +import re +import uuid +from collections.abc import Mapping +from dataclasses import dataclass +from queue import SimpleQueue +from typing import Final + +import pytest +from integration._support.client import eventually +from integration._support.openai_wire import chat_reply +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +import litellm +from litellm import CustomStreamWrapper, Router +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.utils import ModelResponseStream + +_MODEL: Final = "gpt-5.6" +_API_KEY: Final = "synthetic-shadow-key" +_ANSWER: Final = "answered by" +_NONCE: Final = re.compile(r"nonce=([0-9a-f]{32})") +_JSON: Final = TypeAdapter(dict[str, JsonValue]) +_OBJECT: Final = TypeAdapter(dict[str, object]) +_SHADOW_SECONDS: Final = 20 +_TARGETS: Final = ( + pytest.param(("shadow-a",), id="single"), + pytest.param(("shadow-a", "shadow-b"), id="pair"), +) + + +def _prompt(nonce: str) -> str: + return f"which deployment answers this? nonce={nonce}" + + +def _peer(request: Request) -> Reply: + deployment: Final = request.target.lstrip("/").partition("/")[0] + raw: Final = request.body.decode() + found: Final = _NONCE.search(raw) + assert found is not None, raw + return chat_reply( + f"chatcmpl-{deployment}-{found.group(1)}", + _MODEL, + f"{_ANSWER} {deployment}", + stream=_JSON.validate_json(request.body).get("stream") is True, + ) + + +def _deployment(name: str, wire: Wire, **extra: JsonValue) -> dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": {"model": f"openai/{_MODEL}", "api_base": f"{wire.url}/{name}", "api_key": _API_KEY, **extra}, + } + + +@dataclass(frozen=True, slots=True) +class _Posted: + deployment: str + streamed: bool + leaked: bool + + +def _posted(wire: Wire) -> tuple[_Posted, ...]: + return tuple( + _Posted( + request.target.lstrip("/").partition("/")[0], + _JSON.validate_json(request.body).get("stream") is True, + "silent_model" in request.body.decode(), + ) + for request in wire.drain() + if request.method == "POST" + ) + + +class _ShadowRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.shadow_groups: SimpleQueue[str] = SimpleQueue() + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + metadata: Final = _OBJECT.validate_python(_OBJECT.validate_python(kwargs["litellm_params"])["metadata"] or {}) + if metadata.get("is_silent_experiment") is True: + self.shadow_groups.put(str(metadata.get("model_group"))) + + def drained(self) -> tuple[str, ...]: + return tuple(self.shadow_groups.get_nowait() for _ in range(self.shadow_groups.qsize())) + + +def _delta(chunk: object) -> str: + assert isinstance(chunk, ModelResponseStream), chunk + return "".join(str(choice.delta.content or "") for choice in chunk.choices) + + +@pytest.mark.parametrize("targets", _TARGETS) +def test_sync_streamed_completion_streams_and_logs_each_shadow_target( + targets: tuple[str, ...], monkeypatch: pytest.MonkeyPatch +) -> None: + recorder: Final = _ShadowRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + with wire_server(_peer) as wire: + silent_model: Final[JsonValue] = targets[0] if len(targets) == 1 else [*targets] + router: Final = Router( + model_list=[ + _deployment("mirrored", wire, silent_model=silent_model), + *(_deployment(target, wire) for target in targets), + ], + num_retries=0, + disable_cooldowns=True, + ) + response: Final = router.completion( + model="mirrored", + messages=[{"role": "user", "content": _prompt(uuid.uuid4().hex)}], + stream=True, + ) + assert isinstance(response, CustomStreamWrapper), response + text: Final = "".join(_delta(chunk) for chunk in response) + eventually(wire.received.qsize, lambda count: count >= 1 + len(targets), seconds=_SHADOW_SECONDS) + posted: Final = _posted(wire) + eventually(recorder.shadow_groups.qsize, lambda count: count >= len(targets), seconds=_SHADOW_SECONDS) + logged: Final = recorder.drained() + assert text == f"{_ANSWER} mirrored" + assert sorted(item.deployment for item in posted) == sorted(("mirrored", *targets)), posted + assert all(item.streamed for item in posted), posted + assert not any(item.leaked for item in posted), posted + assert sorted(logged) == sorted(targets), logged From 940b865424c5c25a2715f515c911b037f8708fba Mon Sep 17 00:00:00 2001 From: norkelbyte <348963847+norkelbyte@users.noreply.github.com> Date: Wed, 28 Oct 2026 05:17:28 -0700 Subject: [PATCH 029/128] refactor(fallback-generalizations)!: remove the legacy extends and mixed-rule shim (#48737) * refactor(fallback-generalizations): remove the legacy extends and mixed-rule shim Main's remote cost map, the bundled backup, and the root map all ship the split routing/capability rule format, so the shim for the old format is dead code. A rule mixing litellm_provider with other model_info keys is now warned about and skipped. * test(fallback-generalizations): fail when a shipped cost map rule would be skipped at install * test(fallback-generalizations): cover ignored extends and annotate the shipped-rules test * fix(fallback-generalizations): warn about old-format rules with a pointer to a current cost map and skip extends rules * fix(fallback-generalizations): tell local cost map users to replace the bundled backup in the old-format warning --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../fallback_generalizations.py | 117 ++++++--------- .../test_fallback_generalizations.py | 138 ++++++++---------- 2 files changed, 98 insertions(+), 157 deletions(-) diff --git a/litellm/litellm_core_utils/fallback_generalizations.py b/litellm/litellm_core_utils/fallback_generalizations.py index be432d1fdd4..eb4cf421a19 100644 --- a/litellm/litellm_core_utils/fallback_generalizations.py +++ b/litellm/litellm_core_utils/fallback_generalizations.py @@ -19,17 +19,11 @@ ones on key conflicts, and the caller backfills ``litellm_provider`` with the provider it requested. If no capability rule matches, model-info resolution misses as if no rules existed. -LEGACY-SCHEMA SHIM (temporary, until the new-schema JSON reaches main): released -proxies fetch this JSON remotely from main, whose block still ships the old schema -where a rule mixes ``litellm_provider`` with capability keys and may inherit a -parent's ``model_info`` via ``extends``. Such a legacy rule is tolerated rather -than skipped: ``extends`` is resolved once at install time (single level, against -raw parents), and the resolved rule acts as BOTH kinds, a routing rule (its -``litellm_provider`` participates in first-hit inference) and a capability rule -(its full ``model_info``, provider included, participates in the union). New-schema -rules never mix the two and never use ``extends``. A rule whose -``litellm_provider`` is not a string is invalid and is warned about and skipped -(a warning rather than a crash, for the same remote-fetch reason). +A rule whose ``model_info`` carries ``litellm_provider`` alongside any other key, +or carries it as a non-string, is neither kind: it is warned about and skipped +rather than raising, since released proxies fetch this JSON remotely. A rule +carrying ``extends`` is skipped the same way. Mixed rules and ``extends`` are the +old rule format, so their warning points the operator at a current cost map. Rules are only consulted after exact and case-insensitive lookups miss, so an exact cost-map entry always takes precedence over any rule. @@ -62,37 +56,8 @@ NAME_FIELD: Final = "name" PATTERN_FIELD: Final = "pattern" MODEL_INFO_FIELD: Final = "model_info" PROVIDER_KEY: Final = "litellm_provider" -LEGACY_EXTENDS_FIELD: Final = "extends" FILL_MISSING_FOR_PROVIDERS_FIELD: Final = "fill_missing_for_providers" - - -def _resolve_legacy_extends(rules: list) -> list: - """Expand legacy ``extends`` inheritance so each rule's ``model_info`` is self-contained. - - Compatibility shim for the old remote schema: single level, resolved against each - parent's raw ``model_info``, with the child's own keys winning on conflict. Non-dict - rules and dangling parents pass through unchanged; new-schema rules carry no - ``extends`` and are untouched. - """ - base_by_name: Final = { - rule[NAME_FIELD]: rule[MODEL_INFO_FIELD] - for rule in rules - if isinstance(rule, dict) - and isinstance(rule.get(NAME_FIELD), str) - and isinstance(rule.get(MODEL_INFO_FIELD), dict) - } - - def resolved(rule: object) -> object: - if not isinstance(rule, dict): - return rule - parent_name: Final = rule.get(LEGACY_EXTENDS_FIELD) - own_info: Final = rule.get(MODEL_INFO_FIELD) - parent_info: Final = base_by_name.get(parent_name) if isinstance(parent_name, str) else None - if parent_info is None or not isinstance(own_info, dict): - return rule - return {**rule, MODEL_INFO_FIELD: {**parent_info, **own_info}} - - return [resolved(rule) for rule in rules] +EXTENDS_FIELD: Final = "extends" @dataclass(frozen=True, slots=True) @@ -127,19 +92,33 @@ def _parse_fill_missing_for_providers(rule: Mapping[str, object], pattern_label: return frozenset(raw_fill_missing_for_providers) -def _compile_rule(rule: object) -> tuple[_CompiledRule, ...]: +def _warn_old_format_rule(rule_label: str, reason: str) -> None: + verbose_logger.warning( + "LiteLLM: skipping fallback generalization rule %s, written in the old cost map rule format (%s). " + "Point LITELLM_MODEL_COST_MAP_URL at a current cost map, or with LITELLM_LOCAL_MODEL_COST_MAP set, " + "replace the bundled backup with a current one.", + rule_label, + reason, + ) + + +def _compile_rule(rule: object) -> _CompiledRule | None: if not isinstance(rule, dict): - return () + return None pattern: Final = rule.get(PATTERN_FIELD) model_info: Final = rule.get(MODEL_INFO_FIELD) + rule_label: Final = f"{rule.get(NAME_FIELD, pattern)}" if not isinstance(pattern, str) or not isinstance(model_info, dict): verbose_logger.warning( "LiteLLM: skipping malformed fallback generalization rule %s (needs string '%s' and dict '%s').", - rule.get(NAME_FIELD, pattern), + rule_label, PATTERN_FIELD, MODEL_INFO_FIELD, ) - return () + return None + if EXTENDS_FIELD in rule: + _warn_old_format_rule(rule_label, f"'{EXTENDS_FIELD}' is no longer supported") + return None try: compiled: Final = re.compile(pattern, re.IGNORECASE) except re.error as e: @@ -148,37 +127,28 @@ def _compile_rule(rule: object) -> tuple[_CompiledRule, ...]: pattern, e, ) - return () + return None fill_missing_for_providers: Final = _parse_fill_missing_for_providers(rule, pattern) if fill_missing_for_providers is None: - return () + return None if PROVIDER_KEY not in model_info: - return ( - _CapabilityRule( - pattern=compiled, - model_info=model_info, - fill_missing_for_providers=fill_missing_for_providers, - ), + return _CapabilityRule( + pattern=compiled, + model_info=model_info, + fill_missing_for_providers=fill_missing_for_providers, ) provider: Final = model_info[PROVIDER_KEY] if not isinstance(provider, str): verbose_logger.warning( - "LiteLLM: skipping invalid fallback generalization rule %s: '%s' in '%s' must be a string.", - rule.get(NAME_FIELD, pattern), + "LiteLLM: skipping malformed fallback generalization rule %s ('%s' must be a string).", + rule_label, PROVIDER_KEY, - MODEL_INFO_FIELD, ) - return () - if len(model_info) == 1: - return (_RoutingRule(pattern=compiled, provider=provider),) - return ( - _RoutingRule(pattern=compiled, provider=provider), - _CapabilityRule( - pattern=compiled, - model_info=model_info, - fill_missing_for_providers=fill_missing_for_providers, - ), - ) + return None + if len(model_info) != 1: + _warn_old_format_rule(rule_label, f"'{PROVIDER_KEY}' must be the only key in '{MODEL_INFO_FIELD}'") + return None + return _RoutingRule(pattern=compiled, provider=provider) class _FallbackGeneralizations: @@ -192,7 +162,7 @@ class _FallbackGeneralizations: def set_rules(self, rules: list | None) -> None: installed: Final = rules if isinstance(rules, list) else [] - compiled: Final = tuple(kind for rule in _resolve_legacy_extends(installed) for kind in _compile_rule(rule)) + compiled: Final = tuple(kind for kind in map(_compile_rule, installed) if kind is not None) self.rules = installed self.routing_rules = tuple(rule for rule in compiled if isinstance(rule, _RoutingRule)) self.capability_rules = tuple(rule for rule in compiled if isinstance(rule, _CapabilityRule)) @@ -224,10 +194,7 @@ class _FallbackGeneralizations: ) if not matched: return None - fill_missing: Final[Mapping[str, object]] = { - key: value for model_info in matched for key, value in model_info.items() if key != PROVIDER_KEY - } - return fill_missing or None + return {key: value for model_info in matched for key, value in model_info.items()} or None _registry: Final = _FallbackGeneralizations() @@ -236,10 +203,8 @@ _registry: Final = _FallbackGeneralizations() def set_fallback_generalizations(rules: list | None) -> None: """Install the active rule list, compiling and classifying each rule. - Legacy ``extends`` inheritance is resolved here, once, before classification; - a legacy rule mixing ``litellm_provider`` with capability keys installs as both - kinds. Malformed and invalid-regex rules are warned about and skipped. Called - once when the model cost map is loaded (and again on any reload). + Malformed, mixed-kind, and invalid-regex rules are warned about and skipped. + Called once when the model cost map is loaded (and again on any reload). """ _registry.set_rules(rules) diff --git a/tests/unit/litellm_core_utils/test_fallback_generalizations.py b/tests/unit/litellm_core_utils/test_fallback_generalizations.py index e59252e1571..99fa849879c 100644 --- a/tests/unit/litellm_core_utils/test_fallback_generalizations.py +++ b/tests/unit/litellm_core_utils/test_fallback_generalizations.py @@ -7,7 +7,11 @@ its end-to-end wiring into provider routing (get_llm_provider) and model-info resolution (get_model_info) including the shipped rules in the bundled cost map. """ +import json import logging +from collections.abc import Callable +from pathlib import Path +from typing import Final import pytest @@ -21,6 +25,8 @@ from litellm.litellm_core_utils.fallback_generalizations import ( set_fallback_generalizations, ) +REPO_ROOT: Final = Path(__file__).parents[3] + @pytest.fixture def restore_generalizations(): @@ -148,7 +154,7 @@ def test_fill_missing_requires_per_rule_opt_in(restore_generalizations): } ] ) - assert match_fill_missing_generalizations("acme-1", "openai") == {"supports_vision": True} + assert match_fill_missing_generalizations("acme-1", "openai") is None restore_generalizations( [ @@ -203,89 +209,13 @@ def test_reinstalling_rules_replaces_compiled_rules(restore_generalizations): # --------------------------------------------------------------------------- # -# Engine: install-time validation and legacy-schema shim +# Engine: install-time validation # --------------------------------------------------------------------------- # -def test_legacy_mixed_rule_acts_as_both_kinds(restore_generalizations): - """A legacy rule mixing ``litellm_provider`` with capability keys routes AND - contributes its full model_info (provider included) to the capability union.""" - restore_generalizations( - [ - { - "name": "legacy-mixed", - "pattern": r"^acme-", - "model_info": {"litellm_provider": "anthropic", "supports_vision": True}, - }, - {"name": "new-caps", "pattern": r"^acme-pro-", "model_info": {"supports_reasoning": True}}, - ] - ) - assert match_routing_generalization("acme-pro-1") == "anthropic" - assert match_capability_generalizations("acme-pro-1") == { - "litellm_provider": "anthropic", - "supports_vision": True, - "supports_reasoning": True, - } - - -LEGACY_MAIN_RULES = [ - { - "name": "anthropic-claude-adaptive-thinking", - "pattern": "(?:opus|sonnet|haiku)[-._](?:4[-._](?:[6-9]|[1-9]\\d)(?!\\d)|(?:[5-9]|[1-9]\\d{1,})[-._]\\d{1,2}(?!\\d))", - "description": "Claude opus/sonnet/haiku at version 4.6 or higher: 4.6 through 4.99, then any 5.x, 6.x or later major. The minor is capped at two digits so an 8-digit date suffix such as claude-opus-4-20250514 is never read as a >= 4.6 minor. Turns on adaptive thinking for new families with no code change.", - "extends": "anthropic-claude", - "model_info": {"supports_adaptive_thinking": True}, - }, - { - "name": "anthropic-claude", - "pattern": "^claude-[a-z]+-\\d+[-.]\\d+(?:-\\d{8})?$", - "description": "Any Claude family-major-minor id, optionally with an 8-digit date suffix, anchored to the whole name. Version-neutral fallback that gives an unmapped Claude provider routing and baseline capabilities; it carries no pricing, so cost stays on the standard unpriced behavior rather than a guessed number.", - "model_info": { - "litellm_provider": "anthropic", - "mode": "chat", - "max_input_tokens": 200000, - "max_output_tokens": 64000, - "max_tokens": 64000, - "supports_function_calling": True, - "supports_parallel_function_calling": True, - "supports_vision": True, - "supports_tool_choice": True, - "supports_assistant_prefill": True, - "supports_prompt_caching": True, - "supports_response_schema": True, - "supports_reasoning": True, - "supports_pdf_input": True, - "supports_system_messages": True, - }, - }, -] - - -def test_legacy_main_schema_keeps_unmapped_claude_working(restore_generalizations): - """Pins the remote-map transition window: a released proxy running this engine - against main's old-schema block (mixed provider+capability rule plus ``extends``, - copied verbatim above) must keep unmapped-Claude inference and info resolution - working until the new-schema JSON reaches main.""" - restore_generalizations([dict(rule) for rule in LEGACY_MAIN_RULES]) - litellm.get_model_info.cache_clear() - - _, provider, _, _ = litellm.get_llm_provider(model="claude-opus-9-9") - assert provider == "anthropic" - - info = litellm.get_model_info("claude-opus-9-9") - assert info["litellm_provider"] == "anthropic" - assert info["supports_adaptive_thinking"] is True - assert info["supports_function_calling"] is True - assert info["max_input_tokens"] == 200000 - assert not info.get("input_cost_per_token") - - low = litellm.get_model_info("claude-opus-4-0") - assert low["litellm_provider"] == "anthropic" - assert low["supports_function_calling"] is True - assert low.get("supports_adaptive_thinking") is None - - -def test_non_string_provider_rule_warns_and_is_skipped(restore_generalizations, warning_messages): +def test_non_string_provider_rule_warns_and_is_skipped( + restore_generalizations: Callable[[list | None], None], warning_messages: list[str] +) -> None: restore_generalizations([{"name": "bad-provider", "pattern": r"^acme-", "model_info": {"litellm_provider": 42}}]) assert any("bad-provider" in message for message in warning_messages) assert match_routing_generalization("acme-1") is None @@ -306,6 +236,40 @@ def test_malformed_rules_are_skipped_not_fatal(restore_generalizations): assert match_capability_generalizations("claude-opus-9-9") == {"good": True} +@pytest.mark.parametrize( + "rules", + [ + [ + { + "name": "stale-rule", + "pattern": r"^acme-", + "model_info": {"litellm_provider": "anthropic", "supports_vision": True}, + } + ], + [ + {"name": "parent", "pattern": r"^zeta-only$", "model_info": {"supports_vision": True}}, + { + "name": "stale-rule", + "pattern": r"^acme-", + "extends": "parent", + "model_info": {"supports_reasoning": True}, + }, + ], + ], + ids=["provider-mixed-with-capabilities", "extends"], +) +def test_old_format_rule_is_skipped_with_a_warning_pointing_at_a_current_cost_map( + restore_generalizations: Callable[[list | None], None], warning_messages: list[str], rules: list[dict] +) -> None: + restore_generalizations(rules) + assert any( + "stale-rule" in m and "LITELLM_MODEL_COST_MAP_URL" in m and "LITELLM_LOCAL_MODEL_COST_MAP " in m + for m in warning_messages + ) + assert match_routing_generalization("acme-1") is None + assert match_capability_generalizations("acme-1") is None + + # --------------------------------------------------------------------------- # # End-to-end: provider routing + model-info resolution # --------------------------------------------------------------------------- # @@ -435,6 +399,18 @@ def test_exact_entries_fill_only_missing_fields(restore_generalizations, monkeyp # --------------------------------------------------------------------------- # +@pytest.mark.parametrize( + "cost_map_path", + ["model_prices_and_context_window.json", "litellm/model_prices_and_context_window_backup.json"], +) +def test_every_shipped_rule_installs_without_a_skip_warning( + restore_generalizations: Callable[[list | None], None], warning_messages: list[str], cost_map_path: str +) -> None: + rules: Final = json.loads((REPO_ROOT / cost_map_path).read_text())["fallback_generalizations"]["rules"] + restore_generalizations(rules) + assert warning_messages == [] + + @pytest.fixture def shipped_cost_map(monkeypatch): """Activate the bundled cost map so the shipped rules are installed.""" From 11559ed8952d54a39a01efea9595d883ddbdb394 Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Wed, 28 Oct 2026 13:21:31 +0100 Subject: [PATCH 030/128] fix(streaming): exact-match the SSE [DONE] marker in BaseModelResponseIterator (#47832) * fix(streaming): exact-match the SSE [DONE] marker in BaseModelResponseIterator * test(streaming): cover reasoning text and error frames that mention [DONE] --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/llms/base_llm/base_model_iterator.py | 4 +- .../llms/base_llm/test_base_model_iterator.py | 101 +++++++++++++++++- 2 files changed, 103 insertions(+), 2 deletions(-) diff --git a/litellm/llms/base_llm/base_model_iterator.py b/litellm/llms/base_llm/base_model_iterator.py index 611051bb2dd..032951be74a 100644 --- a/litellm/llms/base_llm/base_model_iterator.py +++ b/litellm/llms/base_llm/base_model_iterator.py @@ -3,6 +3,7 @@ from abc import abstractmethod from typing import TYPE_CHECKING, Final, cast import litellm +from litellm.constants import STREAM_SSE_DONE_STRING if TYPE_CHECKING: import httpx @@ -142,7 +143,8 @@ class BaseModelResponseIterator: def _handle_string_chunk(self, str_line: str) -> GenericStreamingChunk | ModelResponseStream: # chunk is a str at this point stripped_json_chunk: Final = BaseModelResponseIterator._string_to_dict_parser(str_line=str_line) - if "[DONE]" in str_line: + stripped_sse_chunk: Final = litellm.CustomStreamWrapper.strip_sse_data_from_chunk(str_line) + if stripped_sse_chunk is not None and stripped_sse_chunk.strip() == STREAM_SSE_DONE_STRING: return GenericStreamingChunk( text="", is_finished=True, diff --git a/tests/unit/llms/base_llm/test_base_model_iterator.py b/tests/unit/llms/base_llm/test_base_model_iterator.py index 204b458cadc..2971405885d 100644 --- a/tests/unit/llms/base_llm/test_base_model_iterator.py +++ b/tests/unit/llms/base_llm/test_base_model_iterator.py @@ -4,15 +4,19 @@ and non-string objects (e.g. Pydantic BaseModel events from the Responses API) p """ import json -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping +from itertools import chain from typing import Final import httpx import pytest from pydantic import BaseModel +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator, convert_model_response_to_streaming from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler +from litellm.llms.openai.common_utils import OpenAIError from litellm.types.utils import Choices, GenericStreamingChunk, Message, ModelResponse @@ -336,3 +340,98 @@ def test_fake_stream_chunk_indexes_thinking_blocks_by_position_without_touching_ {"type": "thinking", "thinking": "Step 2.", "signature": "sig-b", "index": 2}, ] assert [block.get("index") for block in response.choices[0].message.thinking_blocks] == [None, None, None] + + +def _sse_chunk(delta: Mapping[str, object], finish_reason: str | None = None) -> str: + return "data: " + json.dumps( + { + "id": "chatcmpl-done", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "mistral-small-latest", + "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], + } + ) + + +async def _aiter_lines(lines: tuple[str, ...]) -> AsyncIterator[str]: + for line in lines: + yield line + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_done_inside_streamed_json_is_content_not_the_end_of_stream(sync_mode: bool) -> None: + lines: Final = ( + _sse_chunk({"role": "assistant", "reasoning_content": "Mark it [DONE] once tests pass"}), + _sse_chunk({"content": "Checklist: "}), + _sse_chunk({"content": "- [DONE] write tests"}), + _sse_chunk({"content": " - [ ] ship it"}), + _sse_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "set_status", "arguments": ""}, + } + ] + } + ), + _sse_chunk({"tool_calls": [{"index": 0, "function": {"arguments": '{"status": '}}]}), + _sse_chunk({"tool_calls": [{"index": 0, "function": {"arguments": '"[DONE]"'}}]}), + _sse_chunk({"tool_calls": [{"index": 0, "function": {"arguments": ', "n": 1}'}}]}), + _sse_chunk({}, "tool_calls"), + "data: [DONE]", + ) + wrapper: Final = CustomStreamWrapper( + completion_stream=OpenAIChatCompletionStreamingHandler( + streaming_response=iter(lines) if sync_mode else _aiter_lines(lines), sync_stream=sync_mode + ), + model="mistral/mistral-small-latest", + custom_llm_provider="mistral", + logging_obj=Logging( + model="mistral/mistral-small-latest", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=1_790_000_000.0, + litellm_call_id="call-id", + function_id="function-id", + ), + ) + + chunks: Final = list(wrapper) if sync_mode else [chunk async for chunk in wrapper] + + deltas: Final = [chunk.choices[0].delta for chunk in chunks if chunk.choices] + tool_calls: Final = tuple(chain.from_iterable(delta.tool_calls or () for delta in deltas)) + assert "".join(getattr(delta, "reasoning_content", None) or "" for delta in deltas) == ( + "Mark it [DONE] once tests pass" + ) + assert "".join(delta.content or "" for delta in deltas) == "Checklist: - [DONE] write tests - [ ] ship it" + assert json.loads("".join(tool_call.function.arguments or "" for tool_call in tool_calls)) == { + "status": "[DONE]", + "n": 1, + } + assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices and chunk.choices[0].finish_reason] == [ + "tool_calls" + ] + + +@pytest.mark.parametrize("line", ["data: [DONE]", "data:[DONE]", "[DONE]", "data: [DONE]", "data: [DONE] "]) +def test_done_marker_ends_the_stream(line: str) -> None: + chunks: Final = list(BaseModelResponseIterator(streaming_response=iter((line,)), sync_stream=True)) + + assert chunks == [ + GenericStreamingChunk(text="", is_finished=True, finish_reason="stop", usage=None, index=0, tool_use=None) + ] + + +def test_error_frame_mentioning_done_raises_instead_of_ending_the_stream() -> None: + line: Final = 'data: {"error": {"message": "prompt too long [DONE]", "code": 400}}' + + with pytest.raises(OpenAIError, match=r"prompt too long \[DONE\]") as error: + list(OpenAIChatCompletionStreamingHandler(streaming_response=iter((line,)), sync_stream=True)) + + assert error.value.status_code == 400 From ca24b396da9d7fd5264c33da2c84d4cd40335d8c Mon Sep 17 00:00:00 2001 From: quietjules <349048734+quietjules@users.noreply.github.com> Date: Wed, 28 Oct 2026 05:24:08 -0700 Subject: [PATCH 031/128] test(nvidia_riva): audit #41021 (#48801) * test(nvidia_riva): audit #41021 Adds an owned Riva ASR gRPC server and provider integration cells for the riva_offline request-body ban, the offline and streaming paths, the per-deployment caller opt-in, and the 413 and 408 error mappings * test(nvidia_riva): check word timing requests and audio samples in the Riva double --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- tests/integration/_support/riva.py | 143 ++++++++++++++++++ .../providers/test_nvidia_riva_wire.py | 137 +++++++++++++++++ 2 files changed, 280 insertions(+) create mode 100644 tests/integration/_support/riva.py create mode 100644 tests/integration/providers/test_nvidia_riva_wire.py diff --git a/tests/integration/_support/riva.py b/tests/integration/_support/riva.py new file mode 100644 index 00000000000..8552bc2197d --- /dev/null +++ b/tests/integration/_support/riva.py @@ -0,0 +1,143 @@ +from __future__ import annotations + +import io +import threading +import wave +from array import array +from collections.abc import Iterable, Iterator +from concurrent import futures +from contextlib import contextmanager +from dataclasses import dataclass +from enum import Enum +from queue import SimpleQueue +from typing import Final + +import grpc +from riva.client.proto import riva_asr_pb2, riva_asr_pb2_grpc + +SAMPLE_RATE_HZ: Final = 16000 +OFFLINE_ONLY_REFUSAL: Final = ( + "Error: Unavailable model requested given these parameters: language_code=en-US; sample_rate=16000; type=online;" +) +_STALL_LIMIT_SECONDS: Final = 120 + + +class Rpc(Enum): + RECOGNIZE = "Recognize" + STREAMING = "StreamingRecognize" + + +@dataclass(frozen=True, slots=True) +class Word: + word: str + start_ms: int + end_ms: int + + +@dataclass(frozen=True, slots=True) +class RivaCall: + rpc: Rpc + audio: bytes + cancelled: bool + + +@dataclass(frozen=True, slots=True) +class RivaServer: + address: str + calls: SimpleQueue[RivaCall] + + def drain(self) -> tuple[RivaCall, ...]: + return tuple(self.calls.get_nowait() for _ in range(self.calls.qsize())) + + +def _clip_samples(seconds: float) -> array[int]: + return array("h", (((index * 37) % 2001 - 1000) * 16 for index in range(int(seconds * SAMPLE_RATE_HZ)))) + + +def pcm_wav(seconds: float) -> bytes: + buffer: Final = io.BytesIO() + with wave.open(buffer, "wb") as writer: + writer.setnchannels(1) + writer.setsampwidth(2) + writer.setframerate(SAMPLE_RATE_HZ) + writer.writeframes(_clip_samples(seconds).tobytes()) + return buffer.getvalue() + + +def pcm_bytes(seconds: float) -> int: + return int(seconds * SAMPLE_RATE_HZ) * 2 + + +def carries_clip(audio: bytes, seconds: float) -> bool: + sent: Final = _clip_samples(seconds) + if len(audio) != len(sent) * 2: + return False + received: Final = array("h", audio) + return all(abs(left - right) <= 1 for left, right in zip(sent, received, strict=True)) + + +class _Asr(riva_asr_pb2_grpc.RivaSpeechRecognitionServicer): + def __init__(self, words: tuple[Word, ...], offline_only: bool, stall: bool, calls: SimpleQueue[RivaCall]) -> None: + self._words: Final = words + self._offline_only: Final = offline_only + self._stall: Final = stall + self._calls: Final = calls + + def _alternative(self, config: riva_asr_pb2.RecognitionConfig) -> riva_asr_pb2.SpeechRecognitionAlternative: + timed: Final = self._words if config.enable_word_time_offsets else () + return riva_asr_pb2.SpeechRecognitionAlternative( + transcript=" ".join(word.word for word in self._words), + confidence=0.9, + words=[ + riva_asr_pb2.WordInfo(word=word.word, start_time=word.start_ms, end_time=word.end_ms) for word in timed + ], + ) + + def Recognize( + self, request: riva_asr_pb2.RecognizeRequest, context: grpc.ServicerContext + ) -> riva_asr_pb2.RecognizeResponse: + if self._stall: + released: Final = threading.Event() + context.add_callback(released.set) + released.wait(timeout=_STALL_LIMIT_SECONDS) + self._calls.put(RivaCall(Rpc.RECOGNIZE, request.audio, cancelled=not context.is_active())) + return riva_asr_pb2.RecognizeResponse() + self._calls.put(RivaCall(Rpc.RECOGNIZE, request.audio, cancelled=False)) + return riva_asr_pb2.RecognizeResponse( + results=[riva_asr_pb2.SpeechRecognitionResult(alternatives=[self._alternative(request.config)])] + ) + + def StreamingRecognize( + self, request_iterator: Iterable[riva_asr_pb2.StreamingRecognizeRequest], context: grpc.ServicerContext + ) -> Iterator[riva_asr_pb2.StreamingRecognizeResponse]: + if self._offline_only: + self._calls.put(RivaCall(Rpc.STREAMING, b"", cancelled=False)) + context.abort(grpc.StatusCode.INVALID_ARGUMENT, OFFLINE_ONLY_REFUSAL) + requests: Final = tuple(request_iterator) + config: Final = requests[0].streaming_config.config if requests else riva_asr_pb2.RecognitionConfig() + self._calls.put( + RivaCall(Rpc.STREAMING, b"".join(request.audio_content for request in requests), cancelled=False) + ) + yield riva_asr_pb2.StreamingRecognizeResponse( + results=[riva_asr_pb2.StreamingRecognitionResult(alternatives=[self._alternative(config)], is_final=True)] + ) + + +@contextmanager +def riva_server( + words: tuple[Word, ...], + *, + offline_only: bool = False, + stall: bool = False, + max_receive_bytes: int | None = None, +) -> Iterator[RivaServer]: + calls: Final[SimpleQueue[RivaCall]] = SimpleQueue() + options: Final = () if max_receive_bytes is None else (("grpc.max_receive_message_length", max_receive_bytes),) + server: Final = grpc.server(futures.ThreadPoolExecutor(max_workers=8), options=options) + riva_asr_pb2_grpc.add_RivaSpeechRecognitionServicer_to_server(_Asr(words, offline_only, stall, calls), server) + port: Final = server.add_insecure_port("127.0.0.1:0") + server.start() + try: + yield RivaServer(f"127.0.0.1:{port}", calls) + finally: + server.stop(grace=None).wait() diff --git a/tests/integration/providers/test_nvidia_riva_wire.py b/tests/integration/providers/test_nvidia_riva_wire.py new file mode 100644 index 00000000000..ef36c449761 --- /dev/null +++ b/tests/integration/providers/test_nvidia_riva_wire.py @@ -0,0 +1,137 @@ +from collections.abc import Mapping +from typing import Final + +import httpx +import openai +import pytest +from integration._support.client import Gateway, Scenario, eventually +from integration._support.openai_wire import chat_reply +from integration._support.riva import RivaServer, Rpc, Word, carries_clip, pcm_bytes, pcm_wav, riva_server +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +_MODEL: Final = "nvidia_riva/nvidia/parakeet-0.6b-tdt" +_WORDS: Final = (Word("riva", 0, 400), Word("offline", 400, 900), Word("works", 900, 1400)) +_TRANSCRIPT: Final = "riva offline works" +_CLIP_SECONDS: Final = 0.5 +_REFUSED_FIELD: Final = "riva_offline is not allowed in request body" + + +def _riva_model(scenario: Scenario, server: RivaServer, **parameters: JsonValue) -> str: + return scenario.model( + model=_MODEL, + api_base=server.address, + model_info={"mode": "audio_transcription"}, + **parameters, + ) + + +def _transcribe( + gateway: Gateway, key: str, model: str, seconds: float = _CLIP_SECONDS, **fields: str +) -> httpx.Response: + return gateway.request_multipart( + "/v1/audio/transcriptions", + {"model": model, **fields}, + {"file": ("clip.wav", pcm_wav(seconds), "audio/wav")}, + key=key, + ) + + +def _sdk(gateway: Gateway, key: str) -> openai.OpenAI: + return openai.OpenAI(base_url=f"{str(gateway.client.base_url).rstrip('/')}/v1", api_key=key, max_retries=0) + + +def test_a_form_riva_offline_on_transcriptions_is_refused_before_riva_is_called(gateway: Gateway) -> None: + with riva_server(_WORDS) as server, gateway.scenario() as scenario: + model: Final = _riva_model(scenario, server) + refused: Final = _transcribe(gateway, scenario.key(), model, riva_offline="true") + assert refused.status_code == 401, refused.text + assert _REFUSED_FIELD in refused.json()["error"]["message"], refused.text + assert server.drain() == () + + +def test_a_json_riva_offline_on_chat_is_refused_before_the_provider_is_called(gateway: Gateway) -> None: + def respond(_: Request) -> Reply: + return chat_reply("chatcmpl-riva-audit", "gpt-4o-mini", "served", stream=False) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + refused: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "riva flag on chat"}], "riva_offline": True}, + key=scenario.key(), + ) + assert refused.status_code == 401, refused.text + assert _REFUSED_FIELD in refused.json()["error"]["message"], refused.text + assert wire.drain() == () + + +def test_an_offline_deployment_sends_the_whole_clip_in_one_recognize_call(gateway: Gateway) -> None: + with riva_server(_WORDS, offline_only=True) as server, gateway.scenario() as scenario: + model: Final = _riva_model(scenario, server, riva_offline=True) + transcription: Final = _sdk(gateway, scenario.key()).audio.transcriptions.create( + model=model, + file=("clip.wav", pcm_wav(_CLIP_SECONDS), "audio/wav"), + response_format="verbose_json", + timestamp_granularities=["word"], + ) + assert transcription.text == _TRANSCRIPT + assert [(word.word, word.start, word.end) for word in transcription.words or ()] == [ + (word.word, word.start_ms / 1000, word.end_ms / 1000) for word in _WORDS + ] + assert [(call.rpc, carries_clip(call.audio, _CLIP_SECONDS)) for call in server.drain()] == [ + (Rpc.RECOGNIZE, True) + ] + + +def test_a_deployment_without_riva_offline_still_streams_the_clip(gateway: Gateway) -> None: + with riva_server(_WORDS) as server, gateway.scenario() as scenario: + model: Final = _riva_model(scenario, server, timeout=30) + served: Final = _transcribe(gateway, scenario.key(), model) + assert served.status_code == 200, served.text + assert served.json()["text"] == _TRANSCRIPT, served.text + assert [(call.rpc, carries_clip(call.audio, _CLIP_SECONDS)) for call in server.drain()] == [ + (Rpc.STREAMING, True) + ] + + +@pytest.mark.parametrize( + ("sent", "deployment", "status", "rpcs"), + ( + pytest.param("true", {}, 200, (Rpc.RECOGNIZE,), id="true-on-a-streaming-deployment"), + pytest.param("false", {"riva_offline": True}, 200, (Rpc.STREAMING,), id="false-on-an-offline-deployment"), + pytest.param("yes", {}, 400, (), id="malformed"), + pytest.param("yes", {"drop_params": True}, 200, (Rpc.STREAMING,), id="malformed-under-drop-params"), + ), +) +def test_an_opted_in_caller_riva_offline_picks_the_rpc( + gateway: Gateway, sent: str, deployment: Mapping[str, JsonValue], status: int, rpcs: tuple[Rpc, ...] +) -> None: + with riva_server(_WORDS) as server, gateway.scenario() as scenario: + model: Final = _riva_model(scenario, server, configurable_clientside_auth_params=["riva_offline"], **deployment) + answered: Final = _transcribe(gateway, scenario.key(), model, riva_offline=sent) + assert answered.status_code == status, answered.text + assert tuple(call.rpc for call in server.drain()) == rpcs + + +def test_offline_audio_over_the_servers_message_limit_answers_413(gateway: Gateway) -> None: + seconds: Final = 2.0 + limit: Final = pcm_bytes(seconds) // 2 + with ( + riva_server(_WORDS, offline_only=True, max_receive_bytes=limit) as server, + gateway.scenario() as scenario, + ): + model: Final = _riva_model(scenario, server, riva_offline=True) + refused: Final = _transcribe(gateway, scenario.key(), model, seconds) + assert refused.status_code == 413, refused.text + assert server.drain() == () + + +def test_a_stalled_offline_recognize_answers_408_at_the_deployment_timeout(gateway: Gateway) -> None: + with riva_server(_WORDS, offline_only=True, stall=True) as server, gateway.scenario() as scenario: + model: Final = _riva_model(scenario, server, riva_offline=True, timeout=2) + timed_out: Final = _transcribe(gateway, scenario.key(), model) + assert timed_out.status_code == 408, timed_out.text + calls: Final = eventually(server.drain, lambda observed: len(observed) > 0, seconds=30) + assert all(call.rpc is Rpc.RECOGNIZE and call.cancelled for call in calls), calls From 0ceca31081efc6b6f12e7351115a8b12e6ad83e6 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 05:24:41 -0700 Subject: [PATCH 032/128] feat(proxy): per-provider and per-model rpm/tpm requirements for enforce_project_model_quota (#47693) * feat(proxy): per-provider and per-model rpm/tpm requirements for enforce_project_model_quota * test(proxy): cover per-provider and per-model enforce_project_model_quota policies * fix(proxy): keep enforce_project_model_quota reads typed * fix(proxy): apply the quota policy to model add, follow router wildcard rules, keep old setting values * fix(proxy): keep team model names off quota model entries, read the provider behind api_base, match patterns like the router A team admin picks a team model's public name, so a name matching an exempt model entry no longer skips its provider's limits. A mapping with no policy key keeps requiring both limits as before, the project create permission check runs before the setting is read, and an unprefixed model behind a known OpenAI-compatible api_base resolves to that provider * fix(proxy): re-check team projects on team model writes and validate quota provider ids * fix(proxy): check added limits on model updates and match team models by the model they call * fix(proxy): re-check team projects by every name a team model serves and check renamed names * fix(proxy): follow the router's order for tied team patterns and annotate quota tests * fix(proxy): match a team model's called model through JSON provider prefixes too * fix(proxy): keep model updates that keep their route as on main A PATCH or POST model update now checks only the limits its new route adds over the stored one, so removing or zeroing a limit on an unchanged route answers 200 again, under true and under a policy --- .../management_endpoints/project_endpoints.py | 111 +-- .../project_model_quota.py | 49 ++ .../model_management_endpoints.py | 296 +++++++- .../management_helpers/model_quota_policy.py | 327 +++++++++ .../test_project_endpoints_prisma.py | 562 +++++++++++++++- .../test_model_management_endpoints.py | 631 +++++++++++++++++- 6 files changed, 1899 insertions(+), 77 deletions(-) create mode 100644 enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py create mode 100644 litellm/proxy/management_helpers/model_quota_policy.py diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py index b91b4a970e8..c25da7cc187 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py @@ -15,7 +15,11 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Final, cast from fastapi import APIRouter, Depends, HTTPException, Request -from pydantic import TypeAdapter +from litellm_enterprise.proxy.management_endpoints.project_model_quota import ( + describe_missing_quota, + missing_project_model_quota, +) +from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -27,6 +31,15 @@ from litellm.proxy.management.teams.dependencies import get_team_access from litellm.proxy.management.users.service import org_admin_org_ids from litellm.proxy.management_endpoints.common_utils import _set_object_metadata_field from litellm.proxy.management_endpoints.team_admin_field_permissions import team_admin_may_manage_projects +from litellm.proxy.management_helpers.model_quota_policy import ( + ENFORCE_PROJECT_MODEL_QUOTA_SETTING, + REQUIRE_BOTH_LIMITS, + ModelQuotaPolicy, + RouteOf, + no_route, + parse_model_quota_setting, + router_model_routes, +) from litellm.proxy.management_helpers.utils import ( management_endpoint_wrapper, ) @@ -276,26 +289,6 @@ def _check_team_project_limits( ) -def _project_models_missing_positive_quota( - models: list[str] | None, - rpm_limits: Mapping[str, object] | None, - tpm_limits: Mapping[str, object] | None, -) -> list[str]: - """Return the models that lack a positive `rpm` AND `tpm` quota. - - A valid quota is a positive integer; null, zero, and negative are rejected - because downstream rate limiters treat a non-positive limit as immediately - exhausted (every request blocked). - """ - - def _is_positive(value: object) -> bool: - return isinstance(value, int) and not isinstance(value, bool) and value > 0 - - rpm = rpm_limits or {} - tpm = tpm_limits or {} - return [model for model in (models or []) if not _is_positive(rpm.get(model)) or not _is_positive(tpm.get(model))] - - def _router_access_group_names(llm_router: "Router | None") -> frozenset[str]: return frozenset(llm_router.get_model_access_groups()) if llm_router is not None else frozenset() @@ -322,64 +315,79 @@ def _raise_on_project_models_expanding_at_request_time( raise HTTPException( status_code=400, detail={ - "error": f"models {list(expanding)} expand to multiple models at request time, so a per-model rpm/tpm quota cannot be enforced for them while 'enforce_project_model_quota' is enabled. List concrete model names instead." + "error": f"models {list(expanding)} expand to multiple models at request time, so a per-model rpm/tpm quota cannot be enforced for them while '{ENFORCE_PROJECT_MODEL_QUOTA_SETTING}' is enabled. List concrete model names instead." }, ) +def _project_model_quota_policy(general_settings: Mapping[str, object]) -> ModelQuotaPolicy | None: + try: + return parse_model_quota_setting(general_settings.get(ENFORCE_PROJECT_MODEL_QUOTA_SETTING)) + except ValidationError as e: + raise HTTPException( + status_code=500, + detail={"error": f"general_settings.{ENFORCE_PROJECT_MODEL_QUOTA_SETTING} is invalid: {e}"}, + ) from e + + def _raise_on_missing_project_model_quota( - data: NewProjectRequest | UpdateProjectRequest, access_group_names: frozenset[str] = frozenset() + data: NewProjectRequest | UpdateProjectRequest, + access_group_names: frozenset[str] = frozenset(), + policy: ModelQuotaPolicy = REQUIRE_BOTH_LIMITS, + route_of: RouteOf = no_route, ) -> None: - """Require a positive `rpm`/`tpm` quota for every model on project CREATE. + """Require the positive `rpm`/`tpm` quotas the policy asks of each model on project CREATE. `model_rpm_limit`/`model_tpm_limit` are relocated into `metadata` by the request model's `set_model_info` validator, so they are read from there. - - Only invoked when `general_settings.enforce_project_model_quota` is enabled - (default off), so it is opt-in and does not change behavior for existing users. """ _raise_on_project_models_expanding_at_request_time(data.models, access_group_names) metadata = data.metadata or {} - missing = _project_models_missing_positive_quota( - data.models, metadata.get("model_rpm_limit"), metadata.get("model_tpm_limit") + missing: Final = missing_project_model_quota( + data.models or (), metadata.get("model_rpm_limit"), metadata.get("model_tpm_limit"), policy, route_of ) if not missing: return raise HTTPException( status_code=400, detail={ - "error": f"models {missing} added to project without a positive rpm/tpm quota. Set a positive model_rpm_limit and model_tpm_limit for each model." + "error": f"models added to project without a positive rpm/tpm quota: {describe_missing_quota(missing)}. Set a positive value for each listed limit." }, ) def _raise_on_missing_project_model_quota_on_update( - data: UpdateProjectRequest, existing_project: object, access_group_names: frozenset[str] = frozenset() + data: UpdateProjectRequest, + existing_project: object, + access_group_names: frozenset[str] = frozenset(), + policy: ModelQuotaPolicy = REQUIRE_BOTH_LIMITS, + route_of: RouteOf = no_route, ) -> None: - """Require a positive `rpm`/`tpm` quota over the RESULTING state on project UPDATE. + """Require the positive `rpm`/`tpm` quotas the policy asks of each model over the RESULTING state on project UPDATE. `/project/update` replaces `models` and `metadata` when they are provided, so the check runs on what the project WILL look like: a partial update that doesn't touch models/quota keeps the existing values, while one that adds a model or clears a - model's quota must leave every resulting model with a positive limit. - - Only invoked when `general_settings.enforce_project_model_quota` is enabled - (default off), so it is opt-in and does not change behavior for existing users. + model's quota must leave every resulting model with the limits it needs. """ resulting_models = data.models if data.models is not None else (getattr(existing_project, "models", None) or []) resulting_metadata = ( data.metadata if data.metadata is not None else (getattr(existing_project, "metadata", None) or {}) ) _raise_on_project_models_expanding_at_request_time(resulting_models, access_group_names) - missing = _project_models_missing_positive_quota( - resulting_models, resulting_metadata.get("model_rpm_limit"), resulting_metadata.get("model_tpm_limit") + missing: Final = missing_project_model_quota( + resulting_models, + resulting_metadata.get("model_rpm_limit"), + resulting_metadata.get("model_tpm_limit"), + policy, + route_of, ) if not missing: return raise HTTPException( status_code=400, detail={ - "error": f"models {missing} would be left on the project without a positive rpm/tpm quota. Set a positive model_rpm_limit and model_tpm_limit for each model." + "error": f"models would be left on the project without a positive rpm/tpm quota: {describe_missing_quota(missing)}. Set a positive value for each listed limit." }, ) @@ -530,6 +538,7 @@ async def new_project( """ from litellm.proxy.proxy_server import ( general_settings, + general_settings_view, litellm_proxy_admin_name, llm_router, premium_user, @@ -578,10 +587,6 @@ async def new_project( data=data, ) - # Opt-in (default off): require rpm/tpm for every model added to the project. - if general_settings.get("enforce_project_model_quota", False): - _raise_on_missing_project_model_quota(data, _router_access_group_names(llm_router)) - # Check if user has permission to create projects for this team # only team admins can create projects for their team has_permission = await _check_user_permission_for_project( @@ -600,6 +605,15 @@ async def new_project( }, ) + quota_policy: Final = _project_model_quota_policy(general_settings_view()) + if quota_policy is not None: + _raise_on_missing_project_model_quota( + data, + _router_access_group_names(llm_router), + quota_policy, + router_model_routes(llm_router, data.team_id), + ) + # Generate project_id if not provided if data.project_id is None: data.project_id = str(uuid.uuid4()) @@ -721,6 +735,7 @@ async def update_project( """ from litellm.proxy.proxy_server import ( general_settings, + general_settings_view, litellm_proxy_admin_name, llm_router, premium_user, @@ -828,10 +843,14 @@ async def update_project( data=data, ) - # Opt-in (default off): require rpm/tpm for every model the update would leave on the project. - if general_settings.get("enforce_project_model_quota", False): + quota_policy: Final = _project_model_quota_policy(general_settings_view()) + if quota_policy is not None: _raise_on_missing_project_model_quota_on_update( - data, existing_project, _router_access_group_names(llm_router) + data, + existing_project, + _router_access_group_names(llm_router), + quota_policy, + router_model_routes(llm_router, target_team_id), ) # Prepare update data diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py new file mode 100644 index 00000000000..e36ff9b53dc --- /dev/null +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py @@ -0,0 +1,49 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final + +from litellm.proxy.management_helpers.model_quota_policy import ( + ModelQuotaPolicy, + QuotaLimit, + RouteOf, + required_limits, + unmet_limits, +) + +_NO_LIMITS_SET: Final[Mapping[str, object]] = MappingProxyType({}) + + +@dataclass(frozen=True, slots=True) +class MissingModelQuota: + model: str + limits: tuple[QuotaLimit, ...] + + +def _limits_set_for( + model: str, rpm_set: Mapping[str, object], tpm_set: Mapping[str, object] +) -> Mapping[QuotaLimit, object]: + return MappingProxyType({"rpm": rpm_set.get(model), "tpm": tpm_set.get(model)}) + + +def missing_project_model_quota( + models: Sequence[str], + rpm_limits: Mapping[str, object] | None, + tpm_limits: Mapping[str, object] | None, + policy: ModelQuotaPolicy, + route_of: RouteOf, +) -> tuple[MissingModelQuota, ...]: + rpm_set: Final = rpm_limits or _NO_LIMITS_SET + tpm_set: Final = tpm_limits or _NO_LIMITS_SET + gaps: Final = tuple( + MissingModelQuota( + model, + unmet_limits(required_limits(policy, model, route_of(model)), _limits_set_for(model, rpm_set, tpm_set)), + ) + for model in models + ) + return tuple(gap for gap in gaps if gap.limits) + + +def describe_missing_quota(missing: Sequence[MissingModelQuota]) -> str: + return ", ".join(f"{gap.model} ({' and '.join(f'model_{limit}_limit' for limit in gap.limits)})" for gap in missing) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index e625ea6f9f2..1130fb73e9a 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -12,6 +12,7 @@ model/{model_id}/update - PATCH endpoint for model update. import asyncio import datetime +import itertools import json from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence from contextlib import AbstractAsyncContextManager, asynccontextmanager, suppress @@ -42,6 +43,7 @@ from litellm.litellm_core_utils.ptu_pricing import ( ptu_config_error, ) from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload +from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping from litellm.proxy._types import ( BlockModelRequest, CommonProxyErrors, @@ -102,6 +104,18 @@ from litellm.proxy.management_helpers.auto_router_permissions import ( authorize_member_auto_router_write, ) from litellm.proxy.management_helpers.model_allowlist_rename_sync import sync_model_allowlists_for_renamed_model +from litellm.proxy.management_helpers.model_quota_policy import ( + ENFORCE_PROJECT_MODEL_QUOTA_SETTING, + ModelQuotaPolicy, + QuotaLimit, + deployment_route, + exempts_any_route, + limits_added_by_team_write, + parse_model_quota_setting, + required_limits, + team_write_serves, + unmet_limits, +) from litellm.proxy.spend_tracking.ptu_feature_flag import ( PTU_COST_ATTRIBUTION_ENV_VAR, is_ptu_cost_attribution_enabled, @@ -111,6 +125,7 @@ from litellm.repositories.chunked_in import find_many_in from litellm.repositories.credentials_repository import CredentialsRepository from litellm.repositories.model_repository import ModelRepository from litellm.repositories.prisma_protocols import TableActions +from litellm.repositories.project_repository import ProjectRepository from litellm.repositories.table_repositories import ModelTableRepository from litellm.repositories.team_repository import TeamRepository from litellm.router import Router @@ -152,6 +167,7 @@ from litellm.types.router import ( SPECIAL_MODEL_INFO_PARAMS, Deployment, GenericLiteLLMParams, + LiteLLM_Params, ModelInfo, updateDeployment, updateLiteLLMParams, @@ -762,26 +778,93 @@ async def _auto_router_capability_slot( ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING: Final = "enforce_rpm_tpm_on_model_add" -def _raise_if_rate_limits_required_but_missing(*, litellm_params: GenericLiteLLMParams, enforced: bool) -> None: - """Require both rpm and tpm (each a positive value) when the operator opts in via config.yaml. +_ROUTING_KEYS: Final = ("model", "custom_llm_provider", "api_base") +_RATE_LIMIT_KEYS: Final = ("rpm", "tpm") +_NOTHING_SENT: Final[Mapping[str, object]] = MappingProxyType({}) - Off by default, so deployments keep adding models without limits. When - ``enforce_rpm_tpm_on_model_add: true`` is set under general_settings, a model added - without both rpm and tpm set to a positive value is rejected rather than stored - unbounded (or effectively excluded from routing by a zero/negative limit). - """ - if not enforced: + +def _merged_litellm_params( + incoming: GenericLiteLLMParams | None, existing: GenericLiteLLMParams | None, keys: tuple[str, ...] +) -> Mapping[str, object]: + sent: Final = _NOTHING_SENT if incoming is None else _dump_without_none(incoming) + stored: Final = _NOTHING_SENT if existing is None else _decrypted_litellm_params(existing) + return MappingProxyType({key: sent.get(key, stored.get(key)) for key in keys}) + + +def _written_routing_params( + incoming: GenericLiteLLMParams | None, existing: GenericLiteLLMParams | None +) -> Mapping[str, object]: + return _merged_litellm_params(incoming, existing, _ROUTING_KEYS) + + +def _quota_policy(general_settings: Mapping[str, object], setting: str) -> ModelQuotaPolicy | None: + try: + return parse_model_quota_setting(general_settings.get(setting)) + except ValidationError as e: + raise ProxyException( + message=f"general_settings.{setting} is invalid: {e}", + type=ProxyErrorTypes.internal_server_error.value, + code=status.HTTP_500_INTERNAL_SERVER_ERROR, + param=None, + ) from e + + +def _raise_if_rate_limits_required_but_missing( + *, model_name: str, litellm_params: LiteLLM_Params, team_model: bool, policy: ModelQuotaPolicy | None +) -> None: + """Off by default, so deployments keep adding models without limits. Once the operator opts in under + general_settings, the policy decides per model name and provider which of rpm and tpm a new model needs, + and a model missing one, or carrying a zero or negative one, is rejected rather than stored unbounded + (or effectively excluded from routing).""" + if policy is None: return - missing: Final = tuple( - field - for field, value in (("rpm", litellm_params.rpm), ("tpm", litellm_params.tpm)) - if value is None or value <= 0 + _raise_on_missing_rate_limits( + model_name, + unmet_limits( + required_limits( + policy, model_name, deployment_route(_written_routing_params(litellm_params, None), team_model) + ), + MappingProxyType({"rpm": litellm_params.rpm, "tpm": litellm_params.tpm}), + ), ) + + +def _raise_if_model_update_adds_missing_rate_limits( + *, + existing: Deployment, + written_name: str, + written_team_model: bool, + incoming_params: GenericLiteLLMParams | None, + remaining_params: GenericLiteLLMParams, + policy: ModelQuotaPolicy | None, +) -> None: + if policy is None: + return + before: Final = required_limits( + policy, + existing.model_info.team_public_model_name or existing.model_name, + deployment_route( + _written_routing_params(None, existing.litellm_params), existing.model_info.team_id is not None + ), + ) + after: Final = required_limits( + policy, + written_name, + deployment_route(_written_routing_params(incoming_params, remaining_params), written_team_model), + ) + written_limits: Final = _merged_litellm_params(incoming_params, remaining_params, _RATE_LIMIT_KEYS) + checked: Final = after - before + _raise_on_missing_rate_limits( + written_name, unmet_limits(checked, MappingProxyType({limit: written_limits.get(limit) for limit in checked})) + ) + + +def _raise_on_missing_rate_limits(model_name: str, missing: tuple[QuotaLimit, ...]) -> None: if not missing: return raise ProxyException( message=( - f"{' and '.join(missing)} must be set to a positive value when " + f"{' and '.join(missing)} must be set to a positive value for '{model_name}' when " f"'{ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING}' is enabled in general_settings" ), type=ProxyErrorTypes.validation_error.value, @@ -790,6 +873,113 @@ def _raise_if_rate_limits_required_but_missing(*, litellm_params: GenericLiteLLM ) +def _project_limit(metadata: object, limits_key: str, model: str) -> object: + limits: Final = metadata.get(limits_key) if is_str_keyed_mapping(metadata) else None + return limits.get(model) if is_str_keyed_mapping(limits) else None + + +def _projects_short_on( + projects: Sequence["prisma_models.LiteLLM_ProjectTable"], model: str, added: frozenset[QuotaLimit] +) -> tuple[str, ...]: + return tuple( + project.project_alias or project.project_id + for project in projects + if model in project.models + and unmet_limits( + added, + MappingProxyType( + { + "rpm": _project_limit(project.metadata, "model_rpm_limit", model), + "tpm": _project_limit(project.metadata, "model_tpm_limit", model), + } + ), + ) + ) + + +async def _raise_if_team_write_leaves_projects_short( + *, + prisma_client: PrismaClient, + general_settings: Mapping[str, object], + llm_router: Router | None, + team_id: str | None, + public_name: str, + replaced_id: str | None, + written_params: Mapping[str, object] | None, +) -> None: + if team_id is None: + return + policy: Final = _quota_policy(general_settings, ENFORCE_PROJECT_MODEL_QUOTA_SETTING) + if policy is None or not exempts_any_route(policy): + return + projects: Final = await ProjectRepository(prisma_client).table.find_many(where={"team_id": team_id}) + listed: Final = frozenset(itertools.chain.from_iterable(project.models for project in projects)) + served: Final = sorted(model for model in listed if team_write_serves(llm_router, team_id, public_name, model)) + shortfall: Final = next( + ( + (model, added, short) + for model in served + if ( + added := limits_added_by_team_write( + policy, llm_router, team_id, public_name, replaced_id, written_params, model + ) + ) + and (short := _projects_short_on(projects, model, added)) + ), + None, + ) + if shortfall is None: + return + short_model, short_limits, short_projects = shortfall + raise ProxyException( + message=( + f"projects {', '.join(short_projects)} list '{short_model}' without a positive " + f"{' and '.join(sorted(short_limits))} " + f"quota, which '{ENFORCE_PROJECT_MODEL_QUOTA_SETTING}' requires once this team model change applies. " + "Set those limits on the projects first" + ), + type=ProxyErrorTypes.validation_error.value, + code=status.HTTP_400_BAD_REQUEST, + param=None, + ) + + +async def _raise_if_team_update_leaves_projects_short( + *, + prisma_client: PrismaClient, + general_settings: Mapping[str, object], + llm_router: Router | None, + existing: Deployment, + written_team_id: str | None, + written_name: str, + incoming_params: GenericLiteLLMParams | None, + remaining_params: GenericLiteLLMParams, +) -> None: + if written_team_id is None and existing.model_info.team_id is None: + return + stored_name: Final = existing.model_info.team_public_model_name or existing.model_name + await _raise_if_team_write_leaves_projects_short( + prisma_client=prisma_client, + general_settings=general_settings, + llm_router=llm_router, + team_id=written_team_id, + public_name=written_name, + replaced_id=existing.model_info.id, + written_params=_written_routing_params(incoming_params, remaining_params), + ) + if (written_team_id, written_name) == (existing.model_info.team_id, stored_name): + return + await _raise_if_team_write_leaves_projects_short( + prisma_client=prisma_client, + general_settings=general_settings, + llm_router=llm_router, + team_id=existing.model_info.team_id, + public_name=stored_name, + replaced_id=existing.model_info.id, + written_params=None, + ) + + _PTU_PRICED_PAIR: Final = frozenset({"ptu_count", "cost_per_ptu_per_hour"}) @@ -1308,6 +1498,7 @@ async def patch_model( ProxyException: For various error conditions including authentication and database errors """ from litellm.proxy.proxy_server import ( + general_settings_view, litellm_proxy_admin_name, llm_router, premium_user, @@ -1434,6 +1625,29 @@ async def patch_model( existing_params=remaining_params, ) + written_team_id: Final = ( + patch_data.model_info.team_id if patch_data.model_info else None + ) or db_model.model_info.team_id + written_name: Final = _get_public_model_name(effective_patch, db_model) + _raise_if_model_update_adds_missing_rate_limits( + existing=db_model, + written_name=written_name, + written_team_model=written_team_id is not None, + incoming_params=incoming_params, + remaining_params=remaining_params, + policy=_quota_policy(general_settings_view(), ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING), + ) + await _raise_if_team_update_leaves_projects_short( + prisma_client=prisma_client, + general_settings=general_settings_view(), + llm_router=llm_router, + existing=db_model, + written_team_id=written_team_id, + written_name=written_name, + incoming_params=incoming_params, + remaining_params=remaining_params, + ) + effective_params: Final = _effective_complexity_router_params(incoming_params, remaining_params) requested_model_name: Final = patch_data.model_name stored_model_name: str | None = None @@ -2509,6 +2723,7 @@ async def delete_model( from litellm.proxy.proxy_server import ( MODEL_RECONCILE_LOCK, + general_settings_view, llm_router, premium_user, prisma_client, @@ -2542,6 +2757,15 @@ async def delete_model( incoming_params=None, allow_missing_team=True, ) + await _raise_if_team_write_leaves_projects_short( + prisma_client=prisma_client, + general_settings=general_settings_view(), + llm_router=llm_router, + team_id=model_params.model_info.team_id, + public_name=model_params.model_info.team_public_model_name or model_params.model_name, + replaced_id=model_params.model_info.id, + written_params=None, + ) # update DB if store_model_in_db is True: @@ -2708,6 +2932,8 @@ async def add_new_model( """ from litellm.proxy.proxy_server import ( general_settings, + general_settings_view, + llm_router, premium_user, prisma_client, proxy_config, @@ -2766,8 +2992,19 @@ async def add_new_model( ) _raise_if_rate_limits_required_but_missing( + model_name=model_params.model_name, litellm_params=model_params.litellm_params, - enforced=bool(general_settings.get(ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING, False)), + team_model=model_params.model_info.team_id is not None, + policy=_quota_policy(general_settings_view(), ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING), + ) + await _raise_if_team_write_leaves_projects_short( + prisma_client=prisma_client, + general_settings=general_settings_view(), + llm_router=llm_router, + team_id=model_params.model_info.team_id, + public_name=model_params.model_name, + replaced_id=None, + written_params=_written_routing_params(model_params.litellm_params, None), ) clean_model_info: Final = ModelInfo.model_validate( @@ -2906,6 +3143,7 @@ async def update_model( """ from litellm.proxy.proxy_server import ( LITELLM_PROXY_ADMIN_NAME, + general_settings_view, llm_router, premium_user, prisma_client, @@ -2970,6 +3208,30 @@ async def update_model( incoming_params=model_params.litellm_params, existing_params=deployment.litellm_params, ) + renamed_to: Final = ( + model_params.model_name + if model_params.model_name not in (None, deployment.model_name) and deployment.model_info.team_id is None + else None + ) + written_name: Final = renamed_to or deployment.model_info.team_public_model_name or deployment.model_name + _raise_if_model_update_adds_missing_rate_limits( + existing=deployment, + written_name=written_name, + written_team_model=deployment.model_info.team_id is not None, + incoming_params=model_params.litellm_params, + remaining_params=deployment.litellm_params, + policy=_quota_policy(general_settings_view(), ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING), + ) + await _raise_if_team_update_leaves_projects_short( + prisma_client=prisma_client, + general_settings=general_settings_view(), + llm_router=llm_router, + existing=deployment, + written_team_id=deployment.model_info.team_id, + written_name=written_name, + incoming_params=model_params.litellm_params, + remaining_params=deployment.litellm_params, + ) effective_params: Final = _effective_complexity_router_params( model_params.litellm_params, deployment.litellm_params ) @@ -3009,12 +3271,6 @@ async def update_model( if value is not None or _existing_litellm_params_dict.get(key) is not None } - renamed_to: Final = ( - model_params.model_name - if model_params.model_name not in (None, deployment.model_name) - and deployment.model_info.team_id is None - else None - ) base_update: Final[PrismaCompatibleUpdateDBModel] = { "litellm_params": json.dumps(merged_dictionary), "updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, diff --git a/litellm/proxy/management_helpers/model_quota_policy.py b/litellm/proxy/management_helpers/model_quota_policy.py new file mode 100644 index 00000000000..752c78df862 --- /dev/null +++ b/litellm/proxy/management_helpers/model_quota_policy.py @@ -0,0 +1,327 @@ +import re +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal, TypeAlias + +from pydantic import BaseModel, ConfigDict, Field, field_validator + +import litellm +from litellm.exceptions import BadRequestError +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider, is_openai_compatible_endpoint +from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping +from litellm.llms.openai_like.json_loader import JSONProviderRegistry +from litellm.router_utils.common_utils import provider_for_generic_call +from litellm.router_utils.pattern_match_deployments import PatternMatchRouter, PatternUtils +from litellm.secret_managers.main import get_secret_str + +if TYPE_CHECKING: + from litellm import Router + +QuotaLimit: TypeAlias = Literal["rpm", "tpm"] + +ENFORCE_PROJECT_MODEL_QUOTA_SETTING: Final = "enforce_project_model_quota" + +_LIMITS: Final[tuple[QuotaLimit, ...]] = ("rpm", "tpm") +_ROUTER_PATTERNS: Final = PatternMatchRouter() + + +def _is_known_provider(provider: str) -> bool: + return provider in litellm.provider_list or JSONProviderRegistry.exists(provider) + + +class ModelQuotaPolicy(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + default: frozenset[QuotaLimit] = frozenset(_LIMITS) + providers: Mapping[str, frozenset[QuotaLimit]] = Field(default_factory=dict) + models: Mapping[str, frozenset[QuotaLimit]] = Field(default_factory=dict) + + @field_validator("providers") + @classmethod + def _known_providers(cls, providers: Mapping[str, frozenset[QuotaLimit]]) -> Mapping[str, frozenset[QuotaLimit]]: + unknown: Final = sorted(provider for provider in providers if not _is_known_provider(provider)) + if unknown: + raise ValueError(f"unknown provider ids {unknown}") + return providers + + +REQUIRE_BOTH_LIMITS: Final = ModelQuotaPolicy() +_POLICY_KEYS: Final = frozenset(ModelQuotaPolicy.model_fields) + + +@dataclass(frozen=True, slots=True) +class TeamDeployment: + provider: str | None + called_model: str | None + + +@dataclass(frozen=True, slots=True) +class ModelRoute: + providers: frozenset[str] = frozenset() + team_deployments: frozenset[TeamDeployment] = frozenset() + + +RouteOf: TypeAlias = Callable[[str], ModelRoute] +NO_ROUTE: Final = ModelRoute() + + +def parse_model_quota_setting(setting: object) -> ModelQuotaPolicy | None: + """Any falsy value is off and any other value requires both limits, as the boolean setting always + behaved, so only a mapping that names a policy key is read as a per-provider and per-model policy.""" + if not setting: + return None + if is_str_keyed_mapping(setting) and not _POLICY_KEYS.isdisjoint(setting): + return ModelQuotaPolicy.model_validate(setting) + return REQUIRE_BOTH_LIMITS + + +def _wildcard_matches(pattern: str, model: str) -> bool: + return re.match(_ROUTER_PATTERNS.pattern_to_regex(pattern), model) is not None + + +def _router_specificity(pattern: str) -> tuple[int, int]: + return PatternUtils.calculate_pattern_specificity(_ROUTER_PATTERNS.pattern_to_regex(pattern)) + + +def _most_specific_model_pattern(model: str, entries: Mapping[str, frozenset[QuotaLimit]]) -> str | None: + matching: Final = tuple(pattern for pattern in entries if "*" in pattern and _wildcard_matches(pattern, model)) + return max(matching, key=_router_specificity, default=None) + + +def _model_entry(entries: Mapping[str, frozenset[QuotaLimit]], model: str) -> frozenset[QuotaLimit] | None: + exact: Final = entries.get(model) + if exact is not None: + return exact + pattern: Final = _most_specific_model_pattern(model, entries) + return None if pattern is None else entries[pattern] + + +def _provider_limits(policy: ModelQuotaPolicy, provider: str | None) -> frozenset[QuotaLimit]: + return policy.default if provider is None else policy.providers.get(provider, policy.default) + + +def _team_deployment_limits(policy: ModelQuotaPolicy, deployment: TeamDeployment) -> frozenset[QuotaLimit]: + entry: Final = None if deployment.called_model is None else _model_entry(policy.models, deployment.called_model) + return _provider_limits(policy, deployment.provider) if entry is None else entry + + +def required_limits(policy: ModelQuotaPolicy, model: str, route: ModelRoute) -> frozenset[QuotaLimit]: + """A model entry (exact name, then the most specific `*` pattern) wins over the providers of the + model's deployments, which win over the default. A model served by several providers needs every + limit any of them requires. A team picks its model's public name, so a team deployment matches + model entries by the model it calls instead.""" + entry: Final = None if route.team_deployments else _model_entry(policy.models, model) + if entry is not None: + return entry + if not route.providers and not route.team_deployments: + return policy.default + return frozenset[QuotaLimit]().union( + *(_team_deployment_limits(policy, deployment) for deployment in route.team_deployments), + *(_provider_limits(policy, provider) for provider in route.providers), + ) + + +def _is_positive(value: object) -> bool: + return isinstance(value, int) and not isinstance(value, bool) and value > 0 + + +def unmet_limits(required: frozenset[QuotaLimit], values: Mapping[QuotaLimit, object]) -> tuple[QuotaLimit, ...]: + return tuple( + limit + for limit in _LIMITS + if (limit in required or values.get(limit) is not None) and not _is_positive(values.get(limit)) + ) + + +def no_route(_model: str) -> ModelRoute: + return NO_ROUTE + + +def _resolved_api_base(api_base: object) -> object: + if isinstance(api_base, str) and api_base.startswith("os.environ/"): + return get_secret_str(api_base) + return api_base + + +def _provider_behind_api_base(litellm_params: Mapping[str, object]) -> str | None: + model: Final = litellm_params.get("model") + api_base: Final = _resolved_api_base(litellm_params.get("api_base")) + if litellm_params.get("custom_llm_provider") or not isinstance(model, str) or not isinstance(api_base, str): + return None + if model.split("/", 1)[0] in litellm.provider_list or not is_openai_compatible_endpoint(api_base): + return None + try: + _, provider, _, _ = get_llm_provider(model=model, api_base=api_base) + except BadRequestError: + return None + return provider + + +def deployment_provider(litellm_params: object) -> str | None: + if not is_str_keyed_mapping(litellm_params): + return None + return _provider_behind_api_base(litellm_params) or provider_for_generic_call(litellm_params) + + +def _called_model(litellm_params: Mapping[str, object]) -> str | None: + model: Final = litellm_params.get("model") + if not isinstance(model, str): + return None + prefix, _, rest = model.partition("/") + return rest if rest and _is_known_provider(prefix) else model + + +def team_deployment(litellm_params: object) -> TeamDeployment: + if not is_str_keyed_mapping(litellm_params): + return TeamDeployment(provider=None, called_model=None) + return TeamDeployment(provider=deployment_provider(litellm_params), called_model=_called_model(litellm_params)) + + +def deployment_route(litellm_params: Mapping[str, object], team_model: bool) -> ModelRoute: + if team_model: + return ModelRoute(team_deployments=frozenset((team_deployment(litellm_params),))) + provider: Final = deployment_provider(litellm_params) + return ModelRoute(providers=frozenset() if provider is None else frozenset((provider,))) + + +def _named_deployments(llm_router: "Router", model: str, team_id: str | None) -> tuple[Mapping[str, object], ...]: + return ( + *llm_router.get_all_deployments(model_name=model, team_id=team_id), + *llm_router.get_model_list_from_routing_groups(model_name=model), + ) + + +def _routed_deployments(llm_router: "Router", model: str, team_id: str | None) -> Sequence[Mapping[str, object]]: + target: Final = llm_router.routable_model_group(model) + return _named_deployments(llm_router, target, team_id) or llm_router.wildcard_route_deployments(target, team_id) + + +def _is_team_deployment(deployment: Mapping[str, object]) -> bool: + model_info: Final = deployment.get("model_info") + return is_str_keyed_mapping(model_info) and model_info.get("team_id") is not None + + +def router_model_routes(llm_router: "Router | None", team_id: str | None) -> RouteOf: + if llm_router is None: + return no_route + + def route_of(model: str) -> ModelRoute: + deployments: Final = _routed_deployments(llm_router, model, team_id) + shared: Final = tuple(deployment for deployment in deployments if not _is_team_deployment(deployment)) + return ModelRoute( + providers=frozenset( + provider + for deployment in shared + if (provider := deployment_provider(deployment.get("litellm_params"))) is not None + ), + team_deployments=frozenset( + team_deployment(deployment.get("litellm_params")) + for deployment in deployments + if _is_team_deployment(deployment) + ), + ) + + return route_of + + +def _deployment_id(deployment: Mapping[str, object]) -> object: + model_info: Final = deployment.get("model_info") + return model_info.get("id") if is_str_keyed_mapping(model_info) else None + + +def _routable_name(llm_router: "Router | None", model: str) -> str: + return model if llm_router is None else llm_router.routable_model_group(model) + + +def _patterns_tried_before(patterns: Sequence[str], regex: str) -> Sequence[str]: + if regex in patterns: + return patterns[: patterns.index(regex)] + specificity: Final = PatternUtils.calculate_pattern_specificity(regex) + return tuple(rival for rival in patterns if PatternUtils.calculate_pattern_specificity(rival) >= specificity) + + +def _team_pattern_serves(llm_router: "Router", team_id: str, pattern: str, model: str) -> bool: + regex: Final = _ROUTER_PATTERNS.pattern_to_regex(pattern) + if re.match(regex, model) is None or _named_deployments(llm_router, model, team_id): + return False + team_router: Final = llm_router.team_pattern_routers.get(team_id) + patterns: Final = () if team_router is None else tuple(team_router.patterns) + return not any(re.match(rival, model) is not None for rival in _patterns_tried_before(patterns, regex)) + + +def team_write_serves(llm_router: "Router | None", team_id: str, public_name: str, model: str) -> bool: + """Whether a request from the team for ``model`` reaches the team route named ``public_name``: the name + itself or an alias of it, or, for a `*` name, a model no named deployment shadows and no team pattern the + router tries first takes.""" + target: Final = _routable_name(llm_router, model) + if "*" not in public_name or llm_router is None: + return target == public_name + return _team_pattern_serves(llm_router, team_id, public_name, target) + + +def _served_params(litellm_params: object, public_name: str, model: str) -> object: + if "*" not in public_name or not is_str_keyed_mapping(litellm_params): + return litellm_params + matched: Final = re.match(_ROUTER_PATTERNS.pattern_to_regex(public_name), model) + called: Final = litellm_params.get("model") + if matched is None or not isinstance(called, str): + return litellm_params + return MappingProxyType( + { + **litellm_params, + "model": PatternMatchRouter.set_deployment_model_name( + matched_pattern=matched, litellm_deployment_litellm_model=called + ), + } + ) + + +def team_route_after_write( + llm_router: "Router | None", + team_id: str, + public_name: str, + replaced_id: str | None, + written_params: Mapping[str, object] | None, + model: str, +) -> ModelRoute | None: + deployments: Final = () if llm_router is None else llm_router.get_all_deployments(public_name, team_id=team_id) + kept: Final = tuple( + deployment.get("litellm_params") + for deployment in deployments + if _is_team_deployment(deployment) and _deployment_id(deployment) != replaced_id + ) + team_params: Final = kept if written_params is None else (*kept, written_params) + target: Final = _routable_name(llm_router, model) + if team_params: + return ModelRoute( + team_deployments=frozenset( + team_deployment(_served_params(litellm_params, public_name, target)) for litellm_params in team_params + ) + ) + if llm_router is None or not _routed_deployments(llm_router, model, None): + return None + return router_model_routes(llm_router, None)(model) + + +def limits_added_by_team_write( + policy: ModelQuotaPolicy, + llm_router: "Router | None", + team_id: str, + public_name: str, + replaced_id: str | None, + written_params: Mapping[str, object] | None, + model: str, +) -> frozenset[QuotaLimit]: + after: Final = team_route_after_write(llm_router, team_id, public_name, replaced_id, written_params, model) + if after is None: + return frozenset() + now: Final = router_model_routes(llm_router, team_id)(model) + return required_limits(policy, model, after) - required_limits(policy, model, now) + + +def exempts_any_route(policy: ModelQuotaPolicy) -> bool: + return any( + limits != REQUIRE_BOTH_LIMITS.default + for limits in (policy.default, *policy.providers.values(), *policy.models.values()) + ) diff --git a/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py b/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py index 45697723af1..3f14098f6a9 100644 --- a/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py +++ b/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py @@ -1,7 +1,7 @@ import asyncio import os import traceback -from collections.abc import Mapping +from collections.abc import Callable, Mapping from dataclasses import dataclass from datetime import datetime, timezone from typing import Final @@ -30,6 +30,7 @@ from litellm_enterprise.proxy.management_endpoints.project_endpoints import ( delete_project, project_info, ) +from litellm.proxy.management_helpers.model_quota_policy import ModelRoute, TeamDeployment, router_model_routes from litellm.proxy.proxy_server import ( LitellmUserRoles, ) @@ -1205,7 +1206,9 @@ def test_update_quota_existing_quotaless_model_rejected(): _raise_on_missing_project_model_quota_on_update(data, existing) -def _enforced_new_project_mocks(monkeypatch, team_models: list[str], llm_router: mock.MagicMock | None) -> None: +def _enforced_new_project_mocks( + monkeypatch: pytest.MonkeyPatch, team_models: list[str], llm_router: object, quota_setting: object = True +) -> None: from litellm.proxy._types import LiteLLM_TeamTable from litellm_enterprise.proxy.management_endpoints import project_endpoints as pe @@ -1213,7 +1216,7 @@ def _enforced_new_project_mocks(monkeypatch, team_models: list[str], llm_router: monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", mock.MagicMock()) monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True) monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", llm_router) - monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {"enforce_project_model_quota": True}) + monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {"enforce_project_model_quota": quota_setting}) monkeypatch.setattr(pe, "_validate_team_exists", mock.AsyncMock(return_value=team)) monkeypatch.setattr(pe, "_check_user_permission_for_project", mock.AsyncMock(return_value=True)) @@ -1604,6 +1607,559 @@ async def test_new_project_flag_on_access_group_model_returns_400(monkeypatch): assert "expand to multiple models at request time" in str(exc_info.value) +_CUSTOMER_QUOTA_POLICY: Final = { + "providers": {"vertex_ai": ["rpm", "tpm"], "runwayml": []}, + "models": {"veo-*": ["rpm"], "jev-latest": []}, +} + +_CUSTOMER_MODEL_PROVIDERS: Final = { + "gemini-3.8-flash": frozenset({"vertex_ai"}), + "veo-3.1-lite-generate-001": frozenset({"vertex_ai"}), + "runway-gen4-turbo": frozenset({"runwayml"}), +} + + +def _customer_route_of(model: str) -> ModelRoute: + return ModelRoute(providers=_CUSTOMER_MODEL_PROVIDERS.get(model, frozenset())) + + +def _quota_error( + models: list[str], + rpm: dict[str, int] | None = None, + tpm: dict[str, int] | None = None, + setting: object = _CUSTOMER_QUOTA_POLICY, + route_of: Callable[[str], ModelRoute] = _customer_route_of, +) -> str | None: + from litellm_enterprise.proxy.management_endpoints.project_endpoints import ( + _raise_on_missing_project_model_quota, + ) + from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting + + policy: Final = parse_model_quota_setting(setting) + assert policy is not None + data: Final = NewProjectRequest(team_id="test-team", models=models, model_rpm_limit=rpm, model_tpm_limit=tpm) + try: + _raise_on_missing_project_model_quota(data, frozenset(), policy, route_of) + except HTTPException as e: + rejection: Final = e + else: + return None + assert rejection.status_code == 400 + return str(rejection.detail) + + +def test_quota_policy_exempts_provider_and_passthrough_name_but_keeps_vertex_strict(): + assert _quota_error(["runway-gen4-turbo", "jev-latest"]) is None + error: Final = _quota_error(["gemini-3.8-flash"], rpm={"gemini-3.8-flash": 10}) + assert error is not None + assert "gemini-3.8-flash (model_tpm_limit)" in error + + +def test_quota_policy_model_pattern_overrides_its_provider_per_limit(): + assert _quota_error(["veo-3.1-lite-generate-001"], rpm={"veo-3.1-lite-generate-001": 5}) is None + error: Final = _quota_error(["veo-3.1-lite-generate-001"], tpm={"veo-3.1-lite-generate-001": 5000}) + assert error is not None + assert "veo-3.1-lite-generate-001 (model_rpm_limit)" in error + assert "model_tpm_limit" not in error + + +@pytest.mark.parametrize( + ("model", "rpm", "tpm", "expected_missing"), + [ + ("runway-gen4-turbo", {"runway-gen4-turbo": 0}, None, "runway-gen4-turbo (model_rpm_limit)"), + ("jev-latest", None, {"jev-latest": -5}, "jev-latest (model_tpm_limit)"), + ( + "veo-3.1-lite-generate-001", + {"veo-3.1-lite-generate-001": 5}, + {"veo-3.1-lite-generate-001": 0}, + "veo-3.1-lite-generate-001 (model_tpm_limit)", + ), + ], +) +def test_quota_policy_rejects_a_non_positive_limit_set_on_an_exempt_model( + model: str, rpm: dict[str, int] | None, tpm: dict[str, int] | None, expected_missing: str +): + error: Final = _quota_error([model], rpm=rpm, tpm=tpm) + assert error is not None + assert expected_missing in error + + +def test_quota_policy_error_names_each_model_with_the_limits_it_misses(): + error: Final = _quota_error(["gemini-3.8-flash", "veo-3.1-lite-generate-001", "runway-gen4-turbo"]) + assert error is not None + assert "gemini-3.8-flash (model_rpm_limit and model_tpm_limit)" in error + assert "veo-3.1-lite-generate-001 (model_rpm_limit)" in error + assert "runway-gen4-turbo" not in error + + +@pytest.mark.parametrize( + "setting, model, expected_missing", + [ + ({"default": ["tpm"]}, "unrouted-name", "unrouted-name (model_tpm_limit)"), + ({"default": [], "providers": {"runwayml": ["rpm"]}}, "runway-gen4-turbo", "(model_rpm_limit)"), + ( + {"providers": {"runwayml": []}, "models": {"runway-gen4-turbo": ["tpm"]}}, + "runway-gen4-turbo", + "(model_tpm_limit)", + ), + ({"models": {"runway-*": [], "runway-gen4-turbo": ["rpm"]}}, "runway-gen4-turbo", "(model_rpm_limit)"), + ({"models": {"veo-*": [], "veo-3.1-lite-*": ["rpm"]}}, "veo-3.1-lite-generate-001", "(model_rpm_limit)"), + ({"models": {"gpt-5-mini*": [], "gpt-*-*-*": ["rpm"]}}, "gpt-5-mini-fast", "(model_rpm_limit)"), + ], +) +def test_quota_policy_precedence_is_model_then_pattern_then_provider_then_default( + setting: object, model: str, expected_missing: str +): + error: Final = _quota_error([model], setting=setting) + assert error is not None + assert expected_missing in error + + +def test_quota_policy_model_entry_matches_names_literally_apart_from_star(): + setting: Final = {"models": {"claude-opus-5-5[1m]": [], "gpt.5": []}} + assert _quota_error(["claude-opus-5-5[1m]"], setting=setting) is None + assert _quota_error(["claude-opus-5-51"], setting=setting) is not None + assert _quota_error(["gptx5"], setting=setting) is not None + + +def test_quota_policy_model_on_several_providers_needs_every_required_limit(): + def mixed_providers(_model: str) -> ModelRoute: + return ModelRoute(providers=frozenset({"runwayml", "vertex_ai"})) + + error: Final = _quota_error(["video-pool"], rpm={"video-pool": 10}, route_of=mixed_providers) + assert error is not None + assert "video-pool (model_tpm_limit)" in error + + +def test_quota_policy_model_pattern_matches_from_the_start_like_the_router(): + setting: Final = {"providers": {"vertex_ai": ["rpm", "tpm"]}, "models": {"*-lite": []}} + assert _quota_error(["veo-3.1-lite-generate-001"], setting=setting) is None + assert _quota_error(["gemini-3.8-flash"], setting=setting) is not None + + +def _team_vertex_route(called_model: str) -> Callable[[str], ModelRoute]: + def route_of(_model: str) -> ModelRoute: + return ModelRoute(team_deployments=frozenset({TeamDeployment(provider="vertex_ai", called_model=called_model)})) + + return route_of + + +def test_quota_policy_matches_a_team_model_by_the_model_it_calls(): + borrowed_name: Final = _quota_error(["jev-latest"], route_of=_team_vertex_route("gemini-3.8-flash")) + veo_call: Final = _quota_error( + ["video-team"], rpm={"video-team": 5}, route_of=_team_vertex_route("veo-3.1-generate-001") + ) + assert borrowed_name is not None and "jev-latest (model_rpm_limit and model_tpm_limit)" in borrowed_name + assert veo_call is None + + +def test_quota_policy_team_model_beside_a_shared_one_needs_both_routes_limits(): + def team_vertex_beside_shared_runway(_model: str) -> ModelRoute: + return ModelRoute( + providers=frozenset({"runwayml"}), + team_deployments=frozenset({TeamDeployment(provider="vertex_ai", called_model="gemini-3.8-flash")}), + ) + + error: Final = _quota_error(["media-latest"], rpm={"media-latest": 5}, route_of=team_vertex_beside_shared_runway) + assert error is not None and "media-latest (model_tpm_limit)" in error + + +@pytest.mark.parametrize("setting", [True, "true", "false", 1, ["rpm"], {"enabled": True}, {"provider": {"runwayml": []}}]) +def test_quota_setting_truthy_non_policy_value_still_requires_both_limits(setting: object): + error: Final = _quota_error(["runway-gen4-turbo"], setting=setting) + assert error is not None + assert "runway-gen4-turbo (model_rpm_limit and model_tpm_limit)" in error + + +@pytest.mark.parametrize("setting", [False, None, {}, "", 0, []]) +def test_quota_setting_off_enforces_nothing(setting: object): + from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting + + assert parse_model_quota_setting(setting) is None + + +def test_update_quota_policy_resolves_providers_for_the_resulting_models(): + import types + + from litellm_enterprise.proxy.management_endpoints.project_endpoints import ( + _raise_on_missing_project_model_quota_on_update, + ) + from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting + + policy: Final = parse_model_quota_setting(_CUSTOMER_QUOTA_POLICY) + assert policy is not None + existing: Final = types.SimpleNamespace(models=["runway-gen4-turbo"], metadata={}) + assert ( + _raise_on_missing_project_model_quota_on_update( + UpdateProjectRequest(project_id="p", description="unrelated change"), + existing, + frozenset(), + policy, + _customer_route_of, + ) + is None + ) + with pytest.raises(HTTPException) as exc_info: + _raise_on_missing_project_model_quota_on_update( + UpdateProjectRequest(project_id="p", models=["runway-gen4-turbo", "gemini-3.8-flash"]), + existing, + frozenset(), + policy, + _customer_route_of, + ) + assert "gemini-3.8-flash (model_rpm_limit and model_tpm_limit)" in str(exc_info.value.detail) + + +def _customer_router() -> "litellm.Router": + return litellm.Router( + model_list=[ + { + "model_name": "gemini-3.8-flash", + "litellm_params": {"model": "vertex_ai/gemini-3.8-flash", "vertex_project": "qa-project"}, + }, + { + "model_name": "veo-3.1-lite-generate-001", + "litellm_params": {"model": "vertex_ai/veo-3.1-lite-generate-001", "vertex_project": "qa-project"}, + }, + { + "model_name": "runway-gen4-turbo", + "litellm_params": {"model": "runwayml/gen4_turbo", "api_key": "qa-runway-key"}, + }, + ] + ) + + +def test_quota_providers_come_from_the_one_wildcard_route_requests_use(): + router: Final = litellm.Router( + model_list=[ + {"model_name": "gen4*", "litellm_params": {"model": "runwayml/gen4*", "api_key": "qa-runway-key"}}, + { + "model_name": "media-team-gen", + "litellm_params": {"model": "vertex_ai/*", "vertex_project": "qa-project"}, + "model_info": {"team_id": "media-team", "team_public_model_name": "gen*"}, + }, + ] + ) + + assert router_model_routes(router, "media-team")("gen4_turbo") == ModelRoute( + team_deployments=frozenset({TeamDeployment(provider="vertex_ai", called_model="4_turbo")}) + ) + assert router_model_routes(router, None)("gen4_turbo") == ModelRoute(providers=frozenset({"runwayml"})) + + +def test_quota_providers_read_the_provider_behind_a_known_api_base(): + groq_api_base: Final = "https://api.groq.com/openai/v1" + router: Final = litellm.Router( + model_list=[ + { + "model_name": "qa-llama", + "litellm_params": {"model": "qa-llama", "api_base": groq_api_base, "api_key": "qa-groq-key"}, + }, + { + "model_name": "qa-runway", + "litellm_params": {"model": "gen4_turbo", "custom_llm_provider": "runwayml", "api_base": groq_api_base}, + }, + ] + ) + route_of: Final = router_model_routes(router, None) + + assert route_of("qa-llama").providers == frozenset({"groq"}) + assert route_of("qa-runway").providers == frozenset({"runwayml"}) + + +def test_quota_policy_rejects_an_unknown_provider_id(): + from pydantic import ValidationError + + from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting + + with pytest.raises(ValidationError, match=r"unknown provider ids \['runway'\]"): + parse_model_quota_setting({"providers": {"vertex_ai": ["rpm"], "runway": []}}) + + +def test_quota_policy_accepts_a_provider_registered_at_runtime(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting + + monkeypatch.setattr(litellm, "provider_list", [*litellm.provider_list, "qa-custom-llm"]) + policy: Final = parse_model_quota_setting({"providers": {"qa-custom-llm": []}}) + + assert policy is not None + assert policy.providers["qa-custom-llm"] == frozenset() + + +_TEAM_WRITE_POLICY: Final = {"providers": {"vertex_ai": ["rpm", "tpm"], "runwayml": []}, "models": {"jev-latest": []}} +_VERTEX_PARAMS: Final = {"model": "vertex_ai/gemini-3.8-flash", "vertex_project": "qa-project"} +_RUNWAY_PARAMS: Final = {"model": "gen4_turbo", "custom_llm_provider": "runwayml"} + + +def _team_deployment(public_name: str, model_id: str, litellm_params: Mapping[str, object]) -> dict[str, object]: + return { + "model_name": f"model_name_media-team_{model_id}", + "litellm_params": dict(litellm_params), + "model_info": {"id": model_id, "team_id": "media-team", "team_public_model_name": public_name}, + } + + +@pytest.mark.parametrize( + ("setting", "public_name", "proxy_params", "team_routes", "replaced_id", "written_params", "added"), + [ + pytest.param( + _TEAM_WRITE_POLICY, + "jev-latest", + {"model": "openai/jev-latest"}, + (), + None, + _VERTEX_PARAMS, + {"rpm", "tpm"}, + id="team-add-on-vertex-behind-an-exempt-proxy-name", + ), + pytest.param( + _TEAM_WRITE_POLICY, + "jev-latest", + {"model": "openai/jev-latest"}, + (), + None, + _RUNWAY_PARAMS, + set(), + id="team-add-on-an-exempt-provider", + ), + pytest.param( + _TEAM_WRITE_POLICY, + "jev-latest", + {"model": "openai/jev-latest"}, + (), + None, + {"model": "aihubmix/jev-latest"}, + set(), + id="team-add-through-a-json-provider-matches-the-model-it-calls", + ), + pytest.param( + _TEAM_WRITE_POLICY, + "media-latest", + _VERTEX_PARAMS, + (("t1", _RUNWAY_PARAMS),), + "t1", + None, + {"rpm", "tpm"}, + id="team-delete-falls-back-to-the-proxy-vertex-route", + ), + pytest.param( + _TEAM_WRITE_POLICY, + "media-latest", + _VERTEX_PARAMS, + (("t1", _RUNWAY_PARAMS), ("t2", _RUNWAY_PARAMS)), + "t1", + None, + set(), + id="team-delete-keeps-a-sibling-team-route", + ), + pytest.param( + _TEAM_WRITE_POLICY, + "media-latest", + _VERTEX_PARAMS, + (("t1", _RUNWAY_PARAMS),), + "t1", + _VERTEX_PARAMS, + {"rpm", "tpm"}, + id="team-update-moves-the-route-to-vertex", + ), + pytest.param( + True, + "jev-latest", + {"model": "openai/jev-latest"}, + (), + None, + _VERTEX_PARAMS, + set(), + id="true-already-requires-both-limits", + ), + ], +) +def test_team_model_write_adds_only_the_limits_its_new_route_requires( + setting: object, + public_name: str, + proxy_params: Mapping[str, object], + team_routes: tuple[tuple[str, Mapping[str, object]], ...], + replaced_id: str | None, + written_params: Mapping[str, object] | None, + added: set[str], +): + from litellm.proxy.management_helpers.model_quota_policy import ( + limits_added_by_team_write, + parse_model_quota_setting, + ) + + policy: Final = parse_model_quota_setting(setting) + router: Final = litellm.Router( + model_list=[ + {"model_name": public_name, "litellm_params": dict(proxy_params)}, + *(_team_deployment(public_name, model_id, params) for model_id, params in team_routes), + ] + ) + + assert policy is not None + assert limits_added_by_team_write( + policy, router, "media-team", public_name, replaced_id, written_params, public_name + ) == frozenset(added) + + +def test_quota_route_of_an_alias_reads_the_team_deployments_its_target_routes_to(): + router: Final = litellm.Router( + model_list=[ + {"model_name": "media-latest", "litellm_params": dict(_RUNWAY_PARAMS)}, + _team_deployment("media-latest", "t1", _VERTEX_PARAMS), + ], + model_group_alias={"media": "media-latest"}, + ) + route_of: Final = router_model_routes(router, "media-team") + + assert route_of("media") == route_of("media-latest") + assert TeamDeployment(provider="vertex_ai", called_model="gemini-3.8-flash") in route_of("media").team_deployments + + +def _team_pattern_router() -> "litellm.Router": + return litellm.Router( + model_list=[ + {"model_name": "media-latest", "litellm_params": dict(_VERTEX_PARAMS)}, + {"model_name": "video-named", "litellm_params": dict(_VERTEX_PARAMS)}, + _team_deployment("video-*", "t1", {"model": "runwayml/*"}), + _team_deployment("video-veo-*", "t2", {"model": "vertex_ai/veo-*", "vertex_project": "qa-project"}), + _team_deployment("clip-*x", "t3", {"model": "runwayml/*"}), + _team_deployment("clip-x*", "t4", {"model": "vertex_ai/*", "vertex_project": "qa-project"}), + ], + model_group_alias={"media": "media-latest"}, + ) + + +@pytest.mark.parametrize( + ("public_name", "model", "served"), + [ + pytest.param("media-latest", "media-latest", True, id="the-name-itself"), + pytest.param("media-latest", "media", True, id="an-alias-of-the-name"), + pytest.param("media-latest", "video-named", False, id="another-name"), + pytest.param("video-*", "video-gen-1", True, id="a-name-the-pattern-matches"), + pytest.param("video-*", "video-veo-3", False, id="a-more-specific-team-pattern-takes-it"), + pytest.param("video-veo-*", "video-veo-3", True, id="the-most-specific-team-pattern"), + pytest.param("video-*", "video-named", False, id="a-named-deployment-shadows-the-pattern"), + pytest.param("video-*", "image-gen-1", False, id="a-name-the-pattern-misses"), + pytest.param("audio-*", "audio-gen-1", True, id="a-pattern-the-write-adds"), + pytest.param("clip-*x", "clip-xx", True, id="the-first-registered-of-two-tied-team-patterns"), + pytest.param("clip-x*", "clip-xx", False, id="a-tied-team-pattern-registered-first-takes-it"), + pytest.param("cli*-xx", "clip-xx", False, id="a-new-pattern-tied-with-a-registered-one"), + ], +) +def test_team_write_serves_the_project_names_a_team_request_routes_through_it( + public_name: str, model: str, served: bool +): + from litellm.proxy.management_helpers.model_quota_policy import team_write_serves + + assert team_write_serves(_team_pattern_router(), "media-team", public_name, model) is served + + +@pytest.mark.parametrize( + ("model", "added"), + [ + pytest.param("video-gen-1", {"rpm", "tpm"}, id="vertex-requires-both"), + pytest.param("video-veo-3", {"rpm"}, id="the-called-veo-model-entry-wins"), + ], +) +def test_team_pattern_write_reads_each_served_name_by_the_model_it_calls(model: str, added: set[str]): + from litellm.proxy.management_helpers.model_quota_policy import ( + limits_added_by_team_write, + parse_model_quota_setting, + ) + + policy: Final = parse_model_quota_setting( + {"default": [], "providers": {"vertex_ai": ["rpm", "tpm"]}, "models": {"veo-*": ["rpm"]}} + ) + + assert policy is not None + assert limits_added_by_team_write( + policy, + litellm.Router(model_list=[]), + "media-team", + "video-*", + None, + {"model": "vertex_ai/*", "vertex_project": "qa-project"}, + model, + ) == frozenset(added) + + +@pytest.mark.asyncio +async def test_new_project_quota_policy_reads_providers_from_router_deployments(monkeypatch: pytest.MonkeyPatch): + team_models: Final = ["gemini-3.8-flash", "veo-3.1-lite-generate-001", "runway-gen4-turbo", "jev-latest"] + _enforced_new_project_mocks( + monkeypatch, team_models=team_models, llm_router=_customer_router(), quota_setting=_CUSTOMER_QUOTA_POLICY + ) + + with pytest.raises(ProxyException) as exc_info: + await _run_new_project(NewProjectRequest(team_id="test-team", models=team_models)) + + assert exc_info.value.code == "400" + assert ( + "gemini-3.8-flash (model_rpm_limit and model_tpm_limit), veo-3.1-lite-generate-001 (model_rpm_limit)." + in str(exc_info.value.message) + ) + assert "runway-gen4-turbo" not in str(exc_info.value.message) + assert "jev-latest" not in str(exc_info.value.message) + + +@pytest.mark.asyncio +async def test_new_project_rejects_malformed_quota_setting_instead_of_skipping_it(monkeypatch: pytest.MonkeyPatch): + _enforced_new_project_mocks( + monkeypatch, + team_models=["runway-gen4-turbo"], + llm_router=_customer_router(), + quota_setting={"providers": {"runwayml": []}, "model": {"jev-latest": []}}, + ) + + with pytest.raises(ProxyException) as exc_info: + await _run_new_project(NewProjectRequest(team_id="test-team", models=["runway-gen4-turbo"])) + + assert exc_info.value.code == "500" + assert "general_settings.enforce_project_model_quota is invalid" in str(exc_info.value.message) + + +@pytest.mark.asyncio +async def test_new_project_checks_permission_before_reading_the_quota_setting(monkeypatch: pytest.MonkeyPatch): + from litellm_enterprise.proxy.management_endpoints import project_endpoints as pe + + _enforced_new_project_mocks( + monkeypatch, + team_models=["runway-gen4-turbo"], + llm_router=_customer_router(), + quota_setting={"providers": {"runwayml": []}, "model": {"jev-latest": []}}, + ) + monkeypatch.setattr(pe, "_check_user_permission_for_project", mock.AsyncMock(return_value=False)) + + with pytest.raises(ProxyException) as exc_info: + await _run_new_project(NewProjectRequest(team_id="test-team", models=["runway-gen4-turbo"])) + + assert exc_info.value.code == "403" + assert "enforce_project_model_quota" not in str(exc_info.value.message) + + +@pytest.mark.asyncio +async def test_new_project_quota_policy_ignores_model_entries_for_a_team_model_name(monkeypatch: pytest.MonkeyPatch): + router: Final = litellm.Router( + model_list=[ + { + "model_name": "model_name_test-team_qa", + "litellm_params": {"model": "vertex_ai/gemini-3.8-flash", "vertex_project": "qa-project"}, + "model_info": {"team_id": "test-team", "team_public_model_name": "jev-latest"}, + }, + ] + ) + _enforced_new_project_mocks( + monkeypatch, team_models=["jev-latest"], llm_router=router, quota_setting=_CUSTOMER_QUOTA_POLICY + ) + + with pytest.raises(ProxyException) as exc_info: + await _run_new_project(NewProjectRequest(team_id="test-team", models=["jev-latest"])) + + assert exc_info.value.code == "400" + assert "jev-latest (model_rpm_limit and model_tpm_limit)" in str(exc_info.value.message) + + _ACCESS_NOW: Final = datetime.now(timezone.utc) diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index 9ff2a6de9dd..c8f601de280 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -15,6 +15,7 @@ from fastapi.testclient import TestClient from litellm._uuid import uuid from litellm.models.credentials import CredentialItem +from litellm.models.project import LiteLLM_ProjectTable from litellm.proxy._types import ( LiteLLM_ModelTable, @@ -31,6 +32,7 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( KubernetesPodsResponse, ModelManagementAuthChecks, _get_team_deployments, + _raise_if_model_update_adds_missing_rate_limits, _raise_if_rate_limits_required_but_missing, clear_cache, delete_team_models, @@ -38,6 +40,7 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( patch_model, update_model, ) +from litellm.proxy.management_helpers.model_quota_policy import REQUIRE_BOTH_LIMITS, ModelQuotaPolicy from litellm.proxy.utils import PrismaClient from litellm.router import Router from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo, updateDeployment, updateLiteLLMParams @@ -7369,12 +7372,19 @@ class TestAutoRouterClassifierDefaultPrompt: assert response.system_prompt == classification_system_prompt(5) +_CUSTOMER_MODEL_ADD_POLICY: Final = ModelQuotaPolicy.model_validate( + { + "providers": {"vertex_ai": ["rpm", "tpm"], "runwayml": []}, + "models": {"veo-*": ["rpm"], "jev-latest": []}, + } +) + + class TestEnforceRpmTpmOnModelAdd: def test_passes_when_disabled_even_without_limits(self): assert ( _raise_if_rate_limits_required_but_missing( - litellm_params=LiteLLM_Params(model="azure/gpt-5.2"), - enforced=False, + model_name="gpt-5.2", team_model=False, litellm_params=LiteLLM_Params(model="azure/gpt-5.2"), policy=None ) is None ) @@ -7382,8 +7392,9 @@ class TestEnforceRpmTpmOnModelAdd: def test_passes_when_enabled_and_both_set(self): assert ( _raise_if_rate_limits_required_but_missing( + model_name="gpt-5.2", team_model=False, litellm_params=LiteLLM_Params(model="azure/gpt-5.2", rpm=10, tpm=1000), - enforced=True, + policy=REQUIRE_BOTH_LIMITS, ) is None ) @@ -7398,14 +7409,618 @@ class TestEnforceRpmTpmOnModelAdd: (LiteLLM_Params(model="azure/gpt-5.2", rpm=10, tpm=-1), "tpm"), ], ) - def test_raises_when_enabled_and_missing(self, params, expected_missing): - from litellm.proxy._types import ProxyException - + def test_raises_when_enabled_and_missing(self, params: LiteLLM_Params, expected_missing: str): with pytest.raises(ProxyException) as exc_info: - _raise_if_rate_limits_required_but_missing(litellm_params=params, enforced=True) - assert expected_missing in str(exc_info.value.message) + _raise_if_rate_limits_required_but_missing( + model_name="gpt-5.2", team_model=False, litellm_params=params, policy=REQUIRE_BOTH_LIMITS + ) + assert exc_info.value.message.startswith(f"{expected_missing} must be set") assert exc_info.value.code == "400" + @pytest.mark.parametrize( + ("model_name", "params"), + [ + ("runway-gen4", LiteLLM_Params(model="runwayml/gen4_turbo")), + ("my-runway", LiteLLM_Params(model="gen4_turbo", custom_llm_provider="runwayml")), + ("jev-latest", LiteLLM_Params(model="openai/jev-latest")), + ("veo-3.1-generate-001", LiteLLM_Params(model="vertex_ai/veo-3.1-generate-001", rpm=10)), + ("gemini-3-pro", LiteLLM_Params(model="vertex_ai/gemini-3-pro", rpm=10, tpm=1000)), + ], + ) + def test_policy_asks_each_model_only_for_the_limits_its_entry_requires( + self, model_name: str, params: LiteLLM_Params + ): + assert ( + _raise_if_rate_limits_required_but_missing( + model_name=model_name, team_model=False, litellm_params=params, policy=_CUSTOMER_MODEL_ADD_POLICY + ) + is None + ) + + @pytest.mark.parametrize( + ("model_name", "params", "expected_missing"), + [ + ("veo-3.1-generate-001", LiteLLM_Params(model="vertex_ai/veo-3.1-generate-001", tpm=1000), "rpm"), + ("gemini-3-pro", LiteLLM_Params(model="vertex_ai/gemini-3-pro", rpm=10), "tpm"), + ("openai-gpt", LiteLLM_Params(model="openai/gpt-5.2"), "rpm and tpm"), + ("runway-gen4", LiteLLM_Params(model="runwayml/gen4_turbo", rpm=0), "rpm"), + ], + ) + def test_policy_rejects_a_model_missing_a_required_or_carrying_a_non_positive_limit( + self, model_name: str, params: LiteLLM_Params, expected_missing: str + ): + with pytest.raises(ProxyException) as exc_info: + _raise_if_rate_limits_required_but_missing( + model_name=model_name, team_model=False, litellm_params=params, policy=_CUSTOMER_MODEL_ADD_POLICY + ) + assert exc_info.value.message.startswith(f"{expected_missing} must be set to a positive value for '{model_name}'") + assert exc_info.value.code == "400" + + @pytest.mark.parametrize("api_base", ["https://api.groq.com/openai/v1", "os.environ/QA_GROQ_API_BASE"]) + def test_policy_reads_the_provider_behind_a_known_api_base(self, monkeypatch: pytest.MonkeyPatch, api_base: str): + monkeypatch.setenv("QA_GROQ_API_BASE", "https://api.groq.com/openai/v1") + policy: Final = ModelQuotaPolicy.model_validate({"default": [], "providers": {"groq": ["rpm", "tpm"]}}) + with pytest.raises(ProxyException) as exc_info: + _raise_if_rate_limits_required_but_missing( + model_name="qa-llama", team_model=False, + litellm_params=LiteLLM_Params(model="qa-llama", api_base=api_base), + policy=policy, + ) + assert exc_info.value.message.startswith("rpm and tpm must be set to a positive value for 'qa-llama'") + + @pytest.mark.parametrize( + ("model_name", "params", "shared_missing", "team_missing"), + [ + pytest.param( + "jev-latest", LiteLLM_Params(model="vertex_ai/gemini-3.8-flash"), None, "rpm and tpm", + id="team-cannot-borrow-an-exempt-public-name", + ), + pytest.param( + "video-team", LiteLLM_Params(model="vertex_ai/veo-3.1-generate-001", rpm=5), "tpm", None, + id="team-matches-the-model-it-calls", + ), + ], + ) + def test_policy_matches_a_team_model_by_the_model_it_calls( + self, model_name: str, params: LiteLLM_Params, shared_missing: str | None, team_missing: str | None + ): + assert _model_add_error(model_name, params, team_model=False) == _missing_limits_message(shared_missing, model_name) + assert _model_add_error(model_name, params, team_model=True) == _missing_limits_message(team_missing, model_name) + + +def _missing_limits_message(missing: str | None, model_name: str) -> str | None: + if missing is None: + return None + return ( + f"{missing} must be set to a positive value for '{model_name}' when " + "'enforce_rpm_tpm_on_model_add' is enabled in general_settings" + ) + + +def _model_add_error(model_name: str, params: LiteLLM_Params, team_model: bool) -> str | None: + try: + _raise_if_rate_limits_required_but_missing( + model_name=model_name, team_model=team_model, litellm_params=params, policy=_CUSTOMER_MODEL_ADD_POLICY + ) + except ProxyException as error: + return error.message + return None + + +_STORED_RUNWAY: Final = {"model": "runwayml/gen4_turbo"} + + +def _stored_model(litellm_params: Mapping[str, object], team_id: str | None) -> Deployment: + return Deployment( + model_name="media-latest" if team_id is None else "model_name_media-team_team-row", + litellm_params=dict(litellm_params), + model_info={ + "id": "team-row", + "team_id": team_id, + "team_public_model_name": None if team_id is None else "media-latest", + }, + ) + + +@pytest.mark.parametrize( + ("stored", "team_id", "incoming", "missing"), + [ + pytest.param( + _STORED_RUNWAY, "media-team", updateLiteLLMParams(model="vertex_ai/gemini-3.8-flash"), "rpm and tpm", + id="team-model-moves-to-a-provider-that-needs-limits", + ), + pytest.param( + _STORED_RUNWAY, None, updateLiteLLMParams(model="vertex_ai/gemini-3.8-flash", rpm=10), "tpm", + id="proxy-model-moves-with-only-one-limit", + ), + pytest.param( + {**_STORED_RUNWAY, "tpm": 1000}, None, updateLiteLLMParams(model="vertex_ai/gemini-3.8-flash", rpm=10), + None, id="stored-limit-counts-toward-the-move", + ), + pytest.param( + _STORED_RUNWAY, "media-team", updateLiteLLMParams(model="vertex_ai/veo-3.1-generate-001", rpm=5), None, + id="team-model-moves-to-a-model-entry-it-meets", + ), + pytest.param( + {"model": "vertex_ai/gemini-3.8-flash"}, None, updateLiteLLMParams(api_key="qa-rotated-key"), None, + id="unchanged-route-adds-no-limits", + ), + pytest.param( + {"model": "vertex_ai/gemini-3.8-flash", "rpm": 10, "tpm": 1000}, None, updateLiteLLMParams(rpm=0), None, + id="unchanged-route-can-zero-a-limit-as-on-main", + ), + pytest.param( + {"model": "vertex_ai/gemini-3.8-flash"}, None, updateLiteLLMParams(rpm=10), None, + id="legacy-model-without-tpm-stays-editable", + ), + ], +) +def test_model_update_requires_only_the_limits_its_new_route_adds( + stored: Mapping[str, object], team_id: str | None, incoming: updateLiteLLMParams, missing: str | None +) -> None: + assert _model_update_error(_stored_model(stored, team_id), incoming) == _missing_limits_message( + missing, "media-latest" + ) + + +def _model_update_error( + existing: Deployment, + incoming: updateLiteLLMParams, + policy: ModelQuotaPolicy = _CUSTOMER_MODEL_ADD_POLICY, + remaining: LiteLLM_Params | None = None, +) -> str | None: + try: + _raise_if_model_update_adds_missing_rate_limits( + existing=existing, + written_name="media-latest", + written_team_model=existing.model_info.team_id is not None, + incoming_params=incoming, + remaining_params=existing.litellm_params if remaining is None else remaining, + policy=policy, + ) + except ProxyException as error: + return error.message + return None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("setting", "expected_code", "expected_message"), + [ + ({"default": [], "models": {"veo-*": ["rpm"]}}, "400", "rpm must be set to a positive value for 'veo-3.1'"), + ("yes", "400", "rpm and tpm must be set to a positive value for 'veo-3.1'"), + ({"enabled": True}, "400", "rpm and tpm must be set to a positive value for 'veo-3.1'"), + ( + {"providers": {"vertex_ai": ["qps"]}}, + "500", + "general_settings.enforce_rpm_tpm_on_model_add is invalid", + ), + ], +) +async def test_model_add_resolves_the_quota_policy_for_the_public_model_name( + monkeypatch: pytest.MonkeyPatch, setting: object, expected_code: str, expected_message: str +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model + + prisma: Final = MagicMock() + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "general_settings", {"enforce_rpm_tpm_on_model_add": setting}) + with pytest.raises(ProxyException) as error: + await add_new_model( + Deployment(model_name="veo-3.1", litellm_params={"model": "vertex_ai/veo-3.1"}), + UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert error.value.code == expected_code + assert error.value.message.startswith(expected_message) + prisma.db.litellm_proxymodeltable.create.assert_not_called() + + +@pytest.mark.asyncio +async def test_team_model_add_cannot_borrow_a_model_entry_by_its_public_name(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model + + prisma: Final = MagicMock() + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=LiteLLM_TeamTable(team_id="media-team")) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "premium_user", True) + monkeypatch.setattr( + proxy_server, + "general_settings", + {"enforce_rpm_tpm_on_model_add": {"providers": {"vertex_ai": ["rpm", "tpm"]}, "models": {"jev-latest": []}}}, + ) + with pytest.raises(ProxyException) as error: + await add_new_model( + Deployment( + model_name="jev-latest", + litellm_params={"model": "vertex_ai/gemini-3.8-flash"}, + model_info={"team_id": "media-team"}, + ), + UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert error.value.code == "400" + assert error.value.message.startswith("rpm and tpm must be set to a positive value for 'jev-latest'") + prisma.db.litellm_proxymodeltable.create.assert_not_called() + + +_PROJECT_QUOTA_SETTINGS: Final = { + "enforce_project_model_quota": { + "providers": {"vertex_ai": ["rpm", "tpm"], "runwayml": []}, + "models": {"jev-latest": []}, + } +} +_TEAM_VERTEX: Final = {"model": "vertex_ai/gemini-3.8-flash", "vertex_project": "qa-project"} +_TEAM_RUNWAY: Final = {"model": "gen4_turbo", "custom_llm_provider": "runwayml"} + + +def _media_team_db(public_name: str, project_metadata: Mapping[str, object]) -> MagicMock: + prisma: Final = MagicMock() + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=LiteLLM_TeamTable(team_id="media-team")) + prisma.db.litellm_projecttable.find_many = AsyncMock( + return_value=[ + LiteLLM_ProjectTable( + project_id="p-1", + project_alias="video-gen", + team_id="media-team", + models=[public_name], + metadata=dict(project_metadata), + ), + LiteLLM_ProjectTable( + project_id="p-2", project_alias="chat-only", team_id="media-team", models=["chat-only-model"] + ), + ] + ) + return prisma + + +def _team_model_row(public_name: str, litellm_params: Mapping[str, object]) -> LiteLLM_ProxyModelTable: + return LiteLLM_ProxyModelTable( + model_id="team-row", + model_name="model_name_media-team_team-row", + litellm_params=dict(litellm_params), + model_info={"id": "team-row", "team_id": "media-team", "team_public_model_name": public_name}, + created_by="admin", + updated_by="admin", + ) + + +def _quota_router(public_name: str, proxy_params: Mapping[str, object], *team_params: Mapping[str, object]) -> Router: + return Router( + model_list=[ + {"model_name": public_name, "litellm_params": dict(proxy_params)}, + *( + { + "model_name": "model_name_media-team_team-row", + "litellm_params": dict(params), + "model_info": {"id": "team-row", "team_id": "media-team", "team_public_model_name": public_name}, + } + for params in team_params + ), + ] + ) + + +def _use_quota_proxy(monkeypatch: pytest.MonkeyPatch, prisma: MagicMock, router: Router) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "general_settings", _PROJECT_QUOTA_SETTINGS) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "store_model_in_db", False) + monkeypatch.setattr(proxy_server, "premium_user", True) + + +_ADMIN: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + +@pytest.mark.asyncio +async def test_team_model_add_rejects_a_route_its_projects_lack_the_quota_for(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model + + prisma: Final = _media_team_db("jev-latest", {}) + _use_quota_proxy(monkeypatch, prisma, _quota_router("jev-latest", {"model": "openai/jev-latest"})) + + with pytest.raises(ProxyException) as error: + await add_new_model( + Deployment( + model_name="jev-latest", + litellm_params={**_TEAM_VERTEX, "rpm": 10, "tpm": 1000}, + model_info={"team_id": "media-team"}, + ), + _ADMIN, + ) + + assert error.value.code == "400" + assert error.value.message.startswith( + "projects video-gen list 'jev-latest' without a positive rpm and tpm quota, " + "which 'enforce_project_model_quota' requires once this team model change applies" + ) + prisma.db.litellm_projecttable.find_many.assert_awaited_once_with(where={"team_id": "media-team"}) + + +@pytest.mark.parametrize( + ("litellm_params", "project_metadata"), + [ + pytest.param( + _TEAM_VERTEX, + {"model_rpm_limit": {"jev-latest": 5}, "model_tpm_limit": {"jev-latest": 500}}, + id="projects-carry-the-added-limits", + ), + pytest.param(_TEAM_RUNWAY, {}, id="new-route-is-exempt"), + ], +) +@pytest.mark.asyncio +async def test_team_model_add_goes_through_when_its_projects_meet_the_new_route( + monkeypatch: pytest.MonkeyPatch, litellm_params: Mapping[str, object], project_metadata: Mapping[str, object] +) -> None: + from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model + + _use_quota_proxy( + monkeypatch, _media_team_db("jev-latest", project_metadata), _quota_router("jev-latest", {"model": "openai/jev-latest"}) + ) + + with pytest.raises(ProxyException) as error: + await add_new_model( + Deployment(model_name="jev-latest", litellm_params=dict(litellm_params), model_info={"team_id": "media-team"}), + _ADMIN, + ) + + assert "STORE_MODEL_IN_DB" in error.value.message + + +@pytest.mark.parametrize("endpoint", ["patch", "put"]) +@pytest.mark.asyncio +async def test_team_model_update_rejects_moving_its_route_where_its_projects_lack_quota( + monkeypatch: pytest.MonkeyPatch, endpoint: str +) -> None: + from litellm.proxy import proxy_server + + runway_by_prefix: Final = {"model": "runwayml/gen4_turbo"} + prisma: Final = _media_team_db("jev-latest", {}) + prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( + return_value=_team_model_row("jev-latest", runway_by_prefix) + ) + _use_quota_proxy(monkeypatch, prisma, _quota_router("jev-latest", {"model": "openai/jev-latest"}, runway_by_prefix)) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + moved_to_vertex: Final = updateLiteLLMParams(model="vertex_ai/gemini-3.8-flash", vertex_project="qa-project") + + write: Final = ( + patch_model( + model_id="team-row", patch_data=updateDeployment(litellm_params=moved_to_vertex), user_api_key_dict=_ADMIN + ) + if endpoint == "patch" + else update_model( + model_params=updateDeployment(litellm_params=moved_to_vertex, model_info=ModelInfo(id="team-row")), + user_api_key_dict=_ADMIN, + ) + ) + + with pytest.raises(ProxyException) as error: + await write + + assert error.value.code == "400" + assert error.value.message.startswith("projects video-gen list 'jev-latest' without a positive rpm and tpm quota") + + +def _team_route_router(public_name: str, team_params: Mapping[str, object]) -> Router: + return Router( + model_list=[ + {"model_name": "jev-latest", "litellm_params": {"model": "openai/jev-latest"}}, + { + "model_name": "model_name_media-team_team-row", + "litellm_params": dict(team_params), + "model_info": {"id": "team-row", "team_id": "media-team", "team_public_model_name": public_name}, + }, + ], + model_group_alias={"media": "jev-latest"}, + ) + + +@pytest.mark.parametrize( + ("public_name", "project_model", "team_params", "moved_to"), + [ + pytest.param( + "video-*", + "video-gen-1", + {"model": "runwayml/*"}, + updateLiteLLMParams(model="vertex_ai/*", vertex_project="qa-project"), + id="a-name-the-team-wildcard-serves", + ), + pytest.param( + "jev-latest", + "media", + {"model": "runwayml/gen4_turbo"}, + updateLiteLLMParams(model="vertex_ai/gemini-3.8-flash", vertex_project="qa-project"), + id="an-alias-of-the-team-name", + ), + ], +) +@pytest.mark.asyncio +async def test_team_model_patch_rechecks_projects_listing_any_name_it_serves( + monkeypatch: pytest.MonkeyPatch, + public_name: str, + project_model: str, + team_params: Mapping[str, object], + moved_to: updateLiteLLMParams, +) -> None: + from litellm.proxy import proxy_server + + prisma: Final = _media_team_db(project_model, {}) + prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=_team_model_row(public_name, team_params)) + _use_quota_proxy(monkeypatch, prisma, _team_route_router(public_name, team_params)) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + + with pytest.raises(ProxyException) as error: + await patch_model( + model_id="team-row", patch_data=updateDeployment(litellm_params=moved_to), user_api_key_dict=_ADMIN + ) + + assert error.value.code == "400" + assert error.value.message.startswith( + f"projects video-gen list '{project_model}' without a positive rpm and tpm quota" + ) + prisma.db.litellm_proxymodeltable.update.assert_not_called() + + +@pytest.mark.asyncio +async def test_model_update_checks_the_limits_of_the_name_it_renames_to(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + + prisma: Final = MagicMock() + prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( + return_value=LiteLLM_ProxyModelTable( + model_id="vertex-row", + model_name="jev-latest", + litellm_params={"model": "vertex_ai/gemini-3.8-flash"}, + model_info={"id": "vertex-row"}, + created_by="admin", + updated_by="admin", + ) + ) + _use_quota_proxy(monkeypatch, prisma, Router(model_list=[])) + monkeypatch.setattr( + proxy_server, + "general_settings", + {"enforce_rpm_tpm_on_model_add": {"providers": {"vertex_ai": ["rpm", "tpm"]}, "models": {"jev-latest": []}}}, + ) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + + with pytest.raises(ProxyException) as error: + await update_model( + model_params=updateDeployment( + model_name="gemini-3.8-flash", litellm_params=updateLiteLLMParams(), model_info=ModelInfo(id="vertex-row") + ), + user_api_key_dict=_ADMIN, + ) + + assert error.value.code == "400" + assert error.value.message.startswith("rpm and tpm must be set to a positive value for 'gemini-3.8-flash'") + prisma.db.litellm_proxymodeltable.update.assert_not_called() + + +@pytest.mark.parametrize( + ("incoming", "remaining"), + [ + pytest.param(updateLiteLLMParams(rpm=0), {"model": "vertex_ai/gemini-3.8-flash", "tpm": 1000}, id="zeroes-rpm"), + pytest.param(updateLiteLLMParams(), {"model": "vertex_ai/gemini-3.8-flash", "tpm": 1000}, id="removes-rpm"), + pytest.param(updateLiteLLMParams(model="runwayml/gen4_turbo"), {"model": "runwayml/gen4_turbo"}, id="moves"), + ], +) +def test_model_update_under_the_boolean_setting_never_rejects_as_on_main( + incoming: updateLiteLLMParams, remaining: Mapping[str, object] +) -> None: + from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting + + stored: Final = _stored_model({"model": "vertex_ai/gemini-3.8-flash", "rpm": 10, "tpm": 1000}, None) + policy: Final = parse_model_quota_setting(True) + + assert policy is not None + assert _model_update_error(stored, incoming, policy, LiteLLM_Params(**remaining)) is None + + +@pytest.mark.asyncio +async def test_team_model_patch_rejects_a_removal_that_moves_its_route_where_its_projects_lack_quota( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy import proxy_server + from litellm.types.proxy.management_endpoints.model_management_endpoints import PatchModelRequest + + prisma: Final = _media_team_db("jev-latest", {}) + prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=_team_model_row("jev-latest", _TEAM_RUNWAY)) + _use_quota_proxy(monkeypatch, prisma, _quota_router("jev-latest", {"model": "openai/jev-latest"}, _TEAM_RUNWAY)) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + + with pytest.raises(ProxyException) as error: + await patch_model( + model_id="team-row", + patch_data=PatchModelRequest(litellm_params_to_remove=("custom_llm_provider",)), + user_api_key_dict=_ADMIN, + ) + + assert error.value.code == "400" + assert error.value.message.startswith("projects video-gen list 'jev-latest' without a positive rpm and tpm quota") + + +@pytest.mark.asyncio +async def test_team_model_delete_rejects_falling_back_to_a_route_its_projects_lack_quota_for( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy.management_endpoints.model_management_endpoints import ModelInfoDelete, delete_model + + prisma: Final = _media_team_db("media-latest", {}) + prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=_team_model_row("media-latest", _TEAM_RUNWAY)) + _use_quota_proxy(monkeypatch, prisma, _quota_router("media-latest", _TEAM_VERTEX, _TEAM_RUNWAY)) + + with pytest.raises(ProxyException) as error: + await delete_model(model_info=ModelInfoDelete(id="team-row"), user_api_key_dict=_ADMIN) + + assert error.value.code == "400" + assert error.value.message.startswith("projects video-gen list 'media-latest' without a positive rpm and tpm quota") + prisma.db.litellm_proxymodeltable.delete.assert_not_called() + + +@pytest.mark.asyncio +async def test_team_model_delete_goes_through_when_nothing_else_serves_its_name( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy.management_endpoints.model_management_endpoints import ModelInfoDelete, delete_model + + prisma: Final = _media_team_db("media-latest", {}) + prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=_team_model_row("media-latest", _TEAM_RUNWAY)) + only_the_team_model: Final = Router( + model_list=[ + { + "model_name": "model_name_media-team_team-row", + "litellm_params": dict(_TEAM_RUNWAY), + "model_info": {"id": "team-row", "team_id": "media-team", "team_public_model_name": "media-latest"}, + } + ] + ) + _use_quota_proxy(monkeypatch, prisma, only_the_team_model) + + with pytest.raises(ProxyException) as error: + await delete_model(model_info=ModelInfoDelete(id="team-row"), user_api_key_dict=_ADMIN) + + assert "STORE_MODEL_IN_DB" in str(error.value.message) + + +@pytest.mark.parametrize("endpoint", ["patch", "put"]) +@pytest.mark.asyncio +async def test_team_model_update_rejects_moving_to_a_provider_that_needs_limits_it_lacks( + monkeypatch: pytest.MonkeyPatch, endpoint: str +) -> None: + from litellm.proxy import proxy_server + + prisma: Final = _media_team_db("media-latest", {}) + prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=_team_model_row("media-latest", _TEAM_RUNWAY)) + _use_quota_proxy(monkeypatch, prisma, _quota_router("media-latest", _TEAM_RUNWAY, _TEAM_RUNWAY)) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr( + proxy_server, + "general_settings", + {"enforce_rpm_tpm_on_model_add": {"providers": {"vertex_ai": ["rpm", "tpm"], "runwayml": []}}}, + ) + moved_to_vertex: Final = updateLiteLLMParams(model="vertex_ai/gemini-3.8-flash", custom_llm_provider="vertex_ai") + + write: Final = ( + patch_model( + model_id="team-row", patch_data=updateDeployment(litellm_params=moved_to_vertex), user_api_key_dict=_ADMIN + ) + if endpoint == "patch" + else update_model( + model_params=updateDeployment(litellm_params=moved_to_vertex, model_info=ModelInfo(id="team-row")), + user_api_key_dict=_ADMIN, + ) + ) + + with pytest.raises(ProxyException) as error: + await write + + assert error.value.code == "400" + assert error.value.message == _missing_limits_message("rpm and tpm", "media-latest") + prisma.db.litellm_proxymodeltable.update.assert_not_called() + class TestBlockModelResponseSerialization: @pytest.mark.parametrize( From 295d79fc7d9f935165e07c66e150a8579d511d4c Mon Sep 17 00:00:00 2001 From: rj-cairn <349020452+rj-cairn@users.noreply.github.com> Date: Wed, 28 Oct 2026 05:36:26 -0700 Subject: [PATCH 033/128] feat(cli): keep the ChatGPT login with lite codex --subscription (#48741) * feat(cli): keep the ChatGPT login with lite codex --subscription `lite codex --subscription` points Codex at a `litellm_subscription` provider with requires_openai_auth on and no env_key, so Codex keeps its own ChatGPT login and the LiteLLM key is sent only in the x-litellm-api-key header from LITELLM_PROXY_API_KEY. OPENAI_API_KEY is left exactly as the user had it, and the pre-launch key check sends the key the same way `codex_subscription` is a new `lite config` key that turns the flag on by default, like `claude_subscription` does for `lite claude`. The Claude `--settings` overlay now applies to Claude Code only * fix(cli): leave the Codex Stop hook out of subscription launches and load tomli on Python 3.10 * fix(cli): leave OPENAI_BASE_URL as the user had it with lite codex --subscription * fix(cli): give lite codex --subscription its own Codex provider id * test(cli): type the Codex config override helpers * fix(cli): leave the Codex model list to Codex in subscription mode --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../chatgpt_subscription_passthrough.md | 2 + litellm/proxy/client/cli/commands/agents.py | 87 ++++-- litellm/proxy/client/cli/commands/config.py | 14 +- tests/unit/proxy/client/cli/test_agents.py | 270 +++++++++++++++--- .../proxy/client/cli/test_config_commands.py | 49 ++-- 5 files changed, 341 insertions(+), 81 deletions(-) diff --git a/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md b/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md index c79632e9195..e1f606a2eb7 100644 --- a/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md +++ b/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md @@ -40,6 +40,8 @@ env_http_headers = { "x-litellm-api-key" = "LITELLM_API_KEY" } Replace `https://litellm.example.com` with your proxy URL. A local proxy normally uses `http://localhost:4000/v1`. Launch Codex, sign into ChatGPT if needed, and send a prompt with a model available to your subscription. The selected model must match the name exposed by the gateway. +The `lite` CLI can set up this provider each time it starts Codex, with no `config.toml` edit. Export `LITELLM_PROXY_URL` and `LITELLM_PROXY_API_KEY`, then run `lite codex --subscription`. Run `lite config set codex_subscription true` to make it the default. + For a managed workspace rollout, use OpenAI's [gateway configuration guide](https://learn.chatgpt.com/docs/enterprise/sign-in-with-chatgpt-through-a-gateway). ## Native compatibility and Pi diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index 52e36964bc1..4d09e4d1e0f 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -34,7 +34,7 @@ from .claude_settings import ( load_json_or_empty, ) from .cmd_quoting import quote_for_cmd -from .config import claude_subscription_default +from .config import subscription_default from .pi import ( LITELLM_PROXY_API_KEY_ENV, PI_PROVIDER_NAME, @@ -83,9 +83,11 @@ _INSTALL_DOCS: Final[dict[str, str]] = { } _HIDDEN_AGENTS: Final = frozenset({"pi"}) -_SUBSCRIPTION_AGENTS: Final = frozenset({"claude"}) +_SUBSCRIPTION_LOGINS: Final = MappingProxyType({"claude": "claude.ai", "codex": "ChatGPT"}) +_SUBSCRIPTION_AGENTS: Final = frozenset(_SUBSCRIPTION_LOGINS) CODEX_PROXY_PROVIDER: Final = "litellm" +CODEX_SUBSCRIPTION_PROVIDER: Final = "lite_subscription" CODEX_HOME_ENV: Final = "CODEX_HOME" CODEX_MODEL_CATALOG_FILENAME: Final = "litellm-models.json" _CODEX_PREFLIGHT_TIMEOUT_SECONDS: Final = 10.0 @@ -143,6 +145,10 @@ def build_agent_env( travels in the x-litellm-api-key header through ANTHROPIC_CUSTOM_HEADERS, and no ANTHROPIC_AUTH_TOKEN or ANTHROPIC_API_KEY is left in the environment, since either one makes Claude Code stop using the subscription (and its connectors). + With subscription=True the OpenAI profile keeps Codex on its ChatGPT login: the key goes to + LITELLM_PROXY_API_KEY, which Codex sends in the x-litellm-api-key header, and OPENAI_API_KEY and + OPENAI_BASE_URL are left exactly as the caller had them, so the key is never offered to Codex as an + OpenAI credential and a child process never pairs the caller's own OpenAI key with the proxy URL. """ env: Final = dict(base_env) root: Final = base_url.rstrip("/") @@ -161,21 +167,24 @@ def build_agent_env( if ENABLE_GATEWAY_MODEL_DISCOVERY_ENV not in env: env[ENABLE_GATEWAY_MODEL_DISCOVERY_ENV] = ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE if PROFILE_OPENAI in profiles: - env[OPENAI_BASE_URL_ENV] = root + "/v1" - env[OPENAI_API_KEY_ENV] = api_key + if subscription: + env[LITELLM_PROXY_API_KEY_ENV] = api_key + else: + env[OPENAI_BASE_URL_ENV] = root + "/v1" + env[OPENAI_API_KEY_ENV] = api_key if PROFILE_LITELLM in profiles: env[LITELLM_PROXY_API_KEY_ENV] = api_key return env -def codex_proxy_provider(base_url: str) -> Mapping[str, str | bool]: +def codex_proxy_provider(base_url: str, *, subscription: bool = False) -> Mapping[str, str | bool]: return MappingProxyType( { "name": "LiteLLM proxy", "base_url": base_url.rstrip("/") + "/v1", "wire_api": "responses", "supports_websockets": False, - "requires_openai_auth": False, + "requires_openai_auth": subscription, } ) @@ -304,7 +313,7 @@ def subscription_launch_args( return (_SETTINGS_FLAG, str(overlay)), rest -def _codex_proxy_args(base_url: str) -> list[str]: +def _codex_proxy_args(base_url: str, subscription: bool) -> list[str]: """Codex `-c` overrides that point it at the proxy. Codex ignores OPENAI_BASE_URL (it always dials api.openai.com), so the env @@ -312,24 +321,35 @@ def _codex_proxy_args(base_url: str) -> list[str]: inline; supports_websockets=false forces the HTTP/SSE Responses transport because the proxy does not speak the Responses WebSocket protocol. The key is read from OPENAI_API_KEY, which build_agent_env already exports. + With subscription=True the provider is `lite_subscription` with requires_openai_auth, so Codex + signs in with its own ChatGPT login, and the key travels in the x-litellm-api-key header from + LITELLM_PROXY_API_KEY. It sets no env_key, which would win over the login. A `-c` override cannot + remove keys, so the id is one no config.toml defines: `lite configure codex` writes `litellm`, and + the cookbook tells users to write `litellm_subscription`, whose env_key or headers would otherwise apply. """ - provider: Final = f"model_providers.{CODEX_PROXY_PROVIDER}" + provider_id: Final = CODEX_SUBSCRIPTION_PROVIDER if subscription else CODEX_PROXY_PROVIDER + provider: Final = f"model_providers.{provider_id}" + credential: Final = ( + f"{provider}.env_http_headers={{{json.dumps(LITELLM_KEY_HEADER)}={json.dumps(LITELLM_PROXY_API_KEY_ENV)}}}" + if subscription + else f'{provider}.env_key="{OPENAI_API_KEY_ENV}"' + ) return [ "-c", - f'model_provider="{CODEX_PROXY_PROVIDER}"', + f'model_provider="{provider_id}"', *( argument - for key, value in codex_proxy_provider(base_url).items() + for key, value in codex_proxy_provider(base_url, subscription=subscription).items() for argument in ("-c", f"{provider}.{key}={json.dumps(value)}") ), "-c", - f'{provider}.env_key="{OPENAI_API_KEY_ENV}"', + credential, "-c", f"{provider}.http_headers={{}}", ] -_PROXY_ARGS: Final[dict[str, Callable[[str], list[str]]]] = { +_PROXY_ARGS: Final[dict[str, Callable[[str, bool], list[str]]]] = { "codex": _codex_proxy_args, } @@ -410,14 +430,14 @@ _Preparer: TypeAlias = Callable[[str, str, Mapping[str, str]], Sequence[str]] _PREPARERS: Final[Mapping[str, _Preparer]] = MappingProxyType({"pi": prepare_pi, "codex": prepare_codex}) -def agent_launch_args(command: str, base_url: str) -> list[str]: +def agent_launch_args(command: str, base_url: str, *, subscription: bool = False) -> list[str]: """Extra CLI args an agent needs to actually honor the proxy. Claude Code and OpenCode respect the exported env vars, so they get nothing here; Codex needs its provider pointed via config overrides. """ builder: Final = _PROXY_ARGS.get(os.path.basename(command)) - return builder(base_url) if builder else [] + return builder(base_url, subscription) if builder else [] class ListedModel(LiteLLMBaseModel): @@ -889,11 +909,13 @@ def run_agent( (verify_subscription if subscription else verify)(base_url, api_key) env_before_sync: Final = base_env if base_env is not None else os.environ - synced: Final = sync_models(binary, env_before_sync, base_url, api_key, skip_verify) + synced: Final = ( + _NO_EXTRA_ENV if subscription else sync_models(binary, env_before_sync, base_url, api_key, skip_verify) + ) if isinstance(synced, ModelSyncSkipped): warn(f"litellm: not syncing {display_name} models from the proxy: {synced.reason}") - prepare: Final = preparers.get(os.path.basename(command[0])) + prepare: Final = None if subscription else preparers.get(os.path.basename(command[0])) prepared_args: Final = tuple(prepare(base_url, api_key, env_before_sync)) if prepare is not None else () env: Final = MappingProxyType( @@ -903,7 +925,11 @@ def run_agent( } ) synced_args: Final = synced.args if isinstance(synced, ModelSyncArgs) else () - extra_args: Final = (*agent_launch_args(command[0], base_url), *synced_args, *prepared_args) + extra_args: Final = ( + *agent_launch_args(command[0], base_url, subscription=subscription), + *synced_args, + *prepared_args, + ) settings_args, forwarded = ( subscription_launch_args( command, @@ -914,7 +940,7 @@ def run_agent( ), write_overlay, ) - if subscription and os.path.basename(command[0]) in _SUBSCRIPTION_AGENTS + if subscription and os.path.basename(command[0]) == "claude" else ((), tuple(command[1:])) ) if reattach_terminal is not None: @@ -948,10 +974,19 @@ def resolve_api_key(ctx: click.Context) -> str: _SKIP_VERIFY_HELP: Final = "Skip the pre-launch key check against the proxy." -_SUBSCRIPTION_HELP: Final = ( - "Keep using your claude.ai login (and its connectors). The LiteLLM key is sent in the " - "x-litellm-api-key header instead of replacing Claude Code's own credential. " - "Make it the default with `lite config set claude_subscription true`." +_SUBSCRIPTION_HELP: Final = MappingProxyType( + { + "claude": ( + "Keep using your claude.ai login (and its connectors). The LiteLLM key is sent in the " + "x-litellm-api-key header instead of replacing Claude Code's own credential. " + "Make it the default with `lite config set claude_subscription true`." + ), + "codex": ( + "Keep using your ChatGPT login. The LiteLLM key is sent in the " + "x-litellm-api-key header instead of replacing Codex's own credential. " + "Make it the default with `lite config set codex_subscription true`." + ), + } ) @@ -964,10 +999,8 @@ def _launch( api_key: Final = resolve_api_key(ctx) display_name, _profiles = agent_profile(binary) - use_subscription: Final = ( - subscription if subscription is not None else binary in _SUBSCRIPTION_AGENTS and claude_subscription_default() - ) - login_note: Final = " (keeping your claude.ai login)" if use_subscription else "" + use_subscription: Final = subscription if subscription is not None else subscription_default(binary) + login_note: Final = f" (keeping your {_SUBSCRIPTION_LOGINS[binary]} login)" if use_subscription else "" click.echo(f"litellm: routing {display_name} through proxy at {base_url.rstrip('/')}{login_note}") try: @@ -989,7 +1022,7 @@ _Decorated: Final = TypeVar("_Decorated", bound=Callable[..., object] | click.Co def _subscription_option(binary: str) -> Callable[[_Decorated], _Decorated]: if binary not in _SUBSCRIPTION_AGENTS: return lambda func: func - return click.option("--subscription/--no-subscription", default=None, help=_SUBSCRIPTION_HELP) + return click.option("--subscription/--no-subscription", default=None, help=_SUBSCRIPTION_HELP[binary]) def _make_agent_command(binary: str, display_name: str) -> click.Command: diff --git a/litellm/proxy/client/cli/commands/config.py b/litellm/proxy/client/cli/commands/config.py index 539d88dce4a..1689a08c3de 100644 --- a/litellm/proxy/client/cli/commands/config.py +++ b/litellm/proxy/client/cli/commands/config.py @@ -14,6 +14,7 @@ from litellm.litellm_core_utils.private_json import ensure_private_dir, write_pr HIDDEN_COMMANDS_KEY: Final = "hidden_commands" CLAUDE_SUBSCRIPTION_KEY: Final = "claude_subscription" +CODEX_SUBSCRIPTION_KEY: Final = "codex_subscription" PKCE_KEY: Final = "pkce" _config_adapter: Final[TypeAdapter[Mapping[str, str]]] = TypeAdapter(Mapping[str, str]) @@ -114,9 +115,15 @@ def _normalize_boolean(key: str) -> Callable[[str], str]: return normalize -def claude_subscription_default() -> bool: - """Whether `lite claude` keeps the claude.ai login unless told otherwise.""" - return get_config_value(CLAUDE_SUBSCRIPTION_KEY) == "true" +_SUBSCRIPTION_KEYS: Final[Mapping[str, str]] = MappingProxyType( + {"claude": CLAUDE_SUBSCRIPTION_KEY, "codex": CODEX_SUBSCRIPTION_KEY} +) + + +def subscription_default(agent: str) -> bool: + """Whether `lite ` keeps the agent's own subscription login unless told otherwise.""" + key: Final = _SUBSCRIPTION_KEYS.get(agent) + return key is not None and get_config_value(key) == "true" def pkce_default() -> bool: @@ -129,6 +136,7 @@ _NORMALIZERS: Final[Mapping[str, Callable[[str], str]]] = MappingProxyType( "base_url": normalize_base_url, HIDDEN_COMMANDS_KEY: _normalize_hidden_commands, CLAUDE_SUBSCRIPTION_KEY: _normalize_boolean(CLAUDE_SUBSCRIPTION_KEY), + CODEX_SUBSCRIPTION_KEY: _normalize_boolean(CODEX_SUBSCRIPTION_KEY), PKCE_KEY: _normalize_boolean(PKCE_KEY), } ) diff --git a/tests/unit/proxy/client/cli/test_agents.py b/tests/unit/proxy/client/cli/test_agents.py index 3481346076f..3c79b916202 100644 --- a/tests/unit/proxy/client/cli/test_agents.py +++ b/tests/unit/proxy/client/cli/test_agents.py @@ -1,10 +1,13 @@ +import functools import inspect import json import os import stat import subprocess import sys +from collections.abc import Mapping, Sequence from pathlib import Path +from typing import Final from unittest.mock import patch import click @@ -31,8 +34,14 @@ from litellm.proxy.client.cli.commands.agents import ( verify_subscription_key, ) from litellm.proxy.client.cli.commands.claude_settings import ClaudeSettingsError +from litellm.proxy.client.cli.commands.config import save_config from litellm.proxy.client.cli.main import cli as lite +if sys.version_info >= (3, 11): + import tomllib +else: + import tomli as tomllib + AGENTS_MODULE = "litellm.proxy.client.cli.commands.agents" CLAUDE_SETTINGS_MODULE = "litellm.proxy.client.cli.commands.claude_settings" @@ -189,6 +198,30 @@ def _overlay(args): return json.loads(Path(args[2]).read_text()) +def _codex_config(args): + assert set(args[0::2]) == {"-c"} + return tomllib.loads("\n".join(args[1::2])) + + +def _set_codex_key(config: Mapping[str, object], keys: Sequence[str], value: object) -> Mapping[str, object]: + head, *rest = keys + if not rest: + return {**config, head: value} + nested: Final = config.get(head, {}) + assert isinstance(nested, Mapping) + return {**config, head: _set_codex_key(nested, rest, value)} + + +def _with_codex_override(config: Mapping[str, object], override: str) -> Mapping[str, object]: + path, value = override.split("=", 1) + return _set_codex_key(config, path.split("."), tomllib.loads(f"value = {value}")["value"]) + + +def _codex_config_over(user_config_toml: str, args: Sequence[str]) -> Mapping[str, object]: + assert set(args[0::2]) == {"-c"} + return functools.reduce(_with_codex_override, args[1::2], tomllib.loads(user_config_toml)) + + class TestAgentProfile: def test_claude_is_anthropic(self): name, profiles = agent_profile("claude") @@ -254,6 +287,7 @@ class TestBuildAgentEnv: env = build_agent_env({}, "http://localhost:4000/", "sk-key", frozenset({"openai"})) assert env["OPENAI_BASE_URL"] == "http://localhost:4000/v1" assert env["OPENAI_API_KEY"] == "sk-key" + assert "LITELLM_PROXY_API_KEY" not in env assert "ANTHROPIC_BASE_URL" not in env assert "ENABLE_TOOL_SEARCH" not in env assert "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" not in env @@ -317,11 +351,21 @@ class TestBuildAgentEnv: ) assert env["ANTHROPIC_CUSTOM_HEADERS"].splitlines() == ["x-team: blue", "x-litellm-api-key: Bearer sk-new"] - def test_subscription_leaves_openai_agents_on_the_bearer_key(self): - env = build_agent_env({}, "http://localhost:4000", "sk-key", frozenset({"openai"}), subscription=True) - assert env["OPENAI_API_KEY"] == "sk-key" + def test_subscription_routes_openai_agents_without_offering_the_key_as_an_openai_credential(self): + env = build_agent_env({}, "http://localhost:4000/", "sk-key", frozenset({"openai"}), subscription=True) + assert "OPENAI_BASE_URL" not in env + assert env["LITELLM_PROXY_API_KEY"] == "sk-key" + assert "OPENAI_API_KEY" not in env assert "ANTHROPIC_CUSTOM_HEADERS" not in env + def test_subscription_leaves_the_users_own_openai_key_and_base_url_as_they_were(self): + base = {"OPENAI_API_KEY": "sk-users-openai", "OPENAI_BASE_URL": "https://api.openai.com/v1"} + env = build_agent_env(base, "http://localhost:4000", "sk-key", frozenset({"openai"}), subscription=True) + assert env["OPENAI_API_KEY"] == "sk-users-openai" + assert env["OPENAI_BASE_URL"] == "https://api.openai.com/v1" + assert env["LITELLM_PROXY_API_KEY"] == "sk-key" + assert base == {"OPENAI_API_KEY": "sk-users-openai", "OPENAI_BASE_URL": "https://api.openai.com/v1"} + def test_default_mode_does_not_touch_custom_headers(self): env = build_agent_env( {"ANTHROPIC_CUSTOM_HEADERS": "x-team: blue"}, "http://localhost:4000", "sk-key", frozenset({"anthropic"}) @@ -354,6 +398,67 @@ class TestAgentLaunchArgs: agent_launch_args("codex", "http://localhost:4000") ) + def test_codex_default_provider_reads_the_key_from_openai_api_key(self): + assert _codex_config(agent_launch_args("codex", "http://localhost:4000/")) == { + "model_provider": "litellm", + "model_providers": { + "litellm": { + "name": "LiteLLM proxy", + "base_url": "http://localhost:4000/v1", + "wire_api": "responses", + "supports_websockets": False, + "requires_openai_auth": False, + "env_key": "OPENAI_API_KEY", + "http_headers": {}, + } + }, + } + + def test_codex_subscription_provider_keeps_the_chatgpt_login_and_sends_the_key_in_the_litellm_header(self): + assert _codex_config( + agent_launch_args("/usr/local/bin/codex", "http://localhost:4000/", subscription=True) + ) == { + "model_provider": "lite_subscription", + "model_providers": { + "lite_subscription": { + "name": "LiteLLM proxy", + "base_url": "http://localhost:4000/v1", + "wire_api": "responses", + "supports_websockets": False, + "requires_openai_auth": True, + "env_http_headers": {"x-litellm-api-key": "LITELLM_PROXY_API_KEY"}, + "http_headers": {}, + } + }, + } + + def test_codex_subscription_leaves_the_cookbook_provider_table_out_of_the_launch(self): + user_config = """ +model_provider = "litellm_subscription" + +[model_providers.litellm_subscription] +name = "LiteLLM subscription" +base_url = "https://litellm.example.com/v1" +wire_api = "responses" +requires_openai_auth = true +env_key = "OPENAI_API_KEY" +http_headers = { "x-team" = "qa" } +""" + config = _codex_config_over( + user_config, agent_launch_args("codex", "http://localhost:4000/", subscription=True) + ) + launched = config["model_providers"][config["model_provider"]] + assert "env_key" not in launched + assert launched["requires_openai_auth"] is True + assert launched["base_url"] == "http://localhost:4000/v1" + assert config["model_providers"]["litellm_subscription"] == tomllib.loads(user_config)["model_providers"][ + "litellm_subscription" + ] + + @pytest.mark.parametrize("agent", ["claude", "opencode", "pi", "mytool"]) + def test_subscription_adds_no_args_for_agents_without_a_proxy_provider(self, agent): + assert agent_launch_args(agent, "http://localhost:4000", subscription=True) == [] + def test_pi_gets_no_static_args(self): assert agent_launch_args("pi", "http://localhost:4000") == [] @@ -1286,6 +1391,25 @@ class TestRunAgent: with pytest.raises(AgentRunError): run_agent("http://localhost:4000", "sk-key", []) + @pytest.mark.parametrize("subscription", [True, False]) + def test_codex_stop_hook_is_added_only_when_codex_sends_the_key_as_its_openai_key(self, subscription): + launched = {} + run_agent( + "http://localhost:4000", + "sk-key", + ["codex"], + subscription=subscription, + base_env={"OPENAI_API_KEY": "sk-users-own-openai-key"}, + which=lambda name: "/usr/local/bin/codex", + verify=lambda *a: None, + verify_subscription=lambda *a: None, + sync_models=lambda *a: ModelSyncSkipped("not needed"), + warn=lambda message: None, + launcher=lambda p, a, e: launched.update(args=tuple(a)), + preparers={"codex": lambda *a: ("-c", "hooks.Stop=lite")}, + ) + assert ("hooks.Stop=lite" in launched["args"]) is not subscription + def test_subscription_launch_probes_with_the_header_and_exports_it(self): probed = [] launched = {} @@ -1574,6 +1698,56 @@ class TestRunAgent: read_settings=lambda path: {}, ) + def test_codex_subscription_launch_keeps_the_chatgpt_login_and_needs_no_settings_overlay(self): + probed = [] + overlays = [] + launched = {} + run_agent( + "http://localhost:4000", + "sk-key", + ["codex", "exec", "--settings", "kept"], + subscription=True, + base_env={"PATH": "/usr/bin"}, + which=lambda name: "/usr/local/bin/codex", + verify=lambda *a: probed.append(("bearer", a)), + verify_subscription=lambda *a: probed.append(("header", a)), + sync_models=lambda *a: {}, + preparers={}, + write_overlay=lambda text, environ: overlays.append(text), + launcher=lambda p, a, e: launched.update(path=p, args=tuple(a), env=dict(e)), + ) + assert probed == [("header", ("http://localhost:4000", "sk-key"))] + assert overlays == [] + assert launched["path"] == "/usr/local/bin/codex" + assert launched["args"][0] == "codex" + assert launched["args"][-3:] == ("exec", "--settings", "kept") + assert _codex_config(launched["args"][1:-3])["model_provider"] == "lite_subscription" + assert "OPENAI_BASE_URL" not in launched["env"] + assert launched["env"]["LITELLM_PROXY_API_KEY"] == "sk-key" + assert "OPENAI_API_KEY" not in launched["env"] + + @pytest.mark.parametrize( + "sync_result", + [ModelSyncArgs(("-c", 'model_catalog_json="/tmp/c.json"')), ModelSyncSkipped("lists no chat models")], + ) + def test_codex_subscription_launch_leaves_the_model_list_to_codex(self, sync_result): + warnings = [] + launched = {} + run_agent( + "http://localhost:4000", + "sk-key", + ["codex"], + subscription=True, + base_env={}, + which=lambda name: "/usr/local/bin/codex", + verify_subscription=lambda *a: None, + sync_models=lambda *a: sync_result, + warn=warnings.append, + launcher=lambda p, a, e: launched.update(args=tuple(a)), + ) + assert warnings == [] + assert "model_catalog_json" not in json.dumps(launched["args"]) + def test_default_launch_still_probes_with_the_bearer(self): probed = [] run_agent( @@ -1964,7 +2138,11 @@ class TestAgentCommands: assert captured["command"] == ["opencode"] assert "routing OpenCode through proxy at http://localhost:4000" in result.output - def test_claude_subscription_flag_reaches_run_agent_and_is_not_forwarded(self): + @pytest.mark.parametrize( + ("agent", "user_args", "login"), + [("claude", ["--resume"], "claude.ai"), ("codex", ["exec", "hi"], "ChatGPT")], + ) + def test_subscription_flag_reaches_run_agent_and_is_not_forwarded(self, agent, user_args, login): captured = {} def fake_run_agent(base_url, api_key, command, **kwargs): @@ -1973,61 +2151,85 @@ class TestAgentCommands: with patch(f"{AGENTS_MODULE}.run_agent", side_effect=fake_run_agent): result = self.runner.invoke( - _agent_command("claude"), - ["--subscription", "--resume"], + _agent_command(agent), + ["--subscription", *user_args], obj={"base_url": "http://localhost:4000", "api_key": "sk-key"}, ) assert result.exit_code == 0, result.output assert captured["subscription"] is True - assert captured["command"] == ["claude", "--resume"] - assert "keeping your claude.ai login" in result.output + assert captured["command"] == [agent, *user_args] + assert f"keeping your {login} login" in result.output + @pytest.mark.parametrize(("agent", "login"), [("claude", "claude.ai"), ("codex", "ChatGPT")]) @pytest.mark.parametrize( - ("configured", "args", "expected"), + ("saved", "args", "expected"), [ - (False, [], False), - (True, [], True), - (True, ["--no-subscription"], False), - (False, ["--subscription"], True), + (None, [], False), + ("false", [], False), + ("true", [], True), + ("true", ["--no-subscription"], False), + ("false", ["--subscription"], True), ], ) - def test_subscription_comes_from_the_flag_else_the_saved_setting(self, configured, args, expected): + def test_subscription_comes_from_the_flag_else_the_saved_setting(self, agent, login, saved, args, expected): + if saved is not None: + save_config({f"{agent}_subscription": saved}) captured = {} - with ( - patch(f"{AGENTS_MODULE}.claude_subscription_default", return_value=configured), - patch( - f"{AGENTS_MODULE}.run_agent", - side_effect=lambda b, k, c, **kw: captured.update(sub=kw.get("subscription"), command=list(c)), - ), + with patch( + f"{AGENTS_MODULE}.run_agent", + side_effect=lambda b, k, c, **kw: captured.update(sub=kw.get("subscription"), command=list(c)), ): result = self.runner.invoke( - _agent_command("claude"), args, obj={"base_url": "http://localhost:4000", "api_key": "sk-key"} + _agent_command(agent), args, obj={"base_url": "http://localhost:4000", "api_key": "sk-key"} ) assert result.exit_code == 0, result.output assert captured["sub"] is expected - assert captured["command"] == ["claude"] - assert ("keeping your claude.ai login" in result.output) is expected + assert captured["command"] == [agent] + assert (f"keeping your {login} login" in result.output) is expected - @pytest.mark.parametrize("agent", ["codex", "opencode"]) - def test_the_saved_claude_setting_does_not_change_other_agents(self, agent): + @pytest.mark.parametrize( + ("saved_key", "agent"), + [ + ("claude_subscription", "codex"), + ("claude_subscription", "opencode"), + ("codex_subscription", "claude"), + ("codex_subscription", "opencode"), + ], + ) + def test_a_saved_subscription_setting_only_changes_its_own_agent(self, saved_key, agent): + save_config({saved_key: "true"}) captured = {} - with ( - patch(f"{AGENTS_MODULE}.claude_subscription_default", return_value=True), - patch( - f"{AGENTS_MODULE}.run_agent", - side_effect=lambda b, k, c, **kw: captured.update(sub=kw.get("subscription")), - ), + with patch( + f"{AGENTS_MODULE}.run_agent", + side_effect=lambda b, k, c, **kw: captured.update(sub=kw.get("subscription")), ): result = self.runner.invoke( _agent_command(agent), [], obj={"base_url": "http://localhost:4000", "api_key": "sk-key"} ) assert result.exit_code == 0, result.output assert captured["sub"] is False + assert "keeping your" not in result.output - @pytest.mark.parametrize("agent", ["codex", "opencode", "pi"]) - def test_only_claude_offers_the_subscription_flag(self, agent): - assert "--subscription" not in {opt for param in _agent_command(agent).params for opt in param.opts} + def test_only_claude_and_codex_offer_the_subscription_flag(self): + offering = { + command.name + for command in agent_commands() + if "--subscription" in {opt for param in command.params for opt in param.opts} + } + assert offering == {"claude", "codex"} + + @pytest.mark.parametrize( + ("agent", "login", "own_credential"), + [("claude", "claude.ai", "Claude Code's"), ("codex", "ChatGPT", "Codex's")], + ) + def test_subscription_help_names_the_agents_own_login_and_its_setting(self, agent, login, own_credential): + result = self.runner.invoke(lite, [agent, "--help"], prog_name="lite") + assert result.exit_code == 0, result.output + text = " ".join(result.output.split()) + assert f"Keep using your {login} login" in text + assert f"instead of replacing {own_credential} own credential" in text + assert f"lite config set {agent}_subscription true" in text def test_skip_verify_is_consumed_not_forwarded(self): captured = {} diff --git a/tests/unit/proxy/client/cli/test_config_commands.py b/tests/unit/proxy/client/cli/test_config_commands.py index 2cd7a76ab95..e99a4311caf 100644 --- a/tests/unit/proxy/client/cli/test_config_commands.py +++ b/tests/unit/proxy/client/cli/test_config_commands.py @@ -9,12 +9,12 @@ from click.testing import CliRunner from litellm.litellm_core_utils.private_json import write_private_json from litellm.proxy.client.cli import cli from litellm.proxy.client.cli.commands.config import ( - claude_subscription_default, get_config_file_path, get_config_value, load_config, pkce_default, save_config, + subscription_default, ) from litellm.proxy.client.cli.interface import show_commands @@ -362,33 +362,48 @@ class TestWritePrivateJson: assert list(tmp_path.glob(".tmp-*")) == [] -class TestClaudeSubscriptionSetting: +@pytest.mark.parametrize(("agent", "key"), [("claude", "claude_subscription"), ("codex", "codex_subscription")]) +class TestSubscriptionSetting: @pytest.mark.parametrize("word", ["true", "TRUE", "on", "1", "yes"]) - def test_truthy_words_are_stored_as_true(self, cli_runner, isolated_home, word): - result = cli_runner.invoke(cli, ["config", "set", "claude_subscription", word]) + def test_truthy_words_are_stored_as_true(self, cli_runner, isolated_home, agent, key, word): + result = cli_runner.invoke(cli, ["config", "set", key, word]) assert result.exit_code == 0, result.output - assert load_config()["claude_subscription"] == "true" - assert claude_subscription_default() is True + assert load_config()[key] == "true" + assert subscription_default(agent) is True @pytest.mark.parametrize("word", ["false", "Off", "0", "no"]) - def test_falsy_words_are_stored_as_false(self, cli_runner, isolated_home, word): - cli_runner.invoke(cli, ["config", "set", "claude_subscription", "true"]) - result = cli_runner.invoke(cli, ["config", "set", "claude_subscription", word]) + def test_falsy_words_are_stored_as_false(self, cli_runner, isolated_home, agent, key, word): + cli_runner.invoke(cli, ["config", "set", key, "true"]) + result = cli_runner.invoke(cli, ["config", "set", key, word]) assert result.exit_code == 0, result.output - assert load_config()["claude_subscription"] == "false" - assert claude_subscription_default() is False + assert load_config()[key] == "false" + assert subscription_default(agent) is False - def test_anything_else_is_rejected_and_nothing_is_saved(self, cli_runner, isolated_home): - result = cli_runner.invoke(cli, ["config", "set", "claude_subscription", "maybe"]) + def test_anything_else_is_rejected_and_nothing_is_saved(self, cli_runner, isolated_home, agent, key): + result = cli_runner.invoke(cli, ["config", "set", key, "maybe"]) assert result.exit_code != 0 - assert "true or false" in result.output - assert "claude_subscription" not in load_config() + assert f"{key} must be true or false" in result.output + assert key not in load_config() - def test_unset_means_off(self, isolated_home): - assert claude_subscription_default() is False + def test_unset_means_off(self, cli_runner, isolated_home, agent, key): + cli_runner.invoke(cli, ["config", "set", key, "true"]) + result = cli_runner.invoke(cli, ["config", "unset", key]) + + assert result.exit_code == 0, result.output + assert subscription_default(agent) is False + + def test_turning_it_on_leaves_every_other_agent_off(self, cli_runner, isolated_home, agent, key): + cli_runner.invoke(cli, ["config", "set", key, "true"]) + + assert {name: subscription_default(name) for name in ("claude", "codex", "opencode", "pi")} == { + "claude": agent == "claude", + "codex": agent == "codex", + "opencode": False, + "pi": False, + } class TestPkceSetting: From 3507781735115b3136972678de08d84d265f237f Mon Sep 17 00:00:00 2001 From: norkelbyte <348963847+norkelbyte@users.noreply.github.com> Date: Wed, 28 Oct 2026 05:54:40 -0700 Subject: [PATCH 034/128] test(proxy): audit #47352 (#48804) * test(proxy): audit #47352 * test(proxy): audit #47352 lazy load during a cold docs build * test(proxy): route the overlapping chat cell to the rig upstream * test(proxy): keep the root path cell from rewriting the packaged UI * test(proxy): prove the lazy and chat cells overlap the cold docs build * test(proxy): warm the chat path before the overlapping chat cell --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../configuration/test_openapi_docs_route.py | 165 ++++++++++++++++++ 1 file changed, 165 insertions(+) create mode 100644 tests/integration/configuration/test_openapi_docs_route.py diff --git a/tests/integration/configuration/test_openapi_docs_route.py b/tests/integration/configuration/test_openapi_docs_route.py new file mode 100644 index 00000000000..ccb30c2f784 --- /dev/null +++ b/tests/integration/configuration/test_openapi_docs_route.py @@ -0,0 +1,165 @@ +import time +from collections.abc import Iterator +from concurrent.futures import Future, ThreadPoolExecutor +from pathlib import Path +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, list_value, object_value +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from pydantic import JsonValue, TypeAdapter + +from litellm.proxy._lazy_features import LAZY_FEATURES + +DOCS_SECONDS: Final = 120 +COLD_DOCS: Final = 4 +CHATS: Final = 8 +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +OWNED_PROXY_TIMEOUT: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +TWO_OWNED_PROXIES_TIMEOUT: Final = pytest.mark.timeout(2 * (2 * graceful_stop_seconds() + 120)) +LAZY_ROUTES_FLAG: Final = "LITELLM_DISABLE_LAZY_ROUTES" +WARMED_FEATURE: Final = next(feature for feature in LAZY_FEATURES if feature.name == "policies") + + +def _timed_get(candidate: Gateway, path: str) -> tuple[httpx.Response, float]: + started: Final = time.perf_counter() + response: Final = candidate.request("GET", path) + return response, time.perf_counter() - started + + +def _liveliness_until(done: Future[tuple[httpx.Response, float]], candidate: Gateway) -> Iterator[tuple[int, float]]: + while not done.done(): + response, seconds = _timed_get(candidate, "/health/liveliness") + yield response.status_code, seconds + + +def _schema(response: httpx.Response) -> dict[str, JsonValue]: + assert response.status_code == 200, response.text[:500] + return JSON_OBJECT.validate_json(response.content) + + +def _servers(response: httpx.Response) -> list[JsonValue]: + return list(list_value(_schema(response).get("servers", []))) + + +def _chat(candidate: Gateway, model: str, key: str, content: str) -> httpx.Response: + return candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": content}], "cache": {"no-cache": True}}, + key=key, + ) + + +def _warmed_feature_operations(response: httpx.Response) -> dict[str, JsonValue]: + paths: Final = object_value(_schema(response)["paths"]) + return {path: operations for path, operations in paths.items() if WARMED_FEATURE.matches(path)} + + +@OWNED_PROXY_TIMEOUT +def test_liveliness_answers_while_a_cold_openapi_schema_builds(gateway: Gateway, tmp_path: Path) -> None: + with ( + owned_proxy_process(gateway, tmp_path, {}, client_timeout=DOCS_SECONDS) as owned, + ThreadPoolExecutor(1) as pool, + ): + docs: Final = pool.submit(_timed_get, owned.gateway, "/openapi.json") + probes: Final = tuple(_liveliness_until(docs, owned.gateway)) + docs_response, docs_seconds = docs.result() + assert "/v1/chat/completions" in object_value(_schema(docs_response)["paths"]) + assert probes, "no liveliness probe ran while the docs request was pending" + assert {status for status, _ in probes} == {200}, probes + slowest: Final = max(seconds for _, seconds in probes) + assert slowest < docs_seconds / 3, ( + f"a liveliness probe waited {slowest:.2f}s of the {docs_seconds:.2f}s cold docs request" + ) + + +@OWNED_PROXY_TIMEOUT +def test_concurrent_cold_openapi_requests_share_the_schema_later_requests_get(gateway: Gateway, tmp_path: Path) -> None: + with ( + owned_proxy_process(gateway, tmp_path, {}, client_timeout=DOCS_SECONDS) as owned, + ThreadPoolExecutor(COLD_DOCS) as pool, + ): + cold: Final = tuple(pool.map(lambda _: owned.gateway.request("GET", "/openapi.json"), range(COLD_DOCS))) + warm: Final = owned.gateway.request("GET", "/openapi.json") + assert [response.status_code for response in (*cold, warm)] == [200] * (COLD_DOCS + 1) + assert {response.content for response in cold} == {warm.content} + assert "/v1/chat/completions" in object_value(_schema(warm)["paths"]) + + +@OWNED_PROXY_TIMEOUT +def test_chat_completions_overlapping_a_cold_openapi_build_all_succeed(gateway: Gateway, tmp_path: Path) -> None: + with ( + owned_proxy_process(gateway, tmp_path, {}, client_timeout=DOCS_SECONDS) as owned, + owned.gateway.scenario() as scenario, + ThreadPoolExecutor(CHATS + 1) as pool, + ): + model: Final = scenario.model() + key: Final = scenario.key(models=[model]) + first_chat: Final = _chat(owned.gateway, model, key, "before the docs build") + docs: Final = pool.submit(owned.gateway.request, "GET", "/openapi.json") + chats: Final = tuple( + pool.map(lambda index: _chat(owned.gateway, model, key, f"overlapping docs build {index}"), range(CHATS)) + ) + chats_returned_during_build: Final = not docs.done() + docs_response: Final = docs.result() + assert first_chat.status_code == 200, first_chat.text[:300] + assert chats_returned_during_build, "the chat calls returned only after the cold docs build finished" + assert [(response.status_code, response.text[:300]) for response in chats if response.status_code != 200] == [] + replies: Final = tuple(JSON_OBJECT.validate_json(response.content) for response in chats) + assert {reply["object"] for reply in replies} == {"chat.completion"} + assert all(list_value(reply["choices"]) for reply in replies), replies + assert "/v1/chat/completions" in object_value(_schema(docs_response)["paths"]) + + +@OWNED_PROXY_TIMEOUT +def test_docs_filtered_mode_serves_only_llm_routes(gateway: Gateway, tmp_path: Path) -> None: + with owned_proxy_process(gateway, tmp_path, {"DOCS_FILTERED": "True"}, client_timeout=DOCS_SECONDS) as owned: + paths: Final = object_value(_schema(owned.gateway.request("GET", "/openapi.json"))["paths"]) + assert "/v1/chat/completions" in paths + assert {"/key/generate", "/user/new", "/team/new"}.isdisjoint(paths), sorted(paths) + + +@OWNED_PROXY_TIMEOUT +def test_each_root_path_prefix_lists_itself_first_in_servers(gateway: Gateway, tmp_path: Path) -> None: + prefixes: Final = ("/tenant-a", "/tenant-b", "/tenant-a") + with owned_proxy_process( + gateway, tmp_path, {"SERVER_ROOT_PATHS": "/tenant-a,/tenant-b"}, client_timeout=DOCS_SECONDS + ) as owned: + served: Final = tuple( + (prefix, _servers(owned.gateway.request("GET", f"{prefix}/openapi.json"))) for prefix in prefixes + ) + assert served == tuple((prefix, [{"url": prefix}]) for prefix in prefixes) + + +@OWNED_PROXY_TIMEOUT +def test_a_scalar_server_root_path_is_listed_once_in_servers(gateway: Gateway, tmp_path: Path) -> None: + overrides: Final = {"SERVER_ROOT_PATH": "/gateway", "LITELLM_UI_PATH": str(tmp_path / "ui")} + with owned_proxy_process(gateway, tmp_path, overrides, client_timeout=DOCS_SECONDS) as owned: + servers: Final = _servers(owned.gateway.request("GET", "/gateway/openapi.json")) + assert servers == [{"url": "/gateway"}] + + +@TWO_OWNED_PROXIES_TIMEOUT +def test_a_lazy_feature_loaded_during_a_cold_openapi_build_shows_its_live_routes_in_the_next_schema( + gateway: Gateway, tmp_path: Path +) -> None: + with owned_proxy_process(gateway, tmp_path, {LAZY_ROUTES_FLAG: "true"}, client_timeout=DOCS_SECONDS) as eager: + live: Final = _warmed_feature_operations(eager.gateway.request("GET", "/openapi.json")) + with ( + owned_proxy_process( + gateway, tmp_path, {}, remove_environment=(LAZY_ROUTES_FLAG,), client_timeout=DOCS_SECONDS + ) as lazy, + ThreadPoolExecutor(1) as pool, + ): + docs: Final = pool.submit(lazy.gateway.request, "GET", "/openapi.json") + warmed: Final = lazy.gateway.request("POST", f"/lazy/warm/{WARMED_FEATURE.name}") + warmed_during_build: Final = not docs.done() + during: Final = _warmed_feature_operations(docs.result()) + after: Final = _warmed_feature_operations(lazy.gateway.request("GET", "/openapi.json")) + assert warmed.status_code == 200, warmed.text + assert warmed_during_build, "the feature finished loading only after the cold docs build finished" + assert live, "the eager proxy published no route of the warmed feature" + assert during != live, "the cold docs build started after the feature had already loaded" + assert after == live From 356069d2edf87d94882e6e4a539c80b1c68840f7 Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Wed, 28 Oct 2026 06:02:40 -0700 Subject: [PATCH 035/128] test(mcp): audit #47422 (#48821) Adds a real-proxy case where the MCP follow-up call fails on every router retry, checking that the one failure row bills the first call of every attempt on chat completions and /v1/responses Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../integration/mcp/test_mcp_llm_endpoints.py | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/tests/integration/mcp/test_mcp_llm_endpoints.py b/tests/integration/mcp/test_mcp_llm_endpoints.py index 490a8f831bc..394221708bf 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/tests/integration/mcp/test_mcp_llm_endpoints.py @@ -950,6 +950,45 @@ def test_nonstream_failed_follow_up_bills_round_one_on_the_failure_row( ) +NO_CACHE: Final = {"no-cache": True} + + +def _upstream_rounds(rig: Rig) -> tuple[int, int]: + answered: Final = tuple( + _has_tool_result(JSON_OBJECT.validate_json(request.body)) + for request in rig.wire.drain() + if request.method == "POST" + ) + return answered.count(False), answered.count(True) + + +def _billed_rows(call_id: str, call_type: str, billed: float) -> list[dict[str, JsonValue]]: + return eventually( + lambda: read_rows( + 'SELECT status, spend FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s AND call_type=%s', + (call_id, call_type), + ), + lambda found: sum(float(str(row["spend"])) for row in found) >= billed * (1 - 1e-6), + seconds=70, + ) + + +@pytest.mark.parametrize(("surface", "call_type"), (("responses", "aresponses"), ("chat", "acompletion"))) +def test_nonstream_follow_up_failing_on_every_retry_bills_each_attempts_first_round_on_one_failure_row( + gateway: Gateway, surface: Surface, call_type: str +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_fails=True) + with _rig(gateway, surface, billing=billing, prices={**PRICES, "num_retries": 2}) as rig: + response: Final = rig.send(_granted_key(rig), [AUTO], cache=NO_CACHE) + assert response.status_code == 500 and FOLLOW_UP_OUTAGE in response.text, response.text + first_rounds, follow_ups = _upstream_rounds(rig) + assert first_rounds == follow_ups > 1, (first_rounds, follow_ups) + paid: Final = _price(*(FIRST_ROUND,) * first_rounds) + rows: Final = _billed_rows(response.headers["x-litellm-call-id"], call_type, paid) + assert [row["status"] for row in rows] == ["failure"], rows + assert float(str(rows[0]["spend"])) == pytest.approx(paid), rows + + def _tool_rows(call_id: str, expected: int) -> list[dict[str, JsonValue]]: return eventually( lambda: read_rows( From 38b4d3144786d79bfeee7694921f57eb9f997916 Mon Sep 17 00:00:00 2001 From: zoveki <349028329+zoveki@users.noreply.github.com> Date: Wed, 28 Oct 2026 06:08:30 -0700 Subject: [PATCH 036/128] test(scim): audit #47544 (#48791) * test(scim): audit #47544 Add integration cells for the SCIM group delete that now runs through the team delete's bulk path: an admin group delete revokes its members' admin key, a delete sent with the IdP's SCIM token revokes member keys and audits them under that token, a delete that times out on a locked member row changes nothing so a retry works, and two racing team deletes write one tombstone and one team audit row * test(scim): read the racing team delete's audit rows after the proxy drains them --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../management/test_scim_admin_group.py | 3 - .../management/test_scim_group_delete.py | 106 +++++++++++++++++- .../test_team_delete_large_membership.py | 44 +++++++- 3 files changed, 146 insertions(+), 7 deletions(-) diff --git a/tests/integration/management/test_scim_admin_group.py b/tests/integration/management/test_scim_admin_group.py index 3aeed45d186..ffbec93d4a8 100644 --- a/tests/integration/management/test_scim_admin_group.py +++ b/tests/integration/management/test_scim_admin_group.py @@ -691,9 +691,6 @@ def test_scim_group_put_returns_the_updated_members(gateway: Gateway, tmp_path: @pytest.mark.timeout(300) def test_scim_admin_group_delete_revokes_the_members_admin_key(gateway: Gateway, tmp_path: Path) -> None: - pytest.skip( - "BUG: DELETE /scim/v2/Groups/{id} demotes the admin group's members in the DB but their existing keys keep proxy_admin access on /key/generate" - ) group_name: Final = f"scim-admin-delete-{uuid.uuid4().hex}" config: Final = _owned_scim_config(tmp_path / "scim-admin-group-delete.yaml", {"scim_admin_group": group_name}) with owned_proxy_process(gateway, tmp_path, {}, config=config) as owned, owned.gateway.scenario() as scenario: diff --git a/tests/integration/management/test_scim_group_delete.py b/tests/integration/management/test_scim_group_delete.py index 4a4a93f8bd2..53af34c12e0 100644 --- a/tests/integration/management/test_scim_group_delete.py +++ b/tests/integration/management/test_scim_group_delete.py @@ -50,7 +50,7 @@ from integration._support.scim import ( ) from pydantic import JsonValue, TypeAdapter -from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS +from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS, TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS pytestmark: Final = pytest.mark.shares_worker_state @@ -70,6 +70,17 @@ FROM "LiteLLM_AuditLog" WHERE object_id IN (%s, %s) AND action = 'deleted' ORDER BY table_name """ +AUDIT_ROWS_NAMING_SQL: Final = """ +SELECT table_name, object_id, changed_by_api_key +FROM "LiteLLM_AuditLog" +WHERE object_id = ANY(%s::text[]) AND action = 'deleted' +ORDER BY table_name, object_id +""" +OLDEST_TRANSACTION_BLOCKED_BY_SQL: Final = """ +SELECT coalesce(max(extract(epoch FROM clock_timestamp() - xact_start)), 0)::float AS age +FROM pg_stat_activity +WHERE %s::int = ANY(pg_blocking_pids(pid)) +""" _STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") _PIDS: Final = TypeAdapter(tuple[int, ...]) @@ -115,15 +126,28 @@ def _deleted_audit_rows(team: str, hashed_key: str) -> list[dict[str, JsonValue] return read_rows(DELETED_AUDIT_ROWS_SQL, (team, hashed_key)) +def _audit_rows_naming(object_ids: tuple[str, ...]) -> list[dict[str, JsonValue]]: + return read_rows(AUDIT_ROWS_NAMING_SQL, ("{" + ",".join(object_ids) + "}",)) + + +def _oldest_transaction_blocked_by(holder: int) -> float: + age: Final = read_rows(OLDEST_TRANSACTION_BLOCKED_BY_SQL, (str(holder),))[0]["age"] + assert isinstance(age, float), age + return age + + def _worker_startups(log: Path) -> tuple[tuple[int, ...], int]: text: Final = log.read_text() return _PIDS.validate_python(_STARTED_WORKER.findall(text)), text.count("Application startup complete.") @contextmanager -def _fresh_connection(candidate: Gateway) -> Iterator[Gateway]: +def _fresh_connection(candidate: Gateway, timeout: httpx.Timeout | float | None = None) -> Iterator[Gateway]: with httpx.Client( - base_url=candidate.client.base_url, timeout=candidate.client.timeout, trust_env=False, limits=GATEWAY_LIMITS + base_url=candidate.client.base_url, + timeout=candidate.client.timeout if timeout is None else timeout, + trust_env=False, + limits=GATEWAY_LIMITS, ) as client: yield Gateway(client, candidate.key, candidate.upstream_url) @@ -294,6 +318,82 @@ def test_a_group_delete_revokes_member_keys_and_a_re_created_group_gets_its_serv assert _models_status(gateway, member_key) == 401 +@pytest.mark.timeout(OWNED_PROXY_TIMEOUT) +def test_a_group_delete_sent_with_the_idps_scim_token_revokes_member_keys_and_audits_them_under_that_token( + gateway: Gateway, tmp_path: Path +) -> None: + with ( + owned_proxy_process( + gateway, + tmp_path, + {"DATABASE_URL": os.environ["DATABASE_URL"], "LITELLM_STORE_AUDIT_LOGS": "true"}, + remove_environment=("DATABASE_URL_READ_REPLICA",), + ) as owned, + owned.gateway.scenario() as scenario, + ): + candidate: Final = owned.gateway + scim_token: Final = scenario.key(allowed_routes=["/scim/*"], models=[]) + users: Final = seed_users(scenario, 3) + team: Final = created_team(scenario, create_group(candidate, users)) + member_keys: Final = (team_key(candidate, team, users[0]), team_key(candidate, team, users[1])) + service_key: Final = service_account_key(candidate, scenario, team) + + deleted: Final = delete_group(candidate, team, key=scim_token) + assert deleted.status_code == 204, deleted.text + + assert keys_of(team) == frozenset({_hashed(service_key)}) + for key in member_keys: + _assert_revoked_everywhere((candidate,), key) + audited: Final = (team, *(_hashed(key) for key in member_keys), _hashed(service_key)) + rows: Final = eventually(lambda: _audit_rows_naming(audited), lambda found: len(found) >= 3, seconds=30) + token: Final = _hashed(scim_token) + assert rows == [ + {"table_name": "LiteLLM_TeamTable", "object_id": team, "changed_by_api_key": token}, + *( + {"table_name": "LiteLLM_VerificationToken", "object_id": hashed, "changed_by_api_key": token} + for hashed in sorted(_hashed(key) for key in member_keys) + ), + ] + + +@pytest.mark.timeout(TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS + 120) +def test_a_group_delete_that_outlives_its_transaction_on_a_locked_member_row_changes_nothing_and_a_retry_lands( + gateway: Gateway, +) -> None: + patient_timeout: Final = TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS + 90 + with ( + gateway.scenario() as scenario, + _fresh_connection(gateway, timeout=patient_timeout) as patient, + ThreadPoolExecutor(max_workers=1) as pool, + ): + users: Final = seed_users(scenario, 3) + team: Final = created_team(scenario, create_group(gateway, users)) + key: Final = team_key(gateway, team, users[0]) + + with held_user_row(users[0]) as holder: + delete: Final = pool.submit(delete_group, patient, team) + eventually( + lambda: ( + delete.done() + or _oldest_transaction_blocked_by(holder) > TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS + 1 + ), + lambda outlived: outlived, + seconds=TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS + 30, + ) + failed: Final = delete.result(timeout=patient_timeout) + + assert failed.status_code >= 500, failed.text + assert_landed(gateway, team, users) + assert keys_of(team) == frozenset({_hashed(key)}) + assert read_rows(DELETED_TEAM_RECORDS_SQL, (team,)) == [] + assert read_rows(DELETED_KEY_RECORDS_SQL, (_hashed(key),)) == [] + assert _models_status(gateway, key) == 200 + + retried: Final = delete_group(gateway, team) + assert retried.status_code == 204, retried.text + assert_gone(gateway, team, users, key) + + def test_deleting_an_unknown_group_is_404(gateway: Gateway) -> None: missing: Final = f"missing-{uuid.uuid4().hex}" response: Final = delete_group(gateway, missing) diff --git a/tests/integration/management/test_team_delete_large_membership.py b/tests/integration/management/test_team_delete_large_membership.py index e752c835e4f..1a8140373b0 100644 --- a/tests/integration/management/test_team_delete_large_membership.py +++ b/tests/integration/management/test_team_delete_large_membership.py @@ -33,7 +33,7 @@ from tests.integration._support.client import ( string_value, ) from tests.integration._support.database import read_rows, write_rows -from tests.integration._support.process import owned_proxy_process +from tests.integration._support.process import graceful_stop_seconds, owned_proxy_process LARGE_ROSTER: Final = 250 POOL_LIMIT: Final = 5 @@ -131,6 +131,13 @@ def _roster_user_ids(roster: JsonValue) -> list[str]: return sorted(string_value(object_value(member)["user_id"]) for member in roster) +def _deleted_audit_rows(object_id: str, table_name: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT before_value FROM "LiteLLM_AuditLog" WHERE object_id = %s AND table_name = %s AND action = %s', + (object_id, table_name, "deleted"), + ) + + def _waiters_on_lock_held_by(backend_pid: int) -> int: rows: Final = read_rows(WAITERS_ON_HELD_LOCK_SQL, ("%pg_advisory_xact_lock%", str(backend_pid))) waiting: Final = rows[0]["waiting"] @@ -551,6 +558,41 @@ def test_delete_writes_one_audit_row_for_the_team_and_one_per_key(gateway: Gatew ] +@pytest.mark.timeout(int(2 * graceful_stop_seconds() + 120)) +def test_two_deletes_racing_on_one_team_write_one_tombstone_and_one_audit_row(gateway: Gateway, tmp_path: Path) -> None: + with ( + owned_proxy_process( + gateway, + tmp_path, + {"DATABASE_URL": os.environ["DATABASE_URL"], "LITELLM_STORE_AUDIT_LOGS": "true"}, + remove_environment=("DATABASE_URL_READ_REPLICA",), + ) as owned, + owned.gateway.scenario() as scenario, + ThreadPoolExecutor(max_workers=2) as pool, + ): + user: Final = scenario.user() + team: Final = _create_team(scenario) + _bulk_member_add(owned.gateway, team, [user]) + + with psycopg.connect(os.environ["DATABASE_URL"]) as holder: + holder.execute(TAKE_TEAM_LOCK_SQL, (team,)) + deletes: Final = tuple( + _post_in_background(pool, owned.gateway, "/team/delete", {"team_ids": [team]}) for _ in range(2) + ) + eventually( + lambda: _waiters_on_lock_held_by(holder.info.backend_pid), lambda waiting: waiting == 2, seconds=30 + ) + + assert [delete.result(timeout=60).status_code for delete in deletes] == [200, 200] + tombstones: Final = _tombstones(team) + assert [row["team_id"] for row in tombstones] == [team] + assert _roster_user_ids(tombstones[0]["members_with_roles"]) == sorted([PROXY_ADMIN, user]) + + team_audit_rows: Final = _deleted_audit_rows(team, "LiteLLM_TeamTable") + assert len(team_audit_rows) == 1, team_audit_rows + assert object_value(team_audit_rows[0]["before_value"])["team_id"] == team + + def test_second_delete_of_the_same_team_is_404_with_one_tombstone(gateway: Gateway) -> None: with gateway.scenario() as scenario: user: Final = scenario.user() From 6c0585b8f192efb2c07735468537a99824c7a94b Mon Sep 17 00:00:00 2001 From: yujonglee Date: Wed, 28 Oct 2026 07:32:40 -0700 Subject: [PATCH 037/128] build(rust): upgrade reqwest to 0.13 and move MCP out of litellm-http (#48039) * build(rust): upgrade reqwest to 0.13 and move MCP out of litellm-http The workspace moves to reqwest 0.13, so litellm-http no longer carries a second reqwest just for rmcp. The mcp feature, its pool and its proxy conversion go away, and the gateway hands rmcp a pooled NoRedirect client that already carries the TLS, proxy and timeout settings The Client test helpers collapse into for_test, a no-redirect client without TLS or proxy settings, and wrap_for_test for the few tests that need a hand-built reqwest client Co-Authored-By: Claude Opus 5.5 * build(deps): refresh compatible dependencies flagged by OSV * refactor(rust): remove redundant client construction comment * test(rust): cover configured MCP proxy and CA bundle --------- Co-authored-by: Claude Opus 5.5 Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 108 ++---- litellm-rust/Cargo.toml | 2 +- litellm-rust/clippy.toml | 3 + litellm-rust/crates/auth-aws/src/aws.rs | 2 +- .../cache-azure-blob/tests/transport.rs | 2 +- litellm-rust/crates/cache-gcs/tests/cache.rs | 2 +- .../crates/cache-gcs/tests/support/mod.rs | 2 +- .../cache-qdrant-semantic/tests/embedder.rs | 6 +- .../crates/cache-response/tests/response.rs | 2 +- .../crates/cache-s3/tests/support/mod.rs | 2 +- litellm-rust/crates/gateway-mcp/Cargo.toml | 2 +- litellm-rust/crates/gateway/Cargo.toml | 2 +- litellm-rust/crates/gateway/src/lib.rs | 4 +- litellm-rust/crates/gateway/tests/mcp.rs | 86 ++++- litellm-rust/crates/http/Cargo.toml | 3 - litellm-rust/crates/http/src/client.rs | 17 +- litellm-rust/crates/http/src/config.rs | 2 +- litellm-rust/crates/http/src/lib.rs | 2 - litellm-rust/crates/http/src/mcp.rs | 76 ----- litellm-rust/crates/http/src/media.rs | 6 +- litellm-rust/crates/http/src/pool.rs | 66 +--- litellm-rust/crates/http/src/proxy.rs | 13 - .../crates/inference-chat/src/handler.rs | 4 +- .../crates/inference-chat/tests/caching.rs | 2 +- .../crates/inference-messages/src/handler.rs | 2 +- .../crates/inference-ocr/src/handler.rs | 6 +- .../crates/llms/src/base_llm/ocr/document.rs | 8 +- .../crates/llms/src/base_llm/ocr/handler.rs | 24 +- .../llms/src/reducto/ocr/transformation.rs | 5 +- .../python-bridge/src/secrets/callback.rs | 2 +- .../python-bridge/src/secrets/resolved.rs | 6 +- .../crates/secrets-azure/tests/key_vault.rs | 8 +- .../crates/secrets-azure/tests/live.rs | 2 +- .../tests/secret_manager/configuration.rs | 4 +- .../tests/secret_manager/support.rs | 2 +- .../tests/secret_manager/writes.rs | 6 +- .../secrets-google/tests/secret_manager.rs | 8 +- litellm-rust/crates/secrets/tests/aws.rs | 6 +- litellm-rust/crates/secrets/tests/azure.rs | 6 +- .../secrets/tests/common_read_contract.rs | 8 +- litellm-rust/crates/secrets/tests/cyberark.rs | 2 +- litellm-rust/crates/secrets/tests/google.rs | 4 +- .../crates/secrets/tests/hashicorp.rs | 4 +- litellm-rust/crates/secrets/tests/oidc.rs | 28 +- .../crates/secrets/tests/resolution.rs | 14 +- litellm-rust/crates/secrets/tests/source.rs | 4 +- .../crates/spend-clickhouse/tests/storage.rs | 2 +- .../spend-clickhouse/tests/transport.rs | 8 +- .../storage-clickhouse/tests/migrations.rs | 12 +- .../storage-clickhouse/tests/transport.rs | 10 +- pyproject.toml | 2 +- uv.lock | 317 +++++++++++------- vscode-extension/package-lock.json | 6 +- 53 files changed, 433 insertions(+), 499 deletions(-) delete mode 100644 litellm-rust/crates/http/src/mcp.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 63137949f38..d37b65be7a3 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2342,7 +2342,7 @@ dependencies = [ "hmac 0.13.0", "http 1.4.2", "jiff", - "reqwest 0.13.5", + "reqwest", "rustc_version", "rustls 0.23.42", "rustls-pki-types", @@ -2400,7 +2400,7 @@ dependencies = [ "pin-project", "prost", "prost-types", - "reqwest 0.13.5", + "reqwest", "rustc_version", "serde", "serde_json", @@ -2896,7 +2896,6 @@ dependencies = [ "tokio", "tokio-rustls 0.26.4", "tower-service", - "webpki-roots", ] [[package]] @@ -3389,7 +3388,7 @@ dependencies = [ "litellm-auth-types", "litellm-http", "moka", - "reqwest 0.12.28", + "reqwest", "rstest", "serde", "serde_json", @@ -3473,7 +3472,7 @@ dependencies = [ "litellm-cache-response", "litellm-cache-testing", "litellm-http", - "reqwest 0.12.28", + "reqwest", "rstest", "serde_json", "tokio", @@ -3508,7 +3507,7 @@ dependencies = [ "litellm-cache-testing", "litellm-http", "percent-encoding", - "reqwest 0.12.28", + "reqwest", "rstest", "serde_json", "tokio", @@ -3535,7 +3534,7 @@ dependencies = [ "litellm-cache-testing", "litellm-http", "qdrant-client", - "reqwest 0.12.28", + "reqwest", "rstest", "serde", "serde_json", @@ -3611,7 +3610,7 @@ dependencies = [ "litellm-cache", "litellm-cache-testing", "litellm-http", - "reqwest 0.12.28", + "reqwest", "rstest", "serde_json", "tokio", @@ -3978,9 +3977,7 @@ dependencies = [ "hyper-util", "litellm-core-utils", "rcgen", - "reqwest 0.12.28", - "reqwest 0.13.5", - "rmcp", + "reqwest", "rstest", "rustls 0.23.42", "rustls-native-certs", @@ -4013,7 +4010,7 @@ dependencies = [ "litellm-llms", "litellm-secrets", "litellm-tracing", - "reqwest 0.12.28", + "reqwest", "rstest", "serde", "serde_json", @@ -4072,7 +4069,7 @@ dependencies = [ "litellm-router-types", "litellm-secrets", "litellm-tracing", - "reqwest 0.12.28", + "reqwest", "rstest", "serde", "serde_json", @@ -4133,7 +4130,7 @@ dependencies = [ "litellm-llms-types", "litellm-secrets", "litellm-tracing", - "reqwest 0.12.28", + "reqwest", "rstest", "serde_json", "tokio", @@ -4196,7 +4193,7 @@ dependencies = [ "litellm-python-compat", "litellm-router-types", "litellm-secrets", - "reqwest 0.12.28", + "reqwest", "rstest", "serde", "serde_json", @@ -4278,7 +4275,7 @@ dependencies = [ "pyo3", "pyo3-async-runtimes", "redis", - "reqwest 0.12.28", + "reqwest", "rstest", "serde", "serde_json", @@ -4371,7 +4368,7 @@ dependencies = [ "litellm-secrets-hashicorp", "litellm-secrets-types", "moka", - "reqwest 0.12.28", + "reqwest", "rstest", "serde", "serde_json", @@ -4413,7 +4410,7 @@ dependencies = [ "litellm-http", "litellm-secrets-types", "percent-encoding", - "reqwest 0.12.28", + "reqwest", "rstest", "serde", "serde_json", @@ -4436,7 +4433,7 @@ dependencies = [ "moka", "percent-encoding", "rcgen", - "reqwest 0.12.28", + "reqwest", "rstest", "serde", "serde_json", @@ -4463,7 +4460,7 @@ dependencies = [ "litellm-secrets-types", "moka", "percent-encoding", - "reqwest 0.12.28", + "reqwest", "rstest", "serde", "serde_json", @@ -4551,7 +4548,7 @@ version = "0.1.0" dependencies = [ "flate2", "futures-util", - "reqwest 0.12.28", + "reqwest", "rstest", "semver", "serde", @@ -5936,49 +5933,6 @@ version = "1.9.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ba39f3699c378cd8970968dcbff9c43159ea4cfbd88d43c00b22f2ef10a435d2" -[[package]] -name = "reqwest" -version = "0.12.28" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" -dependencies = [ - "base64 0.22.1", - "bytes", - "futures-core", - "futures-util", - "h2 0.4.15", - "http 1.4.2", - "http-body 1.1.0", - "http-body-util", - "hyper 1.10.1", - "hyper-rustls 0.27.9", - "hyper-util", - "js-sys", - "log", - "mime_guess", - "percent-encoding", - "pin-project-lite", - "quinn", - "rustls 0.23.42", - "rustls-pki-types", - "serde", - "serde_json", - "serde_urlencoded", - "sync_wrapper", - "tokio", - "tokio-rustls 0.26.4", - "tokio-util", - "tower", - "tower-http", - "tower-service", - "url", - "wasm-bindgen", - "wasm-bindgen-futures", - "wasm-streams 0.4.2", - "web-sys", - "webpki-roots", -] - [[package]] name = "reqwest" version = "0.13.5" @@ -6000,6 +5954,7 @@ dependencies = [ "js-sys", "log", "mime", + "mime_guess", "percent-encoding", "pin-project-lite", "quinn", @@ -6019,7 +5974,7 @@ dependencies = [ "url", "wasm-bindgen", "wasm-bindgen-futures", - "wasm-streams 0.5.0", + "wasm-streams", "web-sys", ] @@ -6066,7 +6021,7 @@ dependencies = [ "pin-project-lite", "process-wrap", "rand 0.10.2", - "reqwest 0.13.5", + "reqwest", "schemars 1.2.2", "serde", "serde_json", @@ -6189,7 +6144,7 @@ dependencies = [ "async-trait", "bytes", "http 1.4.2", - "reqwest 0.13.5", + "reqwest", "rustify_derive", "serde", "serde_json", @@ -7221,7 +7176,7 @@ dependencies = [ "memchr", "parse-display", "pin-project-lite", - "reqwest 0.13.5", + "reqwest", "serde", "serde_json", "serde_with", @@ -7860,7 +7815,7 @@ dependencies = [ "futures", "pin-project", "rand 0.10.2", - "reqwest 0.13.5", + "reqwest", "serde", "serde_json", "time", @@ -8074,7 +8029,7 @@ dependencies = [ "async-trait", "derive_builder", "http 1.4.2", - "reqwest 0.13.5", + "reqwest", "rustify", "rustify_derive", "serde", @@ -8221,19 +8176,6 @@ dependencies = [ "unicode-ident", ] -[[package]] -name = "wasm-streams" -version = "0.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65" -dependencies = [ - "futures-util", - "js-sys", - "wasm-bindgen", - "wasm-bindgen-futures", - "web-sys", -] - [[package]] name = "wasm-streams" version = "0.5.0" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index dd8d8704370..95f2da8c073 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -90,7 +90,7 @@ pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] } rand = "0.8" macro_rules_attribute = "0.2.3" schemars = "1" -reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls", "http2", "stream"] } +reqwest = { version = "0.13", default-features = false, features = ["json", "multipart", "rustls", "http2", "stream", "query"] } redis = "1.8.0" redis-test = "1.0.4" testcontainers-modules = "0.15.0" diff --git a/litellm-rust/clippy.toml b/litellm-rust/clippy.toml index 18663cfdffb..075e24a7849 100644 --- a/litellm-rust/clippy.toml +++ b/litellm-rust/clippy.toml @@ -10,8 +10,11 @@ disallowed-methods = [ { path = "reqwest::Client::new", reason = "take litellm_http::Client from HttpClientPool" }, { path = "reqwest::Client::builder", reason = "HttpClientConfig owns client construction" }, { path = "reqwest::ClientBuilder::danger_accept_invalid_certs", reason = "set HttpClientConfig::verify instead" }, + { path = "reqwest::ClientBuilder::tls_danger_accept_invalid_certs", reason = "set HttpClientConfig::verify instead" }, + { path = "reqwest::ClientBuilder::tls_danger_accept_invalid_hostnames", reason = "set HttpClientConfig::verify instead" }, { path = "reqwest::ClientBuilder::identity", reason = "set HttpClientConfig::client_certificate instead" }, { path = "reqwest::ClientBuilder::use_preconfigured_tls", reason = "HttpClientConfig owns the TLS configuration" }, + { path = "reqwest::ClientBuilder::tls_backend_preconfigured", reason = "HttpClientConfig owns the TLS configuration" }, { path = "sqlx::query", reason = "use sqlx::query! or query_file! so the SQL is checked against the migrated schema" }, { path = "sqlx::query_as", reason = "use sqlx::query_as! or query_file_as! so the SQL is checked against the migrated schema" }, { path = "sqlx::query_scalar", reason = "use sqlx::query_scalar! so the SQL is checked against the migrated schema" }, diff --git a/litellm-rust/crates/auth-aws/src/aws.rs b/litellm-rust/crates/auth-aws/src/aws.rs index 8fdbba21711..eaa033d631f 100644 --- a/litellm-rust/crates/auth-aws/src/aws.rs +++ b/litellm-rust/crates/auth-aws/src/aws.rs @@ -1080,7 +1080,7 @@ mod tests { &no_env, ) .await?; - let client = litellm_http::Client::plain_for_test(); + let client = litellm_http::Client::for_test(); let mut failures = Vec::new(); for region in ["us-west-2", "us-east-1"] { diff --git a/litellm-rust/crates/cache-azure-blob/tests/transport.rs b/litellm-rust/crates/cache-azure-blob/tests/transport.rs index c8b14cd8543..e600b15af0f 100644 --- a/litellm-rust/crates/cache-azure-blob/tests/transport.rs +++ b/litellm-rust/crates/cache-azure-blob/tests/transport.rs @@ -31,7 +31,7 @@ async fn connect(server: &MockServer) -> AzureBlobCache> { None, ClientOptions { transport: Some(Transport::new(Arc::new(ReqwestTransport( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), )))), ..ClientOptions::default() }, diff --git a/litellm-rust/crates/cache-gcs/tests/cache.rs b/litellm-rust/crates/cache-gcs/tests/cache.rs index 096b691ea16..3e92c3fca19 100644 --- a/litellm-rust/crates/cache-gcs/tests/cache.rs +++ b/litellm-rust/crates/cache-gcs/tests/cache.rs @@ -97,7 +97,7 @@ async fn cache_exposes_its_configuration(#[future(awt)] server: MockServer) { path_service_account: Some("/secrets/sa.json".into()), ..support::config(&server, Some("folder")) }, - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), litellm_cache::JsonCodec::::new(), ); assert_eq!(cache.bucket_name(), "bucket"); diff --git a/litellm-rust/crates/cache-gcs/tests/support/mod.rs b/litellm-rust/crates/cache-gcs/tests/support/mod.rs index beb9aa39d9c..99fb0270a74 100644 --- a/litellm-rust/crates/cache-gcs/tests/support/mod.rs +++ b/litellm-rust/crates/cache-gcs/tests/support/mod.rs @@ -29,7 +29,7 @@ pub fn cache_with_token( ) -> JsonGcsCache { GcsCache::with_token_source( config(server, gcs_path), - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), JsonCodec::new(), token, ) diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs index 24e6e5eba3e..71fc986616e 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs @@ -108,7 +108,7 @@ fn config(base: String, timeout: Option) -> OpenAiEmbedderConfig { async fn posts_embeddings_request_and_parses_vector() { let server = TestHttpServer::response("200 OK", r#"{"data":[{"embedding":[0.1,0.2]}]}"#).await; let embedder = OpenAiEmbedder::new( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), config( format!("{}/", server.base_url()), Some(Duration::from_secs(1)), @@ -161,7 +161,7 @@ async fn status_timeout_and_body_errors_are_unavailable( let server = TestHttpServer::response_after(status, body, Duration::from_millis(delay_ms)).await; let embedder = OpenAiEmbedder::new( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), config(server.base_url(), timeout), ); assert_eq!(embedder.async_embed("hello", None).await, expected); @@ -170,7 +170,7 @@ async fn status_timeout_and_body_errors_are_unavailable( #[rstest] fn sync_embedding_is_unsupported() { let embedder = OpenAiEmbedder::new( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), config("http://127.0.0.1:9".to_owned(), None), ); assert_eq!( diff --git a/litellm-rust/crates/cache-response/tests/response.rs b/litellm-rust/crates/cache-response/tests/response.rs index 655fcb8a46a..86ddf4a72cd 100644 --- a/litellm-rust/crates/cache-response/tests/response.rs +++ b/litellm-rust/crates/cache-response/tests/response.rs @@ -768,7 +768,7 @@ async fn gcs() -> (wiremock::MockServer, Gcs) { path_service_account: None, endpoint: server.uri(), }, - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), litellm_cache_response::ResponseCacheCodec, Arc::new(litellm_cache_gcs::StaticTokenSource("token".into())), ))); diff --git a/litellm-rust/crates/cache-s3/tests/support/mod.rs b/litellm-rust/crates/cache-s3/tests/support/mod.rs index b628c1df431..7378dde21bf 100644 --- a/litellm-rust/crates/cache-s3/tests/support/mod.rs +++ b/litellm-rust/crates/cache-s3/tests/support/mod.rs @@ -34,7 +34,7 @@ pub fn config(endpoint: &str) -> S3CacheConfig { pub fn cache_with(config: S3CacheConfig, runtime: Handle) -> JsonS3Cache { S3Cache::new( config, - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), JsonCodec::new(), runtime, ) diff --git a/litellm-rust/crates/gateway-mcp/Cargo.toml b/litellm-rust/crates/gateway-mcp/Cargo.toml index 07706dcb6a4..1db5257448f 100644 --- a/litellm-rust/crates/gateway-mcp/Cargo.toml +++ b/litellm-rust/crates/gateway-mcp/Cargo.toml @@ -19,7 +19,7 @@ url.workspace = true moka.workspace = true sha2.workspace = true tower = { version = "0.5", features = ["util"] } -rmcp = { version = "=3.4.1", default-features = false, features = ["server", "client", "elicitation", "transport-streamable-http-server", "transport-streamable-http-client", "transport-child-process"] } +rmcp = { version = "=3.4.1", default-features = false, features = ["server", "client", "elicitation", "transport-streamable-http-server", "transport-streamable-http-client-reqwest", "transport-child-process"] } serde.workspace = true serde_json.workspace = true thiserror.workspace = true diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index b2a2f09bfc7..4982c70b18c 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -21,7 +21,7 @@ litellm-auth-types.workspace = true litellm-config.workspace = true litellm-circuit-breaker.workspace = true litellm-redis-client.workspace = true -litellm-http = { workspace = true, features = ["mcp"] } +litellm-http.workspace = true litellm-llms.workspace = true litellm-secrets.workspace = true litellm-tracing.workspace = true diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index 2ed2fc80a51..d7567b83cc1 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -114,11 +114,11 @@ pub async fn build_mcp( Auth::from_config(config, secrets.clone()) .validate() .await?; - let client = pool.mcp_client(http)?; + let client = pool.client(http, ClientVariant::NoRedirect)?; Ok(Some( litellm_gateway_mcp::ConfiguredGateway::connect( &config.mcp_servers, - client, + (*client).clone(), secrets.as_ref(), shutdown, ) diff --git a/litellm-rust/crates/gateway/tests/mcp.rs b/litellm-rust/crates/gateway/tests/mcp.rs index 593fbfeb129..ee8a17ae251 100644 --- a/litellm-rust/crates/gateway/tests/mcp.rs +++ b/litellm-rust/crates/gateway/tests/mcp.rs @@ -78,7 +78,9 @@ async fn upstream() -> Fixture { let transport = StreamableHttpService::new( || Ok(Upstream), Arc::new(LocalSessionManager::default()), - StreamableHttpServerConfig::default().with_cancellation_token(shutdown.clone()), + StreamableHttpServerConfig::default() + .with_allowed_hosts(["127.0.0.1", "upstream.invalid"]) + .with_cancellation_token(shutdown.clone()), ); let app = Router::new().route_service("/mcp", transport); let stop = shutdown.clone(); @@ -222,6 +224,88 @@ mcp_servers: .unwrap(); } +#[rstest] +#[tokio::test] +async fn mcp_tools_are_discovered_through_the_configured_proxy(#[future] upstream: Fixture) { + let upstream = upstream.await; + let config = Config::from_yaml( + "general_settings: {master_key: gateway-key}\nmcp_servers: {docs: {url: 'http://upstream.invalid/mcp', timeout: 2}}", + ) + .unwrap(); + let inference = litellm_gateway::build_inference(&config).unwrap(); + let proxy = upstream.url.trim_end_matches("/mcp"); + let http = litellm_http::HttpClientConfig { + proxies: litellm_http::EnvironmentProxies::from_environment(&|name: &str| { + (name == "HTTP_PROXY").then(|| proxy.to_owned()) + }), + ..inference.http.clone() + }; + let shutdown = CancellationToken::new(); + let mcp = litellm_gateway::build_mcp( + &config, + inference.secrets.clone(), + shutdown.clone(), + &inference.resources.pool, + &http, + ) + .await + .unwrap() + .unwrap(); + let app = litellm_gateway_mcp::router(mcp.operations.clone(), HttpConfig::default()); + let response = app + .oneshot( + Request::get("/mcp-rest/tools/list?server_id=docs") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), 200); + let body: Value = + serde_json::from_slice(&to_bytes(response.into_body(), 65536).await.unwrap()).unwrap(); + assert_eq!(body["tools"][0]["name"], "echo"); + shutdown.cancel(); + tokio::time::timeout(Duration::from_secs(5), mcp.close()) + .await + .unwrap(); + upstream.shutdown.cancel(); + tokio::time::timeout(Duration::from_secs(5), upstream.task) + .await + .unwrap() + .unwrap(); +} + +#[rstest] +#[tokio::test] +async fn mcp_uses_the_configured_ca_bundle() { + let directory = tempfile::tempdir().unwrap(); + let missing = directory.path().join("missing-ca.pem"); + let config = Config::from_yaml( + "general_settings: {master_key: gateway-key}\nmcp_servers: {docs: {url: 'http://127.0.0.1:1/mcp'}}", + ) + .unwrap(); + let inference = litellm_gateway::build_inference(&config).unwrap(); + let http = litellm_http::HttpClientConfig { + verify: litellm_http::Verify::CaBundle(missing.clone()), + ..inference.http.clone() + }; + let error = litellm_gateway::build_mcp( + &config, + inference.secrets.clone(), + CancellationToken::new(), + &inference.resources.pool, + &http, + ) + .await + .err() + .unwrap(); + assert!(matches!( + error, + litellm_gateway::Error::Http(litellm_http::Error::Read { path, .. }) + if path == missing + )); +} + #[rstest] #[case::unsupported_auth("auth_type: oauth2", "auth mode")] #[case::unsupported_policy("allowed_params: {echo: [safe]}", "allowed_params")] diff --git a/litellm-rust/crates/http/Cargo.toml b/litellm-rust/crates/http/Cargo.toml index 7d64467f2f7..3b46825ad9d 100644 --- a/litellm-rust/crates/http/Cargo.toml +++ b/litellm-rust/crates/http/Cargo.toml @@ -7,11 +7,8 @@ repository.workspace = true [features] test-support = [] -mcp = ["dep:rmcp", "dep:reqwest-mcp"] [dependencies] -rmcp = { version = "=3.4.1", default-features = false, features = ["client", "transport-streamable-http-client-reqwest"], optional = true } -reqwest-mcp = { package = "reqwest", version = "0.13.2", default-features = false, features = ["rustls-no-provider", "json", "stream", "http2"], optional = true } http.workspace = true litellm-core-utils.workspace = true hyper-util.workspace = true diff --git a/litellm-rust/crates/http/src/client.rs b/litellm-rust/crates/http/src/client.rs index 1f7017d083b..bb7883f665b 100644 --- a/litellm-rust/crates/http/src/client.rs +++ b/litellm-rust/crates/http/src/client.rs @@ -9,17 +9,7 @@ impl Client { } #[cfg(any(test, feature = "test-support"))] - pub fn for_test(client: reqwest::Client) -> Self { - Self(client) - } - - #[cfg(any(test, feature = "test-support"))] - pub fn plain_for_test() -> Self { - Self(reqwest::Client::new()) - } - - #[cfg(any(test, feature = "test-support"))] - pub fn no_redirect_for_test() -> Self { + pub fn for_test() -> Self { Self( reqwest::Client::builder() .redirect(reqwest::redirect::Policy::none()) @@ -27,6 +17,11 @@ impl Client { .expect("a client without TLS or proxy settings builds"), ) } + + #[cfg(any(test, feature = "test-support"))] + pub fn wrap_for_test(client: reqwest::Client) -> Self { + Self(client) + } } impl Deref for Client { diff --git a/litellm-rust/crates/http/src/config.rs b/litellm-rust/crates/http/src/config.rs index 2f36784bc70..6a2d80e3f85 100644 --- a/litellm-rust/crates/http/src/config.rs +++ b/litellm-rust/crates/http/src/config.rs @@ -98,7 +98,7 @@ impl TryFrom<&HttpClientConfig> for reqwest::ClientBuilder { fn try_from(config: &HttpClientConfig) -> Result { let base = reqwest::Client::builder() - .use_preconfigured_tls(rustls::ClientConfig::try_from(config)?) + .tls_backend_preconfigured(rustls::ClientConfig::try_from(config)?) .connect_timeout(config.connect_timeout) .pool_idle_timeout(config.pool_idle_timeout); let with_keepalive = match config.tcp_keepalive { diff --git a/litellm-rust/crates/http/src/lib.rs b/litellm-rust/crates/http/src/lib.rs index d53da607d39..ef72420af9b 100644 --- a/litellm-rust/crates/http/src/lib.rs +++ b/litellm-rust/crates/http/src/lib.rs @@ -7,8 +7,6 @@ mod client; mod config; mod error; -#[cfg(feature = "mcp")] -mod mcp; pub mod media; pub mod outbound; mod pool; diff --git a/litellm-rust/crates/http/src/mcp.rs b/litellm-rust/crates/http/src/mcp.rs deleted file mode 100644 index 2b101fb9c98..00000000000 --- a/litellm-rust/crates/http/src/mcp.rs +++ /dev/null @@ -1,76 +0,0 @@ -use std::{ - collections::HashMap, - net::{IpAddr, Ipv4Addr}, - sync::{Mutex, PoisonError}, - time::{Duration, Instant}, -}; - -use crate::{Error, HttpClientConfig}; - -struct Entry { - client: reqwest_mcp::Client, - built_at: Instant, -} - -#[derive(Default)] -pub(super) struct Pool(Mutex>); - -impl Pool { - pub fn client( - &self, - config: &HttpClientConfig, - ttl: Duration, - ) -> Result { - let mut clients = self.0.lock().unwrap_or_else(PoisonError::into_inner); - if let Some(entry) = clients.get(config) - && entry.built_at.elapsed() < ttl - { - return Ok(entry.client.clone()); - } - let client = build(config)?; - clients.insert( - config.clone(), - Entry { - client: client.clone(), - built_at: Instant::now(), - }, - ); - Ok(client) - } -} - -fn build(config: &HttpClientConfig) -> Result { - let base = reqwest_mcp::Client::builder() - .tls_backend_preconfigured(rustls::ClientConfig::try_from(config)?) - .connect_timeout(config.connect_timeout) - .pool_idle_timeout(config.pool_idle_timeout) - .redirect(reqwest_mcp::redirect::Policy::none()); - let keepalive = match config.tcp_keepalive { - Some(value) => base - .tcp_keepalive(value.idle) - .tcp_keepalive_interval(value.interval) - .tcp_keepalive_retries(value.retries), - None => base, - }; - let address = if config.force_ipv4 { - keepalive.local_address(IpAddr::V4(Ipv4Addr::UNSPECIFIED)) - } else { - keepalive - }; - let protocol = if config.http2 { - address - } else { - address.http1_only() - }; - let agent = match &config.user_agent { - Some(value) => protocol.user_agent(value), - None => protocol, - }; - config - .proxies - .mcp_proxies() - .into_iter() - .fold(agent.no_proxy(), reqwest_mcp::ClientBuilder::proxy) - .build() - .map_err(|error| Error::Client(error.without_url().to_string())) -} diff --git a/litellm-rust/crates/http/src/media.rs b/litellm-rust/crates/http/src/media.rs index 22a881831cf..960a1d4044a 100644 --- a/litellm-rust/crates/http/src/media.rs +++ b/litellm-rust/crates/http/src/media.rs @@ -522,7 +522,7 @@ mod tests { b"HTTP/1.1 200 OK\r\nContent-Type: application/pdf; charset=binary\r\nContent-Length: 3\r\nConnection: close\r\n\r\nabc", ) .await; - let client = Client::no_redirect_for_test(); + let client = Client::for_test(); let media = MediaFetcher::for_test(client) .fetch(url, policy(3, 0)) .await @@ -538,7 +538,7 @@ mod tests { b"HTTP/1.1 200 OK\r\nContent-Type: application/pdf\r\nContent-Length: 3\r\nConnection: close\r\n\r\nabc", ) .await; - let client = Client::no_redirect_for_test(); + let client = Client::for_test(); let error = MediaFetcher::for_test(client) .fetch(url, policy(2, 0)) .await @@ -553,7 +553,7 @@ mod tests { b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n2\r\nab\r\n2\r\ncd\r\n0\r\n\r\n", ) .await; - let client = Client::no_redirect_for_test(); + let client = Client::for_test(); let error = MediaFetcher::for_test(client) .fetch(url, policy(3, 0)) .await diff --git a/litellm-rust/crates/http/src/pool.rs b/litellm-rust/crates/http/src/pool.rs index 78932f9e29f..dfcc9270fe9 100644 --- a/litellm-rust/crates/http/src/pool.rs +++ b/litellm-rust/crates/http/src/pool.rs @@ -29,8 +29,6 @@ pub struct HttpClientPool { media_resolver: Arc, ttl: Duration, clients: Mutex, - #[cfg(feature = "mcp")] - mcp: crate::mcp::Pool, } impl HttpClientPool { @@ -43,19 +41,9 @@ impl HttpClientPool { media_resolver, ttl, clients: Mutex::default(), - #[cfg(feature = "mcp")] - mcp: crate::mcp::Pool::default(), } } - #[cfg(feature = "mcp")] - pub fn mcp_client( - &self, - config: &HttpClientConfig, - ) -> Result { - self.mcp.client(config, self.ttl) - } - pub fn client( &self, config: &HttpClientConfig, @@ -108,7 +96,7 @@ impl HttpClientPool { } ClientVariant::Media => builder .redirect(reqwest::redirect::Policy::none()) - .dns_resolver2(Arc::clone(&self.media_resolver)), + .dns_resolver(Arc::clone(&self.media_resolver)), } } } @@ -390,56 +378,4 @@ mod tests { .is_err() ); } - #[cfg(feature = "mcp")] - #[rstest::rstest] - #[tokio::test] - async fn mcp_clients_reuse_connections_and_apply_proxy_headers_without_redirects() { - let (address, connections, requests) = serve("HTTP/1.1 302 Found").await; - let pool = pool(); - let config = HttpClientConfig { - proxies: proxied_through(&format!("http://{address}")), - ..config("mcp-pool-test") - }; - for _ in 0..2 { - let response = pool - .mcp - .client(&config, pool.ttl) - .unwrap() - .get("http://upstream.invalid/mcp") - .timeout(Duration::from_secs(2)) - .send() - .await - .unwrap(); - assert_eq!(response.status(), 302); - response.bytes().await.unwrap(); - } - assert_eq!(connections.load(Ordering::SeqCst), 1); - let requests = requests.lock().unwrap(); - assert_eq!(requests.len(), 2); - assert!( - requests - .iter() - .all(|request| request.starts_with("GET http://upstream.invalid/mcp HTTP/1.1")) - ); - assert!( - requests - .iter() - .all(|request| request.contains("user-agent: mcp-pool-test")) - ); - } - - #[cfg(feature = "mcp")] - #[rstest::rstest] - #[tokio::test] - async fn mcp_client_uses_host_tls_configuration() { - let directory = tempfile::tempdir().unwrap(); - let config = HttpClientConfig { - verify: Verify::CaBundle(directory.path().join("missing.pem")), - ..config("mcp-test") - }; - assert!(matches!( - pool().mcp_client(&config), - Err(Error::Read { .. }) - )); - } } diff --git a/litellm-rust/crates/http/src/proxy.rs b/litellm-rust/crates/http/src/proxy.rs index e7e54ed6ae2..eb960d8200d 100644 --- a/litellm-rust/crates/http/src/proxy.rs +++ b/litellm-rust/crates/http/src/proxy.rs @@ -63,19 +63,6 @@ impl EnvironmentProxies { .map(|proxy| proxy.no_proxy(no_proxy.clone())) .collect() } - #[cfg(feature = "mcp")] - pub(crate) fn mcp_proxies(&self) -> Vec { - let no_proxy = reqwest_mcp::NoProxy::from_string(&self.no); - [ - reqwest_mcp::Proxy::http(self.http.as_str()), - reqwest_mcp::Proxy::https(self.https.as_str()), - reqwest_mcp::Proxy::all(self.all.as_str()), - ] - .into_iter() - .filter_map(Result::ok) - .map(|proxy| proxy.no_proxy(no_proxy.clone())) - .collect() - } } #[cfg(test)] diff --git a/litellm-rust/crates/inference-chat/src/handler.rs b/litellm-rust/crates/inference-chat/src/handler.rs index 595e851c6fd..c52ff19ac08 100644 --- a/litellm-rust/crates/inference-chat/src/handler.rs +++ b/litellm-rust/crates/inference-chat/src/handler.rs @@ -283,7 +283,7 @@ mod tests { let (provider, config, request) = prepared(&upstream.uri()); execute( - &Client::plain_for_test(), + &Client::for_test(), &AuthServices::default(), config, provider, @@ -328,7 +328,7 @@ mod tests { let (provider, config, request) = prepared(&upstream.uri()); let error = execute( - &Client::plain_for_test(), + &Client::for_test(), &AuthServices::default(), config, provider, diff --git a/litellm-rust/crates/inference-chat/tests/caching.rs b/litellm-rust/crates/inference-chat/tests/caching.rs index 70da8bc3255..4b3e32f04bc 100644 --- a/litellm-rust/crates/inference-chat/tests/caching.rs +++ b/litellm-rust/crates/inference-chat/tests/caching.rs @@ -239,7 +239,7 @@ async fn chat_cache_identity_follows_resolved_configuration_and_request_callback .store(usize::from(call >= 2), Ordering::SeqCst); let cache = ScopedCache::new(cache.clone(), CacheScope::Shared); ChatCompletionsRoute::new( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), Arc::new(Default::default()), secrets.clone(), ) diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index 55c80d9bddf..8b6154ecbee 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -888,7 +888,7 @@ mod tests { .respond_with(ResponseTemplate::new(200).set_body_raw(body, "text/event-stream")) .mount(&upstream) .await; - let response = litellm_http::Client::plain_for_test() + let response = litellm_http::Client::for_test() .get(upstream.uri()) .send() .await diff --git a/litellm-rust/crates/inference-ocr/src/handler.rs b/litellm-rust/crates/inference-ocr/src/handler.rs index e9cfc9a8989..426d4297f28 100644 --- a/litellm-rust/crates/inference-ocr/src/handler.rs +++ b/litellm-rust/crates/inference-ocr/src/handler.rs @@ -121,11 +121,7 @@ mod tests { }; fn client() -> OcrClient { - OcrClient::for_test( - litellm_http::Client::plain_for_test(), - litellm_http::Client::no_redirect_for_test(), - ) - .with_secrets(litellm_inference_testing::no_secrets()) + OcrClient::for_test().with_secrets(litellm_inference_testing::no_secrets()) } fn request(model: &str, base: &str, document: Value, options: Value) -> LiteLLMOcrRequest { diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/document.rs b/litellm-rust/crates/llms/src/base_llm/ocr/document.rs index b32ea4bf73a..57976240e1b 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/document.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/document.rs @@ -190,16 +190,14 @@ mod tests { clippy::disallowed_types, reason = "the pool has no default-header setting to stand in for provider credentials" )] - let provider_http = litellm_http::Client::for_test( + let provider_http = litellm_http::Client::wrap_for_test( reqwest::Client::builder() .default_headers(provider_headers) .build() .unwrap(), ); - let client = crate::base_llm::ocr::handler::OcrClient::for_test( - provider_http, - litellm_http::Client::no_redirect_for_test(), - ); + let client = + crate::base_llm::ocr::handler::OcrClient::for_test().with_provider_http(provider_http); let converted = inline_remote_document( client.document_fetcher(), OcrDocument::ImageUrl { diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs index 2cf955ff0c6..21ad4ea8025 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs @@ -86,21 +86,27 @@ impl OcrClient { } #[cfg(any(test, feature = "test-support"))] - pub fn for_test(provider_http: Client, no_redirect_http: Client) -> Self { + pub fn for_test() -> Self { Self { secrets: Arc::new( - litellm_secrets::source::EnvironmentSecrets::python_compatible( - provider_http.clone(), - ), + litellm_secrets::source::EnvironmentSecrets::python_compatible(Client::for_test()), ), - provider_http, - polling_http: no_redirect_http.clone(), - document_fetcher: MediaFetcher::for_test(no_redirect_http), + provider_http: Client::for_test(), + polling_http: Client::for_test(), + document_fetcher: MediaFetcher::for_test(Client::for_test()), auth: Arc::new(AuthServices::default()), settings: OcrSettings::default(), } } + #[cfg(any(test, feature = "test-support"))] + pub fn with_provider_http(self, provider_http: Client) -> Self { + Self { + provider_http, + ..self + } + } + #[cfg(any(test, feature = "test-support"))] pub fn with_settings(self, settings: OcrSettings) -> Self { Self { settings, ..self } @@ -322,7 +328,7 @@ mod tests { let _connection = listener.accept().await.unwrap(); tokio::time::sleep(Duration::from_secs(1)).await; }); - let error = litellm_http::Client::plain_for_test() + let error = litellm_http::Client::for_test() .get(format!("http://{address}")) .timeout(Duration::from_millis(10)) .send() @@ -347,7 +353,7 @@ mod tests { socket.write_all(response.as_bytes()).await.unwrap(); std::future::pending::<()>().await; }); - let response = litellm_http::Client::plain_for_test() + let response = litellm_http::Client::for_test() .get(format!("http://{address}")) .send() .await diff --git a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs index 1979b442936..9b991b75e5b 100644 --- a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs @@ -566,10 +566,7 @@ mod tests { let params = ReductoParseV3Config .map_ocr_params(&overrides, "parse-v3") .unwrap(); - let client = OcrClient::for_test( - litellm_http::Client::plain_for_test(), - litellm_http::Client::no_redirect_for_test(), - ); + let client = OcrClient::for_test(); let connection = OcrConnection::default(); let document = serde_json::from_value( json!({"type":"document_url","document_url":"reducto://ready.pdf"}), diff --git a/litellm-rust/crates/python-bridge/src/secrets/callback.rs b/litellm-rust/crates/python-bridge/src/secrets/callback.rs index 71a06bfdc39..669c10e602a 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/callback.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/callback.rs @@ -192,7 +192,7 @@ handler.get_secret_from_manager = get_secret_from_manager KeyManagementSettings::default(), )), Arc::new(move |_: &str| fallback.map(str::to_owned)), - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ) .with_failure_policy(FailurePolicy::EnvironmentFallback); (resolver, locals, handler) diff --git a/litellm-rust/crates/python-bridge/src/secrets/resolved.rs b/litellm-rust/crates/python-bridge/src/secrets/resolved.rs index 5a606ab1039..deb07175e7a 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/resolved.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/resolved.rs @@ -84,7 +84,7 @@ mod tests { } async fn resolve(state: Arc, name: &'static str) -> Option { - ResolvedSecrets::from_state(state, litellm_http::Client::plain_for_test()) + ResolvedSecrets::from_state(state, litellm_http::Client::for_test()) .resolve(&[name]) .await .unwrap() @@ -182,7 +182,7 @@ mod tests { .await; let source = ResolvedSecrets::from_state( state(&server, KeyManagementSettings::default()), - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), ); let snapshot = source.resolve(&[declared]).await.unwrap(); assert_eq!(snapshot.get(undeclared), None); @@ -248,7 +248,7 @@ mod tests { async fn oidc_failures_are_not_converted_to_missing_secrets() { let result = ResolvedSecrets::from_state( Arc::new(SecretManagerState::default()), - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), ) .resolve(&["oidc/"]) .await; diff --git a/litellm-rust/crates/secrets-azure/tests/key_vault.rs b/litellm-rust/crates/secrets-azure/tests/key_vault.rs index fcbc46092e1..05f62d25dce 100644 --- a/litellm-rust/crates/secrets-azure/tests/key_vault.rs +++ b/litellm-rust/crates/secrets-azure/tests/key_vault.rs @@ -11,7 +11,7 @@ use wiremock::{ fn manager(server: &MockServer) -> AzureKeyVault { AzureKeyVault::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), server.uri().parse().unwrap(), Arc::new(|name: &str| (name == "AZURE_AD_TOKEN").then(|| "fake".to_owned())), ) @@ -131,7 +131,7 @@ fn new_validates_vault_environment( #[case] missing_environment: bool, ) { let result = AzureKeyVault::new( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), Arc::new(move |name: &str| { (name == "AZURE_KEY_VAULT_URI") .then(|| uri.map(str::to_owned)) @@ -158,7 +158,7 @@ fn new_validates_vault_environment( #[case::local("http://localhost:8080", "https://localhost/.default")] fn derives_scope_from_vault_host(#[case] uri: &str, #[case] expected: &str) { let manager = AzureKeyVault::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), uri.parse().unwrap(), Arc::new(|_: &str| None), ) @@ -187,7 +187,7 @@ async fn missing_credentials_do_not_request_vault() { fn manager_without_credentials(server: &MockServer) -> AzureKeyVault { AzureKeyVault::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), server.uri().parse().unwrap(), Arc::new(|name: &str| { (name == "AZURE_CREDENTIAL").then(|| "ClientSecretCredential".to_owned()) diff --git a/litellm-rust/crates/secrets-azure/tests/live.rs b/litellm-rust/crates/secrets-azure/tests/live.rs index 429cd3013f7..d7713dc798d 100644 --- a/litellm-rust/crates/secrets-azure/tests/live.rs +++ b/litellm-rust/crates/secrets-azure/tests/live.rs @@ -10,7 +10,7 @@ use rstest::rstest; #[ignore] async fn reads_a_real_secret() { let environment = Arc::new(ProcessEnvironment); - let manager = AzureKeyVault::new(litellm_http::Client::plain_for_test(), environment).unwrap(); + let manager = AzureKeyVault::new(litellm_http::Client::for_test(), environment).unwrap(); let name = std::env::var("AZURE_KEY_VAULT_LIVE_SECRET_NAME").unwrap(); let secret = manager.get_secret(&name).await.unwrap().unwrap(); assert!(matches!(&secret, Secret::String(_))); diff --git a/litellm-rust/crates/secrets-cyberark/tests/secret_manager/configuration.rs b/litellm-rust/crates/secrets-cyberark/tests/secret_manager/configuration.rs index 1bea77e49ae..a3c20a0db3d 100644 --- a/litellm-rust/crates/secrets-cyberark/tests/secret_manager/configuration.rs +++ b/litellm-rust/crates/secrets-cyberark/tests/secret_manager/configuration.rs @@ -77,7 +77,7 @@ async fn authentication_encodes_login(#[case] username: &str, #[case] expected_p .mount(&server) .await; let manager = CyberArkSecretManager::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), server.uri().parse().unwrap(), "acct".into(), username.into(), @@ -433,7 +433,7 @@ async fn trailing_slash_endpoint_preserves_base_path() { .await; let endpoint = format!("{}/prefix/", server.uri()).parse().unwrap(); let manager = CyberArkSecretManager::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), endpoint, "acct".into(), "admin".into(), diff --git a/litellm-rust/crates/secrets-cyberark/tests/secret_manager/support.rs b/litellm-rust/crates/secrets-cyberark/tests/secret_manager/support.rs index 6fa15ecba9b..d7bc71b6503 100644 --- a/litellm-rust/crates/secrets-cyberark/tests/secret_manager/support.rs +++ b/litellm-rust/crates/secrets-cyberark/tests/secret_manager/support.rs @@ -62,7 +62,7 @@ pub(super) fn from_environment( pub(super) fn manager(server: &MockServer, ttl: Duration) -> CyberArkSecretManager { CyberArkSecretManager::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), server.uri().parse().unwrap(), "acct".into(), "admin".into(), diff --git a/litellm-rust/crates/secrets-cyberark/tests/secret_manager/writes.rs b/litellm-rust/crates/secrets-cyberark/tests/secret_manager/writes.rs index 764a591bc71..3db8b6355e1 100644 --- a/litellm-rust/crates/secrets-cyberark/tests/secret_manager/writes.rs +++ b/litellm-rust/crates/secrets-cyberark/tests/secret_manager/writes.rs @@ -233,7 +233,7 @@ async fn writes_match_python_parity_fixture(parity_fixture: ParityFixture) { .mount(&server) .await; let manager = CyberArkSecretManager::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), server.uri().parse().unwrap(), parity_fixture.account, parity_fixture.username, @@ -297,7 +297,7 @@ async fn live_conjur_round_trip() { .as_nanos() ); let manager = CyberArkSecretManager::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), endpoint.clone(), account.clone(), username.clone(), @@ -312,7 +312,7 @@ async fn live_conjur_round_trip() { .await .unwrap(); let verifier = CyberArkSecretManager::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), endpoint.clone(), account.clone(), username.clone(), diff --git a/litellm-rust/crates/secrets-google/tests/secret_manager.rs b/litellm-rust/crates/secrets-google/tests/secret_manager.rs index e9bee7633f5..f2cb5ff3603 100644 --- a/litellm-rust/crates/secrets-google/tests/secret_manager.rs +++ b/litellm-rust/crates/secrets-google/tests/secret_manager.rs @@ -11,7 +11,7 @@ use wiremock::{ fn manager(server: &MockServer, always_read: bool, ttl: Duration) -> GoogleSecretManager { GoogleSecretManager::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), server.uri().parse().unwrap(), "project".into(), Arc::new(|name: &str| (name == "VERTEX_AI_API_KEY").then(|| "token".into())), @@ -215,7 +215,7 @@ async fn always_read_and_expired_cache_fetch_again( fn google_manager_requires_host_license_and_project_configuration() { assert!(matches!( GoogleSecretManager::new( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), Arc::new(|_: &str| None), false ), @@ -223,7 +223,7 @@ fn google_manager_requires_host_license_and_project_configuration() { )); assert!(matches!( GoogleSecretManager::new( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), Arc::new(|_: &str| None), true ), @@ -244,7 +244,7 @@ fn google_manager_rejects_invalid_refresh_intervals(#[case] variable: &'static s }); assert!(matches!( - GoogleSecretManager::new(litellm_http::Client::plain_for_test(), environment, true), + GoogleSecretManager::new(litellm_http::Client::for_test(), environment, true), Err(Error::RefreshInterval) )); } diff --git a/litellm-rust/crates/secrets/tests/aws.rs b/litellm-rust/crates/secrets/tests/aws.rs index 910b1855b77..f28f404c3ff 100644 --- a/litellm-rust/crates/secrets/tests/aws.rs +++ b/litellm-rust/crates/secrets/tests/aws.rs @@ -53,7 +53,7 @@ async fn read_results_follow_the_selected_failure_policy( }, )), Arc::new(move |_: &str| environment.map(str::to_owned)), - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ) .with_failure_policy(policy); let result = resolver @@ -97,7 +97,7 @@ async fn primary_secret_values_other_than_strings_resolve_to_none( let resolver = SecretResolver::new_python_compatible( Arc::new(state(&server, settings)), Arc::new(|_: &str| Some("fallback".into())), - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ); let text = value.as_str(); assert_eq!( @@ -157,7 +157,7 @@ async fn gating_prediction_matches_actual_lookup( let resolver = SecretResolver::new_python_compatible( Arc::new(state), Arc::new(|_: &str| Some("environment".into())), - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ); assert_eq!( resolver diff --git a/litellm-rust/crates/secrets/tests/azure.rs b/litellm-rust/crates/secrets/tests/azure.rs index 60f165add54..6e81017b0f6 100644 --- a/litellm-rust/crates/secrets/tests/azure.rs +++ b/litellm-rust/crates/secrets/tests/azure.rs @@ -22,7 +22,7 @@ async fn azure_handler_reads_missing_and_failed_secrets() { .await; let manager = SecretManager::AzureKeyVault( AzureKeyVault::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), server.uri().parse().unwrap(), std::sync::Arc::new(|name: &str| (name == "AZURE_AD_TOKEN").then(|| "fake".to_owned())), ) @@ -81,7 +81,7 @@ async fn successful_azure_responses_do_not_fall_back_when_the_value_is_empty_or_ .mount(&server) .await; let manager = AzureKeyVault::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), server.uri().parse().unwrap(), Arc::new(|name: &str| (name == "AZURE_AD_TOKEN").then(|| "token".into())), ) @@ -92,7 +92,7 @@ async fn successful_azure_responses_do_not_fall_back_when_the_value_is_empty_or_ Default::default(), )), Arc::new(|_: &str| Some("environment".into())), - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ); assert_eq!( resolver diff --git a/litellm-rust/crates/secrets/tests/common_read_contract.rs b/litellm-rust/crates/secrets/tests/common_read_contract.rs index 2e8fbf46008..e6928038799 100644 --- a/litellm-rust/crates/secrets/tests/common_read_contract.rs +++ b/litellm-rust/crates/secrets/tests/common_read_contract.rs @@ -59,7 +59,7 @@ fn manager(provider: Provider, server: &MockServer) -> SecretManager { ), Provider::Azure => SecretManager::AzureKeyVault( AzureKeyVault::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), server.uri().parse().unwrap(), environment, ) @@ -67,7 +67,7 @@ fn manager(provider: Provider, server: &MockServer) -> SecretManager { ), Provider::Google => SecretManager::GoogleSecretManager( GoogleSecretManager::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), server.uri().parse().unwrap(), "project".into(), environment, @@ -84,7 +84,7 @@ fn manager(provider: Provider, server: &MockServer) -> SecretManager { .unwrap(), ), Provider::Cyberark => SecretManager::Cyberark(CyberArkSecretManager::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), server.uri().parse().unwrap(), "acct".into(), "admin".into(), @@ -240,7 +240,7 @@ async fn python_read_failures_preserve_provider_fallback_rules( KeyManagementSettings::default(), )), Arc::new(move |_: &str| environment_value.map(str::to_owned)), - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ); let expected = if matches!(provider, Provider::Aws) { None diff --git a/litellm-rust/crates/secrets/tests/cyberark.rs b/litellm-rust/crates/secrets/tests/cyberark.rs index b94bd9dc0ad..0a6c394937e 100644 --- a/litellm-rust/crates/secrets/tests/cyberark.rs +++ b/litellm-rust/crates/secrets/tests/cyberark.rs @@ -24,7 +24,7 @@ async fn cyberark_handler_reads_values_and_surfaces_errors() { .mount(&server) .await; let manager = SecretManager::Cyberark(CyberArkSecretManager::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), server.uri().parse().unwrap(), "acct".into(), "admin".into(), diff --git a/litellm-rust/crates/secrets/tests/google.rs b/litellm-rust/crates/secrets/tests/google.rs index 0b45a90c10b..fa1db9916e8 100644 --- a/litellm-rust/crates/secrets/tests/google.rs +++ b/litellm-rust/crates/secrets/tests/google.rs @@ -25,7 +25,7 @@ async fn google_resolver_distinguishes_absence_from_failure(#[case] status: u16) _ => None, }); let manager = GoogleSecretManager::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), server.uri().parse().unwrap(), "project".into(), environment.clone(), @@ -40,7 +40,7 @@ async fn google_resolver_distinguishes_absence_from_failure(#[case] status: u16) let resolver = SecretResolver::new_python_compatible( Arc::new(state), environment, - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ) .with_failure_policy(FailurePolicy::Propagate); let result = resolver.get_secret_str("KEY", None).await; diff --git a/litellm-rust/crates/secrets/tests/hashicorp.rs b/litellm-rust/crates/secrets/tests/hashicorp.rs index e10e903c25a..bfb289cdb97 100644 --- a/litellm-rust/crates/secrets/tests/hashicorp.rs +++ b/litellm-rust/crates/secrets/tests/hashicorp.rs @@ -56,7 +56,7 @@ async fn hashicorp_handler_resolves_found_missing_and_failed_values() { }, )), Arc::new(|_: &str| None), - litellm_secrets::OidcResolver::new(litellm_http::Client::plain_for_test()), + litellm_secrets::OidcResolver::new(litellm_http::Client::for_test()), ); assert_eq!( found_resolver @@ -131,7 +131,7 @@ async fn hashicorp_handler_resolves_found_missing_and_failed_values() { let failed_resolver = SecretResolver::new_python_compatible( Arc::new(failed_state), Arc::new(|_: &str| None), - litellm_secrets::OidcResolver::new(litellm_http::Client::plain_for_test()), + litellm_secrets::OidcResolver::new(litellm_http::Client::for_test()), ) .with_failure_policy(FailurePolicy::Propagate); assert!(matches!( diff --git a/litellm-rust/crates/secrets/tests/oidc.rs b/litellm-rust/crates/secrets/tests/oidc.rs index 6f72bd9d645..867a662cbbb 100644 --- a/litellm-rust/crates/secrets/tests/oidc.rs +++ b/litellm-rust/crates/secrets/tests/oidc.rs @@ -30,7 +30,7 @@ async fn environment_sources_resolve_expected_value( ("CIRCLE_OIDC_TOKEN_V2", "circle-v2"), ]); assert_eq!( - OidcResolver::new(litellm_http::Client::plain_for_test()) + OidcResolver::new(litellm_http::Client::for_test()) .resolve(reference, env.as_ref()) .await .unwrap() @@ -43,7 +43,7 @@ async fn environment_sources_resolve_expected_value( #[tokio::test] async fn environment_sources_bypass_boolean_conversion_and_defaults() { let env = environment(&[("TOKEN", "true")]); - let oidc = OidcResolver::new(litellm_http::Client::plain_for_test()); + let oidc = OidcResolver::new(litellm_http::Client::for_test()); let resolver = SecretResolver::new(Arc::new(SecretManagerState::default()), env, oidc); assert_eq!( resolver @@ -94,7 +94,7 @@ async fn github_requests_are_authenticated_cached_and_revalidate_environment() { ), ("ACTIONS_ID_TOKEN_REQUEST_TOKEN", "request-token"), ]); - let oidc = OidcResolver::new(litellm_http::Client::plain_for_test()); + let oidc = OidcResolver::new(litellm_http::Client::for_test()); for _ in 0..2 { assert_eq!( oidc.resolve("oidc/github/https://service/oidc/path", env.as_ref()) @@ -131,7 +131,7 @@ async fn file_allowlist_resolves_symlinks_while_environment_paths_remain_explici ("PATH_TOKEN", private.to_str().unwrap()), ("AZURE_FEDERATED_TOKEN_FILE", token.to_str().unwrap()), ]); - let oidc = OidcResolver::new(litellm_http::Client::plain_for_test()); + let oidc = OidcResolver::new(litellm_http::Client::for_test()); assert_eq!( oidc.resolve(&format!("oidc/file/{}", token.display()), env.as_ref()) .await @@ -213,7 +213,7 @@ async fn google_expiry_caps_cache_and_preserves_audience( .expect(calls) .mount(&server) .await; - let oidc = OidcResolver::new(litellm_http::Client::plain_for_test()) + let oidc = OidcResolver::new(litellm_http::Client::for_test()) .with_google_identity_endpoint(server.uri().parse().unwrap()) .with_clock(now); for _ in 0..2 { @@ -235,7 +235,7 @@ async fn google_expiry_caps_cache_and_preserves_audience( #[tokio::test] async fn google_oidc_requires_its_build_feature() { assert!(matches!( - OidcResolver::new(litellm_http::Client::plain_for_test()) + OidcResolver::new(litellm_http::Client::for_test()) .resolve("oidc/google/audience", environment(&[]).as_ref()) .await, Err(Error::UnsupportedOidc) @@ -246,7 +246,7 @@ async fn google_oidc_requires_its_build_feature() { #[tokio::test] async fn azure_oidc_without_a_token_file_requires_its_build_feature() { assert!(matches!( - OidcResolver::new(litellm_http::Client::plain_for_test()) + OidcResolver::new(litellm_http::Client::for_test()) .resolve("oidc/azure/scope", environment(&[]).as_ref()) .await, Err(Error::UnsupportedOidc) @@ -262,7 +262,7 @@ async fn invalid_references_fail_before_environment_lookup( #[case] reference: &str, #[case] unsupported: bool, ) { - let error = OidcResolver::new(litellm_http::Client::plain_for_test()) + let error = OidcResolver::new(litellm_http::Client::for_test()) .resolve(reference, &|_: &str| { panic!("invalid reference reached environment lookup") }) @@ -284,7 +284,7 @@ async fn unreadable_expiry_keeps_python_cache_fallback(#[case] token: &str) { .expect(1) .mount(&server) .await; - let resolver = OidcResolver::new(litellm_http::Client::plain_for_test()) + let resolver = OidcResolver::new(litellm_http::Client::for_test()) .with_google_identity_endpoint(server.uri().parse().unwrap()); for _ in 0..2 { assert_eq!( @@ -336,7 +336,7 @@ async fn azure_oidc_acquires_the_requested_scope_and_preserves_failures(#[case] }) } } - let oidc = OidcResolver::new(litellm_http::Client::plain_for_test()) + let oidc = OidcResolver::new(litellm_http::Client::for_test()) .with_azure_token_provider(Arc::new(Provider(failed))); let resolver = SecretResolver::new_python_compatible( Arc::new(SecretManagerState::default()), @@ -364,7 +364,7 @@ async fn azure_oidc_acquires_the_requested_scope_and_preserves_failures(#[case] #[tokio::test] async fn missing_oidc_environment_is_an_error(#[case] reference: &str) { assert!(matches!( - OidcResolver::new(litellm_http::Client::plain_for_test()) + OidcResolver::new(litellm_http::Client::for_test()) .resolve(reference, environment(&[]).as_ref()) .await, Err(Error::MissingEnvironment) @@ -383,7 +383,7 @@ async fn google_oidc_failures_are_not_cached_or_hidden_by_defaults() { let resolver = SecretResolver::new_python_compatible( Arc::new(SecretManagerState::default()), environment(&[]), - OidcResolver::new(litellm_http::Client::plain_for_test()) + OidcResolver::new(litellm_http::Client::for_test()) .with_google_identity_endpoint(server.uri().parse().unwrap()), ); for _ in 0..2 { @@ -424,7 +424,7 @@ async fn google_tokens_expire_at_the_python_cache_deadline( .expect(2) .mount(&server) .await; - let resolver = OidcResolver::new(litellm_http::Client::plain_for_test()) + let resolver = OidcResolver::new(litellm_http::Client::for_test()) .with_google_identity_endpoint(server.uri().parse().unwrap()) .with_clock(|| UNIX_EPOCH + Duration::from_secs(1000)); assert_eq!( @@ -477,7 +477,7 @@ async fn google_cache_uses_payload_expiry_without_requiring_a_jwt_header() { .expect(2) .mount(&server) .await; - let resolver = OidcResolver::new(litellm_http::Client::plain_for_test()) + let resolver = OidcResolver::new(litellm_http::Client::for_test()) .with_google_identity_endpoint(server.uri().parse().unwrap()) .with_clock(|| UNIX_EPOCH + Duration::from_secs(1000)); for _ in 0..2 { diff --git a/litellm-rust/crates/secrets/tests/resolution.rs b/litellm-rust/crates/secrets/tests/resolution.rs index 37557c36fb5..0858089b94f 100644 --- a/litellm-rust/crates/secrets/tests/resolution.rs +++ b/litellm-rust/crates/secrets/tests/resolution.rs @@ -12,7 +12,7 @@ fn resolver(value: Option<&str>) -> SecretResolver { SecretResolver::new_python_compatible( Arc::new(SecretManagerState::default()), Arc::new(move |_: &str| value.clone()), - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ) } @@ -35,7 +35,7 @@ async fn native_reads_preserve_strings_and_report_conversion_errors(#[case] mana let resolver = SecretResolver::new( Arc::new(state), Arc::new(move |_: &str| Some(raw.to_owned())), - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ); assert_eq!( resolver @@ -75,7 +75,7 @@ async fn native_defaults_apply_to_absence_but_never_hide_provider_failures() { KeyManagementSettings::default(), )), Arc::new(|_: &str| None), - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ); let result = resolver .get_secret_str("key", Some(SecretValue::new("default"))) @@ -189,7 +189,7 @@ fn managed(reply: Result, ()>, environment: Option<&'static str>) KeyManagementSettings::default(), )), Arc::new(move |_: &str| environment.map(str::to_owned)), - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ) .with_failure_policy(FailurePolicy::EnvironmentFallback) } @@ -281,7 +281,7 @@ async fn prefix_is_removed_once_and_resolved_from_environment() { let resolver = SecretResolver::new_python_compatible( Arc::new(state), Arc::new(|name: &str| (name == "os.environ/KEY").then(|| "value".into())), - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ); assert_eq!( resolver @@ -340,7 +340,7 @@ async fn excluded_hosted_keys_keep_the_python_manager_conversion_path( }, )), Arc::new(move |_: &str| Some(raw.to_owned())), - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ); assert_eq!( resolver.get_secret("KEY", None).await.unwrap(), @@ -373,7 +373,7 @@ async fn azure_callback_absence_preserves_none_but_errors_fall_back( KeyManagementSettings::default(), )), Arc::new(|_: &str| Some("environment".into())), - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ); assert_eq!( resolver diff --git a/litellm-rust/crates/secrets/tests/source.rs b/litellm-rust/crates/secrets/tests/source.rs index 17210940a67..41f0ecea518 100644 --- a/litellm-rust/crates/secrets/tests/source.rs +++ b/litellm-rust/crates/secrets/tests/source.rs @@ -15,7 +15,7 @@ mod tests { #[case] expected: Option<&str>, ) { unsafe { std::env::set_var(name, value) }; - let secret = EnvironmentSecrets::python_compatible(litellm_http::Client::plain_for_test()) + let secret = EnvironmentSecrets::python_compatible(litellm_http::Client::for_test()) .resolve(&[name]) .await .unwrap() @@ -42,7 +42,7 @@ async fn dynamic_names_use_the_same_resolver_and_snapshots_never_do_fresh_lookup reads.fetch_add(1, Ordering::SeqCst); (name != "missing").then(|| name.to_owned()) }), - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ); let snapshot = source.resolve(&["declared", "missing"]).await.unwrap(); let name = format!("runtime-{}", "key"); diff --git a/litellm-rust/crates/spend-clickhouse/tests/storage.rs b/litellm-rust/crates/spend-clickhouse/tests/storage.rs index 2b3eb75c658..81b6527f9e1 100644 --- a/litellm-rust/crates/spend-clickhouse/tests/storage.rs +++ b/litellm-rust/crates/spend-clickhouse/tests/storage.rs @@ -77,7 +77,7 @@ async fn database() -> Database { 7, ) .unwrap(), - client: Client::no_redirect_for_test(), + client: Client::for_test(), } } diff --git a/litellm-rust/crates/spend-clickhouse/tests/transport.rs b/litellm-rust/crates/spend-clickhouse/tests/transport.rs index ffc3423f65e..49a4b365135 100644 --- a/litellm-rust/crates/spend-clickhouse/tests/transport.rs +++ b/litellm-rust/crates/spend-clickhouse/tests/transport.rs @@ -29,7 +29,7 @@ async fn spend_insert_controls_target_and_retry_settings() { .mount(&server) .await; let config = Config::new("spend_test".into(), &format!("{}?query=DROP&async_insert=0&wait_for_async_insert=0&input_format_skip_unknown_fields=1&database=wrong&readonly=1", server.uri()), 7).unwrap(); - let client = Client::no_redirect_for_test(); + let client = Client::for_test(); let row = BTreeMap::from([ ("request_id".into(), json!("request")), ("start_time".into(), json!(1234)), @@ -77,7 +77,7 @@ async fn spend_insert_preserves_storage_failure(#[case] status: u16) { .await; let config = Config::new("spend_test".into(), &server.uri(), 7).unwrap(); let error = insert_rows( - &Client::no_redirect_for_test(), + &Client::for_test(), config.storage().writer(), "spend_test", vec![BTreeMap::from([("request_id".into(), json!("request"))])], @@ -95,7 +95,7 @@ async fn empty_spend_insert_does_not_contact_storage() { let server = MockServer::start().await; let config = Config::new("spend_test".into(), &server.uri(), 7).unwrap(); insert_rows( - &Client::no_redirect_for_test(), + &Client::for_test(), config.storage().writer(), "spend_test", vec![], @@ -116,7 +116,7 @@ async fn failed_schema_setup_is_reported() { .await; let config = Config::new("spend_test".into(), &server.uri(), 7).unwrap(); assert!(matches!( - ensure_schema(&Client::no_redirect_for_test(), &config).await, + ensure_schema(&Client::for_test(), &config).await, Err(Error::Storage( litellm_storage_clickhouse::Error::SchemaFailed(403) )) diff --git a/litellm-rust/crates/storage-clickhouse/tests/migrations.rs b/litellm-rust/crates/storage-clickhouse/tests/migrations.rs index 29a85a8c4ec..49f481df632 100644 --- a/litellm-rust/crates/storage-clickhouse/tests/migrations.rs +++ b/litellm-rust/crates/storage-clickhouse/tests/migrations.rs @@ -48,7 +48,7 @@ async fn database() -> TestResult { Ok(ClickHouseDatabase { _container: container, url, - client: Client::no_redirect_for_test(), + client: Client::for_test(), }) } @@ -244,7 +244,7 @@ async fn failed_migration_is_not_recorded_and_retains_storage_error( #[rstest] fn invalid_database_identifier_is_rejected() { - let client = Client::no_redirect_for_test(); + let client = Client::for_test(); let connection = Connection::writer("http://127.0.0.1:1").expect("valid URL"); assert!(matches!( ClickHouseMigrate::new( @@ -327,7 +327,7 @@ async fn duplicate_ledger_rows_are_tolerated( #[rstest] #[tokio::test] async fn revert_reports_forward_only_error() { - let client = Client::no_redirect_for_test(); + let client = Client::for_test(); let connection = Connection::writer("http://127.0.0.1:1").expect("valid URL"); let migration = migration( 1, @@ -369,7 +369,7 @@ async fn applied_int64_versions_accept_numeric_and_quoted_json( ) .mount(&mock_server) .await; - let client = Client::no_redirect_for_test(); + let client = Client::for_test(); let connection = Connection::writer(&mock_server.uri()).expect("valid URL"); let migrator = migrator(vec![]); let mut adapter = ClickHouseMigrate::new( @@ -406,7 +406,7 @@ async fn schema_requests_override_unsafe_connection_settings( mock_server.uri() ); execute_statement( - &Client::no_redirect_for_test(), + &Client::for_test(), &Connection::writer(&url).expect("valid URL"), "CREATE DATABASE IF NOT EXISTS trace_test", REQUEST_TIMEOUT, @@ -445,7 +445,7 @@ async fn oversized_ledger_response_is_rejected_before_migrations( ) .mount(&mock_server) .await; - let client = Client::no_redirect_for_test(); + let client = Client::for_test(); let connection = Connection::writer(&mock_server.uri()).expect("valid URL"); let migrator = migrator(vec![]); let mut adapter = ClickHouseMigrate::new( diff --git a/litellm-rust/crates/storage-clickhouse/tests/transport.rs b/litellm-rust/crates/storage-clickhouse/tests/transport.rs index f70fe52e1de..2560c2985f0 100644 --- a/litellm-rust/crates/storage-clickhouse/tests/transport.rs +++ b/litellm-rust/crates/storage-clickhouse/tests/transport.rs @@ -13,7 +13,7 @@ async fn insert_rejects_invalid_identifiers( #[case] table: &str, #[case] invalid_database: bool, ) { - let client = Client::no_redirect_for_test(); + let client = Client::for_test(); let connection = Connection::writer("http://localhost:8123").expect("valid URL"); let result = insert_encoded_rows(&client, &connection, database, table, "token", "{}").await; @@ -24,7 +24,7 @@ async fn insert_rejects_invalid_identifiers( #[rstest] #[tokio::test] async fn read_rejects_empty_sql() { - let client = Client::no_redirect_for_test(); + let client = Client::for_test(); let connection = Connection::reader("http://localhost:8123", "litellm").expect("valid URL"); assert!(matches!( @@ -87,7 +87,7 @@ async fn typed_fetch_encodes_parameters_and_validates_rows( .expect(2) .mount(&server) .await; - let client = Client::no_redirect_for_test(); + let client = Client::for_test(); let connection = Connection::parse(&server.uri()).unwrap(); let params = QueryParams { signed: i64::MIN, @@ -131,7 +131,7 @@ async fn server_result_limits_allow_smaller_pages_without_retrying_other_failure .await; let connection = Connection::parse(&server.uri()).unwrap(); let error = execute_read( - &Client::no_redirect_for_test(), + &Client::for_test(), &connection, "SELECT 1", &BTreeMap::new(), @@ -174,7 +174,7 @@ async fn insert_timeout_environment_child() { .mount(&server) .await; let result = insert_encoded_rows( - &Client::no_redirect_for_test(), + &Client::for_test(), &Connection::parse(&server.uri()).unwrap(), "traces", "otel_traces", diff --git a/pyproject.toml b/pyproject.toml index 0b3480cab23..2d25c1ec737 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -296,7 +296,7 @@ ci = [ "anthropic==0.84.0", "blockbuster==1.5.26", "beautifulsoup4==4.14.3", - "pylint==4.0.5", + "pylint==4.0.9", "langchain-openai==1.1.14", "langgraph>=1.2.4,<1.3.0", "langgraph-prebuilt>=1.1.0,<1.3.0", diff --git a/uv.lock b/uv.lock index d38e96882d0..b40490112cf 100644 --- a/uv.lock +++ b/uv.lock @@ -1077,11 +1077,30 @@ wheels = [ name = "bytecode" version = "0.17.0" source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version == '3.14.*'", + "python_full_version == '3.13.*'", + "python_full_version == '3.12.*'", + "python_full_version == '3.11.*'", + "python_full_version < '3.11'", +] sdist = { url = "https://files.pythonhosted.org/packages/98/c4/4818b392104bd426171fc2ce9c79c8edb4019ba6505747626d0f7107766c/bytecode-0.17.0.tar.gz", hash = "sha256:0c37efa5bd158b1b873f530cceea2c645611d55bd2dc2a4758b09f185749b6fd", size = 105863, upload-time = "2025-09-03T19:55:45.703Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/ce/80/379e685099841f8501a19fb58b496512ef432331fed38276c3938ab09d8e/bytecode-0.17.0-py3-none-any.whl", hash = "sha256:64fb10cde1db7ef5cc39bd414ecebd54ba3b40e1c4cf8121ca5e72f170916ff8", size = 43045, upload-time = "2025-09-03T19:55:43.879Z" }, ] +[[package]] +name = "bytecode" +version = "0.19.1" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.15'", +] +sdist = { url = "https://files.pythonhosted.org/packages/5c/a0/f4499eea31339ba709fad8ff778c6ad9afa2d87434b9771e0b37a1ab6fb1/bytecode-0.19.1.tar.gz", hash = "sha256:f59a6a51a980037f64196448504d12fb04937ae8dcb980943265e799af2943ae", size = 106830, upload-time = "2026-10-07T12:01:11.178Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b5/d2/7b937ec5bcc2fabca017e32541a0b1ef081bb5a1bb3531d382300f7a1ccc/bytecode-0.19.1-py3-none-any.whl", hash = "sha256:211a5091abd6f16f5bb3d34f1b88fc77f8f8539801a02384b51a97bdcee45f0f", size = 43473, upload-time = "2026-10-07T12:01:09.866Z" }, +] + [[package]] name = "cachetools" version = "7.1.1" @@ -1817,64 +1836,71 @@ wheels = [ [[package]] name = "ddtrace" -version = "4.11.0" +version = "4.15.6" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "bytecode" }, + { name = "bytecode", version = "0.17.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.15'" }, + { name = "bytecode", version = "0.19.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.15'" }, + { name = "ddtrace-internal" }, { name = "envier" }, { name = "opentelemetry-api" }, { name = "wrapt" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/81/51/a628f0177274bab5b67c93a1558fd222babb15286a427e4e8c1d65a265b2/ddtrace-4.11.0.tar.gz", hash = "sha256:260c5b46e80565f4fd08cec2650f707627fb57cd6f8951a2d8c9e6a02b490074", size = 2422158, upload-time = "2026-07-10T08:57:53.396Z" } +sdist = { url = "https://files.pythonhosted.org/packages/f6/94/f0f005257dc49a7062c6fed98a6dbcca7691e9d1db5c64e8a2d14f204cac/ddtrace-4.15.6.tar.gz", hash = "sha256:2914f646d4ae338c644b814c4ad80c963e7aadfb39d1e67b0367a15d5c4bc0db", size = 2655968, upload-time = "2026-10-08T09:46:58.392Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/16/1f/05dc12819ca3a3000028a522cab0025c5dd969e97dc472ac38a7d2f1ea71/ddtrace-4.11.0-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:0bc55e48da62b6a287b36e01f5876489023178c09b60e8e01608e2e895ccdc4b", size = 7103767, upload-time = "2026-07-10T08:55:10.348Z" }, - { url = "https://files.pythonhosted.org/packages/e8/8a/efcea448bc15b4d95c48c9b476f2ddc5862a53cd4baf58fe50c64763f9d8/ddtrace-4.11.0-cp310-cp310-macosx_14_0_x86_64.whl", hash = "sha256:5df47873c625993e81fabd8f6ae2ab61f66c6805cfd1a8c551a509d5b5d91062", size = 7447104, upload-time = "2026-07-10T08:55:15.794Z" }, - { url = "https://files.pythonhosted.org/packages/39/28/89aa239e47d28900ab394a0b8e0a5549cab631c0e5977016ad9d578eaee7/ddtrace-4.11.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e857428cc86fe541f5b3768d7d8331e2f1cbde1a2492deedff4619109cbfd5fc", size = 8511839, upload-time = "2026-07-10T08:55:17.609Z" }, - { url = "https://files.pythonhosted.org/packages/30/43/37a5475ade911abae592e121e852922745d842e16294c549cba4ad0f3dbe/ddtrace-4.11.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fc60575f8f47af58e6523f8a802d1f92a4912aa7d4ba71024958421044cba61a", size = 8733381, upload-time = "2026-07-10T08:55:19.66Z" }, - { url = "https://files.pythonhosted.org/packages/10/df/d61f71236a405480c8809522219b184cc486056c32f9cc8bf07cc91e2be9/ddtrace-4.11.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:d65ab60bfb0b2edde6c564bd1f97b99a8ea9bb082c93b265488ba214c606d240", size = 9513521, upload-time = "2026-07-10T08:55:21.921Z" }, - { url = "https://files.pythonhosted.org/packages/79/4f/bff08af986df419179b86fc9c73732836531c2d08ab8dcfe982a6f2e8cb8/ddtrace-4.11.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:7ebfdee1342b5b1cdad8505e040bbfe5bdf6976e820803eaa4a1814417d6bb0e", size = 9788008, upload-time = "2026-07-10T08:55:24.111Z" }, - { url = "https://files.pythonhosted.org/packages/af/b8/2b1fd5b0374edf33ad31adb746d1d1358a44c39a45b1c2a0c4b09651d81e/ddtrace-4.11.0-cp310-cp310-win32.whl", hash = "sha256:6247a431ae3a2622a0e835357a644b1bac81719e28f21748fd7d4034c79c3ed7", size = 5739013, upload-time = "2026-07-10T08:55:27.028Z" }, - { url = "https://files.pythonhosted.org/packages/f2/f5/9bdff59bdce1f80592e70633dc26678c762fff79188c9fb6004c06967708/ddtrace-4.11.0-cp310-cp310-win_amd64.whl", hash = "sha256:9197110982756a97243906ffce7b98e6007c8258a6bbc22451d7b26e8f1b0103", size = 6326099, upload-time = "2026-07-10T08:55:28.962Z" }, - { url = "https://files.pythonhosted.org/packages/76/36/ad0abcc97b63d821ada82cc7452c3070a71f9546f0aa81b50eda1f8ffd41/ddtrace-4.11.0-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:b4074e538e708b0e9d7f812be15db306abbe3a6cb70046f9f786d66daa862344", size = 7105089, upload-time = "2026-07-10T08:55:31.006Z" }, - { url = "https://files.pythonhosted.org/packages/a6/3a/c0f8f70883a70ec9d4050900823ff7b52ad54dc6b2a3d29561727825410f/ddtrace-4.11.0-cp311-cp311-macosx_14_0_x86_64.whl", hash = "sha256:05f701229233ad8a4ad5d52b8ff2bfd6e2f99601a3bc04367615cd1ce4a01a80", size = 7448237, upload-time = "2026-07-10T08:55:32.962Z" }, - { url = "https://files.pythonhosted.org/packages/4e/46/18a42afab21cd718467e7d7ae3e80898569d9969d731b771b5a8ada10294/ddtrace-4.11.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:c3a557891e91a443e6a3a9b7c31dc96b9963acada0986dab5a3d65507efa16a1", size = 8516433, upload-time = "2026-07-10T08:55:35.011Z" }, - { url = "https://files.pythonhosted.org/packages/da/b2/6367a75518f843872e13c062ac3c329f880b2f4d62ff4de0624bdcf64719/ddtrace-4.11.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:b0328bfeb4a6bdc00168a4184aea670eedc073734f6fbc6ca1abfeb6547415d2", size = 8740643, upload-time = "2026-07-10T08:55:37.743Z" }, - { url = "https://files.pythonhosted.org/packages/ae/e8/04f1d7ee4737444d45f81545ce21b01bd9d653a994997ab1e7e04b6ee947/ddtrace-4.11.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:ad6f17111746568a58da89b6db5e5d18eb07db7b8ce3be3396631f37391cff50", size = 9518626, upload-time = "2026-07-10T08:55:40.248Z" }, - { url = "https://files.pythonhosted.org/packages/ed/8b/314ba8a81e562a31655328061afb30d26d46ba04aca4dabd7da85c93a9e7/ddtrace-4.11.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:eb30dd327789ec1851bc42194c8db0b93a7ad708c2963d694e3622b684a085f7", size = 9795392, upload-time = "2026-07-10T08:55:42.609Z" }, - { url = "https://files.pythonhosted.org/packages/9e/6e/a580d53c05a771bc53d33a88174e983146fdf879fe5d4f52e51ceb56ee99/ddtrace-4.11.0-cp311-cp311-win32.whl", hash = "sha256:dfe73130778ca22652ae398fc77de3643d50687deb127bf4cd2566b10778125a", size = 5737732, upload-time = "2026-07-10T08:55:45.605Z" }, - { url = "https://files.pythonhosted.org/packages/12/83/094f5b88690f7801129b0d52a7b6a58a95b34ecbfcb337555d323c639faa/ddtrace-4.11.0-cp311-cp311-win_amd64.whl", hash = "sha256:1241f38d52b40b516d620ed68ecd04afbc715f2c1c9f3ddae8a7b86ecf848c91", size = 6329473, upload-time = "2026-07-10T08:55:47.603Z" }, - { url = "https://files.pythonhosted.org/packages/60/72/036d11ceca6fbdc4ef95a1aba210cb3fc72bc912703cba49cf2bc8b610c5/ddtrace-4.11.0-cp311-cp311-win_arm64.whl", hash = "sha256:3013b67e1a172eb0edd8c37559b81760a9db205a671541925ebd0127f496409b", size = 5990873, upload-time = "2026-07-10T08:55:49.705Z" }, - { url = "https://files.pythonhosted.org/packages/f7/e9/e50295e520f7d14ce2b5d900f7a2fe8f096d8f90e9040f8d7f43108b72ce/ddtrace-4.11.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:d0eadcf7df3143bf0611f09a351040e475b5df7f512410698cd481a076bca056", size = 7101986, upload-time = "2026-07-10T08:55:56.51Z" }, - { url = "https://files.pythonhosted.org/packages/cb/c6/e8d13546a4b2cd72880d8b0823dd487956d9eb291473d7d03b6e742e5bd9/ddtrace-4.11.0-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:57eb0616dc5c8ff04b5c560ba6d3def69d2bc618e69ce65cc015f103fc68649e", size = 7455362, upload-time = "2026-07-10T08:55:58.664Z" }, - { url = "https://files.pythonhosted.org/packages/50/c6/2c22d71290df36b87393794d6d803ec8ea96ee54d42286836a7ddf0dbae3/ddtrace-4.11.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:c311109d1e0f3482b7a44126c855eeb796b186f01604ce74c54800aa5892863d", size = 8502274, upload-time = "2026-07-10T08:56:00.869Z" }, - { url = "https://files.pythonhosted.org/packages/19/84/5afef2c9696b289afa1c608e862930c0e447489776a0133da9f2345edd27/ddtrace-4.11.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:695ec915640072e2917b13d317fe1696086a38924ba02cbe1dcb3918ffb0fa65", size = 8730138, upload-time = "2026-07-10T08:56:03.711Z" }, - { url = "https://files.pythonhosted.org/packages/12/68/a2aff5999fa2b1053b3f50cd78297cbbe7a427f63388fb5f6c500dfb859a/ddtrace-4.11.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:355a288756a7a5d6b73221d28ba4536cb0271d2ea7909adab3e2ee37d62a569f", size = 9507981, upload-time = "2026-07-10T08:56:06.317Z" }, - { url = "https://files.pythonhosted.org/packages/79/12/53bb54b3451f85dbf3c752c2d60b22f672fd55bb08d7495613a70dc87102/ddtrace-4.11.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:2c19888af09f6ad3340f2935ea5c9eb80d9298141e196fd33e544369920aefba", size = 9793730, upload-time = "2026-07-10T08:56:08.893Z" }, - { url = "https://files.pythonhosted.org/packages/9c/7e/d9a52dd0d877ebdce5eb105d9f3d4e3579391ee0106c8335d59aa97006dd/ddtrace-4.11.0-cp312-cp312-win32.whl", hash = "sha256:cd524b299bd15ed14192cb041a35e992003e0d3a9f8831062f9f28bcfacb63e2", size = 5734745, upload-time = "2026-07-10T08:56:11.764Z" }, - { url = "https://files.pythonhosted.org/packages/6d/da/bb7c609ebc368cabb01db38d7022f42a39d30207f484a5e18a306bd9b333/ddtrace-4.11.0-cp312-cp312-win_amd64.whl", hash = "sha256:af663120d3b4116965bebbd79ce8e73a85644816a0917999e8164b6204678cb6", size = 6317182, upload-time = "2026-07-10T08:56:14.156Z" }, - { url = "https://files.pythonhosted.org/packages/30/f3/e3cf49ce0a4c377fdea8203b127984e26ac2c9cad0521fab27d61371d568/ddtrace-4.11.0-cp312-cp312-win_arm64.whl", hash = "sha256:f0b50a1739573318b0a81fb56a8b22ed6ae6dac017b89400cbc76c17b439d9ed", size = 5980198, upload-time = "2026-07-10T08:56:16.458Z" }, - { url = "https://files.pythonhosted.org/packages/d3/f8/89f2b9cb3d413c55a3d1f19a90fda76ec45abdc0ce13a7f6d4aaf10bed1a/ddtrace-4.11.0-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:16b9f06c841858debb3de2c4a7ce6acbd6967eaa6ce264dbd0e2a58e376b8583", size = 7094539, upload-time = "2026-07-10T08:56:18.943Z" }, - { url = "https://files.pythonhosted.org/packages/ed/98/f262ad401cdd1a8a2394299d12f31bcc037d0e2dc61d3a843f1c855c293f/ddtrace-4.11.0-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:1cb352d7df16b864f258d4602addf23b8d29d2d416ccc1d0af8d88c614f70a01", size = 7448347, upload-time = "2026-07-10T08:56:21.359Z" }, - { url = "https://files.pythonhosted.org/packages/1f/17/c9123fc263b89934cc1c831b81d90af299d28512740abf008fd963bd006e/ddtrace-4.11.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:54f4cf656aa452bc37b62fa7939c5bddca751ffd4d315b3c925236d4794767ea", size = 8497865, upload-time = "2026-07-10T08:56:24.143Z" }, - { url = "https://files.pythonhosted.org/packages/39/31/00d529b9ca80c77bf3c6aa7370615e2186fc9ead573ef56cef345e6a4dbf/ddtrace-4.11.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:776d99ccb01abecf7dbb4910f4a570b6ab625ed47e00bb6c0159d88d4e2ef1c8", size = 8721481, upload-time = "2026-07-10T08:56:26.983Z" }, - { url = "https://files.pythonhosted.org/packages/41/56/759eab73938820054faa6db6eda607753d821b6322ce7679a2b9f728f8c7/ddtrace-4.11.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:5533ce94e50412565a36ce0c8fc6ff6f7490e69b4ecff10c550482e49d5d065e", size = 9505774, upload-time = "2026-07-10T08:56:29.901Z" }, - { url = "https://files.pythonhosted.org/packages/5d/a4/6ab67386244c08340c1c30affffc644b14095e59eb8d2e8d8ee1d4d46d84/ddtrace-4.11.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:81d147bd3d6fa067f2c72c7e32aba90d5d61ed379850896f1cc7fd1ea3437bb1", size = 9787192, upload-time = "2026-07-10T08:56:32.824Z" }, - { url = "https://files.pythonhosted.org/packages/e3/8d/44e074f927b35007e8bf3dc2b5aeb5ba5e4c97e76d4b8e994aaabddfd648/ddtrace-4.11.0-cp313-cp313-win32.whl", hash = "sha256:bcf1d3ba8c3771b8c75f0b9ecd9038ae0a44c8bb0c01d21458c405a29341518a", size = 5732179, upload-time = "2026-07-10T08:56:35.958Z" }, - { url = "https://files.pythonhosted.org/packages/7e/ec/ff034010a1fd598863a15606aa7867ee6abaf4a5989b1a22e61753229432/ddtrace-4.11.0-cp313-cp313-win_amd64.whl", hash = "sha256:d465b9f6f049fc32278cbfd74e1e593b0c3460aacb622bef0cee71e50ffc4b05", size = 6314643, upload-time = "2026-07-10T08:56:38.565Z" }, - { url = "https://files.pythonhosted.org/packages/bc/dd/ad7da7ca071c76815b16a0523875c0c061902cfd8e9ae6595577464e923e/ddtrace-4.11.0-cp313-cp313-win_arm64.whl", hash = "sha256:00cc0fbfa6b8741477528c3fe7bc70a4a5ea439221f188e14da6907f03ed0125", size = 5976789, upload-time = "2026-07-10T08:56:41.172Z" }, - { url = "https://files.pythonhosted.org/packages/12/1b/37b9a84cc4a95d74d4358a0f78346376bc03625b6a9fa0d568a1587665ca/ddtrace-4.11.0-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:e6f1ca84237e9f52fd71e3458262cfc407dd43337a56e4a6ba9beea88f0958c5", size = 7099136, upload-time = "2026-07-10T08:56:43.766Z" }, - { url = "https://files.pythonhosted.org/packages/f1/42/f227eeb9c5bbb519c67fe11a5326b531f4c1bbd90257e74ab8b72b5ded75/ddtrace-4.11.0-cp314-cp314-macosx_14_0_x86_64.whl", hash = "sha256:c3ddbf9a5a0cd51b5ab9bca898c85cfce73072054b9f5b577164a95065cb8331", size = 7452151, upload-time = "2026-07-10T08:56:46.822Z" }, - { url = "https://files.pythonhosted.org/packages/eb/c4/35fbc0e8ef585d93f82dadae38246eb3643621f08efe11ca0c20e77e7d42/ddtrace-4.11.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:7c435b708f4956208141c22d941ed60a1551813353f6c1cfc1a5283ec3d80cc5", size = 8507170, upload-time = "2026-07-10T08:56:49.646Z" }, - { url = "https://files.pythonhosted.org/packages/70/99/a9681a53e9b0557497b24a4e0b1495acf2bcb74368e86129e549d6249a4a/ddtrace-4.11.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:26b3fd7ee97cf7fa1296eecacd393bbbef343cf8bb2d0b1ee39d40cb1d4818fa", size = 8725097, upload-time = "2026-07-10T08:56:53.106Z" }, - { url = "https://files.pythonhosted.org/packages/e0/c6/b18f02120c34abf55e28bc180e872ddbe27614559c01bbcfe5be35b5494a/ddtrace-4.11.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:8d23d66ba9525fd9b9a61c0db8c43e64461066b1d49c75c85f639ea65776367b", size = 9516895, upload-time = "2026-07-10T08:56:56.513Z" }, - { url = "https://files.pythonhosted.org/packages/c9/ed/dc73e8317c229300b0ff8513f903d16ff140eede8084b3b6f76c8646a548/ddtrace-4.11.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:85e081e2295a77c9b38e5792a097f1df8169e1dd9b3339f820d00183e10e8b1b", size = 9792739, upload-time = "2026-07-10T08:57:00.333Z" }, - { url = "https://files.pythonhosted.org/packages/7c/8e/b97984d9f43e2f36f5317034711140c3bd1fbb5d96880b1c11b9e1e8c3d3/ddtrace-4.11.0-cp314-cp314-win32.whl", hash = "sha256:c267fe4efceb71019f5e6833fdb4fbdff18a51641bdca9fceabbe579984461e4", size = 5831072, upload-time = "2026-07-10T08:57:03.578Z" }, - { url = "https://files.pythonhosted.org/packages/e8/6a/7fb18925588f645d6d777c90d388f9a640adf7b986be7457943717b74170/ddtrace-4.11.0-cp314-cp314-win_amd64.whl", hash = "sha256:3bebf3cb12ac51c422ad1ed74873746b8398cce845075b234115c2f518a7a652", size = 6455126, upload-time = "2026-07-10T08:57:06.362Z" }, - { url = "https://files.pythonhosted.org/packages/a0/b2/44d7f07ba94c70bd1827669d446a2e979f7efd6e4a409c469a59684d9a7b/ddtrace-4.11.0-cp314-cp314-win_arm64.whl", hash = "sha256:f596cb2da582600e091c8174ab4f8e44bf4f05c50c03cc450e0378c516e0c6e3", size = 6130096, upload-time = "2026-07-10T08:57:09.413Z" }, - { url = "https://files.pythonhosted.org/packages/99/22/dadad1a1213a846499f360344c4be0854635bbd9c5726d2902bc8e92b194/ddtrace-4.11.0-cp315-cp315-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:9d9a7c4644e4ffc858e951eeab37ab86e9b717ef1bc2f2a47de69f1ec373bccb", size = 6454958, upload-time = "2026-07-10T08:57:12.714Z" }, - { url = "https://files.pythonhosted.org/packages/d4/45/4a5bb647d47f01d135c918ad33678d1dba5b0210cc688da5d2f69534830e/ddtrace-4.11.0-cp315-cp315-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:b49433f0e85e16c34f578e01ec4e263ab5181d9636635add8d6194bd47dd163c", size = 6638706, upload-time = "2026-07-10T08:57:15.791Z" }, - { url = "https://files.pythonhosted.org/packages/80/2b/175096efe3784f36d383e09d7bfa1defdba85689008417008f122a3e8d0c/ddtrace-4.11.0-cp315-cp315-musllinux_1_2_aarch64.whl", hash = "sha256:2affbac04583808cba374976d07e59a4fe156d1de122f8def30550bf7717b745", size = 7466814, upload-time = "2026-07-10T08:57:18.873Z" }, - { url = "https://files.pythonhosted.org/packages/80/b3/8fb55c9183ed5a162525ae308e0aaa8237ceb2a3ba78d65a8f1336c16b21/ddtrace-4.11.0-cp315-cp315-musllinux_1_2_x86_64.whl", hash = "sha256:51a28ab8072c741e104d8c6f47d990eb5989a359b7ea4d91bab7c5c1f5c1a03c", size = 7707994, upload-time = "2026-07-10T08:57:21.918Z" }, + { url = "https://files.pythonhosted.org/packages/9b/c8/6b8de8b78503f54a6c19f5b4c9430825976f1a2ef1b8342f21ac47c70e22/ddtrace-4.15.6-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:68c18fa37cd718f2360c980bbd2928a5c2bf502a173e61a83f25e42321d5b4b8", size = 7998372, upload-time = "2026-10-08T09:44:43.566Z" }, + { url = "https://files.pythonhosted.org/packages/67/30/b5ec13494b638453c568ced825b863c086fc7237528aa30498b4fcb6928a/ddtrace-4.15.6-cp310-cp310-macosx_14_0_x86_64.whl", hash = "sha256:7d7696cf05e1675847303701c0a2f7733014f1b2cb5938ed2038e27fa32a18b2", size = 8382818, upload-time = "2026-10-08T09:44:45.73Z" }, + { url = "https://files.pythonhosted.org/packages/fa/dc/57c5d150061cf27935a2e630ca5334f3fa815d82f0585ba694821f793445/ddtrace-4.15.6-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:66647b3c9673468e35fe0294d25015ef6f535bb540a9ce63c76c95ff2679e059", size = 9544001, upload-time = "2026-10-08T09:44:47.366Z" }, + { url = "https://files.pythonhosted.org/packages/f5/7b/6651a225df4438761ddd45c21f624d453aca1508d88ed7fd4c880a3a454a/ddtrace-4.15.6-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4fd260a1339a2df4867f5b0f7ff9111a0f23ba86d6b402e4c3dfe2f71853cd17", size = 9794825, upload-time = "2026-10-08T09:44:49.468Z" }, + { url = "https://files.pythonhosted.org/packages/e4/8d/ffd12a1fa866068c397a1fbbebaebcd6e40f21e618650e95d93860afef44/ddtrace-4.15.6-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:95a04ff00e401caef808689ae6bd70a9fc92d3572647f9f1bbd43dfacfa053f6", size = 10533745, upload-time = "2026-10-08T09:44:51.57Z" }, + { url = "https://files.pythonhosted.org/packages/bf/19/64e52eced9a66cbb61dd40b110dc25297b1185f1027016876ae366fcaa1e/ddtrace-4.15.6-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:4eb19b57f1cc6ab8e13e57b8257b91f6b0e5bede753320b30e3920a28d4b1a4f", size = 10833139, upload-time = "2026-10-08T09:44:53.716Z" }, + { url = "https://files.pythonhosted.org/packages/9d/ae/0d964c6eb264ec0ab797071db48221acca6990b5837f1a6d804b88b818ec/ddtrace-4.15.6-cp310-cp310-win32.whl", hash = "sha256:740a87a6504617ec7584d5afb2a6169d780571ddc2b44270274aa8291a70a03c", size = 6756230, upload-time = "2026-10-08T09:44:55.961Z" }, + { url = "https://files.pythonhosted.org/packages/0a/2c/35a05709d630b3ad6f5c05dd90d34bf784fcde93e0722f9e1cec111ec5c4/ddtrace-4.15.6-cp310-cp310-win_amd64.whl", hash = "sha256:7aa9f30da1832111de778a17a26bc2bc564f4810d573ac2d8633609c3db48306", size = 7387627, upload-time = "2026-10-08T09:44:57.786Z" }, + { url = "https://files.pythonhosted.org/packages/2d/b5/3ce5d64ed92d87388a83b52ed766866177997585e224d018922b0494a88d/ddtrace-4.15.6-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:b3f6e66e3b303252edce525ca573e7092e741049556ae80a395eadf0685a036f", size = 7997045, upload-time = "2026-10-08T09:44:59.612Z" }, + { url = "https://files.pythonhosted.org/packages/f7/25/b7f15af30cf446fa95bf6835e515c135cc801ddc9ec18ecd665b7a82063a/ddtrace-4.15.6-cp311-cp311-macosx_14_0_x86_64.whl", hash = "sha256:6bec8671c1e0bdadd7fad2513f449c07b9bf4f847bc652aef2720f0e2f74aeaa", size = 8381818, upload-time = "2026-10-08T09:45:01.325Z" }, + { url = "https://files.pythonhosted.org/packages/d9/95/1c597a597123eb99fccd0fc0ef64ee25ad46460e158dbf258136126134c3/ddtrace-4.15.6-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:3017a4d4fe8375ae4555b80419d5259bdc444af485bcca6ee52dcd1a41890d35", size = 9541087, upload-time = "2026-10-08T09:45:03.274Z" }, + { url = "https://files.pythonhosted.org/packages/55/88/7e8c66d3db4739599a8092197940d211a006edca56b8fbbf9bbd34869a27/ddtrace-4.15.6-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:220fbb49481c6dd85866110477183b70882cb2578020c3bbbb4db2d89d0b41a8", size = 9793899, upload-time = "2026-10-08T09:45:05.384Z" }, + { url = "https://files.pythonhosted.org/packages/84/f8/08c6b74b9fd13df49cfdb95e6a1a86e06c434994db8f4a4a3395f0fee013/ddtrace-4.15.6-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:2d50b28793f381bb53e2d7cab5ffeb97c1aa0032f436a1fd5984cb3740d75b10", size = 10534489, upload-time = "2026-10-08T09:45:07.549Z" }, + { url = "https://files.pythonhosted.org/packages/44/6c/03cd4a5cc888f7c540fa9282d3c7c3a7675aae2c622c41668d1a765e0699/ddtrace-4.15.6-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:b5a9edbd149d793b134bdaea28f40db5a199d42135f122c2be325519db55fb75", size = 10834145, upload-time = "2026-10-08T09:45:09.661Z" }, + { url = "https://files.pythonhosted.org/packages/25/6d/ab22cf125e6618597c41b55682dc76617cb6e071ca7381aac346840d6404/ddtrace-4.15.6-cp311-cp311-win32.whl", hash = "sha256:6b6a59fc3dfb081fd7708c9431b94dd17c5148625ec84bc425852cdfbf2f214d", size = 6752572, upload-time = "2026-10-08T09:45:11.785Z" }, + { url = "https://files.pythonhosted.org/packages/b6/38/ffc989894b1d7ea7c06fb249185807c574da5b4572adb9c974fcf02ea65d/ddtrace-4.15.6-cp311-cp311-win_amd64.whl", hash = "sha256:fef3cdf314eee555c5f6b831aeeb826a7eec147e3e1ab85591154a0ad96828f9", size = 7388927, upload-time = "2026-10-08T09:45:14.111Z" }, + { url = "https://files.pythonhosted.org/packages/ab/b8/1b8483fac4a4bcf1d7338d6323b2a3bde576553b111c5a288aee71224613/ddtrace-4.15.6-cp311-cp311-win_arm64.whl", hash = "sha256:4f163cc7fd9d76711c9df54273cbc2bb15933d90d1027fbf6ea3e938f042d23a", size = 7391090, upload-time = "2026-10-08T09:45:16.749Z" }, + { url = "https://files.pythonhosted.org/packages/c1/3f/010c67d812aa0950bb44aeabec69c858e5bd7d31d9fe560ae431e2a08d95/ddtrace-4.15.6-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:c4ca75ae638e105e5c28b7f7b3e1d3762f027ab14a315530e62c71d402b49dc1", size = 7993463, upload-time = "2026-10-08T09:45:18.638Z" }, + { url = "https://files.pythonhosted.org/packages/16/e3/c42f708e8506a75621678e5aac5219b03c8d95014b70e5755287cbe81d2a/ddtrace-4.15.6-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:2f533ad60bb9855ab47a429eeb8b7e28667647c1a482899ac6bd5c00aa525407", size = 8387610, upload-time = "2026-10-08T09:45:20.667Z" }, + { url = "https://files.pythonhosted.org/packages/a3/c5/ea86a9859f8e9b69a448c1ccbb70c51c7989c8c7bf68288356b9f53883a1/ddtrace-4.15.6-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e00a2943adf9cc886da3cf377b6c32001828565c877a9f0a7905e60eda527dad", size = 9514941, upload-time = "2026-10-08T09:45:22.929Z" }, + { url = "https://files.pythonhosted.org/packages/8d/af/b3bd0fdacb0d2b6baf8bd033ea00e97b06e74e3e504918ecf2f4d176ad6b/ddtrace-4.15.6-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:dfb419eecf9ab3ad774e0fa612f08814530ebd2e1f0aa10128bb0f98eb7cea2e", size = 9772538, upload-time = "2026-10-08T09:45:25.295Z" }, + { url = "https://files.pythonhosted.org/packages/ce/05/f1cb0e9007a5c0d155948726bfefe97fc5fc9d7e886a951a4d51bc1c0236/ddtrace-4.15.6-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:7f8bf477362e29742a6b282fbbfc9c6985c745ba302e4b6a8a53388911ceacae", size = 10511565, upload-time = "2026-10-08T09:45:27.732Z" }, + { url = "https://files.pythonhosted.org/packages/09/90/9f27058af312b213abfb946ab27b334b58b5482aada8b00ce2c41bee2ee6/ddtrace-4.15.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:bd6c6cd99d6edcffab804755f8246f3c1ecf727bdb241ecfe94f5de31fbb484b", size = 10824106, upload-time = "2026-10-08T09:45:30.202Z" }, + { url = "https://files.pythonhosted.org/packages/8b/9a/f4402bdfb563d103ece19bd56a4622accd0113dfb9ac9b118050757823f9/ddtrace-4.15.6-cp312-cp312-win32.whl", hash = "sha256:903a24868a743e035ff0f7bdc7e397a90c00fac77b551765a9ceae5537c40a44", size = 6756938, upload-time = "2026-10-08T09:45:32.864Z" }, + { url = "https://files.pythonhosted.org/packages/55/fb/0b24d9fbcecfeb72ac2fb24db5b4a4684a0f0ecda73bb2b0c56c14282645/ddtrace-4.15.6-cp312-cp312-win_amd64.whl", hash = "sha256:57cd4ca85c5fef665fc2794c37369e39d6b9f51b1ff37ae8e68ac935360d1561", size = 7375054, upload-time = "2026-10-08T09:45:35.014Z" }, + { url = "https://files.pythonhosted.org/packages/b6/34/09763e279fb865196966cc895fb234cf12350f08fde8a284c796f13d6d34/ddtrace-4.15.6-cp312-cp312-win_arm64.whl", hash = "sha256:be11b8c8fce20b6e02d6069a53d43fc92aec9a50741e0f0dadebd913f2917ae5", size = 7376141, upload-time = "2026-10-08T09:45:37.46Z" }, + { url = "https://files.pythonhosted.org/packages/17/0c/f0b7dfb83917ea3f9c0d072cf8cf9454550ecf26372b19be49b8cf16bbcb/ddtrace-4.15.6-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:8a118c85925d143e88e48e4583500e95a9b88d9c1fabc96df276c086eb515603", size = 7987311, upload-time = "2026-10-08T09:45:40.025Z" }, + { url = "https://files.pythonhosted.org/packages/f5/01/f39c527b10ceef714c16a99d59429ddd8b314c7d5301bff86fab6201dfae/ddtrace-4.15.6-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:1307ee87674c3d2d93f38bcd258b62b70a41081183e4bc94b03594c855751322", size = 8382130, upload-time = "2026-10-08T09:45:42.316Z" }, + { url = "https://files.pythonhosted.org/packages/e2/29/0c46ee3d6dcce28f8bd82c1dc658e420bed935fc0488a1ffac2853b42b0f/ddtrace-4.15.6-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:fabb658b1ec068ab6046fb1ecaa6e6860bd40338235244c3c156c973c4957c00", size = 9507940, upload-time = "2026-10-08T09:45:44.887Z" }, + { url = "https://files.pythonhosted.org/packages/b7/b0/77d18cb89ff7b2652d734ae00e4869ef703945a0393bc465e00878631099/ddtrace-4.15.6-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:2185df0b730e4bb78dfe40cf61b7e242d4b81724cea4f496ca5a0ae0f5427309", size = 9764498, upload-time = "2026-10-08T09:45:47.528Z" }, + { url = "https://files.pythonhosted.org/packages/64/06/7435d3534b646ddafaf16fced76b0dec8d239492eef1e119d4c3b3fe1259/ddtrace-4.15.6-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:7fc21e13b52290d33bcdec36644eb63eac03f49adb0b36c6d84dc4bb02b9b935", size = 10508643, upload-time = "2026-10-08T09:45:50.63Z" }, + { url = "https://files.pythonhosted.org/packages/34/d0/65db05d5b72fcb88378edd4bf44e88b2dc0c3e8e1c3070a4630e5ec2dfdd/ddtrace-4.15.6-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:008a5c93db9489861394190b9a9ecffc464a9a83882792c83c005e938fc69df0", size = 10817707, upload-time = "2026-10-08T09:45:53.417Z" }, + { url = "https://files.pythonhosted.org/packages/df/e0/b2ad40242ce8d17ba4432c4f96a7584d93852c442fcb8ac1285cbd10ae96/ddtrace-4.15.6-cp313-cp313-win32.whl", hash = "sha256:d5ca2d38c37194aee97801775193a034f7f27f2610d5075b6a550ee664b044dd", size = 6754044, upload-time = "2026-10-08T09:45:56.603Z" }, + { url = "https://files.pythonhosted.org/packages/1e/e3/7a549bc05e6b604ef3c5ad8fef19a951178d8d93ee7088660d3743e53b7b/ddtrace-4.15.6-cp313-cp313-win_amd64.whl", hash = "sha256:fc35805f2c5b64bd7a29304a5c54932792a47b1436d2550354c051870cd706b4", size = 7370629, upload-time = "2026-10-08T09:45:58.955Z" }, + { url = "https://files.pythonhosted.org/packages/50/dc/a140b6640dd57f143e9137e8e4d04653ba9047c7ef65ff26ca03721c90ff/ddtrace-4.15.6-cp313-cp313-win_arm64.whl", hash = "sha256:03221f056c03728b01a4295ff4a3854db2dc33bcfab3131757f5eb4ca2a5a41b", size = 7378982, upload-time = "2026-10-08T09:46:01.698Z" }, + { url = "https://files.pythonhosted.org/packages/62/de/448c66ce790d6fcd1abbeb08e8e88ef1a5b7fc990f6ec48b818b4d9a8b29/ddtrace-4.15.6-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:70d5dfc73d970fa1b633e0aad3b9457cb94478958976e139b1c78461a3c48b33", size = 7993200, upload-time = "2026-10-08T09:46:04.284Z" }, + { url = "https://files.pythonhosted.org/packages/2f/87/50bbd2d4fe589aa13038c1fe07defd57e63c1bc83c22b92eac13676b3219/ddtrace-4.15.6-cp314-cp314-macosx_14_0_x86_64.whl", hash = "sha256:26ead4c2956924d0574b6681d063dee6c8741dca98ed3b747c353f6252eae0d7", size = 8388799, upload-time = "2026-10-08T09:46:07.071Z" }, + { url = "https://files.pythonhosted.org/packages/9e/13/d947c88b0fdd7cddb89174f5e9116aead05a935e1406e2921c820cb7bcab/ddtrace-4.15.6-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:03e62ec9c83e9e98122f4544592a6bf7d3338e9ea7ff2e34fe4750ade2b2c3e5", size = 9519810, upload-time = "2026-10-08T09:46:09.975Z" }, + { url = "https://files.pythonhosted.org/packages/48/5a/15711bfd54778a28cca9f7988bcdd03e4d680dfceed7f4db22f6b7a06e73/ddtrace-4.15.6-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:3684507243feb45e7438b68989aae9a7df01b7996112d7fe74846e1d5aeb3726", size = 9771503, upload-time = "2026-10-08T09:46:13.009Z" }, + { url = "https://files.pythonhosted.org/packages/7d/d4/d0e81dff83be9370b90c090828330c4b2f06d74644f331d8f6d9518bad44/ddtrace-4.15.6-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:4882e25808cf6102015d7421d6214aa94a788ce3242fec9a1a726a41dcfd79ac", size = 10521211, upload-time = "2026-10-08T09:46:15.867Z" }, + { url = "https://files.pythonhosted.org/packages/f7/c1/e33b44ebfa4dcb80f66fb2cc99c61d9c7706edd48e5ccb5ad04bcea0d5ad/ddtrace-4.15.6-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:5336630ed8c829333dd4195d817b3f3f8211ddf6074e1c731e884771094fc41c", size = 10826453, upload-time = "2026-10-08T09:46:19.084Z" }, + { url = "https://files.pythonhosted.org/packages/e5/32/2125cf7f03e9e8b68b6a9072e393a3be4af91505ada32bc65811d6965881/ddtrace-4.15.6-cp314-cp314-win32.whl", hash = "sha256:a970bc70c81361c9a27452bfd3335182d722315466109208290a1c3a4c221748", size = 6867770, upload-time = "2026-10-08T09:46:22.66Z" }, + { url = "https://files.pythonhosted.org/packages/b4/a0/c18b1db1cf4843a3254aed03b61c1c9cbd5964bd52c18922f633c7ed2304/ddtrace-4.15.6-cp314-cp314-win_amd64.whl", hash = "sha256:e7e8b948d13737e61a8e0dc1fe082d38431a3d4fb285d23236f1c8f733ee5c67", size = 7538347, upload-time = "2026-10-08T09:46:25.41Z" }, + { url = "https://files.pythonhosted.org/packages/8f/77/6181506d7f7693aceccdc13f5a2dc6b18662ed3eedebc02043f00d502e5e/ddtrace-4.15.6-cp314-cp314-win_arm64.whl", hash = "sha256:11662a7a4e1f51c388450cd1092b57738bb32ef98401c6032754a7e220a5f6d0", size = 7574014, upload-time = "2026-10-08T09:46:28.185Z" }, +] + +[[package]] +name = "ddtrace-internal" +version = "0.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/0f/51/6dc02752dca5920d6627fb92b7e0611f594a823bbd57ed31167571d819bc/ddtrace_internal-0.1.0.tar.gz", hash = "sha256:7390721276ab5f174cb103088be7d0de22d452f2817c949e8bde1182ff5b961e", size = 2504, upload-time = "2026-07-15T18:35:10.062Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/93/ad/3cae600fdfc41f54653c14339ac1f125381a6fc2a823d5e685b627ff152e/ddtrace_internal-0.1.0-py3-none-any.whl", hash = "sha256:8eeac4a812e003cebafc8be35ef05103ae2b132893dfbcee00a67f4af81404a2", size = 2495, upload-time = "2026-07-15T18:35:09.295Z" }, ] [[package]] @@ -1948,11 +1974,30 @@ wheels = [ name = "dnspython" version = "2.8.0" source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version < '3.11'", +] sdist = { url = "https://files.pythonhosted.org/packages/8c/8b/57666417c0f90f08bcafa776861060426765fdb422eb10212086fb811d26/dnspython-2.8.0.tar.gz", hash = "sha256:181d3c6996452cb1189c4046c61599b84a5a86e099562ffde77d26984ff26d0f", size = 368251, upload-time = "2025-09-07T18:58:00.022Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/ba/5a/18ad964b0086c6e62e2e7500f7edc89e3faa45033c71c1893d34eed2b2de/dnspython-2.8.0-py3-none-any.whl", hash = "sha256:01d9bbc4a2d76bf0db7c1f729812ded6d912bd318d3b1cf81d30c0f845dbf3af", size = 331094, upload-time = "2025-09-07T18:57:58.071Z" }, ] +[[package]] +name = "dnspython" +version = "2.9.0" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.15'", + "python_full_version == '3.14.*'", + "python_full_version == '3.13.*'", + "python_full_version == '3.12.*'", + "python_full_version == '3.11.*'", +] +sdist = { url = "https://files.pythonhosted.org/packages/ef/4a/50822184bd67cc6493f0fb6a880749158fcd31ab3fa07409acfd91f9fc85/dnspython-2.9.0.tar.gz", hash = "sha256:b44dc6b18f07a8b1c56676a19fbfdb5209415b046a9cece286baafa87ff3f7f1", size = 423560, upload-time = "2026-10-09T00:07:24.352Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/10/02/cdcc9b7c051786a103c3b09e1003a82fa0c66bcb91ffbdabcfbf7b4163b9/dnspython-2.9.0-py3-none-any.whl", hash = "sha256:9a4aedb833c3c1b49214d04d44d3032ab7a9135f7c1d29a549b4ff78fd82fda9", size = 354822, upload-time = "2026-10-09T00:07:22.622Z" }, +] + [[package]] name = "docker" version = "7.1.0" @@ -2009,7 +2054,8 @@ name = "email-validator" version = "2.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "dnspython" }, + { name = "dnspython", version = "2.8.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "dnspython", version = "2.9.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "idna" }, ] sdist = { url = "https://files.pythonhosted.org/packages/f5/22/900cb125c76b7aaa450ce02fd727f452243f2e91a61af068b40adba60ea9/email_validator-2.3.0.tar.gz", hash = "sha256:9fc05c37f2f6cf439ff414f8fc46d917929974a82244c20eb10231ba60c54426", size = 51238, upload-time = "2025-08-26T13:09:06.831Z" } @@ -5329,7 +5375,7 @@ ci = [ { name = "psycopg2-binary", specifier = "==2.9.11" }, { name = "pyarrow", specifier = "==23.0.1" }, { name = "pygithub", specifier = "==2.8.1" }, - { name = "pylint", specifier = "==4.0.5" }, + { name = "pylint", specifier = "==4.0.9" }, { name = "pytest-codspeed", specifier = "==4.3.0" }, { name = "pytest-retry", specifier = "==1.7.0" }, { name = "tenacity", specifier = "==8.5.0" }, @@ -6111,75 +6157,100 @@ wheels = [ [[package]] name = "msgpack" -version = "1.2.1" +version = "1.2.3" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/31/f9/c0a1c127f9049db9155afc316952ea571720dd01833ff5e4d7e8e6352dbb/msgpack-1.2.1.tar.gz", hash = "sha256:04c721c2c7448767e9e3f2520a475663d8ee0f09c31890f6d2bd70fd636a9647", size = 183960, upload-time = "2026-06-18T16:13:52.594Z" } +sdist = { url = "https://files.pythonhosted.org/packages/0a/e7/bb605a7bab2d8425a64b3fa762b39dc1bf1c7e3f11ba6fb5413d6db0ff8c/msgpack-1.2.3.tar.gz", hash = "sha256:32edb81a2b5eb7cd7c9d941b2bfbbb082fd2cd09e0e725930316af6b708db186", size = 196517, upload-time = "2026-09-29T02:33:52.276Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/5b/16/f70100614b69feb3ade7285f08c9c52d6cda0a5c03f3f5e2facd63acb211/msgpack-1.2.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:8c7b398c56ff125feae96c2737abfec5595f1fa0aa186df60c56040b8accb95c", size = 82926, upload-time = "2026-06-18T16:12:31.531Z" }, - { url = "https://files.pythonhosted.org/packages/e4/3c/08ecd5cdfe4e2de43aec79062028ad0f7b2d9b1fea5430068c198ba570da/msgpack-1.2.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:1548006a91aa93c5da81f3bdcebc1a0d10cea2d25969754fbe848da622b2b895", size = 82730, upload-time = "2026-06-18T16:12:32.894Z" }, - { url = "https://files.pythonhosted.org/packages/19/9f/a70c9cb1a04ecc134005149367dcfe35d167284e8f65035a1e4156ad17b5/msgpack-1.2.1-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1dabedcd0f23559f3596428c6589c1cd8c6eaed3a0d720795b07b0225d769203", size = 400729, upload-time = "2026-06-18T16:12:34.052Z" }, - { url = "https://files.pythonhosted.org/packages/fa/7f/5ce020168cf0439041526e95aa068c722c016aee21624e331aeabeee2e8e/msgpack-1.2.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:83efa1c898e0fc5380fc0cabbf75164c52e3b5cbb45973710d75821928380c73", size = 407625, upload-time = "2026-06-18T16:12:35.239Z" }, - { url = "https://files.pythonhosted.org/packages/79/70/fb7668ce0386819303047057aef6fc1da73b584291d9cff82b821744e2ef/msgpack-1.2.1-cp310-cp310-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:01e2dd6c9b19d333a00282330cc8a73d38d8dabc306dc5b42cd668c3ac82e833", size = 377891, upload-time = "2026-06-18T16:12:36.684Z" }, - { url = "https://files.pythonhosted.org/packages/3d/dc/9ebe654a73c3aed2e40aa6b52e3c2a02b5f53ef0085fa235a45d5b367f87/msgpack-1.2.1-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:350cb813d0af6e65d2f7ef0d729f7ff5be5a8bce03665892f43e5883d4ecc1b8", size = 391987, upload-time = "2026-06-18T16:12:37.839Z" }, - { url = "https://files.pythonhosted.org/packages/42/eb/b67cf64218a2fa25e1c671fe1d3dbb06cbeb973e71bc4b822da079862d0b/msgpack-1.2.1-cp310-cp310-musllinux_1_2_riscv64.whl", hash = "sha256:ee1d9ed27d0497b848923746cf762ed2e7db24f4be7eec8e5cbe8c766aa707b7", size = 374603, upload-time = "2026-06-18T16:12:39.221Z" }, - { url = "https://files.pythonhosted.org/packages/a2/2e/9ee200cde32fd1a0101b4006202fde554c1860adfb9bf7bff31ea4c08df8/msgpack-1.2.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:633727297ed063441fd1cda2288865487f33ad14eeb8831afb5f0c396a62cfce", size = 405121, upload-time = "2026-06-18T16:12:40.524Z" }, - { url = "https://files.pythonhosted.org/packages/43/b6/f10117be7ca7a51e8feed699a907b8e663a8cd66e115ae6b4fb30cc7945c/msgpack-1.2.1-cp310-cp310-win32.whl", hash = "sha256:298872ecf9e61950f1c6af4ca969b859ee91783bb920ef6e6172697d0c8aad74", size = 64088, upload-time = "2026-06-18T16:12:41.762Z" }, - { url = "https://files.pythonhosted.org/packages/ba/93/89976c696fb0224662239d952c47b4d1661b34d79a332ef5584facaa8579/msgpack-1.2.1-cp310-cp310-win_amd64.whl", hash = "sha256:2ff164c1b0bcb740b073b99e945234d0212852fa378e44a208c425379140dbeb", size = 70113, upload-time = "2026-06-18T16:12:42.78Z" }, - { url = "https://files.pythonhosted.org/packages/f4/6b/e9b1cdc042c4458801d2545ed782a95f3d6ba8e270cce8745b8603c7f748/msgpack-1.2.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:29a3f6e9667868429d8240dfd063ea5ffdc1321c13d783aa23827a38de0dcb22", size = 82812, upload-time = "2026-06-18T16:12:45.022Z" }, - { url = "https://files.pythonhosted.org/packages/0c/3a/dd518a1bf78ed1e9ad8afe57307c079a00eafe4b3068932a27ca1ea56b4f/msgpack-1.2.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:aded5bdf32609dc7987a49bbbd15a8ef096193f96dd8bbeb791de729e650acf5", size = 82739, upload-time = "2026-06-18T16:12:46.025Z" }, - { url = "https://files.pythonhosted.org/packages/70/e0/7ba9e1542bf0771a27b8b37c1316e3f95ae9d748fd765284655c476ad4ef/msgpack-1.2.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:146ee4e9ce80b365c6d4c47073da9da7bcec473e58194ceee5dd7620ace77e06", size = 414233, upload-time = "2026-06-18T16:12:47.029Z" }, - { url = "https://files.pythonhosted.org/packages/03/8d/671d81534ea0e2b0e8a121be100020da09eb78861fe3aa8f3ef7dcd3bed1/msgpack-1.2.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a28d076ca7c82b9c8728ad90b7147489449557038bed50e4241eb832395169b4", size = 423843, upload-time = "2026-06-18T16:12:48.19Z" }, - { url = "https://files.pythonhosted.org/packages/d2/b6/e5c737515ed1f166664b87601b532f58cbb73d8aa6a90b99f7c2c5037e8e/msgpack-1.2.1-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:7d31c0ac0c640f877804c67cb2bc9f4e23dc2db97e96c2e67fa27d38283b41f8", size = 390772, upload-time = "2026-06-18T16:12:49.624Z" }, - { url = "https://files.pythonhosted.org/packages/a8/46/62ed8c2e87d7021eab19921594d961ef3aa3794eec76c716dc30f3bfd433/msgpack-1.2.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:8ff92d7feeaf5bc26c51495b69e2f99ed97ab79346fb6555f44be7dd2ac6503b", size = 409559, upload-time = "2026-06-18T16:12:50.936Z" }, - { url = "https://files.pythonhosted.org/packages/70/ff/59aa3887b860bbf43532835e192b1c388a17590d6068ae4f8b2bc74c906e/msgpack-1.2.1-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:779197a6513bab3c3632265e3d0f7cb3227e62510841a6f34f1eaa37efbb345e", size = 387838, upload-time = "2026-06-18T16:12:52.161Z" }, - { url = "https://files.pythonhosted.org/packages/09/11/f8563e471093420cf6478cb3271a0175d8402b82d879783d4035d2d03360/msgpack-1.2.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:67f6dd22fa72a93752643f07889796d62739a13415ee630169a8ce764f86cf9f", size = 421732, upload-time = "2026-06-18T16:12:53.556Z" }, - { url = "https://files.pythonhosted.org/packages/57/cf/e673683c4c6c90c1022b24c65af4b03eda72b182a1176ef6449069d66acc/msgpack-1.2.1-cp311-cp311-win32.whl", hash = "sha256:91054a783328e0ea7954b8771095705c8d2243b814743fbaadf14552c9c52c5d", size = 64091, upload-time = "2026-06-18T16:12:54.821Z" }, - { url = "https://files.pythonhosted.org/packages/3f/07/ca212739d179f9083bff2c7c08c24101c3555a334fadc2b876b18768a3ae/msgpack-1.2.1-cp311-cp311-win_amd64.whl", hash = "sha256:2eda0b7ebb1283a98d3e4492ac933c8af6aff59fd3df1c3ed024f536af4b1dc8", size = 70462, upload-time = "2026-06-18T16:12:55.898Z" }, - { url = "https://files.pythonhosted.org/packages/6d/be/6798347b425e26f35db82e69dd83c09716c856a3714e7bffc4c0860fd830/msgpack-1.2.1-cp311-cp311-win_arm64.whl", hash = "sha256:6ee967f7c7e1df2890c671ff2ee51a28ded0efc95da3e507176dee881ce36c66", size = 65059, upload-time = "2026-06-18T16:12:57.053Z" }, - { url = "https://files.pythonhosted.org/packages/bc/dd/9e8cbd8f5582ca4b590336f2b91ee5662f6a6ca562b565abaf696a0f81ff/msgpack-1.2.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:2ef59c659f289eddf8aa6623823f19fa2f40a4029266889eac7a2505dd210c35", size = 83531, upload-time = "2026-06-18T16:12:58.249Z" }, - { url = "https://files.pythonhosted.org/packages/50/2e/ebdb85a8da151397a2790363676b7ed7c125924fe618e4c6d8befb0cc62c/msgpack-1.2.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d3567748a5107cb40cdf66a275430c2f87c07777698f4bfd25c35f44d533258c", size = 82657, upload-time = "2026-06-18T16:12:59.396Z" }, - { url = "https://files.pythonhosted.org/packages/26/aa/753ad8b007b464e1d8aa0c8e650b9c5f4f725e658fc5ac8a7635c55b7f6e/msgpack-1.2.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:60926b75d00c8e816ef98f3034f484a8bc64242d66839cef4cf7e503142316a0", size = 410634, upload-time = "2026-06-18T16:13:00.383Z" }, - { url = "https://files.pythonhosted.org/packages/6a/fd/6adabd4f6d5e686f97dd02ce7fce3fe4cf672cbac36b8f67ff4040e8ad8b/msgpack-1.2.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:020e881a764b20d8d7ca1a54fc01b8175519d108e3c3f194fddc200bda95951a", size = 419989, upload-time = "2026-06-18T16:13:01.776Z" }, - { url = "https://files.pythonhosted.org/packages/5a/cc/85039b7b0eb168aaad7383a23c97e291a11f08351cb45a606ce865e4e3f1/msgpack-1.2.1-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:4202c74688ca06591f78cb18988228bd4cca2cc75d57b60008372892d2f1e6e6", size = 377544, upload-time = "2026-06-18T16:13:03.637Z" }, - { url = "https://files.pythonhosted.org/packages/ed/bf/35963899493b32030c85fc513b723ae66144ac70c11ebc52e889e16e3d99/msgpack-1.2.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:8b267ce94efb76fbd1b3373511420074ee3187f0f7811bf394531de13294735a", size = 400842, upload-time = "2026-06-18T16:13:05.012Z" }, - { url = "https://files.pythonhosted.org/packages/a6/df/8e2ac970c8f99264cd9997d1c73df5466bc19da3301d7dc5500862a9b089/msgpack-1.2.1-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:e4f1d0f8f98ade9634e01fb704a408f9336c0a8f1117b369f5db83dc7551d8b1", size = 374108, upload-time = "2026-06-18T16:13:06.232Z" }, - { url = "https://files.pythonhosted.org/packages/17/dd/fa8bd265110dfa51c20cb529f9e6d240a16fafe7e645004c6af2d01353ba/msgpack-1.2.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:f02cf17a6ca1abe29b5f980644f7551f94d71f2011509b26d8625ce038f0df64", size = 414939, upload-time = "2026-06-18T16:13:07.478Z" }, - { url = "https://files.pythonhosted.org/packages/2e/b9/8377a5ad8953fc0437c70cc98d9ae29f27fe5ac5109fbec0812085865735/msgpack-1.2.1-cp312-cp312-win32.whl", hash = "sha256:0c0d9802354507bcba62af19c17918e3eb437cc25e6f50657d511b5856a77aac", size = 64504, upload-time = "2026-06-18T16:13:08.822Z" }, - { url = "https://files.pythonhosted.org/packages/57/7f/ce1e377df7e62461fefd9eb23bfb93a4a523f40a517b377b8f844d836828/msgpack-1.2.1-cp312-cp312-win_amd64.whl", hash = "sha256:5c24aa15d5963051e1a5c62b12c50cd705992502b5ec1f3bece6046f33c9fc24", size = 71421, upload-time = "2026-06-18T16:13:09.828Z" }, - { url = "https://files.pythonhosted.org/packages/8f/32/ebfe84c9929f08f188d56c7a2fd913406a9ddad76a634697c1c43b8112e6/msgpack-1.2.1-cp312-cp312-win_arm64.whl", hash = "sha256:4227224aaec8f7fbcbfbd4272319347b2bb4030366502600f8c45588c5187b07", size = 64775, upload-time = "2026-06-18T16:13:11.056Z" }, - { url = "https://files.pythonhosted.org/packages/b0/ac/dcddcab6f6c20ecb387ca5e980371cdb3f87ff69aeca388be97eebc4c074/msgpack-1.2.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:0a70e3cf2804a300d921bb0940426e35f4e489a23adfb77a808892241db0a064", size = 83151, upload-time = "2026-06-18T16:13:12.173Z" }, - { url = "https://files.pythonhosted.org/packages/64/71/fbcfa83a1d6a9c6091942d1cfd070962244664b87427a9a49a6897b1b219/msgpack-1.2.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:491cc39455ca765fad51fb451bf2915eb2cf41192ab5801ce8d67c1d614fe056", size = 82351, upload-time = "2026-06-18T16:13:13.194Z" }, - { url = "https://files.pythonhosted.org/packages/e3/10/ddf7b06db879e8792d13934ddda09ff20bd2a583fd84c9b59aae9b0e650b/msgpack-1.2.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f310233ef7fb9c14e201c93639fe5f5260b005f56f0b29048e999c30935596cc", size = 407518, upload-time = "2026-06-18T16:13:14.233Z" }, - { url = "https://files.pythonhosted.org/packages/79/d3/36a46a8ed992b781acbc05928bd5bee3c810cb0c3563bf81a7b0c04a1a76/msgpack-1.2.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:787c9bebb5833e8f6fc8abca3c0597683d8d87f56a8842b6b89c75a5f3176e2d", size = 416405, upload-time = "2026-06-18T16:13:15.435Z" }, - { url = "https://files.pythonhosted.org/packages/f9/84/e8e9598b557c0ba6ddae901a73780a4c75ac667dddf59414b1e56a42fb34/msgpack-1.2.1-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:dc871b997a9370d855b7394465f2f350e847a5b806dd38dcc9c989e7d87da155", size = 376257, upload-time = "2026-06-18T16:13:17.022Z" }, - { url = "https://files.pythonhosted.org/packages/40/16/738fe6d875ad7e2a9429c165322a4ec088f4f273cdfae63d96a89c467961/msgpack-1.2.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:85f57e960d877f2977f6430896191b04a21f8901b3b4baf2e4604329f4db5402", size = 397469, upload-time = "2026-06-18T16:13:18.287Z" }, - { url = "https://files.pythonhosted.org/packages/ca/be/6d5952df75a7f24f35833af764c3a6860780364cb3a0030beb8099e1b2b4/msgpack-1.2.1-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:1233ee2dd0cefba127583de50ea654677277047d238303521db35def3d7b2e7c", size = 372802, upload-time = "2026-06-18T16:13:19.685Z" }, - { url = "https://files.pythonhosted.org/packages/e1/39/e2ef7dbf0473bcb8dc7c50bf782a892d67414877b63e47fc88eb189ef5e6/msgpack-1.2.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:e3dc2feb0876209d9c38aa56cb1de169bd6c4348f1aa48271f241226590993e6", size = 411273, upload-time = "2026-06-18T16:13:21.028Z" }, - { url = "https://files.pythonhosted.org/packages/ef/c5/133f4512a56e983a93445c836c9d94d88f3bc2e0980ff4b9e577bd8416ce/msgpack-1.2.1-cp313-cp313-win32.whl", hash = "sha256:6d09badf350af2be9d189184e04e64cf54ad93569ab3d96fca58bd3e84aad707", size = 64471, upload-time = "2026-06-18T16:13:22.293Z" }, - { url = "https://files.pythonhosted.org/packages/e2/98/577e10b055096a7dd40732358cabaf7180a20c79ed1dcdbb618e4b9deac7/msgpack-1.2.1-cp313-cp313-win_amd64.whl", hash = "sha256:33f14fba63278b714efe6ad07e50ea5f03d91537aa6a1c5f1ceca4cf44013ca9", size = 71274, upload-time = "2026-06-18T16:13:23.455Z" }, - { url = "https://files.pythonhosted.org/packages/ba/ee/0c0048e7cfbef23c6a94791b8959ab28155232e7956de8a305b5ff588f05/msgpack-1.2.1-cp313-cp313-win_arm64.whl", hash = "sha256:afc5febcd4c99effbc02b528e49d6fd0760b2b7d48c05239e345a5fa6e743d9a", size = 64795, upload-time = "2026-06-18T16:13:24.687Z" }, - { url = "https://files.pythonhosted.org/packages/77/58/cce442852c6b9e1639c7c8ac8fd9143121cb32dab0f308df4d1426a8eb9c/msgpack-1.2.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:05f340e47e7e47d2da8db9b53e1bb1d294369e9ef45a747441309f6650b8351d", size = 83610, upload-time = "2026-06-18T16:13:25.724Z" }, - { url = "https://files.pythonhosted.org/packages/60/5c/15b4c7a0182f75ffa90751958ba36a9c01cafee367d49a3edc10ed140b01/msgpack-1.2.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:810b916696c86ef0deb3b74588480224df4c1b071136c34183e4a2a4284d7ac7", size = 83138, upload-time = "2026-06-18T16:13:26.781Z" }, - { url = "https://files.pythonhosted.org/packages/b8/a6/99e58722feaffc5f2fbcc0c8c0d1451ab9f84097f7af87291b46af2390f4/msgpack-1.2.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ca0dacff965c47afdc3749a8469d7302a8f801d6a28758d55120d75e66ce6889", size = 406090, upload-time = "2026-06-18T16:13:28.072Z" }, - { url = "https://files.pythonhosted.org/packages/19/03/8c63e8cf52958534ef688625965ab04c269a6cadd8caef16758b380a821a/msgpack-1.2.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0e2bf9280bceb5efca998435904b5d3e9fdbcc11d90dc9df30aec7973252b720", size = 412106, upload-time = "2026-06-18T16:13:29.427Z" }, - { url = "https://files.pythonhosted.org/packages/63/d2/155d9e71b40e41fd934bc0c48b9b2770f22263e1ac20aad8e29fdca7be3f/msgpack-1.2.1-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:aa6c4be5d1c02a42b066ca6ddb71adf36432868fdcdb6ee87e634e86e0674190", size = 374851, upload-time = "2026-06-18T16:13:30.631Z" }, - { url = "https://files.pythonhosted.org/packages/98/48/deaf2326262a8d5ea3295ce9649912ecd3f551ba7ec8e33c665d2ba583f3/msgpack-1.2.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:ec0e675d59150a6269ddc9139087c722292664a37d071a849c05c473350f1f2d", size = 396168, upload-time = "2026-06-18T16:13:31.977Z" }, - { url = "https://files.pythonhosted.org/packages/10/2a/b4410f906c2ec0008f1608d3ab5143afc3ad3f4e6da0fed3ea2231d0bef4/msgpack-1.2.1-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:dd3bfe82d53edfe4b7fc9a7ec9761e23a7a5b1dac22264505af428253c29ed24", size = 371959, upload-time = "2026-06-18T16:13:33.282Z" }, - { url = "https://files.pythonhosted.org/packages/59/86/1edc67270099a528fa2093ea60fe191233cd238e4bd30cfacf7db79fc959/msgpack-1.2.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:5ad5467fc3f68b5468e06c5f788d712e9f8ffc8b0cd1bcb160c105c1ee92dae7", size = 408457, upload-time = "2026-06-18T16:13:34.567Z" }, - { url = "https://files.pythonhosted.org/packages/82/90/8b630fef07d8c5ab457b71ff2c217910c83d333c7a68472c186e87cc504a/msgpack-1.2.1-cp314-cp314-win32.whl", hash = "sha256:98b58bdb89c46190e4609bb36abe17c6d4105ad13f9c5f8f6f64d320f8ced3fb", size = 65942, upload-time = "2026-06-18T16:13:36.056Z" }, - { url = "https://files.pythonhosted.org/packages/16/f1/467b81e98b24dd3885d7b1857728797b4ffc76a7a7483af4fb321a07de3c/msgpack-1.2.1-cp314-cp314-win_amd64.whl", hash = "sha256:74847557e28ce71bd3c438a447ca90e4b507e997ddbdef8a12a7b283b86c156b", size = 72627, upload-time = "2026-06-18T16:13:37.079Z" }, - { url = "https://files.pythonhosted.org/packages/a7/1d/5d8c4c89985feb6acefb82a09e501c60392261856d2408d20bfe4f0360b1/msgpack-1.2.1-cp314-cp314-win_arm64.whl", hash = "sha256:b50b727bd652bdc37d950336c848ef20ec54a4cafc38dce19b1cd86ad625d0f7", size = 66908, upload-time = "2026-06-18T16:13:38.23Z" }, - { url = "https://files.pythonhosted.org/packages/1b/02/ad2afb678b4de94496cd432b581759b756a92c1192d8c767edd6b132efdc/msgpack-1.2.1-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:8d00f177ca88a77c1cf848d204a38f249751650b601cb6532acc68805d8a8273", size = 86000, upload-time = "2026-06-18T16:13:39.44Z" }, - { url = "https://files.pythonhosted.org/packages/54/74/0b797484013128837f3b1cbb6cea019277c4de4e377dc512b4d9a0f92940/msgpack-1.2.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:5bb9c386f0a329c035ddbab4b72d1028bf9627add8dda41070288563d57ed1b1", size = 86544, upload-time = "2026-06-18T16:13:40.447Z" }, - { url = "https://files.pythonhosted.org/packages/a9/b4/b774d7eb95561739907fec675582f83203cf41c597a418c2589b4bfb8e9d/msgpack-1.2.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:20466cca18c49c7292a8984bc15d65857b171e7264bdcb5f96baf8be238791fc", size = 427661, upload-time = "2026-06-18T16:13:41.574Z" }, - { url = "https://files.pythonhosted.org/packages/b2/f9/3243191dc9937e00756c8bc1b0272fed8f23758e43df2a3b46f533e5090f/msgpack-1.2.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:196300e7e5d6e74d50f1607ab9c06c4a1484c383cd22defd727902591f7e8dde", size = 426375, upload-time = "2026-06-18T16:13:42.936Z" }, - { url = "https://files.pythonhosted.org/packages/23/c7/1693111db9944ba4ad4b67a1e788400d78a0b6af7a6523dc7e4e58f8274b/msgpack-1.2.1-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:575957e79cd51903a4e8495a242442949641e08f1efd5197b43bebd3ea7682b4", size = 380495, upload-time = "2026-06-18T16:13:44.306Z" }, - { url = "https://files.pythonhosted.org/packages/3e/2b/92f86956a0c13e8662f7e2ad630c4eb4db07497b967589bd5245e018b2c1/msgpack-1.2.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:8c2ed1e48cc0f460bf3c7780e7137ff21a4e18433451916f2442c1b21036cd7d", size = 410897, upload-time = "2026-06-18T16:13:45.629Z" }, - { url = "https://files.pythonhosted.org/packages/da/ea/1479f72d200313a76fc2f823a79d1e07ed052ab7b8a0280640aa7b95de42/msgpack-1.2.1-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:5f6277e5f783c36786a145e0247fc189a03f35f84b251646e53592d2bc12b355", size = 378519, upload-time = "2026-06-18T16:13:46.998Z" }, - { url = "https://files.pythonhosted.org/packages/f5/4d/fa006060ffa1011d32bfae826fe766fe73e02982183601633b7121058ab3/msgpack-1.2.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:f9389552ecf4784886345ead0647e4edc96bee37cbab05b75540f542f766c48c", size = 419815, upload-time = "2026-06-18T16:13:48.205Z" }, - { url = "https://files.pythonhosted.org/packages/2f/e1/aab6c946570496b78e67804721f3d5e2d62a93081b9b37df77764ef56347/msgpack-1.2.1-cp314-cp314t-win32.whl", hash = "sha256:c1c79a604a2969a868a78b6ebd27a887e00c624f14f66b3038e0590cb23332d1", size = 70914, upload-time = "2026-06-18T16:13:49.385Z" }, - { url = "https://files.pythonhosted.org/packages/13/0a/e608956488a2af014cfe6e3d665e090b8ee42aa14b07f8f95b8880d66b09/msgpack-1.2.1-cp314-cp314t-win_amd64.whl", hash = "sha256:f12038a35fabd52e56a3547bab42401af49a45caa6dd00b34c44de235bc93ee2", size = 77999, upload-time = "2026-06-18T16:13:50.467Z" }, - { url = "https://files.pythonhosted.org/packages/d2/8a/27e2e57055176e366a46b85d02d68e7a5bcfbdd8474c9706375d965f24d3/msgpack-1.2.1-cp314-cp314t-win_arm64.whl", hash = "sha256:0adcf06ffde0777c0e1a9b771a2b1c4226ba1bbf748c8efcc02fcdeca3299107", size = 71160, upload-time = "2026-06-18T16:13:51.498Z" }, + { url = "https://files.pythonhosted.org/packages/6d/aa/5b6b09f835791045282dc5d08431db599a5f4743a69fe2f6670045a2cd85/msgpack-1.2.3-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:ec0030361cc861ac699b2ef1c695b741fa145c88f8667fa3d7e3f73deeb648a3", size = 90927, upload-time = "2026-09-29T02:31:28.286Z" }, + { url = "https://files.pythonhosted.org/packages/c9/91/7b288e9133bd1ba92ca0ca4e7f2a4cfc53cf467d99d8d2f57b9939908fac/msgpack-1.2.3-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:5c1efdd9181cb1b719ee46865f368a927f1c0c65d577798340b1194545b7515a", size = 89798, upload-time = "2026-09-29T02:31:30.028Z" }, + { url = "https://files.pythonhosted.org/packages/71/9b/5c3dbc450d14645dcec987970692d6ab24008cc33d2155474b1d818486f9/msgpack-1.2.3-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c309a7abae1d14ba29a8bd0ddbd704a5e469d8e9bd9c3dee0e4ff53d7ae01d56", size = 450687, upload-time = "2026-09-29T02:31:32.407Z" }, + { url = "https://files.pythonhosted.org/packages/2b/21/ea60a8fd0d9e0897fce823e9fd9bf6742567784b35c7eee8f4a18a56eb19/msgpack-1.2.3-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5bf390259cb25a6a1cd197c65810999b811f64cd38683251538bcc5a1e41f7d3", size = 459808, upload-time = "2026-09-29T02:31:34.282Z" }, + { url = "https://files.pythonhosted.org/packages/ee/f7/42140e6afdac8e94bfedae4cfb67ee004b6ad5c4cadd024df42f759bf3b5/msgpack-1.2.3-cp310-cp310-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:39b6986c19e1f2dfa549d185dba6ccf1de2e4c0ba10d8cfc0048935b1c5f9109", size = 423845, upload-time = "2026-09-29T02:31:35.713Z" }, + { url = "https://files.pythonhosted.org/packages/19/7b/cd54f27b59dfbdc438a12361fbb6798b66d377a978f946bc9512598290e9/msgpack-1.2.3-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:fcc6800daac4922960f6eeb7a0dda3dd4105e0bf7bce0e83ebc465a78cb7bdba", size = 445608, upload-time = "2026-09-29T02:31:37.65Z" }, + { url = "https://files.pythonhosted.org/packages/57/38/52bc0dc44cc9f7c2339b632f93d02f8badc78cfb0bb070f2a50a51945e53/msgpack-1.2.3-cp310-cp310-musllinux_1_2_riscv64.whl", hash = "sha256:968583e956d0427878050b371308c5f8647088732ef3e66a117dbe1192ec91e0", size = 421721, upload-time = "2026-09-29T02:31:39.151Z" }, + { url = "https://files.pythonhosted.org/packages/89/e6/451c9a42274fb2be82d8ba8b76a5219c613e20f8de1da521d10cb758a9ef/msgpack-1.2.3-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:1d6bcec3dbbdb89ca385d3a73e63ceae7b841fa0d7ca7c676f1a7bfe7fb2cdb8", size = 460430, upload-time = "2026-09-29T02:31:40.843Z" }, + { url = "https://files.pythonhosted.org/packages/57/bb/663e3100327b58caaa5fb66379e557a2717dac08bb586f22f885756bee47/msgpack-1.2.3-cp310-cp310-win32.whl", hash = "sha256:a6b63917d60d6df451f328bd6afba8565e33c4afe1f62ec4ad758b78731c827b", size = 67987, upload-time = "2026-09-29T02:31:42.157Z" }, + { url = "https://files.pythonhosted.org/packages/28/7a/a00d5d7abc5601099260e0d0af8fadc54fbfac2191315aa56eaee3641d9d/msgpack-1.2.3-cp310-cp310-win_amd64.whl", hash = "sha256:4c0780095871ecc49a58b2ff6b1b43b25214704da67646557ca287a3f49fb2dd", size = 75572, upload-time = "2026-09-29T02:31:43.544Z" }, + { url = "https://files.pythonhosted.org/packages/2a/95/b9c651ccb9d720b2e2c8d537954dff528ab869a03bf89598145716db823c/msgpack-1.2.3-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:ec90a9ae3e1169fa1171147340f0e97d941aa19fcd3b34e8339a55933ed042af", size = 90404, upload-time = "2026-09-29T02:31:44.826Z" }, + { url = "https://files.pythonhosted.org/packages/50/cd/fc9e2e367e80f1493e2ec5f610dda558b344eeede296f88976db133e8f2c/msgpack-1.2.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:9d7e9cbb0998bbfd363fd9a09c330520d5e9cb323c05b5a1a05865d23ccf2226", size = 89683, upload-time = "2026-09-29T02:31:46.413Z" }, + { url = "https://files.pythonhosted.org/packages/19/9e/1028485c6886c1c117f777cc9b053e541eff0fedb3292dfb1da95040edb5/msgpack-1.2.3-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6707d2fa2aa1bb5424ea0b05f44ffc989b15ab41a73ff5855bff4944fec7c8ac", size = 465347, upload-time = "2026-09-29T02:31:47.934Z" }, + { url = "https://files.pythonhosted.org/packages/aa/83/800570e6a22376eb8d599920f70aead4779a63611696f567477c4e85a70f/msgpack-1.2.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:382b219de3d436de3baba0f4b0c6d4336e8f5858d0eb047918b13b69a71c6c55", size = 477820, upload-time = "2026-09-29T02:31:49.479Z" }, + { url = "https://files.pythonhosted.org/packages/ab/ff/817e4a2052f848d3fb67726908d6e4e7c19f68ee7c19553a82ce7b0ed415/msgpack-1.2.3-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:186e6c602b8a9968b8e864c67d622a69279f7d1e55ae25f40e3bff7e815b2b62", size = 436656, upload-time = "2026-09-29T02:31:51.18Z" }, + { url = "https://files.pythonhosted.org/packages/3d/42/040cc55dde6a7d92057baac8d1fc9cfb9f4fd4162900e2ec16dc33917a7d/msgpack-1.2.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:9276ba88891338f2617044429dfd080ae008c9868a25f6f1a7d004a35dc9ac0a", size = 460939, upload-time = "2026-09-29T02:31:53.026Z" }, + { url = "https://files.pythonhosted.org/packages/09/93/4dc007bdef930eed247346773bc0189b710078961d3218d5ee7ba59f322c/msgpack-1.2.3-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:c942c21a93f36b3a69e828c8945bb72c94dc2ffe488a2086950c812f3edf046c", size = 433608, upload-time = "2026-09-29T02:31:54.981Z" }, + { url = "https://files.pythonhosted.org/packages/c0/97/a1b944046f283ec89445cb2a982c42233b5b07cc630f9be739f4f1d469a3/msgpack-1.2.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:18a6ed513023001b28dcd3ba54966f6bb90a38274ba8d2640464bcab3a1b81d4", size = 477373, upload-time = "2026-09-29T02:31:56.713Z" }, + { url = "https://files.pythonhosted.org/packages/59/79/ab411d0d172743732ab2503f4c32a22dd1a7d1436a6feecbb160e4b6376a/msgpack-1.2.3-cp311-cp311-win32.whl", hash = "sha256:d0238cd05dec9ffbe0de1071df685ba63e30a36ac155285b1a094e727c38cbe9", size = 67514, upload-time = "2026-09-29T02:31:58.267Z" }, + { url = "https://files.pythonhosted.org/packages/63/8d/6f0cb2b84e484e96278455c26870196d025bb0cec312b226a663f1fa9000/msgpack-1.2.3-cp311-cp311-win_amd64.whl", hash = "sha256:30e1522e4173230dca4d9ad896f038f73c0da6c1edd42f4dbad88ac583cf5d46", size = 75850, upload-time = "2026-09-29T02:31:59.449Z" }, + { url = "https://files.pythonhosted.org/packages/aa/25/f99e13a2c1d3f5a1dcaa5aab27f474e8c4358188bbc68ad79fecb0d1aefe/msgpack-1.2.3-cp311-cp311-win_arm64.whl", hash = "sha256:8ca67f77938ea6a3663aa9bd22b3e031f6da84d665be850abab910ee90728dfd", size = 72338, upload-time = "2026-09-29T02:32:00.885Z" }, + { url = "https://files.pythonhosted.org/packages/af/12/4d7c6d6203416d9fbf0f59ebaa805e70fb929b93a41b611bc821ec5964a0/msgpack-1.2.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:89c930aece4e972b208ba589c8410b4167b05e411a5ea2cb25fd96f8bc47ee43", size = 91577, upload-time = "2026-09-29T02:32:02.141Z" }, + { url = "https://files.pythonhosted.org/packages/eb/c7/8576ad39f4ca42ddad26f68eb8621d2d0a60501193d480f504bd9d7f36c4/msgpack-1.2.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:905a189853d6bdb204c7ae5f4ab77fb857448abfff574d3d93c62e2815b24b4f", size = 90027, upload-time = "2026-09-29T02:32:03.508Z" }, + { url = "https://files.pythonhosted.org/packages/0a/3a/aa9c580aea1314529a0f3562461479780b0d254b064f0880956bfbcc74a8/msgpack-1.2.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f3d7b3d0018746b5997dd6b14a1870b07cc4c327d9101145d94a1fc264a51a06", size = 460343, upload-time = "2026-09-29T02:32:04.906Z" }, + { url = "https://files.pythonhosted.org/packages/3a/cf/9c2e4d6c179529d5bf4a64cff76fa581486569e9fbdd35bd98f51cb624bf/msgpack-1.2.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ede33b2892ceb976283e009ad12fa1834cfdf1f9c43ee9c97849fc588d00a618", size = 472998, upload-time = "2026-09-29T02:32:06.69Z" }, + { url = "https://files.pythonhosted.org/packages/7b/41/915c81fe6df2d3cbdb0dece4f1a5cd313e1cd2abd9f501d0f50c0582517e/msgpack-1.2.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:666ef5601ab0e6e345e47febc96aa81143cc932201543480cbb9499164f05ffb", size = 423216, upload-time = "2026-09-29T02:32:08.739Z" }, + { url = "https://files.pythonhosted.org/packages/a2/e7/7dda8b1039abfd9bba4c5068172c67135c9e33089f503512db9226f23c24/msgpack-1.2.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:87cf2ef05ff2f2493ba29fcdaef27e960ca64dacfd13460ae29e6f92e0ed05bb", size = 451218, upload-time = "2026-09-29T02:32:10.517Z" }, + { url = "https://files.pythonhosted.org/packages/16/5b/ce995c1ed4a0522b7f2d034bc2034fd63005f240b945961b70fb56fbaf3d/msgpack-1.2.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:b774ff994d844e541439ac5d2d49a14def4104830c3465e9394c153f86200ffb", size = 422453, upload-time = "2026-09-29T02:32:11.956Z" }, + { url = "https://files.pythonhosted.org/packages/d2/3f/ce191fb87e2650d0166b34c437e499ee4a7f9db9c1eb164f41725eb6160e/msgpack-1.2.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:eaf7e82249837e3aa97297b34a0bb9ff562027381631e057cea6e1367f10b438", size = 469003, upload-time = "2026-09-29T02:32:13.663Z" }, + { url = "https://files.pythonhosted.org/packages/42/35/539123407fe200fb16609c835675496fbeb6017ace9fc93909f0613223ae/msgpack-1.2.3-cp312-cp312-win32.whl", hash = "sha256:7c047250096f9fc19dba26e3d1639b5e7a84114003605c94def667149a70ced1", size = 68303, upload-time = "2026-09-29T02:32:15.02Z" }, + { url = "https://files.pythonhosted.org/packages/6f/4c/331b45f9b86fbda6b9e103244d189068e51f726d8c40021ed66e1f2c415e/msgpack-1.2.3-cp312-cp312-win_amd64.whl", hash = "sha256:3ec409b0d6aa8e9eec6eaf881b893caa215dbe68c5319ca96e8a271d81bb111d", size = 76744, upload-time = "2026-09-29T02:32:16.344Z" }, + { url = "https://files.pythonhosted.org/packages/13/9f/fb572dc42b9fac06c7ea848aaee6e140d84469743bd1402bc07089fc4566/msgpack-1.2.3-cp312-cp312-win_arm64.whl", hash = "sha256:59612b4ed48a04cf024584218e813562f3b30a3bafa5f55abe300b15da314751", size = 71580, upload-time = "2026-09-29T02:32:17.617Z" }, + { url = "https://files.pythonhosted.org/packages/1f/8b/3824d65e912e925d09ce30d9130fa9970d6d2855d7888b13639a6604967f/msgpack-1.2.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:21bfa4d2aa0b04c1806ef778a1199e9e53ea2441bcbf284420a32083896320b8", size = 91728, upload-time = "2026-09-29T02:32:18.949Z" }, + { url = "https://files.pythonhosted.org/packages/05/e6/df7f2c9ebb94760113debbcea2bd3afe5fdab88a4f7bec1b618755517460/msgpack-1.2.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:db84203b13aecc222f465061397fdd5b53b7ae73d2c95ffc1c8dc5be0153a709", size = 89955, upload-time = "2026-09-29T02:32:20.224Z" }, + { url = "https://files.pythonhosted.org/packages/08/6a/e5fc57136e8bacccb2b39627dea2cd546540a06181e22fe6db90e15b3ae4/msgpack-1.2.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5e0d7950ca3c1bbae291d0552dd3bb2792fc680629c4c0d44e47e5bab969f3ca", size = 454930, upload-time = "2026-09-29T02:32:21.771Z" }, + { url = "https://files.pythonhosted.org/packages/b0/30/c394d37898db9212d1693456cdf363c7e1a097d0b63e10664007f3df3ec1/msgpack-1.2.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:07c9733089d1b176c3dd2f7fa268452f9d5d784d076473499d754a58e8d1fbbb", size = 466866, upload-time = "2026-09-29T02:32:23.742Z" }, + { url = "https://files.pythonhosted.org/packages/4a/c8/1e4ddf6f6b829b3ee6c530c79dfae89cb609d2b0eedb5e0ae716851c52d1/msgpack-1.2.3-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:f24a43b3560e20f825b807fe1e874bd73d53abaf8bbdcf258a6eb152cddbc1f5", size = 418715, upload-time = "2026-09-29T02:32:25.262Z" }, + { url = "https://files.pythonhosted.org/packages/11/a5/f460ba6d7a12d4301002f3efbb8f841e8bdc9c5fc98d771689677a352885/msgpack-1.2.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:6576f348ed6cc4f31db6fd915a8e94245f042f50eae08d48732425e70638ea37", size = 446489, upload-time = "2026-09-29T02:32:26.988Z" }, + { url = "https://files.pythonhosted.org/packages/49/23/adface88db909bed321c85dd673655152d4a514c67e1f0800eb51c777d07/msgpack-1.2.3-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:cd5a9f9f86a52c24713679aa2631956835f3842512964ff93f736ff76f1f530d", size = 416998, upload-time = "2026-09-29T02:32:28.606Z" }, + { url = "https://files.pythonhosted.org/packages/36/00/5bb3a239ccfc3763c4d0fa49b13b1b7010b00182c499ab3c1fecfe6294bc/msgpack-1.2.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:f9ddd28d3e9bbc602a9dced1591882c7fb9ab776eef8837da2c326fde19e2853", size = 463288, upload-time = "2026-09-29T02:32:30.375Z" }, + { url = "https://files.pythonhosted.org/packages/29/8c/456df77f00d701df9d6980ffb80291bce6e4e2e112e25a4dfae216f0715a/msgpack-1.2.3-cp313-cp313-pyemscripten_2025_0_wasm32.whl", hash = "sha256:62cc1a4ef0e553bac32c8342e1f04834aca7de276b92744eb7307db77759b890", size = 53347, upload-time = "2026-09-29T02:32:31.867Z" }, + { url = "https://files.pythonhosted.org/packages/9d/22/ce780be666f89b77cdb855daa9ec62e87bb7f69e9f403e4a5d83a2b2208f/msgpack-1.2.3-cp313-cp313-win32.whl", hash = "sha256:d2f9c4f85e47a44d26d5baf3b041eef23436e224d44eed273f01bd8a12048d9f", size = 68258, upload-time = "2026-09-29T02:32:33.163Z" }, + { url = "https://files.pythonhosted.org/packages/51/06/c3def9bc4db283103c5901b302ee2a4305cb1e69729244f94d9bd8f8e8e7/msgpack-1.2.3-cp313-cp313-win_amd64.whl", hash = "sha256:bb89b5dc30469c84bbf8684826eb851d82412ca95690e111b9ac5e8fb343961a", size = 76569, upload-time = "2026-09-29T02:32:34.412Z" }, + { url = "https://files.pythonhosted.org/packages/12/9f/cef344073858b80adb92d6ea342e20b0eae7a8f6fe70281b69cf03707270/msgpack-1.2.3-cp313-cp313-win_arm64.whl", hash = "sha256:471e12a6a42498a31490c206e0069e343b6a7c35db540be73a879eb06f5be047", size = 71530, upload-time = "2026-09-29T02:32:35.892Z" }, + { url = "https://files.pythonhosted.org/packages/3f/8e/f777f74e38731c428857933c8011596f2d2f3160c821152f23b6ffba862f/msgpack-1.2.3-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:3a31905206722103a84c1f72633fe30692cff6732c9d262e09a27dbc468797c8", size = 92042, upload-time = "2026-09-29T02:32:37.464Z" }, + { url = "https://files.pythonhosted.org/packages/a0/71/551608543ee5d590f7e8d522267665d6d9946866ad2a2a70a770f7c70793/msgpack-1.2.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:3372475211a9ce1a23acefe512cb3e121d18c95dc74ed56cb1819ef40836ebf4", size = 90578, upload-time = "2026-09-29T02:32:38.883Z" }, + { url = "https://files.pythonhosted.org/packages/ea/11/6d78ce5a9a58bf9ba7b1b6a8f649173b030e6770c8019cf330b91825ee5d/msgpack-1.2.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9324c54995641c3d1f92a9d55093c8cde0ffa2fbc87a467a688ef60428393220", size = 454352, upload-time = "2026-09-29T02:32:40.34Z" }, + { url = "https://files.pythonhosted.org/packages/3d/08/feb9a196269ba7809f44f9117d9e4a601c41c313f6144fd0c337293a5488/msgpack-1.2.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d8ef3a66e4b52d2d7fdd90df2984670124b2ff7546d76bb25dcf68ef47f7df58", size = 462562, upload-time = "2026-09-29T02:32:42.176Z" }, + { url = "https://files.pythonhosted.org/packages/f5/77/3a674f366def24140b103d1ffd4fd27b3d912a13e47da67422afa16bebb3/msgpack-1.2.3-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:902f3490db0e07a7d40b48536a85c9b28fbf1397e7e1658a45a55f958e303620", size = 418134, upload-time = "2026-09-29T02:32:43.693Z" }, + { url = "https://files.pythonhosted.org/packages/48/82/944e71f280577490d99a3951cbce21aa4cbe04e7ab42cb373fd668af883c/msgpack-1.2.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:8e51eca14fbb65c4e0a5a9657346962bd3dca78c08e04e3d4dee70ef48687d30", size = 445937, upload-time = "2026-09-29T02:32:45.739Z" }, + { url = "https://files.pythonhosted.org/packages/b1/ec/feddd629c4a3edf1395313680450c525086cceab56dec0d4de9da9ccb618/msgpack-1.2.3-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:f42f146752eedb6765f07dcc04d72dab0a25779ec8d4a88c0085263ce114f22c", size = 416450, upload-time = "2026-09-29T02:32:47.558Z" }, + { url = "https://files.pythonhosted.org/packages/e4/59/263a10f8c4613ba0713f48cbda7695ac8dd6d6fab2fcbc9168f03f23a94d/msgpack-1.2.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:0ed5823c4efc20fe87d3530665f40ec18a002be003114814c21235cc8d256207", size = 459546, upload-time = "2026-09-29T02:32:49.145Z" }, + { url = "https://files.pythonhosted.org/packages/1e/21/addcfa1e583cfc8a22fbdc57526621b5decd7ad676ae12e9150b7be1be5d/msgpack-1.2.3-cp314-cp314-pyemscripten_2026_0_wasm32.whl", hash = "sha256:2487453ca1b6104442c6442f9a1a8fee1fe8f428a70d99d4cba799108b304150", size = 53462, upload-time = "2026-09-29T02:32:50.708Z" }, + { url = "https://files.pythonhosted.org/packages/8d/2c/3cb5c8524a1335ee27ca952c7ab78d375a16fea8e18ae3767ba0c880416c/msgpack-1.2.3-cp314-cp314-win32.whl", hash = "sha256:6df430419f2338cb71e4a34d6e64f83c88ccd321f91f40ba4513400b36d864ec", size = 70294, upload-time = "2026-09-29T02:32:52.037Z" }, + { url = "https://files.pythonhosted.org/packages/23/f9/9172ff3cdb85d160ad06df5e2708a5fce7682982a5eee8d31869b9f69d2e/msgpack-1.2.3-cp314-cp314-win_amd64.whl", hash = "sha256:84a6616d396ec1bc18a1e83e67c96a393ec35dfe5e17434a5be7b9aa0fe988ab", size = 77778, upload-time = "2026-09-29T02:32:53.429Z" }, + { url = "https://files.pythonhosted.org/packages/04/e8/b4c23178bcf605ae17cec48a75530dd69d49b0a5a6f5f4df5c47d59f746e/msgpack-1.2.3-cp314-cp314-win_arm64.whl", hash = "sha256:7a003b02c6ee2eea6dfe0bb08818631e3597e69f0131f2a8250488a1cc553290", size = 73794, upload-time = "2026-09-29T02:32:54.763Z" }, + { url = "https://files.pythonhosted.org/packages/66/b1/92704be352c4f428b7e0a0e0fb210cb1aa2b1c42c102b8dc22d34b82fac0/msgpack-1.2.3-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:ccea05b5542f6d283fef3f0a8e93a7f0be90af0ddeeef84c25c0216ba76dcae1", size = 93721, upload-time = "2026-09-29T02:32:56.342Z" }, + { url = "https://files.pythonhosted.org/packages/49/78/9c91f1e86cadcbc100b3780fd429c3715648704032a612e77a00646ebe79/msgpack-1.2.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:b1631e12fe572e181cd77e831f69335d6cd5278eac22e3db3f33cf264ac2ac18", size = 94256, upload-time = "2026-09-29T02:32:58.056Z" }, + { url = "https://files.pythonhosted.org/packages/91/4d/270f9725921ae88a29d37a774a77ac24f0ef1411fc960a63f5a4665e81b4/msgpack-1.2.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e54394b7dbe2e12ab032d9d21feef7bb61a90a150a2623633ba3781ba69dcb1f", size = 471673, upload-time = "2026-09-29T02:32:59.886Z" }, + { url = "https://files.pythonhosted.org/packages/48/b8/eaa8d930f72dc1d1dd79511dc2ccf965922b059f2f0ed3b30aebac8c4b11/msgpack-1.2.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:63bb7448a1e9111319ae2430c09a5596140c160422830d6271bc75730ff2ff9a", size = 466257, upload-time = "2026-09-29T02:33:01.517Z" }, + { url = "https://files.pythonhosted.org/packages/5b/5a/97adc805037bc7e24c4e2f711bbcd3b28be8ec9aea3e778f18208cfbdb46/msgpack-1.2.3-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:382bc88fe90f29f5ac8a0b65c7046ff255356f2f2f3186c30e370215736fa1dc", size = 418484, upload-time = "2026-09-29T02:33:03.402Z" }, + { url = "https://files.pythonhosted.org/packages/0d/7e/1c53302606fe436ab48ba539ebafafe4a6a9efe12c4f04dc7eb36912d93e/msgpack-1.2.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:c77e27790ad72989db783d5303825fba0b71550f00a490efba35cde7dc4b719f", size = 454064, upload-time = "2026-09-29T02:33:04.977Z" }, + { url = "https://files.pythonhosted.org/packages/00/2d/9ee0170f638907b396c15c6cd26b3e54f869159efc6206683acfd8f696e1/msgpack-1.2.3-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:700bc0fc9e968a292b9137ee70e7a012f7e115bf0107ce45e3a88202788dfc1e", size = 417901, upload-time = "2026-09-29T02:33:06.489Z" }, + { url = "https://files.pythonhosted.org/packages/cc/d2/905c84490a75cd15a27065407cd085d201f7d392e1e0411f49f03fd31ade/msgpack-1.2.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:5bd5f91ea75c45cafcc5433ba8fae59b708b736ec178d2441c40c499e9e079db", size = 459896, upload-time = "2026-09-29T02:33:08.361Z" }, + { url = "https://files.pythonhosted.org/packages/37/cd/4ce5809b9ab3b114d7cca64863e436820fa1614b49d55ccb93d49824ac2d/msgpack-1.2.3-cp314-cp314t-win32.whl", hash = "sha256:7995a7c6a62a1d6e7df211b4a16de513bd99fd053525050a319f80f44fb8015e", size = 75983, upload-time = "2026-09-29T02:33:10.023Z" }, + { url = "https://files.pythonhosted.org/packages/8a/31/853bb580744c24be0dbd8b090c3e6987dce466a1fc840fe50c0ac2ef9044/msgpack-1.2.3-cp314-cp314t-win_amd64.whl", hash = "sha256:bfe7d5b62cbe7aa664f0b3e2c49077f10fcdd06183d3014f8271ff3c5edbfbf9", size = 83757, upload-time = "2026-09-29T02:33:11.441Z" }, + { url = "https://files.pythonhosted.org/packages/0d/49/9f1b2ee484414eef9e21ee2b2b23b482bb71433ab9bac1da03cbda15ebf5/msgpack-1.2.3-cp314-cp314t-win_arm64.whl", hash = "sha256:1f585407f740a9eac04a3bb82c61d68a0ea78f90e29e670bfb086b9ce3a518dd", size = 78128, upload-time = "2026-09-29T02:33:13.063Z" }, + { url = "https://files.pythonhosted.org/packages/47/b8/50db4235407c3802f622b4ccdf65c6fe1e48d3c3eab6981fa6a9a5e53f11/msgpack-1.2.3-cp315-cp315-macosx_10_15_x86_64.whl", hash = "sha256:13221a6c81ebb8e43ea63a7251c35d54e4175cea37ebf3a62e911bdf42562a3c", size = 92111, upload-time = "2026-09-29T02:33:14.476Z" }, + { url = "https://files.pythonhosted.org/packages/15/56/50cf2a45c6163edafd737e2fd555103a26ce6748e1e241fb56ed445ea835/msgpack-1.2.3-cp315-cp315-macosx_11_0_arm64.whl", hash = "sha256:0955b9000725573d1457c1676944b370dd9643c8d18f25bda5ac72913f850949", size = 90583, upload-time = "2026-09-29T02:33:15.924Z" }, + { url = "https://files.pythonhosted.org/packages/2a/fd/8cc02f767c3bc94d2649c954d28dea935ce9398eb9c93ce2444bb9474cc1/msgpack-1.2.3-cp315-cp315-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0c91762c48cd686dc9cf2b142c0bc544083952de32f5853d6624c956e54b85e5", size = 454751, upload-time = "2026-09-29T02:33:17.475Z" }, + { url = "https://files.pythonhosted.org/packages/80/c9/ddb896767808e3e022453d8dfae26fd52ed404b0aa6fb7f752d39c040208/msgpack-1.2.3-cp315-cp315-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1f4ae8bd4ad9ba085fde95e95d055a896d19210238a4199a771a3cf36dceed49", size = 463597, upload-time = "2026-09-29T02:33:19.309Z" }, + { url = "https://files.pythonhosted.org/packages/4d/a5/e7c261abf75783c07dcac89951cb31dd0c123bf02fbdeda0c67303e698d8/msgpack-1.2.3-cp315-cp315-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:7013534a7163aa4f213c4d9864f1a8a7555daac6fcd48f699a198e29b436bfab", size = 422661, upload-time = "2026-09-29T02:33:21.093Z" }, + { url = "https://files.pythonhosted.org/packages/9d/8e/466d5133f9e1c2e232e15e304f715b62f6f0e28332d18e37d975fe174315/msgpack-1.2.3-cp315-cp315-musllinux_1_2_aarch64.whl", hash = "sha256:6a834097144aabe948b8ca9020a833e8026f7d0abbd0ec54bc7e50f45a8ce012", size = 445188, upload-time = "2026-09-29T02:33:22.877Z" }, + { url = "https://files.pythonhosted.org/packages/d4/b4/33e7ad987ee2f4b3d449a6cbf28f574ed222987ca7f65ad277072646ac5e/msgpack-1.2.3-cp315-cp315-musllinux_1_2_riscv64.whl", hash = "sha256:d31864ba3933a589b6a00249f89c0eb422197f49128fc10da550e57e9cb0f377", size = 420451, upload-time = "2026-09-29T02:33:24.485Z" }, + { url = "https://files.pythonhosted.org/packages/34/2c/9d8be0d6c16e7e6131cd7da20257dd3da65473e3e6df0c00572fb10a195c/msgpack-1.2.3-cp315-cp315-musllinux_1_2_x86_64.whl", hash = "sha256:e15f70588f4db8cd10df0930145b186de70feb9db51710cd378b1399009655bd", size = 460624, upload-time = "2026-09-29T02:33:26.063Z" }, + { url = "https://files.pythonhosted.org/packages/6a/e7/3a04783582c6f44f398cbfcf5f07a111192126ec4e63edf7f5640143bf64/msgpack-1.2.3-cp315-cp315-pyemscripten_2026_5_wasm32.whl", hash = "sha256:b949cc25e4a09252cbcc54e66e507de914d0e94a3a7039bd54c299bf7037c098", size = 53474, upload-time = "2026-09-29T02:33:27.83Z" }, + { url = "https://files.pythonhosted.org/packages/68/fb/db07359851644e258609d84f8e4fe0030ef448c108e20afe73f2a3bf539c/msgpack-1.2.3-cp315-cp315-win32.whl", hash = "sha256:8ec7a1d49ca6c2569d722ab5ec86e90089b0713900aa31905b47b4c4d9e78ce0", size = 70344, upload-time = "2026-09-29T02:33:29.382Z" }, + { url = "https://files.pythonhosted.org/packages/5b/e4/cf5584d2f2a2e4465d5896a855a3e75a34a20ab172360b3d42ad862dd1ce/msgpack-1.2.3-cp315-cp315-win_amd64.whl", hash = "sha256:79dfa38faf92f804aa61beec140d70b18418e1dde1778dbb77a87a4cce85aa8a", size = 77800, upload-time = "2026-09-29T02:33:30.941Z" }, + { url = "https://files.pythonhosted.org/packages/63/f9/518ad4e8a580027b507eafdd26de7aae661a714e43d7c111c212482e4a1b/msgpack-1.2.3-cp315-cp315-win_arm64.whl", hash = "sha256:ed899d73a22f286a72bd9528d63f2ab3030dbad8bf1527fc249319a50d61fb9d", size = 73871, upload-time = "2026-09-29T02:33:32.406Z" }, + { url = "https://files.pythonhosted.org/packages/a4/79/254d4c9ad642b2a3ba84e646787892b34cc815eb36c9976f67a1c4f38515/msgpack-1.2.3-cp315-cp315t-macosx_10_15_x86_64.whl", hash = "sha256:f56fba61b2516be7917cb00151f0d060b5b21184e3499bb57f0f7d9259bea124", size = 93370, upload-time = "2026-09-29T02:33:33.87Z" }, + { url = "https://files.pythonhosted.org/packages/3d/6f/5a2ba167646a25e84eaa8894e12935351e4331b80c28a9237ce6fe8d375f/msgpack-1.2.3-cp315-cp315t-macosx_11_0_arm64.whl", hash = "sha256:69ad12cedb674c73527bed869cddb42b742cac79a207a614202a4abaa24ea173", size = 93959, upload-time = "2026-09-29T02:33:35.503Z" }, + { url = "https://files.pythonhosted.org/packages/e9/a1/2b44612e55f7cf5d5e4b580294959b4429bbbcb1991177888e3e18668137/msgpack-1.2.3-cp315-cp315t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:db9fb67a3a2e75247bae569d34ebb5ff61c0448a4f0d6dbf991dae68af39b007", size = 467921, upload-time = "2026-09-29T02:33:37.023Z" }, + { url = "https://files.pythonhosted.org/packages/0b/6e/3309798ed1c11d7fcfdc7b946642685b0ff1588477925bc0d26bee7dcaae/msgpack-1.2.3-cp315-cp315t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2574ef81c1c8c38b10e330f3f9406fd09198a776b002030fafcf8e7647e9e06e", size = 467310, upload-time = "2026-09-29T02:33:38.799Z" }, + { url = "https://files.pythonhosted.org/packages/6f/79/9c799f489fa4146de4e00cfe9fee17afe33d8012f88ddffffea94f7c4700/msgpack-1.2.3-cp315-cp315t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:fafc3b8898b432b841d30a61082c599fa7f4d06885f9dc58ad72259e12059fa6", size = 420178, upload-time = "2026-09-29T02:33:40.781Z" }, + { url = "https://files.pythonhosted.org/packages/94/c6/5850dc9cafcd2ea315692e65db0e222d20923dd55f44adf35061003de27e/msgpack-1.2.3-cp315-cp315t-musllinux_1_2_aarch64.whl", hash = "sha256:a393e428f6ffb0dcb73308c1fff5593041c16ff42da66e5bac8a83a6107a54b0", size = 450248, upload-time = "2026-09-29T02:33:42.366Z" }, + { url = "https://files.pythonhosted.org/packages/a9/d2/b4c806e3497fe21f0b353568266aec14ff735d092aea672de7b2955db03f/msgpack-1.2.3-cp315-cp315t-musllinux_1_2_riscv64.whl", hash = "sha256:d1c1e8989a855b7f1f2a64ec4a80b23a631822903952770813857b2e4f460471", size = 418431, upload-time = "2026-09-29T02:33:44.178Z" }, + { url = "https://files.pythonhosted.org/packages/b0/f5/f4ecc3ddac4d551bf2f3cdb283ec546dcc826fe7c500074be61aa273e08a/msgpack-1.2.3-cp315-cp315t-musllinux_1_2_x86_64.whl", hash = "sha256:e0bd394e999949c814f7912284243298de1b5a17b6a3dcb6cc8a79b156ffc4fa", size = 457543, upload-time = "2026-09-29T02:33:45.978Z" }, + { url = "https://files.pythonhosted.org/packages/a4/69/1c821d8386fae5cecc5fcaacf3de3947ff0a23f16bb481b5532b5868372a/msgpack-1.2.3-cp315-cp315t-win32.whl", hash = "sha256:3d4c807ed050fe3ddbea5ba7e9f63d7136871ce42861be1f50ff739f0e91047a", size = 75820, upload-time = "2026-09-29T02:33:47.596Z" }, + { url = "https://files.pythonhosted.org/packages/68/9e/41e2f7343a3764a9c1fb10c79f9a6a05db9df93dedd76401d1b511f5a685/msgpack-1.2.3-cp315-cp315t-win_amd64.whl", hash = "sha256:5f304123b90e8b2e49867981b7f6061612c39f50cca51ee88de007c084cf68d3", size = 83345, upload-time = "2026-09-29T02:33:49.325Z" }, + { url = "https://files.pythonhosted.org/packages/80/cd/0c3aa439bc7a7bf24684fef3a0ad776cba170e18ed94445e723bce42fce7/msgpack-1.2.3-cp315-cp315t-win_arm64.whl", hash = "sha256:f41ca154b7737b11893cdce3c78c61d703398a1cd54d4297bdad908392338a8e", size = 77572, upload-time = "2026-09-29T02:33:50.729Z" }, ] [[package]] @@ -8743,7 +8814,7 @@ crypto = [ [[package]] name = "pylint" -version = "4.0.5" +version = "4.0.9" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "astroid" }, @@ -8755,9 +8826,9 @@ dependencies = [ { name = "tomli", marker = "python_full_version < '3.11'" }, { name = "tomlkit" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/e4/b6/74d9a8a68b8067efce8d07707fe6a236324ee1e7808d2eb3646ec8517c7d/pylint-4.0.5.tar.gz", hash = "sha256:8cd6a618df75deb013bd7eb98327a95f02a6fb839205a6bbf5456ef96afb317c", size = 1572474, upload-time = "2026-02-20T09:07:33.621Z" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/5e/d5ec2c861010cf8577a3e5baaf2405fdf1004f1657d9245c4efb006c1dfd/pylint-4.0.9.tar.gz", hash = "sha256:6f1305780229ae720e89b8aff2b3ff857cbb268fbab98c8b586960267667ad34", size = 1604187, upload-time = "2026-09-23T15:26:43.144Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/d5/6f/9ac2548e290764781f9e7e2aaf0685b086379dabfb29ca38536985471eaf/pylint-4.0.5-py3-none-any.whl", hash = "sha256:00f51c9b14a3b3ae08cff6b2cdd43f28165c78b165b628692e428fb1f8dc2cf2", size = 536694, upload-time = "2026-02-20T09:07:31.028Z" }, + { url = "https://files.pythonhosted.org/packages/66/46/12e94ae1f9db89152223bdb17920b12c79e94ac51e19dce1da9fb63627b0/pylint-4.0.9-py3-none-any.whl", hash = "sha256:c4ba5bbaf1e70ecf05f08cd7d6091ea04824b79a2b2d9f5b10dc25fc34cb1124", size = 542068, upload-time = "2026-09-23T15:26:41.626Z" }, ] [[package]] diff --git a/vscode-extension/package-lock.json b/vscode-extension/package-lock.json index 81e9425ef4e..5f6cf7edf4e 100644 --- a/vscode-extension/package-lock.json +++ b/vscode-extension/package-lock.json @@ -2920,9 +2920,9 @@ } }, "node_modules/qs": { - "version": "6.16.0", - "resolved": "https://registry.npmjs.org/qs/-/qs-6.16.0.tgz", - "integrity": "sha512-h6fhOIaRrID2CbEY2fqs+7t+UXZo+MLAnU5gRIq85uFtdiUPCdsApMlHhXogKVM4HM2DVbIjGNTTYH2OcmP1vA==", + "version": "6.17.0", + "resolved": "https://registry.npmjs.org/qs/-/qs-6.17.0.tgz", + "integrity": "sha512-5LK5BZ2XDHMvcUJaf0CQ5HL3aVgKbpklvByC7qLSxYkfE8COEEbJg5xCHWCTqzqi0Gs4Jn5Mv5v3R38W6edscA==", "dev": true, "license": "BSD-3-Clause", "dependencies": { From cfc651dcd5c4daa9f3e12a97b56e7da066680fd0 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Wed, 28 Oct 2026 07:32:41 -0700 Subject: [PATCH 038/128] fix(http): honor transport settings for Rust WebSockets (#47968) * fix(http): honor transport settings for Rust WebSockets The Responses WebSocket connected through tokio-tungstenite with its own native-roots TLS, so it ignored the configured CA bundle, verify flag, client certificate, proxies and timeouts. The handshake now goes through a pooled reqwest client (ClientVariant::WebSocket, HTTP/1.1, no redirects) and upgrades the response, so it shares the settings every other outbound request uses The upgrade response is checked for the Sec-WebSocket-Accept key, the Upgrade and Connection headers and a subprotocol that was offered Co-Authored-By: Claude Opus 5.5 * fix(http): preserve native WebSocket trust and optional protocols * fix(http): allow plain WebSockets without native roots --------- Co-authored-by: Claude Opus 5.5 Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 7 +- litellm-rust/Cargo.toml | 4 +- litellm-rust/crates/http/Cargo.toml | 1 + litellm-rust/crates/http/src/config.rs | 1 + litellm-rust/crates/http/src/pool.rs | 18 +- litellm-rust/crates/http/src/tls.rs | 6 + litellm-rust/crates/http/src/websocket.rs | 179 ++++-- litellm-rust/crates/http/tests/websocket.rs | 550 ++++++++++++++++-- .../crates/inference-responses/Cargo.toml | 1 + .../inference-responses/src/websocket.rs | 49 +- .../inference-responses/tests/responses.rs | 1 + .../src/routes/responses/websocket.rs | 12 +- 12 files changed, 699 insertions(+), 130 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index d37b65be7a3..71ed05c0adf 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3987,6 +3987,7 @@ dependencies = [ "tempfile", "thiserror 2.0.19", "tokio", + "tokio-rustls 0.26.4", "tokio-tungstenite", "tracing", "veil", @@ -7425,11 +7426,7 @@ checksum = "edc5f74e248dc973e0dbb7b74c7e0d6fcc301c694ff50049504004ef4d0cdcd9" dependencies = [ "futures-util", "log", - "rustls 0.23.42", - "rustls-native-certs", - "rustls-pki-types", "tokio", - "tokio-rustls 0.26.4", "tungstenite", ] @@ -7774,8 +7771,6 @@ dependencies = [ "httparse", "log", "rand 0.8.7", - "rustls 0.23.42", - "rustls-pki-types", "sha1 0.10.7", "thiserror 1.0.69", "utf-8", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 95f2da8c073..745b29dd0ff 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -100,10 +100,10 @@ rstest = "0.26.1" wiremock = "0.6.5" rstest_reuse = "0.7.0" rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] } +rustls-native-certs = "0.8" rustify = "=0.7.0" rustify_derive = "=0.5.5" vaultrs = { version = "=0.8.0", default-features = false, features = ["rustls"] } -rustls-native-certs = "0.8" serde = { version = "1.0", features = ["derive"] } serde_json = { version = "1.0", features = ["float_roundtrip"] } serde_with = { version = "=3.21.0", default-features = false, features = ["std", "macros"] } @@ -118,7 +118,7 @@ thiserror = "2.0" tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] } tiktoken-rs = "0.12.0" tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"] } -tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] } +tokio-tungstenite = { version = "0.24", default-features = false, features = ["handshake"] } futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] } base64 = "0.22" flate2 = "1" diff --git a/litellm-rust/crates/http/Cargo.toml b/litellm-rust/crates/http/Cargo.toml index 3b46825ad9d..0552f9775fd 100644 --- a/litellm-rust/crates/http/Cargo.toml +++ b/litellm-rust/crates/http/Cargo.toml @@ -29,5 +29,6 @@ webpki-roots.workspace = true futures-util.workspace = true rcgen = "0.14.10" tempfile.workspace = true +tokio-rustls = { version = "0.26", default-features = false, features = ["ring"] } rstest.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/http/src/config.rs b/litellm-rust/crates/http/src/config.rs index 6a2d80e3f85..b37ea6559b8 100644 --- a/litellm-rust/crates/http/src/config.rs +++ b/litellm-rust/crates/http/src/config.rs @@ -16,6 +16,7 @@ pub enum Verify { Disabled, CaBundle(PathBuf), BuiltInRoots, + NativeRoots, } #[derive(Clone, Debug, PartialEq, Eq, Hash)] diff --git a/litellm-rust/crates/http/src/pool.rs b/litellm-rust/crates/http/src/pool.rs index dfcc9270fe9..bf892b09921 100644 --- a/litellm-rust/crates/http/src/pool.rs +++ b/litellm-rust/crates/http/src/pool.rs @@ -6,7 +6,12 @@ use std::{ use reqwest::dns::Resolve; -use crate::{client::Client, config::HttpClientConfig, error::Error, proxy::EnvironmentProxies}; +use crate::{ + client::Client, + config::{HttpClientConfig, Verify}, + error::Error, + proxy::EnvironmentProxies, +}; #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] pub enum ClientVariant { @@ -14,6 +19,7 @@ pub enum ClientVariant { NoRedirect, Media, UnpinnedMedia, + WebSocket, } const CLIENT_TTL: Duration = Duration::from_secs(3600); @@ -59,6 +65,14 @@ impl HttpClientPool { client_certificate: None, ..config.clone() }, + ClientVariant::WebSocket => HttpClientConfig { + http2: false, + verify: match &config.verify { + Verify::BuiltInRoots => Verify::NativeRoots, + verify => verify.clone(), + }, + ..config.clone() + }, ClientVariant::Provider | ClientVariant::NoRedirect => config.clone(), }; let key = (effective, variant); @@ -91,7 +105,7 @@ impl HttpClientPool { ) -> reqwest::ClientBuilder { match variant { ClientVariant::Provider => builder, - ClientVariant::NoRedirect | ClientVariant::UnpinnedMedia => { + ClientVariant::NoRedirect | ClientVariant::UnpinnedMedia | ClientVariant::WebSocket => { builder.redirect(reqwest::redirect::Policy::none()) } ClientVariant::Media => builder diff --git a/litellm-rust/crates/http/src/tls.rs b/litellm-rust/crates/http/src/tls.rs index 0b1c44bf6fe..a52ea31af59 100644 --- a/litellm-rust/crates/http/src/tls.rs +++ b/litellm-rust/crates/http/src/tls.rs @@ -195,6 +195,12 @@ impl TryFrom<&HttpClientConfig> for ClientConfig { Verify::BuiltInRoots => builder.with_root_certificates(RootCertStore { roots: webpki_roots::TLS_SERVER_ROOTS.to_vec(), }), + Verify::NativeRoots => { + let native = rustls_native_certs::load_native_certs(); + let mut roots = RootCertStore::empty(); + roots.add_parsable_certificates(native.certs); + builder.with_root_certificates(roots) + } Verify::CaBundle(path) => { builder.with_root_certificates(bundle_roots(path, TlsSource::CaBundle)?) } diff --git a/litellm-rust/crates/http/src/websocket.rs b/litellm-rust/crates/http/src/websocket.rs index e904392152a..685c79cf3a1 100644 --- a/litellm-rust/crates/http/src/websocket.rs +++ b/litellm-rust/crates/http/src/websocket.rs @@ -1,62 +1,135 @@ -use std::{ - io, - sync::{Arc, OnceLock}, +use http::{ + HeaderMap, HeaderValue, StatusCode, + header::{ + CONNECTION, SEC_WEBSOCKET_ACCEPT, SEC_WEBSOCKET_KEY, SEC_WEBSOCKET_PROTOCOL, + SEC_WEBSOCKET_VERSION, UPGRADE, + }, }; - -use rustls::{ClientConfig, RootCertStore}; -use tokio::net::TcpStream; +use reqwest::Url; use tokio_tungstenite::{ - Connector, MaybeTlsStream, WebSocketStream, connect_async_tls_with_config, - tungstenite::{client::IntoClientRequest, error::TlsError, handshake::client::Response}, + WebSocketStream, + tungstenite::{ + handshake::{client::generate_key, derive_accept_key}, + protocol::Role, + }, }; -pub type UpstreamWebSocket = WebSocketStream>; +use crate::{client::Client, transport::Error}; -static TLS_CONFIG: OnceLock> = OnceLock::new(); - -fn build_tls_config() -> Result> { - let native = rustls_native_certs::load_native_certs(); - let mut store = RootCertStore::empty(); - let (added, _ignored) = store.add_parsable_certificates(native.certs); - if added == 0 { - return Err(Box::new(tokio_tungstenite::tungstenite::Error::Io( - io::Error::other(format!( - "no usable native root certificates: {:?}", - native.errors - )), - ))); - } - ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider())) - .with_safe_default_protocol_versions() - .map(|builder| builder.with_root_certificates(store).with_no_client_auth()) - .map_err(|error| { - Box::new(tokio_tungstenite::tungstenite::Error::Tls( - TlsError::Rustls(error), - )) - }) -} - -fn tls_config() -> Result, Box> { - if let Some(config) = TLS_CONFIG.get() { - return Ok(Arc::clone(config)); - } - let built = Arc::new(build_tls_config()?); - Ok(Arc::clone(TLS_CONFIG.get_or_init(|| built))) -} +pub type UpstreamWebSocket = WebSocketStream; #[tracing::instrument(name = "litellm.websocket.handshake", level = "debug", skip_all)] -pub async fn connect_upstream( - request: R, -) -> Result<(UpstreamWebSocket, Response), Box> -where - R: IntoClientRequest + Unpin, -{ - let request = request.into_client_request().map_err(Box::new)?; - let connector = match request.uri().scheme_str() { - Some("wss") => Some(Connector::Rustls(tls_config()?)), - _ => None, - }; - connect_async_tls_with_config(request, None, false, connector) +pub async fn connect( + client: &Client, + url: &str, + headers: HeaderMap, +) -> Result { + let offered_protocols = headers.get(SEC_WEBSOCKET_PROTOCOL).cloned(); + let key = generate_key(); + let response = client + .get(http_url(url)?) + .headers(handshake_headers(headers, &key)?) + .send() .await - .map_err(Box::new) + .map_err(Error::from_reqwest_before_dispatch)?; + let status = response.status(); + if status != StatusCode::SWITCHING_PROTOCOLS { + return Err(Error::Http { + status: status.as_u16(), + body: response.text().await.unwrap_or_default(), + }); + } + let expected = derive_accept_key(key.as_bytes()); + if response + .headers() + .get(SEC_WEBSOCKET_ACCEPT) + .map(HeaderValue::as_bytes) + != Some(expected.as_bytes()) + { + return Err(Error::Network( + "WebSocket handshake returned a wrong Sec-WebSocket-Accept".into(), + )); + } + validate_response_headers(response.headers(), offered_protocols.as_ref())?; + let upgraded = response.upgrade().await?; + Ok(WebSocketStream::from_raw_socket(upgraded, Role::Client, None).await) +} + +fn validate_response_headers( + headers: &HeaderMap, + offered_protocols: Option<&HeaderValue>, +) -> Result<(), Error> { + if !headers + .get(UPGRADE) + .is_some_and(|value| value.as_bytes().eq_ignore_ascii_case(b"websocket")) + { + return Err(Error::Network( + "WebSocket handshake returned an invalid Upgrade header".into(), + )); + } + if !headers.get_all(CONNECTION).iter().any(|value| { + value.to_str().is_ok_and(|value| { + value + .split(',') + .any(|token| token.trim().eq_ignore_ascii_case("upgrade")) + }) + }) { + return Err(Error::Network( + "WebSocket handshake returned an invalid Connection header".into(), + )); + } + let protocol_matches = match (offered_protocols, headers.get(SEC_WEBSOCKET_PROTOCOL)) { + (_, None) => true, + (Some(offered), Some(selected)) => offered + .to_str() + .ok() + .zip(selected.to_str().ok()) + .is_some_and(|(offered, selected)| { + offered + .split(',') + .any(|protocol| protocol.trim() == selected) + }), + _ => false, + }; + if !protocol_matches { + return Err(Error::Network( + "WebSocket handshake returned an invalid subprotocol".into(), + )); + } + Ok(()) +} + +fn http_url(url: &str) -> Result { + let mut url = Url::parse(url).map_err(not_dispatched)?; + let scheme = match url.scheme() { + "ws" | "http" => "http", + "wss" | "https" => "https", + other => { + return Err(not_dispatched(format!( + "unsupported WebSocket URL scheme {other:?}" + ))); + } + }; + url.set_scheme(scheme) + .map_err(|()| not_dispatched("WebSocket URL cannot be sent over HTTP"))?; + Ok(url) +} + +fn handshake_headers(headers: HeaderMap, key: &str) -> Result { + let key = HeaderValue::from_str(key).map_err(not_dispatched)?; + Ok([ + (CONNECTION, HeaderValue::from_static("Upgrade")), + (UPGRADE, HeaderValue::from_static("websocket")), + (SEC_WEBSOCKET_VERSION, HeaderValue::from_static("13")), + (SEC_WEBSOCKET_KEY, key), + ] + .into_iter() + .fold(headers, |mut headers, (name, value)| { + headers.insert(name, value); + headers + })) +} + +fn not_dispatched(error: impl std::fmt::Display) -> Error { + Error::Connect(error.to_string()) } diff --git a/litellm-rust/crates/http/tests/websocket.rs b/litellm-rust/crates/http/tests/websocket.rs index 2fbaa3994e8..94aba97a84d 100644 --- a/litellm-rust/crates/http/tests/websocket.rs +++ b/litellm-rust/crates/http/tests/websocket.rs @@ -1,59 +1,539 @@ +use std::{ + net::SocketAddr, + path::PathBuf, + sync::{Arc, Mutex}, +}; + use futures_util::{SinkExt, StreamExt}; -use litellm_http::websocket::connect_upstream; -use rstest::rstest; -use tokio::net::TcpListener; +use http::{HeaderMap, HeaderValue}; +use litellm_http::{ + ClientIdentity, ClientVariant, EnvironmentProxies, HttpClientConfig, HttpClientPool, + HttpSettings, Resolution, Verify, media::PublicDnsResolver, transport::Error, + websocket::connect, +}; +use rcgen::{BasicConstraints, CertificateParams, CertifiedIssuer, IsCa, KeyPair}; +use rstest::{Context, rstest}; +use rustls::{ + RootCertStore, ServerConfig, + crypto::ring, + pki_types::{CertificateDer, PrivateKeyDer}, + server::WebPkiClientVerifier, +}; +use tokio::{ + io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}, + net::{TcpListener, TcpStream}, +}; +use tokio_rustls::TlsAcceptor; use tokio_tungstenite::{ accept_hdr_async, tungstenite::{ Message, - client::IntoClientRequest, handshake::server::{Request, Response}, }, }; +struct Pki { + directory: tempfile::TempDir, + ca: CertifiedIssuer<'static, KeyPair>, + server_chain: Vec>, + server_key: Vec, +} + +impl Pki { + fn new() -> Self { + let mut ca_params = CertificateParams::new(Vec::::new()).unwrap(); + ca_params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained); + let ca = CertifiedIssuer::self_signed(ca_params, KeyPair::generate().unwrap()).unwrap(); + let server_key = KeyPair::generate().unwrap(); + let server = CertificateParams::new(vec!["localhost".into()]) + .unwrap() + .signed_by(&server_key, &ca) + .unwrap(); + let client_key = KeyPair::generate().unwrap(); + let client = CertificateParams::new(vec!["client".into()]) + .unwrap() + .signed_by(&client_key, &ca) + .unwrap(); + let directory = tempfile::tempdir().unwrap(); + std::fs::write(directory.path().join("ca.pem"), ca.pem()).unwrap(); + std::fs::write( + directory.path().join("client.pem"), + format!("{}{}", client.pem(), client_key.serialize_pem()), + ) + .unwrap(); + Self { + directory, + server_chain: vec![server.der().clone()], + server_key: server_key.serialize_der(), + ca, + } + } + + fn path(&self, name: &str) -> PathBuf { + self.directory.path().join(name) + } + + fn acceptor(&self, require_client_certificate: bool) -> TlsAcceptor { + let provider = Arc::new(ring::default_provider()); + let builder = ServerConfig::builder_with_provider(Arc::clone(&provider)) + .with_safe_default_protocol_versions() + .unwrap(); + let builder = if require_client_certificate { + let mut roots = RootCertStore::empty(); + roots.add(self.ca.der().clone()).unwrap(); + builder.with_client_cert_verifier( + WebPkiClientVerifier::builder_with_provider(Arc::new(roots), provider) + .build() + .unwrap(), + ) + } else { + builder.with_no_client_auth() + }; + let mut config = builder + .with_single_cert( + self.server_chain.clone(), + PrivateKeyDer::try_from(self.server_key.clone()).unwrap(), + ) + .unwrap(); + config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()]; + TlsAcceptor::from(Arc::new(config)) + } +} + +fn config(verify: Verify) -> HttpClientConfig { + HttpClientConfig { + verify, + http2: true, + ..Resolution::from(&HttpSettings::default()).config + } +} + +fn proxied_through(proxy: SocketAddr) -> EnvironmentProxies { + EnvironmentProxies::from_environment(&move |name: &str| { + (name == "HTTPS_PROXY").then(|| format!("http://{proxy}")) + }) +} + +async fn open( + config: &HttpClientConfig, + url: &str, + headers: HeaderMap, +) -> Result { + let client = HttpClientPool::new(Arc::new(PublicDnsResolver)) + .client(config, ClientVariant::WebSocket) + .unwrap(); + connect(&client, url, headers).await +} + +async fn echo_once( + stream: S, + seen: Arc>>, +) { + #[expect( + clippy::result_large_err, + reason = "tungstenite requires an unboxed handshake error response" + )] + let Ok(mut socket) = accept_hdr_async(stream, |request: &Request, response: Response| { + *seen.lock().unwrap() = Some(request.clone()); + Ok(response) + }) + .await + else { + return; + }; + if let Some(Ok(message)) = socket.next().await { + let _ = socket.send(message).await; + } +} + +async fn serve_tls( + pki: &Pki, + require_client_certificate: bool, +) -> (SocketAddr, Arc>>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let acceptor = pki.acceptor(require_client_certificate); + let seen = Arc::new(Mutex::new(None)); + let recorded = Arc::clone(&seen); + tokio::spawn(async move { + while let Ok((stream, _)) = listener.accept().await { + if let Ok(stream) = acceptor.accept(stream).await { + echo_once(stream, Arc::clone(&recorded)).await; + } + } + }); + (address, seen) +} + +async fn read_head(stream: &mut TcpStream) -> String { + let mut head = Vec::new(); + while !head.ends_with(b"\r\n\r\n") { + let mut byte = [0u8; 1]; + if stream.read(&mut byte).await.unwrap() == 0 { + break; + } + head.push(byte[0]); + } + String::from_utf8(head).unwrap() +} + +async fn serve_tunnel_proxy(target: SocketAddr) -> (SocketAddr, Arc>>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let requests = Arc::new(Mutex::new(Vec::new())); + let seen = Arc::clone(&requests); + tokio::spawn(async move { + while let Ok((mut client, _)) = listener.accept().await { + let head = read_head(&mut client).await; + seen.lock() + .unwrap() + .push(head.lines().next().unwrap_or_default().to_owned()); + let mut upstream = TcpStream::connect(target).await.unwrap(); + client + .write_all(b"HTTP/1.1 200 Connection established\r\n\r\n") + .await + .unwrap(); + let _ = tokio::io::copy_bidirectional(&mut client, &mut upstream).await; + } + }); + (address, requests) +} + +async fn serve_raw(response: &'static str) -> SocketAddr { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + read_head(&mut stream).await; + stream.write_all(response.as_bytes()).await.unwrap(); + let _ = stream.read(&mut [0u8; 1]).await; + }); + address +} + +async fn round_trip(socket: &mut litellm_http::websocket::UpstreamWebSocket) { + socket.send(Message::Text("ping".into())).await.unwrap(); + assert_eq!( + socket.next().await.unwrap().unwrap(), + Message::Text("ping".into()) + ); +} + #[rstest] #[case::without_query("/responses")] #[case::with_query("/responses?model=test-model")] #[tokio::test] -async fn connects_with_caller_headers_and_exchanges_frames(#[case] path: &'static str) { +async fn plain_websocket_sends_caller_headers_and_path(#[case] path: &'static str) { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let address = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { + let seen = Arc::new(Mutex::new(None)); + let recorded = Arc::clone(&seen); + tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + echo_once(stream, recorded).await; + }); + let headers = HeaderMap::from_iter([( + http::header::AUTHORIZATION, + HeaderValue::from_static("Bearer test-key"), + )]); + + let mut socket = open( + &config(Verify::BuiltInRoots), + &format!("ws://{address}{path}"), + headers, + ) + .await + .unwrap(); + + round_trip(&mut socket).await; + let request = seen.lock().unwrap().take().unwrap(); + assert_eq!(request.uri().path_and_query().unwrap().as_str(), path); + assert_eq!(request.headers()["authorization"], "Bearer test-key"); +} + +#[rstest] +#[case::ca_bundle_trusts_the_private_ca(|pki: &Pki| Verify::CaBundle(pki.path("ca.pem")), true)] +#[case::ssl_verify_off_skips_verification(|_: &Pki| Verify::Disabled, true)] +#[case::built_in_roots_reject_the_private_ca(|_: &Pki| Verify::BuiltInRoots, false)] +#[tokio::test] +async fn wss_honors_ssl_verify(#[case] verify: fn(&Pki) -> Verify, #[case] trusted: bool) { + let pki = Pki::new(); + let (address, _) = serve_tls(&pki, false).await; + + let result = open( + &config(verify(&pki)), + &format!("wss://localhost:{}/responses", address.port()), + HeaderMap::new(), + ) + .await; + + match result { + Ok(mut socket) if trusted => round_trip(&mut socket).await, + Err(Error::Connect(_)) if !trusted => {} + other => panic!("unexpected handshake outcome: {:?}", other.map(drop)), + } +} + +#[rstest] +#[case::plain_without_roots(false, false, true)] +#[case::tls_with_roots(true, true, true)] +#[case::tls_without_roots(true, false, false)] +#[tokio::test] +async fn websocket_uses_system_roots( + #[context] context: Context, + #[case] tls: bool, + #[case] trusted_roots: bool, + #[case] connects: bool, +) { + const CHILD_URL: &str = "LITELLM_TEST_SYSTEM_ROOTS_WSS_URL"; + if let Ok(url) = std::env::var(CHILD_URL) { + let result = open(&config(Verify::BuiltInRoots), &url, HeaderMap::new()).await; + match result { + Ok(mut socket) if connects => round_trip(&mut socket).await, + Err(Error::Connect(_)) if !connects => {} + other => panic!("unexpected handshake outcome: {:?}", other.map(drop)), + } + return; + } + + let pki = Pki::new(); + let url = if tls { + let (address, _) = serve_tls(&pki, false).await; + format!("wss://localhost:{}/responses", address.port()) + } else { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + echo_once(stream, Arc::new(Mutex::new(None))).await; + }); + format!("ws://{address}/responses") + }; + let ca = if trusted_roots { + pki.path("ca.pem") + } else { + let empty = pki.path("empty.pem"); + std::fs::write(&empty, "").unwrap(); + empty + }; + let test = format!( + "{}::case_{}_{}", + context.name, + context.case.unwrap() + 1, + context.description.unwrap() + ); + let output = tokio::task::spawn_blocking(move || { + std::process::Command::new(std::env::current_exe().unwrap()) + .args(["--exact", &test, "--nocapture"]) + .env(CHILD_URL, url) + .env("SSL_CERT_FILE", ca) + .env("SSL_CERT_DIR", "") + .output() + .unwrap() + }) + .await + .unwrap(); + let stdout = String::from_utf8_lossy(&output.stdout); + assert!( + output.status.success() && stdout.contains("1 passed"), + "system-root WebSocket child failed:\n{stdout}\n{}", + String::from_utf8_lossy(&output.stderr) + ); +} + +#[rstest] +#[case::with_certificate(true)] +#[case::without_certificate(false)] +#[tokio::test] +async fn wss_presents_the_configured_client_certificate(#[case] configured: bool) { + let pki = Pki::new(); + let (address, _) = serve_tls(&pki, true).await; + let config = HttpClientConfig { + client_certificate: configured.then(|| ClientIdentity::Pem(pki.path("client.pem"))), + ..config(Verify::CaBundle(pki.path("ca.pem"))) + }; + + let result = open( + &config, + &format!("wss://localhost:{}/responses", address.port()), + HeaderMap::new(), + ) + .await; + + match result { + Ok(mut socket) if configured => round_trip(&mut socket).await, + Err(_) if !configured => {} + other => panic!("unexpected handshake outcome: {:?}", other.map(drop)), + } +} + +#[rstest] +#[tokio::test] +async fn wss_tunnels_through_the_https_proxy() { + let pki = Pki::new(); + let (target, _) = serve_tls(&pki, false).await; + let (proxy, requests) = serve_tunnel_proxy(target).await; + let config = HttpClientConfig { + proxies: proxied_through(proxy), + ..config(Verify::CaBundle(pki.path("ca.pem"))) + }; + + let mut socket = open( + &config, + &format!("wss://localhost:{}/responses", target.port()), + HeaderMap::new(), + ) + .await + .unwrap(); + + round_trip(&mut socket).await; + assert_eq!( + *requests.lock().unwrap(), + [format!("CONNECT localhost:{} HTTP/1.1", target.port())] + ); +} + +#[rstest] +#[tokio::test] +async fn refused_upgrade_reports_the_status_and_body() { + let address = + serve_raw("HTTP/1.1 403 Forbidden\r\nContent-Length: 6\r\nConnection: close\r\n\r\ndenied") + .await; + + let result = open( + &config(Verify::BuiltInRoots), + &format!("ws://{address}/responses"), + HeaderMap::new(), + ) + .await; + + assert_eq!( + result.map(drop), + Err(Error::Http { + status: 403, + body: "denied".into() + }) + ); +} + +#[rstest] +#[tokio::test] +async fn upgrade_with_a_wrong_accept_key_is_rejected() { + let address = serve_raw( + "HTTP/1.1 101 Switching Protocols\r\nConnection: Upgrade\r\nUpgrade: websocket\r\nSec-WebSocket-Accept: wrong\r\n\r\n", + ) + .await; + + let result = open( + &config(Verify::BuiltInRoots), + &format!("ws://{address}/responses"), + HeaderMap::new(), + ) + .await; + + assert!(matches!(result, Err(Error::Network(_)))); +} + +#[rstest] +#[case::non_websocket_scheme("ftp://localhost/responses")] +#[case::unparseable("not a url")] +#[tokio::test] +async fn invalid_url_is_rejected_before_dispatch(#[case] url: &str) { + let result = open(&config(Verify::BuiltInRoots), url, HeaderMap::new()).await; + + assert!(matches!(result, Err(Error::Connect(_)))); +} + +#[rstest] +#[case::missing_upgrade(None, Some("Upgrade"), None, None, false)] +#[case::wrong_upgrade(Some("http"), Some("Upgrade"), None, None, false)] +#[case::missing_connection(Some("websocket"), None, None, None, false)] +#[case::wrong_connection(Some("websocket"), Some("keep-alive"), None, None, false)] +#[case::unsolicited_protocol(Some("websocket"), Some("Upgrade"), None, Some("other"), false)] +#[case::unoffered_protocol( + Some("websocket"), + Some("Upgrade"), + Some("events"), + Some("other"), + false +)] +#[case::declined_optional_protocol(Some("websocket"), Some("Upgrade"), Some("events"), None, true)] +#[case::matching_protocol( + Some("websocket"), + Some("Upgrade"), + Some("events"), + Some("events"), + true +)] +#[case::multiple_offers( + Some("websocket"), + Some("Upgrade"), + Some("events, realtime"), + Some("realtime"), + true +)] +#[case::case_insensitive_upgrade(Some("WebSocket"), Some("uPgRaDe"), None, None, true)] +#[case::connection_tokens(Some("websocket"), Some("keep-alive, Upgrade"), None, None, true)] +#[tokio::test] +async fn upgrade_validates_headers_and_negotiated_protocol( + #[case] upgrade: Option<&'static str>, + #[case] connection: Option<&'static str>, + #[case] offered: Option<&'static str>, + #[case] selected: Option<&'static str>, + #[case] accepted: bool, +) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + tokio::spawn(async move { let (stream, _) = listener.accept().await.unwrap(); #[expect( clippy::result_large_err, reason = "tungstenite requires an unboxed handshake error response" )] - let mut socket = accept_hdr_async(stream, move |request: &Request, response: Response| { - assert_eq!(request.uri().path_and_query().unwrap().as_str(), path); - assert_eq!(request.headers()["authorization"], "Bearer test-key"); + let result = accept_hdr_async(stream, move |_: &Request, mut response: Response| { + response.headers_mut().remove("upgrade"); + response.headers_mut().remove("connection"); + if let Some(value) = upgrade { + response + .headers_mut() + .insert("upgrade", HeaderValue::from_static(value)); + } + if let Some(value) = connection { + response + .headers_mut() + .insert("connection", HeaderValue::from_static(value)); + } + if let Some(value) = selected { + response + .headers_mut() + .insert("sec-websocket-protocol", HeaderValue::from_static(value)); + } Ok(response) }) - .await - .unwrap(); - assert_eq!( - socket.next().await.unwrap().unwrap(), - Message::Text("request".into()) - ); - socket.send(Message::Text("response".into())).await.unwrap(); - socket.close(None).await.unwrap(); + .await; + let Ok(mut socket) = result else { + return; + }; + if let Some(Ok(message)) = socket.next().await { + let _ = socket.send(message).await; + } }); - let mut request = format!("ws://{address}{path}") - .into_client_request() - .unwrap(); - request - .headers_mut() - .insert("authorization", "Bearer test-key".parse().unwrap()); - let (mut socket, response) = connect_upstream(request).await.unwrap(); - assert_eq!(response.status(), 101); - socket.send(Message::Text("request".into())).await.unwrap(); - assert_eq!( - socket.next().await.unwrap().unwrap(), - Message::Text("response".into()) - ); - assert!(matches!( - socket.next().await.unwrap().unwrap(), - Message::Close(_) - )); - server.await.unwrap(); + let headers = offered.map_or_else(HeaderMap::new, |value| { + HeaderMap::from_iter([( + http::header::SEC_WEBSOCKET_PROTOCOL, + HeaderValue::from_static(value), + )]) + }); + let result = open( + &config(Verify::BuiltInRoots), + &format!("ws://{address}/responses"), + headers, + ) + .await; + + match result { + Ok(mut socket) if accepted => round_trip(&mut socket).await, + Err(Error::Network(_)) if !accepted => {} + other => panic!("unexpected handshake outcome: {:?}", other.map(drop)), + } } diff --git a/litellm-rust/crates/inference-responses/Cargo.toml b/litellm-rust/crates/inference-responses/Cargo.toml index 793483ee76d..3dcd9976e0a 100644 --- a/litellm-rust/crates/inference-responses/Cargo.toml +++ b/litellm-rust/crates/inference-responses/Cargo.toml @@ -25,6 +25,7 @@ tokio-tungstenite.workspace = true tracing.workspace = true [dev-dependencies] +litellm-http = { workspace = true, features = ["test-support"] } litellm-cache.workspace = true litellm-cache-memory.workspace = true litellm-host-native.workspace = true diff --git a/litellm-rust/crates/inference-responses/src/websocket.rs b/litellm-rust/crates/inference-responses/src/websocket.rs index 213a9a95831..c2ca89ecbaf 100644 --- a/litellm-rust/crates/inference-responses/src/websocket.rs +++ b/litellm-rust/crates/inference-responses/src/websocket.rs @@ -1,12 +1,14 @@ use std::{collections::HashMap, sync::Arc, time::Duration}; use futures_util::{SinkExt, StreamExt}; -use litellm_http::websocket::{UpstreamWebSocket, connect_upstream}; +use litellm_http::{ + Client, + websocket::{self, UpstreamWebSocket}, +}; use tokio::sync::Mutex; use tokio_tungstenite::tungstenite::{ Message, - client::IntoClientRequest, - http::{HeaderName, HeaderValue}, + http::{HeaderMap, HeaderName, HeaderValue}, }; use super::Error; @@ -24,42 +26,33 @@ impl ResponsesWebSocketConnection { fields(outcome) )] pub async fn connect_url( + client: &Client, url: &str, headers: &HashMap, timeout: Option, ) -> Result { litellm_inference::diagnostic::operation("litellm.websocket.connect_url", async { - let mut request = url.into_client_request().map_err(|error| { - Error::Transport(litellm_http::transport::Error::Network(error.to_string())) - })?; - for (name, value) in headers { - let header_name = name - .parse::() - .map_err(|error| Error::InvalidRequest(error.to_string().into()))?; - let header_value = HeaderValue::from_str(value) - .map_err(|error| Error::InvalidRequest(error.to_string().into()))?; - request.headers_mut().insert(header_name, header_value); - } - let connect = connect_upstream(request); - let result = match timeout { + let headers = headers + .iter() + .map(|(name, value)| { + Ok(( + name.parse::() + .map_err(|error| Error::InvalidRequest(error.to_string().into()))?, + HeaderValue::from_str(value) + .map_err(|error| Error::InvalidRequest(error.to_string().into()))?, + )) + }) + .collect::>()?; + let connect = websocket::connect(client, url, headers); + let socket = match timeout { Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| { Error::Transport(litellm_http::transport::Error::Network( "Responses WebSocket connection timed out".into(), )) })?, None => connect.await, - }; - let (socket, _) = result.map_err(|error| match *error { - tokio_tungstenite::tungstenite::Error::Http(response) => { - Error::Transport(litellm_http::transport::Error::Http { - status: response.status().as_u16(), - body: String::new(), - }) - } - other => { - Error::Transport(litellm_http::transport::Error::Network(other.to_string())) - } - })?; + } + .map_err(Error::Transport)?; Ok(Self { socket: Arc::new(Mutex::new(Some(socket))), }) diff --git a/litellm-rust/crates/inference-responses/tests/responses.rs b/litellm-rust/crates/inference-responses/tests/responses.rs index ae07a71dabd..6d39d55d33c 100644 --- a/litellm-rust/crates/inference-responses/tests/responses.rs +++ b/litellm-rust/crates/inference-responses/tests/responses.rs @@ -538,6 +538,7 @@ async fn websocket_operations_trace_outcomes_without_capturing_frames_or_credent .logger() .instrument(async { let connection = ResponsesWebSocketConnection::connect_url( + &litellm_http::Client::for_test(), &format!("ws://{address}/responses"), &std::collections::HashMap::from([( "authorization".into(), diff --git a/litellm-rust/crates/python-bridge/src/routes/responses/websocket.rs b/litellm-rust/crates/python-bridge/src/routes/responses/websocket.rs index 615a838a597..c89185dc21f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses/websocket.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses/websocket.rs @@ -1,7 +1,9 @@ -use litellm_inference_responses::websocket::ResponsesWebSocketConnection as RustResponsesWebSocketConnection; use pyo3::prelude::*; use serde_json::Value; +use litellm_http::ClientVariant; +use litellm_inference_responses::websocket::ResponsesWebSocketConnection as RustResponsesWebSocketConnection; + use crate::{ errors::route_error_to_pyerr, marshal::{marshal_headers, optional_timeout}, @@ -25,10 +27,12 @@ impl ResponsesWebSocketConnection { ) -> PyResult> { let headers = marshal_headers(headers)?; let timeout = optional_timeout(timeout_seconds); + let client = crate::http::host_client(py, ClientVariant::WebSocket)?; crate::execution::run_async_value(py, async move { - let inner = RustResponsesWebSocketConnection::connect_url(&url, &headers, timeout) - .await - .map_err(route_error_to_pyerr)?; + let inner = + RustResponsesWebSocketConnection::connect_url(&client, &url, &headers, timeout) + .await + .map_err(route_error_to_pyerr)?; Ok(ResponsesWebSocketConnection { inner }) }) } From db5b94014371a52230b20bfa58f79d71a357b6db Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 28 Oct 2026 08:41:50 -0700 Subject: [PATCH 039/128] fix(vertex_ai): send Vertex Claude labels as X-Vertex-AI-Labels header instead of the request body (#48753) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../anthropic/transformation.py | 21 +++ .../test_vertex_partner_labels_wire.py | 120 ++++++++++++++++++ ...partner_models_anthropic_transformation.py | 28 ++++ 3 files changed, 169 insertions(+) create mode 100644 tests/integration/providers/test_vertex_partner_labels_wire.py diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py index 29253359003..363e5596cb9 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py @@ -1,5 +1,7 @@ # What is this? ## Handler file for calling claude-3 on vertex ai +import base64 +import json from collections.abc import Mapping from typing import TYPE_CHECKING, Final @@ -19,6 +21,19 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer _NON_DEFAULT_PARAMS: Final = TypeAdapter(Mapping[str, object]) +_VERTEX_LABELS_HEADER: Final = "X-Vertex-AI-Labels" +_LABELS: Final = TypeAdapter(Mapping[object, object]) + + +def _vertex_labels_header(labels: object, header_names: Mapping[str, object]) -> str | None: + if not isinstance(labels, dict) or any(name.lower() == _VERTEX_LABELS_HEADER.lower() for name in header_names): + return None + string_labels: Final = { + key: value + for key, value in _LABELS.validate_python(labels).items() + if isinstance(key, str) and isinstance(value, str) + } + return base64.b64encode(json.dumps(string_labels).encode()).decode() if string_labels else None class VertexAIError(Exception): @@ -102,6 +117,12 @@ class VertexAIAnthropicConfig(AnthropicConfig): litellm_params: dict, headers: dict, ) -> dict: + params: Final = _NON_DEFAULT_PARAMS.validate_python(optional_params) + labels_header: Final = _vertex_labels_header(params.get("labels"), _NON_DEFAULT_PARAMS.validate_python(headers)) + if "labels" in params: + del optional_params["labels"] # rebind-ok: labels travel in the X-Vertex-AI-Labels header, not the body + if labels_header is not None: + headers[_VERTEX_LABELS_HEADER] = labels_header # rebind-ok: outbound headers are returned by mutation data: Final = super().transform_request( model=model, messages=messages, diff --git a/tests/integration/providers/test_vertex_partner_labels_wire.py b/tests/integration/providers/test_vertex_partner_labels_wire.py new file mode 100644 index 00000000000..ebdf3f2022e --- /dev/null +++ b/tests/integration/providers/test_vertex_partner_labels_wire.py @@ -0,0 +1,120 @@ +import base64 +import json +from typing import Final + +import pytest +from integration._support.client import Gateway, Scenario +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +_BACKEND: Final = "claude-sonnet-5-5" +_PROJECT: Final = "scripted-project" +_LOCATION: Final = "us-east5" +_MODEL_PATH: Final = f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/anthropic/models/{_BACKEND}" +_TARGET: Final = f"{_MODEL_PATH}:rawPredict" +_STREAM_TARGET: Final = f"{_MODEL_PATH}:streamRawPredict?alt=sse" +_LABELS_HEADER: Final = "x-vertex-ai-labels" +_REPLY_TEXT: Final = "labels control" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_LABELS: Final[dict[str, JsonValue]] = {"team": "triage", "env": "prod"} + +_REPLY: Final[dict[str, JsonValue]] = { + "id": "msg_vertex_labels", + "type": "message", + "role": "assistant", + "model": _BACKEND, + "content": [{"type": "text", "text": _REPLY_TEXT}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 3}, +} +_EVENTS: Final[tuple[dict[str, JsonValue], ...]] = ( + {"type": "message_start", "message": {**_REPLY, "content": [], "stop_reason": None}}, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": _REPLY_TEXT}}, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 3}, + }, + {"type": "message_stop"}, +) +_SSE: Final = tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in _EVENTS) + + +def _peer(request: Request) -> Reply: + if request.target == "/_oauth/token": + return Reply( + body=json.dumps({"access_token": "scripted-token", "token_type": "Bearer", "expires_in": 3600}).encode() + ) + if request.target == _STREAM_TARGET: + return Reply(content_type="text/event-stream", chunks=_SSE) + if request.target == _TARGET: + return Reply(body=json.dumps(_REPLY).encode()) + return Reply(status=404, body=json.dumps({"error": f"unscripted target {request.target}"}).encode()) + + +def _deployment(scenario: Scenario, api_base: str) -> str: + return scenario.model( + model=f"vertex_ai/{_BACKEND}", + api_base=api_base, + api_key=None, + vertex_project=_PROJECT, + vertex_location=_LOCATION, + vertex_credentials=service_account_json(_PROJECT, api_base), + num_retries=0, + ) + + +def _chat(gateway: Gateway, wire: Wire, scenario: Scenario, stream: bool, **extra: JsonValue) -> Request: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": "Reply ok"}], + "max_tokens": 16, + "stream": stream, + **extra, + }, + ) + assert response.status_code == 200, response.text + assert _REPLY_TEXT in response.text, response.text + (sent,) = tuple(request for request in wire.drain() if request.target in (_TARGET, _STREAM_TARGET)) + assert sent.target == (_STREAM_TARGET if stream else _TARGET), sent.target + return sent + + +def _decoded_labels(sent: Request) -> JsonValue: + return json.loads(base64.b64decode(sent.headers[_LABELS_HEADER])) + + +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +def test_regression_vertex_claude_labels_move_from_body_to_header(gateway: Gateway, stream: bool) -> None: + with gateway.scenario() as scenario, wire_server(_peer) as wire: + sent: Final = _chat(gateway, wire, scenario, stream, labels=_LABELS) + + assert "labels" not in _JSON_OBJECT.validate_json(sent.body), sent.body + assert _decoded_labels(sent) == _LABELS, sent.headers + + +def test_guard_vertex_claude_without_labels_sends_no_labels_header(gateway: Gateway) -> None: + with gateway.scenario() as scenario, wire_server(_peer) as wire: + sent: Final = _chat(gateway, wire, scenario, False) + + assert "labels" not in _JSON_OBJECT.validate_json(sent.body), sent.body + assert _LABELS_HEADER not in sent.headers, sent.headers + + +def test_regression_vertex_claude_caller_labels_header_wins_over_labels_param(gateway: Gateway) -> None: + own_header: Final = base64.b64encode(json.dumps({"team": "caller"}).encode()).decode() + with gateway.scenario() as scenario, wire_server(_peer) as wire: + sent: Final = _chat( + gateway, wire, scenario, False, labels=_LABELS, extra_headers={"X-Vertex-AI-Labels": own_header} + ) + + assert "labels" not in _JSON_OBJECT.validate_json(sent.body), sent.body + assert sent.headers[_LABELS_HEADER] == own_header, sent.headers diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py index 21bda6f0345..1c432ec79da 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py @@ -1,4 +1,5 @@ +import base64 import copy import json from collections.abc import Mapping @@ -40,6 +41,33 @@ def test_vertex_ai_anthropic_thinking_param(model, expected_thinking): assert "thinking" not in supported_openai_params +@pytest.mark.parametrize( + ("labels", "headers", "expected_header"), + [ + ({"team": "triage", "env": "prod"}, {}, {"team": "triage", "env": "prod"}), + ({"team": "triage", "cost": 3}, {}, {"team": "triage"}), + ({"cost": 3}, {}, None), + ("team=triage", {}, None), + ({"team": "triage"}, {"x-vertex-ai-labels": "e30="}, {}), + ], + ids=["string_labels", "non_string_values_dropped", "no_string_values", "not_a_mapping", "caller_header_wins"], +) +def test_vertex_anthropic_labels_move_from_body_to_header( + labels: object, headers: dict[str, str], expected_header: dict[str, str] | None +) -> None: + body: Final = VertexAIAnthropicConfig().transform_request( + model="claude-sonnet-5-5", + messages=[{"role": "user", "content": "Hi"}], + optional_params={"max_tokens": 16, "anthropic_version": "vertex-2023-10-16", "labels": labels}, + litellm_params={}, + headers=headers, + ) + sent_labels: Final = next((value for name, value in headers.items() if name.lower() == "x-vertex-ai-labels"), None) + + assert "labels" not in body + assert (None if sent_labels is None else json.loads(base64.b64decode(sent_labels))) == expected_header + + @pytest.mark.parametrize( ("stream", "operation"), [(False, "rawPredict"), (True, "streamRawPredict")], From 3afe779125d27448ae4854ba91ee5626cdc497bf Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:07:22 +0000 Subject: [PATCH 040/128] fix(cost-map): rolling provider pricing, capability and retirement audit (#48188) * fix(cost-map): registry audit 2026-10-24, add provider retirement dates and fix mistral-large-4 and azure_ai sonnet 5.5 prices Co-authored-by: Hannes Brantner Co-authored-by: Aasif-Simpplr Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(mistral): derive mistral-large-4 cost expectations from the registry row instead of pinned prices Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(cost-map): add fireworks spot routers and reprice together Qwen3.7-Max Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(cost-map): absorb open registry PRs and update azure_ai mistral-small-2503 retirement date Co-authored-by: kerry Co-authored-by: Aryan Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(cost-map): credit absorbed registry PR authors Co-authored-by: kerry Co-authored-by: Aryan Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(cost-map): add azure gpt-6.1-terra rows and fireworks models/rune path Co-authored-by: ceng-p Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(cost-map): add azure_ai gpt-6.1-terra global standard row Co-authored-by: Ahmet Selim FEDAKAR <122402435+selimfedakar@users.noreply.github.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(cost-map): price mistral-large-4 cache reads at the published $0.14 Co-authored-by: zoveki <349028329+zoveki@users.noreply.github.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * revert(cost-map): leave mistral-large-4 rows to #48825 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(cost-map): refresh provider pricing and retirement metadata Absorb the verified GPT-5.6 Terra and Luna price fields from #48709. Add official Groq caching, xAI batch pricing, Moonshot metadata, and retirement dates without replacing existing model entries. Co-authored-by: devin-ai-integration[bot] <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: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Hannes Brantner Co-authored-by: Aasif-Simpplr Co-authored-by: Aryan Co-authored-by: kerry Co-authored-by: ceng-p Co-authored-by: Ahmet Selim FEDAKAR <122402435+selimfedakar@users.noreply.github.com> Co-authored-by: zoveki <349028329+zoveki@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 1125 ++++++++++++++--- model_prices_and_context_window.json | 1125 ++++++++++++++--- .../test_tool_call_cost_tracking.py | 119 +- .../test_model_prices_and_context_window.py | 40 +- 4 files changed, 1973 insertions(+), 436 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 9c3bd9dacbb..5a7006de2cf 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -4039,6 +4039,7 @@ "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/" }, "azure_ai/gpt-oss-120b": { + "deprecation_date": "2027-09-15", "input_cost_per_token": 1.5e-07, "output_cost_per_token": 6e-07, "litellm_provider": "azure_ai", @@ -4243,6 +4244,54 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure_ai/gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1e-07, + "input_cost_per_token": 1e-06, + "input_cost_per_token_above_272k_tokens": 2e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-06, + "output_cost_per_token_above_272k_tokens": 7.5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/gpt-6-1-terra-now-generally-available-in-microsoft-foundry-build-agents-that-see/4567769", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure_ai/gpt-5.5": { "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5e-07, @@ -4904,6 +4953,11 @@ "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6.6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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, @@ -4922,6 +4976,11 @@ "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6.6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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, @@ -5009,6 +5068,11 @@ "output_cost_per_token": 1.1e-05, "output_cost_per_token_batches": 5.5e-06, "output_cost_per_token_priority": 2.2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -5048,6 +5112,11 @@ "output_cost_per_token": 2.2e-06, "output_cost_per_token_batches": 1.1e-06, "output_cost_per_token_priority": 3.96e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -5113,6 +5182,11 @@ "input_cost_per_token_priority": 2.75e-06, "output_cost_per_token_batches": 5.5e-06, "output_cost_per_token_priority": 2.2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-5.1-codex": { @@ -5188,6 +5262,11 @@ "mode": "chat", "output_cost_per_token": 4.4e-07, "output_cost_per_token_batches": 2.2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -5344,6 +5423,11 @@ "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, @@ -5428,6 +5512,11 @@ "input_cost_per_token_priority": 2.5e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_priority": 2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/global/gpt-5.1-codex": { @@ -5680,6 +5769,11 @@ "output_cost_per_token": 8e-06, "output_cost_per_token_batches": 4e-06, "output_cost_per_token_priority": 1.4e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -5717,6 +5811,11 @@ "output_cost_per_token": 8e-06, "output_cost_per_token_batches": 4e-06, "output_cost_per_token_priority": 1.4e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -5754,6 +5853,11 @@ "output_cost_per_token": 1.6e-06, "output_cost_per_token_batches": 8e-07, "output_cost_per_token_priority": 2.8e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -5791,6 +5895,11 @@ "output_cost_per_token": 1.6e-06, "output_cost_per_token_batches": 8e-07, "output_cost_per_token_priority": 2.8e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -6203,6 +6312,11 @@ "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -6221,6 +6335,11 @@ "mode": "chat", "output_cost_per_token": 6e-07, "output_cost_per_token_batches": 3e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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, @@ -6710,6 +6829,11 @@ "default_reasoning_effort": "none", "input_cost_per_token_batches": 6.25e-07, "output_cost_per_token_batches": 5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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_minimal_reasoning_effort": true, "cache_read_input_token_cost_batches": 6.25e-08 @@ -6797,6 +6921,11 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_priority": 2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -6835,6 +6964,11 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_priority": 2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -6906,6 +7040,11 @@ "output_cost_per_token": 2e-06, "output_cost_per_token_batches": 1e-06, "output_cost_per_token_priority": 3.6e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -6944,6 +7083,11 @@ "output_cost_per_token": 2e-06, "output_cost_per_token_batches": 1e-06, "output_cost_per_token_priority": 3.6e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -6981,6 +7125,11 @@ "mode": "chat", "output_cost_per_token": 4e-07, "output_cost_per_token_batches": 2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -7016,6 +7165,11 @@ "mode": "chat", "output_cost_per_token": 4e-07, "output_cost_per_token_batches": 2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -7051,6 +7205,11 @@ "mode": "responses", "output_cost_per_token": 0.00012, "output_cost_per_token_batches": 6e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/responses" @@ -7112,6 +7271,11 @@ "input_cost_per_token_priority": 2.5e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_priority": 2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-5-chat": { @@ -7459,6 +7623,11 @@ "output_cost_per_token": 1.4e-05, "output_cost_per_token_batches": 7e-06, "output_cost_per_token_priority": 2.8e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -7498,6 +7667,11 @@ "output_cost_per_token": 1.4e-05, "output_cost_per_token_batches": 7e-06, "output_cost_per_token_priority": 2.8e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -7567,6 +7741,11 @@ "mode": "responses", "output_cost_per_token": 1.4e-05, "output_cost_per_token_priority": 2.8e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/responses" @@ -7599,9 +7778,9 @@ "output_cost_per_token": 0.000168, "output_cost_per_token_batches": 8.4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -7637,9 +7816,9 @@ "output_cost_per_token": 0.000168, "output_cost_per_token_batches": 8.4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -7996,9 +8175,9 @@ "output_cost_per_token_batches": 9e-05, "output_cost_per_token_flex": 9e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -8043,9 +8222,9 @@ "output_cost_per_token_batches": 9e-05, "output_cost_per_token_flex": 9e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -8095,9 +8274,9 @@ "output_cost_per_token_priority": 4e-05, "output_cost_per_token_above_272k_tokens_priority": 6e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -8157,9 +8336,9 @@ "output_cost_per_token_above_272k_tokens_priority": 6e-05, "output_cost_per_token_flex": 1e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -8220,9 +8399,9 @@ "output_cost_per_token_above_272k_tokens_priority": 6e-05, "output_cost_per_token_flex": 1e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -8283,9 +8462,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.6e-05, "output_cost_per_token_flex": 6e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -8346,9 +8525,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.6e-05, "output_cost_per_token_flex": 6e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -8409,9 +8588,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.6e-06, "output_cost_per_token_flex": 6e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -8472,9 +8651,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.6e-06, "output_cost_per_token_flex": 6e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -8519,9 +8698,9 @@ "output_cost_per_token": 5e-05, "output_cost_per_token_above_272k_tokens": 7.5e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -8567,9 +8746,9 @@ "output_cost_per_token": 5e-05, "output_cost_per_token_above_272k_tokens": 7.5e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -8615,9 +8794,9 @@ "output_cost_per_token": 5e-07, "output_cost_per_token_above_272k_tokens": 7.5e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ @@ -8663,9 +8842,9 @@ "output_cost_per_token": 5e-07, "output_cost_per_token_above_272k_tokens": 7.5e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ @@ -8718,7 +8897,7 @@ "output_cost_per_token_above_272k_tokens": 1.5e-05, "output_cost_per_token_above_272k_tokens_priority": 3e-05, "output_cost_per_token_priority": 2e-05, - "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ "/v1/chat/completions", @@ -8770,7 +8949,7 @@ "output_cost_per_token_above_272k_tokens": 1.5e-05, "output_cost_per_token_above_272k_tokens_priority": 3e-05, "output_cost_per_token_priority": 2e-05, - "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ "/v1/chat/completions", @@ -8800,6 +8979,53 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1e-07, + "input_cost_per_token": 1e-06, + "input_cost_per_token_above_272k_tokens": 2e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-06, + "output_cost_per_token_above_272k_tokens": 7.5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/gpt-6-1-terra-now-generally-available-in-microsoft-foundry-build-agents-that-see/4567769", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure/gpt-6.1-sol": { "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, @@ -8815,9 +9041,9 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_above_272k_tokens": 1.5e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ @@ -8863,9 +9089,9 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_above_272k_tokens": 1.5e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ @@ -8896,6 +9122,53 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/us/gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.375e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.75e-06, + "cache_read_input_token_cost": 5.5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-07, + "input_cost_per_token": 1.1e-06, + "input_cost_per_token_above_272k_tokens": 2.2e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.5e-06, + "output_cost_per_token_above_272k_tokens": 8.25e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/gpt-6-1-terra-now-generally-available-in-microsoft-foundry-build-agents-that-see/4567769", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure/us/gpt-6.1-sol": { "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_272k_tokens": 5.5e-06, @@ -8911,9 +9184,9 @@ "output_cost_per_token": 1.1e-05, "output_cost_per_token_above_272k_tokens": 1.65e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ @@ -8944,6 +9217,53 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/eu/gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 3e-06, + "cache_read_input_token_cost": 6e-08, + "cache_read_input_token_cost_above_272k_tokens": 1.2e-07, + "input_cost_per_token": 1.2e-06, + "input_cost_per_token_above_272k_tokens": 2.4e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "output_cost_per_token_above_272k_tokens": 9e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/gpt-6-1-terra-now-generally-available-in-microsoft-foundry-build-agents-that-see/4567769", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure/eu/gpt-6.1-sol": { "cache_creation_input_token_cost": 3e-06, "cache_creation_input_token_cost_above_272k_tokens": 6e-06, @@ -8959,9 +9279,9 @@ "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_272k_tokens": 1.8e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ @@ -8992,6 +9312,53 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/apac/gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 3e-06, + "cache_read_input_token_cost": 6e-08, + "cache_read_input_token_cost_above_272k_tokens": 1.2e-07, + "input_cost_per_token": 1.2e-06, + "input_cost_per_token_above_272k_tokens": 2.4e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "output_cost_per_token_above_272k_tokens": 9e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/gpt-6-1-terra-now-generally-available-in-microsoft-foundry-build-agents-that-see/4567769", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure/apac/gpt-6.1-sol": { "cache_creation_input_token_cost": 3e-06, "cache_creation_input_token_cost_above_272k_tokens": 6e-06, @@ -9007,9 +9374,9 @@ "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_272k_tokens": 1.8e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ @@ -9134,9 +9501,9 @@ "output_cost_per_token_above_272k_tokens_priority": 6.6e-05, "output_cost_per_token_priority": 4.4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -9188,9 +9555,9 @@ "output_cost_per_token_above_272k_tokens_priority": 6.6e-05, "output_cost_per_token_priority": 4.4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -9243,9 +9610,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.96e-05, "output_cost_per_token_priority": 2.64e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -9298,9 +9665,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.96e-06, "output_cost_per_token_priority": 2.64e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -9345,9 +9712,9 @@ "output_cost_per_token": 5.5e-05, "output_cost_per_token_above_272k_tokens": 8.25e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -9393,9 +9760,9 @@ "output_cost_per_token": 5.5e-07, "output_cost_per_token_above_272k_tokens": 8.25e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ @@ -9448,7 +9815,7 @@ "output_cost_per_token_above_272k_tokens": 1.65e-05, "output_cost_per_token_above_272k_tokens_priority": 3.3e-05, "output_cost_per_token_priority": 2.2e-05, - "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ "/v1/chat/completions", @@ -9536,9 +9903,9 @@ "output_cost_per_token_above_272k_tokens_priority": 6.6e-05, "output_cost_per_token_priority": 4.4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -9590,9 +9957,9 @@ "output_cost_per_token_above_272k_tokens_priority": 6.6e-05, "output_cost_per_token_priority": 4.4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -9645,9 +10012,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.96e-05, "output_cost_per_token_priority": 2.64e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -9700,9 +10067,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.96e-06, "output_cost_per_token_priority": 2.64e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -9760,9 +10127,9 @@ "output_cost_per_token_above_272k_tokens_priority": 9e-05, "output_cost_per_token_batches": 1.5e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -9814,9 +10181,9 @@ "output_cost_per_token_above_272k_tokens_batches": 2.475e-05, "output_cost_per_token_batches": 1.65e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -9868,9 +10235,9 @@ "output_cost_per_token_above_272k_tokens_batches": 2.475e-05, "output_cost_per_token_batches": 1.65e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -9921,9 +10288,9 @@ "output_cost_per_token_priority": 7.5e-05, "output_cost_per_token_above_272k_tokens_priority": 9e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -9981,9 +10348,9 @@ "output_cost_per_token_priority": 7.5e-05, "output_cost_per_token_above_272k_tokens_priority": 9e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10034,9 +10401,9 @@ "max_output_tokens": 128000, "mode": "chat", "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10086,9 +10453,9 @@ "max_output_tokens": 128000, "mode": "chat", "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10139,9 +10506,9 @@ "max_output_tokens": 128000, "mode": "chat", "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10191,9 +10558,9 @@ "max_output_tokens": 128000, "mode": "chat", "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10239,9 +10606,9 @@ "output_cost_per_token": 0.00018, "output_cost_per_token_above_272k_tokens": 0.00027, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/batch", @@ -10282,9 +10649,9 @@ "output_cost_per_token": 0.00018, "output_cost_per_token_above_272k_tokens": 0.00027, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/batch", @@ -10322,9 +10689,9 @@ "mode": "chat", "output_cost_per_token": 4.5e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10372,9 +10739,9 @@ "mode": "chat", "output_cost_per_token": 4.5e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10423,9 +10790,9 @@ "mode": "chat", "output_cost_per_token": 1.25e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10470,9 +10837,9 @@ "mode": "chat", "output_cost_per_token": 1.25e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10934,6 +11301,11 @@ "max_output_tokens": 100000, "mode": "chat", "output_cost_per_token": 8e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -10965,6 +11337,11 @@ "mode": "chat", "output_cost_per_token": 8e-06, "output_cost_per_token_batches": 4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -10996,9 +11373,9 @@ "mode": "responses", "output_cost_per_token": 4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -11065,6 +11442,11 @@ "mode": "responses", "output_cost_per_token": 8e-05, "output_cost_per_token_batches": 4e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/batch", @@ -11095,6 +11477,11 @@ "mode": "responses", "output_cost_per_token": 8e-05, "output_cost_per_token_batches": 4e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/batch", @@ -11124,6 +11511,11 @@ "max_output_tokens": 100000, "mode": "chat", "output_cost_per_token": 4.4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -11155,6 +11547,11 @@ "mode": "chat", "output_cost_per_token": 4.4e-06, "output_cost_per_token_batches": 2.2e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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, @@ -11268,6 +11665,11 @@ "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, "output_cost_per_token_priority": 1.54e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -11305,6 +11707,11 @@ "output_cost_per_token": 1.76e-06, "output_cost_per_token_batches": 8.8e-07, "output_cost_per_token_priority": 3.08e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -11428,6 +11835,11 @@ "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6.6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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, @@ -11446,6 +11858,11 @@ "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6.6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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, @@ -11533,6 +11950,11 @@ "output_cost_per_token": 1.1e-05, "output_cost_per_token_batches": 5.5e-06, "output_cost_per_token_priority": 2.2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -11572,6 +11994,11 @@ "output_cost_per_token": 2.2e-06, "output_cost_per_token_batches": 1.1e-06, "output_cost_per_token_priority": 3.96e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -11608,6 +12035,11 @@ "mode": "chat", "output_cost_per_token": 4.4e-07, "output_cost_per_token_batches": 2.2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -11673,6 +12105,11 @@ "input_cost_per_token_priority": 2.75e-06, "output_cost_per_token_batches": 5.5e-06, "output_cost_per_token_priority": 2.2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/us/gpt-5.1-codex": { @@ -11800,6 +12237,11 @@ "mode": "chat", "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -11832,6 +12274,11 @@ "mode": "chat", "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -11898,6 +12345,11 @@ "mode": "chat", "output_cost_per_token": 4.84e-06, "output_cost_per_token_batches": 2.42e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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, @@ -11918,6 +12370,11 @@ "mode": "chat", "output_cost_per_token": 4.84e-06, "output_cost_per_token_batches": 2.42e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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, @@ -11935,6 +12392,7 @@ "output_cost_per_second": 0.0001 }, "azure_ai/Cohere-embed-v3-english": { + "deprecation_date": "2027-02-15", "input_cost_per_token": 1e-07, "litellm_provider": "azure_ai", "max_input_tokens": 512, @@ -11945,6 +12403,7 @@ "supports_embedding_image_input": true }, "azure_ai/Cohere-embed-v3-multilingual": { + "deprecation_date": "2027-02-15", "input_cost_per_token": 1e-07, "litellm_provider": "azure_ai", "max_input_tokens": 512, @@ -11955,6 +12414,7 @@ "supports_embedding_image_input": true }, "azure_ai/Codestral-2501": { + "deprecation_date": "2027-02-15", "input_cost_per_token": 3e-07, "litellm_provider": "azure_ai", "max_input_tokens": 256000, @@ -11974,6 +12434,7 @@ "supports_native_streaming": true }, "azure_ai/FLUX-1.1-pro": { + "deprecation_date": "2027-01-05", "litellm_provider": "azure_ai", "max_input_tokens": 5000, "mode": "image_generation", @@ -11984,6 +12445,7 @@ ] }, "azure_ai/FLUX.1-Kontext-pro": { + "deprecation_date": "2027-01-05", "litellm_provider": "azure_ai", "max_input_tokens": 5000, "mode": "image_generation", @@ -11994,6 +12456,7 @@ ] }, "azure_ai/flux.2-pro": { + "deprecation_date": "2027-01-20", "litellm_provider": "azure_ai", "max_input_tokens": 32000, "mode": "image_generation", @@ -12004,6 +12467,7 @@ ] }, "azure_ai/FLUX.2-flex": { + "deprecation_date": "2027-01-20", "input_cost_per_pixel": 5e-08, "litellm_provider": "azure_ai", "max_input_tokens": 32000, @@ -12368,6 +12832,7 @@ "supports_tool_choice": true }, "azure_ai/Llama-3.3-70B-Instruct": { + "deprecation_date": "2027-01-31", "input_cost_per_token": 7.1e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -12379,6 +12844,7 @@ "supports_tool_choice": true }, "azure_ai/Llama-4-Maverick-17B-128E-Instruct-FP8": { + "deprecation_date": "2026-11-19", "input_cost_per_token": 2.5e-07, "litellm_provider": "azure_ai", "max_input_tokens": 1000000, @@ -12391,6 +12857,7 @@ "supports_vision": true }, "azure_ai/Llama-4-Scout-17B-16E-Instruct": { + "deprecation_date": "2027-01-31", "input_cost_per_token": 2e-07, "litellm_provider": "azure_ai", "max_input_tokens": 10000000, @@ -12624,6 +13091,7 @@ ] }, "azure_ai/cohere-command-a": { + "deprecation_date": "2027-02-15", "input_cost_per_token": 2.5e-06, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -12696,6 +13164,7 @@ "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models" }, "azure_ai/cohere-rerank-v4.0-pro": { + "deprecation_date": "2027-09-13", "input_cost_per_query": 0.0025, "input_cost_per_token": 0.0, "litellm_provider": "azure_ai", @@ -12706,6 +13175,7 @@ "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/introducing-cohere-rerank-4-0-in-microsoft-foundry/4477076" }, "azure_ai/cohere-rerank-v4.0-fast": { + "deprecation_date": "2027-09-13", "input_cost_per_query": 0.002, "input_cost_per_token": 0.0, "litellm_provider": "azure_ai", @@ -12716,6 +13186,7 @@ "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/introducing-cohere-rerank-4-0-in-microsoft-foundry/4477076" }, "azure_ai/deepseek-v3.2": { + "deprecation_date": "2027-11-01", "input_cost_per_token": 5.8e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -12730,6 +13201,7 @@ "supports_tool_choice": true }, "azure_ai/deepseek-v3.2-speciale": { + "deprecation_date": "2027-11-01", "input_cost_per_token": 5.8e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -12800,6 +13272,7 @@ "supports_tool_choice": true }, "azure_ai/embed-v-4-0": { + "deprecation_date": "2027-06-01", "input_cost_per_token": 1.2e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -12818,6 +13291,7 @@ "input_cost_per_image_token": 4.7e-07 }, "azure_ai/grok-4": { + "deprecation_date": "2027-01-31", "input_cost_per_token": 3e-06, "litellm_provider": "azure_ai", "max_input_tokens": 262000, @@ -12906,6 +13380,7 @@ "supports_web_search": true }, "azure_ai/grok-4-1-fast-non-reasoning": { + "deprecation_date": "2027-05-19", "input_cost_per_token": 2e-07, "output_cost_per_token": 5e-07, "litellm_provider": "azure_ai", @@ -12919,6 +13394,7 @@ "supports_web_search": true }, "azure_ai/grok-4-1-fast-reasoning": { + "deprecation_date": "2027-05-19", "input_cost_per_token": 2e-07, "output_cost_per_token": 5e-07, "litellm_provider": "azure_ai", @@ -12933,6 +13409,7 @@ "supports_web_search": true }, "azure_ai/grok-code-fast-1": { + "deprecation_date": "2027-01-31", "input_cost_per_token": 2e-07, "litellm_provider": "azure_ai", "max_input_tokens": 256000, @@ -13100,7 +13577,7 @@ "supports_tool_choice": true }, "azure_ai/mistral-small-2503": { - "deprecation_date": "2026-12-09", + "deprecation_date": "2027-02-15", "input_cost_per_token": 1e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -14619,11 +15096,12 @@ "cache_creation_input_token_cost": 4.5e-06, "cache_creation_input_token_cost_above_1hr": 7.2e-06, "cache_read_input_token_cost": 3.6e-07, + "deprecation_date": "2027-04-08", "input_cost_per_token": 3.6e-06, "litellm_provider": "bedrock", "supports_tool_search": true, "max_input_tokens": 200000, - "max_output_tokens": 8192, + "max_output_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.8e-05, "supports_assistant_prefill": true, @@ -14638,6 +15116,7 @@ "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md", "token_count_model": { "model": "anthropic.claude-sonnet-4-5-20250929-v1:0", "litellm_provider": "bedrock" @@ -14647,11 +15126,12 @@ "cache_creation_input_token_cost": 4.5e-06, "cache_creation_input_token_cost_above_1hr": 7.2e-06, "cache_read_input_token_cost": 3.6e-07, + "deprecation_date": "2027-04-08", "input_cost_per_token": 3.6e-06, "litellm_provider": "bedrock", "supports_tool_search": true, "max_input_tokens": 200000, - "max_output_tokens": 8192, + "max_output_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.8e-05, "supports_assistant_prefill": true, @@ -14665,7 +15145,8 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md" }, "bedrock/us-gov-east-1/meta.llama3-70b-instruct-v1:0": { "input_cost_per_token": 2.65e-06, @@ -14787,11 +15268,12 @@ "cache_creation_input_token_cost": 4.5e-06, "cache_creation_input_token_cost_above_1hr": 7.2e-06, "cache_read_input_token_cost": 3.6e-07, + "deprecation_date": "2027-04-08", "input_cost_per_token": 3.6e-06, "litellm_provider": "bedrock", "supports_tool_search": true, "max_input_tokens": 200000, - "max_output_tokens": 8192, + "max_output_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.8e-05, "supports_assistant_prefill": true, @@ -14806,6 +15288,7 @@ "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md", "token_count_model": { "model": "anthropic.claude-sonnet-4-5-20250929-v1:0", "litellm_provider": "bedrock" @@ -14815,11 +15298,12 @@ "cache_creation_input_token_cost": 4.5e-06, "cache_creation_input_token_cost_above_1hr": 7.2e-06, "cache_read_input_token_cost": 3.6e-07, + "deprecation_date": "2027-04-08", "input_cost_per_token": 3.6e-06, "litellm_provider": "bedrock", "supports_tool_search": true, "max_input_tokens": 200000, - "max_output_tokens": 8192, + "max_output_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.8e-05, "supports_assistant_prefill": true, @@ -14833,7 +15317,8 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md" }, "bedrock/us-gov-west-1/meta.llama3-70b-instruct-v1:0": { "input_cost_per_token": 2.65e-06, @@ -15197,6 +15682,7 @@ "supports_tool_choice": true }, "cerebras/llama-3.3-70b": { + "deprecation_date": "2026-02-16", "input_cost_per_token": 8.5e-07, "litellm_provider": "cerebras", "max_input_tokens": 128000, @@ -15207,6 +15693,7 @@ "supports_tool_choice": true }, "cerebras/llama3.1-70b": { + "deprecation_date": "2025-01-17", "input_cost_per_token": 6e-07, "litellm_provider": "cerebras", "max_input_tokens": 128000, @@ -15217,6 +15704,7 @@ "supports_tool_choice": true }, "cerebras/llama3.1-8b": { + "deprecation_date": "2026-05-27", "input_cost_per_token": 1e-07, "litellm_provider": "cerebras", "max_input_tokens": 128000, @@ -15241,6 +15729,7 @@ "supports_tool_choice": true }, "cerebras/qwen-3-32b": { + "deprecation_date": "2026-02-16", "input_cost_per_token": 4e-07, "litellm_provider": "cerebras", "max_input_tokens": 128000, @@ -15552,6 +16041,7 @@ "output_cost_per_token_above_200k_tokens": 2.25e-05, "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, + "deprecation_date": "2027-04-08", "litellm_provider": "bedrock", "supports_tool_search": true, "max_input_tokens": 200000, @@ -15570,7 +16060,8 @@ "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, "input_cost_per_token_batches": 1.5e-06, - "output_cost_per_token_batches": 7.5e-06 + "output_cost_per_token_batches": 7.5e-06, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md" }, "claude-opus-4-5-20251101": { "supports_web_search": true, @@ -21638,10 +22129,10 @@ "supports_vision": true }, "databricks/databricks-gpt-5-6-terra": { - "cache_creation_input_token_cost": 3.12501e-06, - "cache_read_input_token_cost": 2.4997e-07, - "input_cost_per_token": 2.49998e-06, - "input_dbu_cost_per_token": 3.5714e-05, + "cache_creation_input_token_cost": 2.49998e-06, + "cache_read_input_token_cost": 1.9999e-07, + "input_cost_per_token": 1.99997e-06, + "input_dbu_cost_per_token": 2.8571e-05, "litellm_provider": "databricks", "max_input_tokens": 922000, "max_output_tokens": 128000, @@ -21649,8 +22140,8 @@ "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", - "output_cost_per_token": 1.500002e-05, - "output_dbu_cost_per_token": 0.000214286, + "output_cost_per_token": 1.200003e-05, + "output_dbu_cost_per_token": 0.000171429, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -21666,10 +22157,10 @@ "supports_vision": true }, "databricks/databricks-gpt-5-6-luna": { - "cache_creation_input_token_cost": 1.24999e-06, - "cache_read_input_token_cost": 1.0003e-07, - "input_cost_per_token": 1.00002e-06, - "input_dbu_cost_per_token": 1.4286e-05, + "cache_creation_input_token_cost": 2.4997e-07, + "cache_read_input_token_cost": 2.002e-08, + "input_cost_per_token": 1.9999e-07, + "input_dbu_cost_per_token": 2.857e-06, "litellm_provider": "databricks", "max_input_tokens": 922000, "max_output_tokens": 128000, @@ -21677,8 +22168,8 @@ "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", - "output_cost_per_token": 5.99998e-06, - "output_dbu_cost_per_token": 8.5714e-05, + "output_cost_per_token": 1.20001e-06, + "output_dbu_cost_per_token": 1.7143e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -26989,6 +27480,18 @@ "supports_tool_choice": true, "supports_vision": false }, + "fireworks_ai/accounts/fireworks/routers/spot-deepseek-v4-flash": { + "cache_read_input_token_cost": 2.1e-08, + "input_cost_per_token": 1.05e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 2.1e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_vision": false + }, "fireworks_ai/accounts/fireworks/models/deepseek-v4-pro-0813": { "cache_read_input_token_cost": 4.4e-08, "cache_read_input_token_cost_priority": 5.5e-08, @@ -34476,6 +34979,7 @@ "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-4o-mini-tts": { + "deprecation_date": "2027-01-06", "input_cost_per_token": 6e-07, "litellm_provider": "openai", "mode": "audio_speech", @@ -35093,6 +35597,7 @@ "cache_read_input_token_cost_batches": 6.25e-08, "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, + "deprecation_date": "2026-12-11", "input_cost_per_token": 1.25e-06, "input_cost_per_token_batches": 6.25e-07, "input_cost_per_token_flex": 6.25e-07, @@ -35349,6 +35854,7 @@ "cache_read_input_token_cost_batches": 6.25e-08, "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, + "deprecation_date": "2027-04-01", "input_cost_per_token": 1.25e-06, "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", @@ -35648,15 +36154,15 @@ "gpt-6-astra": { "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_12hr": 2.5e-05, + "cache_creation_input_token_cost_above_12hr_ultrafast": 0.00015, + "cache_creation_input_token_cost_above_12hr_batches": 1.25e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 0.0001, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 0.0003, - "cache_creation_input_token_cost_above_12hr_batches": 1.25e-05, "cache_creation_input_token_cost_above_12hr_flex": 1.25e-05, "cache_creation_input_token_cost_above_12hr_priority": 5e-05, - "cache_creation_input_token_cost_above_12hr_ultrafast": 0.00015, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05, @@ -35826,11 +36332,11 @@ "gpt-6-luna": { "cache_creation_input_token_cost": 1.25e-07, "cache_creation_input_token_cost_above_12hr": 2.5e-07, + "cache_creation_input_token_cost_above_12hr_batches": 1.25e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 1e-06, - "cache_creation_input_token_cost_above_12hr_batches": 1.25e-07, "cache_creation_input_token_cost_above_12hr_flex": 1.25e-07, "cache_creation_input_token_cost_above_12hr_priority": 5e-07, "cache_creation_input_token_cost_decisions": 1.25e-07, @@ -37073,6 +37579,7 @@ "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_batches": 1e-08, "cache_read_input_token_cost_flex": 1e-08, + "deprecation_date": "2027-04-01", "input_cost_per_token": 2e-07, "input_cost_per_token_flex": 1e-07, "input_cost_per_token_batches": 1e-07, @@ -37120,6 +37627,7 @@ "supports_minimal_reasoning_effort": false }, "gpt-5-pro": { + "deprecation_date": "2026-12-11", "input_cost_per_token": 1.5e-05, "input_cost_per_token_batches": 7.5e-06, "litellm_provider": "openai", @@ -37399,6 +37907,7 @@ "cache_read_input_token_cost_batches": 1.25e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, + "deprecation_date": "2026-12-11", "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, "input_cost_per_token_flex": 1.25e-07, @@ -37499,6 +38008,7 @@ "cache_read_input_token_cost": 5e-09, "cache_read_input_token_cost_batches": 2.5e-09, "cache_read_input_token_cost_flex": 2.5e-09, + "deprecation_date": "2026-12-11", "input_cost_per_token": 5e-08, "input_cost_per_token_batches": 2.5e-08, "input_cost_per_token_flex": 2.5e-08, @@ -38208,6 +38718,7 @@ }, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -38228,6 +38739,7 @@ }, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -38247,6 +38759,7 @@ "search_context_size_medium": 0.005 }, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -41919,6 +42432,7 @@ }, "moonshot/kimi-k2.5": { "cache_read_input_token_cost": 1e-07, + "deprecation_date": "2026-08-31", "input_cost_per_token": 6e-07, "litellm_provider": "moonshot", "max_input_tokens": 262144, @@ -41949,7 +42463,31 @@ "supports_video_input": true, "supports_vision": true }, + "moonshot/kimi-k3.1-pro": { + "cache_creation_input_token_cost": 2e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "moonshot", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1e-05, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ], + "source": "https://platform.kimi.ai/docs/pricing/chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true + }, "moonshot/kimi-k3": { + "cache_creation_input_token_cost": 3e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, "litellm_provider": "moonshot", @@ -41964,6 +42502,7 @@ ], "source": "https://platform.kimi.ai/docs/pricing/chat-k3", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -41971,6 +42510,7 @@ "supports_vision": true }, "moonshot/moonshot-v1-128k": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 2e-06, "litellm_provider": "moonshot", "max_input_tokens": 131072, @@ -41983,6 +42523,7 @@ "supports_tool_choice": true }, "moonshot/moonshot-v1-128k-vision-preview": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 2e-06, "litellm_provider": "moonshot", "max_input_tokens": 131072, @@ -41996,6 +42537,7 @@ "supports_vision": true }, "moonshot/moonshot-v1-32k": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 1e-06, "litellm_provider": "moonshot", "max_input_tokens": 32768, @@ -42008,6 +42550,7 @@ "supports_tool_choice": true }, "moonshot/moonshot-v1-32k-vision-preview": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 1e-06, "litellm_provider": "moonshot", "max_input_tokens": 32768, @@ -42021,6 +42564,7 @@ "supports_vision": true }, "moonshot/moonshot-v1-8k": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 2e-07, "litellm_provider": "moonshot", "max_input_tokens": 8192, @@ -42033,6 +42577,7 @@ "supports_tool_choice": true }, "moonshot/moonshot-v1-8k-vision-preview": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 2e-07, "litellm_provider": "moonshot", "max_input_tokens": 8192, @@ -42046,6 +42591,7 @@ "supports_vision": true }, "moonshot/moonshot-v1-auto": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 2e-06, "litellm_provider": "moonshot", "max_input_tokens": 131072, @@ -43238,6 +43784,7 @@ "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_flex": 2.5e-07, "cache_read_input_token_cost_priority": 8.75e-07, + "deprecation_date": "2026-12-11", "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, @@ -43428,6 +43975,7 @@ "supports_vision": false }, "o3-pro": { + "deprecation_date": "2026-12-11", "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "openai", @@ -47288,13 +47836,13 @@ "supports_reasoning": true }, "together_ai/Qwen/Qwen3.7-Max": { - "cache_read_input_token_cost": 3e-07, - "input_cost_per_token": 1.5e-06, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 2.5e-06, "litellm_provider": "together_ai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "mode": "chat", - "output_cost_per_token": 4.5e-06, + "output_cost_per_token": 7.5e-06, "source": "https://www.together.ai/pricing", "supports_prompt_caching": true }, @@ -47804,6 +48352,7 @@ "cache_creation_input_token_cost_above_200k_tokens": 9e-06, "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.44e-05, "cache_read_input_token_cost_above_200k_tokens": 7.2e-07, + "deprecation_date": "2027-04-08", "litellm_provider": "bedrock_converse", "supports_tool_search": true, "max_input_tokens": 200000, @@ -47822,7 +48371,7 @@ "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://aws.amazon.com/bedrock/pricing/", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md", "token_count_model": { "model": "anthropic.claude-sonnet-4-5-20250929-v1:0", "litellm_provider": "bedrock" @@ -52899,15 +53448,21 @@ "xai/grok-4.7": { "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_200k_tokens": 1e-06, + "cache_read_input_token_cost_above_200k_tokens_batches": 8e-07, + "cache_read_input_token_cost_batches": 4e-07, "input_cost_per_image_token": 2e-06, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_batches": 3.2e-06, + "input_cost_per_token_batches": 1.6e-06, "litellm_provider": "xai", "max_input_tokens": 500000, "max_output_tokens": 500000, "mode": "chat", "output_cost_per_token": 6e-06, "output_cost_per_token_above_200k_tokens": 1.2e-05, + "output_cost_per_token_above_200k_tokens_batches": 9.6e-06, + "output_cost_per_token_batches": 4.8e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, @@ -56988,6 +57543,7 @@ ] }, "gpt-4o-mini-transcribe-2025-12-15": { + "deprecation_date": "2027-02-26", "input_cost_per_audio_token": 1.25e-06, "input_cost_per_second": 5e-05, "input_cost_per_token": 1.25e-06, @@ -61796,6 +62352,18 @@ "supports_tool_choice": true, "supports_vision": true }, + "fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-spot": { + "cache_read_input_token_cost": 4.2e-09, + "input_cost_per_token": 2.1e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 8.4e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_vision": true + }, "fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-us": { "cache_read_input_token_cost": 9e-09, "cache_read_input_token_cost_priority": 1.125e-08, @@ -61856,6 +62424,18 @@ "supports_tool_choice": true, "supports_vision": true }, + "fireworks_ai/accounts/fireworks/routers/kimi-k3-spot": { + "cache_read_input_token_cost": 2.1e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.05e-05, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_vision": true + }, "fireworks_ai/deepseek-v4-flash-0731": { "cache_read_input_token_cost": 7e-09, "cache_read_input_token_cost_priority": 8.75e-09, @@ -64525,6 +65105,7 @@ ] }, "xai/grok-imagine-image-pro": { + "deprecation_date": "2026-05-15", "input_cost_per_image": 0.05, "litellm_provider": "xai", "mode": "image_generation", @@ -65054,6 +65635,18 @@ "supports_tool_choice": true, "supports_vision": false }, + "fireworks_ai/accounts/fireworks/routers/glm-5p3-spot": { + "cache_read_input_token_cost": 1.82e-07, + "input_cost_per_token": 9.8e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 3.08e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_vision": false + }, "fireworks_ai/accounts/fireworks/routers/glm-5p3-us": { "cache_read_input_token_cost": 3.9e-07, "cache_read_input_token_cost_priority": 4.875e-07, @@ -65362,6 +65955,7 @@ ] }, "scaleway/glm-5.2": { + "deprecation_date": "2026-11-16", "input_cost_per_token": 1.8e-06, "litellm_provider": "scaleway", "max_input_tokens": 256000, @@ -67920,6 +68514,11 @@ "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, "output_cost_per_token_priority": 1.54e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-4.1-mini": { @@ -67934,6 +68533,11 @@ "output_cost_per_token": 1.76e-06, "output_cost_per_token_batches": 8.8e-07, "output_cost_per_token_priority": 3.08e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-4.1-nano": { @@ -67972,6 +68576,11 @@ "output_cost_per_token": 1.1e-05, "output_cost_per_token_batches": 5.5e-06, "output_cost_per_token_priority": 2.2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-5-codex": { @@ -67996,6 +68605,11 @@ "output_cost_per_token": 2.2e-06, "output_cost_per_token_batches": 1.1e-06, "output_cost_per_token_priority": 3.96e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-5-nano": { @@ -68008,6 +68622,11 @@ "mode": "chat", "output_cost_per_token": 4.4e-07, "output_cost_per_token_batches": 2.2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-5-pro": { @@ -68018,6 +68637,11 @@ "mode": "responses", "output_cost_per_token": 0.000132, "output_cost_per_token_batches": 6.6e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-5.1-codex-max": { @@ -68042,6 +68666,11 @@ "output_cost_per_token": 1.54e-05, "output_cost_per_token_batches": 7.7e-06, "output_cost_per_token_priority": 3.08e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-5.2-codex": { @@ -68060,6 +68689,11 @@ "mode": "responses", "output_cost_per_token": 0.0001848, "output_cost_per_token_batches": 9.24e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-5.3-codex": { @@ -68072,6 +68706,11 @@ "mode": "responses", "output_cost_per_token": 1.54e-05, "output_cost_per_token_priority": 3.08e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-5.4-mini": { @@ -68087,6 +68726,11 @@ "output_cost_per_token": 4.95e-06, "output_cost_per_token_batches": 2.475e-06, "output_cost_per_token_priority": 9.9e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-5.4-nano": { @@ -68099,6 +68743,11 @@ "mode": "chat", "output_cost_per_token": 1.375e-06, "output_cost_per_token_batches": 6.875e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-5.4-pro": { @@ -68113,6 +68762,11 @@ "output_cost_per_token_above_272k_tokens": 0.000297, "output_cost_per_token_above_272k_tokens_batches": 0.0001485, "output_cost_per_token_batches": 9.9e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-6-astra": { @@ -68129,6 +68783,11 @@ "mode": "chat", "output_cost_per_token": 6e-05, "output_cost_per_token_above_272k_tokens": 9e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", "supported_endpoints": [ "/v1/chat/completions", @@ -68173,9 +68832,9 @@ "output_cost_per_token": 6e-07, "output_cost_per_token_above_272k_tokens": 9e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ @@ -68221,9 +68880,9 @@ "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_272k_tokens": 1.8e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ @@ -68275,6 +68934,11 @@ "mode": "chat", "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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": { @@ -68288,6 +68952,11 @@ "mode": "chat", "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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-deep-research": { @@ -68310,6 +68979,11 @@ "mode": "chat", "output_cost_per_token": 4.84e-06, "output_cost_per_token_batches": 2.42e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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": { @@ -68323,6 +68997,11 @@ "mode": "chat", "output_cost_per_token": 4.84e-06, "output_cost_per_token_batches": 2.42e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/text-embedding-3-large": { @@ -68416,6 +69095,11 @@ "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, "output_cost_per_token_priority": 1.54e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/us/gpt-4.1-mini": { @@ -68430,6 +69114,11 @@ "output_cost_per_token": 1.76e-06, "output_cost_per_token_batches": 8.8e-07, "output_cost_per_token_priority": 3.08e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/us/gpt-4.1-nano": { @@ -68468,6 +69157,11 @@ "output_cost_per_token": 1.1e-05, "output_cost_per_token_batches": 5.5e-06, "output_cost_per_token_priority": 2.2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/us/gpt-5-codex": { @@ -68492,6 +69186,11 @@ "output_cost_per_token": 2.2e-06, "output_cost_per_token_batches": 1.1e-06, "output_cost_per_token_priority": 3.96e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/us/gpt-5-nano": { @@ -68504,6 +69203,11 @@ "mode": "chat", "output_cost_per_token": 4.4e-07, "output_cost_per_token_batches": 2.2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/us/gpt-5-pro": { @@ -68514,6 +69218,11 @@ "mode": "responses", "output_cost_per_token": 0.000132, "output_cost_per_token_batches": 6.6e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/us/gpt-5.1-codex-max": { @@ -68538,6 +69247,11 @@ "output_cost_per_token": 1.54e-05, "output_cost_per_token_batches": 7.7e-06, "output_cost_per_token_priority": 3.08e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/us/gpt-5.2-codex": { @@ -68556,6 +69270,11 @@ "mode": "responses", "output_cost_per_token": 0.0001848, "output_cost_per_token_batches": 9.24e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/us/gpt-5.3-codex": { @@ -68568,6 +69287,11 @@ "mode": "responses", "output_cost_per_token": 1.54e-05, "output_cost_per_token_priority": 3.08e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/us/gpt-5.4-mini": { @@ -68583,6 +69307,11 @@ "output_cost_per_token": 4.95e-06, "output_cost_per_token_batches": 2.475e-06, "output_cost_per_token_priority": 9.9e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/us/gpt-5.4-nano": { @@ -68595,6 +69324,11 @@ "mode": "chat", "output_cost_per_token": 1.375e-06, "output_cost_per_token_batches": 6.875e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/us/gpt-5.4-pro": { @@ -68609,6 +69343,11 @@ "output_cost_per_token_above_272k_tokens": 0.000297, "output_cost_per_token_above_272k_tokens_batches": 0.0001485, "output_cost_per_token_batches": 9.9e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/us/o1-mini": { @@ -68655,7 +69394,7 @@ "cache_creation_input_audio_token_cost": 4e-07, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, - "deprecation_date": "2026-10-31", + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image_token": 5e-06, "input_cost_per_token": 4e-06, @@ -71000,6 +71739,7 @@ "xai/grok-code-fast-1": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "deprecation_date": "2026-05-15", "input_cost_per_image_token": 1e-06, "input_cost_per_token": 1e-06, "input_cost_per_token_above_200k_tokens": 2e-06, @@ -71896,7 +72636,7 @@ "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, - "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost": 1e-07, "deprecation_date": "2027-09-29", "input_cost_per_token": 2e-06, "litellm_provider": "azure_ai", @@ -72503,15 +73243,15 @@ "gpt-6.1-sol": { "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_12hr": 5e-06, + "cache_creation_input_token_cost_above_12hr_ultrafast": 3e-05, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 6e-05, - "cache_creation_input_token_cost_above_12hr_batches": 2.5e-06, "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, "cache_creation_input_token_cost_above_12hr_priority": 1e-05, - "cache_creation_input_token_cost_above_12hr_ultrafast": 3e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_batches": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, @@ -74472,14 +75212,14 @@ "supports_vision": true }, "wandb/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B": { - "cache_read_input_token_cost": 2e-08, - "input_cost_per_token": 1.5e-07, + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 2.5e-07, "litellm_provider": "wandb", "max_input_tokens": 1049000, "max_output_tokens": 943718, "mode": "chat", - "output_cost_per_token": 7.5e-07, - "source": "https://openrouter.ai/api/v1/models/nvidia/nemotron-3.5-super-vl/endpoints", + "output_cost_per_token": 1.5e-06, + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -74577,6 +75317,18 @@ "supports_tool_choice": true, "supports_vision": true }, + "fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-1017-spot": { + "cache_read_input_token_cost": 4.2e-09, + "input_cost_per_token": 2.1e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 8.4e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_vision": true + }, "github_copilot/claude-sonnet-4.6": { "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, @@ -76331,6 +77083,15 @@ "max" ] }, + "fireworks_ai/accounts/fireworks/models/rune": { + "cache_read_input_token_cost": 8.4e-09, + "input_cost_per_token": 4.2e-08, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 262144, + "mode": "decisions", + "output_cost_per_token": 0.0, + "source": "https://api.fireworks.ai/v1/serverless/models?format=nested" + }, "bedrock_mantle/openai.gpt-6.1-terra": { "cache_creation_input_token_cost": 1.375e-06, "cache_creation_input_token_cost_above_272k_tokens": 2.75e-06, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 9c3bd9dacbb..5a7006de2cf 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -4039,6 +4039,7 @@ "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/" }, "azure_ai/gpt-oss-120b": { + "deprecation_date": "2027-09-15", "input_cost_per_token": 1.5e-07, "output_cost_per_token": 6e-07, "litellm_provider": "azure_ai", @@ -4243,6 +4244,54 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure_ai/gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1e-07, + "input_cost_per_token": 1e-06, + "input_cost_per_token_above_272k_tokens": 2e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-06, + "output_cost_per_token_above_272k_tokens": 7.5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/gpt-6-1-terra-now-generally-available-in-microsoft-foundry-build-agents-that-see/4567769", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure_ai/gpt-5.5": { "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5e-07, @@ -4904,6 +4953,11 @@ "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6.6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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, @@ -4922,6 +4976,11 @@ "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6.6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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, @@ -5009,6 +5068,11 @@ "output_cost_per_token": 1.1e-05, "output_cost_per_token_batches": 5.5e-06, "output_cost_per_token_priority": 2.2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -5048,6 +5112,11 @@ "output_cost_per_token": 2.2e-06, "output_cost_per_token_batches": 1.1e-06, "output_cost_per_token_priority": 3.96e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -5113,6 +5182,11 @@ "input_cost_per_token_priority": 2.75e-06, "output_cost_per_token_batches": 5.5e-06, "output_cost_per_token_priority": 2.2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-5.1-codex": { @@ -5188,6 +5262,11 @@ "mode": "chat", "output_cost_per_token": 4.4e-07, "output_cost_per_token_batches": 2.2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -5344,6 +5423,11 @@ "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, @@ -5428,6 +5512,11 @@ "input_cost_per_token_priority": 2.5e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_priority": 2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/global/gpt-5.1-codex": { @@ -5680,6 +5769,11 @@ "output_cost_per_token": 8e-06, "output_cost_per_token_batches": 4e-06, "output_cost_per_token_priority": 1.4e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -5717,6 +5811,11 @@ "output_cost_per_token": 8e-06, "output_cost_per_token_batches": 4e-06, "output_cost_per_token_priority": 1.4e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -5754,6 +5853,11 @@ "output_cost_per_token": 1.6e-06, "output_cost_per_token_batches": 8e-07, "output_cost_per_token_priority": 2.8e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -5791,6 +5895,11 @@ "output_cost_per_token": 1.6e-06, "output_cost_per_token_batches": 8e-07, "output_cost_per_token_priority": 2.8e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -6203,6 +6312,11 @@ "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -6221,6 +6335,11 @@ "mode": "chat", "output_cost_per_token": 6e-07, "output_cost_per_token_batches": 3e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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, @@ -6710,6 +6829,11 @@ "default_reasoning_effort": "none", "input_cost_per_token_batches": 6.25e-07, "output_cost_per_token_batches": 5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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_minimal_reasoning_effort": true, "cache_read_input_token_cost_batches": 6.25e-08 @@ -6797,6 +6921,11 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_priority": 2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -6835,6 +6964,11 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_priority": 2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -6906,6 +7040,11 @@ "output_cost_per_token": 2e-06, "output_cost_per_token_batches": 1e-06, "output_cost_per_token_priority": 3.6e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -6944,6 +7083,11 @@ "output_cost_per_token": 2e-06, "output_cost_per_token_batches": 1e-06, "output_cost_per_token_priority": 3.6e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -6981,6 +7125,11 @@ "mode": "chat", "output_cost_per_token": 4e-07, "output_cost_per_token_batches": 2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -7016,6 +7165,11 @@ "mode": "chat", "output_cost_per_token": 4e-07, "output_cost_per_token_batches": 2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -7051,6 +7205,11 @@ "mode": "responses", "output_cost_per_token": 0.00012, "output_cost_per_token_batches": 6e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/responses" @@ -7112,6 +7271,11 @@ "input_cost_per_token_priority": 2.5e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_priority": 2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-5-chat": { @@ -7459,6 +7623,11 @@ "output_cost_per_token": 1.4e-05, "output_cost_per_token_batches": 7e-06, "output_cost_per_token_priority": 2.8e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -7498,6 +7667,11 @@ "output_cost_per_token": 1.4e-05, "output_cost_per_token_batches": 7e-06, "output_cost_per_token_priority": 2.8e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -7567,6 +7741,11 @@ "mode": "responses", "output_cost_per_token": 1.4e-05, "output_cost_per_token_priority": 2.8e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/responses" @@ -7599,9 +7778,9 @@ "output_cost_per_token": 0.000168, "output_cost_per_token_batches": 8.4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -7637,9 +7816,9 @@ "output_cost_per_token": 0.000168, "output_cost_per_token_batches": 8.4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -7996,9 +8175,9 @@ "output_cost_per_token_batches": 9e-05, "output_cost_per_token_flex": 9e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -8043,9 +8222,9 @@ "output_cost_per_token_batches": 9e-05, "output_cost_per_token_flex": 9e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -8095,9 +8274,9 @@ "output_cost_per_token_priority": 4e-05, "output_cost_per_token_above_272k_tokens_priority": 6e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -8157,9 +8336,9 @@ "output_cost_per_token_above_272k_tokens_priority": 6e-05, "output_cost_per_token_flex": 1e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -8220,9 +8399,9 @@ "output_cost_per_token_above_272k_tokens_priority": 6e-05, "output_cost_per_token_flex": 1e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -8283,9 +8462,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.6e-05, "output_cost_per_token_flex": 6e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -8346,9 +8525,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.6e-05, "output_cost_per_token_flex": 6e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -8409,9 +8588,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.6e-06, "output_cost_per_token_flex": 6e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -8472,9 +8651,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.6e-06, "output_cost_per_token_flex": 6e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -8519,9 +8698,9 @@ "output_cost_per_token": 5e-05, "output_cost_per_token_above_272k_tokens": 7.5e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -8567,9 +8746,9 @@ "output_cost_per_token": 5e-05, "output_cost_per_token_above_272k_tokens": 7.5e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -8615,9 +8794,9 @@ "output_cost_per_token": 5e-07, "output_cost_per_token_above_272k_tokens": 7.5e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ @@ -8663,9 +8842,9 @@ "output_cost_per_token": 5e-07, "output_cost_per_token_above_272k_tokens": 7.5e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ @@ -8718,7 +8897,7 @@ "output_cost_per_token_above_272k_tokens": 1.5e-05, "output_cost_per_token_above_272k_tokens_priority": 3e-05, "output_cost_per_token_priority": 2e-05, - "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ "/v1/chat/completions", @@ -8770,7 +8949,7 @@ "output_cost_per_token_above_272k_tokens": 1.5e-05, "output_cost_per_token_above_272k_tokens_priority": 3e-05, "output_cost_per_token_priority": 2e-05, - "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ "/v1/chat/completions", @@ -8800,6 +8979,53 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1e-07, + "input_cost_per_token": 1e-06, + "input_cost_per_token_above_272k_tokens": 2e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-06, + "output_cost_per_token_above_272k_tokens": 7.5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/gpt-6-1-terra-now-generally-available-in-microsoft-foundry-build-agents-that-see/4567769", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure/gpt-6.1-sol": { "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, @@ -8815,9 +9041,9 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_above_272k_tokens": 1.5e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ @@ -8863,9 +9089,9 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_above_272k_tokens": 1.5e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ @@ -8896,6 +9122,53 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/us/gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.375e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.75e-06, + "cache_read_input_token_cost": 5.5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-07, + "input_cost_per_token": 1.1e-06, + "input_cost_per_token_above_272k_tokens": 2.2e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.5e-06, + "output_cost_per_token_above_272k_tokens": 8.25e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/gpt-6-1-terra-now-generally-available-in-microsoft-foundry-build-agents-that-see/4567769", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure/us/gpt-6.1-sol": { "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_272k_tokens": 5.5e-06, @@ -8911,9 +9184,9 @@ "output_cost_per_token": 1.1e-05, "output_cost_per_token_above_272k_tokens": 1.65e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ @@ -8944,6 +9217,53 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/eu/gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 3e-06, + "cache_read_input_token_cost": 6e-08, + "cache_read_input_token_cost_above_272k_tokens": 1.2e-07, + "input_cost_per_token": 1.2e-06, + "input_cost_per_token_above_272k_tokens": 2.4e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "output_cost_per_token_above_272k_tokens": 9e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/gpt-6-1-terra-now-generally-available-in-microsoft-foundry-build-agents-that-see/4567769", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure/eu/gpt-6.1-sol": { "cache_creation_input_token_cost": 3e-06, "cache_creation_input_token_cost_above_272k_tokens": 6e-06, @@ -8959,9 +9279,9 @@ "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_272k_tokens": 1.8e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ @@ -8992,6 +9312,53 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/apac/gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 3e-06, + "cache_read_input_token_cost": 6e-08, + "cache_read_input_token_cost_above_272k_tokens": 1.2e-07, + "input_cost_per_token": 1.2e-06, + "input_cost_per_token_above_272k_tokens": 2.4e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "output_cost_per_token_above_272k_tokens": 9e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/gpt-6-1-terra-now-generally-available-in-microsoft-foundry-build-agents-that-see/4567769", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure/apac/gpt-6.1-sol": { "cache_creation_input_token_cost": 3e-06, "cache_creation_input_token_cost_above_272k_tokens": 6e-06, @@ -9007,9 +9374,9 @@ "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_272k_tokens": 1.8e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ @@ -9134,9 +9501,9 @@ "output_cost_per_token_above_272k_tokens_priority": 6.6e-05, "output_cost_per_token_priority": 4.4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -9188,9 +9555,9 @@ "output_cost_per_token_above_272k_tokens_priority": 6.6e-05, "output_cost_per_token_priority": 4.4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -9243,9 +9610,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.96e-05, "output_cost_per_token_priority": 2.64e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -9298,9 +9665,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.96e-06, "output_cost_per_token_priority": 2.64e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -9345,9 +9712,9 @@ "output_cost_per_token": 5.5e-05, "output_cost_per_token_above_272k_tokens": 8.25e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -9393,9 +9760,9 @@ "output_cost_per_token": 5.5e-07, "output_cost_per_token_above_272k_tokens": 8.25e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ @@ -9448,7 +9815,7 @@ "output_cost_per_token_above_272k_tokens": 1.65e-05, "output_cost_per_token_above_272k_tokens_priority": 3.3e-05, "output_cost_per_token_priority": 2.2e-05, - "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ "/v1/chat/completions", @@ -9536,9 +9903,9 @@ "output_cost_per_token_above_272k_tokens_priority": 6.6e-05, "output_cost_per_token_priority": 4.4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -9590,9 +9957,9 @@ "output_cost_per_token_above_272k_tokens_priority": 6.6e-05, "output_cost_per_token_priority": 4.4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -9645,9 +10012,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.96e-05, "output_cost_per_token_priority": 2.64e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -9700,9 +10067,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.96e-06, "output_cost_per_token_priority": 2.64e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -9760,9 +10127,9 @@ "output_cost_per_token_above_272k_tokens_priority": 9e-05, "output_cost_per_token_batches": 1.5e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -9814,9 +10181,9 @@ "output_cost_per_token_above_272k_tokens_batches": 2.475e-05, "output_cost_per_token_batches": 1.65e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -9868,9 +10235,9 @@ "output_cost_per_token_above_272k_tokens_batches": 2.475e-05, "output_cost_per_token_batches": 1.65e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -9921,9 +10288,9 @@ "output_cost_per_token_priority": 7.5e-05, "output_cost_per_token_above_272k_tokens_priority": 9e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -9981,9 +10348,9 @@ "output_cost_per_token_priority": 7.5e-05, "output_cost_per_token_above_272k_tokens_priority": 9e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10034,9 +10401,9 @@ "max_output_tokens": 128000, "mode": "chat", "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10086,9 +10453,9 @@ "max_output_tokens": 128000, "mode": "chat", "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10139,9 +10506,9 @@ "max_output_tokens": 128000, "mode": "chat", "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10191,9 +10558,9 @@ "max_output_tokens": 128000, "mode": "chat", "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10239,9 +10606,9 @@ "output_cost_per_token": 0.00018, "output_cost_per_token_above_272k_tokens": 0.00027, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/batch", @@ -10282,9 +10649,9 @@ "output_cost_per_token": 0.00018, "output_cost_per_token_above_272k_tokens": 0.00027, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/batch", @@ -10322,9 +10689,9 @@ "mode": "chat", "output_cost_per_token": 4.5e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10372,9 +10739,9 @@ "mode": "chat", "output_cost_per_token": 4.5e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10423,9 +10790,9 @@ "mode": "chat", "output_cost_per_token": 1.25e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10470,9 +10837,9 @@ "mode": "chat", "output_cost_per_token": 1.25e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10934,6 +11301,11 @@ "max_output_tokens": 100000, "mode": "chat", "output_cost_per_token": 8e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -10965,6 +11337,11 @@ "mode": "chat", "output_cost_per_token": 8e-06, "output_cost_per_token_batches": 4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -10996,9 +11373,9 @@ "mode": "responses", "output_cost_per_token": 4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "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": [ @@ -11065,6 +11442,11 @@ "mode": "responses", "output_cost_per_token": 8e-05, "output_cost_per_token_batches": 4e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/batch", @@ -11095,6 +11477,11 @@ "mode": "responses", "output_cost_per_token": 8e-05, "output_cost_per_token_batches": 4e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/batch", @@ -11124,6 +11511,11 @@ "max_output_tokens": 100000, "mode": "chat", "output_cost_per_token": 4.4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -11155,6 +11547,11 @@ "mode": "chat", "output_cost_per_token": 4.4e-06, "output_cost_per_token_batches": 2.2e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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, @@ -11268,6 +11665,11 @@ "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, "output_cost_per_token_priority": 1.54e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -11305,6 +11707,11 @@ "output_cost_per_token": 1.76e-06, "output_cost_per_token_batches": 8.8e-07, "output_cost_per_token_priority": 3.08e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -11428,6 +11835,11 @@ "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6.6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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, @@ -11446,6 +11858,11 @@ "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6.6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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, @@ -11533,6 +11950,11 @@ "output_cost_per_token": 1.1e-05, "output_cost_per_token_batches": 5.5e-06, "output_cost_per_token_priority": 2.2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -11572,6 +11994,11 @@ "output_cost_per_token": 2.2e-06, "output_cost_per_token_batches": 1.1e-06, "output_cost_per_token_priority": 3.96e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -11608,6 +12035,11 @@ "mode": "chat", "output_cost_per_token": 4.4e-07, "output_cost_per_token_batches": 2.2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -11673,6 +12105,11 @@ "input_cost_per_token_priority": 2.75e-06, "output_cost_per_token_batches": 5.5e-06, "output_cost_per_token_priority": 2.2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/us/gpt-5.1-codex": { @@ -11800,6 +12237,11 @@ "mode": "chat", "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -11832,6 +12274,11 @@ "mode": "chat", "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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", @@ -11898,6 +12345,11 @@ "mode": "chat", "output_cost_per_token": 4.84e-06, "output_cost_per_token_batches": 2.42e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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, @@ -11918,6 +12370,11 @@ "mode": "chat", "output_cost_per_token": 4.84e-06, "output_cost_per_token_batches": 2.42e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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, @@ -11935,6 +12392,7 @@ "output_cost_per_second": 0.0001 }, "azure_ai/Cohere-embed-v3-english": { + "deprecation_date": "2027-02-15", "input_cost_per_token": 1e-07, "litellm_provider": "azure_ai", "max_input_tokens": 512, @@ -11945,6 +12403,7 @@ "supports_embedding_image_input": true }, "azure_ai/Cohere-embed-v3-multilingual": { + "deprecation_date": "2027-02-15", "input_cost_per_token": 1e-07, "litellm_provider": "azure_ai", "max_input_tokens": 512, @@ -11955,6 +12414,7 @@ "supports_embedding_image_input": true }, "azure_ai/Codestral-2501": { + "deprecation_date": "2027-02-15", "input_cost_per_token": 3e-07, "litellm_provider": "azure_ai", "max_input_tokens": 256000, @@ -11974,6 +12434,7 @@ "supports_native_streaming": true }, "azure_ai/FLUX-1.1-pro": { + "deprecation_date": "2027-01-05", "litellm_provider": "azure_ai", "max_input_tokens": 5000, "mode": "image_generation", @@ -11984,6 +12445,7 @@ ] }, "azure_ai/FLUX.1-Kontext-pro": { + "deprecation_date": "2027-01-05", "litellm_provider": "azure_ai", "max_input_tokens": 5000, "mode": "image_generation", @@ -11994,6 +12456,7 @@ ] }, "azure_ai/flux.2-pro": { + "deprecation_date": "2027-01-20", "litellm_provider": "azure_ai", "max_input_tokens": 32000, "mode": "image_generation", @@ -12004,6 +12467,7 @@ ] }, "azure_ai/FLUX.2-flex": { + "deprecation_date": "2027-01-20", "input_cost_per_pixel": 5e-08, "litellm_provider": "azure_ai", "max_input_tokens": 32000, @@ -12368,6 +12832,7 @@ "supports_tool_choice": true }, "azure_ai/Llama-3.3-70B-Instruct": { + "deprecation_date": "2027-01-31", "input_cost_per_token": 7.1e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -12379,6 +12844,7 @@ "supports_tool_choice": true }, "azure_ai/Llama-4-Maverick-17B-128E-Instruct-FP8": { + "deprecation_date": "2026-11-19", "input_cost_per_token": 2.5e-07, "litellm_provider": "azure_ai", "max_input_tokens": 1000000, @@ -12391,6 +12857,7 @@ "supports_vision": true }, "azure_ai/Llama-4-Scout-17B-16E-Instruct": { + "deprecation_date": "2027-01-31", "input_cost_per_token": 2e-07, "litellm_provider": "azure_ai", "max_input_tokens": 10000000, @@ -12624,6 +13091,7 @@ ] }, "azure_ai/cohere-command-a": { + "deprecation_date": "2027-02-15", "input_cost_per_token": 2.5e-06, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -12696,6 +13164,7 @@ "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models" }, "azure_ai/cohere-rerank-v4.0-pro": { + "deprecation_date": "2027-09-13", "input_cost_per_query": 0.0025, "input_cost_per_token": 0.0, "litellm_provider": "azure_ai", @@ -12706,6 +13175,7 @@ "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/introducing-cohere-rerank-4-0-in-microsoft-foundry/4477076" }, "azure_ai/cohere-rerank-v4.0-fast": { + "deprecation_date": "2027-09-13", "input_cost_per_query": 0.002, "input_cost_per_token": 0.0, "litellm_provider": "azure_ai", @@ -12716,6 +13186,7 @@ "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/introducing-cohere-rerank-4-0-in-microsoft-foundry/4477076" }, "azure_ai/deepseek-v3.2": { + "deprecation_date": "2027-11-01", "input_cost_per_token": 5.8e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -12730,6 +13201,7 @@ "supports_tool_choice": true }, "azure_ai/deepseek-v3.2-speciale": { + "deprecation_date": "2027-11-01", "input_cost_per_token": 5.8e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -12800,6 +13272,7 @@ "supports_tool_choice": true }, "azure_ai/embed-v-4-0": { + "deprecation_date": "2027-06-01", "input_cost_per_token": 1.2e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -12818,6 +13291,7 @@ "input_cost_per_image_token": 4.7e-07 }, "azure_ai/grok-4": { + "deprecation_date": "2027-01-31", "input_cost_per_token": 3e-06, "litellm_provider": "azure_ai", "max_input_tokens": 262000, @@ -12906,6 +13380,7 @@ "supports_web_search": true }, "azure_ai/grok-4-1-fast-non-reasoning": { + "deprecation_date": "2027-05-19", "input_cost_per_token": 2e-07, "output_cost_per_token": 5e-07, "litellm_provider": "azure_ai", @@ -12919,6 +13394,7 @@ "supports_web_search": true }, "azure_ai/grok-4-1-fast-reasoning": { + "deprecation_date": "2027-05-19", "input_cost_per_token": 2e-07, "output_cost_per_token": 5e-07, "litellm_provider": "azure_ai", @@ -12933,6 +13409,7 @@ "supports_web_search": true }, "azure_ai/grok-code-fast-1": { + "deprecation_date": "2027-01-31", "input_cost_per_token": 2e-07, "litellm_provider": "azure_ai", "max_input_tokens": 256000, @@ -13100,7 +13577,7 @@ "supports_tool_choice": true }, "azure_ai/mistral-small-2503": { - "deprecation_date": "2026-12-09", + "deprecation_date": "2027-02-15", "input_cost_per_token": 1e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -14619,11 +15096,12 @@ "cache_creation_input_token_cost": 4.5e-06, "cache_creation_input_token_cost_above_1hr": 7.2e-06, "cache_read_input_token_cost": 3.6e-07, + "deprecation_date": "2027-04-08", "input_cost_per_token": 3.6e-06, "litellm_provider": "bedrock", "supports_tool_search": true, "max_input_tokens": 200000, - "max_output_tokens": 8192, + "max_output_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.8e-05, "supports_assistant_prefill": true, @@ -14638,6 +15116,7 @@ "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md", "token_count_model": { "model": "anthropic.claude-sonnet-4-5-20250929-v1:0", "litellm_provider": "bedrock" @@ -14647,11 +15126,12 @@ "cache_creation_input_token_cost": 4.5e-06, "cache_creation_input_token_cost_above_1hr": 7.2e-06, "cache_read_input_token_cost": 3.6e-07, + "deprecation_date": "2027-04-08", "input_cost_per_token": 3.6e-06, "litellm_provider": "bedrock", "supports_tool_search": true, "max_input_tokens": 200000, - "max_output_tokens": 8192, + "max_output_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.8e-05, "supports_assistant_prefill": true, @@ -14665,7 +15145,8 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md" }, "bedrock/us-gov-east-1/meta.llama3-70b-instruct-v1:0": { "input_cost_per_token": 2.65e-06, @@ -14787,11 +15268,12 @@ "cache_creation_input_token_cost": 4.5e-06, "cache_creation_input_token_cost_above_1hr": 7.2e-06, "cache_read_input_token_cost": 3.6e-07, + "deprecation_date": "2027-04-08", "input_cost_per_token": 3.6e-06, "litellm_provider": "bedrock", "supports_tool_search": true, "max_input_tokens": 200000, - "max_output_tokens": 8192, + "max_output_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.8e-05, "supports_assistant_prefill": true, @@ -14806,6 +15288,7 @@ "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md", "token_count_model": { "model": "anthropic.claude-sonnet-4-5-20250929-v1:0", "litellm_provider": "bedrock" @@ -14815,11 +15298,12 @@ "cache_creation_input_token_cost": 4.5e-06, "cache_creation_input_token_cost_above_1hr": 7.2e-06, "cache_read_input_token_cost": 3.6e-07, + "deprecation_date": "2027-04-08", "input_cost_per_token": 3.6e-06, "litellm_provider": "bedrock", "supports_tool_search": true, "max_input_tokens": 200000, - "max_output_tokens": 8192, + "max_output_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.8e-05, "supports_assistant_prefill": true, @@ -14833,7 +15317,8 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md" }, "bedrock/us-gov-west-1/meta.llama3-70b-instruct-v1:0": { "input_cost_per_token": 2.65e-06, @@ -15197,6 +15682,7 @@ "supports_tool_choice": true }, "cerebras/llama-3.3-70b": { + "deprecation_date": "2026-02-16", "input_cost_per_token": 8.5e-07, "litellm_provider": "cerebras", "max_input_tokens": 128000, @@ -15207,6 +15693,7 @@ "supports_tool_choice": true }, "cerebras/llama3.1-70b": { + "deprecation_date": "2025-01-17", "input_cost_per_token": 6e-07, "litellm_provider": "cerebras", "max_input_tokens": 128000, @@ -15217,6 +15704,7 @@ "supports_tool_choice": true }, "cerebras/llama3.1-8b": { + "deprecation_date": "2026-05-27", "input_cost_per_token": 1e-07, "litellm_provider": "cerebras", "max_input_tokens": 128000, @@ -15241,6 +15729,7 @@ "supports_tool_choice": true }, "cerebras/qwen-3-32b": { + "deprecation_date": "2026-02-16", "input_cost_per_token": 4e-07, "litellm_provider": "cerebras", "max_input_tokens": 128000, @@ -15552,6 +16041,7 @@ "output_cost_per_token_above_200k_tokens": 2.25e-05, "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, + "deprecation_date": "2027-04-08", "litellm_provider": "bedrock", "supports_tool_search": true, "max_input_tokens": 200000, @@ -15570,7 +16060,8 @@ "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, "input_cost_per_token_batches": 1.5e-06, - "output_cost_per_token_batches": 7.5e-06 + "output_cost_per_token_batches": 7.5e-06, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md" }, "claude-opus-4-5-20251101": { "supports_web_search": true, @@ -21638,10 +22129,10 @@ "supports_vision": true }, "databricks/databricks-gpt-5-6-terra": { - "cache_creation_input_token_cost": 3.12501e-06, - "cache_read_input_token_cost": 2.4997e-07, - "input_cost_per_token": 2.49998e-06, - "input_dbu_cost_per_token": 3.5714e-05, + "cache_creation_input_token_cost": 2.49998e-06, + "cache_read_input_token_cost": 1.9999e-07, + "input_cost_per_token": 1.99997e-06, + "input_dbu_cost_per_token": 2.8571e-05, "litellm_provider": "databricks", "max_input_tokens": 922000, "max_output_tokens": 128000, @@ -21649,8 +22140,8 @@ "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", - "output_cost_per_token": 1.500002e-05, - "output_dbu_cost_per_token": 0.000214286, + "output_cost_per_token": 1.200003e-05, + "output_dbu_cost_per_token": 0.000171429, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -21666,10 +22157,10 @@ "supports_vision": true }, "databricks/databricks-gpt-5-6-luna": { - "cache_creation_input_token_cost": 1.24999e-06, - "cache_read_input_token_cost": 1.0003e-07, - "input_cost_per_token": 1.00002e-06, - "input_dbu_cost_per_token": 1.4286e-05, + "cache_creation_input_token_cost": 2.4997e-07, + "cache_read_input_token_cost": 2.002e-08, + "input_cost_per_token": 1.9999e-07, + "input_dbu_cost_per_token": 2.857e-06, "litellm_provider": "databricks", "max_input_tokens": 922000, "max_output_tokens": 128000, @@ -21677,8 +22168,8 @@ "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", - "output_cost_per_token": 5.99998e-06, - "output_dbu_cost_per_token": 8.5714e-05, + "output_cost_per_token": 1.20001e-06, + "output_dbu_cost_per_token": 1.7143e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -26989,6 +27480,18 @@ "supports_tool_choice": true, "supports_vision": false }, + "fireworks_ai/accounts/fireworks/routers/spot-deepseek-v4-flash": { + "cache_read_input_token_cost": 2.1e-08, + "input_cost_per_token": 1.05e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 2.1e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_vision": false + }, "fireworks_ai/accounts/fireworks/models/deepseek-v4-pro-0813": { "cache_read_input_token_cost": 4.4e-08, "cache_read_input_token_cost_priority": 5.5e-08, @@ -34476,6 +34979,7 @@ "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-4o-mini-tts": { + "deprecation_date": "2027-01-06", "input_cost_per_token": 6e-07, "litellm_provider": "openai", "mode": "audio_speech", @@ -35093,6 +35597,7 @@ "cache_read_input_token_cost_batches": 6.25e-08, "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, + "deprecation_date": "2026-12-11", "input_cost_per_token": 1.25e-06, "input_cost_per_token_batches": 6.25e-07, "input_cost_per_token_flex": 6.25e-07, @@ -35349,6 +35854,7 @@ "cache_read_input_token_cost_batches": 6.25e-08, "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, + "deprecation_date": "2027-04-01", "input_cost_per_token": 1.25e-06, "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", @@ -35648,15 +36154,15 @@ "gpt-6-astra": { "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_12hr": 2.5e-05, + "cache_creation_input_token_cost_above_12hr_ultrafast": 0.00015, + "cache_creation_input_token_cost_above_12hr_batches": 1.25e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 0.0001, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 0.0003, - "cache_creation_input_token_cost_above_12hr_batches": 1.25e-05, "cache_creation_input_token_cost_above_12hr_flex": 1.25e-05, "cache_creation_input_token_cost_above_12hr_priority": 5e-05, - "cache_creation_input_token_cost_above_12hr_ultrafast": 0.00015, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05, @@ -35826,11 +36332,11 @@ "gpt-6-luna": { "cache_creation_input_token_cost": 1.25e-07, "cache_creation_input_token_cost_above_12hr": 2.5e-07, + "cache_creation_input_token_cost_above_12hr_batches": 1.25e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 1e-06, - "cache_creation_input_token_cost_above_12hr_batches": 1.25e-07, "cache_creation_input_token_cost_above_12hr_flex": 1.25e-07, "cache_creation_input_token_cost_above_12hr_priority": 5e-07, "cache_creation_input_token_cost_decisions": 1.25e-07, @@ -37073,6 +37579,7 @@ "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_batches": 1e-08, "cache_read_input_token_cost_flex": 1e-08, + "deprecation_date": "2027-04-01", "input_cost_per_token": 2e-07, "input_cost_per_token_flex": 1e-07, "input_cost_per_token_batches": 1e-07, @@ -37120,6 +37627,7 @@ "supports_minimal_reasoning_effort": false }, "gpt-5-pro": { + "deprecation_date": "2026-12-11", "input_cost_per_token": 1.5e-05, "input_cost_per_token_batches": 7.5e-06, "litellm_provider": "openai", @@ -37399,6 +37907,7 @@ "cache_read_input_token_cost_batches": 1.25e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, + "deprecation_date": "2026-12-11", "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, "input_cost_per_token_flex": 1.25e-07, @@ -37499,6 +38008,7 @@ "cache_read_input_token_cost": 5e-09, "cache_read_input_token_cost_batches": 2.5e-09, "cache_read_input_token_cost_flex": 2.5e-09, + "deprecation_date": "2026-12-11", "input_cost_per_token": 5e-08, "input_cost_per_token_batches": 2.5e-08, "input_cost_per_token_flex": 2.5e-08, @@ -38208,6 +38718,7 @@ }, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -38228,6 +38739,7 @@ }, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -38247,6 +38759,7 @@ "search_context_size_medium": 0.005 }, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -41919,6 +42432,7 @@ }, "moonshot/kimi-k2.5": { "cache_read_input_token_cost": 1e-07, + "deprecation_date": "2026-08-31", "input_cost_per_token": 6e-07, "litellm_provider": "moonshot", "max_input_tokens": 262144, @@ -41949,7 +42463,31 @@ "supports_video_input": true, "supports_vision": true }, + "moonshot/kimi-k3.1-pro": { + "cache_creation_input_token_cost": 2e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "moonshot", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1e-05, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ], + "source": "https://platform.kimi.ai/docs/pricing/chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true + }, "moonshot/kimi-k3": { + "cache_creation_input_token_cost": 3e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, "litellm_provider": "moonshot", @@ -41964,6 +42502,7 @@ ], "source": "https://platform.kimi.ai/docs/pricing/chat-k3", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -41971,6 +42510,7 @@ "supports_vision": true }, "moonshot/moonshot-v1-128k": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 2e-06, "litellm_provider": "moonshot", "max_input_tokens": 131072, @@ -41983,6 +42523,7 @@ "supports_tool_choice": true }, "moonshot/moonshot-v1-128k-vision-preview": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 2e-06, "litellm_provider": "moonshot", "max_input_tokens": 131072, @@ -41996,6 +42537,7 @@ "supports_vision": true }, "moonshot/moonshot-v1-32k": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 1e-06, "litellm_provider": "moonshot", "max_input_tokens": 32768, @@ -42008,6 +42550,7 @@ "supports_tool_choice": true }, "moonshot/moonshot-v1-32k-vision-preview": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 1e-06, "litellm_provider": "moonshot", "max_input_tokens": 32768, @@ -42021,6 +42564,7 @@ "supports_vision": true }, "moonshot/moonshot-v1-8k": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 2e-07, "litellm_provider": "moonshot", "max_input_tokens": 8192, @@ -42033,6 +42577,7 @@ "supports_tool_choice": true }, "moonshot/moonshot-v1-8k-vision-preview": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 2e-07, "litellm_provider": "moonshot", "max_input_tokens": 8192, @@ -42046,6 +42591,7 @@ "supports_vision": true }, "moonshot/moonshot-v1-auto": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 2e-06, "litellm_provider": "moonshot", "max_input_tokens": 131072, @@ -43238,6 +43784,7 @@ "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_flex": 2.5e-07, "cache_read_input_token_cost_priority": 8.75e-07, + "deprecation_date": "2026-12-11", "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, @@ -43428,6 +43975,7 @@ "supports_vision": false }, "o3-pro": { + "deprecation_date": "2026-12-11", "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "openai", @@ -47288,13 +47836,13 @@ "supports_reasoning": true }, "together_ai/Qwen/Qwen3.7-Max": { - "cache_read_input_token_cost": 3e-07, - "input_cost_per_token": 1.5e-06, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 2.5e-06, "litellm_provider": "together_ai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "mode": "chat", - "output_cost_per_token": 4.5e-06, + "output_cost_per_token": 7.5e-06, "source": "https://www.together.ai/pricing", "supports_prompt_caching": true }, @@ -47804,6 +48352,7 @@ "cache_creation_input_token_cost_above_200k_tokens": 9e-06, "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.44e-05, "cache_read_input_token_cost_above_200k_tokens": 7.2e-07, + "deprecation_date": "2027-04-08", "litellm_provider": "bedrock_converse", "supports_tool_search": true, "max_input_tokens": 200000, @@ -47822,7 +48371,7 @@ "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://aws.amazon.com/bedrock/pricing/", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md", "token_count_model": { "model": "anthropic.claude-sonnet-4-5-20250929-v1:0", "litellm_provider": "bedrock" @@ -52899,15 +53448,21 @@ "xai/grok-4.7": { "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_200k_tokens": 1e-06, + "cache_read_input_token_cost_above_200k_tokens_batches": 8e-07, + "cache_read_input_token_cost_batches": 4e-07, "input_cost_per_image_token": 2e-06, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_batches": 3.2e-06, + "input_cost_per_token_batches": 1.6e-06, "litellm_provider": "xai", "max_input_tokens": 500000, "max_output_tokens": 500000, "mode": "chat", "output_cost_per_token": 6e-06, "output_cost_per_token_above_200k_tokens": 1.2e-05, + "output_cost_per_token_above_200k_tokens_batches": 9.6e-06, + "output_cost_per_token_batches": 4.8e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, @@ -56988,6 +57543,7 @@ ] }, "gpt-4o-mini-transcribe-2025-12-15": { + "deprecation_date": "2027-02-26", "input_cost_per_audio_token": 1.25e-06, "input_cost_per_second": 5e-05, "input_cost_per_token": 1.25e-06, @@ -61796,6 +62352,18 @@ "supports_tool_choice": true, "supports_vision": true }, + "fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-spot": { + "cache_read_input_token_cost": 4.2e-09, + "input_cost_per_token": 2.1e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 8.4e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_vision": true + }, "fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-us": { "cache_read_input_token_cost": 9e-09, "cache_read_input_token_cost_priority": 1.125e-08, @@ -61856,6 +62424,18 @@ "supports_tool_choice": true, "supports_vision": true }, + "fireworks_ai/accounts/fireworks/routers/kimi-k3-spot": { + "cache_read_input_token_cost": 2.1e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.05e-05, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_vision": true + }, "fireworks_ai/deepseek-v4-flash-0731": { "cache_read_input_token_cost": 7e-09, "cache_read_input_token_cost_priority": 8.75e-09, @@ -64525,6 +65105,7 @@ ] }, "xai/grok-imagine-image-pro": { + "deprecation_date": "2026-05-15", "input_cost_per_image": 0.05, "litellm_provider": "xai", "mode": "image_generation", @@ -65054,6 +65635,18 @@ "supports_tool_choice": true, "supports_vision": false }, + "fireworks_ai/accounts/fireworks/routers/glm-5p3-spot": { + "cache_read_input_token_cost": 1.82e-07, + "input_cost_per_token": 9.8e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 3.08e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_vision": false + }, "fireworks_ai/accounts/fireworks/routers/glm-5p3-us": { "cache_read_input_token_cost": 3.9e-07, "cache_read_input_token_cost_priority": 4.875e-07, @@ -65362,6 +65955,7 @@ ] }, "scaleway/glm-5.2": { + "deprecation_date": "2026-11-16", "input_cost_per_token": 1.8e-06, "litellm_provider": "scaleway", "max_input_tokens": 256000, @@ -67920,6 +68514,11 @@ "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, "output_cost_per_token_priority": 1.54e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-4.1-mini": { @@ -67934,6 +68533,11 @@ "output_cost_per_token": 1.76e-06, "output_cost_per_token_batches": 8.8e-07, "output_cost_per_token_priority": 3.08e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-4.1-nano": { @@ -67972,6 +68576,11 @@ "output_cost_per_token": 1.1e-05, "output_cost_per_token_batches": 5.5e-06, "output_cost_per_token_priority": 2.2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-5-codex": { @@ -67996,6 +68605,11 @@ "output_cost_per_token": 2.2e-06, "output_cost_per_token_batches": 1.1e-06, "output_cost_per_token_priority": 3.96e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-5-nano": { @@ -68008,6 +68622,11 @@ "mode": "chat", "output_cost_per_token": 4.4e-07, "output_cost_per_token_batches": 2.2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-5-pro": { @@ -68018,6 +68637,11 @@ "mode": "responses", "output_cost_per_token": 0.000132, "output_cost_per_token_batches": 6.6e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-5.1-codex-max": { @@ -68042,6 +68666,11 @@ "output_cost_per_token": 1.54e-05, "output_cost_per_token_batches": 7.7e-06, "output_cost_per_token_priority": 3.08e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-5.2-codex": { @@ -68060,6 +68689,11 @@ "mode": "responses", "output_cost_per_token": 0.0001848, "output_cost_per_token_batches": 9.24e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-5.3-codex": { @@ -68072,6 +68706,11 @@ "mode": "responses", "output_cost_per_token": 1.54e-05, "output_cost_per_token_priority": 3.08e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-5.4-mini": { @@ -68087,6 +68726,11 @@ "output_cost_per_token": 4.95e-06, "output_cost_per_token_batches": 2.475e-06, "output_cost_per_token_priority": 9.9e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-5.4-nano": { @@ -68099,6 +68743,11 @@ "mode": "chat", "output_cost_per_token": 1.375e-06, "output_cost_per_token_batches": 6.875e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-5.4-pro": { @@ -68113,6 +68762,11 @@ "output_cost_per_token_above_272k_tokens": 0.000297, "output_cost_per_token_above_272k_tokens_batches": 0.0001485, "output_cost_per_token_batches": 9.9e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/gpt-6-astra": { @@ -68129,6 +68783,11 @@ "mode": "chat", "output_cost_per_token": 6e-05, "output_cost_per_token_above_272k_tokens": 9e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", "supported_endpoints": [ "/v1/chat/completions", @@ -68173,9 +68832,9 @@ "output_cost_per_token": 6e-07, "output_cost_per_token_above_272k_tokens": 9e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ @@ -68221,9 +68880,9 @@ "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_272k_tokens": 1.8e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ @@ -68275,6 +68934,11 @@ "mode": "chat", "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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": { @@ -68288,6 +68952,11 @@ "mode": "chat", "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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-deep-research": { @@ -68310,6 +68979,11 @@ "mode": "chat", "output_cost_per_token": 4.84e-06, "output_cost_per_token_batches": 2.42e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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": { @@ -68323,6 +68997,11 @@ "mode": "chat", "output_cost_per_token": 4.84e-06, "output_cost_per_token_batches": 2.42e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/text-embedding-3-large": { @@ -68416,6 +69095,11 @@ "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, "output_cost_per_token_priority": 1.54e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/us/gpt-4.1-mini": { @@ -68430,6 +69114,11 @@ "output_cost_per_token": 1.76e-06, "output_cost_per_token_batches": 8.8e-07, "output_cost_per_token_priority": 3.08e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/us/gpt-4.1-nano": { @@ -68468,6 +69157,11 @@ "output_cost_per_token": 1.1e-05, "output_cost_per_token_batches": 5.5e-06, "output_cost_per_token_priority": 2.2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/us/gpt-5-codex": { @@ -68492,6 +69186,11 @@ "output_cost_per_token": 2.2e-06, "output_cost_per_token_batches": 1.1e-06, "output_cost_per_token_priority": 3.96e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/us/gpt-5-nano": { @@ -68504,6 +69203,11 @@ "mode": "chat", "output_cost_per_token": 4.4e-07, "output_cost_per_token_batches": 2.2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/us/gpt-5-pro": { @@ -68514,6 +69218,11 @@ "mode": "responses", "output_cost_per_token": 0.000132, "output_cost_per_token_batches": 6.6e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/us/gpt-5.1-codex-max": { @@ -68538,6 +69247,11 @@ "output_cost_per_token": 1.54e-05, "output_cost_per_token_batches": 7.7e-06, "output_cost_per_token_priority": 3.08e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/us/gpt-5.2-codex": { @@ -68556,6 +69270,11 @@ "mode": "responses", "output_cost_per_token": 0.0001848, "output_cost_per_token_batches": 9.24e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/us/gpt-5.3-codex": { @@ -68568,6 +69287,11 @@ "mode": "responses", "output_cost_per_token": 1.54e-05, "output_cost_per_token_priority": 3.08e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/us/gpt-5.4-mini": { @@ -68583,6 +69307,11 @@ "output_cost_per_token": 4.95e-06, "output_cost_per_token_batches": 2.475e-06, "output_cost_per_token_priority": 9.9e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/us/gpt-5.4-nano": { @@ -68595,6 +69324,11 @@ "mode": "chat", "output_cost_per_token": 1.375e-06, "output_cost_per_token_batches": 6.875e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/us/gpt-5.4-pro": { @@ -68609,6 +69343,11 @@ "output_cost_per_token_above_272k_tokens": 0.000297, "output_cost_per_token_above_272k_tokens_batches": 0.0001485, "output_cost_per_token_batches": 9.9e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "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/us/o1-mini": { @@ -68655,7 +69394,7 @@ "cache_creation_input_audio_token_cost": 4e-07, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, - "deprecation_date": "2026-10-31", + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image_token": 5e-06, "input_cost_per_token": 4e-06, @@ -71000,6 +71739,7 @@ "xai/grok-code-fast-1": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "deprecation_date": "2026-05-15", "input_cost_per_image_token": 1e-06, "input_cost_per_token": 1e-06, "input_cost_per_token_above_200k_tokens": 2e-06, @@ -71896,7 +72636,7 @@ "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, - "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost": 1e-07, "deprecation_date": "2027-09-29", "input_cost_per_token": 2e-06, "litellm_provider": "azure_ai", @@ -72503,15 +73243,15 @@ "gpt-6.1-sol": { "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_12hr": 5e-06, + "cache_creation_input_token_cost_above_12hr_ultrafast": 3e-05, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 6e-05, - "cache_creation_input_token_cost_above_12hr_batches": 2.5e-06, "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, "cache_creation_input_token_cost_above_12hr_priority": 1e-05, - "cache_creation_input_token_cost_above_12hr_ultrafast": 3e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_batches": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, @@ -74472,14 +75212,14 @@ "supports_vision": true }, "wandb/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B": { - "cache_read_input_token_cost": 2e-08, - "input_cost_per_token": 1.5e-07, + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 2.5e-07, "litellm_provider": "wandb", "max_input_tokens": 1049000, "max_output_tokens": 943718, "mode": "chat", - "output_cost_per_token": 7.5e-07, - "source": "https://openrouter.ai/api/v1/models/nvidia/nemotron-3.5-super-vl/endpoints", + "output_cost_per_token": 1.5e-06, + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -74577,6 +75317,18 @@ "supports_tool_choice": true, "supports_vision": true }, + "fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-1017-spot": { + "cache_read_input_token_cost": 4.2e-09, + "input_cost_per_token": 2.1e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 8.4e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_vision": true + }, "github_copilot/claude-sonnet-4.6": { "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, @@ -76331,6 +77083,15 @@ "max" ] }, + "fireworks_ai/accounts/fireworks/models/rune": { + "cache_read_input_token_cost": 8.4e-09, + "input_cost_per_token": 4.2e-08, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 262144, + "mode": "decisions", + "output_cost_per_token": 0.0, + "source": "https://api.fireworks.ai/v1/serverless/models?format=nested" + }, "bedrock_mantle/openai.gpt-6.1-terra": { "cache_creation_input_token_cost": 1.375e-06, "cache_creation_input_token_cost_above_272k_tokens": 2.75e-06, diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index 7291c4858f6..e2452a928a4 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -1,4 +1,3 @@ - from typing import Final import pytest @@ -19,9 +18,7 @@ def test_web_search_cost_low(): web_search_options=web_search_options, model_info=model_info ) - assert ( - cost == model_info["search_context_cost_per_query"]["search_context_size_low"] - ) + assert cost == model_info["search_context_cost_per_query"]["search_context_size_low"] def test_web_search_cost_medium(): @@ -32,10 +29,7 @@ def test_web_search_cost_medium(): web_search_options=web_search_options, model_info=model_info ) - assert ( - cost - == model_info["search_context_cost_per_query"]["search_context_size_medium"] - ) + assert cost == model_info["search_context_cost_per_query"]["search_context_size_medium"] def test_web_search_cost_high(): @@ -46,9 +40,7 @@ def test_web_search_cost_high(): web_search_options=web_search_options, model_info=model_info ) - assert ( - cost == model_info["search_context_cost_per_query"]["search_context_size_high"] - ) + assert cost == model_info["search_context_cost_per_query"]["search_context_size_high"] def test_get_web_search_options_reads_tools_when_options_are_null() -> None: @@ -60,9 +52,7 @@ def test_get_web_search_options_reads_tools_when_options_are_null() -> None: def test_get_web_search_options_accepts_non_list_tool_iterables(): - options = StandardBuiltInToolCostTracking.get_web_search_options( - {"tools": ({"type": "web_search_preview"},)} - ) + options = StandardBuiltInToolCostTracking.get_web_search_options({"tools": ({"type": "web_search_preview"},)}) assert options == {"type": "web_search_preview"} @@ -76,25 +66,15 @@ def test_get_file_search_tool_call_does_not_validate_tool_payload(): # Test file search cost calculation def test_file_search_cost(): file_search = FileSearchTool(type="file_search") - cost = StandardBuiltInToolCostTracking.get_cost_for_file_search( - file_search=file_search - ) + cost = StandardBuiltInToolCostTracking.get_cost_for_file_search(file_search=file_search) assert cost == 0.0025 # $2.50/1000 calls = 0.0025 per call # Test edge cases def test_none_inputs(): # Test with None inputs - assert ( - StandardBuiltInToolCostTracking.get_cost_for_web_search( - web_search_options=None, model_info=None - ) - == 0.0 - ) - assert ( - StandardBuiltInToolCostTracking.get_cost_for_file_search(file_search=None) - == 0.0 - ) + assert StandardBuiltInToolCostTracking.get_cost_for_web_search(web_search_options=None, model_info=None) == 0.0 + assert StandardBuiltInToolCostTracking.get_cost_for_file_search(file_search=None) == 0.0 # Test the main get_cost_for_built_in_tools method @@ -119,9 +99,7 @@ def test_get_cost_for_built_in_tools_file_search(): Test that the cost for a file search is 0.00 when no response object is provided """ model = "gpt-4" - standard_built_in_tools_params = StandardBuiltInToolsParams( - file_search=FileSearchTool(type="file_search") - ) + standard_built_in_tools_params = StandardBuiltInToolsParams(file_search=FileSearchTool(type="file_search")) cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( model=model, @@ -144,9 +122,7 @@ def test_get_cost_for_anthropic_web_search_with_server_tool_use_dict(): usage = Usage(server_tool_use={"web_search_requests": 1}) assert isinstance(usage.server_tool_use, ServerToolUse) - assert StandardBuiltInToolCostTracking.response_object_includes_web_search_call( - response_object=None, usage=usage - ) + assert StandardBuiltInToolCostTracking.response_object_includes_web_search_call(response_object=None, usage=usage) def test_anthropic_web_search_zero_requests_from_raw_response_charges_zero(): @@ -257,9 +233,7 @@ def test_completion_cost_includes_web_search_without_standard_built_in_tools_par ) assert web_search_cost > 0, "Web search cost should be non-zero" - assert ( - cost >= web_search_cost - ), f"completion_cost ({cost}) should include web search cost ({web_search_cost})" + assert cost >= web_search_cost, f"completion_cost ({cost}) should include web search cost ({web_search_cost})" @pytest.mark.parametrize( @@ -286,18 +260,14 @@ def test_gemini_3x_web_search_billed_per_query(model, local_model_cost_map): web_search_requests = 2 model_info = litellm.get_model_info(model) assert model_info["web_search_billing_unit"] == "per_query" - per_query_cost = model_info["search_context_cost_per_query"][ - "search_context_size_medium" - ] + per_query_cost = model_info["search_context_cost_per_query"]["search_context_size_medium"] expected_cost = per_query_cost * web_search_requests usage = Usage( prompt_tokens=11, completion_tokens=100, total_tokens=111, - prompt_tokens_details=PromptTokensDetailsWrapper( - text_tokens=11, web_search_requests=web_search_requests - ), + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=11, web_search_requests=web_search_requests), ) cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( @@ -309,8 +279,7 @@ def test_gemini_3x_web_search_billed_per_query(model, local_model_cost_map): ) assert cost == pytest.approx(expected_cost), ( - f"Expected {web_search_requests} x ${per_query_cost} = ${expected_cost} " - f"per_query search fee, got ${cost}" + f"Expected {web_search_requests} x ${per_query_cost} = ${expected_cost} per_query search fee, got ${cost}" ) @@ -353,17 +322,13 @@ def test_gemini_2x_web_search_still_billed_per_prompt(local_model_cost_map): model = "vertex_ai/gemini-2.5-flash" model_info = litellm.get_model_info(model) assert not model_info.get("web_search_billing_unit") - expected_cost = model_info["search_context_cost_per_query"][ - "search_context_size_medium" - ] + expected_cost = model_info["search_context_cost_per_query"]["search_context_size_medium"] usage = Usage( prompt_tokens=11, completion_tokens=100, total_tokens=111, - prompt_tokens_details=PromptTokensDetailsWrapper( - text_tokens=11, web_search_requests=2 - ), + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=11, web_search_requests=2), ) cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( @@ -375,8 +340,7 @@ def test_gemini_2x_web_search_still_billed_per_prompt(local_model_cost_map): ) assert cost == pytest.approx(expected_cost), ( - f"Expected flat ${expected_cost} per_prompt search fee (2 queries clamped to 1), " - f"got ${cost}" + f"Expected flat ${expected_cost} per_prompt search fee (2 queries clamped to 1), got ${cost}" ) @@ -402,9 +366,7 @@ def test_web_search_provider_prefix_fallback_does_not_misprice_non_gemini_model( prompt_tokens=11, completion_tokens=100, total_tokens=111, - prompt_tokens_details=PromptTokensDetailsWrapper( - text_tokens=11, web_search_requests=2 - ), + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=11, web_search_requests=2), ) cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( @@ -427,7 +389,6 @@ def _openai_responses_with_web_search_calls(model, num_calls): ResponseFunctionWebSearch, ) - output = [ ResponseFunctionWebSearch( id=f"ws_{i}", @@ -458,9 +419,7 @@ def test_openai_responses_web_search_multiplied_by_call_count(local_model_cost_m from litellm.types.utils import Usage model = "gpt-4o-search-preview" - per_call = litellm.get_model_info(model)["search_context_cost_per_query"][ - "search_context_size_medium" - ] + per_call = litellm.get_model_info(model)["search_context_cost_per_query"]["search_context_size_medium"] usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) for num_calls in (1, 3): @@ -487,9 +446,7 @@ def test_web_search_call_count_reads_dict_output_items(local_model_cost_map): from litellm.types.utils import Usage model = "gpt-4o-search-preview" - per_call = litellm.get_model_info(model)["search_context_cost_per_query"][ - "search_context_size_medium" - ] + per_call = litellm.get_model_info(model)["search_context_cost_per_query"]["search_context_size_medium"] response = ResponsesAPIResponse.model_validate( { @@ -498,10 +455,7 @@ def test_web_search_call_count_reads_dict_output_items(local_model_cost_map): "model": model, "object": "response", "status": "completed", - "output": [ - {"type": "web_search_call", "id": f"ws_{i}", "status": "completed"} - for i in range(3) - ], + "output": [{"type": "web_search_call", "id": f"ws_{i}", "status": "completed"} for i in range(3)], } ) assert all(isinstance(item, dict) for item in response.output) @@ -514,9 +468,7 @@ def test_web_search_call_count_reads_dict_output_items(local_model_cost_map): standard_built_in_tools_params=None, ) - assert cost == pytest.approx(3 * per_call), ( - f"3 dict-shaped web searches must bill 3 x ${per_call}, got ${cost}" - ) + assert cost == pytest.approx(3 * per_call), f"3 dict-shaped web searches must bill 3 x ${per_call}, got ${cost}" # Note: File search integration test removed due to complex annotation detection logic @@ -596,8 +548,6 @@ _BEDROCK_MANTLE_WEB_SEARCH_MODELS = ( _BEDROCK_MANTLE_WEB_SEARCH_RATE = 0.012 - - def _openai_responses_response(model, output, usage=None, tools=None): return ResponsesAPIResponse.model_validate( { @@ -815,6 +765,31 @@ def test_responses_usage_tool_usage_web_search_billed_without_output_item(local_ assert cost == pytest.approx(num_requests * per_call) +@pytest.mark.parametrize("model", ["azure/gpt-5-mini", "azure/gpt-5.6-luna"]) +def test_azure_responses_web_search_billed_per_bing_request(local_model_cost_map, model): + """Azure OpenAI web_search is Grounding with Bing: $0.014 per reported Bing request, including rows + whose OpenAI twin prices search but the Azure row used to bill it at $0.""" + for num_requests in (1, 5): + response = _openai_responses_response( + model, + [dict(_ASSISTANT_MESSAGE_OUTPUT_ITEM)], + usage={ + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + "tool_usage": {"web_search": {"num_requests": num_requests}}, + }, + ) + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model=model, + response_object=response, + usage=None, + custom_llm_provider="azure", + standard_built_in_tools_params=None, + ) + assert cost == pytest.approx(num_requests * 0.014) + + def _image_gen_token_usage(input_text, input_image, output_image, output_text): return { "input_tokens": 10, @@ -1012,6 +987,8 @@ def test_responses_image_gen_total_without_token_details_falls_back_to_per_image ) ) assert cost > 0 + + @pytest.mark.parametrize( ("model_key", "web_search_requests", "medium_cost_per_query", "expected_cost"), [ diff --git a/tests/unit/test_model_prices_and_context_window.py b/tests/unit/test_model_prices_and_context_window.py index 08a39b94c12..bbfe304429d 100644 --- a/tests/unit/test_model_prices_and_context_window.py +++ b/tests/unit/test_model_prices_and_context_window.py @@ -224,6 +224,41 @@ def test_dated_variants_carry_base_alias_service_tier_pricing(prices: dict): ) +# Azure OpenAI web_search runs on Grounding with Bing and bills each Bing request reported in +# tool_usage.web_search.num_requests: https://learn.microsoft.com/en-us/azure/foundry/openai/how-to/web-search +# Grounding with Bing lists $14 per 1,000 transactions: https://www.microsoft.com/en-us/bing/apis/grounding-pricing +AZURE_WEB_SEARCH_COST_PER_REQUEST: Final = 0.014 +_AZURE_DEPLOYMENT_PREFIX: Final = re.compile(r"^azure/(us/|eu/|global/|global-standard/)?") + + +def _azure_rows(prices: Mapping[str, object]) -> Mapping[str, dict]: + return MappingProxyType( + {key: row for key, row in prices.items() if key.startswith("azure/") and isinstance(row, dict)} + ) + + +def test_azure_web_search_rows_charge_the_grounding_with_bing_price(prices: dict): + wrong = { + key: row["search_context_cost_per_query"] + for key, row in _azure_rows(prices).items() + if row.get("search_context_cost_per_query") + and set(row["search_context_cost_per_query"].values()) != {AZURE_WEB_SEARCH_COST_PER_REQUEST} + } + assert wrong == {} + + +def test_azure_rows_price_web_search_when_their_openai_model_does(prices: dict): + missing = sorted( + key + for key, row in _azure_rows(prices).items() + if not row.get("search_context_cost_per_query") + and row.get("mode") in ("chat", "responses", None) + and isinstance(twin := prices.get(_AZURE_DEPLOYMENT_PREFIX.sub("", key)), dict) + and twin.get("search_context_cost_per_query") + ) + assert missing == [] + + PRIORITY_SUFFIX: Final = "_priority" MIN_PRIORITY_MARKUP: Final = 1.0 MAX_PRIORITY_MARKUP: Final = 5.0 @@ -845,7 +880,10 @@ def test_volcengine_doubao_seed_2_1_pro_260915_row_matches_published_metadata(pa 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")) + assert all( + entry[key] == sibling[key] + for key in ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost") + ) def _bedrock_gpt_oss_rows(prices: Mapping[str, JsonValue]) -> tuple[str, ...]: From ed5fe048ff652242dcc0241b8a81ee253f81d7e2 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Wed, 28 Oct 2026 09:19:36 -0700 Subject: [PATCH 041/128] fix(proxy/auth): apply max_end_user_budget_id to end users without a row on cached keys (#48704) * test(proxy/auth): cover default budgets on cached virtual keys Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy/auth): use in-memory cache in cached budget tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy/auth): apply max_end_user_budget_id to end users without a row on cached keys Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy/auth): share cached end-user parameter update 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> --- litellm/proxy/auth/user_api_key_auth.py | 3 +- .../test_user_api_key_auth_request_flow.py | 91 +++++++++++++++++++ 2 files changed, 93 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index cb490b55352..99edfb940d7 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -3257,8 +3257,9 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e if valid_token.budget_throttle_pct is not None: valid_token_dict["budget_throttle_pct"] = valid_token.budget_throttle_pct - if _end_user_object is not None: + if _end_user_object is not None or litellm.max_end_user_budget_id is not None: valid_token_dict.update(end_user_params) + if _end_user_object is not None: valid_token_dict["end_user_object_permission"] = _end_user_object.object_permission # check if token is from litellm-ui, litellm ui makes keys to allow users to login with sso. These keys can only be used for LiteLLM UI functions diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index 2d9c781b0d1..ae7168ed1e1 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -1537,6 +1537,97 @@ async def test_cached_proxy_admin_key_sets_via_virtual_key_marker(): setattr(_proxy_server_mod, attr, val) +@pytest.mark.parametrize( + ("default_budget_id", "expected_limits"), + [ + pytest.param("proxy-default-end-user-budget", (0, 7), id="default-configured"), + pytest.param(None, (None, None), id="no-default"), + ], +) +@pytest.mark.asyncio +async def test_cached_virtual_key_applies_proxy_default_budget_to_new_end_user( + monkeypatch: pytest.MonkeyPatch, + default_budget_id: str | None, + expected_limits: tuple[int | None, int | None], +) -> None: + """Cached virtual keys apply configured limits to new end users.""" + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy.proxy_server import hash_token + + end_user_id: Final = "cached-virtual-key-end-user" + monkeypatch.setattr(litellm, "max_end_user_budget_id", default_budget_id) + + api_key: Final = "sk-cached-default-end-user-budget-test" + hashed_key: Final = hash_token(api_key) + cached_token: Final = UserAPIKeyAuth( + api_key=api_key, + token=hashed_key, + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + budget: Final = LiteLLM_BudgetTable( + budget_id="proxy-default-end-user-budget", + rpm_limit=0, + tpm_limit=7, + ) + mock_prisma_client: Final = MagicMock() + mock_prisma_client.db.litellm_endusertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=budget) + + user_api_key_cache: Final = UserApiKeyCache() + + mock_proxy_logging_obj: Final = MagicMock() + mock_proxy_logging_obj.internal_usage_cache = MagicMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() + mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + + import litellm.proxy.proxy_server as _proxy_server_mod + + monkeypatch.setattr(_proxy_server_mod, "prisma_client", mock_prisma_client) + monkeypatch.setattr(_proxy_server_mod, "user_api_key_cache", user_api_key_cache) + monkeypatch.setattr(_proxy_server_mod, "proxy_logging_obj", mock_proxy_logging_obj) + monkeypatch.setattr(_proxy_server_mod, "master_key", "sk-master-key") + monkeypatch.setattr(_proxy_server_mod, "general_settings", {}) + monkeypatch.setattr(_proxy_server_mod, "llm_model_list", []) + monkeypatch.setattr(_proxy_server_mod, "llm_router", None) + monkeypatch.setattr(_proxy_server_mod, "open_telemetry_logger", None) + monkeypatch.setattr(_proxy_server_mod, "model_max_budget_limiter", MagicMock()) + monkeypatch.setattr(_proxy_server_mod, "user_custom_auth", None) + monkeypatch.setattr(_proxy_server_mod, "jwt_handler", None) + monkeypatch.setattr(_proxy_server_mod, "litellm_proxy_admin_name", "admin") + + request: Final = Request( + scope={ + "type": "http", + "headers": [(b"x-litellm-end-user-id", end_user_id.encode())], + } + ) + request._url = URL(url="/chat/completions") + + with patch( + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new_callable=AsyncMock, + return_value=cached_token, + ): + result: Final = await user_api_key_auth_builder( + request=request, + api_key=f"Bearer {api_key}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + + assert (result.end_user_rpm_limit, result.end_user_tpm_limit) == expected_limits + if default_budget_id is not None: + assert result.end_user_id == end_user_id + + @pytest.mark.asyncio async def test_master_key_auth_sets_via_virtual_key_marker(): """Master-key requests must also be stamped by overwrite_user_with_key_hash; From 424e5c069c33d18d25b579d2193c4c54ce9d91f9 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Wed, 28 Oct 2026 09:22:10 -0700 Subject: [PATCH 042/128] refactor(http): build outbound requests as reqwest::Request (#48041) * refactor(http): build outbound requests as reqwest::Request OutboundRequest was a hand-rolled copy of reqwest::Request. Routes now build a reqwest::Request directly and RequestSigner reads the body bytes back out of it, so the signer still covers exactly what is sent. SigV4 signing moves onto HeaderMap and takes the method, which drops the BTreeMap plumbing and the header casing table Co-Authored-By: Claude Opus 5.5 * fix(http): normalize provider URLs before building outbound requests --------- Co-authored-by: Claude Opus 5.5 Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 3 + litellm-rust/crates/auth-aws/Cargo.toml | 4 +- litellm-rust/crates/auth-aws/src/aws.rs | 172 +++++----- litellm-rust/crates/auth-aws/src/signer.rs | 152 +++++---- litellm-rust/crates/http/src/error.rs | 2 + litellm-rust/crates/http/src/outbound.rs | 298 ++++++++++-------- litellm-rust/crates/http/src/request.rs | 12 + litellm-rust/crates/inference-chat/Cargo.toml | 4 +- .../crates/inference-chat/src/handler.rs | 7 +- .../crates/inference-messages/src/handler.rs | 6 +- .../inference-ocr/tests/ocr/aws_textract.rs | 32 +- .../crates/inference-responses/src/handler.rs | 2 +- .../inference-transcription/src/handler.rs | 2 +- litellm-rust/crates/inference/src/outbound.rs | 15 +- .../crates/llms/src/base_llm/ocr/handler.rs | 50 ++- .../llms/src/base_llm/ocr/transformation.rs | 16 +- .../llms/src/reducto/ocr/transformation.rs | 17 +- 17 files changed, 446 insertions(+), 348 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 71ed05c0adf..e998fe1abbb 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3385,6 +3385,8 @@ dependencies = [ "aws-sigv4", "aws-smithy-runtime-api", "aws-types", + "bytes", + "futures-util", "litellm-auth-types", "litellm-http", "moka", @@ -4041,6 +4043,7 @@ dependencies = [ "litellm-llms-types", "litellm-secrets", "litellm-tracing", + "reqwest", "rstest", "serde_json", "tokio", diff --git a/litellm-rust/crates/auth-aws/Cargo.toml b/litellm-rust/crates/auth-aws/Cargo.toml index 65b411ffaef..1a7e2cd6782 100644 --- a/litellm-rust/crates/auth-aws/Cargo.toml +++ b/litellm-rust/crates/auth-aws/Cargo.toml @@ -8,6 +8,7 @@ repository.workspace = true [dependencies] litellm-auth-types.workspace = true litellm-http.workspace = true +reqwest.workspace = true moka = { workspace = true, features = ["sync"] } serde.workspace = true @@ -26,5 +27,6 @@ aws-smithy-runtime-api = "1.13.0" [dev-dependencies] rstest.workspace = true litellm-http = { workspace = true, features = ["test-support"] } -reqwest.workspace = true +bytes.workspace = true +futures-util.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/auth-aws/src/aws.rs b/litellm-rust/crates/auth-aws/src/aws.rs index eaa033d631f..d83b01c24ad 100644 --- a/litellm-rust/crates/auth-aws/src/aws.rs +++ b/litellm-rust/crates/auth-aws/src/aws.rs @@ -1,4 +1,3 @@ -use std::collections::BTreeMap; use std::time::Duration; use std::time::{SystemTime, UNIX_EPOCH}; @@ -12,6 +11,7 @@ use aws_sigv4::http_request::{ }; use aws_sigv4::sign::v4; use aws_smithy_runtime_api::client::identity::Identity; +use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; use litellm_auth_types::AwsParams; @@ -435,12 +435,12 @@ fn default_session_name() -> String { /// client header cannot change the canonical request and invalidate the /// signature. Signing everything instead makes the request 403 on a header the /// caller supplied, on a deployment that works on the Python path. -pub fn aws_signature_headers(headers: &BTreeMap) -> BTreeMap { +pub fn aws_signature_headers(headers: &HeaderMap) -> HeaderMap { headers .iter() .filter(|(name, _)| { - let name = name.to_ascii_lowercase(); - AWS_SIGNED_HEADER_NAMES.contains(&name.as_str()) + let name = name.as_str(); + AWS_SIGNED_HEADER_NAMES.contains(&name) || name.starts_with("x-amz-") || name.starts_with("x-amzn-") }) @@ -456,45 +456,63 @@ pub fn is_sigv4_computed_header(name: &str) -> bool { SIGV4_COMPUTED_HEADER_NAMES.contains(&name.to_ascii_lowercase().as_str()) } -pub fn sign_post( +#[derive(Clone, Copy, Debug)] +pub struct SigningScope<'a> { + pub region: &'a str, + pub service: &'a str, + pub credentials: &'a Credentials, + pub time: SystemTime, +} + +pub fn sign_headers( + method: &str, url: &str, body: &[u8], - headers: &BTreeMap, - region: &str, - service: &str, - credentials: &Credentials, - signing_time: SystemTime, -) -> Result, Error> { - let identity: Identity = credentials.clone().into(); + headers: &HeaderMap, + scope: SigningScope<'_>, +) -> Result { + let identity: Identity = scope.credentials.clone().into(); let params = v4::SigningParams::builder() .identity(&identity) - .region(region) - .name(service) - .time(signing_time) + .region(scope.region) + .name(scope.service) + .time(scope.time) .settings(SigningSettings::default()) .build() .map(SigningParams::from) .map_err(|error| Error::AwsSigningParameters(error.to_string()))?; let header_refs = headers .iter() - .map(|(name, value)| (name.as_str(), value.as_str())); - let request = SignableRequest::new("POST", url, header_refs, SignableBody::Bytes(body)) - .map_err(|error| Error::AwsSignableRequest(error.to_string()))?; + .map(|(name, value)| { + value + .to_str() + .map(|value| (name.as_str(), value)) + .map_err(|_| Error::AwsSignableRequest(format!("{name} is not visible ASCII"))) + }) + .collect::, _>>()?; + let request = SignableRequest::new( + method, + url, + header_refs.into_iter(), + SignableBody::Bytes(body), + ) + .map_err(|error| Error::AwsSignableRequest(error.to_string()))?; let (instructions, _) = sign(request, ¶ms) .map_err(|error| Error::AwsSigning(error.to_string()))? .into_parts(); - Ok(instructions + instructions .headers() .map(|(name, value)| { - let normalized_name = match name { - "authorization" => "Authorization", - "x-amz-date" => "X-Amz-Date", - "x-amz-security-token" => "X-Amz-Security-Token", - _ => name, - }; - (normalized_name.to_string(), value.to_string()) + HeaderName::from_bytes(name.as_bytes()) + .map_err(|error| error.to_string()) + .and_then(|name| { + HeaderValue::from_str(value) + .map(|value| (name, value)) + .map_err(|error| error.to_string()) + }) }) - .collect()) + .collect::>() + .map_err(Error::AwsSigning) } /// Model-id and region parsing shared by every Bedrock route. @@ -639,6 +657,7 @@ pub fn host_supplied_credentials(params: &AwsParams) -> Option { mod tests { use super::*; use crate::constants::{AWS_REGION, BEDROCK_SERVICE}; + use reqwest::header::CONTENT_TYPE; fn no_env(_: &str) -> Option { None @@ -731,12 +750,12 @@ mod tests { ); } - fn parity_inputs() -> (String, Vec, BTreeMap) { + fn parity_inputs() -> (String, Vec, HeaderMap) { ( "https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-express-v1/invoke" .to_string(), br#"{"input":"hello"}"#.to_vec(), - BTreeMap::from([("Content-Type".to_string(), "application/json".to_string())]), + HeaderMap::from_iter([(CONTENT_TYPE, HeaderValue::from_static("application/json"))]), ) } @@ -961,9 +980,9 @@ mod tests { // cannot change the canonical request. Signing it instead makes the // request 403 the moment anything on the wire rewrites or drops it. let (url, body, mut headers) = parity_inputs(); - headers.insert("x-request-id".to_string(), "abc-123".to_string()); - headers.insert("Accept-Encoding".to_string(), "gzip".to_string()); - headers.insert("x-amzn-trace-id".to_string(), "Root=1-abc".to_string()); + headers.insert("x-request-id", HeaderValue::from_static("abc-123")); + headers.insert("Accept-Encoding", HeaderValue::from_static("gzip")); + headers.insert("x-amzn-trace-id", HeaderValue::from_static("Root=1-abc")); let signable = aws_signature_headers(&headers); assert!(!signable.contains_key("x-request-id")); @@ -979,18 +998,21 @@ mod tests { None, "test", ); - let signed = sign_post( + let signed = sign_headers( + "POST", &url, &body, &signable, - "us-east-1", - BEDROCK_SERVICE, - &credentials, - SystemTime::UNIX_EPOCH, + SigningScope { + region: "us-east-1", + service: BEDROCK_SERVICE, + credentials: &credentials, + time: SystemTime::UNIX_EPOCH, + }, ) .expect("signs"); - let authorization = signed - .get("Authorization") + let authorization = signed["Authorization"] + .to_str() .expect("carries an authorization header"); assert!( !authorization.contains("x-request-id"), @@ -1012,29 +1034,24 @@ mod tests { None, "test", ); - let signed = sign_post( + let signed = sign_headers( + "POST", &url, &body, &headers, - "us-east-1", - BEDROCK_SERVICE, - &credentials, - UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645), + SigningScope { + region: "us-east-1", + service: BEDROCK_SERVICE, + credentials: &credentials, + time: UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645), + }, ) .expect("golden signature"); + assert_eq!(signed["X-Amz-Date"], "20240102T030405Z"); + assert_eq!(signed["X-Amz-Security-Token"], "session-token"); assert_eq!( - signed.get("X-Amz-Date").map(String::as_str), - Some("20240102T030405Z") - ); - assert_eq!( - signed.get("X-Amz-Security-Token").map(String::as_str), - Some("session-token") - ); - assert_eq!( - signed.get("Authorization").map(String::as_str), - Some( - "AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240102/us-east-1/bedrock/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, Signature=55c027ef47527d3ad63f1735f9d099efdbc99f296ff914bd94e727e24ec0e464" - ) + signed["Authorization"], + "AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240102/us-east-1/bedrock/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, Signature=55c027ef47527d3ad63f1735f9d099efdbc99f296ff914bd94e727e24ec0e464" ); } @@ -1048,14 +1065,17 @@ mod tests { None, "test", ); - let signed = sign_post( + let signed = sign_headers( + "POST", &url, &body, &headers, - "us-east-1", - BEDROCK_SERVICE, - &credentials, - UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645), + SigningScope { + region: "us-east-1", + service: BEDROCK_SERVICE, + credentials: &credentials, + time: UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645), + }, ) .expect("signature"); assert!(!signed.contains_key("X-Amz-Security-Token")); @@ -1068,7 +1088,7 @@ mod tests { let secret_access_key = std::env::var("AWS_BEDROCK_TEST_SECRET_ACCESS_KEY")?; let body = br#"{"anthropic_version":"bedrock-2023-05-31","max_tokens":1,"messages":[{"role":"user","content":[{"type":"text","text":"ping"}]}]}"#.to_vec(); let headers = - BTreeMap::from([("Content-Type".to_string(), "application/json".to_string())]); + HeaderMap::from_iter([(CONTENT_TYPE, HeaderValue::from_static("application/json"))]); let credentials = AwsAuthService::default() .resolve_credentials( AwsAuthConfig { @@ -1087,23 +1107,25 @@ mod tests { let url = format!( "https://bedrock-runtime.{region}.amazonaws.com/model/us.anthropic.claude-opus-4-8/invoke" ); - let signed_headers = sign_post( + let signed_headers = sign_headers( + "POST", &url, &body, &headers, - region, - BEDROCK_SERVICE, - &credentials, - SystemTime::now(), + SigningScope { + region, + service: BEDROCK_SERVICE, + credentials: &credentials, + time: SystemTime::now(), + }, )?; - let mut request = client.post(&url).body(body.clone()); - for (name, value) in &headers { - request = request.header(name, value); - } - for (name, value) in signed_headers { - request = request.header(name, value); - } - let response = request.send().await?; + let response = client + .post(&url) + .headers(headers.clone()) + .headers(signed_headers) + .body(body.clone()) + .send() + .await?; let status = response.status(); let response_body = response.text().await?; let snippet: String = response_body.chars().take(240).collect(); diff --git a/litellm-rust/crates/auth-aws/src/signer.rs b/litellm-rust/crates/auth-aws/src/signer.rs index 3868fdd939b..6c2cb4fd162 100644 --- a/litellm-rust/crates/auth-aws/src/signer.rs +++ b/litellm-rust/crates/auth-aws/src/signer.rs @@ -1,11 +1,13 @@ -use std::{collections::BTreeMap, time::SystemTime}; +use std::time::SystemTime; + +use aws_credential_types::Credentials; +use reqwest::{Request, header::HeaderMap}; use crate::{ - AwsAuthService, AwsCredentialSource, Error, aws_signature_headers, is_sigv4_computed_header, - sign_post, + AwsAuthService, AwsCredentialSource, Error, SigningScope, aws_signature_headers, + is_sigv4_computed_header, sign_headers, }; -use aws_credential_types::Credentials; -use litellm_http::outbound::{RequestSigner, UnsignedRequest}; +use litellm_http::outbound::RequestSigner; #[derive(Clone, Debug)] pub struct SigV4Signer { @@ -45,41 +47,48 @@ impl SigV4Signer { } impl RequestSigner for SigV4Signer { - fn sign( - &self, - request: UnsignedRequest<'_>, - ) -> Result, litellm_http::Error> { - if let Some((name, _)) = request - .headers - .iter() - .find(|(name, _)| is_sigv4_computed_header(name)) + fn sign(&self, request: &Request) -> Result { + if let Some(name) = request + .headers() + .keys() + .find(|name| is_sigv4_computed_header(name.as_str())) { - return Err(litellm_http::Error::ComputedHeader(name.clone())); + return Err(litellm_http::Error::ComputedHeader(name.to_string())); } - let headers: BTreeMap = request.headers.iter().cloned().collect(); - sign_post( - request.url, - request.body, - &aws_signature_headers(&headers), - &self.region, - self.service, - &self.credentials, - (self.clock)(), + let body = match request.body() { + None => &[], + Some(body) => body.as_bytes().ok_or_else(|| { + litellm_http::Error::Signature("a streaming body cannot be signed".into()) + })?, + }; + sign_headers( + request.method().as_str(), + request.url().as_str(), + body, + &aws_signature_headers(request.headers()), + SigningScope { + region: &self.region, + service: self.service, + credentials: &self.credentials, + time: (self.clock)(), + }, ) - .map(|signature| signature.into_iter().collect()) .map_err(|error| litellm_http::Error::Signature(error.to_string())) } } #[cfg(test)] mod tests { + use reqwest::{Body, Method, Url}; + use serde_json::{Value, json as json_value}; use std::time::{Duration, UNIX_EPOCH}; - use litellm_http::outbound::OutboundRequest; - use serde_json::{Value, json}; + use litellm_http::outbound::{json, signed}; use super::*; + const URL: &str = "https://textract.us-east-1.amazonaws.com/"; + fn fixed_clock() -> SystemTime { UNIX_EPOCH + Duration::from_secs(1_700_000_000) } @@ -93,75 +102,83 @@ mod tests { .with_clock(fixed_clock) } + fn target() -> (String, String) { + ("X-Amz-Target".into(), "Textract.DetectDocumentText".into()) + } + fn authorization(body: &Value, service: &'static str) -> String { - OutboundRequest::signed_json( - "https://textract.us-east-1.amazonaws.com/".into(), - vec![("X-Amz-Target".into(), "Textract.DetectDocumentText".into())], - body, - None, + signed( + json(URL, &[target()], body, None).unwrap(), &signer(service), ) .unwrap() - .header("Authorization") - .unwrap() - .to_string() + .headers()["authorization"] + .to_str() + .unwrap() + .to_string() } #[test] fn the_signature_verifies_against_the_bytes_that_are_sent() { - let sent = OutboundRequest::signed_json( - "https://textract.us-east-1.amazonaws.com/".into(), - vec![("X-Amz-Target".into(), "Textract.DetectDocumentText".into())], - &json!({"Document": {"Bytes": "aGk="}}), - None, + let sent = signed( + json( + URL, + &[target()], + &json_value!({"Document": {"Bytes": "aGk="}}), + None, + ) + .unwrap(), &signer("textract"), ) .unwrap(); - let unsigned: BTreeMap = sent + let unsigned: HeaderMap = sent .headers() .iter() - .filter(|(name, _)| !is_sigv4_computed_header(name)) - .cloned() + .filter(|(name, _)| !is_sigv4_computed_header(name.as_str())) + .map(|(name, value)| (name.clone(), value.clone())) .collect(); - let recomputed = sign_post( - sent.url(), - sent.body(), + let recomputed = sign_headers( + "POST", + sent.url().as_str(), + sent.body().unwrap().as_bytes().unwrap(), &aws_signature_headers(&unsigned), - "us-east-1", - "textract", - &Credentials::new("AKIDEXAMPLE", "secret", None, None, "test"), - fixed_clock(), + SigningScope { + region: "us-east-1", + service: "textract", + credentials: &Credentials::new("AKIDEXAMPLE", "secret", None, None, "test"), + time: fixed_clock(), + }, ) .unwrap(); - assert_eq!( - sent.header("Authorization"), - Some(recomputed["Authorization"].as_str()) - ); + assert_eq!(sent.headers()["authorization"], recomputed["authorization"]); } #[test] fn the_signature_depends_on_the_body_and_the_service() { - let original = authorization(&json!({"text": "card 4111"}), "textract"); + let original = authorization(&json_value!({"text": "card 4111"}), "textract"); assert_ne!( original, - authorization(&json!({"text": "card [REDACTED]"}), "textract") + authorization(&json_value!({"text": "card [REDACTED]"}), "textract") ); assert_ne!( original, - authorization(&json!({"text": "card 4111"}), "bedrock") + authorization(&json_value!({"text": "card 4111"}), "bedrock") ); assert!(original.contains("/us-east-1/textract/aws4_request")); } #[test] fn a_forwarded_computed_header_is_refused_instead_of_sent_twice() { - let error = OutboundRequest::signed_json( - "https://textract.us-east-1.amazonaws.com/".into(), - vec![("authorization".into(), "Bearer caller".into())], - &json!({}), - None, + let error = signed( + json( + URL, + &[("authorization".into(), "Bearer caller".into())], + &json_value!({}), + None, + ) + .unwrap(), &signer("textract"), ) .unwrap_err(); @@ -171,4 +188,17 @@ mod tests { litellm_http::Error::ComputedHeader("authorization".into()) ); } + + #[test] + fn a_streaming_body_cannot_be_signed() { + let mut request = Request::new(Method::POST, Url::parse(URL).unwrap()); + *request.body_mut() = Some(Body::wrap_stream(futures_util::stream::once(async { + Ok::<_, std::io::Error>(bytes::Bytes::from_static(b"{}")) + }))); + + assert!(matches!( + signed(request, &signer("textract")), + Err(litellm_http::Error::Signature(_)) + )); + } } diff --git a/litellm-rust/crates/http/src/error.rs b/litellm-rust/crates/http/src/error.rs index eafb4d2976b..ae14756baa3 100644 --- a/litellm-rust/crates/http/src/error.rs +++ b/litellm-rust/crates/http/src/error.rs @@ -24,6 +24,8 @@ pub enum Error { Client(String), #[error("request body could not be serialized: {0}")] RequestBody(String), + #[error("invalid request: {0}")] + Request(String), #[error("request forwards a header the signer computes: {0}")] ComputedHeader(String), #[error("request signing failed: {0}")] diff --git a/litellm-rust/crates/http/src/outbound.rs b/litellm-rust/crates/http/src/outbound.rs index c2cfb00d79b..d51a1994044 100644 --- a/litellm-rust/crates/http/src/outbound.rs +++ b/litellm-rust/crates/http/src/outbound.rs @@ -1,131 +1,69 @@ -//! The request a route hands to the transport. The body is serialized once, -//! when the request is built, and a [`RequestSigner`] sees those exact bytes. +//! Building the request a route hands to the transport. The body is serialized +//! once, into the [`reqwest::Request`] that is sent, and a [`RequestSigner`] +//! reads those exact bytes back out of it. //! //! Host hooks may rewrite the wire request (redaction, guardrails) and a //! signature such as AWS SigV4 covers the body, so a route builds this after -//! its hooks ran and cannot change or re-serialize it afterwards. +//! its hooks ran and signs last. use std::time::Duration; +use reqwest::{ + Request, Url, + header::{CONTENT_TYPE, HeaderMap}, +}; use serde::Serialize; -use crate::{ - Error, - request::{HeaderPolicy, has_header, with_headers}, -}; - -#[derive(Clone, Copy, Debug)] -pub struct UnsignedRequest<'a> { - pub url: &'a str, - pub headers: &'a [(String, String)], - pub body: &'a [u8], -} +use crate::Error; /// Returns the headers to add to the request; it never sees a mutable request. pub trait RequestSigner: Send + Sync { - fn sign(&self, request: UnsignedRequest<'_>) -> Result, Error>; + fn sign(&self, request: &Request) -> Result; } -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct OutboundRequest { - url: String, - headers: Vec<(String, String)>, - body: Vec, +pub fn json( + url: &str, + headers: &[(String, String)], + body: &impl Serialize, timeout: Option, +) -> Result { + let body = serde_json::to_vec(body).map_err(|error| Error::RequestBody(error.to_string()))?; + let url = Url::parse(url).map_err(|error| Error::Request(error.to_string()))?; + let with_headers = headers.iter().fold( + http::Request::post(url.as_str()), + |builder, (name, value)| builder.header(name, value), + ); + let has_content_type = with_headers + .headers_ref() + .is_some_and(|headers| headers.contains_key(CONTENT_TYPE)); + let builder = if has_content_type { + with_headers + } else { + with_headers.header(CONTENT_TYPE, "application/json") + }; + let mut request = builder + .body(body) + .map_err(|error| Error::Request(error.to_string())) + .and_then(|request| { + Request::try_from(request) + .map_err(|error| Error::Request(error.without_url().to_string())) + })?; + *request.timeout_mut() = timeout; + Ok(request) } -impl OutboundRequest { - pub fn json( - url: String, - headers: Vec<(String, String)>, - body: &impl Serialize, - timeout: Option, - ) -> Result { - Self::build(url, headers, body, timeout, None) - } - - pub fn signed_json( - url: String, - headers: Vec<(String, String)>, - body: &impl Serialize, - timeout: Option, - signer: &dyn RequestSigner, - ) -> Result { - Self::build(url, headers, body, timeout, Some(signer)) - } - - fn build( - url: String, - headers: Vec<(String, String)>, - body: &impl Serialize, - timeout: Option, - signer: Option<&dyn RequestSigner>, - ) -> Result { - let body = - serde_json::to_vec(body).map_err(|error| Error::RequestBody(error.to_string()))?; - let content_type = (!has_header(&headers, "content-type")) - .then(|| ("content-type".to_string(), "application/json".to_string())); - let unsigned: Vec<(String, String)> = headers.into_iter().chain(content_type).collect(); - let signature = signer - .map(|signer| { - signer.sign(UnsignedRequest { - url: &url, - headers: &unsigned, - body: &body, - }) - }) - .transpose()? - .unwrap_or_default(); - Ok(Self { - url, - headers: unsigned.into_iter().chain(signature).collect(), - body, - timeout, - }) - } - - pub fn url(&self) -> &str { - &self.url - } - - pub fn headers(&self) -> &[(String, String)] { - &self.headers - } - - pub fn header(&self, name: &str) -> Option<&str> { - self.headers - .iter() - .find(|(key, _)| key.eq_ignore_ascii_case(name)) - .map(|(_, value)| value.as_str()) - } - - pub fn body(&self) -> &[u8] { - &self.body - } - - pub fn timeout(&self) -> Option { - self.timeout - } - - pub async fn send(self, client: &crate::Client) -> Result { - let builder = with_headers( - client.post(&self.url).body(self.body), - &self.headers, - HeaderPolicy::All, - ); - match self.timeout { - Some(timeout) => builder.timeout(timeout), - None => builder, - } - .send() - .await - } +pub fn signed(mut request: Request, signer: &dyn RequestSigner) -> Result { + let signature = signer.sign(&request)?; + request.headers_mut().extend(signature); + Ok(request) } #[cfg(test)] mod tests { use std::sync::Mutex; + use reqwest::header::HeaderValue; + use rstest::rstest; use serde_json::json; use super::*; @@ -134,61 +72,73 @@ mod tests { struct Recording(Mutex>); impl RequestSigner for Recording { - fn sign(&self, request: UnsignedRequest<'_>) -> Result, Error> { - *self.0.lock().unwrap() = request.body.to_vec(); - Ok(vec![("authorization".into(), "signed".into())]) + fn sign(&self, request: &Request) -> Result { + *self.0.lock().unwrap() = request.body().unwrap().as_bytes().unwrap().to_vec(); + Ok(HeaderMap::from_iter([( + reqwest::header::AUTHORIZATION, + HeaderValue::from_static("signed"), + )])) } } #[test] fn the_signer_sees_exactly_the_bytes_that_are_sent() { let signer = Recording::default(); - let request = OutboundRequest::signed_json( - "https://provider.test/".into(), - vec![("x-caller".into(), "kept".into())], - &json!({"b": 1, "a": [true, null]}), - None, + let request = signed( + json( + "https://provider.test/", + &[("x-caller".into(), "kept".into())], + &json!({"b": 1, "a": [true, null]}), + Some(Duration::from_secs(7)), + ) + .unwrap(), &signer, ) .unwrap(); - assert_eq!(request.body(), signer.0.lock().unwrap().as_slice()); - assert_eq!(request.header("authorization"), Some("signed")); - assert_eq!(request.header("x-caller"), Some("kept")); + assert_eq!( + request.body().unwrap().as_bytes().unwrap(), + signer.0.lock().unwrap().as_slice() + ); + assert_eq!(request.headers()["authorization"], "signed"); + assert_eq!(request.headers()["x-caller"], "kept"); + assert_eq!(request.timeout(), Some(&Duration::from_secs(7))); } #[test] fn the_content_type_is_part_of_what_the_signer_sees() { struct RequiresContentType; impl RequestSigner for RequiresContentType { - fn sign(&self, request: UnsignedRequest<'_>) -> Result, Error> { - has_header(request.headers, "content-type") - .then(Vec::new) + fn sign(&self, request: &Request) -> Result { + request + .headers() + .contains_key(CONTENT_TYPE) + .then(HeaderMap::new) .ok_or_else(|| Error::Signature("content-type was not signed".into())) } } - let defaulted = OutboundRequest::signed_json( - "u".into(), - Vec::new(), - &json!({}), - None, + let defaulted = signed( + json("https://u.test/", &[], &json!({}), None).unwrap(), &RequiresContentType, ) .unwrap(); - assert_eq!(defaulted.header("content-type"), Some("application/json")); + assert_eq!(defaulted.headers()[CONTENT_TYPE], "application/json"); - let provider = OutboundRequest::signed_json( - "u".into(), - vec![("Content-Type".into(), "application/x-amz-json-1.1".into())], - &json!({}), - None, + let provider = signed( + json( + "https://u.test/", + &[("Content-Type".into(), "application/x-amz-json-1.1".into())], + &json!({}), + None, + ) + .unwrap(), &RequiresContentType, ) .unwrap(); assert_eq!( - provider.header("content-type"), - Some("application/x-amz-json-1.1") + provider.headers()[CONTENT_TYPE], + "application/x-amz-json-1.1" ); assert_eq!(provider.headers().len(), 1); } @@ -197,14 +147,90 @@ mod tests { fn a_signer_failure_produces_no_request() { struct Refuses; impl RequestSigner for Refuses { - fn sign(&self, _request: UnsignedRequest<'_>) -> Result, Error> { + fn sign(&self, _request: &Request) -> Result { Err(Error::ComputedHeader("authorization".into())) } } + let error = signed( + json("https://u.test/", &[], &json!({}), None).unwrap(), + &Refuses, + ) + .unwrap_err(); + assert_eq!(error, Error::ComputedHeader("authorization".into())); + } + + #[rstest] + #[case::unicode_domain("https://例え.test/v1/chat/completions")] + #[case::unicode_path("https://provider.test/模型/v1/chat/completions")] + #[case::unicode_query("https://provider.test/v1/chat/completions?region=日本")] + #[case::ascii("https://provider.test/v1/chat/completions?region=us-east-1")] + #[case::escaped("https://provider.test/%E6%A8%A1%E5%9E%8B?q=%E6%97%A5%E6%9C%AC")] + fn provider_urls_keep_reqwest_normalization(#[case] url: &str) { + let body = json!({"text": "日本"}); + let timeout = Duration::from_secs(7); + let expected = reqwest::Client::new() + .post(url) + .header(CONTENT_TYPE, "application/json") + .header("x-caller", "kept") + .body(serde_json::to_vec(&body).unwrap()) + .timeout(timeout) + .build() + .unwrap(); + let signer = Recording::default(); + let request = signed( + json( + url, + &[("x-caller".into(), "kept".into())], + &body, + Some(timeout), + ) + .unwrap(), + &signer, + ) + .unwrap(); + + assert_eq!(request.url(), expected.url()); + assert_eq!(request.method(), expected.method()); + assert_eq!(request.timeout(), expected.timeout()); assert_eq!( - OutboundRequest::signed_json("u".into(), Vec::new(), &json!({}), None, &Refuses), - Err(Error::ComputedHeader("authorization".into())) + request.headers()[CONTENT_TYPE], + expected.headers()[CONTENT_TYPE] + ); + assert_eq!( + request.headers()["x-caller"], + expected.headers()["x-caller"] + ); + assert_eq!(request.headers()["authorization"], "signed"); + assert_eq!( + request.body().unwrap().as_bytes(), + expected.body().unwrap().as_bytes() + ); + assert_eq!( + request.body().unwrap().as_bytes().unwrap(), + signer.0.lock().unwrap().as_slice() ); } + + #[rstest] + #[case::relative_url("not a url", "kept")] + #[case::invalid_host("https://[", "kept")] + #[case::invalid_port("https://u.test:invalid", "kept")] + #[case::invalid_header("https://u.test/", "\u{7f}")] + fn a_request_that_cannot_be_sent_fails_when_it_is_built( + #[case] url: &str, + #[case] header: &str, + ) { + assert!( + reqwest::Client::new() + .post(url) + .header("x-caller", header) + .build() + .is_err() + ); + assert!(matches!( + json(url, &[("x-caller".into(), header.into())], &json!({}), None), + Err(Error::Request(_)) + )); + } } diff --git a/litellm-rust/crates/http/src/request.rs b/litellm-rust/crates/http/src/request.rs index cfebe8c6b01..5a3dd19ade2 100644 --- a/litellm-rust/crates/http/src/request.rs +++ b/litellm-rust/crates/http/src/request.rs @@ -53,6 +53,18 @@ pub async fn execute_http_request( client.execute(request).await } +pub fn header_pairs(headers: &reqwest::header::HeaderMap) -> Vec<(String, String)> { + headers + .iter() + .filter_map(|(name, value)| { + value + .to_str() + .ok() + .map(|value| (name.to_string(), value.to_string())) + }) + .collect() +} + pub fn truncate_error_body(body: &str) -> String { if body.chars().count() <= UPSTREAM_ERROR_BODY_MAX_CHARS { return body.to_string(); diff --git a/litellm-rust/crates/inference-chat/Cargo.toml b/litellm-rust/crates/inference-chat/Cargo.toml index 130786e788d..fa3fad067a9 100644 --- a/litellm-rust/crates/inference-chat/Cargo.toml +++ b/litellm-rust/crates/inference-chat/Cargo.toml @@ -6,7 +6,6 @@ license.workspace = true repository.workspace = true [dependencies] -futures-util.workspace = true litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] } litellm-cache-response.workspace = true litellm-core-utils.workspace = true @@ -16,7 +15,10 @@ litellm-inference.workspace = true litellm-llms.workspace = true litellm-llms-types.workspace = true litellm-secrets.workspace = true + serde_json = { workspace = true, features = ["preserve_order"] } +futures-util.workspace = true +reqwest.workspace = true tracing.workspace = true [dev-dependencies] diff --git a/litellm-rust/crates/inference-chat/src/handler.rs b/litellm-rust/crates/inference-chat/src/handler.rs index c52ff19ac08..15a45747997 100644 --- a/litellm-rust/crates/inference-chat/src/handler.rs +++ b/litellm-rust/crates/inference-chat/src/handler.rs @@ -5,7 +5,6 @@ use litellm_auth::AuthServices; use litellm_host::interceptors::{Interceptors, RawResponse, RequestContext, WireRequest}; use litellm_http::{ Client, - outbound::OutboundRequest, request::{truncate_error_body, with_default_headers}, }; use litellm_llms::base_llm::{ @@ -96,7 +95,7 @@ pub(super) async fn execute( headers: wire.headers, signer: authenticated.signer, }, - wire.url, + &wire.url, &wire.body, call.timeout, )?; @@ -170,10 +169,10 @@ pub(super) fn as_response_error(err: Error) -> Error { pub(super) fn outbound_request( authenticated: Authenticated, - url: String, + url: &str, body: &Value, timeout: Option, -) -> Result { +) -> Result { litellm_inference::outbound::outbound_request( authenticated, url, diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index 8b6154ecbee..292995d842c 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -288,7 +288,7 @@ async fn send( ) -> Result { let request = outbound_request( authenticated, - url.to_string(), + url, body, Some(timeout.unwrap_or(Duration::from_secs(MESSAGES_TIMEOUT_SECS))), )?; @@ -859,12 +859,12 @@ mod tests { headers: call.wire.headers, signer: call.signer, }, - call.wire.url, + &call.wire.url, &call.wire.body, None, ) .unwrap(); - let authorization = signed.header("authorization").unwrap_or_default(); + let authorization = signed.headers()["authorization"].to_str().unwrap(); assert!( authorization.contains(&format!("/us-west-2/{service}/aws4_request")), "{authorization}" diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/aws_textract.rs b/litellm-rust/crates/inference-ocr/tests/ocr/aws_textract.rs index 5336d67895d..c4096767188 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/aws_textract.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/aws_textract.rs @@ -1,9 +1,12 @@ -use std::{collections::BTreeMap, time::SystemTime}; +use std::time::SystemTime; -use litellm_auth_aws::{Credentials, aws_signature_headers, sign_post}; +use litellm_auth_aws::{Credentials, SigningScope, aws_signature_headers, sign_headers}; use rstest::rstest; use time::{PrimitiveDateTime, format_description}; -use wiremock::Request; +use wiremock::{ + Request, + http::{HeaderMap, HeaderName}, +}; use super::*; @@ -42,21 +45,26 @@ fn expected_authorization(url: &str, sent: &Request) -> String { .unwrap() .assume_utc() .into(); - let headers: BTreeMap = ["content-type", "x-amz-target"] + let headers: HeaderMap = ["content-type", "x-amz-target"] .into_iter() - .map(|name| (name.to_string(), sent.header(name).unwrap().to_string())) + .map(|name| (HeaderName::from_static(name), sent.headers[name].clone())) .collect(); - sign_post( + sign_headers( + "POST", url, &sent.body, &aws_signature_headers(&headers), - "eu-west-1", - "textract", - &Credentials::new(ACCESS_KEY_ID, SECRET_ACCESS_KEY, None, None, "test"), - signed_at, + SigningScope { + region: "eu-west-1", + service: "textract", + credentials: &Credentials::new(ACCESS_KEY_ID, SECRET_ACCESS_KEY, None, None, "test"), + time: signed_at, + }, ) - .unwrap()["Authorization"] - .clone() + .unwrap()["authorization"] + .to_str() + .unwrap() + .to_string() } /// The recorded URL names wiremock's host, not the address the client signed for. diff --git a/litellm-rust/crates/inference-responses/src/handler.rs b/litellm-rust/crates/inference-responses/src/handler.rs index 02182d23ae3..cbbc8be139d 100644 --- a/litellm-rust/crates/inference-responses/src/handler.rs +++ b/litellm-rust/crates/inference-responses/src/handler.rs @@ -79,7 +79,7 @@ pub(super) async fn execute( headers: wire.headers, signer: authenticated.signer, }, - wire.url, + &wire.url, &wire.body, Some(call.timeout.unwrap_or(Duration::from_secs(600))), )?; diff --git a/litellm-rust/crates/inference-transcription/src/handler.rs b/litellm-rust/crates/inference-transcription/src/handler.rs index 5a7238c5e36..6541ae8263e 100644 --- a/litellm-rust/crates/inference-transcription/src/handler.rs +++ b/litellm-rust/crates/inference-transcription/src/handler.rs @@ -48,7 +48,7 @@ pub(super) async fn execute( let authenticated = resolve_auth(auth, environment, &env_lookup).await?; let outbound = litellm_inference::outbound::outbound_request( authenticated, - url, + &url, &body, Some( call.timeout diff --git a/litellm-rust/crates/inference/src/outbound.rs b/litellm-rust/crates/inference/src/outbound.rs index 41fe3aa62b8..9d9b1e606c2 100644 --- a/litellm-rust/crates/inference/src/outbound.rs +++ b/litellm-rust/crates/inference/src/outbound.rs @@ -1,6 +1,6 @@ use std::time::Duration; -use litellm_http::outbound::OutboundRequest; +use litellm_http::outbound::{json, signed}; use litellm_llms::base_llm::auth::Authenticated; use serde_json::Value; @@ -11,10 +11,10 @@ use serde_json::Value; fields(status) )] pub async fn send( - request: OutboundRequest, + request: reqwest::Request, client: &litellm_http::Client, ) -> Result { - request.send(client).await.inspect(|response| { + client.execute(request).await.inspect(|response| { tracing::Span::current().record("status", response.status().as_u16()); }) } @@ -23,13 +23,14 @@ pub async fn send( /// bytes that are sent. pub fn outbound_request( authenticated: Authenticated, - url: String, + url: &str, body: &Value, timeout: Option, -) -> Result { +) -> Result { let Authenticated { headers, signer } = authenticated; + let request = json(url, &headers, body, timeout)?; match signer { - None => OutboundRequest::json(url, headers, body, timeout), - Some(signer) => OutboundRequest::signed_json(url, headers, body, timeout, &signer), + None => Ok(request), + Some(signer) => signed(request, &signer), } } diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs index 21ad4ea8025..a9ddcb87086 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs @@ -7,7 +7,8 @@ use litellm_host::interceptors::WireRequest; use litellm_http::{ Client, ClientVariant, HttpClientConfig, HttpClientPool, media::{MediaFetcher, UrlPolicy}, - outbound::{OutboundRequest, RequestSigner}, + outbound::{RequestSigner, json, signed}, + request::header_pairs, transport, }; use serde::{Serialize, de::DeserializeOwned}; @@ -128,22 +129,14 @@ pub async fn ocr( ) -> Result { let http = config.prepare_request(request, client, hooks).await?; let url = http.url().to_string(); - let headers = http.headers().to_vec(); - let response = http - .send(client.provider_http()) + let headers = header_pairs(http.headers()); + let response = client + .provider_http() + .execute(http) .await .map_err(transport_error)?; if !response.status().is_success() { - let headers = response - .headers() - .iter() - .filter_map(|(name, value)| { - value - .to_str() - .ok() - .map(|value| (name.to_string(), value.to_string())) - }) - .collect(); + let headers = header_pairs(response.headers()); return match read_response_bytes(response, request.connection.max_response_bytes).await { Err(Error::Transport(transport::Error::Http { status, body })) => { Err(config.get_error_class(body, status, headers)) @@ -225,7 +218,7 @@ pub async fn transform_request_body( body: B, signer: Option<&dyn RequestSigner>, hooks: &dyn CallHooks, -) -> Result { +) -> Result { let composed = litellm_core_utils::call_arguments::compose_body( &request.optional_params, &body, @@ -241,17 +234,16 @@ pub async fn transform_request_body( }); } config.validate_request_body(&changed.body)?; - let timeout = Some(request.connection.timeout); + let built = json( + url, + &changed.headers, + &changed.body, + Some(request.connection.timeout), + )?; Ok(match signer { - Some(signer) => OutboundRequest::signed_json( - url.into(), - changed.headers, - &changed.body, - timeout, - signer, - ), - None => OutboundRequest::json(url.into(), changed.headers, &changed.body, timeout), - }?) + Some(signer) => signed(built, signer)?, + None => built, + }) } fn wire_request(url: &str, headers: &[(String, String)], body: Value) -> WireRequest { @@ -267,10 +259,10 @@ pub fn build_http_request( url: String, headers: Vec<(String, String)>, body: &impl Serialize, -) -> Result { - Ok(OutboundRequest::json( - url, - headers, +) -> Result { + Ok(json( + &url, + &headers, body, Some(request.connection.timeout), )?) diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs index 2da0e1abbfc..73e7b64638d 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs @@ -1,15 +1,15 @@ -use litellm_llms_types::formats::ocr::{LiteLLMOcrResponse, OcrDocument, OcrResponseFormat}; -use std::{collections::BTreeMap, future::Future, sync::Arc, time::Duration}; - -use litellm_auth::{InputSource, SecretValue, Sourced, TokenProviderHandle}; -use litellm_core_utils::{call_arguments::CallArguments, settings::ProcessEnvironment}; -use litellm_http::outbound::{OutboundRequest, RequestSigner}; -use litellm_secrets::source::Secrets; use serde::{ Serialize, de::{DeserializeOwned, IntoDeserializer}, }; use serde_json::{Map, Value}; +use std::{collections::BTreeMap, future::Future, sync::Arc, time::Duration}; + +use litellm_auth::{InputSource, SecretValue, Sourced, TokenProviderHandle}; +use litellm_core_utils::{call_arguments::CallArguments, settings::ProcessEnvironment}; +use litellm_http::outbound::RequestSigner; +use litellm_llms_types::formats::ocr::{LiteLLMOcrResponse, OcrDocument, OcrResponseFormat}; +use litellm_secrets::source::Secrets; use crate::base_llm::ocr::{ error::Error, @@ -391,7 +391,7 @@ pub trait BaseOcrConfig: Send + Sync + Sized + 'static { request: &PreparedOcrRequest, client: &OcrClient, hooks: &dyn CallHooks, - ) -> impl Future> + Send { + ) -> impl Future> + Send { async move { let params = self.map_ocr_params(&request.optional_params, &request.model)?; let environment = self.validate_environment(request, client).await?; diff --git a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs index 9b991b75e5b..fc73ce7c014 100644 --- a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs @@ -1,11 +1,5 @@ use std::collections::BTreeMap; -use litellm_core_utils::{ - call_arguments::{CallArguments, compose_body}, - params::OpaqueParams, - url_utils::ApiUrl, -}; -use litellm_http::outbound::OutboundRequest; use serde::{Deserialize, Deserializer, Serialize}; use serde_json::{Map, Value, json}; @@ -18,6 +12,11 @@ use crate::base_llm::ocr::{ decode_and_normalize_response, }, }; +use litellm_core_utils::{ + call_arguments::{CallArguments, compose_body}, + params::OpaqueParams, + url_utils::ApiUrl, +}; use litellm_llms_types::formats::ocr::{ LiteLLMOcrResponse, OcrDocument, OcrPage, OcrResponseFormat, OcrUsageInfo, }; @@ -173,7 +172,7 @@ impl BaseOcrConfig for ReductoParseV3Config { request: &PreparedOcrRequest, client: &OcrClient, hooks: &dyn CallHooks, - ) -> Result { + ) -> Result { prepare_upload_request(self, request, client, hooks).await } } @@ -262,7 +261,7 @@ impl BaseOcrConfig for ReductoParseLegacyConfig { request: &PreparedOcrRequest, client: &OcrClient, hooks: &dyn CallHooks, - ) -> Result { + ) -> Result { prepare_upload_request(self, request, client, hooks).await } } @@ -275,7 +274,7 @@ async fn prepare_upload_request, -) -> Result { +) -> Result { let params = config.map_ocr_params(&request.optional_params, &request.model)?; let headers = config.validate_environment(request, client).await?; let url = config.get_complete_url(request, ¶ms, &headers)?; From 7b60fd061a5e0500fa075600d3c36fd35be48f43 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Wed, 28 Oct 2026 09:22:11 -0700 Subject: [PATCH 043/128] refactor(rust): adopt one error discipline across litellm-rust with source chains and Report (#48042) * build(rust): upgrade reqwest to 0.13 and move MCP out of litellm-http The workspace moves to reqwest 0.13, so litellm-http no longer carries a second reqwest just for rmcp. The mcp feature, its pool and its proxy conversion go away, and the gateway hands rmcp a pooled NoRedirect client that already carries the TLS, proxy and timeout settings The Client test helpers collapse into for_test, a no-redirect client without TLS or proxy settings, and wrap_for_test for the few tests that need a hand-built reqwest client Co-Authored-By: Claude Opus 5.5 * fix(http): honor transport settings for Rust WebSockets The Responses WebSocket connected through tokio-tungstenite with its own native-roots TLS, so it ignored the configured CA bundle, verify flag, client certificate, proxies and timeouts. The handshake now goes through a pooled reqwest client (ClientVariant::WebSocket, HTTP/1.1, no redirects) and upgrades the response, so it shares the settings every other outbound request uses The upgrade response is checked for the Sec-WebSocket-Accept key, the Upgrade and Connection headers and a subprotocol that was offered Co-Authored-By: Claude Opus 5.5 * refactor(http): build outbound requests as reqwest::Request OutboundRequest was a hand-rolled copy of reqwest::Request. Routes now build a reqwest::Request directly and RequestSigner reads the body bytes back out of it, so the signer still covers exactly what is sent. SigV4 signing moves onto HeaderMap and takes the method, which drops the BTreeMap plumbing and the header casing table Co-Authored-By: Claude Opus 5.5 * refactor(http): make litellm_http::Error the one opaque error transport::Error and media::Error fold into litellm_http::Error, an Arc newtype whose private variants keep the typed reqwest, io, rustls and tungstenite causes and own every message. Callers can no longer build an error with a message of their own. They ask through dispatched, is_request, status, upstream_response and tls_source One From now classifies connect, network and timeout and strips the URL. The WebSocket handshake reports each failed check as its own variant and owns its connect timeout. OCR drops the media variants it mirrored, and inference::Error maps a header the signer computes to Unsupported for every route, not only chat The error rules move out of AGENTS.md into the rust-error skill, which also records the decisions behind this design Co-Authored-By: Claude Opus 5.5 * docs(rust): split rust-error skill into thin SKILL.md with references Co-Authored-By: Claude Opus 5.5 * refactor(rust-error): render causes once and ask predicates at the boundaries The four WebSocket handshake variants collapse into Handshake(HandshakeFault), and media download failures no longer answer dispatched() since they happen before the provider call ErrorDetail in llms and auth-types, the inference error and the OCR error now store their cause in a field named cause with no #[source], so to_string() carries it exactly once and source() is None. Tests that walked the chain now destructure the variant and downcast the cause The gateway asks is_unsupported() and is_missing_credential() on the inference and OCR errors instead of matching through nested enums, and litellm_auth::Error answers is_missing_credential() for both The rust-error skill gains axum and pyo3 references, boundaries.md describes the real predicate order and the known divergences, the testing example asserts dispatched() correctly, and shape, classification, conversions, messages-and-causes and recipes agree with the code again * docs(rust): focus the rust-error gateway page on the error flow and add a rust-gateway skill The former axum page toured gateway-inference: its status match, the envelope table, an extractor, the unsupported route and the Sse closure. Those details belong to the gateway crates and go stale with them. The page is now gateway.md and follows an error from the provider failure through each From into IntoResponse, where it is classified once and rendered by surface and framing. The gateway-structure material moves to a new rust-gateway skill: envelopes per API surface, extractors and path registration, SSE relay, and the router test harness. * test(inference): pin provider URL redaction on every route's transport error Chat, Messages, Responses and transcription each get a route-level test that fails if a reqwest error reaches the caller with its URL. The Messages stream decoder also takes body errors through litellm_http::Error like every other reqwest error Co-Authored-By: Claude Opus 5.5 * docs(rust-error): adopt std source chains and add litellm-error Report Messages describe their own layer and the cause is the source, boundaries render the chain once with litellm_error::Report. Corrects skill claims the audit found stale against the code. Co-Authored-By: Claude Opus 5.5 * refactor(rust-error): keep causes as sources across http, auth, secrets, cache and utility crates Messages describe their own layer and every cause is a typed source, predicates match exhaustively, and boundaries render with litellm_error::Report. Media downloads no longer count as dispatched, Error::upstream bounds the provider body itself, and a forwarded header SigV4 cannot sign is the caller's fault. Co-Authored-By: Claude Opus 5.5 * refactor(rust-error): keep causes as sources in llms and the gateway crates ErrorDetail loses its string variant, OCR provider variants become data, host faults stop counting as the caller's, a tunnelled transport error classifies the same on every stream path, and each gateway crate has one Error with one status() and extractors that reject in its envelope. Co-Authored-By: Claude Opus 5.5 * refactor(rust-error): nest llms errors in inference and ask predicates in route tests litellm_inference::Error nests litellm_llms::Error instead of mirroring its variants, host driver faults get their own variant, every predicate matches exhaustively, and route tests ask predicates instead of matching transport variants. Co-Authored-By: Claude Opus 5.5 * refactor(rust-error): map bridge and gateway-inference errors once and render them with Report RustUpstreamError has one constructor that always sets headers, argument checks raise typed errors mapped once, OCR asks upstream_response(), the Responses stream error event comes from the error type, gateway-inference loses its string variants, and auth faults classify through one litellm_auth predicate on every route. Co-Authored-By: Claude Opus 5.5 * fix(rust-error): carry OCR provider headers as a HeaderMap and keep outages off the caller The OCR provider error and get_error_class take reqwest's HeaderMap, the bridge turns it into Python's (name, value) pairs once. Credential acquisition and provider SDK failures are no longer the caller's fault, the UI hides auth and backend causes, and chat keeps a provider answer as is. Co-Authored-By: Claude Opus 5.5 * docs(rust-error): correct skill claims against the aligned code Co-Authored-By: Claude Opus 5.5 * refactor(rust-error): fix the remaining sweep findings in gateway, llms, auth and cache crates Auth rejections hide their backend cause on every inference renderer, path parameters reject in the surface envelope, Textract's unsupported document and Azure DI's unparseable poll URL get typed variants, the credential chain exposes its last failure as the source, gateway-mcp asks predicates over rmcp errors, and the S3 and Azure blob transports classify through litellm_http::Error. Co-Authored-By: Claude Opus 5.5 * refactor(rust-error): fix the remaining sweep findings on the Python side host-python keeps a native error's chain and returns blocking-pool cancellation instead of panicking, Vault transport failures convert through litellm_http::Error so no URL reaches Python, the bridge asks predicates instead of matching three levels into Vault errors, and the last string-built PyErrs become typed errors mapped once. Co-Authored-By: Claude Opus 5.5 * build(deps): refresh compatible dependencies flagged by OSV * refactor(rust): remove redundant client construction comment * test(rust): cover configured MCP proxy and CA bundle * fix(http): preserve native WebSocket trust and optional protocols * fix(http): allow plain WebSockets without native roots * fix(http): normalize provider URLs before building outbound requests * fix(rust): preserve OCR download status and reduce error size * fix(rust): preserve upstream bodies and bound diagnostics only * test(rust): preserve bridge serialization boundary in body regression --------- Co-authored-by: Claude Opus 5.5 Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../.agents/skills/rust-error/SKILL.md | 35 + .../rust-error/references/boundaries.md | 39 + .../rust-error/references/classification.md | 66 ++ .../rust-error/references/conversions.md | 92 +++ .../skills/rust-error/references/gateway.md | 56 ++ .../references/messages-and-causes.md | 112 +++ .../skills/rust-error/references/pyo3.md | 74 ++ .../skills/rust-error/references/recipes.md | 64 ++ .../skills/rust-error/references/shape.md | 105 +++ .../skills/rust-error/references/testing.md | 37 + .../.agents/skills/rust-gateway/SKILL.md | 20 + .../rust-gateway/references/requests.md | 34 + .../rust-gateway/references/surfaces.md | 20 + .../skills/rust-gateway/references/testing.md | 9 + litellm-rust/AGENTS.md | 12 +- litellm-rust/Cargo.lock | 66 +- litellm-rust/Cargo.toml | 1 + litellm-rust/crates/auth-aws/Cargo.toml | 1 + litellm-rust/crates/auth-aws/src/aws.rs | 94 ++- litellm-rust/crates/auth-aws/src/error.rs | 101 ++- litellm-rust/crates/auth-aws/src/lib.rs | 2 +- litellm-rust/crates/auth-aws/src/signer.rs | 44 +- litellm-rust/crates/auth-azure/Cargo.toml | 1 + litellm-rust/crates/auth-azure/src/native.rs | 43 +- litellm-rust/crates/auth-azure/src/resolve.rs | 119 ++-- litellm-rust/crates/auth-gcp/Cargo.toml | 1 + litellm-rust/crates/auth-gcp/src/auth.rs | 45 +- litellm-rust/crates/auth-gcp/src/sdk.rs | 6 +- litellm-rust/crates/auth-types/Cargo.toml | 1 + .../crates/auth-types/src/credential.rs | 26 +- litellm-rust/crates/auth-types/src/error.rs | 326 +++++++-- litellm-rust/crates/auth-types/src/http.rs | 12 +- litellm-rust/crates/auth-types/src/lib.rs | 4 +- .../crates/auth-types/src/params/aws.rs | 49 +- .../crates/auth-types/src/params/mod.rs | 7 + .../crates/auth-types/src/params/vertex.rs | 5 +- litellm-rust/crates/auth-types/src/policy.rs | 35 +- litellm-rust/crates/auth-types/tests/error.rs | 74 -- .../crates/auth-types/tests/params.rs | 3 +- litellm-rust/crates/auth-types/tests/token.rs | 19 +- .../crates/cache-azure-blob/Cargo.toml | 1 + .../crates/cache-azure-blob/src/cache.rs | 30 +- .../crates/cache-azure-blob/src/credential.rs | 15 +- .../crates/cache-azure-blob/src/transport.rs | 7 +- .../crates/cache-azure-blob/tests/cache.rs | 12 +- litellm-rust/crates/cache-disk/Cargo.toml | 1 + litellm-rust/crates/cache-disk/src/cache.rs | 8 +- .../crates/cache-disk/src/python/value.rs | 18 +- litellm-rust/crates/cache-disk/src/sqlite.rs | 168 +++-- .../crates/cache-disk/tests/python_compat.rs | 2 +- litellm-rust/crates/cache-gcs/Cargo.toml | 2 + litellm-rust/crates/cache-gcs/src/cache.rs | 32 +- litellm-rust/crates/cache-gcs/src/token.rs | 2 +- litellm-rust/crates/cache-gcs/tests/cache.rs | 39 +- litellm-rust/crates/cache-memory/src/cache.rs | 48 +- .../crates/cache-memory/tests/cache.rs | 6 +- .../crates/cache-qdrant-semantic/Cargo.toml | 1 + .../crates/cache-qdrant-semantic/src/cache.rs | 14 +- .../cache-qdrant-semantic/src/embedder.rs | 17 +- .../cache-qdrant-semantic/tests/embedder.rs | 21 +- .../cache-qdrant-semantic/tests/qdrant.rs | 26 +- .../crates/cache-redis-semantic/src/cache.rs | 6 +- .../crates/cache-redis-semantic/src/index.rs | 18 +- .../cache-redis-semantic/tests/cache.rs | 30 +- .../cache-redis-semantic/tests/support/mod.rs | 15 +- litellm-rust/crates/cache-redis/Cargo.toml | 2 + litellm-rust/crates/cache-redis/src/cache.rs | 4 +- litellm-rust/crates/cache-redis/src/claim.rs | 14 +- .../crates/cache-redis/src/connection.rs | 74 +- .../crates/cache-redis/src/counter.rs | 8 +- litellm-rust/crates/cache-redis/src/keys.rs | 4 +- .../crates/cache-redis/src/lifecycle.rs | 9 +- litellm-rust/crates/cache-redis/src/queue.rs | 16 +- litellm-rust/crates/cache-redis/src/script.rs | 2 +- litellm-rust/crates/cache-redis/src/store.rs | 18 +- .../crates/cache-redis/tests/cache.rs | 38 +- .../crates/cache-redis/tests/cluster.rs | 2 +- .../crates/cache-redis/tests/support/mod.rs | 4 +- .../crates/cache-response/src/buffer.rs | 4 +- .../crates/cache-response/src/codec.rs | 28 +- .../crates/cache-response/src/response.rs | 13 +- .../crates/cache-response/tests/codec.rs | 24 +- .../crates/cache-response/tests/response.rs | 12 +- litellm-rust/crates/cache-s3/Cargo.toml | 1 + litellm-rust/crates/cache-s3/src/auth.rs | 2 +- litellm-rust/crates/cache-s3/src/cache.rs | 6 +- litellm-rust/crates/cache-s3/src/transport.rs | 3 +- litellm-rust/crates/cache-s3/tests/cache.rs | 24 +- .../crates/cache-valkey-semantic/Cargo.toml | 2 + .../crates/cache-valkey-semantic/src/index.rs | 39 +- .../cache-valkey-semantic/src/search.rs | 28 +- .../cache-valkey-semantic/tests/cache.rs | 43 +- .../tests/support/mod.rs | 15 +- litellm-rust/crates/cache/Cargo.toml | 1 + litellm-rust/crates/cache/src/capabilities.rs | 4 +- litellm-rust/crates/cache/src/codec.rs | 4 +- litellm-rust/crates/cache/src/dual.rs | 14 +- litellm-rust/crates/cache/src/error.rs | 62 +- litellm-rust/crates/cache/src/lib.rs | 2 +- litellm-rust/crates/cache/tests/caching.rs | 10 +- litellm-rust/crates/cache/tests/codec.rs | 10 +- litellm-rust/crates/cache/tests/dual.rs | 35 +- .../crates/callbacks-legacy-python/Cargo.toml | 2 + .../callbacks-legacy-python/src/adapter.rs | 7 +- .../callbacks-legacy-python/src/error.rs | 17 + .../crates/callbacks-legacy-python/src/lib.rs | 1 + litellm-rust/crates/config/Cargo.toml | 1 + litellm-rust/crates/config/src/error.rs | 11 +- litellm-rust/crates/config/src/includes.rs | 35 +- litellm-rust/crates/config/tests/config.rs | 10 +- litellm-rust/crates/core-utils/Cargo.toml | 1 + .../crates/core-utils/src/call_arguments.rs | 19 +- .../src/exception_mapping_utils/mod.rs | 16 +- litellm-rust/crates/core-utils/src/params.rs | 2 + .../crates/core-utils/src/url_utils.rs | 2 +- litellm-rust/crates/cost/Cargo.toml | 3 + litellm-rust/crates/cost/src/error.rs | 27 + litellm-rust/crates/cost/src/lib.rs | 88 +-- litellm-rust/crates/cost/tests/calculation.rs | 20 +- litellm-rust/crates/error/Cargo.toml | 13 + litellm-rust/crates/error/src/lib.rs | 132 ++++ litellm-rust/crates/framer/src/error.rs | 8 +- litellm-rust/crates/framer/src/framer.rs | 13 - litellm-rust/crates/gateway-auth/AGENTS.md | 4 +- litellm-rust/crates/gateway-auth/Cargo.toml | 1 + litellm-rust/crates/gateway-auth/src/error.rs | 108 +-- litellm-rust/crates/gateway-auth/src/keys.rs | 10 +- litellm-rust/crates/gateway-auth/src/lib.rs | 2 +- litellm-rust/crates/gateway-auth/src/ui.rs | 29 +- .../crates/gateway-auth/tests/auth.rs | 12 +- .../crates/gateway-auth/tests/keys.rs | 22 +- .../crates/gateway-inference/Cargo.toml | 2 + .../src/audio_transcription.rs | 2 +- .../crates/gateway-inference/src/caching.rs | 20 +- .../gateway-inference/src/chat_completions.rs | 9 +- .../crates/gateway-inference/src/error.rs | 348 ++++++--- .../crates/gateway-inference/src/lib.rs | 2 +- .../crates/gateway-inference/src/messages.rs | 2 +- .../crates/gateway-inference/src/ocr.rs | 2 +- .../crates/gateway-inference/src/request.rs | 39 +- .../crates/gateway-inference/src/responses.rs | 9 +- .../gateway-inference/tests/messages.rs | 26 +- .../gateway-inference/tests/responses.rs | 29 + .../crates/gateway-inference/tests/routes.rs | 16 +- .../gateway-inference/tests/support/mod.rs | 24 + .../crates/gateway-management/Cargo.toml | 1 + .../crates/gateway-management/src/error.rs | 2 +- .../crates/gateway-management/tests/keys.rs | 10 +- .../gateway-management/tests/keys/support.rs | 15 +- litellm-rust/crates/gateway-mcp/AGENTS.md | 2 + litellm-rust/crates/gateway-mcp/Cargo.toml | 2 + .../crates/gateway-mcp/src/configured.rs | 120 ++-- litellm-rust/crates/gateway-mcp/src/error.rs | 381 +++++++++- litellm-rust/crates/gateway-mcp/src/legacy.rs | 88 +-- litellm-rust/crates/gateway-mcp/src/lib.rs | 3 +- litellm-rust/crates/gateway-mcp/src/native.rs | 23 +- .../crates/gateway-mcp/src/native/catalog.rs | 15 +- .../crates/gateway-mcp/src/native/registry.rs | 18 +- .../crates/gateway-mcp/src/request.rs | 44 ++ litellm-rust/crates/gateway-mcp/src/routes.rs | 14 +- .../crates/gateway-mcp/src/sessions.rs | 12 +- .../crates/gateway-mcp/tests/gateway/http.rs | 115 ++- .../gateway-mcp/tests/gateway/native.rs | 10 +- litellm-rust/crates/gateway-ui/Cargo.toml | 2 + .../crates/gateway-ui/src/dashboard.rs | 16 +- litellm-rust/crates/gateway-ui/src/error.rs | 119 +++- litellm-rust/crates/gateway-ui/src/session.rs | 35 +- .../crates/gateway-ui/tests/assets.rs | 18 + .../crates/gateway-ui/tests/routes.rs | 53 +- litellm-rust/crates/gateway/Cargo.toml | 1 + litellm-rust/crates/gateway/src/auth.rs | 8 +- litellm-rust/crates/gateway/src/error.rs | 4 +- litellm-rust/crates/gateway/src/lib.rs | 2 +- litellm-rust/crates/gateway/tests/mcp.rs | 146 ++-- litellm-rust/crates/gateway/tests/server.rs | 5 +- litellm-rust/crates/host-python/Cargo.toml | 1 + .../crates/host-python/src/binding.rs | 2 +- litellm-rust/crates/host-python/src/driver.rs | 2 +- litellm-rust/crates/host-python/src/error.rs | 16 +- .../crates/host-python/src/fork_gate.rs | 14 +- litellm-rust/crates/host-python/src/gil.rs | 14 +- litellm-rust/crates/host-python/src/handle.rs | 32 +- .../crates/host-python/src/marshal.rs | 3 +- .../crates/host-python/src/runtime.rs | 31 +- .../crates/host-python/tests/support/mod.rs | 2 + litellm-rust/crates/host/Cargo.toml | 1 + .../crates/host/src/machine/coroutine.rs | 6 +- litellm-rust/crates/http/Cargo.toml | 2 + litellm-rust/crates/http/src/config.rs | 20 +- litellm-rust/crates/http/src/error.rs | 672 +++++++++++++++++- litellm-rust/crates/http/src/lib.rs | 1 - litellm-rust/crates/http/src/media.rs | 122 ++-- litellm-rust/crates/http/src/outbound.rs | 40 +- litellm-rust/crates/http/src/pool.rs | 5 +- litellm-rust/crates/http/src/request.rs | 38 +- litellm-rust/crates/http/src/tls.rs | 39 +- litellm-rust/crates/http/src/transport.rs | 109 --- litellm-rust/crates/http/src/websocket.rs | 64 +- litellm-rust/crates/http/tests/websocket.rs | 53 +- litellm-rust/crates/inference-chat/Cargo.toml | 2 +- .../crates/inference-chat/src/common_utils.rs | 8 +- .../crates/inference-chat/src/handler.rs | 127 ++-- .../inference-chat/tests/chat_completions.rs | 50 +- .../crates/inference-chat/tests/request.rs | 43 +- .../crates/inference-messages/Cargo.toml | 1 + .../inference-messages/src/common_utils.rs | 16 +- .../crates/inference-messages/src/handler.rs | 72 +- .../crates/inference-messages/src/types.rs | 4 +- .../tests/bedrock_native/main.rs | 9 +- .../inference-messages/tests/messages/host.rs | 9 +- .../inference-messages/tests/messages/main.rs | 2 +- .../tests/messages/request.rs | 26 +- .../tests/messages/response.rs | 59 +- .../tests/messages/secrets.rs | 2 +- .../tests/messages/stream.rs | 24 +- .../tests/messages/tencent.rs | 10 +- .../tests/openrouter/main.rs | 15 +- litellm-rust/crates/inference-ocr/Cargo.toml | 2 + .../crates/inference-ocr/src/document.rs | 27 +- .../inference-ocr/src/provider_config.rs | 32 +- .../crates/inference-ocr/src/route.rs | 5 +- .../crates/inference-ocr/src/types.rs | 4 +- litellm-rust/crates/inference-ocr/src/wire.rs | 28 +- .../inference-ocr/tests/ocr/aws_textract.rs | 15 +- .../inference-ocr/tests/ocr/azure_ai.rs | 11 +- .../tests/ocr/azure_document_intelligence.rs | 34 +- .../crates/inference-ocr/tests/ocr/cohere.rs | 11 +- .../inference-ocr/tests/ocr/documents.rs | 7 +- .../inference-ocr/tests/ocr/lifecycle.rs | 7 +- .../crates/inference-ocr/tests/ocr/machine.rs | 35 +- .../crates/inference-ocr/tests/ocr/main.rs | 12 +- .../crates/inference-ocr/tests/ocr/mistral.rs | 25 +- .../crates/inference-ocr/tests/ocr/reducto.rs | 8 +- .../inference-ocr/tests/ocr/vertex_ai.rs | 10 +- .../crates/inference-responses/Cargo.toml | 1 + .../crates/inference-responses/src/handler.rs | 37 +- .../inference-responses/src/websocket.rs | 55 +- .../inference-responses/tests/openrouter.rs | 18 +- .../inference-responses/tests/responses.rs | 30 +- .../inference-responses/tests/support/mod.rs | 3 + .../crates/inference-transcription/Cargo.toml | 1 + .../inference-transcription/src/handler.rs | 19 +- .../tests/audio_transcription.rs | 38 +- litellm-rust/crates/inference/AGENTS.md | 6 +- litellm-rust/crates/inference/Cargo.toml | 1 + litellm-rust/crates/inference/src/error.rs | 340 +++++++-- litellm-rust/crates/inference/src/outbound.rs | 8 +- litellm-rust/crates/inference/src/provider.rs | 6 +- .../crates/inference/tests/caching.rs | 4 +- litellm-rust/crates/llms/Cargo.toml | 1 + .../src/anthropic/batches/transformation.rs | 45 +- .../crates/llms/src/anthropic/chat/handler.rs | 163 ----- .../crates/llms/src/anthropic/chat/mod.rs | 1 - .../llms/src/anthropic/chat/transformation.rs | 34 +- .../llms/src/anthropic/messages/handler.rs | 19 +- .../ocr/analyze_transformation.rs | 2 +- .../llms/src/aws_textract/ocr/common_utils.rs | 90 +-- .../src/aws_textract/ocr/transformation.rs | 4 +- .../ocr/cohere_parse_transformation.rs | 6 +- .../llms/src/azure_ai/ocr/common_utils.rs | 4 +- .../document_intelligence/transformation.rs | 47 +- .../llms/src/azure_ai/ocr/transformation.rs | 11 +- .../llms/src/base_llm/base_model_iterator.rs | 9 + .../crates/llms/src/base_llm/ocr/document.rs | 73 +- .../crates/llms/src/base_llm/ocr/error.rs | 440 +++++++++--- .../crates/llms/src/base_llm/ocr/handler.rs | 89 +-- .../llms/src/base_llm/ocr/transformation.rs | 30 +- litellm-rust/crates/llms/src/base_llm/sse.rs | 55 +- .../src/bedrock/audio_transcription/mod.rs | 33 +- .../bedrock/chat/converse_transformation.rs | 69 +- .../llms/src/bedrock/chat/invoke_handler.rs | 66 +- .../messages/transformation.rs | 3 +- .../llms/src/cohere/ocr/transformation.rs | 50 +- litellm-rust/crates/llms/src/error.rs | 177 +++-- litellm-rust/crates/llms/src/lib.rs | 2 +- .../llms/src/mistral/ocr/transformation.rs | 9 +- .../src/openai/responses/transformation.rs | 18 +- .../src/openai_like/chat/transformation.rs | 19 +- .../llms/src/openai_like/common_utils.rs | 9 +- .../llms/src/reducto/ocr/transformation.rs | 47 +- .../llms/src/vertex_ai/ocr/common_utils.rs | 4 +- .../vertex_ai/ocr/deepseek_transformation.rs | 56 +- .../llms/src/vertex_ai/ocr/transformation.rs | 5 +- litellm-rust/crates/model-catalog/Cargo.toml | 1 + .../crates/model-catalog/src/catalog.rs | 6 +- .../crates/model-catalog/src/error.rs | 8 +- .../crates/model-catalog/src/validation.rs | 114 +-- .../tests/registry_validation.rs | 4 +- litellm-rust/crates/python-bridge/AGENTS.md | 7 +- litellm-rust/crates/python-bridge/Cargo.toml | 2 + .../crates/python-bridge/src/cache/mod.rs | 93 ++- .../python-bridge/src/cache/native/v2.rs | 37 +- .../python-bridge/src/cache/python/host.rs | 34 +- .../python-bridge/src/cache/python/service.rs | 19 +- .../python-bridge/src/cache/selection.rs | 2 +- .../crates/python-bridge/src/callable.rs | 44 +- .../crates/python-bridge/src/coercion.rs | 138 ++-- .../crates/python-bridge/src/credentials.rs | 13 +- .../crates/python-bridge/src/diagnostics.rs | 5 +- .../crates/python-bridge/src/errors.rs | 253 ++++++- litellm-rust/crates/python-bridge/src/http.rs | 79 +- .../crates/python-bridge/src/logger/mod.rs | 8 +- .../crates/python-bridge/src/marshal.rs | 130 ++-- .../src/routes/audio_transcription.rs | 10 +- .../src/routes/chat_completions/host.rs | 2 +- .../src/routes/clickhouse_spend.rs | 58 +- .../python-bridge/src/routes/embeddings.rs | 15 +- .../python-bridge/src/routes/inference.rs | 48 +- .../python-bridge/src/routes/messages/host.rs | 76 +- .../python-bridge/src/routes/ocr/document.rs | 113 +-- .../python-bridge/src/routes/ocr/errors.rs | 315 ++++---- .../python-bridge/src/routes/ocr/host.rs | 13 +- .../python-bridge/src/routes/ocr/project.rs | 19 +- .../src/routes/responses/host.rs | 14 +- .../python-bridge/src/routes/token_counter.rs | 26 +- .../python-bridge/src/secrets/config.rs | 4 +- .../crates/python-bridge/src/secrets/error.rs | 84 ++- .../crates/python-bridge/src/secrets/mod.rs | 4 +- .../python-bridge/src/secrets/mutation.rs | 60 +- .../python-bridge/src/secrets/operations.rs | 25 +- .../python-bridge/src/secrets/provider.rs | 12 +- .../python-bridge/src/secrets/runtime.rs | 46 +- .../crates/python-bridge/src/secrets/vault.rs | 189 ++--- .../src/secrets/vault/operation.rs | 10 +- .../crates/python-bridge/src/tokenizer.rs | 123 +++- .../crates/python-compat/src/error.rs | 8 +- .../crates/python-compat/src/literal.rs | 34 +- .../crates/python-compat/src/pickle.rs | 13 +- litellm-rust/crates/secrets-aws/src/auth.rs | 4 +- litellm-rust/crates/secrets-aws/src/error.rs | 4 +- .../secrets-aws/src/secret_manager/read.rs | 17 +- .../secrets-aws/tests/secret_manager/reads.rs | 6 +- litellm-rust/crates/secrets-azure/Cargo.toml | 1 + .../crates/secrets-azure/src/error.rs | 20 +- .../crates/secrets-azure/src/key_vault.rs | 15 +- .../crates/secrets-azure/tests/key_vault.rs | 18 +- .../crates/secrets-cyberark/Cargo.toml | 1 + .../crates/secrets-cyberark/src/error.rs | 63 +- .../crates/secrets-cyberark/src/lib.rs | 4 +- .../secrets-cyberark/src/secret_manager.rs | 30 +- .../src/secret_manager/client.rs | 32 +- .../src/secret_manager/read.rs | 8 +- .../src/secret_manager/write.rs | 8 +- .../tests/secret_manager/configuration.rs | 12 +- litellm-rust/crates/secrets-google/Cargo.toml | 1 + .../crates/secrets-google/src/error.rs | 38 +- .../secrets-google/src/secret_manager.rs | 20 +- .../secrets-google/tests/secret_manager.rs | 6 +- .../crates/secrets-hashicorp/Cargo.toml | 3 + .../crates/secrets-hashicorp/src/config.rs | 6 +- .../crates/secrets-hashicorp/src/error.rs | 122 +++- .../crates/secrets-hashicorp/src/lib.rs | 4 +- .../secrets-hashicorp/src/secret_manager.rs | 30 +- .../src/secret_manager/client.rs | 12 +- .../src/secret_manager/raw.rs | 132 +++- .../src/secret_manager/write.rs | 2 +- .../tests/secret_manager/configuration.rs | 35 +- .../tests/secret_manager/reads.rs | 51 +- .../crates/secrets-types/src/error.rs | 12 + .../crates/secrets-types/tests/cache.rs | 20 +- litellm-rust/crates/secrets/Cargo.toml | 1 + .../crates/secrets/src/compatibility.rs | 6 +- litellm-rust/crates/secrets/src/error.rs | 14 +- litellm-rust/crates/secrets/src/handler.rs | 10 +- litellm-rust/crates/secrets/src/oidc.rs | 20 +- litellm-rust/crates/secrets/tests/aws.rs | 2 +- litellm-rust/crates/secrets/tests/google.rs | 2 +- .../crates/secrets/tests/resolution.rs | 133 ++-- .../crates/spend-clickhouse/Cargo.toml | 1 + .../crates/spend-clickhouse/src/error.rs | 33 +- .../crates/spend-clickhouse/src/insert.rs | 117 ++- .../crates/spend-clickhouse/src/lib.rs | 4 +- .../crates/storage-clickhouse/Cargo.toml | 2 + .../crates/storage-clickhouse/src/error.rs | 70 +- .../crates/storage-clickhouse/src/insert.rs | 13 +- .../crates/storage-clickhouse/src/lib.rs | 8 +- .../crates/storage-clickhouse/src/migrate.rs | 24 +- .../crates/storage-clickhouse/src/read.rs | 26 +- .../storage-clickhouse/tests/migrations.rs | 29 +- .../storage-clickhouse/tests/transport.rs | 105 ++- litellm-rust/crates/testkit/Cargo.toml | 1 + .../crates/testkit/src/agent/configure.rs | 5 +- litellm-rust/crates/testkit/src/error.rs | 14 +- .../crates/testkit/src/install/fetch.rs | 4 +- .../crates/testkit/src/install/mod.rs | 26 +- litellm-rust/crates/testkit/src/session.rs | 12 +- .../crates/token-counter-fast/Cargo.toml | 1 + .../crates/token-counter-fast/src/error.rs | 19 +- .../crates/token-counter-fast/src/tiktoken.rs | 10 +- .../token-counter-huggingface/src/error.rs | 17 +- .../token-counter-huggingface/src/lib.rs | 4 +- .../crates/token-counter-tiktoken/Cargo.toml | 1 + .../token-counter-tiktoken/src/error.rs | 49 +- .../crates/token-counter-tiktoken/src/lib.rs | 29 +- .../token-counter-tiktoken/src/ranks.rs | 242 ++++--- litellm-rust/crates/token-counter/Cargo.toml | 1 + .../crates/token-counter/src/counter.rs | 10 +- .../crates/token-counter/src/error.rs | 91 ++- litellm-rust/crates/token-counter/src/fast.rs | 12 - .../crates/token-counter/src/huggingface.rs | 11 - litellm-rust/crates/token-counter/src/lib.rs | 2 +- .../crates/token-counter/src/python_json.rs | 16 +- .../crates/token-counter/src/tiktoken.rs | 18 +- .../crates/token-counter/src/tools.rs | 6 +- .../crates/token-counter/src/types.rs | 4 +- .../token-counter/tests/token_counter.rs | 20 +- 406 files changed, 9261 insertions(+), 4747 deletions(-) create mode 100644 litellm-rust/.agents/skills/rust-error/SKILL.md create mode 100644 litellm-rust/.agents/skills/rust-error/references/boundaries.md create mode 100644 litellm-rust/.agents/skills/rust-error/references/classification.md create mode 100644 litellm-rust/.agents/skills/rust-error/references/conversions.md create mode 100644 litellm-rust/.agents/skills/rust-error/references/gateway.md create mode 100644 litellm-rust/.agents/skills/rust-error/references/messages-and-causes.md create mode 100644 litellm-rust/.agents/skills/rust-error/references/pyo3.md create mode 100644 litellm-rust/.agents/skills/rust-error/references/recipes.md create mode 100644 litellm-rust/.agents/skills/rust-error/references/shape.md create mode 100644 litellm-rust/.agents/skills/rust-error/references/testing.md create mode 100644 litellm-rust/.agents/skills/rust-gateway/SKILL.md create mode 100644 litellm-rust/.agents/skills/rust-gateway/references/requests.md create mode 100644 litellm-rust/.agents/skills/rust-gateway/references/surfaces.md create mode 100644 litellm-rust/.agents/skills/rust-gateway/references/testing.md delete mode 100644 litellm-rust/crates/auth-types/tests/error.rs create mode 100644 litellm-rust/crates/callbacks-legacy-python/src/error.rs create mode 100644 litellm-rust/crates/cost/src/error.rs create mode 100644 litellm-rust/crates/error/Cargo.toml create mode 100644 litellm-rust/crates/error/src/lib.rs delete mode 100644 litellm-rust/crates/framer/src/framer.rs create mode 100644 litellm-rust/crates/gateway-mcp/src/request.rs delete mode 100644 litellm-rust/crates/http/src/transport.rs delete mode 100644 litellm-rust/crates/llms/src/anthropic/chat/handler.rs diff --git a/litellm-rust/.agents/skills/rust-error/SKILL.md b/litellm-rust/.agents/skills/rust-error/SKILL.md new file mode 100644 index 00000000000..2f02394c494 --- /dev/null +++ b/litellm-rust/.agents/skills/rust-error/SKILL.md @@ -0,0 +1,35 @@ +--- +name: rust-error +description: Define, wrap, convert, classify, report or test Rust error types in litellm-rust with thiserror. Use when adding an error type or variant, mapping one crate's error into another, turning an error into a gateway response or Python exception, deciding retry or fault from an error, or replacing call sites that build errors from strings +--- + +# Rust errors + +`crates/http/src/error.rs` is the reference implementation. Read the reference for the part you are touching before editing, each one carries the reasoning and worked examples behind its rules + +- [shape](references/shape.md) + - One top-level `Error` per crate in `src/error.rs`. A public enum by default, an opaque struct over a private `Repr` enum when the variants must stay private. No `#[non_exhaustive]` + - A variant is a failure mode a caller handles differently. Provider names and guidance are data, never variants +- [messages and causes](references/messages-and-causes.md) + - Every message is written once, in `#[error(...)]`, and describes its own layer only. Call sites pass typed facts, never a string + - The cause is the `source()`, never part of the message. Keep it typed, never stringify or drop it + - `reqwest::Error` enters our types only through `litellm_http::Error`, which strips its URL + - Anything shown to someone else is rendered with `litellm_error::Report`, never `to_string()` +- [conversions](references/conversions.md) + - One `From` per source and target pair, only when the conversion needs no context. It classifies and sanitizes once + - Nest a lower crate's error in one variant. Translate when the meaning changes, never mirror its variants +- [classification](references/classification.md) + - Callers ask predicates named after their decision, they do not match nested variants. Predicates match exhaustively + - Classification states facts. Retry and fallback policy lives with the caller that knows idempotency +- [boundaries](references/boundaries.md) + - An error becomes text for someone else only in one mapping function per error type, in the gateway crates, the Python bridge and host-python. The inference gateway and the bridge decide by the same predicates +- [gateway](references/gateway.md) + - One `Error` with `IntoResponse` per gateway crate that serves a router. Every failure flows through the same `From` chain, is classified once in `status()` and rendered by surface and framing +- [pyo3](references/pyo3.md) + - Two `create_exception!` classes with the payload in `args`, `ValueError` for the caller's fault and `RuntimeError` otherwise, decoded in `failures.py` +- [testing](references/testing.md) + - Test predicates, chain shape and redaction on errors built the way production builds them +- [recipes](references/recipes.md) + - Common tasks walked through end to end + +The workspace uses thiserror 2 for every error type, and no library crate uses `anyhow` or `Box` in a public signature. The [derive docs](https://docs.rs/thiserror/latest/thiserror/) are the authority on macro behavior, and the [`std::error::Error` docs](https://doc.rust-lang.org/std/error/trait.Error.html) on what `Display` and `source()` mean diff --git a/litellm-rust/.agents/skills/rust-error/references/boundaries.md b/litellm-rust/.agents/skills/rust-error/references/boundaries.md new file mode 100644 index 00000000000..253cf52a7db --- /dev/null +++ b/litellm-rust/.agents/skills/rust-error/references/boundaries.md @@ -0,0 +1,39 @@ +# Boundaries + +A boundary is where an error becomes something another program reads: an HTTP response, a Python exception, a stored or logged message. Each boundary has one mapping function per error type it receives, and nothing outside those functions formats an error for someone else + +| Boundary | Mapping | Reference | +|---|---|---| +| Inference HTTP response | `Error::status()` and its renderers in `crates/gateway-inference/src/error.rs` | [gateway](gateway.md) | +| Other gateway surfaces | the `IntoResponse` of each `crates/gateway-*` crate's `Error` | [gateway](gateway.md) | +| Python bridge, inference and OCR routes | `route_failure` in `crates/python-bridge/src/errors.rs` and in `crates/python-bridge/src/routes/ocr/errors.rs` | [pyo3](pyo3.md) | +| Python bridge, everything else | one `From for PyErr` or mapping function per crate error, all in `crates/python-bridge` | [pyo3](pyo3.md) | +| Python host runtime | `From for PyErr` in `crates/host-python/src/error.rs` | [pyo3](pyo3.md) | + +## Asking the nested error + +The inference gateway and the Python bridge's route mapping both receive a nested route error and decide by the same predicates, so a new failure lands correctly on both once its predicates answer right + +| Predicate | Gateway | Bridge | +|---|---|---| +| `is_unsupported()` | 501 | `ValueError`, since `is_request()` also answers true | +| `upstream_response()` | the provider's status and body | `RustUpstreamError` | +| `is_missing_credential()` | 401 | no class of its own, a `RuntimeError` because `is_request()` answers false | +| `is_request()` | 400 | `ValueError` | +| none of these | 500 | `RuntimeError` | + +If you find yourself adding an arm for your variant at a boundary, check whether a predicate is missing first + +The two boundaries are not identical. The gateway has variants of its own for body parsing, an unknown model and an unsupported path, and the OCR bridge has Python parity arms such as `FileNotFoundError` for a missing file. Those arms match the boundary's own variants or the top level of a public enum, never a nested error's variants. A missing API key is 401 at the gateway and `RuntimeError` at both bridges, pinned by tests to what Python's own path does + +## What reaches the client + +A boundary renders anything that is not a provider answer with `litellm_error::Report`, so the reader sees every layer of the chain once. That is why messages are written for a reader, each layer leaves its cause to `source()`, and reqwest URLs are stripped on entry. A boundary that must not show the chain, such as the bridge's TLS settings error that names the setting but not the file, renders fixed text of its own and says so in its test + +Provider bodies are passed through intact for gateway and Python error mapping. `litellm_http::Error::upstream(status, body)` retains the complete body, but its `Display`, `Debug` and `Report` diagnostics show only the first 256 characters and mark the cut. Do not log the raw body or `RustUpstreamError` args. OCR's `Provider` error also keeps the whole body for Python's `get_error_class`, bounded by the connection's response size limit + +Never format `Debug` of an error into a response or a Python exception. A type whose sources may hold secrets derives `veil::Redact`, so its `Debug` is safe to log + +## Logging + +Propagate with context, and let the terminal handler own the one log or trace event for a failed call. Do not log at every `?`, and do not log an error you also return. Diagnostic tracing excludes raw error strings and bodies, see [`rust-tracing`](../../rust-tracing/SKILL.md) diff --git a/litellm-rust/.agents/skills/rust-error/references/classification.md b/litellm-rust/.agents/skills/rust-error/references/classification.md new file mode 100644 index 00000000000..59c3a1eb3f0 --- /dev/null +++ b/litellm-rust/.agents/skills/rust-error/references/classification.md @@ -0,0 +1,66 @@ +# Classification + +## Ask, don't match + +Downstream crates ask a lower error questions instead of matching its nested variants. `litellm_http::Error` answers these + +| Predicate | Caller decides | +|---|---| +| `dispatched()` | whether retrying elsewhere could repeat a billed call | +| `is_request()` | the caller's input is at fault, so 400 and no reroute | +| `is_response()` | the provider answered with something we cannot use | +| `upstream_response()` | the provider's own status and body to pass through | +| `status()` | an HTTP status to report | +| `is_connect()` | the connection was never made, for a host's own connect-error mapping | +| `tls_source()` | which TLS setting to blame in a config error | +| `is_computed_header()` | the caller forwarded a header the signer owns, which `litellm_inference::Error` turns into `Unsupported` | + +Upper types add the questions their callers ask. `litellm_inference::Error` and the OCR error answer `is_missing_credential()` and `is_unsupported()`, which the gateway asks instead of reaching into `Auth` or `Unsupported` + +```rust +// reaches through two crates' variants and breaks when either changes +Self::Route(InferenceError::Auth(litellm_auth::Error::MissingApiKey { .. })) => StatusCode::UNAUTHORIZED, + +// asks the question it cares about +Self::Route(error) if error.is_missing_credential() => StatusCode::UNAUTHORIZED, +``` + +A predicate is an exhaustive `match` with no `_` arm, never `matches!`, so a new variant does not compile until it has answered every question on purpose + +```rust +pub fn is_request(&self) -> bool { + match &*self.0 { + Repr::Request(_) | Repr::BlockedUrl | Repr::DownloadTooLarge => true, + Repr::Connect(_) | Repr::Network(_) | Repr::Timeout(_) | Repr::Status { .. } => false, + // ... every other variant, listed by name + } +} +``` + +Add a predicate when a real caller needs a new distinction, and name it after the fact the caller needs, in the caller's words, not after the variant's spelling. An upper type's predicates call through to the nested error, as in `Self::Http(error) => error.is_request()` and `Self::Auth(error) => error.is_missing_credential()` + +A public enum may be matched one level down by the type that owns the variant, such as the gateway matching its own `Unsupported` or the OCR bridge matching `FileRead`. Going a second level down, into a variant of a variant, is what to avoid, the gateway used to do that for `MissingApiKey` and `Unsupported` + +## Facts, not policy + +A predicate reports what happened. The decision to retry, reroute or fall back depends on things the error cannot know, such as whether the request is idempotent, whether its body can be replayed, the deadline and the attempt budget. That decision lives with the caller. `dispatched()` exists for the router's reroute decision, so every transport failure answers it even before that caller is wired + +A timeout is the main trap. A request that timed out may already have reached the provider and been billed, so `Timeout` answers `dispatched() == true` and is not safe to retry as if it were a connect failure. Only a failure that provably happened before any byte reached the provider answers `false`: connecting, building or signing the request, and media downloads, which fetch the caller's document before the provider call is made + +```rust +pub fn dispatched(&self) -> bool { + match &*self.0 { + Repr::Network(_) | Repr::Timeout(_) | Repr::Status { .. } => true, + Repr::Connect(_) | Repr::RequestBody(_) | Repr::Download(_) | Repr::BlockedUrl => false, + // ... every other variant, listed by name + } +} +``` + +When the transport cannot tell whether a request reached the provider, it must answer `true`. Wrongly answering `false` risks a duplicate billed call, wrongly answering `true` only costs a reroute + +Keep one predicate per question. `status()` currently mixes the provider's status, the media download's status and a synthesized 408 for timeouts, so a caller that wants "did the provider answer" must use `upstream_response()` instead. Do not add more meanings to an existing predicate, add a new one + +## Equality + +`Arc` errors compare with `Arc::ptr_eq`, so a clone equals its original and two separately built errors never compare equal. That exists only so a wrapping enum such as `litellm_inference::Error` can keep its `PartialEq` derive. Do not give a new type pointer equality for any other reason, and do not rely on it in tests. Assert predicates instead, see [testing](testing.md) diff --git a/litellm-rust/.agents/skills/rust-error/references/conversions.md b/litellm-rust/.agents/skills/rust-error/references/conversions.md new file mode 100644 index 00000000000..db46e4e3aa9 --- /dev/null +++ b/litellm-rust/.agents/skills/rust-error/references/conversions.md @@ -0,0 +1,92 @@ +# Conversions + +## One From per source and target pair + +A `From for Ours` exists when the conversion is obvious, infallible and needs no context from the call site. It does all classification and sanitizing for that pair, so every call site is `?` or `.map_err(Ours::from)` + +```rust +impl From for Error { + fn from(error: reqwest::Error) -> Self { + let error = error.without_url(); + if error.is_timeout() { + Repr::Timeout(error).into() + } else if error.is_connect() || error.is_builder() { + Repr::Connect(error).into() + } else { + Repr::Network(error).into() + } + } +} +``` + +A call site never re-derives that classification + +```rust +// every route used to repeat this, and each drifted a little +.map_err(|err| { + if err.is_connect() || err.is_builder() { + Error::Transport(Connect(err.to_string())) + } else { + Error::Transport(Network(err.to_string())) + } +})? + +// now +.map_err(litellm_http::Error::from)? +``` + +The same foreign type may convert into several of our types, one `From` each. What is not allowed is two places converting the same pair differently + +## When the call site knows something the From cannot + +One `io::Error` can mean "reading the CA bundle failed" or "resolving the media host failed". A global `From` cannot know which, so there is none. Name the operation at the call site instead + +```rust +.map_err(Repr::Resolve)? +.map_err(|source| Repr::Read { path, tls_source, source })? +``` + +Do not drop meaningful context just to use `?` + +## Reclassify once, in the shared From + +When a lower error needs a different meaning for every caller, do it in the upper type's `From`, not per route + +```rust +impl From for Error { + fn from(error: litellm_http::Error) -> Self { + if error.is_computed_header() { + return Self::Unsupported("request forwards a header AWS SigV4 computes"); + } + Self::Http(error) + } +} +``` + +All four inference routes classify a header the signer computes the same way, because none of them does it themselves + +## Nest, translate or add context, never mirror + +| The upper layer | Do | Example | +|---|---|---| +| has nothing to add | nest in one transparent variant and delegate predicates | `Self::Http(error) => error.is_request()` | +| changes what the failure means | translate to an upper variant, keep the cause attached when it helps | computed header -> `Unsupported` | +| knows something the lower one does not | a variant with the extra fields plus the cause as its source | `Read { path, tls_source, source }` | + +Never re-declare a lower crate's variants in the upper type with an upper-layer name and a 1:1 mapping function. OCR once did that for every media download failure, `MediaError::BlockedUrl => Error::BlockedDocumentUrl` and so on, which doubled every change and let the two lists drift. + +## Third-party and host faults + +Converting a fault from a crate we do not own into a string variant is still flattening. Prefer a variant that carries the fault, and when an existing `ErrorDetail` constructor fits, use it + +```rust +// flattens +.map_err(|error| Error::InvalidRequest(error.to_string().into())) + +// keeps the cause typed +.map_err(|error| Error::InvalidRequest(ErrorDetail::invalid("header", error))) +``` + +## Through a foreign error type + +A trait we implement for another crate sometimes forces its error type on us, such as `io::Error` from a body stream or `sqlx::Error` from a driver. Wrap ours as that type's source (`io::Error::other(error)`, `azure_core::Error::new(kind, error)`) and recover it with `downcast_ref` on the other side, never `with_message(error.to_string())` or `from_msg`. The error must come back out as the same value it went in as, so one failure classifies the same way on every path diff --git a/litellm-rust/.agents/skills/rust-error/references/gateway.md b/litellm-rust/.agents/skills/rust-error/references/gateway.md new file mode 100644 index 00000000000..1ec8ecd0acb --- /dev/null +++ b/litellm-rust/.agents/skills/rust-error/references/gateway.md @@ -0,0 +1,56 @@ +# Gateway responses + +The gateway is where a Rust error becomes an HTTP response. This page is about the path an error takes to get there and the one place it is classified. The gateway crates' own conventions, such as which envelope each API surface uses, which extractors exist and how the router is composed, live in the [`rust-gateway`](../../rust-gateway/SKILL.md) skill + +## The flow + +```text +reqwest / provider failure + -> litellm_http::Error From, strips the URL, classifies once + -> litellm_inference::Error::Http(_) From, one transparent variant, predicates delegate + -> gateway_inference::Error::Route(_) From, one variant per nested crate + -> handler returns Result<_, Error> every `?` ends here + -> IntoResponse status() picks the code, a renderer builds the body +``` + +Each hop is one `From` that nests the lower error in one variant, see [conversions](conversions.md). No hop re-derives what the lower crate already decided, and nothing between the provider and `IntoResponse` formats an error for the client. That is what makes "one `Error` with `IntoResponse` per gateway crate that serves a router" work: a handler can `?` any failure from any depth and the response comes out right + +## Classify once + +`status()` is pure and runs once per error. Variants the gateway declares itself, such as a bad body, an unknown model or an oversized body, are matched at the top of its own enum. A nested route error goes to `route_status` and an OCR error to `ocr_status`, which ask the predicates [boundaries](boundaries.md) lists + +| Answer | Status | +|---|---| +| `is_unsupported()` | 501, the gateway matched and parsed the request but core cannot serve it, nothing reached the provider | +| `upstream_response()` | the provider's own status, 502 when it is not a valid code | +| `is_missing_credential()` | 401 | +| `is_request()` | 400, for a route error | +| `http_status_code()` | the transport's status, for an OCR error, such as 408 for a timeout | +| otherwise | 500 | + +Do not reach into a nested enum's variants from the gateway. When a new nested failure needs a different status, add a predicate on the nested error, see [classification](classification.md) + +## Render by surface and framing + +Rendering is separate from classification. The surface picks the envelope and the moment picks the framing, a JSON body before the headers went out and an `event: error` SSE frame after. Both framings of one surface render the same status and `Report` message, so a streaming failure and a JSON failure of the same error agree on everything but framing + +| Surface | JSON | SSE | +|---|---|---| +| Anthropic, `/v1/messages` | `body(request_id)` | `sse_frame()` | +| OpenAI, chat, responses and the rest | `openai_response()` | `responses_sse_frame()`, the Responses API `error` event | + +The message of anything that is not a provider answer is the `litellm_error::Report` of the error, and in the OpenAI and Anthropic JSON envelopes the error `type` string is derived from the status by `error_type()`, so a new failure never names its own type. The Responses `error` event has the fixed type `error` and the status as its `code`. The Anthropic envelope passes a provider body already in its shape through untouched and extracts the message from anything else. The OpenAI envelope always wraps, because that is what the Python proxy returns there + +## Every way in uses the same Error + +A request can fail before the handler runs, in the handler, or after the headers went out. All three end in the same `Error`, so they get the same classification and envelope + +- An extractor sets `type Rejection = Error`, so a bad body is the surface's envelope rather than axum's plain text +- A path the gateway knows but does not serve returns `Error::Unsupported(path)` from one handler, so the client gets a 501 in the surface's envelope rather than a bare 404 +- A stream that fails after its headers went out hands the surface's SSE frame of `Error::from(error)` to `litellm_host_http::Sse::new`. `litellm_host_http::Error` separates `Call(E)`, the route's own error, from `Protocol`, a host bug, and the `From` maps the second to `Internal` + +The axum hooks behind each line are in [`rust-gateway`](../../rust-gateway/SKILL.md) + +## Testing + +`status()` and `body()` are pure, so unit tests call them on errors built the way production builds them, one `#[case]` per fault, as `status_follows_who_is_at_fault`, `an_anthropic_provider_body_passes_through` and `other_failures_get_the_anthropic_envelope_with_their_message` in `crates/gateway-inference/src/error.rs` do. Assert the status and the fields the surface promises, not the whole body, except where passing the body through untouched is the contract. Driving a whole router belongs to the gateway crate's own tests, see [`rust-gateway`](../../rust-gateway/references/testing.md) diff --git a/litellm-rust/.agents/skills/rust-error/references/messages-and-causes.md b/litellm-rust/.agents/skills/rust-error/references/messages-and-causes.md new file mode 100644 index 00000000000..80017423a42 --- /dev/null +++ b/litellm-rust/.agents/skills/rust-error/references/messages-and-causes.md @@ -0,0 +1,112 @@ +# Messages and causes + +## Messages belong to the variant + +Every message template is written once, in the variant's `#[error(...)]`. A caller passes the typed facts the template needs and nothing else + +```rust +// none of these may appear at a call site +Error::Network("WebSocket handshake returned a wrong Sec-WebSocket-Accept".into()) +Error::Request(format!("bad header {name}")) +Error::Signature(error.to_string()) + +// the call site names what happened, the variant owns the words +Repr::Handshake(HandshakeFault::Accept).into() +Error::signature(error) +Error::NonAsciiHeader { name } +``` + +When a check fails for a distinct reason, that reason is a variant or a field, never a string inside a generic variant. A `Message(String)` variant, or a `From` / `From<&str>` into an error, is the hole every stringly call site goes through, so neither exists + +Format arguments can call functions and reach into fields, so a hand-written `Display` is rarely needed + +```rust +#[error("could not read {}", path.display())] +#[error("invalid model: {provider} has no model {model:?} - use one of: {}", supported.join(", "))] +``` + +Keep messages short and lowercase, with no trailing period, unless Python parity pins the exact text. A reporter appends `: {cause}` after them. Never format `Debug` (`{:?}`) of an error, or of a type with no `Display`, into a message + +## The cause is the source, never the message + +This is the rule from the [`std::error::Error` docs](https://doc.rust-lang.org/std/error/trait.Error.html): an error's `Display` describes that layer only, and the error underneath is returned by `source()`. thiserror wires `source()` from a field marked `#[source]`, a `#[from]` field, or a field named `source` + +```rust +// prints the cause twice under any reporter +#[error("could not read {}: {source}", path.display())] +Read { path: PathBuf, source: io::Error }, +#[error("invalid JSON: {0}")] +Json(#[from] serde_json::Error), + +// this layer's words, the cause is the source +#[error("could not read {}", path.display())] +Read { path: PathBuf, source: io::Error }, +#[error("invalid JSON")] +Json(#[from] serde_json::Error), +``` + +| Field | thiserror does | Use for | +|---|---|---| +| `source: E` or `#[source] cause: E` | exposes it as `source()` | a cause with context fields next to it | +| `#[from] E` | implies `#[source]` and adds `From` | context-free wrapping, see [conversions](conversions.md) | +| `#[error(transparent)]` | forwards both `Display` and `source()` to the inner value | nesting another of our errors that adds no text | +| no attribute, any other name | nothing, the field is plain data | facts the message renders, such as a path or a status | + +`#[error(transparent)]` is an invisible layer, not an extra chain node, so a transparent wrapper's `source()` is the inner value's `source()` + +A field that holds an error but is neither rendered nor a source hides that error from every reader. Every error-typed field is a source. A fieldless reason enum, such as `HandshakeFault` or `PagesFault`, is data and not a cause, so the variant renders it in its message + +`source()` is one chain, so an error with several causes, such as `CredentialChain` trying each credential in turn, makes the most relevant one its source and renders the others in its message, each with `Report` + +## Keep causes typed + +Store the lower-level error itself, never its string, and never drop it with `map_err(|_| ...)` + +| Flattening or dropping loses | Because | +|---|---| +| the chain | a `String` has no `source()` | +| downcasting | the type is gone | +| redaction | `reqwest::Error`'s `Display` includes the request URL, which can carry credentials in the path or query | + +If the cause is not `Clone` and the error must be, hold it as `#[source] litellm_error::Source`, a cloneable `Arc` that is transparent to `Display` and `source()`, or make the crate's error opaque. `litellm_error::find::(&error)` walks the chain and looks through `Source`, so it is how anyone reaches a typed cause. A third-party `Box` converts with `Source::from`, and an `anyhow::Error` a dependency returns goes through `Box` first, `Source::from(Box::::from(error))`, and a dependency's own boxed alias such as `tokenizers::Error` is held as it is. A source boxed to keep the variant small is a `Box` node of the chain, so it is reached as `find::>`. Never put `#[error(transparent)]` on a `Source`, the erased error would then be no node of the chain and `find` could not reach it. Do not stringify or drop a cause to get `Clone` or `PartialEq` + +```rust +// loses the type and leaks the URL +.map_err(|error| Error::Network(error.to_string())) +// loses everything +.map_err(|_| Error::Transport) + +// keeps the cause, the From impl strips the URL +.map_err(litellm_http::Error::from) +``` + +A cause that must not be shown at all, such as a secret manager's response, is still a source. The type derives [`veil::Redact`](https://docs.rs/veil) with `#[redact]` on the field, so `Debug` hides it, and the boundary decides how much of the chain to show, see [boundaries](boundaries.md) + +## Strip reqwest URLs on entry + +Every `reqwest::Error` enters our types through `From for litellm_http::Error`, which calls `without_url()`. A crate that talks HTTP holds `litellm_http::Error`, never `reqwest::Error`. When a dependency's own error wraps a `reqwest::Error`, as vaultrs and rustify do, one `From` takes the reqwest error out and converts it, as `secrets-hashicorp` does. That way so every later `Display`, `Debug` or chain walk is safe from that URL + +`without_url()` only removes the URL reqwest attached. It does not scrub headers, response bodies, provider messages, file paths or anything a variant stores next to it. Those need their own policy, see [boundaries](boundaries.md) + +## Rendering + +Anything that shows an error to someone else renders the whole chain with `litellm_error::Report`, which writes each layer once joined by `: `, the single-line form of the unstable [`std::error::Report`](https://doc.rust-lang.org/std/error/struct.Report.html). `to_string()` of an error is only its top layer and is never what a client, a Python exception or a log line shows + +```rust +let message = litellm_error::Report::new(&error).to_string(); +// "could not reach the provider: error sending request: client error (Connect): tcp connect error: Connection refused" +``` + +## #[from] limits + +`#[from]` only works on a variant whose only fields are the source and an optional backtrace. A variant that needs context such as a path, a phase or a `TlsSource` takes its cause through an explicit `map_err` or a named constructor + +```rust +std::fs::read(&path).map_err(|source| Repr::Read { + path: path.clone(), + tls_source: TlsSource::CaBundle, + source, +})?; +``` + +`#[backtrace]` and automatic `provide()` need a nightly compiler and are not used here diff --git a/litellm-rust/.agents/skills/rust-error/references/pyo3.md b/litellm-rust/.agents/skills/rust-error/references/pyo3.md new file mode 100644 index 00000000000..81a7a515726 --- /dev/null +++ b/litellm-rust/.agents/skills/rust-error/references/pyo3.md @@ -0,0 +1,74 @@ +# Python bridge + +`crates/python-bridge` turns a Rust error into a Python exception through one mapping function per error type. A `#[pyfunction]` returns `Err(PyErr)` and pyo3 raises it, so those functions are the whole boundary + +| Error | Mapping | +|---|---| +| `litellm_inference::Error` from a route | `errors::route_failure`, with `route_error_to_pyerr` as its shorthand for routes that set no marker | +| the OCR error | `route_failure` and `to_pyerr` in `routes/ocr/errors.rs` | +| token counter, cache, spend, secrets, argument checks | one `From for PyErr` or mapping function next to the code that registers the functions | + +A route host never builds a `PyErr` from a route error itself. It calls the mapping, and the mapping is the only place that knows the exception classes and their attributes + +## Exception classes + +| Class | Meaning | Payload | +|---|---|---| +| `RustUpstreamError` | the provider call was issued and the provider answered with a failure | `args == (status, body)` and a `headers` attribute of `(name, value)` pairs, always set | +| `RustBridgeDeclined` | the route stopped before calling the provider, so the host may run its own path | `args == (message,)` | +| a standard class from `pyo3::exceptions` | whichever Python class means the same thing, such as `ValueError` or `TypeError` for the caller's input, `FileNotFoundError` or `OSError` for a file, `KeyError` or `OverflowError` where Python raises them, and `RuntimeError` for anything else | the rendered chain, or a parity text where Python pins one | + +The two native classes come from `create_exception!` in `crates/python-bridge/src/errors.rs` and are exported with `#[pymodule_export]`, which is how the [pyo3 guide](https://github.com/PyO3/pyo3/blob/main/guide/src/exception.md) makes an exception importable. Their docstrings are the contract Python reads. `errors::upstream_error(py, status, body, headers)` is the one constructor of `RustUpstreamError`, so the payload cannot differ between routes + +Python decodes a `RustUpstreamError` in `litellm/rust_bridge/failures.py`, which validates `args` and `headers` with a `TypeAdapter` and rebuilds the provider's `httpx.Response`. `litellm/rust_bridge/runtime.py` catches `RustBridgeDeclined` and reads its reason. `_raise_upstream` in `runtime.py` still reads `args` a second time by hand for the `attempt` path, which is the one decode to fold into `failures.py`. A new fact goes into `args` or `headers`, and the decoder learns to read it in the same change + +Two routes also flag a request fault with a boolean attribute, read by their Python route module: `ocr_request_format_error` in `host/routes/ocr.py` and `messages_request_error` in `host/routes/messages.py`. OCR additionally sets `status_code` and `message` for Python parity, and `native_settings_error`, which only the OCR host reads back to tell a settings failure from the caller's exception. Do not add another attribute, nothing on the Python side types them + +## Why args, not a typed subclass + +The bridge builds with pyo3's `abi3-py310` feature. Under abi3, `#[pyclass(extends = PyException)]` is only allowed on Python 3.12 and newer, so a class with typed fields is not available. Even where it is allowed, pyo3 passes an empty tuple to the native base's `tp_new`, so `.args` and `str()` would need manual overrides. The payload therefore travels in `args`, which a `create_exception!` class fills from `new_err` + +## Mapping order + +A Python exception raised inside the call, such as a secret manager callback's, is returned first as it is. Then the provider's answer passes through as `RustUpstreamError`. On `/v1/messages` a request the route rejected becomes a `ValueError` flagged with `messages_request_error`. Everything else is split by `is_request()` into `ValueError` or `RuntimeError`, with the message rendered by `litellm_error::Report`. A missing API key is a `RuntimeError` because `is_request()` answers false for it, and `missing_api_key_stays_a_runtime_error_while_other_auth_failures_are_value_errors` pins that + +OCR puts Python parity arms in front of that order, each matching a top-level variant of its own public enum + +| Arm | Exception | +|---|---| +| the provider answered | `RustUpstreamError` with the provider's headers | +| `RequestFormat` | `ValueError` flagged with `ocr_request_format_error` | +| `FileRead` whose cause is `NotFound` | `FileNotFoundError` naming the path | +| other `FileRead` | `OSError` with the io message | + +Both routes ask `litellm_auth::Error::is_request()` through their own `is_request()`, so a missing key or API base is a `RuntimeError` on both, and a missing setting the caller can pass, such as a region, is a `ValueError` + +Use the standard class whose Python meaning matches. `PyFileNotFoundError` for a missing file is right, a `RuntimeError` whose text says the file is missing is not. `import_exception!` exists for an exception defined in Python, none is used today because `failures.py` maps to LiteLLM's public exceptions through `litellm.exception_type` + +## Declining + +`RustBridgeDeclined` means no provider call was made and the host may retry on its own path. It is raised before dispatch, by the tokenizer, embeddings and token counter routes. Never raise it for a failure after the request went out, the host would repeat a billed call. `route_failures_are_terminal` in `errors.rs` pins that for the inference routes + +## Rules + +- Never format `Debug` of an error into a `PyErr` +- Never build a `PyErr` from a string at a call site. A failed argument check is a variant of a typed error, as `ProjectionError` in `coercion.rs` is, mapped once +- A Python exception raised inside the call stays that exception. Return it as is, never wrap it in a Rust message or stringify it into a new `ValueError`. The one policy exception is a caller-supplied callable, whose failure `callable::wrap_failure` reports as a `RuntimeError` with the original as `__cause__` +- A `panic!`, `unwrap` or `expect` inside a `#[pyfunction]` or `#[getter]` reaches Python as `pyo3_runtime.PanicException`, which is not an `Exception` subclass. Model the failure as a value instead +- When the mapping itself calls Python, as `setattr` does, propagate its failure with `?` rather than discarding it with `.ok()` + +## Testing + +Build the Rust error the way production builds it, map it, and assert the contract Python sees. `Python::initialize()` once, then `Python::attach` + +```rust +Python::initialize(); +Python::attach(|py| { + let failure = route_error_to_pyerr(litellm_http::Error::upstream(429, "slow down".into()).into()); + assert!(failure.is_instance_of::(py)); + let args: (u16, String) = failure.value(py).getattr("args").unwrap().extract().unwrap(); + assert_eq!(args, (429, "slow down".into())); +}); +``` + +One `#[case]` per fault class, asserting `is_instance_of` for the class and `extract` for the payload. The Python side of the contract is tested in `tests/unit/rust_bridge`, through `failures.py` diff --git a/litellm-rust/.agents/skills/rust-error/references/recipes.md b/litellm-rust/.agents/skills/rust-error/references/recipes.md new file mode 100644 index 00000000000..80ba1512f77 --- /dev/null +++ b/litellm-rust/.agents/skills/rust-error/references/recipes.md @@ -0,0 +1,64 @@ +# Recipes + +## A new failure in litellm_http + +Say media downloads must reject a `Content-Type` outside an allow-list + +1. Add a `Repr` variant with typed facts and the message, `#[error("media type {0:?} is not allowed")] DisallowedMediaType(String)` +2. Return it where the check fails, `return Err(Repr::DisallowedMediaType(content_type.to_owned()).into())` +3. Add an arm to every predicate's exhaustive `match`, the compiler lists them, the way its neighbor `DownloadTooLarge` answers. A media download happens before the provider call, so it is not in `dispatched()`. The caller chose the URL, so it is in `is_request()` +4. Leave both boundaries alone. With those answers the gateway returns 400 and Python raises `ValueError` +5. Test it with a mock server answering `Content-Type: application/x-msdownload` and assert `error.is_request()` + +## A provider-specific failure + +A second OCR provider needs a missing-key error whose text names its environment variable and how to set it + +1. Look for an existing provider-neutral variant whose template fits. `litellm_auth::Error::MissingApiKey { provider, environment_variable }` renders `Missing {provider} API Key - Set ...` +2. If none fits, find what varies between providers with the same failure, here the environment variable and the guidance after the dash, and add one neutral variant with those as fields, `#[error("Missing {environment_variable} - {guidance}")] MissingCredential { environment_variable: &'static str, guidance: &'static str }` +3. Pass the provider's values from its transformation. The next provider passes its own and adds no variant, and `is_missing_credential()` lists `MissingCredential { .. }` once + +## Wrapping a new lower crate in litellm_inference + +Say routes start calling `litellm_cache` and its errors must surface + +1. Add one variant, `#[error(transparent)] Cache(#[from] litellm_cache::Error)` +2. Extend every predicate's `match` with an arm that calls through to the cache error's own predicates, adding a predicate there if it has none +3. Do not copy the cache error's variants into `litellm_inference::Error` + +## A foreign error that needs context + +Reading a file fails with `io::Error`, and the message must name the file + +1. A variant with the context and the cause, `Read { path: PathBuf, source: io::Error }` with `#[error("could not read {}", path.display())]` +2. Build it with `map_err(|source| Repr::Read { path: path.clone(), source })`, not a `From` +3. If the error must be `Clone`, the crate's error is opaque or the field is `litellm_error::Source` + +## Migrating a stringly error + +Call sites look like `Error::Network(error.to_string())`, `Error::Request(format!(...))` or `.map_err(|_| Error::Unavailable)` + +1. List every call site of the string variant and group them by what actually failed +2. Each group becomes a variant holding the typed cause or facts. Groups no caller tells apart share one +3. Where a foreign error was stringified with classification logic around it, move that logic into the `From` and turn the call sites into `?` +4. Delete the string variant. If it cannot be deleted yet, nothing new may construct it +5. Re-run the boundary tests, since messages a client sees may change + +## Moving a crate to source chains + +A variant reads `#[error("invalid JSON: {0}")] Json(#[from] serde_json::Error)` or stores `cause` and renders `{cause}` + +1. Drop the cause from the template, `#[error("invalid JSON")]` +2. Make the cause the source, `#[from]`, `#[source]` or a field named `source` +3. Find what renders this error for someone else and make it use `litellm_error::Report` +4. Update tests that asserted the cause in `to_string()` to assert it in the report + +## Making a public enum opaque + +Downstream crates match nested variants, or a variant needs a non-`Clone` cause + +1. Rename the enum to `pub(crate) enum Repr` and add `pub struct Error(Arc)` with `#[error(transparent)]` and `From` +2. For each downstream `match`, find the question it asks and add a predicate. Replace the `match` with the predicate +3. For each downstream construction, add a constructor taking typed data, or move the failure into this crate +4. Swap test assertions to predicates and add the `#[cfg(test)] repr()` accessor for in-crate tests +5. Implement `PartialEq` with `Arc::ptr_eq` only if a wrapping enum's derive needs it diff --git a/litellm-rust/.agents/skills/rust-error/references/shape.md b/litellm-rust/.agents/skills/rust-error/references/shape.md new file mode 100644 index 00000000000..4d51c2cee6b --- /dev/null +++ b/litellm-rust/.agents/skills/rust-error/references/shape.md @@ -0,0 +1,105 @@ +# Error shape + +## One top-level Error per crate + +A crate's errors live in `src/error.rs` and are re-exported from `lib.rs` as `Error`. A caller of any public function in the crate should only ever see `Result<_, crate::Error>`, so one `?` and one set of predicates cover the whole crate + +This deliberately departs from the advice to put an error type next to each unit of fallibility ([Sabrina Jewson, "Modular Errors"](https://sabrinajewson.org/blog/errors)). Our callers do not branch on which function failed, they ask the predicates in [classification](classification.md), and the boundaries map one type per crate. Splitting by function would multiply those predicates and mappings without giving any caller a decision it lacks. Layering, the other half of that advice, is kept: each crate nests the one below it, see [conversions](conversions.md) + +That one type takes one of two forms + +| Form | Use when | Example | +|---|---|---| +| Public enum | Callers outside the crate match variants, and every variant holds our own data or another crate's `Error` that is already `Clone` and comparable | `litellm_inference::Error`, `litellm_auth::Error` | +| Opaque struct over a private `Repr` | The variants hold third-party causes, or the variants must be free to change without breaking callers | `litellm_http::Error` | + +Both forms are still one top-level `Error`. The opaque form just moves the enum behind `Repr`, where it is `pub(crate)` + +```rust +// public enum +#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] +pub enum Error { + #[error("missing required field: {0}")] + MissingField(&'static str), + #[error(transparent)] + Http(litellm_http::Error), +} + +// opaque struct +#[derive(Clone, Debug, thiserror::Error)] +#[error(transparent)] +pub struct Error(Arc); + +#[derive(Debug, thiserror::Error)] +pub(crate) enum Repr { + #[error("could not reach the provider")] + Connect(#[source] reqwest::Error), + #[error("upstream request failed with status {status}: {body}")] + Status { status: u16, body: String }, +} +``` + +Start with the public enum. Move to opaque when a variant needs a non-`Clone` third-party cause, or when downstream crates keep matching variants in ways that duplicate classification. `reqwest::Error`, `std::io::Error` and `serde_json::Error` are built the opaque way + +In the opaque form other crates cannot name a variant, so they cannot build an error with their own message. Public constructors exist only for errors another crate genuinely produces, take typed data and never a message, as in `Error::upstream(status, body)`, `Error::computed_header(name)`, `Error::unsignable_body()` and `Error::signature(cause)` + +`Arc` makes cloning a refcount bump, so `Repr` never has to be `Clone`. `Box` is enough when the goal is only a small `Result`. `Arc` is for when callers clone + +Every crate here is internal to the workspace, so do not mark errors or variants `#[non_exhaustive]`. It only forces a wildcard arm on the other workspace crates and turns off the exhaustiveness check that makes a new variant visit every `match`. When the variants must be free to change, use the opaque form + +## What else may sit next to Error + +These are not a second top-level error and need no folding in + +- Leaf data types whose fields are typed and whose message is in `#[error(...)]`, such as `litellm_http::request::HeaderError`. A wrapping crate nests it as one variant +- Plain data used by predicates, such as `litellm_http::TlsSource` +- A private fault enum one variant carries, such as `PemFault` inside `Repr::InvalidPem` + +## Variants are failure modes + +A public variant is something a caller handles differently: a phase, a status code, retry versus no retry, fault on the request versus the environment, or a message Python parity pins exactly. Failures no caller tells apart share one variant and differ in their data + +```rust +// one variant per provider, so every new provider adds a variant and a match arm +#[error("Missing Azure API Base - set AZURE_API_BASE")] +MissingAzureApiBase, +#[error("Missing Vertex API Base - set VERTEXAI_API_BASE")] +MissingVertexApiBase, + +// one failure mode, the provider and its guidance are data +#[error("Missing {provider} API Base - {guidance}")] +MissingApiBase { provider: &'static str, guidance: &'static str }, +``` + +An exact parity message does not justify a provider-specific variant when caller-supplied data can produce the same text + +Inside a private `Repr` a variant costs no API, so one per distinct reason is fine when each one carries a different message. When several only differ in a word and share every predicate answer, prefer one variant with a fieldless reason enum, so predicates list one arm instead of four, as `Repr::Handshake(HandshakeFault)` does + +```rust +#[error("WebSocket handshake returned {0}")] +Handshake(HandshakeFault), + +#[derive(Debug, thiserror::Error)] +pub(crate) enum HandshakeFault { + #[error("a wrong Sec-WebSocket-Accept")] + Accept, + #[error("an invalid Upgrade header")] + Upgrade, + #[error("an invalid Connection header")] + Connection, + #[error("an invalid subprotocol")] + Subprotocol, +} +``` + +## When to split + +Split into a separate type only when different functions fail in disjoint ways and different callers see them, so no caller matches variants its function can never return. A type per function overdoes it + +Name a split type after what went wrong, not after the function that returns it. `ConnectError` reads better than `OpenSessionError`, and a unit struct is fine for a single failure mode + +A split type still converts into the crate's `Error` when a caller further up needs one `Result` type + +`litellm-llms` has one recorded exception, the OCR error in `src/base_llm/ocr/error.rs`, kept apart until OCR folds into `litellm_llms::Error`, see `crates/inference-ocr/AGENTS.md`. Do not add another + +A utility crate whose modules are unrelated, such as `core-utils`, has no crate-wide failure, so each module that can fail owns one error next to its code instead of a `src/error.rs`. The rest of this skill applies to each of them diff --git a/litellm-rust/.agents/skills/rust-error/references/testing.md b/litellm-rust/.agents/skills/rust-error/references/testing.md new file mode 100644 index 00000000000..4aa22b4a293 --- /dev/null +++ b/litellm-rust/.agents/skills/rust-error/references/testing.md @@ -0,0 +1,37 @@ +# Testing errors + +Assert behavior through predicates and accessors, on errors produced the way production produces them + +```rust +// guesses the message and the variant, and cannot be written once the type is opaque +assert_eq!(error, Error::Transport(TransportError::Http { status: 429, body: "slow down".into() })); + +// asserts what callers rely on +assert_eq!(error.upstream_response(), Some((429, "slow down"))); +assert!(error.dispatched()); +assert!(!error.is_request()); +``` + +## Producing the error + +In order of preference + +1. A real failing call. Bind a `TcpListener`, drop it, then connect to its address for a connect failure. Accept and never answer for a timeout +2. A real cause, such as `tungstenite::Error::ConnectionClosed` or a `serde_json::from_str("{")` error, passed through our `From` +3. A public constructor, such as `Error::upstream(500, "boom".into())` + +Inside the owning crate a test may match the private representation through the `#[cfg(test)] fn repr(&self) -> &Repr` accessor in `crates/http/src/error.rs`, as the `media.rs` and `tls.rs` tests do. Anywhere else, reach the cause the test expects with `litellm_error::find::(&error)`. A transparent layer is no node of the chain, so `find` reaches foreign causes underneath it but never the nested error of ours itself. Ask that error's predicates through the outer type instead, which is what callers do too + +## What to cover + +| Property | Test shape | +|---|---| +| classification | one `#[case]` per input class, asserting the predicate | +| redaction | build from a URL holding `api_key=secret`, assert the message lacks `secret` and the path | +| cause kept | `litellm_error::find::(&error)` returns the expected cause | +| rendered once | the cause's text appears exactly once in `litellm_error::Report::new(&error).to_string()` and not in `to_string()` | +| boundary mapping | the gateway status and body, or the Python exception type and its `status` attribute | + +Assert exact message text only where it is a contract, such as a message Python parity pins. Everywhere else assert that the message contains the facts it must carry. Matching a fieldless variant with `matches!` is fine, full `assert_eq!` on an error that holds a cause is not, since causes compare by pointer + +A test that returns `Result` reports only `Debug` of the error on failure. Prefer `expect` with a short reason, or `unwrap_err` followed by assertions on the predicates diff --git a/litellm-rust/.agents/skills/rust-gateway/SKILL.md b/litellm-rust/.agents/skills/rust-gateway/SKILL.md new file mode 100644 index 00000000000..66b93ee128d --- /dev/null +++ b/litellm-rust/.agents/skills/rust-gateway/SKILL.md @@ -0,0 +1,20 @@ +--- +name: rust-gateway +description: Add or change an HTTP route, extractor, response envelope, stream or router mount in the litellm-rust gateway crates (crates/gateway and crates/gateway-*). Use when adding an endpoint, parsing a request body, choosing an error envelope for an API surface, relaying SSE, or writing a router-level test +--- + +# Rust gateway + +`crates/gateway` mounts endpoint routers and serves the listener. `crates/gateway-inference`, `gateway-mcp` and `gateway-ui` own their endpoints, `gateway-auth` is the inbound authentication middleware and `gateway-management` holds management logic that has no router yet. A crate's `AGENTS.md`, where it has one, says what belongs in it. Read the reference for the part you are touching + +- [surfaces](references/surfaces.md) + - Each API surface mirrors the envelope the Python proxy returns there, so the shapes differ on purpose + - The error type string is derived from the status and provider bodies pass through bounded +- [requests](references/requests.md) + - Extractors reject with the crate's `Error`, handlers take the body as a `Result` when the error needs the request id + - Every alias of a path is registered together, a known but unserved path goes to one handler + - Streams relay through `litellm_host_http::Sse`, with the error frame rendered by the crate's `Error` +- [testing](references/testing.md) + - Build the router through `tests/support` and drive it with `oneshot`, assert the fields the surface promises + +How an error is classified before it reaches a response is the [`rust-error`](../rust-error/SKILL.md) skill, in particular its [gateway](../rust-error/references/gateway.md) page. Axum's own conventions are in its [docs](https://docs.rs/axum/latest/axum/) diff --git a/litellm-rust/.agents/skills/rust-gateway/references/requests.md b/litellm-rust/.agents/skills/rust-gateway/references/requests.md new file mode 100644 index 00000000000..395f29f5c64 --- /dev/null +++ b/litellm-rust/.agents/skills/rust-gateway/references/requests.md @@ -0,0 +1,34 @@ +# Requests + +## Extractors reject with the crate's Error + +A custom extractor sets `type Rejection = Error`, so a bad body gets the surface's envelope rather than axum's plain-text rejection. Never take axum's own `Json`, `Query` or `Path` in a handler, their rejection is plain text. The rejection is the variant's source, never its string. `JsonObject` and `InferenceBody` in `crates/gateway-inference/src/request.rs` are the pattern + +```rust +impl FromRequest for JsonObject { + type Rejection = Error; + + async fn from_request(request: Request, state: &S) -> Result { + let body = Bytes::from_request(request, state).await.map_err(|error| { + if error.status() == StatusCode::PAYLOAD_TOO_LARGE { + Error::BodyTooLarge + } else { + BodyError::Read(error).into() + } + })?; + Ok(Self(object(&body)?)) + } +} +``` + +When the handler needs something from the request in the error body, such as the request id on `/v1/messages`, take the extractor as `Result` so the handler maps the rejection itself, as `messages::create` does. `Infallible` is the right rejection for an extractor that only reads an optional header, as `RequestId` does + +## Paths + +Every alias the Python proxy registers for a route is registered together in `crates/gateway-inference/src/lib.rs`, such as the bare path, its `/v1` form and the `/engines/{model}` and `/openai/deployments/{model}` forms, otherwise one alias silently behaves differently from the others. A path the gateway knows but does not serve, such as `/v1/completions`, is routed to `request::unsupported`, which returns `Error::Unsupported(path)` so the client gets a 501 in the surface's envelope rather than a bare 404 + +`crates/gateway` merges the mounted routers and adds the shared inbound middleware as `route_layer`s, so authentication runs only for matched routes. Endpoint paths and parsing stay in the mounted crate, see the two crates' `AGENTS.md` + +## Streaming + +Once the response headers are sent the status cannot change. A streaming handler builds `litellm_host_http::Sse::new(Json, |error| Bytes::from(Error::from(error).sse_frame()))`, where the closure renders a mid-stream failure as an `event: error` frame with the same status and message a JSON failure on that surface carries, `sse_frame()` for Anthropic and `responses_sse_frame()` for the Responses API's `error` event. `litellm_host_http::Error` separates `Call(E)`, the route's own error, from `Protocol`, a host bug, and the crate's `From` maps the second to `Internal` diff --git a/litellm-rust/.agents/skills/rust-gateway/references/surfaces.md b/litellm-rust/.agents/skills/rust-gateway/references/surfaces.md new file mode 100644 index 00000000000..08d4bbca4ea --- /dev/null +++ b/litellm-rust/.agents/skills/rust-gateway/references/surfaces.md @@ -0,0 +1,20 @@ +# Envelopes follow the API surface + +Each surface mirrors what the Python proxy returns there, so the shapes differ on purpose + +| Surface | Envelope | Where | +|---|---|---| +| `/v1/messages` | Anthropic, `{"type": "error", "error": {"type", "message"}, "request_id"}` | `Error::body()` and `sse_frame()` in `crates/gateway-inference/src/error.rs` | +| OpenAI-shaped routes | `{"error": {"message", "type", "param", "code"}}` | `Error::openai_response()` in the same file | +| `/v1/responses` stream | the Responses `error` event, `{"type": "error", "code", "message", "param"}` | `Error::responses_sse_frame()` in the same file | +| MCP | `{"detail": message}` | `crates/gateway-mcp/src/error.rs` | +| UI | `{"error": {"message"}}`, plus `Retry-After` on 429 | `crates/gateway-ui/src/error.rs` | +| auth rejections | plain text, the top-level message only | `crates/gateway-auth/src/error.rs` | + +Every surface takes its status from one `status()` on its crate's `Error` and renders the message of anything that is not a provider answer with `litellm_error::Report`. In `gateway-inference` the error `type` string is derived from the status by `error_type()`, so a new failure never names its own type. The Anthropic envelope passes a provider body already in its shape through untouched and extracts the message from anything else, and the OpenAI envelope always wraps, as the Python proxy does. A provider body is bounded by `litellm_http::Error::upstream` when the error is built + +Two surfaces show less than the chain on purpose, and a test in each pins it. An auth rejection goes to a caller who has not authenticated, so it never shows what a secret manager or key store said. An MCP upstream or authorization failure shows only its own layer, because rmcp renders the transport error, URL included, into its own `Display`, and an authorizer's cause is the same backend detail auth rejections hide. The inference surfaces show only the top layer of an auth rejection for the same reason + +`/v1/messages` adds the request id to the envelope when the provider did not set one, which is why its handler takes the body as `Result` and renders the rejection itself, see [requests](requests.md) + +A new surface gets its own renderer next to the others, named after the shape it produces, and its status still comes from `Error::status()`. Picking the status is the [`rust-error`](../../rust-error/references/gateway.md) skill's concern, not the surface's diff --git a/litellm-rust/.agents/skills/rust-gateway/references/testing.md b/litellm-rust/.agents/skills/rust-gateway/references/testing.md new file mode 100644 index 00000000000..10c1f01c472 --- /dev/null +++ b/litellm-rust/.agents/skills/rust-gateway/references/testing.md @@ -0,0 +1,9 @@ +# Testing + +Test HTTP contracts in the crate that owns the endpoint: status codes, forwarded headers, error envelopes and streaming behavior. Inference and provider behavior is tested in the `inference-` and `llms` crates, and mounting and server lifecycle in `crates/gateway` + +`crates/gateway-inference/tests/support/mod.rs` builds a `Router` with injected deployments, pools and settings through `app()` and its variants, and `post()` drives it with `tower::ServiceExt::oneshot`. `json()` reads the response body back. Use these instead of building a router or a client by hand in a test + +Assert the status and the fields the surface promises, not the whole body. A test that pins a whole envelope breaks on every unrelated field, one that checks `status`, `error.type` and the presence of `request_id` catches the regressions that matter + +`status()` and `body()` of the error type are pure, so a classification case belongs in the error module's unit tests, see the [`rust-error`](../../rust-error/references/gateway.md) skill. Router tests cover what only the full path shows, such as an extractor rejection arriving in the right envelope or an SSE error frame after the headers went out diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index e62a9463cb5..9e7a55eb232 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -4,6 +4,8 @@ For diagnostic tracing changes, follow [.agents/skills/rust-tracing/SKILL.md](.a For string-valued enums and their Serde conversions, follow [.agents/skills/rust-string-enums/SKILL.md](.agents/skills/rust-string-enums/SKILL.md) +For error types, error variants and conversions between crate errors, follow [.agents/skills/rust-error/SKILL.md](.agents/skills/rust-error/SKILL.md) + For fieldless enums, derive `strum::VariantArray` and use `VARIANTS` instead of a hand-listed `ALL` array. Derive Strum string conversions instead of hand-written variant-to-string matches Use the derived conversions directly (`<&'static str>::from(x)` / `.into()`, `str::parse`) with no `as_str`/`parse` wrapper that only delegates, and spell each variant with explicit `#[strum(serialize = "...")]` instead of `serialize_all` @@ -24,13 +26,3 @@ Use [`#[rstest]`](https://docs.rs/rstest/latest/rstest/attr.rstest.html) for new - Never loop over inputs (`for`, `.iter().for_each`, `.all`) inside a test body. Give each input its own `#[case::name(...)]`, or use `#[values(...)]` for a cross product - Exception: a test pinning a Rust table against a repo-owned data file (for example `include_str!` of a JSON config) may iterate that file's entries - -## Error definitions - -- A crate's errors live in `src/error.rs`, defined with `thiserror`, and re-exported from `lib.rs` -- Put message templates in the variant's `#[error(...)]` declaration. Callers pass only the small typed arguments needed to fill them, never `Error::Variant(format!(...))` or a preformatted message. Keep the smallest set of neutral variants that callers need to distinguish. Different wording or providers do not justify new variants -- Default to one top-level `Error` enum per crate, with one variant per failure mode and a `#[error(...)]` message on each. A failure mode is something a caller handles differently (phase, status code, retry, a message Python parity pins exactly). Failures no caller tells apart share one variant and differ only in its message -- Keep shared error enums minimal and provider-neutral. Provider names, credential types, configuration fields, and setup guidance belong in caller-supplied data, not dedicated variants or hardcoded shared messages. Reuse a variant for the same failure mode across providers, such as `MissingApiBase { provider: "Azure", guidance: "..." }`. An exact parity message does not justify a provider-specific variant when caller-supplied context can preserve it -- Wrap a lower-level error as a variant with `#[from]` or `#[source]` instead of flattening it to a string -- Exception: split into separate types when different functions fail in disjoint ways, especially when different callers see them. A shared enum would force every caller to match variants its function can never return -- Name a split type after what went wrong (a unit struct is fine for a single failure mode), not after the function that returns it diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index e998fe1abbb..9c45c5c4b1c 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3388,6 +3388,7 @@ dependencies = [ "bytes", "futures-util", "litellm-auth-types", + "litellm-error", "litellm-http", "moka", "reqwest", @@ -3407,6 +3408,7 @@ dependencies = [ "azure_core", "azure_identity", "litellm-auth-types", + "litellm-error", "moka", "rstest", "serde_json", @@ -3425,6 +3427,7 @@ dependencies = [ "google-cloud-auth", "http 1.4.2", "litellm-auth-types", + "litellm-error", "moka", "rstest", "serde", @@ -3439,6 +3442,7 @@ dependencies = [ name = "litellm-auth-types" version = "0.1.0" dependencies = [ + "litellm-error", "rstest", "serde", "serde_json", @@ -3452,6 +3456,7 @@ dependencies = [ name = "litellm-cache" version = "0.1.0" dependencies = [ + "litellm-error", "rstest", "serde", "serde_json", @@ -3477,6 +3482,7 @@ dependencies = [ "reqwest", "rstest", "serde_json", + "thiserror 2.0.19", "tokio", "url", "wiremock", @@ -3495,6 +3501,7 @@ dependencies = [ "serde-pickle", "serde_json", "tempfile", + "thiserror 2.0.19", "tokio", ] @@ -3507,11 +3514,13 @@ dependencies = [ "litellm-auth-types", "litellm-cache", "litellm-cache-testing", + "litellm-error", "litellm-http", "percent-encoding", "reqwest", "rstest", "serde_json", + "thiserror 2.0.19", "tokio", "wiremock", ] @@ -3534,6 +3543,7 @@ dependencies = [ "futures-util", "litellm-cache", "litellm-cache-testing", + "litellm-error", "litellm-http", "qdrant-client", "reqwest", @@ -3554,11 +3564,13 @@ version = "0.1.0" dependencies = [ "litellm-cache", "litellm-cache-testing", + "litellm-error", "r2d2", "redis", "redis-test", "rstest", "serde_json", + "thiserror 2.0.19", "tokio", ] @@ -3611,6 +3623,7 @@ dependencies = [ "litellm-auth-aws", "litellm-cache", "litellm-cache-testing", + "litellm-error", "litellm-http", "reqwest", "rstest", @@ -3633,11 +3646,13 @@ dependencies = [ "litellm-cache", "litellm-cache-redis", "litellm-cache-testing", + "litellm-error", "redis", "redis-test", "rstest", "serde_json", "sha2 0.10.9", + "thiserror 2.0.19", "tokio", "uuid", ] @@ -3647,6 +3662,7 @@ name = "litellm-callbacks-legacy-python" version = "0.1.0" dependencies = [ "litellm-auth", + "litellm-error", "litellm-host", "litellm-host-python", "proptest", @@ -3654,6 +3670,7 @@ dependencies = [ "rstest", "serde_json", "strum", + "thiserror 2.0.19", ] [[package]] @@ -3669,6 +3686,7 @@ name = "litellm-config" version = "0.1.0" dependencies = [ "litellm-auth-types", + "litellm-error", "litellm-router-types", "rstest", "serde", @@ -3686,6 +3704,7 @@ name = "litellm-core-utils" version = "0.1.0" dependencies = [ "fancy-regex 0.19.2", + "litellm-error", "litellm-llms-types", "litellm-tracing", "rstest", @@ -3713,6 +3732,7 @@ dependencies = [ "criterion", "proptest", "rstest", + "thiserror 2.0.19", ] [[package]] @@ -3742,6 +3762,14 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-error" +version = "0.1.0" +dependencies = [ + "rstest", + "thiserror 2.0.19", +] + [[package]] name = "litellm-framer" version = "0.1.0" @@ -3769,6 +3797,7 @@ dependencies = [ "litellm-auth-types", "litellm-circuit-breaker", "litellm-config", + "litellm-error", "litellm-gateway-auth", "litellm-gateway-inference", "litellm-gateway-mcp", @@ -3802,6 +3831,7 @@ dependencies = [ "futures-util", "litellm-auth-types", "litellm-config", + "litellm-error", "litellm-secrets", "rstest", "serde", @@ -3825,6 +3855,7 @@ dependencies = [ "litellm-auth", "litellm-cache-memory", "litellm-cache-response", + "litellm-error", "litellm-gateway-auth", "litellm-host", "litellm-host-http", @@ -3844,6 +3875,7 @@ dependencies = [ "serde_json", "thiserror 2.0.19", "tokio", + "tokio-tungstenite", "tower", "wiremock", ] @@ -3854,6 +3886,7 @@ version = "0.1.0" dependencies = [ "base64 0.22.1", "litellm-auth-types", + "litellm-error", "litellm-gateway-auth", "rand 0.8.7", "rstest", @@ -3871,6 +3904,7 @@ dependencies = [ "http 1.4.2", "litellm-auth-types", "litellm-config", + "litellm-error", "litellm-secrets", "moka", "rmcp", @@ -3885,6 +3919,7 @@ dependencies = [ "tower", "url", "uuid", + "veil", ] [[package]] @@ -3894,9 +3929,11 @@ dependencies = [ "axum", "axum-login", "base64 0.22.1", + "futures-util", "governor", "jsonwebtoken", "litellm-auth-types", + "litellm-error", "litellm-gateway-auth", "rand 0.8.7", "rstest", @@ -3921,6 +3958,7 @@ dependencies = [ "litellm-coroutine", "rstest", "serde_json", + "thiserror 2.0.19", "tokio", ] @@ -3958,6 +3996,7 @@ dependencies = [ "bytes", "futures-util", "litellm-coroutine", + "litellm-error", "litellm-host", "parking_lot", "pyo3", @@ -3978,6 +4017,7 @@ dependencies = [ "http 1.4.2", "hyper-util", "litellm-core-utils", + "litellm-error", "rcgen", "reqwest", "rstest", @@ -3992,6 +4032,7 @@ dependencies = [ "tokio-rustls 0.26.4", "tokio-tungstenite", "tracing", + "url", "veil", "webpki-roots", ] @@ -4007,6 +4048,7 @@ dependencies = [ "litellm-cache-memory", "litellm-cache-response", "litellm-core-utils", + "litellm-error", "litellm-framer", "litellm-host", "litellm-http", @@ -4034,6 +4076,7 @@ dependencies = [ "litellm-cache-memory", "litellm-cache-response", "litellm-core-utils", + "litellm-error", "litellm-host", "litellm-host-native", "litellm-http", @@ -4043,7 +4086,6 @@ dependencies = [ "litellm-llms-types", "litellm-secrets", "litellm-tracing", - "reqwest", "rstest", "serde_json", "tokio", @@ -4062,6 +4104,7 @@ dependencies = [ "litellm-cache-memory", "litellm-cache-response", "litellm-core-utils", + "litellm-error", "litellm-framer", "litellm-host", "litellm-host-native", @@ -4093,6 +4136,7 @@ dependencies = [ "litellm-auth-aws", "litellm-auth-gcp", "litellm-core-utils", + "litellm-error", "litellm-host", "litellm-host-native", "litellm-http", @@ -4110,6 +4154,7 @@ dependencies = [ "time", "tokio", "tracing", + "url", "wiremock", ] @@ -4125,6 +4170,7 @@ dependencies = [ "litellm-cache-memory", "litellm-cache-response", "litellm-core-utils", + "litellm-error", "litellm-host", "litellm-host-native", "litellm-http", @@ -4162,6 +4208,7 @@ dependencies = [ "futures-util", "litellm-auth", "litellm-core-utils", + "litellm-error", "litellm-http", "litellm-inference", "litellm-inference-testing", @@ -4190,6 +4237,7 @@ dependencies = [ "litellm-auth-azure", "litellm-auth-gcp", "litellm-core-utils", + "litellm-error", "litellm-framer", "litellm-host", "litellm-http", @@ -4230,6 +4278,7 @@ version = "0.1.0" dependencies = [ "indexmap 2.14.0", "jsonschema", + "litellm-error", "litellm-llms-types", "rstest", "schemars 1.2.2", @@ -4256,6 +4305,7 @@ dependencies = [ "litellm-cache-response", "litellm-callbacks-legacy-python", "litellm-core-utils", + "litellm-error", "litellm-host", "litellm-host-python", "litellm-http", @@ -4286,6 +4336,7 @@ dependencies = [ "serde_with", "sqlx", "strum", + "tempfile", "thiserror 2.0.19", "tokio", "tokio-tungstenite", @@ -4363,6 +4414,7 @@ dependencies = [ "google-cloud-auth", "google-cloud-kms-v1", "litellm-core-utils", + "litellm-error", "litellm-http", "litellm-python-compat", "litellm-secrets-aws", @@ -4421,6 +4473,7 @@ dependencies = [ "sha2 0.10.9", "thiserror 2.0.19", "tokio", + "url", "veil", "wiremock", ] @@ -4444,6 +4497,7 @@ dependencies = [ "tempfile", "thiserror 2.0.19", "tokio", + "url", "veil", "wiremock", ] @@ -4470,6 +4524,7 @@ dependencies = [ "serde_json", "thiserror 2.0.19", "tokio", + "url", "veil", "wiremock", ] @@ -4479,7 +4534,10 @@ name = "litellm-secrets-hashicorp" version = "0.1.0" dependencies = [ "litellm-core-utils", + "litellm-error", + "litellm-http", "litellm-secrets-types", + "reqwest", "rstest", "rustify", "rustify_derive", @@ -4513,6 +4571,7 @@ name = "litellm-spend-clickhouse" version = "0.1.0" dependencies = [ "flate2", + "litellm-error", "litellm-http", "litellm-storage-clickhouse", "rstest", @@ -4533,6 +4592,7 @@ name = "litellm-storage-clickhouse" version = "0.1.0" dependencies = [ "flate2", + "litellm-error", "litellm-http", "rstest", "serde", @@ -4552,6 +4612,7 @@ version = "0.1.0" dependencies = [ "flate2", "futures-util", + "litellm-http", "reqwest", "rstest", "semver", @@ -4575,6 +4636,7 @@ dependencies = [ "criterion", "indexmap 2.14.0", "itoa", + "litellm-error", "litellm-token-counter-fast", "litellm-token-counter-huggingface", "litellm-token-counter-tiktoken", @@ -4591,6 +4653,7 @@ name = "litellm-token-counter-fast" version = "0.1.0" dependencies = [ "base64 0.22.1", + "litellm-error", "rand 0.8.7", "rstest", "rustc-hash", @@ -4616,6 +4679,7 @@ name = "litellm-token-counter-tiktoken" version = "0.1.0" dependencies = [ "base64 0.22.1", + "litellm-error", "once_cell", "rstest", "rustc-hash", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 745b29dd0ff..b7b22cfd32c 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -10,6 +10,7 @@ repository = "https://github.com/BerriAI/litellm" [workspace.dependencies] litellm-config = { path = "crates/config" } +litellm-error = { path = "crates/error" } litellm-router = { path = "crates/router" } litellm-router-types = { path = "crates/router-types" } litellm-tracing = { path = "crates/tracing" } diff --git a/litellm-rust/crates/auth-aws/Cargo.toml b/litellm-rust/crates/auth-aws/Cargo.toml index 1a7e2cd6782..e4578777a08 100644 --- a/litellm-rust/crates/auth-aws/Cargo.toml +++ b/litellm-rust/crates/auth-aws/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true [dependencies] litellm-auth-types.workspace = true +litellm-error.workspace = true litellm-http.workspace = true reqwest.workspace = true diff --git a/litellm-rust/crates/auth-aws/src/aws.rs b/litellm-rust/crates/auth-aws/src/aws.rs index d83b01c24ad..81aa4b068bc 100644 --- a/litellm-rust/crates/auth-aws/src/aws.rs +++ b/litellm-rust/crates/auth-aws/src/aws.rs @@ -15,13 +15,13 @@ use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; use litellm_auth_types::AwsParams; -use super::Error; use super::constants::{ AWS_ACCESS_KEY_ID, AWS_EXTERNAL_ID, AWS_PROFILE_NAME, AWS_REGION_NAME, AWS_ROLE_ARN, AWS_ROLE_NAME, AWS_SECRET_ACCESS_KEY, AWS_SESSION_NAME, AWS_SESSION_TOKEN, AWS_SIGNED_HEADER_NAMES, AWS_STS_ENDPOINT, AWS_WEB_IDENTITY_TOKEN, AWS_WEB_IDENTITY_TOKEN_FILE, DEFAULT_BEDROCK_REGION, DEFAULT_SESSION_NAME_PREFIX, SIGV4_COMPUTED_HEADER_NAMES, }; +use crate::error::{CredentialFlow, Error, SigningStage}; const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60); const AMBIENT_CREDENTIALS_TTL: Duration = Duration::from_secs(600); @@ -271,7 +271,7 @@ impl AwsAuthService { provider .provide_credentials() .await - .map_err(|error| Error::AwsProfile(error.to_string())) + .map_err(Error::credentials(CredentialFlow::Profile)) } AwsAuthFlow::AssumeRole { role, session_name } => { if is_already_running_as_role(&role, &resolved).await? { @@ -287,7 +287,7 @@ impl AwsAuthService { let credentials = provider .provide_credentials() .await - .map_err(|error| Error::AwsDefaultChain(error.to_string()))?; + .map_err(Error::credentials(CredentialFlow::DefaultChain))?; self.set_cached_credentials( key, credentials.clone(), @@ -327,7 +327,7 @@ impl AwsAuthService { provider .provide_credentials() .await - .map_err(|error| Error::AwsAssumeRole(error.to_string())) + .map_err(Error::credentials(CredentialFlow::AssumeRole)) } AwsAuthFlow::WebIdentity { token, @@ -350,12 +350,14 @@ impl AwsAuthService { .web_identity_token(token) .send() .await - .map_err(|error| Error::AwsWebIdentity(error.to_string()))?; + .map_err(Error::credentials(CredentialFlow::WebIdentity))?; let credentials = response .credentials() - .ok_or(Error::AwsMissingWebIdentityCredentials)?; - let expiration = SystemTime::try_from(*credentials.expiration()) - .map_err(|error| Error::AwsWebIdentityExpiration(error.to_string()))?; + .ok_or(Error::MissingWebIdentityCredentials)?; + let expiration = + SystemTime::try_from(*credentials.expiration()).map_err(|source| { + Error::WebIdentityExpiration(litellm_error::Source::new(source)) + })?; Ok(Credentials::new( credentials.access_key_id(), credentials.secret_access_key(), @@ -376,7 +378,7 @@ impl AwsAuthService { let credentials = provider .provide_credentials() .await - .map_err(|error| Error::AwsDefaultChain(error.to_string()))?; + .map_err(Error::credentials(CredentialFlow::DefaultChain))?; self.set_cached_credentials( key, credentials.clone(), @@ -480,14 +482,16 @@ pub fn sign_headers( .settings(SigningSettings::default()) .build() .map(SigningParams::from) - .map_err(|error| Error::AwsSigningParameters(error.to_string()))?; + .map_err(Error::signing(SigningStage::Parameters))?; let header_refs = headers .iter() .map(|(name, value)| { value .to_str() .map(|value| (name.as_str(), value)) - .map_err(|_| Error::AwsSignableRequest(format!("{name} is not visible ASCII"))) + .map_err(|_| Error::NonAsciiHeader { + name: name.to_string(), + }) }) .collect::, _>>()?; let request = SignableRequest::new( @@ -496,23 +500,20 @@ pub fn sign_headers( header_refs.into_iter(), SignableBody::Bytes(body), ) - .map_err(|error| Error::AwsSignableRequest(error.to_string()))?; + .map_err(Error::signing(SigningStage::Request))?; let (instructions, _) = sign(request, ¶ms) - .map_err(|error| Error::AwsSigning(error.to_string()))? + .map_err(Error::signing(SigningStage::Signature))? .into_parts(); instructions .headers() .map(|(name, value)| { - HeaderName::from_bytes(name.as_bytes()) - .map_err(|error| error.to_string()) - .and_then(|name| { - HeaderValue::from_str(value) - .map(|value| (name, value)) - .map_err(|error| error.to_string()) - }) + let name = HeaderName::from_bytes(name.as_bytes()) + .map_err(Error::signing(SigningStage::Output))?; + let value = + HeaderValue::from_str(value).map_err(Error::signing(SigningStage::Output))?; + Ok((name, value)) }) - .collect::>() - .map_err(Error::AwsSigning) + .collect() } /// Model-id and region parsing shared by every Bedrock route. @@ -1081,6 +1082,55 @@ mod tests { assert!(!signed.contains_key("X-Amz-Security-Token")); } + #[test] + fn a_header_value_outside_visible_ascii_is_named_instead_of_signed() { + let (url, body, mut headers) = parity_inputs(); + headers.insert("x-amz-meta-note", HeaderValue::from_str("café").unwrap()); + let credentials = Credentials::new("AKIDEXAMPLE", "secret", None, None, "test"); + + let error = sign_headers( + "POST", + &url, + &body, + &headers, + SigningScope { + region: "us-east-1", + service: BEDROCK_SERVICE, + credentials: &credentials, + time: UNIX_EPOCH, + }, + ) + .unwrap_err(); + + assert_eq!( + error, + Error::NonAsciiHeader { + name: "x-amz-meta-note".into() + } + ); + } + + #[test] + fn the_shared_auth_error_nests_the_aws_failure_and_its_cause() { + let cause = std::io::Error::new(std::io::ErrorKind::NotFound, "profile not found"); + let error = + litellm_auth_types::Error::from(Error::credentials(CredentialFlow::Profile)(cause)); + + assert!(matches!( + litellm_error::find::(&error), + Some(Error::Credentials { + flow: CredentialFlow::Profile, + .. + }) + )); + let report = litellm_error::Report::new(&error).to_string(); + assert!( + report.contains("could not load AWS profile credentials"), + "{report}" + ); + assert_eq!(report.matches("profile not found").count(), 1, "{report}"); + } + #[ignore] #[tokio::test] async fn live_bedrock_invoke_model_returns_200() -> Result<(), Box> { diff --git a/litellm-rust/crates/auth-aws/src/error.rs b/litellm-rust/crates/auth-aws/src/error.rs index d4c6ae8cf6c..5f88fc8a171 100644 --- a/litellm-rust/crates/auth-aws/src/error.rs +++ b/litellm-rust/crates/auth-aws/src/error.rs @@ -1,46 +1,73 @@ -use thiserror::Error as ThisError; +use litellm_error::Source; -#[derive(Clone, Debug, ThisError, PartialEq, Eq)] +#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] pub enum Error { - #[error("AWS profile credentials failed: {0}")] - AwsProfile(String), - #[error("AWS default credentials failed: {0}")] - AwsDefaultChain(String), - #[error("AWS role credentials failed: {0}")] - AwsAssumeRole(String), - #[error("AWS web identity credentials failed: {0}")] - AwsWebIdentity(String), - #[error("AWS web identity expiration was invalid: {0}")] - AwsWebIdentityExpiration(String), - #[error("AWS signing parameters failed: {0}")] - AwsSigningParameters(String), - #[error("AWS signable request failed: {0}")] - AwsSignableRequest(String), - #[error("AWS request signing failed: {0}")] - AwsSigning(String), + #[error("could not load AWS {flow} credentials")] + Credentials { + flow: CredentialFlow, + #[source] + source: Source, + }, #[error("AWS web identity response had no credentials")] - AwsMissingWebIdentityCredentials, + MissingWebIdentityCredentials, + #[error("AWS web identity response had an invalid expiration")] + WebIdentityExpiration(#[source] Source), + #[error("header {name} has a value that is not visible ASCII")] + NonAsciiHeader { name: String }, + #[error("{stage}")] + Signing { + stage: SigningStage, + #[source] + source: Source, + }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] +pub enum CredentialFlow { + #[error("profile")] + Profile, + #[error("default chain")] + DefaultChain, + #[error("assumed role")] + AssumeRole, + #[error("web identity")] + WebIdentity, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] +pub enum SigningStage { + #[error("SigV4 signing parameters are invalid")] + Parameters, + #[error("SigV4 cannot sign this request")] + Request, + #[error("SigV4 could not compute the signature")] + Signature, + #[error("SigV4 produced an invalid header")] + Output, +} + +impl Error { + pub(crate) fn credentials( + flow: CredentialFlow, + ) -> impl FnOnce(E) -> Self { + move |source| Self::Credentials { + flow, + source: Source::new(source), + } + } + + pub(crate) fn signing( + stage: SigningStage, + ) -> impl FnOnce(E) -> Self { + move |source| Self::Signing { + stage, + source: Source::new(source), + } + } } impl From for litellm_auth_types::Error { fn from(error: Error) -> Self { - Self::ProviderAuthentication(error.to_string()) - } -} - -#[cfg(test)] -mod tests { - use super::Error; - - #[test] - fn converts_to_shared_auth_error_without_losing_context() { - let error = litellm_auth_types::Error::from(Error::AwsProfile("profile not found".into())); - - assert_eq!( - error, - litellm_auth_types::Error::ProviderAuthentication( - "AWS profile credentials failed: profile not found".into() - ) - ); + Self::provider(error) } } diff --git a/litellm-rust/crates/auth-aws/src/lib.rs b/litellm-rust/crates/auth-aws/src/lib.rs index 0fe0b390110..bd0f866371f 100644 --- a/litellm-rust/crates/auth-aws/src/lib.rs +++ b/litellm-rust/crates/auth-aws/src/lib.rs @@ -5,5 +5,5 @@ mod signer; pub use aws::*; pub use aws_credential_types::Credentials; -pub use error::Error; +pub use error::{CredentialFlow, Error, SigningStage}; pub use signer::SigV4Signer; diff --git a/litellm-rust/crates/auth-aws/src/signer.rs b/litellm-rust/crates/auth-aws/src/signer.rs index 6c2cb4fd162..dff917d7c13 100644 --- a/litellm-rust/crates/auth-aws/src/signer.rs +++ b/litellm-rust/crates/auth-aws/src/signer.rs @@ -53,13 +53,13 @@ impl RequestSigner for SigV4Signer { .keys() .find(|name| is_sigv4_computed_header(name.as_str())) { - return Err(litellm_http::Error::ComputedHeader(name.to_string())); + return Err(litellm_http::Error::computed_header(name.as_str())); } let body = match request.body() { None => &[], - Some(body) => body.as_bytes().ok_or_else(|| { - litellm_http::Error::Signature("a streaming body cannot be signed".into()) - })?, + Some(body) => body + .as_bytes() + .ok_or_else(litellm_http::Error::unsignable_body)?, }; sign_headers( request.method().as_str(), @@ -73,7 +73,13 @@ impl RequestSigner for SigV4Signer { time: (self.clock)(), }, ) - .map_err(|error| litellm_http::Error::Signature(error.to_string())) + .map_err(|error| match error { + Error::NonAsciiHeader { .. } => litellm_http::Error::unsignable_header(error), + Error::Credentials { .. } + | Error::MissingWebIdentityCredentials + | Error::WebIdentityExpiration(_) + | Error::Signing { .. } => litellm_http::Error::signature(error), + }) } } @@ -183,10 +189,28 @@ mod tests { ) .unwrap_err(); - assert_eq!( - error, - litellm_http::Error::ComputedHeader("authorization".into()) - ); + assert!(error.is_computed_header()); + assert!(error.to_string().contains("authorization"), "{error}"); + } + + #[test] + fn a_forwarded_header_that_cannot_be_signed_is_the_callers_fault() { + let error = signed( + json( + URL, + &[("x-amz-meta-note".into(), "café".into())], + &json_value!({}), + None, + ) + .unwrap(), + &signer("textract"), + ) + .unwrap_err(); + + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains("x-amz-meta-note"), "{report}"); + assert!(error.is_request()); + assert!(!error.dispatched()); } #[test] @@ -198,7 +222,7 @@ mod tests { assert!(matches!( signed(request, &signer("textract")), - Err(litellm_http::Error::Signature(_)) + Err(error) if !error.dispatched() && !error.is_request() )); } } diff --git a/litellm-rust/crates/auth-azure/Cargo.toml b/litellm-rust/crates/auth-azure/Cargo.toml index 1fd9d39d113..f99972f813b 100644 --- a/litellm-rust/crates/auth-azure/Cargo.toml +++ b/litellm-rust/crates/auth-azure/Cargo.toml @@ -18,5 +18,6 @@ azure_core = "1.0.0" azure_identity = { version = "1.0.0", features = ["tokio"] } [dev-dependencies] +litellm-error.workspace = true rstest.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/auth-azure/src/native.rs b/litellm-rust/crates/auth-azure/src/native.rs index 64752162384..9fda67746d1 100644 --- a/litellm-rust/crates/auth-azure/src/native.rs +++ b/litellm-rust/crates/auth-azure/src/native.rs @@ -12,7 +12,7 @@ use azure_identity::{ }; use sha2::{Digest, Sha256}; -use litellm_auth_types::Error; +use litellm_auth_types::{Error, ErrorDetail}; use litellm_auth_types::{InputSource, ResolvedCredential, SecretValue, Sourced}; use super::credential_provider_cache::{ @@ -133,7 +133,9 @@ impl NativeAzureTokenAcquirer { let token = credential .get_token(&[scope.as_str()], None) .await - .map_err(|error| Error::CredentialAcquisition(error.to_string().into()))?; + .map_err(|error| { + Error::CredentialAcquisition(ErrorDetail::failed("Azure token request", error)) + })?; let expires_on = u64::try_from(token.expires_on.unix_timestamp()) .ok() .map(|seconds| UNIX_EPOCH + Duration::from_secs(seconds)); @@ -250,11 +252,8 @@ fn validate_authority(request: &NativeAzureRequest) -> Result<(), Error> { let Some(authority) = authority else { return Ok(()); }; - let url = url::Url::parse(authority.value()).map_err(|_| { - Error::InvalidConfiguration( - "Azure authority must be an HTTPS origin without credentials, query, or fragment" - .into(), - ) + let url = url::Url::parse(authority.value()).map_err(|error| { + Error::InvalidConfiguration(ErrorDetail::failed("Azure authority parsing", error)) })?; if url.scheme() != "https" || url.host_str().is_none() @@ -264,10 +263,10 @@ fn validate_authority(request: &NativeAzureRequest) -> Result<(), Error> { || url.fragment().is_some() || !matches!(url.path(), "" | "/") { - return Err(Error::InvalidConfiguration( - "Azure authority must be an HTTPS origin without credentials, query, or fragment" - .into(), - )); + return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType { + field: "Azure authority".into(), + expected: "an HTTPS origin without credentials, query, or fragment", + })); } Ok(()) } @@ -377,7 +376,7 @@ fn trusted_source(sources: &[InputSource]) -> InputSource { fn mixed_sources() -> Result { Err(Error::InvalidConfiguration( - "request-controlled Azure auth inputs cannot be combined with host credentials".into(), + ErrorDetail::MixedInputSources { provider: "Azure" }, )) } @@ -444,7 +443,7 @@ fn build_credential( .map(|credential| credential as Arc), } .map_err(|error| { - Error::InvalidConfiguration(litellm_auth_types::ErrorDetail::failed( + Error::InvalidConfiguration(ErrorDetail::failed( "Azure credential initialization", error, )) @@ -662,12 +661,11 @@ mod tests { )) .unwrap_err(); - assert_eq!( - error, - litellm_auth_types::Error::InvalidConfiguration( + assert!( + litellm_error::Report::new(&error).to_string().contains( "request-controlled Azure auth inputs cannot be combined with host credentials" - .into() - ) + ), + "{error:?}" ); } @@ -695,12 +693,11 @@ mod tests { authority, )) .unwrap_err(); - assert_eq!( - error, - litellm_auth_types::Error::InvalidConfiguration( + assert!( + litellm_error::Report::new(&error).to_string().contains( "Azure authority must be an HTTPS origin without credentials, query, or fragment" - .into() - ) + ), + "{error:?}" ); } } diff --git a/litellm-rust/crates/auth-azure/src/resolve.rs b/litellm-rust/crates/auth-azure/src/resolve.rs index 2142e22db50..25b71741342 100644 --- a/litellm-rust/crates/auth-azure/src/resolve.rs +++ b/litellm-rust/crates/auth-azure/src/resolve.rs @@ -1,8 +1,8 @@ -use litellm_auth_types::Error; use litellm_auth_types::{ CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential, SecretValue, Sourced, TokenProviderHandle, }; +use litellm_auth_types::{Error, ErrorDetail}; use std::future::Future; use std::pin::Pin; use std::sync::Arc; @@ -91,9 +91,9 @@ impl AzureAuthService { AzureCredentialPlan::Caller(caller) => { let credential = caller.acquire().await?; if credential.secret().expose().is_empty() { - return Err(Error::EmptyCallerCredential( - "Azure AD token provider returned an empty token", - )); + return Err(Error::EmptyCallerCredential { + caller: "Azure AD token provider", + }); } Ok(Some(Sourced::new(credential, InputSource::Deployment))) } @@ -106,11 +106,9 @@ impl AzureAuthService { } => { let assertion = resolve_reference(inputs, env_lookup, reference.value()) .await? - .ok_or_else(|| { - Error::CredentialAcquisition( - "Azure OIDC reference did not resolve to a value".into(), - ) - })?; + .ok_or(Error::CredentialAcquisition(ErrorDetail::Unresolved { + subject: "Azure OIDC reference", + }))?; let request = ValidatedAzureRequest::new(NativeAzureRequest::ClientAssertion { tenant_id, client_id, @@ -142,7 +140,13 @@ impl AzureAuthService { Err(error) => failures.push(error), } } - Err(Error::CredentialChain(failures)) + match failures.pop() { + Some(last) => Err(Error::CredentialChain { + last: Box::new(last), + earlier: failures, + }), + None => Ok(None), + } } AzureCredentialPlan::Missing => Ok(None), } @@ -173,7 +177,12 @@ pub(crate) fn select_auth_plan( .map(|selector| Sourced::new(selector, value.source())) }) .transpose() - .map_err(|_| Error::InvalidConfiguration("invalid Azure credential selector".into()))?; + .map_err(|error| { + Error::InvalidConfiguration(ErrorDetail::failed( + "Azure credential selector parsing", + error, + )) + })?; let federated_token_file = configured_string( &inputs.federated_token_file, AZURE_FEDERATED_TOKEN_FILE_ENV, @@ -263,9 +272,12 @@ fn select_native_plan( let selection_source = selected.source(); match selected.into_value() { - AzureCredentialType::ClientSecretCredential => Err(Error::InvalidConfiguration( - "ClientSecretCredential requires tenant_id, client_id, and client_secret".into(), - )), + AzureCredentialType::ClientSecretCredential => { + Err(Error::InvalidConfiguration(ErrorDetail::MissingSettings { + credential: "ClientSecretCredential", + settings: "tenant_id, client_id, and client_secret", + })) + } AzureCredentialType::WorkloadIdentityCredential => { Ok(AzureCredentialPlan::Native(ValidatedAzureRequest::new( workload_request(tenant_id, client_id, federated_token_file, scope, authority)?, @@ -349,22 +361,21 @@ fn workload_request( authority: Option>, ) -> Result { Ok(NativeAzureRequest::WorkloadIdentity { - tenant_id: tenant_id.ok_or_else(|| { - Error::InvalidConfiguration("WorkloadIdentityCredential requires tenant_id".into()) - })?, - client_id: client_id.ok_or_else(|| { - Error::InvalidConfiguration("WorkloadIdentityCredential requires client_id".into()) - })?, - token_file_path: token_file_path.ok_or_else(|| { - Error::InvalidConfiguration( - "WorkloadIdentityCredential requires azure_federated_token_file".into(), - ) - })?, + tenant_id: tenant_id.ok_or(workload_requires("tenant_id"))?, + client_id: client_id.ok_or(workload_requires("client_id"))?, + token_file_path: token_file_path.ok_or(workload_requires("azure_federated_token_file"))?, scope, authority, }) } +fn workload_requires(settings: &'static str) -> Error { + Error::InvalidConfiguration(ErrorDetail::MissingSettings { + credential: "WorkloadIdentityCredential", + settings, + }) +} + fn configured_string( configured: &ConfigValue, environment_name: &str, @@ -410,11 +421,10 @@ async fn resolve_reference( .map_or(CredentialLookup::Missing, CredentialLookup::Found), CredentialRef::None => return Ok(None), CredentialRef::File(_) | CredentialRef::Request(_) | CredentialRef::Host(_) => { - let resolver = inputs.credential_resolver.as_ref().ok_or_else(|| { - Error::InvalidConfiguration( - "credential reference requires a host credential resolver".into(), - ) - })?; + let resolver = inputs + .credential_resolver + .as_ref() + .ok_or(Error::InvalidConfiguration(ErrorDetail::MissingResolver))?; resolver.resolve(reference).await? } }; @@ -433,7 +443,7 @@ fn oidc_reference( let value = token.value().expose(); if token.source() == InputSource::Request && value.starts_with("oidc/") { return Err(Error::InvalidConfiguration( - "request-controlled Azure credential references are not allowed".into(), + ErrorDetail::RequestControlledReference { provider: "Azure" }, )); } if let Some(name) = value.strip_prefix("oidc/env/") { @@ -457,19 +467,17 @@ fn oidc_reference( } if value.starts_with("oidc/") { return Err(Error::InvalidConfiguration( - "unsupported OIDC reference".into(), + ErrorDetail::UnsupportedReference, )); } Ok(None) } -fn non_empty_reference(value: &str, kind: &str) -> Result { +fn non_empty_reference(value: &str, kind: &'static str) -> Result { if value.is_empty() { - return Err(Error::InvalidConfiguration( - litellm_auth_types::ErrorDetail::Empty { - subject: kind.into(), - }, - )); + return Err(Error::InvalidConfiguration(ErrorDetail::Empty { + subject: kind, + })); } Ok(value.to_string()) } @@ -518,7 +526,9 @@ mod tests { expires_on: None, }) } else { - Err(Error::CredentialAcquisition(kind.into())) + Err(Error::CredentialAcquisition( + litellm_auth_types::ErrorDetail::Unresolved { subject: kind }, + )) } }) } @@ -624,7 +634,11 @@ mod tests { ))) .expect_err("unsupported backend must fail validation"); - assert!(error.to_string().contains("unsupported OIDC reference")); + assert!( + litellm_error::Report::new(&error) + .to_string() + .contains("unsupported OIDC reference") + ); } #[rstest::rstest] @@ -649,11 +663,11 @@ mod tests { }) .unwrap_err(); - assert_eq!( - error, - Error::InvalidConfiguration( - "request-controlled Azure credential references are not allowed".into() - ) + assert!( + litellm_error::Report::new(&error) + .to_string() + .contains("request-controlled Azure credential references are not allowed"), + "{error:?}" ); } @@ -715,7 +729,16 @@ mod tests { .await .unwrap_err(); - assert!(matches!(error, Error::CredentialChain(errors) if errors.len() == 2)); + assert!(matches!(&error, Error::CredentialChain { earlier, .. } if earlier.len() == 1)); + let report = litellm_error::Report::new(&error).to_string(); + assert!( + report.contains("managed-identity did not resolve"), + "{report}" + ); + assert!( + report.contains("developer-tools did not resolve"), + "{report}" + ); } #[derive(Debug)] @@ -763,7 +786,9 @@ mod tests { assert_eq!( error, - Error::EmptyCallerCredential("Azure AD token provider returned an empty token") + Error::EmptyCallerCredential { + caller: "Azure AD token provider" + } ); } } diff --git a/litellm-rust/crates/auth-gcp/Cargo.toml b/litellm-rust/crates/auth-gcp/Cargo.toml index bb4c022580a..7844b35ecc2 100644 --- a/litellm-rust/crates/auth-gcp/Cargo.toml +++ b/litellm-rust/crates/auth-gcp/Cargo.toml @@ -24,5 +24,6 @@ google-cloud-auth = { workspace = true, optional = true } http = { workspace = true, optional = true } [dev-dependencies] +litellm-error.workspace = true rstest.workspace = true tempfile.workspace = true diff --git a/litellm-rust/crates/auth-gcp/src/auth.rs b/litellm-rust/crates/auth-gcp/src/auth.rs index e09a350c012..8ac1e0f54ba 100644 --- a/litellm-rust/crates/auth-gcp/src/auth.rs +++ b/litellm-rust/crates/auth-gcp/src/auth.rs @@ -73,13 +73,17 @@ impl VertexAuth { .or_else(|| non_empty_env(env_lookup, VERTEXAI_API_KEY_ENV)); let project_id = get_vertex_ai_project(config, env_lookup); - if !has_authorization && static_token.is_none() { - let access = self.get_access_token(config, env_lookup).await?; - return Ok(VertexEnvironment { - headers: apply_credential(headers, &access.token, CredentialPlacement::Bearer)?, - project_id: project_id.unwrap_or(access.project_id), - }); - } + let static_token = match (has_authorization, static_token) { + (true, _) => None, + (false, Some(token)) => Some(token), + (false, None) => { + let access = self.get_access_token(config, env_lookup).await?; + return Ok(VertexEnvironment { + headers: apply_credential(headers, &access.token, CredentialPlacement::Bearer)?, + project_id: project_id.unwrap_or(access.project_id), + }); + } + }; let project_id = match project_id { Some(project_id) => project_id, @@ -90,14 +94,9 @@ impl VertexAuth { .await? } }; - let headers = if has_authorization { - headers - } else { - apply_credential( - headers, - static_token.as_deref().expect("static token was checked"), - CredentialPlacement::Bearer, - )? + let headers = match static_token { + Some(token) => apply_credential(headers, &token, CredentialPlacement::Bearer)?, + None => headers, }; Ok(VertexEnvironment { headers, @@ -234,14 +233,13 @@ enum CredentialCacheKey { Adc, } -const BASE64_CREDENTIALS: &str = "The `vertex_credentials` value is base64-encoded JSON, not the JSON itself. \ - Set it to a readable file path, or to the credentials JSON itself."; - fn load_trusted_service_account(configured: &str) -> Result { if is_base64_encoded_json(configured) { - return Err(Error::InvalidConfiguration(ErrorDetail::Message( - BASE64_CREDENTIALS.into(), - ))); + return Err(Error::InvalidConfiguration( + ErrorDetail::Base64EncodedJson { + field: "vertex_credentials", + }, + )); } if names_a_credentials_file(configured) { CustomServiceAccount::from_file(configured) @@ -299,7 +297,8 @@ mod tests { else { panic!("{configured:?} should not load"); }; - assert!(error.to_string().contains(expected), "{error}"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains(expected), "{report}"); } #[rstest::rstest] @@ -321,7 +320,7 @@ mod tests { .await { Ok(_) => panic!("{configured:?} should not load"), - Err(error) => error.to_string(), + Err(error) => litellm_error::Report::new(&error).to_string(), } } diff --git a/litellm-rust/crates/auth-gcp/src/sdk.rs b/litellm-rust/crates/auth-gcp/src/sdk.rs index 566df291b7d..6cfeee64d20 100644 --- a/litellm-rust/crates/auth-gcp/src/sdk.rs +++ b/litellm-rust/crates/auth-gcp/src/sdk.rs @@ -55,7 +55,7 @@ impl CredentialsProvider for GoogleCredentials { entity_tag: EntityTag::new(), data, }) - .map_err(|_| CredentialsError::from_msg(false, "Google authentication failed")) + .map_err(|error| CredentialsError::from_source(false, error)) } async fn universe_domain(&self) -> Option { @@ -102,5 +102,9 @@ mod tests { ); let error = credentials.headers(Extensions::new()).await.unwrap_err(); assert!(!format!("{error:?}").contains("private")); + assert_eq!( + std::error::Error::source(&error).and_then(|cause| cause.downcast_ref::()), + Some(&Error::InvalidHeader) + ); } } diff --git a/litellm-rust/crates/auth-types/Cargo.toml b/litellm-rust/crates/auth-types/Cargo.toml index b54c04f0229..400f2640221 100644 --- a/litellm-rust/crates/auth-types/Cargo.toml +++ b/litellm-rust/crates/auth-types/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-error.workspace = true serde.workspace = true serde_json.workspace = true subtle.workspace = true diff --git a/litellm-rust/crates/auth-types/src/credential.rs b/litellm-rust/crates/auth-types/src/credential.rs index 94fb8a58467..4f69c6fffc8 100644 --- a/litellm-rust/crates/auth-types/src/credential.rs +++ b/litellm-rust/crates/auth-types/src/credential.rs @@ -86,9 +86,9 @@ impl CredentialPlan { Self::Caller(caller) => { let credential = caller.acquire().await?; if credential.secret().expose().is_empty() { - return Err(Error::EmptyCallerCredential( - "credential caller returned an empty credential", - )); + return Err(Error::EmptyCallerCredential { + caller: "credential caller", + }); } Ok(CredentialPlanResolution::Resolved(credential)) } @@ -104,8 +104,8 @@ mod tests { CredentialLookup, CredentialLookupFuture, CredentialPlan, CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle, }; - use crate::Error; use crate::SecretValue; + use crate::{Error, ErrorDetail}; #[derive(Debug)] struct HostResolver; @@ -150,9 +150,10 @@ mod tests { impl CredentialResolver for FailingResolver { fn resolve<'a>(&'a self, _reference: &'a CredentialRef) -> CredentialLookupFuture<'a> { Box::pin(async { - Err(Error::CredentialAcquisition( - "host credential lookup failed".into(), - )) + Err(Error::CredentialAcquisition(ErrorDetail::failed( + "host credential lookup", + std::io::Error::from(std::io::ErrorKind::PermissionDenied), + ))) }) } } @@ -168,9 +169,14 @@ mod tests { .await .expect_err("acquisition errors cannot become fallback"); - assert_eq!( - error, - Error::CredentialAcquisition("host credential lookup failed".into()) + assert!( + matches!(error, Error::CredentialAcquisition(_)), + "{error:?}" + ); + assert!( + litellm_error::Report::new(&error) + .to_string() + .contains("host credential lookup failed") ); } } diff --git a/litellm-rust/crates/auth-types/src/error.rs b/litellm-rust/crates/auth-types/src/error.rs index ff17dcf3810..a8ba7b64e03 100644 --- a/litellm-rust/crates/auth-types/src/error.rs +++ b/litellm-rust/crates/auth-types/src/error.rs @@ -1,17 +1,23 @@ -use thiserror::Error as ThisError; +use litellm_error::{Report, Source}; -#[derive(Clone, Debug, ThisError, PartialEq, Eq)] +use crate::{ParamKind, ParamSpec}; + +#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] pub enum Error { - #[error("invalid authentication configuration: {0}")] + #[error("invalid authentication configuration")] InvalidConfiguration(#[source] ErrorDetail), - #[error("credential acquisition failed: {0}")] + #[error("credential acquisition failed")] CredentialAcquisition(#[source] ErrorDetail), - #[error("credential caller failed: {0}")] - EmptyCallerCredential(&'static str), - #[error("{0}")] - ProviderAuthentication(String), - #[error("credential acquisition failed: {}", .0.iter().map(ToString::to_string).collect::>().join("; "))] - CredentialChain(Vec), + #[error("{caller} returned an empty credential")] + EmptyCallerCredential { caller: &'static str }, + #[error("provider authentication failed")] + Provider(#[source] Source), + #[error("no credential in the chain succeeded{}", earlier_attempts(earlier))] + CredentialChain { + #[source] + last: Box, + earlier: Vec, + }, #[error( "Missing {provider} API Key - Set `api_key` or the {environment_variable} environment variable" )] @@ -22,7 +28,7 @@ pub enum Error { #[error("Missing {provider} {} - {}", .spec.setting, .spec.guidance())] MissingParam { provider: &'static str, - spec: &'static crate::ParamSpec, + spec: &'static ParamSpec, }, #[error("Missing {provider} API Base - {guidance}")] MissingApiBase { @@ -33,78 +39,282 @@ pub enum Error { InvalidHeader, } +fn earlier_attempts(earlier: &[Error]) -> String { + if earlier.is_empty() { + return String::new(); + } + let attempts = earlier + .iter() + .map(|failure| Report::new(failure).to_string()) + .collect::>() + .join("; "); + format!(" (earlier attempts: {attempts})") +} + +impl Error { + pub fn provider(cause: impl std::error::Error + Send + Sync + 'static) -> Self { + Self::Provider(Source::new(cause)) + } + + pub fn is_missing_credential(&self) -> bool { + match self { + Self::MissingApiKey { .. } => true, + Self::MissingParam { spec, .. } => match spec.kind { + ParamKind::Credential => true, + ParamKind::Setting => false, + }, + Self::InvalidConfiguration(_) + | Self::CredentialAcquisition(_) + | Self::EmptyCallerCredential { .. } + | Self::Provider(_) + | Self::CredentialChain { .. } + | Self::MissingApiBase { .. } + | Self::InvalidHeader => false, + } + } + + pub fn is_request(&self) -> bool { + match self { + Self::InvalidConfiguration(_) + | Self::EmptyCallerCredential { .. } + | Self::InvalidHeader => true, + Self::CredentialAcquisition(_) | Self::Provider(_) | Self::CredentialChain { .. } => { + false + } + Self::MissingParam { spec, .. } => match spec.kind { + ParamKind::Setting => true, + ParamKind::Credential => false, + }, + Self::MissingApiKey { .. } | Self::MissingApiBase { .. } => false, + } + } +} + #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub enum ErrorDetail { - #[error("{0}")] - Message(String), #[error("{field} must be {expected}")] InvalidType { field: String, expected: &'static str, }, #[error("{subject} cannot be empty")] - Empty { subject: String }, + Empty { subject: &'static str }, #[error("credential header {0} already exists")] DuplicateHeader(&'static str), - #[error("{operation} failed: {source}")] + #[error("credential plan is not allowed by the provider auth policy")] + PlanNotAllowed, + #[error("{credential} requires {settings}")] + MissingSettings { + credential: &'static str, + settings: &'static str, + }, + #[error("credentials cannot be sent to a request-controlled {provider} endpoint")] + RequestControlledEndpoint { provider: &'static str }, + #[error("request-controlled {provider} credential references are not allowed")] + RequestControlledReference { provider: &'static str }, + #[error("request-controlled {provider} auth inputs cannot be combined with host credentials")] + MixedInputSources { provider: &'static str }, + #[error("unsupported OIDC reference")] + UnsupportedReference, + #[error("credential reference requires a host credential resolver")] + MissingResolver, + #[error("{subject} did not resolve to a value")] + Unresolved { subject: &'static str }, + #[error( + "The `{field}` value is base64-encoded JSON, not the JSON itself. Set it to a readable file path, or to the credentials JSON itself." + )] + Base64EncodedJson { field: &'static str }, + #[error("{operation} failed")] Failed { operation: &'static str, #[source] - source: ErrorSource, + cause: Source, }, } impl ErrorDetail { pub fn failed( operation: &'static str, - source: impl std::error::Error + Send + Sync + 'static, + cause: impl std::error::Error + Send + Sync + 'static, ) -> Self { Self::Failed { operation, - source: ErrorSource::new(source), + cause: Source::new(cause), } } } -impl From for ErrorDetail { - fn from(message: String) -> Self { - Self::Message(message) +#[cfg(test)] +mod tests { + use std::io; + + use litellm_error::{Report, Source}; + use rstest::rstest; + + use super::{Error, ErrorDetail}; + use crate::{AwsParams, VertexParams}; + + #[rstest] + #[case::api_key( + Error::MissingApiKey { provider: "Example", environment_variable: "EXAMPLE_API_KEY" }, + "Missing Example API Key - Set `api_key` or the EXAMPLE_API_KEY environment variable" + )] + #[case::another_api_key( + Error::MissingApiKey { provider: "Custom", environment_variable: "CUSTOM_KEY" }, + "Missing Custom API Key - Set `api_key` or the CUSTOM_KEY environment variable" + )] + #[case::api_base( + Error::MissingApiBase { provider: "Example", guidance: "Pass api_base" }, + "Missing Example API Base - Pass api_base" + )] + #[case::another_api_base( + Error::MissingApiBase { provider: "Custom", guidance: "Set CUSTOM_ENDPOINT" }, + "Missing Custom API Base - Set CUSTOM_ENDPOINT" + )] + fn missing_inputs_keep_the_python_text(#[case] error: Error, #[case] expected: &str) { + assert_eq!(error.to_string(), expected); + } + + #[rstest] + #[case::configuration( + Error::InvalidConfiguration(ErrorDetail::Empty { subject: "credential selector" }), + "invalid authentication configuration: credential selector cannot be empty" + )] + #[case::acquisition( + Error::CredentialAcquisition(ErrorDetail::Unresolved { subject: "OIDC reference" }), + "credential acquisition failed: OIDC reference did not resolve to a value" + )] + #[case::empty_caller( + Error::EmptyCallerCredential { caller: "token provider" }, + "token provider returned an empty credential" + )] + fn the_report_names_every_layer_in_order(#[case] error: Error, #[case] expected: &str) { + assert!( + Report::new(&error).to_string().contains(expected), + "{}", + Report::new(&error) + ); + } + + #[rstest] + #[case::configuration(Error::InvalidConfiguration(ErrorDetail::failed( + "test credential", + denied() + )))] + #[case::acquisition(Error::CredentialAcquisition(ErrorDetail::failed( + "test credential", + denied() + )))] + fn a_failure_keeps_its_cause_as_the_source_and_renders_it_once(#[case] error: Error) { + let cause = + std::iter::successors(std::error::Error::source(&error), |cause| cause.source()) + .find_map(|cause| cause.downcast_ref::()?.downcast_ref::()) + .expect("the io cause is in the source chain"); + assert_eq!(cause.kind(), io::ErrorKind::PermissionDenied); + let report = Report::new(&error).to_string(); + assert_eq!(report.matches("access denied").count(), 1, "{report}"); + } + + #[rstest] + #[case::configuration(Error::InvalidConfiguration(ErrorDetail::failed( + "test credential", + denied() + )))] + #[case::acquisition(Error::CredentialAcquisition(ErrorDetail::failed( + "test credential", + denied() + )))] + fn a_contextual_failure_leaves_its_cause_out_of_its_own_message(#[case] error: Error) { + assert!(!error.to_string().contains("access denied"), "{error}"); + assert!( + Report::new(&error) + .to_string() + .contains("test credential failed: access denied") + ); + } + + #[test] + fn a_provider_failure_keeps_the_provider_error_as_its_source() { + let error = Error::provider(denied()); + assert_eq!( + litellm_error::find::(&error).map(io::Error::kind), + Some(io::ErrorKind::PermissionDenied) + ); + assert!(!error.to_string().contains("access denied"), "{error}"); + let report = Report::new(&error).to_string(); + assert_eq!(report.matches("access denied").count(), 1, "{report}"); + } + + #[test] + fn a_chain_keeps_its_last_failure_as_the_source_and_renders_every_attempt_once() { + let error = Error::CredentialChain { + last: Box::new(Error::provider(denied())), + earlier: vec![ + Error::CredentialAcquisition(ErrorDetail::Unresolved { + subject: "OIDC reference", + }), + Error::EmptyCallerCredential { + caller: "token provider", + }, + ], + }; + + assert_eq!( + litellm_error::find::(&error).map(io::Error::kind), + Some(io::ErrorKind::PermissionDenied) + ); + let report = Report::new(&error).to_string(); + assert_eq!(report.matches("access denied").count(), 1, "{report}"); + assert_eq!( + report.matches("OIDC reference did not resolve").count(), + 1, + "{report}" + ); + assert_eq!( + report + .matches("token provider returned an empty credential") + .count(), + 1, + "{report}" + ); + } + + fn denied() -> io::Error { + io::Error::new(io::ErrorKind::PermissionDenied, "access denied") + } + + #[rstest] + #[case::api_key(Error::MissingApiKey { provider: "Example", environment_variable: "EXAMPLE_API_KEY" }, true)] + #[case::vertex_credentials(VertexParams::CREDENTIALS.missing("Vertex AI"), true)] + #[case::aws_secret_access_key(AwsParams::SECRET_ACCESS_KEY.missing("AWS"), true)] + #[case::vertex_project(VertexParams::PROJECT.missing("Vertex AI"), false)] + #[case::aws_region(AwsParams::REGION.missing("AWS"), false)] + #[case::api_base(Error::MissingApiBase { provider: "Example", guidance: "Pass api_base" }, false)] + #[case::invalid_configuration(Error::InvalidConfiguration(ErrorDetail::PlanNotAllowed), false)] + fn only_an_absent_credential_is_a_missing_credential( + #[case] error: Error, + #[case] expected: bool, + ) { + assert_eq!(error.is_missing_credential(), expected); + } + + #[rstest] + #[case::vertex_project(VertexParams::PROJECT.missing("Vertex AI"), true)] + #[case::aws_region(AwsParams::REGION.missing("AWS"), true)] + #[case::api_base(Error::MissingApiBase { provider: "Example", guidance: "Pass api_base" }, false)] + #[case::vertex_credentials(VertexParams::CREDENTIALS.missing("Vertex AI"), false)] + #[case::api_key(Error::MissingApiKey { provider: "Example", environment_variable: "EXAMPLE_API_KEY" }, false)] + #[case::invalid_configuration(Error::InvalidConfiguration(ErrorDetail::PlanNotAllowed), true)] + #[case::invalid_header(Error::InvalidHeader, true)] + #[case::identity_provider_outage( + Error::CredentialAcquisition(ErrorDetail::MissingResolver), + false + )] + #[case::provider_sdk_failure(Error::provider(denied()), false)] + fn a_setting_the_caller_can_pass_is_theirs_to_fix_but_credentials_and_api_bases_are_not( + #[case] error: Error, + #[case] expected: bool, + ) { + assert_eq!(error.is_request(), expected); } } - -impl From<&str> for ErrorDetail { - fn from(message: &str) -> Self { - Self::Message(message.into()) - } -} - -#[derive(Clone, Debug)] -pub struct ErrorSource(std::sync::Arc); - -impl std::ops::Deref for ErrorSource { - type Target = dyn std::error::Error + Send + Sync; - - fn deref(&self) -> &Self::Target { - self.0.as_ref() - } -} - -impl std::fmt::Display for ErrorSource { - fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - std::fmt::Display::fmt(&self.0, formatter) - } -} - -impl ErrorSource { - pub fn new(error: impl std::error::Error + Send + Sync + 'static) -> Self { - Self(std::sync::Arc::new(error)) - } -} - -impl PartialEq for ErrorSource { - fn eq(&self, other: &Self) -> bool { - std::sync::Arc::ptr_eq(&self.0, &other.0) - } -} - -impl Eq for ErrorSource {} diff --git a/litellm-rust/crates/auth-types/src/http.rs b/litellm-rust/crates/auth-types/src/http.rs index 2de3bad9ecf..2ad07f7d541 100644 --- a/litellm-rust/crates/auth-types/src/http.rs +++ b/litellm-rust/crates/auth-types/src/http.rs @@ -21,9 +21,9 @@ pub fn apply_credential( placement: CredentialPlacement, ) -> Result, Error> { if credential.trim().is_empty() { - return Err(Error::InvalidConfiguration( - "credential cannot be empty".into(), - )); + return Err(Error::InvalidConfiguration(crate::ErrorDetail::Empty { + subject: "credential", + })); } if headers .iter() @@ -71,6 +71,10 @@ mod tests { ) .expect_err("provider policy must handle existing credentials"); - assert!(error.to_string().contains("already exists")); + assert!( + litellm_error::Report::new(&error) + .to_string() + .contains("already exists") + ); } } diff --git a/litellm-rust/crates/auth-types/src/lib.rs b/litellm-rust/crates/auth-types/src/lib.rs index c875bda7979..99651977867 100644 --- a/litellm-rust/crates/auth-types/src/lib.rs +++ b/litellm-rust/crates/auth-types/src/lib.rs @@ -51,9 +51,9 @@ pub use credential::{ CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan, CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle, }; -pub use error::{Error, ErrorDetail, ErrorSource}; +pub use error::{Error, ErrorDetail}; pub use http::CredentialPlacement; -pub use params::{AwsParams, ParamSpec, VertexParams}; +pub use params::{AwsParams, ParamKind, ParamSpec, VertexParams}; pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy}; pub use secret::SecretValue; pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; diff --git a/litellm-rust/crates/auth-types/src/params/aws.rs b/litellm-rust/crates/auth-types/src/params/aws.rs index 25fe52b2ce1..8075fc15e9f 100644 --- a/litellm-rust/crates/auth-types/src/params/aws.rs +++ b/litellm-rust/crates/auth-types/src/params/aws.rs @@ -4,15 +4,17 @@ use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; use veil::Redact; -use super::ParamSpec; +use super::{ParamKind, ParamSpec}; const fn spec( setting: &'static str, + kind: ParamKind, wire: &'static [&'static str], env: &'static [&'static str], ) -> ParamSpec { ParamSpec { setting, + kind, wire, module_global: None, env, @@ -53,40 +55,67 @@ pub struct AwsParams { impl AwsParams { pub const ACCESS_KEY_ID: ParamSpec = spec( "access key id", + ParamKind::Credential, &["aws_access_key_id"], &["AWS_ACCESS_KEY_ID"], ); pub const SECRET_ACCESS_KEY: ParamSpec = spec( "secret access key", + ParamKind::Credential, &["aws_secret_access_key"], &["AWS_SECRET_ACCESS_KEY"], ); pub const SESSION_TOKEN: ParamSpec = spec( "session token", + ParamKind::Credential, &["aws_session_token"], &["AWS_SESSION_TOKEN"], ); pub const REGION: ParamSpec = spec( "region", + ParamKind::Setting, &["aws_region_name"], &["AWS_REGION_NAME", "AWS_REGION"], ); - pub const SESSION_NAME: ParamSpec = - spec("session name", &["aws_session_name"], &["AWS_SESSION_NAME"]); - pub const PROFILE_NAME: ParamSpec = - spec("profile name", &["aws_profile_name"], &["AWS_PROFILE_NAME"]); - pub const ROLE_NAME: ParamSpec = spec("role name", &["aws_role_name"], &["AWS_ROLE_NAME"]); + pub const SESSION_NAME: ParamSpec = spec( + "session name", + ParamKind::Setting, + &["aws_session_name"], + &["AWS_SESSION_NAME"], + ); + pub const PROFILE_NAME: ParamSpec = spec( + "profile name", + ParamKind::Setting, + &["aws_profile_name"], + &["AWS_PROFILE_NAME"], + ); + pub const ROLE_NAME: ParamSpec = spec( + "role name", + ParamKind::Setting, + &["aws_role_name"], + &["AWS_ROLE_NAME"], + ); pub const WEB_IDENTITY_TOKEN: ParamSpec = spec( "web identity token", + ParamKind::Credential, &["aws_web_identity_token"], &["AWS_WEB_IDENTITY_TOKEN"], ); - pub const STS_ENDPOINT: ParamSpec = - spec("STS endpoint", &["aws_sts_endpoint"], &["AWS_STS_ENDPOINT"]); - pub const EXTERNAL_ID: ParamSpec = - spec("external id", &["aws_external_id"], &["AWS_EXTERNAL_ID"]); + pub const STS_ENDPOINT: ParamSpec = spec( + "STS endpoint", + ParamKind::Setting, + &["aws_sts_endpoint"], + &["AWS_STS_ENDPOINT"], + ); + pub const EXTERNAL_ID: ParamSpec = spec( + "external id", + ParamKind::Setting, + &["aws_external_id"], + &["AWS_EXTERNAL_ID"], + ); pub const BEDROCK_RUNTIME_ENDPOINT: ParamSpec = spec( "Bedrock runtime endpoint", + ParamKind::Setting, &["aws_bedrock_runtime_endpoint"], &["AWS_BEDROCK_RUNTIME_ENDPOINT"], ); diff --git a/litellm-rust/crates/auth-types/src/params/mod.rs b/litellm-rust/crates/auth-types/src/params/mod.rs index 88bc64cc2fc..daad1aff8a8 100644 --- a/litellm-rust/crates/auth-types/src/params/mod.rs +++ b/litellm-rust/crates/auth-types/src/params/mod.rs @@ -10,11 +10,18 @@ pub use vertex::VertexParams; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct ParamSpec { pub setting: &'static str, + pub kind: ParamKind, pub wire: &'static [&'static str], pub module_global: Option<&'static str>, pub env: &'static [&'static str], } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ParamKind { + Credential, + Setting, +} + impl ParamSpec { /// The first non-blank value: each wire name through `params`, then each environment /// name through `env`. diff --git a/litellm-rust/crates/auth-types/src/params/vertex.rs b/litellm-rust/crates/auth-types/src/params/vertex.rs index 06ec3008dfb..2a6b8bc0ca6 100644 --- a/litellm-rust/crates/auth-types/src/params/vertex.rs +++ b/litellm-rust/crates/auth-types/src/params/vertex.rs @@ -2,7 +2,7 @@ use serde::{Deserialize, Deserializer, Serialize, de::Unexpected}; use serde_json::Value; use veil::Redact; -use super::ParamSpec; +use super::{ParamKind, ParamSpec}; /// The Vertex AI fields of Python's `GenericLiteLLMParams`, both the current names and the /// `vertex_ai_*` spellings `VertexBase.safe_get_vertex_ai_*` still read. @@ -27,18 +27,21 @@ pub struct VertexParams { impl VertexParams { pub const CREDENTIALS: ParamSpec = ParamSpec { setting: "credentials", + kind: ParamKind::Credential, wire: &["vertex_credentials", "vertex_ai_credentials"], module_global: None, env: &["VERTEXAI_CREDENTIALS"], }; pub const PROJECT: ParamSpec = ParamSpec { setting: "project", + kind: ParamKind::Setting, wire: &["vertex_project", "vertex_ai_project"], module_global: Some("vertex_project"), env: &["VERTEXAI_PROJECT"], }; pub const LOCATION: ParamSpec = ParamSpec { setting: "location", + kind: ParamKind::Setting, wire: &["vertex_location", "vertex_ai_location"], module_global: Some("vertex_location"), env: &["VERTEXAI_LOCATION", "VERTEX_LOCATION"], diff --git a/litellm-rust/crates/auth-types/src/policy.rs b/litellm-rust/crates/auth-types/src/policy.rs index ee2bf08bede..8fff59e284a 100644 --- a/litellm-rust/crates/auth-types/src/policy.rs +++ b/litellm-rust/crates/auth-types/src/policy.rs @@ -1,5 +1,5 @@ use crate::http::apply_credential; -use crate::{CredentialPlacement, Error, ResolvedCredential}; +use crate::{CredentialPlacement, Error, ErrorDetail, ResolvedCredential}; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum CredentialPlanKind { @@ -31,11 +31,18 @@ pub struct ProviderAuthPolicy { impl ProviderAuthPolicy { pub fn has_existing_credential(&self, headers: &[(String, String)]) -> bool { - headers.iter().any(|(name, _)| { - self.accepted_existing_headers - .iter() - .any(|accepted| name.eq_ignore_ascii_case(accepted)) - }) + self.existing_credential_header(headers).is_some() + } + + fn existing_credential_header(&self, headers: &[(String, String)]) -> Option<&'static str> { + self.accepted_existing_headers + .iter() + .copied() + .find(|accepted| { + headers + .iter() + .any(|(name, _)| name.eq_ignore_ascii_case(accepted)) + }) } pub fn apply( @@ -44,11 +51,11 @@ impl ProviderAuthPolicy { kind: CredentialPlanKind, credential: &ResolvedCredential, ) -> Result, Error> { - if self.has_existing_credential(&headers) { + if let Some(existing) = self.existing_credential_header(&headers) { return match self.existing_header_behavior { ExistingHeaderBehavior::Preserve => Ok(headers), ExistingHeaderBehavior::Reject => Err(Error::InvalidConfiguration( - "credential header already exists".into(), + ErrorDetail::DuplicateHeader(existing), )), }; } @@ -56,11 +63,7 @@ impl ProviderAuthPolicy { .rules .iter() .find(|rule| rule.kind == kind) - .ok_or_else(|| { - Error::InvalidConfiguration( - "credential plan is not allowed by the provider auth policy".into(), - ) - })?; + .ok_or(Error::InvalidConfiguration(ErrorDetail::PlanNotAllowed))?; apply_credential(headers, credential.secret().expose(), rule.placement) } } @@ -108,6 +111,10 @@ mod tests { ) .unwrap_err(); - assert!(error.to_string().contains("not allowed")); + assert!( + litellm_error::Report::new(&error) + .to_string() + .contains("not allowed") + ); } } diff --git a/litellm-rust/crates/auth-types/tests/error.rs b/litellm-rust/crates/auth-types/tests/error.rs deleted file mode 100644 index aaedc3814f0..00000000000 --- a/litellm-rust/crates/auth-types/tests/error.rs +++ /dev/null @@ -1,74 +0,0 @@ -use litellm_auth_types::Error; -use rstest::rstest; - -#[rstest] -#[case::api_key( - Error::MissingApiKey { provider: "Example", environment_variable: "EXAMPLE_API_KEY" }, - "Missing Example API Key - Set `api_key` or the EXAMPLE_API_KEY environment variable" -)] -#[case::another_api_key( - Error::MissingApiKey { provider: "Custom", environment_variable: "CUSTOM_KEY" }, - "Missing Custom API Key - Set `api_key` or the CUSTOM_KEY environment variable" -)] -#[case::api_base( - Error::MissingApiBase { provider: "Example", guidance: "Pass api_base" }, - "Missing Example API Base - Pass api_base" -)] -#[case::another_api_base( - Error::MissingApiBase { provider: "Custom", guidance: "Set CUSTOM_ENDPOINT" }, - "Missing Custom API Base - Set CUSTOM_ENDPOINT" -)] -#[case::configuration( - Error::InvalidConfiguration("credential selector is invalid".into()), - "invalid authentication configuration: credential selector is invalid" -)] -#[case::acquisition( - Error::CredentialAcquisition("token expired".into()), - "credential acquisition failed: token expired" -)] -#[case::caller( - Error::EmptyCallerCredential("empty token"), - "credential caller failed: empty token" -)] -#[case::provider( - Error::ProviderAuthentication("provider rejected credentials".into()), - "provider rejected credentials" -)] -#[case::chain( - Error::CredentialChain(vec![ - Error::CredentialAcquisition("token expired".into()), - Error::EmptyCallerCredential("empty token"), - ]), - "credential acquisition failed: credential acquisition failed: token expired; credential caller failed: empty token" -)] -fn display_preserves_failure_phase_and_caller_context( - #[case] error: Error, - #[case] expected: &str, -) { - assert_eq!(error.to_string(), expected); -} - -#[rstest] -#[case::configuration(true)] -#[case::acquisition(false)] -fn contextual_failures_keep_the_original_source(#[case] configuration: bool) { - use litellm_auth_types::ErrorDetail; - - let detail = ErrorDetail::failed( - "test credential", - std::io::Error::from(std::io::ErrorKind::PermissionDenied), - ); - let error = if configuration { - Error::InvalidConfiguration(detail) - } else { - Error::CredentialAcquisition(detail) - }; - let source = std::iter::successors(Some(&error as &dyn std::error::Error), |error| { - error.source() - }) - .find_map(|error| error.downcast_ref::()) - .expect("the original credential error remains available"); - assert_eq!(source.kind(), std::io::ErrorKind::PermissionDenied); - assert!(error.to_string().contains("test credential failed:")); - assert!(error.to_string().ends_with(&source.to_string())); -} diff --git a/litellm-rust/crates/auth-types/tests/params.rs b/litellm-rust/crates/auth-types/tests/params.rs index 9d488b170e1..38890edb39d 100644 --- a/litellm-rust/crates/auth-types/tests/params.rs +++ b/litellm-rust/crates/auth-types/tests/params.rs @@ -1,8 +1,9 @@ -use litellm_auth_types::{AwsParams, ParamSpec}; +use litellm_auth_types::{AwsParams, ParamKind, ParamSpec}; use rstest::rstest; const LOCATION: ParamSpec = ParamSpec { setting: "location", + kind: ParamKind::Setting, wire: &["vertex_location", "vertex_ai_location"], module_global: Some("vertex_location"), env: &["VERTEXAI_LOCATION", "VERTEX_LOCATION"], diff --git a/litellm-rust/crates/auth-types/tests/token.rs b/litellm-rust/crates/auth-types/tests/token.rs index d7b23531596..7f7027942b0 100644 --- a/litellm-rust/crates/auth-types/tests/token.rs +++ b/litellm-rust/crates/auth-types/tests/token.rs @@ -1,5 +1,4 @@ use std::{ - error::Error as StdError, future::{Future, poll_fn}, sync::{ Arc, @@ -53,7 +52,7 @@ async fn callbacks_acquire_fresh_credentials_on_demand(#[case] access_token: boo #[rstest] #[tokio::test] -async fn callback_errors_preserve_the_original_source() { +async fn callback_errors_keep_their_typed_cause() { let provider = TokenProviderHandle::from_callback(|| async { Err(Error::CredentialAcquisition(ErrorDetail::failed( "caller credential", @@ -62,13 +61,15 @@ async fn callback_errors_preserve_the_original_source() { }); let error = provider.acquire().await.unwrap_err(); - assert!(matches!(error, Error::CredentialAcquisition(_))); - let source = std::iter::successors(Some(&error as &(dyn StdError + 'static)), |error| { - (*error).source() - }) - .find_map(|error| error.downcast_ref::()) - .unwrap(); - assert_eq!(source.kind(), std::io::ErrorKind::PermissionDenied); + let Error::CredentialAcquisition(ErrorDetail::Failed { cause, .. }) = &error else { + panic!("the callback's error passes through unchanged: {error:?}"); + }; + assert_eq!( + cause + .downcast_ref::() + .map(std::io::Error::kind), + Some(std::io::ErrorKind::PermissionDenied) + ); } struct Release(Arc); diff --git a/litellm-rust/crates/cache-azure-blob/Cargo.toml b/litellm-rust/crates/cache-azure-blob/Cargo.toml index c28cb90d84a..202300c7ab4 100644 --- a/litellm-rust/crates/cache-azure-blob/Cargo.toml +++ b/litellm-rust/crates/cache-azure-blob/Cargo.toml @@ -16,6 +16,7 @@ azure_core = "1.1.0" azure_storage_blob = "1.1.0" futures-util.workspace = true reqwest.workspace = true +thiserror.workspace = true tokio.workspace = true url.workspace = true diff --git a/litellm-rust/crates/cache-azure-blob/src/cache.rs b/litellm-rust/crates/cache-azure-blob/src/cache.rs index c5c1fdd8ab9..e3b013eba4f 100644 --- a/litellm-rust/crates/cache-azure-blob/src/cache.rs +++ b/litellm-rust/crates/cache-azure-blob/src/cache.rs @@ -18,6 +18,10 @@ use url::Url; use crate::{credential::AzureBlobCredential, transport::ReqwestTransport}; +#[derive(Debug, thiserror::Error)] +#[error("account URL cannot hold a container path")] +struct NotABaseUrl; + pub struct AzureBlobCache { container: BlobContainerClient, codec: C, @@ -57,12 +61,12 @@ impl AzureBlobCache { codec: C, runtime: Handle, ) -> Result { - let parsed = Url::parse(account_url).map_err(|_| Error::Unavailable)?; + let parsed = Url::parse(account_url).map_err(Error::unavailable)?; let account_url = parsed.as_str().trim_end_matches('/').to_string(); let container_url = { let mut url = parsed; url.path_segments_mut() - .map_err(|()| Error::Unavailable)? + .map_err(|()| Error::unavailable(NotABaseUrl))? .pop_if_empty() .push(container); url @@ -75,7 +79,7 @@ impl AzureBlobCache { ..BlobContainerClientOptions::default() }), ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; let cache = Self { container: client, codec, @@ -101,7 +105,7 @@ impl AzureBlobCache { Err(error) if is_storage_error(&error, StorageErrorCode::ContainerAlreadyExists) => { Ok(()) } - Err(_) => Err(Error::Unavailable), + Err(error) => Err(Error::unavailable(error)), } } @@ -116,7 +120,7 @@ impl AzureBlobCache { { Ok(_) => Ok(()), Err(error) if !overwrite && is_already_present(&error) => Ok(()), - Err(_) => Err(Error::Unavailable), + Err(error) => Err(Error::unavailable(error)), } } @@ -126,13 +130,9 @@ impl AzureBlobCache { Err(error) if is_storage_error(&error, StorageErrorCode::BlobNotFound) => { return Ok(None); } - Err(_) => return Err(Error::Unavailable), + Err(error) => return Err(Error::unavailable(error)), }; - let bytes = response - .body - .collect() - .await - .map_err(|_| Error::Unavailable)?; + let bytes = response.body.collect().await.map_err(Error::unavailable)?; self.codec.decode(&bytes).map(Some) } @@ -140,16 +140,16 @@ impl AzureBlobCache { let mut pages = self .container .list_blobs(None) - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? .into_pages(); - while let Some(page) = pages.try_next().await.map_err(|_| Error::Unavailable)? { - let page = page.into_model().map_err(|_| Error::Unavailable)?; + while let Some(page) = pages.try_next().await.map_err(Error::unavailable)? { + let page = page.into_model().map_err(Error::unavailable)?; for name in page.blob_items.into_iter().filter_map(|item| item.name) { self.container .blob_client(&name) .delete(None) .await - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } } Ok(()) diff --git a/litellm-rust/crates/cache-azure-blob/src/credential.rs b/litellm-rust/crates/cache-azure-blob/src/credential.rs index d1a3d0e44ec..f58a599f8e3 100644 --- a/litellm-rust/crates/cache-azure-blob/src/credential.rs +++ b/litellm-rust/crates/cache-azure-blob/src/credential.rs @@ -17,6 +17,10 @@ const LLM_TOKEN_ENV: &str = "AZURE_AD_TOKEN"; type EnvLookup = Arc Option + Send + Sync>; +#[derive(Debug, thiserror::Error)] +#[error("no Azure credential is available for blob storage")] +struct NoCredential; + pub struct AzureBlobCredential { service: AzureAuthService, env_lookup: EnvLookup, @@ -62,15 +66,8 @@ impl TokenCredential for AzureBlobCredential { &lookup, ) .await - .map_err(|error| { - azure_core::Error::with_message(ErrorKind::Credential, error.to_string()) - })? - .ok_or_else(|| { - azure_core::Error::with_message( - ErrorKind::Credential, - "no Azure credential is available for blob storage", - ) - })?; + .map_err(|error| azure_core::Error::new(ErrorKind::Credential, error))? + .ok_or_else(|| azure_core::Error::new(ErrorKind::Credential, NoCredential))?; let (token, expires_on) = match credential.into_value() { ResolvedCredential::AccessToken { token, expires_on } => (token, expires_on), ResolvedCredential::Static(token) => (token, None), diff --git a/litellm-rust/crates/cache-azure-blob/src/transport.rs b/litellm-rust/crates/cache-azure-blob/src/transport.rs index 3914b92365c..62feee8fae0 100644 --- a/litellm-rust/crates/cache-azure-blob/src/transport.rs +++ b/litellm-rust/crates/cache-azure-blob/src/transport.rs @@ -24,6 +24,7 @@ impl HttpClient for ReqwestTransport { Body::SeekableStream(stream) => outgoing.body(reqwest::Body::wrap_stream(stream)), }; let response = outgoing.send().await.map_err(|error| { + let error = litellm_http::Error::from(error); let kind = if error.is_connect() { ErrorKind::Connection } else { @@ -41,9 +42,9 @@ impl HttpClient for ReqwestTransport { ); } } - let body = response - .bytes_stream() - .map_err(|error| azure_core::Error::new(ErrorKind::Io, error)); + let body = response.bytes_stream().map_err(|error| { + azure_core::Error::new(ErrorKind::Io, litellm_http::Error::from(error)) + }); Ok(AsyncRawResponse::new(status, headers, Box::pin(body))) } } diff --git a/litellm-rust/crates/cache-azure-blob/tests/cache.rs b/litellm-rust/crates/cache-azure-blob/tests/cache.rs index b9f50b824b3..3b014f53e82 100644 --- a/litellm-rust/crates/cache-azure-blob/tests/cache.rs +++ b/litellm-rust/crates/cache-azure-blob/tests/cache.rs @@ -128,7 +128,7 @@ fn connect_surfaces_service_failures() { ResponseCacheCodec, runtime.handle().clone(), )); - assert!(matches!(result, Err(Error::Unavailable))); + assert!(matches!(result, Err(Error::Unavailable(_)))); } #[rstest] @@ -263,7 +263,7 @@ fn malformed_blobs_are_invalid_entries(fixture: Fixture, #[case] key: &str, #[ca fixture.service.seed_blob(key, bytes); assert!(matches!( fixture.cache.get_cache(key, &no_ttl()), - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(_)) )); } @@ -378,15 +378,15 @@ fn service_failures_map_to_unavailable(fixture: Fixture) { fixture.service.set_failing(true); assert!(matches!( fixture.cache.get_cache("key", &no_ttl()), - Err(Error::Unavailable) + Err(Error::Unavailable(_)) )); assert!(matches!( fixture.cache.set_cache("key", entry(json!(1)), &no_ttl()), - Err(Error::Unavailable) + Err(Error::Unavailable(_)) )); assert!(matches!( fixture.cache.flush_cache(), - Err(Error::Unavailable) + Err(Error::Unavailable(_)) )); assert!(matches!( fixture.runtime.block_on( @@ -394,7 +394,7 @@ fn service_failures_map_to_unavailable(fixture: Fixture) { .cache .async_set_cache_pipeline(vec![("k".into(), entry(json!(1)))], no_ttl()) ), - Err(Error::Unavailable) + Err(Error::Unavailable(_)) )); } diff --git a/litellm-rust/crates/cache-disk/Cargo.toml b/litellm-rust/crates/cache-disk/Cargo.toml index 825b23899e7..b2f7004910f 100644 --- a/litellm-rust/crates/cache-disk/Cargo.toml +++ b/litellm-rust/crates/cache-disk/Cargo.toml @@ -12,6 +12,7 @@ rand.workspace = true rusqlite = { version = "0.39", features = ["bundled"] } serde-pickle = "1.2" serde_json.workspace = true +thiserror.workspace = true tokio.workspace = true [dev-dependencies] diff --git a/litellm-rust/crates/cache-disk/src/cache.rs b/litellm-rust/crates/cache-disk/src/cache.rs index b8a2c6ed1e4..880a1f86e71 100644 --- a/litellm-rust/crates/cache-disk/src/cache.rs +++ b/litellm-rust/crates/cache-disk/src/cache.rs @@ -64,7 +64,7 @@ impl DiskCache { { tokio::task::spawn_blocking(move || operation(&store)) .await - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? } } @@ -162,7 +162,7 @@ impl BatchCache for DiskCache Ok(BatchEntry::Hit(value)), Ok(None) => Ok(BatchEntry::Miss), - Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid), + Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid), Err(error) => Err(error), }) .collect() @@ -186,7 +186,7 @@ impl BatchCache for DiskCache match self.decode_stored(value) { Ok(Some(value)) => Ok(BatchEntry::Hit(value)), Ok(None) => Ok(BatchEntry::Miss), - Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid), + Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid), Err(error) => Err(error), }, }) @@ -272,7 +272,7 @@ fn increment( Ok((stored, ttl.map(|ttl| unix_now() + ttl.as_secs_f64()))) }; store.update(key, unix_now(), &mut apply)?; - result.ok_or(Error::InvalidEntry) + result.ok_or(Error::InvalidEntry(None)) } fn unix_now() -> f64 { diff --git a/litellm-rust/crates/cache-disk/src/python/value.rs b/litellm-rust/crates/cache-disk/src/python/value.rs index 645eba7b757..6e7c2ee0473 100644 --- a/litellm-rust/crates/cache-disk/src/python/value.rs +++ b/litellm-rust/crates/cache-disk/src/python/value.rs @@ -3,8 +3,8 @@ use py_literal::Value; use serde_json::{Map, Number}; pub(crate) fn from_pickle(bytes: &[u8]) -> Result { - let value = serde_pickle::value_from_slice(bytes, Default::default()) - .map_err(|_| Error::InvalidEntry)?; + let value = + serde_pickle::value_from_slice(bytes, Default::default()).map_err(Error::invalid_entry)?; from_pickle_value(value) } @@ -70,7 +70,7 @@ fn from_pickle_hashable(value: serde_pickle::HashableValue) -> Result Result { - value.parse().map_err(|_| Error::InvalidEntry) + value.parse().map_err(Error::invalid_entry) } pub(crate) fn from_json(value: serde_json::Value) -> Value { @@ -101,7 +101,7 @@ pub(crate) fn from_json(value: serde_json::Value) -> Value { pub(crate) fn from_json_text(value: &str) -> Result { serde_json::from_str(value) .map(from_json) - .map_err(|_| Error::InvalidEntry) + .map_err(Error::invalid_entry) } pub(crate) fn is_truthy(value: &Value) -> bool { @@ -131,7 +131,7 @@ pub(crate) fn to_f64(value: &Value) -> Option { } pub(crate) fn to_json(value: &Value) -> Result, Error> { - serde_json::to_vec(&to_json_value(value)?).map_err(|_| Error::InvalidEntry) + serde_json::to_vec(&to_json_value(value)?).map_err(Error::invalid_entry) } fn to_json_value(value: &Value) -> Result { @@ -142,12 +142,12 @@ fn to_json_value(value: &Value) -> Result { value .to_string() .parse::() - .map_err(|_| Error::InvalidEntry)?, + .map_err(Error::invalid_entry)?, ), Value::Float(value) => { - serde_json::Value::Number(Number::from_f64(*value).ok_or(Error::InvalidEntry)?) + serde_json::Value::Number(Number::from_f64(*value).ok_or(Error::InvalidEntry(None))?) } - Value::Complex(_) | Value::Bytes(_) => return Err(Error::InvalidEntry), + Value::Complex(_) | Value::Bytes(_) => return Err(Error::InvalidEntry(None)), Value::String(value) => serde_json::Value::String(value.clone()), Value::Tuple(values) | Value::List(values) | Value::Set(values) => { serde_json::Value::Array( @@ -162,7 +162,7 @@ fn to_json_value(value: &Value) -> Result { .iter() .map(|(key, value)| { let Value::String(key) = key else { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); }; Ok((key.clone(), to_json_value(value)?)) }) diff --git a/litellm-rust/crates/cache-disk/src/sqlite.rs b/litellm-rust/crates/cache-disk/src/sqlite.rs index 24c2be2e1c5..14cbf18b2ff 100644 --- a/litellm-rust/crates/cache-disk/src/sqlite.rs +++ b/litellm-rust/crates/cache-disk/src/sqlite.rs @@ -48,13 +48,13 @@ struct Row { impl DiskcacheSqliteStore { pub fn open(directory: impl AsRef) -> Result { let directory = directory.as_ref().to_path_buf(); - fs::create_dir_all(&directory).map_err(|_| Error::Unavailable)?; - let directory = std::path::absolute(&directory).map_err(|_| Error::Unavailable)?; + fs::create_dir_all(&directory).map_err(Error::unavailable)?; + let directory = std::path::absolute(&directory).map_err(Error::unavailable)?; let database = directory.join("cache.db"); - let connection = Connection::open(database).map_err(|_| Error::Unavailable)?; + let connection = Connection::open(database).map_err(Error::unavailable)?; connection .busy_timeout(std::time::Duration::from_secs(60)) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; let mut settings = read_settings(&connection)?; for (key, value) in default_settings() { @@ -74,7 +74,7 @@ impl DiskcacheSqliteStore { value )", ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; for (key, value) in &settings { if !matches!(key.as_str(), "count" | "size" | "hits" | "misses") { connection @@ -82,7 +82,7 @@ impl DiskcacheSqliteStore { "INSERT OR REPLACE INTO Settings VALUES (?, ?)", params![key, value], ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } } for (key, value) in [ @@ -96,7 +96,7 @@ impl DiskcacheSqliteStore { "INSERT OR IGNORE INTO Settings VALUES (?, ?)", params![key, value], ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } connection .execute_batch( @@ -117,7 +117,7 @@ impl DiskcacheSqliteStore { CREATE UNIQUE INDEX IF NOT EXISTS Cache_key_raw ON Cache(key, raw); CREATE INDEX IF NOT EXISTS Cache_expire_time ON Cache(expire_time);", ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; let eviction_policy = setting_string(&settings, "eviction_policy") .unwrap_or_else(|| "least-recently-stored".to_string()); @@ -128,23 +128,27 @@ impl DiskcacheSqliteStore { .execute_batch( "CREATE INDEX IF NOT EXISTS Cache_store_time ON Cache(store_time)", ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } "least-recently-used" => { connection .execute_batch( "CREATE INDEX IF NOT EXISTS Cache_access_time ON Cache(access_time)", ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } "least-frequently-used" => { connection .execute_batch( "CREATE INDEX IF NOT EXISTS Cache_access_count ON Cache(access_count)", ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; + } + _ => { + return Err(Error::unavailable(SettingFault::EvictionPolicy( + eviction_policy, + ))); } - _ => return Err(Error::Unavailable), } connection .execute_batch( @@ -170,12 +174,12 @@ impl DiskcacheSqliteStore { UPDATE Settings SET value = value - OLD.size WHERE key = \"size\"; END;", ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; let min_file_size = setting_i64(&settings, "disk_min_file_size") .unwrap_or(DEFAULT_DISK_MIN_FILE_SIZE) .try_into() - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; let size_limit = setting_i64(&settings, "size_limit").unwrap_or(DEFAULT_SIZE_LIMIT); let cull_limit = setting_i64(&settings, "cull_limit").unwrap_or(DEFAULT_CULL_LIMIT); let statistics = setting_i64(&settings, "statistics").unwrap_or_default() != 0; @@ -207,7 +211,7 @@ impl DiskcacheSqliteStore { |row| row.get::<_, Option>(0), ) .optional() - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? .flatten() { cleanup.push(old_filename); @@ -221,7 +225,7 @@ impl DiskcacheSqliteStore { |row| row.get::<_, i64>(0), ) .optional() - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; if let Some(rowid) = rowid { connection .execute( @@ -230,7 +234,7 @@ impl DiskcacheSqliteStore { WHERE rowid = ?", params![now, expire_time, now, size, mode, filename, value, rowid], ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } else { connection .execute( @@ -240,7 +244,7 @@ impl DiskcacheSqliteStore { ) VALUES (?, 1, ?, ?, ?, 0, NULL, ?, ?, ?, ?)", params![key, now, expire_time, now, size, mode, filename, value], ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } cleanup.extend(self.cull(connection, now)?); Ok(cleanup) @@ -257,13 +261,13 @@ impl DiskcacheSqliteStore { WHERE expire_time IS NOT NULL AND expire_time < ? ORDER BY expire_time LIMIT ?", ) - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? .query_map(params![now, self.cull_limit], |row| { Ok((row.get::<_, i64>(0)?, row.get::<_, Option>(1)?)) }) - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? .collect::, _>>() - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; for (_, filename) in &expired { if let Some(filename) = filename { cleanup.push(filename.clone()); @@ -272,7 +276,7 @@ impl DiskcacheSqliteStore { for (rowid, _) in &expired { connection .execute("DELETE FROM Cache WHERE rowid = ?", params![rowid]) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } let remaining = self.cull_limit - i64::try_from(expired.len()).unwrap_or(self.cull_limit); if remaining <= 0 || self.volume(connection)? < self.size_limit { @@ -283,19 +287,23 @@ impl DiskcacheSqliteStore { "least-recently-stored" => "store_time", "least-recently-used" => "access_time", "least-frequently-used" => "access_count", - _ => return Err(Error::Unavailable), + policy => { + return Err(Error::unavailable(SettingFault::EvictionPolicy( + policy.to_owned(), + ))); + } }; let rows = connection .prepare(&format!( "SELECT rowid, filename FROM Cache ORDER BY {order} LIMIT ?" )) - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? .query_map(params![remaining], |row| { Ok((row.get::<_, i64>(0)?, row.get::<_, Option>(1)?)) }) - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? .collect::, _>>() - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; for (_, filename) in &rows { if let Some(filename) = filename { cleanup.push(filename.clone()); @@ -304,7 +312,7 @@ impl DiskcacheSqliteStore { for (rowid, _) in rows { connection .execute("DELETE FROM Cache WHERE rowid = ?", params![rowid]) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } Ok(cleanup) } @@ -312,15 +320,15 @@ impl DiskcacheSqliteStore { fn volume(&self, connection: &Connection) -> Result { let page_count: i64 = connection .query_row("PRAGMA page_count", [], |row| row.get(0)) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; let page_size: i64 = connection .query_row("PRAGMA page_size", [], |row| row.get(0)) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; let size: i64 = connection .query_row("SELECT value FROM Settings WHERE key = 'size'", [], |row| { row.get(0) }) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; Ok(page_count.saturating_mul(page_size).saturating_add(size)) } } @@ -331,13 +339,13 @@ impl DiskStore for DiskcacheSqliteStore { } fn get(&self, key: &str, now: f64) -> Result, Error> { - let connection = self.connection.lock().map_err(|_| Error::Unavailable)?; + let connection = self.connection.lock()?; let select = "SELECT rowid, expire_time, mode, filename, value FROM Cache WHERE key = ? AND raw = 1 AND (expire_time IS NULL OR expire_time > ?)"; let row = connection .query_row(select, params![key, now], row_from_query) .optional() - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; if !self.statistics && !has_get_update(&self.eviction_policy) { return row .map(|row| fetch_row(&self.directory, row)) @@ -348,7 +356,7 @@ impl DiskStore for DiskcacheSqliteStore { let row = connection .query_row(select, params![key, now], row_from_query) .optional() - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; let Some(row) = row else { if self.statistics { connection @@ -356,7 +364,7 @@ impl DiskStore for DiskcacheSqliteStore { "UPDATE Settings SET value = value + 1 WHERE key = 'misses'", [], ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } return Ok(None); }; @@ -369,14 +377,14 @@ impl DiskStore for DiskcacheSqliteStore { "UPDATE Settings SET value = value + 1 WHERE key = 'hits'", [], ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } else if !hit && self.statistics { connection .execute( "UPDATE Settings SET value = value + 1 WHERE key = 'misses'", [], ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } if has_get_update(&self.eviction_policy) && hit { let update = match self.eviction_policy.as_str() { @@ -384,16 +392,20 @@ impl DiskStore for DiskcacheSqliteStore { "least-frequently-used" => { "UPDATE Cache SET access_count = access_count + 1 WHERE rowid = ?" } - _ => return Err(Error::Unavailable), + policy => { + return Err(Error::unavailable(SettingFault::EvictionPolicy( + policy.to_owned(), + ))); + } }; if self.eviction_policy == "least-recently-used" { connection .execute(update, params![now, rowid]) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } else { connection .execute(update, params![rowid]) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } } value @@ -409,7 +421,7 @@ impl DiskStore for DiskcacheSqliteStore { ) -> Result<(), Error> { let columns = store_value(&self.directory, self.min_file_size, value)?; let new_filename = columns.filename.clone(); - let connection = self.connection.lock().map_err(|_| Error::Unavailable)?; + let connection = self.connection.lock()?; let result = transactional(&connection, |connection| { self.set_locked(connection, key, columns, expire_time, now) }); @@ -428,7 +440,7 @@ impl DiskStore for DiskcacheSqliteStore { } fn pop(&self, key: &str, now: f64) -> Result, Error> { - let connection = self.connection.lock().map_err(|_| Error::Unavailable)?; + let connection = self.connection.lock()?; let selected = transactional(&connection, |connection| { let row = connection .query_row( @@ -439,13 +451,13 @@ impl DiskStore for DiskcacheSqliteStore { row_from_query, ) .optional() - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; let Some(row) = row else { return Ok(None); }; connection .execute("DELETE FROM Cache WHERE rowid = ?", params![row.rowid]) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; Ok(Some(row)) })?; let Some(row) = selected else { @@ -460,7 +472,7 @@ impl DiskStore for DiskcacheSqliteStore { } fn clear(&self) -> Result<(), Error> { - let connection = self.connection.lock().map_err(|_| Error::Unavailable)?; + let connection = self.connection.lock()?; let mut last_rowid = 0_i64; loop { let batch = transactional(&connection, |connection| { @@ -469,13 +481,13 @@ impl DiskStore for DiskcacheSqliteStore { "SELECT rowid, filename FROM Cache WHERE rowid > ? ORDER BY rowid LIMIT 100", ) - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? .query_map(params![last_rowid], |row| { Ok((row.get::<_, i64>(0)?, row.get::<_, Option>(1)?)) }) - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? .collect::, _>>() - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; if rows.is_empty() { return Ok(rows); } @@ -486,7 +498,7 @@ impl DiskStore for DiskcacheSqliteStore { .join(","); connection .execute(&format!("DELETE FROM Cache WHERE rowid IN ({ids})"), []) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; Ok(rows) })?; if batch.is_empty() { @@ -509,7 +521,7 @@ impl DiskStore for DiskcacheSqliteStore { now: f64, apply: &mut dyn FnMut(Option) -> Result<(StoredValue, Option), Error>, ) -> Result<(), Error> { - let connection = self.connection.lock().map_err(|_| Error::Unavailable)?; + let connection = self.connection.lock()?; let mut created_filename = None; let result = transactional(&connection, |connection| { let current = connection @@ -521,7 +533,7 @@ impl DiskStore for DiskcacheSqliteStore { row_from_query, ) .optional() - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? .map(|row| fetch_row(&self.directory, row)) .transpose()? .flatten(); @@ -582,21 +594,35 @@ fn read_settings(connection: &Connection) -> Result, Erro }; statement .query_map([], |row| Ok((row.get(0)?, row.get(1)?))) - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? .collect::, _>>() - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) +} + +#[derive(Debug, thiserror::Error)] +enum SettingFault { + #[error("unknown eviction policy {0:?}")] + EvictionPolicy(String), + #[error("setting {0:?} is not an SQLite pragma")] + NotPragma(String), + #[error("setting {0:?} has an unsupported value type")] + PragmaValue(String), } fn apply_pragma(connection: &Connection, key: &str, value: &Value) -> Result<(), Error> { - let pragma = key.strip_prefix("sqlite_").ok_or(Error::Unavailable)?; + let pragma = key + .strip_prefix("sqlite_") + .ok_or_else(|| Error::unavailable(SettingFault::NotPragma(key.to_owned())))?; match value { Value::Integer(value) => connection .pragma_update(None, pragma, value) - .map_err(|_| Error::Unavailable), + .map_err(Error::unavailable), Value::Text(value) => connection .pragma_update(None, pragma, value) - .map_err(|_| Error::Unavailable), - _ => Err(Error::Unavailable), + .map_err(Error::unavailable), + _ => Err(Error::unavailable(SettingFault::PragmaValue( + key.to_owned(), + ))), } } @@ -634,7 +660,7 @@ fn fetch_row(directory: &Path, row: Row) -> Result, Error> { Value::Text(value) => Ok(Some(StoredValue::Text(value))), Value::Integer(value) => Ok(Some(StoredValue::Integer(value))), Value::Real(value) => Ok(Some(StoredValue::Float(value))), - Value::Null => Err(Error::InvalidEntry), + Value::Null => Err(Error::InvalidEntry(None)), }, MODE_BINARY | MODE_PICKLE => { let bytes = match row.value { @@ -645,7 +671,7 @@ fn fetch_row(directory: &Path, row: Row) -> Result, Error> { }; value } - _ => return Err(Error::InvalidEntry), + _ => return Err(Error::InvalidEntry(None)), }; Ok(Some(if row.mode == MODE_BINARY { StoredValue::Bytes(bytes) @@ -663,24 +689,24 @@ fn fetch_row(directory: &Path, row: Row) -> Result, Error> { } Value::Blob(value) => value, Value::Text(value) => value.into_bytes(), - _ => return Err(Error::InvalidEntry), + _ => return Err(Error::InvalidEntry(None)), }; Ok(Some(StoredValue::Text( - String::from_utf8(bytes).map_err(|_| Error::InvalidEntry)?, + String::from_utf8(bytes).map_err(Error::invalid_entry)?, ))) } - _ => Err(Error::InvalidEntry), + _ => Err(Error::InvalidEntry(None)), } } fn read_file(directory: &Path, filename: Option<&str>) -> Result>, Error> { let Some(filename) = filename else { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); }; match fs::read(directory.join(filename)) { Ok(value) => Ok(Some(value)), Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None), - Err(_) => Err(Error::Unavailable), + Err(error) => Err(Error::unavailable(error)), } } @@ -712,7 +738,7 @@ fn store_value( let bytes = value.into_bytes(); let filename = write_file(directory, &bytes)?; Ok(StoredColumns { - size: i64::try_from(bytes.len()).map_err(|_| Error::Unavailable)?, + size: i64::try_from(bytes.len()).map_err(Error::unavailable)?, mode: MODE_TEXT, filename: Some(filename), value: None, @@ -727,7 +753,7 @@ fn store_value( StoredValue::Bytes(value) => { let filename = write_file(directory, &value)?; Ok(StoredColumns { - size: i64::try_from(value.len()).map_err(|_| Error::Unavailable)?, + size: i64::try_from(value.len()).map_err(Error::unavailable)?, mode: MODE_BINARY, filename: Some(filename), value: None, @@ -742,7 +768,7 @@ fn store_value( StoredValue::Pickle(value) => { let filename = write_file(directory, &value)?; Ok(StoredColumns { - size: i64::try_from(value.len()).map_err(|_| Error::Unavailable)?, + size: i64::try_from(value.len()).map_err(Error::unavailable)?, mode: MODE_PICKLE, filename: Some(filename), value: None, @@ -761,14 +787,14 @@ fn write_file(directory: &Path, bytes: &[u8]) -> Result { let filename = format!("{}/{}/{}.val", &hex[..2], &hex[2..4], &hex[4..]); let path = directory.join(&filename); if let Some(parent) = path.parent() { - fs::create_dir_all(parent).map_err(|_| Error::Unavailable)?; + fs::create_dir_all(parent).map_err(Error::unavailable)?; } let mut file = OpenOptions::new() .write(true) .create_new(true) .open(path) - .map_err(|_| Error::Unavailable)?; - file.write_all(bytes).map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; + file.write_all(bytes).map_err(Error::unavailable)?; Ok(filename) } @@ -789,12 +815,12 @@ fn transactional( ) -> Result { connection .execute_batch("BEGIN IMMEDIATE") - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; match operation(connection) { Ok(value) => { connection .execute_batch("COMMIT") - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; Ok(value) } Err(error) => { diff --git a/litellm-rust/crates/cache-disk/tests/python_compat.rs b/litellm-rust/crates/cache-disk/tests/python_compat.rs index 9cbef8573bd..e167b0adb37 100644 --- a/litellm-rust/crates/cache-disk/tests/python_compat.rs +++ b/litellm-rust/crates/cache-disk/tests/python_compat.rs @@ -84,7 +84,7 @@ fn python_read_cases(#[case] row: StoredValue, #[case] expected: ReadExpectation match expected { ReadExpectation::Bytes(expected) => assert_eq!(result.unwrap().unwrap(), expected), ReadExpectation::Miss => assert_eq!(result.unwrap(), None), - ReadExpectation::Invalid => assert!(matches!(result, Err(Error::InvalidEntry))), + ReadExpectation::Invalid => assert!(matches!(result, Err(Error::InvalidEntry(_)))), } } diff --git a/litellm-rust/crates/cache-gcs/Cargo.toml b/litellm-rust/crates/cache-gcs/Cargo.toml index 91630879cbe..dbd8c9ac908 100644 --- a/litellm-rust/crates/cache-gcs/Cargo.toml +++ b/litellm-rust/crates/cache-gcs/Cargo.toml @@ -13,9 +13,11 @@ litellm-auth-types.workspace = true litellm-cache.workspace = true percent-encoding.workspace = true reqwest.workspace = true +thiserror.workspace = true tokio.workspace = true [dev-dependencies] +litellm-error.workspace = true litellm-http = { workspace = true, features = ["test-support"] } litellm-cache-testing.workspace = true rstest.workspace = true diff --git a/litellm-rust/crates/cache-gcs/src/cache.rs b/litellm-rust/crates/cache-gcs/src/cache.rs index bad81573cb4..a6b08445bc8 100644 --- a/litellm-rust/crates/cache-gcs/src/cache.rs +++ b/litellm-rust/crates/cache-gcs/src/cache.rs @@ -25,6 +25,14 @@ pub fn key_prefix(gcs_path: Option<&str>) -> String { } } +#[derive(Debug, thiserror::Error)] +#[error("GCS answered {0}")] +struct UnexpectedStatus(reqwest::StatusCode); + +#[derive(Debug, thiserror::Error)] +#[error("GCS cache worker thread panicked")] +struct WorkerPanicked; + #[derive(Clone, Debug)] pub struct GcsConfig { pub bucket_name: String, @@ -114,9 +122,9 @@ impl GcsCache { .body(payload) .send() .await - .map_err(|_| Error::Unavailable)?; + .map_err(|error| Error::unavailable(litellm_http::Error::from(error)))?; if !response.status().is_success() { - return Err(Error::Unavailable); + return Err(Error::unavailable(UnexpectedStatus(response.status()))); } Ok(()) } @@ -135,18 +143,18 @@ impl GcsCache { .header(reqwest::header::CONTENT_TYPE, "application/json") .send() .await - .map_err(|_| Error::Unavailable)?; + .map_err(|error| Error::unavailable(litellm_http::Error::from(error)))?; if response.status() == reqwest::StatusCode::NOT_FOUND { return Ok(None); } if !response.status().is_success() { - return Err(Error::Unavailable); + return Err(Error::unavailable(UnexpectedStatus(response.status()))); } - let body = response.bytes().await.map_err(|_| Error::Unavailable)?; - self.codec - .decode(&body) - .map(Some) - .map_err(|_| Error::InvalidEntry) + let body = response + .bytes() + .await + .map_err(|error| Error::unavailable(litellm_http::Error::from(error)))?; + self.codec.decode(&body).map(Some) } fn run_sync(future: F) -> Result @@ -158,7 +166,7 @@ impl GcsCache { tokio::runtime::Builder::new_current_thread() .enable_all() .build() - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) .and_then(|runtime| runtime.block_on(future)) }; match tokio::runtime::Handle::try_current() { @@ -169,7 +177,7 @@ impl GcsCache { scope .spawn(|| run(future)) .join() - .map_err(|_| Error::Unavailable) + .map_err(|_| Error::unavailable(WorkerPanicked)) .and_then(|result| result) }), Err(_) => run(future), @@ -242,7 +250,7 @@ impl BatchCache for GcsCache { match self.async_get_cache(&key, &context).await { Ok(Some(value)) => Ok(BatchEntry::Hit(value)), Ok(None) => Ok(BatchEntry::Miss), - Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid), + Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid), Err(error) => Err(error), } } diff --git a/litellm-rust/crates/cache-gcs/src/token.rs b/litellm-rust/crates/cache-gcs/src/token.rs index adb601c276c..079ec44dbce 100644 --- a/litellm-rust/crates/cache-gcs/src/token.rs +++ b/litellm-rust/crates/cache-gcs/src/token.rs @@ -30,7 +30,7 @@ impl TokenSource for GcpTokenSource { self.auth .access_token(&self.config, &|name| std::env::var(name).ok()) .await - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) }) } } diff --git a/litellm-rust/crates/cache-gcs/tests/cache.rs b/litellm-rust/crates/cache-gcs/tests/cache.rs index 3e92c3fca19..f8dba94ea8e 100644 --- a/litellm-rust/crates/cache-gcs/tests/cache.rs +++ b/litellm-rust/crates/cache-gcs/tests/cache.rs @@ -55,19 +55,27 @@ async fn set_writes_encoded_object_and_headers(#[future(awt)] server: MockServer Ok(Some(json!({"value": "entry"}))) )] #[case::missing("missing", ResponseTemplate::new(404), Ok(None))] -#[case::server_error("server-error", ResponseTemplate::new(500), Err(Error::Unavailable))] -#[case::unauthorized("unauthorized", ResponseTemplate::new(401), Err(Error::Unavailable))] +#[case::server_error( + "server-error", + ResponseTemplate::new(500), + Err("cache is unavailable: GCS answered 500 Internal Server Error") +)] +#[case::unauthorized( + "unauthorized", + ResponseTemplate::new(401), + Err("cache is unavailable: GCS answered 401 Unauthorized") +)] #[case::invalid( "invalid", ResponseTemplate::new(200).set_body_string("not json"), - Err(Error::InvalidEntry) + Err("invalid cache entry: ") )] #[tokio::test] async fn get_maps_statuses_and_decode_failures( #[future(awt)] server: MockServer, #[case] key: &str, #[case] response: ResponseTemplate, - #[case] expected: Result, Error>, + #[case] expected: Result, &str>, ) { Mock::given(method("GET")) .and(path(format!("/storage/v1/b/bucket/o/{key}"))) @@ -76,8 +84,21 @@ async fn get_maps_statuses_and_decode_failures( .mount(&server) .await; let cache = support::cache(&server, None); - assert_eq!(cache.get_cache(key, &context()), expected); - assert_eq!(cache.async_get_cache(key, &context()).await, expected); + let outcome = |result: Result, Error>| { + result.map_err(|error| litellm_error::Report::new(&error).to_string()) + }; + let sync = outcome(cache.get_cache(key, &context())); + let asynchronous = outcome(cache.async_get_cache(key, &context()).await); + match expected { + Ok(value) => { + assert_eq!(sync, Ok(value.clone())); + assert_eq!(asynchronous, Ok(value)); + } + Err(prefix) => { + assert!(sync.unwrap_err().starts_with(prefix)); + assert!(asynchronous.unwrap_err().starts_with(prefix)); + } + } } #[rstest] @@ -257,7 +278,7 @@ struct FailingTokenSource; impl TokenSource for FailingTokenSource { fn bearer_token(&self) -> Pin> + Send + '_>> { - Box::pin(async { Err(Error::Unavailable) }) + Box::pin(async { Err(Error::Unavailable(None)) }) } } @@ -267,14 +288,14 @@ async fn token_source_failure_skips_http(#[future(awt)] server: MockServer) { let cache = support::cache_with_token(&server, None, Arc::new(FailingTokenSource)); assert_eq!( cache.get_cache("key", &context()).unwrap_err(), - Error::Unavailable + Error::Unavailable(None) ); assert_eq!( cache .async_set_cache("key", json!(1), context()) .await .unwrap_err(), - Error::Unavailable + Error::Unavailable(None) ); assert_eq!(server.received_requests().await.unwrap().len(), 0); } diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs index 74476559cbe..73fbf14b07a 100644 --- a/litellm-rust/crates/cache-memory/src/cache.rs +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -2,7 +2,7 @@ use std::{ cmp::Reverse, collections::{BinaryHeap, HashMap, HashSet}, hash::Hash, - sync::{Arc, Mutex}, + sync::{Arc, Mutex, MutexGuard}, time::{Duration, SystemTime, UNIX_EPOCH}, }; @@ -84,6 +84,10 @@ impl InMemoryCache { } } + fn lock(&self) -> Result>, Error> { + Ok(self.state.lock()?) + } + pub fn set_cache( &self, key: impl Into, @@ -94,13 +98,13 @@ impl InMemoryCache { return Ok(CacheWrite::Disabled); } let now = (self.now)(); - let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; + let mut state = self.lock()?; self.store(&mut state, key.into(), value, ttl, now) } pub fn get_cache(&self, key: &str) -> Result, Error> { let now = (self.now)(); - let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; + let mut state = self.lock()?; if state .expirations .get(key) @@ -124,7 +128,7 @@ impl InMemoryCache { /// fits. pub fn evict_cache(&self) -> Result<(), Error> { let now = (self.now)(); - let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; + let mut state = self.lock()?; Self::evict(&mut state, self.max_size_in_memory, now, None); Ok(()) } @@ -132,7 +136,7 @@ impl InMemoryCache { /// `evict_element_if_expired`: `true` when `key` had expired and was removed. pub fn evict_element_if_expired(&self, key: &str) -> Result { let now = (self.now)(); - let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; + let mut state = self.lock()?; let expired = state .expirations .get(key) @@ -153,12 +157,7 @@ impl InMemoryCache { /// The number of stored entries, expired ones included until they are evicted. pub fn len(&self) -> Result { - Ok(self - .state - .lock() - .map_err(|_| Error::Unavailable)? - .values - .len()) + Ok(self.lock()?.values.len()) } pub fn is_empty(&self) -> Result { @@ -167,12 +166,7 @@ impl InMemoryCache { /// Entries in the expiration heap, stale ones included; bounded by eviction. pub fn expiration_heap_len(&self) -> Result { - Ok(self - .state - .lock() - .map_err(|_| Error::Unavailable)? - .expiration_heap - .len()) + Ok(self.lock()?.expiration_heap.len()) } pub fn max_size_in_memory(&self) -> usize { @@ -184,13 +178,7 @@ impl InMemoryCache { } pub fn expires_at(&self, key: &str) -> Result, Error> { - Ok(self - .state - .lock() - .map_err(|_| Error::Unavailable)? - .expirations - .get(key) - .copied()) + Ok(self.lock()?.expirations.get(key).copied()) } pub async fn async_get_ttl(&self, key: &str) -> Result, Error> { @@ -198,7 +186,7 @@ impl InMemoryCache { } pub async fn async_get_oldest_n_keys(&self, count: usize) -> Result, Error> { - let state = self.state.lock().map_err(|_| Error::Unavailable)?; + let state = self.lock()?; let mut expirations = state .expirations .iter() @@ -213,13 +201,13 @@ impl InMemoryCache { } pub fn delete_cache(&self, key: &str) -> Result<(), Error> { - let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; + let mut state = self.lock()?; Self::remove(&mut state, key); Ok(()) } pub fn flush_cache(&self) -> Result<(), Error> { - let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; + let mut state = self.lock()?; state.values.clear(); state.expirations.clear(); state.expiration_heap.clear(); @@ -316,7 +304,7 @@ where return Ok(candidate); } let now = (self.now)(); - let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; + let mut state = self.lock()?; Self::evict(&mut state, self.max_size_in_memory, now, Some(key)); let existing = state .values @@ -351,7 +339,7 @@ impl CounterCache for InMemoryCache { return Ok(amount); } let now = (self.now)(); - let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; + let mut state = self.lock()?; let value = Self::live(&mut state, key, now).unwrap_or_default() + amount; self.store(&mut state, key.into(), value, self.get_ttl(&context), now)?; Ok(value) @@ -424,7 +412,7 @@ where return Ok(values); } let now = (self.now)(); - let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; + let mut state = self.lock()?; let mut stored = Self::live(&mut state, key, now).unwrap_or_default(); stored.extend(values.iter().cloned()); self.store(&mut state, key.into(), stored, ttl, now)?; diff --git a/litellm-rust/crates/cache-memory/tests/cache.rs b/litellm-rust/crates/cache-memory/tests/cache.rs index ce7d6ac8a98..e6c74e5ccda 100644 --- a/litellm-rust/crates/cache-memory/tests/cache.rs +++ b/litellm-rust/crates/cache-memory/tests/cache.rs @@ -52,7 +52,7 @@ fn measured(capacity: usize) -> InMemoryCache { Some(4), Some(Arc::new(|value: &String| { if value.is_empty() { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); } Ok(value.len()) })), @@ -275,7 +275,7 @@ fn evict_element_if_expired_reports_removal(clock: Clock) { #[case::fits("ok", Ok(true))] #[case::at_limit("four", Ok(true))] #[case::too_large("oversized", Ok(false))] -#[case::measure_error("", Err(Error::InvalidEntry))] +#[case::measure_error("", Err(Error::InvalidEntry(None)))] fn check_value_size_applies_the_entry_limit( #[case] value: &str, #[case] expected: Result, @@ -310,7 +310,7 @@ fn disabled_size_limited_and_validated_writes_are_observable() { assert_eq!(cache.get_cache("small").unwrap(), Some("ok".into())); assert_eq!( cache.set_cache("invalid", String::new(), None), - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(None)) ); assert_eq!(cache.get_cache("invalid").unwrap(), None); cache.delete_cache("small").unwrap(); diff --git a/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml b/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml index a44bef0a731..8a4caa8def7 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml +++ b/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml @@ -18,6 +18,7 @@ tokio.workspace = true uuid.workspace = true [dev-dependencies] +litellm-error.workspace = true litellm-http = { workspace = true, features = ["test-support"] } futures-executor = "0.3" litellm-cache-testing.workspace = true diff --git a/litellm-rust/crates/cache-qdrant-semantic/src/cache.rs b/litellm-rust/crates/cache-qdrant-semantic/src/cache.rs index a140e0af174..9c93cc45386 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/src/cache.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/src/cache.rs @@ -36,7 +36,7 @@ impl QdrantSemanticCache { let exists = client .collection_exists(config.collection_name.clone()) .await - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; if !exists { client .create_collection( @@ -48,7 +48,7 @@ impl QdrantSemanticCache { .quantization_config(quantization(&config.quantization)), ) .await - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } let _ = client .create_field_index(CreateFieldIndexCollectionBuilder::new( @@ -99,13 +99,13 @@ impl QdrantSemanticCache { .async_embed(&prompt, context.metadata.as_ref()) .await?; let response = - String::from_utf8(self.codec.encode(&value)?).map_err(|_| Error::InvalidEntry)?; + String::from_utf8(self.codec.encode(&value)?).map_err(Error::invalid_entry)?; let payload = Payload::try_from(json!({ "litellm_cache_key": key, "text": prompt, "response": response, })) - .map_err(|_| Error::InvalidEntry)?; + .map_err(Error::invalid_entry)?; self.client .upsert_points( UpsertPointsBuilder::new( @@ -119,7 +119,7 @@ impl QdrantSemanticCache { .wait(true), ) .await - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; Ok(()) } @@ -152,7 +152,7 @@ impl QdrantSemanticCache { ), ) .await - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; let Some(point) = result.result.into_iter().next() else { return Ok(SemanticLookup::miss(Some(0.0))); }; @@ -170,7 +170,7 @@ impl QdrantSemanticCache { let response = payload .get("response") .and_then(Value::as_str) - .ok_or(Error::InvalidEntry)?; + .ok_or(Error::InvalidEntry(None))?; Ok(SemanticLookup { value: Some(self.codec.decode(response.as_bytes())?), similarity: Some(similarity), diff --git a/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs b/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs index b7fbcd9b02d..bad780bc57f 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs @@ -55,10 +55,9 @@ impl Embedder for OpenAiEmbedder { } .send() .await - .map_err(|_| Error::Unavailable)? - .error_for_status() - .map_err(|_| Error::Unavailable)?; - let body: Value = response.json().await.map_err(|_| Error::Unavailable)?; + .and_then(reqwest::Response::error_for_status) + .map_err(http_error)?; + let body: Value = response.json().await.map_err(http_error)?; body.get("data") .and_then(Value::as_array) .and_then(|data| data.first()) @@ -70,6 +69,14 @@ impl Embedder for OpenAiEmbedder { .map(|value| value.as_f64().map(|value| value as f32)) .collect::>>() }) - .ok_or(Error::Unavailable) + .ok_or(Error::unavailable(MissingEmbedding)) } } + +fn http_error(error: reqwest::Error) -> Error { + Error::unavailable(litellm_http::Error::from(error)) +} + +#[derive(Debug, thiserror::Error)] +#[error("embeddings response has no embedding")] +struct MissingEmbedding; diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs index 71fc986616e..42261d21524 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs @@ -134,13 +134,19 @@ async fn posts_embeddings_request_and_parses_vector() { } #[rstest] -#[case::error_status("500 Internal Server Error", "{}", 0, None, Err(Error::Unavailable))] +#[case::error_status( + "500 Internal Server Error", + "{}", + 0, + None, + Err(Error::Unavailable(None)) +)] #[case::timed_out( "200 OK", r#"{"data":[{"embedding":[0.1,0.2]}]}"#, 500, Some(Duration::from_millis(200)), - Err(Error::Unavailable) + Err(Error::Unavailable(None)) )] #[case::within_timeout( "200 OK", @@ -149,7 +155,7 @@ async fn posts_embeddings_request_and_parses_vector() { Some(Duration::from_secs(1)), Ok(vec![0.1, 0.2]) )] -#[case::missing_embedding("200 OK", r#"{"data":[]}"#, 0, None, Err(Error::Unavailable))] +#[case::missing_embedding("200 OK", r#"{"data":[]}"#, 0, None, Err(Error::Unavailable(None)))] #[tokio::test] async fn status_timeout_and_body_errors_are_unavailable( #[case] status: &str, @@ -164,7 +170,14 @@ async fn status_timeout_and_body_errors_are_unavailable( litellm_http::Client::for_test(), config(server.base_url(), timeout), ); - assert_eq!(embedder.async_embed("hello", None).await, expected); + let result = embedder.async_embed("hello", None).await; + if let Err(error) = &result { + let report = litellm_error::Report::new(error).to_string(); + assert!(!report.contains(&server.base_url()), "{report}"); + } + let kind = + |result: Result, Error>| result.map_err(|error| std::mem::discriminant(&error)); + assert_eq!(kind(result), kind(expected)); } #[rstest] diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs index fe25c503a36..b0764b4470b 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs @@ -53,7 +53,10 @@ impl Embedder for FixedEmbedder { .lock() .unwrap() .push((input.to_owned(), metadata.cloned())); - self.vectors.get(input).cloned().ok_or(Error::Unavailable) + self.vectors + .get(input) + .cloned() + .ok_or(Error::Unavailable(None)) } } @@ -396,9 +399,9 @@ async fn lookup_reports_python_semantic_similarity( #[rstest] #[case::codec_decodes_the_payload(Some(json!("{\"a\":1}")), Ok(Some(json!({"a": 1}))))] -#[case::undecodable_response(Some(json!("not json")), Err(Error::InvalidEntry))] -#[case::non_string_response(Some(json!(1)), Err(Error::InvalidEntry))] -#[case::missing_response(None, Err(Error::InvalidEntry))] +#[case::undecodable_response(Some(json!("not json")), Err(Error::InvalidEntry(None)))] +#[case::non_string_response(Some(json!(1)), Err(Error::InvalidEntry(None)))] +#[case::missing_response(None, Err(Error::InvalidEntry(None)))] #[tokio::test(flavor = "multi_thread")] async fn stored_responses_go_through_the_codec( #[case] response: Option, @@ -419,9 +422,12 @@ async fn stored_responses_go_through_the_codec( .into(), }); + let kind = |result: Result, Error>| { + result.map_err(|error| std::mem::discriminant(&error)) + }; assert_eq!( - cache.async_get_cache("key", &context("hello")).await, - expected + kind(cache.async_get_cache("key", &context("hello")).await), + kind(expected) ); server.stop(); } @@ -434,7 +440,7 @@ async fn embedding_failures_propagate() { assert_eq!( cache.async_get_cache("key", &context("unknown")).await, - Err(Error::Unavailable) + Err(Error::Unavailable(None)) ); server.stop(); } @@ -493,10 +499,10 @@ async fn stopped_qdrant_server_maps_to_unavailable() { let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await; server.stop(); tokio::time::sleep(Duration::from_millis(50)).await; - assert_eq!( + assert!(matches!( cache.async_get_cache("key", &context("hello")).await, - Err(Error::Unavailable) - ); + Err(Error::Unavailable(Some(_))) + )); } /// `_payload_matches_cache_key` compares `str(cached_key) == str(key)`, so a point whose stored diff --git a/litellm-rust/crates/cache-redis-semantic/src/cache.rs b/litellm-rust/crates/cache-redis-semantic/src/cache.rs index 02feb3cbdf4..42fc6ae585f 100644 --- a/litellm-rust/crates/cache-redis-semantic/src/cache.rs +++ b/litellm-rust/crates/cache-redis-semantic/src/cache.rs @@ -63,13 +63,13 @@ impl Inner { .arg(CACHE_KEY_FIELD) .arg(tag) .query::<()>(connection) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; if let Some(ttl) = ttl { redis::cmd("EXPIRE") .arg(&hash_key) .arg(ttl_seconds(ttl)) .query::<()>(connection) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } Ok(()) } @@ -111,7 +111,7 @@ impl Inner { .arg("vector") .arg(vector_buffer(vector)) .query::(connection) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; let Some(fields) = first_document(&result) else { return Ok(SemanticLookup::miss(Some(0.0))); }; diff --git a/litellm-rust/crates/cache-redis-semantic/src/index.rs b/litellm-rust/crates/cache-redis-semantic/src/index.rs index c141e47003a..23401938a0e 100644 --- a/litellm-rust/crates/cache-redis-semantic/src/index.rs +++ b/litellm-rust/crates/cache-redis-semantic/src/index.rs @@ -40,10 +40,10 @@ impl Index { Some(false) => self.isolated(connection, dims)?, None => match create_index(connection, &self.name, dims) { Ok(()) => self.name.clone(), - Err(_) => match index_compatible(connection, &self.name, dims)? { + Err(error) => match index_compatible(connection, &self.name, dims)? { Some(true) => self.name.clone(), Some(false) => self.isolated(connection, dims)?, - None => return Err(Error::Unavailable), + None => return Err(error), }, }, }; @@ -59,7 +59,7 @@ impl Index { redis::cmd("FT.DROPINDEX") .arg(&name) .query::<()>(connection) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; create_index(connection, &name, dims)?; Ok(name) } @@ -109,7 +109,7 @@ fn create_index(connection: &mut ConnectionRef<'_>, name: &str, dims: usize) -> .arg("SEPARATOR") .arg(",") .query::<()>(connection) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) } fn index_compatible( @@ -123,14 +123,18 @@ fn index_compatible( { Ok(info) => info, Err(error) if unknown_index(&error) => return Ok(None), - Err(_) => return Err(Error::Unavailable), + Err(error) => return Err(Error::unavailable(error)), }; Ok(Some(schema_compatible(&info, dims))) } fn unknown_index(error: &redis::RedisError) -> bool { - let message = error.to_string().to_lowercase(); - message.contains("unknown") && message.contains("index") + error + .code() + .is_some_and(|code| code.eq_ignore_ascii_case("unknown")) + && error + .detail() + .is_some_and(|detail| detail.to_ascii_lowercase().starts_with("index")) } struct Attribute { diff --git a/litellm-rust/crates/cache-redis-semantic/tests/cache.rs b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs index fa4ac00e767..cc42f4c3924 100644 --- a/litellm-rust/crates/cache-redis-semantic/tests/cache.rs +++ b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs @@ -79,7 +79,7 @@ fn s(value: &str) -> redis::Value { } fn unknown_index_error() -> redis::RedisError { - redis::RedisError::from((redis::ErrorKind::Extension, "Unknown index name")) + support::server_error("Unknown index name") } fn info_missing(index: &str) -> MockCmd { @@ -320,7 +320,7 @@ fn store_writes_the_redisvl_hash( #[case::below_distance_threshold("key1", "0.05", None, Ok(Some(entry())))] #[case::above_distance_threshold("key1", "0.5", None, Ok(None))] #[case::other_cache_key("other", "0.05", None, Ok(None))] -#[case::malformed_response("key1", "0.05", Some(b"not json!".as_slice()), Err(Error::InvalidEntry))] +#[case::malformed_response("key1", "0.05", Some(b"not json!".as_slice()), Err(Error::InvalidEntry(None)))] fn lookup_applies_threshold_scope_and_codec( #[case] stored_tag: &str, #[case] distance: &str, @@ -343,7 +343,10 @@ fn lookup_applies_threshold_scope_and_codec( FakeEmbedder::new(&[]), ); - assert_eq!(cache.get_cache("key1", &context), expected); + let kind = |result: Result, Error>| { + result.map_err(|error| std::mem::discriminant(&error)) + }; + assert_eq!(kind(cache.get_cache("key1", &context)), kind(expected)); } #[rstest] @@ -451,10 +454,7 @@ fn create_index_race_rechecks_schema_and_stores(entry: Value, context: SemanticC info_missing(INDEX), MockCmd::new( create_index_command(INDEX, 3), - Err::<&str, _>(redis::RedisError::from(( - redis::ErrorKind::Extension, - "Index already exists", - ))), + Err::<&str, _>(support::server_error("Index already exists")), ), info(INDEX, compatible_info(3)), hset(INDEX, PROMPT, "key1", &VECTOR, &entry), @@ -673,10 +673,10 @@ fn shared_base_index_across_dimensions_replaces_the_isolated_index(entry: Value) .set_cache("key1", entry.clone(), &context()) .unwrap(); - assert_eq!( - worker_b.get_cache("key1", &context()).unwrap_err(), - Error::Unavailable - ); + assert!(matches!( + worker_b.get_cache("key1", &context()), + Err(Error::Unavailable(Some(_))) + )); } #[fixture] @@ -732,10 +732,10 @@ fn live_shared_index_is_replaced_across_dimensions(redis_stack_url: Option { let name = text(1); if state.indexes.contains_key(&name) { - return Err(error("Index already exists")); + return Err(server_error("Index already exists")); } let position = |token: &str| args.iter().position(|arg| arg == token.as_bytes()); let prefix = args[position("PREFIX").unwrap() + 2].clone(); @@ -94,7 +94,7 @@ impl FakeSearch { let index = state .indexes .get(&text(1)) - .ok_or_else(|| error("Unknown index name"))?; + .ok_or_else(|| server_error("Unknown index name"))?; Ok(index_info(index)) } "FT.DROPINDEX" => { @@ -118,7 +118,7 @@ impl FakeSearch { let index = state .indexes .get(&text(1)) - .ok_or_else(|| error("no such index"))?; + .ok_or_else(|| server_error("no such index"))?; let query = text(2); let tag = query_tag(&query); let params = args.iter().position(|arg| arg == b"PARAMS").unwrap(); @@ -155,7 +155,7 @@ impl FakeSearch { ])) } "PING" => Ok(redis::Value::SimpleString("PONG".into())), - _ => Err(error("unsupported command")), + _ => Err(server_error("unsupported command")), } } } @@ -192,8 +192,11 @@ impl redis::ConnectionLike for FakeSearch { } } -fn error(message: &'static str) -> redis::RedisError { - redis::RedisError::from((redis::ErrorKind::Extension, message)) +pub fn server_error(reply: &str) -> redis::RedisError { + redis::parse_redis_value(format!("-{reply}\r\n").as_bytes()) + .expect("error reply") + .extract_error() + .expect_err("server error") } fn bulk(bytes: &[u8]) -> redis::Value { diff --git a/litellm-rust/crates/cache-redis/Cargo.toml b/litellm-rust/crates/cache-redis/Cargo.toml index bebfcafd98a..f7621299506 100644 --- a/litellm-rust/crates/cache-redis/Cargo.toml +++ b/litellm-rust/crates/cache-redis/Cargo.toml @@ -7,8 +7,10 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true +litellm-error.workspace = true redis = { workspace = true, features = ["cluster", "tls-rustls"] } r2d2 = "0.8.10" +thiserror.workspace = true tokio.workspace = true [dev-dependencies] diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index ebcf0b6916b..b873a2c0282 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -141,7 +141,7 @@ where redis::Value::Nil => Ok(None), redis::Value::BulkString(bytes) => self.codec.decode(&bytes).map(Some), redis::Value::SimpleString(text) => self.codec.decode(text.as_bytes()).map(Some), - _ => Err(Error::InvalidEntry), + _ => Err(Error::InvalidEntry(None)), } } @@ -152,7 +152,7 @@ where match self.decode_response(value) { Ok(Some(value)) => Ok(BatchEntry::Hit(value)), Ok(None) => Ok(BatchEntry::Miss), - Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid), + Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid), Err(error) => Err(error), } } diff --git a/litellm-rust/crates/cache-redis/src/claim.rs b/litellm-rust/crates/cache-redis/src/claim.rs index e4cbe8c8d5c..c57a1b03395 100644 --- a/litellm-rust/crates/cache-redis/src/claim.rs +++ b/litellm-rust/crates/cache-redis/src/claim.rs @@ -17,7 +17,7 @@ fn stored_bytes(value: redis::Value) -> Result>, Error> { redis::Value::Nil => Ok(None), redis::Value::BulkString(bytes) => Ok(Some(bytes)), redis::Value::SimpleString(text) => Ok(Some(text.into_bytes())), - _ => Err(Error::InvalidEntry), + _ => Err(Error::InvalidEntry(None)), } } @@ -37,13 +37,13 @@ where { let payload = codec.encode(&candidate)?; if payload.is_empty() { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); } for _ in 0..CLAIM_ATTEMPTS { let current = stored_bytes( connection .get::<_, redis::Value>(key) - .map_err(|_| Error::Unavailable)?, + .map_err(Error::unavailable)?, )? .filter(|bytes| !bytes.is_empty()); let existing = current @@ -63,14 +63,18 @@ where .arg(write) .arg(u8::from(refresh)) .query::(connection) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; if applied { return Ok(existing.unwrap_or(candidate)); } } - Err(Error::Unavailable) + Err(Error::unavailable(ClaimContended)) } +#[derive(Debug, thiserror::Error)] +#[error("claim kept losing the race after {CLAIM_ATTEMPTS} attempts")] +struct ClaimContended; + impl ClaimCache for RedisCache where S: CacheCodec + Clone + 'static, diff --git a/litellm-rust/crates/cache-redis/src/connection.rs b/litellm-rust/crates/cache-redis/src/connection.rs index 2f58e2a9b80..99befa16ffa 100644 --- a/litellm-rust/crates/cache-redis/src/connection.rs +++ b/litellm-rust/crates/cache-redis/src/connection.rs @@ -36,19 +36,19 @@ where ) -> Result { match self { Self::Pool(pool) => { - let mut pooled = pool.get().map_err(|_| Error::Unavailable)?; + let mut pooled = pool.get().map_err(Error::unavailable)?; let result = operation(&mut ConnectionRef::Node(&mut pooled.connection)); - pooled.failed = matches!(result, Err(Error::Unavailable)); + pooled.failed = matches!(result, Err(Error::Unavailable(_))); result } Self::Cluster(pool) => { - let mut pooled = pool.get().map_err(|_| Error::Unavailable)?; + let mut pooled = pool.get().map_err(Error::unavailable)?; let result = operation(&mut ConnectionRef::Cluster(&mut pooled.connection)); - pooled.failed = matches!(result, Err(Error::Unavailable)); + pooled.failed = matches!(result, Err(Error::Unavailable(_))); result } Self::Fixed(connection) => { - let mut connection = connection.lock().map_err(|_| Error::Unavailable)?; + let mut connection = connection.lock()?; operation(&mut ConnectionRef::Node(&mut *connection)) } } @@ -61,7 +61,7 @@ where { tokio::task::spawn_blocking(move || connections.execute(operation)) .await - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? } pub fn fixed(connection: C) -> Self { @@ -95,7 +95,7 @@ fn pool(manager: M) -> Result, Error> { .connection_timeout(REDIS_TIMEOUT) .test_on_check_out(false) .build(manager) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) } fn close_idle(pool: &r2d2::Pool) @@ -123,7 +123,7 @@ impl ConnectionManager { fn open(url: &str) -> Result { redis::Client::open(url) .map(Self) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) } } @@ -151,14 +151,26 @@ impl r2d2::ManageConnection for ConnectionManager { } } +#[derive(Debug, thiserror::Error)] +enum TopologyFault { + #[error("Redis Cluster needs at least one startup node")] + NoStartupNodes, + #[error("Redis Cluster nodes need a TCP address")] + NonTcpAddress, + #[error("Redis Cluster answered a fan-out command without a per-node map")] + FanOutReply, + #[error("scan target does not match the connection topology")] + ScanTarget, +} + pub(crate) struct ClusterConnectionManager(ClusterClient); impl ClusterConnectionManager { fn open(url: &str, startup_nodes: &[RedisNode]) -> Result { if startup_nodes.is_empty() { - return Err(Error::Unavailable); + return Err(Error::unavailable(TopologyFault::NoStartupNodes)); } - let info = url.into_connection_info().map_err(|_| Error::Unavailable)?; + let info = url.into_connection_info().map_err(Error::unavailable)?; let nodes = startup_nodes .iter() .map(|node| node_info(&info, node)) @@ -168,7 +180,7 @@ impl ClusterConnectionManager { .response_timeout(REDIS_TIMEOUT) .build() .map(Self) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) } } @@ -185,7 +197,7 @@ fn node_info(info: &ConnectionInfo, node: &RedisNode) -> Result return Err(Error::Unavailable), + _ => return Err(Error::unavailable(TopologyFault::NonTcpAddress)), }; Ok(info.clone().set_addr(addr)) } @@ -298,7 +310,7 @@ impl ConnectionRef<'_> { .and_then(|replies| pipeline.query(&mut Replies(Some(replies)))) } } - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) } pub(crate) fn scan( @@ -311,7 +323,7 @@ impl ConnectionRef<'_> { Self::Node(connection) => { let page = scan_command(0, pattern, count) .query::(*connection) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; vec![(None, page)] } Self::Cluster(connection) => connection @@ -322,7 +334,7 @@ impl ConnectionRef<'_> { Some(ResponsePolicy::Special), )), ) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) .and_then(primary_pages)? .into_iter() .map(|(node, page)| (Some(node), page)) @@ -362,7 +374,7 @@ impl ConnectionRef<'_> { pub(crate) fn node_text(&mut self, command: &redis::Cmd) -> Result { match self { - Self::Node(connection) => command.query(*connection).map_err(|_| Error::Unavailable), + Self::Node(connection) => command.query(*connection).map_err(Error::unavailable), Self::Cluster(connection) => { let value = connection .route_command( @@ -372,18 +384,16 @@ impl ConnectionRef<'_> { Some(ResponsePolicy::Special), )), ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; let redis::Value::Map(entries) = value else { - return Err(Error::Unavailable); + return Err(Error::unavailable(TopologyFault::FanOutReply)); }; let mut replies = entries .into_iter() .map(|(node, reply)| { Ok(( - redis::from_redis_value::(node) - .map_err(|_| Error::Unavailable)?, - redis::from_redis_value::(reply) - .map_err(|_| Error::Unavailable)?, + redis::from_redis_value::(node).map_err(Error::unavailable)?, + redis::from_redis_value::(reply).map_err(Error::unavailable)?, )) }) .collect::, Error>>()?; @@ -400,7 +410,7 @@ impl ConnectionRef<'_> { pub(crate) fn flushall(&mut self) -> Result<(), Error> { let command = redis::cmd("FLUSHALL"); match self { - Self::Node(connection) => command.query(*connection).map_err(|_| Error::Unavailable), + Self::Node(connection) => command.query(*connection).map_err(Error::unavailable), Self::Cluster(connection) => connection .route_command( &command, @@ -410,7 +420,7 @@ impl ConnectionRef<'_> { )), ) .map(|_| ()) - .map_err(|_| Error::Unavailable), + .map_err(Error::unavailable), } } @@ -424,7 +434,7 @@ impl ConnectionRef<'_> { let command = scan_command(cursor, pattern, count); match (self, node) { (Self::Node(connection), None) => { - command.query(*connection).map_err(|_| Error::Unavailable) + command.query(*connection).map_err(Error::unavailable) } (Self::Cluster(connection), Some(node)) => connection .route_command( @@ -434,9 +444,9 @@ impl ConnectionRef<'_> { port: node.port(), }), ) - .map_err(|_| Error::Unavailable) - .and_then(|value| redis::from_redis_value(value).map_err(|_| Error::Unavailable)), - _ => Err(Error::Unavailable), + .map_err(Error::unavailable) + .and_then(|value| redis::from_redis_value(value).map_err(Error::unavailable)), + _ => Err(Error::unavailable(TopologyFault::ScanTarget)), } } } @@ -445,14 +455,14 @@ type ScanPage = (u64, Vec); fn primary_pages(value: redis::Value) -> Result, Error> { let redis::Value::Map(entries) = value else { - return Err(Error::Unavailable); + return Err(Error::unavailable(TopologyFault::FanOutReply)); }; entries .into_iter() .map(|(node, page)| { - let node = redis::from_redis_value::(node).map_err(|_| Error::Unavailable)?; - let node = NodeAddress::try_from(node.as_str()).map_err(|_| Error::Unavailable)?; - let page = redis::from_redis_value::(page).map_err(|_| Error::Unavailable)?; + let node = redis::from_redis_value::(node).map_err(Error::unavailable)?; + let node = NodeAddress::try_from(node.as_str()).map_err(Error::unavailable)?; + let page = redis::from_redis_value::(page).map_err(Error::unavailable)?; Ok((node, page)) }) .collect() diff --git a/litellm-rust/crates/cache-redis/src/counter.rs b/litellm-rust/crates/cache-redis/src/counter.rs index adcd597b2f4..a4beb9f5712 100644 --- a/litellm-rust/crates/cache-redis/src/counter.rs +++ b/litellm-rust/crates/cache-redis/src/counter.rs @@ -100,7 +100,7 @@ fn increment( .arg(amount) .arg(ttl) .query(connection) - .map_err(|_| Error::Unavailable); + .map_err(Error::unavailable); } connection .query_pipeline( @@ -140,7 +140,7 @@ where } fn count(value: redis::Value) -> Result, Error> { - redis::from_redis_value(value).map_err(|_| Error::InvalidEntry) + redis::from_redis_value(value).map_err(Error::invalid_entry) } impl BoundedCounterCache for RedisCache @@ -182,7 +182,7 @@ where .arg(value) .arg(ttl) .query(connection) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) }) .await } @@ -201,5 +201,5 @@ fn increment_with_floor( .arg(amount) .arg(ttl) .query(connection) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) } diff --git a/litellm-rust/crates/cache-redis/src/keys.rs b/litellm-rust/crates/cache-redis/src/keys.rs index b9d7bb20ff3..2e8a9e60e65 100644 --- a/litellm-rust/crates/cache-redis/src/keys.rs +++ b/litellm-rust/crates/cache-redis/src/keys.rs @@ -16,7 +16,7 @@ where redis::cmd("TTL") .arg(key) .query::(connection) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) }) .await?; Ok(u64::try_from(ttl).ok().map(Duration::from_secs)) @@ -36,7 +36,7 @@ where .arg(key) .arg(ttl) .query(connection) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) }) .await } diff --git a/litellm-rust/crates/cache-redis/src/lifecycle.rs b/litellm-rust/crates/cache-redis/src/lifecycle.rs index 6ab8c31509e..54fa8ecc3bf 100644 --- a/litellm-rust/crates/cache-redis/src/lifecycle.rs +++ b/litellm-rust/crates/cache-redis/src/lifecycle.rs @@ -11,11 +11,11 @@ where C: redis::ConnectionLike + Send + 'static, { fn sync_ping(&self) -> Result { - self.execute(|connection| connection.ping().map_err(|_| Error::Unavailable)) + self.execute(|connection| connection.ping().map_err(Error::unavailable)) } async fn ping(&self) -> Result { - self.run(|connection| connection.ping().map_err(|_| Error::Unavailable)) + self.run(|connection| connection.ping().map_err(Error::unavailable)) .await } } @@ -33,9 +33,10 @@ where RedisTopology::Cluster { .. } => "Redis Cluster", }; let ping = self - .run(|connection| Ok(connection.ping().map_err(|error| error.to_string()))) + .run(|connection| Ok(connection.ping().map_err(Error::unavailable))) .await - .unwrap_or_else(|error| Err(error.to_string())); + .and_then(|ping| ping) + .map_err(|error| litellm_error::Report::new(&error).to_string()); Ok(match ping { Ok(true) => CacheConnectionResult { status: CacheConnectionStatus::Success, diff --git a/litellm-rust/crates/cache-redis/src/queue.rs b/litellm-rust/crates/cache-redis/src/queue.rs index 623bce23640..d55e4f49187 100644 --- a/litellm-rust/crates/cache-redis/src/queue.rs +++ b/litellm-rust/crates/cache-redis/src/queue.rs @@ -29,7 +29,7 @@ where ttl: Option, ) -> Result { if values.is_empty() { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); } let key = self.namespaced_key(key); let ttl = self.ttl_or_default(ttl); @@ -58,7 +58,7 @@ where async fn async_rpush(&self, key: &str, values: Vec) -> Result { if values.is_empty() { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); } let key = self.namespaced_key(key); self.run(move |connection| { @@ -66,7 +66,7 @@ where .arg(key) .arg(values) .query(connection) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) }) .await } @@ -78,7 +78,7 @@ where max_len: usize, ) -> Result { if values.is_empty() { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); } let key = self.namespaced_key(key); let start = i64::try_from(max_len).map_or(i64::MIN, |max_len| -max_len); @@ -108,7 +108,7 @@ where let mut pipeline = redis::pipe(); for operation in operations { if operation.values.is_empty() { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); } pipeline .cmd("RPUSH") @@ -130,7 +130,7 @@ where .run(move |connection| { command .query::(connection) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) }) .await?; lpop_result(value, count.is_some()) @@ -182,7 +182,7 @@ fn redis_bytes(value: redis::Value) -> Result, Error> { match value { redis::Value::BulkString(bytes) => Ok(bytes), redis::Value::SimpleString(text) => Ok(text.into_bytes()), - _ => Err(Error::InvalidEntry), + _ => Err(Error::InvalidEntry(None)), } } @@ -221,6 +221,6 @@ fn lpop_result(value: redis::Value, multiple: bool) -> Result, _>>() .map(RedisLpopResult::Values), value if !multiple => redis_bytes(value).map(RedisLpopResult::Value), - _ => Err(Error::InvalidEntry), + _ => Err(Error::InvalidEntry(None)), } } diff --git a/litellm-rust/crates/cache-redis/src/script.rs b/litellm-rust/crates/cache-redis/src/script.rs index e8a16bcc2b5..09987437ed4 100644 --- a/litellm-rust/crates/cache-redis/src/script.rs +++ b/litellm-rust/crates/cache-redis/src/script.rs @@ -116,7 +116,7 @@ fn eval( .arg(keys) .arg(arguments) .query(connection) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) } impl ScriptCache for RedisCache diff --git a/litellm-rust/crates/cache-redis/src/store.rs b/litellm-rust/crates/cache-redis/src/store.rs index 111d481ad02..daadc940ca9 100644 --- a/litellm-rust/crates/cache-redis/src/store.rs +++ b/litellm-rust/crates/cache-redis/src/store.rs @@ -32,7 +32,7 @@ where self.execute(|connection| { connection .set_ex::<_, _, ()>(key, payload, ttl) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) }) } @@ -41,7 +41,7 @@ where let value = self.execute(|connection| { connection .get::<_, redis::Value>(key) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) })?; self.decode_response(value) } @@ -58,7 +58,7 @@ where self.run(move |connection| { connection .set_ex::<_, _, ()>(key, payload, ttl) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) }) .await } @@ -73,7 +73,7 @@ where .run(move |connection| { connection .get::<_, redis::Value>(key) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) }) .await?; self.decode_response(value) @@ -158,7 +158,7 @@ pub(crate) fn mget( redis::cmd("MGET") .arg(keys) .query(connection) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) } impl DeleteCache for RedisCache @@ -168,12 +168,12 @@ where { fn delete_cache(&self, key: &str) -> Result<(), Error> { let key = self.namespaced_key(key); - self.execute(|connection| connection.del::<_, ()>(key).map_err(|_| Error::Unavailable)) + self.execute(|connection| connection.del::<_, ()>(key).map_err(Error::unavailable)) } async fn async_delete_cache(&self, key: &str) -> Result<(), Error> { let key = self.namespaced_key(key); - self.run(move |connection| connection.del::<_, ()>(key).map_err(|_| Error::Unavailable)) + self.run(move |connection| connection.del::<_, ()>(key).map_err(Error::unavailable)) .await } } @@ -188,7 +188,7 @@ where return Ok(0); } let keys = self.namespaced_keys(&keys); - self.run(move |connection| connection.del(keys).map_err(|_| Error::Unavailable)) + self.run(move |connection| connection.del(keys).map_err(Error::unavailable)) .await } } @@ -225,7 +225,7 @@ fn flush_matching(connection: &mut ConnectionRef<'_>, pattern: &str) -> Result<( if !keys.is_empty() { connection .del::<_, usize>(keys) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } Ok(true) }) diff --git a/litellm-rust/crates/cache-redis/tests/cache.rs b/litellm-rust/crates/cache-redis/tests/cache.rs index 70baeffd572..9ec85251733 100644 --- a/litellm-rust/crates/cache-redis/tests/cache.rs +++ b/litellm-rust/crates/cache-redis/tests/cache.rs @@ -178,11 +178,11 @@ async fn codec_errors_propagate_without_writing_partial_batches(context: ExactCa ]); assert_eq!( cache.set_cache("invalid", 255, &context), - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(None)) ); assert_eq!( cache.async_set_cache("invalid", 255, context.clone()).await, - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(None)) ); assert_eq!( cache @@ -191,7 +191,7 @@ async fn codec_errors_propagate_without_writing_partial_batches(context: ExactCa context.clone(), ) .await, - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(None)) ); assert_eq!( cache @@ -200,15 +200,15 @@ async fn codec_errors_propagate_without_writing_partial_batches(context: ExactCa ("invalid".into(), 255, None), ]) .await, - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(None)) ); assert_eq!( cache.get_cache("invalid", &context), - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(None)) ); assert_eq!( cache.async_get_cache("invalid", &context).await, - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(None)) ); } @@ -516,7 +516,7 @@ async fn sets_add_members_and_arm_the_default_ttl() { cache .async_set_cache_sadd("members", Vec::new(), None) .await, - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(None)) ); } @@ -547,7 +547,7 @@ async fn queues_push_and_pop_namespaced_lists() { ); assert_eq!( cache.async_rpush("queue", Vec::new()).await, - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(None)) ); assert_eq!( cache.async_lpop("queue", Some(2)).await.unwrap(), @@ -673,15 +673,15 @@ async fn rpush_and_trim_raises_when_a_queued_command_fails() { )]) .with_namespace(Some("ns".into())); - assert_eq!( + assert!(matches!( cache .async_rpush_and_trim("buf", vec!["c".into(), "d".into()], 3) .await, - Err(Error::Unavailable) - ); + Err(Error::Unavailable(Some(_))) + )); assert_eq!( cache.async_rpush_and_trim("buf", Vec::new(), 3).await, - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(None)) ); } @@ -1122,24 +1122,30 @@ async fn test_connection_reports_success_with_the_python_message() { } #[rstest] -#[case::unexpected_reply(Ok("NOPE"), "Redis ping returned False", false)] +#[case::unexpected_reply(Ok("NOPE"), "Redis ping returned False", None)] #[case::connection_refused( Err(redis::RedisError::from((redis::ErrorKind::Io, "connection refused"))), "Redis connection failed:", - true + Some("connection refused") )] #[tokio::test] async fn test_connection_failures_use_the_python_result_contract( #[case] reply: redis::RedisResult<&'static str>, #[case] message: &str, - #[case] has_error: bool, + #[case] cause: Option<&str>, ) { let cache = json_cache(vec![MockCmd::new(redis::cmd("PING"), reply)]); let result = cache.test_connection().await.unwrap(); assert_eq!(result.status, CacheConnectionStatus::Failed); assert!(result.message.starts_with(message), "{}", result.message); - assert_eq!(result.error.is_some(), has_error); + assert_eq!(result.error.is_some(), cause.is_some()); + assert!(cause.is_none_or(|cause| { + result + .error + .as_deref() + .is_some_and(|error| error.contains(cause)) + })); } #[rstest] diff --git a/litellm-rust/crates/cache-redis/tests/cluster.rs b/litellm-rust/crates/cache-redis/tests/cluster.rs index c1a9a70a5d2..e1ba01bcb04 100644 --- a/litellm-rust/crates/cache-redis/tests/cluster.rs +++ b/litellm-rust/crates/cache-redis/tests/cluster.rs @@ -66,7 +66,7 @@ fn constructor_rejects_unusable_cluster_configs( JsonCodec::new(), ) .err(); - assert!(matches!(error, Some(Error::Unavailable))); + assert!(matches!(error, Some(Error::Unavailable(_)))); } #[rstest] diff --git a/litellm-rust/crates/cache-redis/tests/support/mod.rs b/litellm-rust/crates/cache-redis/tests/support/mod.rs index bff1e618582..b04e1c29f3b 100644 --- a/litellm-rust/crates/cache-redis/tests/support/mod.rs +++ b/litellm-rust/crates/cache-redis/tests/support/mod.rs @@ -16,7 +16,7 @@ impl CacheCodec for TaggedByteCodec { fn encode(&self, value: &u8) -> Result, Error> { if *value > 127 { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); } Ok(vec![self.0, *value]) } @@ -24,7 +24,7 @@ impl CacheCodec for TaggedByteCodec { fn decode(&self, bytes: &[u8]) -> Result { match bytes { [tag, value] if *tag == self.0 => Ok(*value), - _ => Err(Error::InvalidEntry), + _ => Err(Error::InvalidEntry(None)), } } } diff --git a/litellm-rust/crates/cache-response/src/buffer.rs b/litellm-rust/crates/cache-response/src/buffer.rs index 68f2348e28b..deb1a56b503 100644 --- a/litellm-rust/crates/cache-response/src/buffer.rs +++ b/litellm-rust/crates/cache-response/src/buffer.rs @@ -26,7 +26,7 @@ impl WriteBuffer { now: Duration, ) -> Result<(), Error> { let pending = { - let mut entries = self.entries.lock().map_err(|_| Error::Unavailable)?; + let mut entries = self.entries.lock()?; entries.push((request.clone(), response, now)); (entries.len() >= self.flush_size).then(|| std::mem::take(&mut *entries)) }; @@ -39,7 +39,7 @@ impl WriteBuffer { } pub fn clear(&self) -> Result<(), Error> { - self.entries.lock().map_err(|_| Error::Unavailable)?.clear(); + self.entries.lock()?.clear(); Ok(()) } } diff --git a/litellm-rust/crates/cache-response/src/codec.rs b/litellm-rust/crates/cache-response/src/codec.rs index 6b0f29e0a58..1adebe65bc3 100644 --- a/litellm-rust/crates/cache-response/src/codec.rs +++ b/litellm-rust/crates/cache-response/src/codec.rs @@ -14,24 +14,24 @@ impl CacheCodec for ResponseCacheCodec { .timestamp .is_some_and(|timestamp| !timestamp.is_finite()) { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); } // Python reads a `response` that is either a dict or a serialized string, so every // other shape is written serialized. A string on the wire is therefore always a // serialized response, which keeps string-valued responses unambiguous. if value.timestamp.is_none() || value.response.is_object() { - return serde_json::to_vec(value).map_err(|_| Error::InvalidEntry); + return serde_json::to_vec(value).map_err(Error::invalid_entry); } - let response = serde_json::to_string(&value.response).map_err(|_| Error::InvalidEntry)?; + let response = serde_json::to_string(&value.response).map_err(Error::invalid_entry)?; serde_json::to_vec(&CacheEntry { timestamp: value.timestamp, response: Value::String(response), }) - .map_err(|_| Error::InvalidEntry) + .map_err(Error::invalid_entry) } fn decode(&self, bytes: &[u8]) -> Result { - let text = std::str::from_utf8(bytes).map_err(|_| Error::InvalidEntry)?; + let text = std::str::from_utf8(bytes).map_err(Error::invalid_entry)?; let value = decode_value(text)?; let Some(timestamp) = value.get("timestamp") else { return Ok(CacheEntry { @@ -40,9 +40,9 @@ impl CacheCodec for ResponseCacheCodec { }); }; let Some(timestamp) = timestamp.as_f64().filter(|timestamp| timestamp.is_finite()) else { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); }; - let response = match value.get("response").ok_or(Error::InvalidEntry)? { + let response = match value.get("response").ok_or(Error::InvalidEntry(None))? { Value::String(text) => decode_value(text)?, response => response.clone(), }; @@ -58,25 +58,25 @@ fn decode_value(text: &str) -> Result { return Ok(value); } check_literal_depth(text)?; - let literal: py_literal::Value = text.parse().map_err(|_| Error::InvalidEntry)?; + let literal: py_literal::Value = text.parse().map_err(Error::invalid_entry)?; literal_value(literal, 0) } fn literal_value(value: py_literal::Value, depth: usize) -> Result { use py_literal::Value as Literal; if depth > 128 { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); } match value { Literal::String(text) => Ok(Value::String(text)), Literal::Boolean(value) => Ok(Value::Bool(value)), Literal::None => Ok(Value::Null), Literal::Integer(value) => { - serde_json::from_str(&value.to_string()).map_err(|_| Error::InvalidEntry) + serde_json::from_str(&value.to_string()).map_err(Error::invalid_entry) } Literal::Float(value) => serde_json::Number::from_f64(value) .map(Value::Number) - .ok_or(Error::InvalidEntry), + .ok_or(Error::InvalidEntry(None)), Literal::List(values) | Literal::Tuple(values) => values .into_iter() .map(|value| literal_value(value, depth + 1)) @@ -86,13 +86,13 @@ fn literal_value(value: py_literal::Value, depth: usize) -> Result .into_iter() .map(|(key, value)| { let Literal::String(key) = key else { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); }; Ok((key, literal_value(value, depth + 1)?)) }) .collect::, _>>() .map(Value::Object), - _ => Err(Error::InvalidEntry), + _ => Err(Error::InvalidEntry(None)), } } @@ -118,7 +118,7 @@ fn check_literal_depth(text: &str) -> Result<(), Error> { '[' | '{' | '(' => { depth += 1; if depth > 128 { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); } } ']' | '}' | ')' => depth = depth.saturating_sub(1), diff --git a/litellm-rust/crates/cache-response/src/response.rs b/litellm-rust/crates/cache-response/src/response.rs index a5bbef99a3a..1505297315f 100644 --- a/litellm-rust/crates/cache-response/src/response.rs +++ b/litellm-rust/crates/cache-response/src/response.rs @@ -2,7 +2,7 @@ use std::{sync::Arc, time::Duration}; use litellm_cache::{ BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, ConnectionCache, Error, - FlushCache, + FlushCache, ShortBatch, semantic::{SemanticCache, SemanticLookup}, }; use serde_json::Value; @@ -114,7 +114,7 @@ where .get_cache(&cache_key(&request.key), &request.context) { Ok(entry) => entry, - Err(Error::InvalidEntry) => None, + Err(Error::InvalidEntry(_)) => None, Err(error) => return Err(error), }; Ok(Self::fresh_or_miss(entry, now, request.max_age)) @@ -134,7 +134,7 @@ where .await { Ok(entry) => entry, - Err(Error::InvalidEntry) => None, + Err(Error::InvalidEntry(_)) => None, Err(error) => return Err(error), }; Ok(Self::fresh_or_miss(entry, now, request.max_age)) @@ -338,7 +338,10 @@ where now: Duration, ) -> Result { if readable.len() != entries.len() { - return Err(Error::Unavailable); + return Err(Error::unavailable(ShortBatch { + requested: readable.len(), + returned: entries.len(), + })); } let mut values = vec![None; requests.len()]; for ((index, request), entry) in readable.into_iter().zip(entries) { @@ -361,7 +364,7 @@ where value: Self::fresh_or_miss(lookup.value, now, max_age), similarity: lookup.similarity, }), - Err(Error::InvalidEntry) => Ok(SemanticLookup::miss(None)), + Err(Error::InvalidEntry(_)) => Ok(SemanticLookup::miss(None)), Err(error) => Err(error), } } diff --git a/litellm-rust/crates/cache-response/tests/codec.rs b/litellm-rust/crates/cache-response/tests/codec.rs index 4fb5a5094cb..289465fcda9 100644 --- a/litellm-rust/crates/cache-response/tests/codec.rs +++ b/litellm-rust/crates/cache-response/tests/codec.rs @@ -44,25 +44,23 @@ fn decode_reads_every_python_envelope(#[case] bytes: &[u8], #[case] response: Va #[case::non_utf8(vec![0xff, 0xfe])] #[case::deep_nesting(format!("{}None{}", "[".repeat(1000), "]".repeat(1000)).into_bytes())] fn decode_rejects_invalid_entries(#[case] bytes: Vec) { - assert_eq!( - ResponseCacheCodec.decode(&bytes).unwrap_err(), - Error::InvalidEntry - ); + assert!(matches!( + ResponseCacheCodec.decode(&bytes), + Err(Error::InvalidEntry(_)) + )); } #[rstest] #[case::nan(f64::NAN)] #[case::infinity(f64::INFINITY)] fn encode_rejects_non_finite_timestamps(#[case] timestamp: f64) { - assert_eq!( - ResponseCacheCodec - .encode(&CacheEntry { - timestamp: Some(timestamp), - response: json!({}), - }) - .unwrap_err(), - Error::InvalidEntry - ); + assert!(matches!( + ResponseCacheCodec.encode(&CacheEntry { + timestamp: Some(timestamp), + response: json!({}), + }), + Err(Error::InvalidEntry(_)) + )); } #[rstest] diff --git a/litellm-rust/crates/cache-response/tests/response.rs b/litellm-rust/crates/cache-response/tests/response.rs index 86ddf4a72cd..70ae259d33a 100644 --- a/litellm-rust/crates/cache-response/tests/response.rs +++ b/litellm-rust/crates/cache-response/tests/response.rs @@ -158,8 +158,16 @@ fn scored(timestamp: f64, similarity: f64) -> Result, )] #[case::no_search(Ok(SemanticLookup::miss(None)), true, Ok(SemanticLookup::miss(None)))] #[case::disabled_reads_skip_the_backend(scored(95.0, 0.95), false, Ok(SemanticLookup::miss(None)))] -#[case::invalid_entry_is_a_miss(Err(Error::InvalidEntry), true, Ok(SemanticLookup::miss(None)))] -#[case::backend_errors_propagate(Err(Error::Unavailable), true, Err(Error::Unavailable))] +#[case::invalid_entry_is_a_miss( + Err(Error::InvalidEntry(None)), + true, + Ok(SemanticLookup::miss(None)) +)] +#[case::backend_errors_propagate( + Err(Error::Unavailable(None)), + true, + Err(Error::Unavailable(None)) +)] #[tokio::test] async fn semantic_lookup_applies_freshness_to_the_value_only( #[case] backend: Result, Error>, diff --git a/litellm-rust/crates/cache-s3/Cargo.toml b/litellm-rust/crates/cache-s3/Cargo.toml index eb3a2fff1ac..d51c282d4d7 100644 --- a/litellm-rust/crates/cache-s3/Cargo.toml +++ b/litellm-rust/crates/cache-s3/Cargo.toml @@ -20,6 +20,7 @@ reqwest.workspace = true tokio.workspace = true [dev-dependencies] +litellm-error.workspace = true litellm-http = { workspace = true, features = ["test-support"] } litellm-cache-testing.workspace = true rstest.workspace = true diff --git a/litellm-rust/crates/cache-s3/src/auth.rs b/litellm-rust/crates/cache-s3/src/auth.rs index f4f06e5371d..659390a5659 100644 --- a/litellm-rust/crates/cache-s3/src/auth.rs +++ b/litellm-rust/crates/cache-s3/src/auth.rs @@ -46,7 +46,7 @@ impl ProvideCredentials for S3Credentials { self.auth .resolve_credentials(self.config.clone(), &self.env) .await - .map_err(|_| CredentialsError::provider_error("S3 cache authentication failed")) + .map_err(CredentialsError::provider_error) }) } } diff --git a/litellm-rust/crates/cache-s3/src/cache.rs b/litellm-rust/crates/cache-s3/src/cache.rs index ced948e80c6..5369bca0f18 100644 --- a/litellm-rust/crates/cache-s3/src/cache.rs +++ b/litellm-rust/crates/cache-s3/src/cache.rs @@ -125,7 +125,7 @@ impl S3Cache { } None => request.cache_control("immutable, max-age=31536000, s-maxage=31536000"), }; - request.send().await.map_err(|_| Error::Unavailable)?; + request.send().await.map_err(Error::unavailable)?; Ok(()) } @@ -152,7 +152,7 @@ impl S3Cache { return Ok(None); } } - return Err(Error::Unavailable); + return Err(Error::unavailable(error)); } }; if let Some(expires) = output.expires_string() @@ -165,7 +165,7 @@ impl S3Cache { .body .collect() .await - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? .into_bytes(); self.codec.decode(&bytes).map(Some) } diff --git a/litellm-rust/crates/cache-s3/src/transport.rs b/litellm-rust/crates/cache-s3/src/transport.rs index eabc54ac9e2..bf1711276f2 100644 --- a/litellm-rust/crates/cache-s3/src/transport.rs +++ b/litellm-rust/crates/cache-s3/src/transport.rs @@ -29,9 +29,10 @@ impl HttpConnector for ReqwestHttpClient { .try_into_http1x() .map_err(|error| ConnectorError::other(error.into(), None))? .map(reqwest::Body::wrap); - let request = reqwest::Request::try_from(request) + let request = litellm_http::outbound::from_http(request) .map_err(|error| ConnectorError::other(error.into(), None))?; let response = client.execute(request).await.map_err(|error| { + let error = litellm_http::Error::from(error); if error.is_timeout() { ConnectorError::timeout(error.into()) } else { diff --git a/litellm-rust/crates/cache-s3/tests/cache.rs b/litellm-rust/crates/cache-s3/tests/cache.rs index 53a51f8a289..d0469620493 100644 --- a/litellm-rust/crates/cache-s3/tests/cache.rs +++ b/litellm-rust/crates/cache-s3/tests/cache.rs @@ -147,15 +147,15 @@ async fn async_set_signs_with_the_configured_keys(#[future(awt)] server: MockSer #[case::malformed( "malformed", ResponseTemplate::new(200).set_body_string("not a cache entry"), - Err(Error::InvalidEntry) + Err("invalid cache entry: ") )] -#[case::server_error("broken", ResponseTemplate::new(500), Err(Error::Unavailable))] +#[case::server_error("broken", ResponseTemplate::new(500), Err("cache is unavailable: "))] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn get_maps_s3_responses( #[future(awt)] server: MockServer, #[case] key: &str, #[case] response: ResponseTemplate, - #[case] expected: Result, Error>, + #[case] expected: Result, &str>, ) { Mock::given(method("GET")) .and(path(format!("/cache-bucket/team/{key}"))) @@ -164,9 +164,21 @@ async fn get_maps_s3_responses( .await; let cache = support::cache(&server.uri()); let context = ExactCacheContext::default(); - - assert_eq!(cache.get_cache(key, &context), expected); - assert_eq!(cache.async_get_cache(key, &context).await, expected); + let outcome = |result: Result, Error>| { + result.map_err(|error| litellm_error::Report::new(&error).to_string()) + }; + let sync = outcome(cache.get_cache(key, &context)); + let asynchronous = outcome(cache.async_get_cache(key, &context).await); + match expected { + Ok(value) => { + assert_eq!(sync, Ok(value.clone())); + assert_eq!(asynchronous, Ok(value)); + } + Err(prefix) => { + assert!(sync.unwrap_err().starts_with(prefix)); + assert!(asynchronous.unwrap_err().starts_with(prefix)); + } + } } #[rstest] diff --git a/litellm-rust/crates/cache-valkey-semantic/Cargo.toml b/litellm-rust/crates/cache-valkey-semantic/Cargo.toml index ec755e61b8c..5b4e99de382 100644 --- a/litellm-rust/crates/cache-valkey-semantic/Cargo.toml +++ b/litellm-rust/crates/cache-valkey-semantic/Cargo.toml @@ -10,9 +10,11 @@ litellm-cache.workspace = true litellm-cache-redis.workspace = true redis = { workspace = true, features = ["tls-rustls"] } sha2.workspace = true +thiserror.workspace = true uuid = { version = "1", features = ["v4"] } [dev-dependencies] +litellm-error.workspace = true litellm-cache-testing.workspace = true redis-test.workspace = true rstest.workspace = true diff --git a/litellm-rust/crates/cache-valkey-semantic/src/index.rs b/litellm-rust/crates/cache-valkey-semantic/src/index.rs index a0b7be4ca7b..3dfff32ff7c 100644 --- a/litellm-rust/crates/cache-valkey-semantic/src/index.rs +++ b/litellm-rust/crates/cache-valkey-semantic/src/index.rs @@ -23,8 +23,7 @@ pub(crate) fn ensure_index( ) -> Result<(), Error> { if index .dimension - .lock() - .map_err(|_| Error::Unavailable)? + .lock()? .is_some_and(|existing| existing == dimension) { return Ok(()); @@ -50,24 +49,42 @@ pub(crate) fn ensure_index( .arg("DISTANCE_METRIC") .arg("COSINE") .query::(connection) - .map(|_| ()) - .map_err(|error| error.to_string()); - if let Err(message) = create { - if !message.to_ascii_lowercase().contains("already exists") { - return Err(Error::Unavailable); + .map(|_| ()); + if let Err(error) = create { + if !index_exists(&error) { + return Err(Error::unavailable(error)); } let info = redis::cmd("FT.INFO") .arg(&index.name) .query::(connection) - .map_err(|_| Error::Unavailable)?; - if index_dimension_from_info(&info).is_some_and(|existing| existing != dimension) { - return Err(Error::Unavailable); + .map_err(Error::unavailable)?; + if let Some(existing) = + index_dimension_from_info(&info).filter(|existing| *existing != dimension) + { + return Err(Error::unavailable(DimensionMismatch { + existing, + dimension, + })); } } - *index.dimension.lock().map_err(|_| Error::Unavailable)? = Some(dimension); + *index.dimension.lock()? = Some(dimension); Ok(()) } +fn index_exists(error: &redis::RedisError) -> bool { + error.code().is_some() + && error + .detail() + .is_some_and(|detail| detail.to_ascii_lowercase().contains("already exists")) +} + +#[derive(Debug, thiserror::Error)] +#[error("index exists with dimension {existing}, not {dimension}")] +struct DimensionMismatch { + existing: usize, + dimension: usize, +} + /// `_extract_index_dim`: flatten each attribute one level and read the value after /// `dimensions`. fn index_dimension_from_info(value: &redis::Value) -> Option { diff --git a/litellm-rust/crates/cache-valkey-semantic/src/search.rs b/litellm-rust/crates/cache-valkey-semantic/src/search.rs index 3b4fae84317..c5e4e748b58 100644 --- a/litellm-rust/crates/cache-valkey-semantic/src/search.rs +++ b/litellm-rust/crates/cache-valkey-semantic/src/search.rs @@ -53,9 +53,7 @@ pub(crate) fn write_document( .arg(ttl.as_secs()) .ignore(); } - pipeline - .query::<()>(connection) - .map_err(|_| Error::Unavailable) + pipeline.query::<()>(connection).map_err(Error::unavailable) } /// The KNN-1 search within `scope`: the closest document's similarity, and its stored response @@ -83,7 +81,7 @@ pub(crate) fn search_document( .arg("DIALECT") .arg(2) .query::(connection) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; let Some(fields) = search_fields(response)? else { return Ok(SemanticLookup::miss(Some(0.0))); }; @@ -91,7 +89,7 @@ pub(crate) fn search_document( fields .iter() .find_map(|(field, value)| (field == name).then(|| value.clone())) - .ok_or(Error::InvalidEntry) + .ok_or(Error::InvalidEntry(None)) }; let response = field("response")?; let similarity = 1.0 - parse_f64(&field("vector_distance")?)?; @@ -105,24 +103,24 @@ type SearchFields = Vec<(String, Vec)>; fn search_fields(value: redis::Value) -> Result, Error> { let redis::Value::Array(values) = value else { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); }; - let total = parse_i64(values.first().ok_or(Error::InvalidEntry)?)?; + let total = parse_i64(values.first().ok_or(Error::InvalidEntry(None))?)?; if total <= 0 || values.len() < 3 { return Ok(None); } let redis::Value::Array(fields) = &values[2] else { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); }; let (pairs, remainder) = fields.as_chunks::<2>(); if !remainder.is_empty() { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); } pairs .iter() .map(|pair| { Ok(( - value_text(&pair[0]).ok_or(Error::InvalidEntry)?, + value_text(&pair[0]).ok_or(Error::InvalidEntry(None))?, value_bytes(&pair[1])?, )) }) @@ -132,16 +130,16 @@ fn search_fields(value: redis::Value) -> Result, Error> { fn parse_i64(value: &redis::Value) -> Result { value_text(value) - .ok_or(Error::InvalidEntry)? + .ok_or(Error::InvalidEntry(None))? .parse() - .map_err(|_| Error::InvalidEntry) + .map_err(Error::invalid_entry) } fn parse_f64(value: &[u8]) -> Result { std::str::from_utf8(value) - .map_err(|_| Error::InvalidEntry)? + .map_err(Error::invalid_entry)? .parse() - .map_err(|_| Error::InvalidEntry) + .map_err(Error::invalid_entry) } pub(crate) fn value_text(value: &redis::Value) -> Option { @@ -158,6 +156,6 @@ fn value_bytes(value: &redis::Value) -> Result, Error> { redis::Value::BulkString(bytes) => Ok(bytes.clone()), redis::Value::SimpleString(value) => Ok(value.as_bytes().to_vec()), redis::Value::Int(value) => Ok(value.to_string().into_bytes()), - _ => Err(Error::InvalidEntry), + _ => Err(Error::InvalidEntry(None)), } } diff --git a/litellm-rust/crates/cache-valkey-semantic/tests/cache.rs b/litellm-rust/crates/cache-valkey-semantic/tests/cache.rs index 1c693d0875d..10e32c0183b 100644 --- a/litellm-rust/crates/cache-valkey-semantic/tests/cache.rs +++ b/litellm-rust/crates/cache-valkey-semantic/tests/cache.rs @@ -83,10 +83,7 @@ fn ok() -> redis::RedisResult { } fn already_exists() -> redis::RedisResult { - Err(redis::RedisError::from(( - redis::ErrorKind::Io, - "already exists", - ))) + Err(support::server_error("Index already exists")) } fn bulk(value: &[u8]) -> redis::Value { @@ -250,19 +247,29 @@ fn second_set_skips_create_after_dimension_is_cached(entry: Value, context: Sema #[rstest] #[case::nested_matching(nested_dimension_info(3), Ok(()))] -#[case::nested_mismatch(nested_dimension_info(2), Err(Error::Unavailable))] +#[case::nested_mismatch( + nested_dimension_info(2), + Err("cache is unavailable: index exists with dimension 2, not 3") +)] #[case::flat_matching(flat_dimension_info("3"), Ok(()))] -#[case::flat_mismatch(flat_dimension_info("2"), Err(Error::Unavailable))] +#[case::flat_mismatch( + flat_dimension_info("2"), + Err("cache is unavailable: index exists with dimension 2, not 3") +)] #[case::unreported_dimension_is_accepted(redis::Value::Array(vec![]), Ok(()))] fn existing_index_dimension_must_match_embedding( #[case] info: redis::Value, - #[case] expected: Result<(), Error>, + #[case] expected: Result<(), &str>, entry: Value, context: SemanticCacheContext, ) { let recording = recording([already_exists(), Ok(info)]); - assert_eq!(recording.cache.set_cache("key", entry, &context), expected); + let result = recording.cache.set_cache("key", entry, &context); + assert_eq!( + result.map_err(|error| litellm_error::Report::new(&error).to_string()), + expected.map_err(str::to_owned) + ); } #[rstest] @@ -274,10 +281,10 @@ fn index_creation_failures_are_unavailable( ) { let recording = recording([reply]); - assert_eq!( + assert!(matches!( recording.cache.set_cache("key", entry, &context), - Err(Error::Unavailable) - ); + Err(Error::Unavailable(Some(_))) + )); } #[rstest] @@ -291,10 +298,10 @@ fn index_creation_failures_are_unavailable( bulk(b"document"), redis::Value::Array(vec![bulk(b"vector_distance"), bulk(b"0.1")]), ]), - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(None)) )] -#[case::unparsable_distance(search_hit(b"not-json", "abc"), Err(Error::InvalidEntry))] -#[case::undecodable_response(search_hit(b"not-json", "0.1"), Err(Error::InvalidEntry))] +#[case::unparsable_distance(search_hit(b"not-json", "abc"), Err(Error::InvalidEntry(None)))] +#[case::undecodable_response(search_hit(b"not-json", "0.1"), Err(Error::InvalidEntry(None)))] fn get_applies_threshold_and_decodes_entry( #[case] reply: redis::Value, #[case] expected: Result, Error>, @@ -302,7 +309,13 @@ fn get_applies_threshold_and_decodes_entry( ) { let recording = recording([ok(), Ok(reply)]); - assert_eq!(recording.cache.get_cache("key", &context), expected); + let kind = |result: Result, Error>| { + result.map_err(|error| std::mem::discriminant(&error)) + }; + assert_eq!( + kind(recording.cache.get_cache("key", &context)), + kind(expected) + ); } #[rstest] diff --git a/litellm-rust/crates/cache-valkey-semantic/tests/support/mod.rs b/litellm-rust/crates/cache-valkey-semantic/tests/support/mod.rs index 2c6ff06b134..2af05943c98 100644 --- a/litellm-rust/crates/cache-valkey-semantic/tests/support/mod.rs +++ b/litellm-rust/crates/cache-valkey-semantic/tests/support/mod.rs @@ -74,7 +74,7 @@ impl FakeSearch { "FT.CREATE" => { let name = text(1); if state.indexes.contains_key(&name) { - return Err(error("Index already exists")); + return Err(server_error("Index already exists")); } let position = |token: &str| args.iter().position(|arg| arg == token.as_bytes()); let prefix = args[position("PREFIX").unwrap() + 2].clone(); @@ -94,7 +94,7 @@ impl FakeSearch { let index = state .indexes .get(&text(1)) - .ok_or_else(|| error("Unknown index name"))?; + .ok_or_else(|| server_error("Unknown index name"))?; Ok(index_info(index)) } "FT.DROPINDEX" => { @@ -118,7 +118,7 @@ impl FakeSearch { let index = state .indexes .get(&text(1)) - .ok_or_else(|| error("no such index"))?; + .ok_or_else(|| server_error("no such index"))?; let query = text(2); let tag = query_tag(&query); let params = args.iter().position(|arg| arg == b"PARAMS").unwrap(); @@ -155,7 +155,7 @@ impl FakeSearch { ])) } "PING" => Ok(redis::Value::SimpleString("PONG".into())), - _ => Err(error("unsupported command")), + _ => Err(server_error("unsupported command")), } } } @@ -192,8 +192,11 @@ impl redis::ConnectionLike for FakeSearch { } } -fn error(message: &'static str) -> redis::RedisError { - redis::RedisError::from((redis::ErrorKind::Extension, message)) +pub fn server_error(reply: &str) -> redis::RedisError { + redis::parse_redis_value(format!("-{reply}\r\n").as_bytes()) + .expect("error reply") + .extract_error() + .expect_err("server error") } fn bulk(bytes: &[u8]) -> redis::Value { diff --git a/litellm-rust/crates/cache/Cargo.toml b/litellm-rust/crates/cache/Cargo.toml index 51bbe12a85e..282c34e6b02 100644 --- a/litellm-rust/crates/cache/Cargo.toml +++ b/litellm-rust/crates/cache/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-error.workspace = true serde.workspace = true serde_json = { workspace = true, features = ["preserve_order"] } strum.workspace = true diff --git a/litellm-rust/crates/cache/src/capabilities.rs b/litellm-rust/crates/cache/src/capabilities.rs index ac9d8fc8764..b30f9e162f2 100644 --- a/litellm-rust/crates/cache/src/capabilities.rs +++ b/litellm-rust/crates/cache/src/capabilities.rs @@ -48,7 +48,7 @@ pub trait BatchCache: BaseCache { .map(|key| match self.get_cache(key, context) { Ok(Some(value)) => Ok(BatchEntry::Hit(value)), Ok(None) => Ok(BatchEntry::Miss), - Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid), + Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid), Err(error) => Err(error), }) .collect() @@ -65,7 +65,7 @@ pub trait BatchCache: BaseCache { entries.push(match self.async_get_cache(&key, &context).await { Ok(Some(value)) => BatchEntry::Hit(value), Ok(None) => BatchEntry::Miss, - Err(Error::InvalidEntry) => BatchEntry::Invalid, + Err(Error::InvalidEntry(_)) => BatchEntry::Invalid, Err(error) => return Err(error), }); } diff --git a/litellm-rust/crates/cache/src/codec.rs b/litellm-rust/crates/cache/src/codec.rs index 6d47c682406..b802e40a8b3 100644 --- a/litellm-rust/crates/cache/src/codec.rs +++ b/litellm-rust/crates/cache/src/codec.rs @@ -41,10 +41,10 @@ where type Value = V; fn encode(&self, value: &Self::Value) -> Result, Error> { - serde_json::to_vec(value).map_err(|_| Error::InvalidEntry) + serde_json::to_vec(value).map_err(Error::invalid_entry) } fn decode(&self, bytes: &[u8]) -> Result { - serde_json::from_slice(bytes).map_err(|_| Error::InvalidEntry) + serde_json::from_slice(bytes).map_err(Error::invalid_entry) } } diff --git a/litellm-rust/crates/cache/src/dual.rs b/litellm-rust/crates/cache/src/dual.rs index c4bc953239d..3deee2560f4 100644 --- a/litellm-rust/crates/cache/src/dual.rs +++ b/litellm-rust/crates/cache/src/dual.rs @@ -2,7 +2,7 @@ use std::{sync::Arc, time::Duration}; use crate::{ BaseCache, BatchCache, BatchEntry, BulkDeleteCache, CacheContext, ClaimCache, CounterCache, - DeleteCache, Error, FlushCache, IncrementOperation, SetCache, TtlCache, + DeleteCache, Error, FlushCache, IncrementOperation, SetCache, ShortBatch, TtlCache, }; #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] @@ -98,7 +98,7 @@ impl DualCache { fn remote(&self, result: Result) -> Result, Error> { match result { Ok(value) => Ok(Some(value)), - Err(Error::Unavailable) + Err(Error::Unavailable(_)) if self.remote_failure_policy == RemoteFailurePolicy::UseLocal => { Ok(None) @@ -136,7 +136,10 @@ where remote: Vec>, ) -> Result>, Error> { if missing.len() != remote.len() { - return Err(Error::Unavailable); + return Err(Error::unavailable(ShortBatch { + requested: missing.len(), + returned: remote.len(), + })); } for (index, entry) in missing.into_iter().zip(remote) { if let BatchEntry::Hit(value) = &entry { @@ -389,7 +392,10 @@ where } let values = self.l2.async_increment_pipeline(operations.clone()).await?; if values.len() != operations.len() { - return Err(Error::Unavailable); + return Err(Error::unavailable(ShortBatch { + requested: operations.len(), + returned: values.len(), + })); } for (operation, value) in operations.iter().zip(&values) { self.l1 diff --git a/litellm-rust/crates/cache/src/error.rs b/litellm-rust/crates/cache/src/error.rs index 79ef9cd18b1..8861651a1c8 100644 --- a/litellm-rust/crates/cache/src/error.rs +++ b/litellm-rust/crates/cache/src/error.rs @@ -1,9 +1,11 @@ +use litellm_error::Source; + #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub enum Error { #[error("cache is unavailable")] - Unavailable, + Unavailable(#[source] Option), #[error("invalid cache entry")] - InvalidEntry, + InvalidEntry(#[source] Option), #[error("flushing Redis requires an explicit namespace")] UnscopedFlush, #[error("operation is not supported by this cache")] @@ -11,3 +13,59 @@ pub enum Error { #[error("semantic cache requires request messages")] MissingPrompt, } + +#[derive(Debug, thiserror::Error)] +#[error("cache backend returned {returned} results for {requested} keys")] +pub struct ShortBatch { + pub requested: usize, + pub returned: usize, +} + +#[derive(Debug, thiserror::Error)] +#[error("cache lock is poisoned")] +struct Poisoned; + +impl From> for Error { + fn from(_: std::sync::PoisonError) -> Self { + Self::unavailable(Poisoned) + } +} + +impl Error { + pub fn unavailable(cause: impl std::error::Error + Send + Sync + 'static) -> Self { + Self::Unavailable(Some(Source::new(cause))) + } + + pub fn invalid_entry(cause: impl std::error::Error + Send + Sync + 'static) -> Self { + Self::InvalidEntry(Some(Source::new(cause))) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::Error; + + #[rstest] + #[case::unavailable( + Error::unavailable(std::io::Error::other("connection refused")), + "cache is unavailable" + )] + #[case::invalid_entry( + Error::invalid_entry(std::io::Error::other("connection refused")), + "invalid cache entry" + )] + fn the_backend_cause_is_the_source_and_rendered_once( + #[case] error: Error, + #[case] message: &str, + ) { + let report = litellm_error::Report::new(&error).to_string(); + assert_eq!(report, format!("{message}: connection refused")); + assert!(!error.to_string().contains("connection refused")); + let cause = std::error::Error::source(&error) + .and_then(|cause| cause.downcast_ref::()) + .expect("cause"); + assert!(cause.downcast_ref::().is_some()); + } +} diff --git a/litellm-rust/crates/cache/src/lib.rs b/litellm-rust/crates/cache/src/lib.rs index 55d8a5fa28b..4e07f3b4e8c 100644 --- a/litellm-rust/crates/cache/src/lib.rs +++ b/litellm-rust/crates/cache/src/lib.rs @@ -23,4 +23,4 @@ pub use codec::{CacheCodec, JsonCodec}; pub use dual::{ DEFAULT_DELETE_BATCH_SIZE, DualCache, ReadPolicy, RemoteFailurePolicy, WritePolicy, }; -pub use error::Error; +pub use error::{Error, ShortBatch}; diff --git a/litellm-rust/crates/cache/tests/caching.rs b/litellm-rust/crates/cache/tests/caching.rs index baf968f4659..84a3b8c4d2a 100644 --- a/litellm-rust/crates/cache/tests/caching.rs +++ b/litellm-rust/crates/cache/tests/caching.rs @@ -67,7 +67,7 @@ impl BaseCache for TestCache { } fn set_cache(&self, _: &str, _: Self::Value, _: &ExactCacheContext) -> Result<(), Error> { - Err(Error::Unavailable) + Err(Error::Unavailable(None)) } async fn async_set_cache( @@ -77,7 +77,7 @@ impl BaseCache for TestCache { context: ExactCacheContext, ) -> Result<(), Error> { if key == "unavailable" { - return Err(Error::Unavailable); + return Err(Error::Unavailable(None)); } self.writes .lock() @@ -118,7 +118,7 @@ impl BaseCache for RecordingCounter { impl CounterCache for RecordingCounter { fn increment_cache(&self, key: &str, amount: f64, _: ExactCacheContext) -> Result { if key == "unavailable" { - return Err(Error::Unavailable); + return Err(Error::Unavailable(None)); } let mut total = self.total.lock().unwrap(); *total += amount; @@ -226,7 +226,7 @@ async fn default_batch_operations_use_async_writes_and_stop_on_failure(cache: Te context.clone(), ) .await, - Err(Error::Unavailable) + Err(Error::Unavailable(None)) ); assert_eq!( *cache.writes.lock().unwrap(), @@ -323,7 +323,7 @@ async fn default_increment_pipeline_stops_at_the_first_failure() { operation("skipped", 1.0, None), ]) .await, - Err(Error::Unavailable) + Err(Error::Unavailable(None)) ); let keys = cache .increments diff --git a/litellm-rust/crates/cache/tests/codec.rs b/litellm-rust/crates/cache/tests/codec.rs index 5320545fd49..ff4aeed87e6 100644 --- a/litellm-rust/crates/cache/tests/codec.rs +++ b/litellm-rust/crates/cache/tests/codec.rs @@ -31,12 +31,18 @@ fn json_codec_round_trips_typed_domain_values() { #[case::wrongly_typed(br#"{"deployment":12}"#)] fn json_codec_rejects_malformed_and_wrongly_typed_entries(#[case] bytes: &[u8]) { let codec = JsonCodec::::new(); - assert_eq!(codec.decode(bytes).unwrap_err(), Error::InvalidEntry); + assert!(matches!( + codec.decode(bytes), + Err(Error::InvalidEntry(Some(_))) + )); } #[rstest] fn json_codec_propagates_encoding_errors() { let codec = JsonCodec::>::new(); let value = BTreeMap::from([((1, 2), "invalid JSON object key".into())]); - assert_eq!(codec.encode(&value).unwrap_err(), Error::InvalidEntry); + assert!(matches!( + codec.encode(&value), + Err(Error::InvalidEntry(Some(_))) + )); } diff --git a/litellm-rust/crates/cache/tests/dual.rs b/litellm-rust/crates/cache/tests/dual.rs index 87dd76388ef..86af5e80abb 100644 --- a/litellm-rust/crates/cache/tests/dual.rs +++ b/litellm-rust/crates/cache/tests/dual.rs @@ -71,7 +71,7 @@ where impl CounterCache for TestCache { fn increment_cache(&self, _: &str, amount: f64, _: ExactCacheContext) -> Result { if self.fail { - return Err(Error::Unavailable); + return Err(Error::Unavailable(None)); } let mut value = self.value.lock().unwrap(); let incremented = value.unwrap_or_default() + amount; @@ -92,7 +92,7 @@ where _: ExactCacheContext, ) -> Result { if self.fail { - return Err(Error::Unavailable); + return Err(Error::Unavailable(None)); } let mut value = self.value.lock().unwrap(); let winner = match value.as_ref() { @@ -113,7 +113,7 @@ fn failed_l2_increment_leaves_l1_unchanged() { assert_eq!( cache.increment_cache("counter", 2.0, ExactCacheContext::default()), - Err(Error::Unavailable) + Err(Error::Unavailable(None)) ); assert_eq!( l1.get_cache("counter", &ExactCacheContext::default()) @@ -276,11 +276,11 @@ impl BaseCache for Unavailable { } fn set_cache(&self, _: &str, _: String, _: &ExactCacheContext) -> Result<(), Error> { - Err(Error::Unavailable) + Err(Error::Unavailable(None)) } fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result, Error> { - Err(Error::Unavailable) + Err(Error::Unavailable(None)) } } @@ -288,13 +288,13 @@ impl BatchCache for Unavailable {} impl DeleteCache for Unavailable { fn delete_cache(&self, _: &str) -> Result<(), Error> { - Err(Error::Unavailable) + Err(Error::Unavailable(None)) } } impl FlushCache for Unavailable { fn flush_cache(&self) -> Result<(), Error> { - Err(Error::Unavailable) + Err(Error::Unavailable(None)) } } @@ -306,7 +306,7 @@ impl ClaimCache for Unavailable { _: &[String], _: ExactCacheContext, ) -> Result { - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(None)) } } @@ -316,9 +316,12 @@ fn remote_failure_policy_selects_propagation_or_the_local_tier() { let strict = DualCache::new(Arc::new(TestCache::new(None, false)), Arc::new(Unavailable)); assert_eq!( strict.set_cache("key", "value".into(), &context), - Err(Error::Unavailable) + Err(Error::Unavailable(None)) + ); + assert_eq!( + strict.get_cache("key", &context), + Err(Error::Unavailable(None)) ); - assert_eq!(strict.get_cache("key", &context), Err(Error::Unavailable)); let l1 = Arc::new(TestCache::new(None, false)); let degraded = DualCache::new(l1.clone(), Arc::new(Unavailable)) @@ -347,7 +350,7 @@ fn claim_fallback_does_not_hide_non_availability_errors() { &[], ExactCacheContext::default() ), - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(None)) ); } @@ -419,7 +422,7 @@ impl Tier { fn check(&self) -> Result<(), Error> { if self.fail { - return Err(Error::Unavailable); + return Err(Error::Unavailable(None)); } Ok(()) } @@ -598,7 +601,7 @@ async fn failed_async_l2_increment_leaves_l1_unchanged(log: Log) { cache .async_increment("counter", 2.0, seconds(30), true) .await, - Err(Error::Unavailable) + Err(Error::Unavailable(None)) ); assert_eq!( l1.counter("counter"), @@ -669,7 +672,7 @@ async fn failed_l2_increment_pipeline_leaves_l1_unchanged(log: Log) { ttl: None, }]) .await, - Err(Error::Unavailable) + Err(Error::Unavailable(None)) ); assert_eq!(events(&log), ["l2 pipeline a"]); assert_eq!(l1.counter("a"), Some((1.0, ExactCacheContext::default()))); @@ -713,7 +716,7 @@ async fn failed_local_set_add_never_reaches_l2(log: Log) { cache .async_set_cache_sadd("members", vec!["a".into()], None) .await, - Err(Error::Unavailable) + Err(Error::Unavailable(None)) ); assert!(events(&log).is_empty()); } @@ -768,7 +771,7 @@ async fn bulk_delete_stops_before_l2_when_the_local_delete_fails(log: Log) { ); assert_eq!( cache.delete_cache_keys(vec!["a".into()]).await, - Err(Error::Unavailable) + Err(Error::Unavailable(None)) ); assert!(events(&log).is_empty()); } diff --git a/litellm-rust/crates/callbacks-legacy-python/Cargo.toml b/litellm-rust/crates/callbacks-legacy-python/Cargo.toml index ed5e0fb9691..1cd05710366 100644 --- a/litellm-rust/crates/callbacks-legacy-python/Cargo.toml +++ b/litellm-rust/crates/callbacks-legacy-python/Cargo.toml @@ -6,12 +6,14 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-error.workspace = true litellm-host.workspace = true litellm-host-python.workspace = true pyo3.workspace = true strum.workspace = true serde_json.workspace = true +thiserror.workspace = true [dev-dependencies] litellm-auth.workspace = true diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index d891b5a1f0c..493e91bb5ea 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -134,9 +134,10 @@ impl LegacyLogging { } fn logger(&self) -> PyResult<&PythonLogger> { - self.logger.as_ref().ok_or_else(|| { - pyo3::exceptions::PyRuntimeError::new_err("call logging is not initialized") - }) + Ok(self + .logger + .as_ref() + .ok_or(crate::error::Error::LoggingNotInitialized)?) } /// The keyword view the rest of the call reads: a copy, so the deployment hook's own diff --git a/litellm-rust/crates/callbacks-legacy-python/src/error.rs b/litellm-rust/crates/callbacks-legacy-python/src/error.rs new file mode 100644 index 00000000000..97ce62ae99c --- /dev/null +++ b/litellm-rust/crates/callbacks-legacy-python/src/error.rs @@ -0,0 +1,17 @@ +use litellm_error::Report; +use pyo3::{PyErr, exceptions::PyRuntimeError}; + +#[derive(Debug, thiserror::Error)] +pub(crate) enum Error { + #[error("call logging is not initialized")] + LoggingNotInitialized, +} + +impl From for PyErr { + fn from(error: Error) -> Self { + let message = Report::new(&error).to_string(); + match error { + Error::LoggingNotInitialized => PyRuntimeError::new_err(message), + } + } +} diff --git a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs index 38c6b1aedbd..f453237b469 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs @@ -11,6 +11,7 @@ mod adapter; mod call; mod callbacks; mod deferred; +mod error; mod logger; mod mapping; mod python; diff --git a/litellm-rust/crates/config/Cargo.toml b/litellm-rust/crates/config/Cargo.toml index 9746a002bb9..5b0cd22332f 100644 --- a/litellm-rust/crates/config/Cargo.toml +++ b/litellm-rust/crates/config/Cargo.toml @@ -16,6 +16,7 @@ thiserror.workspace = true tracing.workspace = true [dev-dependencies] +litellm-error.workspace = true rstest.workspace = true tempfile.workspace = true tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] } diff --git a/litellm-rust/crates/config/src/error.rs b/litellm-rust/crates/config/src/error.rs index 95c33967cf0..e5219ac992c 100644 --- a/litellm-rust/crates/config/src/error.rs +++ b/litellm-rust/crates/config/src/error.rs @@ -1,9 +1,16 @@ +use std::{io, path::PathBuf}; + #[derive(Debug, thiserror::Error)] pub enum Error { - #[error("could not read config")] - Read(#[from] std::io::Error), + #[error("could not read config {}", path.display())] + Read { path: PathBuf, source: io::Error }, #[error("invalid YAML config")] Parse(#[from] serde_yaml_ng::Error), + #[error("invalid YAML config {}", path.display())] + ParseFile { + path: PathBuf, + source: serde_yaml_ng::Error, + }, #[error("Redis host does not yet support configured setting {0}")] RedisSetting(String), #[error("invalid Redis setting {setting}: {reason}")] diff --git a/litellm-rust/crates/config/src/includes.rs b/litellm-rust/crates/config/src/includes.rs index 78df8e9a33a..27b908650df 100644 --- a/litellm-rust/crates/config/src/includes.rs +++ b/litellm-rust/crates/config/src/includes.rs @@ -15,11 +15,31 @@ struct Includes { } fn read(path: &Path) -> Result { - Ok(serde_yaml_ng::from_str(&std::fs::read_to_string(path)?)?) + let text = std::fs::read_to_string(path).map_err(|source| Error::Read { + path: path.to_owned(), + source, + })?; + serde_yaml_ng::from_str(&text).map_err(|source| Error::ParseFile { + path: path.to_owned(), + source, + }) +} + +fn canonical(path: &Path) -> Result { + path.canonicalize().map_err(|source| Error::Read { + path: path.to_owned(), + source, + }) } fn entries(config: &Mapping, path: &Path) -> Result, Error> { - let includes: Includes = serde_yaml_ng::from_value(Value::Mapping(config.clone()))?; + let includes: Includes = + serde_yaml_ng::from_value(Value::Mapping(config.clone())).map_err(|source| { + Error::ParseFile { + path: path.to_owned(), + source, + } + })?; Ok(includes .include .into_iter() @@ -28,7 +48,7 @@ fn entries(config: &Mapping, path: &Path) -> Result, Erro } pub(super) fn load(path: &Path) -> Result { - let root = path.canonicalize()?; + let root = canonical(path)?; let mut merged = read(&root)?; let mut pending: VecDeque<_> = entries(&merged, &root)?.into(); let mut loaded = BTreeSet::from([root.clone()]); @@ -36,12 +56,11 @@ pub(super) fn load(path: &Path) -> Result { while let Some((entry, declaring)) = pending.pop_front() { let declared = declaring.parent().unwrap_or(Path::new(".")).join(&entry); let fallback = root.parent().unwrap_or(Path::new(".")).join(&entry); - let location = if declared.exists() { - declared + let location = canonical(if declared.exists() { + &declared } else { - fallback - } - .canonicalize()?; + &fallback + })?; if !loaded.insert(location.clone()) { continue; } diff --git a/litellm-rust/crates/config/tests/config.rs b/litellm-rust/crates/config/tests/config.rs index d425b62efc2..3c7d02203e5 100644 --- a/litellm-rust/crates/config/tests/config.rs +++ b/litellm-rust/crates/config/tests/config.rs @@ -84,9 +84,15 @@ fn rejects_malformed_and_incomplete_config(#[case] yaml: &str) { #[rstest] fn distinguishes_read_errors_from_parse_errors(directory: TempDir) { + let error = Config::load(directory.path().join("missing.yaml")).unwrap_err(); + assert!( + litellm_error::Report::new(&error) + .to_string() + .contains("missing.yaml") + ); assert!(matches!( - Config::load(directory.path().join("missing.yaml")), - Err(Error::Read(error)) if error.kind() == std::io::ErrorKind::NotFound + error, + Error::Read { source, .. } if source.kind() == std::io::ErrorKind::NotFound )); } diff --git a/litellm-rust/crates/core-utils/Cargo.toml b/litellm-rust/crates/core-utils/Cargo.toml index 1feea5fcc08..2e5857c79f3 100644 --- a/litellm-rust/crates/core-utils/Cargo.toml +++ b/litellm-rust/crates/core-utils/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true [dependencies] fancy-regex.workspace = true +litellm-error.workspace = true litellm-tracing.workspace = true litellm-llms-types.workspace = true serde.workspace = true diff --git a/litellm-rust/crates/core-utils/src/call_arguments.rs b/litellm-rust/crates/core-utils/src/call_arguments.rs index 062dad0a92a..3277ca26c7f 100644 --- a/litellm-rust/crates/core-utils/src/call_arguments.rs +++ b/litellm-rust/crates/core-utils/src/call_arguments.rs @@ -1,5 +1,7 @@ use std::ops::Deref; +use litellm_error::Source; + use serde::{Deserialize, Serialize, de::DeserializeOwned}; use serde_json::{Map, Value}; @@ -17,9 +19,11 @@ impl CallArguments { } #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] -#[error("invalid argument: {path}")] +#[error("invalid argument {path}")] pub struct ArgumentError { pub path: String, + #[source] + source: Source, } pub fn parse_options(arguments: &CallArguments) -> Result { @@ -28,6 +32,7 @@ pub fn parse_options(arguments: &CallArguments) -> Result( body: &B, consumed: &[&str], ) -> Result { - let Value::Object(fields) = - serde_json::to_value(body).map_err(|_| crate::params::Error::Body)? + let Value::Object(fields) = serde_json::to_value(body) + .map_err(|error| crate::params::Error::Serialize(Source::new(error)))? else { return Err(crate::params::Error::Body); }; @@ -183,9 +188,9 @@ mod tests { assert_eq!(arguments.get("enabled"), Some(&Value::Null)); assert_eq!(arguments.get("missing"), None); let invalid = serde_json::from_value(json!({"enabled":0})).unwrap(); - assert_eq!( - parse_options::(&invalid).err().unwrap().path, - "enabled" - ); + let error = parse_options::(&invalid).err().unwrap(); + assert_eq!(error.path, "enabled"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.starts_with("invalid argument enabled: invalid type")); } } diff --git a/litellm-rust/crates/core-utils/src/exception_mapping_utils/mod.rs b/litellm-rust/crates/core-utils/src/exception_mapping_utils/mod.rs index 8b63fadf5a6..c99bbf78e02 100644 --- a/litellm-rust/crates/core-utils/src/exception_mapping_utils/mod.rs +++ b/litellm-rust/crates/core-utils/src/exception_mapping_utils/mod.rs @@ -337,15 +337,15 @@ mod tests { #[case::request_timed_out("Request timed out")] #[case::timed_out_generating("Timed out generating response")] #[case::read_operation("The read operation timed out")] - fn timeout_markers_win_over_every_family(#[case] marker: &str) { + fn timeout_markers_win_over_every_family( + #[case] marker: &str, + #[values("mistral", "vertex_ai", "cohere", "reducto")] provider: &str, + ) { let body = format!("rate limit invalid api token {marker}"); - for provider in ["mistral", "vertex_ai", "cohere", "reducto"] { - assert_eq!( - mapped(provider, &http(429, &body)).error, - PublicError::Timeout { status: 408 }, - "{provider}" - ); - } + assert_eq!( + mapped(provider, &http(429, &body)).error, + PublicError::Timeout { status: 408 } + ); } #[test] diff --git a/litellm-rust/crates/core-utils/src/params.rs b/litellm-rust/crates/core-utils/src/params.rs index 9545a3ef17b..f163399b878 100644 --- a/litellm-rust/crates/core-utils/src/params.rs +++ b/litellm-rust/crates/core-utils/src/params.rs @@ -4,6 +4,8 @@ pub enum Error { ExtraBody, #[error("invalid request: body must be a JSON object")] Body, + #[error("invalid request: body could not be serialized")] + Serialize(#[source] litellm_error::Source), } use std::ops::{Deref, DerefMut}; diff --git a/litellm-rust/crates/core-utils/src/url_utils.rs b/litellm-rust/crates/core-utils/src/url_utils.rs index ebbe7fca00b..63e4f041a72 100644 --- a/litellm-rust/crates/core-utils/src/url_utils.rs +++ b/litellm-rust/crates/core-utils/src/url_utils.rs @@ -4,7 +4,7 @@ use url::Url; #[derive(Debug, thiserror::Error)] pub enum ApiUrlError { - #[error("invalid URL: {0}")] + #[error("invalid URL")] Parse(#[from] url::ParseError), #[error("URL cannot be used as a base")] CannotBeBase, diff --git a/litellm-rust/crates/cost/Cargo.toml b/litellm-rust/crates/cost/Cargo.toml index 3ed27e82c9a..6921e6151b5 100644 --- a/litellm-rust/crates/cost/Cargo.toml +++ b/litellm-rust/crates/cost/Cargo.toml @@ -5,6 +5,9 @@ edition.workspace = true license.workspace = true repository.workspace = true +[dependencies] +thiserror.workspace = true + [dev-dependencies] criterion.workspace = true proptest.workspace = true diff --git a/litellm-rust/crates/cost/src/error.rs b/litellm-rust/crates/cost/src/error.rs new file mode 100644 index 00000000000..1753080617f --- /dev/null +++ b/litellm-rust/crates/cost/src/error.rs @@ -0,0 +1,27 @@ +#[derive(Clone, Copy, Debug, Eq, PartialEq, thiserror::Error)] +pub enum Error { + #[error("model has no input rate")] + MissingInputRate, + #[error("model has no output rate")] + MissingOutputRate, + #[error("rate must be finite and nonnegative")] + InvalidRate, + #[error("region multiplier must be finite and positive")] + InvalidRegionMultiplier, + #[error("off-peak pricing needs a billing minute within the day")] + InvalidBillingTime, + #[error("off-peak window must start before it ends within the day")] + InvalidOffPeakWindow, + #[error("cached tokens exceed prompt tokens")] + CacheExceedsPrompt, + #[error("cache write details do not add up to the cache write tokens")] + InvalidCacheWriteDetails, + #[error("token count or cost overflows")] + TokenCountOverflow, + #[error("service tier is repeated or cannot carry its own rates")] + DuplicateTier, + #[error("prompt token threshold is repeated")] + DuplicateThreshold, + #[error("service tier is repeated or cannot carry its own rates within a threshold")] + DuplicateThresholdTier, +} diff --git a/litellm-rust/crates/cost/src/lib.rs b/litellm-rust/crates/cost/src/lib.rs index 32cc755ce15..24c69afd9d4 100644 --- a/litellm-rust/crates/cost/src/lib.rs +++ b/litellm-rust/crates/cost/src/lib.rs @@ -1,3 +1,7 @@ +mod error; + +pub use error::Error; + #[derive(Clone, Copy, Debug, PartialEq)] pub enum Rate { Missing, @@ -156,22 +160,6 @@ pub struct EffectiveRates { pub cache_write_1h: f64, } -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub enum PricingError { - MissingInputRate, - MissingOutputRate, - InvalidRate, - InvalidRegionMultiplier, - InvalidBillingTime, - InvalidOffPeakWindow, - CacheExceedsPrompt, - InvalidCacheWriteDetails, - TokenCountOverflow, - DuplicateTier, - DuplicateThreshold, - DuplicateThresholdTier, -} - fn selected_tier(tier: ServiceTier) -> ServiceTier { if tier == ServiceTier::Fast { ServiceTier::Priority @@ -210,9 +198,9 @@ fn valid_rates(rates: Rates) -> bool { }) } -fn validate_tiers(tiers: &[TierRates], duplicate: PricingError) -> Result<(), PricingError> { +fn validate_tiers(tiers: &[TierRates], duplicate: Error) -> Result<(), Error> { if tiers.iter().any(|entry| !valid_rates(entry.rates)) { - return Err(PricingError::InvalidRate); + return Err(Error::InvalidRate); } if tiers.iter().enumerate().any(|(index, entry)| { matches!( @@ -227,20 +215,20 @@ fn validate_tiers(tiers: &[TierRates], duplicate: PricingError) -> Result<(), Pr Ok(()) } -pub fn compile(pricing: &Pricing<'_>) -> Result { +pub fn compile(pricing: &Pricing<'_>) -> Result { if !valid_rates(pricing.standard) { - return Err(PricingError::InvalidRate); + return Err(Error::InvalidRate); } - validate_tiers(pricing.tiers, PricingError::DuplicateTier)?; + validate_tiers(pricing.tiers, Error::DuplicateTier)?; if let Some(window) = pricing.off_peak { if window.start_utc_minute >= 1440 || window.end_utc_minute > 1440 || window.start_utc_minute >= window.end_utc_minute { - return Err(PricingError::InvalidOffPeakWindow); + return Err(Error::InvalidOffPeakWindow); } if !valid_rates(window.rates) { - return Err(PricingError::InvalidRate); + return Err(Error::InvalidRate); } } let mut thresholds: Vec<_> = pricing @@ -248,9 +236,9 @@ pub fn compile(pricing: &Pricing<'_>) -> Result { .iter() .map(|entry| { if !valid_rates(entry.standard) { - return Err(PricingError::InvalidRate); + return Err(Error::InvalidRate); } - validate_tiers(entry.tiers, PricingError::DuplicateThresholdTier)?; + validate_tiers(entry.tiers, Error::DuplicateThresholdTier)?; Ok(CompiledThreshold { above_prompt_tokens: entry.above_prompt_tokens, standard: entry.standard, @@ -263,7 +251,7 @@ pub fn compile(pricing: &Pricing<'_>) -> Result { .windows(2) .any(|pair| pair[0].above_prompt_tokens == pair[1].above_prompt_tokens) { - return Err(PricingError::DuplicateThreshold); + return Err(Error::DuplicateThreshold); } Ok(PricingPlan { standard: pricing.standard, @@ -274,11 +262,7 @@ pub fn compile(pricing: &Pricing<'_>) -> Result { } impl PricingPlan { - fn resolve_rates( - &self, - request: &Request, - threshold_tokens: u64, - ) -> Result { + fn resolve_rates(&self, request: &Request, threshold_tokens: u64) -> Result { let tier = selected_tier(request.service_tier); let base = self .tiers @@ -305,13 +289,13 @@ impl PricingPlan { || window.end_utc_minute > 1440 || window.start_utc_minute >= window.end_utc_minute { - return Err(PricingError::InvalidOffPeakWindow); + return Err(Error::InvalidOffPeakWindow); } let minute = request .billed_at_utc_minute - .ok_or(PricingError::InvalidBillingTime)?; + .ok_or(Error::InvalidBillingTime)?; if minute >= 1440 { - return Err(PricingError::InvalidBillingTime); + return Err(Error::InvalidBillingTime); } if (window.start_utc_minute..window.end_utc_minute).contains(&minute) { Ok(window.rates.overlay(selected)) @@ -322,26 +306,26 @@ impl PricingPlan { } } - fn checked_rate(rate: Rate, missing: PricingError) -> Result { + fn checked_rate(rate: Rate, missing: Error) -> Result { let value = rate.value().ok_or(missing)?; if !value.is_finite() || value < 0.0 { - return Err(PricingError::InvalidRate); + return Err(Error::InvalidRate); } Ok(value) } - pub fn calculate(&self, request: &Request) -> Result { + pub fn calculate(&self, request: &Request) -> Result { let usage = request.usage; let cached = usage .cache_read_tokens .checked_add(usage.cache_write_tokens) - .ok_or(PricingError::TokenCountOverflow)?; + .ok_or(Error::TokenCountOverflow)?; let (regular, threshold_tokens) = match usage.prompt_convention { PromptConvention::IncludesCache => ( usage .prompt_tokens .checked_sub(cached) - .ok_or(PricingError::CacheExceedsPrompt)?, + .ok_or(Error::CacheExceedsPrompt)?, usage.prompt_tokens, ), PromptConvention::ExcludesCache => ( @@ -349,7 +333,7 @@ impl PricingPlan { usage .prompt_tokens .checked_add(cached) - .ok_or(PricingError::TokenCountOverflow)?, + .ok_or(Error::TokenCountOverflow)?, ), }; let writes = match (usage.cache_write_5m_tokens, usage.cache_write_1h_tokens) { @@ -357,26 +341,20 @@ impl PricingPlan { (Some(five), Some(one)) if five.checked_add(one) == Some(usage.cache_write_tokens) => { (five, one) } - _ => return Err(PricingError::InvalidCacheWriteDetails), + _ => return Err(Error::InvalidCacheWriteDetails), }; let rates = self.resolve_rates(request, threshold_tokens)?; - let input = Self::checked_rate(rates.input, PricingError::MissingInputRate)?; - let output = Self::checked_rate(rates.output, PricingError::MissingOutputRate)?; - let read = Self::checked_rate( - rates.cache_read.or(rates.input), - PricingError::MissingInputRate, - )?; - let write = Self::checked_rate( - rates.cache_write.or(rates.input), - PricingError::MissingInputRate, - )?; + let input = Self::checked_rate(rates.input, Error::MissingInputRate)?; + let output = Self::checked_rate(rates.output, Error::MissingOutputRate)?; + let read = Self::checked_rate(rates.cache_read.or(rates.input), Error::MissingInputRate)?; + let write = Self::checked_rate(rates.cache_write.or(rates.input), Error::MissingInputRate)?; let write_1h = Self::checked_rate( rates.cache_write_1h.or(rates.cache_write).or(rates.input), - PricingError::MissingInputRate, + Error::MissingInputRate, )?; let multiplier = request.region_multiplier.unwrap_or(1.0); if !multiplier.is_finite() || multiplier <= 0.0 { - return Err(PricingError::InvalidRegionMultiplier); + return Err(Error::InvalidRegionMultiplier); } let cost = Cost { uncached_input: regular as f64 * input, @@ -394,12 +372,12 @@ impl PricingPlan { }, }; if !cost.total().is_finite() { - return Err(PricingError::TokenCountOverflow); + return Err(Error::TokenCountOverflow); } Ok(cost) } } -pub fn calculate(pricing: &Pricing<'_>, request: &Request) -> Result { +pub fn calculate(pricing: &Pricing<'_>, request: &Request) -> Result { compile(pricing)?.calculate(request) } diff --git a/litellm-rust/crates/cost/tests/calculation.rs b/litellm-rust/crates/cost/tests/calculation.rs index f8152483a7c..af1a73d8c61 100644 --- a/litellm-rust/crates/cost/tests/calculation.rs +++ b/litellm-rust/crates/cost/tests/calculation.rs @@ -1,5 +1,5 @@ use litellm_cost::{ - OffPeakRates, Pricing, PricingError, PromptConvention, Rate, Rates, Request, ServiceTier, + Error, OffPeakRates, Pricing, PromptConvention, Rate, Rates, Request, ServiceTier, ThresholdPolicy, ThresholdRates, TierRates, Usage, calculate, compile, }; use rstest::rstest; @@ -132,7 +132,7 @@ fn split_writes_and_invalid_accounting() { }; assert_eq!( calculate(&pricing(standard), &overlapping), - Err(PricingError::CacheExceedsPrompt) + Err(Error::CacheExceedsPrompt) ); let incomplete = Request { usage: Usage { @@ -143,7 +143,7 @@ fn split_writes_and_invalid_accounting() { }; assert_eq!( calculate(&pricing(standard), &incomplete), - Err(PricingError::InvalidCacheWriteDetails) + Err(Error::InvalidCacheWriteDetails) ); } @@ -297,10 +297,10 @@ fn compile_rejects_ambiguous_rates() { }; assert_eq!( compile(&specification).err(), - Some(PricingError::DuplicateThreshold) + Some(Error::DuplicateThreshold) ); let invalid = pricing(rates(Rate::Value(f64::NAN), Rate::Value(1.0))); - assert_eq!(compile(&invalid).err(), Some(PricingError::InvalidRate)); + assert_eq!(compile(&invalid).err(), Some(Error::InvalidRate)); } #[test] @@ -326,7 +326,7 @@ fn off_peak_is_one_non_wrapping_utc_window() { }; assert_eq!( calculate(&specification, &base), - Err(PricingError::InvalidBillingTime) + Err(Error::InvalidBillingTime) ); assert_eq!(calculate(&specification, &start).unwrap().input(), 100.0); assert_eq!(calculate(&specification, &end).unwrap().input(), 200.0); @@ -347,11 +347,11 @@ fn missing_rates_and_free_rates_remain_distinct() { }; assert_eq!( calculate(&pricing(Rates::EMPTY), &empty), - Err(PricingError::MissingInputRate) + Err(Error::MissingInputRate) ); assert_eq!( calculate(&pricing(rates(Rate::Value(0.0), Rate::Missing)), &empty), - Err(PricingError::MissingOutputRate) + Err(Error::MissingOutputRate) ); assert_eq!( calculate(&pricing(rates(Rate::Value(0.0), Rate::Value(0.0))), &empty) @@ -481,7 +481,7 @@ fn regional_multiplier_applies_after_input_components_are_summed() { }; assert_eq!( calculate(&pricing(standard), &invalid), - Err(PricingError::InvalidRegionMultiplier) + Err(Error::InvalidRegionMultiplier) ); } @@ -531,5 +531,5 @@ fn compilation_sorts_thresholds_and_rejects_duplicate_tiers() { thresholds: &[], ..specification }; - assert_eq!(compile(&invalid).err(), Some(PricingError::DuplicateTier)); + assert_eq!(compile(&invalid).err(), Some(Error::DuplicateTier)); } diff --git a/litellm-rust/crates/error/Cargo.toml b/litellm-rust/crates/error/Cargo.toml new file mode 100644 index 00000000000..fa2d26a1b99 --- /dev/null +++ b/litellm-rust/crates/error/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "litellm-error" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true +description = "Shared error rendering and type-erased causes for litellm-rust boundaries" + +[dependencies] +thiserror.workspace = true + +[dev-dependencies] +rstest.workspace = true diff --git a/litellm-rust/crates/error/src/lib.rs b/litellm-rust/crates/error/src/lib.rs new file mode 100644 index 00000000000..ec8d5c3e1f4 --- /dev/null +++ b/litellm-rust/crates/error/src/lib.rs @@ -0,0 +1,132 @@ +use std::{error::Error, fmt, ops::Deref, sync::Arc}; + +pub struct Report<'a>(&'a (dyn Error + 'static)); + +impl<'a> Report<'a> { + pub fn new(error: &'a (dyn Error + 'static)) -> Self { + Self(error) + } +} + +impl fmt::Display for Report<'_> { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(formatter, "{}", self.0)?; + std::iter::successors(self.0.source(), |&cause| cause.source()) + .try_for_each(|cause| write!(formatter, ": {cause}")) + } +} + +pub fn find<'a, T: Error + 'static>(error: &'a (dyn Error + 'static)) -> Option<&'a T> { + std::iter::successors(Some(error), |&cause| cause.source()).find_map(|cause| { + cause + .downcast_ref::() + .or_else(|| cause.downcast_ref::()?.downcast_ref::()) + }) +} + +#[derive(Clone, Debug, thiserror::Error)] +#[error(transparent)] +pub struct Source(Arc); + +impl Source { + pub fn new(error: impl Error + Send + Sync + 'static) -> Self { + Self(Arc::new(error)) + } +} + +impl From> for Source { + fn from(error: Box) -> Self { + Self(Arc::from(error)) + } +} + +impl Deref for Source { + type Target = dyn Error + Send + Sync; + + fn deref(&self) -> &Self::Target { + self.0.as_ref() + } +} + +impl PartialEq for Source { + fn eq(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.0, &other.0) + } +} + +impl Eq for Source {} + +#[cfg(test)] +mod tests { + use std::io; + + use rstest::rstest; + + use super::{Report, Source, find}; + + #[derive(Debug, thiserror::Error)] + enum Leaf { + #[error("could not read the config")] + Read(#[source] io::Error), + #[error("could not parse the config")] + Parse(#[source] Source), + #[error("config is empty")] + Empty, + } + + #[derive(Debug, thiserror::Error)] + enum Outer { + #[error("startup failed")] + Startup(#[from] Leaf), + #[error(transparent)] + Nested(Leaf), + } + + fn denied() -> io::Error { + io::Error::new(io::ErrorKind::PermissionDenied, "permission denied") + } + + #[rstest] + #[case::leaf(Outer::Nested(Leaf::Empty), "config is empty")] + #[case::context_over_a_cause( + Outer::Startup(Leaf::Read(denied())), + "startup failed: could not read the config: permission denied" + )] + #[case::transparent_adds_no_text( + Outer::Nested(Leaf::Read(denied())), + "could not read the config: permission denied" + )] + #[case::erased_cause_keeps_its_chain( + Outer::Nested(Leaf::Parse(Source::new(Leaf::Read(denied())))), + "could not parse the config: could not read the config: permission denied" + )] + fn report_renders_every_layer_once(#[case] error: Outer, #[case] expected: &str) { + assert_eq!(Report::new(&error).to_string(), expected); + } + + #[rstest] + #[case::direct_source(Outer::Startup(Leaf::Read(denied())))] + #[case::behind_an_erased_source(Outer::Startup(Leaf::Parse(Source::new(Leaf::Read( + denied() + )))))] + #[case::behind_a_boxed_source(Outer::Startup(Leaf::Parse(Source::from( + Box::::from(denied()) + ))))] + #[case::behind_a_transparent_layer(Outer::Nested(Leaf::Parse(Source::new(denied()))))] + fn find_reaches_the_typed_cause_anywhere_in_the_chain(#[case] error: Outer) { + let cause = find::(&error).expect("io cause"); + assert_eq!(cause.kind(), io::ErrorKind::PermissionDenied); + } + + #[test] + fn find_answers_none_when_the_chain_has_no_such_cause() { + assert!(find::(&Outer::Nested(Leaf::Empty)).is_none()); + } + + #[test] + fn a_clone_shares_its_cause() { + let source = Source::new(denied()); + assert_eq!(source.clone(), source); + assert_ne!(Source::new(denied()), source); + } +} diff --git a/litellm-rust/crates/framer/src/error.rs b/litellm-rust/crates/framer/src/error.rs index 879d7557671..dde421d71af 100644 --- a/litellm-rust/crates/framer/src/error.rs +++ b/litellm-rust/crates/framer/src/error.rs @@ -1,21 +1,21 @@ #[cfg(feature = "sse")] #[derive(Debug, thiserror::Error)] pub enum SseError { - #[error("body stream failed: {0}")] + #[error("body stream failed")] Body(#[from] std::io::Error), - #[error("SSE field is not UTF-8: {0}")] + #[error("SSE field is not UTF-8")] InvalidUtf8(#[from] std::str::Utf8Error), } #[cfg(feature = "aws")] #[derive(Debug, thiserror::Error)] pub enum EventStreamError { - #[error("body stream failed: {0}")] + #[error("body stream failed")] Body(#[from] std::io::Error), #[error("invalid AWS EventStream frame length: {0}")] InvalidLength(usize), #[error("truncated AWS EventStream frame")] Truncated, - #[error("malformed AWS EventStream frame: {0}")] + #[error("malformed AWS EventStream frame")] Malformed(#[from] aws_smithy_eventstream::error::Error), } diff --git a/litellm-rust/crates/framer/src/framer.rs b/litellm-rust/crates/framer/src/framer.rs deleted file mode 100644 index 507aed54700..00000000000 --- a/litellm-rust/crates/framer/src/framer.rs +++ /dev/null @@ -1,13 +0,0 @@ -use futures_util::Stream; - -use crate::Error; - -pub trait Framer: Send { - type Frame: Send; - - fn frame(self, input: S) -> impl Stream> + Send - where - S: Stream> + Send, - B: bytes::Buf + Send, - E: std::error::Error + Send + Sync + 'static; -} diff --git a/litellm-rust/crates/gateway-auth/AGENTS.md b/litellm-rust/crates/gateway-auth/AGENTS.md index 58457a3151b..0a4bc6450ae 100644 --- a/litellm-rust/crates/gateway-auth/AGENTS.md +++ b/litellm-rust/crates/gateway-auth/AGENTS.md @@ -20,7 +20,9 @@ Local UI login continues using axum-login and tower-sessions. Only after session Authentication evidence and principals contain no raw token, password, request body, or mutable accounting state. Session ownership is scoped by principal authority, subject, verifier, and credential ID, so separate credentials do not silently share an MCP session. Credential rotation may retain ownership when the verifier preserves a stable credential ID. Scope and expiry checks still run for each operation -Failures distinguish invalid or expired credentials, forbidden operations, unavailable authentication services, and missing server configuration/context. HTTP adapters map those outcomes to status codes. Inference keeps its API-specific error envelopes and MCP translates policy failures to its own protocol +Failures distinguish invalid or expired credentials, forbidden operations, unavailable authentication services, and missing server configuration/context. Bearer tokens, virtual keys and UI sessions all fail with the one `Error`, classified once by `status()`, and `is_denied()` tells a rejected caller from an auth failure. Inference keeps its API-specific error envelopes, the UI renders `Error` in its own envelope and MCP translates policy failures to its own protocol + +The plain-text rejection body is the top-level message only, never the `Report` chain, because the caller is unauthenticated and the sources carry secret-manager and backend detail. That stays reachable through `source()` for logs Virtual-key storage, JWT verification, OAuth2 introspection, trusted-proxy validation, SSO, and custom Python hook adapters are not implemented here yet. They should supply these contracts rather than bypassing the shared authorization boundary. Request-body-dependent custom hooks will need a bounded endpoint adapter after parsing diff --git a/litellm-rust/crates/gateway-auth/Cargo.toml b/litellm-rust/crates/gateway-auth/Cargo.toml index cf040c6eef3..27f20203843 100644 --- a/litellm-rust/crates/gateway-auth/Cargo.toml +++ b/litellm-rust/crates/gateway-auth/Cargo.toml @@ -10,6 +10,7 @@ axum = { workspace = true, features = ["matched-path"] } axum-login.workspace = true litellm-auth-types.workspace = true litellm-config.workspace = true +litellm-error.workspace = true litellm-secrets.workspace = true sha2.workspace = true subtle.workspace = true diff --git a/litellm-rust/crates/gateway-auth/src/error.rs b/litellm-rust/crates/gateway-auth/src/error.rs index 262a666df80..7bb6292b18c 100644 --- a/litellm-rust/crates/gateway-auth/src/error.rs +++ b/litellm-rust/crates/gateway-auth/src/error.rs @@ -2,44 +2,59 @@ use axum::{ http::StatusCode, response::{IntoResponse, Response}, }; - -#[derive(Debug, thiserror::Error)] -pub enum KeyError { - #[error("missing or invalid virtual key")] - Invalid, - #[error("virtual key has expired")] - Expired, - #[error("virtual key lookup unavailable")] - Lookup(#[source] Box), -} +use litellm_error::Source; #[derive(Debug, thiserror::Error)] pub enum Error { - #[error("operation is not permitted")] - Forbidden, - #[error("credential has expired")] - Expired, - #[error("authenticated request context is missing")] - MissingIdentity, - #[error("authentication service unavailable")] - Unavailable, - #[error("gateway auth not configured")] - Unconfigured, #[error("missing or invalid bearer token")] InvalidToken, + #[error("invalid UI credentials or session")] + InvalidSession, + #[error("credential has expired")] + Expired, + #[error("operation is not permitted")] + Forbidden, + #[error("authentication service unavailable")] + Unavailable(#[source] Source), #[error("gateway authentication unavailable")] Secret(#[from] litellm_secrets::Error), + #[error("UI session unavailable")] + Session(#[from] tower_sessions::session::Error), + #[error("gateway auth not configured")] + Unconfigured, + #[error("UI_USERNAME and UI_PASSWORD must be nonempty")] + UiUnconfigured, + #[error("authenticated request context is missing")] + MissingIdentity, + #[error("UI session layer is not installed")] + MissingSessionLayer, } impl Error { pub fn status(&self) -> StatusCode { match self { - Self::InvalidToken | Self::Expired => StatusCode::UNAUTHORIZED, + Self::InvalidToken | Self::InvalidSession | Self::Expired => StatusCode::UNAUTHORIZED, Self::Forbidden => StatusCode::FORBIDDEN, - Self::Unavailable => StatusCode::SERVICE_UNAVAILABLE, - Self::MissingIdentity | Self::Unconfigured | Self::Secret(_) => { - StatusCode::INTERNAL_SERVER_ERROR - } + Self::Unavailable(_) => StatusCode::SERVICE_UNAVAILABLE, + Self::Secret(_) + | Self::Session(_) + | Self::Unconfigured + | Self::UiUnconfigured + | Self::MissingIdentity + | Self::MissingSessionLayer => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + pub fn is_denied(&self) -> bool { + match self { + Self::InvalidToken | Self::InvalidSession | Self::Expired | Self::Forbidden => true, + Self::Unavailable(_) + | Self::Secret(_) + | Self::Session(_) + | Self::Unconfigured + | Self::UiUnconfigured + | Self::MissingIdentity + | Self::MissingSessionLayer => false, } } } @@ -50,24 +65,31 @@ impl IntoResponse for Error { } } -#[derive(Debug, thiserror::Error)] -pub enum UiAuthError { - #[error("UI_USERNAME and UI_PASSWORD must be nonempty")] - Unconfigured, - #[error("invalid UI credentials or session")] - Unauthorized, - #[error("UI authentication unavailable")] - Unavailable, - #[error("UI session unavailable")] - Session(#[from] tower_sessions::session::Error), -} +#[cfg(test)] +mod tests { + use rstest::rstest; -impl IntoResponse for UiAuthError { - fn into_response(self) -> Response { - let status = match self { - Self::Unauthorized => StatusCode::UNAUTHORIZED, - _ => StatusCode::INTERNAL_SERVER_ERROR, - }; - (status, self.to_string()).into_response() + use super::*; + + #[rstest] + #[tokio::test] + async fn rejection_body_hides_backend_causes() { + let error = + Error::Unavailable(Source::new(std::io::Error::other("private-backend-detail"))); + assert!( + litellm_error::Report::new(&error) + .to_string() + .contains("private-backend-detail") + ); + let response = error.into_response(); + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + let body = axum::body::to_bytes(response.into_body(), 4096) + .await + .unwrap(); + assert!( + !std::str::from_utf8(&body) + .unwrap() + .contains("private-backend-detail") + ); } } diff --git a/litellm-rust/crates/gateway-auth/src/keys.rs b/litellm-rust/crates/gateway-auth/src/keys.rs index 46d8c02253f..1fe8fa60e0e 100644 --- a/litellm-rust/crates/gateway-auth/src/keys.rs +++ b/litellm-rust/crates/gateway-auth/src/keys.rs @@ -2,7 +2,7 @@ use std::{future::Future, pin::Pin, time::SystemTime}; use litellm_auth_types::SecretValue; -use crate::{KeyError, hash_token}; +use crate::{Error, hash_token}; #[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)] pub struct KeyHash(String); @@ -24,7 +24,7 @@ pub enum KeyStatus { } pub type LookupFuture<'a> = - Pin, KeyError>> + Send + 'a>>; + Pin, Error>> + Send + 'a>>; pub trait KeyLookup: Send + Sync { fn lookup<'a>(&'a self, hash: &'a KeyHash) -> LookupFuture<'a>; @@ -34,15 +34,15 @@ pub async fn verify_key( lookup: &dyn KeyLookup, token: &SecretValue, now: impl FnOnce() -> SystemTime, -) -> Result { +) -> Result { let hash = KeyHash::from_token(token); match lookup.lookup(&hash).await? { Some(KeyStatus::Active { expires_at }) => { if expires_at.is_some_and(|expiry| expiry <= now()) { - return Err(KeyError::Expired); + return Err(Error::Expired); } Ok(hash) } - Some(KeyStatus::Revoked) | None => Err(KeyError::Invalid), + Some(KeyStatus::Revoked) | None => Err(Error::InvalidToken), } } diff --git a/litellm-rust/crates/gateway-auth/src/lib.rs b/litellm-rust/crates/gateway-auth/src/lib.rs index 1b400a94746..31ce17b3aa4 100644 --- a/litellm-rust/crates/gateway-auth/src/lib.rs +++ b/litellm-rust/crates/gateway-auth/src/lib.rs @@ -16,7 +16,7 @@ pub use authorization::{ AccessRequest, AuthenticatedRequest, AuthorizedOperation, Authorizer, McpAction, NoAdditionalPolicy, Permissions, UiAction, }; -pub use error::{Error, KeyError, UiAuthError}; +pub use error::Error; pub use http::{Bearer, CredentialExtractor, RequireMasterKey, authenticate}; pub use identity::{ AuthenticatedCaller, Authentication, AuthenticationMethod, Principal, PrincipalKind, diff --git a/litellm-rust/crates/gateway-auth/src/ui.rs b/litellm-rust/crates/gateway-auth/src/ui.rs index 9e6b77a17d0..3d655bde187 100644 --- a/litellm-rust/crates/gateway-auth/src/ui.rs +++ b/litellm-rust/crates/gateway-auth/src/ui.rs @@ -11,9 +11,9 @@ use tower_sessions::Session; use veil::Redact; use crate::{ - AccessRequest, AuthenticatedRequest, Authentication, AuthenticationMethod, LocalAdministrator, - NoAdditionalPolicy, Permissions, Principal, PrincipalKind, SystemClock, UiAction, UiAuthError, - VerifiedIdentity, + AccessRequest, AuthenticatedRequest, Authentication, AuthenticationMethod, Error, + LocalAdministrator, NoAdditionalPolicy, Permissions, Principal, PrincipalKind, SystemClock, + UiAction, VerifiedIdentity, }; pub const UI_CSRF_KEY: &str = "litellm.ui.csrf"; @@ -50,9 +50,9 @@ pub struct UiBackend { } impl UiBackend { - pub fn new(username: String, password: SecretValue) -> Result { + pub fn new(username: String, password: SecretValue) -> Result { if username.trim().is_empty() || password.expose().trim().is_empty() { - return Err(UiAuthError::Unconfigured); + return Err(Error::UiUnconfigured); } Ok(Self { user: UiUser { @@ -89,28 +89,28 @@ pub struct UiSession { } impl FromRequestParts for UiSession { - type Rejection = UiAuthError; + type Rejection = Error; async fn from_request_parts(parts: &mut Parts, state: &S) -> Result { let auth = UiAuthSession::from_request_parts(parts, state) .await - .map_err(|_| UiAuthError::Unavailable)?; - let user = auth.user.ok_or(UiAuthError::Unauthorized)?; + .map_err(|_| Error::MissingSessionLayer)?; + let user = auth.user.ok_or(Error::InvalidSession)?; let session = Session::from_request_parts(parts, state) .await - .map_err(|_| UiAuthError::Unavailable)?; + .map_err(|_| Error::MissingSessionLayer)?; let expected = session .get::(UI_CSRF_KEY) .await? - .ok_or(UiAuthError::Unauthorized)?; + .ok_or(Error::InvalidSession)?; let provided = parts .headers .get(AUTHORIZATION) .and_then(|value| value.to_str().ok()) .and_then(|value| value.strip_prefix("Bearer ")) - .ok_or(UiAuthError::Unauthorized)?; + .ok_or(Error::InvalidSession)?; if !bool::from(Sha256::digest(expected).ct_eq(&Sha256::digest(provided))) { - return Err(UiAuthError::Unauthorized); + return Err(Error::InvalidSession); } let identity = session_identity(&user).await?; parts.extensions.insert(identity.clone()); @@ -118,7 +118,7 @@ impl FromRequestParts for UiSession { } } -async fn session_identity(user: &UiUser) -> Result { +async fn session_identity(user: &UiUser) -> Result { crate::authentication::resolve( VerifiedIdentity { principal: Principal::new( @@ -142,7 +142,6 @@ async fn session_identity(user: &UiUser) -> Result Ok(Some(SecretValue::new("resolved-key"))), "EMPTY" => Ok(Some(SecretValue::new(""))), - "ERROR" => Err(litellm_secrets::Error::ExternalRead(Box::new( - std::io::Error::other("private-backend-detail"), - ))), + "ERROR" => Err(litellm_secrets::Error::ExternalRead( + litellm_error::Source::new(std::io::Error::other("private-backend-detail")), + )), _ => Ok(None), } }) @@ -132,7 +132,11 @@ impl Authenticator for Verifier { Box::pin(async move { match self.failure { Some(401) => return Err(Error::InvalidToken), - Some(_) => return Err(Error::Unavailable), + Some(_) => { + return Err(Error::Unavailable(litellm_error::Source::new( + std::io::Error::other("verifier offline"), + ))); + } None => (), } let litellm_gateway_auth::Credential::Token(credential) = credential else { diff --git a/litellm-rust/crates/gateway-auth/tests/keys.rs b/litellm-rust/crates/gateway-auth/tests/keys.rs index 51ef23a859f..4d16d3573bc 100644 --- a/litellm-rust/crates/gateway-auth/tests/keys.rs +++ b/litellm-rust/crates/gateway-auth/tests/keys.rs @@ -2,7 +2,7 @@ use std::time::{Duration, SystemTime}; use litellm_auth_types::SecretValue; use litellm_gateway_auth::{ - KeyError, + Error, keys::{KeyHash, KeyLookup, KeyStatus, LookupFuture, verify_key}, }; use rstest::{fixture, rstest}; @@ -59,7 +59,7 @@ async fn checks_expiration( match valid { true => assert_eq!(result.unwrap(), hash), - false => assert!(matches!(result, Err(KeyError::Expired))), + false => assert!(matches!(result, Err(Error::Expired))), } } @@ -79,7 +79,7 @@ async fn rejects_unusable_keys( assert!(matches!( verify_key(&lookup, &token, || now).await, - Err(KeyError::Invalid) + Err(Error::InvalidToken) )); } @@ -93,7 +93,7 @@ async fn rejects_a_different_token(token: SecretValue, now: SystemTime) { assert!(matches!( verify_key(&lookup, &SecretValue::new("sk-other-credential"), || now).await, - Err(KeyError::Invalid) + Err(Error::InvalidToken) )); } @@ -102,9 +102,9 @@ struct Unavailable; impl KeyLookup for Unavailable { fn lookup<'a>(&'a self, _: &'a KeyHash) -> LookupFuture<'a> { Box::pin(async { - Err(KeyError::Lookup(Box::new(std::io::Error::other( - "backend detail", - )))) + Err(Error::Unavailable(litellm_error::Source::new( + std::io::Error::other("backend detail"), + ))) }) } } @@ -114,10 +114,8 @@ impl KeyLookup for Unavailable { async fn lookup_failure_preserves_source_without_exposing_it(token: SecretValue, now: SystemTime) { let error = verify_key(&Unavailable, &token, || now).await.unwrap_err(); - assert!(matches!(&error, KeyError::Lookup(_))); - assert_eq!( - std::error::Error::source(&error).unwrap().to_string(), - "backend detail" - ); + assert_eq!(error.status(), axum::http::StatusCode::SERVICE_UNAVAILABLE); + let cause = litellm_error::find::(&error).expect("lookup cause"); + assert_eq!(cause.to_string(), "backend detail"); assert!(!error.to_string().contains("backend detail")); } diff --git a/litellm-rust/crates/gateway-inference/Cargo.toml b/litellm-rust/crates/gateway-inference/Cargo.toml index 312d44b9a4b..034c3a9826e 100644 --- a/litellm-rust/crates/gateway-inference/Cargo.toml +++ b/litellm-rust/crates/gateway-inference/Cargo.toml @@ -11,6 +11,7 @@ axum = { workspace = true, features = ["json", "multipart", "original-uri"] } base64.workspace = true bytes.workspace = true litellm-auth.workspace = true +litellm-error.workspace = true litellm-gateway-auth.workspace = true litellm-inference.workspace = true litellm-inference-transcription.workspace = true @@ -34,5 +35,6 @@ litellm-cache-memory.workspace = true futures-util.workspace = true tokio = { workspace = true, features = ["io-util"] } rstest.workspace = true +tokio-tungstenite.workspace = true tower = { version = "0.5.3", features = ["util"] } wiremock.workspace = true diff --git a/litellm-rust/crates/gateway-inference/src/audio_transcription.rs b/litellm-rust/crates/gateway-inference/src/audio_transcription.rs index 8d571eccc5b..d4cb467f6fd 100644 --- a/litellm-rust/crates/gateway-inference/src/audio_transcription.rs +++ b/litellm-rust/crates/gateway-inference/src/audio_transcription.rs @@ -27,7 +27,7 @@ async fn handle( upload, }: InferenceBody, ) -> Result { - let deployment = request::resolve_deployment(gateway, &body)?; + let deployment = request::resolve_deployment(gateway, "/audio/transcriptions", &body)?; request::authorize_model(identity, deployment, &body).await?; let audio = match upload { Some(upload) => { diff --git a/litellm-rust/crates/gateway-inference/src/caching.rs b/litellm-rust/crates/gateway-inference/src/caching.rs index 426b8b91182..72f07d5d621 100644 --- a/litellm-rust/crates/gateway-inference/src/caching.rs +++ b/litellm-rust/crates/gateway-inference/src/caching.rs @@ -5,7 +5,7 @@ use litellm_gateway_auth::AuthenticatedRequest; use serde::Deserialize; use serde_json::{Map, Value}; -use crate::Error; +use crate::BodyError; #[derive(Default, Deserialize)] #[serde(default, deny_unknown_fields)] @@ -24,10 +24,9 @@ type Prepared = (Map, CacheOptions); pub(crate) fn prepare( identity: &AuthenticatedRequest, body: Map, -) -> Result { +) -> Result { let controls: Controls = match body.get("cache").filter(|value| !value.is_null()) { - Some(value) => serde_json::from_value(value.clone()) - .map_err(|error| Error::InvalidBody(error.to_string()))?, + Some(value) => serde_json::from_value(value.clone()).map_err(BodyError::CacheControls)?, None => Controls::default(), }; let caching: Option = body @@ -35,7 +34,7 @@ pub(crate) fn prepare( .filter(|value| !value.is_null()) .map(|value| serde_json::from_value(value.clone())) .transpose() - .map_err(|error| Error::InvalidBody(error.to_string()))?; + .map_err(BodyError::Caching)?; let caller = identity.caller(); let options = CacheOptions { policy: CachePolicy { @@ -62,11 +61,12 @@ pub(crate) fn prepare( )) } -fn duration(seconds: f64) -> Result { - Duration::try_from_secs_f64(seconds) - .ok() - .filter(|duration| !duration.is_zero()) - .ok_or_else(|| Error::InvalidBody("cache durations must be finite and positive".into())) +fn duration(seconds: f64) -> Result { + match Duration::try_from_secs_f64(seconds) { + Ok(duration) if !duration.is_zero() => Ok(duration), + Ok(_) => Err(BodyError::CacheDuration(None)), + Err(error) => Err(BodyError::CacheDuration(Some(error))), + } } #[derive(Clone, Default)] diff --git a/litellm-rust/crates/gateway-inference/src/chat_completions.rs b/litellm-rust/crates/gateway-inference/src/chat_completions.rs index 742a373b286..fb8a6f4609d 100644 --- a/litellm-rust/crates/gateway-inference/src/chat_completions.rs +++ b/litellm-rust/crates/gateway-inference/src/chat_completions.rs @@ -3,13 +3,16 @@ use std::sync::Arc; use axum::{ Json, - extract::{Path, State}, + extract::State, response::{IntoResponse, Response}, }; use litellm_inference_chat::types::ChatCompletionsCall; use serde_json::{Map, Value}; -use crate::{Error, Gateway, JsonObject, request}; +use crate::{ + Error, Gateway, JsonObject, + request::{self, Path}, +}; pub(crate) async fn create( State(gateway): State>, @@ -40,7 +43,7 @@ async fn handle( identity: &AuthenticatedRequest, body: Map, ) -> Result { - let deployment = request::resolve_deployment(gateway, &body)?; + let deployment = request::resolve_deployment(gateway, "/chat/completions", &body)?; request::authorize_model(identity, deployment, &body).await?; let (body, cache_options) = crate::caching::prepare(identity, body)?; let route = gateway.chat_completions.clone(); diff --git a/litellm-rust/crates/gateway-inference/src/error.rs b/litellm-rust/crates/gateway-inference/src/error.rs index ec82a6ede12..505483b426c 100644 --- a/litellm-rust/crates/gateway-inference/src/error.rs +++ b/litellm-rust/crates/gateway-inference/src/error.rs @@ -3,7 +3,6 @@ use axum::{ Json, response::{IntoResponse, Response}, }; -use litellm_http::transport::Error as TransportError; use litellm_inference::Error as InferenceError; use litellm_llms::base_llm::ocr::error::Error as OcrError; use serde_json::{Map, Value, json}; @@ -12,12 +11,14 @@ use serde_json::{Map, Value, json}; pub enum Error { #[error(transparent)] Auth(#[from] litellm_gateway_auth::Error), - #[error("invalid request body: {0}")] - InvalidBody(String), + #[error("invalid request body")] + InvalidBody(#[from] BodyError), + #[error("invalid path parameter")] + InvalidPath(#[source] axum::extract::rejection::PathRejection), #[error( - "/v1/messages: Invalid model name passed in model={0}. Call `/v1/models` to view available models for your key." + "{route}: Invalid model name passed in model={model}. Call `/v1/models` to view available models for your key." )] - UnknownModel(String), + UnknownModel { route: &'static str, model: String }, #[error(transparent)] Route(#[from] InferenceError), #[error(transparent)] @@ -26,8 +27,30 @@ pub enum Error { Unsupported(String), #[error("request body exceeds the size limit")] BodyTooLarge, - #[error("{0}")] - Internal(String), + #[error("unexpected HTTP host operation")] + Internal, +} + +#[derive(Debug, thiserror::Error)] +pub enum BodyError { + #[error("could not read the body")] + Read(#[source] axum::extract::rejection::BytesRejection), + #[error("invalid JSON")] + Json(#[source] serde_json::Error), + #[error("expected a JSON object")] + NotAnObject, + #[error("invalid multipart body")] + Multipart(#[source] axum::extract::multipart::MultipartRejection), + #[error("invalid multipart field")] + MultipartField(#[source] axum::extract::multipart::MultipartError), + #[error("model is required")] + MissingModel, + #[error("invalid cache controls")] + CacheControls(#[source] serde_json::Error), + #[error("invalid caching flag")] + Caching(#[source] serde_json::Error), + #[error("cache durations must be finite and positive")] + CacheDuration(#[source] Option), } impl IntoResponse for Error { @@ -40,7 +63,7 @@ impl From> for Error { fn from(error: litellm_host_http::Error) -> Self { match error { litellm_host_http::Error::Call(error) => Self::Route(error), - litellm_host_http::Error::Protocol => Self::Internal(error.to_string()), + litellm_host_http::Error::Protocol => Self::Internal, } } } @@ -49,58 +72,50 @@ impl Error { pub fn status(&self) -> StatusCode { match self { Self::Auth(error) => error.status(), - Self::Unsupported(_) - | Self::Route(InferenceError::Unsupported(_)) - | Self::Ocr(OcrError::Unsupported(_)) => StatusCode::NOT_IMPLEMENTED, + Self::InvalidPath(rejection) => rejection.status(), + Self::Unsupported(_) => StatusCode::NOT_IMPLEMENTED, Self::BodyTooLarge => StatusCode::PAYLOAD_TOO_LARGE, - Self::Ocr( - OcrError::Auth(litellm_auth::Error::MissingApiKey { .. }) - | OcrError::MissingAzureAiCredentials - | OcrError::MissingAzureDocumentIntelligenceCredentials - | OcrError::MissingReductoApiKey, - ) => StatusCode::UNAUTHORIZED, - Self::Ocr(error) => error - .http_status_code() - .and_then(|status| StatusCode::from_u16(status).ok()) - .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR), - Self::InvalidBody(_) | Self::UnknownModel(_) => StatusCode::BAD_REQUEST, - Self::Route(InferenceError::Transport(TransportError::Http { status, .. })) => { - StatusCode::from_u16(*status).unwrap_or(StatusCode::BAD_GATEWAY) - } - Self::Route(InferenceError::Auth(litellm_auth::Error::MissingApiKey { .. })) => { - StatusCode::UNAUTHORIZED - } - Self::Route(error) if error.is_request() => StatusCode::BAD_REQUEST, - Self::Route(_) | Self::Internal(_) => StatusCode::INTERNAL_SERVER_ERROR, + Self::InvalidBody(_) | Self::UnknownModel { .. } => StatusCode::BAD_REQUEST, + Self::Internal => StatusCode::INTERNAL_SERVER_ERROR, + Self::Ocr(error) => ocr_status(error), + Self::Route(error) => route_status(error), } } pub fn openai_response(self) -> Response { + (self.status(), Json(self.openai_body())).into_response() + } + + /// The OpenAI envelope as an `event: error` frame, for a stream that fails after its + /// headers went out. + pub fn responses_sse_frame(&self) -> String { + let event = json!({ + "type": "error", + "code": self.status().as_u16().to_string(), + "message": self.message(), + "param": null, + }); + format!("event: error\ndata: {event}\n\n") + } + + fn openai_body(&self) -> Value { let status = self.status(); - let message = match &self { - Self::UnknownModel(model) => format!("Invalid model name passed in model={model}"), - _ => self.to_string(), - }; - ( - status, - Json(json!({"error": { - "message": message, - "type": error_type(status), - "param": null, - "code": status.as_u16(), - }})), - ) - .into_response() + json!({"error": { + "message": self.message(), + "type": error_type(status), + "param": null, + "code": status.as_u16(), + }}) } /// The Anthropic error envelope Python's `AnthropicExceptionMapping` builds: an upstream /// body already in that shape passes through, any other has its message extracted. pub fn body(&self, request_id: Option<&str>) -> Value { let raw = match self { - Self::Route(InferenceError::Transport(TransportError::Http { body, .. })) => { - body.clone() + Self::Route(error) if let Some((_, body)) = error.upstream_response() => { + body.to_owned() } - other => other.to_string(), + other => other.message(), }; let parsed = serde_json::from_str::(&raw).ok(); let envelope = match parsed { @@ -117,6 +132,52 @@ impl Error { pub fn sse_frame(&self) -> String { format!("event: error\ndata: {}\n\n", self.body(None)) } + + fn message(&self) -> String { + match self { + Self::Auth(error) => error.to_string(), + Self::InvalidBody(_) + | Self::InvalidPath(_) + | Self::UnknownModel { .. } + | Self::Route(_) + | Self::Ocr(_) + | Self::Unsupported(_) + | Self::BodyTooLarge + | Self::Internal => litellm_error::Report::new(self).to_string(), + } + } +} + +fn route_status(error: &InferenceError) -> StatusCode { + if error.is_unsupported() { + return StatusCode::NOT_IMPLEMENTED; + } + if let Some((status, _)) = error.upstream_response() { + return StatusCode::from_u16(status).unwrap_or(StatusCode::BAD_GATEWAY); + } + if error.is_missing_credential() { + return StatusCode::UNAUTHORIZED; + } + if error.is_request() { + return StatusCode::BAD_REQUEST; + } + StatusCode::INTERNAL_SERVER_ERROR +} + +fn ocr_status(error: &OcrError) -> StatusCode { + if error.is_unsupported() { + return StatusCode::NOT_IMPLEMENTED; + } + if let Some((status, ..)) = error.upstream_response() { + return StatusCode::from_u16(status).unwrap_or(StatusCode::BAD_GATEWAY); + } + if error.is_missing_credential() { + return StatusCode::UNAUTHORIZED; + } + error + .http_status_code() + .and_then(|status| StatusCode::from_u16(status).ok()) + .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR) } fn error_type(status: StatusCode) -> &'static str { @@ -133,13 +194,13 @@ fn error_type(status: StatusCode) -> &'static str { } fn envelope(status: StatusCode, message: &str) -> Map { - let Value::Object(envelope) = json!({ - "type": "error", - "error": {"type": error_type(status), "message": message}, - }) else { - unreachable!("a json object literal is an object") - }; - envelope + Map::from_iter([ + ("type".to_owned(), Value::from("error")), + ( + "error".to_owned(), + json!({"type": error_type(status), "message": message}), + ), + ]) } fn is_anthropic_error(object: &Map) -> bool { @@ -177,14 +238,22 @@ mod tests { use super::*; fn upstream(status: u16, body: &str) -> Error { - Error::Route(InferenceError::Transport(TransportError::Http { - status, - body: body.into(), - })) + Error::Route(litellm_http::Error::upstream(status, body.into()).into()) + } + + fn invalid_json() -> Error { + BodyError::Json(serde_json::from_str::("{").unwrap_err()).into() + } + + fn unknown_model() -> Error { + Error::UnknownModel { + route: "/v1/messages", + model: "nope".into(), + } } #[rstest] - #[case::anthropic_body_passes_through( + #[case::anthropic_body( upstream(529, r#"{"type":"error","error":{"type":"overloaded_error","message":"busy","extra":1}}"#), Some("req_1"), json!({"type": "error", "error": {"type": "overloaded_error", "message": "busy", "extra": 1}, "request_id": "req_1"}), @@ -194,30 +263,7 @@ mod tests { Some("caller"), json!({"type": "error", "error": {"type": "x", "message": "m"}, "request_id": "upstream"}), )] - #[case::bedrock_detail( - upstream(403, r#"{"detail":{"message":"denied"}}"#), - None, - json!({"type": "error", "error": {"type": "permission_error", "message": "denied"}}), - )] - #[case::aws_message( - upstream(429, r#"{"Message":"slow down"}"#), - None, - json!({"type": "error", "error": {"type": "rate_limit_error", "message": "slow down"}}), - )] - #[case::plain_text_with_unmapped_status( - upstream(502, "bad gateway"), - None, - json!({"type": "error", "error": {"type": "api_error", "message": "bad gateway"}}), - )] - #[case::unknown_model( - Error::UnknownModel("nope".into()), - None, - json!({"type": "error", "error": { - "type": "invalid_request_error", - "message": "/v1/messages: Invalid model name passed in model=nope. Call `/v1/models` to view available models for your key.", - }}), - )] - fn body_follows_the_anthropic_exception_mapping( + fn an_anthropic_provider_body_passes_through( #[case] error: Error, #[case] request_id: Option<&str>, #[case] expected: Value, @@ -225,9 +271,68 @@ mod tests { assert_eq!(error.body(request_id), expected); } + #[rstest] + #[case::ascii("busy ".repeat(80))] + #[case::unicode("é".repeat(300))] + fn long_provider_errors_preserve_the_messages_envelope(#[case] message: String) { + let body = json!({ + "type": "error", + "error": {"type": "overloaded_error", "message": message}, + "request_id": "provider-request", + }); + let error = upstream(529, &body.to_string()); + let report = litellm_error::Report::new(&error).to_string(); + let diagnostic_body: String = body.to_string().chars().take(256).collect(); + assert_eq!(report.matches(&diagnostic_body).count(), 1, "{report}"); + assert!(report.ends_with("... (truncated)"), "{report}"); + assert!(!report.contains(&message), "{report}"); + assert!(!format!("{error:?}").contains(&message)); + assert_eq!(error.status(), StatusCode::from_u16(529).unwrap()); + assert_eq!(error.body(Some("gateway-request")), body); + } + + #[rstest] + #[case::bedrock_detail( + upstream(403, r#"{"detail":{"message":"denied"}}"#), + "permission_error", + "denied" + )] + #[case::aws_message( + upstream(429, r#"{"Message":"slow down"}"#), + "rate_limit_error", + "slow down" + )] + #[case::plain_text_with_unmapped_status( + upstream(502, "bad gateway"), + "api_error", + "bad gateway" + )] + #[case::unknown_model( + unknown_model(), + "invalid_request_error", + "/v1/messages: Invalid model name passed in model=nope. Call `/v1/models` to view available models for your key." + )] + fn other_failures_get_the_anthropic_envelope_with_their_message( + #[case] error: Error, + #[case] error_type: &str, + #[case] message: &str, + ) { + let body = error.body(Some("req_1")); + assert_eq!(body["type"], "error"); + assert_eq!(body["error"]["type"], error_type); + assert_eq!(body["error"]["message"], message); + assert_eq!(body["request_id"], "req_1"); + } + #[rstest] #[case::upstream_status(upstream(429, ""), StatusCode::TOO_MANY_REQUESTS)] - #[case::rejected_request(Error::Route(InferenceError::InvalidRequest("top_k".into())), StatusCode::BAD_REQUEST)] + #[case::rejected_request( + Error::Route(InferenceError::invalid_request(litellm_llms::ErrorDetail::Expected { + subject: "top_k", + expected: "an integer", + })), + StatusCode::BAD_REQUEST + )] #[case::missing_key( Error::Route(InferenceError::Auth(litellm_auth::Error::MissingApiKey { provider: "Anthropic", @@ -236,10 +341,83 @@ mod tests { StatusCode::UNAUTHORIZED, )] #[case::lost_connection( - Error::Route(InferenceError::Transport(TransportError::Network("reset".into()))), + Error::Route( + litellm_http::Error::from(tokio_tungstenite::tungstenite::Error::ConnectionClosed) + .into() + ), StatusCode::INTERNAL_SERVER_ERROR, )] + #[case::invalid_body(invalid_json(), StatusCode::BAD_REQUEST)] + #[case::missing_model(BodyError::MissingModel.into(), StatusCode::BAD_REQUEST)] + #[case::unknown_model(unknown_model(), StatusCode::BAD_REQUEST)] + #[case::oversized_body(Error::BodyTooLarge, StatusCode::PAYLOAD_TOO_LARGE)] + #[case::unserved_path(Error::Unsupported("/v1/completions".into()), StatusCode::NOT_IMPLEMENTED)] + #[case::host_protocol( + litellm_host_http::Error::::Protocol.into(), + StatusCode::INTERNAL_SERVER_ERROR + )] + #[case::ocr_provider_answer( + Error::Ocr(OcrError::Provider { status: 429, body: "slow".into(), headers: Default::default() }), + StatusCode::TOO_MANY_REQUESTS + )] + #[case::ocr_missing_credential( + Error::Ocr(OcrError::MissingCredential { credential: "Azure AI credentials", guidance: "set a key" }), + StatusCode::UNAUTHORIZED + )] + #[case::ocr_request_fault(Error::Ocr(OcrError::MissingDocumentUrl), StatusCode::BAD_REQUEST)] + #[case::ocr_response_fault( + Error::Ocr(OcrError::EmptyContent), + StatusCode::INTERNAL_SERVER_ERROR + )] fn status_follows_who_is_at_fault(#[case] error: Error, #[case] status: StatusCode) { assert_eq!(error.status(), status); } + + #[test] + fn a_body_fault_names_the_body_and_keeps_its_cause_once() { + let cause = serde_json::from_str::("{").unwrap_err().to_string(); + let body = invalid_json().openai_body(); + let message = body["error"]["message"].as_str().unwrap(); + assert!( + message.starts_with("invalid request body: invalid JSON: "), + "{message}" + ); + assert_eq!(message.matches(&cause).count(), 1, "{message}"); + assert_eq!(body["error"]["type"], "invalid_request_error"); + assert_eq!(body["error"]["code"], 400); + } + + fn auth_outage() -> Error { + Error::Auth(litellm_gateway_auth::Error::Unavailable( + litellm_error::Source::new(std::io::Error::other("private-backend-detail")), + )) + } + + #[rstest] + #[case::openai(auth_outage().openai_body()["error"]["message"].clone())] + #[case::anthropic(auth_outage().body(None)["error"]["message"].clone())] + #[case::responses_sse(Value::from(auth_outage().responses_sse_frame()))] + #[case::anthropic_sse(Value::from(auth_outage().sse_frame()))] + fn an_auth_failure_hides_its_backend_cause(#[case] rendered: Value) { + let top = auth_outage().to_string(); + let rendered = rendered.to_string(); + assert!(rendered.contains(&top), "{rendered}"); + assert!(!rendered.contains("private-backend-detail"), "{rendered}"); + } + + #[rstest] + #[case::provider_answer(upstream(429, "slow down"))] + #[case::caller_fault(invalid_json())] + #[case::host_bug(litellm_host_http::Error::::Protocol.into())] + fn the_responses_sse_frame_is_the_responses_error_event(#[case] error: Error) { + let frame = error.responses_sse_frame(); + let data = frame + .strip_prefix("event: error\ndata: ") + .and_then(|data| data.strip_suffix("\n\n")) + .unwrap_or_else(|| panic!("an event: error frame, got {frame:?}")); + let event = serde_json::from_str::(data).unwrap(); + assert_eq!(event["type"], "error"); + assert_eq!(event["code"], error.status().as_u16().to_string()); + assert_eq!(event["message"], error.openai_body()["error"]["message"]); + } } diff --git a/litellm-rust/crates/gateway-inference/src/lib.rs b/litellm-rust/crates/gateway-inference/src/lib.rs index 4663151f1de..43966f2bee7 100644 --- a/litellm-rust/crates/gateway-inference/src/lib.rs +++ b/litellm-rust/crates/gateway-inference/src/lib.rs @@ -25,7 +25,7 @@ use litellm_inference_transcription::AudioTranscriptionRoute; use litellm_llms::base_llm::ocr::{handler::OcrClient, settings::OcrSettings}; use litellm_secrets::source::SecretSource; -pub use error::Error; +pub use error::{BodyError, Error}; pub use litellm_router::{Deployment, Router as ModelRouter}; pub use request::{JsonObject, RequestId}; diff --git a/litellm-rust/crates/gateway-inference/src/messages.rs b/litellm-rust/crates/gateway-inference/src/messages.rs index d612a042e03..58ecd40204d 100644 --- a/litellm-rust/crates/gateway-inference/src/messages.rs +++ b/litellm-rust/crates/gateway-inference/src/messages.rs @@ -46,7 +46,7 @@ async fn handle( headers: &HeaderMap, body: Map, ) -> Result { - let deployment = request::resolve_deployment(gateway, &body)?; + let deployment = request::resolve_deployment(gateway, "/v1/messages", &body)?; request::authorize_model(identity, deployment, &body).await?; let (body, cache_options) = crate::caching::prepare(identity, body)?; let route = gateway.messages.clone(); diff --git a/litellm-rust/crates/gateway-inference/src/ocr.rs b/litellm-rust/crates/gateway-inference/src/ocr.rs index 03d064b6bae..f200fd7b141 100644 --- a/litellm-rust/crates/gateway-inference/src/ocr.rs +++ b/litellm-rust/crates/gateway-inference/src/ocr.rs @@ -35,7 +35,7 @@ async fn handle( .get("x-req-format") .and_then(|value| value.to_str().ok()) .map(str::to_owned); - let deployment = request::resolve_deployment(gateway, &body)?; + let deployment = request::resolve_deployment(gateway, "/ocr", &body)?; request::authorize_model(identity, deployment, &body).await?; let document = match upload { Some(upload) => OcrDocumentInput::Bytes { diff --git a/litellm-rust/crates/gateway-inference/src/request.rs b/litellm-rust/crates/gateway-inference/src/request.rs index b398da2013b..d41794242ff 100644 --- a/litellm-rust/crates/gateway-inference/src/request.rs +++ b/litellm-rust/crates/gateway-inference/src/request.rs @@ -6,7 +6,7 @@ use axum::{ }; use serde_json::{Map, Value}; -use crate::{Deployment, Error, Gateway}; +use crate::{BodyError, Deployment, Error, Gateway}; pub(crate) const MAX_FILE_BYTES: usize = 50 * 1024 * 1024; pub(crate) const MAX_BODY_BYTES: usize = MAX_FILE_BYTES + 1024 * 1024; @@ -21,6 +21,19 @@ pub struct JsonObject(pub Map); pub struct RequestId(pub Option); +pub(crate) struct Path(pub T); + +impl FromRequestParts for Path { + type Rejection = Error; + + async fn from_request_parts(parts: &mut Parts, state: &S) -> Result { + axum::extract::Path::::from_request_parts(parts, state) + .await + .map(|axum::extract::Path(value)| Self(value)) + .map_err(Error::InvalidPath) + } +} + impl FromRequestParts for RequestId { type Rejection = std::convert::Infallible; @@ -43,10 +56,10 @@ impl FromRequest for JsonObject { if error.status() == axum::http::StatusCode::PAYLOAD_TOO_LARGE { Error::BodyTooLarge } else { - Error::InvalidBody(error.to_string()) + BodyError::Read(error).into() } })?; - object(&body).map(Self) + Ok(Self(object(&body)?)) } } @@ -77,16 +90,16 @@ impl FromRequest for InferenceBody { } let multipart = Multipart::from_request(request, state) .await - .map_err(|error| Error::InvalidBody(error.to_string()))?; + .map_err(BodyError::Multipart)?; parse_multipart(multipart).await } } -fn object(body: &[u8]) -> Result, Error> { +fn object(body: &[u8]) -> Result, BodyError> { match serde_json::from_slice(body) { Ok(Value::Object(body)) => Ok(body), - Ok(_) => Err(Error::InvalidBody("expected a JSON object".into())), - Err(error) => Err(Error::InvalidBody(error.to_string())), + Ok(_) => Err(BodyError::NotAnObject), + Err(error) => Err(BodyError::Json(error)), } } @@ -98,7 +111,7 @@ pub(crate) async fn authorize_model( let name = body .get("model") .and_then(Value::as_str) - .ok_or_else(|| Error::InvalidBody("model is required".into()))?; + .ok_or(BodyError::MissingModel)?; let access = litellm_gateway_auth::AccessRequest::Model { name: name.into(), deployment: deployment.model.clone(), @@ -112,16 +125,20 @@ pub(crate) async fn authorize_model( pub(crate) fn resolve_deployment<'a>( gateway: &'a Gateway, + route: &'static str, body: &Map, ) -> Result<&'a Deployment, Error> { let model = body .get("model") .and_then(Value::as_str) - .ok_or_else(|| Error::InvalidBody("model is required".into()))?; + .ok_or(BodyError::MissingModel)?; gateway .models .get(model) - .ok_or_else(|| Error::UnknownModel(model.to_owned())) + .ok_or_else(|| Error::UnknownModel { + route, + model: model.to_owned(), + }) } async fn parse_multipart(mut multipart: Multipart) -> Result { @@ -159,7 +176,7 @@ fn multipart_error(error: axum::extract::multipart::MultipartError) -> Error { if error.status() == axum::http::StatusCode::PAYLOAD_TOO_LARGE { return Error::BodyTooLarge; } - Error::InvalidBody(error.to_string()) + BodyError::MultipartField(error).into() } pub(crate) async fn unsupported(OriginalUri(uri): OriginalUri) -> impl IntoResponse { diff --git a/litellm-rust/crates/gateway-inference/src/responses.rs b/litellm-rust/crates/gateway-inference/src/responses.rs index 77bebf610e1..df272460d21 100644 --- a/litellm-rust/crates/gateway-inference/src/responses.rs +++ b/litellm-rust/crates/gateway-inference/src/responses.rs @@ -4,7 +4,6 @@ use axum::{Json, body::Bytes, extract::State, response::Response}; use litellm_gateway_auth::AuthenticatedRequest; use litellm_host_http::Sse; use litellm_inference_responses::{route::Responses, types::ResponsesCall}; -use serde_json::json; use crate::{Error, Gateway, JsonObject, request}; @@ -13,7 +12,7 @@ pub(crate) async fn create( identity: AuthenticatedRequest, JsonObject(body): JsonObject, ) -> Result { - let deployment = request::resolve_deployment(&gateway, &body)?; + let deployment = request::resolve_deployment(&gateway, "/responses", &body)?; request::authorize_model(&identity, deployment, &body).await?; let (body, cache_options) = crate::caching::prepare(&identity, body)?; let route = gateway.responses.clone(); @@ -40,11 +39,7 @@ pub(crate) async fn create( }; let machine = route.machine(call, cache_options.policy); let stream = Sse::::new(Json, |error| { - let error = Error::from(error); - Bytes::from(format!( - "event: error\ndata: {}\n\n", - json!({"type": "error", "code": error.status().as_u16().to_string(), "message": error.to_string(), "param": null}) - )) + Bytes::from(Error::from(error).responses_sse_frame()) }); let headers = crate::caching::CacheHeaders::default(); let response = litellm_host_http::serve(machine, (), headers.clone(), stream, None).await?; diff --git a/litellm-rust/crates/gateway-inference/tests/messages.rs b/litellm-rust/crates/gateway-inference/tests/messages.rs index bf11e67e0fc..7d6d1f160d1 100644 --- a/litellm-rust/crates/gateway-inference/tests/messages.rs +++ b/litellm-rust/crates/gateway-inference/tests/messages.rs @@ -6,7 +6,6 @@ use axum::{ }; use rstest::rstest; use serde_json::json; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tower::ServiceExt; use wiremock::{ Mock, MockServer, ResponseTemplate, @@ -90,34 +89,11 @@ async fn invalid_messages_stays_an_anthropic_error() { assert_eq!(body["error"]["type"], "invalid_request_error"); } -/// Answers with the SSE head and one event, then drops the connection short of the -/// announced body length. -async fn truncating_upstream() -> String { - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://{}", listener.local_addr().unwrap()); - tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut request = vec![0; 4096]; - let _ = socket.read(&mut request).await; - socket - .write_all( - format!( - "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\n\r\n{FIRST_EVENT}", - FIRST_EVENT.len() * 2 - ) - .as_bytes(), - ) - .await - .unwrap(); - }); - base -} - const FIRST_EVENT: &str = "event: message_start\ndata: {}\n\n"; #[tokio::test] async fn a_stream_that_fails_after_opening_ends_with_an_sse_error_frame() { - let base = truncating_upstream().await; + let base = support::truncating_upstream(FIRST_EVENT).await; let request = Request::post("/v1/messages") .header("content-type", "application/json") .body(Body::from( diff --git a/litellm-rust/crates/gateway-inference/tests/responses.rs b/litellm-rust/crates/gateway-inference/tests/responses.rs index ded7d73c5f3..ed7c8e6fe76 100644 --- a/litellm-rust/crates/gateway-inference/tests/responses.rs +++ b/litellm-rust/crates/gateway-inference/tests/responses.rs @@ -116,3 +116,32 @@ async fn responses_authorize_the_public_model_before_provider_execution() { assert_eq!(response.status(), 403); assert!(upstream.received_requests().await.unwrap().is_empty()); } + +const RESPONSE_CREATED: &str = "event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"response-1\",\"model\":\"test-model\",\"output\":[]}}\n\n"; + +#[rstest] +#[tokio::test] +async fn a_stream_that_fails_after_opening_ends_with_the_responses_error_event() { + let base = support::truncating_upstream(RESPONSE_CREATED).await; + let response = support::post( + support::app("openai/test-model", &base), + "/v1/responses", + json!({"model": "public/model", "input": "hello", "stream": true}), + ) + .await; + assert_eq!(response.status(), 200); + let body = to_bytes(response.into_body(), 4096).await.unwrap(); + let text = std::str::from_utf8(&body).unwrap(); + let frame = text + .strip_prefix(RESPONSE_CREATED) + .and_then(|rest| rest.strip_prefix("event: error\ndata: ")) + .unwrap_or_else(|| panic!("the delivered event then one error frame, got {text:?}")); + let error: serde_json::Value = serde_json::from_str(frame.trim_end()).unwrap(); + assert_eq!(error["type"], "error"); + assert_eq!(error["code"], "500"); + assert!( + error["message"] + .as_str() + .is_some_and(|message| !message.is_empty()) + ); +} diff --git a/litellm-rust/crates/gateway-inference/tests/routes.rs b/litellm-rust/crates/gateway-inference/tests/routes.rs index b0699ff10b5..24b8c5cc694 100644 --- a/litellm-rust/crates/gateway-inference/tests/routes.rs +++ b/litellm-rust/crates/gateway-inference/tests/routes.rs @@ -125,7 +125,10 @@ async fn chat_errors_come_from_core( .await; assert_eq!(response.status(), status); let body = support::json(response).await; - assert_eq!(body["error"]["message"], error.to_string()); + assert_eq!( + body["error"]["message"], + litellm_error::Report::new(&error).to_string() + ); assert_eq!(body["error"]["code"], status); assert!(upstream.received_requests().await.unwrap().is_empty()); } @@ -193,17 +196,20 @@ async fn transcription_aliases_reach_core_validation(#[case] path: &str) { } #[rstest] -#[case::no_extension("test")] -#[case::unsupported_extension("test.invalid")] +#[case::no_extension("test", json!(null))] +#[case::unsupported_extension("test.invalid", json!("invalid"))] #[tokio::test] -async fn upload_audio_format_validation_matches_core(#[case] filename: &str) { +async fn upload_audio_format_validation_matches_core( + #[case] filename: &str, + #[case] format: Value, +) { let upstream = MockServer::start().await; let app = support::app("bedrock/test-model", &upstream.uri()); let path = "/v1/audio/transcriptions"; let expected = support::post( app.clone(), path, - json!({"model": "public/model", "audio": {"data": "YWJj", "format": null}}), + json!({"model": "public/model", "audio": {"data": "YWJj", "format": format}}), ) .await; let payload = format!( diff --git a/litellm-rust/crates/gateway-inference/tests/support/mod.rs b/litellm-rust/crates/gateway-inference/tests/support/mod.rs index 0e37bb0aff9..737dc0c830c 100644 --- a/litellm-rust/crates/gateway-inference/tests/support/mod.rs +++ b/litellm-rust/crates/gateway-inference/tests/support/mod.rs @@ -15,6 +15,7 @@ use litellm_http::{HttpClientPool, HttpSettings, Resolution, media::PublicDnsRes use litellm_inference::resources::CoreResources; use litellm_secrets::{SecretValue, source::SecretSource}; use serde_json::Value; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tower::ServiceExt; pub struct NoSecrets; @@ -166,3 +167,26 @@ impl litellm_gateway_auth::IdentityResolver for TestPermissions { }) } } + +/// Answers with the SSE head and one event, then drops the connection short of the +/// announced body length. +pub async fn truncating_upstream(first_event: &'static str) -> String { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = vec![0; 4096]; + let _ = socket.read(&mut request).await; + socket + .write_all( + format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\n\r\n{first_event}", + first_event.len() * 2 + ) + .as_bytes(), + ) + .await + .unwrap(); + }); + base +} diff --git a/litellm-rust/crates/gateway-management/Cargo.toml b/litellm-rust/crates/gateway-management/Cargo.toml index 53ff2ea59aa..502760f8fbe 100644 --- a/litellm-rust/crates/gateway-management/Cargo.toml +++ b/litellm-rust/crates/gateway-management/Cargo.toml @@ -8,6 +8,7 @@ repository.workspace = true [dependencies] base64.workspace = true litellm-auth-types.workspace = true +litellm-error.workspace = true litellm-gateway-auth.workspace = true rand.workspace = true thiserror.workspace = true diff --git a/litellm-rust/crates/gateway-management/src/error.rs b/litellm-rust/crates/gateway-management/src/error.rs index 5da7977c7cb..7194395d668 100644 --- a/litellm-rust/crates/gateway-management/src/error.rs +++ b/litellm-rust/crates/gateway-management/src/error.rs @@ -5,7 +5,7 @@ pub enum Error { #[error("key already exists")] AlreadyExists, #[error("key storage unavailable")] - Storage(#[source] Box), + Storage(#[source] litellm_error::Source), #[error("key generation unavailable")] Entropy(#[source] rand::Error), } diff --git a/litellm-rust/crates/gateway-management/tests/keys.rs b/litellm-rust/crates/gateway-management/tests/keys.rs index dbfedb7ac4d..bf0221cf1be 100644 --- a/litellm-rust/crates/gateway-management/tests/keys.rs +++ b/litellm-rust/crates/gateway-management/tests/keys.rs @@ -9,7 +9,7 @@ use std::{ use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; use litellm_auth_types::SecretValue; use litellm_gateway_auth::{ - KeyError, hash_token, + hash_token, keys::{KeyHash, KeyLookup, KeyStatus, verify_key}, }; use litellm_gateway_management::{ @@ -73,7 +73,7 @@ async fn generated_key_can_be_retrieved_verified_and_revoked( ); assert!(matches!( verify_key(lookup.as_ref(), &generated.token, || now).await, - Err(KeyError::Invalid) + Err(litellm_gateway_auth::Error::InvalidToken) )); assert_eq!(keys.revoke(hash).await.unwrap(), Revocation::Revoked); } @@ -137,9 +137,7 @@ async fn storage_failures_are_returned_for_all_operations(now: SystemTime) { assert!(matches!(keys.get(&hash).await, Err(Error::Storage(_)))); let error = keys.revoke(&hash).await.unwrap_err(); assert!(matches!(&error, Error::Storage(_))); - assert_eq!( - std::error::Error::source(&error).unwrap().to_string(), - "backend detail" - ); + let cause = litellm_error::find::(&error).expect("storage cause"); + assert_eq!(cause.to_string(), "backend detail"); assert!(!error.to_string().contains("backend detail")); } diff --git a/litellm-rust/crates/gateway-management/tests/keys/support.rs b/litellm-rust/crates/gateway-management/tests/keys/support.rs index d0efec4c780..71511c7eaba 100644 --- a/litellm-rust/crates/gateway-management/tests/keys/support.rs +++ b/litellm-rust/crates/gateway-management/tests/keys/support.rs @@ -1,10 +1,7 @@ use std::{collections::BTreeMap, sync::Mutex}; use litellm_gateway_auth::keys::KeyStatus; -use litellm_gateway_auth::{ - KeyError, - keys::{KeyHash, KeyLookup, LookupFuture}, -}; +use litellm_gateway_auth::keys::{KeyHash, KeyLookup, LookupFuture}; use litellm_gateway_management::{ Error, keys::{KeyRecord, KeyStore, NewKey, Revocation, StoreFuture}, @@ -56,7 +53,9 @@ impl KeyLookup for MemoryStore { self.get(hash) .await .map(|record| record.map(|record| record.status)) - .map_err(|error| KeyError::Lookup(Box::new(error))) + .map_err(|error| { + litellm_gateway_auth::Error::Unavailable(litellm_error::Source::new(error)) + }) }) } } @@ -64,9 +63,9 @@ impl KeyLookup for MemoryStore { pub struct Unavailable; fn unavailable() -> Result { - Err(Error::Storage(Box::new(std::io::Error::other( - "backend detail", - )))) + Err(Error::Storage(litellm_error::Source::new( + std::io::Error::other("backend detail"), + ))) } impl KeyStore for Unavailable { diff --git a/litellm-rust/crates/gateway-mcp/AGENTS.md b/litellm-rust/crates/gateway-mcp/AGENTS.md index aa608f1e720..c4b00e3acbf 100644 --- a/litellm-rust/crates/gateway-mcp/AGENTS.md +++ b/litellm-rust/crates/gateway-mcp/AGENTS.md @@ -6,6 +6,8 @@ The HTTP router serves `/mcp`, `/{server}/mcp`, `/mcp/{server}`, `/mcp/{server}/ `NativeGateway` aggregates upstream tools, prompts, resources, and resource templates. Tool and prompt calls resolve registered server prefixes, including prefixes containing hyphens. Resource names receive prefixes while their URIs remain unchanged, and reads require exactly one selected server. Tool allowlists apply to discovery and execution. Upstream pagination is drained with cycle detection. Aggregate tool listings retain healthy servers and expose sanitized outcomes under `_meta["litellm.ai/server_outcomes"]`, while scoped REST listings preserve upstream failures +Every HTTP failure, including extractor rejections and the legacy SSE and session transport checks, is an `Error` with its status from `status()` and the `{"detail"}` envelope. The detail is the `litellm_error::Report` of the error, except for upstream and authorization failures, whose causes can carry upstream URLs and backend detail. Startup failures are `ConnectError`, whose `Debug` redacts the SDK's initialization cause for the same reason + ## Integration Construct `Server` entries from connected SDK peers and pass a `Registry` to `NativeGateway::new`. Keep their `RunningService` handles alive in the host and cancel them during shutdown. SDK peers can use stdio, Streamable HTTP, or another SDK transport. Connection establishment, outbound credentials, HTTP client pools, and reconnection remain owned by the host diff --git a/litellm-rust/crates/gateway-mcp/Cargo.toml b/litellm-rust/crates/gateway-mcp/Cargo.toml index 1db5257448f..bed48bad58d 100644 --- a/litellm-rust/crates/gateway-mcp/Cargo.toml +++ b/litellm-rust/crates/gateway-mcp/Cargo.toml @@ -12,6 +12,7 @@ futures-util.workspace = true litellm-config.workspace = true litellm-secrets.workspace = true litellm-auth-types.workspace = true +litellm-error.workspace = true base64.workspace = true http.workspace = true uuid.workspace = true @@ -25,6 +26,7 @@ serde_json.workspace = true thiserror.workspace = true tokio.workspace = true tokio-util = "0.7" +veil.workspace = true [dev-dependencies] rstest.workspace = true diff --git a/litellm-rust/crates/gateway-mcp/src/configured.rs b/litellm-rust/crates/gateway-mcp/src/configured.rs index bb4f5594f1f..cc7a52fc029 100644 --- a/litellm-rust/crates/gateway-mcp/src/configured.rs +++ b/litellm-rust/crates/gateway-mcp/src/configured.rs @@ -22,25 +22,25 @@ use rmcp::{ use sha2::{Digest, Sha256}; use tokio_util::sync::CancellationToken; -use crate::{ConnectError, Limits, NativeGateway, Registry, Server, ServerInfo}; +use crate::{ConfigFault, ConnectError, Limits, NativeGateway, Registry, Server, ServerInfo}; pub struct ConfiguredGateway { pub operations: Arc, services: Vec>, } -fn invalid(server: &str, message: impl Into) -> ConnectError { +fn invalid(server: &str, fault: ConfigFault) -> ConnectError { ConnectError::Configuration { server: server.into(), - message: message.into(), + fault, } } fn timeout(name: &str, config: &McpServer) -> Result { let duration = Duration::try_from_secs_f64(config.timeout.unwrap_or(30.0)) - .map_err(|_| invalid(name, "timeout must be finite and positive"))?; + .map_err(|source| invalid(name, ConfigFault::InvalidTimeout(source)))?; if duration.is_zero() { - return Err(invalid(name, "timeout must be positive")); + return Err(invalid(name, ConfigFault::ZeroTimeout)); } Ok(duration) } @@ -75,27 +75,24 @@ fn validate(name: &str, config: &McpServer) -> Result<(), ConnectError> { .as_ref() .is_some_and(|alias| alias.trim().is_empty()) { - return Err(invalid(name, "server identifiers must not be blank")); + return Err(invalid(name, ConfigFault::BlankIdentifier)); } if let Some(field) = config.unsupported.keys().next() { return Err(invalid( name, - format!("setting '{field}' is not supported by the Rust MCP host"), + ConfigFault::UnsupportedSetting(field.clone()), )); } if config .max_concurrent_requests .is_some_and(|limit| limit == 0 || limit > tokio::sync::Semaphore::MAX_PERMITS) { - return Err(invalid(name, "max_concurrent_requests is out of range")); + return Err(invalid(name, ConfigFault::ConcurrencyOutOfRange)); } if config.mcp_info.contains_key("protocol_version") || config.mcp_info.contains_key("mcp_server_cost_info") { - return Err(invalid( - name, - "mcp_info protocol and cost settings are not supported", - )); + return Err(invalid(name, ConfigFault::UnsupportedMcpInfo)); } if !matches!( config.auth_type, @@ -108,17 +105,11 @@ fn validate(name: &str, config: &McpServer) -> Result<(), ConnectError> { | McpAuth::Token ) ) { - return Err(invalid( - name, - "configured upstream auth mode is not supported by the Rust MCP host", - )); + return Err(invalid(name, ConfigFault::UnsupportedAuth)); } match config.transport { McpTransport::Sse => { - return Err(invalid( - name, - "legacy SSE upstreams are not supported by the pinned SDK; use transport: http", - )); + return Err(invalid(name, ConfigFault::SseTransport)); } McpTransport::Http => { if config.url.is_none() @@ -126,10 +117,7 @@ fn validate(name: &str, config: &McpServer) -> Result<(), ConnectError> { || !config.args.is_empty() || !config.env.is_empty() { - return Err(invalid( - name, - "HTTP transport requires url and does not accept command, args or env", - )); + return Err(invalid(name, ConfigFault::HttpShape)); } } McpTransport::Stdio => { @@ -139,53 +127,49 @@ fn validate(name: &str, config: &McpServer) -> Result<(), ConnectError> { .is_none_or(|command| command.trim().is_empty()) || config.url.is_some() { - return Err(invalid( - name, - "stdio transport requires command and does not accept url", - )); + return Err(invalid(name, ConfigFault::StdioShape)); } if config.authentication_token.is_some() || !config.static_headers.is_empty() || config.upstream_token_header.is_some() || !matches!(config.auth_type, None | Some(McpAuth::None)) { - return Err(invalid( - name, - "stdio credentials must be supplied through env", - )); + return Err(invalid(name, ConfigFault::StdioCredentials)); } } } if matches!(config.auth_type, None | Some(McpAuth::None)) { if config.authentication_token.is_some() || config.upstream_token_header.is_some() { - return Err(invalid( - name, - "authentication_token and upstream_token_header require an auth_type", - )); + return Err(invalid(name, ConfigFault::TokenWithoutAuth)); } } else if config.authentication_token.is_none() { - return Err(invalid( - name, - "configured auth_type requires authentication_token", - )); + return Err(invalid(name, ConfigFault::AuthWithoutToken)); } Ok(()) } -async fn resolve(value: &SecretValue, secrets: &dyn SecretSource) -> Result { +async fn resolve( + name: &str, + value: &SecretValue, + secrets: &dyn SecretSource, +) -> Result { match value.expose().strip_prefix("os.environ/") { Some(key) => secrets .get_secret_str(key) .await? .map(|secret| secret.expose().to_owned()) - .ok_or_else(|| invalid("configuration", "referenced secret is unavailable")), + .ok_or_else(|| invalid(name, ConfigFault::MissingSecret)), None => Ok(value.expose().to_owned()), } } -fn header(name: &str, value: &str) -> Result<(HeaderName, HeaderValue), ConnectError> { +fn header( + server: &str, + name: &str, + value: &str, +) -> Result<(HeaderName, HeaderValue), ConnectError> { let key = HeaderName::try_from(name) - .map_err(|_| invalid("configuration", "invalid upstream header name"))?; + .map_err(|source| invalid(server, ConfigFault::HeaderName(source)))?; if matches!( key.as_str(), "host" @@ -198,13 +182,10 @@ fn header(name: &str, value: &str) -> Result<(HeaderName, HeaderValue), ConnectE | "mcp-protocol-version" | "last-event-id" ) { - return Err(invalid( - "configuration", - "upstream header conflicts with HTTP framing", - )); + return Err(invalid(server, ConfigFault::FramingHeader)); } let mut value = HeaderValue::try_from(value) - .map_err(|_| invalid("configuration", "invalid upstream header value"))?; + .map_err(|source| invalid(server, ConfigFault::HeaderValue(source)))?; value.set_sensitive(true); Ok((key, value)) } @@ -231,21 +212,19 @@ fn basic(value: &str) -> String { } async fn headers( + name: &str, config: &McpServer, secrets: &dyn SecretSource, ) -> Result, ConnectError> { - let entries = try_join_all( - config - .static_headers - .iter() - .map(|(key, value)| async move { header(key, &resolve(value, secrets).await?) }), - ) + let entries = try_join_all(config.static_headers.iter().map(|(key, value)| async move { + header(name, key, &resolve(name, value, secrets).await?) + })) .await?; let auth = match (&config.authentication_token, config.auth_type) { (Some(token), Some(kind)) => { - let token = resolve(token, secrets).await?; + let token = resolve(name, token, secrets).await?; if token.trim().is_empty() { - return Err(invalid("configuration", "upstream credential is empty")); + return Err(invalid(name, ConfigFault::EmptyCredential)); } let (slot, value) = match kind { McpAuth::ApiKey => ("x-api-key", token), @@ -253,9 +232,10 @@ async fn headers( McpAuth::Token => ("authorization", scheme(&token, "token")), McpAuth::Basic => ("authorization", basic(&token)), McpAuth::Authorization => ("authorization", token), - _ => return Err(invalid("configuration", "unsupported credential mode")), + _ => return Err(invalid(name, ConfigFault::UnsupportedAuth)), }; Some(header( + name, config.upstream_token_header.as_deref().unwrap_or(slot), &value, )?) @@ -349,36 +329,34 @@ async fn connect_service( let configured = config .url .as_ref() - .ok_or_else(|| invalid(name, "HTTP transport requires url"))?; - let url = resolve(configured, secrets).await?; - let parsed = url::Url::parse(&url).map_err(|_| invalid(name, "invalid HTTP URL"))?; + .ok_or_else(|| invalid(name, ConfigFault::HttpShape))?; + let url = resolve(name, configured, secrets).await?; + let parsed = url::Url::parse(&url) + .map_err(|source| invalid(name, ConfigFault::InvalidUrl(source)))?; if !matches!(parsed.scheme(), "http" | "https") || parsed.host_str().is_none() || !parsed.username().is_empty() || parsed.password().is_some() || parsed.fragment().is_some() { - return Err(invalid( - name, - "HTTP URL must use http or https without user credentials or fragments", - )); + return Err(invalid(name, ConfigFault::UrlShape)); } let transport = StreamableHttpClientTransport::with_client( client, StreamableHttpClientTransportConfig::with_uri(url) - .custom_headers(headers(config, secrets).await?), + .custom_headers(headers(name, config, secrets).await?), ); ().serve_with_ct(transport, shutdown).await } McpTransport::Stdio => { let environment = try_join_all(config.env.iter().map(|(key, value)| async move { - Ok::<_, ConnectError>((key, resolve(value, secrets).await?)) + Ok::<_, ConnectError>((key, resolve(name, value, secrets).await?)) })) .await?; let program = config .command .as_deref() - .ok_or_else(|| invalid(name, "stdio requires command"))?; + .ok_or_else(|| invalid(name, ConfigFault::StdioShape))?; let mut command = tokio::process::Command::new(program); command .args(&config.args) @@ -387,7 +365,7 @@ async fn connect_service( ().serve_with_ct(TokioChildProcess::new(command)?, shutdown) .await } - McpTransport::Sse => return Err(invalid(name, "unsupported transport")), + McpTransport::Sse => return Err(invalid(name, ConfigFault::SseTransport)), }; initialized.map_err(|source| ConnectError::Initialize { server: name.into(), @@ -452,7 +430,7 @@ mod tests { "mcp_servers:\n test:\n auth_type: {mode}\n authentication_token: '{token}'\n" )) .unwrap(); - let built = headers(&config.mcp_servers["test"], &Secrets) + let built = headers("test", &config.mcp_servers["test"], &Secrets) .await .unwrap(); let key = HeaderName::try_from(slot).unwrap(); @@ -464,7 +442,7 @@ mod tests { #[tokio::test] async fn explicit_credential_slot_overrides_static_header() { let config = litellm_config::Config::from_yaml("mcp_servers: {test: {auth_type: api_key, authentication_token: os.environ/KEY, upstream_token_header: x-custom, static_headers: {X-Custom: old, x-extra: extra}}}").unwrap(); - let built = headers(&config.mcp_servers["test"], &Secrets) + let built = headers("test", &config.mcp_servers["test"], &Secrets) .await .unwrap(); assert_eq!( diff --git a/litellm-rust/crates/gateway-mcp/src/error.rs b/litellm-rust/crates/gateway-mcp/src/error.rs index 1aa72e34703..9f08abaf26c 100644 --- a/litellm-rust/crates/gateway-mcp/src/error.rs +++ b/litellm-rust/crates/gateway-mcp/src/error.rs @@ -1,58 +1,231 @@ +use std::time::TryFromFloatSecsError; + use axum::{ Json, + extract::rejection::{JsonRejection, PathRejection, QueryRejection}, http::StatusCode, response::{IntoResponse, Response}, }; +use http::header::{InvalidHeaderName, InvalidHeaderValue, ToStrError}; +use litellm_error::{Report, Source}; use rmcp::{ model::{ErrorCode, ErrorData}, - service::ServiceError, + service::{ClientInitializeError, ServiceError}, }; use serde_json::json; #[derive(Debug, thiserror::Error)] pub enum Error { - #[error("{0}")] - InvalidRequest(String), + #[error("{0} is required in request body")] + MissingField(&'static str), + #[error("Gateway listings do not accept an upstream cursor")] + UpstreamCursor, + #[error("read_resource requires exactly one allowed MCP server")] + AmbiguousResourceServer, + #[error("Invalid x-mcp-servers header")] + InvalidServersHeader(#[source] ToStrError), + #[error("invalid JSON body")] + Json(#[source] JsonRejection), + #[error("invalid query string")] + Query(#[source] QueryRejection), + #[error("invalid path")] + Path(#[source] PathRejection), + #[error("request body is too large")] + BodyTooLarge(#[source] axum::Error), #[error("User not allowed to access this MCP server or operation.")] Forbidden, - #[error("{0}")] - Configuration(String), + #[error("request host is not allowed")] + HostNotAllowed, + #[error("request origin is not allowed")] + OriginNotAllowed, + #[error("MCP session belongs to another caller")] + ForeignSession, + #[error("MCP session not found")] + UnknownSession, + #[error("MCP session is busy")] + SessionBusy, #[error("MCP upstream request failed")] Upstream(#[from] ServiceError), + #[error("MCP upstream returned invalid pagination")] + InvalidPagination, #[error("MCP operation returned an unexpected result")] UnexpectedResult, + #[error("MCP authorization unavailable")] + AuthorizationUnavailable(#[source] Source), #[error("MCP request cancelled")] Cancelled, } impl Error { - pub fn into_mcp(self) -> ErrorData { + pub fn status(&self) -> StatusCode { match self { - Self::Upstream(ServiceError::McpError(error)) => error, - Self::InvalidRequest(message) => ErrorData::invalid_params(message, None), - Self::Forbidden => ErrorData::new(ErrorCode(-32003), self.to_string(), None), - _ => ErrorData::internal_error(self.to_string(), None), + Self::MissingField(_) + | Self::UpstreamCursor + | Self::AmbiguousResourceServer + | Self::InvalidServersHeader(_) => StatusCode::BAD_REQUEST, + Self::Json(rejection) => rejection.status(), + Self::Query(rejection) => rejection.status(), + Self::Path(rejection) => rejection.status(), + Self::BodyTooLarge(_) => StatusCode::PAYLOAD_TOO_LARGE, + Self::Forbidden + | Self::HostNotAllowed + | Self::OriginNotAllowed + | Self::ForeignSession => StatusCode::FORBIDDEN, + Self::UnknownSession => StatusCode::NOT_FOUND, + Self::SessionBusy => StatusCode::TOO_MANY_REQUESTS, + Self::Upstream(_) | Self::InvalidPagination => StatusCode::BAD_GATEWAY, + Self::Cancelled => StatusCode::REQUEST_TIMEOUT, + Self::UnexpectedResult | Self::AuthorizationUnavailable(_) => { + StatusCode::INTERNAL_SERVER_ERROR + } + } + } + + fn is_request(&self) -> bool { + match self { + Self::MissingField(_) + | Self::UpstreamCursor + | Self::AmbiguousResourceServer + | Self::InvalidServersHeader(_) + | Self::Json(_) + | Self::Query(_) + | Self::Path(_) + | Self::BodyTooLarge(_) => true, + Self::Forbidden + | Self::HostNotAllowed + | Self::OriginNotAllowed + | Self::ForeignSession + | Self::UnknownSession + | Self::SessionBusy + | Self::Upstream(_) + | Self::InvalidPagination + | Self::UnexpectedResult + | Self::AuthorizationUnavailable(_) + | Self::Cancelled => false, + } + } + + fn detail(&self) -> String { + match self { + Self::Upstream(_) | Self::AuthorizationUnavailable(_) => self.to_string(), + Self::MissingField(_) + | Self::UpstreamCursor + | Self::AmbiguousResourceServer + | Self::InvalidServersHeader(_) + | Self::Json(_) + | Self::Query(_) + | Self::Path(_) + | Self::BodyTooLarge(_) + | Self::Forbidden + | Self::HostNotAllowed + | Self::OriginNotAllowed + | Self::ForeignSession + | Self::UnknownSession + | Self::SessionBusy + | Self::InvalidPagination + | Self::UnexpectedResult + | Self::Cancelled => Report::new(self).to_string(), + } + } + + fn upstream(&self) -> Option<&ServiceError> { + match self { + Self::Upstream(error) => Some(error), + Self::MissingField(_) + | Self::UpstreamCursor + | Self::AmbiguousResourceServer + | Self::InvalidServersHeader(_) + | Self::Json(_) + | Self::Query(_) + | Self::Path(_) + | Self::BodyTooLarge(_) + | Self::Forbidden + | Self::HostNotAllowed + | Self::OriginNotAllowed + | Self::ForeignSession + | Self::UnknownSession + | Self::SessionBusy + | Self::InvalidPagination + | Self::UnexpectedResult + | Self::AuthorizationUnavailable(_) + | Self::Cancelled => None, + } + } + + pub(crate) fn is_upstream(&self) -> bool { + self.upstream().is_some() + } + + pub(crate) fn is_timeout(&self) -> bool { + self.upstream().is_some_and(|error| match error { + ServiceError::Timeout { .. } => true, + ServiceError::McpError(_) + | ServiceError::TransportSend(_) + | ServiceError::TransportClosed + | ServiceError::UnexpectedResponse + | ServiceError::SubscriptionLagged { .. } + | ServiceError::Cancelled { .. } + | ServiceError::InputRequiredRoundsExceeded { .. } => false, + _ => false, + }) + } + + pub(crate) fn is_unreachable(&self) -> bool { + self.upstream().is_some_and(|error| match error { + ServiceError::TransportSend(_) | ServiceError::TransportClosed => true, + ServiceError::McpError(_) + | ServiceError::Timeout { .. } + | ServiceError::UnexpectedResponse + | ServiceError::SubscriptionLagged { .. } + | ServiceError::Cancelled { .. } + | ServiceError::InputRequiredRoundsExceeded { .. } => false, + _ => false, + }) + } + + pub(crate) fn mcp_error(&self) -> Option<&ErrorData> { + self.upstream().and_then(|error| match error { + ServiceError::McpError(error) => Some(error), + ServiceError::TransportSend(_) + | ServiceError::TransportClosed + | ServiceError::Timeout { .. } + | ServiceError::UnexpectedResponse + | ServiceError::SubscriptionLagged { .. } + | ServiceError::Cancelled { .. } + | ServiceError::InputRequiredRoundsExceeded { .. } => None, + _ => None, + }) + } + + pub fn into_mcp(self) -> ErrorData { + if let Some(error) = self.mcp_error() { + return error.clone(); + } + let message = self.detail(); + if self.status() == StatusCode::FORBIDDEN { + ErrorData::new(ErrorCode(-32003), message, None) + } else if self.is_request() { + ErrorData::invalid_params(message, None) + } else { + ErrorData::internal_error(message, None) } } } impl IntoResponse for Error { fn into_response(self) -> Response { - let status = match &self { - Self::InvalidRequest(_) => StatusCode::BAD_REQUEST, - Self::Forbidden => StatusCode::FORBIDDEN, - Self::Upstream(_) => StatusCode::BAD_GATEWAY, - Self::Cancelled => StatusCode::REQUEST_TIMEOUT, - Self::Configuration(_) | Self::UnexpectedResult => StatusCode::INTERNAL_SERVER_ERROR, - }; - (status, Json(json!({"detail": self.to_string()}))).into_response() + (self.status(), Json(json!({"detail": self.detail()}))).into_response() } } -#[derive(thiserror::Error)] +#[derive(thiserror::Error, veil::Redact)] pub enum ConnectError { - #[error("MCP server {server}: {message}")] - Configuration { server: String, message: String }, + #[error("MCP server {server} is misconfigured")] + Configuration { + server: String, + #[source] + fault: ConfigFault, + }, #[error("could not resolve MCP credentials")] Secret(#[from] litellm_secrets::Error), #[error("could not start MCP child process")] @@ -61,38 +234,174 @@ pub enum ConnectError { Initialize { server: String, #[source] - source: Box, + #[redact] + source: Box, }, #[error("MCP upstream {0} initialization timed out")] Timeout(String), #[error("MCP upstream startup cancelled")] Cancelled, - #[error(transparent)] - Registry(#[from] Error), + #[error("MCP server identifiers must not be empty")] + EmptyIdentifier, + #[error("MCP server identifiers and prefixes must be unique")] + DuplicateIdentifier, } -impl std::fmt::Debug for ConnectError { - fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - std::fmt::Display::fmt(self, formatter) - } +#[derive(Debug, thiserror::Error)] +pub enum ConfigFault { + #[error("timeout must be finite and positive")] + InvalidTimeout(#[source] TryFromFloatSecsError), + #[error("timeout must be positive")] + ZeroTimeout, + #[error("server identifiers must not be blank")] + BlankIdentifier, + #[error("setting '{0}' is not supported by the Rust MCP host")] + UnsupportedSetting(String), + #[error("max_concurrent_requests is out of range")] + ConcurrencyOutOfRange, + #[error("mcp_info protocol and cost settings are not supported")] + UnsupportedMcpInfo, + #[error("configured upstream auth mode is not supported by the Rust MCP host")] + UnsupportedAuth, + #[error("legacy SSE upstreams are not supported by the pinned SDK; use transport: http")] + SseTransport, + #[error("HTTP transport requires url and does not accept command, args or env")] + HttpShape, + #[error("stdio transport requires command and does not accept url")] + StdioShape, + #[error("stdio credentials must be supplied through env")] + StdioCredentials, + #[error("authentication_token and upstream_token_header require an auth_type")] + TokenWithoutAuth, + #[error("configured auth_type requires authentication_token")] + AuthWithoutToken, + #[error("referenced secret is unavailable")] + MissingSecret, + #[error("upstream credential is empty")] + EmptyCredential, + #[error("invalid upstream header name")] + HeaderName(#[source] InvalidHeaderName), + #[error("invalid upstream header value")] + HeaderValue(#[source] InvalidHeaderValue), + #[error("upstream header conflicts with HTTP framing")] + FramingHeader, + #[error("invalid HTTP URL")] + InvalidUrl(#[source] url::ParseError), + #[error("HTTP URL must use http or https without user credentials or fragments")] + UrlShape, } #[cfg(test)] mod tests { + use rstest::rstest; + use super::*; - #[rstest::rstest] + const SECRET_URL: &str = "https://example.test/mcp?token=private-secret"; + + #[rstest] fn startup_diagnostics_do_not_expose_upstream_details() { let error = ConnectError::Initialize { server: "docs".into(), - source: Box::new(rmcp::service::ClientInitializeError::ConnectionClosed( - "https://example.test/mcp?token=private-secret".into(), - )), + source: Box::new(ClientInitializeError::ConnectionClosed(SECRET_URL.into())), }; - assert_eq!( - format!("{error:?}"), - "could not initialize MCP upstream docs" - ); + let debug = format!("{error:?}"); + assert!(debug.contains("docs")); + assert!(!debug.contains("private-secret")); assert!(std::error::Error::source(&error).is_some()); } + + #[rstest] + fn configuration_faults_reach_the_startup_diagnostic() { + let error = ConnectError::Configuration { + server: "docs".into(), + fault: ConfigFault::ZeroTimeout, + }; + assert!(format!("{error:?}").contains("ZeroTimeout")); + let report = Report::new(&error).to_string(); + assert!(report.contains("docs"), "{report}"); + assert_eq!( + report.matches("timeout must be positive").count(), + 1, + "{report}" + ); + } + + #[rstest] + #[case::missing_field(Error::MissingField("server_id"), StatusCode::BAD_REQUEST)] + #[case::upstream_cursor(Error::UpstreamCursor, StatusCode::BAD_REQUEST)] + #[case::forbidden(Error::Forbidden, StatusCode::FORBIDDEN)] + #[case::foreign_session(Error::ForeignSession, StatusCode::FORBIDDEN)] + #[case::unknown_session(Error::UnknownSession, StatusCode::NOT_FOUND)] + #[case::busy_session(Error::SessionBusy, StatusCode::TOO_MANY_REQUESTS)] + #[case::invalid_pagination(Error::InvalidPagination, StatusCode::BAD_GATEWAY)] + #[case::upstream( + Error::Upstream(ServiceError::TransportClosed), + StatusCode::BAD_GATEWAY + )] + #[case::cancelled(Error::Cancelled, StatusCode::REQUEST_TIMEOUT)] + fn status_follows_who_is_at_fault(#[case] error: Error, #[case] status: StatusCode) { + assert_eq!(error.status(), status); + } + + fn mcp_failure() -> Error { + Error::Upstream(ServiceError::McpError(ErrorData::invalid_params( + "bad", None, + ))) + } + + #[rstest] + #[case::timeout( + Error::Upstream(ServiceError::Timeout { timeout: std::time::Duration::from_secs(1) }), + true, + false, + false + )] + #[case::closed(Error::Upstream(ServiceError::TransportClosed), false, true, false)] + #[case::mcp_error(mcp_failure(), false, false, true)] + #[case::unexpected_response( + Error::Upstream(ServiceError::UnexpectedResponse), + false, + false, + false + )] + #[case::gateway_cancelled(Error::Cancelled, false, false, false)] + fn upstream_failures_answer_what_went_wrong( + #[case] error: Error, + #[case] timeout: bool, + #[case] unreachable: bool, + #[case] mcp_error: bool, + ) { + assert_eq!(error.is_timeout(), timeout); + assert_eq!(error.is_unreachable(), unreachable); + assert_eq!(error.mcp_error().is_some(), mcp_error); + } + + #[rstest] + fn an_upstream_mcp_error_passes_through_unchanged() { + let mapped = mcp_failure().into_mcp(); + assert_eq!(mapped.code, ErrorCode::INVALID_PARAMS); + assert_eq!(mapped.message, "bad"); + } + + #[rstest] + #[case::request(Error::MissingField("name"), ErrorCode::INVALID_PARAMS)] + #[case::forbidden(Error::Forbidden, ErrorCode(-32003))] + #[case::upstream_fault(Error::InvalidPagination, ErrorCode::INTERNAL_ERROR)] + fn mcp_error_code_follows_the_fault(#[case] error: Error, #[case] code: ErrorCode) { + assert_eq!(error.into_mcp().code, code); + } + + #[rstest] + #[case::upstream(Error::Upstream(ServiceError::Cancelled { + reason: Some(SECRET_URL.into()) + }))] + #[case::authorization(Error::AuthorizationUnavailable(Source::new(std::io::Error::other( + SECRET_URL + ))))] + fn private_causes_stay_out_of_client_messages(#[case] error: Error) { + assert!(Report::new(&error).to_string().contains("private-secret")); + assert!(!error.detail().contains("private-secret")); + assert!(!error.into_mcp().message.contains("private-secret")); + } } diff --git a/litellm-rust/crates/gateway-mcp/src/legacy.rs b/litellm-rust/crates/gateway-mcp/src/legacy.rs index 429fd933135..1673e1b7b19 100644 --- a/litellm-rust/crates/gateway-mcp/src/legacy.rs +++ b/litellm-rust/crates/gateway-mcp/src/legacy.rs @@ -1,8 +1,8 @@ use std::{sync::Arc, time::Duration}; use axum::{ - Json, Router, - extract::{Query, Request, State}, + Router, + extract::{FromRequest, Request, State}, http::{StatusCode, request::Parts}, response::{ IntoResponse, Response, Sse, @@ -18,7 +18,11 @@ use tokio::sync::mpsc; use tokio_util::sync::CancellationToken; use uuid::Uuid; -use crate::{Context, HttpConfig, McpServer, Operations, sessions::fingerprint}; +use crate::{ + Context, Error, HttpConfig, McpServer, Operations, + request::{Json, Query}, + sessions::fingerprint, +}; #[derive(Clone)] struct Session { @@ -63,13 +67,13 @@ pub(crate) fn router( } impl Legacy { - fn validate(&self, parts: &Parts) -> Result<(), StatusCode> { + fn validate(&self, parts: &Parts) -> Result<(), Error> { let host = parts .headers .get("host") .and_then(|value| value.to_str().ok()) .and_then(|value| value.parse::().ok()) - .ok_or(StatusCode::FORBIDDEN)?; + .ok_or(Error::HostNotAllowed)?; let allowed = self.hosts.iter().any(|allowed| { host.as_str().eq_ignore_ascii_case(allowed) || host @@ -78,17 +82,16 @@ impl Legacy { .eq_ignore_ascii_case(allowed) }); if !allowed { - return Err(StatusCode::FORBIDDEN); + return Err(Error::HostNotAllowed); } - if let Some(origin) = parts.headers.get("origin") { - let origin = origin.to_str().map_err(|_| StatusCode::FORBIDDEN)?; - if !self - .origins - .iter() - .any(|allowed| same_origin(origin, allowed)) - { - return Err(StatusCode::FORBIDDEN); - } + if let Some(origin) = parts.headers.get("origin") + && !origin.to_str().is_ok_and(|origin| { + self.origins + .iter() + .any(|allowed| same_origin(origin, allowed)) + }) + { + return Err(Error::OriginNotAllowed); } Ok(()) } @@ -111,27 +114,22 @@ fn same_origin(value: &str, allowed: &str) -> bool { } } -async fn connect(State(state): State, request: Request) -> Response { +async fn connect(State(state): State, request: Request) -> Result { let owner = fingerprint(&request); let (parts, _) = request.into_parts(); - if let Err(status) = state.validate(&parts) { - return status.into_response(); - } - if let Err(error) = state + state.validate(&parts)?; + state .operations .authorize(Context { parts, server: None, mcp: None, }) - .await - { - return error.into_response(); - } + .await?; let connection = start_session(state, owner).await; - Sse::new(connection.events()) + Ok(Sse::new(connection.events()) .keep_alive(KeepAlive::default()) - .into_response() + .into_response()) } struct Connection { @@ -216,29 +214,21 @@ async fn message( State(state): State, Query(query): Query, request: Request, -) -> Response { +) -> Result { let owner = fingerprint(&request); let (parts, body) = request.into_parts(); - if let Err(status) = state.validate(&parts) { - return status.into_response(); - } - let Some(session) = state.sessions.get(&query.session_id).await else { - return StatusCode::NOT_FOUND.into_response(); - }; - if session.owner != owner { - return StatusCode::FORBIDDEN.into_response(); - } - let request = Request::from_parts(parts.clone(), body); - let Json(mut message) = - match as axum::extract::FromRequest<()>>::from_request( - request, - &(), - ) + state.validate(&parts)?; + let session = state + .sessions + .get(&query.session_id) .await - { - Ok(message) => message, - Err(error) => return error.into_response(), - }; + .ok_or(Error::UnknownSession)?; + if session.owner != owner { + return Err(Error::ForeignSession); + } + let Json(mut message) = + Json::::from_request(Request::from_parts(parts.clone(), body), &()) + .await?; match &mut message { ClientJsonRpcMessage::Request(request) => { request.request.extensions_mut().insert(parts); @@ -249,8 +239,8 @@ async fn message( _ => (), } match session.sender.try_send(message) { - Ok(()) => (StatusCode::ACCEPTED, "Accepted").into_response(), - Err(mpsc::error::TrySendError::Full(_)) => StatusCode::TOO_MANY_REQUESTS.into_response(), - Err(mpsc::error::TrySendError::Closed(_)) => StatusCode::NOT_FOUND.into_response(), + Ok(()) => Ok((StatusCode::ACCEPTED, "Accepted").into_response()), + Err(mpsc::error::TrySendError::Full(_)) => Err(Error::SessionBusy), + Err(mpsc::error::TrySendError::Closed(_)) => Err(Error::UnknownSession), } } diff --git a/litellm-rust/crates/gateway-mcp/src/lib.rs b/litellm-rust/crates/gateway-mcp/src/lib.rs index dd72940e24b..fcb7fd98383 100644 --- a/litellm-rust/crates/gateway-mcp/src/lib.rs +++ b/litellm-rust/crates/gateway-mcp/src/lib.rs @@ -3,6 +3,7 @@ mod error; mod legacy; mod native; mod relay; +mod request; mod routes; mod server; mod sessions; @@ -13,7 +14,7 @@ use http::request::Parts; use rmcp::{RoleServer, model::*, service::RequestContext}; pub use configured::ConfiguredGateway; -pub use error::{ConnectError, Error}; +pub use error::{ConfigFault, ConnectError, Error}; pub use native::{Limits, NativeGateway, Registry, Server, ServerResolver}; pub use relay::RelayClient; pub use rmcp; diff --git a/litellm-rust/crates/gateway-mcp/src/native.rs b/litellm-rust/crates/gateway-mcp/src/native.rs index 3a60ce5aca7..7a83f1316e9 100644 --- a/litellm-rust/crates/gateway-mcp/src/native.rs +++ b/litellm-rust/crates/gateway-mcp/src/native.rs @@ -178,14 +178,15 @@ impl NativeGateway { .ok_or(Error::Forbidden)?; let mut upstream = request; upstream.name = tool.tool.name.clone(); - match self + let result = self .send(server, CallToolRequest::new(upstream).into(), &context) - .await - { - Err(Error::Upstream(rmcp::service::ServiceError::McpError(error))) => { - Ok(CallToolResult::error(vec![ContentBlock::text(error.message)]).into()) - } - result => result, + .await; + match result.as_ref().err().and_then(Error::mcp_error) { + Some(error) => Ok(CallToolResult::error(vec![ContentBlock::text( + error.message.clone(), + )]) + .into()), + None => result, } } Operation::ListPrompts(params) => { @@ -321,9 +322,7 @@ impl NativeGateway { return Err(Error::Forbidden); } let [server] = servers.as_ref() else { - return Err(Error::InvalidRequest( - "read_resource requires exactly one allowed MCP server".into(), - )); + return Err(Error::AmbiguousResourceServer); }; self.send(server, ReadResourceRequest::new(request).into(), &context) .await @@ -363,9 +362,7 @@ impl Operations for NativeGateway { fn reject_cursor(params: &Option) -> Result<(), Error> { match params.as_ref().and_then(|params| params.cursor.as_ref()) { - Some(_) => Err(Error::InvalidRequest( - "Gateway listings do not accept an upstream cursor".into(), - )), + Some(_) => Err(Error::UpstreamCursor), None => Ok(()), } } diff --git a/litellm-rust/crates/gateway-mcp/src/native/catalog.rs b/litellm-rust/crates/gateway-mcp/src/native/catalog.rs index 448729ec686..6dd16d1fe63 100644 --- a/litellm-rust/crates/gateway-mcp/src/native/catalog.rs +++ b/litellm-rust/crates/gateway-mcp/src/native/catalog.rs @@ -28,9 +28,7 @@ impl NativeGateway { let state = match next { None => None, Some(cursor) if seen.contains(&cursor) || seen.len() >= 1000 => { - return Err(Error::InvalidRequest( - "MCP upstream returned invalid pagination".into(), - )); + return Err(Error::InvalidPagination); } Some(cursor) => Some(( Some(cursor.clone()), @@ -95,14 +93,9 @@ impl NativeGateway { Ok(tools) => ServerOutcome::Ok { tool_count: tools.len(), }, - Err(Error::Upstream(rmcp::service::ServiceError::Timeout { .. })) => { - ServerOutcome::Timeout - } - Err(Error::Upstream( - rmcp::service::ServiceError::TransportClosed - | rmcp::service::ServiceError::TransportSend(_), - )) => ServerOutcome::Unreachable, - Err(Error::Upstream(_)) => ServerOutcome::UpstreamError, + Err(error) if error.is_timeout() => ServerOutcome::Timeout, + Err(error) if error.is_unreachable() => ServerOutcome::Unreachable, + Err(error) if error.is_upstream() => ServerOutcome::UpstreamError, Err(_) => ServerOutcome::Internal, }; (server.prefix(), outcome) diff --git a/litellm-rust/crates/gateway-mcp/src/native/registry.rs b/litellm-rust/crates/gateway-mcp/src/native/registry.rs index 0aa051d6f32..231737e9772 100644 --- a/litellm-rust/crates/gateway-mcp/src/native/registry.rs +++ b/litellm-rust/crates/gateway-mcp/src/native/registry.rs @@ -1,5 +1,5 @@ use super::NativeGateway; -use crate::{Context, Error, GatewayFuture, ServerInfo}; +use crate::{ConnectError, Context, Error, GatewayFuture, ServerInfo}; use rmcp::{Peer, RoleClient}; use std::sync::Arc; @@ -39,14 +39,12 @@ pub trait ServerResolver: Send + Sync + 'static { pub struct Registry(Arc<[Server]>); impl Registry { - pub fn new(servers: impl Into>) -> Result { + pub fn new(servers: impl Into>) -> Result { let servers = servers.into(); for (index, server) in servers.iter().enumerate() { let identifiers = [&server.info.server_id, &server.info.server_name]; if identifiers.iter().any(|value| value.is_empty()) || server.prefix().is_empty() { - return Err(Error::Configuration( - "MCP server identifiers must not be empty".into(), - )); + return Err(ConnectError::EmptyIdentifier); } if servers[..index].iter().any(|previous| { previous.prefix() == server.prefix() @@ -57,9 +55,7 @@ impl Registry { .as_deref() .is_some_and(|alias| previous.matches(alias)) }) { - return Err(Error::Configuration( - "MCP server identifiers and prefixes must be unique".into(), - )); + return Err(ConnectError::DuplicateIdentifier); } } Ok(Self(servers)) @@ -79,11 +75,7 @@ impl NativeGateway { .parts .headers .get("x-mcp-servers") - .map(|value| { - value - .to_str() - .map_err(|_| Error::InvalidRequest("Invalid x-mcp-servers header".into())) - }) + .map(|value| value.to_str().map_err(Error::InvalidServersHeader)) .transpose()?; let scope = context.server.as_deref().or(header); let Some(scope) = scope else { diff --git a/litellm-rust/crates/gateway-mcp/src/request.rs b/litellm-rust/crates/gateway-mcp/src/request.rs new file mode 100644 index 00000000000..eaa9a6f1840 --- /dev/null +++ b/litellm-rust/crates/gateway-mcp/src/request.rs @@ -0,0 +1,44 @@ +use axum::extract::{FromRequest, FromRequestParts, Request}; +use http::request::Parts; +use serde::de::DeserializeOwned; + +use crate::Error; + +pub(crate) struct Json(pub T); + +impl FromRequest for Json { + type Rejection = Error; + + async fn from_request(request: Request, state: &S) -> Result { + axum::Json::::from_request(request, state) + .await + .map(|axum::Json(value)| Self(value)) + .map_err(Error::Json) + } +} + +pub(crate) struct Query(pub T); + +impl FromRequestParts for Query { + type Rejection = Error; + + async fn from_request_parts(parts: &mut Parts, state: &S) -> Result { + axum::extract::Query::::from_request_parts(parts, state) + .await + .map(|axum::extract::Query(value)| Self(value)) + .map_err(Error::Query) + } +} + +pub(crate) struct Path(pub T); + +impl FromRequestParts for Path { + type Rejection = Error; + + async fn from_request_parts(parts: &mut Parts, state: &S) -> Result { + axum::extract::Path::::from_request_parts(parts, state) + .await + .map(|axum::extract::Path(value)| Self(value)) + .map_err(Error::Path) + } +} diff --git a/litellm-rust/crates/gateway-mcp/src/routes.rs b/litellm-rust/crates/gateway-mcp/src/routes.rs index b5d314e4882..56d67a5bc77 100644 --- a/litellm-rust/crates/gateway-mcp/src/routes.rs +++ b/litellm-rust/crates/gateway-mcp/src/routes.rs @@ -2,7 +2,7 @@ use std::sync::Arc; use axum::{ Json, Router, - extract::{Path, Query, Request, State}, + extract::{Request, State}, http::request::Parts, middleware::{self, Next}, response::{IntoResponse, Response}, @@ -12,7 +12,11 @@ use rmcp::{model::*, transport::streamable_http_server::StreamableHttpServerConf use serde::{Deserialize, Serialize}; use tokio_util::sync::CancellationToken; -use crate::{Context, Error, McpServer, Operation, Operations, ToolCatalog, server::ServerScope}; +use crate::{ + Context, Error, McpServer, Operation, Operations, ToolCatalog, + request::{self, Path, Query}, + server::ServerScope, +}; pub struct HttpConfig { pub allowed_hosts: Vec, @@ -132,16 +136,16 @@ struct ToolCall { async fn rest_call( State(operations): State>, parts: Parts, - Json(call): Json, + request::Json(call): request::Json, ) -> Result { let server = call .server_id .filter(|id| !id.is_empty()) - .ok_or_else(|| Error::InvalidRequest("server_id is required in request body".into()))?; + .ok_or(Error::MissingField("server_id"))?; let name = call .name .filter(|name| !name.is_empty()) - .ok_or_else(|| Error::InvalidRequest("name is required in request body".into()))?; + .ok_or(Error::MissingField("name"))?; let params = CallToolRequestParams::new(name).with_arguments(call.arguments.unwrap_or_default()); let context = Context { diff --git a/litellm-rust/crates/gateway-mcp/src/sessions.rs b/litellm-rust/crates/gateway-mcp/src/sessions.rs index 54d0634bd3d..583a8b86341 100644 --- a/litellm-rust/crates/gateway-mcp/src/sessions.rs +++ b/litellm-rust/crates/gateway-mcp/src/sessions.rs @@ -4,7 +4,7 @@ use axum::{ Router, body::{Body, to_bytes}, extract::{Request, State}, - http::{Method, StatusCode}, + http::Method, response::IntoResponse, routing::any, }; @@ -19,7 +19,7 @@ use rmcp::{ use sha2::{Digest, Sha256}; use tower::ServiceExt; -use crate::{McpServer, server::ServerScope}; +use crate::{Error, McpServer, server::ServerScope}; #[derive(Clone)] pub struct SessionOwner(pub String); @@ -65,7 +65,7 @@ pub(crate) fn transport(server: McpServer, config: StreamableHttpServerConfig) - async fn dispatch( State(transport): State, request: Request, -) -> Result { +) -> Result { let owner = fingerprint(&request); let session = request .headers() @@ -75,8 +75,8 @@ async fn dispatch( if let Some(session) = &session { match transport.owners.get(session).await { Some(expected) if expected == owner => (), - Some(_) => return Err(StatusCode::FORBIDDEN), - None => return Err(StatusCode::NOT_FOUND), + Some(_) => return Err(Error::ForeignSession), + None => return Err(Error::UnknownSession), } } let method = request.method().clone(); @@ -84,7 +84,7 @@ async fn dispatch( let (parts, body) = request.into_parts(); let bytes = to_bytes(body, 4 * 1024 * 1024) .await - .map_err(|_| StatusCode::PAYLOAD_TOO_LARGE)?; + .map_err(Error::BodyTooLarge)?; let initialize = matches!(serde_json::from_slice::(&bytes), Ok(ClientJsonRpcMessage::Request(request)) if matches!(request.request, ClientRequest::InitializeRequest(_))); (Request::from_parts(parts, Body::from(bytes)), initialize) diff --git a/litellm-rust/crates/gateway-mcp/tests/gateway/http.rs b/litellm-rust/crates/gateway-mcp/tests/gateway/http.rs index 5feab0194d3..c5060ebe8f3 100644 --- a/litellm-rust/crates/gateway-mcp/tests/gateway/http.rs +++ b/litellm-rust/crates/gateway-mcp/tests/gateway/http.rs @@ -171,43 +171,44 @@ async fn resolver_receives_each_requests_identity( Duration::from_secs(2), )); let app = router(operations, HttpConfig::default()); - for token in ["allowed", "denied", "allowed"] { - let response = app - .clone() - .oneshot( - Request::builder() - .method(method) - .uri(path) - .header("host", "localhost") - .header("authorization", format!("Bearer {token}")) - .header("content-type", "application/json") - .header("accept", "application/json, text/event-stream") - .header("mcp-protocol-version", "2025-11-25") - .body(Body::from( - json!({"jsonrpc":"2.0","id":1,"method":"tools/list"}).to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - let status = response.status(); - let body = String::from_utf8( - to_bytes(response.into_body(), 65536) - .await - .unwrap() - .to_vec(), + let (status, body) = list_as(&app, method, path, "allowed").await; + assert_eq!(status, StatusCode::OK); + assert!(body.contains("echo-value"), "{body}"); + let (status, body) = list_as(&app, method, path, "denied").await; + assert_eq!(status == StatusCode::FORBIDDEN, method == "GET"); + assert!(body.contains("not allowed"), "{body}"); + let (status, body) = list_as(&app, method, path, "allowed").await; + assert_eq!(status, StatusCode::OK); + assert!(body.contains("echo-value"), "{body}"); +} + +async fn list_as( + app: &axum::Router, + method: &str, + path: &str, + token: &str, +) -> (StatusCode, String) { + let response = app + .clone() + .oneshot( + Request::builder() + .method(method) + .uri(path) + .header("host", "localhost") + .header("authorization", format!("Bearer {token}")) + .header("content-type", "application/json") + .header("accept", "application/json, text/event-stream") + .header("mcp-protocol-version", "2025-11-25") + .body(Body::from( + json!({"jsonrpc":"2.0","id":1,"method":"tools/list"}).to_string(), + )) + .unwrap(), ) + .await .unwrap(); - if token == "denied" { - if method == "GET" { - assert_eq!(status, StatusCode::FORBIDDEN); - } - assert!(body.contains("not allowed"), "{body}"); - } else { - assert_eq!(status, StatusCode::OK); - assert!(body.contains("echo-value"), "{body}"); - } - } + let status = response.status(); + let body = to_bytes(response.into_body(), 65536).await.unwrap(); + (status, String::from_utf8(body.to_vec()).unwrap()) } fn rpc_request(path: &str, token: &str, session: Option<&str>, body: Value) -> Request { @@ -352,10 +353,14 @@ async fn header_scope_cannot_broaden_path_scope( } #[rstest] -#[case::missing_server(json!({"name":"echo-value"}))] -#[case::missing_name(json!({"server_id":"alpha"}))] +#[case::missing_server(json!({"name":"echo-value"}), "server_id")] +#[case::missing_name(json!({"server_id":"alpha"}), "name")] #[tokio::test] -async fn rest_requires_server_and_tool(#[future] harness: Harness, #[case] body: Value) { +async fn rest_requires_server_and_tool( + #[future] harness: Harness, + #[case] body: Value, + #[case] field: &str, +) { let harness = harness.await; let response = harness .app @@ -368,6 +373,40 @@ async fn rest_requires_server_and_tool(#[future] harness: Harness, #[case] body: .await .unwrap(); assert_eq!(response.status(), StatusCode::BAD_REQUEST); + let body: Value = + serde_json::from_slice(&to_bytes(response.into_body(), 65536).await.unwrap()).unwrap(); + assert!(body["detail"].as_str().unwrap().contains(field), "{body}"); +} + +#[rstest] +#[case::malformed_call("/mcp-rest/tools/call", "POST", "{")] +#[case::unknown_list_parameter("/mcp-rest/tools/list?unknown=1", "GET", "")] +#[case::legacy_message_without_session("/mcp/sse/messages", "POST", "{}")] +#[tokio::test] +async fn rejected_requests_use_the_detail_envelope( + #[future] harness: Harness, + #[case] path: &str, + #[case] method: &str, + #[case] body: &'static str, +) { + let harness = harness.await; + let response = harness + .app + .oneshot( + Request::builder() + .method(method) + .uri(path) + .header("host", "localhost") + .header("content-type", "application/json") + .body(Body::from(body)) + .unwrap(), + ) + .await + .unwrap(); + assert!(response.status().is_client_error(), "{}", response.status()); + let body: Value = + serde_json::from_slice(&to_bytes(response.into_body(), 65536).await.unwrap()).unwrap(); + assert!(body["detail"].is_string(), "{body}"); } #[rstest] diff --git a/litellm-rust/crates/gateway-mcp/tests/gateway/native.rs b/litellm-rust/crates/gateway-mcp/tests/gateway/native.rs index aca19b38d6e..54fdabc0da7 100644 --- a/litellm-rust/crates/gateway-mcp/tests/gateway/native.rs +++ b/litellm-rust/crates/gateway-mcp/tests/gateway/native.rs @@ -148,10 +148,12 @@ async fn routes_prompts_and_resources_without_rewriting_uris(#[future] harness: "fixture://{file}" ); let read = Operation::ReadResource(ReadResourceRequestParams::new("fixture://document")); - assert!(matches!( - harness.gateway.execute(read.clone(), context(None)).await, - Err(Error::InvalidRequest(_)) - )); + let ambiguous = harness + .gateway + .execute(read.clone(), context(None)) + .await + .expect_err("read without a server scope"); + assert_eq!(ambiguous.status(), axum::http::StatusCode::BAD_REQUEST); let result = harness .gateway .execute(read, context(Some("alpha"))) diff --git a/litellm-rust/crates/gateway-ui/Cargo.toml b/litellm-rust/crates/gateway-ui/Cargo.toml index 94546ef9118..1c65665be72 100644 --- a/litellm-rust/crates/gateway-ui/Cargo.toml +++ b/litellm-rust/crates/gateway-ui/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-error.workspace = true axum = { workspace = true, features = ["json", "original-uri", "query"] } axum-login.workspace = true base64.workspace = true @@ -22,6 +23,7 @@ tower-http = { version = "0.6.11", features = ["fs", "set-header"] } tower-sessions.workspace = true [dev-dependencies] +futures-util.workspace = true serde_json.workspace = true tower = { version = "0.5", features = ["util"] } litellm-auth-types.workspace = true diff --git a/litellm-rust/crates/gateway-ui/src/dashboard.rs b/litellm-rust/crates/gateway-ui/src/dashboard.rs index 2642506fde6..04ec8490411 100644 --- a/litellm-rust/crates/gateway-ui/src/dashboard.rs +++ b/litellm-rust/crates/gateway-ui/src/dashboard.rs @@ -2,7 +2,8 @@ use std::path::Path; use axum::{ Router, - extract::{Query, Request}, + extract::{FromRequestParts, Query, Request}, + http::request::Parts, routing::get, }; use serde::{Deserialize, Serialize}; @@ -34,6 +35,17 @@ struct LogoQuery { variant: Option, } +impl FromRequestParts for LogoQuery { + type Rejection = crate::Error; + + async fn from_request_parts(parts: &mut Parts, state: &S) -> Result { + Query::::from_request_parts(parts, state) + .await + .map(|Query(query)| query) + .map_err(crate::Error::Query) + } +} + fn logo_file(query: &LogoQuery) -> &'static str { match (query.variant, query.theme) { (Some(Variant::Monogram), Some(Theme::Dark)) => "assets/logos/litellm_monogram_dark.svg", @@ -59,7 +71,7 @@ pub fn dashboard_assets(directory: impl AsRef) -> Router { ) .route( "/get_image", - get(move |Query(query): Query, request: Request| { + get(move |query: LogoQuery, request: Request| { ServeFile::new(logos.join(logo_file(&query))).oneshot(request) }), ) diff --git a/litellm-rust/crates/gateway-ui/src/error.rs b/litellm-rust/crates/gateway-ui/src/error.rs index 16e6e205d99..cce01ba76ce 100644 --- a/litellm-rust/crates/gateway-ui/src/error.rs +++ b/litellm-rust/crates/gateway-ui/src/error.rs @@ -1,9 +1,9 @@ use axum::{ Json, + extract::rejection::{JsonRejection, QueryRejection}, http::{StatusCode, header}, response::{IntoResponse, Response}, }; -use litellm_gateway_auth::UiAuthError; #[derive(serde::Serialize)] struct ErrorBody { @@ -21,39 +21,108 @@ pub enum Error { InvalidCredentials, #[error("Too many login attempts, try again in one minute")] RateLimited, - #[error("UI authentication unavailable")] - Auth(#[from] axum_login::Error), + #[error("invalid login request")] + Body(#[source] JsonRejection), + #[error("invalid query string")] + Query(#[source] QueryRejection), + #[error("UI login unavailable")] + Login(#[from] axum_login::Error), #[error("UI session unavailable")] Session(#[from] tower_sessions::session::Error), #[error("UI token unavailable")] Token(#[from] jsonwebtoken::errors::Error), - #[error("Invalid UI session")] - Unauthorized(#[from] UiAuthError), + #[error(transparent)] + Auth(#[from] litellm_gateway_auth::Error), +} + +impl Error { + pub fn status(&self) -> StatusCode { + match self { + Self::InvalidCredentials => StatusCode::UNAUTHORIZED, + Self::RateLimited => StatusCode::TOO_MANY_REQUESTS, + Self::Body(rejection) => rejection.status(), + Self::Query(rejection) => rejection.status(), + Self::Auth(error) => error.status(), + Self::Login(_) | Self::Session(_) | Self::Token(_) => StatusCode::INTERNAL_SERVER_ERROR, + } + } +} + +impl Error { + fn message(&self) -> String { + match self { + Self::Body(_) | Self::Query(_) => litellm_error::Report::new(self).to_string(), + Self::InvalidCredentials + | Self::RateLimited + | Self::Login(_) + | Self::Session(_) + | Self::Token(_) + | Self::Auth(_) => self.to_string(), + } + } } impl IntoResponse for Error { fn into_response(self) -> Response { - let status = match &self { - Self::InvalidCredentials | Self::Unauthorized(UiAuthError::Unauthorized) => { - StatusCode::UNAUTHORIZED + let status = self.status(); + let body = Json(ErrorBody { + error: ErrorMessage { + message: self.message(), + }, + }); + match status { + StatusCode::TOO_MANY_REQUESTS => { + (status, [(header::RETRY_AFTER, "60")], body).into_response() } - Self::RateLimited => StatusCode::TOO_MANY_REQUESTS, - _ => StatusCode::INTERNAL_SERVER_ERROR, - }; - let response = ( - status, - Json(ErrorBody { - error: ErrorMessage { - message: self.to_string(), - }, - }), - ) - .into_response(); - - if matches!(self, Self::RateLimited) { - return ([(header::RETRY_AFTER, "60")], response).into_response(); + _ => (status, body).into_response(), } - - response + } +} + +#[cfg(test)] +mod tests { + use litellm_gateway_auth::Error as AuthError; + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::rejected_session(Error::Auth(AuthError::InvalidSession), StatusCode::UNAUTHORIZED)] + #[case::session_layer_missing( + Error::Auth(AuthError::MissingSessionLayer), + StatusCode::INTERNAL_SERVER_ERROR + )] + #[case::operation_outside_session_scope( + Error::Auth(AuthError::Forbidden), + StatusCode::FORBIDDEN + )] + #[case::rate_limited(Error::RateLimited, StatusCode::TOO_MANY_REQUESTS)] + fn status_follows_who_is_at_fault(#[case] error: Error, #[case] status: StatusCode) { + assert_eq!(error.status(), status); + } + + #[rstest] + #[case::rate_limited(Error::RateLimited, true)] + #[case::invalid_credentials(Error::InvalidCredentials, false)] + fn only_rate_limits_ask_to_retry_later(#[case] error: Error, #[case] retry_after: bool) { + assert_eq!( + error + .into_response() + .headers() + .contains_key(header::RETRY_AFTER), + retry_after + ); + } + + #[tokio::test] + async fn a_body_never_shows_what_the_backend_said() { + let backend = std::io::Error::other("vault token rejected"); + let response = Error::Auth(AuthError::Unavailable(litellm_error::Source::new(backend))) + .into_response(); + let body = axum::body::to_bytes(response.into_body(), 4096) + .await + .unwrap(); + let body = String::from_utf8(body.to_vec()).unwrap(); + assert!(!body.contains("vault token rejected"), "{body}"); } } diff --git a/litellm-rust/crates/gateway-ui/src/session.rs b/litellm-rust/crates/gateway-ui/src/session.rs index b80f40bb376..b5d847ad353 100644 --- a/litellm-rust/crates/gateway-ui/src/session.rs +++ b/litellm-rust/crates/gateway-ui/src/session.rs @@ -4,14 +4,18 @@ use serde::Serialize; use std::sync::Arc; use time::{Duration, OffsetDateTime}; -use axum::{Json, extract::State as ExtractState, response::IntoResponse}; +use axum::{ + Json, + extract::{FromRequest, Request, State as ExtractState}, + response::IntoResponse, +}; use governor::DefaultDirectRateLimiter; use tower_cookies::{Cookie, Cookies}; use tower_sessions::{Expiry, Session, cookie::SameSite}; use crate::Error; use litellm_gateway_auth::{ - AccessRequest, UI_CSRF_KEY, UiAction, UiAuthError, UiAuthSession, UiCredentials, UiSession, + AccessRequest, UI_CSRF_KEY, UiAction, UiAuthSession, UiCredentials, UiSession, }; pub struct State { @@ -60,12 +64,25 @@ struct Discovery { is_control_plane: bool, } +pub struct Credentials(UiCredentials); + +impl FromRequest for Credentials { + type Rejection = Error; + + async fn from_request(request: Request, state: &S) -> Result { + Json::::from_request(request, state) + .await + .map(|Json(credentials)| Self(credentials)) + .map_err(Error::Body) + } +} + pub async fn login( ExtractState(state): ExtractState>, mut auth: UiAuthSession, session: Session, cookies: Cookies, - Json(credentials): Json, + Credentials(credentials): Credentials, ) -> Result { state.login_limit.check().map_err(|_| Error::RateLimited)?; let user = auth @@ -107,13 +124,14 @@ pub async fn login( })) } -pub async fn info(session: Result) -> Result { +pub async fn info( + session: Result, +) -> Result { let session = session?; session .identity .authorize(AccessRequest::Ui(UiAction::SessionInfo)) - .await - .map_err(|_| UiAuthError::Unauthorized)?; + .await?; Ok(Json(SessionInfo { user_id: session.user.username, user_role: "proxy_admin", @@ -121,7 +139,7 @@ pub async fn info(session: Result) -> Result, + validated: Result, mut auth: UiAuthSession, cookies: Cookies, ) -> Result { @@ -129,8 +147,7 @@ pub async fn logout( validated .identity .authorize(AccessRequest::Ui(UiAction::Logout)) - .await - .map_err(|_| UiAuthError::Unauthorized)?; + .await?; auth.logout().await?; cookies.add( Cookie::build(("token", "")) diff --git a/litellm-rust/crates/gateway-ui/tests/assets.rs b/litellm-rust/crates/gateway-ui/tests/assets.rs index ba741fbf2ec..e2819033f95 100644 --- a/litellm-rust/crates/gateway-ui/tests/assets.rs +++ b/litellm-rust/crates/gateway-ui/tests/assets.rs @@ -200,6 +200,24 @@ async fn logo_discovery_points_to_served_image(dashboard: App) { ); } +#[rstest] +#[tokio::test] +async fn unknown_logo_theme_is_rejected_in_the_ui_envelope(dashboard: App) { + let response = dashboard + .router + .oneshot( + Request::get("/get_image?theme=purple") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + let body: Value = + serde_json::from_slice(&to_bytes(response.into_body(), 65536).await.unwrap()).unwrap(); + assert!(body["error"]["message"].as_str().unwrap().contains("theme")); +} + #[rstest] #[case::logo("/get_image")] #[case::logo_dark("/get_image?theme=dark")] diff --git a/litellm-rust/crates/gateway-ui/tests/routes.rs b/litellm-rust/crates/gateway-ui/tests/routes.rs index 4a3238f79c2..95246f33a35 100644 --- a/litellm-rust/crates/gateway-ui/tests/routes.rs +++ b/litellm-rust/crates/gateway-ui/tests/routes.rs @@ -330,24 +330,45 @@ async fn reauthentication_rotates_and_revokes_the_previous_session(app: App) { ); } +async fn failed_login(app: &App, username: &str) -> Response { + app.router + .clone() + .oneshot( + Request::post("/v2/login") + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from( + json!({"username": username, "password": "wrong-password"}).to_string(), + )) + .unwrap(), + ) + .await + .unwrap() +} + #[rstest] #[tokio::test] async fn failed_logins_are_generic_and_rate_limited(app: App) { - for attempt in 0..11 { - let response = app.router.clone().oneshot(Request::post("/v2/login") - .header(header::CONTENT_TYPE, "application/json") - .body(Body::from(json!({"username": format!("unknown-{attempt}"), "password": "wrong-password"}).to_string())).unwrap()).await.unwrap(); - if attempt < 10 { - assert_eq!(response.status(), StatusCode::UNAUTHORIZED); - assert_eq!( - body(response).await["error"]["message"], - "Invalid username or password" - ); - } else { - assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); - assert!(response.headers().contains_key(header::RETRY_AFTER)); - } - } + let first = failed_login(&app, "unknown-0").await; + assert_eq!(first.status(), StatusCode::UNAUTHORIZED); + assert_eq!( + body(first).await["error"]["message"], + "Invalid username or password" + ); + let usernames: Vec = (1..10) + .map(|attempt| format!("unknown-{attempt}")) + .collect(); + let burst = futures_util::future::join_all( + usernames + .iter() + .map(|username| failed_login(&app, username)), + ) + .await; + let statuses: Vec = burst.iter().map(Response::status).collect(); + assert_eq!(statuses, vec![StatusCode::UNAUTHORIZED; 9]); + let limited = failed_login(&app, "unknown-10").await; + assert_eq!(limited.status(), StatusCode::TOO_MANY_REQUESTS); + assert!(limited.headers().contains_key(header::RETRY_AFTER)); + assert!(body(limited).await["error"]["message"].is_string()); } #[rstest] @@ -365,6 +386,7 @@ async fn login_rejects_browser_form_posts(app: App) { .unwrap(); assert_eq!(response.status(), StatusCode::UNSUPPORTED_MEDIA_TYPE); assert!(!response.headers().contains_key(header::SET_COOKIE)); + assert!(body(response).await["error"]["message"].is_string()); } #[rstest] @@ -468,6 +490,7 @@ async fn invalid_login_payloads_do_not_create_sessions( .unwrap(); assert_eq!(response.status(), status); assert!(!response.headers().contains_key(header::SET_COOKIE)); + assert!(body(response).await["error"]["message"].is_string()); } #[rstest] diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index 4982c70b18c..4d831aa2314 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -19,6 +19,7 @@ litellm-gateway-auth.workspace = true litellm-gateway-ui.workspace = true litellm-auth-types.workspace = true litellm-config.workspace = true +litellm-error.workspace = true litellm-circuit-breaker.workspace = true litellm-redis-client.workspace = true litellm-http.workspace = true diff --git a/litellm-rust/crates/gateway/src/auth.rs b/litellm-rust/crates/gateway/src/auth.rs index 72892be783c..d5d8a229d0d 100644 --- a/litellm-rust/crates/gateway/src/auth.rs +++ b/litellm-rust/crates/gateway/src/auth.rs @@ -64,10 +64,8 @@ impl OperationAuthorizer for McpPolicy { } fn mcp_error(error: Error) -> litellm_gateway_mcp::Error { - match error { - Error::InvalidToken | Error::Expired | Error::Forbidden => { - litellm_gateway_mcp::Error::Forbidden - } - _ => litellm_gateway_mcp::Error::Configuration("MCP authorization unavailable".into()), + if error.is_denied() { + return litellm_gateway_mcp::Error::Forbidden; } + litellm_gateway_mcp::Error::AuthorizationUnavailable(litellm_error::Source::new(error)) } diff --git a/litellm-rust/crates/gateway/src/error.rs b/litellm-rust/crates/gateway/src/error.rs index 170a8284c18..e2f646c5448 100644 --- a/litellm-rust/crates/gateway/src/error.rs +++ b/litellm-rust/crates/gateway/src/error.rs @@ -6,6 +6,8 @@ pub enum Error { Environment, #[error("MCP host does not yet support configured setting {0}")] McpSetting(String), + #[error("MCP host does not yet support guardrails or policies")] + McpGuardrails, #[error(transparent)] Mcp(#[from] litellm_gateway_mcp::ConnectError), #[error(transparent)] @@ -17,6 +19,6 @@ pub enum Error { }, #[error(transparent)] RedisConnect(#[from] litellm_redis_client::ConnectError), - #[error("MCP requires a configured master key")] + #[error("MCP requires working gateway authentication")] Auth(#[from] litellm_gateway_auth::Error), } diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index d7567b83cc1..b8ff8e46033 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -109,7 +109,7 @@ pub async fn build_mcp( || !config.policies.is_empty() || !config.policy_attachments.is_empty() { - return Err(Error::McpSetting("guardrails or policies".into())); + return Err(Error::McpGuardrails); } Auth::from_config(config, secrets.clone()) .validate() diff --git a/litellm-rust/crates/gateway/tests/mcp.rs b/litellm-rust/crates/gateway/tests/mcp.rs index ee8a17ae251..3d12efd730e 100644 --- a/litellm-rust/crates/gateway/tests/mcp.rs +++ b/litellm-rust/crates/gateway/tests/mcp.rs @@ -97,11 +97,29 @@ async fn upstream() -> Fixture { } } -#[rstest] -#[tokio::test] -async fn serves_configured_upstream_with_separate_admission_and_outbound_credentials( - #[future] upstream: Fixture, -) { +struct Served { + app: Router, + mcp: litellm_gateway_mcp::ConfiguredGateway, + shutdown: CancellationToken, + upstream: Fixture, +} + +impl Served { + async fn close(self) { + self.shutdown.cancel(); + tokio::time::timeout(Duration::from_secs(5), self.mcp.close()) + .await + .unwrap(); + self.upstream.shutdown.cancel(); + tokio::time::timeout(Duration::from_secs(5), self.upstream.task) + .await + .unwrap() + .unwrap(); + } +} + +#[fixture] +async fn served(#[future] upstream: Fixture) -> Served { let upstream = upstream.await; let config = Config::from_yaml(&format!( r#" @@ -148,30 +166,53 @@ mcp_servers: }, ); let app = litellm_gateway::router(inference, &config, None, Some(mcp_router)); - - for (path, method) in [ - ("/mcp", "POST"), - ("/mcp", "GET"), - ("/mcp", "DELETE"), - ("/mcp/sse", "GET"), - ("/mcp/sse/messages", "POST"), - ("/mcp-rest/tools/list", "GET"), - ("/mcp/enabled", "GET"), - ] { - let response = app - .clone() - .oneshot( - Request::builder() - .method(method) - .uri(path) - .header("host", "localhost") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), 401, "{method} {path}"); + Served { + app, + mcp, + shutdown, + upstream, } +} + +#[rstest] +#[case::streamable_post("/mcp", "POST")] +#[case::streamable_get("/mcp", "GET")] +#[case::streamable_delete("/mcp", "DELETE")] +#[case::legacy_sse("/mcp/sse", "GET")] +#[case::legacy_messages("/mcp/sse/messages", "POST")] +#[case::rest_list("/mcp-rest/tools/list", "GET")] +#[case::enabled("/mcp/enabled", "GET")] +#[tokio::test] +async fn every_mcp_route_requires_admission_credentials( + #[future] served: Served, + #[case] path: &str, + #[case] method: &str, +) { + let served = served.await; + let response = served + .app + .clone() + .oneshot( + Request::builder() + .method(method) + .uri(path) + .header("host", "localhost") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), 401); + served.close().await; +} + +#[rstest] +#[tokio::test] +async fn serves_configured_upstream_with_separate_admission_and_outbound_credentials( + #[future] served: Served, +) { + let served = served.await; + let app = served.app.clone(); let list = app .clone() .oneshot( @@ -213,15 +254,7 @@ mcp_servers: .unwrap(); assert_eq!(blocked.status(), 403); - shutdown.cancel(); - tokio::time::timeout(Duration::from_secs(5), mcp.close()) - .await - .unwrap(); - upstream.shutdown.cancel(); - tokio::time::timeout(Duration::from_secs(5), upstream.task) - .await - .unwrap() - .unwrap(); + served.close().await; } #[rstest] @@ -275,6 +308,7 @@ async fn mcp_tools_are_discovered_through_the_configured_proxy(#[future] upstrea .unwrap(); } +#[rstest] #[rstest] #[tokio::test] async fn mcp_uses_the_configured_ca_bundle() { @@ -299,11 +333,20 @@ async fn mcp_uses_the_configured_ca_bundle() { .await .err() .unwrap(); - assert!(matches!( - error, - litellm_gateway::Error::Http(litellm_http::Error::Read { path, .. }) - if path == missing - )); + let litellm_gateway::Error::Http(error) = error else { + panic!("expected an HTTP configuration error"); + }; + assert_eq!(error.tls_source(), Some(litellm_http::TlsSource::CaBundle)); + assert_eq!( + error.to_string(), + format!("could not read {}", missing.display()) + ); + assert_eq!( + litellm_error::find::(&error) + .unwrap() + .kind(), + std::io::ErrorKind::NotFound, + ); } #[rstest] @@ -336,9 +379,10 @@ async fn rejects_unsupported_or_invalid_settings_before_serving( &inference.http, ) .await; - let error = result.err().expect("startup must fail").to_string(); - assert!(error.contains(expected), "{error}"); - assert!(!error.contains("secret-value")); + let error = result.err().expect("startup must fail"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains(expected), "{report}"); + assert!(!report.contains("secret-value")); } #[rstest] @@ -397,9 +441,10 @@ fn applies_configured_http_environment_before_building_clients() { .unwrap(); let error = litellm_gateway::build_inference(&config).err().unwrap(); assert!(matches!( - error, - litellm_gateway::Error::Http(litellm_http::Error::Read { path, .. }) - if path == missing + &error, + litellm_gateway::Error::Http(http) + if http.tls_source() == Some(litellm_http::TlsSource::CaBundle) + && http.to_string().contains(&missing.display().to_string()) )); } @@ -424,5 +469,8 @@ async fn refuses_unimplemented_global_mcp_policy(#[case] settings: &str) { .await .err() .unwrap(); - assert!(matches!(error, litellm_gateway::Error::McpSetting(_))); + assert!(matches!( + error, + litellm_gateway::Error::McpSetting(_) | litellm_gateway::Error::McpGuardrails + )); } diff --git a/litellm-rust/crates/gateway/tests/server.rs b/litellm-rust/crates/gateway/tests/server.rs index 53a97e7098c..c876c1cbcbb 100644 --- a/litellm-rust/crates/gateway/tests/server.rs +++ b/litellm-rust/crates/gateway/tests/server.rs @@ -94,7 +94,10 @@ async fn authenticates_before_serving_mounted_inference_routes( let response = request.send().await.unwrap(); assert_eq!(response.status().as_u16(), status); if status == 400 { - let expected = Error::UnknownModel("unconfigured-model".into()); + let expected = Error::UnknownModel { + route: "/v1/messages", + model: "unconfigured-model".into(), + }; assert_eq!( response.json::().await.unwrap(), expected.body(Some("gateway-test")) diff --git a/litellm-rust/crates/host-python/Cargo.toml b/litellm-rust/crates/host-python/Cargo.toml index 8d220b19def..0b0107d0548 100644 --- a/litellm-rust/crates/host-python/Cargo.toml +++ b/litellm-rust/crates/host-python/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true [dependencies] litellm-coroutine.workspace = true +litellm-error.workspace = true litellm-host.workspace = true bytes.workspace = true diff --git a/litellm-rust/crates/host-python/src/binding.rs b/litellm-rust/crates/host-python/src/binding.rs index 91fbc8c0e52..e36b7d70c28 100644 --- a/litellm-rust/crates/host-python/src/binding.rs +++ b/litellm-rust/crates/host-python/src/binding.rs @@ -5,7 +5,7 @@ use pyo3::types::PyDict; /// Converts requests, responses, stream values and errors at the Python boundary. pub trait PythonBinding: PythonOwned { - type Protocol: Protocol; + type Protocol: Protocol; /// The public exception a native failure maps to, kept as a value until the driver /// raises it. diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index f8e35aa09ba..5dd9a72eb4a 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -112,7 +112,7 @@ enum Setup { /// The route's public exception for a native failure. When classification itself fails, /// that failure is raised with the native error's text as its `__context__`. fn classified(py: Python<'_>, binding: &H, error: ErrorOf) -> PyErr { - let native = error.to_string(); + let native = litellm_error::Report::new(&error).to_string(); let classifier_error = match binding.map_error(py, error) { Ok(failure) => return failure.into(), Err(classifier_error) => classifier_error, diff --git a/litellm-rust/crates/host-python/src/error.rs b/litellm-rust/crates/host-python/src/error.rs index b06a4a2fd41..8d47d97df25 100644 --- a/litellm-rust/crates/host-python/src/error.rs +++ b/litellm-rust/crates/host-python/src/error.rs @@ -1,3 +1,4 @@ +use litellm_error::Report; use pyo3::{ PyErr, exceptions::{PyRuntimeError, PyTypeError}, @@ -24,16 +25,19 @@ pub enum Error { SyncInsideTokio, #[error("the context body produced no value")] ContextBodyProducedNothing, + #[error("the context body ran more than once")] + ContextBodyRanTwice, + #[error("the blocking pool dropped a Python call")] + BlockingCallCancelled(#[source] tokio::task::JoinError), #[error("file read must return bytes or str, got {got}")] FileContentType { got: String }, } impl From for PyErr { fn from(error: Error) -> Self { + let message = Report::new(&error).to_string(); match error { - Error::SendBeforeStart | Error::FileContentType { .. } => { - PyTypeError::new_err(error.to_string()) - } + Error::SendBeforeStart | Error::FileContentType { .. } => PyTypeError::new_err(message), Error::MissingState | Error::ExecutionRunning | Error::ExecutionClosed @@ -41,7 +45,9 @@ impl From for PyErr { | Error::ExecutionClosedWhileRunning | Error::NativeExecutionClosed | Error::SyncInsideTokio - | Error::ContextBodyProducedNothing => PyRuntimeError::new_err(error.to_string()), + | Error::ContextBodyProducedNothing + | Error::ContextBodyRanTwice + | Error::BlockingCallCancelled(_) => PyRuntimeError::new_err(message), } } } @@ -51,7 +57,7 @@ impl From for PyErr { /// raised, which reaches the caller as it was raised. #[derive(Debug, thiserror::Error)] pub enum InvokeError { - #[error("{0}")] + #[error(transparent)] Native(E), #[error(transparent)] Python(#[from] PyErr), diff --git a/litellm-rust/crates/host-python/src/fork_gate.rs b/litellm-rust/crates/host-python/src/fork_gate.rs index c4842dd9223..82698931f8f 100644 --- a/litellm-rust/crates/host-python/src/fork_gate.rs +++ b/litellm-rust/crates/host-python/src/fork_gate.rs @@ -12,13 +12,23 @@ pub(crate) struct ForkGate { fork_only_pid: AtomicU32, } -#[derive(Debug, PartialEq, Eq)] +#[derive(Debug, PartialEq, Eq, thiserror::Error)] pub(crate) enum Refused { + #[error( + "this process is reserved for forking workers and cannot run native routes; \ + move the call into a worker, after the fork" + )] ReservedForForking, + #[error( + "this process was forked after the native runtime started, and runtime threads \ + do not survive fork(); start workers with spawn or forkserver, or fork before \ + the first native call" + )] ForkedAfterStart, } -#[derive(Debug, PartialEq, Eq)] +#[derive(Debug, PartialEq, Eq, thiserror::Error)] +#[error("the native runtime already started in this process")] pub struct RuntimeAlreadyStarted; impl ForkGate { diff --git a/litellm-rust/crates/host-python/src/gil.rs b/litellm-rust/crates/host-python/src/gil.rs index aacecb3af81..5f54fae81ed 100644 --- a/litellm-rust/crates/host-python/src/gil.rs +++ b/litellm-rust/crates/host-python/src/gil.rs @@ -1,8 +1,9 @@ use std::sync::{ - Arc, Mutex, + Arc, atomic::{AtomicU64, Ordering}, }; +use parking_lot::Mutex; use pyo3::prelude::*; use crate::Error; @@ -36,12 +37,8 @@ impl PythonContext { let value = Arc::clone(&value); pyo3::types::PyCFunction::new_closure(py, None, None, move |_, _| { Python::attach(|py| -> PyResult<()> { - let body = body - .lock() - .expect("context body slot poisoned") - .take() - .expect("the context body ran more than once"); - *value.lock().expect("context value slot poisoned") = Some(body(py)); + let body = body.lock().take().ok_or(Error::ContextBodyRanTwice)?; + *value.lock() = Some(body(py)); Ok(()) }) })? @@ -49,7 +46,6 @@ impl PythonContext { copy.call_method1("run", (callback,))?; value .lock() - .expect("context value slot poisoned") .take() .ok_or_else(|| Error::ContextBodyProducedNothing.into()) } @@ -84,7 +80,7 @@ where match tokio::task::spawn_blocking(move || Python::attach(|py| context.enter(py, f))).await { Ok(value) => value, Err(error) if error.is_panic() => std::panic::resume_unwind(error.into_panic()), - Err(error) => panic!("the blocking pool dropped a Python call: {error}"), + Err(error) => Err(Error::BlockingCallCancelled(error).into()), } } diff --git a/litellm-rust/crates/host-python/src/handle.rs b/litellm-rust/crates/host-python/src/handle.rs index 2361128bf43..cb4de2e4763 100644 --- a/litellm-rust/crates/host-python/src/handle.rs +++ b/litellm-rust/crates/host-python/src/handle.rs @@ -56,25 +56,27 @@ impl Execution { fn advance(slf: &Bound<'_, Self>, py: Python<'_>, input: Resume) -> PyResult> { let (mut body, started) = { let mut execution = slf.borrow_mut(); - let started = match (&execution.state, &input) { - (ExecutionState::Created(_), Resume::Send(value)) if value.is_none(py) => false, - (ExecutionState::Created(_), Resume::Send(_)) => { + match ( + std::mem::replace(&mut execution.state, ExecutionState::Running), + &input, + ) { + (ExecutionState::Created(body), Resume::Send(value)) if value.is_none(py) => { + (body, false) + } + (ExecutionState::Suspended(body), _) => (body, true), + (ExecutionState::Running, _) => return Err(Error::ExecutionRunning.into()), + (state @ ExecutionState::Created(_), Resume::Send(_)) => { + execution.state = state; return Err(Error::SendBeforeStart.into()); } - (ExecutionState::Suspended(_), _) => true, - (ExecutionState::Running, _) => { - return Err(Error::ExecutionRunning.into()); - } - (ExecutionState::Closed, _) => { - return Err(Error::ExecutionClosed.into()); - } - (ExecutionState::Created(_), Resume::Throw(_)) => { + (state @ ExecutionState::Created(_), Resume::Throw(_)) => { + execution.state = state; return Err(Error::ExecutionOutOfOrder.into()); } - }; - match std::mem::replace(&mut execution.state, ExecutionState::Running) { - ExecutionState::Created(body) | ExecutionState::Suspended(body) => (body, started), - _ => unreachable!(), + (ExecutionState::Closed, _) => { + execution.state = ExecutionState::Closed; + return Err(Error::ExecutionClosed.into()); + } } }; let result = match input { diff --git a/litellm-rust/crates/host-python/src/marshal.rs b/litellm-rust/crates/host-python/src/marshal.rs index ae9066379ae..465521cf9d2 100644 --- a/litellm-rust/crates/host-python/src/marshal.rs +++ b/litellm-rust/crates/host-python/src/marshal.rs @@ -14,7 +14,8 @@ pub fn from_py_argument(value: &Bound<'_, PyAny>) -> PyResult where T: DeserializeOwned, { - pythonize::depythonize(value).map_err(|error| PyValueError::new_err(error.to_string())) + pythonize::depythonize(value) + .map_err(|error| PyValueError::new_err(litellm_error::Report::new(&error).to_string())) } pub fn from_py(value: &Bound<'_, PyAny>) -> PyResult diff --git a/litellm-rust/crates/host-python/src/runtime.rs b/litellm-rust/crates/host-python/src/runtime.rs index c128242ef14..e8872c9340a 100644 --- a/litellm-rust/crates/host-python/src/runtime.rs +++ b/litellm-rust/crates/host-python/src/runtime.rs @@ -7,6 +7,7 @@ use std::time::Duration; use crate::fork_gate::{ForkGate, Refused, RuntimeAlreadyStarted}; use crate::{Error, Pythonized, panic_to_pyerr, release_gil}; use futures_util::FutureExt; +use litellm_error::Report; use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; use serde::Serialize; @@ -42,19 +43,23 @@ pub fn reserve_process_for_forking() -> Result<(), RuntimeAlreadyStarted> { /// Claims process-bound native state before runtime startup or tokenizer execution. pub fn enter_native() -> PyResult<()> { - FORK_GATE - .enter(std::process::id()) - .map_err(|refused| match refused { - Refused::ReservedForForking => ProcessReservedForForking::new_err( - "this process is reserved for forking workers and cannot run native routes; \ - move the call into a worker, after the fork", - ), - Refused::ForkedAfterStart => ForkedAfterNativeRuntimeStarted::new_err( - "this process was forked after the native runtime started, and runtime threads \ - do not survive fork(); start workers with spawn or forkserver, or fork before \ - the first native call", - ), - }) + Ok(FORK_GATE.enter(std::process::id())?) +} + +impl From for PyErr { + fn from(refused: Refused) -> Self { + let message = Report::new(&refused).to_string(); + match refused { + Refused::ReservedForForking => ProcessReservedForForking::new_err(message), + Refused::ForkedAfterStart => ForkedAfterNativeRuntimeStarted::new_err(message), + } + } +} + +impl From for PyErr { + fn from(error: RuntimeAlreadyStarted) -> Self { + PyRuntimeError::new_err(Report::new(&error).to_string()) + } } #[expect(clippy::disallowed_methods, reason = "this is the gated door")] diff --git a/litellm-rust/crates/host-python/tests/support/mod.rs b/litellm-rust/crates/host-python/tests/support/mod.rs index ef09135698e..003fce74309 100644 --- a/litellm-rust/crates/host-python/tests/support/mod.rs +++ b/litellm-rust/crates/host-python/tests/support/mod.rs @@ -92,6 +92,8 @@ impl std::fmt::Display for Error { } } +impl std::error::Error for Error {} + impl From for Error { fn from(fault: MachineFault) -> Self { Self(format!("{fault:?}")) diff --git a/litellm-rust/crates/host/Cargo.toml b/litellm-rust/crates/host/Cargo.toml index 51015e99906..d3354fa8622 100644 --- a/litellm-rust/crates/host/Cargo.toml +++ b/litellm-rust/crates/host/Cargo.toml @@ -10,6 +10,7 @@ futures-util.workspace = true litellm-auth.workspace = true litellm-coroutine.workspace = true serde_json.workspace = true +thiserror.workspace = true tokio = { workspace = true, features = ["sync"] } [dev-dependencies] diff --git a/litellm-rust/crates/host/src/machine/coroutine.rs b/litellm-rust/crates/host/src/machine/coroutine.rs index 825536e1d9b..db69f1dceb2 100644 --- a/litellm-rust/crates/host/src/machine/coroutine.rs +++ b/litellm-rust/crates/host/src/machine/coroutine.rs @@ -10,12 +10,14 @@ use super::{ use crate::protocol::{HostRequest, Protocol}; /// The machine's own failures, distinct from anything the provider call reports. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] pub enum MachineFault { /// The host dropped an op's reply unanswered, or went away while the call waited. + #[error("the host abandoned the call")] Abandoned, /// The host resumed the call out of turn. - Protocol(ResumeError), + #[error("the host resumed the call out of turn")] + Protocol(#[source] ResumeError), } pub type ExecuteFuture::Response> = diff --git a/litellm-rust/crates/http/Cargo.toml b/litellm-rust/crates/http/Cargo.toml index 0552f9775fd..6b3e3cc7852 100644 --- a/litellm-rust/crates/http/Cargo.toml +++ b/litellm-rust/crates/http/Cargo.toml @@ -11,6 +11,7 @@ test-support = [] [dependencies] http.workspace = true litellm-core-utils.workspace = true +litellm-error.workspace = true hyper-util.workspace = true reqwest.workspace = true rustls.workspace = true @@ -22,6 +23,7 @@ strum.workspace = true thiserror.workspace = true tokio.workspace = true tracing.workspace = true +url.workspace = true veil.workspace = true webpki-roots.workspace = true diff --git a/litellm-rust/crates/http/src/config.rs b/litellm-rust/crates/http/src/config.rs index b37ea6559b8..5e5f2d66fa2 100644 --- a/litellm-rust/crates/http/src/config.rs +++ b/litellm-rust/crates/http/src/config.rs @@ -136,7 +136,7 @@ mod tests { use rstest::rstest; use super::*; - use crate::TlsSource; + use crate::{TlsSource, error::Repr}; fn settings(ssl_verify: Option, ssl_cert_file: Option<&str>) -> HttpSettings { HttpSettings { @@ -306,11 +306,10 @@ mod tests { }; assert!(matches!( reqwest::ClientBuilder::try_from(&config), - Err(Error::Read { - path: reported, - tls_source: TlsSource::CaBundle, - .. - }) if reported == path + Err(error) if matches!( + error.repr(), + Repr::Read { path: reported, tls_source: TlsSource::CaBundle, .. } if *reported == path + ) )); } @@ -327,11 +326,10 @@ mod tests { std::fs::remove_file(&path).unwrap(); assert!(matches!( result, - Err(Error::InvalidPem { - path: reported, - tls_source: TlsSource::CaBundle, - .. - }) if reported == path + Err(error) if matches!( + error.repr(), + Repr::InvalidPem { path: reported, tls_source: TlsSource::CaBundle, .. } if *reported == path + ) )); } } diff --git a/litellm-rust/crates/http/src/error.rs b/litellm-rust/crates/http/src/error.rs index ae14756baa3..29c8b988c0b 100644 --- a/litellm-rust/crates/http/src/error.rs +++ b/litellm-rust/crates/http/src/error.rs @@ -1,4 +1,10 @@ -use std::path::PathBuf; +use std::{borrow::Cow, fmt, io, path::PathBuf, sync::Arc}; + +use tokio::time::error::Elapsed; +use tokio_tungstenite::tungstenite; + +const UPSTREAM_BODY_MAX_CHARS: usize = 256; +const TRUNCATION_MARKER: &str = "... (truncated)"; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum TlsSource { @@ -6,34 +12,666 @@ pub enum TlsSource { ClientIdentity, } -#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] -pub enum Error { - #[error("could not read {}: {message}", path.display())] +#[derive(Clone, Debug, thiserror::Error)] +#[error(transparent)] +pub struct Error(Arc); + +#[derive(Debug, thiserror::Error)] +pub(crate) enum Repr { + #[error("could not read {}", path.display())] Read { path: PathBuf, - message: String, tls_source: TlsSource, + source: io::Error, }, - #[error("{} is not a PEM file: {message}", path.display())] + #[error("could not load {}", path.display())] InvalidPem { path: PathBuf, - message: String, tls_source: TlsSource, + source: PemFault, }, - #[error("could not build the HTTP client: {0}")] - Client(String), - #[error("request body could not be serialized: {0}")] - RequestBody(String), - #[error("invalid request: {0}")] - Request(String), - #[error("request forwards a header the signer computes: {0}")] + #[error("could not build the HTTP client")] + Client(#[source] reqwest::Error), + #[error("could not build the TLS configuration")] + TlsConfig(#[source] rustls::Error), + + #[error("request body could not be serialized")] + RequestBody(#[source] serde_json::Error), + #[error("invalid request")] + Request(#[source] http::Error), + #[error("invalid request")] + RequestUrl(#[source] reqwest::Error), + #[error("invalid request URL")] + ParseUrl(#[source] url::ParseError), + #[error("request forwards the header {0:?}, which the signer computes")] ComputedHeader(String), - #[error("request signing failed: {0}")] - Signature(String), + #[error("request signing failed: a streaming body cannot be signed")] + UnsignableBody, + #[error("request forwards a header that cannot be signed")] + UnsignableHeader(#[source] litellm_error::Source), + #[error("request signing failed")] + Signature(#[source] litellm_error::Source), + + #[error("could not reach the provider")] + Connect(#[source] reqwest::Error), + #[error("upstream network error")] + Network(#[source] reqwest::Error), + #[error("upstream request timed out")] + Timeout(#[source] reqwest::Error), + #[error("upstream request failed with status {status}: {body}")] + Status { status: u16, body: UpstreamBody }, + + #[error("invalid WebSocket URL")] + WebSocketUrl(#[source] url::ParseError), + #[error("unsupported WebSocket URL scheme {0:?}")] + WebSocketScheme(String), + #[error("could not rewrite the WebSocket URL scheme to {0}")] + WebSocketSchemeRewrite(&'static str), + #[error("WebSocket handshake returned {0}")] + Handshake(HandshakeFault), + #[error("WebSocket connection timed out")] + HandshakeTimeout(#[source] Elapsed), + #[error("upstream WebSocket error")] + WebSocket(#[source] tungstenite::Error), + #[error("the WebSocket connection is closed")] + WebSocketClosed, + + #[error("media URL rejected by network policy")] + BlockedUrl, + #[error("could not resolve the media host")] + Resolve(#[source] io::Error), + #[error("media download is disabled")] + DownloadDisabled, + #[error("media download exceeds the maximum size")] + DownloadTooLarge, + #[error("too many redirects while fetching media")] + TooManyRedirects, + #[error("media redirect is missing a Location header")] + MissingRedirectLocation, + #[error("invalid media redirect")] + InvalidRedirect(#[source] RedirectFault), + #[error("media download failed with status {0}")] + DownloadStatus(u16), + #[error("media download failed")] + Download(#[source] reqwest::Error), + #[error("media download timed out")] + DownloadTimeout(#[source] Elapsed), +} + +#[derive(Debug, thiserror::Error)] +pub(crate) enum HandshakeFault { + #[error("a wrong Sec-WebSocket-Accept")] + Accept, + #[error("an invalid Upgrade header")] + Upgrade, + #[error("an invalid Connection header")] + Connection, + #[error("an invalid subprotocol")] + Subprotocol, +} + +#[derive(Debug, thiserror::Error)] +pub(crate) enum PemFault { + #[error("no certificates found")] + NoCertificates, + #[error("could not parse PEM")] + Parse(#[from] rustls::pki_types::pem::Error), + #[error("rustls rejected the certificate or key")] + Rejected(#[from] rustls::Error), +} + +#[derive(Debug, thiserror::Error)] +pub(crate) enum RedirectFault { + #[error("Location header is not valid text")] + Location(#[from] http::header::ToStrError), + #[error("Location is not a valid URL")] + Url(#[from] url::ParseError), +} + +impl Repr { + pub(crate) fn download(error: reqwest::Error) -> Self { + Self::Download(error.without_url()) + } +} + +impl Error { + pub fn upstream(status: u16, body: String) -> Self { + Repr::Status { + status, + body: UpstreamBody(body), + } + .into() + } + + pub fn computed_header(name: impl Into) -> Self { + Repr::ComputedHeader(name.into()).into() + } + + pub fn unsignable_body() -> Self { + Repr::UnsignableBody.into() + } + + pub fn unsignable_header(cause: impl std::error::Error + Send + Sync + 'static) -> Self { + Repr::UnsignableHeader(litellm_error::Source::new(cause)).into() + } + + pub fn websocket_closed() -> Self { + Repr::WebSocketClosed.into() + } + + pub fn signature(cause: impl std::error::Error + Send + Sync + 'static) -> Self { + Repr::Signature(litellm_error::Source::new(cause)).into() + } + + /// The request may have reached the provider, so retrying elsewhere could + /// repeat a billed call. + pub fn dispatched(&self) -> bool { + match &*self.0 { + Repr::Network(_) + | Repr::Timeout(_) + | Repr::Status { .. } + | Repr::Handshake(_) + | Repr::HandshakeTimeout(_) + | Repr::WebSocket(_) => true, + Repr::WebSocketClosed + | Repr::Read { .. } + | Repr::InvalidPem { .. } + | Repr::Client(_) + | Repr::TlsConfig(_) + | Repr::RequestBody(_) + | Repr::Request(_) + | Repr::RequestUrl(_) + | Repr::ParseUrl(_) + | Repr::ComputedHeader(_) + | Repr::UnsignableBody + | Repr::UnsignableHeader(_) + | Repr::Signature(_) + | Repr::Connect(_) + | Repr::WebSocketUrl(_) + | Repr::WebSocketScheme(_) + | Repr::WebSocketSchemeRewrite(_) + | Repr::BlockedUrl + | Repr::Resolve(_) + | Repr::DownloadDisabled + | Repr::DownloadTooLarge + | Repr::TooManyRedirects + | Repr::MissingRedirectLocation + | Repr::InvalidRedirect(_) + | Repr::DownloadStatus(_) + | Repr::Download(_) + | Repr::DownloadTimeout(_) => false, + } + } + + pub fn is_request(&self) -> bool { + match &*self.0 { + Repr::Request(_) + | Repr::RequestUrl(_) + | Repr::ParseUrl(_) + | Repr::ComputedHeader(_) + | Repr::UnsignableHeader(_) + | Repr::WebSocketUrl(_) + | Repr::WebSocketScheme(_) + | Repr::WebSocketSchemeRewrite(_) + | Repr::BlockedUrl + | Repr::DownloadDisabled + | Repr::DownloadTooLarge + | Repr::TooManyRedirects => true, + Repr::Read { .. } + | Repr::InvalidPem { .. } + | Repr::Client(_) + | Repr::TlsConfig(_) + | Repr::RequestBody(_) + | Repr::UnsignableBody + | Repr::Signature(_) + | Repr::Connect(_) + | Repr::Network(_) + | Repr::Timeout(_) + | Repr::Status { .. } + | Repr::Handshake(_) + | Repr::HandshakeTimeout(_) + | Repr::WebSocket(_) + | Repr::WebSocketClosed + | Repr::Resolve(_) + | Repr::MissingRedirectLocation + | Repr::InvalidRedirect(_) + | Repr::DownloadStatus(_) + | Repr::Download(_) + | Repr::DownloadTimeout(_) => false, + } + } + + pub fn is_response(&self) -> bool { + match &*self.0 { + Repr::MissingRedirectLocation | Repr::InvalidRedirect(_) => true, + Repr::Read { .. } + | Repr::InvalidPem { .. } + | Repr::Client(_) + | Repr::TlsConfig(_) + | Repr::RequestBody(_) + | Repr::Request(_) + | Repr::RequestUrl(_) + | Repr::ParseUrl(_) + | Repr::ComputedHeader(_) + | Repr::UnsignableBody + | Repr::UnsignableHeader(_) + | Repr::Signature(_) + | Repr::Connect(_) + | Repr::Network(_) + | Repr::Timeout(_) + | Repr::Status { .. } + | Repr::WebSocketUrl(_) + | Repr::WebSocketScheme(_) + | Repr::WebSocketSchemeRewrite(_) + | Repr::Handshake(_) + | Repr::HandshakeTimeout(_) + | Repr::WebSocket(_) + | Repr::WebSocketClosed + | Repr::BlockedUrl + | Repr::Resolve(_) + | Repr::DownloadDisabled + | Repr::DownloadTooLarge + | Repr::TooManyRedirects + | Repr::DownloadStatus(_) + | Repr::Download(_) + | Repr::DownloadTimeout(_) => false, + } + } + + pub fn is_computed_header(&self) -> bool { + match &*self.0 { + Repr::ComputedHeader(_) => true, + Repr::Read { .. } + | Repr::InvalidPem { .. } + | Repr::Client(_) + | Repr::TlsConfig(_) + | Repr::RequestBody(_) + | Repr::Request(_) + | Repr::RequestUrl(_) + | Repr::ParseUrl(_) + | Repr::UnsignableBody + | Repr::UnsignableHeader(_) + | Repr::Signature(_) + | Repr::Connect(_) + | Repr::Network(_) + | Repr::Timeout(_) + | Repr::Status { .. } + | Repr::WebSocketUrl(_) + | Repr::WebSocketScheme(_) + | Repr::WebSocketSchemeRewrite(_) + | Repr::Handshake(_) + | Repr::HandshakeTimeout(_) + | Repr::WebSocket(_) + | Repr::WebSocketClosed + | Repr::BlockedUrl + | Repr::Resolve(_) + | Repr::DownloadDisabled + | Repr::DownloadTooLarge + | Repr::TooManyRedirects + | Repr::MissingRedirectLocation + | Repr::InvalidRedirect(_) + | Repr::DownloadStatus(_) + | Repr::Download(_) + | Repr::DownloadTimeout(_) => false, + } + } + + pub fn is_connect(&self) -> bool { + match &*self.0 { + Repr::Connect(_) => true, + Repr::ComputedHeader(_) + | Repr::Read { .. } + | Repr::InvalidPem { .. } + | Repr::Client(_) + | Repr::TlsConfig(_) + | Repr::RequestBody(_) + | Repr::Request(_) + | Repr::RequestUrl(_) + | Repr::ParseUrl(_) + | Repr::UnsignableBody + | Repr::UnsignableHeader(_) + | Repr::Signature(_) + | Repr::Network(_) + | Repr::Timeout(_) + | Repr::Status { .. } + | Repr::WebSocketUrl(_) + | Repr::WebSocketScheme(_) + | Repr::WebSocketSchemeRewrite(_) + | Repr::Handshake(_) + | Repr::HandshakeTimeout(_) + | Repr::WebSocket(_) + | Repr::WebSocketClosed + | Repr::BlockedUrl + | Repr::Resolve(_) + | Repr::DownloadDisabled + | Repr::DownloadTooLarge + | Repr::TooManyRedirects + | Repr::MissingRedirectLocation + | Repr::InvalidRedirect(_) + | Repr::DownloadStatus(_) + | Repr::Download(_) + | Repr::DownloadTimeout(_) => false, + } + } + + pub fn is_timeout(&self) -> bool { + match &*self.0 { + Repr::Timeout(_) | Repr::HandshakeTimeout(_) | Repr::DownloadTimeout(_) => true, + Repr::Read { .. } + | Repr::InvalidPem { .. } + | Repr::Client(_) + | Repr::TlsConfig(_) + | Repr::RequestBody(_) + | Repr::Request(_) + | Repr::RequestUrl(_) + | Repr::ParseUrl(_) + | Repr::ComputedHeader(_) + | Repr::UnsignableBody + | Repr::UnsignableHeader(_) + | Repr::Signature(_) + | Repr::Connect(_) + | Repr::Network(_) + | Repr::Status { .. } + | Repr::WebSocketUrl(_) + | Repr::WebSocketScheme(_) + | Repr::WebSocketSchemeRewrite(_) + | Repr::Handshake(_) + | Repr::WebSocket(_) + | Repr::WebSocketClosed + | Repr::BlockedUrl + | Repr::Resolve(_) + | Repr::DownloadDisabled + | Repr::DownloadTooLarge + | Repr::TooManyRedirects + | Repr::MissingRedirectLocation + | Repr::InvalidRedirect(_) + | Repr::DownloadStatus(_) + | Repr::Download(_) => false, + } + } + + pub fn upstream_response(&self) -> Option<(u16, &str)> { + match &*self.0 { + Repr::Status { status, body } => Some((*status, &body.0)), + Repr::Read { .. } + | Repr::InvalidPem { .. } + | Repr::Client(_) + | Repr::TlsConfig(_) + | Repr::RequestBody(_) + | Repr::Request(_) + | Repr::RequestUrl(_) + | Repr::ParseUrl(_) + | Repr::ComputedHeader(_) + | Repr::UnsignableBody + | Repr::UnsignableHeader(_) + | Repr::Signature(_) + | Repr::Connect(_) + | Repr::Network(_) + | Repr::Timeout(_) + | Repr::WebSocketUrl(_) + | Repr::WebSocketScheme(_) + | Repr::WebSocketSchemeRewrite(_) + | Repr::Handshake(_) + | Repr::HandshakeTimeout(_) + | Repr::WebSocket(_) + | Repr::WebSocketClosed + | Repr::BlockedUrl + | Repr::Resolve(_) + | Repr::DownloadDisabled + | Repr::DownloadTooLarge + | Repr::TooManyRedirects + | Repr::MissingRedirectLocation + | Repr::InvalidRedirect(_) + | Repr::DownloadStatus(_) + | Repr::Download(_) + | Repr::DownloadTimeout(_) => None, + } + } + + pub fn status(&self) -> Option { + match &*self.0 { + Repr::Status { status, .. } | Repr::DownloadStatus(status) => Some(*status), + Repr::Timeout(_) | Repr::DownloadTimeout(_) => Some(408), + Repr::Read { .. } + | Repr::InvalidPem { .. } + | Repr::Client(_) + | Repr::TlsConfig(_) + | Repr::RequestBody(_) + | Repr::Request(_) + | Repr::RequestUrl(_) + | Repr::ParseUrl(_) + | Repr::ComputedHeader(_) + | Repr::UnsignableBody + | Repr::UnsignableHeader(_) + | Repr::Signature(_) + | Repr::Connect(_) + | Repr::Network(_) + | Repr::WebSocketUrl(_) + | Repr::WebSocketScheme(_) + | Repr::WebSocketSchemeRewrite(_) + | Repr::Handshake(_) + | Repr::HandshakeTimeout(_) + | Repr::WebSocket(_) + | Repr::WebSocketClosed + | Repr::BlockedUrl + | Repr::Resolve(_) + | Repr::DownloadDisabled + | Repr::DownloadTooLarge + | Repr::TooManyRedirects + | Repr::MissingRedirectLocation + | Repr::InvalidRedirect(_) + | Repr::Download(_) => None, + } + } + + pub fn tls_source(&self) -> Option { + match &*self.0 { + Repr::Read { tls_source, .. } | Repr::InvalidPem { tls_source, .. } => { + Some(*tls_source) + } + Repr::Client(_) + | Repr::TlsConfig(_) + | Repr::RequestBody(_) + | Repr::Request(_) + | Repr::RequestUrl(_) + | Repr::ParseUrl(_) + | Repr::ComputedHeader(_) + | Repr::UnsignableBody + | Repr::UnsignableHeader(_) + | Repr::Signature(_) + | Repr::Connect(_) + | Repr::Network(_) + | Repr::Timeout(_) + | Repr::Status { .. } + | Repr::WebSocketUrl(_) + | Repr::WebSocketScheme(_) + | Repr::WebSocketSchemeRewrite(_) + | Repr::Handshake(_) + | Repr::HandshakeTimeout(_) + | Repr::WebSocket(_) + | Repr::WebSocketClosed + | Repr::BlockedUrl + | Repr::Resolve(_) + | Repr::DownloadDisabled + | Repr::DownloadTooLarge + | Repr::TooManyRedirects + | Repr::MissingRedirectLocation + | Repr::InvalidRedirect(_) + | Repr::DownloadStatus(_) + | Repr::Download(_) + | Repr::DownloadTimeout(_) => None, + } + } + + #[cfg(test)] + pub(crate) fn repr(&self) -> &Repr { + &self.0 + } +} + +pub(crate) struct UpstreamBody(String); + +impl fmt::Display for UpstreamBody { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Display::fmt(&bounded(&self.0), f) + } +} + +impl fmt::Debug for UpstreamBody { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Debug::fmt(&bounded(&self.0), f) + } +} + +fn bounded(body: &str) -> Cow<'_, str> { + match body.char_indices().nth(UPSTREAM_BODY_MAX_CHARS) { + None => Cow::Borrowed(body), + Some((cut, _)) => Cow::Owned(format!("{}{TRUNCATION_MARKER}", &body[..cut])), + } +} + +impl PartialEq for Error { + fn eq(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.0, &other.0) + } +} + +impl Eq for Error {} + +impl From for Error { + fn from(repr: Repr) -> Self { + Self(Arc::new(repr)) + } } impl From for Error { fn from(error: reqwest::Error) -> Self { - Self::Client(error.without_url().to_string()) + let error = error.without_url(); + if error.is_timeout() { + Repr::Timeout(error).into() + } else if error.is_connect() || error.is_builder() { + Repr::Connect(error).into() + } else { + Repr::Network(error).into() + } + } +} + +impl From for Error { + fn from(error: tungstenite::Error) -> Self { + Repr::WebSocket(error).into() + } +} + +#[cfg(test)] +mod tests { + use litellm_error::Report; + use rstest::rstest; + + use super::{TRUNCATION_MARKER, UPSTREAM_BODY_MAX_CHARS}; + + #[tokio::test] + async fn transport_errors_remove_urls_and_keep_dispatch_context() { + let error = reqwest::Client::builder() + .no_proxy() + .build() + .expect("client") + .get("http://localhost:invalid/private?api_key=secret") + .send() + .await + .expect_err("invalid port"); + let error = crate::Error::from(error); + let report = Report::new(&error).to_string(); + assert!(!error.dispatched()); + assert!(!report.contains("secret"), "{report}"); + assert!(!report.contains("private"), "{report}"); + } + + fn root_cause<'a>( + error: &'a (dyn std::error::Error + 'static), + ) -> &'a (dyn std::error::Error + 'static) { + error.source().map_or(error, root_cause) + } + + #[tokio::test] + async fn connect_failure_reports_its_root_cause_once() { + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind"); + let address = listener.local_addr().expect("address"); + drop(listener); + let error = reqwest::Client::builder() + .no_proxy() + .build() + .expect("client") + .get(format!("http://{address}/private?api_key=secret")) + .send() + .await + .expect_err("nothing listens on the port"); + let error = crate::Error::from(error); + let root_cause = root_cause(&error).to_string(); + let report = Report::new(&error).to_string(); + assert_eq!(report.matches(&root_cause).count(), 1, "{report}"); + assert!(!error.to_string().contains(&root_cause)); + assert!(error.is_connect()); + assert!(!error.is_timeout()); + assert!(!report.contains("secret"), "{report}"); + assert!(!report.contains("private"), "{report}"); + } + + #[tokio::test] + async fn request_timeout_is_not_safe_to_retry_as_a_connect_failure() { + use std::time::Duration; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind"); + let address = listener.local_addr().expect("address"); + let request = reqwest::Client::builder() + .no_proxy() + .build() + .expect("client") + .get(format!("http://{address}")) + .timeout(Duration::from_millis(200)) + .send(); + let (response, accepted) = tokio::join!( + request, + tokio::time::timeout(Duration::from_secs(2), listener.accept()) + ); + let _connection = accepted + .expect("accept deadline") + .expect("accepted connection"); + let error = response.expect_err("server does not respond"); + assert!(error.is_timeout()); + let error = crate::Error::from(error); + assert!(error.dispatched()); + assert!(error.is_timeout()); + assert!(!error.is_connect()); + assert_eq!(error.status(), Some(408)); + } + + #[rstest] + #[case::empty(String::new(), String::new())] + #[case::short("slow down".to_owned(), "slow down".to_owned())] + #[case::ascii_over_the_limit( + format!("{}private-tail", "a".repeat(UPSTREAM_BODY_MAX_CHARS)), + format!("{}{TRUNCATION_MARKER}", "a".repeat(UPSTREAM_BODY_MAX_CHARS)) + )] + #[case::at_the_limit("\u{00e9}".repeat(UPSTREAM_BODY_MAX_CHARS), "\u{00e9}".repeat(UPSTREAM_BODY_MAX_CHARS))] + #[case::over_the_limit( + format!("{}private-tail", "\u{00e9}".repeat(UPSTREAM_BODY_MAX_CHARS)), + format!("{}{TRUNCATION_MARKER}", "\u{00e9}".repeat(UPSTREAM_BODY_MAX_CHARS)) + )] + fn upstream_preserves_the_body_and_bounds_only_diagnostics( + #[case] body: String, + #[case] expected: String, + ) { + let error = crate::Error::upstream(429, body.clone()); + let expected_message = format!("upstream request failed with status 429: {expected}"); + assert_eq!(error.to_string(), expected_message); + assert_eq!(Report::new(&error).to_string(), expected_message); + let debug = format!("{error:?}"); + assert!(debug.contains(&format!("{expected:?}")), "{debug}"); + assert!(!debug.contains("private-tail"), "{debug}"); + assert!(!format!("{error:#?}").contains("private-tail")); + assert_eq!(error.upstream_response(), Some((429, body.as_str()))); + assert_eq!(error.status(), Some(429)); } } diff --git a/litellm-rust/crates/http/src/lib.rs b/litellm-rust/crates/http/src/lib.rs index ef72420af9b..02129c62083 100644 --- a/litellm-rust/crates/http/src/lib.rs +++ b/litellm-rust/crates/http/src/lib.rs @@ -14,7 +14,6 @@ mod proxy; pub mod request; mod settings; mod tls; -pub mod transport; pub mod websocket; pub use client::Client; diff --git a/litellm-rust/crates/http/src/media.rs b/litellm-rust/crates/http/src/media.rs index 960a1d4044a..4acaca5e502 100644 --- a/litellm-rust/crates/http/src/media.rs +++ b/litellm-rust/crates/http/src/media.rs @@ -12,29 +12,7 @@ use reqwest::{ dns::{Addrs, Name, Resolve, Resolving}, }; -use crate::{Client, ClientVariant, HttpClientConfig, HttpClientPool}; - -#[derive(Debug, thiserror::Error)] -pub enum Error { - #[error("media URL rejected by network policy")] - BlockedUrl, - #[error("media download is disabled")] - DownloadDisabled, - #[error("media download exceeds the maximum size")] - DownloadTooLarge, - #[error("too many redirects while fetching media")] - TooManyRedirects, - #[error("media redirect is missing a Location header")] - MissingRedirectLocation, - #[error("invalid media redirect")] - InvalidRedirect, - #[error("media download failed with status {0}")] - Http(u16), - #[error("media download timed out")] - Timeout, - #[error("{0}")] - Transport(#[from] crate::transport::Error), -} +use crate::{Client, ClientVariant, Error, HttpClientConfig, HttpClientPool, error::Repr}; #[derive(Clone, Debug, PartialEq, Eq)] pub struct UrlPolicy { @@ -125,7 +103,7 @@ impl MediaFetcher { pool: &HttpClientPool, config: &HttpClientConfig, url_policy: UrlPolicy, - ) -> Result { + ) -> Result { let uses_proxy: ProxyMatch = Arc::new(config.proxies.matcher()); Self::with_resolution( pool, @@ -142,7 +120,7 @@ impl MediaFetcher { url_policy: UrlPolicy, address_resolver: Arc, uses_proxy: ProxyMatch, - ) -> Result { + ) -> Result { Ok(Self { pinned: pool.client(config, ClientVariant::Media)?, unpinned: pool.client(config, ClientVariant::UnpinnedMedia)?, @@ -167,11 +145,11 @@ impl MediaFetcher { pub async fn fetch(&self, url: Url, policy: DownloadPolicy) -> Result { if policy.max_bytes == 0 { - return Err(Error::DownloadDisabled); + return Err(Repr::DownloadDisabled.into()); } tokio::time::timeout(policy.timeout, self.fetch_before_deadline(url, policy)) .await - .map_err(|_| Error::Timeout)? + .map_err(Repr::DownloadTimeout)? } async fn fetch_before_deadline( @@ -187,22 +165,25 @@ impl MediaFetcher { .get(url.clone()) .send() .await - .map_err(crate::transport::Error::from)?; + .map_err(Repr::download)?; if response.status().is_redirection() { if redirects_followed == policy.max_redirects { - return Err(Error::TooManyRedirects); + return Err(Repr::TooManyRedirects.into()); } let location = response .headers() .get(reqwest::header::LOCATION) - .and_then(|value| value.to_str().ok()) - .ok_or(Error::MissingRedirectLocation)?; - url = url.join(location).map_err(|_| Error::InvalidRedirect)?; + .ok_or(Repr::MissingRedirectLocation)? + .to_str() + .map_err(|error| Repr::InvalidRedirect(error.into()))?; + url = url + .join(location) + .map_err(|error| Repr::InvalidRedirect(error.into()))?; redirects_followed += 1; continue; } if !response.status().is_success() { - return Err(Error::Http(response.status().as_u16())); + return Err(Repr::DownloadStatus(response.status().as_u16()).into()); } enforce_download_size(response.content_length().unwrap_or(0), policy.max_bytes)?; let content_type = response @@ -215,11 +196,7 @@ impl MediaFetcher { .unwrap_or("application/octet-stream") .to_string(); let mut bytes = Vec::new(); - while let Some(chunk) = response - .chunk() - .await - .map_err(crate::transport::Error::from)? - { + while let Some(chunk) = response.chunk().await.map_err(Repr::download)? { enforce_download_size(bytes.len() as u64 + chunk.len() as u64, policy.max_bytes)?; bytes.extend_from_slice(&chunk); } @@ -238,13 +215,13 @@ impl MediaFetcher { || !url.username().is_empty() || url.password().is_some() { - return Err(Error::BlockedUrl); + return Err(Repr::BlockedUrl.into()); } - let host = url.host_str().ok_or(Error::BlockedUrl)?; + let host = url.host_str().ok_or(Repr::BlockedUrl)?; if self.allow_private_network { return Ok(&self.pinned); } - let port = url.port_or_known_default().ok_or(Error::BlockedUrl)?; + let port = url.port_or_known_default().ok_or(Repr::BlockedUrl)?; if self.url_policy.allows(host, port) { return Ok(&self.unpinned); } @@ -262,27 +239,29 @@ impl MediaFetcher { .trim_end_matches(']') .parse::() { - return (!is_blocked_ip(ip)).then_some(()).ok_or(Error::BlockedUrl); + return (!is_blocked_ip(ip)) + .then_some(()) + .ok_or(Repr::BlockedUrl.into()); } let addresses = self .address_resolver .resolve(host, port) .await - .map_err(|error| crate::transport::Error::Network(error.to_string()))?; + .map_err(Repr::Resolve)?; validate_addresses(&addresses) } } fn enforce_download_size(length: u64, max_bytes: u64) -> Result<(), Error> { if length > max_bytes { - return Err(Error::DownloadTooLarge); + return Err(Repr::DownloadTooLarge.into()); } Ok(()) } fn validate_addresses(addresses: &[SocketAddr]) -> Result<(), Error> { if addresses.is_empty() || addresses.iter().any(|address| is_blocked_ip(address.ip())) { - return Err(Error::BlockedUrl); + return Err(Repr::BlockedUrl.into()); } Ok(()) } @@ -356,10 +335,8 @@ impl Resolve for PublicDnsResolver { .await .map_err(|error| Box::new(error) as Box)? .collect::>(); - validate_addresses(&addresses).map_err(|_| { - Box::new(io::Error::other("destination rejected by network policy")) - as Box - })?; + validate_addresses(&addresses) + .map_err(|error| Box::new(error) as Box)?; Ok(Box::new(addresses.into_iter()) as Addrs) }) } @@ -544,7 +521,7 @@ mod tests { .await .expect_err("oversize body is rejected"); server.await.expect("server completes"); - assert!(matches!(error, Error::DownloadTooLarge)); + assert!(matches!(error.repr(), Repr::DownloadTooLarge)); } #[tokio::test] @@ -559,7 +536,7 @@ mod tests { .await .expect_err("stream crossing limit is rejected"); server.await.expect("server completes"); - assert!(matches!(error, Error::DownloadTooLarge)); + assert!(matches!(error.repr(), Repr::DownloadTooLarge)); } #[tokio::test] @@ -595,7 +572,7 @@ mod tests { .expect_err("private redirect is rejected"); let requests = server.await.expect("server completes"); assert_eq!(requests.len(), 1); - assert!(matches!(error, Error::BlockedUrl)); + assert!(matches!(error.repr(), Repr::BlockedUrl)); } #[tokio::test] @@ -622,7 +599,27 @@ mod tests { .await .expect_err("fetch times out"); server.await.expect("server completes"); - assert!(matches!(error, Error::Timeout)); + assert!(matches!(error.repr(), Repr::DownloadTimeout(_))); + assert!(!error.dispatched()); + } + + #[rstest::rstest] + #[case::closed_before_answering(b"")] + #[case::reset_mid_body(b"HTTP/1.1 200 OK\r\nContent-Length: 10\r\n\r\nab")] + #[tokio::test] + async fn a_failed_download_never_reached_the_provider(#[case] response: &'static [u8]) { + let (url, server) = serve(response).await; + let url = url.join("?sig=secret").expect("valid test URL"); + let error = MediaFetcher::for_test(Client::for_test()) + .fetch(url, policy(10, 0)) + .await + .expect_err("download fails"); + server.await.expect("server completes"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(!error.dispatched()); + assert!(!error.is_request()); + assert!(!report.contains("secret"), "{report}"); + assert!(!report.contains("document"), "{report}"); } #[tokio::test] @@ -652,8 +649,13 @@ mod tests { let url = Url::parse("https://user:password@8.8.8.8/document").expect("credentialed URL parses"); assert!(matches!( - fetcher.fetch(url, policy(1, 0)).await, - Err(Error::BlockedUrl) + fetcher + .fetch(url, policy(1, 0)) + .await + .err() + .as_ref() + .map(Error::repr), + Some(Repr::BlockedUrl) )); } @@ -683,7 +685,10 @@ mod tests { let result = fetcher(UNROUTABLE, HashSet::from(["localhost"]), other_port, false) .fetch(url, policy(2, 0)) .await; - assert!(matches!(result, Err(Error::BlockedUrl))); + assert!(matches!( + result.as_ref().err().map(Error::repr), + Some(Repr::BlockedUrl) + )); } #[test] @@ -742,6 +747,9 @@ mod tests { ) .fetch(url, policy(2, 0)) .await; - assert!(matches!(blocked, Err(Error::BlockedUrl))); + assert!(matches!( + blocked.as_ref().err().map(Error::repr), + Some(Repr::BlockedUrl) + )); } } diff --git a/litellm-rust/crates/http/src/outbound.rs b/litellm-rust/crates/http/src/outbound.rs index d51a1994044..20b0b1be972 100644 --- a/litellm-rust/crates/http/src/outbound.rs +++ b/litellm-rust/crates/http/src/outbound.rs @@ -14,7 +14,7 @@ use reqwest::{ }; use serde::Serialize; -use crate::Error; +use crate::{Error, error::Repr}; /// Returns the headers to add to the request; it never sees a mutable request. pub trait RequestSigner: Send + Sync { @@ -27,8 +27,8 @@ pub fn json( body: &impl Serialize, timeout: Option, ) -> Result { - let body = serde_json::to_vec(body).map_err(|error| Error::RequestBody(error.to_string()))?; - let url = Url::parse(url).map_err(|error| Error::Request(error.to_string()))?; + let body = serde_json::to_vec(body).map_err(Repr::RequestBody)?; + let url = Url::parse(url).map_err(Repr::ParseUrl)?; let with_headers = headers.iter().fold( http::Request::post(url.as_str()), |builder, (name, value)| builder.header(name, value), @@ -41,17 +41,15 @@ pub fn json( } else { with_headers.header(CONTENT_TYPE, "application/json") }; - let mut request = builder - .body(body) - .map_err(|error| Error::Request(error.to_string())) - .and_then(|request| { - Request::try_from(request) - .map_err(|error| Error::Request(error.without_url().to_string())) - })?; + let mut request = from_http(builder.body(body).map_err(Repr::Request)?)?; *request.timeout_mut() = timeout; Ok(request) } +pub fn from_http>(request: http::Request) -> Result { + Request::try_from(request).map_err(|error| Repr::RequestUrl(error.without_url()).into()) +} + pub fn signed(mut request: Request, signer: &dyn RequestSigner) -> Result { let signature = signer.sign(&request)?; request.headers_mut().extend(signature); @@ -114,7 +112,7 @@ mod tests { .headers() .contains_key(CONTENT_TYPE) .then(HeaderMap::new) - .ok_or_else(|| Error::Signature("content-type was not signed".into())) + .ok_or_else(|| Error::signature(std::fmt::Error)) } } @@ -148,7 +146,7 @@ mod tests { struct Refuses; impl RequestSigner for Refuses { fn sign(&self, _request: &Request) -> Result { - Err(Error::ComputedHeader("authorization".into())) + Err(Error::computed_header("authorization")) } } @@ -157,7 +155,7 @@ mod tests { &Refuses, ) .unwrap_err(); - assert_eq!(error, Error::ComputedHeader("authorization".into())); + assert!(error.is_computed_header()); } #[rstest] @@ -230,7 +228,21 @@ mod tests { ); assert!(matches!( json(url, &[("x-caller".into(), header.into())], &json!({}), None), - Err(Error::Request(_)) + Err(error) if error.is_request() )); } + + #[test] + fn a_relative_http_request_is_the_callers_fault_and_never_reached_the_provider() { + let request = http::Request::get("/private?api_key=secret") + .body(Vec::::new()) + .unwrap(); + + let error = from_http(request).unwrap_err(); + + assert!(error.is_request()); + assert!(!error.dispatched()); + let report = litellm_error::Report::new(&error).to_string(); + assert!(!report.contains("secret"), "{report}"); + } } diff --git a/litellm-rust/crates/http/src/pool.rs b/litellm-rust/crates/http/src/pool.rs index bf892b09921..3bdc63453ef 100644 --- a/litellm-rust/crates/http/src/pool.rs +++ b/litellm-rust/crates/http/src/pool.rs @@ -9,7 +9,7 @@ use reqwest::dns::Resolve; use crate::{ client::Client, config::{HttpClientConfig, Verify}, - error::Error, + error::{Error, Repr}, proxy::EnvironmentProxies, }; @@ -83,7 +83,8 @@ impl HttpClientPool { } let client = self .apply(variant, reqwest::ClientBuilder::try_from(&key.0)?) - .build()?; + .build() + .map_err(Repr::Client)?; self.lock().insert( key, PooledClient { diff --git a/litellm-rust/crates/http/src/request.rs b/litellm-rust/crates/http/src/request.rs index 5a3dd19ade2..11e19f7fa24 100644 --- a/litellm-rust/crates/http/src/request.rs +++ b/litellm-rust/crates/http/src/request.rs @@ -9,9 +9,7 @@ pub struct HeaderError { use litellm_core_utils::core_helpers::json_type_name; use serde_json::{Map, Value}; -/// Max characters of an upstream error body echoed across the call boundary -/// before truncation, so provider bodies are bounded and data-minimized. -const UPSTREAM_ERROR_BODY_MAX_CHARS: usize = 256; +use crate::Error; pub enum HeaderPolicy<'a> { All, @@ -40,17 +38,15 @@ pub fn with_headers( }) } -pub async fn http_request( - request: reqwest::RequestBuilder, -) -> Result { - request.send().await +pub async fn http_request(request: reqwest::RequestBuilder) -> Result { + Ok(request.send().await?) } pub async fn execute_http_request( client: &reqwest::Client, request: reqwest::Request, -) -> Result { - client.execute(request).await +) -> Result { + Ok(client.execute(request).await?) } pub fn header_pairs(headers: &reqwest::header::HeaderMap) -> Vec<(String, String)> { @@ -65,14 +61,6 @@ pub fn header_pairs(headers: &reqwest::header::HeaderMap) -> Vec<(String, String .collect() } -pub fn truncate_error_body(body: &str) -> String { - if body.chars().count() <= UPSTREAM_ERROR_BODY_MAX_CHARS { - return body.to_string(); - } - let truncated: String = body.chars().take(UPSTREAM_ERROR_BODY_MAX_CHARS).collect(); - format!("{truncated}... (truncated)") -} - pub fn string_headers( context: &'static str, extra_headers: Option>, @@ -215,22 +203,6 @@ mod tests { assert_ne!(request.headers()["content-length"], "0"); } - #[test] - fn truncate_leaves_short_bodies_untouched() { - assert_eq!(truncate_error_body("short"), "short"); - } - - #[test] - fn truncate_bounds_long_bodies_by_characters() { - let body = "\u{00e9}".repeat(UPSTREAM_ERROR_BODY_MAX_CHARS + 10); - let truncated = truncate_error_body(&body); - assert!(truncated.ends_with("... (truncated)")); - assert_eq!( - truncated.chars().count(), - UPSTREAM_ERROR_BODY_MAX_CHARS + "... (truncated)".chars().count() - ); - } - #[test] fn string_headers_rejects_non_string_values() { let headers = Map::from_iter([("x-trace".to_string(), json!(7))]); diff --git a/litellm-rust/crates/http/src/tls.rs b/litellm-rust/crates/http/src/tls.rs index a52ea31af59..a83b0018b2b 100644 --- a/litellm-rust/crates/http/src/tls.rs +++ b/litellm-rust/crates/http/src/tls.rs @@ -1,4 +1,4 @@ -use std::{fmt, path::Path, str::FromStr, sync::Arc}; +use std::{path::Path, str::FromStr, sync::Arc}; use rustls::{ CipherSuite, ClientConfig, DigitallySignedStruct, RootCertStore, SignatureScheme, @@ -9,7 +9,7 @@ use rustls::{ use crate::{ config::{ClientIdentity, HttpClientConfig, Verify}, - error::{Error, TlsSource}, + error::{Error, PemFault, Repr, TlsSource}, }; #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] @@ -187,7 +187,7 @@ impl TryFrom<&HttpClientConfig> for ClientConfig { }); let builder = ClientConfig::builder_with_provider(Arc::clone(&provider)) .with_safe_default_protocol_versions() - .map_err(|error| Error::Client(error.to_string()))?; + .map_err(Repr::TlsConfig)?; let verified = match &config.verify { Verify::Disabled => builder .dangerous() @@ -230,7 +230,7 @@ impl TryFrom<&HttpClientConfig> for ClientConfig { fn bundle_roots(path: &Path, source: TlsSource) -> Result { let certificates = certificates(path, source)?; if certificates.is_empty() { - return Err(invalid_pem(path, source, "no certificates found")); + return Err(invalid_pem(path, source, PemFault::NoCertificates)); } let mut store = RootCertStore::empty(); for certificate in certificates { @@ -248,7 +248,7 @@ fn client_identity( let source = TlsSource::ClientIdentity; let chain = certificates(certificate, source)?; if chain.is_empty() { - return Err(invalid_pem(certificate, source, "no certificates found")); + return Err(invalid_pem(certificate, source, PemFault::NoCertificates)); } let private_key = PrivateKeyDer::from_pem_slice(&read(key, source)?) .map_err(|error| invalid_pem(key, source, error))?; @@ -262,19 +262,23 @@ fn certificates(path: &Path, source: TlsSource) -> Result Result, Error> { - std::fs::read(path).map_err(|error| Error::Read { - path: path.to_path_buf(), - message: error.to_string(), - tls_source: source, + std::fs::read(path).map_err(|error| { + Repr::Read { + path: path.to_path_buf(), + tls_source: source, + source: error, + } + .into() }) } -fn invalid_pem(path: &Path, source: TlsSource, message: impl fmt::Display) -> Error { - Error::InvalidPem { +fn invalid_pem(path: &Path, tls_source: TlsSource, fault: impl Into) -> Error { + Repr::InvalidPem { path: path.to_path_buf(), - message: message.to_string(), - tls_source: source, + tls_source, + source: fault.into(), } + .into() } #[derive(Debug)] @@ -454,11 +458,10 @@ mod tests { std::fs::remove_file(&path).unwrap(); assert!(matches!( result, - Err(Error::InvalidPem { - path: reported, - tls_source: TlsSource::ClientIdentity, - .. - }) if reported == path + Err(error) if matches!( + error.repr(), + Repr::InvalidPem { path: reported, tls_source: TlsSource::ClientIdentity, .. } if *reported == path + ) )); } diff --git a/litellm-rust/crates/http/src/transport.rs b/litellm-rust/crates/http/src/transport.rs deleted file mode 100644 index 8814925bbf2..00000000000 --- a/litellm-rust/crates/http/src/transport.rs +++ /dev/null @@ -1,109 +0,0 @@ -#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] -pub enum Error { - #[error("upstream request failed with status {status}: {body}")] - Http { status: u16, body: String }, - #[error("upstream network error: {0}")] - Network(String), - #[error("could not reach the provider: {0}")] - Connect(String), -} - -impl Error { - pub fn from_reqwest_before_dispatch(error: reqwest::Error) -> Self { - let before_dispatch = !error.is_timeout() && (error.is_connect() || error.is_builder()); - let message = describe(error); - if before_dispatch { - Self::Connect(message) - } else { - Self::Network(message) - } - } -} - -impl From for Error { - fn from(error: reqwest::Error) -> Self { - Self::Network(describe(error)) - } -} - -fn describe(error: reqwest::Error) -> String { - let error = error.without_url(); - std::iter::successors(std::error::Error::source(&error), |cause| cause.source()) - .fold(error.to_string(), |message, cause| { - format!("{message}: {cause}") - }) -} - -#[cfg(test)] -mod tests { - #[tokio::test] - async fn transport_errors_remove_urls_and_keep_dispatch_context() { - let error = reqwest::Client::builder() - .no_proxy() - .build() - .expect("client") - .get("http://localhost:invalid/private?api_key=secret") - .send() - .await - .expect_err("invalid port"); - let error = crate::transport::Error::from_reqwest_before_dispatch(error); - assert!(matches!(error, crate::transport::Error::Connect(_))); - assert!(!error.to_string().contains("secret")); - assert!(!error.to_string().contains("private")); - } - - fn root_cause(error: &dyn std::error::Error) -> Option { - match error.source() { - Some(cause) => root_cause(cause).or_else(|| Some(cause.to_string())), - None => None, - } - } - - #[tokio::test] - async fn network_error_message_names_the_underlying_cause() { - let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind"); - let address = listener.local_addr().expect("address"); - drop(listener); - let error = reqwest::Client::builder() - .no_proxy() - .build() - .expect("client") - .get(format!("http://{address}/private?api_key=secret")) - .send() - .await - .expect_err("nothing listens on the port"); - let root_cause = root_cause(&error).expect("reqwest reports a cause"); - let message = crate::transport::Error::from(error).to_string(); - assert!(message.contains(&root_cause), "{message}"); - assert!(!message.contains("secret")); - } - - #[tokio::test] - async fn request_timeout_is_not_safe_to_retry_as_a_connect_failure() { - use std::time::Duration; - let listener = tokio::net::TcpListener::bind("127.0.0.1:0") - .await - .expect("bind"); - let address = listener.local_addr().expect("address"); - let request = reqwest::Client::builder() - .no_proxy() - .build() - .expect("client") - .get(format!("http://{address}")) - .timeout(Duration::from_millis(200)) - .send(); - let (response, accepted) = tokio::join!( - request, - tokio::time::timeout(Duration::from_secs(2), listener.accept()) - ); - let _connection = accepted - .expect("accept deadline") - .expect("accepted connection"); - let error = response.expect_err("server does not respond"); - assert!(error.is_timeout()); - assert!(matches!( - crate::transport::Error::from_reqwest_before_dispatch(error), - crate::transport::Error::Network(_) - )); - } -} diff --git a/litellm-rust/crates/http/src/websocket.rs b/litellm-rust/crates/http/src/websocket.rs index 685c79cf3a1..b2a6571e2b8 100644 --- a/litellm-rust/crates/http/src/websocket.rs +++ b/litellm-rust/crates/http/src/websocket.rs @@ -1,3 +1,5 @@ +use std::time::Duration; + use http::{ HeaderMap, HeaderValue, StatusCode, header::{ @@ -14,7 +16,10 @@ use tokio_tungstenite::{ }, }; -use crate::{client::Client, transport::Error}; +use crate::{ + client::Client, + error::{Error, HandshakeFault, Repr}, +}; pub type UpstreamWebSocket = WebSocketStream; @@ -23,6 +28,20 @@ pub async fn connect( client: &Client, url: &str, headers: HeaderMap, + timeout: Option, +) -> Result { + match timeout { + Some(timeout) => tokio::time::timeout(timeout, handshake(client, url, headers)) + .await + .map_err(Repr::HandshakeTimeout)?, + None => handshake(client, url, headers).await, + } +} + +async fn handshake( + client: &Client, + url: &str, + headers: HeaderMap, ) -> Result { let offered_protocols = headers.get(SEC_WEBSOCKET_PROTOCOL).cloned(); let key = generate_key(); @@ -30,14 +49,13 @@ pub async fn connect( .get(http_url(url)?) .headers(handshake_headers(headers, &key)?) .send() - .await - .map_err(Error::from_reqwest_before_dispatch)?; + .await?; let status = response.status(); if status != StatusCode::SWITCHING_PROTOCOLS { - return Err(Error::Http { - status: status.as_u16(), - body: response.text().await.unwrap_or_default(), - }); + return Err(Error::upstream( + status.as_u16(), + response.text().await.unwrap_or_default(), + )); } let expected = derive_accept_key(key.as_bytes()); if response @@ -46,9 +64,7 @@ pub async fn connect( .map(HeaderValue::as_bytes) != Some(expected.as_bytes()) { - return Err(Error::Network( - "WebSocket handshake returned a wrong Sec-WebSocket-Accept".into(), - )); + return Err(Repr::Handshake(HandshakeFault::Accept).into()); } validate_response_headers(response.headers(), offered_protocols.as_ref())?; let upgraded = response.upgrade().await?; @@ -63,9 +79,7 @@ fn validate_response_headers( .get(UPGRADE) .is_some_and(|value| value.as_bytes().eq_ignore_ascii_case(b"websocket")) { - return Err(Error::Network( - "WebSocket handshake returned an invalid Upgrade header".into(), - )); + return Err(Repr::Handshake(HandshakeFault::Upgrade).into()); } if !headers.get_all(CONNECTION).iter().any(|value| { value.to_str().is_ok_and(|value| { @@ -74,9 +88,7 @@ fn validate_response_headers( .any(|token| token.trim().eq_ignore_ascii_case("upgrade")) }) }) { - return Err(Error::Network( - "WebSocket handshake returned an invalid Connection header".into(), - )); + return Err(Repr::Handshake(HandshakeFault::Connection).into()); } let protocol_matches = match (offered_protocols, headers.get(SEC_WEBSOCKET_PROTOCOL)) { (_, None) => true, @@ -92,31 +104,25 @@ fn validate_response_headers( _ => false, }; if !protocol_matches { - return Err(Error::Network( - "WebSocket handshake returned an invalid subprotocol".into(), - )); + return Err(Repr::Handshake(HandshakeFault::Subprotocol).into()); } Ok(()) } fn http_url(url: &str) -> Result { - let mut url = Url::parse(url).map_err(not_dispatched)?; + let mut url = Url::parse(url).map_err(Repr::WebSocketUrl)?; let scheme = match url.scheme() { "ws" | "http" => "http", "wss" | "https" => "https", - other => { - return Err(not_dispatched(format!( - "unsupported WebSocket URL scheme {other:?}" - ))); - } + other => return Err(Repr::WebSocketScheme(other.to_owned()).into()), }; url.set_scheme(scheme) - .map_err(|()| not_dispatched("WebSocket URL cannot be sent over HTTP"))?; + .map_err(|()| Repr::WebSocketSchemeRewrite(scheme))?; Ok(url) } fn handshake_headers(headers: HeaderMap, key: &str) -> Result { - let key = HeaderValue::from_str(key).map_err(not_dispatched)?; + let key = HeaderValue::from_str(key).map_err(|error| Repr::Request(error.into()))?; Ok([ (CONNECTION, HeaderValue::from_static("Upgrade")), (UPGRADE, HeaderValue::from_static("websocket")), @@ -129,7 +135,3 @@ fn handshake_headers(headers: HeaderMap, key: &str) -> Result headers })) } - -fn not_dispatched(error: impl std::fmt::Display) -> Error { - Error::Connect(error.to_string()) -} diff --git a/litellm-rust/crates/http/tests/websocket.rs b/litellm-rust/crates/http/tests/websocket.rs index 94aba97a84d..a5758d45e56 100644 --- a/litellm-rust/crates/http/tests/websocket.rs +++ b/litellm-rust/crates/http/tests/websocket.rs @@ -7,9 +7,8 @@ use std::{ use futures_util::{SinkExt, StreamExt}; use http::{HeaderMap, HeaderValue}; use litellm_http::{ - ClientIdentity, ClientVariant, EnvironmentProxies, HttpClientConfig, HttpClientPool, - HttpSettings, Resolution, Verify, media::PublicDnsResolver, transport::Error, - websocket::connect, + ClientIdentity, ClientVariant, EnvironmentProxies, Error, HttpClientConfig, HttpClientPool, + HttpSettings, Resolution, Verify, media::PublicDnsResolver, websocket::connect, }; use rcgen::{BasicConstraints, CertificateParams, CertifiedIssuer, IsCa, KeyPair}; use rstest::{Context, rstest}; @@ -122,7 +121,7 @@ async fn open( let client = HttpClientPool::new(Arc::new(PublicDnsResolver)) .client(config, ClientVariant::WebSocket) .unwrap(); - connect(&client, url, headers).await + connect(&client, url, headers, None).await } async fn echo_once( @@ -269,7 +268,7 @@ async fn wss_honors_ssl_verify(#[case] verify: fn(&Pki) -> Verify, #[case] trust match result { Ok(mut socket) if trusted => round_trip(&mut socket).await, - Err(Error::Connect(_)) if !trusted => {} + Err(error) if !trusted && !error.dispatched() => {} other => panic!("unexpected handshake outcome: {:?}", other.map(drop)), } } @@ -290,7 +289,7 @@ async fn websocket_uses_system_roots( let result = open(&config(Verify::BuiltInRoots), &url, HeaderMap::new()).await; match result { Ok(mut socket) if connects => round_trip(&mut socket).await, - Err(Error::Connect(_)) if !connects => {} + Err(error) if !connects && error.is_connect() => {} other => panic!("unexpected handshake outcome: {:?}", other.map(drop)), } return; @@ -407,13 +406,8 @@ async fn refused_upgrade_reports_the_status_and_body() { ) .await; - assert_eq!( - result.map(drop), - Err(Error::Http { - status: 403, - body: "denied".into() - }) - ); + let error = result.map(drop).unwrap_err(); + assert_eq!(error.upstream_response(), Some((403, "denied"))); } #[rstest] @@ -431,7 +425,7 @@ async fn upgrade_with_a_wrong_accept_key_is_rejected() { ) .await; - assert!(matches!(result, Err(Error::Network(_)))); + assert!(matches!(result, Err(error) if error.dispatched() && !error.is_request())); } #[rstest] @@ -441,7 +435,7 @@ async fn upgrade_with_a_wrong_accept_key_is_rejected() { async fn invalid_url_is_rejected_before_dispatch(#[case] url: &str) { let result = open(&config(Verify::BuiltInRoots), url, HeaderMap::new()).await; - assert!(matches!(result, Err(Error::Connect(_)))); + assert!(matches!(result, Err(error) if !error.dispatched() && error.is_request())); } #[rstest] @@ -533,7 +527,34 @@ async fn upgrade_validates_headers_and_negotiated_protocol( match result { Ok(mut socket) if accepted => round_trip(&mut socket).await, - Err(Error::Network(_)) if !accepted => {} + Err(error) if !accepted && error.dispatched() => {} other => panic!("unexpected handshake outcome: {:?}", other.map(drop)), } } + +#[rstest] +#[tokio::test] +async fn a_handshake_that_never_answers_times_out_as_dispatched() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + read_head(&mut stream).await; + let _ = stream.read(&mut [0u8; 1]).await; + }); + let client = HttpClientPool::new(Arc::new(PublicDnsResolver)) + .client(&config(Verify::BuiltInRoots), ClientVariant::WebSocket) + .unwrap(); + + let error = connect( + &client, + &format!("ws://{address}/responses"), + HeaderMap::new(), + Some(std::time::Duration::from_millis(100)), + ) + .await + .map(drop) + .unwrap_err(); + + assert!(error.dispatched()); +} diff --git a/litellm-rust/crates/inference-chat/Cargo.toml b/litellm-rust/crates/inference-chat/Cargo.toml index fa3fad067a9..8a779d44976 100644 --- a/litellm-rust/crates/inference-chat/Cargo.toml +++ b/litellm-rust/crates/inference-chat/Cargo.toml @@ -18,11 +18,11 @@ litellm-secrets.workspace = true serde_json = { workspace = true, features = ["preserve_order"] } futures-util.workspace = true -reqwest.workspace = true tracing.workspace = true [dev-dependencies] litellm-cache-memory.workspace = true +litellm-error.workspace = true litellm-host-native.workspace = true litellm-http = { workspace = true, features = ["test-support"] } litellm-inference-testing.workspace = true diff --git a/litellm-rust/crates/inference-chat/src/common_utils.rs b/litellm-rust/crates/inference-chat/src/common_utils.rs index 4063f33432d..b55aba22325 100644 --- a/litellm-rust/crates/inference-chat/src/common_utils.rs +++ b/litellm-rust/crates/inference-chat/src/common_utils.rs @@ -50,14 +50,16 @@ pub(super) fn resolve_request( config: &'static dyn BaseConfig, ) -> Result { let messages = serde_json::from_value::>(call.messages).map_err(|err| { - Error::InvalidRequest(litellm_llms::ErrorDetail::invalid( + Error::invalid_request(litellm_llms::ErrorDetail::invalid( "chat completions messages", err, )) })?; if messages.is_empty() { - return Err(Error::InvalidRequest( - "chat completions requires at least one message".into(), + return Err(Error::invalid_request( + litellm_llms::ErrorDetail::NoMessages { + route: "chat completions", + }, )); } if let Some(reason) = config.unsupported_reason(&messages, &call.optional_params) { diff --git a/litellm-rust/crates/inference-chat/src/handler.rs b/litellm-rust/crates/inference-chat/src/handler.rs index 15a45747997..c34d60e1b8a 100644 --- a/litellm-rust/crates/inference-chat/src/handler.rs +++ b/litellm-rust/crates/inference-chat/src/handler.rs @@ -3,10 +3,7 @@ use std::time::Duration; use litellm_auth::AuthServices; use litellm_host::interceptors::{Interceptors, RawResponse, RequestContext, WireRequest}; -use litellm_http::{ - Client, - request::{truncate_error_body, with_default_headers}, -}; +use litellm_http::{Client, request::with_default_headers}; use litellm_llms::base_llm::{ auth::{Authenticated, ValidatedEnvironment, resolve_auth}, chat::transformation::BaseConfig, @@ -90,39 +87,26 @@ pub(super) async fn execute( interceptors, observers, || async move { - let outbound = outbound_request( + let outbound = litellm_inference::outbound::outbound_request( Authenticated { headers: wire.headers, signer: authenticated.signer, }, &wire.url, &wire.body, - call.timeout, + Some( + call.timeout + .unwrap_or(Duration::from_secs(CHAT_COMPLETIONS_TIMEOUT_SECS)), + ), )?; - let response = litellm_inference::outbound::send(outbound, http) - .await - .map_err(|err| { - // Failing to establish the connection means the request never went out, - // so the host can still serve it. Everything else here, a timeout - // above all, may have reached the provider and been answered. - if err.is_connect() || err.is_builder() { - Error::Transport(litellm_http::transport::Error::Connect(err.to_string())) - } else { - Error::Transport(litellm_http::transport::Error::Network(err.to_string())) - } - })?; + let response = litellm_inference::outbound::send(outbound, http).await?; let status = response.status(); - let text = response.text().await.map_err(|err| { - Error::Transport(litellm_http::transport::Error::Network(err.to_string())) - })?; + let text = response.text().await.map_err(litellm_http::Error::from)?; if !status.is_success() { - return Err(Error::Transport(litellm_http::transport::Error::Http { - status: status.as_u16(), - body: truncate_error_body(&text), - })); + return Err(litellm_http::Error::upstream(status.as_u16(), text).into()); } let raw = RawResponse { body: text.clone() }; if let Some(observers) = observers { @@ -136,7 +120,7 @@ pub(super) async fn execute( .map_err(Error::post_call)?; let body: Value = serde_json::from_str(&text).map_err(|err| { - Error::InvalidResponse(litellm_llms::ErrorDetail::invalid( + Error::invalid_response(litellm_llms::ErrorDetail::invalid( "chat completions response JSON", err, )) @@ -160,33 +144,13 @@ pub(super) async fn execute( /// implementation must not retry those, so collapse them to one variant that /// can only mean the provider was already called. pub(super) fn as_response_error(err: Error) -> Error { - match err { - already @ (Error::InvalidResponse(_) - | Error::Transport(litellm_http::transport::Error::Http { .. })) => already, - other => Error::InvalidResponse(other.to_string().into()), + if err.is_response() || err.upstream_response().is_some() { + return err; } -} - -pub(super) fn outbound_request( - authenticated: Authenticated, - url: &str, - body: &Value, - timeout: Option, -) -> Result { - litellm_inference::outbound::outbound_request( - authenticated, - url, - body, - Some(timeout.unwrap_or(Duration::from_secs(CHAT_COMPLETIONS_TIMEOUT_SECS))), - ) - .map_err(|error| match error { - // Python drops the caller's copy and prefers a forwarded Authorization - // over the signature, so leave the request to it. - litellm_http::Error::ComputedHeader(_) => { - Error::Unsupported("request forwards a header AWS SigV4 computes") - } - other => Error::Http(other), - }) + Error::invalid_response(litellm_llms::ErrorDetail::invalid( + "chat completions response", + err, + )) } #[cfg(test)] @@ -341,31 +305,48 @@ mod tests { .await .expect_err("the upstream failure fails the call"); - assert!(matches!( - error, - Error::Transport(litellm_http::transport::Error::Http { status: 500, .. }) - )); + assert!( + error + .upstream_response() + .is_some_and(|(status, _)| status == 500), + "{error:?}" + ); assert!(interceptors.raw.into_inner().unwrap().is_empty()); } #[rstest::rstest] - fn response_errors_collapse_to_one_variant_that_can_only_mean_already_sent() { - for original in [ - Error::MissingField("usage"), - Error::Unsupported("non-text response content block"), - Error::InvalidRequest("whatever".to_string().into()), - Error::Auth(litellm_auth::Error::InvalidHeader), - ] { - let label = format!("{original:?}"); - assert!( - matches!(as_response_error(original), Error::InvalidResponse(_)), - "{label} must not stay retryable once the provider has answered" - ); - } - let upstream = Error::Transport(litellm_http::transport::Error::Http { - status: 500, - body: "boom".to_string(), + #[case::missing_field(litellm_llms::Error::MissingField("usage").into())] + #[case::unsupported(Error::Unsupported("non-text response content block"))] + #[case::invalid_request(Error::invalid_request(litellm_llms::ErrorDetail::Expected { + subject: "whatever", + expected: "anything else", + }))] + #[case::auth(Error::Auth(litellm_auth::Error::InvalidHeader))] + fn response_errors_collapse_to_a_fault_that_can_only_mean_already_sent( + #[case] original: Error, + ) { + let collapsed = as_response_error(original); + assert!(collapsed.is_response(), "{collapsed:?}"); + assert!(!collapsed.is_request(), "{collapsed:?}"); + assert!(!collapsed.is_unsupported(), "{collapsed:?}"); + } + + #[rstest::rstest] + fn a_response_fault_is_not_wrapped_twice() { + let original = Error::invalid_response(litellm_llms::ErrorDetail::MissingField { + subject: "response", + field: "usage", }); - assert_eq!(as_response_error(upstream.clone()), upstream); + let report = litellm_error::Report::new(&as_response_error(original)).to_string(); + assert!(!report.contains("chat completions response"), "{report}"); + } + + #[test] + fn a_provider_answer_keeps_its_status_and_body() { + let upstream = Error::from(litellm_http::Error::upstream(429, "slow down".into())); + assert_eq!( + as_response_error(upstream).upstream_response(), + Some((429, "slow down")) + ); } } diff --git a/litellm-rust/crates/inference-chat/tests/chat_completions.rs b/litellm-rust/crates/inference-chat/tests/chat_completions.rs index 461d4146a63..6ed513ba32a 100644 --- a/litellm-rust/crates/inference-chat/tests/chat_completions.rs +++ b/litellm-rust/crates/inference-chat/tests/chat_completions.rs @@ -5,7 +5,6 @@ use litellm_host::{ }; use std::time::Duration; -use litellm_http::transport::Error as TransportError; use litellm_inference_chat::{Error, types::ChatCompletionsRequest}; use litellm_llms_types::formats::chat_completions::ChatCompletionsResponse; use rstest::{fixture, rstest}; @@ -180,7 +179,7 @@ async fn a_response_it_cannot_normalize_is_reported_as_already_sent( .await .expect_err("response cannot be normalized"); - assert!(matches!(error, Error::InvalidResponse(_)), "{error:?}"); + assert!(error.is_response(), "{error:?}"); } #[rstest] @@ -201,13 +200,7 @@ async fn an_upstream_error_status_keeps_its_code_and_body( .await .expect_err("upstream rejects"); - assert_eq!( - error, - Error::Transport(TransportError::Http { - status, - body: "slow down".into() - }) - ); + assert_eq!(error.upstream_response(), Some((status, "slow down"))); } #[rstest] @@ -222,10 +215,25 @@ async fn a_connection_that_is_never_established_returns_a_connect_error( .await .expect_err("nothing is listening"); - assert!( - matches!(error, Error::Transport(TransportError::Connect(_))), - "{error:?}" - ); + assert!(!error.dispatched(), "{error:?}"); +} + +#[rstest] +#[tokio::test] +async fn a_transport_error_does_not_echo_the_provider_url( + request: ChatCompletionsRequest<'static>, +) { + let base = format!("{UNREACHABLE_BASE}/tenant-token-1234"); + + let error = complete(ChatCompletionsRequest { + api_base: Some(&base), + ..request + }) + .await + .expect_err("nothing is listening"); + + let report = litellm_error::Report::new(&error).to_string(); + assert!(!report.contains("tenant-token-1234"), "{report}"); } #[rstest] @@ -243,10 +251,8 @@ async fn a_timeout_after_sending_returns_a_network_error(request: ChatCompletion .await .expect_err("the call times out"); - assert!( - matches!(error, Error::Transport(TransportError::Network(_))), - "{error:?}" - ); + assert!(error.dispatched(), "{error:?}"); + assert!(!error.is_request(), "{error:?}"); } #[rstest] @@ -342,7 +348,7 @@ async fn a_post_call_hook_failure_never_looks_safe_to_retry( Ok(wire) } async fn after_provider_response(&self, _: RawResponse) -> Result<(), Error> { - Err(Error::InvalidRequest("callback rejected".into())) + Err(litellm_llms::Error::MissingField("callback rejected").into()) } } let upstream = upstream([anthropic_response(ANTHROPIC_MESSAGE)]).await; @@ -358,10 +364,10 @@ async fn a_post_call_hook_failure_never_looks_safe_to_retry( ) .await .unwrap_err(); - let Error::PostCallHook(source) = error else { - panic!("expected retained callback error") - }; - assert_eq!(*source, Error::InvalidRequest("callback rejected".into())); + assert!(error.dispatched(), "{error:?}"); + assert!(!error.is_request(), "{error:?}"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains("callback rejected"), "{report}"); assert_eq!(received(&upstream).await.len(), 1); } diff --git a/litellm-rust/crates/inference-chat/tests/request.rs b/litellm-rust/crates/inference-chat/tests/request.rs index 5988a7baee5..4025c137cdc 100644 --- a/litellm-rust/crates/inference-chat/tests/request.rs +++ b/litellm-rust/crates/inference-chat/tests/request.rs @@ -211,7 +211,7 @@ async fn rejects_empty_messages_before_resolving_credentials() { }) .await .expect_err("empty messages are rejected"); - assert!(matches!(error, Error::InvalidRequest(_))); + assert!(error.is_request(), "{error:?}"); } #[rstest] @@ -233,26 +233,27 @@ async fn rejects_an_unknown_provider() { #[rstest] #[tokio::test] async fn rejects_a_model_with_no_resolvable_provider() { - assert!(matches!( - complete(request( - "claude-sonnet-4-5", - None, - json!([{"role": "user", "content": "hi"}]), - json!({}), - )) - .await, - Err(Error::InvalidProvider(_)) - )); + let error = complete(request( + "claude-sonnet-4-5", + None, + json!([{"role": "user", "content": "hi"}]), + json!({}), + )) + .await + .unwrap_err(); + assert!(error.is_request(), "{error:?}"); + let report = litellm_error::Report::new(&error).to_string(); + assert!( + report.contains("unable to resolve custom_llm_provider"), + "{report}" + ); } #[rstest] -#[case::empty(json!([]), Some("chat completions requires at least one message"))] -#[case::malformed(json!("not a list"), None)] +#[case::empty(json!([]), "chat completions requires at least one message")] +#[case::malformed(json!("not a list"), "invalid chat completions messages")] #[tokio::test] -async fn rejects_an_empty_or_malformed_message_list( - #[case] messages: Value, - #[case] expected: Option<&str>, -) { +async fn rejects_an_empty_or_malformed_message_list(#[case] messages: Value, #[case] reason: &str) { let error = complete(request( "anthropic/claude-sonnet-4-5", None, @@ -261,11 +262,9 @@ async fn rejects_an_empty_or_malformed_message_list( )) .await .unwrap_err(); - if let Some(expected) = expected { - assert_eq!(error, Error::InvalidRequest(expected.to_string().into())); - } else { - assert!(matches!(error, Error::InvalidRequest(_))); - } + assert!(error.is_request(), "{error:?}"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains(reason), "{report}"); } #[rstest] diff --git a/litellm-rust/crates/inference-messages/Cargo.toml b/litellm-rust/crates/inference-messages/Cargo.toml index e310a0be9b1..694d9b4de31 100644 --- a/litellm-rust/crates/inference-messages/Cargo.toml +++ b/litellm-rust/crates/inference-messages/Cargo.toml @@ -26,6 +26,7 @@ serde_json = { workspace = true, features = ["preserve_order"] } tracing.workspace = true [dev-dependencies] +litellm-error.workspace = true litellm-framer = { workspace = true, features = ["sse"] } litellm-cache-memory.workspace = true litellm-host-native.workspace = true diff --git a/litellm-rust/crates/inference-messages/src/common_utils.rs b/litellm-rust/crates/inference-messages/src/common_utils.rs index 9be7ee6ffc3..a57eca1c6a9 100644 --- a/litellm-rust/crates/inference-messages/src/common_utils.rs +++ b/litellm-rust/crates/inference-messages/src/common_utils.rs @@ -1,6 +1,5 @@ use litellm_core_utils::get_llm_provider_logic::LlmProviders; use litellm_http::request::string_headers as shared_string_headers; -pub(super) use litellm_http::request::truncate_error_body; use litellm_llms::tencent::messages::transformation::TENCENT_MESSAGES_CONFIG; use litellm_llms::{ anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG, @@ -77,7 +76,7 @@ mod tests { use rstest::rstest; - use super::{messages_provider, string_headers, truncate_error_body}; + use super::{messages_provider, string_headers}; use crate::Error; use litellm_core_utils::get_llm_provider_logic::LlmProviders; @@ -106,19 +105,6 @@ mod tests { assert!(messages_provider(provider, model).is_none()); } - #[test] - fn truncate_error_body_caps_long_payloads() { - let body = "x".repeat(400); - let truncated = truncate_error_body(&body); - assert!(truncated.ends_with("... (truncated)")); - let prefix_chars = truncated - .strip_suffix("... (truncated)") - .expect("truncated marker present") - .chars() - .count(); - assert_eq!(prefix_chars, 256); - } - #[test] fn string_headers_rejects_non_string_values() { let headers = json!({"x-count": 3}).as_object().unwrap().clone(); diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index 292995d842c..67c3b8f98b5 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -10,7 +10,7 @@ use litellm_core_utils::{ use litellm_host::interceptors::{ Cost, ExecutionFacts, Interceptors, ProviderIdentity, RequestContext, ResultSource, WireRequest, }; -use litellm_http::{request::with_default_headers, transport::Error as TransportError}; +use litellm_http::request::with_default_headers; use litellm_llms::base_llm::{ auth::{Authenticated, ValidatedEnvironment, resolve_auth}, messages::{ @@ -26,7 +26,7 @@ use serde_json::Value; use super::{ Error, MessagesCall, MessagesCallResponse, MessagesRoute, - common_utils::{MessagesProvider, string_headers, truncate_error_body}, + common_utils::{MessagesProvider, string_headers}, constants::MESSAGES_TIMEOUT_SECS, route::Messages, types::invalid_request, @@ -188,7 +188,7 @@ async fn provider_call( None | Some(Value::Null) => false, Some(Value::Bool(stream)) => *stream, Some(value) => { - return Err(Error::InvalidRequest( + return Err(Error::invalid_request( litellm_llms::ErrorDetail::InvalidValue { field: "stream", expected: "a boolean", @@ -261,7 +261,7 @@ async fn call_provider( provider_name, )); } - let text = response.text().await.map_err(network)?; + let text = response.text().await.map_err(litellm_http::Error::from)?; log_response_body(&text); context.response_received(&text).await?; decode_response(config, &identity.model, &text) @@ -269,16 +269,12 @@ async fn call_provider( } fn serialize_failure(err: serde_json::Error) -> Error { - Error::InvalidRequest(litellm_llms::ErrorDetail::failed( + Error::invalid_request(litellm_llms::ErrorDetail::failed( "Anthropic messages request serialization", err, )) } -fn network(error: reqwest::Error) -> Error { - Error::Transport(TransportError::Network(error.to_string())) -} - async fn send( http: &litellm_http::Client, authenticated: Authenticated, @@ -292,9 +288,7 @@ async fn send( body, Some(timeout.unwrap_or(Duration::from_secs(MESSAGES_TIMEOUT_SECS))), )?; - litellm_inference::outbound::send(request, http) - .await - .map_err(network) + Ok(litellm_inference::outbound::send(request, http).await?) } async fn provider_error(response: reqwest::Response) -> Error { @@ -302,12 +296,9 @@ async fn provider_error(response: reqwest::Response) -> Error { match response.text().await { Ok(text) => { log_error_body(status, &text); - Error::Transport(TransportError::Http { - status, - body: truncate_error_body(&text), - }) + litellm_http::Error::upstream(status, text).into() } - Err(error) => network(error), + Err(error) => litellm_http::Error::from(error).into(), } } @@ -317,7 +308,7 @@ fn decode_response( text: &str, ) -> Result { let response = serde_json::from_str(text).map_err(|err| { - Error::InvalidResponse(litellm_llms::ErrorDetail::invalid( + Error::invalid_response(litellm_llms::ErrorDetail::invalid( "messages response JSON", err, )) @@ -339,7 +330,7 @@ fn streaming_response( .collect(); let chunks = match decoder { None => futures_util::stream::try_unfold(response, move |mut response| async move { - let chunk = response.chunk().await.map_err(network)?; + let chunk = response.chunk().await.map_err(litellm_http::Error::from)?; Ok(chunk.map(|chunk| { log_chunk(provider, "provider_response", &chunk); (chunk, response) @@ -362,7 +353,7 @@ fn decoded_chunks( let bytes: ByteStream = response .bytes_stream() .inspect_ok(move |chunk| log_chunk(provider, "provider_response", chunk)) - .map_err(std::io::Error::other) + .map_err(|error| std::io::Error::other(litellm_http::Error::from(error))) .boxed(); futures_util::stream::try_unfold(decode(bytes), move |mut events| async move { let Some(event) = events.try_next().await? else { @@ -541,14 +532,10 @@ mod tests { shaping, }) .await; - assert!( - matches!( - result, - Err(Error::Auth(litellm_auth::Error::MissingApiKey { .. })) - ), - "{:?}", - result.err() - ); + let Err(error) = result else { + panic!("a call without credentials fails"); + }; + assert!(error.is_missing_credential(), "{error:?}"); } #[rstest] @@ -771,19 +758,18 @@ mod tests { }, ..shaping }; - assert!(matches!( - wire_body( - json!({ - "model": "claude-test", - "messages": [{"role": "user", "content": "hi"}], - "max_tokens": 16, - "metadata": {"user_id": 123} - }), - shaping, - ) - .await, - Err(Error::InvalidRequest(_)) - )); + let error = wire_body( + json!({ + "model": "claude-test", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 16, + "metadata": {"user_id": 123} + }), + shaping, + ) + .await + .unwrap_err(); + assert!(error.is_request(), "{error:?}"); } #[rstest] @@ -800,7 +786,7 @@ mod tests { shaping, ) .await, - Err(Error::InvalidRequest( + Err(Error::invalid_request( litellm_llms::ErrorDetail::InvalidValue { field: "metadata.user_id", expected: "a string", @@ -902,7 +888,7 @@ mod tests { let chunk = chunks.next().await.unwrap(); match expected { Some(expected) => assert_eq!(chunk.unwrap().as_ref(), expected.as_bytes()), - None => assert!(matches!(chunk, Err(Error::InvalidResponse(_))), "{chunk:?}"), + None => assert!(chunk.unwrap_err().is_response()), } assert!(chunks.next().await.is_none()); } diff --git a/litellm-rust/crates/inference-messages/src/types.rs b/litellm-rust/crates/inference-messages/src/types.rs index 21eac0574ed..ae346259559 100644 --- a/litellm-rust/crates/inference-messages/src/types.rs +++ b/litellm-rust/crates/inference-messages/src/types.rs @@ -34,11 +34,11 @@ pub fn messages_body(body: Map) -> Result /// Python's `GenericLiteLLMParams(**kwargs)`. pub fn litellm_params(fields: Map) -> Result { serde_json::from_value(Value::Object(fields)) - .map_err(|err| Error::InvalidRequest(ErrorDetail::invalid("litellm params", err))) + .map_err(|err| Error::invalid_request(ErrorDetail::invalid("litellm params", err))) } pub(super) fn invalid_request(err: serde_json::Error) -> Error { - Error::InvalidRequest(ErrorDetail::invalid("Anthropic messages request", err)) + Error::invalid_request(ErrorDetail::invalid("Anthropic messages request", err)) } pub type MessagesCallResponse = diff --git a/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs b/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs index 5420477c34a..f91d5c79f7c 100644 --- a/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs +++ b/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs @@ -309,7 +309,14 @@ async fn deployment_workspace_is_required_and_sent_as_a_header( }) .await; if !configured { - assert!(matches!(result, Err(Error::Auth(_)))); + let Err(error) = result else { + panic!("the call fails"); + }; + assert!(!error.dispatched(), "{error:?}"); + assert!( + error.is_request() || error.is_missing_credential(), + "{error:?}" + ); assert!(received(&upstream).await.is_empty()); return; } diff --git a/litellm-rust/crates/inference-messages/tests/messages/host.rs b/litellm-rust/crates/inference-messages/tests/messages/host.rs index ba20fb9eec4..48afae99cb5 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/host.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/host.rs @@ -301,7 +301,7 @@ async fn response_mode_follows_the_intercepted_request( let summaries = traces.summaries("litellm.route"); assert_eq!(summaries.len(), 1); let Some(expected_stream) = expected_stream else { - assert!(matches!(result, Err(Error::InvalidRequest(_)))); + assert!(result.unwrap_err().is_request()); assert!(received(&upstream).await.is_empty()); assert_eq!(summaries[0]["outcome"], "failure"); return; @@ -322,14 +322,17 @@ async fn a_before_send_failure_never_sends(call: MessagesCall) { let upstream = upstream([message_response()]).await; let host = RecordingHost::new( authenticated(call, upstream.uri()), - Box::new(|_| Err(Error::InvalidRequest("vetoed by the host".into()))), + Box::new(|_| Err(litellm_llms::Error::MissingField("vetoed by the host").into())), ); let error = run_through(&host) .await .expect_err("the host failure fails the call"); - assert_eq!(error, Error::InvalidRequest("vetoed by the host".into())); + assert_eq!( + error, + Error::from(litellm_llms::Error::MissingField("vetoed by the host")) + ); assert!(received(&upstream).await.is_empty()); assert!(host.raw_responses().is_empty()); } diff --git a/litellm-rust/crates/inference-messages/tests/messages/main.rs b/litellm-rust/crates/inference-messages/tests/messages/main.rs index fcc432acb76..a01e571ec75 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/main.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/main.rs @@ -144,7 +144,7 @@ impl LocalMessagesHost { .lock() .unwrap_or_else(|error| error.into_inner()) .take() - .ok_or_else(|| Error::InvalidRequest("messages request was already projected".into())) + .ok_or(litellm_llms::Error::MissingField("messages request").into()) } pub fn runtime(&self) -> litellm_host_native::in_process::Host<'_, (), Self, ()> { litellm_host_native::in_process::Host { diff --git a/litellm-rust/crates/inference-messages/tests/messages/request.rs b/litellm-rust/crates/inference-messages/tests/messages/request.rs index 93bbbb071ba..e53bdb8d252 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/request.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/request.rs @@ -108,13 +108,7 @@ async fn a_call_without_credentials_fails_before_sending( .await .expect_err("a call without credentials fails"); - assert!( - matches!( - error, - Error::Auth(litellm_auth::Error::MissingApiKey { .. }) - ), - "{error:?}" - ); + assert!(error.is_missing_credential(), "{error:?}"); assert!(received(&upstream).await.is_empty()); } @@ -185,7 +179,9 @@ async fn unsupported_providers_are_rejected_before_sending( .await .expect_err("unsupported provider errors"); - assert_eq!(error, Error::InvalidProvider(reported.into())); + assert!(error.is_request(), "{error:?}"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains(reported), "{report}"); } #[rstest] @@ -513,10 +509,9 @@ async fn unsupported_params_are_dropped_under_drop_params_and_rejected_without_i let error = run(shaped(false)) .await .expect_err("an unsupported param is rejected without drop_params"); - assert!( - matches!(&error, Error::InvalidRequest(message) if message.to_string().contains(rejected_as)), - "{error:?}" - ); + assert!(error.is_request(), "{error:?}"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains(rejected_as), "{report}"); assert!(received(&upstream).await.is_empty()); run_message(shaped(true)).await; @@ -957,11 +952,10 @@ async fn provider_validation_runs_before_caller_parameter_removal( )) .await; let error = result.expect_err("metadata is validated before removal"); + let report = litellm_error::Report::new(&error).to_string(); assert!( - error - .to_string() - .contains("metadata.user_id must be a string"), - "{error}" + report.contains("metadata.user_id must be a string"), + "{report}" ); assert!(received(&upstream).await.is_empty()); } diff --git a/litellm-rust/crates/inference-messages/tests/messages/response.rs b/litellm-rust/crates/inference-messages/tests/messages/response.rs index b5d623a5582..7aff1286858 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/response.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/response.rs @@ -2,7 +2,6 @@ use litellm_host::{ interceptors::{ExecutionFacts, ResultSource}, lifecycle::ExecutionEvent, }; -use litellm_http::transport::Error as TransportError; use litellm_inference_messages::{MessagesCallResponse, messages_body}; use litellm_inference_testing::{ RecordingSecrets, http_config, no_secrets, provider_http, resources, @@ -166,17 +165,21 @@ async fn a_json_error_envelope_is_kept_verbatim(call: MessagesCall) { .await .expect_err("upstream error propagates"); - let Error::Transport(TransportError::Http { status, body }) = error else { + let Some((status, body)) = error.upstream_response() else { panic!("{error:?}"); }; assert_eq!(status, 400); - assert_eq!(serde_json::from_str::(&body).unwrap(), envelope); + assert_eq!(serde_json::from_str::(body).unwrap(), envelope); } #[rstest] +#[case::ascii(format!("{}private-tail", "x".repeat(600)))] +#[case::unicode(format!("{}private-tail", "é".repeat(600)))] #[tokio::test] -async fn a_long_error_body_is_truncated_at_the_documented_cap(call: MessagesCall) { - let long = "x".repeat(600); +async fn a_long_error_body_is_preserved_with_bounded_diagnostics( + call: MessagesCall, + #[case] long: String, +) { let upstream = upstream([ResponseTemplate::new(500).set_body_string(long.clone())]).await; let error = run(MessagesCall { @@ -187,13 +190,13 @@ async fn a_long_error_body_is_truncated_at_the_documented_cap(call: MessagesCall .await .expect_err("upstream error propagates"); - assert_eq!( - error, - Error::Transport(TransportError::Http { - status: 500, - body: format!("{}... (truncated)", &long[..256]) - }) - ); + assert_eq!(error.upstream_response(), Some((500, long.as_str()))); + let diagnostic_body: String = long.chars().take(256).collect(); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.ends_with(&format!("{diagnostic_body}... (truncated)"))); + assert_eq!(report.matches(&diagnostic_body).count(), 1, "{report}"); + assert!(!report.contains("private-tail"), "{report}"); + assert!(!format!("{error:?}").contains("private-tail")); } #[rstest] @@ -216,11 +219,8 @@ async fn an_upstream_error_keeps_its_status_and_body(call: MessagesCall, #[case] .expect_err("upstream error propagates"); assert_eq!( - error, - Error::Transport(TransportError::Http { - status, - body: "upstream said no".into() - }) + error.upstream_response(), + Some((status, "upstream said no")) ); } @@ -242,7 +242,7 @@ async fn an_unreadable_success_body_is_an_invalid_response( .await .expect_err("an unreadable body fails"); - assert!(matches!(error, Error::InvalidResponse(_)), "{error:?}"); + assert!(error.is_response(), "{error:?}"); } #[rstest] @@ -259,7 +259,23 @@ async fn a_provider_slower_than_the_timeout_fails_the_call(call: MessagesCall) { .await .expect_err("the call times out"); - assert!(matches!(error, Error::Transport(_)), "{error:?}"); + assert!(error.dispatched(), "{error:?}"); + assert!(!error.is_request(), "{error:?}"); + assert_eq!(error.upstream_response(), None); +} + +#[rstest] +#[tokio::test] +async fn a_transport_error_does_not_echo_the_provider_url(call: MessagesCall) { + let error = run(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(format!("{UNREACHABLE_BASE}/tenant-token-1234")), + ..call + }) + .await + .expect_err("nothing is listening"); + + assert!(!error.to_string().contains("tenant-token-1234"), "{error}"); } #[rstest] @@ -305,8 +321,11 @@ async fn the_facade_sends_through_the_injected_http_pool_configuration(call: Mes fn a_body_that_does_not_parse_is_an_invalid_request(#[case] raw: Value) { let error = messages_body(object(raw)).expect_err("the body is rejected"); + assert!(error.is_request(), "{error:?}"); assert!( - matches!(&error, Error::InvalidRequest(message) if message.to_string().starts_with("invalid Anthropic messages request: ")), + litellm_error::Report::new(&error) + .to_string() + .contains("invalid Anthropic messages request"), "{error:?}" ); } diff --git a/litellm-rust/crates/inference-messages/tests/messages/secrets.rs b/litellm-rust/crates/inference-messages/tests/messages/secrets.rs index d4f7d2bf3a7..b4a3018ab93 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/secrets.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/secrets.rs @@ -109,7 +109,7 @@ async fn a_secret_manager_failure_fails_the_call_before_sending(call: MessagesCa .expect_err("a secret manager failure fails the call"); assert!( - matches!(&error, Error::Secret(source) if matches!(source.source_error(), litellm_secrets::Error::ManagedSecretMissing)), + matches!(&error, Error::Secret(_)) && !error.is_request(), "{error:?}" ); assert!(received(&upstream).await.is_empty()); diff --git a/litellm-rust/crates/inference-messages/tests/messages/stream.rs b/litellm-rust/crates/inference-messages/tests/messages/stream.rs index cda688107b5..4a8a061f764 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/stream.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/stream.rs @@ -260,13 +260,7 @@ async fn an_upstream_error_fails_the_call_without_opening_the_stream( .await .expect_err("upstream error propagates"); - assert_eq!( - error, - Error::Transport(litellm_http::transport::Error::Http { - status: 429, - body: body.into() - }) - ); + assert_eq!(error.upstream_response(), Some((429, body))); assert!(host.seen.into_inner().unwrap().is_empty()); } @@ -332,7 +326,9 @@ async fn the_timeout_covers_a_stalled_stream_body(call: MessagesCall) { .expect("the stalled stream gives up within the timeout") .expect_err("a stalled body fails the call"); - assert!(matches!(error, Error::Transport(_)), "{error:?}"); + assert!(error.dispatched(), "{error:?}"); + assert!(!error.is_request(), "{error:?}"); + assert_eq!(error.upstream_response(), None); let seen = host.seen.into_inner().unwrap(); assert!( matches!(seen.as_slice(), [Seen::Open(_), Seen::Deliver(chunk)] if chunk.as_ref() == b"event: message_start\ndata: {}\n\n"), @@ -390,13 +386,7 @@ async fn the_sdk_returns_http_errors_before_opening_a_stream(call: MessagesCall) .err() .expect("upstream failure is returned by messages()"); - assert_eq!( - error, - Error::Transport(litellm_http::transport::Error::Http { - status: 429, - body: "slow down".into(), - }) - ); + assert_eq!(error.upstream_response(), Some((429, "slow down"))); } #[rstest] @@ -470,7 +460,9 @@ async fn the_sdk_yields_a_body_error_once_after_delivered_chunks(call: MessagesC .expect("the stalled body times out") .unwrap() .unwrap_err(); - assert!(matches!(error, Error::Transport(_)), "{error:?}"); + assert!(error.dispatched(), "{error:?}"); + assert!(!error.is_request(), "{error:?}"); + assert_eq!(error.upstream_response(), None); assert!(chunks.next().await.is_none()); tokio::time::timeout(Duration::from_secs(5), connection) .await diff --git a/litellm-rust/crates/inference-messages/tests/messages/tencent.rs b/litellm-rust/crates/inference-messages/tests/messages/tencent.rs index dce187f4ea2..4bab7de65e1 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/tencent.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/tencent.rs @@ -79,13 +79,9 @@ async fn missing_credentials_fail_before_sending(tencent_call: MessagesCall) { }) .await .unwrap_err(); - assert!(matches!( - error, - Error::Auth(litellm_auth::Error::MissingApiKey { - environment_variable: "TENCENT_API_KEY", - .. - }) - )); + assert!(error.is_missing_credential(), "{error:?}"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains("TENCENT_API_KEY"), "{report}"); assert!(received(&upstream).await.is_empty()); } diff --git a/litellm-rust/crates/inference-messages/tests/openrouter/main.rs b/litellm-rust/crates/inference-messages/tests/openrouter/main.rs index 74444142d37..80243718768 100644 --- a/litellm-rust/crates/inference-messages/tests/openrouter/main.rs +++ b/litellm-rust/crates/inference-messages/tests/openrouter/main.rs @@ -1,7 +1,7 @@ use std::{sync::Arc, time::Duration}; use futures_util::TryStreamExt; -use litellm_inference_messages::{Error, MessagesCall, MessagesCallResponse, MessagesShaping}; +use litellm_inference_messages::{MessagesCall, MessagesCallResponse, MessagesShaping}; use litellm_inference_testing::{RecordingSecrets, no_secrets}; use rstest::{fixture, rstest}; use serde_json::json; @@ -92,13 +92,12 @@ async fn an_anthropic_key_alone_is_not_an_openrouter_credential(call: MessagesCa None, ) .await; - assert!(matches!( - result, - Err(Error::Auth(litellm_auth::Error::MissingApiKey { - provider: "OpenRouter", - .. - })) - )); + let Err(error) = result else { + panic!("the call fails"); + }; + assert!(error.is_missing_credential(), "{error:?}"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains("OpenRouter"), "{report}"); assert!(received(&upstream).await.is_empty()); } diff --git a/litellm-rust/crates/inference-ocr/Cargo.toml b/litellm-rust/crates/inference-ocr/Cargo.toml index 21ec013f873..bbbdaddb050 100644 --- a/litellm-rust/crates/inference-ocr/Cargo.toml +++ b/litellm-rust/crates/inference-ocr/Cargo.toml @@ -25,6 +25,7 @@ tokio.workspace = true tracing.workspace = true [dev-dependencies] +litellm-error.workspace = true litellm-auth-aws.workspace = true litellm-auth-gcp.workspace = true litellm-host-native.workspace = true @@ -34,4 +35,5 @@ litellm-llms = { workspace = true, features = ["test-support"] } rand.workspace = true rstest.workspace = true time.workspace = true +url.workspace = true wiremock.workspace = true diff --git a/litellm-rust/crates/inference-ocr/src/document.rs b/litellm-rust/crates/inference-ocr/src/document.rs index 0d9a296d289..44c1049102a 100644 --- a/litellm-rust/crates/inference-ocr/src/document.rs +++ b/litellm-rust/crates/inference-ocr/src/document.rs @@ -46,9 +46,9 @@ pub fn read_path_document(path: &Path, mime_type: Option<&str>) -> Result( model: &'a str, custom_llm_provider: Option<&'a str>, ) -> Result<(ResolvedProvider<'a>, OcrConfigKind), Error> { - let provider = match resolve_llm_provider(model, custom_llm_provider, "ocr") { - Ok(provider) => provider, - Err(_) if get_custom_llm_provider(model, custom_llm_provider).is_none() => { - ResolvedProvider { - model, - provider: LlmProviders::Mistral, - } - } - Err(litellm_inference::Error::InvalidProvider(provider)) => { - return Err(Error::InvalidProvider(provider)); - } - Err(error) => return Err(Error::InvalidProvider(error.to_string())), + let provider = match get_custom_llm_provider(model, custom_llm_provider) { + None => ResolvedProvider { + model, + provider: LlmProviders::Mistral, + }, + Some(CustomLlmProvider { + model, + custom_llm_provider, + }) => ResolvedProvider { + model, + provider: custom_llm_provider + .parse() + .map_err(|_| Error::InvalidProvider(custom_llm_provider.to_owned()))?, + }, }; let config = match provider.provider { LlmProviders::AwsTextract => match TextractOperation::from_model(provider.model)? { diff --git a/litellm-rust/crates/inference-ocr/src/route.rs b/litellm-rust/crates/inference-ocr/src/route.rs index 85125452bda..74ff829e8ad 100644 --- a/litellm-rust/crates/inference-ocr/src/route.rs +++ b/litellm-rust/crates/inference-ocr/src/route.rs @@ -43,7 +43,10 @@ fn caller_token_provider(services: HostServices) -> TokenProviderHandle { .call(OcrOp::AcquireAzureAdToken) .await .map_err(|error| { - litellm_auth::Error::CredentialAcquisition(error.to_string().into()) + litellm_auth::Error::CredentialAcquisition(litellm_auth::ErrorDetail::failed( + "host Azure AD token request", + error, + )) }) } }) diff --git a/litellm-rust/crates/inference-ocr/src/types.rs b/litellm-rust/crates/inference-ocr/src/types.rs index 1bd70c06116..98c609a158f 100644 --- a/litellm-rust/crates/inference-ocr/src/types.rs +++ b/litellm-rust/crates/inference-ocr/src/types.rs @@ -67,6 +67,7 @@ impl OcrConnectionInputs { .map(|value| (name.clone(), value.to_string())) .ok_or_else(|| Error::RequestField { path: format!("extra_headers.{name}"), + cause: None, }) }) .collect() @@ -105,6 +106,7 @@ impl LiteLLMOcrRequest { .filter(|value| *value > 0 && *value <= default_transport.max_response_bytes) .ok_or_else(|| Error::RequestField { path: "max_response_bytes".into(), + cause: None, }) }) .transpose()? @@ -313,7 +315,7 @@ mod tests { }; assert!(matches!( error, - Error::RequestField { ref path } if path == "extra_headers.x-a" + Error::RequestField { ref path, .. } if path == "extra_headers.x-a" )); } } diff --git a/litellm-rust/crates/inference-ocr/src/wire.rs b/litellm-rust/crates/inference-ocr/src/wire.rs index a8f5a77a1cd..56d22be6523 100644 --- a/litellm-rust/crates/inference-ocr/src/wire.rs +++ b/litellm-rust/crates/inference-ocr/src/wire.rs @@ -73,9 +73,8 @@ pub fn decode_request_input>( let timeout = wire .timeout_seconds .map(|seconds| { - Duration::try_from_secs_f64(seconds).map_err(|_| Error::RequestField { - path: "timeout_seconds".into(), - }) + Duration::try_from_secs_f64(seconds) + .map_err(|error| Error::request_field("timeout_seconds", error)) }) .transpose()?; LiteLLMOcrRequest::from_inputs( @@ -123,8 +122,8 @@ mod tests { #[case::non_object(json!([]), "document")] #[case::missing_type(json!({"document_url":"https://example.com/a.pdf"}), "document")] #[case::unsupported_type(json!({"type":"text"}), "type")] - #[case::missing_document_url(json!({"type":"document_url"}), "Document URL")] - #[case::missing_image_url(json!({"type":"image_url"}), "Document URL")] + #[case::missing_document_url(json!({"type":"document_url"}), "document URL")] + #[case::missing_image_url(json!({"type":"image_url"}), "document URL")] fn ocr_contract_malformed_document_is_bad_request( #[case] document: Value, #[case] field: &str, @@ -191,16 +190,13 @@ mod tests { assert!(is_supported_request("deepseek-ocr", Some("vertex_ai"))); } - #[test] - fn missing_document_source_has_a_typed_public_error() { - for document in [ - serde_json::json!({"type": "document_url"}), - serde_json::json!({"type": "image_url"}), - ] { - assert!(matches!( - decode_document(document), - Err(Error::MissingDocumentUrl) - )); - } + #[rstest] + #[case::document_url(serde_json::json!({"type": "document_url"}))] + #[case::image_url(serde_json::json!({"type": "image_url"}))] + fn missing_document_source_has_a_typed_public_error(#[case] document: serde_json::Value) { + assert!(matches!( + decode_document(document), + Err(Error::MissingDocumentUrl) + )); } } diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/aws_textract.rs b/litellm-rust/crates/inference-ocr/tests/ocr/aws_textract.rs index c4096767188..f0e72b19fc1 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/aws_textract.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/aws_textract.rs @@ -191,12 +191,15 @@ async fn a_multi_page_rejection_reaches_the_caller_with_the_single_page_limit(#[ .await .unwrap_err(); - let Error::Provider { status, body, .. } = error else { - panic!("expected a provider error, got {error:?}"); - }; - assert_eq!(status, 400); + let message = litellm_error::Report::new(&error).to_string(); + assert!(error.is_request()); + assert_eq!(error.http_status_code(), Some(400)); assert!( - body.contains("multi-page documents are not supported"), - "{body}" + message.contains("Request has unsupported document format"), + "{message}" + ); + assert!( + message.contains("multi-page documents are not supported"), + "{message}" ); } diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/azure_ai.rs b/litellm-rust/crates/inference-ocr/tests/ocr/azure_ai.rs index 9bff7490d58..6f591af73e0 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/azure_ai.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/azure_ai.rs @@ -222,24 +222,23 @@ async fn credential_precedence( false, json!({}), numbered_token, - |error: &Error| matches!(error, Error::Auth(litellm_auth::Error::MissingApiBase { - provider: "Azure AI", - guidance: "Set AZURE_AI_API_BASE environment variable or pass api_base parameter", - })), + |error: &Error| litellm_error::Report::new(error).to_string() + == "Missing Azure AI API Base - Set AZURE_AI_API_BASE environment variable or pass api_base parameter", 0 )] #[case::unsupported_oidc_reference( true, json!({"azure_ad_token": "oidc/assertion", "client_id": "client", "tenant_id": "tenant"}), numbered_token, - |error: &Error| matches!(error, Error::Auth(litellm_auth::Error::InvalidConfiguration(detail)) if detail.to_string() == "unsupported OIDC reference"), + |error: &Error| litellm_error::Report::new(error).to_string().contains("unsupported OIDC reference"), 0 )] #[case::empty_provider_token_ignores_static_token( true, json!({"azure_ad_token": "static-token"}), |_| String::new(), - |error: &Error| matches!(error, Error::MissingAzureAiCredentials), + |error: &Error| error.is_missing_credential() + && litellm_error::Report::new(error).to_string().contains("Azure AI credentials"), 1 )] #[tokio::test] diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/azure_document_intelligence.rs b/litellm-rust/crates/inference-ocr/tests/ocr/azure_document_intelligence.rs index cec43e2ed0e..98875cc777c 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/azure_document_intelligence.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/azure_document_intelligence.rs @@ -5,7 +5,7 @@ use std::{ use litellm_host::lifecycle::CallEvent; use litellm_host::lifecycle::ExecutionEvent; -use litellm_llms::base_llm::ocr::settings::OcrSettings; +use litellm_llms::base_llm::ocr::{error::PagesFault, settings::OcrSettings}; use rstest::rstest; @@ -64,10 +64,10 @@ async fn pages_features_and_extra_options_map_to_the_analyze_call() { } #[rstest] -#[case(json!({"pages": [true]}), Error::Pages("expected only integers or only strings".into()))] -#[case(json!({"pages": [1, "2"]}), Error::Pages("expected only integers or only strings".into()))] -#[case(json!({"pages": [-1]}), Error::Pages("negative page index".into()))] -#[case(json!({"pages": "1&&features=bad"}), Error::Pages("invalid native page range".into()))] +#[case(json!({"pages": [true]}), Error::Pages(PagesFault::Mixed))] +#[case(json!({"pages": [1, "2"]}), Error::Pages(PagesFault::Mixed))] +#[case(json!({"pages": [-1]}), Error::Pages(PagesFault::Negative))] +#[case(json!({"pages": "1&&features=bad"}), Error::Pages(PagesFault::Range))] #[case(json!({"features": "languages&pages=1"}), Error::Features)] #[case(json!({"req_format": "azure"}), Error::RequestFormat)] #[tokio::test] @@ -315,7 +315,11 @@ async fn polling_does_not_follow_redirects() { .await .unwrap_err(); - assert!(error.to_string().contains("status 302"), "{error}"); + assert_eq!( + error.upstream_response().map(|(status, ..)| status), + Some(302), + "{error}" + ); assert_eq!(received(&upstream).await.len(), 2); } @@ -389,12 +393,15 @@ async fn malformed_provider_pages_report_the_response_path( } #[rstest] -#[case::missing(None)] -#[case::relative(Some("/relative"))] -#[case::cross_origin(Some("http://example.com/operation"))] -#[case::with_userinfo(Some("http://user:password@127.0.0.1/operation"))] +#[case::missing(None, false)] +#[case::relative(Some("/relative"), true)] +#[case::cross_origin(Some("http://example.com/operation"), false)] +#[case::with_userinfo(Some("http://user:password@127.0.0.1/operation"), false)] #[tokio::test] -async fn an_unusable_operation_location_is_rejected(#[case] location: Option<&str>) { +async fn an_unusable_operation_location_is_rejected( + #[case] location: Option<&str>, + #[case] unparseable: bool, +) { let response = location .into_iter() .fold(ResponseTemplate::new(202), |response, location| { @@ -407,6 +414,11 @@ async fn an_unusable_operation_location_is_rejected(#[case] location: Option<&st .unwrap_err(); assert!(error.to_string().contains("operation-location"), "{error}"); + assert!(error.is_response()); + assert_eq!( + litellm_error::find::(&error).is_some(), + unparseable + ); assert_eq!(received(&upstream).await.len(), 1); } diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/cohere.rs b/litellm-rust/crates/inference-ocr/tests/ocr/cohere.rs index 007aa49a2fd..a5076cf0112 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/cohere.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/cohere.rs @@ -37,6 +37,15 @@ async fn a_non_image_document_is_rejected_before_sending( .await .unwrap_err(); - assert!(matches!(error, Error::CohereImageOnly), "{error:?}"); + assert!( + matches!( + error, + Error::UnsupportedDocument { + provider: "Cohere Parse", + .. + } + ), + "{error:?}" + ); assert!(received(&upstream).await.is_empty()); } diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/documents.rs b/litellm-rust/crates/inference-ocr/tests/ocr/documents.rs index 338a47c5f57..35de1326386 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/documents.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/documents.rs @@ -176,8 +176,8 @@ async fn a_missing_path_document_fails_before_sending() { assert!( matches!( &error, - Error::FileRead { path: failed, source } - if *failed == path && source.kind() == std::io::ErrorKind::NotFound + Error::FileRead { path: failed, cause } + if *failed == path && cause.kind() == std::io::ErrorKind::NotFound ), "{error:?}" ); @@ -216,7 +216,8 @@ async fn configured_client_preserves_document_url_policy(#[case] allowed: bool) .await; if !allowed { - assert!(matches!(result, Err(Error::BlockedDocumentUrl))); + let error = result.unwrap_err(); + assert!(error.is_request(), "{error:?}"); assert!(received(&documents).await.is_empty()); assert!(received(&upstream).await.is_empty()); return; diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs b/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs index e413781f1a3..75a91092daf 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs @@ -38,7 +38,7 @@ fn recording_host( .with_before_send(move |wire, _| { interceptions.fetch_add(1, Ordering::SeqCst); match block { - true => Err(Error::InvalidRequest("blocked".into())), + true => Err(Error::MissingField("blocked")), false => Ok(wire), } }) @@ -95,10 +95,7 @@ async fn a_blocking_before_send_prevents_the_call_and_emits_one_failure() { .await .unwrap_err(); - assert!( - matches!(&error, Error::InvalidRequest(message) if message == "blocked"), - "{error:?}" - ); + assert!(matches!(error, Error::MissingField("blocked")), "{error:?}"); assert_eq!(interceptions.load(Ordering::SeqCst), 1); assert_eq!(*events.lock().unwrap(), ["started", "failure"]); assert!(received(&upstream).await.is_empty()); diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/machine.rs b/litellm-rust/crates/inference-ocr/tests/ocr/machine.rs index 4f9266c4a8d..9a1c67fab0f 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/machine.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/machine.rs @@ -117,10 +117,14 @@ async fn a_hand_driven_machine_performs_the_same_call() { assert_eq!(outcome.unwrap().pages[0].markdown, "native"); assert_eq!(received(&upstream).await.len(), 1); assert_eq!(ops, ["BeforeSend", "response"]); - assert!(matches!( - machine.resume().await, - Err(Error::InvalidRequest(_)) - )); + let Err(error) = machine.resume().await else { + panic!("resuming a finished call is a host protocol fault"); + }; + assert!(!error.is_request(), "{error:?}"); + assert!( + litellm_error::find::(&error).is_some(), + "{error:?}" + ); } #[rstest::rstest] @@ -153,8 +157,11 @@ async fn a_path_document_is_read_by_core_without_a_host_operation() { } #[rstest] -#[case::failed(HostFailure::Error(Error::InvalidRequest("before_provider_request failed".into())), "before_provider_request failed")] -#[case::cancelled(HostFailure::Cancelled(Error::InvalidRequest("cancelled".into())), "cancelled")] +#[case::failed( + HostFailure::Error(Error::MissingField("before_provider_request failed")), + "before_provider_request failed" +)] +#[case::cancelled(HostFailure::Cancelled(Error::MissingField("cancelled")), "cancelled")] #[tokio::test] async fn a_before_send_failure_ends_the_call_without_reaching_transport( #[case] failure: HostFailure, @@ -174,7 +181,7 @@ async fn a_before_send_failure_ends_the_call_without_reaching_transport( .await; assert!( - matches!(&outcome, Err(Error::InvalidRequest(actual)) if actual == message), + matches!(&outcome, Err(Error::MissingField(actual)) if *actual == message), "{outcome:?}" ); assert_eq!(ops, ["BeforeSend"]); @@ -255,17 +262,17 @@ async fn interrupt_drops_provider_captures_before_returning() { drive_until_notified(&mut machine, &host, &entered).await; assert!(!dropped.load(Ordering::SeqCst)); - let acknowledgement = machine.interrupt(HostFailure::Cancelled(Error::InvalidRequest( - "cancelled".into(), - ))); + let acknowledgement = + machine.interrupt(HostFailure::Cancelled(Error::MissingField("cancelled"))); assert!( dropped.load(Ordering::SeqCst), "interrupt returned while provider captures were still alive" ); - assert!( - matches!(acknowledgement.await, Err(Error::InvalidRequest(message)) if message == "cancelled") - ); + assert!(matches!( + acknowledgement.await, + Err(Error::MissingField("cancelled")) + )); } #[rstest::rstest] @@ -290,7 +297,7 @@ async fn interrupting_an_in_flight_provider_request_closes_its_connection() { let mut machine = ocr_route().machine(host.request().unwrap(), None); drive_until_notified(&mut machine, &host, &received).await; - let cancelled = Error::InvalidRequest("cancelled".into()); + let cancelled = Error::MissingField("cancelled"); assert!( machine diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/main.rs b/litellm-rust/crates/inference-ocr/tests/ocr/main.rs index 5a905f907d2..5a48bba9e1a 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/main.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/main.rs @@ -209,7 +209,7 @@ impl LocalOcrHost { request, caller_token: false, }) - .ok_or_else(|| Error::InvalidRequest("OCR request was already projected".into())) + .ok_or(Error::MissingField("OCR request")) } pub fn runtime(&self) -> litellm_host_native::in_process::Host<'_, Self, Self, ()> { litellm_host_native::in_process::Host { @@ -223,11 +223,11 @@ impl LocalOcrHost { impl litellm_host_native::services::HostCallHandler for LocalOcrHost { async fn handle_host_call(&self, op: OcrOp) -> Result<(), Error> { match op { - OcrOp::AcquireAzureAdToken(_) => { - Err(Error::Auth(litellm_auth::Error::CredentialAcquisition( - "OCR host has no Azure AD token provider".into(), - ))) - } + OcrOp::AcquireAzureAdToken(_) => Err(Error::Auth( + litellm_auth::Error::CredentialAcquisition(litellm_auth::ErrorDetail::Unresolved { + subject: "OCR host Azure AD token provider", + }), + )), } } } diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/mistral.rs b/litellm-rust/crates/inference-ocr/tests/ocr/mistral.rs index 528adc7feaf..21cf4ce455d 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/mistral.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/mistral.rs @@ -92,28 +92,13 @@ async fn an_upstream_error_keeps_its_status_whole_body_and_headers( .unwrap_err(); assert_eq!(received(&upstream).await.len(), 1); - let Error::Provider { - status, - body, - headers, - } = error - else { - panic!("expected provider error, got {error:?}"); - }; + let (status, body, headers) = error.upstream_response().expect("the provider answered"); assert_eq!(status, 422); - for (name, value) in [ - ("retry-after", "17"), - ("x-request-id", "request-123"), - ("x-future-header", "retained"), - ] { - assert!( - headers - .iter() - .any(|(key, actual)| key.eq_ignore_ascii_case(name) && actual == value), - "{name} missing from {headers:?}" - ); - } assert_eq!(body, expected_body); + let header = |name: &str| headers.get(name).and_then(|value| value.to_str().ok()); + assert_eq!(header("retry-after"), Some("17")); + assert_eq!(header("x-request-id"), Some("request-123")); + assert_eq!(header("x-future-header"), Some("retained")); } #[rstest] diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/reducto.rs b/litellm-rust/crates/inference-ocr/tests/ocr/reducto.rs index 649dca3d690..f6f458a4916 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/reducto.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/reducto.rs @@ -175,8 +175,11 @@ async fn a_failed_upload_stops_before_parse(#[case] upload: ResponseTemplate) { } #[rstest] -#[case::remote_url("https://example.com/a.pdf", Error::ReductoSource)] -#[case::empty_file_id("reducto://", Error::RequestField { path: "document file id".into() })] +#[case::remote_url( + "https://example.com/a.pdf", + Error::UnsupportedDocument { provider: "Reducto", guidance: "" } +)] +#[case::empty_file_id("reducto://", Error::RequestField { path: "document file id".into(), cause: None })] #[case::data_uri_without_payload("data:application/pdf;base64", Error::InvalidDataUri)] #[case::invalid_base64("data:application/pdf;base64,INVALID!", Error::InvalidDataUri)] #[tokio::test] @@ -202,7 +205,6 @@ async fn invalid_document_sources_are_rejected_before_sending( std::mem::discriminant(&expected) ); assert_eq!(error.http_status_code(), Some(400)); - assert_eq!(error.to_string(), expected.to_string()); } #[tokio::test] diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/vertex_ai.rs b/litellm-rust/crates/inference-ocr/tests/ocr/vertex_ai.rs index 81b52d076a5..dba12de60f6 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/vertex_ai.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/vertex_ai.rs @@ -71,7 +71,8 @@ async fn invalid_credentials_fail_before_sending() { .await .unwrap_err(); - assert!(error.to_string().contains("vertex_credentials"), "{error}"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains("vertex_credentials"), "{report}"); } #[rstest] @@ -91,11 +92,10 @@ async fn a_request_controlled_api_base_is_rejected_before_vertex_auth( let error = perform(request).await.unwrap_err(); + let report = litellm_error::Report::new(&error).to_string(); assert!( - error - .to_string() - .contains("request-controlled Vertex AI endpoint"), - "{error}" + report.contains("request-controlled Vertex AI endpoint"), + "{report}" ); } diff --git a/litellm-rust/crates/inference-responses/Cargo.toml b/litellm-rust/crates/inference-responses/Cargo.toml index 3dcd9976e0a..083179c50c8 100644 --- a/litellm-rust/crates/inference-responses/Cargo.toml +++ b/litellm-rust/crates/inference-responses/Cargo.toml @@ -28,6 +28,7 @@ tracing.workspace = true litellm-http = { workspace = true, features = ["test-support"] } litellm-cache.workspace = true litellm-cache-memory.workspace = true +litellm-error.workspace = true litellm-host-native.workspace = true litellm-inference-testing.workspace = true litellm-tracing.workspace = true diff --git a/litellm-rust/crates/inference-responses/src/handler.rs b/litellm-rust/crates/inference-responses/src/handler.rs index cbbc8be139d..8d3d1075e8e 100644 --- a/litellm-rust/crates/inference-responses/src/handler.rs +++ b/litellm-rust/crates/inference-responses/src/handler.rs @@ -72,7 +72,14 @@ pub(super) async fn execute( let stream = match wire.body.get("stream") { None => false, Some(serde_json::Value::Bool(value)) => *value, - Some(_) => return Err(Error::InvalidRequest("stream must be a boolean".into())), + Some(_) => { + return Err(Error::invalid_request( + litellm_llms::ErrorDetail::Expected { + subject: "stream", + expected: "a boolean", + }, + )); + } }; let outbound = litellm_inference::outbound::outbound_request( Authenticated { @@ -83,17 +90,11 @@ pub(super) async fn execute( &wire.body, Some(call.timeout.unwrap_or(Duration::from_secs(600))), )?; - let response = litellm_inference::outbound::send(outbound, &route.http) - .await - .map_err(network)?; + let response = litellm_inference::outbound::send(outbound, &route.http).await?; let status = response.status().as_u16(); if !response.status().is_success() { - let body = response.text().await.map_err(network)?; - return Err(litellm_http::transport::Error::Http { - status, - body: litellm_http::request::truncate_error_body(&body), - } - .into()); + let body = response.text().await.map_err(litellm_http::Error::from)?; + return Err(litellm_http::Error::upstream(status, body).into()); } if stream { let headers = response @@ -105,14 +106,14 @@ pub(super) async fn execute( .collect(); let chunks = response .bytes_stream() - .map(|chunk| chunk.map_err(network)) + .map(|chunk| chunk.map_err(|error| litellm_http::Error::from(error).into())) .boxed(); return Ok(ResponsesOutput::Stream { head: ResponsesStreamHead { headers }, chunks, }); } - let body = response.text().await.map_err(network)?; + let body = response.text().await.map_err(litellm_http::Error::from)?; let raw = RawResponse { body: body.clone() }; if let Some(observers) = observers { observers.emit(litellm_host::lifecycle::CallEvent::Execution( @@ -123,8 +124,12 @@ pub(super) async fn execute( .after_provider_response(raw) .await .map_err(Error::post_call)?; - let value = serde_json::from_str(&body) - .map_err(|error| Error::InvalidResponse(error.to_string().into()))?; + let value = serde_json::from_str(&body).map_err(|error| { + Error::invalid_response(litellm_llms::ErrorDetail::invalid( + "Responses API response JSON", + error, + )) + })?; config .transform_response_api_response(value) .map(ResponsesOutput::Complete) @@ -133,7 +138,3 @@ pub(super) async fn execute( ) .await } - -fn network(error: reqwest::Error) -> Error { - litellm_http::transport::Error::Network(error.to_string()).into() -} diff --git a/litellm-rust/crates/inference-responses/src/websocket.rs b/litellm-rust/crates/inference-responses/src/websocket.rs index c2ca89ecbaf..2afcde0a7a4 100644 --- a/litellm-rust/crates/inference-responses/src/websocket.rs +++ b/litellm-rust/crates/inference-responses/src/websocket.rs @@ -11,6 +11,8 @@ use tokio_tungstenite::tungstenite::{ http::{HeaderMap, HeaderName, HeaderValue}, }; +use litellm_llms::ErrorDetail; + use super::Error; #[derive(Clone)] @@ -36,23 +38,22 @@ impl ResponsesWebSocketConnection { .iter() .map(|(name, value)| { Ok(( - name.parse::() - .map_err(|error| Error::InvalidRequest(error.to_string().into()))?, - HeaderValue::from_str(value) - .map_err(|error| Error::InvalidRequest(error.to_string().into()))?, + name.parse::().map_err(|error| { + Error::invalid_request(ErrorDetail::invalid( + "WebSocket header name", + error, + )) + })?, + HeaderValue::from_str(value).map_err(|error| { + Error::invalid_request(ErrorDetail::invalid( + "WebSocket header value", + error, + )) + })?, )) }) .collect::>()?; - let connect = websocket::connect(client, url, headers); - let socket = match timeout { - Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| { - Error::Transport(litellm_http::transport::Error::Network( - "Responses WebSocket connection timed out".into(), - )) - })?, - None => connect.await, - } - .map_err(Error::Transport)?; + let socket = websocket::connect(client, url, headers, timeout).await?; Ok(Self { socket: Arc::new(Mutex::new(Some(socket))), }) @@ -70,13 +71,12 @@ impl ResponsesWebSocketConnection { litellm_inference::diagnostic::operation("litellm.websocket.send_text", async { let mut socket = self.socket.lock().await; let Some(socket) = socket.as_mut() else { - return Err(Error::Transport(litellm_http::transport::Error::Network( - "Responses WebSocket is closed".into(), - ))); + return Err(litellm_http::Error::websocket_closed().into()); }; - socket.send(Message::Text(text)).await.map_err(|error| { - Error::Transport(litellm_http::transport::Error::Network(error.to_string())) - }) + socket + .send(Message::Text(text)) + .await + .map_err(|error| litellm_http::Error::from(error).into()) }) .await } @@ -97,12 +97,12 @@ impl ResponsesWebSocketConnection { Some(Ok(Message::Text(text))) => Ok(Some(text)), Some(Ok(Message::Binary(bytes))) => String::from_utf8(bytes.to_vec()) .map(Some) - .map_err(|error| Error::InvalidResponse(error.to_string().into())), + .map_err(|error| { + Error::invalid_response(ErrorDetail::invalid("WebSocket text frame", error)) + }), Some(Ok(Message::Close(_))) | None => Ok(None), Some(Ok(_)) => Ok(None), - Some(Err(error)) => Err(Error::Transport(litellm_http::transport::Error::Network( - error.to_string(), - ))), + Some(Err(error)) => Err(litellm_http::Error::from(error).into()), } }) .await @@ -118,9 +118,10 @@ impl ResponsesWebSocketConnection { litellm_inference::diagnostic::operation("litellm.websocket.close", async { let mut socket = self.socket.lock().await; if let Some(socket) = socket.as_mut() { - socket.close(None).await.map_err(|error| { - Error::Transport(litellm_http::transport::Error::Network(error.to_string())) - })?; + socket + .close(None) + .await + .map_err(litellm_http::Error::from)?; } *socket = None; Ok(()) diff --git a/litellm-rust/crates/inference-responses/tests/openrouter.rs b/litellm-rust/crates/inference-responses/tests/openrouter.rs index e52cdba8ce8..8cdb68c376f 100644 --- a/litellm-rust/crates/inference-responses/tests/openrouter.rs +++ b/litellm-rust/crates/inference-responses/tests/openrouter.rs @@ -1,10 +1,7 @@ use std::sync::Arc; use futures_util::TryStreamExt; -use litellm_inference_responses::{ - Error, - types::{ResponsesCall, ResponsesOutput}, -}; +use litellm_inference_responses::types::{ResponsesCall, ResponsesOutput}; use litellm_inference_testing::{RecordingSecrets, no_secrets}; use rstest::{fixture, rstest}; use serde_json::json; @@ -108,13 +105,12 @@ async fn a_missing_key_fails_before_transport(call: ResponsesCall) { None, ) .await; - assert!(matches!( - result, - Err(Error::Auth(litellm_auth::Error::MissingApiKey { - provider: "OpenRouter", - .. - })) - )); + let Err(error) = result else { + panic!("the call fails"); + }; + assert!(error.is_missing_credential(), "{error:?}"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains("OpenRouter"), "{report}"); assert!(received(&upstream).await.is_empty()); } diff --git a/litellm-rust/crates/inference-responses/tests/responses.rs b/litellm-rust/crates/inference-responses/tests/responses.rs index 6d39d55d33c..66402e595eb 100644 --- a/litellm-rust/crates/inference-responses/tests/responses.rs +++ b/litellm-rust/crates/inference-responses/tests/responses.rs @@ -304,6 +304,25 @@ async fn provider_failures_emit_failure_once( ); } +#[rstest] +#[tokio::test] +async fn a_transport_error_does_not_echo_the_provider_url(call: ResponsesCall) { + let host = RecordingCall::::new(ResponsesCall { + api_base: Some(format!("{UNREACHABLE_BASE}/tenant-token-1234")), + ..call + }); + let call = host.request.lock().unwrap().take().unwrap(); + + let Err(error) = responses_route(no_secrets()) + .execute(call, &host, None) + .await + else { + panic!("nothing is listening"); + }; + + assert!(!error.to_string().contains("tenant-token-1234"), "{error}"); +} + #[rstest] #[case::explicit(true)] #[case::from_secrets(false)] @@ -557,12 +576,11 @@ async fn websocket_operations_trace_outcomes_without_capturing_frames_or_credent Some("private-frame-sentinel") ); connection.close().await.unwrap(); - assert!( - connection - .send_text("private-frame-sentinel".into()) - .await - .is_err() - ); + let error = connection + .send_text("private-frame-sentinel".into()) + .await + .expect_err("a closed connection sends nothing"); + assert!(!error.dispatched(), "{error:?}"); }) .await; server.await.unwrap(); diff --git a/litellm-rust/crates/inference-responses/tests/support/mod.rs b/litellm-rust/crates/inference-responses/tests/support/mod.rs index 81f133a0471..984b878eb21 100644 --- a/litellm-rust/crates/inference-responses/tests/support/mod.rs +++ b/litellm-rust/crates/inference-responses/tests/support/mod.rs @@ -13,6 +13,9 @@ use litellm_secrets::source::SecretSource; use serde_json::Value; use wiremock::{Mock, MockServer, Request, ResponseTemplate, matchers::any}; +/// A port nothing listens on, for calls that must fail before any request is sent. +pub const UNREACHABLE_BASE: &str = "http://127.0.0.1:1"; + pub fn responses_route( secrets: Arc, ) -> litellm_inference_responses::ResponsesRoute { diff --git a/litellm-rust/crates/inference-transcription/Cargo.toml b/litellm-rust/crates/inference-transcription/Cargo.toml index c501d329852..cf50b3f84fd 100644 --- a/litellm-rust/crates/inference-transcription/Cargo.toml +++ b/litellm-rust/crates/inference-transcription/Cargo.toml @@ -17,6 +17,7 @@ serde_json.workspace = true tracing.workspace = true [dev-dependencies] +litellm-error.workspace = true litellm-inference-testing.workspace = true litellm-tracing.workspace = true rstest.workspace = true diff --git a/litellm-rust/crates/inference-transcription/src/handler.rs b/litellm-rust/crates/inference-transcription/src/handler.rs index 6541ae8263e..ba7eb740e2f 100644 --- a/litellm-rust/crates/inference-transcription/src/handler.rs +++ b/litellm-rust/crates/inference-transcription/src/handler.rs @@ -2,7 +2,7 @@ use std::time::Duration; use litellm_http::{ Client, - request::{string_headers, truncate_error_body, with_default_headers}, + request::{string_headers, with_default_headers}, }; use litellm_llms::base_llm::{ audio_transcription::transformation::BaseAudioTranscriptionConfig, @@ -55,23 +55,14 @@ pub(super) async fn execute( .unwrap_or(Duration::from_secs(AUDIO_TRANSCRIPTION_TIMEOUT_SECS)), ), )?; - let response = litellm_inference::outbound::send(outbound, http) - .await - .map_err(|error| { - Error::Transport(litellm_http::transport::Error::Network(error.to_string())) - })?; + let response = litellm_inference::outbound::send(outbound, http).await?; let status = response.status(); - let text = response.text().await.map_err(|error| { - Error::Transport(litellm_http::transport::Error::Network(error.to_string())) - })?; + let text = response.text().await.map_err(litellm_http::Error::from)?; if !status.is_success() { - return Err(Error::Transport(litellm_http::transport::Error::Http { - status: status.as_u16(), - body: truncate_error_body(&text), - })); + return Err(litellm_http::Error::upstream(status.as_u16(), text).into()); } let response_json = serde_json::from_str(&text).map_err(|error| { - Error::InvalidResponse(litellm_llms::ErrorDetail::invalid( + Error::invalid_response(litellm_llms::ErrorDetail::invalid( "audio response JSON", error, )) diff --git a/litellm-rust/crates/inference-transcription/tests/audio_transcription.rs b/litellm-rust/crates/inference-transcription/tests/audio_transcription.rs index eb252bdff7e..9e48edce0c0 100644 --- a/litellm-rust/crates/inference-transcription/tests/audio_transcription.rs +++ b/litellm-rust/crates/inference-transcription/tests/audio_transcription.rs @@ -152,13 +152,7 @@ async fn invalid_audio_is_rejected_before_sending( .await .expect_err("invalid audio is rejected"); - assert!( - matches!( - error, - Error::InvalidRequest(_) | Error::MissingField(_) | Error::InvalidType { .. } - ), - "{error:?}" - ); + assert!(error.is_request(), "{error:?}"); assert!(received(&upstream).await.is_empty()); } @@ -185,7 +179,9 @@ async fn unsupported_providers_are_rejected_before_sending( .await .expect_err("unsupported provider errors"); - assert_eq!(error, Error::InvalidProvider(reported.into())); + assert!(error.is_request(), "{error:?}"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains(reported), "{report}"); } #[rstest] @@ -222,14 +218,28 @@ async fn an_upstream_error_keeps_its_status_and_body( .expect_err("upstream error propagates"); assert_eq!( - error, - Error::Transport(litellm_http::transport::Error::Http { - status, - body: "upstream said no".into() - }) + error.upstream_response(), + Some((status, "upstream said no")) ); } +#[rstest] +#[tokio::test] +async fn a_transport_error_does_not_echo_the_provider_url( + request: AudioTranscriptionRequest<'static>, +) { + let base = format!("{UNREACHABLE_BASE}/tenant-token-1234"); + + let error = transcribe(AudioTranscriptionRequest { + api_base: Some(&base), + ..request + }) + .await + .expect_err("nothing is listening"); + + assert!(!error.to_string().contains("tenant-token-1234"), "{error}"); +} + #[rstest] #[case::not_json(ResponseTemplate::new(200).set_body_string("not json"))] #[case::no_output(json_response(json!({"unexpected": true})))] @@ -248,7 +258,7 @@ async fn an_unreadable_success_body_is_an_invalid_response( .await .expect_err("an unreadable body fails"); - assert!(matches!(error, Error::InvalidResponse(_)), "{error:?}"); + assert!(error.is_response(), "{error:?}"); } #[rstest] diff --git a/litellm-rust/crates/inference/AGENTS.md b/litellm-rust/crates/inference/AGENTS.md index 7f2bd0387be..8cc8fc8c26b 100644 --- a/litellm-rust/crates/inference/AGENTS.md +++ b/litellm-rust/crates/inference/AGENTS.md @@ -53,10 +53,10 @@ ## Error placement -- The workspace `Error definitions` rules shape each crate's error. This section decides which crate and module a failure belongs to +- The [rust-error skill](../../.agents/skills/rust-error/SKILL.md) shapes each crate's error. This section decides which crate and module a failure belongs to - A failure is declared once, by the lowest crate that raises it - - every crate above nests it unchanged (`#[error(transparent)] Auth(#[from] litellm_auth::Error)`) or maps it once at its boundary, as `src/error.rs` does for `litellm_llms::Error` - - `Error` collects route failures for every format crate and never re-declares a variant a lower crate raises + - every crate above nests it in one variant and delegates its predicates, as `Error::Llm`, `Error::Auth` and `Error::Http` do, or translates it once in the shared `From` when its meaning changes, as `From` does for `Unsupported` and `From` does for a header the signer computes + - `Error` collects route failures for every format crate and never re-declares a variant a lower crate raises. Routes build transformation failures with `Error::invalid_request` and `Error::invalid_response`, and callers ask `is_unsupported()`, `is_request()`, `is_response()`, `is_missing_credential()`, `dispatched()`, `upstream_response()` and `rejected_request()` instead of matching nested variants - Scope follows the concept, not the first caller - an error type under `litellm-llms`'s `/` is private to that provider: no other provider and nothing in `base_llm` may import it - a failure two providers or two routes can hit (wire framing, stream event decoding, a malformed provider response) belongs to the crate that owns the concept: `litellm-framer` for framing, `litellm_llms::Error` for the transformation layer diff --git a/litellm-rust/crates/inference/Cargo.toml b/litellm-rust/crates/inference/Cargo.toml index acc211f214b..807a3ea85d3 100644 --- a/litellm-rust/crates/inference/Cargo.toml +++ b/litellm-rust/crates/inference/Cargo.toml @@ -29,6 +29,7 @@ time.workspace = true url.workspace = true [dev-dependencies] +litellm-error.workspace = true litellm-cache-memory.workspace = true litellm-http = { workspace = true, features = ["test-support"] } litellm-llms = { workspace = true, features = ["test-support"] } diff --git a/litellm-rust/crates/inference/src/error.rs b/litellm-rust/crates/inference/src/error.rs index 895c8c55e01..2b42d2fced1 100644 --- a/litellm-rust/crates/inference/src/error.rs +++ b/litellm-rust/crates/inference/src/error.rs @@ -1,57 +1,47 @@ -//! One error for every route in this crate. OCR still carries its own, richer enum. -//! -//! A variant is declared by the layer that produces it and nested here as is: -//! credentials by `litellm_auth` (AWS folds into it at that crate's boundary), the wire by -//! `litellm_http`, secrets by `litellm_secrets`. The transformation layer's [`LlmError`] -//! maps onto the same-named variants once, here, so no route re-declares them. - use std::sync::Arc; -use litellm_http::transport::Error as TransportError; -use litellm_llms::Error as LlmError; +use litellm_host::machine::MachineFault; +use litellm_llms::{Error as LlmError, ErrorDetail}; #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub enum Error { - #[error("expected {expected}, got {actual}")] - InvalidType { - expected: &'static str, - actual: &'static str, - }, - #[error("missing required field: {0}")] - MissingField(&'static str), #[error("invalid provider: {0}")] InvalidProvider(String), - #[error("invalid request: {0}")] - InvalidRequest(#[source] litellm_llms::ErrorDetail), - #[error("invalid response: {0}")] - InvalidResponse(#[source] litellm_llms::ErrorDetail), + #[error("invalid provider: unable to resolve custom_llm_provider for {route} request")] + UnresolvedProvider { route: &'static str }, #[error("unsupported by the rust path: {0}")] Unsupported(&'static str), #[error(transparent)] - Auth(#[from] litellm_auth::Error), + Llm(LlmError), #[error(transparent)] - Transport(#[from] TransportError), + Auth(#[from] litellm_auth::Error), #[error(transparent)] Headers(#[from] litellm_http::request::HeaderError), #[error(transparent)] - Http(#[from] litellm_http::Error), + Http(litellm_http::Error), #[error(transparent)] Secret(#[from] SecretError), - #[error("post-call hook failed: {0}")] + #[error("the host driver failed")] + Host(#[source] MachineFault), + #[error("post-call hook failed")] PostCallHook(#[source] Arc), } -impl From for Error { - fn from(fault: litellm_host::machine::MachineFault) -> Self { - use litellm_host::machine::MachineFault; - Self::InvalidRequest(match fault { - MachineFault::Abandoned => "host driver was abandoned".into(), - MachineFault::Protocol(message) => format!("host {message}").into(), - }) +impl From for Error { + fn from(fault: MachineFault) -> Self { + Self::Host(fault) } } impl Error { + pub fn invalid_request(detail: ErrorDetail) -> Self { + Self::Llm(LlmError::InvalidRequest(detail)) + } + + pub fn invalid_response(detail: ErrorDetail) -> Self { + Self::Llm(LlmError::InvalidResponse(detail)) + } + pub fn post_call(error: Self) -> Self { Self::PostCallHook(Arc::new(error)) } @@ -60,31 +50,128 @@ impl Error { /// the provider's answer. pub fn is_request(&self) -> bool { match self { - Self::InvalidType { .. } - | Self::MissingField(_) - | Self::InvalidProvider(_) - | Self::InvalidRequest(_) + Self::InvalidProvider(_) + | Self::UnresolvedProvider { .. } | Self::Unsupported(_) | Self::Headers(_) => true, - Self::Auth(error) => !matches!(error, litellm_auth::Error::MissingApiKey { .. }), - Self::InvalidResponse(_) - | Self::Transport(_) - | Self::Http(_) + Self::Llm(error) => error.is_request(), + Self::Auth(error) => error.is_request(), + Self::Http(error) => error.is_request(), + Self::Secret(_) | Self::Host(_) | Self::PostCallHook(_) => false, + } + } + + pub fn is_response(&self) -> bool { + match self { + Self::Llm(error) => error.is_response(), + Self::Http(error) => error.is_response(), + Self::InvalidProvider(_) + | Self::UnresolvedProvider { .. } + | Self::Unsupported(_) + | Self::Auth(_) + | Self::Headers(_) | Self::Secret(_) + | Self::Host(_) | Self::PostCallHook(_) => false, } } + + pub fn is_missing_credential(&self) -> bool { + match self { + Self::Llm(error) => error.is_missing_credential(), + Self::Auth(error) => error.is_missing_credential(), + Self::InvalidProvider(_) + | Self::UnresolvedProvider { .. } + | Self::Unsupported(_) + | Self::Headers(_) + | Self::Http(_) + | Self::Secret(_) + | Self::Host(_) + | Self::PostCallHook(_) => false, + } + } + + pub fn is_unsupported(&self) -> bool { + match self { + Self::Unsupported(_) => true, + Self::InvalidProvider(_) + | Self::UnresolvedProvider { .. } + | Self::Llm(_) + | Self::Auth(_) + | Self::Headers(_) + | Self::Http(_) + | Self::Secret(_) + | Self::Host(_) + | Self::PostCallHook(_) => false, + } + } + + pub fn dispatched(&self) -> bool { + match self { + Self::Llm(error) => error.dispatched(), + Self::Http(error) => error.dispatched(), + Self::Host(_) | Self::PostCallHook(_) => true, + Self::InvalidProvider(_) + | Self::UnresolvedProvider { .. } + | Self::Unsupported(_) + | Self::Auth(_) + | Self::Headers(_) + | Self::Secret(_) => false, + } + } + + pub fn rejected_request(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + Self::Llm(error) => error.rejected_request(), + Self::InvalidProvider(_) + | Self::UnresolvedProvider { .. } + | Self::Unsupported(_) + | Self::Auth(_) + | Self::Headers(_) + | Self::Http(_) + | Self::Secret(_) + | Self::Host(_) + | Self::PostCallHook(_) => None, + } + } + + pub fn upstream_response(&self) -> Option<(u16, &str)> { + match self { + Self::Llm(error) => error.upstream_response(), + Self::Http(error) => error.upstream_response(), + Self::InvalidProvider(_) + | Self::UnresolvedProvider { .. } + | Self::Unsupported(_) + | Self::Auth(_) + | Self::Headers(_) + | Self::Secret(_) + | Self::Host(_) + | Self::PostCallHook(_) => None, + } + } +} + +impl From for Error { + fn from(error: litellm_http::Error) -> Self { + // Python drops the caller's copy and prefers a forwarded Authorization + // over the signature, so leave the request to it. + if error.is_computed_header() { + return Self::Unsupported("request forwards a header AWS SigV4 computes"); + } + Self::Http(error) + } } impl From for Error { fn from(error: LlmError) -> Self { match error { - LlmError::InvalidType { expected, actual } => Self::InvalidType { expected, actual }, - LlmError::MissingField(field) => Self::MissingField(field), - LlmError::InvalidRequest(message) => Self::InvalidRequest(message), - LlmError::InvalidResponse(message) => Self::InvalidResponse(message), LlmError::Unsupported(reason) => Self::Unsupported(reason), + LlmError::Http(error) => Self::from(error), LlmError::Auth(error) => Self::Auth(error), + error @ (LlmError::InvalidType { .. } + | LlmError::MissingField(_) + | LlmError::InvalidRequest(_) + | LlmError::InvalidResponse(_)) => Self::Llm(error), } } } @@ -116,26 +203,151 @@ impl Eq for SecretError {} #[cfg(test)] mod tests { use super::Error; + use litellm_host::machine::MachineFault; use litellm_llms::{Error as LlmError, ErrorDetail}; use rstest::rstest; - #[test] - fn a_missing_api_key_is_the_environment_not_the_request() { - assert!( - !Error::Auth(litellm_auth::Error::MissingApiKey { - provider: "Anthropic", - environment_variable: "ANTHROPIC_API_KEY", - }) - .is_request() - ); - assert!(Error::Auth(litellm_auth::Error::InvalidHeader).is_request()); - assert!(Error::InvalidRequest("top_k".into()).is_request()); - assert!(!Error::InvalidResponse("bad json".into()).is_request()); + #[rstest] + #[case::missing_api_key( + litellm_auth::Error::MissingApiKey { + provider: "Anthropic", + environment_variable: "ANTHROPIC_API_KEY", + } + .into(), + false + )] + #[case::invalid_auth_header(litellm_auth::Error::InvalidHeader.into(), true)] + #[case::transformation_missing_api_key( + LlmError::Auth(litellm_auth::Error::MissingApiKey { + provider: "Anthropic", + environment_variable: "ANTHROPIC_API_KEY", + }) + .into(), + false + )] + #[case::invalid_request( + LlmError::InvalidRequest(ErrorDetail::Expected { + subject: "top_k", + expected: "an integer", + }) + .into(), + true + )] + #[case::invalid_response( + LlmError::InvalidResponse(ErrorDetail::MissingField { + subject: "response", + field: "usage", + }) + .into(), + false + )] + #[case::host_driver(MachineFault::Abandoned.into(), false)] + #[case::post_call_hook_over_a_request_fault( + Error::post_call(LlmError::MissingField("usage").into()), + false + )] + fn only_the_callers_input_is_a_request_fault(#[case] error: Error, #[case] is_request: bool) { + assert_eq!(error.is_request(), is_request); } + + #[rstest] + #[case::own_auth(litellm_auth::Error::MissingApiKey { + provider: "Anthropic", + environment_variable: "ANTHROPIC_API_KEY", + } + .into())] + #[case::transformation_auth( + LlmError::Auth(litellm_auth::Error::MissingApiKey { + provider: "Anthropic", + environment_variable: "ANTHROPIC_API_KEY", + }) + .into() + )] + fn a_missing_api_key_is_a_missing_credential_wherever_it_is_raised(#[case] error: Error) { + assert!(error.is_missing_credential()); + } + + #[rstest] + #[case::direct(litellm_http::Error::computed_header("authorization").into())] + #[case::through_a_transformation( + LlmError::Http(litellm_http::Error::computed_header("authorization")).into() + )] + fn a_header_the_signer_computes_is_left_to_the_host_on_every_route(#[case] error: Error) { + assert!(error.is_unsupported(), "{error:?}"); + assert!(error.is_request()); + } + + #[test] + fn a_transformation_unsupported_stays_a_decline() { + assert!(Error::from(LlmError::Unsupported("test capability")).is_unsupported()); + } + + #[test] + fn an_upstream_failure_through_a_transformation_keeps_the_provider_answer() { + let error = Error::from(LlmError::Http(litellm_http::Error::upstream( + 429, + "slow down".into(), + ))); + assert_eq!(error.upstream_response(), Some((429, "slow down"))); + assert!(error.dispatched()); + } + + #[rstest] + #[case::invalid_request( + Error::invalid_request(ErrorDetail::Expected { + subject: "messages", + expected: "non-empty", + }), + Some("messages must be non-empty") + )] + #[case::missing_field( + LlmError::MissingField("max_tokens").into(), + Some("missing required field: max_tokens") + )] + #[case::invalid_response( + Error::invalid_response(ErrorDetail::MissingField { + subject: "response", + field: "usage", + }), + None + )] + #[case::provider_rejection(litellm_http::Error::upstream(400, "bad".into()).into(), None)] + #[case::unresolved_provider(Error::InvalidProvider("unknown".into()), None)] + fn a_rejected_request_names_only_what_the_caller_got_wrong( + #[case] error: Error, + #[case] expected: Option<&str>, + ) { + let rejection = error + .rejected_request() + .map(|rejection| litellm_error::Report::new(rejection).to_string()); + assert_eq!(rejection.as_deref(), expected); + } + + #[test] + fn a_post_call_hook_failure_is_not_the_providers_answer() { + let error = + Error::post_call(litellm_http::Error::upstream(500, "hook backend down".into()).into()); + assert_eq!(error.upstream_response(), None); + assert!(!error.is_request()); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.starts_with("post-call hook failed: "), "{report}"); + assert!(report.contains("hook backend down"), "{report}"); + } + + #[test] + fn a_host_driver_fault_keeps_the_fault_as_its_source() { + let error = Error::from(MachineFault::Abandoned); + assert_eq!( + litellm_error::find::(&error), + Some(&MachineFault::Abandoned) + ); + assert!(error.dispatched()); + } + #[rstest] #[case::request(true)] #[case::response(false)] - fn contextual_errors_preserve_sources_and_route_classification(#[case] request: bool) { + fn contextual_errors_keep_their_cause_and_route_classification(#[case] request: bool) { let source = serde_json::from_str::("{").unwrap_err(); let source_message = source.to_string(); let detail = ErrorDetail::invalid("test payload", source); @@ -145,16 +357,10 @@ mod tests { LlmError::InvalidResponse(detail) }); assert_eq!(error.is_request(), request); - let category = if request { "request" } else { "response" }; - assert_eq!( - error.to_string(), - format!("invalid {category}: invalid test payload: {source_message}") - ); - let source = std::iter::successors(Some(&error as &dyn std::error::Error), |error| { - error.source() - }) - .find_map(|error| error.downcast_ref::()) - .expect("the original JSON error remains available"); - assert_eq!(source.to_string(), source_message); + assert_eq!(error.is_response(), !request); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains("test payload"), "{report}"); + assert_eq!(report.matches(&source_message).count(), 1, "{report}"); + assert!(litellm_error::find::(&error).is_some()); } } diff --git a/litellm-rust/crates/inference/src/outbound.rs b/litellm-rust/crates/inference/src/outbound.rs index 9d9b1e606c2..b0cc71d425d 100644 --- a/litellm-rust/crates/inference/src/outbound.rs +++ b/litellm-rust/crates/inference/src/outbound.rs @@ -13,10 +13,10 @@ use serde_json::Value; pub async fn send( request: reqwest::Request, client: &litellm_http::Client, -) -> Result { - client.execute(request).await.inspect(|response| { - tracing::Span::current().record("status", response.status().as_u16()); - }) +) -> Result { + let response = client.execute(request).await?; + tracing::Span::current().record("status", response.status().as_u16()); + Ok(response) } /// Header credentials are already in `headers`; SigV4 is applied here, over the diff --git a/litellm-rust/crates/inference/src/provider.rs b/litellm-rust/crates/inference/src/provider.rs index 305aca85f97..fa226a4ef8d 100644 --- a/litellm-rust/crates/inference/src/provider.rs +++ b/litellm-rust/crates/inference/src/provider.rs @@ -25,11 +25,7 @@ pub fn resolve_llm_provider<'a>( custom_llm_provider: provider, }) }) - .ok_or_else(|| { - Error::InvalidProvider(format!( - "unable to resolve custom_llm_provider for {route} request" - )) - })?; + .ok_or(Error::UnresolvedProvider { route })?; let provider = custom_llm_provider .parse() .map_err(|_| Error::InvalidProvider(custom_llm_provider.to_string()))?; diff --git a/litellm-rust/crates/inference/tests/caching.rs b/litellm-rust/crates/inference/tests/caching.rs index a0a6a757565..02c2edd95dd 100644 --- a/litellm-rust/crates/inference/tests/caching.rs +++ b/litellm-rust/crates/inference/tests/caching.rs @@ -417,7 +417,7 @@ impl litellm_cache::BaseCache for UnavailableCache { _: &str, _: &Self::Context, ) -> Result, litellm_cache::Error> { - Err(litellm_cache::Error::Unavailable) + Err(litellm_cache::Error::Unavailable(None)) } fn set_cache( @@ -426,7 +426,7 @@ impl litellm_cache::BaseCache for UnavailableCache { _: Self::Value, _: &Self::Context, ) -> Result<(), litellm_cache::Error> { - Err(litellm_cache::Error::Unavailable) + Err(litellm_cache::Error::Unavailable(None)) } } diff --git a/litellm-rust/crates/llms/Cargo.toml b/litellm-rust/crates/llms/Cargo.toml index 2d3ea85d1fc..962f61ef4ec 100644 --- a/litellm-rust/crates/llms/Cargo.toml +++ b/litellm-rust/crates/llms/Cargo.toml @@ -11,6 +11,7 @@ test-support = ["litellm-http/test-support"] [dependencies] litellm-llms-types.workspace = true litellm-core-utils.workspace = true +litellm-error.workspace = true litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] } litellm-auth-aws.workspace = true litellm-auth-azure.workspace = true diff --git a/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs b/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs index b8ea701af99..66658802dd5 100644 --- a/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs @@ -149,7 +149,12 @@ impl AnthropicBatchesConfig for AnthropicBatchesTransformation { } let mut url = batches_base_url(api_base, env_lookup)?; url.path_segments_mut() - .map_err(|_| Error::InvalidRequest("Anthropic API base cannot be a base URL".into()))? + .map_err(|()| { + Error::InvalidRequest(crate::ErrorDetail::Expected { + subject: "Anthropic API base", + expected: "a base URL", + }) + })? .push(batch_id); Ok(url.into()) } @@ -217,11 +222,11 @@ impl AnthropicBatchesConfig for AnthropicBatchesTransformation { .map(|(index, line)| { let record: BatchResultRecord = serde_json::from_str(line.trim()).map_err(|error| { - Error::InvalidResponse(crate::ErrorDetail::InvalidLine { - subject: "Anthropic batch result", - line: index + 1, - source: crate::ErrorSource::new(error), - }) + Error::InvalidResponse(crate::ErrorDetail::invalid_line( + "Anthropic batch result", + index + 1, + error, + )) })?; match record.result { BatchResult::Succeeded { message } => Ok(*message), @@ -314,18 +319,22 @@ mod tests { assert_eq!(messages[0].id, "msg_1"); } - #[test] - fn reports_malformed_and_unsuccessful_batch_results() { - assert!(matches!( - ANTHROPIC_BATCHES_TRANSFORMATION.transform_batch_results("not-json"), - Err(Error::InvalidResponse(message)) if message.to_string().contains("line 1") - )); - assert!(matches!( - ANTHROPIC_BATCHES_TRANSFORMATION.transform_batch_results( - r#"{"result":{"type":"errored","error":{"type":"invalid_request_error"}}}"# - ), - Err(Error::InvalidResponse(message)) if message.to_string().contains("request failed") - )); + #[rstest::rstest] + #[case::malformed("not-json", "line 1")] + #[case::errored( + r#"{"result":{"type":"errored","error":{"type":"invalid_request_error"}}}"#, + "request failed" + )] + fn reports_malformed_and_unsuccessful_batch_results(#[case] body: &str, #[case] fact: &str) { + let error = ANTHROPIC_BATCHES_TRANSFORMATION + .transform_batch_results(body) + .unwrap_err(); + assert!(matches!(error, Error::InvalidResponse(_)), "{error:?}"); + assert!( + litellm_error::Report::new(&error) + .to_string() + .contains(fact) + ); } #[test] diff --git a/litellm-rust/crates/llms/src/anthropic/chat/handler.rs b/litellm-rust/crates/llms/src/anthropic/chat/handler.rs deleted file mode 100644 index eaee006f228..00000000000 --- a/litellm-rust/crates/llms/src/anthropic/chat/handler.rs +++ /dev/null @@ -1,163 +0,0 @@ -use std::collections::HashMap; - -use litellm_llms_types::formats::{ - chat_completions::{ - ChatCompletionChunk, ChatCompletionThinkingBlock, ChatCompletionToolCallChunk, - ChatCompletionsUsage, - }, - messages::streaming::{ - MessagesContentBlock, MessagesContentBlockDelta, MessagesStreamEvent, MessagesStreamUsage, - }, -}; -use serde_json::Value; - -use crate::{ - Error, - base_llm::{base_model_iterator::StreamTransformer, chat::streaming::StreamShape}, -}; - -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub enum AnthropicJsonChunkType { - ValidJson, - AccumulatedJson, -} - -#[derive(Clone, Debug, Eq, PartialEq)] -pub enum AnthropicContentBlockType { - Text, - ToolUse, - ServerToolUse, - Thinking, - RedactedThinking, - Compaction, - ToolResult(String), - Other(String), -} - -#[derive(Clone, Debug, PartialEq)] -pub struct AnthropicContentBlockDeltaEvent { - pub index: u64, - pub delta: MessagesContentBlockDelta, -} - -pub struct ModelResponseIterator { - pub content_blocks: Vec, - pub tool_index: i64, - pub json_mode: bool, - pub speed: Option, - pub tool_name_reverse_map: HashMap, - pub response_id: String, - pub served_model: Option, - pub is_response_format_tool: bool, - pub converted_response_format_tool: bool, - pub accumulated_json: String, - pub chunk_type: AnthropicJsonChunkType, - pub current_content_block_type: Option, - pub web_search_results: Vec, - pub web_search_calls: HashMap, - pub compaction_blocks: Vec, - pub reasoning_content_chunks: Vec, - pub server_tool_inputs: HashMap, - pub tool_results: Vec, - pub current_server_tool_id: Option, - pub container_id: Option, -} - -impl ModelResponseIterator { - pub fn new(_shape: StreamShape) -> Self { - todo!() - } - - pub fn check_empty_tool_call_args(&self) -> bool { - todo!() - } - - pub fn handle_usage(&mut self, _usage: MessagesStreamUsage) -> ChatCompletionsUsage { - todo!() - } - - pub fn handle_content_block_delta( - &mut self, - _index: u64, - _delta: MessagesContentBlockDelta, - ) -> ( - String, - Option, - Vec, - Option, - Option, - ) { - todo!() - } - - pub fn handle_content_block_start( - &mut self, - _index: u64, - _content_block: MessagesContentBlock, - ) -> Result { - todo!() - } - - pub fn handle_json_mode_chunk( - &mut self, - _text: String, - _tool_use: Option, - ) -> (String, Option) { - todo!() - } - - pub fn handle_accumulated_json_chunk( - &mut self, - _data: &str, - _is_final: bool, - ) -> Result, Error> { - todo!() - } - - pub fn handle_redacted_thinking_content( - &mut self, - _content_block: &MessagesContentBlock, - ) -> Vec { - todo!() - } - - pub fn web_search_call_snapshot(&self) -> HashMap { - todo!() - } - - pub fn complete_web_search_call(&mut self, _result: Value) { - todo!() - } - - pub fn build_code_interpreter_results(&self) -> Vec { - todo!() - } - - pub fn handle_message_delta( - &mut self, - _event: MessagesStreamEvent, - ) -> (Option, Option, Option) { - todo!() - } - - pub fn chunk_parser( - &mut self, - _event: MessagesStreamEvent, - ) -> Result { - todo!() - } -} - -impl StreamTransformer for ModelResponseIterator { - type Input = MessagesStreamEvent; - type Output = ChatCompletionChunk; - type Error = Error; - - fn transform(&mut self, _input: Self::Input) -> Result, Self::Error> { - todo!() - } - - fn finish(&mut self) -> Result, Self::Error> { - todo!() - } -} diff --git a/litellm-rust/crates/llms/src/anthropic/chat/mod.rs b/litellm-rust/crates/llms/src/anthropic/chat/mod.rs index f0050b7dc71..f239b6921fa 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/mod.rs +++ b/litellm-rust/crates/llms/src/anthropic/chat/mod.rs @@ -1,2 +1 @@ -pub mod handler; pub mod transformation; diff --git a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs index 91baa247ac8..4d4f110fbe1 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs @@ -14,23 +14,19 @@ use serde_json::{Map, Value, json}; use crate::{ Error, - anthropic::{ - chat::handler::ModelResponseIterator, - common_utils::{ - API_KEY_PLACEMENT, complete_anthropic_url, forwarded_oauth_bearer, - resolve_anthropic_api_key, - }, + anthropic::common_utils::{ + API_KEY_PLACEMENT, complete_anthropic_url, forwarded_oauth_bearer, + resolve_anthropic_api_key, }, base_llm::{ auth::AuthScheme, chat::{ - streaming::{ChatStream, StreamShape}, + streaming::StreamShape, transformation::{ BaseConfig, Headers, ProviderChatRequestData, ProviderChatResponseData, Unsupported, ValidatedEnvironment, unsupported_message, unsupported_param, }, }, - messages::streaming::event_stream as messages_event_stream, }, }; @@ -194,13 +190,6 @@ impl BaseConfig for AnthropicConfig { Ok(ValidatedEnvironment { headers, auth }) } - fn model_response_iterator(&self, shape: StreamShape) -> Option { - Some(ChatStream::new( - messages_event_stream, - ModelResponseIterator::new(shape), - )) - } - fn default_headers(&self) -> &'static [(&'static str, &'static str)] { DEFAULT_HEADERS } @@ -676,16 +665,11 @@ mod tests { #[case::missing_model(json!({"content": [], "usage": {"input_tokens": 1, "output_tokens": 1}}))] fn errors_on_a_response_missing_required_fields(#[case] body: Value) { let error = transform_response(body).expect_err("invalid response"); - assert!(matches!(error, Error::InvalidResponse(_))); - let source = std::iter::successors(Some(&error as &dyn std::error::Error), |error| { - error.source() - }) - .find_map(|error| error.downcast_ref::()) - .expect("the JSON decoding source is preserved"); - assert_eq!( - error.to_string(), - format!("invalid response: invalid messages response: {source}") - ); + assert!(matches!(error, Error::InvalidResponse(_)), "{error:?}"); + let cause = litellm_error::find::(&error).expect("JSON cause"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains("invalid messages response"), "{report}"); + assert_eq!(report.matches(&cause.to_string()).count(), 1, "{report}"); } #[test] diff --git a/litellm-rust/crates/llms/src/anthropic/messages/handler.rs b/litellm-rust/crates/llms/src/anthropic/messages/handler.rs index 6e8de65d6fa..54f924ddec6 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/handler.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/handler.rs @@ -208,24 +208,29 @@ mod tests { #[case::empty(json!({}), Ok(json!({})))] #[case::numeric_user_id( json!({"user_id": 123}), - Err(Error::InvalidRequest("metadata.user_id must be a string, got 123".to_string().into())), + Err(invalid_value("metadata.user_id", "a string", json!(123))), )] #[case::boolean_user_id( json!({"user_id": true}), - Err(Error::InvalidRequest("metadata.user_id must be a string, got true".to_string().into())), + Err(invalid_value("metadata.user_id", "a string", json!(true))), )] #[case::not_an_object( json!(["u-1"]), - Err(Error::InvalidRequest(r#"metadata must be an object, got ["u-1"]"#.to_string().into())), + Err(invalid_value("metadata", "an object", json!(["u-1"]))), )] fn validate_anthropic_api_metadata_passes_only_a_string_user_id( #[case] metadata: Value, #[case] expected: Result, ) { - assert_eq!( - validate_anthropic_api_metadata(&metadata).map_err(|error| error.to_string()), - expected.map_err(|error| error.to_string()), - ); + assert_eq!(validate_anthropic_api_metadata(&metadata), expected); + } + + fn invalid_value(field: &'static str, expected: &'static str, actual: Value) -> Error { + Error::InvalidRequest(crate::ErrorDetail::InvalidValue { + field, + expected, + actual, + }) } #[rstest] diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs index e813d4ac8cc..82119cecfa5 100644 --- a/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs @@ -123,7 +123,7 @@ impl BaseOcrConfig for TextractAnalyzeDocumentConfig { &self, error_message: String, status_code: u16, - headers: Vec<(String, String)>, + headers: reqwest::header::HeaderMap, ) -> Error { error_class(error_message, status_code, headers) } diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs index bc199f40938..e3a4f947433 100644 --- a/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs @@ -17,6 +17,7 @@ const AWS_JSON_CONTENT_TYPE: &str = "application/x-amz-json-1.1"; const TARGET_HEADER: &str = "X-Amz-Target"; const CONTENT_TYPE_HEADER: &str = "Content-Type"; const UNSUPPORTED_DOCUMENT: &str = "UnsupportedDocumentException"; +const SYNC_API_GUIDANCE: &str = "aws_textract uses Textract's synchronous API, which reads a JPEG, PNG, or a single-page PDF or TIFF; other formats and multi-page documents are not supported"; const SYNC_DOCUMENT_MAX_BYTES: usize = 10 * 1024 * 1024; const HEALTH_CHECK_IMAGE_DATA_URI: &str = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4//8/AAX+Av4N70a4AAAAAElFTkSuQmCC"; @@ -321,20 +322,18 @@ struct AwsError { /// Textract answers both an unsupported format and a multi-page PDF or TIFF /// with a bare "unsupported document format", which reads like a corrupt file. /// Say what the synchronous API accepts. -pub(super) fn error_class(body: String, status: u16, headers: Vec<(String, String)>) -> Error { - let unsupported = serde_json::from_str::(&body) - .ok() - .filter(|error| error.kind.ends_with(UNSUPPORTED_DOCUMENT)); - Error::Provider { - status, - body: match unsupported { - Some(error) => format!( - "{UNSUPPORTED_DOCUMENT}: {}. aws_textract uses Textract's synchronous API, which reads a JPEG, PNG, or a single-page PDF or TIFF; other formats and multi-page documents are not supported", - error.message - ), - None => body, +pub(super) fn error_class(body: String, status: u16, headers: reqwest::header::HeaderMap) -> Error { + match serde_json::from_str::(&body) { + Ok(error) if error.kind.ends_with(UNSUPPORTED_DOCUMENT) => Error::RejectedDocument { + provider: "AWS Textract", + message: error.message, + guidance: SYNC_API_GUIDANCE, + }, + Ok(_) | Err(_) => Error::Provider { + status, + body, + headers: Box::new(headers), }, - headers, } } @@ -408,11 +407,11 @@ mod tests { fn a_model_outside_the_operations_is_refused_with_the_supported_names(#[case] model: &str) { let error = TextractOperation::from_model(model).unwrap_err(); - assert_eq!( - error.to_string(), - format!( - "invalid model: aws_textract has no model {model:?} - use one of: detect-document-text, analyze-document" - ) + let message = error.to_string(); + assert!(message.contains(&format!("{model:?}")), "{message}"); + assert!( + message.contains("detect-document-text, analyze-document"), + "{message}" ); assert_eq!(error.http_status_code(), Some(400)); } @@ -542,43 +541,46 @@ mod tests { #[rstest] #[case::bare_type( r#"{"__type":"UnsupportedDocumentException","Message":"Request has unsupported document format"}"#, - Some("Request has unsupported document format") + "Request has unsupported document format" )] #[case::namespaced_type( r#"{"__type":"com.amazonaws.textract#UnsupportedDocumentException","Message":"bad"}"#, - Some("bad") + "bad" )] #[case::lowercase_message( r#"{"__type":"UnsupportedDocumentException","message":"bad"}"#, - Some("bad") + "bad" )] - #[case::other_exception(r#"{"__type":"AccessDeniedException","Message":"no"}"#, None)] - #[case::json_without_a_type(r#"{"Message":"no"}"#, None)] - #[case::not_json("bad gateway", None)] - fn only_an_unsupported_document_gains_the_sync_api_hint( + fn an_unsupported_document_is_the_callers_fault_with_the_sync_api_hint( #[case] body: &str, - #[case] hinted_message: Option<&str>, + #[case] provider_message: &str, ) { - let response_headers = vec![("x-amzn-requestid".to_string(), "abc".to_string())]; + let error = error_class(body.into(), 400, reqwest::header::HeaderMap::new()); - let Error::Provider { - status, - body: reported, - headers, - } = error_class(body.into(), 400, response_headers.clone()) - else { - panic!("expected a provider error"); - }; + let message = error.to_string(); + assert!(error.is_request()); + assert_eq!(error.http_status_code(), Some(400)); + assert!(error.upstream_response().is_none()); + assert!(message.contains(provider_message), "{message}"); + assert!(message.contains(HINT), "{message}"); + } - assert_eq!(status, 400); - assert_eq!(headers, response_headers); - match hinted_message { - Some(message) => { - assert!(reported.contains(message), "{reported}"); - assert!(reported.contains(HINT), "{reported}"); - } - None => assert_eq!(reported, body), - } + #[rstest] + #[case::other_exception(r#"{"__type":"AccessDeniedException","Message":"no"}"#)] + #[case::json_without_a_type(r#"{"Message":"no"}"#)] + #[case::not_json("bad gateway")] + fn any_other_rejection_passes_through_untouched(#[case] body: &str) { + let response_headers = reqwest::header::HeaderMap::from_iter([( + reqwest::header::HeaderName::from_static("x-amzn-requestid"), + reqwest::header::HeaderValue::from_static("abc"), + )]); + + let error = error_class(body.into(), 403, response_headers.clone()); + + assert_eq!( + error.upstream_response(), + Some((403, body, &response_headers)) + ); } #[rstest] diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs index 4b27a191734..34213926942 100644 --- a/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs @@ -104,7 +104,7 @@ impl BaseOcrConfig for TextractDetectTextConfig { &self, error_message: String, status_code: u16, - headers: Vec<(String, String)>, + headers: reqwest::header::HeaderMap, ) -> Error { error_class(error_message, status_code, headers) } @@ -237,7 +237,7 @@ mod tests { let error = TextractDetectTextConfig.get_error_class( r#"{"__type":"UnsupportedDocumentException","Message":"Request has unsupported document format"}"#.into(), 400, - Vec::new(), + reqwest::header::HeaderMap::new(), ); assert!( diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/cohere_parse_transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/cohere_parse_transformation.rs index 72e321547e3..bec8bee3f34 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/cohere_parse_transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/cohere_parse_transformation.rs @@ -120,7 +120,8 @@ impl BaseOcrConfig for AzureAICohereParseConfig { impl AzureAICohereParseConfig { fn get_complete_url(&self, base: &str) -> Result { - let mut url = reqwest::Url::parse(base).map_err(|_| invalid_api_base())?; + let mut url = + reqwest::Url::parse(base).map_err(|error| Error::request_field("api_base", error))?; if !matches!(url.scheme(), "http" | "https") { return Err(invalid_api_base()); } @@ -133,13 +134,14 @@ impl AzureAICohereParseConfig { ApiUrl::parse(url.as_str()) .and_then(|url| url.complete_path(&AZURE_COHERE_PARSE_PATH)) .map(|url| url.into_string()) - .map_err(|_| invalid_api_base()) + .map_err(|error| Error::request_field("api_base", error)) } } fn invalid_api_base() -> Error { Error::RequestField { path: "api_base".into(), + cause: None, } } diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/common_utils.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/common_utils.rs index 8cbf2455e01..7da685febf9 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/common_utils.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/common_utils.rs @@ -26,7 +26,7 @@ pub(super) async fn resolve_entra( .get_azure_ad_token(config, env_lookup) .await .or_else(|error| match error { - litellm_auth::Error::EmptyCallerCredential(_) => Ok(None), + litellm_auth::Error::EmptyCallerCredential { .. } => Ok(None), other => Err(other), }) .map(|credential| { @@ -48,7 +48,7 @@ pub(super) fn validate_destination( && credential_source != InputSource::Request { return Err(litellm_auth::Error::InvalidConfiguration( - "host credentials cannot be sent to a request-controlled Azure endpoint".into(), + litellm_auth::ErrorDetail::RequestControlledEndpoint { provider: "Azure" }, ) .into()); } diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs index 2a50989f4e4..d7233db0a76 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs @@ -13,7 +13,7 @@ use tokio::time::Instant; use crate::base_llm::ocr::{ document::InlineDocument, - error::Error, + error::{Error, PagesFault, PollUrl}, handler::{CallHooks, OcrClient, read_json_response}, settings::OcrSettings, transformation::{ @@ -185,7 +185,10 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig { ) -> Result { let endpoint = nonblank(request.connection.api_base.clone()) .or_else(|| nonblank(request.connection.secret(AZURE_DI_ENDPOINT_ENV))) - .ok_or_else(|| Error::Auth(litellm_auth::Error::ProviderAuthentication("Missing Azure Document Intelligence API Base - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT or pass api_base".into())))?; + .ok_or(Error::Auth(litellm_auth::Error::MissingApiBase { + provider: "Azure Document Intelligence", + guidance: "Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT or pass api_base", + }))?; self.build_ocr_url( &endpoint, &request.model, @@ -256,14 +259,12 @@ fn normalize_pages_param(pages: Option<&Value>) -> Result, Error> Some(Value::Array(pages)) if pages.iter().all(Value::is_number) => pages .iter() .map(|page| { - let page = page - .as_i64() - .ok_or_else(|| Error::Pages("page index is out of range".into()))?; + let page = page.as_i64().ok_or(Error::Pages(PagesFault::OutOfRange))?; if page < 0 { - return Err(Error::Pages("negative page index".into())); + return Err(Error::Pages(PagesFault::Negative)); } page.checked_add(1) - .ok_or_else(|| Error::Pages("page index is out of range".into())) + .ok_or(Error::Pages(PagesFault::OutOfRange)) }) .collect::, _>>()? .into_iter() @@ -276,7 +277,7 @@ fn normalize_pages_param(pages: Option<&Value>) -> Result, Error> token .as_str() .map(str::trim) - .ok_or_else(|| Error::Pages("expected only integers or only strings".into())) + .ok_or(Error::Pages(PagesFault::Mixed)) }) .collect::, _>>()? .join(","), @@ -286,13 +287,11 @@ fn normalize_pages_param(pages: Option<&Value>) -> Result, Error> .collect::>() .join(","), Some(_) => { - return Err(Error::Pages( - "expected an array of integers or strings, or a native page range".into(), - )); + return Err(Error::Pages(PagesFault::Shape)); } }; if !normalized.split(',').all(valid_page_token) { - return Err(Error::Pages("invalid native page range".into())); + return Err(Error::Pages(PagesFault::Range)); } Ok(Some(normalized)) } @@ -458,8 +457,14 @@ async fn read_operation_response( .and_then(|value| value.to_str().ok()) .ok_or(Error::PollLocation)? .to_string(); - let original = Url::parse(original_url).map_err(|_| Error::PollOrigin)?; - let operation = Url::parse(&location).map_err(|_| Error::PollOrigin)?; + let original = Url::parse(original_url).map_err(|source| Error::InvalidPollUrl { + url: PollUrl::Submission, + source, + })?; + let operation = Url::parse(&location).map_err(|source| Error::InvalidPollUrl { + url: PollUrl::OperationLocation, + source, + })?; if original.origin() != operation.origin() || !operation.username().is_empty() || operation.password().is_some() @@ -504,8 +509,7 @@ async fn poll_operation( let response = tokio::time::timeout_at(deadline, litellm_http::request::http_request(builder)) .await - .map_err(|_| Error::PollTimeout)? - .map_err(litellm_http::transport::Error::from)?; + .map_err(|_| Error::PollTimeout)??; let retry = response .headers() .get(reqwest::header::RETRY_AFTER) @@ -578,9 +582,7 @@ impl AzureDocumentIntelligenceOcrConfig { ) .into_string() }) - .map_err(|_| Error::RequestField { - path: "api_base".into(), - }) + .map_err(|error| Error::request_field("api_base", error)) } async fn resolve_headers( @@ -623,7 +625,10 @@ impl AzureDocumentIntelligenceOcrConfig { } let token = super::super::common_utils::resolve_entra(auth, config, env_lookup) .await? - .ok_or(Error::MissingAzureDocumentIntelligenceCredentials)?; + .ok_or(Error::MissingCredential { + credential: "Azure Document Intelligence credentials", + guidance: "set AZURE_DOCUMENT_INTELLIGENCE_API_KEY or configure Entra ID", + })?; super::super::common_utils::validate_destination(connection, token.source())?; Ok( std::iter::once(("Authorization".into(), format!("Bearer {}", token.value()))) @@ -842,7 +847,7 @@ mod tests { .unwrap_err(); assert!( - error + litellm_error::Report::new(&error) .to_string() .contains("request-controlled Azure endpoint") ); diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs index 83d88bbb3cb..2f7957fe56e 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs @@ -169,7 +169,10 @@ impl AzureAiOcrConfig { } let key = super::common_utils::resolve_entra(auth, config, env_lookup) .await? - .ok_or(Error::MissingAzureAiCredentials)?; + .ok_or(Error::MissingCredential { + credential: "Azure AI credentials", + guidance: "set AZURE_AI_API_KEY or configure Entra ID", + })?; super::common_utils::validate_destination(connection, key.source())?; Ok(bearer_headers(connection, key.value())) } @@ -183,9 +186,7 @@ impl AzureAiOcrConfig { ApiUrl::parse(&base) .and_then(|url| url.complete_path(&AZURE_AI_OCR_PATH)) .map(|url| url.into_string()) - .map_err(|_| Error::RequestField { - path: "api_base".into(), - }) + .map_err(|error| Error::request_field("api_base", error)) } } @@ -302,7 +303,7 @@ mod tests { .unwrap_err(); assert!( - error + litellm_error::Report::new(&error) .to_string() .contains("request-controlled Azure endpoint") ); diff --git a/litellm-rust/crates/llms/src/base_llm/base_model_iterator.rs b/litellm-rust/crates/llms/src/base_llm/base_model_iterator.rs index a283ca84089..6df321b4b16 100644 --- a/litellm-rust/crates/llms/src/base_llm/base_model_iterator.rs +++ b/litellm-rust/crates/llms/src/base_llm/base_model_iterator.rs @@ -5,6 +5,15 @@ use futures_util::{Stream, StreamExt, stream, stream::BoxStream}; pub type ByteStream = BoxStream<'static, Result>; +/// The transport error a body stream carried as an `io::Error`, however many times framing +/// wrapped it on the way. +pub(crate) fn tunnelled_http(error: &io::Error) -> Option<&litellm_http::Error> { + let inner = error.get_ref()?; + inner + .downcast_ref::() + .or_else(|| tunnelled_http(inner.downcast_ref::()?)) +} + pub trait StreamTransformer { type Input; type Output; diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/document.rs b/litellm-rust/crates/llms/src/base_llm/ocr/document.rs index 57976240e1b..898c24369c3 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/document.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/document.rs @@ -1,9 +1,6 @@ use base64::{Engine, engine::general_purpose::STANDARD}; use data_url::{DataUrl, DataUrlError, forgiving_base64::DecodeError, mime::Mime}; -use litellm_http::{ - media::{DownloadPolicy, Error as MediaError, MediaFetcher}, - transport::Error as TransportError, -}; +use litellm_http::media::{DownloadPolicy, MediaFetcher}; use reqwest::Url; use crate::base_llm::ocr::{ @@ -61,9 +58,7 @@ pub async fn inline_remote_document( validate_inline_document(&document)?; return Ok(document); } - let url = Url::parse(source).map_err(|_| Error::RequestField { - path: "document URL".into(), - })?; + let url = Url::parse(source).map_err(|error| Error::request_field("document URL", error))?; let downloaded = fetcher .fetch( url, @@ -74,7 +69,7 @@ pub async fn inline_remote_document( }, ) .await - .map_err(map_media_error)?; + .map_err(Error::DocumentDownload)?; let result = document.with_source(format!( "data:{};base64,{}", downloaded.content_type, @@ -84,28 +79,6 @@ pub async fn inline_remote_document( Ok(result) } -fn map_media_error(error: MediaError) -> Error { - match error { - MediaError::BlockedUrl => Error::BlockedDocumentUrl, - MediaError::DownloadDisabled => Error::DownloadDisabled, - MediaError::DownloadTooLarge => Error::DownloadTooLarge, - MediaError::TooManyRedirects => Error::TooManyRedirects, - MediaError::MissingRedirectLocation => Error::MissingRedirectLocation, - MediaError::InvalidRedirect => Error::InvalidRedirect, - MediaError::Http(status) => TransportError::Http { - status, - body: "OCR document download failed".into(), - } - .into(), - MediaError::Timeout => TransportError::Http { - status: 408, - body: "OCR document download timed out".into(), - } - .into(), - MediaError::Transport(error) => error.into(), - } -} - #[cfg(test)] mod tests { use std::collections::BTreeMap as Map; @@ -119,20 +92,17 @@ mod tests { } } - #[test] - fn decodes_data_urls_and_limits_decoded_size() { - for (source, expected) in [ - ("data:application/pdf;base64,YWJj", b"abc".as_slice()), - ("DATA:application/pdf;BASE64,YWI", b"ab".as_slice()), - ("data:,a%20b%00%FF", b"a b\0\xff".as_slice()), - ] { - let inline = InlineDocument::parse(source).unwrap().unwrap(); - assert_eq!(inline.decode(expected.len()).unwrap(), expected); - assert!(matches!( - inline.decode(expected.len() - 1), - Err(Error::InlineDocumentTooLarge) - )); - } + #[rstest::rstest] + #[case::base64("data:application/pdf;base64,YWJj", b"abc")] + #[case::uppercase_unpadded("DATA:application/pdf;BASE64,YWI", b"ab")] + #[case::percent_encoded("data:,a%20b%00%FF", b"a b\0\xff")] + fn decodes_data_urls_and_limits_decoded_size(#[case] source: &str, #[case] expected: &[u8]) { + let inline = InlineDocument::parse(source).unwrap().unwrap(); + assert_eq!(inline.decode(expected.len()).unwrap(), expected); + assert!(matches!( + inline.decode(expected.len() - 1), + Err(Error::InlineDocumentTooLarge) + )); } #[test] @@ -150,15 +120,12 @@ mod tests { ); } - #[test] - fn rejects_invalid_inline_documents() { - for source in [ - "https://example.com/document.pdf", - "data:application/pdf;base64", - "data:application/pdf;base64,INVALID!", - ] { - assert!(validate_inline_document(&document(source)).is_err()); - } + #[rstest::rstest] + #[case::remote("https://example.com/document.pdf")] + #[case::no_payload("data:application/pdf;base64")] + #[case::invalid_base64("data:application/pdf;base64,INVALID!")] + fn rejects_invalid_inline_documents(#[case] source: &str) { + assert!(validate_inline_document(&document(source)).is_err()); } #[tokio::test] diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs index 392dd56412c..9a0e0519f12 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs @@ -1,53 +1,62 @@ +use std::{io, path::PathBuf, sync::Arc}; + +use litellm_core_utils::call_arguments::ArgumentError; +use litellm_error::Source; +use reqwest::header::HeaderMap; + +static NO_HEADERS: std::sync::LazyLock = std::sync::LazyLock::new(HeaderMap::new); + #[derive(Clone, Debug, thiserror::Error)] pub enum Error { #[error("upstream OCR error ({status}): {body}")] Provider { status: u16, body: String, - headers: Vec<(String, String)>, + headers: Box, }, #[error("File is empty or could not be read")] EmptyFile, - #[error("Failed to read OCR file {}: {source}", path.display())] + #[error("could not read OCR file {}", path.display())] FileRead { - path: std::path::PathBuf, + path: PathBuf, #[source] - source: std::sync::Arc, + cause: Arc, }, - #[error("OCR document preparation task failed: {0}")] - DocumentTask(#[source] std::sync::Arc), - #[error("Invalid MIME type: {0}")] + #[error("OCR document preparation task failed")] + DocumentTask(#[source] Arc), + #[error("invalid MIME type {0:?}")] InvalidMimeType(String), - #[error( - "Cohere Parse only accepts `image_url` documents; document_url and PDF inputs are not supported" - )] - CohereImageOnly, + #[error("{provider} cannot read this document - {guidance}")] + UnsupportedDocument { + provider: &'static str, + guidance: &'static str, + }, + #[error("{provider} cannot read this document: {message} - {guidance}")] + RejectedDocument { + provider: &'static str, + message: String, + guidance: &'static str, + }, #[error("Invalid `req_format`. Expected 'native' or 'litellm'.")] RequestFormat, #[error("invalid OCR request field: {path}")] - RequestField { path: String }, + RequestField { + path: String, + #[source] + cause: Option, + }, + #[error("invalid OCR optional params")] + Arguments(#[from] ArgumentError), #[error("missing required field: {0}")] MissingField(&'static str), - #[error("Document URL is required")] + #[error("document URL is required")] MissingDocumentUrl, #[error("invalid OCR document data URI")] InvalidDataUri, - #[error( - "Reducto requires a reducto:// id or a data URI; plain HTTP URLs are not supported, upload the file first" - )] - ReductoSource, #[error("inline OCR document exceeds the size limit")] InlineDocumentTooLarge, - #[error("OCR document URL is blocked by network policy")] - BlockedDocumentUrl, - #[error("OCR document downloads are disabled")] - DownloadDisabled, - #[error("OCR document download exceeds the size limit")] - DownloadTooLarge, - #[error("OCR document download exceeded the redirect limit")] - TooManyRedirects, #[error("invalid OCR pages: {0}")] - Pages(String), + Pages(PagesFault), #[error("invalid OCR features")] Features, #[error("OCR model cannot be a dot segment")] @@ -55,19 +64,25 @@ pub enum Error { #[error("OCR response exceeds the size limit of {limit} bytes")] TooLarge { limit: usize }, #[error("invalid OCR response field: {path}")] - ResponseField { path: String }, + ResponseField { + path: String, + #[source] + cause: Option, + }, #[error("OCR response is missing non-empty content")] EmptyContent, - #[error("OCR document redirect is missing a location")] - MissingRedirectLocation, - #[error("OCR document redirect location is invalid")] - InvalidRedirect, #[error("OCR operation ended with status {0}")] OperationStatus(String), #[error("OCR response numeric value is out of range: {0}")] NumericRange(&'static str), #[error("OCR accepted response is missing a valid operation-location")] PollLocation, + #[error("OCR {url} is not a valid URL")] + InvalidPollUrl { + url: PollUrl, + #[source] + source: url::ParseError, + }, #[error("OCR operation-location must use the submission origin without credentials")] PollOrigin, #[error("OCR polling timed out")] @@ -82,108 +97,313 @@ pub enum Error { model: String, supported: &'static [&'static str], }, - #[error("invalid request: {0}")] - InvalidRequest(String), - #[error("invalid response: {0}")] - InvalidResponse(String), - #[error( - "invalid authentication configuration: Missing Azure AI credentials - set AZURE_AI_API_KEY or configure Entra ID" - )] - MissingAzureAiCredentials, - #[error( - "invalid authentication configuration: Missing Azure Document Intelligence credentials - set AZURE_DOCUMENT_INTELLIGENCE_API_KEY or configure Entra ID" - )] - MissingAzureDocumentIntelligenceCredentials, - #[error( - "Missing REDUCTO_API_KEY - set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()" - )] - MissingReductoApiKey, - #[error("secret resolution failed: {0}")] - Secret(#[source] std::sync::Arc), + #[error("Missing {credential} - {guidance}")] + MissingCredential { + credential: &'static str, + guidance: &'static str, + }, + #[error("secret resolution failed")] + Secret(#[source] Arc), + #[error("the OCR host driver failed")] + Host(#[source] litellm_host::machine::MachineFault), + #[error("the OCR host call raised an exception")] + HostCall, #[error(transparent)] Auth(#[from] litellm_auth::Error), #[error(transparent)] - Transport(#[from] litellm_http::transport::Error), - #[error(transparent)] Params(#[from] litellm_core_utils::params::Error), #[error(transparent)] Headers(#[from] litellm_http::request::HeaderError), #[error(transparent)] Http(#[from] litellm_http::Error), + #[error(transparent)] + DocumentDownload(litellm_http::Error), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] +pub enum PagesFault { + #[error("page index is out of range")] + OutOfRange, + #[error("negative page index")] + Negative, + #[error("expected only integers or only strings")] + Mixed, + #[error("expected an array of integers or strings, or a native page range")] + Shape, + #[error("invalid native page range")] + Range, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] +pub enum PollUrl { + #[error("submission URL")] + Submission, + #[error("operation-location")] + OperationLocation, } impl From for Error { fn from(fault: litellm_host::machine::MachineFault) -> Self { - use litellm_host::machine::MachineFault; - Self::InvalidRequest(match fault { - MachineFault::Abandoned => "OCR host driver was abandoned".into(), - MachineFault::Protocol(message) => format!("OCR {message}"), - }) - } -} - -impl From for Error { - fn from(error: litellm_core_utils::call_arguments::ArgumentError) -> Self { - Self::RequestField { - path: format!("optional_params.{}", error.path), - } + Self::Host(fault) } } impl Error { - pub fn http_status_code(&self) -> Option { - match self { - Self::Provider { status, .. } - | Self::Transport(litellm_http::transport::Error::Http { status, .. }) => Some(*status), - error if error.is_request() => Some(400), - _ => None, + pub fn request_field( + path: impl Into, + cause: impl std::error::Error + Send + Sync + 'static, + ) -> Self { + Self::RequestField { + path: path.into(), + cause: Some(Source::new(cause)), } } - pub fn is_request(&self) -> bool { - matches!( - self, + pub fn response_field( + path: impl Into, + cause: impl std::error::Error + Send + Sync + 'static, + ) -> Self { + Self::ResponseField { + path: path.into(), + cause: Some(Source::new(cause)), + } + } + + pub fn upstream_response(&self) -> Option<(u16, &str, &HeaderMap)> { + match self { + Self::Provider { + status, + body, + headers, + } => Some((*status, body, headers)), + Self::DocumentDownload(error) => error.status().map(|status| { + let body = if error.is_timeout() { + "OCR document download timed out" + } else { + "OCR document download failed" + }; + (status, body, &*NO_HEADERS) + }), + Self::Http(error) => error + .upstream_response() + .map(|(status, body)| (status, body, &*NO_HEADERS)), Self::EmptyFile - | Self::InvalidMimeType(_) - | Self::CohereImageOnly - | Self::RequestFormat - | Self::RequestField { .. } - | Self::MissingField(_) - | Self::MissingDocumentUrl - | Self::InvalidDataUri - | Self::ReductoSource - | Self::InlineDocumentTooLarge - | Self::BlockedDocumentUrl - | Self::DownloadDisabled - | Self::DownloadTooLarge - | Self::TooManyRedirects - | Self::Pages(_) - | Self::Features - | Self::DotModel - | Self::InvalidRequest(_) - | Self::InvalidProvider(_) - | Self::InvalidModel { .. } - | Self::Params(_) - | Self::Headers(_) - | Self::Http(_) - | Self::Auth(litellm_auth::Error::MissingParam { .. }) - ) + | Self::FileRead { .. } + | Self::DocumentTask(_) + | Self::InvalidMimeType(_) + | Self::UnsupportedDocument { .. } + | Self::RejectedDocument { .. } + | Self::RequestFormat + | Self::RequestField { .. } + | Self::Arguments(_) + | Self::MissingField(_) + | Self::MissingDocumentUrl + | Self::InvalidDataUri + | Self::InlineDocumentTooLarge + | Self::Pages(_) + | Self::Features + | Self::DotModel + | Self::TooLarge { .. } + | Self::ResponseField { .. } + | Self::EmptyContent + | Self::OperationStatus(_) + | Self::NumericRange(_) + | Self::PollLocation + | Self::InvalidPollUrl { .. } + | Self::PollOrigin + | Self::PollTimeout + | Self::Unsupported(_) + | Self::InvalidProvider(_) + | Self::InvalidModel { .. } + | Self::MissingCredential { .. } + | Self::Secret(_) + | Self::Host(_) + | Self::HostCall + | Self::Auth(_) + | Self::Params(_) + | Self::Headers(_) => None, + } + } + + pub fn is_missing_credential(&self) -> bool { + match self { + Self::MissingCredential { .. } => true, + Self::Auth(error) => error.is_missing_credential(), + Self::Provider { .. } + | Self::EmptyFile + | Self::FileRead { .. } + | Self::DocumentTask(_) + | Self::InvalidMimeType(_) + | Self::UnsupportedDocument { .. } + | Self::RejectedDocument { .. } + | Self::RequestFormat + | Self::RequestField { .. } + | Self::Arguments(_) + | Self::MissingField(_) + | Self::MissingDocumentUrl + | Self::InvalidDataUri + | Self::InlineDocumentTooLarge + | Self::Pages(_) + | Self::Features + | Self::DotModel + | Self::TooLarge { .. } + | Self::ResponseField { .. } + | Self::EmptyContent + | Self::OperationStatus(_) + | Self::NumericRange(_) + | Self::PollLocation + | Self::InvalidPollUrl { .. } + | Self::PollOrigin + | Self::PollTimeout + | Self::Unsupported(_) + | Self::InvalidProvider(_) + | Self::InvalidModel { .. } + | Self::Secret(_) + | Self::Host(_) + | Self::HostCall + | Self::Params(_) + | Self::Headers(_) + | Self::Http(_) + | Self::DocumentDownload(_) => false, + } + } + + pub fn is_unsupported(&self) -> bool { + match self { + Self::Unsupported(_) => true, + Self::Provider { .. } + | Self::EmptyFile + | Self::FileRead { .. } + | Self::DocumentTask(_) + | Self::InvalidMimeType(_) + | Self::UnsupportedDocument { .. } + | Self::RejectedDocument { .. } + | Self::RequestFormat + | Self::RequestField { .. } + | Self::Arguments(_) + | Self::MissingField(_) + | Self::MissingDocumentUrl + | Self::InvalidDataUri + | Self::InlineDocumentTooLarge + | Self::Pages(_) + | Self::Features + | Self::DotModel + | Self::TooLarge { .. } + | Self::ResponseField { .. } + | Self::EmptyContent + | Self::OperationStatus(_) + | Self::NumericRange(_) + | Self::PollLocation + | Self::InvalidPollUrl { .. } + | Self::PollOrigin + | Self::PollTimeout + | Self::InvalidProvider(_) + | Self::InvalidModel { .. } + | Self::MissingCredential { .. } + | Self::Secret(_) + | Self::Host(_) + | Self::HostCall + | Self::Auth(_) + | Self::Params(_) + | Self::Headers(_) + | Self::Http(_) + | Self::DocumentDownload(_) => false, + } + } + + pub fn http_status_code(&self) -> Option { + if let Some((status, ..)) = self.upstream_response() { + return Some(status); + } + if let Self::Http(error) = self + && let Some(status) = error.status() + { + return Some(status); + } + self.is_request().then_some(400) + } + + pub fn is_request(&self) -> bool { + match self { + Self::EmptyFile + | Self::InvalidMimeType(_) + | Self::UnsupportedDocument { .. } + | Self::RejectedDocument { .. } + | Self::RequestFormat + | Self::RequestField { .. } + | Self::Arguments(_) + | Self::MissingField(_) + | Self::MissingDocumentUrl + | Self::InvalidDataUri + | Self::InlineDocumentTooLarge + | Self::Pages(_) + | Self::Features + | Self::DotModel + | Self::InvalidProvider(_) + | Self::InvalidModel { .. } + | Self::Params(_) + | Self::Headers(_) => true, + Self::Auth(error) => error.is_request(), + Self::Http(error) | Self::DocumentDownload(error) => error.is_request(), + Self::Provider { .. } + | Self::FileRead { .. } + | Self::DocumentTask(_) + | Self::TooLarge { .. } + | Self::ResponseField { .. } + | Self::EmptyContent + | Self::OperationStatus(_) + | Self::NumericRange(_) + | Self::PollLocation + | Self::InvalidPollUrl { .. } + | Self::PollOrigin + | Self::PollTimeout + | Self::Unsupported(_) + | Self::MissingCredential { .. } + | Self::Secret(_) + | Self::Host(_) + | Self::HostCall => false, + } } pub fn is_response(&self) -> bool { - matches!( - self, + match self { Self::TooLarge { .. } - | Self::ResponseField { .. } - | Self::EmptyContent - | Self::MissingRedirectLocation - | Self::InvalidRedirect - | Self::OperationStatus(_) - | Self::NumericRange(_) - | Self::PollLocation - | Self::PollOrigin - | Self::PollTimeout - | Self::InvalidResponse(_) - ) + | Self::ResponseField { .. } + | Self::EmptyContent + | Self::OperationStatus(_) + | Self::NumericRange(_) + | Self::PollLocation + | Self::InvalidPollUrl { .. } + | Self::PollOrigin + | Self::PollTimeout => true, + Self::Http(error) | Self::DocumentDownload(error) => error.is_response(), + Self::Provider { .. } + | Self::EmptyFile + | Self::FileRead { .. } + | Self::DocumentTask(_) + | Self::InvalidMimeType(_) + | Self::UnsupportedDocument { .. } + | Self::RejectedDocument { .. } + | Self::RequestFormat + | Self::RequestField { .. } + | Self::Arguments(_) + | Self::MissingField(_) + | Self::MissingDocumentUrl + | Self::InvalidDataUri + | Self::InlineDocumentTooLarge + | Self::Pages(_) + | Self::Features + | Self::DotModel + | Self::Unsupported(_) + | Self::InvalidProvider(_) + | Self::InvalidModel { .. } + | Self::MissingCredential { .. } + | Self::Secret(_) + | Self::Host(_) + | Self::HostCall + | Self::Auth(_) + | Self::Params(_) + | Self::Headers(_) => false, + } } } diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs index a9ddcb87086..7bbc751c614 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs @@ -9,7 +9,6 @@ use litellm_http::{ media::{MediaFetcher, UrlPolicy}, outbound::{RequestSigner, json, signed}, request::header_pairs, - transport, }; use serde::{Serialize, de::DeserializeOwned}; use serde_json::Value; @@ -134,16 +133,12 @@ pub async fn ocr( .provider_http() .execute(http) .await - .map_err(transport_error)?; + .map_err(litellm_http::Error::from)?; if !response.status().is_success() { - let headers = header_pairs(response.headers()); - return match read_response_bytes(response, request.connection.max_response_bytes).await { - Err(Error::Transport(transport::Error::Http { status, body })) => { - Err(config.get_error_class(body, status, headers)) - } - Err(error) => Err(error), - Ok(_) => unreachable!("non-success response produces an HTTP error"), - }; + let status = response.status().as_u16(); + let headers = response.headers().clone(); + let body = read_error_body(response, request.connection.max_response_bytes).await?; + return Err(config.get_error_class(body, status, headers)); } let context = OcrResponseContext { client, @@ -171,43 +166,41 @@ pub async fn read_response_bytes( mut response: reqwest::Response, limit: usize, ) -> Result { - let status = response.status(); - if status.is_success() - && response - .content_length() - .is_some_and(|length| length > limit as u64) + if !response.status().is_success() { + let status = response.status().as_u16(); + let headers = response.headers().clone(); + let body = read_error_body(response, limit).await?; + return Err(Error::Provider { + status, + body, + headers: Box::new(headers), + }); + } + if response + .content_length() + .is_some_and(|length| length > limit as u64) { return Err(Error::TooLarge { limit }); } let mut bytes = BytesMut::new(); - while let Some(chunk) = response.chunk().await.map_err(transport_error)? { - let remaining = limit.saturating_sub(bytes.len()); - if status.is_success() && chunk.len() > remaining { + while let Some(chunk) = response.chunk().await.map_err(litellm_http::Error::from)? { + if chunk.len() > limit.saturating_sub(bytes.len()) { return Err(Error::TooLarge { limit }); } - bytes.extend_from_slice(&chunk[..chunk.len().min(remaining)]); - if !status.is_success() && bytes.len() == limit { - break; - } - } - if !status.is_success() { - return Err(transport::Error::Http { - status: status.as_u16(), - body: String::from_utf8_lossy(&bytes).into_owned(), - } - .into()); + bytes.extend_from_slice(&chunk); } Ok(bytes.freeze()) } -pub fn transport_error(error: reqwest::Error) -> Error { - if error.is_timeout() { - return Error::Transport(transport::Error::Http { - status: 408, - body: "OCR request timed out".into(), - }); +async fn read_error_body(mut response: reqwest::Response, limit: usize) -> Result { + let mut bytes = BytesMut::new(); + while bytes.len() < limit + && let Some(chunk) = response.chunk().await.map_err(litellm_http::Error::from)? + { + let remaining = limit - bytes.len(); + bytes.extend_from_slice(&chunk[..chunk.len().min(remaining)]); } - transport::Error::from(error).into() + Ok(String::from_utf8_lossy(&bytes).into_owned()) } pub async fn transform_request_body( @@ -231,6 +224,7 @@ pub async fn transform_request_body( if !changed.body.is_object() { return Err(Error::RequestField { path: "guardrail.body".into(), + cause: None, }); } config.validate_request_body(&changed.body)?; @@ -274,9 +268,8 @@ pub async fn guardrail_document( headers: &[(String, String)], hooks: &dyn CallHooks, ) -> Result<(OcrDocument, Vec<(String, String)>), Error> { - let body = serde_json::to_value(&request.document).map_err(|_| Error::RequestField { - path: "document".into(), - })?; + let body = serde_json::to_value(&request.document) + .map_err(|error| Error::request_field("document", error))?; let changed = hooks .before_provider_request(wire_request(url, headers, body)) .await?; @@ -290,6 +283,7 @@ pub fn body_document(body: &Value) -> Result { .and_then(Value::as_object) .ok_or_else(|| Error::RequestField { path: "body.document".into(), + cause: None, })?; let source = document .iter() @@ -309,8 +303,6 @@ mod tests { net::TcpListener, }; - use super::*; - #[rstest::rstest] #[tokio::test] async fn request_timeout_has_an_http_408_status() { @@ -326,10 +318,10 @@ mod tests { .send() .await .unwrap_err(); - assert!(matches!( - transport_error(error), - Error::Transport(transport::Error::Http { status: 408, .. }) - )); + assert_eq!( + Error::from(litellm_http::Error::from(error)).http_status_code(), + Some(408) + ); server.abort(); } @@ -400,10 +392,7 @@ mod tests { .await .unwrap_err(); - let Error::Transport(litellm_http::transport::Error::Http { status, body }) = error else { - panic!("unexpected error: {error}"); - }; - assert_eq!(status, 429); - assert_eq!(body, prefix); + let (status, body, _) = error.upstream_response().expect("the provider answered"); + assert_eq!((status, body), (429, prefix.as_str())); } } diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs index 73e7b64638d..19dd8e24226 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs @@ -194,17 +194,13 @@ pub struct DecodedOcrResponse { pub fn decode_request_value(value: Value, prefix: &str) -> Result { serde_path_to_error::deserialize(value.into_deserializer()).map_err(|error| { - Error::RequestField { - path: format!("{prefix}.{}", error.path()), - } + Error::request_field(format!("{prefix}.{}", error.path()), error.into_inner()) }) } pub fn decode_response_value(value: Value, prefix: &str) -> Result { serde_path_to_error::deserialize(value.into_deserializer()).map_err(|error| { - Error::ResponseField { - path: format!("{prefix}.{}", error.path()), - } + Error::response_field(format!("{prefix}.{}", error.path()), error.into_inner()) }) } @@ -213,19 +209,15 @@ pub fn decode_response( native: bool, ) -> Result, Error> { let mut deserializer = serde_json::Deserializer::from_slice(bytes); - let data = serde_path_to_error::deserialize(&mut deserializer).map_err(|error| { - Error::ResponseField { - path: error.path().to_string(), - } - })?; - deserializer.end().map_err(|_| Error::ResponseField { - path: "response".into(), - })?; + let data = serde_path_to_error::deserialize(&mut deserializer) + .map_err(|error| Error::response_field(error.path().to_string(), error.into_inner()))?; + deserializer + .end() + .map_err(|error| Error::response_field("response", error))?; let native = if native { Some( - serde_json::from_slice(bytes).map_err(|_| Error::ResponseField { - path: "response".into(), - })?, + serde_json::from_slice(bytes) + .map_err(|error| Error::response_field("response", error))?, ) } else { None @@ -369,12 +361,12 @@ pub trait BaseOcrConfig: Send + Sync + Sized + 'static { &self, error_message: String, status_code: u16, - headers: Vec<(String, String)>, + headers: reqwest::header::HeaderMap, ) -> Error { Error::Provider { status: status_code, body: error_message, - headers, + headers: Box::new(headers), } } diff --git a/litellm-rust/crates/llms/src/base_llm/sse.rs b/litellm-rust/crates/llms/src/base_llm/sse.rs index 5ff2e3adddd..23f0c1b9389 100644 --- a/litellm-rust/crates/llms/src/base_llm/sse.rs +++ b/litellm-rust/crates/llms/src/base_llm/sse.rs @@ -15,12 +15,15 @@ use bytes::Bytes; use futures_util::{StreamExt, TryStreamExt, future::ready, stream::BoxStream}; use litellm_framer::{ - frames, + SseError, frames, sse::{SseCodec, SseEvent, decode_all}, }; use serde::{Serialize, de::DeserializeOwned}; -use crate::{Error, ErrorDetail, base_llm::base_model_iterator::ByteStream}; +use crate::{ + Error, ErrorDetail, + base_llm::base_model_iterator::{ByteStream, tunnelled_http}, +}; pub type EventStream = BoxStream<'static, Result>; @@ -70,8 +73,13 @@ pub fn invalid(error: serde_json::Error) -> Error { Error::InvalidResponse(ErrorDetail::invalid("stream event", error)) } -fn framing(error: litellm_framer::SseError) -> Error { - Error::InvalidResponse(ErrorDetail::failed("stream framing", error)) +fn framing(error: SseError) -> Error { + match &error { + SseError::Body(body) if let Some(http) = tunnelled_http(body) => Error::Http(http.clone()), + SseError::Body(_) | SseError::InvalidUtf8(_) => { + Error::InvalidResponse(ErrorDetail::failed("stream framing", error)) + } + } } #[cfg(test)] @@ -94,6 +102,45 @@ mod tests { assert_eq!(is_done(&event), done); } + #[rstest] + #[tokio::test] + async fn a_transport_failure_tunnelled_through_the_body_comes_back_as_itself() { + use futures_util::StreamExt; + + let failure = litellm_http::Error::upstream(502, "gateway went away".into()); + let bytes: ByteStream = futures_util::stream::iter([ + Ok(Bytes::from_static(b"data: {}\n\n")), + Err(std::io::Error::other(failure.clone())), + ]) + .boxed(); + + let mut events = event_stream::(bytes, is_done); + assert!(events.next().await.expect("first event").is_ok()); + let error = events.next().await.expect("failure").unwrap_err(); + assert_eq!(error, Error::Http(failure)); + } + + #[rstest] + #[tokio::test] + async fn a_body_failure_that_is_not_transport_is_an_invalid_response() { + use futures_util::StreamExt; + + let bytes: ByteStream = + futures_util::stream::iter([Err(std::io::Error::other("disk full"))]).boxed(); + + let error = event_stream::(bytes, is_done) + .next() + .await + .expect("failure") + .unwrap_err(); + assert!(matches!(error, Error::InvalidResponse(_)), "{error:?}"); + assert!( + litellm_error::Report::new(&error) + .to_string() + .contains("disk full") + ); + } + #[rstest] fn the_done_frame_we_send_decodes_as_a_control_frame() { let events: Vec = decode(DONE_FRAME, is_done).unwrap(); diff --git a/litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs b/litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs index 61dd264336e..0f8703f223b 100644 --- a/litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs +++ b/litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs @@ -14,7 +14,7 @@ use serde_json::{Map, Value, json}; use strum::IntoStaticStr; use crate::{ - Error, + Error, ErrorDetail, base_llm::{ audio_transcription::transformation::{ AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, Headers, @@ -59,20 +59,12 @@ fn audio_fields(audio: Value) -> Result { .and_then(Value::as_str) .filter(|value| !value.is_empty()) .ok_or(Error::MissingField("audio.data"))?; - let format = object.get("format").cloned().ok_or_else(|| { - Error::InvalidRequest( - "audio.format must be wav, mp3, flac, or ogg" - .to_string() - .into(), - ) - })?; - let format: AudioFormat = serde_json::from_value(format).map_err(|_| { - Error::InvalidRequest( - "audio.format must be wav, mp3, flac, or ogg" - .to_string() - .into(), - ) - })?; + let format = object + .get("format") + .cloned() + .ok_or(Error::MissingField("audio.format"))?; + let format: AudioFormat = serde_json::from_value(format) + .map_err(|error| Error::InvalidRequest(ErrorDetail::invalid("audio.format", error)))?; Ok(AudioInput { data: data.to_string(), format, @@ -135,14 +127,13 @@ impl BaseAudioTranscriptionConfig for BedrockAudioTranscriptionConfig { ) -> Result { let response: ConverseResponse = serde_json::from_value(response_json).map_err(|error| { - Error::InvalidResponse(format!("invalid Converse response: {error}").into()) + Error::InvalidResponse(ErrorDetail::invalid("Converse response", error)) })?; if response.message_content_is_non_text() { - return Err(Error::InvalidResponse( - "Bedrock response contains non-text transcript content" - .to_string() - .into(), - )); + return Err(Error::InvalidResponse(ErrorDetail::Expected { + subject: "Bedrock transcript content", + expected: "text", + })); } Ok(AudioTranscriptionResponseData { text: response.content_text(), diff --git a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs index caa345a7024..acaf7b56f8b 100644 --- a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs @@ -19,7 +19,7 @@ use litellm_llms_types::{ use serde_json::{Map, Value, json}; use crate::{ - Error, + Error, ErrorDetail, base_llm::{ auth::AuthScheme, chat::{ @@ -154,21 +154,16 @@ impl BaseConfig for AmazonConverseConfig { response: ProviderChatResponseData, ) -> Result { let body = response.body; - if !body.is_object() { - return Err(Error::InvalidResponse( - "converse response is not an object".into(), - )); - } - for field in ["output", "usage"] { - if body.get(field).is_none() { - return Err(Error::InvalidResponse( - format!("invalid Converse response: missing field `{field}`").into(), - )); - } - } + let has_usage = body.get("usage").is_some(); let response: ConverseResponse = serde_json::from_value(body).map_err(|error| { - Error::InvalidResponse(format!("invalid Converse response: {error}").into()) + Error::InvalidResponse(ErrorDetail::invalid("Converse response", error)) })?; + if !has_usage { + return Err(Error::InvalidResponse(ErrorDetail::MissingField { + subject: "Converse response", + field: "usage", + })); + } // The route declines tool requests, so anything other than a text block // is something this path never asked for. Decline; the host falls back. if response.message_content_is_non_text() { @@ -858,34 +853,24 @@ mod tests { assert_eq!(err, Error::Unsupported("non-text response content block")); } - #[test] - fn errors_on_a_response_missing_required_fields() { - assert_eq!( - transform_response(json!("nope")).expect_err("not an object"), - Error::InvalidResponse("converse response is not an object".to_string().into()) - ); - assert_eq!( - transform_response(json!({"usage": {}})).expect_err("no output"), - Error::InvalidResponse("invalid Converse response: missing field `output`".into()) - ); - assert_eq!( - transform_response(json!({"output": {"message": {"content": []}}})) - .expect_err("no usage"), - Error::InvalidResponse("invalid Converse response: missing field `usage`".into()) - ); - } - - #[test] - fn rejects_malformed_text_and_token_counts() { - for response in [ - json!({"output": {"message": {"content": [{"text": 123}]}}, "usage": {"inputTokens": 1, "outputTokens": 1}}), - json!({"output": {"message": {"content": [{"text": "x"}]}}, "usage": {"inputTokens": "1", "outputTokens": 1}}), - ] { - assert!(matches!( - transform_response(response), - Err(Error::InvalidResponse(message)) if message.to_string().starts_with("invalid Converse response:") - )); - } + #[rstest::rstest] + #[case::not_an_object(json!("nope"), "invalid type")] + #[case::no_output(json!({"usage": {"inputTokens": 1, "outputTokens": 1}}), "missing field `output`")] + #[case::no_usage(json!({"output": {"message": {"content": []}}}), "missing usage")] + #[case::non_string_text( + json!({"output": {"message": {"content": [{"text": 123}]}}, "usage": {"inputTokens": 1, "outputTokens": 1}}), + "invalid type" + )] + #[case::non_integer_tokens( + json!({"output": {"message": {"content": [{"text": "x"}]}}, "usage": {"inputTokens": "1", "outputTokens": 1}}), + "invalid type" + )] + fn rejects_a_malformed_response(#[case] response: Value, #[case] fact: &str) { + let error = transform_response(response).expect_err("malformed response"); + assert!(matches!(error, Error::InvalidResponse(_)), "{error:?}"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains("Converse response"), "{report}"); + assert!(report.contains(fact), "{report}"); } #[test] diff --git a/litellm-rust/crates/llms/src/bedrock/chat/invoke_handler.rs b/litellm-rust/crates/llms/src/bedrock/chat/invoke_handler.rs index 878c424ae2d..15a6ff38822 100644 --- a/litellm-rust/crates/llms/src/bedrock/chat/invoke_handler.rs +++ b/litellm-rust/crates/llms/src/bedrock/chat/invoke_handler.rs @@ -2,6 +2,7 @@ use base64::Engine; use bytes::Buf; use futures_util::{Stream, StreamExt}; use litellm_framer::{ + EventStreamError, aws_event_stream::{AwsEventStreamCodec, Message}, frames, }; @@ -11,9 +12,8 @@ use serde_json::Value; use crate::{ Error, - anthropic::chat::handler::ModelResponseIterator, base_llm::{ - chat::streaming::{ChatStream, StreamShape}, + base_model_iterator::tunnelled_http, messages::streaming::{ByteStream, EventStream}, }, }; @@ -47,11 +47,21 @@ where B: Buf + Send, E: std::error::Error + Send + Sync + 'static, { - frames(input, AwsEventStreamCodec).map(|message| { - decode_invoke_chunk(message.map_err(|error| { + frames(input, AwsEventStreamCodec).map(|message| decode_invoke_chunk(message.map_err(framing)?)) +} + +fn framing(error: EventStreamError) -> Error { + match &error { + EventStreamError::Body(body) if let Some(http) = tunnelled_http(body) => { + Error::Http(http.clone()) + } + EventStreamError::Body(_) + | EventStreamError::InvalidLength(_) + | EventStreamError::Truncated + | EventStreamError::Malformed(_) => { Error::InvalidResponse(crate::ErrorDetail::failed("stream framing", error)) - })?) - }) + } + } } pub fn decode_invoke_anthropic_chunk(chunk: Value) -> Result { @@ -64,34 +74,6 @@ pub fn invoke_anthropic_event_stream(bytes: ByteStream) -> EventStream { Box::pin(invoke_chunk_stream(bytes).map(|chunk| decode_invoke_anthropic_chunk(chunk?))) } -#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::EnumString)] -enum InvokeProvider { - #[strum(serialize = "anthropic")] - Anthropic, - #[strum(serialize = "deepseek_r1")] - DeepseekR1, - #[strum(serialize = "moonshot")] - Moonshot, - #[strum(disabled)] - Unsupported, -} - -pub fn invoke_chat_stream(invoke_provider: &str, shape: StreamShape) -> Result { - match invoke_provider - .parse() - .unwrap_or(InvokeProvider::Unsupported) - { - InvokeProvider::Anthropic => Ok(ChatStream::new( - invoke_anthropic_event_stream, - ModelResponseIterator::new(shape), - )), - InvokeProvider::DeepseekR1 | InvokeProvider::Moonshot => Err(Error::Unsupported( - "Bedrock invoke streaming for this model family", - )), - InvokeProvider::Unsupported => Err(Error::Unsupported("Bedrock invoke streaming")), - } -} - #[cfg(test)] mod tests { use aws_smithy_eventstream::frame::write_message_to; @@ -112,22 +94,6 @@ mod tests { const TEXT_DELTA: &str = r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}"#; - #[rstest::rstest] - #[case::anthropic("anthropic", InvokeProvider::Anthropic)] - #[case::deepseek_r1("deepseek_r1", InvokeProvider::DeepseekR1)] - #[case::moonshot("moonshot", InvokeProvider::Moonshot)] - #[case::unknown("qwen", InvokeProvider::Unsupported)] - #[case::case_sensitive("Anthropic", InvokeProvider::Unsupported)] - fn invoke_provider_maps_each_model_family_name( - #[case] name: &str, - #[case] expected: InvokeProvider, - ) { - assert_eq!( - name.parse().unwrap_or(InvokeProvider::Unsupported), - expected - ); - } - fn aws_wire(chunk: &str) -> Vec { let payload = serde_json::json!({"bytes": STANDARD.encode(chunk)}); let message = Message::new(Bytes::from(serde_json::to_vec(&payload).unwrap())).add_header( diff --git a/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs index 72958fd6cb9..79c30df9b19 100644 --- a/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs @@ -1,6 +1,6 @@ use std::sync::LazyLock; -use litellm_auth::{AwsParams, CredentialPlacement, ParamSpec, SecretValue}; +use litellm_auth::{AwsParams, CredentialPlacement, ParamKind, ParamSpec, SecretValue}; use litellm_auth_aws::{AwsCredentialSource, resolve_aws_region}; use litellm_core_utils::settings::resolve_non_empty; use litellm_http::request::{header_value, with_header, without_headers}; @@ -36,6 +36,7 @@ const WORKSPACE_FIELDS: &[&str] = &[ const WORKSPACE_HEADER: &str = "anthropic-workspace-id"; static WORKSPACE: ParamSpec = ParamSpec { setting: "workspace_id", + kind: ParamKind::Setting, wire: WORKSPACE_FIELDS, module_global: None, env: WORKSPACE_ENVS, diff --git a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs index dadab413fe4..76ff27b56fc 100644 --- a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs @@ -195,11 +195,10 @@ impl CohereParseConfig { .and_then(env_lookup) .filter(|key| !key.trim().is_empty()) }) - .ok_or_else(|| { - Error::Auth(litellm_auth::Error::ProviderAuthentication( - "Missing COHERE_API_KEY - set it in the environment or pass api_key".into(), - )) - })?; + .ok_or(Error::Auth(litellm_auth::Error::MissingApiKey { + provider: "Cohere", + environment_variable: COHERE_API_KEY_ENV, + }))?; Ok( std::iter::once(("Authorization".into(), format!("Bearer {key}"))) .chain(connection.extra_headers.clone()) @@ -208,27 +207,28 @@ impl CohereParseConfig { } fn build_ocr_url(&self, api_base: &str) -> Result { - let parsed = reqwest::Url::parse(api_base).map_err(|_| invalid_api_base())?; + let parsed = reqwest::Url::parse(api_base) + .map_err(|error| Error::request_field("api_base", error))?; if !matches!(parsed.scheme(), "http" | "https") { return Err(invalid_api_base()); } ApiUrl::parse(api_base) .and_then(|url| url.complete_path(&["v2", "parse"])) .map(|url| url.into_string()) - .map_err(|_| invalid_api_base()) + .map_err(|error| Error::request_field("api_base", error)) } } pub fn validate_document(document: &OcrDocument) -> Result<(), Error> { let OcrDocument::ImageUrl { image_url, .. } = document else { - return Err(Error::CohereImageOnly); + return Err(image_only()); }; if image_url.is_empty() { - return Err(Error::CohereImageOnly); + return Err(image_only()); } if let Some(inline) = InlineDocument::parse(image_url)? { if !inline.mime_type().type_.eq_ignore_ascii_case("image") { - return Err(Error::CohereImageOnly); + return Err(image_only()); } inline.decode(OCR_INLINE_MAX_BYTES)?; } @@ -260,7 +260,7 @@ pub fn normalize_response( fn image_url(document: OcrDocument) -> Result { validate_document(&document)?; let OcrDocument::ImageUrl { image_url, .. } = document else { - return Err(Error::CohereImageOnly); + return Err(image_only()); }; Ok(image_url) } @@ -329,9 +329,17 @@ fn billed_pages(response: &CohereResponse) -> Option { response.meta.as_ref()?.billed_units.as_ref()?.pages } +fn image_only() -> Error { + Error::UnsupportedDocument { + provider: "Cohere Parse", + guidance: "it only accepts `image_url` documents, not document_url or PDF inputs", + } +} + fn invalid_api_base() -> Error { Error::RequestField { path: "api_base".into(), + cause: None, } } @@ -368,11 +376,14 @@ mod tests { #[test] fn options_reject_invalid_output_format() { let invalid = serde_json::from_value(json!({"output_format":"html"})).unwrap(); - assert!(matches!( - CohereParseConfig.map_ocr_params(&invalid, "parse"), - Err(Error::RequestField { path }) - if path == "optional_params.output_format" - )); + let Err(error) = CohereParseConfig.map_ocr_params(&invalid, "parse") else { + panic!("html is not a Cohere output format"); + }; + assert!(error.is_request()); + let argument = + litellm_error::find::(&error) + .expect("the argument failure is kept"); + assert_eq!(argument.path, "output_format"); } #[test] @@ -432,7 +443,7 @@ mod tests { .unwrap(); assert!(matches!( normalize_response("parse", response).unwrap_err(), - Error::ResponseField { path } + Error::ResponseField { path, .. } if path == "pages[0].markdown.images[0].image_base64" )); } @@ -657,7 +668,10 @@ mod tests { fn request_requires_image(#[case] value: Value) { assert!(matches!( validate_document(&serde_json::from_value(value).unwrap()), - Err(Error::CohereImageOnly) + Err(Error::UnsupportedDocument { + provider: "Cohere Parse", + .. + }) )); } diff --git a/litellm-rust/crates/llms/src/error.rs b/litellm-rust/crates/llms/src/error.rs index 10a4f696406..eddcb27d1d9 100644 --- a/litellm-rust/crates/llms/src/error.rs +++ b/litellm-rust/crates/llms/src/error.rs @@ -1,3 +1,5 @@ +use litellm_error::Source; + #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub enum Error { #[error("expected {expected}, got {actual}")] @@ -7,31 +9,117 @@ pub enum Error { }, #[error("missing required field: {0}")] MissingField(&'static str), - #[error("invalid request: {0}")] + #[error("invalid request")] InvalidRequest(#[source] ErrorDetail), - #[error("invalid response: {0}")] + #[error("invalid response")] InvalidResponse(#[source] ErrorDetail), #[error("unsupported: {0}")] Unsupported(&'static str), #[error(transparent)] Auth(#[from] litellm_auth::Error), + #[error(transparent)] + Http(#[from] litellm_http::Error), +} + +impl Error { + pub fn is_request(&self) -> bool { + match self { + Self::InvalidType { .. } + | Self::MissingField(_) + | Self::InvalidRequest(_) + | Self::Unsupported(_) => true, + Self::Auth(error) => error.is_request(), + Self::Http(error) => error.is_request(), + Self::InvalidResponse(_) => false, + } + } + + pub fn is_response(&self) -> bool { + match self { + Self::InvalidResponse(_) => true, + Self::Http(error) => error.is_response(), + Self::InvalidType { .. } + | Self::MissingField(_) + | Self::InvalidRequest(_) + | Self::Unsupported(_) + | Self::Auth(_) => false, + } + } + + pub fn is_missing_credential(&self) -> bool { + match self { + Self::Auth(error) => error.is_missing_credential(), + Self::InvalidType { .. } + | Self::MissingField(_) + | Self::InvalidRequest(_) + | Self::InvalidResponse(_) + | Self::Unsupported(_) + | Self::Http(_) => false, + } + } + + /// A transformation runs on both sides of the send, so only the transport and credential + /// resolution, which always runs before it, can tell whether the request reached the provider. + pub fn dispatched(&self) -> bool { + match self { + Self::Http(error) => error.dispatched(), + Self::Auth(_) => false, + Self::InvalidType { .. } + | Self::MissingField(_) + | Self::InvalidRequest(_) + | Self::InvalidResponse(_) + | Self::Unsupported(_) => true, + } + } + + /// Why the transformation rejected the caller's request, rendered as Python's own + /// validation renders it, without this layer's "invalid request" prefix. + pub fn rejected_request(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + Self::InvalidRequest(detail) => Some(detail), + Self::InvalidType { .. } | Self::MissingField(_) => Some(self), + Self::InvalidResponse(_) | Self::Unsupported(_) | Self::Auth(_) | Self::Http(_) => None, + } + } + + pub fn upstream_response(&self) -> Option<(u16, &str)> { + match self { + Self::Http(error) => error.upstream_response(), + Self::InvalidType { .. } + | Self::MissingField(_) + | Self::InvalidRequest(_) + | Self::InvalidResponse(_) + | Self::Unsupported(_) + | Self::Auth(_) => None, + } + } } #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub enum ErrorDetail { - #[error("{0}")] - Message(String), - #[error("invalid {subject}: {source}")] + #[error("invalid {subject}")] Invalid { subject: &'static str, #[source] - source: ErrorSource, + cause: Source, }, - #[error("{operation} failed: {source}")] + #[error("{operation} failed")] Failed { operation: &'static str, #[source] - source: ErrorSource, + cause: Source, + }, + #[error("{route} requires at least one message")] + NoMessages { route: &'static str }, + #[error("{subject} must be {expected}")] + Expected { + subject: &'static str, + expected: &'static str, + }, + #[error("{subject} is missing {field}")] + MissingField { + subject: &'static str, + field: &'static str, }, #[error("{field} must be {expected}, got {actual}")] InvalidValue { @@ -60,12 +148,18 @@ pub enum ErrorDetail { value: String, hint: String, }, - #[error("invalid {subject} on line {line}: {source}")] + #[error( + "Missing API Base - A call is being made to LLM Provider but no api base is set either in the environment variables ({{LLM_PROVIDER}}_API_KEY) or via params" + )] + MissingApiBase, + #[error("the host call raised an exception")] + HostCall, + #[error("invalid {subject} on line {line}")] InvalidLine { subject: &'static str, line: usize, #[source] - source: ErrorSource, + cause: Source, }, #[error("{operation} failed: {detail}")] RemoteFailure { @@ -77,64 +171,33 @@ pub enum ErrorDetail { impl ErrorDetail { pub fn invalid( subject: &'static str, - source: impl std::error::Error + Send + Sync + 'static, + cause: impl std::error::Error + Send + Sync + 'static, ) -> Self { Self::Invalid { subject, - source: ErrorSource::new(source), + cause: Source::new(cause), } } pub fn failed( operation: &'static str, - source: impl std::error::Error + Send + Sync + 'static, + cause: impl std::error::Error + Send + Sync + 'static, ) -> Self { Self::Failed { operation, - source: ErrorSource::new(source), + cause: Source::new(cause), + } + } + + pub fn invalid_line( + subject: &'static str, + line: usize, + cause: impl std::error::Error + Send + Sync + 'static, + ) -> Self { + Self::InvalidLine { + subject, + line, + cause: Source::new(cause), } } } - -impl From for ErrorDetail { - fn from(message: String) -> Self { - Self::Message(message) - } -} - -impl From<&str> for ErrorDetail { - fn from(message: &str) -> Self { - Self::Message(message.into()) - } -} - -#[derive(Clone, Debug)] -pub struct ErrorSource(std::sync::Arc); - -impl std::ops::Deref for ErrorSource { - type Target = dyn std::error::Error + Send + Sync; - - fn deref(&self) -> &Self::Target { - self.0.as_ref() - } -} - -impl std::fmt::Display for ErrorSource { - fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - std::fmt::Display::fmt(&self.0, formatter) - } -} - -impl ErrorSource { - pub fn new(error: impl std::error::Error + Send + Sync + 'static) -> Self { - Self(std::sync::Arc::new(error)) - } -} - -impl PartialEq for ErrorSource { - fn eq(&self, other: &Self) -> bool { - std::sync::Arc::ptr_eq(&self.0, &other.0) - } -} - -impl Eq for ErrorSource {} diff --git a/litellm-rust/crates/llms/src/lib.rs b/litellm-rust/crates/llms/src/lib.rs index eb232404d93..68e6e4010c2 100644 --- a/litellm-rust/crates/llms/src/lib.rs +++ b/litellm-rust/crates/llms/src/lib.rs @@ -17,4 +17,4 @@ pub mod reducto; pub mod tencent; pub mod vertex_ai; -pub use error::{Error, ErrorDetail, ErrorSource}; +pub use error::{Error, ErrorDetail}; diff --git a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs index 29d4d5c1610..fa4e2254e86 100644 --- a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs @@ -171,9 +171,7 @@ impl MistralOcrConfig { ApiUrl::parse(base) .and_then(|url| url.complete_path(&["v1", "ocr"])) .map(|url| url.into_string()) - .map_err(|_| Error::RequestField { - path: "api_base".into(), - }) + .map_err(|error| Error::request_field("api_base", error)) } } @@ -186,6 +184,7 @@ pub fn normalize_response( Some(None) => { return Err(Error::ResponseField { path: "model".into(), + cause: None, }); } None => model.to_string(), @@ -231,7 +230,7 @@ mod tests { let response = serde_json::from_value(json!({"model":null})).unwrap(); assert!(matches!( normalize_response("fallback", response).unwrap_err(), - Error::ResponseField { path } if path == "model" + Error::ResponseField { path, .. } if path == "model" )); } @@ -263,7 +262,7 @@ mod tests { .unwrap_err(); assert!(matches!( error, - Error::ResponseField { path: actual } if actual == path + Error::ResponseField { path: actual, .. } if actual == path )); } diff --git a/litellm-rust/crates/llms/src/openai/responses/transformation.rs b/litellm-rust/crates/llms/src/openai/responses/transformation.rs index 99cc3559514..cf74871894b 100644 --- a/litellm-rust/crates/llms/src/openai/responses/transformation.rs +++ b/litellm-rust/crates/llms/src/openai/responses/transformation.rs @@ -6,7 +6,7 @@ use serde_json::{Map, Value}; use litellm_auth::{CredentialPlacement, SecretValue}; use crate::{ - Error, + Error, ErrorDetail, base_llm::{ auth::{AuthScheme, ValidatedEnvironment}, responses::transformation::{ @@ -180,15 +180,19 @@ pub fn responses_request( params: Map, ) -> Result { if !input.is_string() && !input.is_array() { - return Err(Error::InvalidRequest( - "responses input must be a string or an array".into(), - )); + return Err(Error::InvalidRequest(ErrorDetail::Expected { + subject: "responses input", + expected: "a string or an array", + })); } if params .get("stream") .is_some_and(|stream| !stream.is_boolean()) { - return Err(Error::InvalidRequest("stream must be a boolean".into())); + return Err(Error::InvalidRequest(ErrorDetail::Expected { + subject: "stream", + expected: "a boolean", + })); } Ok(Value::Object( params @@ -202,7 +206,9 @@ pub fn responses_request( } pub fn responses_response(body: Value) -> Result { - serde_json::from_value(body).map_err(|error| Error::InvalidResponse(error.to_string().into())) + serde_json::from_value(body).map_err(|error| { + Error::InvalidResponse(ErrorDetail::invalid("Responses API response", error)) + }) } #[cfg(test)] diff --git a/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs b/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs index 0a89adf90f7..9f208db2471 100644 --- a/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs @@ -123,9 +123,12 @@ impl BaseConfig for OpenAILikeChatConfig { ) -> Result { let mut body = response.body; sanitize_usage(&mut body); - let body = body - .as_object() - .ok_or_else(|| Error::InvalidResponse("chat response is not an object".into()))?; + let body = + body.as_object() + .ok_or(Error::InvalidResponse(crate::ErrorDetail::Expected { + subject: "chat response", + expected: "an object", + }))?; let choices = body .get("choices") @@ -414,9 +417,13 @@ mod tests { #[rstest] fn a_missing_api_base_is_an_error() { assert!(matches!( - OPENAI_LIKE_CHAT_COMPLETIONS_CONFIG - .get_complete_url(None, "my-model", ¶ms(json!({})), &no_env), - Err(Error::InvalidRequest(message)) if message.to_string().starts_with("Missing API Base") + OPENAI_LIKE_CHAT_COMPLETIONS_CONFIG.get_complete_url( + None, + "my-model", + ¶ms(json!({})), + &no_env + ), + Err(Error::InvalidRequest(crate::ErrorDetail::MissingApiBase)) )); } diff --git a/litellm-rust/crates/llms/src/openai_like/common_utils.rs b/litellm-rust/crates/llms/src/openai_like/common_utils.rs index b0f5adafe00..d798f851d79 100644 --- a/litellm-rust/crates/llms/src/openai_like/common_utils.rs +++ b/litellm-rust/crates/llms/src/openai_like/common_utils.rs @@ -1,7 +1,7 @@ //! Shared OpenAI-like credential and endpoint resolution, mirroring //! `litellm/llms/openai_like/common_utils.py`. -use crate::Error; +use crate::{Error, ErrorDetail}; /// `OpenAILikeChatConfig._get_openai_compatible_provider_info`: the deployment's /// `api_base` wins over `OPENAI_LIKE_API_BASE`, and the deployment key over @@ -31,12 +31,7 @@ pub fn complete_openai_like_url( env_lookup: &dyn Fn(&str) -> Option, ) -> Result { let (api_base, _) = openai_compatible_provider_info(api_base, None, env_lookup); - let api_base = api_base.ok_or_else(|| { - Error::InvalidRequest( - "Missing API Base - A call is being made to LLM Provider but no api base is set either in the environment variables ({LLM_PROVIDER}_API_KEY) or via params" - .into(), - ) - })?; + let api_base = api_base.ok_or(Error::InvalidRequest(ErrorDetail::MissingApiBase))?; if custom_endpoint { return Ok(api_base); } diff --git a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs index fc73ce7c014..bf0ddc1c7a8 100644 --- a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs @@ -299,9 +299,16 @@ async fn prepare_upload_request Error { + Error::UnsupportedDocument { + provider: "Reducto", + guidance: "pass a reducto:// id or a data URI, or upload the file first for plain HTTP URLs", + } +} + fn uploaded_file_id(document: OcrDocument) -> Result { if !document.source().starts_with(REDUCTO_ID_PREFIX) { - return Err(Error::ReductoSource); + return Err(unsupported_source()); } if document.source()[REDUCTO_ID_PREFIX.len()..] .trim() @@ -309,6 +316,7 @@ fn uploaded_file_id(document: OcrDocument) -> Result { { return Err(Error::RequestField { path: "document file id".into(), + cause: None, }); } Ok(ReductoFileId(document.source().into())) @@ -392,6 +400,7 @@ fn build_pages_from_reducto(chunks: Vec) -> Result, E Some(Value::String(content)) => Ok(Some(content.as_str())), Some(_) => Err(Error::ResponseField { path: "result.chunks.blocks.content".into(), + cause: None, }), }) .collect::, _>>()?; @@ -436,9 +445,7 @@ fn complete_endpoint_url(api_base: Option<&str>, path: &str) -> Result( response, false, @@ -546,6 +555,7 @@ async fn upload_bytes_async( let Some(file_id) = file_id else { return Err(Error::ResponseField { path: "file_id".into(), + cause: None, }); }; Ok(ReductoFileId(file_id.to_string())) @@ -553,6 +563,8 @@ async fn upload_bytes_async( #[cfg(test)] mod tests { + use rstest::rstest; + use super::*; #[tokio::test] @@ -625,15 +637,16 @@ mod tests { ); } + #[rstest] + #[case::fractional_pages(json!({"num_pages":1.5}))] + #[case::non_numeric_pages(json!({"num_pages":[]}))] + #[case::non_numeric_credits(json!({"credits":{}}))] + fn invalid_usage_is_rejected(#[case] usage: Value) { + assert!(serde_json::from_value::(json!({"usage":usage})).is_err()); + } + #[test] - fn usage_uses_shared_validation_while_block_page_numbers_are_best_effort() { - for usage in [ - json!({"num_pages":1.5}), - json!({"num_pages":[]}), - json!({"credits":{}}), - ] { - assert!(serde_json::from_value::(json!({"usage":usage})).is_err()); - } + fn block_page_numbers_are_best_effort() { let response = serde_json::from_value(json!({"result":{"chunks":[{"blocks":[ {"content":"ignored", "bbox":{"page":"invalid"}}, {"content":"kept", "bbox":{"page":2.5}, "extra":null} diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/common_utils.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/common_utils.rs index 810462c57f8..d760da9073e 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/common_utils.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/common_utils.rs @@ -16,7 +16,9 @@ pub(super) fn vertex_config(request: &PreparedOcrRequest) -> Result Result<(), Error> { if connection.api_base.is_some() && connection.api_base_source == InputSource::Request { return Err(litellm_auth::Error::InvalidConfiguration( - "credentials cannot be sent to a request-controlled Vertex AI endpoint".into(), + litellm_auth::ErrorDetail::RequestControlledEndpoint { + provider: "Vertex AI", + }, ) .into()); } diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs index 67f9122ed67..5d680588612 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs @@ -235,8 +235,9 @@ pub fn normalize_response( &mut output, PythonJsonFormatter, )) - .map_err(|_| response_field("content"))?; - String::from_utf8(output).map_err(|_| response_field("content"))? + .map_err(|error| Error::response_field(content_path("content"), error))?; + String::from_utf8(output) + .map_err(|error| Error::response_field(content_path("content"), error))? }; (data, fallback) } @@ -261,7 +262,12 @@ pub fn normalize_response( }) }) .collect::, Error>>()?, - Some(_) => return Err(response_field("pages")), + Some(_) => { + return Err(Error::ResponseField { + path: content_path("pages"), + cause: None, + }); + } None => Vec::new(), }; let usage = ocr_data @@ -273,7 +279,12 @@ pub fn normalize_response( .transpose()?; let model = match ocr_data.get("model") { Some(Value::String(model)) => model.clone(), - Some(_) => return Err(response_field("model")), + Some(_) => { + return Err(Error::ResponseField { + path: content_path("model"), + cause: None, + }); + } None => model.to_string(), }; Ok(LiteLLMOcrResponse { @@ -369,10 +380,8 @@ impl serde_json::ser::Formatter for PythonJsonFormatter { } } -fn response_field(field: &str) -> Error { - Error::ResponseField { - path: format!("choices[0].message.content.{field}"), - } +fn content_path(field: &str) -> String { + format!("choices[0].message.content.{field}") } pub fn provider_model(model: &str) -> Result { @@ -380,6 +389,7 @@ pub fn provider_model(model: &str) -> Result { if local_model.is_empty() { return Err(Error::RequestField { path: "model".into(), + cause: None, }); } Ok(format!("{MODEL_PREFIX}{local_model}")) @@ -411,9 +421,7 @@ impl VertexAIDeepSeekOCRConfig { ]) }) .map(|url| url.into_string()) - .map_err(|_| Error::RequestField { - path: "api_base".into(), - }) + .map_err(|error| Error::request_field("api_base", error)) } } @@ -625,20 +633,18 @@ mod tests { assert_eq!(result["future"], "kept"); } - #[test] - fn response_transform_rejects_missing_empty_and_malformed_content() { - for value in [ - json!({"choices":[]}), - json!({"choices":[{"message":{"content":{}}}]}), - json!({"choices":[{"message":{"content":""}}]}), - json!({"choices":[{"message":{"content":"{\"pages\":[{\"markdown\":42}]}"}}]}), - json!({"choices":[{"message":{"content":{"pages":[{"markdown":42}]}}}]}), - ] { - let result = serde_json::from_value::(value) - .map_err(|_| ()) - .and_then(|response| normalize_response("model", response).map_err(|_| ())); - assert!(result.is_err()); - } + #[rstest::rstest] + #[case::no_choices(json!({"choices":[]}))] + #[case::empty_object(json!({"choices":[{"message":{"content":{}}}]}))] + #[case::empty_string(json!({"choices":[{"message":{"content":""}}]}))] + #[case::malformed_string(json!({"choices":[{"message":{"content":"{\"pages\":[{\"markdown\":42}]}"}}]}))] + #[case::malformed_object(json!({"choices":[{"message":{"content":{"pages":[{"markdown":42}]}}}]}))] + fn response_transform_rejects_missing_empty_and_malformed_content(#[case] value: Value) { + let rejected = match serde_json::from_value::(value) { + Ok(response) => normalize_response("model", response).is_err(), + Err(_) => true, + }; + assert!(rejected); } #[test] diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs index ab28ee8bff4..77881cd2e78 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs @@ -174,9 +174,7 @@ impl VertexAiOcrConfig { ]) }) .map(|url| url.into_string()) - .map_err(|_| Error::RequestField { - path: "api_base".into(), - }) + .map_err(|error| Error::request_field("api_base", error)) } } @@ -198,6 +196,7 @@ fn validate_location(location: &str) -> Result<(), Error> { } Err(Error::RequestField { path: "vertex_location".into(), + cause: None, }) } diff --git a/litellm-rust/crates/model-catalog/Cargo.toml b/litellm-rust/crates/model-catalog/Cargo.toml index 570e68ca4fd..2c2866d7090 100644 --- a/litellm-rust/crates/model-catalog/Cargo.toml +++ b/litellm-rust/crates/model-catalog/Cargo.toml @@ -20,4 +20,5 @@ time.workspace = true [dev-dependencies] jsonschema = { version = "0.55.1", default-features = false } +litellm-error.workspace = true rstest.workspace = true diff --git a/litellm-rust/crates/model-catalog/src/catalog.rs b/litellm-rust/crates/model-catalog/src/catalog.rs index 0b113436a5b..37412ef1e14 100644 --- a/litellm-rust/crates/model-catalog/src/catalog.rs +++ b/litellm-rust/crates/model-catalog/src/catalog.rs @@ -20,11 +20,15 @@ pub struct IntegrityLimits { pub min_reference_ratio: f64, } -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub enum AliasIssue { + #[error("model {model:?} has aliases that are not an array")] InvalidList { model: String }, + #[error("model {model:?} has an alias that is not a string")] InvalidName { model: String }, + #[error("alias {alias:?} of model {model:?} is already a model name")] CanonicalCollision { model: String, alias: String }, + #[error("alias {alias:?} of model {model:?} is already an alias of another model")] AliasCollision { model: String, alias: String }, } diff --git a/litellm-rust/crates/model-catalog/src/error.rs b/litellm-rust/crates/model-catalog/src/error.rs index edb6ba1eb11..00902858be4 100644 --- a/litellm-rust/crates/model-catalog/src/error.rs +++ b/litellm-rust/crates/model-catalog/src/error.rs @@ -2,26 +2,20 @@ use thiserror::Error; #[derive(Debug, Error)] pub enum Error { - /// The body is not valid JSON, or a model entry fails typed deserialization. - #[error("invalid JSON: {0}")] + #[error("invalid catalog JSON")] Json(#[from] serde_json::Error), - /// The catalog has no entries at all. #[error("catalog is empty")] Empty, - /// A non-reserved top level value is not a JSON object. #[error("model {model:?} must be an object")] EntryNotObject { model: String }, - /// Canonical entry count is under the configured minimum. #[error("catalog has {actual} models, below minimum {minimum}")] BelowMinimum { actual: usize, minimum: usize }, - /// Canonical entry count is under the configured reference ratio. #[error("catalog has {actual} models, below {ratio} of reference count {reference}")] Shrunk { actual: usize, reference: usize, ratio: f64, }, - /// The configured minimum reference ratio is not finite or outside `[0, 1]`. #[error("minimum reference ratio must be finite and between zero and one")] InvalidRatio, } diff --git a/litellm-rust/crates/model-catalog/src/validation.rs b/litellm-rust/crates/model-catalog/src/validation.rs index 73f08a865c1..6dd4174148c 100644 --- a/litellm-rust/crates/model-catalog/src/validation.rs +++ b/litellm-rust/crates/model-catalog/src/validation.rs @@ -5,24 +5,57 @@ use thiserror::Error; use crate::{AliasIssue, Catalog, ModelInfo, UtcHours, Weekday}; -/// A registry entry violates the checked-in catalog contract. #[derive(Debug, Error)] pub enum RegistryValidationError { - #[error("{reason}")] - Entry { model: String, reason: String }, - #[error("alias issue: {0:?}")] + #[error("invalid registry entry {model:?}")] + Entry { + model: String, + #[source] + reason: EntryViolation, + }, + #[error(transparent)] Alias(AliasIssue), } -/// Validate one registry entry without restricting the tolerant catalog reader. +#[derive(Debug, Error)] +pub enum EntryViolation { + #[error("entry must be an object")] + NotObject, + #[error("aliases must be an array")] + AliasesNotArray, + #[error("aliases must contain strings")] + AliasNotString, + #[error("entry does not match the model schema")] + Schema(#[source] serde_json::Error), + #[error("litellm_provider is required")] + MissingProvider, + #[error("invalid deprecation_date")] + DeprecationDate(#[source] time::error::Parse), + #[error("off_peak_pricing needs hours or windows")] + OffPeakWithoutHours, + #[error("off_peak_pricing.hours_utc is invalid")] + InvalidHours, + #[error("off_peak_pricing.windows is empty")] + EmptyWindows, + #[error("off_peak_pricing.weekdays is invalid")] + InvalidWeekdays, + #[error("typed entry could not be serialized")] + Serialize(#[source] serde_json::Error), + #[error("entry has an unknown field, null, or changed value: {}", fields.join(", "))] + ChangedFields { fields: Vec }, + #[error("{field} must be nonnegative")] + Negative { field: String }, + #[error("{field} must be at least one")] + UpliftBelowOne { field: String }, +} + pub fn validate_model_entry(model: &str, value: &Value) -> Result<(), RegistryValidationError> { - validate_entry_inner(model, value).map_err(|reason| RegistryValidationError::Entry { + validate_entry_inner(value).map_err(|reason| RegistryValidationError::Entry { model: model.to_owned(), reason, }) } -/// Check every model and alias in a parsed catalog against registry rules. pub fn validate_registry(catalog: &Catalog) -> Result<(), RegistryValidationError> { if let Some(issue) = catalog.alias_issues().first() { return Err(RegistryValidationError::Alias(issue.clone())); @@ -53,70 +86,61 @@ fn keys(value: &Map) -> BTreeSet { value.keys().cloned().collect() } -fn symmetric_difference(left: &BTreeSet, right: &BTreeSet) -> BTreeSet { +fn symmetric_difference(left: &BTreeSet, right: &BTreeSet) -> Vec { left.symmetric_difference(right).cloned().collect() } -fn validate_entry_inner(model_name: &str, value: &Value) -> Result<(), String> { - let object = value - .as_object() - .ok_or_else(|| format!("{model_name} must be an object"))?; +fn validate_entry_inner(value: &Value) -> Result<(), EntryViolation> { + let object = value.as_object().ok_or(EntryViolation::NotObject)?; if let Some(aliases) = object.get("aliases") { - let names = aliases - .as_array() - .ok_or_else(|| format!("{model_name}.aliases must be an array"))?; + let names = aliases.as_array().ok_or(EntryViolation::AliasesNotArray)?; if names.iter().any(|name| !name.is_string()) { - return Err(format!("{model_name}.aliases must contain strings")); + return Err(EntryViolation::AliasNotString); } } - let info: ModelInfo = - serde_json::from_value(value.clone()).map_err(|error| format!("{model_name}: {error}"))?; + let info: ModelInfo = serde_json::from_value(value.clone()).map_err(EntryViolation::Schema)?; if info.litellm_provider.is_none() { - return Err(format!("{model_name}.litellm_provider is required")); + return Err(EntryViolation::MissingProvider); } - validate_dates_and_windows(model_name, &info)?; - let serialized = serde_json::to_value(info).map_err(|error| error.to_string())?; + validate_dates_and_windows(&info)?; + let serialized = serde_json::to_value(info).map_err(EntryViolation::Serialize)?; let mut expected = object.clone(); expected.remove("aliases"); if !json_eq(&Value::Object(expected.clone()), &serialized) { let actual = serialized .as_object() .expect("ModelInfo serializes as an object"); - return Err(format!( - "{model_name} has an unknown field, null, or changed value: {:?}", - symmetric_difference(&keys(&expected), &keys(actual)) - )); + return Err(EntryViolation::ChangedFields { + fields: symmetric_difference(&keys(&expected), &keys(actual)), + }); } - check_prices(model_name, value) + check_prices("", value) } -fn validate_dates_and_windows(model_name: &str, info: &ModelInfo) -> Result<(), String> { +fn validate_dates_and_windows(info: &ModelInfo) -> Result<(), EntryViolation> { if let Some(date) = &info.deprecation_date { let format = time::format_description::parse_borrowed::<2>("[year]-[month]-[day]").unwrap(); - time::Date::parse(date, &format) - .map_err(|error| format!("{model_name}.deprecation_date: {error}"))?; + time::Date::parse(date, &format).map_err(EntryViolation::DeprecationDate)?; } let Some(pricing) = &info.off_peak_pricing else { return Ok(()); }; if pricing.hours_utc.is_none() && pricing.windows.is_none() { - return Err(format!( - "{model_name}.off_peak_pricing needs hours or windows" - )); + return Err(EntryViolation::OffPeakWithoutHours); } if let Some(hours) = &pricing.hours_utc { validate_hours(hours)?; } if let Some(windows) = &pricing.windows { if windows.is_empty() { - return Err(format!("{model_name}.off_peak_pricing.windows is empty")); + return Err(EntryViolation::EmptyWindows); } windows.iter().try_for_each(|window| { validate_hours(&window.hours_utc)?; if let Some(days) = &window.weekdays && (days.is_empty() || days.iter().any(|day| !valid_weekday(day))) { - return Err(format!("{model_name}.off_peak_pricing.weekdays is invalid")); + return Err(EntryViolation::InvalidWeekdays); } Ok(()) })?; @@ -124,13 +148,13 @@ fn validate_dates_and_windows(model_name: &str, info: &ModelInfo) -> Result<(), Ok(()) } -fn validate_hours(hours: &UtcHours) -> Result<(), String> { +fn validate_hours(hours: &UtcHours) -> Result<(), EntryViolation> { let values = match hours { UtcHours::Single(value) => std::slice::from_ref(value), UtcHours::Multiple(values) => values.as_slice(), }; if values.is_empty() || values.iter().any(|value| !valid_utc_window(value)) { - return Err("off_peak_pricing.hours_utc is invalid".into()); + return Err(EntryViolation::InvalidHours); } Ok(()) } @@ -176,7 +200,7 @@ fn valid_weekday(day: &Weekday) -> bool { } } -fn check_prices(path: &str, value: &Value) -> Result<(), String> { +fn check_prices(path: &str, value: &Value) -> Result<(), EntryViolation> { let Some(object) = value.as_object() else { return Ok(()); }; @@ -188,13 +212,17 @@ fn check_prices(path: &str, value: &Value) -> Result<(), String> { && let Some(number) = field.as_f64() && number < 0.0 { - return Err(format!("{field_path} must be nonnegative")); + return Err(EntryViolation::Negative { + field: relative(&field_path), + }); } if key.contains("uplift_multiplier") && let Some(number) = field.as_f64() && number < 1.0 { - return Err(format!("{field_path} must be at least one")); + return Err(EntryViolation::UpliftBelowOne { + field: relative(&field_path), + }); } if matches!(key.as_str(), "range" | "max_results_range") && field.as_array().is_some_and(|values| { @@ -203,7 +231,9 @@ fn check_prices(path: &str, value: &Value) -> Result<(), String> { .any(|value| value.as_f64().is_some_and(|n| n < 0.0)) }) { - return Err(format!("{field_path} must be nonnegative")); + return Err(EntryViolation::Negative { + field: relative(&field_path), + }); } if matches!(key.as_str(), "metadata" | "provider_specific_entry") { return Ok(()); @@ -216,3 +246,7 @@ fn check_prices(path: &str, value: &Value) -> Result<(), String> { } }) } + +fn relative(field_path: &str) -> String { + field_path.trim_start_matches('.').to_owned() +} diff --git a/litellm-rust/crates/model-catalog/tests/registry_validation.rs b/litellm-rust/crates/model-catalog/tests/registry_validation.rs index 67589010aba..3b65ef13d78 100644 --- a/litellm-rust/crates/model-catalog/tests/registry_validation.rs +++ b/litellm-rust/crates/model-catalog/tests/registry_validation.rs @@ -52,9 +52,9 @@ fn fallback_generalizations_are_typed(repo_root: PathBuf) { #[case::null_aliases(serde_json::json!({"litellm_provider": "test", "aliases": null}), "aliases must be an array")] #[ignore] fn registry_validation_rejects_malformed_entries(#[case] entry: Value, #[case] expected: &str) { + let error = validate_model_entry("test", &entry).unwrap_err(); assert!( - validate_model_entry("test", &entry) - .unwrap_err() + litellm_error::Report::new(&error) .to_string() .contains(expected) ); diff --git a/litellm-rust/crates/python-bridge/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index c435ba37e78..d90fe073b20 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -111,8 +111,11 @@ GIL handling to `litellm-host-python`. - Avoid copying large payloads more than needed. The current JSON round-trip is acceptable for the first scaffold, but future performance work should evaluate direct PyO3 conversion before expanding Rust coverage to image-heavy paths. -- Do not expose raw Rust errors that include document contents or upstream - bodies. +- Render a Rust error for Python only in its mapping function, with + `litellm_error::Report`, never `Debug` or a document's contents. An upstream + body reaches Python intact as the `RustUpstreamError` args for provider error + mapping. The 256-character cap applies only to HTTP error diagnostics, not + response bodies. Do not log those args. ## Tests diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 8afde8362f6..9d69cd8c14b 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -27,6 +27,7 @@ litellm-host.workspace = true bytes.workspace = true futures-util.workspace = true litellm-cache.workspace = true +litellm-error.workspace = true litellm-cache-memory.workspace = true litellm-cache-redis.workspace = true litellm-cache-response.workspace = true @@ -70,6 +71,7 @@ criterion.workspace = true futures-util.workspace = true rstest.workspace = true sqlx = { workspace = true, features = ["migrate"] } +tempfile.workspace = true tokio-tungstenite.workspace = true wiremock.workspace = true aws-sdk-secretsmanager = "1.117.0" diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index e190b12e665..87dba3efcb1 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -12,10 +12,93 @@ use pyo3::{ prelude::*, }; -fn cache_error(error: Error) -> PyErr { - match error { - Error::InvalidEntry => PyValueError::new_err(error.to_string()), - Error::UnsupportedOperation => PyNotImplementedError::new_err(error.to_string()), - _ => PyRuntimeError::new_err(error.to_string()), +/// A cache call the bridge rejects itself, before or after the cache ran. +#[derive(Debug, thiserror::Error)] +pub(crate) enum HandleError { + #[error("recreate the v2 cache after fork")] + Forked, + #[error("cache limits must be positive")] + Limits, + #[error("namespace and a positive cache limit are required")] + RedisLimits, + #[error("cache durations must be finite and positive")] + Duration(#[source] Option), + #[error("unsupported v2 cache control: {0}")] + Control(String), + #[error("native inference needs a v2 cache, the catalog keeps Python caches on Python")] + NotNative, + #[error("cache operation without configured cache")] + Unconfigured, + #[error("cache arguments unavailable")] + MissingArguments, + #[error("cache reply without pending operation")] + NothingPending, +} + +impl From for PyErr { + fn from(error: HandleError) -> Self { + let message = litellm_error::Report::new(&error).to_string(); + match error { + HandleError::Limits + | HandleError::RedisLimits + | HandleError::Duration(_) + | HandleError::Control(_) => PyValueError::new_err(message), + HandleError::Forked + | HandleError::NotNative + | HandleError::Unconfigured + | HandleError::MissingArguments + | HandleError::NothingPending => PyRuntimeError::new_err(message), + } + } +} + +fn cache_error(error: Error) -> PyErr { + let message = litellm_error::Report::new(&error).to_string(); + match error { + Error::InvalidEntry(_) => PyValueError::new_err(message), + Error::UnsupportedOperation => PyNotImplementedError::new_err(message), + Error::Unavailable(_) | Error::UnscopedFlush | Error::MissingPrompt => { + PyRuntimeError::new_err(message) + } + } +} + +#[cfg(test)] +mod tests { + use pyo3::{ + exceptions::{PyRuntimeError, PyValueError}, + prelude::*, + }; + use rstest::rstest; + + use super::{Error, cache_error}; + + fn refused() -> std::io::Error { + std::io::Error::other("connection refused") + } + + #[rstest] + #[case::unavailable( + Error::unavailable(refused()), + false, + "cache is unavailable: connection refused" + )] + #[case::invalid_entry( + Error::invalid_entry(refused()), + true, + "invalid cache entry: connection refused" + )] + fn python_sees_the_class_and_the_whole_chain( + #[case] error: Error, + #[case] caller_fault: bool, + #[case] message: &str, + ) { + Python::initialize(); + Python::attach(|py| { + let failure = cache_error(error); + assert_eq!(failure.is_instance_of::(py), caller_fault); + assert_eq!(failure.is_instance_of::(py), !caller_fault); + assert_eq!(failure.value(py).to_string(), message); + }); } } diff --git a/litellm-rust/crates/python-bridge/src/cache/native/v2.rs b/litellm-rust/crates/python-bridge/src/cache/native/v2.rs index eb09b8131c3..c54cb641843 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native/v2.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native/v2.rs @@ -1,4 +1,4 @@ -use crate::cache::cache_error; +use crate::cache::{HandleError, cache_error}; use std::{ sync::Arc, time::{Duration, SystemTime, UNIX_EPOCH}, @@ -14,7 +14,7 @@ use litellm_cache_response::{ CacheEntry, CacheKeyInput, ExactResponseCache, ResponseCache, ResponseCacheCodec, ResponseCacheConfig, ResponseCacheRequest, ResponseCacheService, }; -use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; +use pyo3::{prelude::*, types::PyDict}; #[pyclass( frozen, @@ -37,9 +37,7 @@ enum Storage { impl NativeCacheHandle { fn check_process(&self) -> PyResult<()> { if self.pid != std::process::id() { - return Err(pyo3::exceptions::PyRuntimeError::new_err( - "recreate the v2 cache after fork", - )); + return Err(HandleError::Forked.into()); } Ok(()) } @@ -61,7 +59,7 @@ impl NativeCacheHandle { fn memory(ttl: f64, capacity: usize, max_entry_bytes: usize) -> PyResult { let ttl = duration(ttl)?; if capacity == 0 || max_entry_bytes == 0 { - return Err(PyValueError::new_err("cache limits must be positive")); + return Err(HandleError::Limits.into()); } let storage = Arc::new(InMemoryCache::new(Some(capacity), Some(ttl))); let backend = Arc::new(ResponseCache::new(storage.clone()).with_config( @@ -89,9 +87,7 @@ impl NativeCacheHandle { ) -> PyResult { let ttl = duration(ttl)?; if namespace.is_empty() || max_entry_bytes == 0 { - return Err(PyValueError::new_err( - "namespace and a positive cache limit are required", - )); + return Err(HandleError::RedisLimits.into()); } let storage = Arc::new( release_gil(py, || { @@ -252,11 +248,12 @@ impl NativeCacheHandle { } } -fn duration(seconds: f64) -> PyResult { - Duration::try_from_secs_f64(seconds) - .ok() - .filter(|value| !value.is_zero()) - .ok_or_else(|| PyValueError::new_err("cache durations must be finite and positive")) +fn duration(seconds: f64) -> Result { + match Duration::try_from_secs_f64(seconds) { + Ok(duration) if !duration.is_zero() => Ok(duration), + Ok(_) => Err(HandleError::Duration(None)), + Err(error) => Err(HandleError::Duration(Some(error))), + } } pub(in crate::cache) fn native_handle<'py>( @@ -277,11 +274,7 @@ pub(in crate::cache) fn configured( Option>, litellm_cache_response::CacheOptions, )> { - let handle = native_handle(configured)?.ok_or_else(|| { - pyo3::exceptions::PyRuntimeError::new_err( - "native inference needs a v2 cache, the catalog keeps Python caches on Python", - ) - })?; + let handle = native_handle(configured)?.ok_or(HandleError::NotNative)?; let cache = handle.extract::>()?; cache.check_process()?; let controls = kwargs.get_item("cache")?.filter(|value| !value.is_none()); @@ -296,9 +289,7 @@ pub(in crate::cache) fn configured( name.as_str(), "no-cache" | "no-store" | "ttl" | "s-maxage" | "s-max-age" | "use-cache" ) { - return Err(PyValueError::new_err(format!( - "unsupported v2 cache control: {name}" - ))); + return Err(HandleError::Control(name).into()); } } } @@ -316,7 +307,7 @@ pub(in crate::cache) fn configured( .map(|values| values.get_item(name)) .transpose()? .flatten() - .map(|value| duration(value.extract()?)) + .map(|value| Ok(duration(value.extract()?)?)) .transpose() }; Ok(( diff --git a/litellm-rust/crates/python-bridge/src/cache/python/host.rs b/litellm-rust/crates/python-bridge/src/cache/python/host.rs index d8f8ac7c181..82b7d881df6 100644 --- a/litellm-rust/crates/python-bridge/src/cache/python/host.rs +++ b/litellm-rust/crates/python-bridge/src/cache/python/host.rs @@ -9,6 +9,7 @@ use pyo3::{ use serde_json::Value; use super::service::CacheCall; +use crate::cache::HandleError; enum Pending { Lookup(Reply, Error>>), @@ -43,16 +44,12 @@ impl PythonCache { pub fn begin(&mut self, py: Python<'_>, call: CacheCall) -> PyResult>> { let Some(cache) = self.cache.as_ref() else { - return Err(pyo3::exceptions::PyRuntimeError::new_err( - "cache operation without configured cache", - )); + return Err(HandleError::Unconfigured.into()); }; let arguments = self .arguments .as_ref() - .ok_or_else(|| { - pyo3::exceptions::PyRuntimeError::new_err("cache arguments unavailable") - })? + .ok_or(HandleError::MissingArguments)? .bind(py) .copy()?; let (method, result) = match call { @@ -97,33 +94,30 @@ impl PythonCache { py: Python<'_>, result: PyResult>, ) -> PyResult>> { - if let Err(error) = &result - && !error.is_instance_of::(py) - { - self.pending = None; - return Err(result.err().unwrap()); - } + let result = match result { + Err(error) if !error.is_instance_of::(py) => { + self.pending = None; + return Err(error); + } + result => result, + }; match self.pending.take() { Some(Pending::Lookup(reply)) => { - let value = result.map_err(|_| Error::Unavailable).and_then(|value| { + let value = result.map_err(Error::unavailable).and_then(|value| { if value.bind(py).is_none() { Ok(None) } else { from_py(value.bind(py)) .map(Some) - .map_err(|_| Error::InvalidEntry) + .map_err(Error::invalid_entry) } }); reply.send(value); } Some(Pending::Store(reply)) => { - reply.send(result.map(|_| ()).map_err(|_| Error::Unavailable)); - } - None => { - return Err(pyo3::exceptions::PyRuntimeError::new_err( - "cache reply without pending operation", - )); + reply.send(result.map(|_| ()).map_err(Error::unavailable)); } + None => return Err(HandleError::NothingPending.into()), } Ok(None) } diff --git a/litellm-rust/crates/python-bridge/src/cache/python/service.rs b/litellm-rust/crates/python-bridge/src/cache/python/service.rs index f215e5e97f5..38872b85d30 100644 --- a/litellm-rust/crates/python-bridge/src/cache/python/service.rs +++ b/litellm-rust/crates/python-bridge/src/cache/python/service.rs @@ -17,18 +17,21 @@ fn from_python(value: Value) -> Result { let output = match value.get(STREAM_EVENTS_KEY) { Some(events) => { let events: Vec = - serde_json::from_value(events.clone()).map_err(|_| Error::InvalidEntry)?; + serde_json::from_value(events.clone()).map_err(Error::invalid_entry)?; CachedOutput::Stream(events.concat()) } None => CachedOutput::Response(value), }; - serde_json::to_value(ResponseEnvelope::new("messages", output)).map_err(|_| Error::InvalidEntry) + serde_json::to_value(ResponseEnvelope::new("messages", output)).map_err(Error::invalid_entry) } fn to_python(value: Value) -> Result { let envelope: ResponseEnvelope> = - serde_json::from_value(value).map_err(|_| Error::InvalidEntry)?; - match envelope.decode("messages").ok_or(Error::InvalidEntry)? { + serde_json::from_value(value).map_err(Error::invalid_entry)?; + match envelope + .decode("messages") + .ok_or(Error::InvalidEntry(None))? + { CachedOutput::Response(response) => Ok(response), CachedOutput::Stream(text) => Ok(serde_json::json!({ STREAM_EVENTS_KEY: text.split_inclusive("\n\n").collect::>() @@ -56,7 +59,7 @@ pub(in crate::cache) fn service>( namespace: String, ) -> std::sync::Arc where - P::Error: From, + P::Error: From + std::error::Error, { std::sync::Arc::new(PythonCacheService { services, @@ -69,7 +72,7 @@ where impl> ResponseCacheService for PythonCacheService

where - P::Error: From, + P::Error: From + std::error::Error, { fn config(&self) -> &ResponseCacheConfig { &self.config @@ -84,7 +87,7 @@ where self.services .call(|reply| CacheCall::Lookup { reply }) .await - .map_err(|_| Error::Unavailable)?? + .map_err(Error::unavailable)?? .map(from_python) .transpose() }) @@ -101,7 +104,7 @@ where self.services .call(|reply| CacheCall::Store { value, reply }) .await - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? }) } } diff --git a/litellm-rust/crates/python-bridge/src/cache/selection.rs b/litellm-rust/crates/python-bridge/src/cache/selection.rs index cd19533fc73..00fc7cd822f 100644 --- a/litellm-rust/crates/python-bridge/src/cache/selection.rs +++ b/litellm-rust/crates/python-bridge/src/cache/selection.rs @@ -39,7 +39,7 @@ impl Selection { services: HostServices

, ) -> (Option, CacheOptions) where - P::Error: From, + P::Error: From + std::error::Error, { let service = match self.backend { Backend::Disabled => None, diff --git a/litellm-rust/crates/python-bridge/src/callable.rs b/litellm-rust/crates/python-bridge/src/callable.rs index d0ed76a3e8b..605b649e1f3 100644 --- a/litellm-rust/crates/python-bridge/src/callable.rs +++ b/litellm-rust/crates/python-bridge/src/callable.rs @@ -48,14 +48,8 @@ mod tests { Err(PyErr::from_value(error.clone())) } - #[rstest::rstest] - fn only_ordinary_exceptions_are_reported_under_the_template() { - Python::initialize(); - Python::attach(|py| { - let locals = PyDict::new(py); - py.run( - pyo3::ffi::c_str!( - r#" + const FAILURES: &std::ffi::CStr = pyo3::ffi::c_str!( + r#" class CallerError(Exception): def __format__(self, specification): return 'unavailable' @@ -63,12 +57,19 @@ ordinary = CallerError('must use __format__') type_error = TypeError('signature') abort = KeyboardInterrupt('cancelled') "# - ), - Some(&locals), - Some(&locals), - ) - .unwrap(); + ); + fn failures(py: Python<'_>) -> Bound<'_, PyDict> { + let locals = PyDict::new(py); + py.run(FAILURES, Some(&locals), Some(&locals)).unwrap(); + locals + } + + #[test] + fn ordinary_exceptions_are_reported_under_the_template() { + Python::initialize(); + Python::attach(|py| { + let locals = failures(py); let original = raised(&locals, "ordinary"); let wrapped = wrap_failure(py, TEMPLATE, failure(&original)).unwrap_err(); assert!(wrapped.is_instance_of::(py)); @@ -78,12 +79,19 @@ abort = KeyboardInterrupt('cancelled') wrapped.value(py).str().unwrap().to_str().unwrap(), "Failed to reach the caller: unavailable" ); + }); + } - for name in ["type_error", "abort"] { - let original = raised(&locals, name); - let error = wrap_failure(py, TEMPLATE, failure(&original)).unwrap_err(); - assert!(error.value(py).is(&original)); - } + #[rstest::rstest] + #[case::type_error("type_error")] + #[case::keyboard_interrupt("abort")] + fn other_exceptions_pass_through_unwrapped(#[case] name: &str) { + Python::initialize(); + Python::attach(|py| { + let locals = failures(py); + let original = raised(&locals, name); + let error = wrap_failure(py, TEMPLATE, failure(&original)).unwrap_err(); + assert!(error.value(py).is(&original)); }); } diff --git a/litellm-rust/crates/python-bridge/src/coercion.rs b/litellm-rust/crates/python-bridge/src/coercion.rs index 77c8cabc22b..8cdf1e6c5bc 100644 --- a/litellm-rust/crates/python-bridge/src/coercion.rs +++ b/litellm-rust/crates/python-bridge/src/coercion.rs @@ -1,4 +1,5 @@ use litellm_core_utils::serde_compat::parse_str_bool; +use litellm_error::Report; use pyo3::{ exceptions::{PyAttributeError, PyRuntimeError, PyValueError}, prelude::*, @@ -9,15 +10,20 @@ use pyo3::{ pub(crate) enum ProjectionError { Python(PyErr), InvalidConfiguration(ProjectionDetail), - UnsupportedLiveObject(String), + UnsupportedLiveObject(ProjectionDetail), InternalSchemaFailure(ProjectionDetail), } #[derive(Debug, thiserror::Error)] pub(crate) enum ProjectionDetail { - #[error("{0}")] - Message(String), - #[error("{field}: {source}")] + #[error("{group}.{name}: expected {expected}, got {actual}")] + Expected { + group: &'static str, + name: &'static str, + expected: &'static str, + actual: String, + }, + #[error("{field}")] InvalidField { field: &'static str, #[source] @@ -30,12 +36,6 @@ pub(crate) enum ProjectionDetail { }, } -impl From for ProjectionDetail { - fn from(message: String) -> Self { - Self::Message(message) - } -} - impl From for ProjectionError { fn from(error: PyErr) -> Self { Self::Python(error) @@ -46,12 +46,12 @@ impl From for PyErr { fn from(error: ProjectionError) -> Self { match error { ProjectionError::Python(error) => error, - ProjectionError::InvalidConfiguration(detail) => { - PyValueError::new_err(detail.to_string()) + ProjectionError::InvalidConfiguration(detail) + | ProjectionError::UnsupportedLiveObject(detail) => { + PyValueError::new_err(Report::new(&detail).to_string()) } - ProjectionError::UnsupportedLiveObject(message) => PyValueError::new_err(message), ProjectionError::InternalSchemaFailure(detail) => { - PyRuntimeError::new_err(detail.to_string()) + PyRuntimeError::new_err(Report::new(&detail).to_string()) } } } @@ -124,26 +124,26 @@ impl<'py> Field<'py> { && getter.is(object.getattr("__getattribute__")?)) } - pub(crate) fn path(&self) -> String { - format!("{}.{}", self.group, self.name) - } - /// A member of this field's collection, reported under the same path. pub(crate) fn member(&self, value: Bound<'py, PyAny>) -> Self { Self::new(self.group, self.name, value) } - pub(crate) fn expected(&self, expected: &str) -> Result { - Ok(format!( - "{}: expected {expected}, got {}", - self.path(), - self.value.get_type().name()? - )) + pub(crate) fn expected( + &self, + expected: &'static str, + ) -> Result { + Ok(ProjectionDetail::Expected { + group: self.group, + name: self.name, + expected, + actual: self.value.get_type().name()?.to_string(), + }) } - pub(crate) fn invalid(&self, expected: &str) -> ProjectionError { + pub(crate) fn invalid(&self, expected: &'static str) -> ProjectionError { match self.expected(expected) { - Ok(message) => ProjectionError::InvalidConfiguration(message.into()), + Ok(detail) => ProjectionError::InvalidConfiguration(detail), Err(error) => error, } } @@ -163,24 +163,23 @@ impl<'py> Field<'py> { pub(crate) fn schema_bool(&self) -> Result { if !self.value.is_instance_of::() { return Err(ProjectionError::InternalSchemaFailure( - self.expected("a Boolean")?.into(), + self.expected("a Boolean")?, )); } Ok(self.exact_true()) } pub(crate) fn strict_string(&self) -> Result { - let value = self - .value - .cast::() - .map_err(|_| self.invalid("a string"))?; + let Ok(value) = self.value.cast::() else { + return Err(self.invalid("a string")); + }; Ok(value.to_str()?.to_owned()) } pub(crate) fn schema_string(&self) -> Result { if !self.value.is_instance_of::() { return Err(ProjectionError::InternalSchemaFailure( - self.expected("a string")?.into(), + self.expected("a string")?, )); } self.strict_string() @@ -474,12 +473,7 @@ text = Text(' False ') }); } - #[test] - fn missing_snapshot_fields_and_descriptor_attribute_errors_are_distinct() { - with_python(|py| { - let locals = PyDict::new(py); - py.run( - c" + const RAISING_ATTRIBUTES: &std::ffi::CStr = c" failure = AttributeError('descriptor failed') class Snapshot: @property @@ -491,46 +485,56 @@ class Intercepted: def __getattribute__(self, name): raise failure dynamic = Dynamic() intercepted = Intercepted() -", - Some(&locals), - Some(&locals), - ) - .unwrap(); - let snapshot = locals.get_item("snapshot").unwrap().unwrap(); - let descriptor = PyErr::from(Field::read(&snapshot, "test", "flag").err().unwrap()); +"; + + #[rstest] + #[case::descriptor("snapshot")] + #[case::dynamic("dynamic")] + #[case::intercepted("intercepted")] + fn attribute_errors_raised_by_a_snapshot_surface_as_raised(#[case] name: &str) { + with_python(|py| { + let locals = PyDict::new(py); + py.run(RAISING_ATTRIBUTES, Some(&locals), Some(&locals)) + .unwrap(); + let value = locals.get_item(name).unwrap().unwrap(); + let error = PyErr::from(Field::read(&value, "test", "flag").err().unwrap()); assert!( - descriptor + error .value(py) .is(locals.get_item("failure").unwrap().unwrap()) ); - for name in ["dynamic", "intercepted"] { - let value = locals.get_item(name).unwrap().unwrap(); - let error = PyErr::from(Field::read(&value, "test", "flag").err().unwrap()); - assert!( - error - .value(py) - .is(locals.get_item("failure").unwrap().unwrap()) - ); - } + }); + } + + #[test] + fn missing_snapshot_fields_are_runtime_errors_naming_the_field() { + with_python(|py| { + let locals = PyDict::new(py); + py.run(RAISING_ATTRIBUTES, Some(&locals), Some(&locals)) + .unwrap(); + let snapshot = locals.get_item("snapshot").unwrap().unwrap(); let missing = PyErr::from(Field::read(&snapshot, "test", "missing").err().unwrap()); assert!(missing.is_instance_of::(py)); assert!(missing.to_string().contains("test.missing")); }); } - #[test] - fn configuration_errors_name_fields_without_exposing_values() { + #[rstest] + #[case::mapping("{'secret': 'do-not-print'}")] + #[case::nested("['host.test', {'secret': 'do-not-print'}]")] + fn configuration_errors_name_fields_without_exposing_values(#[case] source: &str) { + with_python(|py| { + let field = Field::new("test", "setting", evaluate(py, source)); + let error = PyErr::from(field.falsy_optional_string().err().unwrap()); + assert!(error.is_instance_of::(py)); + assert!(error.to_string().contains("test.setting")); + assert!(!error.to_string().contains("do-not-print")); + }); + } + + #[test] + fn invalid_collections_and_flags_are_configuration_errors() { with_python(|py| { - for source in [ - "{'secret': 'do-not-print'}", - "['host.test', {'secret': 'do-not-print'}]", - ] { - let field = Field::new("test", "setting", evaluate(py, source)); - let error = PyErr::from(field.falsy_optional_string().err().unwrap()); - assert!(error.is_instance_of::(py)); - assert!(error.to_string().contains("test.setting")); - assert!(!error.to_string().contains("do-not-print")); - } let hosts = Field::new( "url_policy", "user_url_allowed_hosts", diff --git a/litellm-rust/crates/python-bridge/src/credentials.rs b/litellm-rust/crates/python-bridge/src/credentials.rs index 04d8e94c8fe..a4bb064aae0 100644 --- a/litellm-rust/crates/python-bridge/src/credentials.rs +++ b/litellm-rust/crates/python-bridge/src/credentials.rs @@ -10,10 +10,19 @@ use pyo3::{ types::{PyDict, PyString}, }; -const NOT_CALLABLE: &str = "Azure AD token provider must be callable"; const NOT_A_STRING: &str = "Azure AD token must be a string, got {}"; const FAILED: &str = "Failed to get Azure AD token: {}"; +#[derive(Debug, thiserror::Error)] +#[error("Azure AD token provider must be callable")] +struct NotCallable; + +impl From for PyErr { + fn from(error: NotCallable) -> Self { + PyTypeError::new_err(litellm_error::Report::new(&error).to_string()) + } +} + /// The `azure_ad_token_provider` keyword argument, kept alive for the rest of the call. pub(crate) struct CallerTokenProvider { provider: Py, @@ -36,7 +45,7 @@ impl CallerTokenProvider { pub(crate) fn acquire(&self, py: Python<'_>) -> PyResult { let provider = self.provider.bind(py); if !provider.is_callable() { - return Err(PyTypeError::new_err(NOT_CALLABLE)); + return Err(NotCallable.into()); } let token = wrap_failure( py, diff --git a/litellm-rust/crates/python-bridge/src/diagnostics.rs b/litellm-rust/crates/python-bridge/src/diagnostics.rs index 687a090e768..5368da4f930 100644 --- a/litellm-rust/crates/python-bridge/src/diagnostics.rs +++ b/litellm-rust/crates/python-bridge/src/diagnostics.rs @@ -1,5 +1,5 @@ use litellm_host_python::{release_count, runtime_started}; -use pyo3::{exceptions::PyRuntimeError, prelude::*, types::PyDict}; +use pyo3::{prelude::*, types::PyDict}; #[pyfunction] pub(crate) fn gil_stats(py: Python<'_>) -> PyResult> { @@ -18,8 +18,7 @@ pub(crate) fn process_state_started() -> bool { /// so the runtime can never start. Raises if it already has. Forked workers are unaffected. #[pyfunction] pub(crate) fn reserve_process_for_forking() -> PyResult<()> { - litellm_host_python::reserve_process_for_forking() - .map_err(|_| PyRuntimeError::new_err("the native runtime already started in this process")) + Ok(litellm_host_python::reserve_process_for_forking()?) } #[cfg(feature = "panic-test")] diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index 4dab0297e73..a883ec608e5 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -1,9 +1,12 @@ -use litellm_http::transport::Error as TransportError; +use std::error::Error as StdError; + +use litellm_error::Report; use litellm_inference::Error; use pyo3::{ exceptions::{PyRuntimeError, PyValueError}, prelude::*, }; +use reqwest::header::HeaderMap; pyo3::create_exception!( _native, @@ -16,20 +19,34 @@ pyo3::create_exception!( _native, RustUpstreamError, pyo3::exceptions::PyException, - "The provider call was already issued and failed. Args are (status, message); status is 0 when there was no HTTP response." + "The provider call was already issued and the provider answered with a failure. Args are (status, body); the headers attribute holds the provider's (name, value) response headers." ); -pub(crate) fn route_error_to_pyerr(error: Error) -> PyErr { - match error { - Error::Transport(TransportError::Http { status, body }) => { - RustUpstreamError::new_err((status, body)) - } - other => by_fault(other.is_request(), other.to_string()), - } +const MESSAGES_REQUEST_ERROR: &str = "messages_request_error"; + +pub(crate) fn upstream_error( + py: Python<'_>, + status: u16, + body: &str, + headers: &HeaderMap, +) -> PyResult { + let error = RustUpstreamError::new_err((status, body.to_owned())); + let pairs: Vec<(&str, String)> = headers + .iter() + .map(|(name, value)| { + ( + name.as_str(), + String::from_utf8_lossy(value.as_bytes()).into_owned(), + ) + }) + .collect(); + error.value(py).setattr("headers", pairs)?; + Ok(error) } /// A request the caller got wrong is a `ValueError`; anything else is a `RuntimeError`. -pub(crate) fn by_fault(is_request: bool, message: String) -> PyErr { +pub(crate) fn by_fault(is_request: bool, error: &(dyn StdError + 'static)) -> PyErr { + let message = Report::new(error).to_string(); if is_request { PyValueError::new_err(message) } else { @@ -37,17 +54,82 @@ pub(crate) fn by_fault(is_request: bool, message: String) -> PyErr { } } +/// How a route flags a rejected request for its Python route module, which raises the public +/// exception for it instead of mapping the native one. +#[derive(Clone, Copy)] +pub(crate) enum RequestMarker { + Unmarked, + Messages, +} + +/// Maps a route failure once. A Python exception raised inside the call is returned as it is, +/// everything else is mapped and handed to `public`, the route's public exception mapping. +pub(crate) fn route_failure( + py: Python<'_>, + error: Error, + marker: RequestMarker, + public: impl FnOnce(PyErr) -> PyResult, +) -> PyResult { + if let Some(raised) = crate::secrets::python_error(py, &error) { + return Ok(raised); + } + if let Some((status, body)) = error.upstream_response() { + return public(upstream_error(py, status, body, &HeaderMap::new())?); + } + let native = match (marker, error.rejected_request()) { + (RequestMarker::Messages, Some(rejection)) => { + let native = PyValueError::new_err(Report::new(rejection).to_string()); + native.value(py).setattr(MESSAGES_REQUEST_ERROR, true)?; + native + } + (RequestMarker::Messages, None) | (RequestMarker::Unmarked, _) => { + by_fault(error.is_request(), &error) + } + }; + public(native) +} + +pub(crate) fn route_error_to_pyerr(error: Error) -> PyErr { + Python::attach(|py| { + route_failure(py, error, RequestMarker::Unmarked, Ok).unwrap_or_else(std::convert::identity) + }) +} + #[cfg(test)] mod tests { + use reqwest::header::{HeaderValue, RETRY_AFTER}; + use super::*; + fn upstream_payload(py: Python<'_>, failure: &PyErr) -> ((u16, String), Vec<(String, String)>) { + let value = failure.value(py); + ( + value.getattr("args").unwrap().extract().unwrap(), + value.getattr("headers").unwrap().extract().unwrap(), + ) + } + #[rstest::rstest] #[case::unsupported(Error::Unsupported("test capability"), true)] #[case::invalid_provider(Error::InvalidProvider("unknown".into()), true)] - #[case::invalid_request(Error::InvalidRequest("empty messages".into()), true)] - #[case::connection(TransportError::Connect("unreachable".into()).into(), false)] - #[case::network(TransportError::Network("timed out".into()).into(), false)] - #[case::invalid_response(Error::InvalidResponse("missing usage".into()), false)] + #[case::invalid_request( + Error::invalid_request(litellm_llms::ErrorDetail::Expected { + subject: "messages", + expected: "non-empty", + }), + true + )] + #[case::lost_connection( + litellm_http::Error::from(tokio_tungstenite::tungstenite::Error::ConnectionClosed).into(), + false + )] + #[case::invalid_response( + Error::invalid_response(litellm_llms::ErrorDetail::MissingField { + subject: "response", + field: "usage", + }), + false + )] fn route_failures_are_terminal(#[case] error: Error, #[case] is_request: bool) { Python::initialize(); Python::attach(|py| { @@ -59,29 +141,69 @@ mod tests { } #[rstest::rstest] - fn transport_status_survives_python_mapping() { + #[case::unmarked(RequestMarker::Unmarked)] + #[case::messages(RequestMarker::Messages)] + fn the_provider_answer_passes_through_with_a_headers_attribute(#[case] marker: RequestMarker) { Python::initialize(); Python::attach(|py| { - let upstream = route_error_to_pyerr( - TransportError::Http { - status: 429, - body: "slow down".into(), - } - .into(), - ); + let error = litellm_http::Error::upstream(429, "slow down".into()).into(); + let upstream = route_failure(py, error, marker, Ok).unwrap(); assert!(upstream.is_instance_of::(py)); assert_eq!( + upstream_payload(py, &upstream), + ((429, "slow down".into()), Vec::new()) + ); + assert!( upstream .value(py) - .getattr("args") + .getattr_opt(MESSAGES_REQUEST_ERROR) .unwrap() - .extract::<(u16, String)>() - .unwrap(), - (429, "slow down".into()) + .is_none() ); }); } + #[test] + fn upstream_error_keeps_the_provider_headers() { + Python::initialize(); + Python::attach(|py| { + let headers = HeaderMap::from_iter([(RETRY_AFTER, HeaderValue::from_static("17"))]); + let upstream = upstream_error(py, 503, "busy", &headers).unwrap(); + assert_eq!( + upstream_payload(py, &upstream), + ( + (503, "busy".into()), + vec![("retry-after".into(), "17".into())] + ) + ); + }); + } + + #[rstest::rstest] + #[case::ascii("busy ".repeat(80))] + #[case::unicode("é".repeat(300))] + fn long_provider_bodies_reach_python_intact( + #[case] message: String, + #[values(RequestMarker::Unmarked, RequestMarker::Messages)] marker: RequestMarker, + ) { + let envelope = serde_json::json!({ + "type": "error", + "error": {"type": "overloaded_error", "message": message}, + "request_id": "provider-request", + }); + let body = envelope.to_string(); + let error = litellm_http::Error::upstream(529, body.clone()).into(); + Python::initialize(); + Python::attach(|py| { + let upstream = route_failure(py, error, marker, Ok).unwrap(); + assert!(upstream.is_instance_of::(py)); + let ((status, returned_body), headers) = upstream_payload(py, &upstream); + assert_eq!(status, 529); + assert_eq!(returned_body, body); + assert!(headers.is_empty()); + }); + } + #[test] fn missing_api_key_stays_a_runtime_error_while_other_auth_failures_are_value_errors() { Python::initialize(); @@ -95,4 +217,83 @@ mod tests { assert!(invalid.is_instance_of::(py)); }); } + + #[test] + fn a_wrapped_failure_renders_every_layer() { + Python::initialize(); + Python::attach(|py| { + let error = Error::post_call( + litellm_http::Error::upstream(500, "hook backend down".into()).into(), + ); + let report = Report::new(&error).to_string(); + let failure = route_error_to_pyerr(error); + assert_eq!(failure.value(py).to_string(), report); + assert!(report.contains("hook backend down"), "{report}"); + }); + } + + #[rstest::rstest] + #[case::rejected_request( + Error::invalid_request(litellm_llms::ErrorDetail::UnsupportedParameter { + model: "claude".into(), + param: "top_k".into(), + value: "5".into(), + hint: String::new(), + }), + Some("claude does not support top_k=5. To drop unsupported params, set `litellm.drop_params = True`.") + )] + #[case::missing_field( + litellm_llms::Error::MissingField("max_tokens").into(), + Some("missing required field: max_tokens") + )] + #[case::unresolvable_provider(Error::InvalidProvider("openai".into()), None)] + #[case::invalid_response( + Error::invalid_response(litellm_llms::ErrorDetail::MissingField { + subject: "response", + field: "usage", + }), + None + )] + fn only_messages_request_rejections_carry_the_marker_and_python_text( + #[case] error: Error, + #[case] rejection: Option<&str>, + ) { + Python::initialize(); + Python::attach(|py| { + let unmarked = route_failure(py, error.clone(), RequestMarker::Unmarked, Ok).unwrap(); + assert!( + unmarked + .value(py) + .getattr_opt(MESSAGES_REQUEST_ERROR) + .unwrap() + .is_none() + ); + let native = route_failure(py, error, RequestMarker::Messages, Ok).unwrap(); + let marked = native + .value(py) + .getattr_opt(MESSAGES_REQUEST_ERROR) + .unwrap() + .map(|value| value.extract::().unwrap()); + assert_eq!(marked.unwrap_or(false), rejection.is_some()); + if let Some(rejection) = rejection { + assert!(native.is_instance_of::(py)); + assert_eq!(native.value(py).to_string(), rejection); + } + }); + } + + #[test] + fn a_python_secret_manager_exception_skips_the_public_mapping() { + Python::initialize(); + Python::attach(|py| { + let raised = pyo3::exceptions::PyKeyError::new_err("vault down"); + let identity = raised.value(py).clone().unbind(); + let error = Error::from(crate::secrets::read_error(py, raised)); + let surfaced = route_failure(py, error, RequestMarker::Unmarked, |_| { + panic!("the raised exception must not be remapped") + }) + .unwrap(); + assert!(surfaced.value(py).is(identity.bind(py))); + }); + } } diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index 57bf8aafc4e..b4e20a8f750 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -132,26 +132,14 @@ pub(crate) fn host_client(py: Python<'_>, variant: ClientVariant) -> PyResult PyErr { - match error { - litellm_http::Error::Read { - tls_source: TlsSource::ClientIdentity, - .. - } - | litellm_http::Error::InvalidPem { - tls_source: TlsSource::ClientIdentity, - .. - } => PyValueError::new_err( + match error.tls_source() { + Some(TlsSource::ClientIdentity) => PyValueError::new_err( "http_settings.ssl_certificate: expected a readable PEM certificate and private key", ), - litellm_http::Error::Read { - tls_source: TlsSource::CaBundle, - .. + Some(TlsSource::CaBundle) => { + PyValueError::new_err("http_settings.ssl_verify: expected a readable PEM CA bundle") } - | litellm_http::Error::InvalidPem { - tls_source: TlsSource::CaBundle, - .. - } => PyValueError::new_err("http_settings.ssl_verify: expected a readable PEM CA bundle"), - _ => PyValueError::new_err("http_settings: native HTTP client configuration is invalid"), + None => PyValueError::new_err(litellm_error::Report::new(&error).to_string()), } } @@ -274,21 +262,28 @@ settings = types.SimpleNamespace(**defaults) fn client_error_uses_tls_source_when_paths_match() { Python::initialize(); Python::attach(|py| { - let path = PathBuf::from("/shared.pem"); - let ca_error = client_error(litellm_http::Error::InvalidPem { - path: path.clone(), - message: "invalid".into(), - tls_source: TlsSource::CaBundle, - }); + let file = tempfile::NamedTempFile::new().unwrap(); + std::fs::write(file.path(), b"not a certificate").unwrap(); + let path = file.path().to_path_buf(); + let defaults = Resolution::from(&HttpSettings::default()).config; + let build = |config: HttpClientConfig| { + HttpClientPool::new(Arc::new(PublicDnsResolver)) + .client(&config, ClientVariant::Provider) + .map(drop) + .unwrap_err() + }; + let ca_error = client_error(build(HttpClientConfig { + verify: Verify::CaBundle(path.clone()), + ..defaults.clone() + })); assert_eq!( ca_error.to_string(), "ValueError: http_settings.ssl_verify: expected a readable PEM CA bundle" ); - let client_error = client_error(litellm_http::Error::InvalidPem { - path, - message: "invalid".into(), - tls_source: TlsSource::ClientIdentity, - }); + let client_error = client_error(build(HttpClientConfig { + client_certificate: Some(litellm_http::ClientIdentity::Pem(path)), + ..defaults + })); assert!(client_error.is_instance_of::(py)); assert_eq!( client_error.to_string(), @@ -297,6 +292,34 @@ settings = types.SimpleNamespace(**defaults) }); } + #[test] + fn client_error_without_a_tls_source_reports_the_build_failure() { + Python::initialize(); + Python::attach(|py| { + let error = HttpClientPool::new(Arc::new(PublicDnsResolver)) + .client( + &HttpClientConfig { + user_agent: Some("bad\nagent".into()), + ..Resolution::from(&HttpSettings::default()).config + }, + ClientVariant::Provider, + ) + .map(drop) + .unwrap_err(); + assert_eq!(error.tls_source(), None); + let report = litellm_error::Report::new(&error).to_string(); + + let failure = client_error(error); + + assert!(failure.is_instance_of::(py)); + assert_eq!(failure.value(py).str().unwrap().to_string(), report); + assert!( + report.contains("could not build the HTTP client: "), + "{report}" + ); + }); + } + #[test] fn python_settings_flow_into_the_configured_layer() { Python::initialize(); diff --git a/litellm-rust/crates/python-bridge/src/logger/mod.rs b/litellm-rust/crates/python-bridge/src/logger/mod.rs index bf5c735360b..8f40d94f458 100644 --- a/litellm-rust/crates/python-bridge/src/logger/mod.rs +++ b/litellm-rust/crates/python-bridge/src/logger/mod.rs @@ -81,8 +81,12 @@ impl NativeDiagnosticProcessor { } } -fn processing_error(_: fancy_regex::Error) -> PyErr { - PyRuntimeError::new_err("diagnostic processing failed") +#[derive(Debug, thiserror::Error)] +#[error("diagnostic processing failed")] +struct ProcessingError(#[source] fancy_regex::Error); + +fn processing_error(error: fancy_regex::Error) -> PyErr { + PyRuntimeError::new_err(litellm_error::Report::new(&ProcessingError(error)).to_string()) } struct PythonSink { diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs index 038fa8c9630..dec8ff0b0ce 100644 --- a/litellm-rust/crates/python-bridge/src/marshal.rs +++ b/litellm-rust/crates/python-bridge/src/marshal.rs @@ -21,10 +21,27 @@ pub(crate) struct RouteOptions { pub(crate) timeout: Option, } -fn required_object(name: &'static str, value: Value) -> PyResult> { +/// A caller argument the bridge cannot use, raised as a `ValueError`. +#[derive(Debug, thiserror::Error)] +pub(crate) enum ArgumentError { + #[error("{0} is required")] + Required(&'static str), + #[error("{0} must be a dict")] + NotADict(&'static str), + #[error("header values must be strings")] + HeaderValue, +} + +impl From for PyErr { + fn from(error: ArgumentError) -> Self { + PyValueError::new_err(litellm_error::Report::new(&error).to_string()) + } +} + +fn required_object(name: &'static str, value: Value) -> Result, ArgumentError> { match value { Value::Object(values) => Ok(values), - _ => Err(PyValueError::new_err(format!("{name} must be a dict"))), + _ => Err(ArgumentError::NotADict(name)), } } @@ -35,7 +52,7 @@ fn optional_object( if value.is_none() { return Ok(None); } - required_object(name, from_py_argument(value)?).map(Some) + Ok(Some(required_object(name, from_py_argument(value)?)?)) } pub(crate) fn optional_timeout(timeout_seconds: Option) -> Option { @@ -57,11 +74,11 @@ pub(crate) fn python_timeout_seconds(py: Python<'_>, timeout: Py) -> PyRe pub(crate) fn required_field<'py>( fields: &Bound<'py, PyDict>, - name: &str, + name: &'static str, ) -> PyResult> { - fields + Ok(fields .get_item(name)? - .ok_or_else(|| PyValueError::new_err(format!("{name} is required"))) + .ok_or(ArgumentError::Required(name))?) } pub(crate) fn optional_field( @@ -171,21 +188,18 @@ pub(crate) fn request_input_sources<'a>( .collect()) } -pub(crate) fn marshal_headers(headers: Option) -> PyResult> { - let value = match headers { - Some(headers) => headers, - None => Value::Object(Map::new()), - }; - let Value::Object(headers) = value else { - return Err(PyValueError::new_err("headers must be a dict")); +pub(crate) fn marshal_headers( + headers: Option, +) -> Result, ArgumentError> { + let headers = match headers { + Some(headers) => required_object("headers", headers)?, + None => Map::new(), }; headers .into_iter() - .map(|(name, value)| { - value - .as_str() - .map(|value| (name, value.to_string())) - .ok_or_else(|| PyValueError::new_err("header values must be strings")) + .map(|(name, value)| match value { + Value::String(value) => Ok((name, value)), + _ => Err(ArgumentError::HeaderValue), }) .collect() } @@ -395,48 +409,52 @@ module.__getattr__ = fail float: Option, } - #[test] - fn numeric_adapters_agree_across_json_and_python_boundaries() { + #[rstest::rstest] + #[case::empty(json!({}))] + #[case::explicit_nulls(json!({"integers": null, "float": null}))] + #[case::lax_values(json!({ + "integers": [i64::MIN, i64::MAX, "9007199254740993.0", " +1_000.00 ", true, 3.0], + "float": " 1.25 ", + }))] + #[case::integer_overflow(json!({"integers": [u64::MAX]}))] + #[case::imprecise_integer_string(json!({"integers": ["1.0000000000000001"]}))] + #[case::fractional_integer(json!({"integers": [2.5]}))] + #[case::nan_string(json!({"float": "NaN"}))] + #[case::infinity_string(json!({"float": "inf"}))] + #[case::overflowing_float_string(json!({"float": "1e999"}))] + #[case::boolean_float(json!({"float": true}))] + #[case::large_integer_float(json!({"float": u64::MAX}))] + fn numeric_adapters_agree_across_json_and_python_boundaries(#[case] input: Value) { Python::initialize(); Python::attach(|py| { - for input in [ - json!({}), - json!({"integers": null, "float": null}), - json!({"integers": [i64::MIN, i64::MAX, "9007199254740993.0", " +1_000.00 ", true, 3.0], "float": " 1.25 "}), - json!({"integers": [u64::MAX]}), - json!({"integers": ["1.0000000000000001"]}), - json!({"integers": [2.5]}), - json!({"float": "NaN"}), - json!({"float": "inf"}), - json!({"float": "1e999"}), - json!({"float": true}), - json!({"float": u64::MAX}), - ] { - let expected = serde_json::from_value::(input.clone()); - let python = litellm_host_python::to_py(py, &input).unwrap(); - let actual = from_py::(python.bind(py)); - match (expected, actual) { - (Ok(expected), Ok(actual)) => { - assert_eq!(actual, expected); - let serialized = litellm_host_python::to_py(py, &actual).unwrap(); - assert_eq!( - from_py::(serialized.bind(py)).unwrap(), - serde_json::to_value(expected).unwrap() - ); - } - (Err(_), Err(_)) => {} - mismatch => panic!("boundary mismatch for {input}: {mismatch:?}"), + let expected = serde_json::from_value::(input.clone()); + let python = litellm_host_python::to_py(py, &input).unwrap(); + let actual = from_py::(python.bind(py)); + match (expected, actual) { + (Ok(expected), Ok(actual)) => { + assert_eq!(actual, expected); + let serialized = litellm_host_python::to_py(py, &actual).unwrap(); + assert_eq!( + from_py::(serialized.bind(py)).unwrap(), + serde_json::to_value(expected).unwrap() + ); } + (Err(_), Err(_)) => {} + mismatch => panic!("boundary mismatch for {input}: {mismatch:?}"), } - for source in [ - c"{'float': float('nan')}", - c"{'float': float('inf')}", - c"{'integers': [float('inf')]}", - c"{'integers': [2 ** 100]}", - ] { - let value = py.eval(source, None, None).unwrap(); - assert!(from_py::(&value).is_err()); - } + }); + } + + #[rstest::rstest] + #[case::nan(c"{'float': float('nan')}")] + #[case::infinity(c"{'float': float('inf')}")] + #[case::infinite_integer(c"{'integers': [float('inf')]}")] + #[case::integer_overflow(c"{'integers': [2 ** 100]}")] + fn non_finite_and_overflowing_python_numbers_are_rejected(#[case] source: &std::ffi::CStr) { + Python::initialize(); + Python::attach(|py| { + let value = py.eval(source, None, None).unwrap(); + assert!(from_py::(&value).is_err()); }); } diff --git a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs index e62673d4d4f..21a57123921 100644 --- a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs +++ b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs @@ -13,7 +13,7 @@ use crate::{ }; async fn execute( - http: Result, + http: litellm_http::Client, secrets: std::sync::Arc, audio: Value, optional_params: Map, @@ -27,7 +27,7 @@ async fn execute( extra_headers, timeout, } = options; - AudioTranscriptionRoute::new(http?, crate::http::resources().auth.clone(), secrets) + AudioTranscriptionRoute::new(http, crate::http::resources().auth.clone(), secrets) .execute(AudioTranscriptionRequest { model: &model, audio, @@ -48,7 +48,8 @@ pub(crate) fn transcription(py: Python<'_>, call: NativeCall<'_>) -> PyResult( litellm_host_python::from_py_argument(&required_field(&arguments, "audio")?)?; let options = value_route_options(&arguments)?; let optional_params = optional_object_field(&arguments, "optional_params")?.unwrap_or_default(); - let http = crate::http::provider_client(py, &call.kwargs, true)?; + let http = + crate::http::provider_client(py, &call.kwargs, true)?.map_err(crate::http::client_error)?; let secrets = crate::secrets::source(py)?; run_async( py, diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs index 5ad9f8df8b8..63ccfc2433f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs @@ -79,7 +79,7 @@ impl PythonBinding for ChatCompletionsPythonHost { self.0.error(py, error) } fn host_error(error: &PyErr) -> Error { - Error::InvalidRequest(error.to_string().into()) + crate::routes::inference::host_call_error(error) } } diff --git a/litellm-rust/crates/python-bridge/src/routes/clickhouse_spend.rs b/litellm-rust/crates/python-bridge/src/routes/clickhouse_spend.rs index f8df7a9de23..80d7bd5360e 100644 --- a/litellm-rust/crates/python-bridge/src/routes/clickhouse_spend.rs +++ b/litellm-rust/crates/python-bridge/src/routes/clickhouse_spend.rs @@ -8,24 +8,13 @@ use pyo3::{ }; fn map_error(error: Error) -> PyErr { - use litellm_storage_clickhouse::Error as StorageError; - match error { - Error::InsertTooLarge | Error::Storage(StorageError::InsertTooLarge) => { - PyOverflowError::new_err(error.to_string()) - } - Error::InvalidRow - | Error::InvalidLimit(_) - | Error::InvalidSchema - | Error::Storage( - StorageError::InvalidRow - | StorageError::InvalidLimit(_) - | StorageError::InvalidTable - | StorageError::InvalidSchema - | StorageError::EmptySql - | StorageError::InvalidParameters - | StorageError::InvalidQuery, - ) => PyValueError::new_err(error.to_string()), - Error::Migration(_) | Error::Storage(_) => PyRuntimeError::new_err(error.to_string()), + let message = litellm_error::Report::new(&error).to_string(); + if error.is_insert_too_large() { + PyOverflowError::new_err(message) + } else if error.is_request() { + PyValueError::new_err(message) + } else { + PyRuntimeError::new_err(message) } } @@ -99,26 +88,22 @@ mod tests { use super::*; use rstest::rstest; + use litellm_storage_clickhouse::Error as StorageError; + #[rstest] - #[case::row(Error::InvalidRow, "ValueError")] - #[case::insert_limit(Error::InvalidLimit("CLICKHOUSE_TRACE_MAX_INSERT_BYTES"), "ValueError")] - #[case::insert_timeout( - Error::Storage(litellm_storage_clickhouse::Error::InvalidLimit( - "CLICKHOUSE_INSERT_TIMEOUT_SECONDS" - )), + #[case::timestamp(Error::NonIntegerTimestamp { field: "start_time" }, "ValueError")] + #[case::insert_limit( + StorageError::InvalidLimit("CLICKHOUSE_TRACE_MAX_INSERT_BYTES").into(), "ValueError" )] - #[case::insert_budget(Error::InsertTooLarge, "OverflowError")] - #[case::schema( - Error::Storage(litellm_storage_clickhouse::Error::SchemaFailed(503)), - "RuntimeError" - )] + #[case::insert_budget(StorageError::InsertTooLarge.into(), "OverflowError")] + #[case::schema(StorageError::SchemaFailed(503).into(), "RuntimeError")] #[case::migration( Error::Migration(sqlx::migrate::MigrateError::VersionMismatch(6)), "RuntimeError" )] #[case::storage( - Error::Storage(litellm_storage_clickhouse::Error::InvalidUrl), + Config::new("spend".into(), "ftp://localhost", 1).err().unwrap(), "RuntimeError" )] fn spend_failures_preserve_public_exception_types( @@ -127,7 +112,7 @@ mod tests { ) { Python::initialize(); Python::attach(|py| { - let message = error.to_string(); + let message = litellm_error::Report::new(&error).to_string(); let exception = map_error(error); assert_eq!(exception.get_type(py).name().unwrap(), exception_name); assert_eq!( @@ -136,4 +121,15 @@ mod tests { ); }); } + + #[test] + fn spend_failures_render_their_cause() { + let error = Config::new("spend".into(), "not a url", 1).err().unwrap(); + let cause = std::error::Error::source(&error).unwrap().to_string(); + Python::initialize(); + Python::attach(|py| { + let message = map_error(error).value(py).str().unwrap().to_string(); + assert!(message.ends_with(&cause), "{message}"); + }); + } } diff --git a/litellm-rust/crates/python-bridge/src/routes/embeddings.rs b/litellm-rust/crates/python-bridge/src/routes/embeddings.rs index 0f73eefedd5..0935f42f85a 100644 --- a/litellm-rust/crates/python-bridge/src/routes/embeddings.rs +++ b/litellm-rust/crates/python-bridge/src/routes/embeddings.rs @@ -1,14 +1,23 @@ +use litellm_error::Report; use pyo3::prelude::*; use super::NativeCall; use crate::errors::RustBridgeDeclined; +#[derive(Debug, thiserror::Error)] +#[error("native embeddings route is not implemented")] +struct NotImplemented; + +impl From for PyErr { + fn from(error: NotImplemented) -> Self { + RustBridgeDeclined::new_err(Report::new(&error).to_string()) + } +} + #[pyfunction] pub(crate) fn embedding(call: NativeCall<'_>) -> PyResult> { drop(call); - Err(RustBridgeDeclined::new_err( - "native embeddings route is not implemented", - )) + Err(NotImplemented.into()) } #[pyfunction] diff --git a/litellm-rust/crates/python-bridge/src/routes/inference.rs b/litellm-rust/crates/python-bridge/src/routes/inference.rs index c9eb60e47ed..ca30c94b02f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/inference.rs +++ b/litellm-rust/crates/python-bridge/src/routes/inference.rs @@ -6,11 +6,9 @@ use litellm_callbacks_legacy_python::LoggingOperation; use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider; use litellm_host::{call::HostedMachine, protocol::Protocol}; use litellm_host_python::{PythonBinding, PythonHostCalls, from_py}; -use litellm_http::transport::Error as TransportError; use litellm_inference::Error; use litellm_secrets::source::SecretSource; use pyo3::{ - exceptions::PyValueError, gc::{PyTraverseError, PyVisit}, prelude::*, types::PyDict, @@ -20,8 +18,10 @@ use serde_json::{Map, Value}; use super::{NativeCall, RequestView}; use crate::{ - errors::{RustUpstreamError, route_error_to_pyerr}, - marshal::{RouteOptions, optional_timeout, public_response, python_timeout_seconds}, + errors::{RequestMarker, route_failure}, + marshal::{ + ArgumentError, RouteOptions, optional_timeout, public_response, python_timeout_seconds, + }, }; pub(super) struct InferenceHost { @@ -44,7 +44,7 @@ impl InferenceHost { &mut self, py: Python<'_>, arguments: &Bound<'_, PyDict>, - input: &str, + input: &'static str, ) -> PyResult { self.view.resolve(py, arguments)?; let argument = |name: &str| self.argument(py, arguments, name); @@ -57,7 +57,7 @@ impl InferenceHost { .map(|value| python_timeout_seconds(py, value.unbind())) .transpose()? .flatten(); - let model = string("model")?.ok_or_else(|| PyValueError::new_err("model is required"))?; + let model = string("model")?.ok_or(ArgumentError::Required("model"))?; let custom_llm_provider = string("custom_llm_provider")?; let provider = get_custom_llm_provider(&model, custom_llm_provider.as_deref()) .map_or("", |resolved| resolved.custom_llm_provider); @@ -82,10 +82,7 @@ impl InferenceHost { .transpose()?, timeout: optional_timeout(timeout), }, - input: from_py( - &argument(input)? - .ok_or_else(|| PyValueError::new_err(format!("{input} is required")))?, - )?, + input: from_py(&argument(input)?.ok_or(ArgumentError::Required(input))?)?, params, }) } @@ -120,29 +117,20 @@ impl InferenceHost { } pub fn error(&self, py: Python<'_>, error: Error) -> PyResult { - if let Error::Secret(source) = &error - && let Some(original) = crate::secrets::python_error(py, source.source_error()) - { - return Ok(original); - } - let native = match error { - Error::Transport(TransportError::Http { status, body }) => { - let error = RustUpstreamError::new_err((status, body)); - error - .value(py) - .setattr("headers", Vec::<(String, String)>::new())?; - error - } - other => route_error_to_pyerr(other), - }; - let mapped = py - .import(self.module)? - .getattr("map_failure")? - .call1((native.value(py), self.view.request(py)?))?; - Ok(PyErr::from_value(mapped)) + route_failure(py, error, RequestMarker::Unmarked, |native| { + let mapped = py + .import(self.module)? + .getattr("map_failure")? + .call1((native.value(py), self.view.request(py)?))?; + Ok(PyErr::from_value(mapped)) + }) } } +pub(super) fn host_call_error(_: &PyErr) -> Error { + Error::invalid_request(litellm_llms::ErrorDetail::HostCall) +} + pub(super) trait InferenceRoute: Sized + 'static { type Protocol: Protocol; const OPERATION: LoggingOperation; diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index 5e445fcaa7d..27ce8c943f0 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -3,7 +3,6 @@ use litellm_host_python::{PythonHostCalls, PythonOwned}; use bytes::Bytes; use litellm_host_python::{InvokeError, PythonBinding, from_py, to_py}; -use litellm_http::transport::Error as TransportError; use litellm_inference_messages::{ Error, MessagesCall, MessagesSettings, MessagesShaping, litellm_params, messages_body, route::{Messages, MessagesStreamHead}, @@ -12,7 +11,7 @@ use litellm_llms::base_llm::messages::context::MessagesModelCapabilities; use litellm_llms_types::headers::ProviderSpecificHeaders; use litellm_router_types::LitellmParams; use pyo3::{ - exceptions::{PyException, PyValueError}, + exceptions::PyException, gc::{PyTraverseError, PyVisit}, prelude::*, types::{PyBytes, PyDict}, @@ -20,14 +19,16 @@ use pyo3::{ use serde_json::{Map, Value}; use crate::{ - errors::{RustUpstreamError, route_error_to_pyerr}, - marshal::{optional_timeout, project_optional_fields, public_response, python_timeout_seconds}, + errors::{RequestMarker, route_failure}, + marshal::{ + ArgumentError, optional_timeout, project_optional_fields, public_response, + python_timeout_seconds, + }, python_settings::missing_module, routes::RequestView, }; const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.host.routes.messages"; -const REQUEST_ERROR_MARKER: &str = "messages_request_error"; fn merge_headers( forwarded: Option>, @@ -68,29 +69,6 @@ fn project_litellm_params<'py>( Ok(litellm_params(fields.into_iter().chain(folded).collect())) } -fn native_error(py: Python<'_>, error: Error) -> PyResult { - match error { - Error::Transport(TransportError::Http { status, body }) => { - let error = RustUpstreamError::new_err((status, body)); - error - .value(py) - .setattr("headers", Vec::<(String, String)>::new())?; - Ok(error) - } - Error::InvalidRequest(message) => { - let error = PyValueError::new_err(message.to_string()); - error.value(py).setattr(REQUEST_ERROR_MARKER, true)?; - Ok(error) - } - Error::MissingField(field) => { - let error = PyValueError::new_err(format!("missing required field: {field}")); - error.value(py).setattr(REQUEST_ERROR_MARKER, true)?; - Ok(error) - } - other => Ok(route_error_to_pyerr(other)), - } -} - /// The Python side of the Messages route: projects the prepared arguments and builds the /// public response, chunks and exceptions. pub(super) struct MessagesPythonHost { @@ -115,9 +93,8 @@ impl MessagesPythonHost { let string = |name: &str| -> PyResult> { argument(name)?.map(|value| value.extract()).transpose() }; - let model = string("model")?.ok_or_else(|| PyValueError::new_err("model is required"))?; - let messages = - argument("messages")?.ok_or_else(|| PyValueError::new_err("messages is required"))?; + let model = string("model")?.ok_or(ArgumentError::Required("model"))?; + let messages = argument("messages")?.ok_or(ArgumentError::Required("messages"))?; let fields = self.view.body_params(py, arguments, ROUTE_HOST_MODULE)?; let body = [ ("model".to_string(), Value::String(model.clone())), @@ -276,16 +253,13 @@ impl PythonBinding for MessagesPythonHost { } fn map_error(&self, py: Python<'_>, error: Error) -> PyResult { - if let Error::Secret(source) = &error - && let Some(original) = crate::secrets::python_error(py, source.source_error()) - { - return Ok(original); - } - Ok(self.map_failure(py, native_error(py, error)?)) + route_failure(py, error, RequestMarker::Messages, |native| { + Ok(self.map_failure(py, native)) + }) } fn host_error(error: &PyErr) -> Error { - Error::InvalidRequest(error.to_string().into()) + crate::routes::inference::host_call_error(error) } } @@ -459,33 +433,9 @@ mod tests { projected, litellm_params(serde_json::from_value(fields).unwrap()).unwrap() ), - (Err(error), Err(())) => assert!(matches!(error, Error::InvalidRequest(_))), + (Err(error), Err(())) => assert!(error.is_request(), "{error:?}"), (projected, expected) => panic!("got {projected:?}, expected {expected:?}"), } }); } - - #[rstest] - #[case::rejected_request(Error::InvalidRequest("does not support top_k=5".into()), true)] - #[case::missing_field(Error::MissingField("max_tokens"), true)] - #[case::unresolvable_provider(Error::InvalidProvider("openai".into()), false)] - #[case::upstream_failure( - Error::Transport(TransportError::Http { status: 400, body: "bad".into() }), - false, - )] - fn only_request_rejections_carry_the_request_error_marker( - #[case] error: Error, - #[case] marked: bool, - ) { - Python::initialize(); - Python::attach(|py| { - let native = native_error(py, error).unwrap(); - let marker = native - .value(py) - .getattr_opt(REQUEST_ERROR_MARKER) - .unwrap() - .map(|value| value.extract::().unwrap()); - assert_eq!(marker.unwrap_or(false), marked); - }); - } } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs index f85d7d4c7cb..71eb743138d 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs @@ -9,6 +9,29 @@ use pyo3::{ types::{PyBytes, PyString, PyType}, }; +/// A `type='file'` document the bridge cannot read, with the text Python's OCR path raises. +#[derive(Debug, thiserror::Error)] +enum DocumentError { + #[error( + "document with type='file' must include a 'file' field containing a pathlib.Path, file-like object, or bytes" + )] + MissingFile, + #[error( + "OCR file input does not accept bare str values. Pass bytes, a pathlib.Path, or a file-like object." + )] + BareString, + #[error( + "Unsupported file input type: {0}. Expected pathlib.Path, bytes, or a file-like object." + )] + UnsupportedType(String), +} + +impl From for PyErr { + fn from(error: DocumentError) -> Self { + PyValueError::new_err(litellm_error::Report::new(&error).to_string()) + } +} + /// A `type='file'` document as projected: paths and bytes are typed inputs already; a /// file-like object is a reader the projection consumes once every other field is read. pub(super) enum FileDocumentInput { @@ -45,25 +68,18 @@ impl FromPyObject<'_, '_> for FileDocumentInput { Err(error) if error.is_instance_of::(py) => None, Err(error) => return Err(error), }; - let missing = || { - PyValueError::new_err( - "document with type='file' must include a 'file' field containing a pathlib.Path, file-like object, or bytes", - ) - }; - let file = document.get_item("file").map_err(|error| { - if error.is_instance_of::(py) { - missing() - } else { - error + let file = match document.get_item("file") { + Ok(file) => file, + Err(error) if error.is_instance_of::(py) => { + return Err(DocumentError::MissingFile.into()); } - })?; + Err(error) => return Err(error), + }; if file.is_none() { - return Err(missing()); + return Err(DocumentError::MissingFile.into()); } if file.is_instance_of::() { - return Err(PyValueError::new_err( - "OCR file input does not accept bare str values. Pass bytes, a pathlib.Path, or a file-like object.", - )); + return Err(DocumentError::BareString.into()); } static PATH_LIKE: PyOnceLock> = PyOnceLock::new(); if file.is_instance(PATH_LIKE.import(py, "os", "PathLike")?)? { @@ -81,10 +97,7 @@ impl FromPyObject<'_, '_> for FileDocumentInput { } match PythonFileReader::from_file_like(&file)? { Some(reader) => Ok(Self::Deferred { reader, mime_type }), - None => Err(PyValueError::new_err(format!( - "Unsupported file input type: {}. Expected pathlib.Path, bytes, or a file-like object.", - file.get_type(), - ))), + None => Err(DocumentError::UnsupportedType(file.get_type().to_string()).into()), } } } @@ -92,6 +105,7 @@ impl FromPyObject<'_, '_> for FileDocumentInput { #[cfg(test)] mod tests { use pyo3::{exceptions::PyTypeError, types::PyDict}; + use rstest::rstest; use super::*; @@ -108,40 +122,47 @@ mod tests { } } - #[test] - fn extraction_validates_required_file_and_optional_mime_type() { + enum Rejection { + Value(&'static str), + Type, + } + + #[rstest] + #[case::missing_file(c"{}", Rejection::Value("must include a 'file' field"))] + #[case::none_file(c"{'file': None}", Rejection::Value("must include a 'file' field"))] + #[case::none_mime_type(c"{'file': b'abc', 'mime_type': None}", Rejection::Type)] + #[case::numeric_mime_type(c"{'file': b'abc', 'mime_type': 7}", Rejection::Type)] + #[case::bare_str(c"{'file': 'scan.pdf'}", Rejection::Value("bare str"))] + #[case::unsupported_type( + c"{'file': object()}", + Rejection::Value("Unsupported file input type") + )] + fn extraction_rejects_invalid_documents( + #[case] expression: &std::ffi::CStr, + #[case] rejection: Rejection, + ) { Python::initialize(); Python::attach(|py| { - for expression in [c"{}", c"{'file': None}"] { - let document = py.eval(expression, None, None).unwrap(); - let error = document.extract::().err().unwrap(); - assert!(error.is_instance_of::(py)); - assert!(error.to_string().contains("must include a 'file' field")); - } - for expression in [ - c"{'file': b'abc', 'mime_type': None}", - c"{'file': b'abc', 'mime_type': 7}", - ] { - let document = py.eval(expression, None, None).unwrap(); - let error = document.extract::().err().unwrap(); - assert!(error.is_instance_of::(py)); - } let error = py - .eval(c"{'file': 'scan.pdf'}", None, None) + .eval(expression, None, None) .unwrap() .extract::() .err() .unwrap(); - assert!(error.is_instance_of::(py)); - assert!(error.to_string().contains("bare str")); - let error = py - .eval(c"{'file': object()}", None, None) - .unwrap() - .extract::() - .err() - .unwrap(); - assert!(error.is_instance_of::(py)); - assert!(error.to_string().contains("Unsupported file input type")); + match rejection { + Rejection::Value(fragment) => { + assert!(error.is_instance_of::(py)); + assert!(error.to_string().contains(fragment), "{error}"); + } + Rejection::Type => assert!(error.is_instance_of::(py)), + } + }); + } + + #[test] + fn extraction_keeps_bytes_and_their_mime_type() { + Python::initialize(); + Python::attach(|py| { let document = py .eval(c"{'file': b'abc', 'mime_type': 'image/png'}", None, None) .unwrap(); diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs index 2068b6a6e4b..31736c7b654 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs @@ -1,128 +1,203 @@ +use std::io::ErrorKind; + +use litellm_error::Report; use litellm_llms::base_llm::ocr::error::Error; use pyo3::{ exceptions::{PyFileNotFoundError, PyOSError}, prelude::*, }; -use crate::errors::{RustUpstreamError, by_fault}; +use crate::errors::{by_fault, upstream_error}; + +const REQUEST_FORMAT_ERROR: &str = "ocr_request_format_error"; pub(super) fn to_pyerr(error: Error) -> PyErr { - let status = error.http_status_code(); - let mapped = Python::attach(|py| -> PyResult { - Ok(match error { - Error::Provider { - status, - body, - headers, - } => upstream_error(py, status, body, headers)?, - Error::Transport(litellm_http::transport::Error::Http { status, body }) => { - upstream_error(py, status, body, Vec::new())? - } - Error::RequestFormat => { - let error = by_fault(true, Error::RequestFormat.to_string()); - error - .value(py) - .setattr("ocr_request_format_error", true) - .ok(); - error - } - Error::FileRead { path, source } if source.kind() == std::io::ErrorKind::NotFound => { - PyFileNotFoundError::new_err(format!("File not found: {}", path.display())) - } - Error::FileRead { source, .. } => PyOSError::new_err(source.to_string()), - other => by_fault(is_request(&other), other.to_string()), - }) - }) - .unwrap_or_else(|error| error); - attach_status(mapped, status) + Python::attach(|py| native(py, error).unwrap_or_else(std::convert::identity)) } -fn is_request(error: &Error) -> bool { - error.is_request() - || matches!( - error, - Error::Auth(_) - | Error::InvalidProvider(_) - | Error::InvalidRequest(_) - | Error::MissingField(_) - | Error::MissingDocumentUrl - ) -} - -fn upstream_error( +/// Maps an OCR route failure once. A Python exception raised inside the call is returned as it +/// is, everything else is mapped and handed to `public`, the route's public exception mapping. +pub(super) fn route_failure( py: Python<'_>, - status: u16, - body: String, - headers: Vec<(String, String)>, + error: Error, + public: impl FnOnce(PyErr) -> PyResult, ) -> PyResult { - let error = RustUpstreamError::new_err((status, body)); - error.value(py).setattr("headers", headers)?; - Ok(error) + match crate::secrets::python_error(py, &error) { + Some(raised) => Ok(raised), + None => public(native(py, error)?), + } } -fn attach_status(error: PyErr, status: Option) -> PyErr { - if let Some(status) = status { - Python::attach(|py| { - let value = error.value(py); - value.setattr("status_code", status).ok(); - value.setattr("message", value.to_string()).ok(); - }); +fn native(py: Python<'_>, error: Error) -> PyResult { + let mapped = match (&error, error.upstream_response()) { + (_, Some((status, body, headers))) => upstream_error(py, status, body, headers)?, + (Error::RequestFormat, None) => { + let mapped = by_fault(true, &error); + mapped.value(py).setattr(REQUEST_FORMAT_ERROR, true)?; + mapped + } + (Error::FileRead { path, cause }, None) if cause.kind() == ErrorKind::NotFound => { + PyFileNotFoundError::new_err(format!("File not found: {}", path.display())) + } + (Error::FileRead { cause, .. }, None) => { + PyOSError::new_err(Report::new(cause.as_ref()).to_string()) + } + (_, None) => by_fault(error.is_request(), &error), + }; + if let Some(status) = error.http_status_code() { + let value = mapped.value(py); + value.setattr("status_code", status)?; + value.setattr("message", value.str()?)?; } - error + Ok(mapped) } #[cfg(test)] mod tests { - use pyo3::exceptions::PyValueError; + use pyo3::exceptions::{PyRuntimeError, PyValueError}; + use rstest::rstest; use super::*; + use crate::errors::RustUpstreamError; - #[test] - fn preserves_python_validation_and_provider_details() { + fn status_code(py: Python<'_>, mapped: &PyErr) -> Option { + mapped + .value(py) + .getattr_opt("status_code") + .unwrap() + .map(|status| status.extract().unwrap()) + } + + fn file_read(kind: ErrorKind) -> Error { + Error::FileRead { + path: "/missing/scan.pdf".into(), + cause: std::sync::Arc::new(std::io::Error::new(kind, "disk said no")), + } + } + + #[rstest] + #[case::provider( + Error::Provider { + status: 429, + body: r#"{"message":"rate limited"}"#.into(), + headers: Box::new(reqwest::header::HeaderMap::from_iter([( + reqwest::header::RETRY_AFTER, + reqwest::header::HeaderValue::from_static("17"), + )])), + }, + vec![("retry-after".into(), "17".into())] + )] + #[case::transport(litellm_http::Error::upstream(429, r#"{"message":"rate limited"}"#.into()).into(), vec![])] + fn a_provider_answer_keeps_its_status_body_and_headers( + #[case] error: Error, + #[case] headers: Vec<(String, String)>, + ) { Python::initialize(); Python::attach(|py| { - let mapped = to_pyerr(Error::MissingDocumentUrl); - assert!(mapped.is_instance_of::(py)); - assert_eq!(mapped.value(py).to_string(), "Document URL is required"); - assert_eq!( - mapped - .value(py) - .getattr("status_code") - .unwrap() - .extract::() - .unwrap(), - 400 - ); - let mapped = to_pyerr(Error::Provider { - status: 429, - body: r#"{"message":"rate limited"}"#.to_string(), - headers: vec![("Retry-After".to_string(), "17".to_string())], - }); + let mapped = to_pyerr(error); assert!(mapped.is_instance_of::(py)); - let headers: Vec<(String, String)> = mapped - .value(py) - .getattr("headers") - .and_then(|headers| headers.extract()) - .expect("OCR failures retain provider headers"); - assert_eq!(headers, vec![("Retry-After".to_string(), "17".to_string())]); - let args: (u16, String) = mapped - .value(py) - .getattr("args") - .and_then(|args| args.extract()) - .expect("OCR failures retain status and unprefixed provider message"); - assert_eq!(args, (429, r#"{"message":"rate limited"}"#.to_string())); + let value = mapped.value(py); + let args: (u16, String) = value.getattr("args").unwrap().extract().unwrap(); + assert_eq!(args, (429, r#"{"message":"rate limited"}"#.into())); + let kept: Vec<(String, String)> = value.getattr("headers").unwrap().extract().unwrap(); + assert_eq!(kept, headers); + assert_eq!(status_code(py, &mapped), Some(429)); + }); + } - let mapped = to_pyerr(Error::InvalidRequest("invalid format".into())); + #[rstest] + #[case::not_found(Some(404), 404, "OCR document download failed")] + #[case::rate_limited(Some(429), 429, "OCR document download failed")] + #[case::timeout(None, 408, "OCR document download timed out")] + #[tokio::test] + async fn document_download_failures_keep_the_upstream_exception( + #[case] response_status: Option, + #[case] expected_status: u16, + #[case] expected_body: &str, + ) { + use std::time::Duration; + + use litellm_http::{Client, media::MediaFetcher}; + use litellm_llms::base_llm::ocr::{ + document::inline_remote_document, transformation::OcrConnection, + }; + use litellm_llms_types::formats::ocr::OcrDocument; + use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpListener, + }; + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0_u8; 1]; + socket.read_exact(&mut request).await.unwrap(); + if let Some(status) = response_status { + socket.write_all(format!("HTTP/1.1 {status} Error\r\nContent-Length: 0\r\nConnection: close\r\n\r\n").as_bytes()).await.unwrap(); + } else { + std::future::pending::<()>().await; + } + }); + let connection = OcrConnection { + timeout: Duration::from_millis(100), + ..OcrConnection::default() + }; + let error = inline_remote_document( + &MediaFetcher::for_test(Client::for_test()), + OcrDocument::DocumentUrl { + document_url: format!("http://{address}/document.pdf"), + extra_fields: Default::default(), + }, + &connection, + ) + .await + .unwrap_err(); + server.abort(); + if response_status.is_none() { + assert!(litellm_error::find::(&error).is_some()); + } + Python::initialize(); + Python::attach(|py| { + let mapped = to_pyerr(error); + assert!(mapped.is_instance_of::(py), "{mapped}"); + let args: (u16, String) = mapped.value(py).getattr("args").unwrap().extract().unwrap(); + assert_eq!(args, (expected_status, expected_body.into())); + assert_eq!(status_code(py, &mapped), Some(expected_status)); + }); + } + + #[rstest] + #[case::missing_document_url(Error::MissingDocumentUrl)] + #[case::invalid_features(Error::Features)] + #[case::invalid_provider(Error::InvalidProvider("unknown".into()))] + fn request_faults_are_value_errors_with_status_400(#[case] error: Error) { + Python::initialize(); + Python::attach(|py| { + let report = Report::new(&error).to_string(); + let mapped = to_pyerr(error); assert!(mapped.is_instance_of::(py)); - assert_eq!( - mapped - .value(py) - .getattr("status_code") - .unwrap() - .extract::() - .unwrap(), - 400 - ); + assert_eq!(mapped.value(py).to_string(), report); + assert_eq!(status_code(py, &mapped), Some(400)); + }); + } + + #[rstest] + #[case::missing_api_key(litellm_auth::Error::MissingApiKey { + provider: "Mistral", + environment_variable: "MISTRAL_API_KEY", + } + .into())] + #[case::missing_credential(Error::MissingCredential { + credential: "Azure AI credentials", + guidance: "set AZURE_AI_API_KEY", + })] + fn a_missing_credential_is_a_statusless_runtime_error(#[case] error: Error) { + Python::initialize(); + Python::attach(|py| { + let mapped = to_pyerr(error); + assert!(mapped.is_instance_of::(py)); + assert_eq!(status_code(py, &mapped), None); }); } @@ -135,19 +210,12 @@ mod tests { assert!(mapped.is_instance_of::(py)); assert!( value - .getattr("ocr_request_format_error") + .getattr(REQUEST_FORMAT_ERROR) .unwrap() .extract::() .unwrap() ); - assert_eq!( - value - .getattr("status_code") - .unwrap() - .extract::() - .unwrap(), - 400 - ); + assert_eq!(status_code(py, &mapped), Some(400)); assert_eq!( value .getattr("message") @@ -159,18 +227,11 @@ mod tests { }); } - fn file_read(kind: std::io::ErrorKind) -> Error { - Error::FileRead { - path: "/missing/scan.pdf".into(), - source: std::sync::Arc::new(std::io::Error::new(kind, "disk said no")), - } - } - #[test] fn missing_files_map_to_file_not_found_naming_the_path() { Python::initialize(); Python::attach(|py| { - let mapped = to_pyerr(file_read(std::io::ErrorKind::NotFound)); + let mapped = to_pyerr(file_read(ErrorKind::NotFound)); assert!(mapped.is_instance_of::(py)); assert_eq!( mapped.value(py).to_string(), @@ -183,28 +244,32 @@ mod tests { fn other_file_read_failures_map_to_os_error_with_the_io_message() { Python::initialize(); Python::attach(|py| { - let mapped = to_pyerr(file_read(std::io::ErrorKind::PermissionDenied)); + let mapped = to_pyerr(file_read(ErrorKind::PermissionDenied)); assert!(mapped.is_instance_of::(py)); assert!(!mapped.is_instance_of::(py)); assert_eq!(mapped.value(py).to_string(), "disk said no"); }); } - #[rstest::rstest] + #[rstest] #[case::oversized(Error::TooLarge { limit: 7 })] - #[case::malformed_field(Error::ResponseField { path: "pages[0].index".into() })] - fn response_failures_are_statusless_runtime_errors(#[case] error: Error) { + #[case::malformed_field(Error::response_field( + "pages[0].index", + "-1".parse::().unwrap_err() + ))] + fn response_failures_are_statusless_runtime_errors( + #[case] error: Error, + #[values("status_code", "ocr_request_format_error", "headers")] attribute: &str, + ) { Python::initialize(); Python::attach(|py| { - let message = error.to_string(); + let report = Report::new(&error).to_string(); let mapped = to_pyerr(error); let value = mapped.value(py); - assert!(mapped.is_instance_of::(py)); + assert!(mapped.is_instance_of::(py)); assert!(!mapped.is_instance_of::(py)); - assert_eq!(value.to_string(), message); - for attribute in ["status_code", "ocr_request_format_error", "headers"] { - assert!(!value.hasattr(attribute).unwrap(), "{attribute}"); - } + assert_eq!(value.to_string(), report); + assert!(!value.hasattr(attribute).unwrap()); }); } } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs index fc1b70392e2..a1a56baef8f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs @@ -12,7 +12,7 @@ use pyo3::{ }; use super::{ - errors::to_pyerr as ocr_error_to_pyerr, + errors::route_failure, project::{OcrHostHandles, SETTINGS_ERROR_MARKER, project_request}, }; use crate::{marshal::public_response, routes::RequestView}; @@ -135,16 +135,11 @@ impl PythonBinding for OcrPythonHost { } fn map_error(&self, py: Python<'_>, error: Error) -> PyResult { - if let Error::Secret(source) = &error - && let Some(original) = crate::secrets::python_error(py, source) - { - return Ok(original); - } - Ok(self.map_failure(py, ocr_error_to_pyerr(error))) + route_failure(py, error, |native| Ok(self.map_failure(py, native))) } - fn host_error(error: &PyErr) -> Error { - Error::InvalidRequest(error.to_string()) + fn host_error(_: &PyErr) -> Error { + Error::HostCall } } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs index fe2f83bff46..40287c669eb 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs @@ -6,14 +6,16 @@ use litellm_inference_ocr::{ wire::{OcrWireRequest, consumed_optional_params, decode_document, decode_request_input}, }; use litellm_llms::base_llm::ocr::error::Error; -use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; +use pyo3::{prelude::*, types::PyDict}; use serde_json::{Map, Value}; use super::{document::FileDocumentInput, errors::to_pyerr as ocr_error_to_pyerr}; use crate::{ coercion::FieldSpec, credentials::{self, CallerTokenProvider}, - marshal::{project_optional_fields, python_timeout_seconds, request_input_sources}, + marshal::{ + ArgumentError, project_optional_fields, python_timeout_seconds, request_input_sources, + }, python_settings::{PythonSettings, Snapshot}, }; @@ -30,9 +32,9 @@ struct OcrArguments<'a, 'py> { } impl<'py> OcrArguments<'_, 'py> { - fn lookup(&self, name: &str) -> PyResult> { - litellm_host_python::lookup(self.kwargs, self.bound, name)? - .ok_or_else(|| PyValueError::new_err(format!("missing argument: {name}"))) + fn lookup(&self, name: &'static str) -> PyResult> { + Ok(litellm_host_python::lookup(self.kwargs, self.bound, name)? + .ok_or(ArgumentError::Required(name))?) } fn model(&self) -> PyResult { @@ -92,6 +94,7 @@ impl ProjectedDocument { { ocr_error_to_pyerr(Error::RequestField { path: "document.type".into(), + cause: None, }) } else { error @@ -218,8 +221,10 @@ fn fold_module_globals( pub(super) const SETTINGS_ERROR_MARKER: &str = "native_settings_error"; fn settings_error(py: Python<'_>, error: PyErr) -> PyErr { - error.value(py).setattr(SETTINGS_ERROR_MARKER, true).ok(); - error + match error.value(py).setattr(SETTINGS_ERROR_MARKER, true) { + Ok(()) => error, + Err(failure) => failure, + } } #[cfg(test)] diff --git a/litellm-rust/crates/python-bridge/src/routes/responses/host.rs b/litellm-rust/crates/python-bridge/src/routes/responses/host.rs index 9c430ba7868..eb19d299625 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses/host.rs @@ -141,9 +141,9 @@ impl PythonBinding for ResponsesPythonHost { head: ::StreamHead, ) -> PyResult> { let _ = head; - Err(pyo3::exceptions::PyRuntimeError::new_err( - "native Python Responses streaming is not supported", - )) + Err(crate::errors::route_error_to_pyerr(Error::Unsupported( + "native Python responses streaming", + ))) } fn encode_chunk( @@ -152,16 +152,16 @@ impl PythonBinding for ResponsesPythonHost { chunk: ::Chunk, ) -> PyResult> { let _ = chunk; - Err(pyo3::exceptions::PyRuntimeError::new_err( - "native Python Responses streaming is not supported", - )) + Err(crate::errors::route_error_to_pyerr(Error::Unsupported( + "native Python responses streaming", + ))) } fn map_error(&self, py: Python<'_>, error: Error) -> PyResult { self.host.error(py, error) } fn host_error(error: &PyErr) -> Error { - Error::InvalidRequest(error.to_string().into()) + crate::routes::inference::host_call_error(error) } } diff --git a/litellm-rust/crates/python-bridge/src/routes/token_counter.rs b/litellm-rust/crates/python-bridge/src/routes/token_counter.rs index 2c26311231f..4b10eed087f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/token_counter.rs +++ b/litellm-rust/crates/python-bridge/src/routes/token_counter.rs @@ -2,6 +2,7 @@ use crate::execution::run_async; use std::sync::Arc; use std::{num::NonZero, thread::available_parallelism}; +use litellm_error::{Report, Source}; use litellm_host_python::enter_native; use litellm_token_counter::{ CountableRequest, Error, InputTokenCount, TokenCounter as CoreTokenCounter, @@ -50,10 +51,10 @@ impl TokenCounter { let _slot = encode_slots .acquire_owned() .await - .map_err(|error| Error::Task(error.to_string()))?; + .map_err(|error| Error::Task(Source::new(error)))?; tokio::task::spawn_blocking(move || count_body(&counter, &body)) .await - .map_err(|error| Error::Task(error.to_string()))? + .map_err(|error| Error::Task(Source::new(error)))? }, token_count_error_to_pyerr, ) @@ -70,19 +71,12 @@ fn count_body(counter: &CoreTokenCounter, body: &[u8]) -> Result PyErr { - let message = error.to_string(); - match error { - Error::Load(_) - | Error::Ranks(_) - | Error::UnicodeClasses - | Error::UnsupportedTokenizer(_) => PyValueError::new_err(message), - Error::RequestParse(_) - | Error::MissingInput - | Error::FloatText - | Error::ContentBlock - | Error::ArrayItems - | Error::JsonSerialization(_) - | Error::JsonUtf8(_) => RustBridgeDeclined::new_err(message), - Error::Encode(_) | Error::Decode(_) | Error::Task(_) => PyRuntimeError::new_err(message), + let message = Report::new(&error).to_string(); + if error.is_unsupported() { + RustBridgeDeclined::new_err(message) + } else if error.is_load() { + PyValueError::new_err(message) + } else { + PyRuntimeError::new_err(message) } } diff --git a/litellm-rust/crates/python-bridge/src/secrets/config.rs b/litellm-rust/crates/python-bridge/src/secrets/config.rs index d1a1eafd468..7afd06e2373 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/config.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/config.rs @@ -128,9 +128,7 @@ pub(crate) fn read(py: Python<'_>) -> PyResult { .system .is_some_and(|system| system != backend.system()) { - return Err(pyo3::exceptions::PyValueError::new_err( - "native secret manager system does not match configuration", - )); + return Err(super::ManagerError::SystemMismatch.into()); } Ok(SecretManagerSnapshot { client: SecretManagerClient::Native(Box::new(backend)), diff --git a/litellm-rust/crates/python-bridge/src/secrets/error.rs b/litellm-rust/crates/python-bridge/src/secrets/error.rs index 5ccfd0a1554..ca51ee79921 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/error.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/error.rs @@ -1,7 +1,66 @@ use std::fmt; +use litellm_error::Report; use litellm_secrets::Error; -use pyo3::{exceptions::PyBaseException, prelude::*}; +use pyo3::{ + exceptions::{PyBaseException, PyRuntimeError, PyValueError}, + prelude::*, +}; + +/// A secret manager call the bridge rejects, with the text Python's own managers raise. +#[derive(Debug, thiserror::Error)] +pub(crate) enum ManagerError { + #[error("native secret manager must be recreated after fork")] + Forked, + #[error("unknown secret manager system")] + UnknownSystem(#[source] serde_json::Error), + #[error("invalid secret manager settings")] + Settings(#[source] litellm_error::Source), + #[error("native secret manager system does not match configuration")] + SystemMismatch, + #[error("secret manager does not support provider reads")] + ProviderReads, + #[error("invalid timeout")] + Timeout(#[source] std::time::TryFromFloatSecsError), + #[error("No secret found in CyberArk Secret Manager for {0}")] + CyberarkSecretMissing(String), + #[error("Current secret {0} not found")] + CurrentMissing(String), + #[error("Failed to verify new secret {0}")] + ReplacementMissing(String), + #[error("New secret value mismatch")] + ReplacementMismatch, + #[error("native secret manager mutation is unavailable")] + MutationUnsupported, + #[error(transparent)] + Secret(Error), +} + +impl ManagerError { + pub(crate) fn settings(error: impl std::error::Error + Send + Sync + 'static) -> Self { + Self::Settings(litellm_error::Source::new(error)) + } +} + +impl From for PyErr { + fn from(error: ManagerError) -> Self { + let message = Report::new(&error).to_string(); + match error { + ManagerError::Forked => PyRuntimeError::new_err(message), + ManagerError::UnknownSystem(_) + | ManagerError::Settings(_) + | ManagerError::SystemMismatch + | ManagerError::ProviderReads + | ManagerError::Timeout(_) + | ManagerError::CyberarkSecretMissing(_) + | ManagerError::CurrentMissing(_) + | ManagerError::ReplacementMissing(_) + | ManagerError::ReplacementMismatch + | ManagerError::MutationUnsupported + | ManagerError::Secret(_) => PyValueError::new_err(message), + } + } +} #[derive(thiserror::Error)] #[error("Python secret manager failed")] @@ -14,18 +73,23 @@ impl fmt::Debug for PythonSecretError { } pub(super) fn external_error(py: Python<'_>, error: PyErr) -> Error { - Error::ExternalManager(Box::new(PythonSecretError(error.into_value(py)))) + Error::ExternalManager(litellm_error::Source::new(PythonSecretError( + error.into_value(py), + ))) } -pub(super) fn read_error(py: Python<'_>, error: PyErr) -> Error { - Error::ExternalRead(Box::new(PythonSecretError(error.into_value(py)))) +pub(crate) fn read_error(py: Python<'_>, error: PyErr) -> Error { + Error::ExternalRead(litellm_error::Source::new(PythonSecretError( + error.into_value(py), + ))) } -pub(crate) fn python_error(py: Python<'_>, error: &Error) -> Option { - let (Error::ExternalManager(source) | Error::ExternalRead(source)) = error else { - return None; - }; - source - .downcast_ref::() +/// The exception a Python secret manager raised anywhere in the chain, so a route surfaces it +/// unchanged instead of wrapping it. +pub(crate) fn python_error( + py: Python<'_>, + error: &(dyn std::error::Error + 'static), +) -> Option { + litellm_error::find::(error) .map(|error| PyErr::from_value(error.0.clone_ref(py).into_bound(py).into_any())) } diff --git a/litellm-rust/crates/python-bridge/src/secrets/mod.rs b/litellm-rust/crates/python-bridge/src/secrets/mod.rs index cddd47d08db..57042e6a32e 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/mod.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/mod.rs @@ -11,7 +11,9 @@ mod vault; use std::sync::Arc; -pub(crate) use error::python_error; +#[cfg(test)] +pub(crate) use error::read_error; +pub(crate) use error::{ManagerError, python_error}; use litellm_secrets::source::{EnvironmentSecrets, SecretSource}; use pyo3::prelude::*; use python::PythonSecrets; diff --git a/litellm-rust/crates/python-bridge/src/secrets/mutation.rs b/litellm-rust/crates/python-bridge/src/secrets/mutation.rs index eaab06a5b12..b04060651ac 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/mutation.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/mutation.rs @@ -1,7 +1,8 @@ use super::operations::{PythonMutationError, PythonMutationResponse}; +use litellm_error::Report; use litellm_host_python::{json_loads, to_py}; use litellm_secrets::cyberark; -use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; +use pyo3::{prelude::*, types::PyDict}; pub(super) fn mutation_value( result: Result, @@ -23,18 +24,7 @@ pub(super) fn mutation_value( &serde_json::json!({"status": "error", "message": message}), ) } - Err(PythonMutationError::CurrentMissing(name)) => Err(PyValueError::new_err(format!( - "Current secret {name} not found" - ))), - Err(PythonMutationError::ReplacementMissing(name)) => Err(PyValueError::new_err(format!( - "Failed to verify new secret {name}" - ))), - Err(PythonMutationError::ReplacementMismatch) => { - Err(PyValueError::new_err("New secret value mismatch")) - } - Err(PythonMutationError::Unsupported) => Err(PyValueError::new_err( - "native secret manager mutation is unavailable", - )), + Err(PythonMutationError::Manager(error)) => Err(error.into()), }) } @@ -43,18 +33,35 @@ fn cyberark_failure( name: &str, failure: cyberark::WriteFailure, ) -> PyResult { - let message = match failure.source { + let message = if failure.source.is_unsafe_secret_name() { + format!("Invalid secret_name {}", name.into_pyobject(py)?.repr()?) + } else { + cyberark_message( + py, + failure.source, + failure.authentication, + failure.request_url.as_ref(), + )? + }; + Ok(if failure.authentication { + format!("Could not authenticate to CyberArk Conjur: {message}") + } else { + message + }) +} + +fn cyberark_message( + py: Python<'_>, + source: cyberark::Error, + authentication: bool, + request_url: Option<&reqwest::Url>, +) -> PyResult { + Ok(match source { cyberark::Error::Status(status) | cyberark::Error::AuthStatus(status) => { - let url = failure - .request_url - .as_ref() - .map_or("", reqwest::Url::as_str); + let url = request_url.map_or("", reqwest::Url::as_str); http_message(py, "POST", url, status)? } - cyberark::Error::Operation(litellm_secrets_types::Error::UnsafeSecretName) => { - format!("Invalid secret_name {}", name.into_pyobject(py)?.repr()?) - } - cyberark::Error::Http(source) if failure.authentication => match os_error_code(&source) { + cyberark::Error::Http(source) if authentication => match os_error_code(&source) { Some(code) => { let reason = py.import("os")?.getattr("strerror")?.call1((code,))?; py.import("builtins")? @@ -63,17 +70,12 @@ fn cyberark_failure( .str()? .extract()? } - None => cyberark::Error::Http(source).to_string(), + None => Report::new(&cyberark::Error::Http(source)).to_string(), }, cyberark::Error::Http(source) if source.is_connect() => { "All connection attempts failed".to_owned() } - source => source.to_string(), - }; - Ok(if failure.authentication { - format!("Could not authenticate to CyberArk Conjur: {message}") - } else { - message + source => Report::new(&source).to_string(), }) } diff --git a/litellm-rust/crates/python-bridge/src/secrets/operations.rs b/litellm-rust/crates/python-bridge/src/secrets/operations.rs index 8e338741aff..b458e3f34e8 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/operations.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/operations.rs @@ -4,6 +4,8 @@ use litellm_secrets::cyberark::AuthenticationRetry; use litellm_secrets::{Error, SecretManager}; use litellm_secrets_types::{PythonSecretRead, SecretOperationContext}; +use super::ManagerError; + pub(super) struct PythonReadRequest { pub secret_name: String, pub primary_secret_name: Option, @@ -62,15 +64,18 @@ pub(super) async fn read_python_provider( #[derive(Debug)] pub(super) enum PythonMutationError { - Unsupported, + Manager(ManagerError), Vault(Box), CyberarkWrite { name: String, failure: Box, }, - CurrentMissing(String), - ReplacementMissing(String), - ReplacementMismatch, +} + +impl From for PythonMutationError { + fn from(error: ManagerError) -> Self { + Self::Manager(error) + } } pub(super) async fn write_python_provider( @@ -89,7 +94,7 @@ pub(super) async fn write_python_provider( })?; Ok(write_success(name)) } - _ => Err(PythonMutationError::Unsupported), + _ => Err(ManagerError::MutationUnsupported.into()), } } @@ -115,7 +120,7 @@ pub(super) async fn delete_python_provider( "message": "CyberArk Conjur does not support direct secret deletion. Use policy updates to remove variables.", })) } - _ => Err(PythonMutationError::Unsupported), + _ => Err(ManagerError::MutationUnsupported.into()), } } @@ -138,7 +143,7 @@ pub(super) async fn rotate_python_provider( .flatten() .is_none() { - return Err(PythonMutationError::CurrentMissing(current_name.to_owned())); + return Err(ManagerError::CurrentMissing(current_name.to_owned()).into()); } client .write_with_retry( @@ -157,16 +162,16 @@ pub(super) async fn rotate_python_provider( .await .ok() .flatten() - .ok_or_else(|| PythonMutationError::ReplacementMissing(new_name.to_owned()))?; + .ok_or_else(|| ManagerError::ReplacementMissing(new_name.to_owned()))?; if actual != *value { - return Err(PythonMutationError::ReplacementMismatch); + return Err(ManagerError::ReplacementMismatch.into()); } if current_name != new_name { client.invalidate_cached_secret(current_name).await; } Ok(write_success(new_name)) } - _ => Err(PythonMutationError::Unsupported), + _ => Err(ManagerError::MutationUnsupported.into()), } } fn write_success(name: &str) -> serde_json::Value { diff --git a/litellm-rust/crates/python-bridge/src/secrets/provider.rs b/litellm-rust/crates/python-bridge/src/secrets/provider.rs index 5f4cd193fcf..ff10f596311 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/provider.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/provider.rs @@ -6,7 +6,9 @@ use litellm_secrets_types::{ AwsOperationContext, CyberarkOperationContext, GoogleOperationContext, HashicorpOperationContext, SecretOperationContext, }; -use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; +use pyo3::{prelude::*, types::PyDict}; + +use super::ManagerError; pub(super) fn read_request( system: KeyManagementSystem, @@ -39,9 +41,7 @@ pub(super) fn read_request( SecretOperationContext::Google(GoogleOperationContext::default()) } _ => { - return Err(PyValueError::new_err( - "secret manager does not support provider reads", - )); + return Err(ManagerError::ProviderReads.into()); } }; Ok(PythonReadRequest { @@ -89,9 +89,7 @@ fn aws_context( fn read_timeout(value: Option<&Bound<'_, PyAny>>) -> PyResult> { read_timeout_seconds(value)? - .map(|value| { - Duration::try_from_secs_f64(value).map_err(|_| PyValueError::new_err("invalid timeout")) - }) + .map(|value| Ok(Duration::try_from_secs_f64(value).map_err(ManagerError::Timeout)?)) .transpose() } diff --git a/litellm-rust/crates/python-bridge/src/secrets/runtime.rs b/litellm-rust/crates/python-bridge/src/secrets/runtime.rs index 72ba91ded16..b956ec297b6 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/runtime.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/runtime.rs @@ -7,12 +7,9 @@ use litellm_secrets::{ read_secret_from_python_manager, }; use litellm_secrets_types::PythonSecretRead; -use pyo3::{ - exceptions::{PyAttributeError, PyRuntimeError, PyValueError}, - prelude::*, - types::PyDict, -}; +use pyo3::{exceptions::PyAttributeError, prelude::*, types::PyDict}; +use super::ManagerError; use crate::execution::{run_async_value, run_sync_value}; const SECRET_MANAGER_MODULE: &str = "litellm.rust_bridge.host.secret_manager"; @@ -35,9 +32,7 @@ pub(crate) struct NativeSecretManager { impl NativeSecretManager { pub(super) fn backend(&self) -> PyResult { if self.pid != std::process::id() { - return Err(PyRuntimeError::new_err( - "native secret manager must be recreated after fork", - )); + return Err(ManagerError::Forked.into()); } Ok(self.backend.clone()) } @@ -60,7 +55,7 @@ impl NativeSecretManager { enterprise_enabled, ) .await - .map_err(|error| PyValueError::new_err(error.to_string())) + .map_err(|error| ManagerError::Secret(error).into()) })?; Ok(Self { backend, @@ -82,7 +77,7 @@ impl NativeSecretManager { enterprise_enabled: bool, ) -> PyResult { let system = serde_json::from_value(serde_json::Value::String(system.to_owned())) - .map_err(|_| PyValueError::new_err("unknown secret manager system"))?; + .map_err(ManagerError::UnknownSystem)?; let settings = parse_settings(settings)?; Self::build( py, @@ -146,7 +141,7 @@ impl NativeSecretManager { .collect::>>()?; let settings = from_py::>(&config.getattr("settings")?) - .map_err(|_| PyValueError::new_err("invalid secret manager settings"))?; + .map_err(ManagerError::settings)?; let attributes = config .getattr("settings_attributes")? .extract::>()?; @@ -161,11 +156,11 @@ impl NativeSecretManager { system: serde_json::from_value(serde_json::Value::String( config.getattr("system")?.extract()?, )) - .map_err(|_| PyValueError::new_err("unknown secret manager system"))?, + .map_err(ManagerError::UnknownSystem)?, settings: serde_json::from_value(serde_json::Value::Object( settings.into_iter().chain(setting_overrides).collect(), )) - .map_err(|_| PyValueError::new_err("invalid secret manager settings"))?, + .map_err(ManagerError::settings)?, environment: captured .into_iter() .filter(|(key, _)| !environment_attributes.contains_key(key)) @@ -190,12 +185,8 @@ impl NativeSecretManager { } #[getter] - fn system(&self) -> String { - serde_json::to_value(self.configuration.system) - .expect("serializable system") - .as_str() - .expect("string system") - .to_owned() + fn system(&self) -> &'static str { + self.configuration.system.into() } #[pyo3(signature = (name, settings=None))] @@ -238,7 +229,7 @@ impl NativeSecretManager { run_sync_value(py, async move { super::operations::read_python_provider(&backend, &request, &ProcessEnvironment) .await - .map_err(|error| PyValueError::new_err(error.to_string())) + .map_err(|error| ManagerError::Secret(error).into()) .and_then(|value| python_secret_value(value, &request.secret_name)) }) } @@ -264,7 +255,7 @@ impl NativeSecretManager { run_async_value(py, async move { super::operations::read_python_provider(&backend, &request, &ProcessEnvironment) .await - .map_err(|error| PyValueError::new_err(error.to_string())) + .map_err(|error| ManagerError::Secret(error).into()) .and_then(|value| python_secret_value(value, &request.secret_name)) }) } @@ -438,8 +429,8 @@ fn cached( fn parse_settings(value: Option<&Bound<'_, PyAny>>) -> PyResult { value .map(|value| { - serde_json::from_value(from_py::(value)?) - .map_err(|_| PyValueError::new_err("invalid secret manager settings")) + Ok(serde_json::from_value(from_py::(value)?) + .map_err(ManagerError::settings)?) }) .transpose() .map(Option::unwrap_or_default) @@ -466,11 +457,10 @@ fn python_read_error( name: &str, error: litellm_secrets::Error, ) -> PyErr { - let message = match (system, error) { + match (system, error) { (KeyManagementSystem::Cyberark, litellm_secrets::Error::ManagedSecretMissing) => { - format!("No secret found in CyberArk Secret Manager for {name}") + ManagerError::CyberarkSecretMissing(name.to_owned()).into() } - (_, error) => error.to_string(), - }; - PyValueError::new_err(message) + (_, error) => ManagerError::Secret(error).into(), + } } diff --git a/litellm-rust/crates/python-bridge/src/secrets/vault.rs b/litellm-rust/crates/python-bridge/src/secrets/vault.rs index 79281662848..30e3bd2c413 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/vault.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/vault.rs @@ -2,8 +2,11 @@ mod operation; pub(super) use operation::{Failure, FailureKind, FailureStage, delete, rotate, write}; -use litellm_core_utils::exception_mapping_utils::{TimeoutPhase, timeout_message}; -use litellm_secrets::hashicorp::{Error, RawOperationError}; +use litellm_core_utils::exception_mapping_utils::{ + TimeoutPhase, timeout_message as transport_timeout_message, +}; +use litellm_error::Report; +use litellm_secrets::hashicorp::RawOperationError; use pyo3::prelude::*; use super::mutation::{error_value, http_message, json_value}; @@ -16,40 +19,39 @@ pub(super) fn failure_value( failure: Failure, context: &ErrorContext, ) -> PyResult> { - let message = match *failure.kind { - FailureKind::Native(RawOperationError::Http { - method, - url, - status, - body, - }) => match failure.stage { - FailureStage::Current(name) if status == 404 => { - format!("Current secret {name} not found") + let message = match (&*failure.kind, failure.stage) { + (FailureKind::Native(error), stage) if let Some(response) = error.http_failure() => { + match stage { + FailureStage::Current(name) if response.status == 404 => { + format!("Current secret {name} not found") + } + FailureStage::Replacement(name) if response.status == 404 => { + format!("Failed to verify new secret {name}") + } + FailureStage::Current(_) => format!( + "HTTP error occurred while checking current secret: {}", + response_text(py, &response.body)? + ), + FailureStage::Replacement(_) => format!( + "HTTP error occurred while verifying new secret: {}", + response_text(py, &response.body)? + ), + FailureStage::Mutation => { + http_message(py, &response.method, &response.url, response.status)? + } } - FailureStage::Replacement(name) if status == 404 => { - format!("Failed to verify new secret {name}") - } - FailureStage::Current(_) => format!( - "HTTP error occurred while checking current secret: {}", - response_text(py, &body)? - ), - FailureStage::Replacement(_) => format!( - "HTTP error occurred while verifying new secret: {}", - response_text(py, &body)? - ), - FailureStage::Mutation => http_message(py, &method, &url, status)?, - }, - FailureKind::ValueMismatch { expected, actual } => { - let actual = json_value(py, &actual)?; + } + (FailureKind::ValueMismatch { expected, actual }, _) => { + let actual = json_value(py, actual)?; format!( "New secret value mismatch. Expected: {}, Got: {}", expected.expose(), actual.bind(py).str()? ) } - kind => { + (kind, stage) => { let message = cause_message(py, kind, context)?; - match failure.stage { + match stage { FailureStage::Current(_) => { format!("Error checking current secret: {message}") } @@ -63,86 +65,91 @@ pub(super) fn failure_value( error_value(py, message) } -fn cause_message(py: Python<'_>, kind: FailureKind, context: &ErrorContext) -> PyResult { +fn cause_message(py: Python<'_>, kind: &FailureKind, context: &ErrorContext) -> PyResult { Ok(match kind { - FailureKind::Native(RawOperationError::Local(error)) => error.to_string(), + FailureKind::Native(error) => native_message(py, error, context)?, FailureKind::UnsafeName(name) => { format!("Invalid secret_name {}", name.into_pyobject(py)?.repr()?) } - FailureKind::Native(RawOperationError::Timeout { method, elapsed }) => { - if method == "POST" { - let kwargs = pyo3::types::PyDict::new(py); - kwargs.set_item( - "message", - timeout_message( - TimeoutPhase::Read, - context.timeout, - Some(elapsed.as_secs_f64()), - ), - )?; - kwargs.set_item("model", "default-model-name")?; - kwargs.set_item("llm_provider", "litellm-httpx-handler")?; - py.import("litellm.exceptions")? - .getattr("Timeout")? - .call((), Some(&kwargs))? - .str()? - .extract()? - } else if context.aiohttp { - "Timeout on reading data from socket".to_owned() - } else { - String::new() - } - } - FailureKind::Native(RawOperationError::Transport(source)) => { - if let Some(error) = request_error(&source) { - if error.is_timeout() { - String::new() - } else if error.is_connect() { - "All connection attempts failed".to_owned() - } else { - "HashiCorp Vault request failed".to_owned() - } - } else { - "HashiCorp Vault request failed".to_owned() - } - } FailureKind::MissingGet(value) => { - let value = json_value(py, &value)?; + let value = json_value(py, value)?; match value.bind(py).getattr("get") { Err(error) => error.value(py).str()?.extract()?, Ok(_) => "HashiCorp Vault response payload is malformed".to_owned(), } } - FailureKind::Json(body) => match json_value(py, &body) { + FailureKind::Json(body) => match json_value(py, body) { Err(error) => error.value(py).str()?.extract()?, Ok(_) => "HashiCorp Vault response payload is malformed".to_owned(), }, - FailureKind::Native(RawOperationError::Authentication { - source, - url, - certificate, - }) => { - let message = match source { - Error::LoginStatus { status } => http_message(py, "POST", &url, status)?, - error => error.to_string(), - }; - let mechanism = if certificate { "TLS cert" } else { "AppRole" }; - format!("Could not authenticate to Vault via {mechanism}: {message}") - } - FailureKind::Native(RawOperationError::Http { - method, - url, - status, - .. - }) => http_message(py, &method, &url, status)?, FailureKind::ValueMismatch { .. } => "New secret value mismatch".to_owned(), }) } -fn request_error<'a>(error: &'a (dyn std::error::Error + 'static)) -> Option<&'a reqwest::Error> { - error - .downcast_ref::() - .or_else(|| error.source().and_then(request_error)) +fn native_message( + py: Python<'_>, + error: &RawOperationError, + context: &ErrorContext, +) -> PyResult { + if let Some(error) = error.local() { + return Ok(Report::new(error).to_string()); + } + if let Some((method, elapsed)) = error.timeout() { + return timeout_message(py, method, elapsed, context); + } + if let Some(response) = error.http_failure() { + return http_message(py, &response.method, &response.url, response.status); + } + if let Some(login) = error.authentication() { + let message = match login.source.login_status() { + Some(status) => http_message(py, "POST", login.url, status)?, + None => Report::new(login.source).to_string(), + }; + let mechanism = if login.certificate { + "TLS cert" + } else { + "AppRole" + }; + return Ok(format!( + "Could not authenticate to Vault via {mechanism}: {message}" + )); + } + Ok(match error.transport() { + Some(error) if error.status() == Some(408) => String::new(), + Some(error) if error.is_connect() => "All connection attempts failed".to_owned(), + Some(_) | None => "HashiCorp Vault request failed".to_owned(), + }) +} + +fn timeout_message( + py: Python<'_>, + method: &str, + elapsed: std::time::Duration, + context: &ErrorContext, +) -> PyResult { + if method != "POST" { + return Ok(if context.aiohttp { + "Timeout on reading data from socket".to_owned() + } else { + String::new() + }); + } + let kwargs = pyo3::types::PyDict::new(py); + kwargs.set_item( + "message", + transport_timeout_message( + TimeoutPhase::Read, + context.timeout, + Some(elapsed.as_secs_f64()), + ), + )?; + kwargs.set_item("model", "default-model-name")?; + kwargs.set_item("llm_provider", "litellm-httpx-handler")?; + py.import("litellm.exceptions")? + .getattr("Timeout")? + .call((), Some(&kwargs))? + .str()? + .extract() } fn response_text(py: Python<'_>, body: &[u8]) -> PyResult { diff --git a/litellm-rust/crates/python-bridge/src/secrets/vault/operation.rs b/litellm-rust/crates/python-bridge/src/secrets/vault/operation.rs index cf56beb5efe..41658f0baf8 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/vault/operation.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/vault/operation.rs @@ -2,7 +2,7 @@ use std::collections::HashMap; use litellm_secrets::{ SecretValue, - hashicorp::{Error, HashicorpVault, RawOperationError}, + hashicorp::{HashicorpVault, RawOperationError}, }; use litellm_secrets_types::{HashicorpOperationContext, SecretWriteContext}; use serde_json::value::RawValue; @@ -53,12 +53,10 @@ impl Failure { } fn native_failure(name: &str, error: RawOperationError) -> Failure { - match error { - RawOperationError::Local(Error::InvalidSecretName(_)) => { - FailureKind::UnsafeName(name.to_owned()).into() - } - error => FailureKind::Native(error).into(), + if error.is_unsafe_secret_name() { + return FailureKind::UnsafeName(name.to_owned()).into(); } + FailureKind::Native(error).into() } pub(crate) async fn write( diff --git a/litellm-rust/crates/python-bridge/src/tokenizer.rs b/litellm-rust/crates/python-bridge/src/tokenizer.rs index df219d55eb6..2026b2fdc42 100644 --- a/litellm-rust/crates/python-bridge/src/tokenizer.rs +++ b/litellm-rust/crates/python-bridge/src/tokenizer.rs @@ -8,16 +8,18 @@ use std::sync::Arc; #[cfg(feature = "fast")] use std::sync::OnceLock; +use litellm_error::Report; use litellm_host_python::{enter_native, release_gil}; #[cfg(feature = "fast")] use litellm_token_counter::fast::{FastCounter, FastTokenizer}; use litellm_token_counter::{Error, TextCodec}; use pyo3::{exceptions::PyUnicodeEncodeError, prelude::*, types::PyString}; -#[cfg(any(feature = "tiktoken", feature = "huggingface"))] -use pyo3::exceptions::PyValueError; #[cfg(feature = "huggingface")] -use pyo3::{exceptions::PyIOError, types::PyDict}; +use pyo3::{ + exceptions::{PyIOError, PyTypeError, PyValueError}, + types::PyDict, +}; #[cfg(feature = "tiktoken")] use pyo3::{ exceptions::{PyKeyError, PyRuntimeError}, @@ -51,6 +53,70 @@ pub(crate) fn load_tiktoken(py: Python<'_>, encoding: &str) -> PyResult for PyErr { + fn from(error: TokenizerError) -> Self { + let message = Report::new(&error).to_string(); + match error { + #[cfg(not(all(feature = "tiktoken", feature = "huggingface")))] + TokenizerError::MissingFeature(_) => RustBridgeDeclined::new_err(message), + #[cfg(all(feature = "tiktoken", feature = "huggingface"))] + TokenizerError::NotTiktoken | TokenizerError::NotHuggingFace => { + PyValueError::new_err(message) + } + #[cfg(feature = "tiktoken")] + TokenizerError::NoVocabulary | TokenizerError::EncodeSpecial(_) => { + PyRuntimeError::new_err(message) + } + #[cfg(feature = "tiktoken")] + TokenizerError::DecodeBytes(_) => PyKeyError::new_err(message), + #[cfg(feature = "huggingface")] + TokenizerError::Direction { .. } | TokenizerError::EncodingJson(_) => { + PyValueError::new_err(message) + } + #[cfg(feature = "huggingface")] + TokenizerError::PretokenizedMismatch => PyTypeError::new_err(message), + #[cfg(feature = "huggingface")] + TokenizerError::Download(_) => PyIOError::new_err(message), + } + } +} + pub(crate) enum Codec { #[cfg(feature = "tiktoken")] Tiktoken(TiktokenTokenizer), @@ -100,9 +166,7 @@ impl Tokenizer { #[cfg(not(feature = "tiktoken"))] { let _ = (py, encoding); - Err(RustBridgeDeclined::new_err( - "tokenizer backend requires the tiktoken feature", - )) + Err(TokenizerError::MissingFeature("tiktoken").into()) } } @@ -118,9 +182,7 @@ impl Tokenizer { #[cfg(not(feature = "huggingface"))] { let _ = (py, tokenizer_json); - Err(RustBridgeDeclined::new_err( - "tokenizer backend requires the huggingface feature", - )) + Err(TokenizerError::MissingFeature("huggingface").into()) } } @@ -144,16 +206,14 @@ impl Tokenizer { .getattr("hf_hub_download")? .call((), Some(&kwargs))? .extract()?; - let json = - release_gil(py, || std::fs::read_to_string(path)).map_err(PyIOError::new_err)?; + let json = release_gil(py, || std::fs::read_to_string(path)) + .map_err(TokenizerError::Download)?; Self::from_json(py, &json) } #[cfg(not(feature = "huggingface"))] { let _ = (py, identifier, revision, token); - Err(RustBridgeDeclined::new_err( - "tokenizer backend requires the huggingface feature", - )) + Err(TokenizerError::MissingFeature("huggingface").into()) } } @@ -199,7 +259,7 @@ impl Tokenizer { let tokenizer = self.tiktoken()?; let text = self.text(text)?; release_gil(py, || tokenizer.encode_special(&text, &allowed)) - .map_err(PyRuntimeError::new_err) + .map_err(|error| TokenizerError::EncodeSpecial(error.into()).into()) } /// tiktoken's `encode_with_unstable`: `(stable_tokens, completions)`. @@ -258,8 +318,8 @@ impl Tokenizer { fn decode_bytes<'py>(&self, py: Python<'py>, ids: Vec) -> PyResult> { enter_native()?; let tokenizer = self.tiktoken()?; - let bytes = - release_gil(py, || tokenizer.decode_bytes(&ids)).map_err(PyKeyError::new_err)?; + let bytes = release_gil(py, || tokenizer.decode_bytes(&ids)) + .map_err(|error| TokenizerError::DecodeBytes(error.into()))?; Ok(PyBytes::new(py, &bytes)) } @@ -480,15 +540,16 @@ impl Tokenizer { match *self.inner { Codec::Tiktoken(ref tokenizer) => Ok(tokenizer), #[cfg(feature = "huggingface")] - Codec::HuggingFace(_) => Err(PyValueError::new_err("requires a tiktoken encoding")), + Codec::HuggingFace(_) => Err(TokenizerError::NotTiktoken.into()), } } #[cfg(feature = "tiktoken")] fn vocabulary(&self) -> PyResult<&Vocabulary> { - self.tiktoken()?.vocabulary().ok_or_else(|| { - PyRuntimeError::new_err("this encoding was built without its vocabulary") - }) + Ok(self + .tiktoken()? + .vocabulary() + .ok_or(TokenizerError::NoVocabulary)?) } #[cfg(feature = "huggingface")] @@ -496,7 +557,7 @@ impl Tokenizer { match *self.inner { Codec::HuggingFace(ref tokenizer) => Ok(tokenizer), #[cfg(feature = "tiktoken")] - Codec::Tiktoken(_) => Err(PyValueError::new_err("requires a Hugging Face tokenizer")), + Codec::Tiktoken(_) => Err(TokenizerError::NotHuggingFace.into()), } } } @@ -530,21 +591,21 @@ impl Sequence { match (self, is_pretokenized) { (Self::Text(text), false) => Ok(text.into()), (Self::Words(words), true) => Ok(words.into()), - _ => Err(pyo3::exceptions::PyTypeError::new_err( - "input must match is_pretokenized", - )), + _ => Err(TokenizerError::PretokenizedMismatch.into()), } } } #[cfg(feature = "huggingface")] -fn direction(value: &str, left: T, right: T, what: &str) -> PyResult { +fn direction(value: &str, left: T, right: T, what: &'static str) -> PyResult { match value { "left" => Ok(left), "right" => Ok(right), - other => Err(PyValueError::new_err(format!( - "invalid {what} direction {other:?}: expected 'left' or 'right'" - ))), + other => Err(TokenizerError::Direction { + what, + value: other.to_owned(), + } + .into()), } } @@ -564,7 +625,7 @@ impl HuggingFaceEncoding { fn new(json: Option<&str>) -> PyResult { let inner = match json { Some(json) => encoding_from_json(json) - .map_err(|error| PyValueError::new_err(error.to_string()))?, + .map_err(|error| TokenizerError::EncodingJson(error.into()))?, None => Encoding::default(), }; Ok(Self { inner }) @@ -586,7 +647,7 @@ impl HuggingFaceEncoding { py: Python<'py>, ) -> PyResult<(Bound<'py, pyo3::types::PyType>, (String,))> { let json = encoding_to_json(&self.inner) - .map_err(|error| PyValueError::new_err(error.to_string()))?; + .map_err(|error| TokenizerError::EncodingJson(error.into()))?; Ok((py.get_type::(), (json,))) } diff --git a/litellm-rust/crates/python-compat/src/error.rs b/litellm-rust/crates/python-compat/src/error.rs index 43990eb7602..6547907056e 100644 --- a/litellm-rust/crates/python-compat/src/error.rs +++ b/litellm-rust/crates/python-compat/src/error.rs @@ -10,15 +10,17 @@ pub enum Error { Unhashable(&'static str), #[error("value nests deeper than {MAX_DEPTH} levels")] TooDeep, - #[error("invalid pickle: {0}")] - InvalidPickle(String), + #[error("invalid pickle")] + InvalidPickle(#[source] serde_pickle::Error), + #[error("could not write pickle")] + PickleWrite(#[source] serde_pickle::Error), #[error("Object of type {0} is not JSON serializable")] NotJsonSerializable(&'static str), #[error("keys must be str, int, float, bool or None, not {0}")] InvalidJsonKey(&'static str), #[error("Out of range float values are not JSON compliant")] NonFiniteFloat, - #[error("integer does not fit in the JSON number range")] + #[error("integer does not fit in 64 bits")] IntegerOutOfRange, #[error("Object of type {0} cannot be pickled as plain data")] NotPicklable(&'static str), diff --git a/litellm-rust/crates/python-compat/src/literal.rs b/litellm-rust/crates/python-compat/src/literal.rs index 12e9036572f..336c6dfcb2e 100644 --- a/litellm-rust/crates/python-compat/src/literal.rs +++ b/litellm-rust/crates/python-compat/src/literal.rs @@ -58,13 +58,6 @@ impl Term { kind: Kind::Other, } } - - fn is_number(&self) -> bool { - matches!( - self.value, - Value::Int(_) | Value::Float(_) | Value::Complex { .. } - ) - } } struct Parser<'a> { @@ -204,18 +197,25 @@ impl Parser<'_> { self.enter()?; let operand = self.unary()?; self.depth -= 1; - if operand.kind != Kind::Constant || !operand.is_number() { + if operand.kind != Kind::Constant { return Err(Error::InvalidLiteral(self.offset + at)); } - let value = if !negative { - operand.value - } else { - match operand.value { - Value::Int(value) => Value::Int(-value), - Value::Float(value) => Value::Float(-value), - Value::Complex { re, im } => Value::Complex { re: -re, im: -im }, - _ => unreachable!("checked numeric above"), - } + let value = match (operand.value, negative) { + (value @ (Value::Int(_) | Value::Float(_) | Value::Complex { .. }), false) => value, + (Value::Int(value), true) => Value::Int(-value), + (Value::Float(value), true) => Value::Float(-value), + (Value::Complex { re, im }, true) => Value::Complex { re: -re, im: -im }, + ( + Value::None + | Value::Bool(_) + | Value::Str(_) + | Value::Bytes(_) + | Value::List(_) + | Value::Tuple(_) + | Value::Dict(_) + | Value::Set(_), + _, + ) => return Err(Error::InvalidLiteral(self.offset + at)), }; Ok(Term { value, diff --git a/litellm-rust/crates/python-compat/src/pickle.rs b/litellm-rust/crates/python-compat/src/pickle.rs index bff7442dc6b..dd7ac6a0e4f 100644 --- a/litellm-rust/crates/python-compat/src/pickle.rs +++ b/litellm-rust/crates/python-compat/src/pickle.rs @@ -22,10 +22,8 @@ pub fn loads(data: &[u8]) -> Result { let mut deserializer = serde_pickle::Deserializer::new(data, DeOptions::new()); let value = Seed { depth: 0 } .deserialize(&mut deserializer) - .map_err(|error| Error::InvalidPickle(error.to_string()))?; - deserializer - .end() - .map_err(|error| Error::InvalidPickle(error.to_string()))?; + .map_err(Error::InvalidPickle)?; + deserializer.end().map_err(Error::InvalidPickle)?; Ok(value) } @@ -33,8 +31,7 @@ pub fn loads(data: &[u8]) -> Result { /// default. Sets and complex numbers are rejected rather than silently changing type. pub fn dumps(value: &Value) -> Result, Error> { check_picklable(value, 0)?; - serde_pickle::to_vec(&Pickled(value), SerOptions::new()) - .map_err(|error| Error::InvalidPickle(error.to_string())) + serde_pickle::to_vec(&Pickled(value), SerOptions::new()).map_err(Error::PickleWrite) } fn check_picklable(value: &Value, depth: usize) -> Result<(), Error> { @@ -66,7 +63,7 @@ impl Serialize for Pickled<'_> { Value::Bool(value) => serializer.serialize_bool(*value), Value::Int(value) => { let value = i64::try_from(value) - .map_err(|_| serde::ser::Error::custom("integer out of i64 range"))?; + .map_err(|_| serde::ser::Error::custom(Error::IntegerOutOfRange))?; serializer.serialize_i64(value) } Value::Float(value) => serializer.serialize_f64(*value), @@ -94,7 +91,7 @@ impl Serialize for Pickled<'_> { map.end() } value @ (Value::Set(_) | Value::Complex { .. }) => Err(serde::ser::Error::custom( - format!("{} cannot be pickled as plain data", value.type_name()), + Error::NotPicklable(value.type_name()), )), } } diff --git a/litellm-rust/crates/secrets-aws/src/auth.rs b/litellm-rust/crates/secrets-aws/src/auth.rs index 06019cc3c9d..66f4f51bda9 100644 --- a/litellm-rust/crates/secrets-aws/src/auth.rs +++ b/litellm-rust/crates/secrets-aws/src/auth.rs @@ -79,9 +79,7 @@ impl ProvideCredentials for Credentials { self.auth .resolve_credentials(self.config.clone(), &|name| self.environment.get(name)) .await - .map_err(|_| { - CredentialsError::provider_error("secret manager authentication failed") - }) + .map_err(CredentialsError::provider_error) }) } } diff --git a/litellm-rust/crates/secrets-aws/src/error.rs b/litellm-rust/crates/secrets-aws/src/error.rs index c8ca671e9b5..c537da204dd 100644 --- a/litellm-rust/crates/secrets-aws/src/error.rs +++ b/litellm-rust/crates/secrets-aws/src/error.rs @@ -11,7 +11,7 @@ pub enum Error { #[error("KMS response has no plaintext")] MissingPlaintext, #[error("AWS request timed out")] - Timeout, + Timeout(#[source] #[redact] Box>), #[error("AWS KMS decrypt failed")] Decrypt(#[from] #[redact] Box>), #[error("AWS Secrets Manager read failed")] @@ -33,7 +33,7 @@ pub enum Error { #[error("AWS Secrets Manager response has no string payload")] MissingString, #[error("primary secret is not a JSON object")] - PrimarySecret, + PrimarySecret(#[source] #[redact] Option), #[error(transparent)] Operation(#[from] litellm_secrets_types::Error), } diff --git a/litellm-rust/crates/secrets-aws/src/secret_manager/read.rs b/litellm-rust/crates/secrets-aws/src/secret_manager/read.rs index b7583f9785f..d7f10a9847e 100644 --- a/litellm-rust/crates/secrets-aws/src/secret_manager/read.rs +++ b/litellm-rust/crates/secrets-aws/src/secret_manager/read.rs @@ -72,7 +72,7 @@ impl AwsSecretsManagerV2 { } let client = self.client_for_context(context)?; let value = match Self::read_with_client(&client, name, ReadPolicy::Python).await { - Err(Error::Read(_) | Error::MissingString | Error::Timeout) => None, + Err(Error::Read(_) | Error::MissingString | Error::Timeout(_)) => None, result => result?, }; Ok(PythonSecretRead::Value(value.map(Secret::String))) @@ -124,9 +124,10 @@ impl AwsSecretsManagerV2 { Self::read_with_client(&self.client, name, policy).await, policy, ) { - (Err(Error::Read(_) | Error::MissingString | Error::Timeout), ReadPolicy::Python) => { - Ok(None) - } + ( + Err(Error::Read(_) | Error::MissingString | Error::Timeout(_)), + ReadPolicy::Python, + ) => Ok(None), (result, _) => result, } } @@ -173,7 +174,7 @@ impl AwsSecretsManagerV2 { aws_sdk_secretsmanager::error::SdkError::TimeoutError(_) ) || matches!(&error, aws_sdk_secretsmanager::error::SdkError::DispatchFailure(failure) if failure.is_timeout()) => { - Err(Error::Timeout) + Err(Error::Timeout(Box::new(error))) } Err(error) if error @@ -216,9 +217,9 @@ fn resolve_payload(payload: PythonSecretRead, name: &str) -> Result Ok(value), PythonSecretRead::PrimaryJson(document) => { - let object: Value = - serde_json::from_str(document.expose()).map_err(|_| Error::PrimarySecret)?; - let object = object.as_object().ok_or(Error::PrimarySecret)?; + let object: Value = serde_json::from_str(document.expose()) + .map_err(|error| Error::PrimarySecret(Some(error)))?; + let object = object.as_object().ok_or(Error::PrimarySecret(None))?; Ok(object.get(name).cloned().map(Secret::from_json)) } } diff --git a/litellm-rust/crates/secrets-aws/tests/secret_manager/reads.rs b/litellm-rust/crates/secrets-aws/tests/secret_manager/reads.rs index 3c39924f09f..30aec467543 100644 --- a/litellm-rust/crates/secrets-aws/tests/secret_manager/reads.rs +++ b/litellm-rust/crates/secrets-aws/tests/secret_manager/reads.rs @@ -88,7 +88,7 @@ async fn failed_read_returns_none_but_invalid_primary_json_is_an_error( manager .read_secret_for_resolver("KEY", Some("invalid"), &|_: &str| None) .await, - Err(Error::PrimarySecret) + Err(Error::PrimarySecret(_)) )); assert!(matches!( manager.async_read_secret("no-string").await, @@ -143,7 +143,7 @@ async fn trait_read_applies_the_aws_operation_timeout() { }; assert!(matches!( BaseSecretManager::async_read_secret(&loaded_manager(&server), "key", &context).await, - Err(Error::Timeout) + Err(Error::Timeout(_)) )); } @@ -169,7 +169,7 @@ async fn read_timeout_is_an_error_and_cannot_be_mistaken_for_missing() { let manager = AwsSecretsManagerV2::new(Client::from_conf(config), Default::default()); assert!(matches!( manager.async_read_secret("key").await, - Err(Error::Timeout) + Err(Error::Timeout(_)) )); } diff --git a/litellm-rust/crates/secrets-azure/Cargo.toml b/litellm-rust/crates/secrets-azure/Cargo.toml index 7ec03fb98da..2aff6129c37 100644 --- a/litellm-rust/crates/secrets-azure/Cargo.toml +++ b/litellm-rust/crates/secrets-azure/Cargo.toml @@ -15,6 +15,7 @@ litellm-core-utils.workspace = true reqwest.workspace = true serde.workspace = true thiserror.workspace = true +url.workspace = true veil.workspace = true percent-encoding = "2.3" diff --git a/litellm-rust/crates/secrets-azure/src/error.rs b/litellm-rust/crates/secrets-azure/src/error.rs index 2862f08a56c..097f42cb89c 100644 --- a/litellm-rust/crates/secrets-azure/src/error.rs +++ b/litellm-rust/crates/secrets-azure/src/error.rs @@ -7,7 +7,7 @@ pub enum Error { #[error("{0} environment variable is missing")] MissingEnvironment(&'static str), #[error("AZURE_KEY_VAULT_URI is not a valid https vault URL")] - VaultUri, + VaultUri(#[source] Option), #[error("Azure Key Vault credentials are not configured")] MissingCredentials, #[error(transparent)] @@ -20,10 +20,26 @@ pub enum Error { Http( #[source] #[redact] - reqwest::Error, + litellm_http::Error, ), #[error("Azure Key Vault returned HTTP {0}")] Status(u16), #[error("Azure Key Vault response is missing the secret value")] MissingValue, } + +impl Error { + pub fn is_missing_value(&self) -> bool { + match self { + Self::MissingValue => true, + Self::Operation(_) + | Self::Timeout + | Self::MissingEnvironment(_) + | Self::VaultUri(_) + | Self::MissingCredentials + | Self::Auth(_) + | Self::Http(_) + | Self::Status(_) => false, + } + } +} diff --git a/litellm-rust/crates/secrets-azure/src/key_vault.rs b/litellm-rust/crates/secrets-azure/src/key_vault.rs index 095c451927c..d97b7a22cfc 100644 --- a/litellm-rust/crates/secrets-azure/src/key_vault.rs +++ b/litellm-rust/crates/secrets-azure/src/key_vault.rs @@ -38,7 +38,7 @@ impl AzureKeyVault { environment: Arc, ) -> Result { if vault.host_str().is_none() { - return Err(Error::VaultUri); + return Err(Error::VaultUri(None)); } let inputs = AzureAuthInputs { azure_scope: ConfigValue::Value(Sourced::new( @@ -64,9 +64,9 @@ impl AzureKeyVault { let value = environment .get(AZURE_KEY_VAULT_URI) .ok_or(Error::MissingEnvironment(AZURE_KEY_VAULT_URI))?; - let vault = reqwest::Url::parse(&value).map_err(|_| Error::VaultUri)?; + let vault = reqwest::Url::parse(&value).map_err(|error| Error::VaultUri(Some(error)))?; if vault.scheme() != "https" || vault.host_str().is_none() { - return Err(Error::VaultUri); + return Err(Error::VaultUri(None)); } Self::with_client(client, vault, environment) } @@ -95,7 +95,7 @@ impl AzureKeyVault { let url = self .vault .join(&format!("secrets/{encoded_name}?api-version={API_VERSION}")) - .map_err(|_| Error::VaultUri)?; + .map_err(|error| Error::VaultUri(Some(error)))?; let response = self .client .get(url) @@ -103,14 +103,17 @@ impl AzureKeyVault { .header(reqwest::header::ACCEPT, "application/json") .send() .await - .map_err(Error::Http)?; + .map_err(|error| Error::Http(error.into()))?; if response.status() == reqwest::StatusCode::NOT_FOUND { return Ok(None); } if response.status() != reqwest::StatusCode::OK { return Err(Error::Status(response.status().as_u16())); } - let payload: SecretResponse = response.json().await.map_err(Error::Http)?; + let payload: SecretResponse = response + .json() + .await + .map_err(|error| Error::Http(error.into()))?; let value = payload.value.ok_or(Error::MissingValue)?; Ok(Some(SecretValue::new(value))) } diff --git a/litellm-rust/crates/secrets-azure/tests/key_vault.rs b/litellm-rust/crates/secrets-azure/tests/key_vault.rs index 05f62d25dce..a0ba9e6ddf8 100644 --- a/litellm-rust/crates/secrets-azure/tests/key_vault.rs +++ b/litellm-rust/crates/secrets-azure/tests/key_vault.rs @@ -122,13 +122,14 @@ async fn missing_value_is_an_error() { } #[rstest] -#[case::missing(None, true)] -#[case::http(Some("http://vault.example"), false)] -#[case::relative(Some("vault.example"), false)] -#[case::malformed(Some("://"), false)] +#[case::missing(None, true, false)] +#[case::http(Some("http://vault.example"), false, false)] +#[case::relative(Some("vault.example"), false, true)] +#[case::malformed(Some("://"), false, true)] fn new_validates_vault_environment( #[case] uri: Option<&'static str>, #[case] missing_environment: bool, + #[case] unparsable: bool, ) { let result = AzureKeyVault::new( litellm_http::Client::for_test(), @@ -145,7 +146,14 @@ fn new_validates_vault_environment( Err(Error::MissingEnvironment("AZURE_KEY_VAULT_URI")) )); } else { - assert!(matches!(result, Err(Error::VaultUri))); + let error = result.err().expect("invalid vault URI"); + assert!(matches!(error, Error::VaultUri(_))); + assert_eq!( + std::error::Error::source(&error) + .and_then(|cause| cause.downcast_ref::()) + .is_some(), + unparsable + ); } } diff --git a/litellm-rust/crates/secrets-cyberark/Cargo.toml b/litellm-rust/crates/secrets-cyberark/Cargo.toml index f630d5857d8..2ae2c1ea5c3 100644 --- a/litellm-rust/crates/secrets-cyberark/Cargo.toml +++ b/litellm-rust/crates/secrets-cyberark/Cargo.toml @@ -14,6 +14,7 @@ moka.workspace = true reqwest.workspace = true serde_json.workspace = true thiserror.workspace = true +url.workspace = true veil.workspace = true litellm-tracing.workspace = true percent-encoding = "2.3" diff --git a/litellm-rust/crates/secrets-cyberark/src/error.rs b/litellm-rust/crates/secrets-cyberark/src/error.rs index 4f21647225b..b5dbd7342fa 100644 --- a/litellm-rust/crates/secrets-cyberark/src/error.rs +++ b/litellm-rust/crates/secrets-cyberark/src/error.rs @@ -6,7 +6,7 @@ pub enum Error { Http( #[from] #[redact] - reqwest::Error, + litellm_http::Error, ), #[error("CyberArk Conjur authentication returned HTTP {0}")] AuthStatus(u16), @@ -17,15 +17,68 @@ pub enum Error { )] MissingCredentials, #[error("CyberArk client certificate could not be loaded")] - ClientCertificate, + ClientCertificate( + #[source] + #[redact] + litellm_http::Error, + ), #[error("CyberArk Conjur HTTP client could not be built")] - Client(#[redact] Box), + Client( + #[source] + #[redact] + litellm_http::Error, + ), #[error("invalid refresh interval")] - RefreshInterval, + RefreshInterval(#[source] std::num::ParseIntError), #[error("invalid CyberArk Conjur endpoint")] - Endpoint, + Endpoint(#[source] url::ParseError), #[error("CyberArk secret manager requires an enterprise license")] EnterpriseRequired, #[error(transparent)] Operation(#[from] litellm_secrets_types::Error), } + +impl Error { + pub fn is_unsafe_secret_name(&self) -> bool { + match self { + Self::Operation(error) => error.is_unsafe_secret_name(), + Self::Timeout + | Self::Http(_) + | Self::AuthStatus(_) + | Self::Status(_) + | Self::MissingCredentials + | Self::ClientCertificate(_) + | Self::Client(_) + | Self::RefreshInterval(_) + | Self::Endpoint(_) + | Self::EnterpriseRequired => false, + } + } +} + +#[derive(thiserror::Error, veil::Redact)] +#[error("CyberArk Conjur write failed")] +pub struct WriteFailure { + pub source: Error, + #[redact] + pub request_url: Option, + pub authentication: bool, +} + +impl WriteFailure { + pub(crate) fn local(source: Error) -> Self { + Self { + source, + request_url: None, + authentication: false, + } + } + + pub(crate) fn request(source: Error, url: reqwest::Url) -> Self { + Self { + source, + request_url: Some(url), + authentication: false, + } + } +} diff --git a/litellm-rust/crates/secrets-cyberark/src/lib.rs b/litellm-rust/crates/secrets-cyberark/src/lib.rs index 74a8b6febf2..d780fb43921 100644 --- a/litellm-rust/crates/secrets-cyberark/src/lib.rs +++ b/litellm-rust/crates/secrets-cyberark/src/lib.rs @@ -3,5 +3,5 @@ mod error; mod secret_manager; -pub use error::Error; -pub use secret_manager::{AuthenticationRetry, CyberArkSecretManager, DeleteOutcome, WriteFailure}; +pub use error::{Error, WriteFailure}; +pub use secret_manager::{AuthenticationRetry, CyberArkSecretManager, DeleteOutcome}; diff --git a/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs b/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs index 44473c9fe90..c5cb0308b13 100644 --- a/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs @@ -17,7 +17,7 @@ use litellm_secrets_types::{ use moka::future::Cache; use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, utf8_percent_encode}; -use crate::Error; +use crate::{Error, WriteFailure}; const CYBERARK_API_BASE: &str = "CYBERARK_API_BASE"; const CYBERARK_ACCOUNT: &str = "CYBERARK_ACCOUNT"; @@ -62,38 +62,12 @@ pub enum AuthenticationRetry { Unauthorized, } -#[derive(veil::Redact)] -pub struct WriteFailure { - pub source: Error, - #[redact] - pub request_url: Option, - pub authentication: bool, -} - -impl WriteFailure { - fn local(source: Error) -> Self { - Self { - source, - request_url: None, - authentication: false, - } - } - - fn request(source: Error, url: reqwest::Url) -> Self { - Self { - source, - request_url: Some(url), - authentication: false, - } - } -} - impl CyberArkSecretManager { fn secret_url(&self, name: &str) -> Result { let encoded = utf8_percent_encode(name, SECRET_NAME_SAFE); self.endpoint .join(&format!("secrets/{}/variable/{}", self.account, encoded)) - .map_err(|_| Error::Endpoint) + .map_err(Error::Endpoint) } } diff --git a/litellm-rust/crates/secrets-cyberark/src/secret_manager/client.rs b/litellm-rust/crates/secrets-cyberark/src/secret_manager/client.rs index ea2262392d5..1cda0ae5306 100644 --- a/litellm-rust/crates/secrets-cyberark/src/secret_manager/client.rs +++ b/litellm-rust/crates/secrets-cyberark/src/secret_manager/client.rs @@ -64,25 +64,18 @@ impl CyberArkSecretManager { }), ..config.clone() }; - let client = - pool.client(&config, ClientVariant::Provider) - .map_err(|error| match error { - litellm_http::Error::Read { - tls_source: TlsSource::ClientIdentity, - .. - } - | litellm_http::Error::InvalidPem { - tls_source: TlsSource::ClientIdentity, - .. - } => Error::ClientCertificate, - other => Error::Client(Box::new(other)), - })?; + let client = pool + .client(&config, ClientVariant::Provider) + .map_err(|error| match error.tls_source() { + Some(TlsSource::ClientIdentity) => Error::ClientCertificate(error), + Some(TlsSource::CaBundle) | None => Error::Client(error), + })?; let endpoint = reqwest::Url::parse( &environment .get(CYBERARK_API_BASE) .unwrap_or_else(|| DEFAULT_API_BASE.to_owned()), ) - .map_err(|_| Error::Endpoint)?; + .map_err(Error::Endpoint)?; let account = environment .get(CYBERARK_ACCOUNT) .unwrap_or_else(|| DEFAULT_ACCOUNT.to_owned()); @@ -95,7 +88,7 @@ impl CyberArkSecretManager { value .parse::() .map(Duration::from_secs) - .map_err(|_| Error::RefreshInterval) + .map_err(Error::RefreshInterval) }) .transpose()?; Ok(Self::with_client( @@ -115,7 +108,7 @@ impl CyberArkSecretManager { self.account, utf8_percent_encode(&self.username, SECRET_NAME_SAFE) )) - .map_err(|_| Error::Endpoint) + .map_err(Error::Endpoint) } pub(super) async fn authenticate( @@ -135,11 +128,14 @@ impl CyberArkSecretManager { context, ) .send() - .await?; + .await + .map_err(litellm_http::Error::from)?; if !response.status().is_success() { return Err(Error::AuthStatus(response.status().as_u16())); } - let token = SecretValue::new(STANDARD.encode(response.text().await?)); + let token = SecretValue::new( + STANDARD.encode(response.text().await.map_err(litellm_http::Error::from)?), + ); self.token.insert((), token.clone()).await; Ok(token) } diff --git a/litellm-rust/crates/secrets-cyberark/src/secret_manager/read.rs b/litellm-rust/crates/secrets-cyberark/src/secret_manager/read.rs index daa9ed1114d..2862c6a956d 100644 --- a/litellm-rust/crates/secrets-cyberark/src/secret_manager/read.rs +++ b/litellm-rust/crates/secrets-cyberark/src/secret_manager/read.rs @@ -76,7 +76,8 @@ impl CyberArkSecretManager { context, ) .send() - .await?; + .await + .map_err(litellm_http::Error::from)?; let response = if matches!(retry, AuthenticationRetry::Unauthorized) && had_cached_token && response.status() == reqwest::StatusCode::UNAUTHORIZED @@ -89,7 +90,8 @@ impl CyberArkSecretManager { context, ) .send() - .await? + .await + .map_err(litellm_http::Error::from)? } else { response }; @@ -99,7 +101,7 @@ impl CyberArkSecretManager { if !response.status().is_success() { return Err(Error::Status(response.status().as_u16())); } - let value = SecretValue::new(response.text().await?); + let value = SecretValue::new(response.text().await.map_err(litellm_http::Error::from)?); Ok(Some(value)) } } diff --git a/litellm-rust/crates/secrets-cyberark/src/secret_manager/write.rs b/litellm-rust/crates/secrets-cyberark/src/secret_manager/write.rs index 87f6cea3619..7a2c79f506a 100644 --- a/litellm-rust/crates/secrets-cyberark/src/secret_manager/write.rs +++ b/litellm-rust/crates/secrets-cyberark/src/secret_manager/write.rs @@ -83,7 +83,9 @@ impl CyberArkSecretManager { ) .send() .await - .map_err(|source| WriteFailure::request(source.into(), url.clone())) + .map_err(|source| { + WriteFailure::request(litellm_http::Error::from(source).into(), url.clone()) + }) } pub(super) async fn ensure_variable_exists( @@ -141,8 +143,8 @@ impl CyberArkSecretManager { ); return; } - Err(error) => { - litellm_tracing::warn!("Error ensuring CyberArk variable exists: {error}"); + Err(_) => { + litellm_tracing::warn!("Error ensuring CyberArk variable exists"); return; } } diff --git a/litellm-rust/crates/secrets-cyberark/tests/secret_manager/configuration.rs b/litellm-rust/crates/secrets-cyberark/tests/secret_manager/configuration.rs index a3c20a0db3d..8a0746f8002 100644 --- a/litellm-rust/crates/secrets-cyberark/tests/secret_manager/configuration.rs +++ b/litellm-rust/crates/secrets-cyberark/tests/secret_manager/configuration.rs @@ -201,7 +201,7 @@ async fn trait_read_applies_cyberark_operation_timeout_to_authentication() { match result { Err(Error::Timeout) => {} - Err(Error::Http(error)) => assert!(error.is_timeout()), + Err(Error::Http(error)) => assert_eq!(error.status(), Some(408)), other => panic!("expected timeout, got {other:?}"), } } @@ -237,7 +237,7 @@ fn new_validates_credentials_before_license_and_configuration() { }), true ), - Err(Error::RefreshInterval) + Err(Error::RefreshInterval(_)) )); assert!(matches!( from_environment( @@ -248,7 +248,7 @@ fn new_validates_credentials_before_license_and_configuration() { }), true ), - Err(Error::Endpoint) + Err(Error::Endpoint(_)) )); } @@ -263,7 +263,7 @@ fn certificate_only_credentials_are_validated_as_a_client_identity() { true, ); - assert!(matches!(result, Err(Error::ClientCertificate))); + assert!(matches!(result, Err(Error::ClientCertificate(_)))); } #[rstest] @@ -347,7 +347,7 @@ fn invalid_client_identity_is_not_ignored( true, ); - assert!(matches!(result, Err(Error::ClientCertificate))); + assert!(matches!(result, Err(Error::ClientCertificate(_)))); } #[rstest] @@ -413,7 +413,7 @@ fn new_reports_missing_client_certificate_files() { }), true ), - Err(Error::ClientCertificate) + Err(Error::ClientCertificate(_)) )); } diff --git a/litellm-rust/crates/secrets-google/Cargo.toml b/litellm-rust/crates/secrets-google/Cargo.toml index 3ce14fe7a12..06b3747def5 100644 --- a/litellm-rust/crates/secrets-google/Cargo.toml +++ b/litellm-rust/crates/secrets-google/Cargo.toml @@ -17,6 +17,7 @@ base64.workspace = true crc32c = "0.6.8" serde_json.workspace = true thiserror.workspace = true +url.workspace = true veil.workspace = true google-cloud-kms-v1 = "1.14.0" google-cloud-gax = { version = "1.14.0", default-features = false } diff --git a/litellm-rust/crates/secrets-google/src/error.rs b/litellm-rust/crates/secrets-google/src/error.rs index 159e4983ddc..486c7a3f35e 100644 --- a/litellm-rust/crates/secrets-google/src/error.rs +++ b/litellm-rust/crates/secrets-google/src/error.rs @@ -26,7 +26,7 @@ pub enum Error { Http( #[from] #[redact] - reqwest::Error, + litellm_http::Error, ), #[error("Google Secret Manager returned HTTP {0}")] Status(u16), @@ -35,15 +35,41 @@ pub enum Error { #[error("required environment variable is missing: {0}")] MissingEnvironment(&'static str), #[error("invalid refresh interval")] - RefreshInterval, + RefreshInterval(#[source] std::num::ParseIntError), #[error("payload is not valid base64")] - Base64(#[from] base64::DecodeError), + Base64( + #[from] + #[redact] + base64::DecodeError, + ), #[error("Google Secret Manager payload checksum mismatch")] - Checksum, + Checksum(#[source] Option), #[error("decrypted value is not UTF-8")] - Utf8, + Utf8(#[source] std::str::Utf8Error), #[error("invalid Google Secret Manager endpoint")] - Endpoint, + Endpoint(#[source] url::ParseError), #[error("Google Secret Manager requires an enterprise license")] EnterpriseRequired, } + +impl Error { + pub fn is_not_found(&self) -> bool { + match self { + Self::Status(status) => *status == 404, + Self::Operation(_) + | Self::Timeout + | Self::Client(_) + | Self::Auth(_) + | Self::Kms(_) + | Self::Http(_) + | Self::MissingPayload + | Self::MissingEnvironment(_) + | Self::RefreshInterval(_) + | Self::Base64(_) + | Self::Checksum(_) + | Self::Utf8(_) + | Self::Endpoint(_) + | Self::EnterpriseRequired => false, + } + } +} diff --git a/litellm-rust/crates/secrets-google/src/secret_manager.rs b/litellm-rust/crates/secrets-google/src/secret_manager.rs index 08ba466b799..463d1f6fd73 100644 --- a/litellm-rust/crates/secrets-google/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-google/src/secret_manager.rs @@ -92,12 +92,12 @@ impl GoogleSecretManager { let ttl = environment .get(GOOGLE_SECRET_MANAGER_REFRESH_INTERVAL) .filter(|v| !v.is_empty()) - .map(|v| v.parse::().map_err(|_| Error::RefreshInterval)) + .map(|v| v.parse::().map_err(Error::RefreshInterval)) .transpose()? .unwrap_or( environment .get(SECRET_MANAGER_REFRESH_INTERVAL) - .map(|v| v.parse::().map_err(|_| Error::RefreshInterval)) + .map(|v| v.parse::().map_err(Error::RefreshInterval)) .transpose()? .unwrap_or(DEFAULT_REFRESH_INTERVAL.as_secs() as i64), ); @@ -164,20 +164,21 @@ impl GoogleSecretManager { ), percent_encoding::utf8_percent_encode(name, percent_encoding::NON_ALPHANUMERIC) )) - .map_err(|_| Error::Endpoint)?; + .map_err(Error::Endpoint)?; let response = self .client .get(url) .headers(self.credentials.request_headers().await?) .send() - .await?; + .await + .map_err(litellm_http::Error::from)?; if response.status() == reqwest::StatusCode::NOT_FOUND { return Ok(None); } if response.status() != reqwest::StatusCode::OK { return Err(Error::Status(response.status().as_u16())); } - let response: Response = response.json().await?; + let response: Response = response.json().await.map_err(litellm_http::Error::from)?; let Some(payload) = response.payload else { return Err(Error::MissingPayload); }; @@ -186,12 +187,15 @@ impl GoogleSecretManager { }; let bytes = STANDARD.decode(data)?; if let Some(expected) = payload.data_crc32c { - let expected = expected.parse::().map_err(|_| Error::Checksum)?; + let expected = expected + .parse::() + .map_err(|error| Error::Checksum(Some(error)))?; if crc32c::crc32c(&bytes) != expected { - return Err(Error::Checksum); + return Err(Error::Checksum(None)); } } - let plaintext = String::from_utf8(bytes).map_err(|_| Error::Utf8)?; + let plaintext = + String::from_utf8(bytes).map_err(|error| Error::Utf8(error.utf8_error()))?; let value = SecretValue::new(plaintext); Ok(Some(value)) } diff --git a/litellm-rust/crates/secrets-google/tests/secret_manager.rs b/litellm-rust/crates/secrets-google/tests/secret_manager.rs index f2cb5ff3603..615e331cb96 100644 --- a/litellm-rust/crates/secrets-google/tests/secret_manager.rs +++ b/litellm-rust/crates/secrets-google/tests/secret_manager.rs @@ -152,8 +152,8 @@ async fn failed_or_missing_reads_are_not_cached( assert!(matches!(result, Err(Error::MissingPayload))); } ExpectedReadFailure::Base64 => assert!(matches!(result, Err(Error::Base64(_)))), - ExpectedReadFailure::Utf8 => assert!(matches!(result, Err(Error::Utf8))), - ExpectedReadFailure::Checksum => assert!(matches!(result, Err(Error::Checksum))), + ExpectedReadFailure::Utf8 => assert!(matches!(result, Err(Error::Utf8(_)))), + ExpectedReadFailure::Checksum => assert!(matches!(result, Err(Error::Checksum(None)))), } drop(failing); Mock::given(path( @@ -245,7 +245,7 @@ fn google_manager_rejects_invalid_refresh_intervals(#[case] variable: &'static s assert!(matches!( GoogleSecretManager::new(litellm_http::Client::for_test(), environment, true), - Err(Error::RefreshInterval) + Err(Error::RefreshInterval(_)) )); } diff --git a/litellm-rust/crates/secrets-hashicorp/Cargo.toml b/litellm-rust/crates/secrets-hashicorp/Cargo.toml index 7dd3d3c674f..8731a8deb7a 100644 --- a/litellm-rust/crates/secrets-hashicorp/Cargo.toml +++ b/litellm-rust/crates/secrets-hashicorp/Cargo.toml @@ -6,8 +6,11 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-error.workspace = true litellm-core-utils.workspace = true +litellm-http.workspace = true litellm-secrets-types.workspace = true +reqwest.workspace = true rustify.workspace = true rustify_derive.workspace = true serde.workspace = true diff --git a/litellm-rust/crates/secrets-hashicorp/src/config.rs b/litellm-rust/crates/secrets-hashicorp/src/config.rs index d32491199c4..4232667ae16 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/config.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/config.rs @@ -153,9 +153,11 @@ fn refresh_interval(environment: &dyn Lookup) -> Result { let Some(value) = value else { return Ok(DEFAULT_REFRESH_INTERVAL); }; - let seconds: i64 = value.parse().map_err(|_| Error::RefreshInterval)?; + let seconds: i64 = value + .parse() + .map_err(|error| Error::RefreshInterval(Some(error)))?; if seconds < 0 { - return Err(Error::RefreshInterval); + return Err(Error::RefreshInterval(None)); } Ok(Duration::from_secs(seconds as u64)) } diff --git a/litellm-rust/crates/secrets-hashicorp/src/error.rs b/litellm-rust/crates/secrets-hashicorp/src/error.rs index 1813cb486bc..73399a44eaa 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/error.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/error.rs @@ -1,42 +1,144 @@ +use rustify::errors::ClientError as RustifyClientError; +use vaultrs::error::ClientError; + #[derive(thiserror::Error, veil::Redact)] pub enum Error { #[error("HashiCorp Vault requires an enterprise license")] EnterpriseRequired, #[error("invalid secret name")] - InvalidSecretName(litellm_secrets_types::Error), + InvalidSecretName(#[source] litellm_secrets_types::Error), #[error(transparent)] Operation(#[from] litellm_secrets_types::Error), #[error("HashiCorp Vault client failed")] Client( - #[from] + #[source] #[redact] vaultrs::error::ClientError, ), - #[error("HashiCorp Vault client settings are invalid: {message}")] - ClientSettings { message: String }, - #[error("HashiCorp Vault TLS identity could not be configured for {path}: {message}")] + #[error("HashiCorp Vault request failed")] + Transport( + #[source] + #[redact] + litellm_http::Error, + ), + #[error("HashiCorp Vault client settings are invalid")] + ClientSettings(#[source] vaultrs::client::VaultClientSettingsBuilderError), + #[error("HashiCorp Vault TLS identity could not be configured for {path}")] TlsIdentity { path: std::path::PathBuf, - message: String, + source: litellm_error::Source, }, #[error("HashiCorp Vault login returned HTTP {status}")] LoginStatus { status: u16 }, #[error("HashiCorp Vault login response is malformed")] - MalformedLogin, + MalformedLogin( + #[source] + #[redact] + vaultrs::error::ClientError, + ), #[error("HashiCorp Vault authentication is not configured")] NoAuthConfigured, #[error("HashiCorp Vault returned HTTP {status}")] Status { status: u16 }, #[error("HashiCorp Vault response payload is malformed")] - MalformedPayload, + MalformedPayload( + #[source] + #[redact] + vaultrs::error::ClientError, + ), #[error("HashiCorp Vault secret value is not a string")] NonStringValue, #[error("HashiCorp Vault data key conflicts with description")] DataKeyConflictsWithDescription, #[error("HashiCorp Vault secret version exceeds CAS range")] - CasVersionOverflow, + CasVersionOverflow(#[source] std::num::TryFromIntError), #[error("HashiCorp Vault operation timed out")] Timeout, #[error("invalid HashiCorp Vault refresh interval")] - RefreshInterval, + RefreshInterval(#[source] Option), +} + +impl Error { + pub fn login_status(&self) -> Option { + match self { + Self::LoginStatus { status } => Some(*status), + Self::EnterpriseRequired + | Self::InvalidSecretName(_) + | Self::Operation(_) + | Self::Client(_) + | Self::Transport(_) + | Self::ClientSettings(_) + | Self::TlsIdentity { .. } + | Self::MalformedLogin(_) + | Self::NoAuthConfigured + | Self::Status { .. } + | Self::MalformedPayload(_) + | Self::NonStringValue + | Self::DataKeyConflictsWithDescription + | Self::CasVersionOverflow(_) + | Self::Timeout + | Self::RefreshInterval(_) => None, + } + } + + pub fn is_unsafe_secret_name(&self) -> bool { + match self { + Self::InvalidSecretName(error) | Self::Operation(error) => { + error.is_unsafe_secret_name() + } + Self::EnterpriseRequired + | Self::Client(_) + | Self::Transport(_) + | Self::ClientSettings(_) + | Self::TlsIdentity { .. } + | Self::LoginStatus { .. } + | Self::MalformedLogin(_) + | Self::NoAuthConfigured + | Self::Status { .. } + | Self::MalformedPayload(_) + | Self::NonStringValue + | Self::DataKeyConflictsWithDescription + | Self::CasVersionOverflow(_) + | Self::Timeout + | Self::RefreshInterval(_) => false, + } + } +} + +impl From for Error { + fn from(error: ClientError) -> Self { + match error { + ClientError::RestClientError { source } => match transport(source) { + Ok(error) => Self::Transport(error), + Err(source) => Self::Client(ClientError::RestClientError { source }), + }, + ClientError::RestClientBuildError { source } => Self::Transport(source.into()), + error => Self::Client(error), + } + } +} + +pub(crate) fn transport( + error: RustifyClientError, +) -> Result { + match error { + RustifyClientError::RequestError { + source, + url, + method, + } => source + .downcast::() + .map(litellm_http::Error::from) + .map_err(|source| RustifyClientError::RequestError { + source, + url, + method, + }), + RustifyClientError::ResponseError { source } => source + .downcast::() + .map(litellm_http::Error::from) + .map_err(|source| RustifyClientError::ResponseError { source }), + RustifyClientError::ReqwestBuildError { source } => Ok(source.into()), + error => Err(error), + } } diff --git a/litellm-rust/crates/secrets-hashicorp/src/lib.rs b/litellm-rust/crates/secrets-hashicorp/src/lib.rs index 7aab5a38b01..15c03724f72 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/lib.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/lib.rs @@ -7,4 +7,6 @@ pub mod secret_manager; pub use config::{AppRoleAuth, HashicorpVaultConfig, TlsCertAuth}; pub use error::Error; -pub use secret_manager::{HashicorpVault, RawOperationError, SecretLocation}; +pub use secret_manager::{ + HashicorpVault, HttpFailure, LoginFailure, RawOperationError, SecretLocation, +}; diff --git a/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs index f0d4fe8c0b4..447beacda85 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs @@ -31,7 +31,7 @@ use vaultrs::{ use crate::{Error, HashicorpVaultConfig, TlsCertAuth, cert_login::CertLoginRequest}; -pub use raw::RawOperationError; +pub use raw::{HttpFailure, LoginFailure, RawOperationError}; const CACHE_CAPACITY: u64 = 200; @@ -178,23 +178,25 @@ fn map_api_error(error: ClientError, context: ErrorContext) -> Error { ErrorContext::Login => Error::LoginStatus { status: code }, ErrorContext::Read | ErrorContext::Secret => Error::Status { status: code }, }, - ClientError::JsonParseError { source } => match context { - ErrorContext::Login => Error::MalformedLogin, - ErrorContext::Read => Error::MalformedPayload, - ErrorContext::Secret => Error::Client(ClientError::JsonParseError { source }), + error @ ClientError::JsonParseError { .. } => match context { + ErrorContext::Login => Error::MalformedLogin(error), + ErrorContext::Read => Error::MalformedPayload(error), + ErrorContext::Secret => Error::Client(error), }, - ClientError::ResponseEmptyError | ClientError::ResponseDataEmptyError => { - malformed_response(context) + error @ (ClientError::ResponseEmptyError | ClientError::ResponseDataEmptyError) => { + malformed_response(context, error) } ClientError::RestClientError { source } => match source { RustifyClientError::ServerResponseError { code, .. } => match context { ErrorContext::Login => Error::LoginStatus { status: code }, ErrorContext::Read | ErrorContext::Secret => Error::Status { status: code }, }, - RustifyClientError::ResponseParseError { .. } => malformed_response(context), - source => Error::Client(ClientError::RestClientError { source }), + source @ RustifyClientError::ResponseParseError { .. } => { + malformed_response(context, ClientError::RestClientError { source }) + } + source => ClientError::RestClientError { source }.into(), }, - error => Error::Client(error), + error => error.into(), } } @@ -208,10 +210,10 @@ fn api_status(error: &ClientError) -> Option { } } -fn malformed_response(context: ErrorContext) -> Error { +fn malformed_response(context: ErrorContext, error: ClientError) -> Error { match context { - ErrorContext::Login => Error::MalformedLogin, - ErrorContext::Read => Error::MalformedPayload, - ErrorContext::Secret => Error::Client(ClientError::ResponseDataEmptyError), + ErrorContext::Login => Error::MalformedLogin(error), + ErrorContext::Read => Error::MalformedPayload(error), + ErrorContext::Secret => Error::Client(error), } } diff --git a/litellm-rust/crates/secrets-hashicorp/src/secret_manager/client.rs b/litellm-rust/crates/secrets-hashicorp/src/secret_manager/client.rs index c1dc073dcf8..01d3807d290 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/secret_manager/client.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/secret_manager/client.rs @@ -83,10 +83,8 @@ impl HashicorpVault { .ca_certs(Vec::new()) .verify(true) .build() - .map_err(|message| Error::ClientSettings { - message: message.to_string(), - })?; - VaultClient::new(settings).map_err(Error::Client) + .map_err(Error::ClientSettings)?; + VaultClient::new(settings).map_err(Error::from) } } @@ -94,16 +92,16 @@ fn identity_for(tls: Option<&TlsCertAuth>) -> Result, Error> { tls.map(|tls| { let cert: Vec = std::fs::read(&tls.cert_path).map_err(|source| Error::TlsIdentity { path: tls.cert_path.clone(), - message: source.to_string(), + source: litellm_error::Source::new(source), })?; let key: Vec = std::fs::read(&tls.key_path).map_err(|source| Error::TlsIdentity { path: tls.key_path.clone(), - message: source.to_string(), + source: litellm_error::Source::new(source), })?; Identity::from_pem(&[cert.as_slice(), key.as_slice()].concat()).map_err(|source| { Error::TlsIdentity { path: tls.cert_path.clone(), - message: source.to_string(), + source: litellm_error::Source::new(source), } }) }) diff --git a/litellm-rust/crates/secrets-hashicorp/src/secret_manager/raw.rs b/litellm-rust/crates/secrets-hashicorp/src/secret_manager/raw.rs index 0325561a7a0..9c722c0b6c8 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/secret_manager/raw.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/secret_manager/raw.rs @@ -4,26 +4,123 @@ use vaultrs::api::kv2::requests::{ DeleteLatestSecretVersionRequest, ReadSecretRequest, SetSecretRequest, }; -#[derive(veil::Redact)] +#[derive(thiserror::Error, veil::Redact)] pub enum RawOperationError { + #[error(transparent)] Local(Error), + #[error("HashiCorp Vault authentication failed")] Authentication { source: Error, url: String, certificate: bool, }, - Http { - method: String, - url: String, - status: u16, + #[error("HashiCorp Vault {} request returned HTTP {}", .0.method, .0.status)] + Http(HttpFailure), + #[error("HashiCorp Vault request failed")] + Transport( + #[source] #[redact] - body: Vec, - }, - Transport(#[redact] RustifyClientError), - Timeout { - method: String, - elapsed: Duration, - }, + litellm_http::Error, + ), + #[error("HashiCorp Vault client failed")] + Client( + #[source] + #[redact] + RustifyClientError, + ), + #[error("HashiCorp Vault {method} request timed out")] + Timeout { method: String, elapsed: Duration }, +} + +#[derive(veil::Redact)] +pub struct HttpFailure { + pub method: String, + pub url: String, + pub status: u16, + #[redact] + pub body: Vec, +} + +pub struct LoginFailure<'a> { + pub source: &'a Error, + pub url: &'a str, + pub certificate: bool, +} + +impl RawOperationError { + pub fn is_unsafe_secret_name(&self) -> bool { + match self { + Self::Local(error) => error.is_unsafe_secret_name(), + Self::Authentication { .. } + | Self::Http(_) + | Self::Transport(_) + | Self::Client(_) + | Self::Timeout { .. } => false, + } + } + + pub fn local(&self) -> Option<&Error> { + match self { + Self::Local(error) => Some(error), + Self::Authentication { .. } + | Self::Http(_) + | Self::Transport(_) + | Self::Client(_) + | Self::Timeout { .. } => None, + } + } + + pub fn authentication(&self) -> Option> { + match self { + Self::Authentication { + source, + url, + certificate, + } => Some(LoginFailure { + source, + url, + certificate: *certificate, + }), + Self::Local(_) + | Self::Http(_) + | Self::Transport(_) + | Self::Client(_) + | Self::Timeout { .. } => None, + } + } + + pub fn http_failure(&self) -> Option<&HttpFailure> { + match self { + Self::Http(failure) => Some(failure), + Self::Local(_) + | Self::Authentication { .. } + | Self::Transport(_) + | Self::Client(_) + | Self::Timeout { .. } => None, + } + } + + pub fn timeout(&self) -> Option<(&str, Duration)> { + match self { + Self::Timeout { method, elapsed } => Some((method, *elapsed)), + Self::Local(_) + | Self::Authentication { .. } + | Self::Http(_) + | Self::Transport(_) + | Self::Client(_) => None, + } + } + + pub fn transport(&self) -> Option<&litellm_http::Error> { + match self { + Self::Transport(error) => Some(error), + Self::Local(_) + | Self::Authentication { .. } + | Self::Http(_) + | Self::Client(_) + | Self::Timeout { .. } => None, + } + } } impl HashicorpVault { @@ -118,7 +215,7 @@ impl HashicorpVault { let request = endpoint .with_middleware(&client.middle) .request(client.http.base()) - .map_err(RawOperationError::Transport)?; + .map_err(RawOperationError::Client)?; let method = request.method().to_string(); let url = format!( "{}/v1/{}{}/data/{}", @@ -141,14 +238,17 @@ impl HashicorpVault { })?, None => client.http.send(request).await, } - .map_err(RawOperationError::Transport)?; + .map_err(|error| match crate::error::transport(error) { + Ok(error) => RawOperationError::Transport(error), + Err(error) => RawOperationError::Client(error), + })?; if !response.status().is_success() { - return Err(RawOperationError::Http { + return Err(RawOperationError::Http(HttpFailure { method, url, status: response.status().as_u16(), body: response.into_body(), - }); + })); } Ok(response.into_body()) } diff --git a/litellm-rust/crates/secrets-hashicorp/src/secret_manager/write.rs b/litellm-rust/crates/secrets-hashicorp/src/secret_manager/write.rs index 0b1acb3a179..15a0e4c7dff 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/secret_manager/write.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/secret_manager/write.rs @@ -37,7 +37,7 @@ impl HashicorpVault { .await { Ok(metadata) => u32::try_from(metadata.current_version) - .map_err(|_| Error::CasVersionOverflow)?, + .map_err(Error::CasVersionOverflow)?, Err(error) if api_status(&error) == Some(404) => 0, Err(_) => return Err(map_api_error(error, ErrorContext::Secret)), }; diff --git a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager/configuration.rs b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager/configuration.rs index be3059ce986..b45b843fde8 100644 --- a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager/configuration.rs +++ b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager/configuration.rs @@ -32,7 +32,7 @@ fn invalid_refresh_intervals_are_rejected(#[case] value: &str) { assert!(matches!( HashicorpVaultConfig::from_environment(environment.as_ref()), - Err(Error::RefreshInterval) + Err(Error::RefreshInterval(_)) )); } @@ -79,6 +79,39 @@ async fn approle_login_uses_namespace_and_reuses_the_token() { assert!(manager.async_read_secret("name-2").await.unwrap().is_none()); } +#[rstest] +#[tokio::test] +async fn an_unreadable_client_certificate_keeps_its_io_cause() { + let server: MockServer = MockServer::start().await; + let directory: tempfile::TempDir = tempfile::tempdir().unwrap(); + let cert_path = directory.path().join("missing.crt"); + let values: HashMap = HashMap::from([ + ("HCP_VAULT_ADDR".to_owned(), server.uri()), + ( + "HCP_VAULT_CLIENT_CERT".to_owned(), + cert_path.to_str().unwrap().to_owned(), + ), + ( + "HCP_VAULT_CLIENT_KEY".to_owned(), + cert_path.to_str().unwrap().to_owned(), + ), + ]); + let environment: Arc = + Arc::new(move |name: &str| values.get(name).cloned()); + let manager: HashicorpVault = HashicorpVault::new(environment, true).unwrap(); + + let error = manager.async_read_secret("name").await.unwrap_err(); + + let cause = std::error::Error::source(&error) + .and_then(|cause| cause.downcast_ref::()) + .and_then(|cause| cause.downcast_ref::()) + .expect("io cause"); + assert_eq!(cause.kind(), std::io::ErrorKind::NotFound); + let report = litellm_error::Report::new(&error).to_string(); + assert_eq!(report.matches(cert_path.to_str().unwrap()).count(), 1); + assert!(report.ends_with(&cause.to_string())); +} + #[rstest] #[tokio::test] async fn tls_login_posts_the_role_and_uses_the_client_identity() { diff --git a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager/reads.rs b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager/reads.rs index 0c662eb5b55..c31b29fb536 100644 --- a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager/reads.rs +++ b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager/reads.rs @@ -138,7 +138,7 @@ async fn read_responses_distinguish_absence_and_malformed_payloads( .await; match expected { ExpectedRead::Missing => assert!(result.unwrap().is_none()), - ExpectedRead::Malformed => assert!(matches!(result, Err(Error::MalformedPayload))), + ExpectedRead::Malformed => assert!(matches!(result, Err(Error::MalformedPayload(_)))), ExpectedRead::NonString => assert!(matches!(result, Err(Error::NonStringValue))), } } @@ -383,3 +383,52 @@ async fn concurrent_reads_share_a_load_but_verification_fetches_fresh( "new" ); } + +fn unreachable_manager() -> HashicorpVault { + let address = std::net::TcpListener::bind("127.0.0.1:0") + .unwrap() + .local_addr() + .unwrap(); + let environment: Arc = Arc::new(move |name: &str| match name { + "HCP_VAULT_ADDR" => Some(format!("http://{address}")), + "HCP_VAULT_TOKEN" => Some("token".to_owned()), + _ => None, + }); + HashicorpVault::new(environment, true).unwrap() +} + +#[tokio::test] +async fn unreachable_vault_reads_are_connect_failures_without_the_url() { + let error = unreachable_manager() + .async_read_secret("do-not-leak") + .await + .unwrap_err(); + + assert!(matches!(&error, Error::Transport(transport) if transport.is_connect())); + let report = litellm_error::Report::new(&error).to_string(); + assert!( + report.starts_with("HashiCorp Vault request failed: "), + "{report}" + ); + assert!(!report.contains("do-not-leak"), "{report}"); +} + +#[tokio::test] +async fn unreachable_vault_raw_reads_are_connect_failures_without_the_url() { + let error = unreachable_manager() + .read_raw("do-not-leak", &HashicorpOperationContext::default()) + .await + .unwrap_err(); + + assert!( + error + .transport() + .is_some_and(litellm_http::Error::is_connect) + ); + let report = litellm_error::Report::new(&error).to_string(); + assert!( + report.starts_with("HashiCorp Vault request failed: "), + "{report}" + ); + assert!(!report.contains("do-not-leak"), "{report}"); +} diff --git a/litellm-rust/crates/secrets-types/src/error.rs b/litellm-rust/crates/secrets-types/src/error.rs index 027dd54d8cc..891d8a0fba9 100644 --- a/litellm-rust/crates/secrets-types/src/error.rs +++ b/litellm-rust/crates/secrets-types/src/error.rs @@ -11,3 +11,15 @@ pub enum Error { #[error("secret manager received an incompatible operation context")] InvalidOperationContext, } + +impl Error { + pub fn is_unsafe_secret_name(&self) -> bool { + match self { + Self::UnsafeSecretName => true, + Self::CurrentSecretMissing + | Self::NewSecretMissing + | Self::NewSecretMismatch + | Self::InvalidOperationContext => false, + } + } +} diff --git a/litellm-rust/crates/secrets-types/tests/cache.rs b/litellm-rust/crates/secrets-types/tests/cache.rs index cec4f5198e3..18efde4289b 100644 --- a/litellm-rust/crates/secrets-types/tests/cache.rs +++ b/litellm-rust/crates/secrets-types/tests/cache.rs @@ -110,7 +110,7 @@ async fn concurrent_misses_share_a_load_and_cancellation_allows_a_retry() { } #[tokio::test] -async fn refresh_bypasses_cached_values_and_errors_and_absence_are_retried() { +async fn refresh_bypasses_cached_values() { let cache = SecretCache::new(10, Duration::from_secs(60)); cache.insert("key", "old").await; assert_eq!( @@ -127,9 +127,21 @@ async fn refresh_bypasses_cached_values_and_errors_and_absence_are_retried() { .unwrap(), Some("fresh") ); - for result in [Err("failure"), Ok(None), Ok(Some("recovered"))] { - assert_eq!(cache.read("retry", async { result }).await, result); - } +} + +#[rstest] +#[case::failure(Err("failure"))] +#[case::absence(Ok(None))] +#[tokio::test] +async fn errors_and_absence_are_retried(#[case] first: Result, &'static str>) { + let cache = SecretCache::new(10, Duration::from_secs(60)); + assert_eq!(cache.read("retry", async { first }).await, first); + assert_eq!( + cache + .read("retry", async { Ok::<_, &str>(Some("recovered")) }) + .await, + Ok(Some("recovered")) + ); } #[tokio::test] diff --git a/litellm-rust/crates/secrets/Cargo.toml b/litellm-rust/crates/secrets/Cargo.toml index 3655ce8bbc2..023d61136e4 100644 --- a/litellm-rust/crates/secrets/Cargo.toml +++ b/litellm-rust/crates/secrets/Cargo.toml @@ -23,6 +23,7 @@ litellm-secrets-hashicorp = { workspace = true, optional = true } litellm-secrets-azure = { workspace = true, optional = true } litellm-secrets-cyberark = { workspace = true, optional = true } litellm-core-utils.workspace = true +litellm-error.workspace = true litellm-http.workspace = true base64.workspace = true serde.workspace = true diff --git a/litellm-rust/crates/secrets/src/compatibility.rs b/litellm-rust/crates/secrets/src/compatibility.rs index 621f7a3f432..368c9931b51 100644 --- a/litellm-rust/crates/secrets/src/compatibility.rs +++ b/litellm-rust/crates/secrets/src/compatibility.rs @@ -40,11 +40,9 @@ pub async fn get_secret_from_python_manager( }; match result { #[cfg(feature = "google")] - Err(Error::Google(litellm_secrets_google::Error::Status(404))) => { - Err(Error::ManagedSecretMissing) - } + Err(Error::Google(error)) if error.is_not_found() => Err(Error::ManagedSecretMissing), #[cfg(feature = "azure")] - Err(Error::Azure(litellm_secrets_azure::Error::MissingValue)) => Ok(None), + Err(Error::Azure(error)) if error.is_missing_value() => Ok(None), Ok(None) if matches!(manager, SecretManager::External(_)) && manager.system() == KeyManagementSystem::AzureKeyVault => diff --git a/litellm-rust/crates/secrets/src/error.rs b/litellm-rust/crates/secrets/src/error.rs index 7fc756c5c7d..ade8555242d 100644 --- a/litellm-rust/crates/secrets/src/error.rs +++ b/litellm-rust/crates/secrets/src/error.rs @@ -7,9 +7,9 @@ pub enum Error { #[error("encrypted environment value is missing")] MissingCiphertext, #[error("ciphertext is not valid base64 for the configured manager")] - InvalidCiphertext, + InvalidCiphertext(#[source] Option), #[error("decrypted value is not UTF-8")] - Utf8, + Utf8(#[source] std::str::Utf8Error), #[error(transparent)] Client(#[from] litellm_http::Error), #[error("unsupported OIDC provider or missing build feature")] @@ -19,21 +19,21 @@ pub enum Error { #[error("OIDC environment variable is missing")] MissingEnvironment, #[error("OIDC request failed")] - OidcHttp, + OidcHttp(#[source] litellm_http::Error), #[error("OIDC provider returned HTTP {0}")] OidcStatus(u16), #[error("OIDC response is invalid")] - OidcResponse, + OidcResponse(#[source] litellm_http::Error), #[error("OIDC file path must be absolute and within the credential allowlist")] UnsafeOidcPath, #[error("OIDC file could not be read")] - OidcFile, + OidcFile(#[source] std::io::Error), #[error("secret cannot be converted to {expected}")] TypeMismatch { expected: &'static str }, #[error("external secret manager failed")] - ExternalManager(#[source] Box), + ExternalManager(#[source] litellm_error::Source), #[error("external secret manager read failed")] - ExternalRead(#[source] Box), + ExternalRead(#[source] litellm_error::Source), #[cfg(feature = "aws")] #[error(transparent)] Aws(#[from] litellm_secrets_aws::Error), diff --git a/litellm-rust/crates/secrets/src/handler.rs b/litellm-rust/crates/secrets/src/handler.rs index db37d793384..f42e09c6bdb 100644 --- a/litellm-rust/crates/secrets/src/handler.rs +++ b/litellm-rust/crates/secrets/src/handler.rs @@ -87,7 +87,8 @@ pub async fn get_secret_from_manager( let plaintext = client .decrypt(decode_ciphertext(&ciphertext, Base64Mode::Permissive)?) .await?; - let value = String::from_utf8(plaintext).map_err(|_| Error::Utf8)?; + let value = + String::from_utf8(plaintext).map_err(|error| Error::Utf8(error.utf8_error()))?; Ok(Some(Secret::String(SecretValue::new(value.trim())))) } #[cfg(feature = "google")] @@ -98,7 +99,8 @@ pub async fn get_secret_from_manager( let plaintext = client .decrypt(decode_ciphertext(&ciphertext, Base64Mode::Canonical)?) .await?; - let value = String::from_utf8(plaintext).map_err(|_| Error::Utf8)?; + let value = + String::from_utf8(plaintext).map_err(|error| Error::Utf8(error.utf8_error()))?; Ok(Some(Secret::String(SecretValue::new(value)))) } #[cfg(feature = "aws")] @@ -149,9 +151,9 @@ fn decode_ciphertext(value: &str, mode: Base64Mode) -> Result, Error> { }; let ciphertext = STANDARD .decode(&encoded) - .map_err(|_| Error::InvalidCiphertext)?; + .map_err(|error| Error::InvalidCiphertext(Some(error)))?; if canonical && STANDARD.encode(&ciphertext) != encoded { - return Err(Error::InvalidCiphertext); + return Err(Error::InvalidCiphertext(None)); } Ok(ciphertext) } diff --git a/litellm-rust/crates/secrets/src/oidc.rs b/litellm-rust/crates/secrets/src/oidc.rs index adeb6c3db00..24e6912c797 100644 --- a/litellm-rust/crates/secrets/src/oidc.rs +++ b/litellm-rust/crates/secrets/src/oidc.rs @@ -191,7 +191,7 @@ impl OidcResolver { .header("Accept", "application/json; api-version=2.0") .send() .await - .map_err(|_| Error::OidcHttp)?; + .map_err(|error| Error::OidcHttp(error.into()))?; if response.status() != reqwest::StatusCode::OK { return Err(Error::OidcStatus(response.status().as_u16())); } @@ -199,7 +199,10 @@ impl OidcResolver { struct Token { value: Option, } - let token: Token = response.json().await.map_err(|_| Error::OidcResponse)?; + let token: Token = response + .json() + .await + .map_err(|error| Error::OidcResponse(error.into()))?; if let Some(value) = &token.value { self.cache .insert( @@ -225,11 +228,14 @@ impl OidcResolver { .header("Metadata-Flavor", "Google") .send() .await - .map_err(|_| Error::OidcHttp)?; + .map_err(|error| Error::OidcHttp(error.into()))?; if response.status() != reqwest::StatusCode::OK { return Err(Error::OidcStatus(response.status().as_u16())); } - let token = response.text().await.map_err(|_| Error::OidcResponse)?; + let token = response + .text() + .await + .map_err(|error| Error::OidcResponse(error.into()))?; let now = (self.clock)(); let ttl = oidc_token_cache_ttl(&token, now, GOOGLE_TOKEN_MAX_TTL); let value = SecretValue::new(token); @@ -259,7 +265,7 @@ async fn read_file(path: &str) -> Result { tokio::fs::read_to_string(path) .await .map(|value| SecretValue::new(value.replace("\r\n", "\n").replace('\r', "\n"))) - .map_err(|_| Error::OidcFile) + .map_err(Error::OidcFile) } async fn read_allowed_file( @@ -271,7 +277,7 @@ async fn read_allowed_file( } let resolved = tokio::fs::canonicalize(path) .await - .map_err(|_| Error::OidcFile)?; + .map_err(Error::OidcFile)?; let allowed = environment .get(OIDC_ALLOWED_CREDENTIAL_DIRS) .filter(|v| !v.is_empty()) @@ -283,7 +289,7 @@ async fn read_allowed_file( return tokio::fs::read_to_string(&resolved) .await .map(|value| SecretValue::new(value.replace("\r\n", "\n").replace('\r', "\n"))) - .map_err(|_| Error::OidcFile); + .map_err(Error::OidcFile); } } Err(Error::UnsafeOidcPath) diff --git a/litellm-rust/crates/secrets/tests/aws.rs b/litellm-rust/crates/secrets/tests/aws.rs index f28f404c3ff..abad04f31c5 100644 --- a/litellm-rust/crates/secrets/tests/aws.rs +++ b/litellm-rust/crates/secrets/tests/aws.rs @@ -218,6 +218,6 @@ async fn aws_handler_reads_ciphertext_decodes_trims_and_redacts() { )); assert!(matches!( get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| Some("abc".into())).await, - Err(Error::InvalidCiphertext) + Err(Error::InvalidCiphertext(_)) )); } diff --git a/litellm-rust/crates/secrets/tests/google.rs b/litellm-rust/crates/secrets/tests/google.rs index fa1db9916e8..5542c864dfd 100644 --- a/litellm-rust/crates/secrets/tests/google.rs +++ b/litellm-rust/crates/secrets/tests/google.rs @@ -108,7 +108,7 @@ async fn google_handler_requires_canonical_base64_and_preserves_plaintext_whites STANDARD.encode("encrypted") ))) .await, - Err(Error::InvalidCiphertext) + Err(Error::InvalidCiphertext(_)) )); assert!(matches!( get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| None).await, diff --git a/litellm-rust/crates/secrets/tests/resolution.rs b/litellm-rust/crates/secrets/tests/resolution.rs index 0858089b94f..8435e051315 100644 --- a/litellm-rust/crates/secrets/tests/resolution.rs +++ b/litellm-rust/crates/secrets/tests/resolution.rs @@ -17,74 +17,81 @@ fn resolver(value: Option<&str>) -> SecretResolver { } #[rstest::rstest] -#[case::environment(false)] -#[case::manager(true)] +#[case::capitalized_true("True", Some(true))] +#[case::padded_false(" FALSE ", Some(false))] +#[case::empty("", None)] +#[case::parenthesized("(True)", None)] +#[case::json("{\"key\":1}", None)] #[tokio::test] -async fn native_reads_preserve_strings_and_report_conversion_errors(#[case] managed: bool) { - for raw in ["True", " FALSE ", "", "(True)", "{\"key\":1}"] { - let state = if managed { - SecretManagerState::new( - SecretManager::External(Arc::new(FixedManager::custom(Ok(Some(Secret::String( - SecretValue::new(raw), - )))))), - KeyManagementSettings::default(), - ) - } else { - SecretManagerState::default() - }; - let resolver = SecretResolver::new( - Arc::new(state), - Arc::new(move |_: &str| Some(raw.to_owned())), - OidcResolver::new(litellm_http::Client::for_test()), - ); - assert_eq!( - resolver - .get_secret_str("key", None) - .await - .unwrap() - .unwrap() - .expose(), - raw - ); - assert_eq!( - resolver.get_secret("key", None).await.unwrap(), - Some(Secret::String(SecretValue::new(raw))) - ); - if raw == "True" || raw == " FALSE " { - assert_eq!( - resolver.get_secret_bool("key", None).await.unwrap(), - Some(raw == "True") - ); - } else { - assert!(matches!( - resolver.get_secret_bool("key", None).await, - Err(Error::TypeMismatch { - expected: "boolean" - }) - )); - } +async fn native_reads_preserve_strings_and_report_conversion_errors( + #[case] raw: &'static str, + #[case] boolean: Option, + #[values(false, true)] managed: bool, +) { + let state = if managed { + SecretManagerState::new( + SecretManager::External(Arc::new(FixedManager::custom(Ok(Some(Secret::String( + SecretValue::new(raw), + )))))), + KeyManagementSettings::default(), + ) + } else { + SecretManagerState::default() + }; + let resolver = SecretResolver::new( + Arc::new(state), + Arc::new(move |_: &str| Some(raw.to_owned())), + OidcResolver::new(litellm_http::Client::for_test()), + ); + assert_eq!( + resolver + .get_secret_str("key", None) + .await + .unwrap() + .unwrap() + .expose(), + raw + ); + assert_eq!( + resolver.get_secret("key", None).await.unwrap(), + Some(Secret::String(SecretValue::new(raw))) + ); + match boolean { + Some(expected) => assert_eq!( + resolver.get_secret_bool("key", None).await.unwrap(), + Some(expected) + ), + None => assert!(matches!( + resolver.get_secret_bool("key", None).await, + Err(Error::TypeMismatch { + expected: "boolean" + }) + )), } } +#[rstest::rstest] +#[case::absent(Ok(None))] +#[case::failed(Err(()))] #[tokio::test] -async fn native_defaults_apply_to_absence_but_never_hide_provider_failures() { - for reply in [Ok(None), Err(())] { - let resolver = SecretResolver::new( - Arc::new(SecretManagerState::new( - SecretManager::External(Arc::new(FixedManager::custom(reply.clone()))), - KeyManagementSettings::default(), - )), - Arc::new(|_: &str| None), - OidcResolver::new(litellm_http::Client::for_test()), - ); - let result = resolver - .get_secret_str("key", Some(SecretValue::new("default"))) - .await; - match reply { - Ok(None) => assert_eq!(result.unwrap().unwrap().expose(), "default"), - Err(()) => assert!(matches!(result, Err(Error::MissingCiphertext))), - Ok(Some(_)) => unreachable!(), - } +async fn native_defaults_apply_to_absence_but_never_hide_provider_failures( + #[case] reply: Result, ()>, +) { + let resolver = SecretResolver::new( + Arc::new(SecretManagerState::new( + SecretManager::External(Arc::new(FixedManager::custom(reply.clone()))), + KeyManagementSettings::default(), + )), + Arc::new(|_: &str| None), + OidcResolver::new(litellm_http::Client::for_test()), + ); + let result = resolver + .get_secret_str("key", Some(SecretValue::new("default"))) + .await; + match reply { + Ok(None) => assert_eq!(result.unwrap().unwrap().expose(), "default"), + Err(()) => assert!(matches!(result, Err(Error::MissingCiphertext))), + Ok(Some(_)) => unreachable!(), } } diff --git a/litellm-rust/crates/spend-clickhouse/Cargo.toml b/litellm-rust/crates/spend-clickhouse/Cargo.toml index 893c5daded3..ef7f4a0efdf 100644 --- a/litellm-rust/crates/spend-clickhouse/Cargo.toml +++ b/litellm-rust/crates/spend-clickhouse/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true [dependencies] flate2.workspace = true +litellm-error.workspace = true litellm-http.workspace = true litellm-storage-clickhouse.workspace = true serde.workspace = true diff --git a/litellm-rust/crates/spend-clickhouse/src/error.rs b/litellm-rust/crates/spend-clickhouse/src/error.rs index 0d037ef06f8..c5a9f542ca6 100644 --- a/litellm-rust/crates/spend-clickhouse/src/error.rs +++ b/litellm-rust/crates/spend-clickhouse/src/error.rs @@ -1,15 +1,32 @@ +use litellm_error::Source; + #[derive(Debug, thiserror::Error)] pub enum Error { - #[error("invalid ClickHouse insert row")] - InvalidRow, - #[error("{0} must be a positive integer")] - InvalidLimit(&'static str), - #[error("database must be a nonempty SQL identifier and retention must be positive")] - InvalidSchema, - #[error("ClickHouse insert exceeds the encoded size limit")] - InsertTooLarge, + #[error("{field} must be an integer millisecond timestamp")] + NonIntegerTimestamp { field: &'static str }, + #[error("{field} is outside the supported timestamp range")] + TimestampOutOfRange { field: &'static str, source: Source }, #[error(transparent)] Storage(#[from] litellm_storage_clickhouse::Error), #[error(transparent)] Migration(#[from] sqlx::migrate::MigrateError), } + +impl Error { + pub fn is_request(&self) -> bool { + match self { + Self::NonIntegerTimestamp { .. } | Self::TimestampOutOfRange { .. } => true, + Self::Storage(error) => error.is_request(), + Self::Migration(_) => false, + } + } + + pub fn is_insert_too_large(&self) -> bool { + match self { + Self::Storage(error) => error.is_insert_too_large(), + Self::NonIntegerTimestamp { .. } + | Self::TimestampOutOfRange { .. } + | Self::Migration(_) => false, + } + } +} diff --git a/litellm-rust/crates/spend-clickhouse/src/insert.rs b/litellm-rust/crates/spend-clickhouse/src/insert.rs index 495f0d8ecb2..66d0ea59536 100644 --- a/litellm-rust/crates/spend-clickhouse/src/insert.rs +++ b/litellm-rust/crates/spend-clickhouse/src/insert.rs @@ -13,7 +13,8 @@ use sha2::{Digest, Sha256}; use time::{OffsetDateTime, format_description::well_known::Rfc3339}; use crate::Error; -use litellm_storage_clickhouse::Connection; +use litellm_error::Source; +use litellm_storage_clickhouse::{Connection, Error as StorageError}; fn max_insert_bytes() -> Result { let name = "CLICKHOUSE_TRACE_MAX_INSERT_BYTES"; @@ -22,9 +23,9 @@ fn max_insert_bytes() -> Result { .parse::() .ok() .filter(|value| *value > 0) - .ok_or(Error::InvalidLimit(name)), + .ok_or(StorageError::InvalidLimit(name).into()), Err(std::env::VarError::NotPresent) => Ok(64 * 1024 * 1024), - Err(_) => Err(Error::InvalidLimit(name)), + Err(_) => Err(StorageError::InvalidLimit(name).into()), } } @@ -55,7 +56,11 @@ pub async fn insert_rows( pub fn encode_rows(rows: Vec) -> Result { let body = write_rows(&rows, None, Vec::new(), usize::MAX)?; - String::from_utf8(body).map_err(|_| Error::InvalidRow) + String::from_utf8(body).map_err(invalid_row) +} + +fn invalid_row(error: impl std::error::Error + Send + Sync + 'static) -> Error { + StorageError::InvalidRow(Source::new(error)).into() } fn prepare_insert( @@ -73,9 +78,9 @@ fn prepare_insert( )?; let body = encoder .into_inner() - .map_err(|_| Error::InvalidRow)? + .map_err(|error| invalid_row(error.into_error()))? .finish() - .map_err(|_| Error::InvalidRow)?; + .map_err(invalid_row)?; Ok((token, body)) } @@ -102,7 +107,7 @@ impl Write for LimitedWriter { fn write(&mut self, bytes: &[u8]) -> std::io::Result { if bytes.len() > self.remaining { self.exceeded = true; - return Err(std::io::Error::other(Error::InsertTooLarge)); + return Err(std::io::Error::other(StorageError::InsertTooLarge)); } let written = self.inner.write(bytes)?; self.remaining -= written; @@ -126,17 +131,18 @@ fn write_rows( exceeded: false, }; for (index, row) in rows.iter().enumerate() { + let row = EncodedRow::new(row, received_ms)?; let result = (|| { if index != 0 { writer.write_all(b"\n").map_err(serde_json::Error::io)?; } - serde_json::to_writer(&mut writer, &EncodedRow { row, received_ms }) + serde_json::to_writer(&mut writer, &row) })(); - if result.is_err() { + if let Err(error) = result { return Err(if writer.exceeded { - Error::InsertTooLarge + StorageError::InsertTooLarge.into() } else { - Error::InvalidRow + invalid_row(error) }); } } @@ -144,25 +150,37 @@ fn write_rows( } struct EncodedRow<'a> { - row: &'a InsertRow, + fields: Vec<(&'a str, Cow<'a, Value>)>, received_ms: Option, } +impl<'a> EncodedRow<'a> { + fn new(row: &'a InsertRow, received_ms: Option) -> Result { + let fields = row + .iter() + .map(|(name, value)| Ok((name.as_str(), insert_value(name, value)?))) + .collect::>()?; + Ok(Self { + fields, + received_ms, + }) + } +} + impl Serialize for EncodedRow<'_> { fn serialize(&self, serializer: S) -> Result { let mut map = serializer.serialize_map(None)?; let mut received_ms = self.received_ms; - for (name, value) in self.row { - if name.as_str() >= "EngineReceivedMs" + for (name, value) in &self.fields { + if *name >= "EngineReceivedMs" && let Some(timestamp) = received_ms.take() { map.serialize_entry("EngineReceivedMs", ×tamp)?; } - if name == "EngineReceivedMs" && self.received_ms.is_some() { + if *name == "EngineReceivedMs" && self.received_ms.is_some() { continue; } - let value = insert_value(name, value).map_err(serde::ser::Error::custom)?; - map.serialize_entry(name, &value)?; + map.serialize_entry(name, value)?; } if let Some(timestamp) = received_ms { map.serialize_entry("EngineReceivedMs", ×tamp)?; @@ -171,21 +189,23 @@ impl Serialize for EncodedRow<'_> { } } +const TIMESTAMP_FIELDS: [&str; 3] = ["start_time", "end_time", "completion_start_time"]; + fn insert_value<'a>(name: &str, value: &'a Value) -> Result, Error> { - let multiplier = match name { - "start_time" | "end_time" | "completion_start_time" => 1_000_000, - _ => return Ok(Cow::Borrowed(value)), + let Some(&field) = TIMESTAMP_FIELDS.iter().find(|field| **field == name) else { + return Ok(Cow::Borrowed(value)); }; - if name == "completion_start_time" && value.is_null() { + if field == "completion_start_time" && value.is_null() { return Ok(Cow::Borrowed(value)); } - let timestamp = value.as_i64().ok_or(Error::InvalidRow)?; - let datetime = OffsetDateTime::from_unix_timestamp_nanos(i128::from(timestamp) * multiplier) - .map_err(|_| Error::InvalidRow)?; + let out_of_range = |source| Error::TimestampOutOfRange { field, source }; + let timestamp = value.as_i64().ok_or(Error::NonIntegerTimestamp { field })?; + let datetime = OffsetDateTime::from_unix_timestamp_nanos(i128::from(timestamp) * 1_000_000) + .map_err(|error| out_of_range(Source::new(error)))?; datetime .format(&Rfc3339) .map(|value| Cow::Owned(Value::String(value))) - .map_err(|_| Error::InvalidRow) + .map_err(|error| out_of_range(Source::new(error))) } #[cfg(test)] @@ -204,10 +224,11 @@ mod tests { ]; let encoded = write_rows(&rows, None, Vec::new(), usize::MAX).unwrap(); assert!(write_rows(&rows, None, Vec::new(), encoded.len()).is_ok()); - assert!(matches!( - write_rows(&rows, None, Vec::new(), encoded.len() - 1), - Err(Error::InsertTooLarge) - )); + assert!( + write_rows(&rows, None, Vec::new(), encoded.len() - 1) + .unwrap_err() + .is_insert_too_large() + ); } #[rstest] @@ -264,10 +285,11 @@ mod tests { fn stamped_insert_enforces_the_encoded_limit() { let rows = vec![BTreeMap::new()]; assert!(prepare_insert(&rows, 1, 22).is_ok()); - assert!(matches!( - prepare_insert(&rows, 1, 21), - Err(Error::InsertTooLarge) - )); + assert!( + prepare_insert(&rows, 1, 21) + .unwrap_err() + .is_insert_too_large() + ); } #[rstest] @@ -282,15 +304,36 @@ mod tests { assert_eq!(insert_value(field, &input).unwrap().as_ref(), &expected); } + fn rejected_start_time(input: Value) -> Error { + prepare_insert(&[BTreeMap::from([("start_time".into(), input)])], 456, 4096).unwrap_err() + } + #[rstest] #[case::fractional(json!(1.5))] #[case::string(json!("1234"))] #[case::null(Value::Null)] - #[case::out_of_range(json!(i64::MAX))] - fn invalid_spend_dates_fail_before_transport(#[case] input: Value) { + fn non_integer_spend_dates_fail_before_transport(#[case] input: Value) { + let error = rejected_start_time(input); + assert!(error.is_request()); assert!(matches!( - prepare_insert(&[BTreeMap::from([("start_time".into(), input)])], 456, 4096), - Err(Error::InvalidRow) + error, + Error::NonIntegerTimestamp { + field: "start_time" + } )); } + + #[rstest] + fn out_of_range_spend_dates_keep_their_cause() { + let error = rejected_start_time(json!(i64::MAX)); + assert!(error.is_request()); + let cause = std::error::Error::source(&error) + .and_then(|cause| cause.downcast_ref::()) + .expect("time cause"); + assert!( + cause + .downcast_ref::() + .is_some() + ); + } } diff --git a/litellm-rust/crates/spend-clickhouse/src/lib.rs b/litellm-rust/crates/spend-clickhouse/src/lib.rs index f4c108b1e7f..9da658aa3e8 100644 --- a/litellm-rust/crates/spend-clickhouse/src/lib.rs +++ b/litellm-rust/crates/spend-clickhouse/src/lib.rs @@ -4,7 +4,7 @@ mod schema; pub use error::Error; pub use insert::{encode_rows, insert_rows}; -use litellm_storage_clickhouse::Storage; +use litellm_storage_clickhouse::{Error as StorageError, Storage}; pub use schema::ensure_schema; #[derive(Clone)] @@ -16,7 +16,7 @@ pub struct Config { impl Config { pub fn new(database: String, url: &str, retention_days: u32) -> Result { if retention_days == 0 { - return Err(Error::InvalidSchema); + return Err(StorageError::InvalidSchema.into()); } Ok(Self { storage: Storage::new(database, url)?, diff --git a/litellm-rust/crates/storage-clickhouse/Cargo.toml b/litellm-rust/crates/storage-clickhouse/Cargo.toml index 39d04f2de5d..8b985fc40d8 100644 --- a/litellm-rust/crates/storage-clickhouse/Cargo.toml +++ b/litellm-rust/crates/storage-clickhouse/Cargo.toml @@ -8,6 +8,7 @@ repository.workspace = true [dependencies] flate2.workspace = true +litellm-error.workspace = true litellm-http.workspace = true serde.workspace = true serde_json.workspace = true @@ -17,6 +18,7 @@ thiserror.workspace = true url.workspace = true [dev-dependencies] +litellm-error.workspace = true litellm-http = { workspace = true, features = ["test-support"] } rstest.workspace = true testcontainers-modules = { workspace = true, features = ["clickhouse"] } diff --git a/litellm-rust/crates/storage-clickhouse/src/error.rs b/litellm-rust/crates/storage-clickhouse/src/error.rs index fa0c5eabd73..ef7b5046455 100644 --- a/litellm-rust/crates/storage-clickhouse/src/error.rs +++ b/litellm-rust/crates/storage-clickhouse/src/error.rs @@ -1,21 +1,23 @@ +use litellm_error::Source; + #[derive(Clone, Debug, thiserror::Error)] pub enum Error { #[error("invalid ClickHouse insert row")] - InvalidRow, + InvalidRow(#[source] Source), #[error("{0} must be a positive integer")] InvalidLimit(&'static str), #[error("invalid ClickHouse insert table")] InvalidTable, #[error("invalid ClickHouse HTTP URL")] - InvalidUrl, + InvalidUrl(#[source] url::ParseError), + #[error("ClickHouse URL must use http or https and name a host")] + UnsupportedUrl, #[error("database must be a nonempty SQL identifier and retention must be positive")] InvalidSchema, #[error("SQL query must not be empty")] EmptySql, #[error("invalid ClickHouse query parameters")] - InvalidParameters, - #[error("unknown ClickHouse read query")] - InvalidQuery, + InvalidParameters(#[source] Source), #[error("ClickHouse query failed with HTTP status {0}")] QueryFailed(u16), #[error("ClickHouse insert failed with HTTP status {0}")] @@ -24,10 +26,60 @@ pub enum Error { InsertTooLarge, #[error("ClickHouse schema setup failed with HTTP status {0}")] SchemaFailed(u16), + #[error("ClickHouse migrations are forward-only")] + Irreversible, #[error("ClickHouse query exceeded the response size limit")] ResponseTooLarge, - #[error("ClickHouse returned an invalid or failed JSON query response")] - InvalidResponse, - #[error("ClickHouse query transport failed")] - Transport, + #[error("ClickHouse returned an invalid response")] + InvalidResponse(#[source] Source), + #[error("ClickHouse returned a failed or unexpected response")] + UnexpectedResponse, + #[error("ClickHouse transport failed")] + Transport(#[from] litellm_http::Error), +} + +impl Error { + pub fn is_request(&self) -> bool { + match self { + Self::InvalidRow(_) + | Self::InvalidLimit(_) + | Self::InvalidTable + | Self::InvalidSchema + | Self::EmptySql + | Self::InvalidParameters(_) => true, + Self::InvalidUrl(_) + | Self::UnsupportedUrl + | Self::QueryFailed(_) + | Self::InsertFailed(_) + | Self::InsertTooLarge + | Self::SchemaFailed(_) + | Self::Irreversible + | Self::ResponseTooLarge + | Self::InvalidResponse(_) + | Self::UnexpectedResponse + | Self::Transport(_) => false, + } + } + + pub fn is_insert_too_large(&self) -> bool { + match self { + Self::InsertTooLarge => true, + Self::InvalidRow(_) + | Self::InvalidLimit(_) + | Self::InvalidTable + | Self::InvalidUrl(_) + | Self::UnsupportedUrl + | Self::InvalidSchema + | Self::EmptySql + | Self::InvalidParameters(_) + | Self::QueryFailed(_) + | Self::InsertFailed(_) + | Self::SchemaFailed(_) + | Self::Irreversible + | Self::ResponseTooLarge + | Self::InvalidResponse(_) + | Self::UnexpectedResponse + | Self::Transport(_) => false, + } + } } diff --git a/litellm-rust/crates/storage-clickhouse/src/insert.rs b/litellm-rust/crates/storage-clickhouse/src/insert.rs index 1cdad1897c1..29837359f0b 100644 --- a/litellm-rust/crates/storage-clickhouse/src/insert.rs +++ b/litellm-rust/crates/storage-clickhouse/src/insert.rs @@ -1,6 +1,7 @@ use std::{io::Write, time::Duration}; use flate2::{Compression, write::GzEncoder}; +use litellm_error::Source; use litellm_http::Client; use crate::{Connection, Error, valid_identifier}; @@ -34,13 +35,15 @@ pub async fn insert_encoded_rows( return Err(Error::InvalidTable); } let mut encoder = GzEncoder::new(Vec::new(), Compression::default()); - encoder - .write_all(encoded.as_bytes()) - .map_err(|_| Error::InvalidRow)?; - let body = encoder.finish().map_err(|_| Error::InvalidRow)?; + encoder.write_all(encoded.as_bytes()).map_err(invalid_row)?; + let body = encoder.finish().map_err(invalid_row)?; insert_compressed_rows(client, connection, database, table, token, body).await } +fn invalid_row(error: std::io::Error) -> Error { + Error::InvalidRow(Source::new(error)) +} + pub async fn insert_compressed_rows( client: &Client, connection: &Connection, @@ -91,7 +94,7 @@ pub async fn insert_compressed_rows( .body(body) .send() .await - .map_err(|_| Error::Transport)?; + .map_err(litellm_http::Error::from)?; if !response.status().is_success() { return Err(Error::InsertFailed(response.status().as_u16())); } diff --git a/litellm-rust/crates/storage-clickhouse/src/lib.rs b/litellm-rust/crates/storage-clickhouse/src/lib.rs index e2ef62d1bc3..60a4e77a36c 100644 --- a/litellm-rust/crates/storage-clickhouse/src/lib.rs +++ b/litellm-rust/crates/storage-clickhouse/src/lib.rs @@ -16,9 +16,9 @@ pub struct Connection { impl Connection { pub fn parse(value: &str) -> Result { - let url = Url::parse(value).map_err(|_| Error::InvalidUrl)?; + let url = Url::parse(value).map_err(Error::InvalidUrl)?; if !matches!(url.scheme(), "http" | "https") || url.host().is_none() { - return Err(Error::InvalidUrl); + return Err(Error::UnsupportedUrl); } Ok(Self { url }) } @@ -33,11 +33,11 @@ impl Connection { connection .url .set_username(user) - .map_err(|_| Error::InvalidUrl)?; + .map_err(|()| Error::UnsupportedUrl)?; connection .url .set_password(Some(password)) - .map_err(|_| Error::InvalidUrl)?; + .map_err(|()| Error::UnsupportedUrl)?; let pairs: Vec<_> = connection .url .query_pairs() diff --git a/litellm-rust/crates/storage-clickhouse/src/migrate.rs b/litellm-rust/crates/storage-clickhouse/src/migrate.rs index c249020faaf..4da26baa681 100644 --- a/litellm-rust/crates/storage-clickhouse/src/migrate.rs +++ b/litellm-rust/crates/storage-clickhouse/src/migrate.rs @@ -12,7 +12,7 @@ use sqlx::{ migrate::{AppliedMigration, Migrate, MigrateError, Migration}, }; -use crate::{Connection, Error, READ_LIMITS, valid_identifier}; +use crate::{Connection, Error, READ_LIMITS, read::invalid_response, valid_identifier}; pub async fn execute_statement( client: &Client, @@ -22,7 +22,7 @@ pub async fn execute_statement( ) -> Result<(), Error> { let body = execute_sql(client, connection, sql, timeout).await?; if !body.trim().is_empty() { - return Err(Error::InvalidResponse); + return Err(Error::UnexpectedResponse); } Ok(()) } @@ -56,18 +56,18 @@ async fn execute_sql( .body(sql.to_owned()) .send() .await - .map_err(|_| Error::Transport)?; + .map_err(litellm_http::Error::from)?; if !response.status().is_success() { return Err(Error::SchemaFailed(response.status().as_u16())); } let mut body = Vec::new(); - while let Some(chunk) = response.chunk().await.map_err(|_| Error::Transport)? { + while let Some(chunk) = response.chunk().await.map_err(litellm_http::Error::from)? { if body.len() + chunk.len() > READ_LIMITS.response_bytes { return Err(Error::ResponseTooLarge); } body.extend_from_slice(&chunk); } - String::from_utf8(body).map_err(|_| Error::InvalidResponse) + String::from_utf8(body).map_err(invalid_response) } /// The startup runner records success after execution, never marks migrations dirty, and skips locks @@ -131,13 +131,13 @@ fn encode_hex(bytes: &[u8]) -> String { fn decode_hex(value: &str) -> Result, Error> { let (pairs, remainder) = value.as_bytes().as_chunks::<2>(); if !remainder.is_empty() { - return Err(Error::InvalidResponse); + return Err(Error::UnexpectedResponse); } pairs .iter() .map(|pair| { - let high = decode_hex_digit(pair[0]).ok_or(Error::InvalidResponse)?; - let low = decode_hex_digit(pair[1]).ok_or(Error::InvalidResponse)?; + let high = decode_hex_digit(pair[0]).ok_or(Error::UnexpectedResponse)?; + let low = decode_hex_digit(pair[1]).ok_or(Error::UnexpectedResponse)?; Ok((high << 4) | low) }) .collect() @@ -231,7 +231,7 @@ where .filter(|line| !line.trim().is_empty()) .map(|line| { let row = serde_json::from_str::(line) - .map_err(|_| migrate_error(Error::InvalidResponse))?; + .map_err(|error| migrate_error(invalid_response(error)))?; let checksum = decode_hex(&row.checksum).map_err(migrate_error)?; Ok(AppliedMigration { version: row.version, @@ -288,11 +288,7 @@ where _table_name: &'e str, _migration: &'e Migration, ) -> MigrateFuture<'e, Result> { - Box::pin(async { - Err(MigrateError::Execute(SqlxError::AnyDriverError(Box::new( - std::io::Error::other("ClickHouse migrations are forward-only"), - )))) - }) + Box::pin(async { Err(migrate_error(Error::Irreversible)) }) } } diff --git a/litellm-rust/crates/storage-clickhouse/src/read.rs b/litellm-rust/crates/storage-clickhouse/src/read.rs index 03a5522e595..40d4ad0a073 100644 --- a/litellm-rust/crates/storage-clickhouse/src/read.rs +++ b/litellm-rust/crates/storage-clickhouse/src/read.rs @@ -1,5 +1,6 @@ use std::{collections::BTreeMap, time::Duration}; +use litellm_error::Source; use litellm_http::Client; use serde::{Deserialize, Serialize, de::DeserializeOwned}; @@ -115,7 +116,7 @@ async fn execute_read_with_limits( .post(url) .timeout(Duration::from_secs(15)) .body(sql.to_owned()); - let mut response = request.send().await.map_err(|_| Error::Transport)?; + let mut response = request.send().await.map_err(litellm_http::Error::from)?; if !response.status().is_success() { if response .headers() @@ -128,20 +129,19 @@ async fn execute_read_with_limits( } let mut body = Vec::new(); - while let Some(chunk) = response.chunk().await.map_err(|_| Error::Transport)? { + while let Some(chunk) = response.chunk().await.map_err(litellm_http::Error::from)? { if body.len() + chunk.len() > limits.response_bytes { return Err(Error::ResponseTooLarge); } body.extend_from_slice(&chunk); } - let json: serde_json::Value = - serde_json::from_slice(&body).map_err(|_| Error::InvalidResponse)?; + let json: serde_json::Value = serde_json::from_slice(&body).map_err(invalid_response)?; if json.get("exception").is_some() || !json.get("data").is_some_and(serde_json::Value::is_array) { - return Err(Error::InvalidResponse); + return Err(Error::UnexpectedResponse); } - String::from_utf8(body).map_err(|_| Error::InvalidResponse) + String::from_utf8(body).map_err(invalid_response) } pub trait Query { @@ -158,8 +158,8 @@ struct Rows { } fn parameters(params: &T) -> Result, Error> { - let value = serde_json::to_value(params).map_err(|_| Error::InvalidParameters)?; - serde_json::from_value(value).map_err(|_| Error::InvalidParameters) + let value = serde_json::to_value(params).map_err(invalid_parameters)?; + serde_json::from_value(value).map_err(invalid_parameters) } pub async fn fetch( @@ -198,5 +198,13 @@ pub async fn fetch_json( fn decode_rows(body: &str) -> Result, Error> { serde_json::from_str::>(body) .map(|rows| rows.data) - .map_err(|_| Error::InvalidResponse) + .map_err(invalid_response) +} + +pub(crate) fn invalid_response(error: impl std::error::Error + Send + Sync + 'static) -> Error { + Error::InvalidResponse(Source::new(error)) +} + +fn invalid_parameters(error: serde_json::Error) -> Error { + Error::InvalidParameters(Source::new(error)) } diff --git a/litellm-rust/crates/storage-clickhouse/tests/migrations.rs b/litellm-rust/crates/storage-clickhouse/tests/migrations.rs index 49f481df632..09e56766891 100644 --- a/litellm-rust/crates/storage-clickhouse/tests/migrations.rs +++ b/litellm-rust/crates/storage-clickhouse/tests/migrations.rs @@ -346,11 +346,7 @@ async fn revert_reports_forward_only_error() { .revert("_sqlx_migrations", &migration) .await .expect_err("ClickHouse migrations cannot be reverted"); - assert!( - error - .to_string() - .contains("ClickHouse migrations are forward-only") - ); + assert!(matches!(storage_error(&error), Some(Error::Irreversible))); } #[rstest] @@ -391,6 +387,7 @@ async fn applied_int64_versions_accept_numeric_and_quoted_json( #[tokio::test] async fn schema_requests_override_unsafe_connection_settings( #[future(awt)] mock_server: MockServer, + #[values("wait_end_of_query", "send_progress_in_http_headers", "async_insert")] name: &str, ) { Mock::given(method("POST")) .and(query_param("wait_end_of_query", "1")) @@ -417,20 +414,14 @@ async fn schema_requests_override_unsafe_connection_settings( .received_requests() .await .expect("requests recorded"); - for name in [ - "wait_end_of_query", - "send_progress_in_http_headers", - "async_insert", - ] { - assert_eq!( - requests[0] - .url - .query_pairs() - .filter(|(key, _)| key == name) - .count(), - 1 - ); - } + assert_eq!( + requests[0] + .url + .query_pairs() + .filter(|(key, _)| key == name) + .count(), + 1 + ); } #[rstest] diff --git a/litellm-rust/crates/storage-clickhouse/tests/transport.rs b/litellm-rust/crates/storage-clickhouse/tests/transport.rs index 2560c2985f0..6b2c3fb8ef1 100644 --- a/litellm-rust/crates/storage-clickhouse/tests/transport.rs +++ b/litellm-rust/crates/storage-clickhouse/tests/transport.rs @@ -55,18 +55,24 @@ impl litellm_storage_clickhouse::Query for TypedQuery { const SQL: &'static str = "SELECT typed_parameters"; } +enum Expected { + Rows, + Undecodable, + Failed, +} + #[rstest] #[case::valid( r#"{"meta":[],"data":[{"answer":"ok"}],"rows":1,"statistics":{"elapsed":0.1}}"#, - true + Expected::Rows )] -#[case::wrong_type(r#"{"data":[{"answer":1}]}"#, false)] -#[case::missing_column(r#"{"data":[{}]}"#, false)] -#[case::exception(r#"{"data":[],"exception":"failed"}"#, false)] +#[case::wrong_type(r#"{"data":[{"answer":1}]}"#, Expected::Undecodable)] +#[case::missing_column(r#"{"data":[{}]}"#, Expected::Undecodable)] +#[case::exception(r#"{"data":[],"exception":"failed"}"#, Expected::Failed)] #[tokio::test] async fn typed_fetch_encodes_parameters_and_validates_rows( #[case] body: &str, - #[case] valid: bool, + #[case] expected: Expected, ) { use litellm_storage_clickhouse::{fetch, fetch_json}; use wiremock::{ @@ -98,17 +104,24 @@ async fn typed_fetch_encodes_parameters_and_validates_rows( }; let rows = fetch::(&client, &connection, ¶ms).await; let envelope = fetch_json::(&client, &connection, ¶ms).await; - if valid { - assert_eq!( - rows.unwrap(), - vec![QueryRow { - answer: "ok".into() - }] - ); - assert_eq!(envelope.unwrap(), body); - } else { - assert!(matches!(rows, Err(Error::InvalidResponse))); - assert!(matches!(envelope, Err(Error::InvalidResponse))); + match expected { + Expected::Rows => { + assert_eq!( + rows.unwrap(), + vec![QueryRow { + answer: "ok".into() + }] + ); + assert_eq!(envelope.unwrap(), body); + } + Expected::Undecodable => { + assert!(matches!(rows, Err(Error::InvalidResponse(_)))); + assert!(matches!(envelope, Err(Error::InvalidResponse(_)))); + } + Expected::Failed => { + assert!(matches!(rows, Err(Error::UnexpectedResponse))); + assert!(matches!(envelope, Err(Error::UnexpectedResponse))); + } } } @@ -145,21 +158,23 @@ async fn server_result_limits_allow_smaller_pages_without_retrying_other_failure } } -#[test] -fn insert_timeout_environment_controls_transport() { - for value in ["1", "3", "0", "invalid"] { - let result = std::process::Command::new(std::env::current_exe().unwrap()) - .args(["--exact", "insert_timeout_environment_child"]) - .env("LITELLM_TEST_INSERT_TIMEOUT", value) - .env("CLICKHOUSE_INSERT_TIMEOUT_SECONDS", value) - .output() - .unwrap(); - assert!( - result.status.success(), - "{}", - String::from_utf8_lossy(&result.stdout) - ); - } +#[rstest] +#[case::times_out("1")] +#[case::completes("3")] +#[case::zero("0")] +#[case::not_a_number("invalid")] +fn insert_timeout_environment_controls_transport(#[case] value: &str) { + let result = std::process::Command::new(std::env::current_exe().unwrap()) + .args(["--exact", "insert_timeout_environment_child"]) + .env("LITELLM_TEST_INSERT_TIMEOUT", value) + .env("CLICKHOUSE_INSERT_TIMEOUT_SECONDS", value) + .output() + .unwrap(); + assert!( + result.status.success(), + "{}", + String::from_utf8_lossy(&result.stdout) + ); } #[tokio::test] @@ -183,7 +198,7 @@ async fn insert_timeout_environment_child() { ) .await; match value.as_str() { - "1" => assert!(matches!(result, Err(Error::Transport))), + "1" => assert!(matches!(result, Err(Error::Transport(error)) if error.dispatched())), "3" => assert!(result.is_ok()), _ => assert!(matches!( result, @@ -191,3 +206,29 @@ async fn insert_timeout_environment_child() { )), } } + +#[rstest] +#[tokio::test] +async fn unreachable_server_keeps_the_transport_cause() { + let address = std::net::TcpListener::bind("127.0.0.1:0") + .unwrap() + .local_addr() + .unwrap(); + let connection = Connection::parse(&format!("http://{address}")).unwrap(); + let error = execute_read( + &Client::for_test(), + &connection, + "SELECT 1", + &BTreeMap::new(), + ) + .await + .unwrap_err(); + let cause = std::error::Error::source(&error) + .and_then(|cause| cause.downcast_ref::()) + .expect("transport cause"); + assert!(!cause.dispatched()); + assert_ne!( + litellm_error::Report::new(&error).to_string(), + error.to_string() + ); +} diff --git a/litellm-rust/crates/testkit/Cargo.toml b/litellm-rust/crates/testkit/Cargo.toml index 97efb753384..ef3c3e6cd7e 100644 --- a/litellm-rust/crates/testkit/Cargo.toml +++ b/litellm-rust/crates/testkit/Cargo.toml @@ -8,6 +8,7 @@ publish = false [dependencies] flate2.workspace = true +litellm-http.workspace = true reqwest.workspace = true serde.workspace = true semver.workspace = true diff --git a/litellm-rust/crates/testkit/src/agent/configure.rs b/litellm-rust/crates/testkit/src/agent/configure.rs index a41bec9379a..6f1b970e8fb 100644 --- a/litellm-rust/crates/testkit/src/agent/configure.rs +++ b/litellm-rust/crates/testkit/src/agent/configure.rs @@ -5,10 +5,13 @@ use semver::Version; use crate::Error; -#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray, strum::Display)] pub enum Wire { + #[strum(serialize = "chat completions")] ChatCompletions, + #[strum(serialize = "messages")] Messages, + #[strum(serialize = "responses")] Responses, } diff --git a/litellm-rust/crates/testkit/src/error.rs b/litellm-rust/crates/testkit/src/error.rs index 03520d827f1..8b747c28bda 100644 --- a/litellm-rust/crates/testkit/src/error.rs +++ b/litellm-rust/crates/testkit/src/error.rs @@ -15,7 +15,7 @@ pub enum Error { Request { url: String, #[source] - source: reqwest::Error, + source: litellm_http::Error, }, #[error("{url} answered with status {status}")] Status { url: String, status: u16 }, @@ -47,10 +47,14 @@ pub enum Error { expected: String, reported: String, }, - #[error("{agent} cannot talk to the gateway over {wire:?}")] + #[error("{agent} cannot talk to the gateway over {wire}")] UnsupportedWire { agent: &'static str, wire: Wire }, - #[error("agent did not finish within {0:?}")] + #[error("agent did not finish within {} seconds", .0.as_secs_f64())] Timeout(std::time::Duration), - #[error("io failure")] - Io(#[from] io::Error), + #[error("could not install the binary into {}", dir.display())] + Install { dir: PathBuf, source: io::Error }, + #[error("could not run {}", binary.display())] + Spawn { binary: PathBuf, source: io::Error }, + #[error("could not write agent files under {}", home.display())] + WriteFiles { home: PathBuf, source: io::Error }, } diff --git a/litellm-rust/crates/testkit/src/install/fetch.rs b/litellm-rust/crates/testkit/src/install/fetch.rs index 73008f7a0da..59a61bfaa0a 100644 --- a/litellm-rust/crates/testkit/src/install/fetch.rs +++ b/litellm-rust/crates/testkit/src/install/fetch.rs @@ -38,9 +38,9 @@ impl Fetch for HttpFetch { (Some(token), true) => request.bearer_auth(token), _ => request, }; - let request_error = |source| Error::Request { + let request_error = |source: reqwest::Error| Error::Request { url: url.to_owned(), - source, + source: source.into(), }; let response = request.send().await.map_err(request_error)?; let status = response.status(); diff --git a/litellm-rust/crates/testkit/src/install/mod.rs b/litellm-rust/crates/testkit/src/install/mod.rs index 1104bcec102..73f7deec864 100644 --- a/litellm-rust/crates/testkit/src/install/mod.rs +++ b/litellm-rust/crates/testkit/src/install/mod.rs @@ -52,7 +52,13 @@ impl Installer { version: version.clone(), binary: binary.clone(), }; - if fs::try_exists(&binary).await? && probe_version(&binary, version).await.is_ok() { + let install_error = |source| Error::Install { + dir: dir.clone(), + source, + }; + if fs::try_exists(&binary).await.map_err(install_error)? + && probe_version(&binary, version).await.is_ok() + { return Ok(installed); } @@ -61,21 +67,23 @@ impl Installer { verify_sha256(&release.asset, &release.sha256, &archive)?; let contents = extract_binary(&release.packaging, &archive)?; - fs::create_dir_all(&dir).await?; + fs::create_dir_all(&dir).await.map_err(install_error)?; let staging = dir.join(format!( ".{}.{}.{}.partial", agent.binary(), std::process::id(), STAGING_COUNTER.fetch_add(1, Ordering::Relaxed) )); - fs::write(&staging, contents).await?; - fs::set_permissions(&staging, std::fs::Permissions::from_mode(0o755)).await?; - fs::rename(&staging, &binary).await?; + fs::write(&staging, contents).await.map_err(install_error)?; + fs::set_permissions(&staging, std::fs::Permissions::from_mode(0o755)) + .await + .map_err(install_error)?; + fs::rename(&staging, &binary).await.map_err(install_error)?; match probe_version(&binary, version).await { Ok(()) => Ok(installed), Err(error) => { - fs::remove_file(&binary).await?; + fs::remove_file(&binary).await.map_err(install_error)?; Err(error) } } @@ -98,7 +106,11 @@ async fn probe_version(binary: &Path, expected: &Version) -> Result<(), Error> { .env("DISABLE_AUTOUPDATER", "1") .stdin(Stdio::null()) .output() - .await?; + .await + .map_err(|source| Error::Spawn { + binary: binary.to_owned(), + source, + })?; let stdout = String::from_utf8_lossy(&output.stdout); if stdout .split_whitespace() diff --git a/litellm-rust/crates/testkit/src/session.rs b/litellm-rust/crates/testkit/src/session.rs index 6b06e756cca..cd5eb08ae2b 100644 --- a/litellm-rust/crates/testkit/src/session.rs +++ b/litellm-rust/crates/testkit/src/session.rs @@ -28,7 +28,11 @@ impl Session { ) -> Result { let home = home.into(); let spec = agent.configure(&installed.version, &settings, &home)?; - spec.write_files(&home)?; + spec.write_files(&home) + .map_err(|source| Error::WriteFiles { + home: home.clone(), + source, + })?; Ok(Self { binary: installed.binary.clone(), home, @@ -55,7 +59,11 @@ impl Session { .output(); let output = timeout(limit, child) .await - .map_err(|_| Error::Timeout(limit))??; + .map_err(|_| Error::Timeout(limit))? + .map_err(|source| Error::Spawn { + binary: self.binary.clone(), + source, + })?; let parsed = agent.parse(&self.version, &String::from_utf8_lossy(&output.stdout)); let failed_silently = !output.status.success() && parsed.errors.is_empty(); Ok(Outcome { diff --git a/litellm-rust/crates/token-counter-fast/Cargo.toml b/litellm-rust/crates/token-counter-fast/Cargo.toml index 5127dc17f22..070bec58519 100644 --- a/litellm-rust/crates/token-counter-fast/Cargo.toml +++ b/litellm-rust/crates/token-counter-fast/Cargo.toml @@ -13,6 +13,7 @@ tokenizers.workspace = true unicode-normalization-alignments = "0.1.12" [dev-dependencies] +litellm-error.workspace = true rand.workspace = true rstest.workspace = true serde.workspace = true diff --git a/litellm-rust/crates/token-counter-fast/src/error.rs b/litellm-rust/crates/token-counter-fast/src/error.rs index d2a706191ac..2cd8d5c810c 100644 --- a/litellm-rust/crates/token-counter-fast/src/error.rs +++ b/litellm-rust/crates/token-counter-fast/src/error.rs @@ -2,16 +2,25 @@ use thiserror::Error as ThisError; #[derive(Debug, ThisError)] pub enum Error { - #[error("failed to load tokenizer: {0}")] + #[error("failed to load tokenizer")] Load(#[source] tokenizers::Error), - #[error("failed to load tokenizer: tiktoken rank file: {0}")] + #[error("failed to load tokenizer from a tiktoken rank file")] Ranks(#[source] RankError), #[error("failed to load tokenizer: Unicode character classes are unavailable")] UnicodeClasses, - #[error("tokenization failed: {0}")] + #[error("tokenization failed")] Encode(#[source] tokenizers::Error), } +impl Error { + pub fn is_load(&self) -> bool { + match self { + Self::Load(_) | Self::Ranks(_) | Self::UnicodeClasses => true, + Self::Encode(_) => false, + } + } +} + #[derive(Debug, ThisError)] pub enum RankError { #[error("rank {0} is reserved")] @@ -20,8 +29,8 @@ pub enum RankError { MissingByte(u8), #[error("line without a rank: {0:?}")] MissingRank(String), - #[error("token is not base64: {0}")] + #[error("token is not base64")] InvalidToken(#[source] base64::DecodeError), - #[error("rank is not an integer: {0}")] + #[error("rank is not an integer")] InvalidRank(#[source] std::num::ParseIntError), } diff --git a/litellm-rust/crates/token-counter-fast/src/tiktoken.rs b/litellm-rust/crates/token-counter-fast/src/tiktoken.rs index f2d4de38f8c..11182e84e8d 100644 --- a/litellm-rust/crates/token-counter-fast/src/tiktoken.rs +++ b/litellm-rust/crates/token-counter-fast/src/tiktoken.rs @@ -254,8 +254,12 @@ mod tests { .err() .expect("invalid rank file"); assert!(matches!(error, Error::Ranks(_))); - assert!(error.to_string().starts_with(&format!( - "failed to load tokenizer: tiktoken rank file: {detail}" - ))); + assert!( + litellm_error::Report::new(&error) + .to_string() + .starts_with(&format!( + "failed to load tokenizer from a tiktoken rank file: {detail}" + )) + ); } } diff --git a/litellm-rust/crates/token-counter-huggingface/src/error.rs b/litellm-rust/crates/token-counter-huggingface/src/error.rs index e7f6260321b..b4e09db253f 100644 --- a/litellm-rust/crates/token-counter-huggingface/src/error.rs +++ b/litellm-rust/crates/token-counter-huggingface/src/error.rs @@ -2,10 +2,21 @@ use thiserror::Error as ThisError; #[derive(Debug, ThisError)] pub enum Error { - #[error("failed to load tokenizer: {0}")] + #[error("failed to load tokenizer")] Load(#[source] tokenizers::Error), - #[error("tokenization failed: {0}")] + #[error("invalid encoding JSON")] + EncodingJson(#[source] serde_json::Error), + #[error("tokenization failed")] Encode(#[source] tokenizers::Error), - #[error("token decoding failed: {0}")] + #[error("token decoding failed")] Decode(#[source] tokenizers::Error), } + +impl Error { + pub fn is_load(&self) -> bool { + match self { + Self::Load(_) | Self::EncodingJson(_) => true, + Self::Encode(_) | Self::Decode(_) => false, + } + } +} diff --git a/litellm-rust/crates/token-counter-huggingface/src/lib.rs b/litellm-rust/crates/token-counter-huggingface/src/lib.rs index 0acea0532f8..c3af9b42767 100644 --- a/litellm-rust/crates/token-counter-huggingface/src/lib.rs +++ b/litellm-rust/crates/token-counter-huggingface/src/lib.rs @@ -13,11 +13,11 @@ pub use tokenizers::{ }; pub fn encoding_from_json(json: &str) -> Result { - serde_json::from_str(json).map_err(|error| Error::Load(error.into())) + serde_json::from_str(json).map_err(Error::EncodingJson) } pub fn encoding_to_json(encoding: &Encoding) -> Result { - serde_json::to_string(encoding).map_err(|error| Error::Load(error.into())) + serde_json::to_string(encoding).map_err(Error::EncodingJson) } pub struct HuggingFaceTokenizer { diff --git a/litellm-rust/crates/token-counter-tiktoken/Cargo.toml b/litellm-rust/crates/token-counter-tiktoken/Cargo.toml index d1789a693a3..bca7345dcd4 100644 --- a/litellm-rust/crates/token-counter-tiktoken/Cargo.toml +++ b/litellm-rust/crates/token-counter-tiktoken/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true [dependencies] base64.workspace = true +litellm-error.workspace = true once_cell = "1.21.3" rustc-hash = "2.1.3" thiserror.workspace = true diff --git a/litellm-rust/crates/token-counter-tiktoken/src/error.rs b/litellm-rust/crates/token-counter-tiktoken/src/error.rs index e28cbbfb620..024be00f048 100644 --- a/litellm-rust/crates/token-counter-tiktoken/src/error.rs +++ b/litellm-rust/crates/token-counter-tiktoken/src/error.rs @@ -1,5 +1,50 @@ +use std::{io, num::ParseIntError}; + +use litellm_error::Source; use thiserror::Error as ThisError; +use tiktoken_rs::{DecodeKeyError, EncodeError}; #[derive(Debug, ThisError)] -#[error("unsupported tokenizer: {0}")] -pub struct UnsupportedTokenizer(pub String); +pub enum Error { + #[error("unsupported tokenizer: {0}")] + UnsupportedTokenizer(String), + #[error("could not read tiktoken rank file {file}")] + Read { + file: &'static str, + source: io::Error, + }, + #[error("invalid tiktoken rank file")] + Ranks(#[from] RankError), + #[error("could not build the tiktoken encoder")] + Encoder(#[source] Source), + #[error("tokenization failed")] + Encode(#[source] EncodeError), + #[error("token decoding failed")] + Decode(#[source] DecodeKeyError), +} + +impl Error { + pub fn is_load(&self) -> bool { + match self { + Self::UnsupportedTokenizer(_) + | Self::Read { .. } + | Self::Ranks(_) + | Self::Encoder(_) => true, + Self::Encode(_) | Self::Decode(_) => false, + } + } +} + +#[derive(Debug, ThisError)] +pub enum RankError { + #[error("line without a rank: {0:?}")] + MissingRank(String), + #[error("token is not base64")] + InvalidToken(#[source] base64::DecodeError), + #[error("rank is not an integer")] + InvalidRank(#[source] ParseIntError), + #[error("a rank is assigned to more than one token")] + DuplicateRank, + #[error("byte 0x{0:02X} has no token")] + MissingByte(u8), +} diff --git a/litellm-rust/crates/token-counter-tiktoken/src/lib.rs b/litellm-rust/crates/token-counter-tiktoken/src/lib.rs index 30d0f820da3..c4cb786d9e4 100644 --- a/litellm-rust/crates/token-counter-tiktoken/src/lib.rs +++ b/litellm-rust/crates/token-counter-tiktoken/src/lib.rs @@ -5,8 +5,8 @@ mod ranks; use std::collections::HashSet; -pub use error::UnsupportedTokenizer; -pub use ranks::{LoadError, Vocabulary}; +pub use error::{Error, RankError}; +pub use ranks::Vocabulary; pub struct TiktokenTokenizer { encoder: &'static tiktoken_rs::CoreBPE, @@ -22,7 +22,7 @@ impl TiktokenTokenizer { pub fn from_cached_ranks( name: &str, load: impl FnOnce(&str) -> std::io::Result, - ) -> Result { + ) -> Result { let (loaded, name) = ranks::load(name, load)?; Ok(Self { encoder: &loaded.bpe, @@ -32,7 +32,7 @@ impl TiktokenTokenizer { } /// The encodings tiktoken-rs embeds, for hosts without the packaged rank files. - pub fn from_name(name: &str) -> Result { + pub fn from_name(name: &str) -> Result { let (encoder, name) = match name { "cl100k_base" => (tiktoken_rs::cl100k_base_singleton(), "cl100k_base"), "o200k_base" => (tiktoken_rs::o200k_base_singleton(), "o200k_base"), @@ -41,7 +41,7 @@ impl TiktokenTokenizer { "p50k_edit" => (tiktoken_rs::p50k_edit_singleton(), "p50k_edit"), "r50k_base" => (tiktoken_rs::r50k_base_singleton(), "r50k_base"), "gpt2" => (tiktoken_rs::r50k_base_singleton(), "gpt2"), - _ => return Err(UnsupportedTokenizer(name.to_owned())), + _ => return Err(Error::UnsupportedTokenizer(name.to_owned())), }; Ok(Self { encoder, @@ -62,12 +62,12 @@ impl TiktokenTokenizer { self.encoder.encode_ordinary(text) } - pub fn encode_special(&self, text: &str, allowed: &[String]) -> Result, String> { + pub fn encode_special(&self, text: &str, allowed: &[String]) -> Result, Error> { let allowed = allowed.iter().map(String::as_str).collect(); self.encoder .encode(text, &allowed) .map(|(ids, _)| ids) - .map_err(|error| error.to_string()) + .map_err(Error::Encode) } pub fn special_tokens(&self) -> HashSet { @@ -92,17 +92,13 @@ impl TiktokenTokenizer { (stable, completions) } - pub fn decode_bytes(&self, ids: &[u32]) -> Result, String> { - self.encoder - .decode_bytes(ids) - .map_err(|error| error.to_string()) + pub fn decode_bytes(&self, ids: &[u32]) -> Result, Error> { + self.encoder.decode_bytes(ids).map_err(Error::Decode) } - pub fn decode(&self, ids: &[u32]) -> Result { - self.encoder - .decode_bytes(ids) + pub fn decode(&self, ids: &[u32]) -> Result { + self.decode_bytes(ids) .map(|bytes| String::from_utf8_lossy(&bytes).into_owned()) - .map_err(|error| error.to_string()) } pub fn name(&self) -> &str { @@ -161,7 +157,8 @@ mod tests { #[test] fn unsupported_encoding_preserves_its_name() { - let Err(UnsupportedTokenizer(name)) = TiktokenTokenizer::from_name("unknown-encoding") + let Err(Error::UnsupportedTokenizer(name)) = + TiktokenTokenizer::from_name("unknown-encoding") else { panic!("unknown encoding must be rejected"); }; diff --git a/litellm-rust/crates/token-counter-tiktoken/src/ranks.rs b/litellm-rust/crates/token-counter-tiktoken/src/ranks.rs index 192e6f65a0b..36c30e711c1 100644 --- a/litellm-rust/crates/token-counter-tiktoken/src/ranks.rs +++ b/litellm-rust/crates/token-counter-tiktoken/src/ranks.rs @@ -1,10 +1,10 @@ use base64::{Engine, engine::general_purpose::STANDARD}; +use litellm_error::Source; use once_cell::sync::OnceCell; use rustc_hash::FxHashMap; -use thiserror::Error; use tiktoken_rs::{CoreBPE, O200K_BASE_PAT_STR, Rank}; -use crate::UnsupportedTokenizer; +use crate::error::{Error, RankError}; const CL100K: &str = "9b5ad71b2ce5302211f9c61530b329a4922fc6a4"; const O200K: &str = "fb374d419588a4632f3f557e76b4b70aebbca790"; @@ -76,22 +76,30 @@ impl Vocabulary { } } -#[derive(Debug, Error)] -pub enum LoadError { - #[error(transparent)] - Unsupported(#[from] UnsupportedTokenizer), - #[error("failed to load tiktoken ranks: {0}")] - Ranks(String), -} - /// Loads `name` once per process. The returned name is the one requested (`gpt2` stays /// `gpt2`, as `tiktoken.get_encoding("gpt2").name` does), while `gpt2` and `r50k_base` share /// one cached encoder. pub(super) fn load( name: &str, load_file: impl FnOnce(&str) -> std::io::Result, -) -> Result<(&'static Loaded, &'static str), LoadError> { - let (requested, canonical, file, cache) = match name { +) -> Result<(&'static Loaded, &'static str), Error> { + let (requested, canonical, file, cache) = encoding(name)?; + let loaded = cache.get_or_try_init(|| { + let ranks = load_file(file).map_err(|source| Error::Read { file, source })?; + build(canonical, &ranks) + })?; + Ok((loaded, requested)) +} + +type Encoding = ( + &'static str, + &'static str, + &'static str, + &'static OnceCell, +); + +fn encoding(name: &str) -> Result { + Ok(match name { "cl100k_base" => ("cl100k_base", "cl100k_base", CL100K, &CL100K_ENCODER), "o200k_base" => ("o200k_base", "o200k_base", O200K, &O200K_ENCODER), "o200k_harmony" => ("o200k_harmony", "o200k_harmony", O200K, &HARMONY_ENCODER), @@ -99,16 +107,11 @@ pub(super) fn load( "p50k_edit" => ("p50k_edit", "p50k_edit", P50K, &EDIT_ENCODER), "r50k_base" => ("r50k_base", "r50k_base", P50K, &R50K_ENCODER), "gpt2" => ("gpt2", "r50k_base", P50K, &R50K_ENCODER), - _ => return Err(UnsupportedTokenizer(name.to_owned()).into()), - }; - let loaded = cache.get_or_try_init(|| { - let ranks = load_file(file).map_err(|error| LoadError::Ranks(error.to_string()))?; - build(canonical, &ranks) - })?; - Ok((loaded, requested)) + _ => return Err(Error::UnsupportedTokenizer(name.to_owned())), + }) } -fn build(name: &str, ranks: &str) -> Result { +fn build(name: &str, ranks: &str) -> Result { let parsed = ranks .lines() .map(parse_rank) @@ -122,9 +125,11 @@ fn build(name: &str, ranks: &str) -> Result { .collect::>() .len() != encoder.len() - || (0..=u8::MAX).any(|byte| !encoder.contains_key(&[byte][..])) { - return Err(LoadError::Ranks("invalid vocabulary ranks".into())); + return Err(RankError::DuplicateRank.into()); + } + if let Some(byte) = (0..=u8::MAX).find(|byte| !encoder.contains_key(&[*byte][..])) { + return Err(RankError::MissingByte(byte).into()); } let (pattern, specials): (&str, &[(&str, Rank)]) = match name { "cl100k_base" => ( @@ -185,9 +190,12 @@ fn build(name: &str, ranks: &str) -> Result { .chain(special_tokens.values()) .copied() .max() - .ok_or_else(|| LoadError::Ranks("empty vocabulary".into()))?; - let bpe = CoreBPE::new(encoder.clone(), special_tokens.clone(), pattern) - .map_err(|error| LoadError::Ranks(error.to_string()))?; + .unwrap_or_default(); + let bpe = CoreBPE::new(encoder.clone(), special_tokens.clone(), pattern).map_err(|error| { + Error::Encoder(Source::from( + Box::::from(error), + )) + })?; Ok(Loaded { bpe, vocabulary: Vocabulary { @@ -198,16 +206,12 @@ fn build(name: &str, ranks: &str) -> Result { }) } -fn parse_rank(line: &str) -> Result<(Vec, Rank), LoadError> { +fn parse_rank(line: &str) -> Result<(Vec, Rank), RankError> { let (token, rank) = line .split_once(' ') - .ok_or_else(|| LoadError::Ranks("missing rank".into()))?; - let bytes = STANDARD - .decode(token) - .map_err(|error| LoadError::Ranks(error.to_string()))?; - let rank = rank - .parse() - .map_err(|error: std::num::ParseIntError| LoadError::Ranks(error.to_string()))?; + .ok_or_else(|| RankError::MissingRank(line.to_owned()))?; + let bytes = STANDARD.decode(token).map_err(RankError::InvalidToken)?; + let rank = rank.parse().map_err(RankError::InvalidRank)?; Ok((bytes, rank)) } @@ -226,79 +230,117 @@ mod tests { ) } - #[test] - fn packaged_encodings_match_embedded_encodings_and_reuse_successful_loads() { - for name in [ + #[rstest] + #[case::cl100k_base("cl100k_base", None)] + #[case::o200k_base("o200k_base", None)] + #[case::o200k_harmony("o200k_harmony", None)] + #[case::p50k_base("p50k_base", None)] + #[case::p50k_edit("p50k_edit", None)] + #[case::r50k_base_shared_with_gpt2("r50k_base", Some("gpt2"))] + fn packaged_encodings_load_once_and_match_embedded_encodings( + #[case] name: &str, + #[case] alias: Option<&str>, + ) { + assert!( + TiktokenTokenizer::from_cached_ranks(name, |_| { + Err(std::io::Error::other("unreadable vocabulary")) + }) + .is_err() + ); + let loads = std::sync::atomic::AtomicUsize::new(0); + let barrier = std::sync::Barrier::new(4); + let encoders = std::thread::scope(|scope| { + let tasks: Vec<_> = (0..4) + .map(|_| { + scope.spawn(|| { + barrier.wait(); + TiktokenTokenizer::from_cached_ranks(name, |file| { + loads.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + read_packaged_ranks(file) + }) + .unwrap() + }) + }) + .collect(); + tasks + .into_iter() + .map(|task| task.join().unwrap()) + .collect::>() + }); + assert_eq!(loads.into_inner(), 1); + let actual = &encoders[0]; + let expected = TiktokenTokenizer::from_name(name).unwrap(); + assert_eq!(actual.special_tokens(), expected.special_tokens()); + let specials: Vec<_> = expected.special_tokens().into_iter().collect(); + let special_text = specials.join(" "); + assert_eq!( + actual.encode_special(&special_text, &specials).unwrap(), + expected.encode_special(&special_text, &specials).unwrap() + ); + assert_eq!(actual.encode(&special_text), expected.encode(&special_text)); + let cached = + TiktokenTokenizer::from_cached_ranks(name, |_| panic!("reloaded cached vocabulary")) + .unwrap(); + assert_eq!(cached.encode("cached"), expected.encode("cached")); + assert_eq!(cached.name(), name); + assert!(expected.vocabulary().is_none()); + assert_vocabulary_lookups(name, actual); + if let Some(alias) = alias { + let shared = TiktokenTokenizer::from_cached_ranks(alias, |_| { + panic!("reloaded the shared vocabulary") + }) + .unwrap(); + assert_eq!(shared.name(), alias); + assert_vocabulary_lookups(alias, &shared); + } + } + + fn uncached(name: &str) -> TiktokenTokenizer { + static BUILT: std::sync::Mutex> = + std::sync::Mutex::new(Vec::new()); + let (requested, canonical, file, _) = encoding(name).unwrap(); + let mut built = BUILT.lock().unwrap(); + let loaded = match built.iter().find(|(built, _)| *built == requested) { + Some((_, loaded)) => *loaded, + None => { + let ranks = read_packaged_ranks(file).unwrap(); + let loaded: &'static Loaded = + Box::leak(Box::new(build(canonical, &ranks).unwrap())); + built.push((requested, loaded)); + loaded + } + }; + TiktokenTokenizer { + encoder: &loaded.bpe, + vocabulary: Some(&loaded.vocabulary), + name: requested, + } + } + + #[rstest] + fn packaged_ranks_encode_like_embedded_encodings( + #[values( "cl100k_base", "o200k_base", "o200k_harmony", "p50k_base", "p50k_edit", "r50k_base", - "gpt2", - ] { - if name != "gpt2" { - assert!( - TiktokenTokenizer::from_cached_ranks(name, |_| { - Err(std::io::Error::other("unreadable vocabulary")) - }) - .is_err() - ); - } - let loads = std::sync::atomic::AtomicUsize::new(0); - let barrier = std::sync::Barrier::new(4); - let encoders = std::thread::scope(|scope| { - let tasks: Vec<_> = (0..4) - .map(|_| { - scope.spawn(|| { - barrier.wait(); - TiktokenTokenizer::from_cached_ranks(name, |file| { - loads.fetch_add(1, std::sync::atomic::Ordering::Relaxed); - read_packaged_ranks(file) - }) - .unwrap() - }) - }) - .collect(); - tasks - .into_iter() - .map(|task| task.join().unwrap()) - .collect::>() - }); - assert_eq!(loads.into_inner(), usize::from(name != "gpt2")); - let actual = &encoders[0]; - let expected = TiktokenTokenizer::from_name(name).unwrap(); - assert_eq!(actual.special_tokens(), expected.special_tokens()); - let specials: Vec<_> = expected.special_tokens().into_iter().collect(); - let special_text = specials.join(" "); - assert_eq!( - actual.encode_special(&special_text, &specials).unwrap(), - expected.encode_special(&special_text, &specials).unwrap() - ); - for text in [ - "", - "café 漢字 ع 🙂", - "a\r\nb\t ", - " hello 123456789", - &special_text, - ] { - let ids = expected.encode(text); - assert_eq!(actual.encode(text), ids, "{name}: {text:?}"); - assert_eq!(actual.count_tokens(text), ids.len(), "{name}: {text:?}"); - assert_eq!( - actual.decode_bytes(&ids).unwrap(), - expected.decode_bytes(&ids).unwrap() - ); - } - let cached = TiktokenTokenizer::from_cached_ranks(name, |_| { - panic!("reloaded cached vocabulary") - }) - .unwrap(); - assert_eq!(cached.encode("cached"), expected.encode("cached")); - assert_eq!(cached.name(), name); - assert!(expected.vocabulary().is_none()); - assert_vocabulary_lookups(name, actual); - } + "gpt2" + )] + name: &str, + #[values("", "café 漢字 ع 🙂", "a\r\nb\t ", " hello 123456789")] text: &str, + ) { + let actual = uncached(name); + let expected = TiktokenTokenizer::from_name(name).unwrap(); + + let ids = expected.encode(text); + assert_eq!(actual.encode(text), ids); + assert_eq!(actual.count_tokens(text), ids.len()); + assert_eq!( + actual.decode_bytes(&ids).unwrap(), + expected.decode_bytes(&ids).unwrap() + ); } /// The token-level lookups tiktoken's Python `Encoding` exposes, checked against the diff --git a/litellm-rust/crates/token-counter/Cargo.toml b/litellm-rust/crates/token-counter/Cargo.toml index 67e5bd60537..01ce4b6bac8 100644 --- a/litellm-rust/crates/token-counter/Cargo.toml +++ b/litellm-rust/crates/token-counter/Cargo.toml @@ -14,6 +14,7 @@ tiktoken = ["dep:litellm-token-counter-tiktoken"] [dependencies] indexmap = { version = "2.14.0", features = ["serde"] } itoa = "1.0" +litellm-error.workspace = true litellm-token-counter-fast = { workspace = true, optional = true } litellm-token-counter-huggingface = { workspace = true, optional = true } litellm-token-counter-tiktoken = { workspace = true, optional = true } diff --git a/litellm-rust/crates/token-counter/src/counter.rs b/litellm-rust/crates/token-counter/src/counter.rs index 6370d9d74f8..f3c16f93b84 100644 --- a/litellm-rust/crates/token-counter/src/counter.rs +++ b/litellm-rust/crates/token-counter/src/counter.rs @@ -1,12 +1,12 @@ use serde::Serialize; -use crate::Error; use crate::python_json; use crate::tools::format_function_definitions; use crate::types::{ ContentItem, CountableContentBlock, CountableRequest, Message, MessageContent, TextValue, ToolChoice, ToolDefinition, }; +use crate::{Error, Unsupported}; const TOKENS_PER_MESSAGE: usize = 3; const TOKENS_PER_NAME: usize = 1; @@ -53,7 +53,7 @@ impl TokenCounter { self.count_optional_text_value(request.query.as_ref())? + self.count_optional_text_value(request.documents.as_ref())? } else { - return Err(Error::MissingInput); + return Err(Unsupported::MissingInput.into()); }; Ok(InputTokenCount { model: request.model.clone(), @@ -96,7 +96,7 @@ impl TokenCounter { TextValue::Number(number) => match (number.as_i64(), number.as_u64()) { (Some(number), _) => self.count_text(itoa::Buffer::new().format(number)), (_, Some(number)) => self.count_text(itoa::Buffer::new().format(number)), - _ => Err(Error::FloatText), + _ => Err(Unsupported::FloatText.into()), }, TextValue::Text(text) => self.count_text(text), TextValue::List(items) => items @@ -143,7 +143,9 @@ impl TokenCounter { None => Ok(0), } } - ContentItem::Block(CountableContentBlock::Unsupported) => Err(Error::ContentBlock), + ContentItem::Block(CountableContentBlock::Unsupported) => { + Err(Unsupported::ContentBlock.into()) + } } } diff --git a/litellm-rust/crates/token-counter/src/error.rs b/litellm-rust/crates/token-counter/src/error.rs index 7a74138c074..89a97423c48 100644 --- a/litellm-rust/crates/token-counter/src/error.rs +++ b/litellm-rust/crates/token-counter/src/error.rs @@ -1,39 +1,66 @@ use std::string::FromUtf8Error; +use litellm_error::Source; use thiserror::Error as ThisError; #[derive(Debug, ThisError)] pub enum Error { - #[error("unsupported tokenizer: {0}")] - UnsupportedTokenizer(String), - #[error("failed to load tokenizer: {0}")] - Load(#[source] Box), - #[error("failed to load tokenizer: tiktoken rank file: {0}")] - Ranks(#[source] Box), - #[error("failed to load tokenizer: Unicode character classes are unavailable")] - UnicodeClasses, - #[error("unsupported by the rust token counter: request body could not be parsed: {0}")] - RequestParse(#[source] serde_json::Error), - #[error("unsupported by the rust token counter: request has no countable input")] - MissingInput, - #[error( - "unsupported by the rust token counter: float text values are counted by the python path" - )] - FloatText, - #[error( - "unsupported by the rust token counter: content block type is counted by the python path" - )] - ContentBlock, - #[error("unsupported by the rust token counter: array parameter without items")] - ArrayItems, - #[error("unsupported by the rust token counter: text value could not be serialized: {0}")] - JsonSerialization(#[source] serde_json::Error), - #[error("unsupported by the rust token counter: serialized text value is not UTF-8: {0}")] - JsonUtf8(#[source] FromUtf8Error), - #[error("tokenization failed: {0}")] - Encode(#[source] Box), - #[error("token decoding failed: {0}")] - Decode(String), - #[error("token counting task failed: {0}")] - Task(String), + #[cfg(feature = "fast")] + #[error(transparent)] + Fast(#[from] litellm_token_counter_fast::Error), + #[cfg(feature = "huggingface")] + #[error(transparent)] + HuggingFace(#[from] litellm_token_counter_huggingface::Error), + #[cfg(feature = "tiktoken")] + #[error(transparent)] + Tiktoken(#[from] litellm_token_counter_tiktoken::Error), + #[error("unsupported by the rust token counter")] + Unsupported(#[from] Unsupported), + #[error("token counting task failed")] + Task(#[source] Source), +} + +impl Error { + pub fn is_unsupported(&self) -> bool { + match self { + Self::Unsupported(_) => true, + #[cfg(feature = "fast")] + Self::Fast(_) => false, + #[cfg(feature = "huggingface")] + Self::HuggingFace(_) => false, + #[cfg(feature = "tiktoken")] + Self::Tiktoken(_) => false, + Self::Task(_) => false, + } + } + + pub fn is_load(&self) -> bool { + match self { + #[cfg(feature = "fast")] + Self::Fast(error) => error.is_load(), + #[cfg(feature = "huggingface")] + Self::HuggingFace(error) => error.is_load(), + #[cfg(feature = "tiktoken")] + Self::Tiktoken(error) => error.is_load(), + Self::Unsupported(_) | Self::Task(_) => false, + } + } +} + +#[derive(Debug, ThisError)] +pub enum Unsupported { + #[error("request body could not be parsed")] + RequestParse(#[source] serde_json::Error), + #[error("request has no countable input")] + MissingInput, + #[error("float text values are counted by the python path")] + FloatText, + #[error("content block type is counted by the python path")] + ContentBlock, + #[error("array parameter without items")] + ArrayItems, + #[error("text value could not be serialized")] + JsonSerialization(#[source] serde_json::Error), + #[error("serialized text value is not UTF-8")] + JsonUtf8(#[source] FromUtf8Error), } diff --git a/litellm-rust/crates/token-counter/src/fast.rs b/litellm-rust/crates/token-counter/src/fast.rs index d1f270dcdc9..3ef0021e355 100644 --- a/litellm-rust/crates/token-counter/src/fast.rs +++ b/litellm-rust/crates/token-counter/src/fast.rs @@ -1,4 +1,3 @@ -use litellm_token_counter_fast::Error as BackendError; pub use litellm_token_counter_fast::FastTokenizer; use crate::{Error, TextCodec, TokenCounter, Tokenizer}; @@ -53,17 +52,6 @@ impl Tokenizer for FastTokenizer { } } -impl From for Error { - fn from(error: BackendError) -> Self { - match error { - BackendError::Load(source) => Self::Load(source), - BackendError::Ranks(source) => Self::Ranks(source.into()), - BackendError::UnicodeClasses => Self::UnicodeClasses, - BackendError::Encode(source) => Self::Encode(source), - } - } -} - #[cfg(all(test, feature = "huggingface", feature = "tiktoken"))] mod tests { use super::*; diff --git a/litellm-rust/crates/token-counter/src/huggingface.rs b/litellm-rust/crates/token-counter/src/huggingface.rs index 43613fac8ea..5730e9ce873 100644 --- a/litellm-rust/crates/token-counter/src/huggingface.rs +++ b/litellm-rust/crates/token-counter/src/huggingface.rs @@ -1,4 +1,3 @@ -use litellm_token_counter_huggingface::Error as BackendError; pub use litellm_token_counter_huggingface::{ AddedToken, EncodeInput, Encoding, HuggingFaceTokenizer, InputSequence, PaddingDirection, PaddingParams, PaddingStrategy, TruncationDirection, TruncationParams, encoding_from_json, @@ -34,13 +33,3 @@ impl TextCodec for HuggingFaceTokenizer { HuggingFaceTokenizer::name(self) } } - -impl From for Error { - fn from(error: BackendError) -> Self { - match error { - BackendError::Load(source) => Self::Load(source), - BackendError::Encode(source) => Self::Encode(source), - BackendError::Decode(source) => Self::Decode(source.to_string()), - } - } -} diff --git a/litellm-rust/crates/token-counter/src/lib.rs b/litellm-rust/crates/token-counter/src/lib.rs index 84bf35ca2ba..c58b8cf372f 100644 --- a/litellm-rust/crates/token-counter/src/lib.rs +++ b/litellm-rust/crates/token-counter/src/lib.rs @@ -19,6 +19,6 @@ pub mod huggingface; pub mod tiktoken; pub use counter::{InputTokenCount, TokenCounter}; -pub use error::Error; +pub use error::{Error, Unsupported}; pub use tokenizer::{TextCodec, Tokenizer}; pub use types::CountableRequest; diff --git a/litellm-rust/crates/token-counter/src/python_json.rs b/litellm-rust/crates/token-counter/src/python_json.rs index 00a55ad300c..b5c074a24c7 100644 --- a/litellm-rust/crates/token-counter/src/python_json.rs +++ b/litellm-rust/crates/token-counter/src/python_json.rs @@ -6,8 +6,8 @@ use std::io::{self, Write}; use serde::Serialize; use serde_json::ser::{Formatter, Serializer}; -use super::Error; use super::types::TextValue; +use super::{Error, Unsupported}; pub(super) fn dumps(value: &TextValue) -> Result { let mut output = Vec::with_capacity(serialized_len(value)?); @@ -16,9 +16,9 @@ pub(super) fn dumps(value: &TextValue) -> Result { &mut output, PythonFormatter, )) - .map_err(Error::JsonSerialization)?; + .map_err(Unsupported::JsonSerialization)?; debug_assert_eq!(output.len(), output.capacity()); - String::from_utf8(output).map_err(Error::JsonUtf8) + String::from_utf8(output).map_err(|error| Unsupported::JsonUtf8(error).into()) } fn serialized_len(value: &TextValue) -> Result { @@ -29,7 +29,7 @@ fn serialized_len(value: &TextValue) -> Result { TextValue::Number(number) => match (number.as_i64(), number.as_u64()) { (Some(number), _) => Ok(unsigned_len(number.unsigned_abs()) + usize::from(number < 0)), (_, Some(number)) => Ok(unsigned_len(number)), - _ => Err(Error::FloatText), + _ => Err(Unsupported::FloatText.into()), }, TextValue::Text(text) => Ok(quoted_len(text)), TextValue::List(items) => items @@ -147,9 +147,9 @@ mod tests { let value = serde_json::from_str(input).expect("fixture parses"); let error = dumps(&value).expect_err("floats are declined"); - assert_eq!( - error.to_string(), - "unsupported by the rust token counter: float text values are counted by the python path" - ); + assert!(matches!( + error, + crate::Error::Unsupported(crate::Unsupported::FloatText) + )); } } diff --git a/litellm-rust/crates/token-counter/src/tiktoken.rs b/litellm-rust/crates/token-counter/src/tiktoken.rs index b95621f3ccc..136280b6115 100644 --- a/litellm-rust/crates/token-counter/src/tiktoken.rs +++ b/litellm-rust/crates/token-counter/src/tiktoken.rs @@ -1,4 +1,3 @@ -use litellm_token_counter_tiktoken::{LoadError, UnsupportedTokenizer}; pub use litellm_token_counter_tiktoken::{TiktokenTokenizer, Vocabulary, encoding_for_model}; use crate::{Error, TextCodec, TokenCounter, Tokenizer}; @@ -23,25 +22,10 @@ impl TextCodec for TiktokenTokenizer { } fn decode(&self, ids: &[u32], _skip_special_tokens: bool) -> Result { - TiktokenTokenizer::decode(self, ids).map_err(|error| Error::Decode(error.to_string())) + TiktokenTokenizer::decode(self, ids).map_err(Error::from) } fn name(&self) -> &str { TiktokenTokenizer::name(self) } } - -impl From for Error { - fn from(error: UnsupportedTokenizer) -> Self { - Self::UnsupportedTokenizer(error.0) - } -} - -impl From for Error { - fn from(error: LoadError) -> Self { - match error { - LoadError::Unsupported(error) => error.into(), - LoadError::Ranks(source) => Self::Ranks(source.into()), - } - } -} diff --git a/litellm-rust/crates/token-counter/src/tools.rs b/litellm-rust/crates/token-counter/src/tools.rs index 22150a4e4b1..a5a3fdbfabf 100644 --- a/litellm-rust/crates/token-counter/src/tools.rs +++ b/litellm-rust/crates/token-counter/src/tools.rs @@ -3,8 +3,8 @@ use std::fmt::Write; -use super::Error; use super::types::{EnumValue, Schema, SchemaType, ToolDefinition}; +use super::{Error, Unsupported}; pub(super) fn format_function_definitions(tools: &[ToolDefinition]) -> Result { ToolFormatter::new(tools.len()).format(tools) @@ -117,7 +117,7 @@ impl ToolFormatter { None => self.output.push_str("number"), }, "array" => { - let items = props.items.as_deref().ok_or(Error::ArrayItems)?; + let items = props.items.as_deref().ok_or(Unsupported::ArrayItems)?; self.write_type(items, indent)?; self.output.push_str("[]"); } @@ -298,7 +298,7 @@ mod tests { assert!(matches!( format_function_definitions(&tools), - Err(Error::ArrayItems) + Err(Error::Unsupported(Unsupported::ArrayItems)) )); } } diff --git a/litellm-rust/crates/token-counter/src/types.rs b/litellm-rust/crates/token-counter/src/types.rs index c1236f94d9a..7e1513d7706 100644 --- a/litellm-rust/crates/token-counter/src/types.rs +++ b/litellm-rust/crates/token-counter/src/types.rs @@ -5,7 +5,7 @@ use serde::de::{MapAccess, SeqAccess, Visitor}; use serde::{Deserialize, Deserializer, Serialize}; use serde_json::Number; -use super::Error; +use super::{Error, Unsupported}; /// The parts of a request body the host's budget counter reads. Chat and /// Anthropic Messages bodies carry `messages`; completions carry `prompt`; @@ -33,7 +33,7 @@ pub struct CountableRequest { impl CountableRequest { pub fn parse(body: &[u8]) -> Result { - serde_json::from_slice(body).map_err(Error::RequestParse) + serde_json::from_slice(body).map_err(|error| Unsupported::RequestParse(error).into()) } } diff --git a/litellm-rust/crates/token-counter/tests/token_counter.rs b/litellm-rust/crates/token-counter/tests/token_counter.rs index 542bd4a1fc4..cc18bad5402 100644 --- a/litellm-rust/crates/token-counter/tests/token_counter.rs +++ b/litellm-rust/crates/token-counter/tests/token_counter.rs @@ -2,7 +2,7 @@ use rstest::rstest; #[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))] use litellm_token_counter::TokenCounter; -use litellm_token_counter::{CountableRequest, Error}; +use litellm_token_counter::{CountableRequest, Error, Unsupported}; #[cfg(any(feature = "fast", feature = "huggingface"))] mod json { @@ -100,7 +100,12 @@ mod json { let request = CountableRequest::parse(body).expect("shape parses"); assert!(matches!( counter(load).count_request(&request), - Err(Error::MissingInput | Error::FloatText | Error::ContentBlock | Error::ArrayItems) + Err(Error::Unsupported( + Unsupported::MissingInput + | Unsupported::FloatText + | Unsupported::ContentBlock + | Unsupported::ArrayItems + )) )); } @@ -139,7 +144,7 @@ mod json { } fn assert_loading_a_bad_tokenizer_is_a_load_error(load: JsonLoader) { - assert!(matches!(load("{}"), Err(Error::Load(_)))); + assert!(load("{}").is_err_and(|error| error.is_load())); } macro_rules! json_backend_tests { @@ -206,7 +211,7 @@ mod json { let counter = $loader(&tokenizer.to_string(false).unwrap()).unwrap(); let request = CountableRequest::parse(br#"{"prompt":"hello"}"#).unwrap(); let error = counter.count_request(&request).unwrap_err(); - assert!(matches!(error, Error::Encode(_))); + assert!(!error.is_load() && !error.is_unsupported()); assert_eq!(error.source().unwrap().to_string(), expected.to_string()); } @@ -250,7 +255,7 @@ mod json { fn shapes_outside_the_mirror_are_declined_at_parse(#[case] body: &[u8]) { assert!(matches!( CountableRequest::parse(body), - Err(Error::RequestParse(_)) + Err(Error::Unsupported(Unsupported::RequestParse(_))) )); } @@ -443,7 +448,7 @@ mod tiktoken { #[case] rank_file: &str, #[values(CL100K, O200K)] encoding: TiktokenEncoding, ) { - assert!(matches!((encoding.load)(rank_file), Err(Error::Ranks(_)))); + assert!((encoding.load)(rank_file).is_err_and(|error| error.is_load())); } } @@ -452,6 +457,7 @@ mod tiktoken { fn unsupported_encoding_reaches_the_counter_caller() { assert!(matches!( TokenCounter::from_tiktoken("unknown-encoding"), - Err(Error::UnsupportedTokenizer(name)) if name == "unknown-encoding" + Err(error) if error.is_load() + && litellm_error::Report::new(&error).to_string().contains("unknown-encoding") )); } From 323d3956d640772332f65c5dda63c0c351f41aef Mon Sep 17 00:00:00 2001 From: yujonglee Date: Wed, 28 Oct 2026 09:22:11 -0700 Subject: [PATCH 044/128] feat(rust): one typed decline from core up to each host (#47772) * feat(rust): one typed decline from core up to each host Rust core says it cannot serve a request with one Decline that travels unchanged through llms and inference up to the host. The Python bridge raises RustBridgeDeclined for it in route_failure before the public exception mapping, so dispatch falls back to legacy Python instead of failing the call, and the gateway answers 501. Every route collapses errors after the send with Error::after_send, which keeps a provider answer and a response fault as they are and nests anything else as the source of an invalid response, so a billed call can never read as a decline. Response-side cases that used to decline are now ErrorDetail::Unhandled response faults. The rust-error skill documents declined() as the first question both boundaries ask, and the Python suites pin that a native decline reaches the dispatcher before any provider IO. Co-Authored-By: Claude Opus 5.5 * fix(rust): keep post-provider hook failures terminal --------- Co-authored-by: Claude Opus 5.5 Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../.agents/skills/rust-error/SKILL.md | 1 + .../rust-error/references/boundaries.md | 10 +- .../rust-error/references/classification.md | 6 +- .../rust-error/references/conversions.md | 4 +- .../skills/rust-error/references/gateway.md | 2 +- .../skills/rust-error/references/pyo3.md | 6 +- litellm-rust/crates/core-utils/src/decline.rs | 16 +++ litellm-rust/crates/core-utils/src/lib.rs | 1 + .../crates/gateway-inference/src/error.rs | 11 +- litellm-rust/crates/host-native/src/driver.rs | 15 +-- litellm-rust/crates/host-python/src/driver.rs | 4 +- litellm-rust/crates/host/src/interceptors.rs | 2 +- .../crates/host/src/machine/context.rs | 4 +- litellm-rust/crates/host/src/protocol.rs | 8 +- .../crates/inference-chat/src/common_utils.rs | 2 +- .../crates/inference-chat/src/handler.rs | 57 +-------- .../crates/inference-chat/tests/request.rs | 2 +- .../crates/inference-messages/src/handler.rs | 1 + .../inference-messages/tests/messages/host.rs | 119 ++++++++++++------ .../crates/inference-ocr/tests/ocr/machine.rs | 19 ++- .../crates/inference-responses/src/handler.rs | 1 + .../crates/inference-responses/src/lib.rs | 10 +- .../inference-responses/tests/responses.rs | 5 +- .../inference-transcription/src/handler.rs | 7 +- litellm-rust/crates/inference/AGENTS.md | 2 +- litellm-rust/crates/inference/src/caching.rs | 63 +++++++++- litellm-rust/crates/inference/src/context.rs | 14 ++- litellm-rust/crates/inference/src/error.rs | 109 ++++++++++++---- .../crates/inference/tests/caching.rs | 16 +-- .../src/anthropic/batches/transformation.rs | 35 ++++-- .../llms/src/anthropic/chat/transformation.rs | 46 ++++--- .../llms/src/base_llm/chat/transformation.rs | 31 ++--- .../crates/llms/src/base_llm/ocr/error.rs | 49 -------- .../bedrock/chat/converse_transformation.rs | 45 ++++--- .../anthropic_claude_transformation.rs | 5 +- litellm-rust/crates/llms/src/error.rs | 38 ++++-- .../src/openai_like/chat/transformation.rs | 41 +++--- .../crates/python-bridge/src/errors.rs | 42 ++++++- .../src/routes/chat_completions/host.rs | 5 +- .../src/routes/responses/host.rs | 21 ++-- .../chat_completions/test_failures.py | 13 ++ tests/test_litellm_rust/responses/data.py | 2 +- .../responses/test_failures.py | 13 ++ 43 files changed, 558 insertions(+), 345 deletions(-) create mode 100644 litellm-rust/crates/core-utils/src/decline.rs diff --git a/litellm-rust/.agents/skills/rust-error/SKILL.md b/litellm-rust/.agents/skills/rust-error/SKILL.md index 2f02394c494..fd5d6222287 100644 --- a/litellm-rust/.agents/skills/rust-error/SKILL.md +++ b/litellm-rust/.agents/skills/rust-error/SKILL.md @@ -23,6 +23,7 @@ description: Define, wrap, convert, classify, report or test Rust error types in - Classification states facts. Retry and fallback policy lives with the caller that knows idempotency - [boundaries](references/boundaries.md) - An error becomes text for someone else only in one mapping function per error type, in the gateway crates, the Python bridge and host-python. The inference gateway and the bridge decide by the same predicates + - A failure before the provider was called is a decline. Both ask `declined()` first and hand it back to the host unchanged - [gateway](references/gateway.md) - One `Error` with `IntoResponse` per gateway crate that serves a router. Every failure flows through the same `From` chain, is classified once in `status()` and rendered by surface and framing - [pyo3](references/pyo3.md) diff --git a/litellm-rust/.agents/skills/rust-error/references/boundaries.md b/litellm-rust/.agents/skills/rust-error/references/boundaries.md index 253cf52a7db..876ebfdc184 100644 --- a/litellm-rust/.agents/skills/rust-error/references/boundaries.md +++ b/litellm-rust/.agents/skills/rust-error/references/boundaries.md @@ -12,11 +12,11 @@ A boundary is where an error becomes something another program reads: an HTTP re ## Asking the nested error -The inference gateway and the Python bridge's route mapping both receive a nested route error and decide by the same predicates, so a new failure lands correctly on both once its predicates answer right +The inference gateway and the Python bridge's route mapping both receive a nested route error and decide by the same predicates in the same order, so a new failure lands correctly on both once its predicates answer right | Predicate | Gateway | Bridge | |---|---|---| -| `is_unsupported()` | 501 | `ValueError`, since `is_request()` also answers true | +| `declined()` | 501 | `RustBridgeDeclined`, without the route's public mapping | | `upstream_response()` | the provider's status and body | `RustUpstreamError` | | `is_missing_credential()` | 401 | no class of its own, a `RuntimeError` because `is_request()` answers false | | `is_request()` | 400 | `ValueError` | @@ -26,6 +26,12 @@ If you find yourself adding an arm for your variant at a boundary, check whether The two boundaries are not identical. The gateway has variants of its own for body parsing, an unknown model and an unsupported path, and the OCR bridge has Python parity arms such as `FileNotFoundError` for a missing file. Those arms match the boundary's own variants or the top level of a public enum, never a nested error's variants. A missing API key is 401 at the gateway and `RuntimeError` at both bridges, pinned by tests to what Python's own path does +## Declines + +A failure before anything reached the provider is a decline, and `declined()` on the route error names its reason. It means another host may serve the same request as it was received, so each boundary hands it to the host unchanged: the gateway answers 501 and the Python bridge raises `RustBridgeDeclined`, never a public exception or a `ValueError`. See [gateway](gateway.md) and [pyo3](pyo3.md) for each host's side + +Anything after the send is terminal, never a decline. `declined()` answers `None` for a response fault and for a post-call hook failure even when the wrapped error was a decline, and every route that transforms the provider's answer collapses its errors with `Error::after_send`, which keeps a provider answer and a response fault as they are and nests anything else as the source of an invalid response, so a decline can never describe a call that was billed. Ask `declined()`, do not match the variant behind it, because that after-send cut lives in the predicate + ## What reaches the client A boundary renders anything that is not a provider answer with `litellm_error::Report`, so the reader sees every layer of the chain once. That is why messages are written for a reader, each layer leaves its cause to `source()`, and reqwest URLs are stripped on entry. A boundary that must not show the chain, such as the bridge's TLS settings error that names the setting but not the file, renders fixed text of its own and says so in its test diff --git a/litellm-rust/.agents/skills/rust-error/references/classification.md b/litellm-rust/.agents/skills/rust-error/references/classification.md index 59c3a1eb3f0..e32463847db 100644 --- a/litellm-rust/.agents/skills/rust-error/references/classification.md +++ b/litellm-rust/.agents/skills/rust-error/references/classification.md @@ -13,9 +13,9 @@ Downstream crates ask a lower error questions instead of matching its nested var | `status()` | an HTTP status to report | | `is_connect()` | the connection was never made, for a host's own connect-error mapping | | `tls_source()` | which TLS setting to blame in a config error | -| `is_computed_header()` | the caller forwarded a header the signer owns, which `litellm_inference::Error` turns into `Unsupported` | +| `is_computed_header()` | the caller forwarded a header the signer owns, which `litellm_inference::Error` turns into a decline | -Upper types add the questions their callers ask. `litellm_inference::Error` and the OCR error answer `is_missing_credential()` and `is_unsupported()`, which the gateway asks instead of reaching into `Auth` or `Unsupported` +Upper types add the questions their callers ask. `litellm_inference::Error` and the OCR error answer `is_missing_credential()`, and `litellm_llms::Error` and `litellm_inference::Error` answer `declined()`, which the boundaries ask first instead of reaching into `Auth` or `Declined`, see [boundaries](boundaries.md) ```rust // reaches through two crates' variants and breaks when either changes @@ -39,7 +39,7 @@ pub fn is_request(&self) -> bool { Add a predicate when a real caller needs a new distinction, and name it after the fact the caller needs, in the caller's words, not after the variant's spelling. An upper type's predicates call through to the nested error, as in `Self::Http(error) => error.is_request()` and `Self::Auth(error) => error.is_missing_credential()` -A public enum may be matched one level down by the type that owns the variant, such as the gateway matching its own `Unsupported` or the OCR bridge matching `FileRead`. Going a second level down, into a variant of a variant, is what to avoid, the gateway used to do that for `MissingApiKey` and `Unsupported` +A public enum may be matched one level down by the type that owns the variant, such as the gateway matching its own `Unsupported` or the OCR bridge matching `FileRead`. Going a second level down, into a variant of a variant, is what to avoid, the gateway used to do that for `MissingApiKey` and the old `Unsupported` route variant ## Facts, not policy diff --git a/litellm-rust/.agents/skills/rust-error/references/conversions.md b/litellm-rust/.agents/skills/rust-error/references/conversions.md index db46e4e3aa9..e40b76b123f 100644 --- a/litellm-rust/.agents/skills/rust-error/references/conversions.md +++ b/litellm-rust/.agents/skills/rust-error/references/conversions.md @@ -56,7 +56,7 @@ When a lower error needs a different meaning for every caller, do it in the uppe impl From for Error { fn from(error: litellm_http::Error) -> Self { if error.is_computed_header() { - return Self::Unsupported("request forwards a header AWS SigV4 computes"); + return Self::Declined(Decline::new("request forwards a header AWS SigV4 computes")); } Self::Http(error) } @@ -70,7 +70,7 @@ All four inference routes classify a header the signer computes the same way, be | The upper layer | Do | Example | |---|---|---| | has nothing to add | nest in one transparent variant and delegate predicates | `Self::Http(error) => error.is_request()` | -| changes what the failure means | translate to an upper variant, keep the cause attached when it helps | computed header -> `Unsupported` | +| changes what the failure means | translate to an upper variant, keep the cause attached when it helps | computed header -> `Declined` | | knows something the lower one does not | a variant with the extra fields plus the cause as its source | `Read { path, tls_source, source }` | Never re-declare a lower crate's variants in the upper type with an upper-layer name and a 1:1 mapping function. OCR once did that for every media download failure, `MediaError::BlockedUrl => Error::BlockedDocumentUrl` and so on, which doubled every change and let the two lists drift. diff --git a/litellm-rust/.agents/skills/rust-error/references/gateway.md b/litellm-rust/.agents/skills/rust-error/references/gateway.md index 1ec8ecd0acb..5da34aa84b2 100644 --- a/litellm-rust/.agents/skills/rust-error/references/gateway.md +++ b/litellm-rust/.agents/skills/rust-error/references/gateway.md @@ -21,7 +21,7 @@ Each hop is one `From` that nests the lower error in one variant, see [conversio | Answer | Status | |---|---| -| `is_unsupported()` | 501, the gateway matched and parsed the request but core cannot serve it, nothing reached the provider | +| `declined()` | 501, for a route error, the gateway matched and parsed the request but core cannot serve it, nothing reached the provider | | `upstream_response()` | the provider's own status, 502 when it is not a valid code | | `is_missing_credential()` | 401 | | `is_request()` | 400, for a route error | diff --git a/litellm-rust/.agents/skills/rust-error/references/pyo3.md b/litellm-rust/.agents/skills/rust-error/references/pyo3.md index 81a7a515726..1c6ffb96a82 100644 --- a/litellm-rust/.agents/skills/rust-error/references/pyo3.md +++ b/litellm-rust/.agents/skills/rust-error/references/pyo3.md @@ -30,7 +30,7 @@ The bridge builds with pyo3's `abi3-py310` feature. Under abi3, `#[pyclass(exten ## Mapping order -A Python exception raised inside the call, such as a secret manager callback's, is returned first as it is. Then the provider's answer passes through as `RustUpstreamError`. On `/v1/messages` a request the route rejected becomes a `ValueError` flagged with `messages_request_error`. Everything else is split by `is_request()` into `ValueError` or `RuntimeError`, with the message rendered by `litellm_error::Report`. A missing API key is a `RuntimeError` because `is_request()` answers false for it, and `missing_api_key_stays_a_runtime_error_while_other_auth_failures_are_value_errors` pins that +A core decline is returned first, as `RustBridgeDeclined` with the decline's reason as its only arg, without calling the route's public mapping. Then a Python exception raised inside the call, such as a secret manager callback's, is returned as it is. Then the provider's answer passes through as `RustUpstreamError`. On `/v1/messages` a request the route rejected becomes a `ValueError` flagged with `messages_request_error`. Everything else is split by `is_request()` into `ValueError` or `RuntimeError`, with the message rendered by `litellm_error::Report`. A missing API key is a `RuntimeError` because `is_request()` answers false for it, and `missing_api_key_stays_a_runtime_error_while_other_auth_failures_are_value_errors` pins that OCR puts Python parity arms in front of that order, each matching a top-level variant of its own public enum @@ -47,7 +47,9 @@ Use the standard class whose Python meaning matches. `PyFileNotFoundError` for a ## Declining -`RustBridgeDeclined` means no provider call was made and the host may retry on its own path. It is raised before dispatch, by the tokenizer, embeddings and token counter routes. Never raise it for a failure after the request went out, the host would repeat a billed call. `route_failures_are_terminal` in `errors.rs` pins that for the inference routes +`RustBridgeDeclined` means no provider call was made and the host may run its own path for the same call. The tokenizer, embeddings and token counter routes raise it from their own pre-checks, and `route_failure` raises it for a route error whose `declined()` answers, so a decline from deep inside a transformation reaches Python the same way. Never raise it for a failure after the request went out, the host would repeat a billed call. `a_core_decline_hands_the_call_back_to_python_before_the_public_mapping` and `route_failures_are_terminal` in `errors.rs` pin both sides + +Route hosts pass a terminal failure through the Python route module's `map_failure`, which runs `litellm.exception_type` and returns the public exception a Python caller expects. A decline must skip that step. Once wrapped it looks like a provider failure to the dispatcher and the host no longer falls back, so `route_failure` returns it before it calls `public` ## Rules diff --git a/litellm-rust/crates/core-utils/src/decline.rs b/litellm-rust/crates/core-utils/src/decline.rs new file mode 100644 index 00000000000..6f398b16a06 --- /dev/null +++ b/litellm-rust/crates/core-utils/src/decline.rs @@ -0,0 +1,16 @@ +/// Rust core cannot serve this request, and nothing was sent to the provider yet. +/// +/// Every error on the call path carries it unchanged, so the host can hand the same request +/// to another implementation: the SDK runs legacy Python, the gateway forwards to the proxy. +/// Raise it only before the provider call. Anything after the send is terminal. +#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] +#[error("declined by the rust path: {reason}")] +pub struct Decline { + pub reason: &'static str, +} + +impl Decline { + pub const fn new(reason: &'static str) -> Self { + Self { reason } + } +} diff --git a/litellm-rust/crates/core-utils/src/lib.rs b/litellm-rust/crates/core-utils/src/lib.rs index a937f55654e..cfbf451bf1b 100644 --- a/litellm-rust/crates/core-utils/src/lib.rs +++ b/litellm-rust/crates/core-utils/src/lib.rs @@ -1,5 +1,6 @@ pub mod call_arguments; pub mod core_helpers; +pub mod decline; pub mod dot_notation_indexing; pub mod exception_mapping_utils; pub mod get_llm_provider_logic; diff --git a/litellm-rust/crates/gateway-inference/src/error.rs b/litellm-rust/crates/gateway-inference/src/error.rs index 505483b426c..190561079f6 100644 --- a/litellm-rust/crates/gateway-inference/src/error.rs +++ b/litellm-rust/crates/gateway-inference/src/error.rs @@ -149,7 +149,7 @@ impl Error { } fn route_status(error: &InferenceError) -> StatusCode { - if error.is_unsupported() { + if error.declined().is_some() { return StatusCode::NOT_IMPLEMENTED; } if let Some((status, _)) = error.upstream_response() { @@ -165,9 +165,6 @@ fn route_status(error: &InferenceError) -> StatusCode { } fn ocr_status(error: &OcrError) -> StatusCode { - if error.is_unsupported() { - return StatusCode::NOT_IMPLEMENTED; - } if let Some((status, ..)) = error.upstream_response() { return StatusCode::from_u16(status).unwrap_or(StatusCode::BAD_GATEWAY); } @@ -326,6 +323,12 @@ mod tests { #[rstest] #[case::upstream_status(upstream(429, ""), StatusCode::TOO_MANY_REQUESTS)] + #[case::declined( + Error::Route(InferenceError::from(litellm_http::Error::computed_header( + "authorization" + ))), + StatusCode::NOT_IMPLEMENTED + )] #[case::rejected_request( Error::Route(InferenceError::invalid_request(litellm_llms::ErrorDetail::Expected { subject: "top_k", diff --git a/litellm-rust/crates/host-native/src/driver.rs b/litellm-rust/crates/host-native/src/driver.rs index 7254721da87..af3248848a2 100644 --- a/litellm-rust/crates/host-native/src/driver.rs +++ b/litellm-rust/crates/host-native/src/driver.rs @@ -66,11 +66,10 @@ where MachineStep::Suspended(request) => request, }; let answered = match request { - HostRequest::Intercept(InterceptRequest::ResultReady { facts, reply }) => self - .interceptors - .result_ready(facts) - .await - .map(|()| reply.send(())), + HostRequest::Intercept(InterceptRequest::ResultReady { facts, reply }) => { + reply.send(self.interceptors.result_ready(facts).await); + Ok(()) + } HostRequest::HostCall(call) => self.services.handle_host_call(call).await, HostRequest::Intercept(InterceptRequest::BeforeProviderRequest { wire, @@ -82,10 +81,8 @@ where .await .map(|wire| reply.send(wire)), HostRequest::Intercept(InterceptRequest::AfterProviderResponse { raw, reply }) => { - self.interceptors - .after_provider_response(raw) - .await - .map(|()| reply.send(())) + reply.send(self.interceptors.after_provider_response(raw).await); + Ok(()) } HostRequest::Stream(StreamDelivery::Open(head, reply)) => { self.demand = Some(reply); diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index 5dd9a72eb4a..dbf843a5295 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -299,7 +299,7 @@ where facts, })) .await?; - reply.send(()); + reply.send(Ok(())); Ok(()) } HostRequest::Intercept(InterceptRequest::AfterProviderResponse { raw, reply }) => { @@ -307,7 +307,7 @@ where ExecutionEvent::ProviderResponseReceived { raw: &raw }, )) .await?; - reply.send(()); + reply.send(Ok(())); Ok(()) } HostRequest::Stream(StreamDelivery::Open(head, reply)) => { diff --git a/litellm-rust/crates/host/src/interceptors.rs b/litellm-rust/crates/host/src/interceptors.rs index f0fd7434c6d..3cff781792f 100644 --- a/litellm-rust/crates/host/src/interceptors.rs +++ b/litellm-rust/crates/host/src/interceptors.rs @@ -212,7 +212,7 @@ mod tests { panic!("on_event yields Emit"); }; assert_eq!(raw.body, "raw"); - reply.send(()); + reply.send(Ok(())); let Ok(MachineStep::Complete((sent, ()))) = machine.resume().await else { panic!("the call completes with the answers"); diff --git a/litellm-rust/crates/host/src/machine/context.rs b/litellm-rust/crates/host/src/machine/context.rs index b2c9e95c563..1e53605ed7b 100644 --- a/litellm-rust/crates/host/src/machine/context.rs +++ b/litellm-rust/crates/host/src/machine/context.rs @@ -109,7 +109,7 @@ where .request_reply(|reply| { HostRequest::Intercept(InterceptRequest::ResultReady { facts, reply }) }) - .await + .await? } async fn after_provider_response(&self, raw: RawResponse) -> Result<(), P::Error> { @@ -117,7 +117,7 @@ where .request_reply(|reply| { HostRequest::Intercept(InterceptRequest::AfterProviderResponse { raw, reply }) }) - .await + .await? } } diff --git a/litellm-rust/crates/host/src/protocol.rs b/litellm-rust/crates/host/src/protocol.rs index 9e5aff8aae8..c24557b0348 100644 --- a/litellm-rust/crates/host/src/protocol.rs +++ b/litellm-rust/crates/host/src/protocol.rs @@ -15,14 +15,14 @@ pub trait Protocol: Send + Sync + 'static { pub enum HostRequest { HostCall(P::HostCall), - Intercept(InterceptRequest), + Intercept(InterceptRequest), Stream(StreamDelivery

), } -pub enum InterceptRequest { +pub enum InterceptRequest { ResultReady { facts: crate::interceptors::ExecutionFacts, - reply: Reply<()>, + reply: Reply>, }, BeforeProviderRequest { wire: Box, @@ -31,7 +31,7 @@ pub enum InterceptRequest { }, AfterProviderResponse { raw: RawResponse, - reply: Reply<()>, + reply: Reply>, }, } diff --git a/litellm-rust/crates/inference-chat/src/common_utils.rs b/litellm-rust/crates/inference-chat/src/common_utils.rs index b55aba22325..1c177d58552 100644 --- a/litellm-rust/crates/inference-chat/src/common_utils.rs +++ b/litellm-rust/crates/inference-chat/src/common_utils.rs @@ -63,7 +63,7 @@ pub(super) fn resolve_request( )); } if let Some(reason) = config.unsupported_reason(&messages, &call.optional_params) { - return Err(Error::Unsupported(reason.0)); + return Err(Error::Declined(reason)); } Ok(ResolvedChatCompletionsRequest { messages, diff --git a/litellm-rust/crates/inference-chat/src/handler.rs b/litellm-rust/crates/inference-chat/src/handler.rs index c34d60e1b8a..1c724a47956 100644 --- a/litellm-rust/crates/inference-chat/src/handler.rs +++ b/litellm-rust/crates/inference-chat/src/handler.rs @@ -128,31 +128,12 @@ pub(super) async fn execute( config .transform_response(provider.model, ProviderChatResponseData { body }) .map_err(Error::from) - .map_err(as_response_error) + .map_err(Error::after_send) }, ) .await } -/// Re-tag an error raised while normalizing a response the provider already -/// returned. -/// -/// A config reports the same variants on either side of the call: a missing -/// field or an unsupported block can mean "this request cannot be translated" -/// before the request is sent and "this response cannot be normalized" here. Only the -/// second kind has already been billed, and a host that keeps a reference -/// implementation must not retry those, so collapse them to one variant that -/// can only mean the provider was already called. -pub(super) fn as_response_error(err: Error) -> Error { - if err.is_response() || err.upstream_response().is_some() { - return err; - } - Error::invalid_response(litellm_llms::ErrorDetail::invalid( - "chat completions response", - err, - )) -} - #[cfg(test)] mod tests { use std::sync::Mutex; @@ -313,40 +294,4 @@ mod tests { ); assert!(interceptors.raw.into_inner().unwrap().is_empty()); } - - #[rstest::rstest] - #[case::missing_field(litellm_llms::Error::MissingField("usage").into())] - #[case::unsupported(Error::Unsupported("non-text response content block"))] - #[case::invalid_request(Error::invalid_request(litellm_llms::ErrorDetail::Expected { - subject: "whatever", - expected: "anything else", - }))] - #[case::auth(Error::Auth(litellm_auth::Error::InvalidHeader))] - fn response_errors_collapse_to_a_fault_that_can_only_mean_already_sent( - #[case] original: Error, - ) { - let collapsed = as_response_error(original); - assert!(collapsed.is_response(), "{collapsed:?}"); - assert!(!collapsed.is_request(), "{collapsed:?}"); - assert!(!collapsed.is_unsupported(), "{collapsed:?}"); - } - - #[rstest::rstest] - fn a_response_fault_is_not_wrapped_twice() { - let original = Error::invalid_response(litellm_llms::ErrorDetail::MissingField { - subject: "response", - field: "usage", - }); - let report = litellm_error::Report::new(&as_response_error(original)).to_string(); - assert!(!report.contains("chat completions response"), "{report}"); - } - - #[test] - fn a_provider_answer_keeps_its_status_and_body() { - let upstream = Error::from(litellm_http::Error::upstream(429, "slow down".into())); - assert_eq!( - as_response_error(upstream).upstream_response(), - Some((429, "slow down")) - ); - } } diff --git a/litellm-rust/crates/inference-chat/tests/request.rs b/litellm-rust/crates/inference-chat/tests/request.rs index 4025c137cdc..121f7347c24 100644 --- a/litellm-rust/crates/inference-chat/tests/request.rs +++ b/litellm-rust/crates/inference-chat/tests/request.rs @@ -354,7 +354,7 @@ async fn rejects_a_forwarded_header_the_signer_computes(#[case] forwarded: &str) let error = complete(call) .await .expect_err("conflicting signing headers fail"); - assert!(matches!(error, Error::Unsupported(_)), "{error:?}"); + assert!(error.declined().is_some(), "{error:?}"); assert!(received(&upstream).await.is_empty()); } diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index 67c3b8f98b5..fa802d20587 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -316,6 +316,7 @@ fn decode_response( config .transform_anthropic_messages_response(model, response) .map_err(Error::from) + .map_err(Error::after_send) } fn streaming_response( diff --git a/litellm-rust/crates/inference-messages/tests/messages/host.rs b/litellm-rust/crates/inference-messages/tests/messages/host.rs index 48afae99cb5..c3d06027c32 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/host.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/host.rs @@ -1,3 +1,4 @@ +use litellm_core_utils::decline::Decline; use litellm_host::lifecycle::ExecutionEvent; use std::sync::Mutex; @@ -23,6 +24,7 @@ struct RecordingHost { optional_params: Mutex>, facts: Mutex>, reject_result: bool, + reject_response: bool, } impl RecordingHost { @@ -34,6 +36,7 @@ impl RecordingHost { optional_params: Mutex::new(Vec::new()), facts: Mutex::new(Vec::new()), reject_result: false, + reject_response: false, } } @@ -81,7 +84,7 @@ impl litellm_host::interceptors::Interceptors< Result<(), Error> { self.facts.lock().unwrap().push(facts); if self.reject_result { - return Err(Error::Unsupported("result rejected")); + return Err(Error::Declined(Decline::new("result rejected"))); } Ok(()) } @@ -101,6 +104,9 @@ impl litellm_host::interceptors::Interceptors< Result<(), ::Error> { + if self.reject_response { + return Err(Error::Declined(Decline::new("response rejected"))); + } litellm_host::lifecycle::CallObserver::observe( self, litellm_host::lifecycle::CallEvent::Execution( @@ -121,6 +127,7 @@ async fn rejected_results_are_not_delivered_or_cached( call: MessagesCall, #[case] streaming: bool, #[case] hosted: bool, + #[values(false, true)] cached: bool, ) { use futures_util::TryStreamExt; use litellm_cache_memory::InMemoryCache; @@ -142,13 +149,8 @@ async fn rejected_results_are_not_delivered_or_cached( )))), CacheScope::Shared, )); - for (reject, expected_requests, cached) in [ - (true, 1, false), - (false, 2, false), - (true, 2, true), - (false, 2, true), - ] { - let request = authenticated( + let request = || { + authenticated( with_fields( MessagesCall { body: call.body.clone(), @@ -157,36 +159,41 @@ async fn rejected_results_are_not_delivered_or_cached( json!({"stream": streaming}), ), upstream.uri(), - ); - let host = RecordingHost { - reject_result: reject, - ..RecordingHost::passthrough(request) - }; - let result = if hosted { - litellm_host_native::in_process::run_hosted( - route.clone().machine(host.request().unwrap(), None), - host.runtime(), - ) + ) + }; + let host = RecordingHost { + reject_result: true, + ..RecordingHost::passthrough(request()) + }; + if cached { + let output = route.execute(request(), &(), None).await.unwrap(); + if let MessagesCallResponse::Stream { chunks, .. } = output { + chunks.try_collect::>().await.unwrap(); + } + } + let result = if hosted { + litellm_host_native::in_process::run_hosted( + route.clone().machine(host.request().unwrap(), None), + host.runtime(), + ) + .await + .map(|_| ()) + } else { + route + .execute(host.request().unwrap(), &host, None) .await .map(|_| ()) - } else { - match route.execute(host.request().unwrap(), &host, None).await { - Ok(MessagesCallResponse::Complete(_)) => Ok(()), - Ok(MessagesCallResponse::Stream { chunks, .. }) => { - chunks.try_collect::>().await.map(|_| ()) - } - Err(error) => Err(error), - } - }; - assert_eq!( - result, - if reject { - Err(Error::Unsupported("result rejected")) - } else { - Ok(()) - } - ); - assert_eq!(received(&upstream).await.len(), expected_requests); + }; + let error = result.unwrap_err(); + assert_eq!(error.declined().is_some(), cached); + assert_eq!(error.dispatched(), !cached); + assert!( + litellm_error::Report::new(&error) + .to_string() + .contains("result rejected") + ); + assert_eq!(received(&upstream).await.len(), 1); + { let facts = host.facts.lock().unwrap(); assert_eq!(facts.len(), 1); assert_eq!( @@ -194,6 +201,46 @@ async fn rejected_results_are_not_delivered_or_cached( cached ); } + let output = route.execute(request(), &(), None).await.unwrap(); + if let MessagesCallResponse::Stream { chunks, .. } = output { + chunks.try_collect::>().await.unwrap(); + } + assert_eq!(received(&upstream).await.len(), if cached { 1 } else { 2 }); +} + +#[rstest] +#[case::native(false)] +#[case::hosted(true)] +#[tokio::test] +async fn response_hook_declines_are_terminal(call: MessagesCall, #[case] hosted: bool) { + let upstream = upstream([message_response()]).await; + let host = RecordingHost { + reject_response: true, + ..RecordingHost::passthrough(authenticated(call, upstream.uri())) + }; + let route = messages_route(no_secrets()); + let result = if hosted { + litellm_host_native::in_process::run_hosted( + route.machine(host.request().unwrap(), None), + host.runtime(), + ) + .await + .map(|_| ()) + } else { + route + .execute(host.request().unwrap(), &host, None) + .await + .map(|_| ()) + }; + let error = result.unwrap_err(); + assert!(error.declined().is_none()); + assert!(error.dispatched()); + assert!( + litellm_error::Report::new(&error) + .to_string() + .contains("response rejected") + ); + assert_eq!(received(&upstream).await.len(), 1); } async fn run_through(host: &RecordingHost) -> Result { diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/machine.rs b/litellm-rust/crates/inference-ocr/tests/ocr/machine.rs index 9a1c67fab0f..c7a5d58ad3d 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/machine.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/machine.rs @@ -41,11 +41,10 @@ async fn drive_until( Err(error) => break Err(error), }; let answer = match op { - HostRequest::Intercept(InterceptRequest::ResultReady { facts, reply }) => host - .result_ready(facts) - .await - .map(|()| reply.send(())) - .map_err(HostFailure::Error), + HostRequest::Intercept(InterceptRequest::ResultReady { facts, reply }) => { + reply.send(host.result_ready(facts).await); + Ok(()) + } HostRequest::Stream(stream) => match stream { litellm_host::protocol::StreamDelivery::Open(head, _) => match head {}, litellm_host::protocol::StreamDelivery::Chunk(chunk, _) => match chunk {}, @@ -64,10 +63,8 @@ async fn drive_until( } HostRequest::Intercept(InterceptRequest::AfterProviderResponse { raw, reply }) => { ops.push("response"); - host.after_provider_response(raw) - .await - .map(|()| reply.send(())) - .map_err(HostFailure::Error) + reply.send(host.after_provider_response(raw).await); + Ok(()) } }; if let Err(failure) = answer { @@ -85,10 +82,10 @@ async fn drive_until_notified(machine: &mut OcrMachine, host: &LocalOcrHost, sto _ = stop.notified() => break, step = machine.resume() => { match step.unwrap() { - MachineStep::Suspended(HostRequest::Intercept(InterceptRequest::ResultReady { reply, .. })) => reply.send(()), + MachineStep::Suspended(HostRequest::Intercept(InterceptRequest::ResultReady { reply, .. })) => reply.send(Ok(())), MachineStep::Suspended(HostRequest::HostCall(op)) => host.handle_host_call(op).await.unwrap(), MachineStep::Suspended(HostRequest::Intercept(InterceptRequest::BeforeProviderRequest { wire, reply, .. })) => reply.send(*wire), - MachineStep::Suspended(HostRequest::Intercept(InterceptRequest::AfterProviderResponse { reply, .. })) => reply.send(()), + MachineStep::Suspended(HostRequest::Intercept(InterceptRequest::AfterProviderResponse { reply, .. })) => reply.send(Ok(())), MachineStep::Suspended(HostRequest::Stream(stream)) => match stream { litellm_host::protocol::StreamDelivery::Open(head, _) => match head {}, litellm_host::protocol::StreamDelivery::Chunk(chunk, _) => match chunk {}, diff --git a/litellm-rust/crates/inference-responses/src/handler.rs b/litellm-rust/crates/inference-responses/src/handler.rs index 8d3d1075e8e..19092a63a14 100644 --- a/litellm-rust/crates/inference-responses/src/handler.rs +++ b/litellm-rust/crates/inference-responses/src/handler.rs @@ -134,6 +134,7 @@ pub(super) async fn execute( .transform_response_api_response(value) .map(ResponsesOutput::Complete) .map_err(Error::from) + .map_err(Error::after_send) }, ) .await diff --git a/litellm-rust/crates/inference-responses/src/lib.rs b/litellm-rust/crates/inference-responses/src/lib.rs index b6e151ecec3..9af5dea1fe8 100644 --- a/litellm-rust/crates/inference-responses/src/lib.rs +++ b/litellm-rust/crates/inference-responses/src/lib.rs @@ -11,6 +11,7 @@ mod handler; use std::sync::Arc; use litellm_auth::AuthServices; +use litellm_core_utils::decline::Decline; use litellm_core_utils::get_llm_provider_logic::{ CustomLlmProvider, LlmProviders, get_custom_llm_provider, }; @@ -24,6 +25,8 @@ use litellm_llms::{ use litellm_secrets::source::SecretSource; use types::{ResponsesCall, ResponsesOutput, ResponsesProvider}; +const NOT_NATIVE: Decline = Decline::new("native HTTP responses provider"); + fn responses_provider(provider: LlmProviders) -> Option<&'static dyn BaseResponsesApiConfig> { match provider { LlmProviders::Openai => Some(&OpenAiResponsesApiConfig), @@ -111,9 +114,8 @@ impl ResponsesRoute { } let provider = provider_name .parse::() - .map_err(|_| Error::Unsupported("native HTTP responses provider"))?; - let config = responses_provider(provider) - .ok_or(Error::Unsupported("native HTTP responses provider"))?; + .map_err(|_| Error::Declined(NOT_NATIVE))?; + let config = responses_provider(provider).ok_or(Error::Declined(NOT_NATIVE))?; let wire_model = match provider { LlmProviders::Openai if resolved_model.contains('/') => { return Err(Error::InvalidProvider(call.model)); @@ -123,7 +125,7 @@ impl ResponsesRoute { format!("openrouter/{resolved_model}") } LlmProviders::Openrouter => resolved_model.into(), - _ => return Err(Error::Unsupported("native HTTP responses provider")), + _ => return Err(Error::Declined(NOT_NATIVE)), }; let provider = ResponsesProvider { identity: ProviderIdentity { diff --git a/litellm-rust/crates/inference-responses/tests/responses.rs b/litellm-rust/crates/inference-responses/tests/responses.rs index 66402e595eb..89368a941ef 100644 --- a/litellm-rust/crates/inference-responses/tests/responses.rs +++ b/litellm-rust/crates/inference-responses/tests/responses.rs @@ -103,10 +103,7 @@ async fn unsupported_providers_fail_before_transport( ) .await; - assert!(matches!( - result, - Err(Error::Unsupported("native HTTP responses provider")) - )); + assert!(matches!(result, Err(ref error) if error.declined().is_some())); assert!(received(&upstream).await.is_empty()); } diff --git a/litellm-rust/crates/inference-transcription/src/handler.rs b/litellm-rust/crates/inference-transcription/src/handler.rs index ba7eb740e2f..2ed3b1eba7b 100644 --- a/litellm-rust/crates/inference-transcription/src/handler.rs +++ b/litellm-rust/crates/inference-transcription/src/handler.rs @@ -67,7 +67,8 @@ pub(super) async fn execute( error, )) })?; - Ok(config - .transform_audio_transcription_response(provider.model, response_json)? - .into_json()) + config + .transform_audio_transcription_response(provider.model, response_json) + .map(|response| response.into_json()) + .map_err(|error| Error::from(error).after_send()) } diff --git a/litellm-rust/crates/inference/AGENTS.md b/litellm-rust/crates/inference/AGENTS.md index 8cc8fc8c26b..bcbc98a209e 100644 --- a/litellm-rust/crates/inference/AGENTS.md +++ b/litellm-rust/crates/inference/AGENTS.md @@ -56,7 +56,7 @@ - The [rust-error skill](../../.agents/skills/rust-error/SKILL.md) shapes each crate's error. This section decides which crate and module a failure belongs to - A failure is declared once, by the lowest crate that raises it - every crate above nests it in one variant and delegates its predicates, as `Error::Llm`, `Error::Auth` and `Error::Http` do, or translates it once in the shared `From` when its meaning changes, as `From` does for `Unsupported` and `From` does for a header the signer computes - - `Error` collects route failures for every format crate and never re-declares a variant a lower crate raises. Routes build transformation failures with `Error::invalid_request` and `Error::invalid_response`, and callers ask `is_unsupported()`, `is_request()`, `is_response()`, `is_missing_credential()`, `dispatched()`, `upstream_response()` and `rejected_request()` instead of matching nested variants + - `Error` collects route failures for every format crate and never re-declares a variant a lower crate raises. Routes build transformation failures with `Error::invalid_request` and `Error::invalid_response` and collapse anything raised after the send with `Error::after_send`, and callers ask `declined()`, `is_request()`, `is_response()`, `is_missing_credential()`, `dispatched()`, `upstream_response()` and `rejected_request()` instead of matching nested variants - Scope follows the concept, not the first caller - an error type under `litellm-llms`'s `/` is private to that provider: no other provider and nothing in `base_llm` may import it - a failure two providers or two routes can hit (wire framing, stream event decoding, a malformed provider response) belongs to the crate that owns the concept: `litellm-framer` for framing, `litellm_llms::Error` for the transformation layer diff --git a/litellm-rust/crates/inference/src/caching.rs b/litellm-rust/crates/inference/src/caching.rs index 09d03c06215..d9b5b891015 100644 --- a/litellm-rust/crates/inference/src/caching.rs +++ b/litellm-rust/crates/inference/src/caching.rs @@ -377,5 +377,66 @@ async fn publish( facts: facts.clone(), })); } - interceptors.result_ready(facts).await + let from_provider = facts.source == ResultSource::Provider; + interceptors.result_ready(facts).await.map_err(|error| { + if from_provider { + Error::post_call(error) + } else { + error + } + }) +} + +#[cfg(test)] +mod tests { + use litellm_core_utils::decline::Decline; + use litellm_host::interceptors::{RawResponse, RequestContext, WireRequest}; + + use super::*; + + struct RejectResult; + + impl Interceptors for RejectResult { + async fn before_provider_request( + &self, + wire: WireRequest, + _: RequestContext, + ) -> Result { + Ok(wire) + } + + async fn after_provider_response(&self, _: RawResponse) -> Result<(), Error> { + Ok(()) + } + + async fn result_ready(&self, _: ExecutionFacts) -> Result<(), Error> { + Err(Error::Declined(Decline::new("result rejected"))) + } + } + + #[rstest::rstest] + #[case::provider(ResultSource::Provider, false)] + #[case::cache(ResultSource::Cache { key: "cached".into() }, true)] + #[tokio::test] + async fn result_rejection_only_declines_without_a_provider_call( + #[case] source: ResultSource, + #[case] may_decline: bool, + ) { + let facts = ExecutionFacts { + provider: ProviderIdentity { + model: "model".into(), + provider: "provider".into(), + }, + cost: Cost::settle(&source, None), + source, + }; + let error = publish(facts, &RejectResult, None).await.unwrap_err(); + assert_eq!(error.declined().is_some(), may_decline); + assert_eq!(error.dispatched(), !may_decline); + assert!( + litellm_error::Report::new(&error) + .to_string() + .contains("result rejected") + ); + } } diff --git a/litellm-rust/crates/inference/src/context.rs b/litellm-rust/crates/inference/src/context.rs index 655b768dde1..963aa7b90c5 100644 --- a/litellm-rust/crates/inference/src/context.rs +++ b/litellm-rust/crates/inference/src/context.rs @@ -1,6 +1,6 @@ use litellm_cache_response::CachePolicy; use litellm_host::{ - interceptors::{ExecutionFacts, Interceptors, RawResponse}, + interceptors::{ExecutionFacts, Interceptors, RawResponse, ResultSource}, lifecycle::{CallEvent, ExecutionEvent}, observation::ObservationSender, }; @@ -28,7 +28,17 @@ impl<'a, I: Interceptors> CallContext<'a, I> { facts: facts.clone(), })); } - self.interceptors.result_ready(facts).await + let from_provider = facts.source == ResultSource::Provider; + self.interceptors + .result_ready(facts) + .await + .map_err(|error| { + if from_provider { + Error::post_call(error) + } else { + error + } + }) } pub async fn response_received(&self, body: &str) -> Result<(), Error> { diff --git a/litellm-rust/crates/inference/src/error.rs b/litellm-rust/crates/inference/src/error.rs index 2b42d2fced1..a8e6fb4e1b3 100644 --- a/litellm-rust/crates/inference/src/error.rs +++ b/litellm-rust/crates/inference/src/error.rs @@ -1,5 +1,6 @@ use std::sync::Arc; +use litellm_core_utils::decline::Decline; use litellm_host::machine::MachineFault; use litellm_llms::{Error as LlmError, ErrorDetail}; @@ -9,8 +10,8 @@ pub enum Error { InvalidProvider(String), #[error("invalid provider: unable to resolve custom_llm_provider for {route} request")] UnresolvedProvider { route: &'static str }, - #[error("unsupported by the rust path: {0}")] - Unsupported(&'static str), + #[error(transparent)] + Declined(Decline), #[error(transparent)] Llm(LlmError), #[error(transparent)] @@ -50,14 +51,11 @@ impl Error { /// the provider's answer. pub fn is_request(&self) -> bool { match self { - Self::InvalidProvider(_) - | Self::UnresolvedProvider { .. } - | Self::Unsupported(_) - | Self::Headers(_) => true, + Self::InvalidProvider(_) | Self::UnresolvedProvider { .. } | Self::Headers(_) => true, Self::Llm(error) => error.is_request(), Self::Auth(error) => error.is_request(), Self::Http(error) => error.is_request(), - Self::Secret(_) | Self::Host(_) | Self::PostCallHook(_) => false, + Self::Declined(_) | Self::Secret(_) | Self::Host(_) | Self::PostCallHook(_) => false, } } @@ -67,7 +65,7 @@ impl Error { Self::Http(error) => error.is_response(), Self::InvalidProvider(_) | Self::UnresolvedProvider { .. } - | Self::Unsupported(_) + | Self::Declined(_) | Self::Auth(_) | Self::Headers(_) | Self::Secret(_) @@ -82,7 +80,7 @@ impl Error { Self::Auth(error) => error.is_missing_credential(), Self::InvalidProvider(_) | Self::UnresolvedProvider { .. } - | Self::Unsupported(_) + | Self::Declined(_) | Self::Headers(_) | Self::Http(_) | Self::Secret(_) @@ -91,21 +89,34 @@ impl Error { } } - pub fn is_unsupported(&self) -> bool { + pub fn declined(&self) -> Option<&Decline> { match self { - Self::Unsupported(_) => true, + Self::Declined(decline) => Some(decline), + Self::Llm(error) => error.declined(), Self::InvalidProvider(_) | Self::UnresolvedProvider { .. } - | Self::Llm(_) | Self::Auth(_) | Self::Headers(_) | Self::Http(_) | Self::Secret(_) | Self::Host(_) - | Self::PostCallHook(_) => false, + | Self::PostCallHook(_) => None, } } + /// Re-tag an error raised while handling an answer the provider already returned. + /// + /// A config reports the same variants on either side of the call, so a missing field or + /// a decline can mean "this request cannot be translated" before the send and "this + /// response cannot be normalized" after it. Only the second kind was billed, and a host + /// that keeps a reference implementation must not retry it. + pub fn after_send(self) -> Self { + if self.is_response() || self.upstream_response().is_some() { + return self; + } + Self::invalid_response(ErrorDetail::invalid("provider response", self)) + } + pub fn dispatched(&self) -> bool { match self { Self::Llm(error) => error.dispatched(), @@ -113,7 +124,7 @@ impl Error { Self::Host(_) | Self::PostCallHook(_) => true, Self::InvalidProvider(_) | Self::UnresolvedProvider { .. } - | Self::Unsupported(_) + | Self::Declined(_) | Self::Auth(_) | Self::Headers(_) | Self::Secret(_) => false, @@ -125,7 +136,7 @@ impl Error { Self::Llm(error) => error.rejected_request(), Self::InvalidProvider(_) | Self::UnresolvedProvider { .. } - | Self::Unsupported(_) + | Self::Declined(_) | Self::Auth(_) | Self::Headers(_) | Self::Http(_) @@ -141,7 +152,7 @@ impl Error { Self::Http(error) => error.upstream_response(), Self::InvalidProvider(_) | Self::UnresolvedProvider { .. } - | Self::Unsupported(_) + | Self::Declined(_) | Self::Auth(_) | Self::Headers(_) | Self::Secret(_) @@ -156,7 +167,7 @@ impl From for Error { // Python drops the caller's copy and prefers a forwarded Authorization // over the signature, so leave the request to it. if error.is_computed_header() { - return Self::Unsupported("request forwards a header AWS SigV4 computes"); + return Self::Declined(Decline::new("request forwards a header AWS SigV4 computes")); } Self::Http(error) } @@ -165,7 +176,7 @@ impl From for Error { impl From for Error { fn from(error: LlmError) -> Self { match error { - LlmError::Unsupported(reason) => Self::Unsupported(reason), + LlmError::Declined(decline) => Self::Declined(decline), LlmError::Http(error) => Self::from(error), LlmError::Auth(error) => Self::Auth(error), error @ (LlmError::InvalidType { .. } @@ -203,6 +214,7 @@ impl Eq for SecretError {} #[cfg(test)] mod tests { use super::Error; + use litellm_core_utils::decline::Decline; use litellm_host::machine::MachineFault; use litellm_llms::{Error as LlmError, ErrorDetail}; use rstest::rstest; @@ -273,13 +285,66 @@ mod tests { LlmError::Http(litellm_http::Error::computed_header("authorization")).into() )] fn a_header_the_signer_computes_is_left_to_the_host_on_every_route(#[case] error: Error) { - assert!(error.is_unsupported(), "{error:?}"); - assert!(error.is_request()); + assert!(error.declined().is_some(), "{error:?}"); + assert!(!error.dispatched()); } #[test] - fn a_transformation_unsupported_stays_a_decline() { - assert!(Error::from(LlmError::Unsupported("test capability")).is_unsupported()); + fn a_transformation_decline_survives_into_the_route_error() { + let decline = Decline::new("unrecognized request parameter"); + let error = Error::from(LlmError::Declined(decline)); + assert_eq!(error.declined(), Some(&decline)); + assert!(!error.is_request()); + assert!(!error.dispatched()); + } + + #[test] + fn a_post_call_hook_failure_is_never_a_decline() { + let hook = Error::post_call(Error::Declined(Decline::new("hook"))); + assert_eq!(hook.declined(), None); + } + + #[rstest] + #[case::declined(Error::Declined(Decline::new("non-text response content block")))] + #[case::transformation_decline(LlmError::Declined(Decline::new("streaming")).into())] + #[case::missing_field(LlmError::MissingField("usage").into())] + #[case::invalid_request(Error::invalid_request(ErrorDetail::Expected { + subject: "whatever", + expected: "anything else", + }))] + #[case::auth(Error::Auth(litellm_auth::Error::InvalidHeader))] + #[case::host_driver(MachineFault::Abandoned.into())] + fn errors_after_the_send_are_response_faults_and_never_a_decline(#[case] error: Error) { + let report = litellm_error::Report::new(&error).to_string(); + let collapsed = error.after_send(); + assert!(collapsed.is_response(), "{collapsed:?}"); + assert!(!collapsed.is_request(), "{collapsed:?}"); + assert!(collapsed.dispatched(), "{collapsed:?}"); + assert_eq!(collapsed.declined(), None); + let collapsed_report = litellm_error::Report::new(&collapsed).to_string(); + assert!(collapsed_report.ends_with(&report), "{collapsed_report}"); + } + + #[test] + fn a_response_fault_is_not_wrapped_twice_after_the_send() { + let original = Error::invalid_response(ErrorDetail::MissingField { + subject: "response", + field: "usage", + }); + let report = litellm_error::Report::new(&original).to_string(); + assert_eq!( + litellm_error::Report::new(&original.after_send()).to_string(), + report + ); + } + + #[test] + fn a_provider_answer_stays_as_is_after_the_send() { + let upstream = Error::from(litellm_http::Error::upstream(429, "slow down".into())); + assert_eq!( + upstream.after_send().upstream_response(), + Some((429, "slow down")) + ); } #[test] diff --git a/litellm-rust/crates/inference/tests/caching.rs b/litellm-rust/crates/inference/tests/caching.rs index 02c2edd95dd..4e8b2349ac8 100644 --- a/litellm-rust/crates/inference/tests/caching.rs +++ b/litellm-rust/crates/inference/tests/caching.rs @@ -14,6 +14,7 @@ use litellm_cache_memory::InMemoryCache; use litellm_cache_response::{ CacheOptions, CachePolicy, CacheScope, ResponseCache, ResponseCacheConfig, ResponseCacheService, }; +use litellm_core_utils::decline::Decline; use litellm_host::{ call::{CallOutput, OutputOf}, interceptors::{ @@ -204,7 +205,8 @@ async fn streamed( .chunks(3) .map(|bytes| Ok(Bytes::copy_from_slice(bytes))) .collect::>(); - let ending = fail.then_some(Err(Error::Unsupported("test transport failure"))); + let ending = + fail.then_some(Err(Error::Declined(Decline::new("test transport failure")))); Ok(CallOutput::Stream { head: (), chunks: stream::iter(chunks.into_iter().chain(ending)).boxed(), @@ -315,7 +317,7 @@ async fn a_provider_failure_never_populates_the_cache(cache: Arc for CacheHitAccounting { ); self.key.set(key).unwrap(); if self.reject { - return Err(Error::Unsupported("cache accounting rejected")); + return Err(Error::Declined(Decline::new("cache accounting rejected"))); } Ok(()) } @@ -569,10 +571,10 @@ async fn cache_hits_notify_accounting_once_and_propagate_its_failure( .await }; if reject { - assert!(matches!( - result, - Err(Error::Unsupported("cache accounting rejected")) - )); + assert_eq!( + result.unwrap_err().declined(), + Some(&Decline::new("cache accounting rejected")) + ); } else { assert_eq!(result.unwrap(), expected); } diff --git a/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs b/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs index 66658802dd5..0513f88038d 100644 --- a/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs @@ -1,3 +1,4 @@ +use litellm_core_utils::decline::Decline; use litellm_llms_types::formats::batches::{BatchRequestCounts, BatchResponse, BatchStatus}; use litellm_llms_types::formats::messages::MessagesResponse; use litellm_llms_types::providers::anthropic::{BATCHES_PATH, MESSAGES_PATH}; @@ -127,7 +128,9 @@ impl AnthropicBatchesConfig for AnthropicBatchesTransformation { } fn transform_create_batch_request(&self) -> Result { - Err(Error::Unsupported("Anthropic message batch creation")) + Err(Error::Declined(Decline { + reason: "Anthropic message batch creation", + })) } fn transform_create_batch_response( @@ -135,7 +138,9 @@ impl AnthropicBatchesConfig for AnthropicBatchesTransformation { _response: AnthropicMessageBatch, _now: i64, ) -> Result { - Err(Error::Unsupported("Anthropic message batch creation")) + Err(Error::InvalidResponse(crate::ErrorDetail::Unhandled( + "Anthropic message batch creation", + ))) } fn retrieve_batch_url( @@ -337,16 +342,22 @@ mod tests { ); } - #[test] - fn preserves_python_placeholder_for_batch_creation() { - assert!(matches!( - ANTHROPIC_BATCHES_TRANSFORMATION.transform_create_batch_request(), - Err(Error::Unsupported("Anthropic message batch creation")) - )); + #[rstest::rstest] + fn batch_creation_only_declines_before_a_provider_response() { + let error = ANTHROPIC_BATCHES_TRANSFORMATION + .transform_create_batch_request() + .unwrap_err(); + assert!(!error.dispatched()); + assert_eq!( + error.declined().map(|decline| decline.reason), + Some("Anthropic message batch creation") + ); let response: AnthropicMessageBatch = serde_json::from_value(json!({})).unwrap(); - assert!(matches!( - ANTHROPIC_BATCHES_TRANSFORMATION.transform_create_batch_response(response, 0), - Err(Error::Unsupported("Anthropic message batch creation")) - )); + let error = ANTHROPIC_BATCHES_TRANSFORMATION + .transform_create_batch_response(response, 0) + .unwrap_err(); + assert!(error.dispatched()); + assert!(error.is_response()); + assert!(error.declined().is_none()); } } diff --git a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs index 4d4f110fbe1..eb796a900f5 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs @@ -1,4 +1,5 @@ use litellm_auth::SecretValue; +use litellm_core_utils::decline::Decline; use litellm_core_utils::{ core_helpers::{finish_reason_for, unix_now, usage_from_parts}, prompt_templates::factory::{Conversation, build_conversation}, @@ -24,7 +25,7 @@ use crate::{ streaming::StreamShape, transformation::{ BaseConfig, Headers, ProviderChatRequestData, ProviderChatResponseData, - Unsupported, ValidatedEnvironment, unsupported_message, unsupported_param, + ValidatedEnvironment, unsupported_message, unsupported_param, }, }, }, @@ -126,15 +127,14 @@ impl BaseConfig for AnthropicConfig { serde_json::from_value(response.body).map_err(|error| { Error::InvalidResponse(crate::ErrorDetail::invalid("messages response", error)) })?; - // The route declines tool and thinking requests, so a non-text block - // means the response carries something this path never asked for. - // Decline rather than silently dropping it; the host falls back. if body .content .iter() .any(|block| matches!(block, TextResponseBlock::Other)) { - return Err(Error::Unsupported("non-text response content block")); + return Err(Error::InvalidResponse(crate::ErrorDetail::Unhandled( + "non-text response content block", + ))); } let text: String = body .content @@ -203,7 +203,7 @@ impl BaseConfig for AnthropicConfig { &self, messages: &[ChatMessage], optional_params: &Map, - ) -> Option { + ) -> Option { unsupported_param(self.supported_openai_param_mappings(), &[], optional_params) .or_else(|| messages.iter().find_map(unsupported_message)) // Anthropic rejects a request whose first turn is not a user turn. @@ -211,7 +211,7 @@ impl BaseConfig for AnthropicConfig { // core cannot observe, so decline instead of guessing. .or_else(|| { (!build_conversation(messages).opens_on_user_turn()) - .then_some(Unsupported("conversation does not open on a user turn")) + .then_some(Decline::new("conversation does not open on a user turn")) }) } } @@ -260,9 +260,10 @@ mod tests { anthropic::chat::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG, base_llm::{ auth::AuthScheme, - chat::transformation::{BaseConfig, ProviderChatResponseData, Unsupported}, + chat::transformation::{BaseConfig, ProviderChatResponseData}, }, }; + use litellm_core_utils::decline::Decline; use litellm_llms_types::formats::chat_completions::{ChatCompletionsResponse, ChatMessage}; use rstest::rstest; use serde_json::{Map, Value, json}; @@ -290,7 +291,7 @@ mod tests { .transform_response("claude-sonnet-4-5", ProviderChatResponseData { body }) } - fn reason(msgs: Value, opts: Value) -> Option { + fn reason(msgs: Value, opts: Value) -> Option { ANTHROPIC_CHAT_COMPLETIONS_CONFIG.unsupported_reason(&messages(msgs), ¶ms(opts)) } @@ -398,7 +399,7 @@ mod tests { json!([{"role": "user", "content": "hi"}]), json!({"top_k": 40}) ), - Some(Unsupported("unrecognized request parameter")) + Some(Decline::new("unrecognized request parameter")) ); } @@ -409,7 +410,7 @@ mod tests { json!([{"role": "user", "content": "hi"}]), json!({"stream": true, "max_tokens": 16}) ), - Some(Unsupported("streaming")) + Some(Decline::new("streaming")) ); } @@ -434,7 +435,7 @@ mod tests { fn declines_any_param_outside_the_allowlist(#[case] param: Value) { assert_eq!( reason(json!([{"role": "user", "content": "hi"}]), param.clone()), - Some(Unsupported("unrecognized request parameter")), + Some(Decline::new("unrecognized request parameter")), "expected {param} to decline" ); } @@ -452,7 +453,7 @@ mod tests { ]), json!({}) ), - Some(Unsupported("unrecognized message field")) + Some(Decline::new("unrecognized message field")) ); assert_eq!( reason( @@ -462,7 +463,7 @@ mod tests { ]), json!({}) ), - Some(Unsupported("unrecognized message field")) + Some(Decline::new("unrecognized message field")) ); assert_eq!( reason( @@ -472,7 +473,7 @@ mod tests { ]}]), json!({}) ), - Some(Unsupported("non-text message content")) + Some(Decline::new("non-text message content")) ); assert_eq!( reason( @@ -482,7 +483,7 @@ mod tests { ]}]), json!({}) ), - Some(Unsupported("non-text message content")) + Some(Decline::new("non-text message content")) ); } @@ -493,7 +494,7 @@ mod tests { // than declining to Python before it. assert_eq!( reason(json!([{"role": "user", "content": []}]), json!({})), - Some(Unsupported("message without content")) + Some(Decline::new("message without content")) ); assert_eq!( reason( @@ -514,7 +515,7 @@ mod tests { ]), json!({}) ), - Some(Unsupported("conversation does not open on a user turn")) + Some(Decline::new("conversation does not open on a user turn")) ); } @@ -647,7 +648,7 @@ mod tests { } #[test] - fn declines_a_response_carrying_a_non_text_block() { + fn a_response_carrying_a_non_text_block_is_invalid() { let err = transform_response(json!({ "model": "claude-sonnet-4-5", "content": [{"type": "tool_use", "id": "t1", "name": "f", "input": {}}], @@ -655,7 +656,12 @@ mod tests { "usage": {"input_tokens": 1, "output_tokens": 1} })) .expect_err("non-text block"); - assert_eq!(err, Error::Unsupported("non-text response content block")); + assert_eq!( + err, + Error::InvalidResponse(crate::ErrorDetail::Unhandled( + "non-text response content block" + )) + ); } #[rstest::rstest] diff --git a/litellm-rust/crates/llms/src/base_llm/chat/transformation.rs b/litellm-rust/crates/llms/src/base_llm/chat/transformation.rs index 30bff76651e..8ba4030540c 100644 --- a/litellm-rust/crates/llms/src/base_llm/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/chat/transformation.rs @@ -1,3 +1,4 @@ +use litellm_core_utils::decline::Decline; use litellm_llms_types::formats::chat_completions::{ ChatCompletionsResponse, ChatMessage, ChatMessageContent, }; @@ -29,16 +30,6 @@ const IGNORABLE_MESSAGE_FIELDS: &[&str] = &["name"]; pub use crate::base_llm::auth::{Headers, ValidatedEnvironment}; -/// Why a request cannot be served by the Rust path. -/// -/// The core declines rather than guessing: the host turns this into a -/// transparent fallback to the Python implementation, which covers the full -/// surface. Acceptance is an allowlist, so a parameter or message shape the -/// core has never seen declines by construction instead of being translated -/// wrong. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct Unsupported(pub &'static str); - pub trait BaseConfig: Sync { fn secret_names(&self) -> Vec<&'static str>; @@ -97,7 +88,7 @@ pub trait BaseConfig: Sync { &self, messages: &[ChatMessage], optional_params: &Map, - ) -> Option { + ) -> Option { unsupported_param( self.supported_openai_param_mappings(), self.config_params(), @@ -111,13 +102,13 @@ pub fn unsupported_param( supported: &'static [(&'static str, &'static str)], config: &'static [&'static str], optional_params: &Map, -) -> Option { +) -> Option { if optional_params .get(STREAM_PARAM) .and_then(Value::as_bool) .unwrap_or(false) { - return Some(Unsupported("streaming")); + return Some(Decline::new("streaming")); } optional_params .keys() @@ -128,29 +119,29 @@ pub fn unsupported_param( .any(|(_, provider_name)| *provider_name == key) && !config.contains(&key.as_str()) }) - .then_some(Unsupported("unrecognized request parameter")) + .then_some(Decline::new("unrecognized request parameter")) } /// Message shapes the core can translate faithfully: text content, either a /// plain string or a non-empty list of parts that are all /// `{"type": "text", "text": ...}`. Tool calls, tool results, and multimodal /// parts decline so Python's fuller translation handles them. -pub fn unsupported_message(message: &ChatMessage) -> Option { +pub fn unsupported_message(message: &ChatMessage) -> Option { if message .extra .keys() .any(|key| !IGNORABLE_MESSAGE_FIELDS.contains(&key.as_str())) { - return Some(Unsupported("unrecognized message field")); + return Some(Decline::new("unrecognized message field")); } if !matches!(message.role.as_str(), "system" | "user" | "assistant") { - return Some(Unsupported("unrecognized message role")); + return Some(Decline::new("unrecognized message role")); } match &message.content { - None => Some(Unsupported("message without content")), + None => Some(Decline::new("message without content")), Some(ChatMessageContent::Text(_)) => None, Some(ChatMessageContent::Parts(parts)) if parts.is_empty() => { - Some(Unsupported("message without content")) + Some(Decline::new("message without content")) } Some(ChatMessageContent::Parts(parts)) => parts .iter() @@ -159,6 +150,6 @@ pub fn unsupported_message(message: &ChatMessage) -> Option { || part.get("text").and_then(Value::as_str).is_none() || part.as_object().is_some_and(|object| object.len() != 2) }) - .then_some(Unsupported("non-text message content")), + .then_some(Decline::new("non-text message content")), } } diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs index 9a0e0519f12..d51fe996dc4 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs @@ -87,8 +87,6 @@ pub enum Error { PollOrigin, #[error("OCR polling timed out")] PollTimeout, - #[error("unsupported by the rust path: {0}")] - Unsupported(&'static str), #[error("invalid provider: {0}")] InvalidProvider(String), #[error("invalid model: {provider} has no model {model:?} - use one of: {}", supported.join(", "))] @@ -212,7 +210,6 @@ impl Error { | Self::InvalidPollUrl { .. } | Self::PollOrigin | Self::PollTimeout - | Self::Unsupported(_) | Self::InvalidProvider(_) | Self::InvalidModel { .. } | Self::MissingCredential { .. } @@ -255,7 +252,6 @@ impl Error { | Self::InvalidPollUrl { .. } | Self::PollOrigin | Self::PollTimeout - | Self::Unsupported(_) | Self::InvalidProvider(_) | Self::InvalidModel { .. } | Self::Secret(_) @@ -268,49 +264,6 @@ impl Error { } } - pub fn is_unsupported(&self) -> bool { - match self { - Self::Unsupported(_) => true, - Self::Provider { .. } - | Self::EmptyFile - | Self::FileRead { .. } - | Self::DocumentTask(_) - | Self::InvalidMimeType(_) - | Self::UnsupportedDocument { .. } - | Self::RejectedDocument { .. } - | Self::RequestFormat - | Self::RequestField { .. } - | Self::Arguments(_) - | Self::MissingField(_) - | Self::MissingDocumentUrl - | Self::InvalidDataUri - | Self::InlineDocumentTooLarge - | Self::Pages(_) - | Self::Features - | Self::DotModel - | Self::TooLarge { .. } - | Self::ResponseField { .. } - | Self::EmptyContent - | Self::OperationStatus(_) - | Self::NumericRange(_) - | Self::PollLocation - | Self::InvalidPollUrl { .. } - | Self::PollOrigin - | Self::PollTimeout - | Self::InvalidProvider(_) - | Self::InvalidModel { .. } - | Self::MissingCredential { .. } - | Self::Secret(_) - | Self::Host(_) - | Self::HostCall - | Self::Auth(_) - | Self::Params(_) - | Self::Headers(_) - | Self::Http(_) - | Self::DocumentDownload(_) => false, - } - } - pub fn http_status_code(&self) -> Option { if let Some((status, ..)) = self.upstream_response() { return Some(status); @@ -357,7 +310,6 @@ impl Error { | Self::InvalidPollUrl { .. } | Self::PollOrigin | Self::PollTimeout - | Self::Unsupported(_) | Self::MissingCredential { .. } | Self::Secret(_) | Self::Host(_) @@ -394,7 +346,6 @@ impl Error { | Self::Pages(_) | Self::Features | Self::DotModel - | Self::Unsupported(_) | Self::InvalidProvider(_) | Self::InvalidModel { .. } | Self::MissingCredential { .. } diff --git a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs index acaf7b56f8b..976c9161dd2 100644 --- a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs @@ -5,6 +5,7 @@ use litellm_auth_aws::{ constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE}, resolve_bedrock_region, }; +use litellm_core_utils::decline::Decline; use litellm_core_utils::{ core_helpers::{finish_reason_for, unix_now, usage_from_parts}, prompt_templates::factory::{Conversation, TurnRole, build_conversation}, @@ -26,7 +27,7 @@ use crate::{ streaming::StreamShape, transformation::{ BaseConfig, Headers, ProviderChatRequestData, ProviderChatResponseData, - Unsupported, ValidatedEnvironment, unsupported_message, unsupported_param, + ValidatedEnvironment, unsupported_message, unsupported_param, }, }, }, @@ -164,10 +165,10 @@ impl BaseConfig for AmazonConverseConfig { field: "usage", })); } - // The route declines tool requests, so anything other than a text block - // is something this path never asked for. Decline; the host falls back. if response.message_content_is_non_text() { - return Err(Error::Unsupported("non-text response content block")); + return Err(Error::InvalidResponse(ErrorDetail::Unhandled( + "non-text response content block", + ))); } let text = response.content_text(); let computed = usage_from_parts( @@ -264,7 +265,7 @@ impl BaseConfig for AmazonConverseConfig { &self, messages: &[ChatMessage], optional_params: &Map, - ) -> Option { + ) -> Option { unsupported_param( self.supported_openai_param_mappings(), CONFIG_PARAMS, @@ -278,7 +279,7 @@ impl BaseConfig for AmazonConverseConfig { messages .iter() .any(has_blank_text) - .then_some(Unsupported("blank message text")) + .then_some(Decline::new("blank message text")) }) // Converse has no assistant prefill: Python inserts a continue turn // when a conversation opens or closes on an assistant message, and @@ -291,7 +292,7 @@ impl BaseConfig for AmazonConverseConfig { .turns .last() .is_some_and(|turn| turn.role == TurnRole::Assistant); - (!conversation.opens_on_user_turn() || ends_on_assistant).then_some(Unsupported( + (!conversation.opens_on_user_turn() || ends_on_assistant).then_some(Decline::new( "conversation does not run user turn to user turn", )) }) @@ -353,11 +354,12 @@ mod tests { Error, base_llm::{ auth::AuthScheme, - chat::transformation::{BaseConfig, ProviderChatResponseData, Unsupported}, + chat::transformation::{BaseConfig, ProviderChatResponseData}, }, bedrock::chat::converse_transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG, }; use litellm_auth::CredentialPlacement; + use litellm_core_utils::decline::Decline; use litellm_llms_types::formats::chat_completions::{ChatCompletionsResponse, ChatMessage}; use rstest::rstest; use serde_json::{Map, Value, json}; @@ -391,7 +393,7 @@ mod tests { ) } - fn reason(msgs: Value, opts: Value) -> Option { + fn reason(msgs: Value, opts: Value) -> Option { BEDROCK_CHAT_COMPLETIONS_CONFIG.unsupported_reason(&messages(msgs), ¶ms(opts)) } @@ -467,7 +469,7 @@ mod tests { json!([{"role": "user", "content": "hi"}]), json!({"stream": true}) ), - Some(Unsupported("streaming")) + Some(Decline::new("streaming")) ); } @@ -478,7 +480,7 @@ mod tests { json!([{"role": "user", "content": "hi"}]), json!({"topK": 40}) ), - Some(Unsupported("unrecognized request parameter")) + Some(Decline::new("unrecognized request parameter")) ); } @@ -492,7 +494,7 @@ mod tests { fn declines_tools_and_other_params_outside_the_allowlist(#[case] param: Value) { assert_eq!( reason(json!([{"role": "user", "content": "hi"}]), param.clone()), - Some(Unsupported("unrecognized request parameter")), + Some(Decline::new("unrecognized request parameter")), "expected {param} to decline" ); } @@ -509,7 +511,7 @@ mod tests { json!([{"role": "user", "content": content}, {"role": "user", "content": "hi"}]), json!({}) ), - Some(Unsupported("blank message text")), + Some(Decline::new("blank message text")), "expected blank content {content} to decline" ); } @@ -521,7 +523,7 @@ mod tests { // host before the call rather than an error after it. assert_eq!( reason(json!([{"role": "user", "content": []}]), json!({})), - Some(Unsupported("message without content")) + Some(Decline::new("message without content")) ); assert_eq!( reason( @@ -542,7 +544,7 @@ mod tests { ]), json!({}) ), - Some(Unsupported( + Some(Decline::new( "conversation does not run user turn to user turn" )) ); @@ -554,7 +556,7 @@ mod tests { ]), json!({}) ), - Some(Unsupported( + Some(Decline::new( "conversation does not run user turn to user turn" )) ); @@ -814,7 +816,7 @@ mod tests { ]}]), json!({}) ), - Some(Unsupported("non-text message content")) + Some(Decline::new("non-text message content")) ); } @@ -841,7 +843,7 @@ mod tests { } #[test] - fn declines_a_response_carrying_a_tool_use_block() { + fn a_response_carrying_a_tool_use_block_is_invalid() { let err = transform_response(json!({ "output": {"message": {"content": [ {"toolUse": {"toolUseId": "t1", "name": "f", "input": {}}} @@ -850,7 +852,12 @@ mod tests { "usage": {"inputTokens": 1, "outputTokens": 1} })) .expect_err("tool use block"); - assert_eq!(err, Error::Unsupported("non-text response content block")); + assert_eq!( + err, + Error::InvalidResponse(crate::ErrorDetail::Unhandled( + "non-text response content block" + )) + ); } #[rstest::rstest] diff --git a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs index 1f6b6edf28d..3ec60f12021 100644 --- a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs @@ -1,3 +1,4 @@ +use litellm_core_utils::decline::Decline; use std::convert::Infallible; use crate::{ @@ -130,9 +131,9 @@ impl BaseMessagesConfig for AmazonAnthropicClaudeMessagesConfig { _request: MessagesRequest, _context: &MessagesTransformContext, ) -> Result { - Err(Error::Unsupported( + Err(Error::Declined(Decline::new( "Bedrock invoke messages request shaping", - )) + ))) } fn secret_names(&self) -> &'static [&'static str] { diff --git a/litellm-rust/crates/llms/src/error.rs b/litellm-rust/crates/llms/src/error.rs index eddcb27d1d9..b3d1421083a 100644 --- a/litellm-rust/crates/llms/src/error.rs +++ b/litellm-rust/crates/llms/src/error.rs @@ -1,3 +1,4 @@ +use litellm_core_utils::decline::Decline; use litellm_error::Source; #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] @@ -13,8 +14,8 @@ pub enum Error { InvalidRequest(#[source] ErrorDetail), #[error("invalid response")] InvalidResponse(#[source] ErrorDetail), - #[error("unsupported: {0}")] - Unsupported(&'static str), + #[error(transparent)] + Declined(Decline), #[error(transparent)] Auth(#[from] litellm_auth::Error), #[error(transparent)] @@ -24,13 +25,10 @@ pub enum Error { impl Error { pub fn is_request(&self) -> bool { match self { - Self::InvalidType { .. } - | Self::MissingField(_) - | Self::InvalidRequest(_) - | Self::Unsupported(_) => true, + Self::InvalidType { .. } | Self::MissingField(_) | Self::InvalidRequest(_) => true, Self::Auth(error) => error.is_request(), Self::Http(error) => error.is_request(), - Self::InvalidResponse(_) => false, + Self::InvalidResponse(_) | Self::Declined(_) => false, } } @@ -41,7 +39,7 @@ impl Error { Self::InvalidType { .. } | Self::MissingField(_) | Self::InvalidRequest(_) - | Self::Unsupported(_) + | Self::Declined(_) | Self::Auth(_) => false, } } @@ -53,7 +51,7 @@ impl Error { | Self::MissingField(_) | Self::InvalidRequest(_) | Self::InvalidResponse(_) - | Self::Unsupported(_) + | Self::Declined(_) | Self::Http(_) => false, } } @@ -67,8 +65,8 @@ impl Error { Self::InvalidType { .. } | Self::MissingField(_) | Self::InvalidRequest(_) - | Self::InvalidResponse(_) - | Self::Unsupported(_) => true, + | Self::InvalidResponse(_) => true, + Self::Declined(_) => false, } } @@ -78,7 +76,19 @@ impl Error { match self { Self::InvalidRequest(detail) => Some(detail), Self::InvalidType { .. } | Self::MissingField(_) => Some(self), - Self::InvalidResponse(_) | Self::Unsupported(_) | Self::Auth(_) | Self::Http(_) => None, + Self::InvalidResponse(_) | Self::Declined(_) | Self::Auth(_) | Self::Http(_) => None, + } + } + + pub fn declined(&self) -> Option<&Decline> { + match self { + Self::Declined(decline) => Some(decline), + Self::InvalidType { .. } + | Self::MissingField(_) + | Self::InvalidRequest(_) + | Self::InvalidResponse(_) + | Self::Auth(_) + | Self::Http(_) => None, } } @@ -89,7 +99,7 @@ impl Error { | Self::MissingField(_) | Self::InvalidRequest(_) | Self::InvalidResponse(_) - | Self::Unsupported(_) + | Self::Declined(_) | Self::Auth(_) => None, } } @@ -154,6 +164,8 @@ pub enum ErrorDetail { MissingApiBase, #[error("the host call raised an exception")] HostCall, + #[error("{0} is not handled by the rust path")] + Unhandled(&'static str), #[error("invalid {subject} on line {line}")] InvalidLine { subject: &'static str, diff --git a/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs b/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs index 9f208db2471..e998f8785cb 100644 --- a/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs @@ -246,14 +246,20 @@ fn normalize_choice(position: usize, choice: &Value) -> Result Option { + fn reason(msgs: Value, opts: Value) -> Option { OPENAI_LIKE_CHAT_COMPLETIONS_CONFIG.unsupported_reason(&messages(msgs), ¶ms(opts)) } @@ -535,9 +542,7 @@ mod tests { } #[rstest] - fn a_tool_call_response_declines_instead_of_dropping_the_calls() { - // The `json_mode` rewrite needs a request flag the route does not carry, so - // a tool-call answer falls back to Python rather than losing the calls. + fn a_tool_call_response_fails_instead_of_dropping_the_calls() { assert_eq!( transform_response(json!({ "model": "m", @@ -554,12 +559,14 @@ mod tests { "finish_reason": "tool_calls", }], })), - Err(Error::Unsupported("tool call response")) + Err(Error::InvalidResponse(crate::ErrorDetail::Unhandled( + "tool call response" + ))) ); } #[rstest] - fn a_refusal_declines_instead_of_returning_an_empty_reply() { + fn a_refusal_fails_instead_of_returning_an_empty_reply() { assert_eq!( transform_response(json!({ "model": "m", @@ -568,12 +575,14 @@ mod tests { "finish_reason": "stop", }], })), - Err(Error::Unsupported("refusal response")) + Err(Error::InvalidResponse(crate::ErrorDetail::Unhandled( + "refusal response" + ))) ); } #[rstest] - fn a_non_text_response_content_declines() { + fn a_non_text_response_content_is_invalid() { assert_eq!( transform_response(json!({ "model": "m", @@ -582,7 +591,9 @@ mod tests { "finish_reason": "stop", }], })), - Err(Error::Unsupported("non-text response content")) + Err(Error::InvalidResponse(crate::ErrorDetail::Unhandled( + "non-text response content" + ))) ); } @@ -592,7 +603,7 @@ mod tests { fn declines(#[case] opts: Value, #[case] expected: &'static str) { assert_eq!( reason(json!([{"role": "user", "content": "hi"}]), opts), - Some(Unsupported(expected)) + Some(Decline::new(expected)) ); } @@ -622,7 +633,7 @@ mod tests { json!([{"role": "user", "content": "hi"}]), json!({"tools": [{"type": "function", "function": {"name": "f"}}]}), ), - Some(Unsupported("unrecognized request parameter")) + Some(Decline::new("unrecognized request parameter")) ); } } diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index a883ec608e5..b42c08b7e78 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -62,14 +62,18 @@ pub(crate) enum RequestMarker { Messages, } -/// Maps a route failure once. A Python exception raised inside the call is returned as it is, -/// everything else is mapped and handed to `public`, the route's public exception mapping. +/// Maps a route failure once. A core decline and a Python exception raised inside the call are +/// returned as they are, everything else is mapped and handed to `public`, the route's public +/// exception mapping. pub(crate) fn route_failure( py: Python<'_>, error: Error, marker: RequestMarker, public: impl FnOnce(PyErr) -> PyResult, ) -> PyResult { + if let Some(decline) = error.declined() { + return Ok(RustBridgeDeclined::new_err(decline.reason)); + } if let Some(raised) = crate::secrets::python_error(py, &error) { return Ok(raised); } @@ -97,6 +101,7 @@ pub(crate) fn route_error_to_pyerr(error: Error) -> PyErr { #[cfg(test)] mod tests { + use litellm_core_utils::decline::Decline; use reqwest::header::{HeaderValue, RETRY_AFTER}; use super::*; @@ -110,7 +115,6 @@ mod tests { } #[rstest::rstest] - #[case::unsupported(Error::Unsupported("test capability"), true)] #[case::invalid_provider(Error::InvalidProvider("unknown".into()), true)] #[case::invalid_request( Error::invalid_request(litellm_llms::ErrorDetail::Expected { @@ -130,6 +134,10 @@ mod tests { }), false )] + #[case::post_call_decline( + Error::post_call(Error::Declined(Decline::new("result rejected"))), + false + )] fn route_failures_are_terminal(#[case] error: Error, #[case] is_request: bool) { Python::initialize(); Python::attach(|py| { @@ -140,6 +148,34 @@ mod tests { }); } + #[rstest::rstest] + #[case::unmarked(RequestMarker::Unmarked)] + #[case::messages(RequestMarker::Messages)] + fn a_core_decline_hands_the_call_back_to_python_before_the_public_mapping( + #[case] marker: RequestMarker, + ) { + Python::initialize(); + Python::attach(|py| { + let error = Error::from(litellm_llms::Error::Declined(Decline::new( + "unrecognized request parameter", + ))); + let declined = route_failure(py, error, marker, |_| { + panic!("a decline must not reach the public mapping") + }) + .unwrap(); + assert!(declined.is_instance_of::(py)); + assert_eq!( + declined + .value(py) + .getattr("args") + .unwrap() + .extract::<(String,)>() + .unwrap(), + ("unrecognized request parameter".into(),) + ); + }); + } + #[rstest::rstest] #[case::unmarked(RequestMarker::Unmarked)] #[case::messages(RequestMarker::Messages)] diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs index 63ccfc2433f..4db9b37cbb8 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs @@ -1,3 +1,4 @@ +use litellm_core_utils::decline::Decline; use std::convert::Infallible; use crate::routes::inference::InferenceHost; @@ -44,9 +45,9 @@ impl PythonBinding for ChatCompletionsPythonHost { .get("stream") .is_some_and(|value| value == &serde_json::Value::Bool(true)) { - return Err(InvokeError::Native(Error::Unsupported( + return Err(InvokeError::Native(Error::Declined(Decline::new( "native Python chat_completions streaming", - ))); + )))); } Ok(call) } diff --git a/litellm-rust/crates/python-bridge/src/routes/responses/host.rs b/litellm-rust/crates/python-bridge/src/routes/responses/host.rs index eb19d299625..b797e334c6a 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses/host.rs @@ -1,3 +1,4 @@ +use litellm_core_utils::decline::Decline; use std::convert::Infallible; use crate::routes::inference::InferenceHost; @@ -118,9 +119,9 @@ impl PythonBinding for ResponsesPythonHost { .get("stream") .is_some_and(|value| value == &serde_json::Value::Bool(true)) { - return Err(InvokeError::Native(Error::Unsupported( + return Err(InvokeError::Native(Error::Declined(Decline::new( "native Python responses streaming", - ))); + )))); } Ok(call) } @@ -141,9 +142,11 @@ impl PythonBinding for ResponsesPythonHost { head: ::StreamHead, ) -> PyResult> { let _ = head; - Err(crate::errors::route_error_to_pyerr(Error::Unsupported( - "native Python responses streaming", - ))) + Err(crate::errors::route_error_to_pyerr( + Error::invalid_response(litellm_llms::ErrorDetail::Unhandled( + "native Python responses streaming", + )), + )) } fn encode_chunk( @@ -152,9 +155,11 @@ impl PythonBinding for ResponsesPythonHost { chunk: ::Chunk, ) -> PyResult> { let _ = chunk; - Err(crate::errors::route_error_to_pyerr(Error::Unsupported( - "native Python responses streaming", - ))) + Err(crate::errors::route_error_to_pyerr( + Error::invalid_response(litellm_llms::ErrorDetail::Unhandled( + "native Python responses streaming", + )), + )) } fn map_error(&self, py: Python<'_>, error: Error) -> PyResult { diff --git a/tests/test_litellm_rust/chat_completions/test_failures.py b/tests/test_litellm_rust/chat_completions/test_failures.py index 68a81096a7c..dda5a1bad90 100644 --- a/tests/test_litellm_rust/chat_completions/test_failures.py +++ b/tests/test_litellm_rust/chat_completions/test_failures.py @@ -67,3 +67,16 @@ async def test_empty_messages_fail_before_provider_io(server: RecordingServer, a await call(server, asynchronous=asynchronous, messages=[]) assert not server.requests + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +async def test_a_native_decline_reaches_the_dispatcher_before_provider_io( + server: RecordingServer, asynchronous: bool +) -> None: + server.expected_requests = 0 + + with pytest.raises(RuntimeError, match="declined the request: native Python chat_completions streaming"): + await call(server, asynchronous=asynchronous, stream=True) + + assert not server.requests diff --git a/tests/test_litellm_rust/responses/data.py b/tests/test_litellm_rust/responses/data.py index 0347fd7a57b..3379b0f80b1 100644 --- a/tests/test_litellm_rust/responses/data.py +++ b/tests/test_litellm_rust/responses/data.py @@ -70,6 +70,6 @@ async def call(server: RecordingServer, *, asynchronous: bool = True, **kwargs: SYNC_DEPLOYMENT_HOOK: Final = gap("native sync calls skip the async pre-call deployment hook") STREAMING: Final = gap( - "native Responses streaming is unsupported, the catalog keeps it on Python", raises=litellm.APIConnectionError + "native Responses streaming is unsupported, the catalog keeps it on Python", raises=RuntimeError ) BODY_HOOKS: Final = gap("native Responses does not run the handler body hooks") diff --git a/tests/test_litellm_rust/responses/test_failures.py b/tests/test_litellm_rust/responses/test_failures.py index 4c4bb96c6ac..bfd6629fb33 100644 --- a/tests/test_litellm_rust/responses/test_failures.py +++ b/tests/test_litellm_rust/responses/test_failures.py @@ -56,3 +56,16 @@ async def test_provider_failure_is_terminal_and_shared_with_callbacks( assert len(failure) == 1 assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value assert not any("success" in name for name in recorder.names) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +async def test_a_native_decline_reaches_the_dispatcher_before_provider_io( + server: RecordingServer, asynchronous: bool +) -> None: + server.expected_requests = 0 + + with pytest.raises(RuntimeError, match="declined the request: native Python responses streaming"): + await call(server, asynchronous=asynchronous, stream=True) + + assert not server.requests From 4ce26025481044922ea93e2f6d933d72612b7af7 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Wed, 28 Oct 2026 09:27:26 -0700 Subject: [PATCH 045/128] fix(huggingface): stop forcing truncate=false on TEI rerank and pass truncate params through (#48706) * test(huggingface): cover TEI rerank truncate passthrough and server default Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(huggingface): stop forcing truncate=false on TEI rerank and pass truncate params through Co-authored-by: Luke Schaefer <225965040+luke-e-schaefer@users.noreply.github.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(huggingface): keep TEI truncate options provider-specific Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(huggingface): use MockTransport for rerank truncation 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: Luke Schaefer <225965040+luke-e-schaefer@users.noreply.github.com> --- .../llms/huggingface/rerank/transformation.py | 7 +- .../test_huggingface_rerank_truncate_wire.py | 74 +++++++++++++++ .../test_huggingface_rerank_transformation.py | 91 +++++++++++++------ 3 files changed, 145 insertions(+), 27 deletions(-) create mode 100644 tests/integration/providers/test_huggingface_rerank_truncate_wire.py diff --git a/litellm/llms/huggingface/rerank/transformation.py b/litellm/llms/huggingface/rerank/transformation.py index 1f130f84283..476e22fc088 100644 --- a/litellm/llms/huggingface/rerank/transformation.py +++ b/litellm/llms/huggingface/rerank/transformation.py @@ -78,6 +78,8 @@ class HuggingFaceRerankConfig(BaseRerankConfig): "documents", "top_n", "return_documents", + "truncate", + "truncation_direction", ] def map_cohere_rerank_params( @@ -104,6 +106,10 @@ class HuggingFaceRerankConfig(BaseRerankConfig): optional_rerank_params["return_text"] = v elif k == "top_n" and v is not None: optional_rerank_params["top_n"] = v + elif k == "truncate" and isinstance(v, bool): + optional_rerank_params["truncate"] = v + elif k == "truncation_direction" and v is not None: + optional_rerank_params["truncation_direction"] = v elif k == "documents" and v is not None: optional_rerank_params["texts"] = v elif k == "query" and v is not None: @@ -151,7 +157,6 @@ class HuggingFaceRerankConfig(BaseRerankConfig): # HuggingFace API expects return_text parameter, corresponding to our return_documents parameter request_body: Final = { "raw_scores": False, - "truncate": False, "truncation_direction": "Right", } diff --git a/tests/integration/providers/test_huggingface_rerank_truncate_wire.py b/tests/integration/providers/test_huggingface_rerank_truncate_wire.py new file mode 100644 index 00000000000..0d29dbf0ded --- /dev/null +++ b/tests/integration/providers/test_huggingface_rerank_truncate_wire.py @@ -0,0 +1,74 @@ +import json +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +import pytest +from integration._support.client import JSON_OBJECT, Gateway +from integration._support.wire import Reply, Request, wire_server + +MODEL: Final = "huggingface/BAAI/bge-reranker-base" +QUERY: Final = "what is a fox" +TEXT: Final = "a short doc about foxes" +RESPONSE: Final = json.dumps([{"index": 0, "score": 0.82}]).encode() +TRUNCATION_CASES: Final = ( + pytest.param( + MappingProxyType({}), + MappingProxyType({"truncation_direction": "Right"}), + id="omitted-truncate-uses-server-default", + ), + pytest.param( + MappingProxyType({"truncate": True, "truncation_direction": "Left"}), + MappingProxyType({"truncate": True, "truncation_direction": "Left"}), + id="explicit-true-and-left-are-forwarded", + ), + pytest.param( + MappingProxyType({"truncate": False}), + MappingProxyType({"truncate": False, "truncation_direction": "Right"}), + id="explicit-false-is-forwarded", + ), +) + + +def rerank_peer(request: Request, expected_body: Mapping[str, object]) -> Reply: + assert request.method == "POST" and request.target == "/rerank", request.target + assert request.headers["authorization"] == "Bearer integration-provider-key" + body: Final = JSON_OBJECT.validate_json(request.body) + assert body == expected_body, body + return Reply(body=RESPONSE) + + +@pytest.mark.parametrize(("request_params", "expected_truncation_params"), TRUNCATION_CASES) +def test_huggingface_rerank_forwards_truncation_options( + gateway: Gateway, + request_params: Mapping[str, bool | str], + expected_truncation_params: Mapping[str, bool | str], +) -> None: + expected_body: Final = { + "raw_scores": False, + "truncation_direction": "Right", + "query": QUERY, + "texts": [TEXT], + "return_text": False, + **expected_truncation_params, + } + with wire_server(lambda request: rerank_peer(request, expected_body)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=MODEL, + api_base=wire.url, + model_info={"mode": "rerank"}, + ) + response: Final = gateway.request( + "POST", + "/v1/rerank", + { + "model": model, + "query": QUERY, + "documents": [TEXT], + "return_documents": False, + **request_params, + }, + ) + assert response.status_code == 200, response.text + assert response.json()["results"] == [{"index": 0, "relevance_score": 0.82}], response.text + assert len(wire.drain()) == 1, "Expected exactly one provider rerank call" diff --git a/tests/unit/llms/huggingface/rerank/test_huggingface_rerank_transformation.py b/tests/unit/llms/huggingface/rerank/test_huggingface_rerank_transformation.py index 6fd2b006fef..00a3cbfdb7f 100644 --- a/tests/unit/llms/huggingface/rerank/test_huggingface_rerank_transformation.py +++ b/tests/unit/llms/huggingface/rerank/test_huggingface_rerank_transformation.py @@ -5,11 +5,14 @@ Based on the test patterns from other rerank providers and the current HuggingFa import asyncio import json +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler def assert_response_shape(response, custom_llm_provider): @@ -219,11 +222,11 @@ def test_huggingface_rerank_return_documents(mock_post): assert "text" in result["document"] -def test_huggingface_rerank_config(): +def test_huggingface_rerank_config() -> None: """Test HuggingFaceRerankConfig class functionality.""" from litellm.llms.huggingface.rerank.transformation import HuggingFaceRerankConfig - config = HuggingFaceRerankConfig() + config: Final = HuggingFaceRerankConfig() # Test complete URL generation assert config.get_complete_url(None, "test") == "https://api-inference.huggingface.co/rerank" @@ -233,14 +236,16 @@ def test_huggingface_rerank_config(): assert custom_url == "https://custom.huggingface.co/rerank" # Test supported parameters - supported_params = config.get_supported_cohere_rerank_params("test") + supported_params: Final = config.get_supported_cohere_rerank_params("test") assert "query" in supported_params assert "documents" in supported_params assert "top_n" in supported_params assert "return_documents" in supported_params + assert "truncate" in supported_params + assert "truncation_direction" in supported_params # Test parameter mapping - params = config.map_cohere_rerank_params( + params: Final = config.map_cohere_rerank_params( non_default_params={ "query": "hello", "documents": ["hello", "world"], @@ -259,24 +264,28 @@ def test_huggingface_rerank_config(): assert params["return_text"] is True -def test_request_transformation(): +def test_request_transformation() -> None: """Test request transformation logic.""" from litellm.llms.huggingface.rerank.transformation import HuggingFaceRerankConfig from litellm.types.rerank import OptionalRerankParams - config = HuggingFaceRerankConfig() + config: Final = HuggingFaceRerankConfig() - optional_params = OptionalRerankParams(query="hello", texts=["hello", "world"], top_n=2, return_text=True) + optional_params: Final = OptionalRerankParams(query="hello", texts=["hello", "world"], top_n=2, return_text=True) - request_body = config.transform_rerank_request(model="test", optional_rerank_params=optional_params, headers={}) + request_body: Final = config.transform_rerank_request( + model="test", optional_rerank_params=optional_params, headers={} + ) - assert request_body["query"] == "hello" - assert request_body["texts"] == ["hello", "world"] - assert request_body["top_n"] == 2 - assert request_body["return_text"] is True - assert request_body["raw_scores"] is False - assert request_body["truncate"] is False - assert request_body["truncation_direction"] == "Right" + expected_request: Final = { + "query": "hello", + "texts": ["hello", "world"], + "top_n": 2, + "return_text": True, + "raw_scores": False, + "truncation_direction": "Right", + } + assert request_body == expected_request def test_response_transformation(): @@ -345,20 +354,20 @@ def test_validate_environment(): @patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") -def test_huggingface_rerank_request_payload(mock_post): +def test_huggingface_rerank_request_payload(mock_post: MagicMock) -> None: """Test that the request payload is correctly formatted for HuggingFace API.""" - mock_response_data = [{"index": 0, "score": 0.9}, {"index": 1, "score": 0.1}] + mock_response_data: Final = [{"index": 0, "score": 0.9}, {"index": 1, "score": 0.1}] def return_val(): return mock_response_data - mock_response = MagicMock() + mock_response: Final = MagicMock() mock_response.json = return_val mock_response.status_code = 200 mock_response.headers = {"content-type": "application/json"} mock_post.return_value = mock_response - response = litellm.rerank( + response: Final = litellm.rerank( model="huggingface/BAAI/bge-reranker-base", query="hello", documents=["hello", "world"], @@ -370,28 +379,58 @@ def test_huggingface_rerank_request_payload(mock_post): mock_post.assert_called_once() # Verify URL - call_url = mock_post.call_args.kwargs["url"] + call_url: Final = mock_post.call_args.kwargs["url"] assert call_url == "https://api-inference.huggingface.co/rerank" # Verify headers - headers = mock_post.call_args.kwargs["headers"] + headers: Final = mock_post.call_args.kwargs["headers"] assert "Bearer test_api_key" in headers["Authorization"] assert headers["content-type"] == "application/json" # Verify request body - request_data = json.loads(mock_post.call_args.kwargs["data"]) - expected_request = { + request_data: Final = json.loads(mock_post.call_args.kwargs["data"]) + expected_request: Final = { "query": "hello", "texts": ["hello", "world"], "raw_scores": False, "return_text": True, - "truncate": False, "truncation_direction": "Right", "top_n": 2, } - for key, value in expected_request.items(): - assert request_data[key] == value + assert request_data == expected_request assert response.results is not None assert len(response.results) == 2 + + +def test_huggingface_rerank_forwards_truncate_options() -> None: + expected_request: Final = { + "query": "hello", + "texts": ["hello", "world"], + "raw_scores": False, + "return_text": False, + "truncation_direction": "Left", + "truncate": True, + } + + def respond(request: httpx.Request) -> httpx.Response: + assert request.url == "https://api-inference.huggingface.co/rerank" + assert request.headers["Authorization"] == "Bearer test_api_key" + assert json.loads(request.content) == expected_request + return httpx.Response(200, json=[{"index": 0, "score": 0.9}]) + + with httpx.Client(transport=httpx.MockTransport(respond)) as http_client: + response: Final = litellm.rerank( + model="huggingface/BAAI/bge-reranker-base", + query="hello", + documents=["hello", "world"], + api_key="test_api_key", + return_documents=False, + truncate=True, + truncation_direction="Left", + client=HTTPHandler(client=http_client), + ) + + assert response.results is not None + assert response.results[0]["relevance_score"] == 0.9 From 3d7b94b34f15c56123790436b1880f62291370bf Mon Sep 17 00:00:00 2001 From: moe-berri Date: Wed, 28 Oct 2026 09:32:47 -0700 Subject: [PATCH 046/128] fix(ollama): forward embedding credentials, headers and timeout (#48707) * test(integration): cover ollama embedding key, header and timeout forwarding * fix(ollama): forward api_key, headers and timeout on native embeddings Co-authored-by: Mrun25 <165297479+Mrun25@users.noreply.github.com> * test(ollama): disable fixture key for keyless proxy guard --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Mrun25 <165297479+Mrun25@users.noreply.github.com> --- litellm/llms/ollama/completion/handler.py | 38 ++- litellm/main.py | 4 + tests/integration/_support/wire.py | 3 +- .../test_api_base_trailing_slash_wire.py | 3 +- .../providers/test_ollama_embedding_wire.py | 134 +++++++++++ .../sdk/test_ollama_embedding_wire.py | 218 ++++++++++++++++++ .../ollama/test_ollama_chat_transformation.py | 6 + .../unit/llms/ollama/test_ollama_embedding.py | 180 ++++++++++++++- 8 files changed, 577 insertions(+), 9 deletions(-) create mode 100644 tests/integration/providers/test_ollama_embedding_wire.py create mode 100644 tests/integration/sdk/test_ollama_embedding_wire.py diff --git a/litellm/llms/ollama/completion/handler.py b/litellm/llms/ollama/completion/handler.py index 0c13b1cd1d3..1e70141f939 100644 --- a/litellm/llms/ollama/completion/handler.py +++ b/litellm/llms/ollama/completion/handler.py @@ -5,12 +5,15 @@ Ollama /chat/completion calls handled in llm_http_handler.py """ from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import Any, Final, Protocol, TypedDict +import httpx from typing_extensions import NotRequired, ReadOnly import litellm from litellm.litellm_core_utils.url_utils import ensure_api_route +from litellm.types.mcp import has_header from litellm.types.utils import EmbeddingResponse @@ -82,6 +85,15 @@ def _process_ollama_embedding_response( return model_response +def _ollama_embedding_headers(api_key: str | None, headers: Mapping[str, str] | None) -> Mapping[str, str]: + return MappingProxyType( + { + **({"Authorization": f"Bearer {api_key}"} if api_key and not has_header(headers, "Authorization") else {}), + **(headers or {}), + } + ) + + async def ollama_aembeddings( api_base: str, model: str, @@ -90,10 +102,19 @@ async def ollama_aembeddings( optional_params: Mapping[str, object], logging_obj: object, encoding: TokenEncoder | None, -): + *, + api_key: str | None = None, + headers: Mapping[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, +) -> EmbeddingResponse: data: Final = _prepare_ollama_embedding_payload(model, prompts, optional_params) - response: Final = await litellm.module_level_aclient.post(url=ensure_api_route(api_base, "api/embed"), json=data) + response: Final = await litellm.module_level_aclient.post( + url=ensure_api_route(api_base, "api/embed"), + json=data, + headers=dict(_ollama_embedding_headers(api_key, headers)), + timeout=timeout, + ) response_json: Final[OllamaEmbeddingResponse] = response.json() return _process_ollama_embedding_response( @@ -114,10 +135,19 @@ def ollama_embeddings( model_response: EmbeddingResponse, logging_obj: object, encoding: TokenEncoder | None = None, -): + *, + api_key: str | None = None, + headers: Mapping[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, +) -> EmbeddingResponse: data: Final = _prepare_ollama_embedding_payload(model, prompts, optional_params) - response: Final = litellm.module_level_client.post(url=ensure_api_route(api_base, "api/embed"), json=data) + response: Final = litellm.module_level_client.post( + url=ensure_api_route(api_base, "api/embed"), + json=data, + headers=dict(_ollama_embedding_headers(api_key, headers)), + timeout=timeout, + ) response_json: Final[OllamaEmbeddingResponse] = response.json() return _process_ollama_embedding_response( diff --git a/litellm/main.py b/litellm/main.py index 39a0a0cd3c1..79d2d0370f7 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -4229,6 +4229,7 @@ def _complete_sagemaker(ctx: CompletionDispatchContext) -> _CompletionDispatchRe _ADDITIONAL_DROP_PARAMS_ADAPTER: Final = TypeAdapter(list[str]) _OPTIONAL_PARAMS_ADAPTER: Final = TypeAdapter(dict[str, object]) +_OLLAMA_EMBEDDING_HEADERS_ADAPTER: Final = TypeAdapter(dict[str, str]) _REQUEST_KWARGS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _PROVIDER_ROUTING_HINTS: Final = TypeAdapter(ProviderRoutingHints) @@ -7320,6 +7321,9 @@ def embedding( logging_obj=logging, optional_params=optional_params, model_response=EmbeddingResponse(), + api_key=api_key, + headers=_OLLAMA_EMBEDDING_HEADERS_ADAPTER.validate_python(headers), + timeout=timeout, ) elif custom_llm_provider == "sagemaker": response = sagemaker_llm.embedding( diff --git a/tests/integration/_support/wire.py b/tests/integration/_support/wire.py index df6036ff8d7..e278a796146 100644 --- a/tests/integration/_support/wire.py +++ b/tests/integration/_support/wire.py @@ -19,6 +19,7 @@ class Request: target: str headers: Mapping[str, str] body: bytes + header_lines: tuple[tuple[str, str], ...] = () class AbortedBody(Exception): @@ -136,7 +137,7 @@ def wire_server( self.close_connection = True disconnected.put(self.path) return - request: Final = Request(self.command, self.path, headers, body) + request: Final = Request(self.command, self.path, headers, body, tuple(self.headers.items())) received.put(request) try: reply = respond(request) diff --git a/tests/integration/providers/test_api_base_trailing_slash_wire.py b/tests/integration/providers/test_api_base_trailing_slash_wire.py index e914199e825..01b68f52032 100644 --- a/tests/integration/providers/test_api_base_trailing_slash_wire.py +++ b/tests/integration/providers/test_api_base_trailing_slash_wire.py @@ -42,7 +42,7 @@ _MESSAGES_ROWS: Final = ( Row("ollama_chat/llama3.1", "/ol", "/api/chat"), ) _EMBEDDING_ROWS: Final = ( - Row("ollama/nomic-embed-text", "/ol", "/api/embed", keyed=False), + Row("ollama/nomic-embed-text", "/ol", "/api/embed"), Row("jina_ai/jina-embeddings-v3", "/jn/v1", "/embeddings"), Row("voyage/voyage-3", "/vy/v1", "/embeddings"), Row("voyage/voyage-multimodal-3", "/vy/v1", "/multimodalembeddings"), @@ -65,6 +65,7 @@ _FULL_ROUTE_EMBEDDING_ROWS: Final = ( Row("jina_ai/jina-embeddings-v3", "/jn/v1/embeddings", ""), Row("voyage/voyage-3", "/vy/v1/embeddings", ""), Row("perplexity/pplx-embed-v1-0.6b", "/px/v1/embeddings", ""), + Row("ollama/nomic-embed-text", "/ol/api/embed", ""), ) diff --git a/tests/integration/providers/test_ollama_embedding_wire.py b/tests/integration/providers/test_ollama_embedding_wire.py new file mode 100644 index 00000000000..fd38260b21c --- /dev/null +++ b/tests/integration/providers/test_ollama_embedding_wire.py @@ -0,0 +1,134 @@ +from __future__ import annotations + +import json +import threading +import time +import uuid +from collections.abc import Mapping +from typing import Final + +import httpx +import pytest +from integration._support.client import JSON_OBJECT, Gateway, list_value, object_value +from integration._support.timing import assert_cut_short +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +_MODEL: Final = "nomic-embed-text" +_DEPLOYMENT_KEY: Final = "synthetic-ollama-deployment-key" +_CALLER_TOKEN: Final = "Bearer synthetic-deployment-header-token" +_DELAY_SECONDS: Final = 5.0 + + +def _vector(index: int) -> list[float]: + return [0.5, float(index)] + + +def _embeddings(request: Request) -> Reply: + body: Final = JSON_OBJECT.validate_json(request.body) + count: Final = len(body["input"]) if isinstance(body["input"], list) else 1 + return Reply( + body=json.dumps({"embeddings": [_vector(index) for index in range(count)], "prompt_eval_count": 3}).encode() + ) + + +def _vectors(response: httpx.Response) -> list[JsonValue]: + data: Final = list_value(JSON_OBJECT.validate_json(response.content)["data"]) + return [object_value(row)["embedding"] for row in data] + + +def _only(wire: Wire) -> Request: + received: Final = wire.drain() + assert [(request.method, request.target) for request in received] == [("POST", "/api/embed")], [ + (request.method, request.target) for request in received + ] + return received[0] + + +def _values_of(request: Request, name: str) -> list[str]: + return [value for header, value in request.header_lines if header.lower() == name.lower()] + + +def _texts(count: int) -> list[str]: + return [f"ollama proxy embedding {uuid.uuid4().hex}" for _ in range(count)] + + +@pytest.mark.parametrize( + ("route", "batched"), (("/v1/embeddings", False), ("/embeddings", True)), ids=("v1-string", "alias-batch") +) +@pytest.mark.parametrize( + ("credentials", "authorization"), + ( + ({"api_key": _DEPLOYMENT_KEY}, f"Bearer {_DEPLOYMENT_KEY}"), + ( + {"api_key": _DEPLOYMENT_KEY, "extra_headers": {"authorization": _CALLER_TOKEN, "X-Embed-Tenant": "t-1"}}, + _CALLER_TOKEN, + ), + ), + ids=("api-key", "lowercase-header-overrides-key"), +) +def test_regression_proxy_ollama_embedding_sends_deployment_credentials( + gateway: Gateway, route: str, batched: bool, credentials: Mapping[str, JsonValue], authorization: str +) -> None: + texts: Final = _texts(2 if batched else 1) + with wire_server(_embeddings) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_MODEL}", api_base=wire.url, **credentials) + response: Final = gateway.request("POST", route, {"model": model, "input": texts if batched else texts[0]}) + request: Final = _only(wire) + assert response.status_code == 200, response.text + assert _vectors(response) == [_vector(i) for i in range(len(texts))] + assert _values_of(request, "authorization") == [authorization], request.header_lines + assert JSON_OBJECT.validate_json(request.body) == {"model": _MODEL, "input": texts}, request.body + expected_tenant: Final = ["t-1"] if "extra_headers" in credentials else [] + assert _values_of(request, "x-embed-tenant") == expected_tenant, request.header_lines + + +@pytest.mark.parametrize("caller_timeout", ({}, {"timeout": 0}), ids=("omitted", "zero")) +def test_regression_proxy_ollama_deployment_timeout_cuts_a_held_read( + gateway: Gateway, caller_timeout: Mapping[str, JsonValue] +) -> None: + release: Final = threading.Event() + + def respond(request: Request) -> Reply: + release.wait(timeout=_DELAY_SECONDS) + return _embeddings(request) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_MODEL}", api_base=wire.url, timeout=0.5) + started: Final = time.monotonic() + try: + response: Final = gateway.request( + "POST", "/v1/embeddings", {"model": model, "input": _texts(1)[0], **caller_timeout} + ) + finally: + elapsed: Final = time.monotonic() - started + release.set() + _only(wire) + assert response.status_code == 408, response.text + assert_cut_short(elapsed, injected_delay=_DELAY_SECONDS) + + +def test_guard_proxy_keyless_ollama_embedding_sends_no_authorization(gateway: Gateway) -> None: + texts: Final = _texts(2) + with wire_server(_embeddings) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_MODEL}", api_base=f"{wire.url}/", api_key=None) + response: Final = gateway.request("POST", "/v1/embeddings", {"model": model, "input": texts}) + request: Final = _only(wire) + assert response.status_code == 200, response.text + assert _vectors(response) == [_vector(0), _vector(1)] + assert _values_of(request, "authorization") == [], request.header_lines + assert JSON_OBJECT.validate_json(request.body) == {"model": _MODEL, "input": texts}, request.body + + +def test_guard_proxy_ollama_upstream_error_reaches_the_caller(gateway: Gateway) -> None: + marker: Final = f"scripted-bad-request-{uuid.uuid4().hex}" + + def respond(_: Request) -> Reply: + return Reply(status=400, body=json.dumps({"error": marker}).encode()) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_MODEL}", api_base=wire.url) + response: Final = gateway.request("POST", "/v1/embeddings", {"model": model, "input": _texts(1)[0]}) + _only(wire) + assert response.status_code == 400, response.text + assert marker in response.text, response.text diff --git a/tests/integration/sdk/test_ollama_embedding_wire.py b/tests/integration/sdk/test_ollama_embedding_wire.py new file mode 100644 index 00000000000..21b803b8728 --- /dev/null +++ b/tests/integration/sdk/test_ollama_embedding_wire.py @@ -0,0 +1,218 @@ +import json +import threading +import time +import uuid +from collections.abc import Callable, Iterator, Mapping +from contextlib import contextmanager +from typing import Final, Literal + +import httpx +import pytest +from integration._support.client import JSON_OBJECT +from integration._support.timing import assert_cut_short +from integration._support.wire import Reply, Request, Wire, wire_server + +import litellm +from litellm.types.utils import EmbeddingResponse + +_MODEL: Final = "nomic-embed-text" +_EXPLICIT_KEY: Final = "synthetic-ollama-explicit-key" +_ENV_KEY: Final = "synthetic-ollama-env-key" +_GLOBAL_KEY: Final = "synthetic-global-key-never-sent" +_CALLER_TOKEN: Final = "Bearer synthetic-caller-token" +_DELAY_SECONDS: Final = 5.0 +_SHORT_DELAY_SECONDS: Final = 1.0 + +Mode = Literal["sync", "async"] +_MODES: Final[tuple[Mode, ...]] = ("sync", "async") + + +def _vector(index: int) -> list[float]: + return [0.25, float(index)] + + +def _embeddings(request: Request) -> Reply: + body: Final = JSON_OBJECT.validate_json(request.body) + count: Final = len(body["input"]) if isinstance(body["input"], list) else 1 + return Reply( + body=json.dumps({"embeddings": [_vector(index) for index in range(count)], "prompt_eval_count": 3}).encode() + ) + + +def _held(release: threading.Event, seconds: float) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + release.wait(timeout=seconds) + return _embeddings(request) + + return respond + + +@pytest.fixture(autouse=True) +def _no_ambient_ollama_settings(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("OLLAMA_API_KEY", raising=False) + monkeypatch.delenv("OLLAMA_API_BASE", raising=False) + + +@contextmanager +def _peer(respond: Callable[[Request], Reply] = _embeddings) -> Iterator[Wire]: + with wire_server(respond) as wire: + yield wire + + +async def _call(mode: Mode, **kwargs: object) -> EmbeddingResponse: + if mode == "sync": + return litellm.embedding(model=f"ollama/{_MODEL}", **kwargs) + return await litellm.aembedding(model=f"ollama/{_MODEL}", **kwargs) + + +async def _seconds_until_timeout(mode: Mode, **kwargs: object) -> float: + started: Final = time.monotonic() + with pytest.raises(litellm.Timeout): + await _call(mode, **kwargs) + return time.monotonic() - started + + +def _only(wire: Wire) -> Request: + received: Final = wire.drain() + assert [(request.method, request.target) for request in received] == [("POST", "/api/embed")], [ + (request.method, request.target) for request in received + ] + return received[0] + + +def _values_of(request: Request, name: str) -> list[str]: + return [value for header, value in request.header_lines if header.lower() == name.lower()] + + +def _text() -> str: + return f"ollama embedding {uuid.uuid4().hex}" + + +@pytest.mark.parametrize("mode", _MODES) +@pytest.mark.parametrize( + ("explicit", "env", "expected"), + ( + (_EXPLICIT_KEY, None, _EXPLICIT_KEY), + (None, _ENV_KEY, _ENV_KEY), + (_EXPLICIT_KEY, _ENV_KEY, _EXPLICIT_KEY), + ), + ids=("explicit-key", "env-key", "explicit-beats-env"), +) +async def test_regression_embedding_sends_the_resolved_ollama_key_as_one_bearer_header( + monkeypatch: pytest.MonkeyPatch, mode: Mode, explicit: str | None, env: str | None, expected: str +) -> None: + if env is not None: + monkeypatch.setenv("OLLAMA_API_KEY", env) + texts: Final = [_text(), _text()] + with _peer() as wire: + response: Final = await _call(mode, input=texts, api_base=wire.url, api_key=explicit) + request: Final = _only(wire) + assert [row["embedding"] for row in response.data] == [_vector(0), _vector(1)], response + assert _values_of(request, "authorization") == [f"Bearer {expected}"], request.header_lines + assert JSON_OBJECT.validate_json(request.body) == {"model": _MODEL, "input": texts}, request.body + + +@pytest.mark.parametrize("mode", _MODES) +@pytest.mark.parametrize("auth_name", ("Authorization", "authorization"), ids=("canonical", "lowercase")) +@pytest.mark.parametrize("header_kwarg", ("extra_headers", "headers")) +async def test_regression_caller_headers_reach_ollama_and_override_the_key_once( + mode: Mode, auth_name: str, header_kwarg: str +) -> None: + caller: Final[dict[str, str]] = {auth_name: _CALLER_TOKEN, "X-Embed-Tenant": "tenant-a"} + snapshot: Final = dict(caller) + text: Final = _text() + with _peer() as wire: + response: Final = await _call( + mode, input=text, api_base=wire.url, api_key=_EXPLICIT_KEY, **{header_kwarg: caller} + ) + request: Final = _only(wire) + assert [row["embedding"] for row in response.data] == [_vector(0)], response + assert _values_of(request, "authorization") == [_CALLER_TOKEN], request.header_lines + assert _values_of(request, "x-embed-tenant") == ["tenant-a"], request.header_lines + assert caller == snapshot, "the caller's header map was changed by the call" + assert JSON_OBJECT.validate_json(request.body) == {"model": _MODEL, "input": [text]}, request.body + + +@pytest.mark.parametrize("mode", _MODES) +async def test_regression_extra_headers_win_over_headers_and_both_reach_ollama(mode: Mode) -> None: + extra: Final[dict[str, str]] = {"X-Embed-Tenant": "from-extra-headers"} + extra_snapshot: Final = dict(extra) + with _peer() as wire: + await _call( + mode, + input=_text(), + api_base=wire.url, + headers={"X-Embed-Tenant": "from-headers", "X-Embed-Trace": "trace-1"}, + extra_headers=extra, + ) + request: Final = _only(wire) + assert _values_of(request, "x-embed-tenant") == ["from-extra-headers"], request.header_lines + assert _values_of(request, "x-embed-trace") == ["trace-1"], request.header_lines + assert extra == extra_snapshot, "the caller's extra_headers map was changed by the call" + + +@pytest.mark.parametrize("mode", _MODES) +@pytest.mark.parametrize("timeout", (0.5, httpx.Timeout(0.5)), ids=("float", "httpx-timeout")) +async def test_regression_short_timeout_cuts_a_held_ollama_read(mode: Mode, timeout: float | httpx.Timeout) -> None: + release: Final = threading.Event() + with _peer(_held(release, _DELAY_SECONDS)) as wire: + try: + elapsed: Final = await _seconds_until_timeout(mode, input=_text(), api_base=wire.url, timeout=timeout) + finally: + release.set() + _only(wire) + assert_cut_short(elapsed, injected_delay=_DELAY_SECONDS) + + +@pytest.mark.parametrize("mode", _MODES) +@pytest.mark.parametrize("timeout_kwargs", ({}, {"timeout": None}), ids=("omitted", "none")) +async def test_guard_default_timeout_does_not_cut_a_slow_ollama_reply( + mode: Mode, timeout_kwargs: Mapping[str, None] +) -> None: + release: Final = threading.Event() + with _peer(_held(release, _SHORT_DELAY_SECONDS)) as wire: + response: Final = await _call(mode, input=_text(), api_base=wire.url, **timeout_kwargs) + _only(wire) + assert [row["embedding"] for row in response.data] == [_vector(0)], response + + +@pytest.mark.parametrize("mode", _MODES) +async def test_guard_no_ollama_key_sends_no_authorization_even_with_global_keys( + monkeypatch: pytest.MonkeyPatch, mode: Mode +) -> None: + monkeypatch.setattr(litellm, "api_key", _GLOBAL_KEY) + monkeypatch.setattr(litellm, "ollama_key", _GLOBAL_KEY, raising=False) + monkeypatch.setattr(litellm, "openai_key", _GLOBAL_KEY) + texts: Final = [_text(), _text()] + with _peer() as wire: + response: Final = await _call(mode, input=texts, api_base=f"{wire.url}/") + request: Final = _only(wire) + assert [row["embedding"] for row in response.data] == [_vector(0), _vector(1)], response + assert _values_of(request, "authorization") == [], request.header_lines + assert not any(_GLOBAL_KEY in value for _, value in request.header_lines), request.header_lines + assert JSON_OBJECT.validate_json(request.body) == {"model": _MODEL, "input": texts}, request.body + + +@pytest.mark.parametrize("mode", _MODES) +async def test_guard_non_text_input_is_rejected_before_ollama_is_called(mode: Mode) -> None: + valid: Final = _text() + with _peer() as wire: + with pytest.raises(litellm.BadRequestError): + await _call(mode, input=[[1, 2, 3]], api_base=wire.url, api_key=_EXPLICIT_KEY) + await _call(mode, input=valid, api_base=wire.url) + received: Final = wire.drain() + assert [JSON_OBJECT.validate_json(request.body)["input"] for request in received] == [[valid]] + + +@pytest.mark.parametrize("mode", _MODES) +async def test_guard_ollama_error_reaches_the_caller_as_its_status(mode: Mode) -> None: + marker: Final = f"scripted-unauthorized-{uuid.uuid4().hex}" + + def respond(_: Request) -> Reply: + return Reply(status=401, body=json.dumps({"error": marker}).encode()) + + with _peer(respond) as wire: + with pytest.raises(litellm.AuthenticationError) as caught: + await _call(mode, input=_text(), api_base=wire.url) + _only(wire) + assert marker in str(caught.value), caught.value diff --git a/tests/unit/llms/ollama/test_ollama_chat_transformation.py b/tests/unit/llms/ollama/test_ollama_chat_transformation.py index 1a0bbe171a9..df9f5fd9d14 100644 --- a/tests/unit/llms/ollama/test_ollama_chat_transformation.py +++ b/tests/unit/llms/ollama/test_ollama_chat_transformation.py @@ -1064,6 +1064,9 @@ def test_ollama_embeddings(mock_embeddings): logging_obj=mock.ANY, model_response=mock.ANY, encoding=mock.ANY, + api_key=None, + headers={}, + timeout=600, ) except Exception as e: pytest.fail(f"Error occurred: {e}") @@ -1088,6 +1091,9 @@ def test_ollama_aembeddings(mock_aembeddings): logging_obj=mock.ANY, model_response=mock.ANY, encoding=mock.ANY, + api_key=None, + headers={}, + timeout=600, ) except Exception as e: pytest.fail(f"Error occurred: {e}") diff --git a/tests/unit/llms/ollama/test_ollama_embedding.py b/tests/unit/llms/ollama/test_ollama_embedding.py index 6a1dfebc236..97f6f8468f4 100644 --- a/tests/unit/llms/ollama/test_ollama_embedding.py +++ b/tests/unit/llms/ollama/test_ollama_embedding.py @@ -1,3 +1,6 @@ +import json +from collections.abc import Mapping +from types import MappingProxyType from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -127,12 +130,183 @@ def test_prompt_eval_fallback_when_missing(mock_embedding_response, mock_encodin ], ) @respx.mock -def test_ollama_embeddings_ignore_a_trailing_slash_on_api_base(api_base: str) -> None: +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_ollama_embeddings_ignore_a_trailing_slash_on_api_base( + api_base: str, asynchronous: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") route: Final = respx.post("http://gateway.example.com/ollama/api/embed").mock( return_value=httpx.Response(200, json={"embeddings": [[0.1, 0.2], [0.3, 0.4]], "prompt_eval_count": 2}) ) - response: Final = litellm.embedding(model="ollama/all-minilm", api_base=api_base, input=["hello", "world"]) + response: Final = ( + await litellm.aembedding(model="ollama/all-minilm", api_base=api_base, input=["hello", "world"]) + if asynchronous + else litellm.embedding(model="ollama/all-minilm", api_base=api_base, input=["hello", "world"]) + ) - assert route.called + assert route.call_count == 1 assert [item["embedding"] for item in response.data] == [[0.1, 0.2], [0.3, 0.4]] + + +@pytest.fixture +def embedding_route(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch) -> respx.Route: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + return respx_mock.post("http://ollama.test/api/embed").respond( + 200, json={"embeddings": [[0.1, 0.2]], "prompt_eval_count": 2} + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize( + ("api_key", "env_key", "authorization"), + [ + ("explicit-key", None, "Bearer explicit-key"), + (None, "env-key", "Bearer env-key"), + ("explicit-key", "env-key", "Bearer explicit-key"), + (None, None, None), + ], +) +async def test_sdk_embedding_forwards_resolved_auth_and_merged_headers( + embedding_route: respx.Route, + monkeypatch: pytest.MonkeyPatch, + asynchronous: bool, + api_key: str | None, + env_key: str | None, + authorization: str | None, +) -> None: + if env_key is None: + monkeypatch.delenv("OLLAMA_API_KEY", raising=False) + else: + monkeypatch.setenv("OLLAMA_API_KEY", env_key) + response: Final = ( + await litellm.aembedding( + model="ollama/nomic-embed-text", + input="hello", + api_base="http://ollama.test", + api_key=api_key, + headers={"X-Caller": "caller", "X-Shared": "original"}, + extra_headers={"X-Extra": "extra", "X-Shared": "override"}, + ) + if asynchronous + else litellm.embedding( + model="ollama/nomic-embed-text", + input="hello", + api_base="http://ollama.test", + api_key=api_key, + headers={"X-Caller": "caller", "X-Shared": "original"}, + extra_headers={"X-Extra": "extra", "X-Shared": "override"}, + ) + ) + request: Final = embedding_route.calls.last.request + + assert embedding_route.call_count == 1 + assert request.headers.get("Authorization") == authorization + assert (request.headers["X-Caller"], request.headers["X-Extra"], request.headers["X-Shared"]) == ( + "caller", + "extra", + "override", + ) + assert request.extensions["timeout"] == httpx.Timeout(600).as_dict() + assert json.loads(request.content) == {"model": "nomic-embed-text", "input": ["hello"]} + assert isinstance(response, EmbeddingResponse) + assert response.data[0]["embedding"] == [0.1, 0.2] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("timeout", [0, 2.5, httpx.Timeout(connect=1, read=2, write=3, pool=4)]) +async def test_sdk_embedding_applies_timeout_to_http_request( + embedding_route: respx.Route, asynchronous: bool, timeout: float | httpx.Timeout +) -> None: + response: Final = ( + await litellm.aembedding( + model="ollama/nomic-embed-text", input=["hello"], api_base="http://ollama.test", timeout=timeout + ) + if asynchronous + else litellm.embedding( + model="ollama/nomic-embed-text", input=["hello"], api_base="http://ollama.test", timeout=timeout + ) + ) + + assert embedding_route.call_count == 1 + assert embedding_route.calls.last.request.extensions["timeout"] == httpx.Timeout(timeout).as_dict() + assert isinstance(response, EmbeddingResponse) + assert response.data[0]["embedding"] == [0.1, 0.2] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize( + ("api_key", "headers", "authorization"), + [ + ("explicit-key", {"X-Custom": "custom"}, "Bearer explicit-key"), + ("explicit-key", {"Authorization": "Basic caller"}, "Basic caller"), + ("explicit-key", {"authorization": "Bearer caller"}, "Bearer caller"), + (None, {}, None), + ("", {}, None), + ], +) +@pytest.mark.parametrize("timeout_params", [{}, {"timeout": None}], ids=["omitted", "none"]) +async def test_handler_embedding_preserves_caller_headers_and_client_timeout( + embedding_route: respx.Route, + asynchronous: bool, + api_key: str | None, + headers: Mapping[str, str], + authorization: str | None, + timeout_params: Mapping[str, None], +) -> None: + caller_headers: Final = MappingProxyType(dict(headers)) + expected_timeout: Final = ( + litellm.module_level_aclient.client.timeout if asynchronous else litellm.module_level_client.client.timeout + ) + response: Final = ( + await ollama_aembeddings( + "http://ollama.test", + "nomic-embed-text", + ["hello"], + EmbeddingResponse(), + {}, + None, + None, + api_key=api_key, + headers=caller_headers, + **timeout_params, + ) + if asynchronous + else ollama_embeddings( + "http://ollama.test", + "nomic-embed-text", + ["hello"], + {}, + EmbeddingResponse(), + None, + api_key=api_key, + headers=caller_headers, + **timeout_params, + ) + ) + request: Final = embedding_route.calls.last.request + + assert embedding_route.call_count == 1 + assert request.headers.get_list("Authorization") == ([] if authorization is None else [authorization]) + assert dict(caller_headers) == headers + assert all(request.headers[name] == value for name, value in caller_headers.items()) + assert request.extensions["timeout"] == expected_timeout.as_dict() + assert json.loads(request.content) == {"model": "nomic-embed-text", "input": ["hello"]} + assert response.data[0]["embedding"] == [0.1, 0.2] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_sdk_embedding_rejects_non_string_batch(embedding_route: respx.Route, asynchronous: bool) -> None: + if asynchronous: + with pytest.raises(litellm.BadRequestError, match="Invalid input for ollama embeddings"): + await litellm.aembedding(model="ollama/nomic-embed-text", input=[123], api_base="http://ollama.test") + else: + with pytest.raises(litellm.BadRequestError, match="Invalid input for ollama embeddings"): + litellm.embedding(model="ollama/nomic-embed-text", input=[123], api_base="http://ollama.test") + + assert not embedding_route.called From 97a13ce6763126f127ade1c060492ce93a61c8e9 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 28 Oct 2026 13:00:48 -0400 Subject: [PATCH 047/128] fix(router): isolate child tier pins in user turn mode (#48701) --- litellm/constants.py | 1 + litellm/proxy/litellm_pre_call_utils.py | 57 +++++- .../complexity_router/README.md | 6 + .../complexity_router/complexity_router.py | 24 ++- .../unit/proxy/test_litellm_pre_call_utils.py | 105 +++++++++++ .../router_strategy/test_complexity_router.py | 170 ++++++++++++++++++ 6 files changed, 352 insertions(+), 11 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 0d6dfb5c0d4..7c6d6deba66 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1669,6 +1669,7 @@ MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: Final = "_managed_background_response" SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated" SESSION_ID_OMITTED_METADATA_KEY: Final = "litellm_session_id_omitted" CALLER_SESSION_ID_METADATA_KEY: Final = "litellm_caller_session_id" +COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY: Final = "_complexity_router_child_session_id" LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated" LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = ( "Truncation is a DB storage safeguard. " diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 314f14b6279..24a31723faa 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -32,6 +32,7 @@ from litellm.constants import ( BILLED_RETRY_ATTEMPT_METADATA_KEY, CALLER_SESSION_ID_METADATA_KEY, CLIENT_OUTPUT_CEILING_METADATA_KEY, + COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY, CONSUMED_REQUEST_TAGS_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, LITELLM_PROXY_MASTER_KEY_ALIAS, @@ -210,13 +211,46 @@ def _caller_session_beside_trace_header( return get_chain_id_from_headers(session_headers) or _caller_body_session_id(data, metadata_variable_name) -def _drop_client_caller_session(request_body: Mapping[str, object]) -> None: +def _child_tier_session_metadata( + normalized_headers: Mapping[str, str], data: Mapping[str, object], metadata_variable_name: str +) -> Mapping[str, str | None]: + parent: Final = normalized_headers.get("x-parent-session-id") + if not parent or not _SESSION_ID_VALUE_RE.fullmatch(parent): + return {} + own_headers: Final = { + key: value + for key, value in normalized_headers.items() + if key not in ("x-parent-session-id", "x-litellm-trace-id") + } + active: Final = data.get(metadata_variable_name) + inactive: Final = data.get("litellm_metadata") + inactive_session: Final = ( + _caller_body_session_id(data, "litellm_metadata") + if metadata_variable_name == "metadata" + and not (isinstance(active, Mapping) and "session_id" in active) + and isinstance(inactive, Mapping) + else None + ) + own_session: Final = ( + get_chain_id_from_headers(own_headers) + or _get_anthropic_session_id_from_metadata(data.get("metadata")) + or _caller_body_session_id(data, metadata_variable_name) + or inactive_session + or _session_id_from_baggage(normalized_headers.get("baggage", "")) + ) + return {COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY: own_session if own_session != parent else None} + + +def _drop_client_session_metadata(request_body: Mapping[str, object]) -> None: for metadata_key in ("metadata", "litellm_metadata"): if isinstance(metadata := request_body.get(metadata_key), dict): - metadata.pop(CALLER_SESSION_ID_METADATA_KEY, None) + for key in (CALLER_SESSION_ID_METADATA_KEY, COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY): + cast(dict[str, object], metadata).pop(key, None) # cast-ok: request metadata is a str-keyed JSON object -_PROXY_SET_METADATA_KEYS: Final = frozenset({CALLER_SESSION_ID_METADATA_KEY, BAGGAGE_MAPPED_TAGS_METADATA_KEY}) +_PROXY_SET_METADATA_KEYS: Final = frozenset( + {CALLER_SESSION_ID_METADATA_KEY, BAGGAGE_MAPPED_TAGS_METADATA_KEY, COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY} +) _METADATA_TAGS_ADAPTER: Final = TypeAdapter(tuple[object, ...]) @@ -1855,7 +1889,7 @@ class LiteLLMProxyRequestSetup: metadata_from_headers: Final = LitellmMetadataFromRequestHeaders() request_body: Final = cast(Mapping[str, object], data) # cast-ok: request body is a str-keyed JSON object - _drop_client_caller_session(request_body) + _drop_client_session_metadata(request_body) spend_logs_metadata: Final = LiteLLMProxyRequestSetup._get_spend_logs_metadata_from_request_headers(headers) if spend_logs_metadata is not None: metadata_from_headers["spend_logs_metadata"] = spend_logs_metadata @@ -1870,6 +1904,9 @@ class LiteLLMProxyRequestSetup: agent_id_from_header: Final = headers.get("x-litellm-agent-id") # Explicit litellm headers take precedence; fall back to any x-*-session-id header. chain_id: Final = get_chain_id_from_headers(dict(headers)) + child_tier_metadata: Final = _child_tier_session_metadata( + normalized_headers, request_body, _metadata_variable_name + ) if agent_id_from_header: metadata_from_headers["agent_id"] = agent_id_from_header @@ -1939,9 +1976,15 @@ class LiteLLMProxyRequestSetup: if isinstance(data[_metadata_variable_name], dict): data[_metadata_variable_name].update( - {**metadata_from_headers, CALLER_SESSION_ID_METADATA_KEY: caller_session_beside_trace} - if caller_session_beside_trace - else metadata_from_headers + { + **metadata_from_headers, + **child_tier_metadata, + **( + {CALLER_SESSION_ID_METADATA_KEY: caller_session_beside_trace} + if caller_session_beside_trace + else {} + ), + } ) return data diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index 3df1803d264..d06248a555f 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -477,6 +477,12 @@ unless `modality_routing` is also on. ### Session pin retention +With `classification_mode: user_turn`, a child request carrying `x-parent-session-id` classifies +its own new asks and stores the tier under its own session ID. Its continuations reuse that tier +without replacing the parent's tier. Existing trace grouping, per-tier model affinity and deployment +affinity retain their session identities. A child without a distinct, usable session ID classifies +without storing a tier pin. `session_affinity: true` retains precedence and the existing shared-pin behavior + `session_affinity_ttl_seconds` is the idle window for both the model pin selected by session affinity and the deployment pin. Every request that reuses a pin refreshes its TTL, so a session actively sending requests stays pinned. After the window passes with no pin reuse, the next request classifies again and creates a fresh pin. Omit the setting to track the default of 3600 seconds. ### Mid-task stall escalation diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 82aaa221d31..2b0187fd667 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -37,6 +37,7 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger from litellm.caching.affinity_cache import ROUTER_SESSION_PINS_TARGET, claim_affinity_pin, read_affinity_pin from litellm.constants import ( + COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY, EMPTY_MAPPING, INTERNAL_CALL_ORIGIN_METADATA_KEY, OUTPUT_TOKEN_CEILING_PARAMS, @@ -4358,17 +4359,17 @@ class ComplexityRouter(CustomLogger): return newest_ask is not None and newest_ask.startswith(_CODEX_COMPACTION_SUMMARY_PREFIX) @staticmethod - def _iter_metadata_dicts(request_kwargs: dict) -> list[dict]: + def _iter_metadata_dicts(request_kwargs: Mapping[str, object]) -> list[Mapping[str, object]]: """Metadata may land on `metadata` or `litellm_metadata` depending on the endpoint, mirroring DeploymentAffinityCheck's precedence.""" return [ - metadata + cast(Mapping[str, object], metadata) # cast-ok: request metadata is a str-keyed JSON object for metadata_key in ("litellm_metadata", "metadata") if isinstance(metadata := request_kwargs.get(metadata_key), dict) ] @staticmethod - def _get_session_id_from_request_kwargs(request_kwargs: dict) -> str | None: + def _get_session_id_from_request_kwargs(request_kwargs: Mapping[str, object]) -> str | None: """Resolve a client-supplied session_id.""" for metadata in ComplexityRouter._iter_metadata_dicts(request_kwargs): session_id = metadata.get("session_id") @@ -4376,6 +4377,21 @@ class ComplexityRouter(CustomLogger): return str(session_id) return None + def _get_tier_pin_session_id(self, request_kwargs: Mapping[str, object]) -> str | None: + if self.config.classification_mode == "user_turn" and not self.config.session_affinity: + child_metadata: Final = next( + ( + metadata + for metadata in self._iter_metadata_dicts(request_kwargs) + if COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY in metadata + ), + None, + ) + if child_metadata is not None: + child_session: Final = child_metadata[COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY] + return child_session if isinstance(child_session, str) and child_session else None + return self._get_session_id_from_request_kwargs(request_kwargs) + @staticmethod def _get_user_api_key_hash_from_request_kwargs(request_kwargs: dict) -> str | None: """Resolve the proxy-derived API key hash, the same trust boundary @@ -4460,7 +4476,7 @@ class ComplexityRouter(CustomLogger): conversation_continuing: Final = _conversation_is_continuing(resolved_messages) use_session_affinity: Final = self._uses_tier_pin - session_id: Final = self._get_session_id_from_request_kwargs(request_kwargs) if use_session_affinity else None + session_id: Final = self._get_tier_pin_session_id(request_kwargs) if use_session_affinity else None cache_key = self._get_session_affinity_cache_key(session_id, request_kwargs) if session_id is not None else None # In 'user_turn' mode a held pin is replayed only on continuation turns; a new human diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index a32c40e5a43..81a196af676 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -19,6 +19,7 @@ from starlette.datastructures import Headers import litellm from litellm.constants import ( CALLER_SESSION_ID_METADATA_KEY, + COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, SERVER_STREAMING_CLASSIFICATION_KEY, SERVER_STREAMING_CLASSIFICATION_MARKER, @@ -3684,6 +3685,110 @@ async def test_spend_log_caller_session_stays_out_of_requester_metadata(monkeypa assert CALLER_SESSION_ID_METADATA_KEY not in updated["metadata"]["requester_metadata"] +@pytest.mark.asyncio +@pytest.mark.parametrize( + "surface,headers,body,child_id,trace_id", + [ + ( + "chat/completions", {"X-Session-Id": "child-01"}, {"metadata": {"session_id": "body-001"}}, + "child-01", "parent-01", + ), + ("responses", {"x-vendor-session-id": "child-01"}, {}, "child-01", "parent-01"), + ("messages", {"session-id": "child-01", "user-agent": "codex_cli_rs/1.0"}, {}, "child-01", "parent-01"), + ("chat/completions", {"x-litellm-session-id": "short", "x-session-id": "child-01"}, {}, "short", "short"), + ("responses", {"x-litellm-trace-id": "trace-01", "x-session-id": "child-01"}, {}, "child-01", "trace-01"), + ( + "messages", {"x-litellm-trace-id": "trace-01", "x-litellm-session-id": "child-01"}, + {}, "child-01", "trace-01", + ), + ("responses", {}, {"metadata": {"session_id": "body-001"}}, "body-001", "parent-01"), + ("messages", {}, {"litellm_metadata": {"session_id": "body-001"}}, "body-001", "parent-01"), + ("chat/completions", {}, {"litellm_metadata": {"session_id": "body-001"}}, "body-001", "parent-01"), + ("messages", {"baggage": "session.id=child-01"}, {}, "child-01", "parent-01"), + ( + "responses", {"baggage": "session.id=child-01"}, {"metadata": {"session_id": "body-001"}}, + "body-001", "parent-01", + ), + ( + "chat/completions", {}, {"litellm_session_id": "root-001", "metadata": {"session_id": "body-001"}}, + "root-001", "parent-01", + ), + ( + "responses", {}, {"metadata": {"session_id": "body-001"}, "litellm_metadata": {"session_id": "active-1"}}, + "active-1", "parent-01", + ), + ( + "chat/completions", {}, + {"metadata": {"session_id": "active-1"}, "litellm_metadata": {"session_id": "body-001"}}, + "active-1", "parent-01", + ), + ( + "chat/completions", {}, {"metadata": {"session_id": ""}, "litellm_metadata": {"session_id": "body-001"}}, + None, "parent-01", + ), + ( + "messages", {}, {"metadata": {"session_id": "body-001"}, "litellm_metadata": {"session_id": ""}}, + None, "parent-01", + ), + ( + "messages", {}, {"litellm_session_id": "root-001", "litellm_metadata": {"session_id": "active-1"}, + "metadata": {"session_id": "body-001", "user_id": "user_abc_account__session_anthro-1"}}, + "anthro-1", "parent-01", + ), + ("messages", {"x-session-id": "child-01"}, {"metadata": {"user_id": {"session_id": "body-001"}}}, + "child-01", "parent-01"), + ("chat/completions", {}, {}, None, "parent-01"), + ("responses", {"x-session-id": "bad"}, {}, None, "parent-01"), + ("messages", {}, {"metadata": {"session_id": 123}}, None, "parent-01"), + ("chat/completions", {"x-session-id": "parent-01"}, {}, None, "parent-01"), + ], +) +async def test_child_tier_identity_preserves_trace_and_body_session_precedence( + surface: str, headers: dict[str, str], body: dict[str, object], child_id: str | None, trace_id: str +) -> None: + request: Final = _make_request_mock(f"/v1/{surface}", {"x-parent-session-id": "parent-01", **headers}) + updated: Final = await add_litellm_data_to_request( + data={"model": "test-router", "messages": [], **copy.deepcopy(body)}, + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + ) + metadata: Final = updated[get_metadata_variable_name(request)] + assert metadata[COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY] == child_id + assert metadata["session_id"] == metadata["trace_id"] == trace_id + assert updated["litellm_session_id"] == updated["litellm_trace_id"] == trace_id + expected_caller: Final = ( + headers.get("x-litellm-session-id", "parent-01") if "x-litellm-trace-id" in headers else None + ) + assert metadata.get(CALLER_SESSION_ID_METADATA_KEY) == expected_caller + assert COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY not in metadata.get("requester_metadata", {}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("surface,parent_id", [("chat/completions", None), ("responses", None), ("messages", "bad")]) +async def test_non_child_requests_cannot_forge_a_child_tier_identity(surface: str, parent_id: str | None) -> None: + headers: Final = {"x-session-id": "child-01", **({"x-parent-session-id": parent_id} if parent_id else {})} + request: Final = _make_request_mock(f"/v1/{surface}", headers) + updated: Final = await add_litellm_data_to_request( + data={ + "model": "test-router", + "messages": [], + "metadata": {COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY: "forged-1"}, + "litellm_metadata": {COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY: "forged-2"}, + }, + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + ) + for bucket in ("metadata", "litellm_metadata"): + metadata: Final = updated.get(bucket, {}) + assert COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY not in metadata + assert COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY not in metadata.get("requester_metadata", {}) + assert updated[get_metadata_variable_name(request)]["session_id"] == "child-01" + + def test_add_litellm_metadata_from_request_headers_x_litellm_session_id_sets_chain_id(): """x-litellm-session-id sets both metadata and top-level litellm_session_id/litellm_trace_id for call chaining.""" headers = {"x-litellm-session-id": "bar"} diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index cb3faf89358..ca39c7f5c8f 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -7714,6 +7714,176 @@ class _DummyPlugin: return context +class TestChildSessionTierPins: + @staticmethod + def _router(session_affinity: bool = False) -> Router: + return Router( + model_list=[ + { + "model_name": "child-tier-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": {"SIMPLE": "simple", "REASONING": ["reasoning", "reasoning-peer"]}, + "classification_mode": "user_turn", + "session_affinity": session_affinity, + "keyword_tier_rules": [{"keywords": ["reason carefully"], "tier": "REASONING"}], + }, + }, + }, + *[ + { + "model_name": group, + "litellm_params": {"model": "openai/test-model", "api_key": "test-only"}, + "model_info": {"id": f"{group}-{index}"}, + } + for group, index in ( + ("simple", 0), ("simple", 1), ("reasoning", 0), + ("reasoning", 1), ("reasoning-peer", 0), ("reasoning-peer", 1), + ) + ], + ], + num_retries=0, + ) + + @staticmethod + def _normalized_request( + surface: str, session_id: str | None, caller: str = "caller-a", child: bool = True + ) -> tuple[dict[str, object], str]: + from litellm.proxy.litellm_pre_call_utils import ( + LiteLLMProxyRequestSetup, + metadata_variable_name_for_route, + ) + + route: Final = { + "chat": "/v1/chat/completions", "responses": "/v1/responses", "messages": "/v1/messages" + }[surface] + bucket: Final = metadata_variable_name_for_route(route) + data: Final[dict[str, object]] = {bucket: {"user_api_key_hash": caller, "user_api_key_request_route": route}} + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers={ + **({"X-Session-Id": session_id} if session_id is not None else {}), + **({"x-parent-session-id": "parent-session"} if child else {}), + }, + data=data, + _metadata_variable_name=bucket, + ) + return data, bucket + + async def _route( + self, router: Router, surface: str, session_id: str | None, prompt: str | None, + expected_tier: str, expected_cause: str, caller: str = "caller-a", child: bool = True, + expected_deployment: str | None = None, + ) -> str: + from litellm.router_utils.pre_call_checks.deployment_affinity_check import DeploymentAffinityCheck + + data, bucket = self._normalized_request(surface, session_id, caller, child) + turn: Final[dict[str, object]] = ( + {"role": "user", "content": prompt} if prompt is not None else + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_1", "content": "done"}]} + if surface == "messages" else + {"type": "function_call_output", "call_id": "call_1", "output": "done"} + if surface == "responses" else + {"role": "tool", "tool_call_id": "call_1", "content": "done"} + ) + tool_call: Final = ( + {"type": "function_call", "call_id": "call_1", "name": "read_file", "arguments": "{}"} + if surface == "responses" else + {"role": "assistant", "content": [{"type": "tool_use", "id": "call_1", "name": "read_file", "input": {}}]} + if surface == "messages" else TestClassificationMode.TOOL_CALL_1 + ) + turns: Final = [turn] if prompt is not None else [TestClassificationMode.SIMPLE_ASK, tool_call, turn] + if surface == "responses": + data["input"] = turns + selected: Final = await router.async_get_available_deployment( + model="child-tier-router", request_kwargs=data, + messages=None if surface == "responses" else turns, + input=turns if surface == "responses" else None, + ) + metadata: Final = data[bucket] + assert isinstance(metadata, dict) + assert (metadata["session_id"], metadata["trace_id"], data["litellm_session_id"], data["litellm_trace_id"]) == ( + "parent-session", "parent-session", "parent-session", "parent-session" + ) + assert (metadata["routing_decision"]["tier"], metadata["routing_decision"]["cause"]) == ( + expected_tier, expected_cause + ) + affinity: Final = next(cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentAffinityCheck)) + if expected_deployment is not None: + filtered: Final = await affinity.async_filter_deployments( + model=selected["model_name"], healthy_deployments=router.get_model_list(selected["model_name"]) or [], + messages=None, request_kwargs=data, + ) + assert [deployment["model_info"]["id"] for deployment in filtered] == [expected_deployment] + assert selected["model_info"]["id"] == expected_deployment + await affinity.async_pre_call_deployment_hook( + kwargs={ + **data, "model_info": selected["model_info"], + bucket: {**metadata, "deployment_model_name": selected["model_name"]}, + }, + call_type=None, + ) + return str(selected["model_info"]["id"]) + + @pytest.mark.asyncio + @pytest.mark.parametrize("surface", ["chat", "responses", "messages"]) + async def test_children_keep_their_tier_without_replacing_the_parent(self, surface: str) -> None: + router: Final = self._router() + try: + parent: Final = await self._route( + router, surface, "parent-session", "reason carefully", "REASONING", "literal_keyword_match", child=False + ) + for session, prompt, tier, cause, caller, child in ( + ("child-session-a", "Hello!", "SIMPLE", "heuristic_scorer", "caller-a", True), + ("child-session-a", None, "SIMPLE", "user_turn_continuation", "caller-a", True), + ("parent-session", None, "REASONING", "user_turn_continuation", "caller-a", False), + ("child-session-b", None, "SIMPLE", "heuristic_scorer", "caller-a", True), + ("child-session-b", "reason carefully", "REASONING", "literal_keyword_match", "caller-a", True), + ("child-session-a", None, "SIMPLE", "user_turn_continuation", "caller-a", True), + ("child-session-a", "reason carefully", "REASONING", "literal_keyword_match", "caller-b", True), + ("child-session-a", None, "SIMPLE", "user_turn_continuation", "caller-a", True), + (None, "Hello!", "SIMPLE", "heuristic_scorer", "caller-a", True), + (None, None, "SIMPLE", "heuristic_scorer", "caller-a", True), + ): + await self._route( + router, surface, session, prompt, tier, cause, caller, child, + expected_deployment=parent if tier == "REASONING" and caller == "caller-a" else None, + ) + assert await self._route( + router, surface, "parent-session", None, "REASONING", "user_turn_continuation", + child=False, expected_deployment=parent, + ) == parent + finally: + router.discard() + + @pytest.mark.asyncio + async def test_session_affinity_still_shares_the_parent_tier(self) -> None: + router: Final = self._router(session_affinity=True) + try: + parent: Final = await self._route( + router, "chat", "parent-session", "reason carefully", "REASONING", "literal_keyword_match", child=False + ) + assert await self._route( + router, "chat", "child-session-a", "Hello!", "REASONING", "session_affinity_pin", + expected_deployment=parent, + ) == parent + finally: + router.discard() + + @pytest.mark.asyncio + async def test_child_tier_model_affinity_keeps_using_the_parent_identity(self) -> None: + router: Final = self._router() + try: + strategy: Final = router.complexity_routers["child-tier-router"][0].strategy + parent, _ = self._normalized_request("chat", "parent-session", child=False) + child, _ = self._normalized_request("chat", "child-session-a") + candidates: Final = ("reasoning", "reasoning-peer") + assert await strategy._pin_model_for_tier("REASONING", "reasoning", candidates, parent) == "reasoning" + assert await strategy._pin_model_for_tier("REASONING", "reasoning-peer", candidates, child) == "reasoning" + finally: + router.discard() + + class TestClassificationMode: """Test classification_mode='user_turn': classify only requests whose newest turn is a new human ask; tool-loop continuation turns replay the session's held routing decision.""" From 471918c30b246d98be0c233be3c923ebd278ad23 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Wed, 28 Oct 2026 10:08:24 -0700 Subject: [PATCH 048/128] fix(refuseless): update GLM 5.3 and DeepSeek V4.1 Flash prices (#48844) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...model_prices_and_context_window_backup.json | 18 +++++++++--------- model_prices_and_context_window.json | 18 +++++++++--------- 2 files changed, 18 insertions(+), 18 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 5a7006de2cf..43bf7a7d029 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -76979,13 +76979,13 @@ "supports_vision": true }, "refuseless/glm-5.3-abliterated": { - "cache_read_input_token_cost": 2.5e-07, - "input_cost_per_token": 1.25e-06, + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1e-06, "litellm_provider": "refuseless", "max_input_tokens": 1048576, "max_output_tokens": 131072, "mode": "chat", - "output_cost_per_token": 2.25e-06, + "output_cost_per_token": 2e-06, "source": "https://refuseless.com/pricing", "supports_function_calling": true, "supports_prompt_caching": true, @@ -77000,13 +77000,13 @@ ] }, "refuseless/glm-5.3-flash-abliterated": { - "cache_read_input_token_cost": 1.5e-07, - "input_cost_per_token": 2.5e-07, + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 1.2e-07, "litellm_provider": "refuseless", "max_input_tokens": 1048576, "max_output_tokens": 131072, "mode": "chat", - "output_cost_per_token": 8e-07, + "output_cost_per_token": 5e-07, "source": "https://refuseless.com/pricing", "supports_function_calling": true, "supports_prompt_caching": true, @@ -77021,13 +77021,13 @@ ] }, "refuseless/deepseek-v4.1-flash-abliterated": { - "cache_read_input_token_cost": 4.5e-08, - "input_cost_per_token": 2.2e-07, + "cache_read_input_token_cost": 3.5e-08, + "input_cost_per_token": 1.4e-07, "litellm_provider": "refuseless", "max_input_tokens": 1048576, "max_output_tokens": 130000, "mode": "chat", - "output_cost_per_token": 7.5e-07, + "output_cost_per_token": 5e-07, "source": "https://refuseless.com/pricing", "supports_function_calling": true, "supports_prompt_caching": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 5a7006de2cf..43bf7a7d029 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -76979,13 +76979,13 @@ "supports_vision": true }, "refuseless/glm-5.3-abliterated": { - "cache_read_input_token_cost": 2.5e-07, - "input_cost_per_token": 1.25e-06, + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1e-06, "litellm_provider": "refuseless", "max_input_tokens": 1048576, "max_output_tokens": 131072, "mode": "chat", - "output_cost_per_token": 2.25e-06, + "output_cost_per_token": 2e-06, "source": "https://refuseless.com/pricing", "supports_function_calling": true, "supports_prompt_caching": true, @@ -77000,13 +77000,13 @@ ] }, "refuseless/glm-5.3-flash-abliterated": { - "cache_read_input_token_cost": 1.5e-07, - "input_cost_per_token": 2.5e-07, + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 1.2e-07, "litellm_provider": "refuseless", "max_input_tokens": 1048576, "max_output_tokens": 131072, "mode": "chat", - "output_cost_per_token": 8e-07, + "output_cost_per_token": 5e-07, "source": "https://refuseless.com/pricing", "supports_function_calling": true, "supports_prompt_caching": true, @@ -77021,13 +77021,13 @@ ] }, "refuseless/deepseek-v4.1-flash-abliterated": { - "cache_read_input_token_cost": 4.5e-08, - "input_cost_per_token": 2.2e-07, + "cache_read_input_token_cost": 3.5e-08, + "input_cost_per_token": 1.4e-07, "litellm_provider": "refuseless", "max_input_tokens": 1048576, "max_output_tokens": 130000, "mode": "chat", - "output_cost_per_token": 7.5e-07, + "output_cost_per_token": 5e-07, "source": "https://refuseless.com/pricing", "supports_function_calling": true, "supports_prompt_caching": true, From fd9a37d2330f2e836b077238891bf1cebbdfc488 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Wed, 28 Oct 2026 10:16:37 -0700 Subject: [PATCH 049/128] fix(ui): match simple auto router models to advanced setup (#48851) --- .../add_model/ComplexityRouterConfig.tsx | 12 +++--- .../add_auto_router_tab.integration.test.tsx | 21 ++++++++++ .../components/add_model/auto_setup.test.ts | 38 +++++++++++++++++-- .../src/components/add_model/auto_setup.ts | 37 ++++++++---------- 4 files changed, 77 insertions(+), 31 deletions(-) diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index b59aa051e38..ac13af14a6b 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -37,6 +37,7 @@ import { } from "./tier_rows"; import React from "react"; import { ModelGroup } from "@/components/llm_calls/fetch_models"; +import { routingModelGroups } from "./auto_setup"; import ClassificationMethodConfig, { InactiveHeuristicV2Threshold } from "./ClassificationMethodConfig"; import ComplexityRouterAdvancedSections from "./ComplexityRouterAdvancedSections"; import { type TierSetAction, applyTierSetAction, setFallbackTier } from "./tier_set_actions"; @@ -657,13 +658,10 @@ const ComplexityRouterConfig: React.FC = ({ ); const classifierEffortOptionsByModel = classifierEffortOptionsForModels(modelInfo); - // Embedding models can't serve a chat-completion role, so they're excluded here. - const modelOptions = modelInfo - .filter((model) => model.mode !== "embedding") - .map((model) => ({ - value: model.model_group, - label: model.model_group, - })); + const modelOptions = routingModelGroups(modelInfo).map((model) => ({ + value: model.model_group, + label: model.model_group, + })); const handleTierModelParamChange = (tier: string, model: string, change: TierModelParamChange) => onChange({ diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.integration.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.integration.test.tsx index 52c4e9093a4..78782e755ac 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.integration.test.tsx @@ -319,6 +319,27 @@ describe("AddAutoRouterTab", () => { }, ); + it("carries Responses model choices from simple setup into Advanced tiers", async () => { + const models = ["gpt-6-luna", "gpt-6-astra"]; + mockFetchAvailableModels.mockResolvedValue(models.map((model_group) => ({ model_group, mode: "responses" }))); + mockFetchAllModelDeployments.mockResolvedValue( + models.map((model_name) => ({ + model_name, + litellm_params: { model: `chatgpt/${model_name}` }, + model_info: { mode: "responses" }, + })), + ); + renderWithProviders(); + await waitFor(() => + expect(screen.getByRole("combobox", { name: "Everyday model" })).toHaveTextContent("gpt-6-luna"), + ); + expect(screen.getByRole("combobox", { name: "Hard tasks model" })).toHaveTextContent("gpt-6-astra"); + expect(screen.queryByText(/Two deployed models are needed/)).not.toBeInTheDocument(); + fireEvent.click(screen.getByRole("button", { name: "Advanced configuration", exact: true })); + expectTierModel("Simple", "gpt-6-luna"); + expectTierModel("Reasoning", "gpt-6-astra"); + }); + it("keeps manual solver choices on goal changes and carries them into Advanced", async () => { mockFetchAvailableModels.mockResolvedValue( ["claude-haiku-5-5", "custom-fast", "claude-opus-5-5"].map((model_group) => ({ diff --git a/ui/litellm-dashboard/src/components/add_model/auto_setup.test.ts b/ui/litellm-dashboard/src/components/add_model/auto_setup.test.ts index 4097afc17fe..c0d4931873f 100644 --- a/ui/litellm-dashboard/src/components/add_model/auto_setup.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/auto_setup.test.ts @@ -7,6 +7,7 @@ import { solverPairConfig, nextRouterName, simpleSetupError, + routingModelGroups, } from "./auto_setup"; const deployment = (name: string, model: string, extra: Partial = {}): AutoRouterDeployment => ({ @@ -76,7 +77,7 @@ describe("deployedRouterModels", () => { ]); }); - it("uses only caller-visible chat models and excludes blocked and recursive targets", () => { + it("matches Advanced choices while keeping blocked and recursive targets out of suggestions", () => { const rows = [ deployment("hidden", "anthropic/claude-opus-6"), deployment("blocked", "anthropic/claude-opus-5", { model_info: { blocked: true } }), @@ -94,9 +95,40 @@ describe("deployedRouterModels", () => { ], rows, ).map((model) => model.name), - ).toEqual(["ready"]); + ).toEqual(["blocked", "ready", "recursive"]); + expect( + recommendSolverPair( + deployedRouterModels( + rows + .filter((row) => row.model_name !== "hidden" && row.model_name !== "embedding") + .map((row) => ({ model_group: row.model_name! })), + rows, + ), + "balanced", + "auto", + ), + ).toBeNull(); }); + it.each(["chat", "responses", "completion", "unknown", undefined])( + "offers the same routing choices as Advanced for mode %s", + (mode) => { + const groups = [ + { model_group: "everyday", mode }, + { model_group: "hard", mode }, + { model_group: "embedding", mode: "embedding" }, + ]; + const rows = [ + deployment("everyday", "chatgpt/gpt-6-luna", { model_info: { mode } }), + deployment("hard", "chatgpt/gpt-6-astra", { model_info: { mode } }), + ]; + const models = deployedRouterModels(groups, rows); + expect(models.map((model) => model.name)).toEqual(routingModelGroups(groups).map((group) => group.model_group)); + expect(models.map((model) => model.name)).toEqual(["everyday", "hard"]); + expect(names(recommendSolverPair(models, "balanced", "openai"))).toEqual(["everyday", "hard"]); + }, + ); + it("does not claim a mixed-version or unknown model group is a known version", () => { const rows = [ deployment("mixed", "anthropic/claude-opus-5"), @@ -112,7 +144,7 @@ describe("deployedRouterModels", () => { ], rows, ); - expect(models.map((model) => model.name)).toEqual(["claude-sonnet-5", "mixed"]); + expect(models.map((model) => model.name)).toEqual(["claude-sonnet-5", "mixed", "wildcard"]); expect(models.every((model) => model.family === undefined)).toBe(true); expect(recommendSolverPair(models, "balanced", "auto")).toBeNull(); }); diff --git a/ui/litellm-dashboard/src/components/add_model/auto_setup.ts b/ui/litellm-dashboard/src/components/add_model/auto_setup.ts index 39476470928..eb21e17c265 100644 --- a/ui/litellm-dashboard/src/components/add_model/auto_setup.ts +++ b/ui/litellm-dashboard/src/components/add_model/auto_setup.ts @@ -146,24 +146,21 @@ const compareVersion = (left: readonly number[], right: readonly number[]): numb const compareRelease = (left: ModelIdentity, right: ModelIdentity): number => compareVersion(left.version, right.version) || left.snapshot - right.snapshot; +export const routingModelGroups = (groups: readonly ModelGroup[]): ModelGroup[] => + groups.filter((group) => group.mode !== "embedding"); + export const deployedRouterModels = ( groups: readonly ModelGroup[], deployments: readonly AutoRouterDeployment[], ): SolverModel[] => - groups - .flatMap((group) => { - if (group.mode && group.mode !== "chat") return []; + routingModelGroups(groups) + .map((group) => { const rows = deployments.filter((row) => row.model_name === group.model_group); - if ( - rows.some(isAutoRouterDeployment) || - group.model_group.includes("*") || - group.model_group.startsWith("auto_router/") - ) - return []; const active = rows.filter((row) => !row.model_info?.blocked); - if (rows.length > 0 && active.length === 0) return []; - if (active.some((row) => row.model_info?.mode && row.model_info.mode !== "chat")) return []; - if (active.some((row) => row.litellm_params?.model?.includes("*"))) return []; + const isRouter = rows.some(isAutoRouterDeployment) || group.model_group.startsWith("auto_router/"); + const hasWildcard = + group.model_group.includes("*") || active.some((row) => row.litellm_params?.model?.includes("*")); + const manualOnly = isRouter || hasWildcard || (rows.length > 0 && active.length === 0); const identities = active.length > 0 ? active.map(identityOf) : [parseModel(group.model_group)]; const first = identities[0]; const homogeneous = @@ -173,19 +170,17 @@ export const deployedRouterModels = ( const sameLine = identity.family === first.family && identity.line === first.line; return sameLine && compareRelease(identity, first) === 0; }); - const identity = homogeneous ? first : undefined; + const identity = homogeneous && !manualOnly ? first : undefined; const providers = active.length > 0 ? active.map((row) => providerOf(row, group, identity?.family)) : [providerOf({}, group, identity?.family)]; - return [ - { - ...identity, - name: group.model_group, - model: identity?.model ?? group.model_group, - provider: [...new Set(providers)].toSorted().join(", "), - }, - ]; + return { + ...identity, + name: group.model_group, + model: identity?.model ?? group.model_group, + provider: [...new Set(providers)].toSorted().join(", "), + }; }) .toSorted((a, b) => a.name.localeCompare(b.name)); From 9da926b90f0d08954b729e54b70907d7e56058c0 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Wed, 28 Oct 2026 10:16:52 -0700 Subject: [PATCH 050/128] fix(ui): show reasoning in the Playground for /v1/responses (#48850) * test(ui): cover Responses API reasoning events in the playground Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): show reasoning for current Responses API reasoning events 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> --- .../llm_calls/responses_api.test.tsx | 131 ++++++++++++++++++ .../components/llm_calls/responses_api.tsx | 37 ++++- 2 files changed, 161 insertions(+), 7 deletions(-) diff --git a/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx b/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx index 47a1e3608c5..f9c5e3b0181 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx @@ -490,6 +490,137 @@ describe("responses_api", () => { }); }); +describe("responses_api reasoning", () => { + type Run = { reasoning: string; answer: string; timeToFirstToken: number[]; reasoningBeforeFirstToken: string }; + + const run = async (streamingEnabled: boolean): Promise => { + const reasoningChunks: string[] = []; + const answerChunks: string[] = []; + const timeToFirstToken: number[] = []; + const reasoningAtFirstToken: string[] = []; + await makeOpenAIResponsesRequest( + [{ role: "user", content: "What is 3+4?" }], + (_role, delta) => answerChunks.push(delta), + "qwen3", + "test-token", + undefined, + undefined, + (content) => reasoningChunks.push(content), + (ttft) => { + timeToFirstToken.push(ttft); + reasoningAtFirstToken.push(reasoningChunks.join("") + "|" + answerChunks.join("")); + }, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + streamingEnabled, + ); + return { + reasoning: reasoningChunks.join(""), + answer: answerChunks.join(""), + timeToFirstToken, + reasoningBeforeFirstToken: reasoningAtFirstToken[0] ?? "", + }; + }; + + const streamOf = (events: Record[]) => { + async function* mockStream() { + for (const event of events) { + yield event; + } + yield { type: "response.completed", response: { id: "resp_reasoning" } }; + } + mockResponsesCreate.mockResolvedValue(mockStream()); + }; + + afterEach(() => { + vi.clearAllMocks(); + }); + + it("shows streamed reasoning_text deltas and starts the first-token timer on them", async () => { + streamOf([ + { type: "response.reasoning_text.delta", item_id: "rs_1", content_index: 0, delta: "Add " }, + { type: "response.reasoning_text.delta", item_id: "rs_1", content_index: 0, delta: "them" }, + { type: "response.output_text.delta", delta: "7" }, + ]); + + const result = await run(true); + + expect(result.reasoning).toBe("Add them"); + expect(result.answer).toBe("7"); + expect(result.timeToFirstToken).toHaveLength(1); + expect(result.reasoningBeforeFirstToken).toBe("Add |"); + }); + + it("separates streamed reasoning summary parts and items with a blank line", async () => { + streamOf([ + { type: "response.reasoning_summary_text.delta", item_id: "rs_1", summary_index: 0, delta: "First" }, + { type: "response.reasoning_summary_text.delta", item_id: "rs_1", summary_index: 0, delta: " part" }, + { type: "response.reasoning_summary_text.delta", item_id: "rs_1", summary_index: 1, delta: "Second" }, + { type: "response.reasoning_summary_text.delta", item_id: "rs_2", summary_index: 0, delta: "Third" }, + ]); + + await expect(run(true)).resolves.toMatchObject({ reasoning: "First part\n\nSecond\n\nThird" }); + }); + + it("still shows the legacy reasoning.delta event", async () => { + streamOf([{ type: "response.reasoning.delta", delta: "Legacy thought" }]); + + await expect(run(true)).resolves.toMatchObject({ reasoning: "Legacy thought" }); + }); + + it("shows non-streamed reasoning content when the summary is empty", async () => { + mockResponsesCreate.mockReturnValueOnce( + nonStreamingResponse({ + id: "resp_content", + output: [ + { + type: "reasoning", + summary: [], + content: [ + { type: "reasoning_text", text: "Add" }, + { type: "reasoning_text", text: "them" }, + ], + }, + { type: "message", content: [{ type: "output_text", text: "7" }] }, + ], + }), + ); + + await expect(run(false)).resolves.toMatchObject({ reasoning: "Add\n\nthem", answer: "7" }); + }); + + it("prefers the non-streamed reasoning summary over its content", async () => { + mockResponsesCreate.mockReturnValueOnce( + nonStreamingResponse({ + id: "resp_summary", + output: [ + { + type: "reasoning", + summary: [{ type: "summary_text", text: "Short summary" }], + content: [{ type: "reasoning_text", text: "Raw chain" }], + }, + { type: "reasoning", content: [{ type: "output_text", text: "Second item" }] }, + ], + }), + ); + + await expect(run(false)).resolves.toMatchObject({ reasoning: "Short summary\n\nSecond item" }); + }); +}); + describe("responses_api prompt cache usage", () => { const captureUsage = async (usage: Record): Promise => { async function* mockStream() { diff --git a/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx b/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx index 426abd4e172..ce1a8aab0c5 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx @@ -44,6 +44,22 @@ type SynthesizedResponseEvent = | { type: "response.output_text.delta"; delta: string } | { type: "response.completed"; response: NonStreamedResponse }; +const REASONING_PART_SEPARATOR = "\n\n"; + +const REASONING_DELTA_EVENT_TYPES: ReadonlySet = new Set([ + "response.reasoning_text.delta", + "response.reasoning_summary_text.delta", + "response.reasoning.delta", +]); + +const nonEmptyTexts = (parts: ResponseOutputPart[] | undefined): string[] => + (parts ?? []).map((part) => part.text ?? "").filter((text) => text.length > 0); + +const reasoningItemText = (item: ResponseOutputItem): string => { + const summary = nonEmptyTexts(item.summary); + return (summary.length > 0 ? summary : nonEmptyTexts(item.content)).join(REASONING_PART_SEPARATOR); +}; + const responseAsEvents = (response: NonStreamedResponse): SynthesizedResponseEvent[] => { const outputItems = response.output ?? []; const outputText = outputItems @@ -54,9 +70,9 @@ const responseAsEvents = (response: NonStreamedResponse): SynthesizedResponseEve .join(""); const reasoningText = outputItems .filter((item) => item.type === "reasoning") - .flatMap((item) => item.summary ?? []) - .map((part) => part.text ?? "") - .join(""); + .map(reasoningItemText) + .filter((text) => text.length > 0) + .join(REASONING_PART_SEPARATOR); return [ ...outputItems.map((item) => ({ type: "response.output_item.done" as const, item })), @@ -124,6 +140,7 @@ export async function makeOpenAIResponsesRequest( try { const startTime = Date.now(); let firstTokenReceived = false; + let reasoningPartKey: string | null = null; let servedFromResponseCache = false; const formattedInput = messages.flatMap((message) => { @@ -298,10 +315,16 @@ export async function makeOpenAIResponsesRequest( } // Handle reasoning content - if (event.type === "response.reasoning.delta" && "delta" in event) { - const delta = event.delta; - if (typeof delta === "string" && onReasoningContent) { - onReasoningContent(delta); + if (REASONING_DELTA_EVENT_TYPES.has(event.type) && typeof event.delta === "string" && event.delta.length > 0) { + const partKey = `${event.item_id ?? ""}:${event.summary_index ?? event.content_index ?? 0}`; + const separator = reasoningPartKey !== null && reasoningPartKey !== partKey ? REASONING_PART_SEPARATOR : ""; + reasoningPartKey = partKey; + onReasoningContent?.(separator + event.delta); + + const isFirstToken = !firstTokenReceived; + firstTokenReceived = true; + if (isFirstToken && onTimingData && streamingEnabled) { + onTimingData(Date.now() - startTime); } } From 12deb60e270fffe8de4832c26b4fcd6634affff4 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Wed, 28 Oct 2026 10:33:11 -0700 Subject: [PATCH 051/128] fix(responses): forward top-level moderation param to OpenAI (#48708) * test: cover Responses moderation forwarding Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(responses): forward moderation parameter Co-authored-by: Poppy Huang Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: remove Perplexity drop_params moderation case Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(responses): preserve typed moderation forwarding Co-authored-by: Poppy Huang 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: Poppy Huang --- .../llms/openai/responses/transformation.py | 18 ++++- litellm/types/llms/openai.py | 5 ++ .../e2e/llm_translation/test_responses_e2e.py | 42 ++++++++++ ...est_perplexity_responses_transformation.py | 57 +++++++++---- .../test_responses_api_request_body.py | 80 ++++++++++++++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 9 +++ 6 files changed, 191 insertions(+), 20 deletions(-) diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 07822bb17b3..9cb074a177e 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -499,8 +499,13 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): ) if sanitized_tools is not None: response_api_optional_request_params["tools"] = sanitized_tools - final_request_params: Final = dict( - ResponsesAPIRequestParams(model=model, input=replay_safe_input, **response_api_optional_request_params) + final_request_params: Final = cast( # cast-ok: typed dict spread is inferred as Unknown by basedpyright + dict[str, object], + { + "model": model, + "input": replay_safe_input, + **response_api_optional_request_params, + }, ) return final_request_params @@ -1209,8 +1214,13 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): ) if sanitized_tools is not None: response_api_optional_request_params["tools"] = sanitized_tools - data: Final = dict( - ResponsesAPIRequestParams(model=model, input=replay_safe_input, **response_api_optional_request_params) + data: Final = cast( # cast-ok: typed dict spread is inferred as Unknown by basedpyright + dict[str, object], + { + "model": model, + "input": replay_safe_input, + **response_api_optional_request_params, + }, ) return url, data diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 4912f9f1375..8e1db87c16f 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1342,6 +1342,10 @@ class ResponsesAPIStreamOptions(TypedDict, total=False): include_obfuscation: bool +class ResponsesAPIModeration(TypedDict, total=False): + model: Required[ReadOnly[str]] + + class ResponsesAPIOptionalRequestParams(TypedDict, total=False): """TypedDict for Optional parameters supported by the responses API.""" @@ -1369,6 +1373,7 @@ class ResponsesAPIOptionalRequestParams(TypedDict, total=False): prompt_cache_key: str | None prompt_cache_retention: str | None prompt_cache_options: ReadOnly[PromptCacheOptions | None] + moderation: ReadOnly[ResponsesAPIModeration | None] stream_options: ResponsesAPIStreamOptions | None top_logprobs: int | None partial_images: int | None # Number of partial images to generate (1-3) for streaming image generation diff --git a/tests/e2e/llm_translation/test_responses_e2e.py b/tests/e2e/llm_translation/test_responses_e2e.py index bb75569ad99..a0f185c4077 100644 --- a/tests/e2e/llm_translation/test_responses_e2e.py +++ b/tests/e2e/llm_translation/test_responses_e2e.py @@ -183,6 +183,48 @@ class TestResponses: ) assert response.output_text.strip(), f"/responses returned no output text: {response.output!r}" + @pytest.mark.covers("llm.responses.openai.basic.nonstream.works") + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.RESPONSES, + providers=(Provider.OPENAI,), + models=(OPENAI_RESPONSES_BACKEND,), + mode=Mode.NONSTREAM, + ) + ) + def test_responses_forwards_moderation( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients + ) -> None: + model: Final = _register(proxy, resources, _openai_responses_params()) + client: Final = sdk.openai(resources.key()) + + response: Final = client.responses.create( + model=model, + input="hi", + extra_body={**NO_PROXY_CACHE, "moderation": {"model": "omni-moderation-latest"}}, + ) + response_body: Final = TypeAdapter(dict[str, object]).validate_python(cast(object, response.model_dump())) + moderation_value: Final = response_body.get("moderation") + assert moderation_value is not None, f"response had no moderation results: {response_body!r}" + moderation: Final = TypeAdapter(dict[str, object]).validate_python(moderation_value) + moderation_input: Final = TypeAdapter(dict[str, object]).validate_python(moderation["input"]) + moderation_output: Final = TypeAdapter(dict[str, object]).validate_python(moderation["output"]) + assert moderation_input.get("type") == "moderation_result", ( + f"unexpected input moderation: {moderation_input!r}" + ) + assert ( + moderation_output.get("type") == "moderation_result" + ), f"unexpected output moderation: {moderation_output!r}" + + with pytest.raises(openai.BadRequestError) as exc_info: + client.responses.create( + model=model, + input="hi", + extra_body={**NO_PROXY_CACHE, "moderation": {"model": "does-not-exist"}}, + ) + assert exc_info.value.status_code == 400 + @pytest.mark.covers("llm.responses.openai.basic.stream.works") @meta( Subject( diff --git a/tests/unit/llms/perplexity/responses/test_perplexity_responses_transformation.py b/tests/unit/llms/perplexity/responses/test_perplexity_responses_transformation.py index 534176e381a..b6d26aa1de6 100644 --- a/tests/unit/llms/perplexity/responses/test_perplexity_responses_transformation.py +++ b/tests/unit/llms/perplexity/responses/test_perplexity_responses_transformation.py @@ -8,10 +8,13 @@ Source: litellm/llms/perplexity/responses/transformation.py """ import json +from typing import Final import httpx import pytest +import respx +import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -21,6 +24,17 @@ from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager +_SUCCESS_RESPONSE: Final[dict[str, object]] = { + "id": "resp_moderation_test", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "openai/gpt-5.2", + "output": [], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, +} + + class TestPerplexityResponsesTransformation: """Test Perplexity Responses API configuration and transformations""" @@ -60,9 +74,7 @@ class TestPerplexityResponsesTransformation: assert len(result["tools"]) == 1 assert result["tools"][0]["type"] == "function" assert result["tools"][0]["function"]["name"] == "get_weather" - assert ( - result["tools"][0]["function"]["description"] == "Get the current weather" - ) + assert result["tools"][0]["function"]["description"] == "Get the current weather" assert "parameters" in result["tools"][0]["function"] def test_web_search_tool_passthrough(self): @@ -130,9 +142,7 @@ class TestPerplexityResponsesTransformation: """tool_choice passes through""" config = PerplexityResponsesConfig() - params = ResponsesAPIOptionalRequestParams( - tool_choice="required", temperature=0.7 - ) + params = ResponsesAPIOptionalRequestParams(tool_choice="required", temperature=0.7) result = config.map_openai_params( response_api_optional_params=params, @@ -146,9 +156,7 @@ class TestPerplexityResponsesTransformation: """parallel_tool_calls passes through""" config = PerplexityResponsesConfig() - params = ResponsesAPIOptionalRequestParams( - parallel_tool_calls=True, temperature=0.7 - ) + params = ResponsesAPIOptionalRequestParams(parallel_tool_calls=True, temperature=0.7) result = config.map_openai_params( response_api_optional_params=params, @@ -492,9 +500,7 @@ class TestPerplexityResponsesTransformation: "id": "msg_123", "role": "assistant", "status": "completed", - "content": [ - {"type": "output_text", "text": "Hello!", "annotations": []} - ], + "content": [{"type": "output_text", "text": "Hello!", "annotations": []}], } ], "usage": { @@ -549,9 +555,7 @@ class TestPerplexityResponsesTransformation: "id": "msg_123", "role": "assistant", "status": "completed", - "content": [ - {"type": "output_text", "text": "Hello!", "annotations": []} - ], + "content": [{"type": "output_text", "text": "Hello!", "annotations": []}], } ], "usage": { @@ -587,3 +591,26 @@ class TestPerplexityResponsesTransformation: assert result.type == "response.completed" assert result.response.usage.cost == 0.0003 assert isinstance(result.response.usage.cost, float) + + +@pytest.mark.asyncio +async def test_aresponses_rejects_moderation_for_perplexity_without_drop_params( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "drop_params", False) + litellm.in_memory_llm_clients_cache.flush_cache() + upstream: Final = respx_mock.post("https://api.perplexity.ai/v1/responses").mock( + return_value=httpx.Response(200, json=_SUCCESS_RESPONSE) + ) + + with pytest.raises(litellm.UnsupportedParamsError): + await litellm.aresponses( + model="perplexity/openai/gpt-5.2", + api_base="https://api.perplexity.ai", + api_key="fake-api-key", + input="hi", + moderation={"model": "omni-moderation-latest"}, + ) + + assert upstream.call_count == 0 diff --git a/tests/unit/responses/test_responses_api_request_body.py b/tests/unit/responses/test_responses_api_request_body.py index bf571c74e22..b926b824fc6 100644 --- a/tests/unit/responses/test_responses_api_request_body.py +++ b/tests/unit/responses/test_responses_api_request_body.py @@ -12,7 +12,7 @@ import math from collections.abc import Awaitable, Callable, Mapping from pathlib import Path from importlib import import_module -from typing import Final +from typing import Final, cast from unittest.mock import AsyncMock, patch import httpx @@ -254,6 +254,84 @@ async def test_aresponses_forwards_non_enum_reasoning_effort( assert response.output[0].content[0].text == "Done." +@pytest.mark.asyncio +async def test_aresponses_forwards_moderation( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + upstream: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, json=_minimal_responses_api_payload("resp_moderation", "gpt-5.4")) + ) + moderation: Final = {"model": "omni-moderation-latest"} + + await litellm.aresponses(model="openai/gpt-5.4", input="hi", moderation=moderation) + + request_body: Final = TypeAdapter(dict[str, object]).validate_python( + cast(object, json.loads(upstream.calls[0].request.read())) + ) + assert upstream.call_count == 1 + assert request_body.get("moderation") == moderation + + +@pytest.mark.asyncio +async def test_aresponses_omits_null_moderation( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + upstream: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, json=_minimal_responses_api_payload("resp_null_moderation", "gpt-5.4")) + ) + + await litellm.aresponses(model="openai/gpt-5.4", input="hi", moderation=None) + + request_body: Final = TypeAdapter(dict[str, object]).validate_python( + cast(object, json.loads(upstream.calls[0].request.read())) + ) + assert upstream.call_count == 1 + assert "moderation" not in request_body + + +@pytest.mark.asyncio +async def test_aresponses_stream_forwards_moderation(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + received: Final[asyncio.Queue[httpx.Request]] = asyncio.Queue() + payload: Final = _minimal_responses_api_payload("resp_stream_moderation", "gpt-5.4") + moderation: Final = {"model": "omni-moderation-latest"} + + def respond(request: httpx.Request) -> httpx.Response: + received.put_nowait(request) + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text="".join( + f"event: {kind}\ndata: {json.dumps({'type': kind, 'response': payload})}\n\n" + for kind in ("response.created", "response.completed") + ), + ) + + async_client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(respond)) + try: + response: Final = await litellm.aresponses( + model="openai/gpt-5.4", input="hi", moderation=moderation, stream=True, client=async_client + ) + events: Final = tuple([event.type async for event in response]) + assert "response.completed" in events + finally: + await async_client.close() + + request_body: Final = TypeAdapter(dict[str, object]).validate_python( + cast(object, json.loads(received.get_nowait().content)) + ) + assert request_body.get("moderation") == moderation + assert received.empty() + + @pytest.mark.asyncio async def test_acompletion_with_tools_forwards_non_enum_reasoning_effort_over_the_bridge( monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 8af76e4d16f..76e3c95fb75 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -44163,6 +44163,11 @@ export interface components { */ type: "tool_search_output"; }; + /** ResponsesAPIModeration */ + ResponsesAPIModeration: { + /** Model */ + model: string; + }; /** * ResponsesAPIRequestParams * @description TypedDict for request parameters supported by the responses API. @@ -44188,6 +44193,7 @@ export interface components { } | null; /** Model */ model: string; + moderation?: components["schemas"]["ResponsesAPIModeration"] | null; /** Parallel Tool Calls */ parallel_tool_calls?: boolean | null; /** Partial Images */ @@ -66975,6 +66981,7 @@ export interface operations { } | null; /** Model */ model: string; + moderation?: components["schemas"]["ResponsesAPIModeration"] | null; /** Parallel Tool Calls */ parallel_tool_calls?: boolean | null; /** Partial Images */ @@ -70439,6 +70446,7 @@ export interface operations { } | null; /** Model */ model: string; + moderation?: components["schemas"]["ResponsesAPIModeration"] | null; /** Parallel Tool Calls */ parallel_tool_calls?: boolean | null; /** Partial Images */ @@ -81237,6 +81245,7 @@ export interface operations { } | null; /** Model */ model: string; + moderation?: components["schemas"]["ResponsesAPIModeration"] | null; /** Parallel Tool Calls */ parallel_tool_calls?: boolean | null; /** Partial Images */ From 0f49db5c3b979092f6d9c67dc6053130d926801f Mon Sep 17 00:00:00 2001 From: shrey Date: Wed, 28 Oct 2026 10:39:01 -0700 Subject: [PATCH 052/128] fix(proxy): store boolean settings as JSON booleans on /config/field/update (#48595) --- litellm/proxy/proxy_server.py | 8 +++++--- .../authorization/test_agent_least_privilege.py | 2 +- .../unit/proxy/proxy_server/test_routes_config.py | 14 ++++++++------ 3 files changed, 14 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 04f84f1f7b5..4e7c4601803 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -19899,12 +19899,14 @@ async def update_config_general_settings( ) try: - ConfigGeneralSettings.model_validate({data.field_name: requested_value}) + validated: Final = ConfigGeneralSettings.model_validate({data.field_name: requested_value}) except Exception: raise HTTPException( status_code=400, detail={"error": f"Invalid type of field value={type(requested_value)} passed in."}, ) + parsed_value: Final[object] = _UNTYPED_VALUE.validate_python(getattr(validated, data.field_name)) + stored_value: Final = parsed_value if isinstance(parsed_value, bool) else requested_value if data.field_name == "alerting_args": try: @@ -19933,10 +19935,10 @@ async def update_config_general_settings( proxy_config.reject_config_owned_writes( section_name="general_settings", - changed_keys={data.field_name: cast(JsonValue, requested_value)}, # cast-ok: validated above + changed_keys={data.field_name: cast(JsonValue, stored_value)}, # cast-ok: validated above ) - field_value = requested_value + field_value = stored_value if data.field_name == "plugins": field_value = _preserve_redacted_plugin_keys(field_value, general_settings.get("plugins")) diff --git a/tests/integration/authorization/test_agent_least_privilege.py b/tests/integration/authorization/test_agent_least_privilege.py index 498a53c3ca9..d993920a32e 100644 --- a/tests/integration/authorization/test_agent_least_privilege.py +++ b/tests/integration/authorization/test_agent_least_privilege.py @@ -313,7 +313,7 @@ def test_agent_access_default_deny_toggles_at_runtime_through_config_api( for field in config_fields.json() if object_value(field)["field_name"] == "agent_access_default_deny" ) - assert setting["field_type"] == "Boolean" and setting["field_value"] == "true", setting + assert setting["field_type"] == "Boolean" and setting["field_value"] is True, setting eventually( lambda: tuple(_card(candidate, agent_id, ungranted_key) for _ in range(8)), lambda responses: all(response.status_code == 403 for response in responses), diff --git a/tests/unit/proxy/proxy_server/test_routes_config.py b/tests/unit/proxy/proxy_server/test_routes_config.py index 7d0a7f3f261..e3b8aff8fcb 100644 --- a/tests/unit/proxy/proxy_server/test_routes_config.py +++ b/tests/unit/proxy/proxy_server/test_routes_config.py @@ -716,12 +716,16 @@ def test_config_list_exposes_config_reload_interval(client, auth_as, mock_prisma assert entry["field_default_value"] == 30 -@pytest.mark.parametrize("enabled", (True, False)) +@pytest.mark.parametrize( + ("sent", "enabled"), + ((True, True), (False, False), ("true", True), ("yes", True), (1, True), ("off", False)), +) def test_config_list_exposes_and_saves_wildcard_routes_toggle( client: TestClient, auth_as: Callable[..., contextlib.AbstractContextManager[object]], mock_prisma: MagicMock, monkeypatch: pytest.MonkeyPatch, + sent: object, enabled: bool, ) -> None: from pydantic import TypeAdapter @@ -747,13 +751,11 @@ def test_config_list_exposes_and_saves_wildcard_routes_toggle( saved: Final = client.post( "/config/field/update", - json={"field_name": "return_wildcard_routes", "field_value": enabled, "config_type": "general_settings"}, + json={"field_name": "return_wildcard_routes", "field_value": sent, "config_type": "general_settings"}, ) assert saved.status_code == 200, saved.text - persisted: Final = TypeAdapter(dict[str, bool]).validate_json( - table.upsert.call_args.kwargs["data"]["create"]["param_value"] - ) - assert persisted == {"return_wildcard_routes": enabled} + persisted: Final = json.loads(table.upsert.call_args.kwargs["data"]["create"]["param_value"]) + assert persisted["return_wildcard_routes"] is enabled after: Final = client.get("/config/list", params={"config_type": "general_settings"}) assert after.status_code == 200, after.text From 2f84296ad271bb73cbf07b2529e34dd054483a4e Mon Sep 17 00:00:00 2001 From: shrey Date: Wed, 28 Oct 2026 10:43:27 -0700 Subject: [PATCH 053/128] fix(azure): map context_length_exceeded error code to ContextWindowExceededError (#48616) --- litellm/exceptions.py | 8 +- .../exception_mapping_utils.py | 20 ++- .../azure/test_azure_exception_mapping.py | 144 +++++++++++++++++- 3 files changed, 160 insertions(+), 12 deletions(-) diff --git a/litellm/exceptions.py b/litellm/exceptions.py index b8e933bf046..c8eee9ad1f8 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -232,7 +232,7 @@ class BadRequestError(openai.BadRequestError): litellm_debug_info: str | None = None, max_retries: int | None = None, num_retries: int | None = None, - body: dict | None = None, + body: Mapping[str, object] | None = None, ): self.status_code = 400 self.message = f"litellm.BadRequestError: {message}" @@ -276,7 +276,7 @@ class ImageFetchError(BadRequestError): litellm_debug_info: str | None = None, max_retries: int | None = None, num_retries: int | None = None, - body: dict | None = None, + body: Mapping[str, object] | None = None, ): super().__init__( message=message, @@ -556,6 +556,7 @@ class ContextWindowExceededError(BadRequestError): llm_provider, response: httpx.Response | None = None, litellm_debug_info: str | None = None, + body: Mapping[str, object] | None = None, ): self.status_code = 400 self.model = model @@ -567,6 +568,7 @@ class ContextWindowExceededError(BadRequestError): llm_provider=self.llm_provider, response=response, litellm_debug_info=self.litellm_debug_info, + body=body, ) # Call the base class constructor with the parameters it needs # set after, to make it clear the raised error is a context window exceeded error @@ -685,7 +687,7 @@ class ContentPolicyViolationError(BadRequestError): response: httpx.Response | None = None, litellm_debug_info: str | None = None, provider_specific_fields: dict | None = None, - body: dict | None = None, + body: Mapping[str, object] | None = None, ): self.status_code = 400 self.message = f"litellm.ContentPolicyViolationError: {message}" diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 1c816b22ef1..47ec8331c83 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -2036,14 +2036,6 @@ def _map_azure_exception( litellm_debug_info=extra_information, response=getattr(original_exception, "response", None), ) - elif "This model's maximum context length is" in error_str: - raise ContextWindowExceededError( - message=f"AzureException ContextWindowExceededError - {message}", - llm_provider="azure", - model=model, - litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), - ) elif "DeploymentNotFound" in error_str: raise NotFoundError( message=f"AzureException NotFoundError - {message}", @@ -2085,6 +2077,18 @@ def _map_azure_exception( response=getattr(original_exception, "response", None), body=getattr(original_exception, "body", None), ) + elif getattr(original_exception, "status_code", None) in (None, 400) and ( + azure_error_code == "context_length_exceeded" + or ExceptionCheckers.is_error_str_context_window_exceeded(error_str) + ): + raise ContextWindowExceededError( + message=f"AzureException ContextWindowExceededError - {message}", + llm_provider="azure", + model=model, + litellm_debug_info=extra_information, + response=getattr(original_exception, "response", None), + body=getattr(original_exception, "body", None), + ) elif "invalid_request_error" in error_str and getattr(original_exception, "status_code", None) in (None, 400): raise BadRequestError( message=f"AzureException BadRequestError - {message}", diff --git a/tests/unit/llms/azure/test_azure_exception_mapping.py b/tests/unit/llms/azure/test_azure_exception_mapping.py index 8670a9df2ea..cbbf70c691b 100644 --- a/tests/unit/llms/azure/test_azure_exception_mapping.py +++ b/tests/unit/llms/azure/test_azure_exception_mapping.py @@ -1,13 +1,16 @@ import json +from collections.abc import Awaitable, Callable, Mapping +from typing import Final from unittest.mock import MagicMock, patch import pytest import httpx +import openai import respx import litellm -from litellm.exceptions import ContentPolicyViolationError +from litellm.exceptions import BadRequestError, ContentPolicyViolationError, ContextWindowExceededError from litellm.litellm_core_utils.exception_mapping_utils import exception_type @@ -490,3 +493,142 @@ class TestAzureExceptionMapping: error = exc_info.value assert "encrypted_content_affinity" in error.message assert "enable_pre_call_checks" in error.message + + +_AZURE_API_BASE: Final = "https://azure-context.openai.azure.com" +_AZURE_CALL_KWARGS: Final = { + "api_base": _AZURE_API_BASE, + "api_version": "2024-10-21", + "api_key": "test-azure-key", + "max_retries": 0, +} + + +_CONTEXT_LENGTH_ERROR: Final = { + "message": "Your input exceeds the context window of this model. Please adjust your input and try again.", + "type": "invalid_request_error", + "param": "input", + "code": "context_length_exceeded", +} + + +def _azure_error_response(status_code: int, error: Mapping[str, object]) -> httpx.Response: + return httpx.Response( + status_code, text=json.dumps({"error": error}, indent=2), headers={"content-type": "application/json"} + ) + + +async def _azure_chat(*, stream: bool) -> None: + messages: Final = [{"role": "user", "content": "hi"}] + if stream: + stream_response: Final = await litellm.acompletion( + model="azure/gpt-4o", messages=messages, stream=True, **_AZURE_CALL_KWARGS + ) + assert isinstance(stream_response, litellm.CustomStreamWrapper) + async for _chunk in stream_response: + pass + else: + litellm.completion(model="azure/gpt-4o", messages=messages, **_AZURE_CALL_KWARGS) + + +async def _azure_responses(*, is_async: bool) -> None: + if is_async: + await litellm.aresponses(model="azure/gpt-4o", input="hi", **_AZURE_CALL_KWARGS) + else: + litellm.responses(model="azure/gpt-4o", input="hi", **_AZURE_CALL_KWARGS) + + +async def _azure_embedding() -> None: + litellm.embedding(model="azure/text-embedding-3-small", input="hi", **_AZURE_CALL_KWARGS) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "call", + [ + pytest.param(lambda: _azure_chat(stream=False), id="chat"), + pytest.param(lambda: _azure_chat(stream=True), id="async-chat-stream"), + pytest.param(lambda: _azure_responses(is_async=False), id="responses"), + pytest.param(lambda: _azure_responses(is_async=True), id="async-responses"), + pytest.param(_azure_embedding, id="embedding"), + ], +) +@pytest.mark.parametrize( + ("code", "expected"), + [("context_length_exceeded", ContextWindowExceededError), ("invalid_value", BadRequestError)], +) +@respx.mock +async def test_azure_400_is_classified_by_error_code( + monkeypatch: pytest.MonkeyPatch, call: Callable[[], Awaitable[None]], code: str, expected: type[Exception] +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx.post(url__startswith=_AZURE_API_BASE).mock( + return_value=_azure_error_response(400, {**_CONTEXT_LENGTH_ERROR, "code": code}) + ) + + with pytest.raises(BadRequestError) as exc_info: + await call() + + assert type(exc_info.value) is expected + + +@pytest.mark.asyncio +@respx.mock +async def test_azure_context_length_code_triggers_context_window_fallback(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx.post(url__startswith=_AZURE_API_BASE).mock( + return_value=_azure_error_response(400, _CONTEXT_LENGTH_ERROR) + ) + router: Final = litellm.Router( + model_list=[ + {"model_name": "small", "litellm_params": {"model": "azure/gpt-4o", **_AZURE_CALL_KWARGS}}, + {"model_name": "large", "litellm_params": {"model": "gpt-4.1", "mock_response": "served by large"}}, + ], + context_window_fallbacks=[{"small": ["large"]}], + num_retries=0, + ) + + response: Final = await router.acompletion( # pyright: ignore[reportUnknownMemberType] # untyped **kwargs + model="small", messages=[{"role": "user", "content": "hi"}] + ) + + assert isinstance(response, litellm.ModelResponse) + assert response.choices[0].message.content == "served by large" + assert len(route.calls) == 1 + + +def _azure_chat_error(monkeypatch: pytest.MonkeyPatch, response: httpx.Response) -> openai.APIError: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx.post(url__startswith=_AZURE_API_BASE).mock(return_value=response) + with pytest.raises(openai.APIError) as exc_info: + litellm.completion(model="azure/gpt-4o", messages=[{"role": "user", "content": "hi"}], **_AZURE_CALL_KWARGS) + return exc_info.value + + +@respx.mock +def test_azure_context_window_error_keeps_upstream_code_and_param(monkeypatch: pytest.MonkeyPatch) -> None: + error: Final = _azure_chat_error(monkeypatch, _azure_error_response(400, _CONTEXT_LENGTH_ERROR)) + + assert isinstance(error, ContextWindowExceededError) + assert error.code == "context_length_exceeded" + assert error.param == "input" + + +@pytest.mark.parametrize("status_code", [429, 500]) +@respx.mock +def test_azure_non_400_with_context_length_code_keeps_its_status( + monkeypatch: pytest.MonkeyPatch, status_code: int +) -> None: + error: Final = _azure_chat_error(monkeypatch, _azure_error_response(status_code, _CONTEXT_LENGTH_ERROR)) + + assert not isinstance(error, ContextWindowExceededError) + assert getattr(error, "status_code", None) == status_code + + +@respx.mock +def test_azure_content_policy_wins_over_context_length_code(monkeypatch: pytest.MonkeyPatch) -> None: + policy_error: Final = {**_CONTEXT_LENGTH_ERROR, "inner_error": {"code": "ResponsibleAIPolicyViolation"}} + + error: Final = _azure_chat_error(monkeypatch, _azure_error_response(400, policy_error)) + + assert type(error) is ContentPolicyViolationError From 208ef99b88690f9abb28251d87565f8887fba6ad Mon Sep 17 00:00:00 2001 From: moe-berri Date: Wed, 28 Oct 2026 10:46:47 -0700 Subject: [PATCH 054/128] fix(images): honor max_retries and prevent Router retry amplification (#48705) * test(images): cover SDK retry budgets and router ownership Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(images): forward max_retries to the OpenAI image generation client Per-call, deployment and default_litellm_params max_retries were filtered out before reaching the OpenAI-compatible image handler, which then always used 2 SDK retries. Forward the caller's value, and pin it to 0 for routed calls so the Router stays the only retry owner, matching completion() Co-authored-by: Ericz Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(images): read image max_retries through a typed kwargs view Co-authored-by: Ericz Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(images): assert the public invalid retry error mapping Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(images): preserve the invalid retry status code Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(images): group the proxy test imports Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(images): discard proxy retry test routers 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: Ericz --- litellm/images/main.py | 11 +- .../test_image_gen_drop_params_wire.py | 87 ++++++- tests/unit/images/test_main.py | 214 +++++++++++++++++- .../test_router_per_deployment_num_retries.py | 69 +++++- 4 files changed, 376 insertions(+), 5 deletions(-) diff --git a/litellm/images/main.py b/litellm/images/main.py index 9c88e737a40..3312ca417d5 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -284,6 +284,12 @@ def image_generation( extra_headers: Final = kwargs.get("extra_headers", None) headers: Final[dict] = kwargs.get("headers", None) or {} base_model: Final = kwargs.get("base_model", None) + call_kwargs: Final = cast(Mapping[str, object], kwargs) # cast-ok: image_generation's **kwargs is untyped + max_retries: Final = call_kwargs.get("max_retries") + is_router_call: Final = any( + isinstance(call_metadata := call_kwargs.get(k), Mapping) and "model_group" in call_metadata + for k in ("metadata", "litellm_metadata") + ) if extra_headers is not None: headers.update(extra_headers) model_response: ImageResponse = litellm.utils.ImageResponse() @@ -550,6 +556,7 @@ def image_generation( optional_params["extra_headers"] = extra_headers # Forward OpenAI organization if present (set by proxy pre-call utils) organization: Final[str | None] = kwargs.get("organization", None) + sdk_max_retries: Final = 0 if is_router_call else max_retries model_response = openai_chat_completions.image_generation( model=model, prompt=prompt, @@ -557,7 +564,9 @@ def image_generation( api_key=api_key or dynamic_api_key, api_base=api_base, logging_obj=litellm_logging_obj, - optional_params=optional_params, + optional_params=( + optional_params if sdk_max_retries is None else {**optional_params, "max_retries": sdk_max_retries} + ), model_response=model_response, organization=organization, aimg_generation=aimg_generation, diff --git a/tests/integration/providers/test_image_gen_drop_params_wire.py b/tests/integration/providers/test_image_gen_drop_params_wire.py index 4307d20c598..f61d7d6dc17 100644 --- a/tests/integration/providers/test_image_gen_drop_params_wire.py +++ b/tests/integration/providers/test_image_gen_drop_params_wire.py @@ -1,16 +1,101 @@ import json -from typing import Final +import threading +from types import MappingProxyType +from typing import Final, Literal +import httpx import openai import pytest from integration._support.client import Gateway from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue, TypeAdapter +import litellm + _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _PROMPT: Final = "a transparent lighthouse at dawn" +@pytest.mark.asyncio +@pytest.mark.parametrize("router_retries", (0, 1)) +@pytest.mark.parametrize("sdk_retries", (0, 5)) +@pytest.mark.parametrize("outcome", ("503", "timeout", "success")) +async def test_proxy_image_retry_attempts_are_owned_only_by_router( + monkeypatch: pytest.MonkeyPatch, + router_retries: int, + sdk_retries: int, + outcome: Literal["503", "timeout", "success"], +) -> None: + from litellm.proxy import proxy_server + + release: Final = threading.Event() + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/v1/images/generations" + assert request.headers["authorization"] == "Bearer synthetic-image-key" + assert _JSON_OBJECT.validate_json(request.body) == {"model": "gpt-image-1-mini", "prompt": "a red circle"} + if outcome == "timeout": + return Reply(chunks=(b"{}",), gate_before_first=release) + if outcome == "503": + return Reply( + status=503, + headers=MappingProxyType({"retry-after-ms": "1"}), + body=b'{"error":{"message":"scripted unavailable","type":"server_error"}}', + ) + return Reply(body=b'{"created":1700000000,"data":[{"b64_json":"aW1n"}]}') + + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "num_retries", None) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + with wire_server(respond) as wire: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "image-retries", + "litellm_params": { + "model": "openai/gpt-image-1-mini", + "api_base": f"{wire.url}/v1", + "api_key": "synthetic-image-key", + "max_retries": sdk_retries, + "timeout": 0.1, + }, + } + ], + num_retries=router_retries, + disable_cooldowns=True, + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "master_key", None) + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(proxy_server, "general_settings", {}) + try: + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=proxy_server.app), base_url="http://proxy.test" + ) as client: + response: Final = await client.post( + "/v1/images/generations", + json={"model": "image-retries", "prompt": "a red circle", "max_retries": sdk_retries}, + ) + finally: + release.set() + router.discard() + litellm.in_memory_llm_clients_cache.flush_cache() + requests: Final = wire.drain() + status: Final = {"503": 503, "timeout": 408, "success": 200}[outcome] + attempts: Final = 1 if outcome == "success" else router_retries + 1 + assert response.status_code == status, response.text + assert len(requests) == attempts, tuple(request.headers.get("x-stainless-retry-count") for request in requests) + if outcome == "success": + data: Final = TypeAdapter(tuple[dict[str, JsonValue], ...]).validate_python( + _JSON_OBJECT.validate_json(response.content)["data"] + ) + assert data[0]["b64_json"] == "aW1n" + + def test_image_generation_additional_drop_params_reaches_provider_body(gateway: Gateway) -> None: def respond(request: Request) -> Reply: assert request.method == "POST" diff --git a/tests/unit/images/test_main.py b/tests/unit/images/test_main.py index 38af99ff505..88d5ae1fc2e 100644 --- a/tests/unit/images/test_main.py +++ b/tests/unit/images/test_main.py @@ -3,14 +3,18 @@ import base64 import json from datetime import datetime from pathlib import Path -from typing import Final +from collections.abc import Generator +from queue import SimpleQueue +from typing import Final, Literal, TypedDict from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx -from openai import AsyncOpenAI +from openai import AsyncOpenAI, OpenAI from openai.types.image import Image +from pydantic import JsonValue, TypeAdapter +from typing_extensions import ReadOnly import litellm from litellm.integrations.custom_logger import CustomLogger @@ -19,6 +23,212 @@ from litellm.types.utils import CallTypes, StandardLoggingPayload import os +class ImageRetryParameters(TypedDict, total=False): + max_retries: ReadOnly[int | None] + + +@pytest.fixture +def image_retry_delays() -> Generator[None, None, None]: + with patch("time.sleep"), patch("asyncio.sleep", new_callable=AsyncMock): + yield + + +def _image_retry_transport( + requests: SimpleQueue[httpx.Request], outcome: Literal["503", "timeout", "success"] +) -> httpx.MockTransport: + def respond(request: httpx.Request) -> httpx.Response: + requests.put(request) + assert request.method == "POST" + assert request.url.path == "/v1/images/generations" + body: Final = TypeAdapter(dict[str, JsonValue]).validate_json(request.content) + assert body == {"model": "gpt-image-1-mini", "prompt": "a red circle"} + if outcome == "timeout": + raise httpx.ReadTimeout("scripted image timeout", request=request) + if outcome == "503": + return httpx.Response(503, json={"error": {"message": "scripted unavailable", "type": "server_error"}}) + return httpx.Response(200, json={"created": 1700000000, "data": [{"b64_json": "aW1n"}]}) + + return httpx.MockTransport(respond) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("image_retry_delays") +@pytest.mark.parametrize("is_async", (False, True), ids=("sync", "async")) +@pytest.mark.parametrize("outcome", ("503", "timeout")) +@pytest.mark.parametrize( + ("retries", "attempts"), + ((0, 1), (1, 2), (None, 3), ("unset", 3), ("global", 3)), + ids=("zero", "one", "none", "unset", "global"), +) +async def test_image_generation_honors_sdk_retry_budget( + monkeypatch: pytest.MonkeyPatch, + is_async: bool, + outcome: Literal["503", "timeout"], + retries: int | None | Literal["unset", "global"], + attempts: int, +) -> None: + requests: Final[SimpleQueue[httpx.Request]] = SimpleQueue() + transport: Final = _image_retry_transport(requests, outcome) + if retries == "global": + monkeypatch.setattr(litellm, "max_retries", 0, raising=False) + retry_params: Final[ImageRetryParameters] = {} if isinstance(retries, str) else {"max_retries": retries} + error: Final = litellm.Timeout if outcome == "timeout" else litellm.ServiceUnavailableError + if is_async: + async with AsyncOpenAI( + api_key="synthetic-image-key", + base_url="https://image.test/v1", + max_retries=5, + http_client=httpx.AsyncClient(transport=transport), + ) as client: + with pytest.raises(error): + await litellm.aimage_generation( + model="openai/gpt-image-1-mini", + prompt="a red circle", + api_key="synthetic-image-key", + api_base="https://image.test/v1", + client=client, + **retry_params, + ) + else: + with OpenAI( + api_key="synthetic-image-key", + base_url="https://image.test/v1", + max_retries=5, + http_client=httpx.Client(transport=transport), + ) as client: + with pytest.raises(error): + litellm.image_generation( + model="openai/gpt-image-1-mini", + prompt="a red circle", + api_key="synthetic-image-key", + api_base="https://image.test/v1", + client=client, + **retry_params, + ) + assert requests.qsize() == attempts + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("image_retry_delays") +@pytest.mark.parametrize("is_async", (False, True), ids=("sync", "async")) +@pytest.mark.parametrize("provider", ("litellm_proxy", "hosted_vllm")) +async def test_compatible_image_generation_honors_zero_retries(provider: str, is_async: bool) -> None: + requests: Final[SimpleQueue[httpx.Request]] = SimpleQueue() + transport: Final = _image_retry_transport(requests, "503") + if is_async: + async with AsyncOpenAI( + api_key="synthetic-image-key", + base_url="https://image.test/v1", + http_client=httpx.AsyncClient(transport=transport), + ) as client: + with pytest.raises(litellm.ServiceUnavailableError): + await litellm.aimage_generation( + model=f"{provider}/gpt-image-1-mini", + prompt="a red circle", + api_key="synthetic-image-key", + api_base="https://image.test/v1", + client=client, + max_retries=0, + ) + else: + with OpenAI( + api_key="synthetic-image-key", + base_url="https://image.test/v1", + http_client=httpx.Client(transport=transport), + ) as client: + with pytest.raises(litellm.ServiceUnavailableError): + litellm.image_generation( + model=f"{provider}/gpt-image-1-mini", + prompt="a red circle", + api_key="synthetic-image-key", + api_base="https://image.test/v1", + client=client, + max_retries=0, + ) + assert requests.qsize() == 1 + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("image_retry_delays") +@pytest.mark.parametrize("is_async", (False, True), ids=("sync", "async")) +async def test_image_generation_rejects_non_integer_retries_before_sending(is_async: bool) -> None: + requests: Final[SimpleQueue[httpx.Request]] = SimpleQueue() + transport: Final = _image_retry_transport(requests, "success") + if is_async: + async with AsyncOpenAI( + api_key="synthetic-image-key", + base_url="https://image.test/v1", + http_client=httpx.AsyncClient(transport=transport), + ) as client: + with pytest.raises(litellm.BadRequestError) as raised: + await litellm.aimage_generation( + model="openai/gpt-image-1-mini", + prompt="a red circle", + api_key="synthetic-image-key", + api_base="https://image.test/v1", + client=client, + max_retries="invalid", + ) + else: + with OpenAI( + api_key="synthetic-image-key", + base_url="https://image.test/v1", + http_client=httpx.Client(transport=transport), + ) as client: + with pytest.raises(litellm.BadRequestError) as raised: + litellm.image_generation( + model="openai/gpt-image-1-mini", + prompt="a red circle", + api_key="synthetic-image-key", + api_base="https://image.test/v1", + client=client, + max_retries="invalid", + ) + assert raised.value.status_code == 422 + assert "max retries must be an int" in raised.value.message + assert requests.qsize() == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", (False, True), ids=("sync", "async")) +async def test_image_generation_returns_success_without_forwarding_retry_control(is_async: bool) -> None: + requests: Final[SimpleQueue[httpx.Request]] = SimpleQueue() + transport: Final = _image_retry_transport(requests, "success") + if is_async: + async with AsyncOpenAI( + api_key="synthetic-image-key", + base_url="https://image.test/v1", + http_client=httpx.AsyncClient(transport=transport), + ) as client: + async_response: Final = await litellm.aimage_generation( + model="openai/gpt-image-1-mini", + prompt="a red circle", + api_key="synthetic-image-key", + api_base="https://image.test/v1", + client=client, + max_retries=0, + ) + async_data: Final = TypeAdapter(tuple[Image, ...]).validate_python(async_response.data) + assert async_data[0].b64_json == "aW1n" + else: + with OpenAI( + api_key="synthetic-image-key", + base_url="https://image.test/v1", + http_client=httpx.Client(transport=transport), + ) as client: + sync_response: Final = litellm.image_generation( + model="openai/gpt-image-1-mini", + prompt="a red circle", + api_key="synthetic-image-key", + api_base="https://image.test/v1", + client=client, + max_retries=0, + ) + sync_data: Final = TypeAdapter(tuple[Image, ...]).validate_python(sync_response.data) + assert sync_data[0].b64_json == "aW1n" + assert requests.qsize() == 1 + + def test_image_generation_keeps_an_internal_prefixed_kwarg_out_of_the_provider_request( respx_mock: respx.MockRouter, ) -> None: diff --git a/tests/unit/test_router_per_deployment_num_retries.py b/tests/unit/test_router_per_deployment_num_retries.py index 99ad7c224f8..b08ae25c438 100644 --- a/tests/unit/test_router_per_deployment_num_retries.py +++ b/tests/unit/test_router_per_deployment_num_retries.py @@ -6,7 +6,10 @@ GitHub Issue: #18968 - Per-deployment max_retries/num_retries in litellm_params import httpx import pytest import pytest_asyncio -from unittest.mock import patch +from queue import SimpleQueue +from typing import Final, Literal +from unittest.mock import AsyncMock, patch +from pydantic import JsonValue, TypeAdapter import litellm from litellm import Router @@ -14,6 +17,70 @@ from litellm.types.router import RetryPolicy from litellm.integrations.custom_logger import CustomLogger +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", (False, True), ids=("sync", "async")) +@pytest.mark.parametrize("source", ("deployment", "default", "request")) +@pytest.mark.parametrize(("sdk_retries", "router_retries"), ((0, 0), (0, 1), (5, 1))) +async def test_routed_images_do_not_multiply_sdk_and_router_retries( + monkeypatch: pytest.MonkeyPatch, + is_async: bool, + source: Literal["deployment", "default", "request"], + sdk_retries: int, + router_retries: int, +) -> None: + requests: Final[SimpleQueue[httpx.Request]] = SimpleQueue() + + def respond(request: httpx.Request) -> httpx.Response: + requests.put(request) + assert request.url.path == "/v1/images/generations" + body: Final = TypeAdapter(dict[str, JsonValue]).validate_json(request.content) + assert body == {"model": "gpt-image-1-mini", "prompt": "a red circle"} + return httpx.Response( + 503, + headers={"retry-after": "0"}, + json={"error": {"message": "scripted unavailable", "type": "server_error"}}, + ) + + transport: Final = httpx.MockTransport(respond) + litellm.in_memory_llm_clients_cache.flush_cache() + with httpx.Client(transport=transport) as sync_session: + async with httpx.AsyncClient(transport=transport) as async_session: + monkeypatch.setattr(litellm, "client_session", sync_session) + monkeypatch.setattr(litellm, "aclient_session", async_session) + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "num_retries", None) + router: Final = Router( + model_list=[ + { + "model_name": "image-retries", + "litellm_params": { + "model": "openai/gpt-image-1-mini", + "api_key": "synthetic-image-key", + "api_base": "https://router-image.test/v1", + **({"max_retries": sdk_retries} if source == "deployment" else {}), + }, + } + ], + default_litellm_params={"max_retries": sdk_retries} if source == "default" else {}, + num_retries=router_retries, + disable_cooldowns=True, + ) + call_params: Final = {"max_retries": sdk_retries} if source == "request" else {} + try: + with patch("time.sleep"), patch("asyncio.sleep", new_callable=AsyncMock): + if is_async: + with pytest.raises(litellm.ServiceUnavailableError): + await router.aimage_generation(model="image-retries", prompt="a red circle", **call_params) + else: + with pytest.raises(litellm.ServiceUnavailableError): + router.image_generation(model="image-retries", prompt="a red circle", **call_params) + finally: + litellm.in_memory_llm_clients_cache.flush_cache() + assert requests.qsize() == router_retries + 1 + + class TestPerDeploymentNumRetries: """Test that per-deployment num_retries in litellm_params is correctly used.""" From 4fa860b64c948373a18830145583711432bf594c Mon Sep 17 00:00:00 2001 From: shrey Date: Wed, 28 Oct 2026 10:47:47 -0700 Subject: [PATCH 055/128] fix(vertex-ai): preserve image size aspect ratios (#48463) --- litellm/llms/gemini/common_utils.py | 57 +++++++-- litellm/llms/vertex_ai/common_utils.py | 31 +++++ .../vertex_gemini_transformation.py | 25 ++-- .../vertex_imagen_transformation.py | 24 ++-- .../vertex_gemini_transformation.py | 32 ++--- .../vertex_imagen_transformation.py | 18 +-- .../llms/gemini/test_gemini_common_utils.py | 34 +++++ ...test_vertex_and_google_ai_studio_gemini.py | 5 +- ...est_vertex_ai_image_edit_transformation.py | 100 ++++++++++++++- ...rtex_ai_image_generation_transformation.py | 117 ++++++++++++++++-- 10 files changed, 349 insertions(+), 94 deletions(-) diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 9634c580858..88c6494e3f7 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -18,6 +18,25 @@ from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat, TokenCountResponse +# 2.5 Flash has no configurable imageSize and supports ten aspect ratios, excluding these four +# https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/gemini/2-5-flash-image +_GEMINI_2_5_FLASH_MODEL_FAMILY: Final = "2.5-flash" +_GEMINI_2_5_FLASH_UNSUPPORTED_ASPECT_RATIOS: Final = frozenset(("1:4", "4:1", "1:8", "8:1")) + +# 3.1 Flash supports the 512 tier, while other Gemini image models start at 1K +# https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/gemini/3-1-flash-image +_GEMINI_512_IMAGE_SIZE_MODEL: Final = "gemini-3.1-flash-image" + +# Flash Lite supports only 1K, even when the requested dimensions are larger +# https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/gemini/3-1-flash-lite-image +_GEMINI_1K_ONLY_IMAGE_SIZE_MODEL: Final = "gemini-3.1-flash-lite-image" + +# LiteLLM chooses the nearest resolution tier using the side length of a square with the same total number of pixels +# as the requested image +_GEMINI_IMAGE_SIZE_1K_THRESHOLD: Final = (512 + 1024) // 2 +_GEMINI_IMAGE_SIZE_2K_THRESHOLD: Final = (1024 + 2048) // 2 +_GEMINI_IMAGE_SIZE_4K_THRESHOLD: Final = (2048 + 4096) // 2 + GEMINI_IMAGE_ASPECT_RATIOS: Final[dict[str, float]] = { "1:1": 1 / 1, "1:4": 1 / 4, @@ -105,8 +124,8 @@ def map_openai_size_to_gemini_image_config(size: str, model: str) -> dict[str, s return None width, height = dimensions - image_config: Final = {"aspectRatio": _map_dimensions_to_gemini_aspect_ratio(width, height)} - image_size: Final = _map_dimensions_to_gemini_image_size(width, height) + image_config: Final = {"aspectRatio": _map_dimensions_to_gemini_aspect_ratio(width, height, model)} + image_size: Final = _map_dimensions_to_gemini_image_size(width, height, model) if is_gemini_image_model(model): if supports_gemini_image_size(model): image_config["imageSize"] = image_size @@ -123,7 +142,7 @@ def supports_gemini_image_size(model: str) -> bool: return bool(value) except Exception: pass - return "2.5-flash" not in model + return _GEMINI_2_5_FLASH_MODEL_FAMILY not in model def is_gemini_image_model(model: str) -> bool: @@ -333,24 +352,36 @@ def _parse_openai_image_size(size: str) -> tuple[int, int] | None: return width, height -def _map_dimensions_to_gemini_aspect_ratio(width: int, height: int) -> str: - if (width, height) in GEMINI_IMAGE_SIZE_TO_ASPECT_RATIO: - return GEMINI_IMAGE_SIZE_TO_ASPECT_RATIO[(width, height)] +def _map_dimensions_to_gemini_aspect_ratio(width: int, height: int, model: str) -> str: + aspect_ratios: Final = tuple( + ratio + for ratio in GEMINI_IMAGE_ASPECT_RATIOS + if _GEMINI_2_5_FLASH_MODEL_FAMILY not in model or ratio not in _GEMINI_2_5_FLASH_UNSUPPORTED_ASPECT_RATIOS + ) + known_ratio: Final = GEMINI_IMAGE_SIZE_TO_ASPECT_RATIO.get((width, height)) + if known_ratio is not None and known_ratio in aspect_ratios: + return known_ratio requested_ratio: Final = width / height return min( - GEMINI_IMAGE_ASPECT_RATIOS, + aspect_ratios, key=lambda aspect_ratio: abs(math.log(GEMINI_IMAGE_ASPECT_RATIOS[aspect_ratio] / requested_ratio)), ) -def _map_dimensions_to_gemini_image_size(width: int, height: int) -> str: - effective_square_side: Final = math.sqrt(width * height) - if effective_square_side < 768: - return "512" - if effective_square_side < 1536: +def _map_dimensions_to_gemini_image_size(width: int, height: int, model: str) -> str: + if _GEMINI_1K_ONLY_IMAGE_SIZE_MODEL in model: return "1K" - if effective_square_side < 3072: + + effective_square_side: Final = math.sqrt(width * height) + if effective_square_side < _GEMINI_IMAGE_SIZE_1K_THRESHOLD: + if is_gemini_image_model(model) and _GEMINI_512_IMAGE_SIZE_MODEL not in model: + return "1K" + else: + return "512" + if effective_square_side < _GEMINI_IMAGE_SIZE_2K_THRESHOLD: + return "1K" + if effective_square_side < _GEMINI_IMAGE_SIZE_4K_THRESHOLD: return "2K" return "4K" diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 27f2aedb300..0fa80cb2cad 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1,3 +1,4 @@ +import math import re from collections.abc import Callable, Iterable, Mapping, Sequence from copy import deepcopy @@ -18,6 +19,11 @@ from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.litellm_core_utils.prompt_templates.common_utils import unpack_defs from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.gemini.common_utils import ( + GEMINI_IMAGE_ASPECT_RATIOS, + GEMINI_IMAGE_SIZE_TO_ASPECT_RATIO, + map_openai_size_to_gemini_image_config, +) from litellm.types.llms.openai import AllMessageValues from litellm.types.llms.vertex_ai import ( VERTEX_AI_PROVIDER_METADATA_FIELDS, @@ -35,6 +41,31 @@ GEMINI_VIDEO_METADATA_KEYS: Final = MappingProxyType( GEMINI_FILES_API_URI_PREFIX: Final = "https://generativelanguage.googleapis.com/v1beta/files/" +# https://cloud.google.com/vertex-ai/generative-ai/docs/image/configure-aspect-ratio +_IMAGEN_ASPECT_RATIOS: Final = MappingProxyType( + {"1:1": 1.0, "3:4": 3 / 4, "4:3": 4 / 3, "9:16": 9 / 16, "16:9": 16 / 9} +) + + +def map_vertex_image_size(size: object, model: str, *, imagen: bool = False) -> Mapping[str, str] | None: + if not isinstance(size, str) or (image_config := map_openai_size_to_gemini_image_config(size, model)) is None: + return None + + width, height = map(int, size.lower().split("x")) + requested_ratio: Final = width / height + aspect_ratio: Final = ( + min(_IMAGEN_ASPECT_RATIOS, key=lambda ratio: abs(math.log(_IMAGEN_ASPECT_RATIOS[ratio] / requested_ratio))) + if imagen + else image_config["aspectRatio"] + ) + if ( + not math.isclose(GEMINI_IMAGE_ASPECT_RATIOS[aspect_ratio], requested_ratio) + and GEMINI_IMAGE_SIZE_TO_ASPECT_RATIO.get((width, height)) != aspect_ratio + ): + verbose_logger.warning("Vertex AI size %s is approximated with aspectRatio %s", size, aspect_ratio) + return MappingProxyType({"aspectRatio": aspect_ratio} if imagen else image_config) + + def gemini_video_metadata_from_openai(video_metadata: Mapping[str, object]) -> dict[str, object]: return { gemini_key: video_metadata[openai_key] diff --git a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py index 6d071363823..0b90f74ea6d 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py @@ -10,7 +10,7 @@ from httpx._types import RequestFiles import litellm from litellm.images.utils import ImageEditRequestUtils from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig -from litellm.llms.vertex_ai.common_utils import get_vertex_base_url +from litellm.llms.vertex_ai.common_utils import get_vertex_base_url, map_vertex_image_size from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams @@ -67,7 +67,7 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): return {} size: Final = image_edit_optional_params.get("size") - return {"aspectRatio": self._map_size_to_aspect_ratio(size or "")} + return dict(map_vertex_image_size(size, model) or {"aspectRatio": "1:1"}) def _resolve_vertex_project(self) -> str | None: return ( @@ -176,11 +176,11 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): contents: Final[dict[str, object]] = {"role": "USER", "parts": parts} # Add image-specific configuration - image_config: Final = ( - {"aspect_ratio": image_edit_optional_request_params["aspectRatio"]} - if "aspectRatio" in image_edit_optional_request_params - else None - ) + image_config: Final = { + field: image_edit_optional_request_params[param] + for param, field in (("aspectRatio", "aspect_ratio"), ("imageSize", "image_size")) + if param in image_edit_optional_request_params + } generation_config: Final[dict[str, object]] = { key: value for key, value in (("response_modalities", ["IMAGE"]), ("image_config", image_config)) if value @@ -222,17 +222,6 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): model_response.data = cast(list[OpenAIImage], data_list) return model_response - def _map_size_to_aspect_ratio(self, size: str) -> str: - """Map OpenAI size format to Gemini aspect ratio format""" - aspect_ratio_map: Final = { - "1024x1024": "1:1", - "1792x1024": "16:9", - "1024x1792": "9:16", - "1280x896": "4:3", - "896x1280": "3:4", - } - return aspect_ratio_map.get(size, "1:1") - def _prepare_inline_image_parts(self, image: FileTypes | list[FileTypes]) -> list[HttpxPartType]: images: Final[list[FileTypes]] = image if isinstance(image, list) else [image] return [ diff --git a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py index deaa32f7308..87b0dfdbcbc 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py @@ -15,7 +15,7 @@ import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig -from litellm.llms.vertex_ai.common_utils import get_vertex_base_url +from litellm.llms.vertex_ai.common_utils import get_vertex_base_url, map_vertex_image_size from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams @@ -70,13 +70,11 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): if "n" in filtered_params: mapped_params["sampleCount"] = filtered_params["n"] - if "size" in filtered_params: - mapped_params["aspectRatio"] = self._map_size_to_aspect_ratio(filtered_params["size"]) - if "mask" in filtered_params: mapped_params["mask"] = filtered_params["mask"] - return mapped_params + size_params: Final = map_vertex_image_size(filtered_params.get("size"), model, imagen=True) or {} + return {**mapped_params, **size_params} def _resolve_vertex_project(self) -> str | None: return ( @@ -189,6 +187,11 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): "sampleCount": sample_count, "editMode": edit_mode, "editConfig": {"baseSteps": base_steps}, + **( + {"aspectRatio": image_edit_optional_request_params["aspectRatio"]} + if "aspectRatio" in image_edit_optional_request_params + else {} + ), } # Set default values for Vertex AI-specific parameters (not configurable by users via OpenAI API) @@ -239,17 +242,6 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): model_response.data = cast(list[OpenAIImage], data_list) return model_response - def _map_size_to_aspect_ratio(self, size: str) -> str: - """Map OpenAI size format to Imagen aspect ratio format""" - aspect_ratio_map: Final = { - "1024x1024": "1:1", - "1792x1024": "16:9", - "1024x1792": "9:16", - "1280x896": "4:3", - "896x1280": "3:4", - } - return aspect_ratio_map.get(size, "1:1") - def _prepare_reference_images( self, image: FileTypes | list[FileTypes], diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py index f6df772dd56..e1049ef8ece 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -14,7 +14,7 @@ from litellm.llms.gemini.common_utils import ( get_gemini_image_web_search_requests, map_gemini_image_tools_params, ) -from litellm.llms.vertex_ai.common_utils import get_vertex_base_url +from litellm.llms.vertex_ai.common_utils import get_vertex_base_url, map_vertex_image_size from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -89,8 +89,7 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): if k == "n": mapped_params["candidate_count"] = v elif k == "size": - # Map OpenAI size format to Gemini aspectRatio - mapped_params["aspectRatio"] = self._map_size_to_aspect_ratio(v) + continue elif k in ("aspectRatio", "aspect_ratio"): mapped_params["aspectRatio"] = v elif k in ("imageSize", "image_size"): @@ -103,22 +102,23 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): elif k not in ("tools", "web_search_options", "imageConfig"): mapped_params[k] = v + if "size" in non_default_params and "size" not in optional_params: + image_params: Final = _JSON_OBJECT.validate_python(mapped_params) + image_config: Final = _JSON_OBJECT.validate_python(image_params.get("imageConfig", {})) + size_params: Final = { + key: value + for key, value in ( + map_vertex_image_size(non_default_params["size"], model) or {"aspectRatio": "1:1"} + ).items() + if key not in image_config + } + return map_gemini_image_tools_params( + _JSON_OBJECT.validate_python(non_default_params), {**size_params, **image_params} + ) + mapped_params = map_gemini_image_tools_params(non_default_params, mapped_params) return mapped_params - def _map_size_to_aspect_ratio(self, size: str) -> str: - """ - Map OpenAI size format to Gemini aspect ratio format - """ - aspect_ratio_map: Final = { - "1024x1024": "1:1", - "1792x1024": "16:9", - "1024x1792": "9:16", - "1280x896": "4:3", - "896x1280": "3:4", - } - return aspect_ratio_map.get(size, "1:1") - def _resolve_vertex_project(self) -> str | None: return ( getattr(self, "_vertex_project", None) diff --git a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py index 927c7b74e67..daf87b1ec58 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py @@ -11,6 +11,7 @@ from litellm.llms.base_llm.image_generation.transformation import ( ) from litellm.llms.vertex_ai.common_utils import ( get_vertex_base_url, + map_vertex_image_size, pop_vertex_request_labels, ) from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM @@ -72,25 +73,14 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): mapped_params["sampleCount"] = v elif k == "size": # Map OpenAI size format to Imagen aspectRatio - mapped_params["aspectRatio"] = self._map_size_to_aspect_ratio(v) + mapped_params["aspectRatio"] = ( + map_vertex_image_size(v, model, imagen=True) or {"aspectRatio": "1:1"} + )["aspectRatio"] else: mapped_params[k] = v return mapped_params - def _map_size_to_aspect_ratio(self, size: str) -> str: - """ - Map OpenAI size format to Imagen aspect ratio format - """ - aspect_ratio_map: Final = { - "1024x1024": "1:1", - "1792x1024": "16:9", - "1024x1792": "9:16", - "1280x896": "4:3", - "896x1280": "3:4", - } - return aspect_ratio_map.get(size, "1:1") - def _resolve_vertex_project(self) -> str | None: return ( getattr(self, "_vertex_project", None) diff --git a/tests/unit/llms/gemini/test_gemini_common_utils.py b/tests/unit/llms/gemini/test_gemini_common_utils.py index 20c36878fef..5f2af34e513 100644 --- a/tests/unit/llms/gemini/test_gemini_common_utils.py +++ b/tests/unit/llms/gemini/test_gemini_common_utils.py @@ -512,6 +512,40 @@ def test_map_openai_image_params_to_gemini_forwards_a_json_object_image_config() assert _map_image_config('{"aspectRatio": "1:1"}') == {"imageConfig": {"aspectRatio": "1:1"}} +@pytest.mark.parametrize("prefix", ("", "gemini/", "vertex_ai/")) +@pytest.mark.parametrize( + "model, size, expected", + ( + ("gemini-3-pro-image", "256x256", {"aspectRatio": "1:1", "imageSize": "1K"}), + ("gemini-3-pro-image", "512x512", {"aspectRatio": "1:1", "imageSize": "1K"}), + ("gemini-3-pro-image-preview", "512x512", {"aspectRatio": "1:1", "imageSize": "1K"}), + ("gemini-3-pro-image", "2048x2048", {"aspectRatio": "1:1", "imageSize": "2K"}), + ("gemini-3-pro-image", "4096x4096", {"aspectRatio": "1:1", "imageSize": "4K"}), + ("gemini-3.1-flash-image", "512x512", {"aspectRatio": "1:1", "imageSize": "512"}), + ("gemini-3.1-flash-image-preview", "512x512", {"aspectRatio": "1:1", "imageSize": "512"}), + ("gemini-3.1-flash-image", "2048x2048", {"aspectRatio": "1:1", "imageSize": "2K"}), + ("gemini-3.1-flash-image", "4096x4096", {"aspectRatio": "1:1", "imageSize": "4K"}), + ("gemini-3.1-flash-lite-image", "256x256", {"aspectRatio": "1:1", "imageSize": "1K"}), + ("gemini-3.1-flash-lite-image", "512x512", {"aspectRatio": "1:1", "imageSize": "1K"}), + ("gemini-3.1-flash-lite-image", "2048x2048", {"aspectRatio": "1:1", "imageSize": "1K"}), + ("gemini-3.1-flash-lite-image", "4096x4096", {"aspectRatio": "1:1", "imageSize": "1K"}), + ("gemini-2.5-flash-image", "512x512", {"aspectRatio": "1:1"}), + ("gemini-2.5-flash-image", "1024x256", {"aspectRatio": "21:9"}), + ("gemini-2.5-flash-image", "256x1024", {"aspectRatio": "9:16"}), + ("gemini-2.5-flash-image", "8000x1000", {"aspectRatio": "21:9"}), + ("gemini-2.5-flash-image", "1000x8000", {"aspectRatio": "9:16"}), + ("gemini-3.1-flash-image", "1024x256", {"aspectRatio": "4:1", "imageSize": "512"}), + ), +) +def test_size_mapping_respects_model_image_limits( + prefix: str, model: str, size: str, expected: dict[str, str] +) -> None: + """Limits as of 2026-10-26, https://ai.google.dev/gemini-api/docs/image-generation#aspect_ratios_and_image_size.""" + assert map_openai_image_params_to_gemini({"size": size}, prefix + model, ["size"]) == { + "imageConfig": expected + } + + @pytest.mark.parametrize("raw_image_config", ["[1, 2]", '"16:9"', "null"]) def test_map_openai_image_params_to_gemini_drops_or_refuses_a_non_object_image_config(raw_image_config: str): with contextlib.suppress(litellm.UnsupportedParamsError): diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 0f5c21d08c7..065d894e658 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -8289,11 +8289,12 @@ def test_gemini_image_generation_ignores_non_dict_image_config(): def test_gemini_image_generation_openai_size_maps_to_google_table( size: str, expected_aspect_ratio: str, expected_image_size: str ): + """Flash table as of 2026-10-26, https://ai.google.dev/gemini-api/docs/image-generation#aspect_ratios_and_image_size.""" from litellm.llms.gemini.common_utils import ( map_openai_size_to_gemini_image_config, ) - assert map_openai_size_to_gemini_image_config(size, "gemini-3-pro-image-preview") == { + assert map_openai_size_to_gemini_image_config(size, "gemini-3.1-flash-image") == { "aspectRatio": expected_aspect_ratio, "imageSize": expected_image_size, } @@ -8317,7 +8318,7 @@ def test_gemini_image_generation_openai_size_snaps_to_nearest_option( map_openai_size_to_gemini_image_config, ) - assert map_openai_size_to_gemini_image_config(size, "gemini-3-pro-image-preview") == { + assert map_openai_size_to_gemini_image_config(size, "gemini-3.1-flash-image") == { "aspectRatio": expected_aspect_ratio, "imageSize": expected_image_size, } diff --git a/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py b/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py index 756e6441f1f..6ae0bbd0d72 100644 --- a/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py +++ b/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py @@ -3,22 +3,118 @@ import json import os from datetime import datetime from io import BytesIO -from typing import Dict +from typing import Dict, Final from unittest.mock import MagicMock, patch import httpx import pytest -from pydantic import ValidationError +from pydantic import JsonValue, TypeAdapter, ValidationError import litellm from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.gemini.common_utils import map_openai_size_to_gemini_image_config from litellm.llms.vertex_ai.image_edit.vertex_gemini_transformation import ( VertexAIGeminiImageEditConfig, ) from litellm.llms.vertex_ai.image_edit.vertex_imagen_transformation import ( VertexAIImagenImageEditConfig, ) +from litellm.types.router import GenericLiteLLMParams + + +@pytest.mark.parametrize("size", ("1536x1024", "1024x1536", "3072x2048", "256x256", "512x512", "4096x4096")) +@pytest.mark.parametrize("model", ("gemini-3.1-flash-image", "gemini-3-pro-image", "gemini-3.1-flash-lite-image")) +def test_vertex_gemini_edit_size_reaches_request(size: str, model: str) -> None: + config: Final = VertexAIGeminiImageEditConfig() + payload, _ = config.transform_image_edit_request( + model, + "mountains", + BytesIO(b"image"), + config.map_openai_params({"size": size}, model, False), + GenericLiteLLMParams(), + {}, + ) + request: Final = TypeAdapter(dict[str, JsonValue]).validate_json(TypeAdapter(str).validate_python(payload)) + expected: Final = map_openai_size_to_gemini_image_config(size, model) or {} + assert request["generationConfig"] == { + "response_modalities": ["IMAGE"], + **( + {"image_config": {"aspect_ratio": expected["aspectRatio"], "image_size": expected["imageSize"]}} + if expected + else {} + ), + } + + +@pytest.mark.parametrize( + "size, expected_ratio", + (("1536x1024", "4:3"), ("1024x1536", "3:4"), ("100000x1", "16:9")), +) +def test_vertex_imagen_edit_size_reaches_request(size: str, expected_ratio: str | None) -> None: + model: Final = "imagen-3.0-capability-001" + config: Final = VertexAIImagenImageEditConfig() + payload, _ = config.transform_image_edit_request( + model, + "mountains", + BytesIO(b"image"), + config.map_openai_params({"size": size, "n": 2}, model, False), + GenericLiteLLMParams(), + {}, + ) + request: Final = TypeAdapter(dict[str, JsonValue]).validate_json(TypeAdapter(str).validate_python(payload)) + assert request["parameters"] == { + "sampleCount": 2, + "editMode": "EDIT_MODE_INPAINT_INSERTION", + "editConfig": {"baseSteps": 50}, + "guidanceScale": 7.5, + "seed": None, + **({"aspectRatio": expected_ratio} if expected_ratio is not None else {}), + } + + +@pytest.mark.parametrize("size", ("auto", "unknown", "0x1024", "-1x1024", "1024x", "3:2")) +@pytest.mark.parametrize( + "config_type, model, config_key, expected", + ( + ( + VertexAIGeminiImageEditConfig, + "gemini-3.1-flash-image", + "generationConfig", + {"response_modalities": ["IMAGE"], "image_config": {"aspect_ratio": "1:1"}}, + ), + ( + VertexAIImagenImageEditConfig, + "imagen-3.0-capability-001", + "parameters", + { + "sampleCount": 1, + "editMode": "EDIT_MODE_INPAINT_INSERTION", + "editConfig": {"baseSteps": 50}, + "guidanceScale": 7.5, + "seed": None, + }, + ), + ), +) +def test_vertex_edit_preserves_legacy_size_fallback( + size: str, + config_type: type[VertexAIGeminiImageEditConfig] | type[VertexAIImagenImageEditConfig], + model: str, + config_key: str, + expected: dict[str, object], +) -> None: + config: Final = config_type() + payload, _ = config.transform_image_edit_request( + model, + "mountains", + BytesIO(b"image"), + config.map_openai_params({"size": size}, model, False), + GenericLiteLLMParams(), + {}, + ) + request: Final = TypeAdapter(dict[str, JsonValue]).validate_json(TypeAdapter(str).validate_python(payload)) + assert request[config_key] == expected class TestVertexAIGeminiImageEditTransformation: diff --git a/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py b/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py index 5ac16283094..10ce24e969c 100644 --- a/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py +++ b/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py @@ -11,6 +11,7 @@ from pydantic import ValidationError import litellm from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.gemini.image_generation.transformation import GoogleImageGenConfig from litellm.llms.vertex_ai.image_generation import ( get_vertex_ai_image_generation_config, ) @@ -23,6 +24,98 @@ from litellm.llms.vertex_ai.image_generation.vertex_imagen_transformation import from litellm.types.utils import ImageResponse +@pytest.mark.parametrize("size", ("1536x1024", "1024x1536", "3072x2048", "256x256", "512x512", "4096x4096")) +@pytest.mark.parametrize("model", ("gemini-3.1-flash-image", "gemini-3-pro-image", "gemini-3.1-flash-lite-image")) +def test_vertex_gemini_size_matches_direct_gemini_request(size: str, model: str) -> None: + vertex: Final = VertexAIGeminiImageGenerationConfig() + direct: Final = GoogleImageGenConfig() + vertex_request: Final = vertex.transform_image_generation_request( + model, "mountains", vertex.map_openai_params({"size": size}, {}, model, False), {}, {} + ) + direct_request: Final = direct.transform_image_generation_request( + model, "mountains", direct.map_openai_params({"size": size}, {}, model, False), {}, {} + ) + assert vertex_request["generationConfig"].get("imageConfig") == direct_request["generationConfig"].get( + "imageConfig" + ) + + +@pytest.mark.parametrize( + "size, expected_ratio", + (("1536x1024", "4:3"), ("1024x1536", "3:4"), ("100000x1", "16:9")), +) +def test_vertex_imagen_generation_size_reaches_request(size: str, expected_ratio: str | None) -> None: + model: Final = "imagen-4.0-generate-001" + config: Final = VertexAIImagenImageGenerationConfig() + request: Final = config.transform_image_generation_request( + model, "mountains", config.map_openai_params({"size": size, "n": 2}, {}, model, False), {}, {} + ) + assert request["parameters"] == { + "sampleCount": 2, + **({"aspectRatio": expected_ratio} if expected_ratio is not None else {}), + } + + +@pytest.mark.parametrize("size", ("auto", "unknown", "0x1024", "-1x1024", "1024x", "3:2")) +@pytest.mark.parametrize( + "config_type, model, config_key, expected", + ( + ( + VertexAIGeminiImageGenerationConfig, + "gemini-3.1-flash-image", + "generationConfig", + {"responseModalities": ["IMAGE"], "imageConfig": {"aspectRatio": "1:1"}}, + ), + ( + VertexAIImagenImageGenerationConfig, + "imagen-4.0-generate-001", + "parameters", + {"sampleCount": 1, "aspectRatio": "1:1"}, + ), + ), +) +def test_vertex_generation_preserves_legacy_size_fallback( + size: str, + config_type: type[VertexAIGeminiImageGenerationConfig] | type[VertexAIImagenImageGenerationConfig], + model: str, + config_key: str, + expected: dict[str, object], +) -> None: + config: Final = config_type() + request: Final = config.transform_image_generation_request( + model, "mountains", config.map_openai_params({"size": size}, {}, model, False), {}, {} + ) + assert request[config_key] == expected + + +@pytest.mark.parametrize("config_type", (VertexAIGeminiImageGenerationConfig, VertexAIImagenImageGenerationConfig)) +def test_vertex_generation_warns_when_size_ratio_is_approximated( + config_type: type[VertexAIGeminiImageGenerationConfig] | type[VertexAIImagenImageGenerationConfig], + caplog: pytest.LogCaptureFixture, +) -> None: + mapped: Final = config_type().map_openai_params({"size": "100000x1"}, {}, "gemini-3.1-flash-image", False) + assert "100000x1" in caplog.text + assert str(mapped["aspectRatio"]) in caplog.text + + +@pytest.mark.parametrize("size_first", (True, False)) +@pytest.mark.parametrize( + "overrides", + ( + {"aspect_ratio": "9:16", "image_size": "2K"}, + {"imageConfig": {"aspectRatio": "9:16", "imageSize": "2K"}}, + ), +) +def test_vertex_gemini_explicit_image_settings_override_size(size_first: bool, overrides: dict[str, object]) -> None: + model: Final = "gemini-3.1-flash-image" + config: Final = VertexAIGeminiImageGenerationConfig() + params: Final = {"size": "1536x1024", **overrides} if size_first else {**overrides, "size": "1536x1024"} + request: Final = config.transform_image_generation_request( + model, "mountains", config.map_openai_params(params, {}, model, False), {}, {} + ) + assert request["generationConfig"]["imageConfig"] == {"aspectRatio": "9:16", "imageSize": "2K"} + + class TestVertexAIGeminiImageGenerationConfig: def setup_method(self): """Set up test fixtures""" @@ -55,14 +148,13 @@ class TestVertexAIGeminiImageGenerationConfig: result = self.config.map_openai_params(non_default_params, optional_params, "gemini-2.5-flash-image", False) assert result.get("aspectRatio") == "16:9" - def test_map_size_to_aspect_ratio(self): - """Test size to aspect ratio mapping""" - assert self.config._map_size_to_aspect_ratio("1024x1024") == "1:1" - assert self.config._map_size_to_aspect_ratio("1792x1024") == "16:9" - assert self.config._map_size_to_aspect_ratio("1024x1792") == "9:16" - assert self.config._map_size_to_aspect_ratio("1280x896") == "4:3" - assert self.config._map_size_to_aspect_ratio("896x1280") == "3:4" - assert self.config._map_size_to_aspect_ratio("unknown") == "1:1" # default + @pytest.mark.parametrize( + "size, expected_ratio", + (("1024x1024", "1:1"), ("1792x1024", "16:9"), ("1024x1792", "9:16"), ("1280x896", "4:3"), ("896x1280", "3:4")), + ) + def test_map_openai_params_legacy_sizes(self, size: str, expected_ratio: str) -> None: + mapped: Final = self.config.map_openai_params({"size": size}, {}, "gemini-3.1-flash-image", False) + assert mapped["aspectRatio"] == expected_ratio def test_get_supported_openai_params_includes_native_gemini_params(self): """Test that native Gemini imageConfig params are supported""" @@ -474,11 +566,10 @@ class TestVertexAIImagenImageGenerationConfig: result = self.config.map_openai_params(non_default_params, optional_params, "imagegeneration@006", False) assert result.get("aspectRatio") == "1:1" - def test_map_size_to_aspect_ratio(self): - """Test size to aspect ratio mapping""" - assert self.config._map_size_to_aspect_ratio("1024x1024") == "1:1" - assert self.config._map_size_to_aspect_ratio("1792x1024") == "16:9" - assert self.config._map_size_to_aspect_ratio("unknown") == "1:1" # default + @pytest.mark.parametrize("size, expected_ratio", (("1024x1024", "1:1"), ("1792x1024", "16:9"))) + def test_map_openai_params_legacy_sizes(self, size: str, expected_ratio: str) -> None: + mapped: Final = self.config.map_openai_params({"size": size}, {}, "imagen-4.0-generate-001", False) + assert mapped["aspectRatio"] == expected_ratio def test_transform_image_generation_request_basic(self): """Test basic request transformation""" From 74392efbd4292dd79e9cc769c8a56e44d4569346 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Wed, 28 Oct 2026 10:48:34 -0700 Subject: [PATCH 056/128] fix(otel): label configured pass-through routes litellm.passthrough and read the operation from the upstream (#48675) --- .../opentelemetry_utils/gen_ai_semconv.py | 6 +- litellm/integrations/otel/logger.py | 1 + litellm/integrations/otel/model/metadata.py | 6 +- litellm/integrations/otel/model/payloads.py | 7 +- litellm/integrations/otel/model/semconv.py | 10 +- litellm/integrations/otel/plumbing/metrics.py | 5 +- .../pass_through_endpoints.py | 12 ++ .../upstream_usage_headers.py | 23 ++- .../test_passthrough_operation_name.py | 163 ++++++++++++++++++ .../integrations/otel/test_otel_v2_metrics.py | 16 ++ .../otel/test_otel_v2_sources_of_truth.py | 11 ++ tests/unit/integrations/test_opentelemetry.py | 10 ++ .../test_upstream_usage_headers.py | 23 +++ 13 files changed, 284 insertions(+), 9 deletions(-) create mode 100644 tests/integration/observability/test_passthrough_operation_name.py diff --git a/litellm/integrations/opentelemetry_utils/gen_ai_semconv.py b/litellm/integrations/opentelemetry_utils/gen_ai_semconv.py index 035e5c5db10..84c4b2e58ab 100644 --- a/litellm/integrations/opentelemetry_utils/gen_ai_semconv.py +++ b/litellm/integrations/opentelemetry_utils/gen_ai_semconv.py @@ -35,6 +35,7 @@ from enum import Enum from typing import TYPE_CHECKING, Final, Protocol from litellm.integrations.otel.mappers.utils import drop_none +from litellm.integrations.otel.model.semconv import PASSTHROUGH_OPERATION_KEY, GenAIOperation from litellm.litellm_core_utils.safe_json_dumps import safe_dumps if TYPE_CHECKING: @@ -148,8 +149,11 @@ class OTELGenAISemconvMixin: """Map a LiteLLM ``call_type`` to spec ``gen_ai.operation.name``. Substring match (e.g. ``aembedding`` -> ``embeddings``); defaults to - ``chat``. + ``chat``. A pass-through call's operation, stored under ``PASSTHROUGH_OPERATION_KEY``, wins. """ + passthrough_operation: Final = kwargs.get(PASSTHROUGH_OPERATION_KEY) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # kwargs is an untyped dict + if isinstance(passthrough_operation, GenAIOperation): + return passthrough_operation.value call_type: Final = kwargs.get("call_type", "") or "" match call_type: case s if "embedding" in s: diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index e09226ac9cd..4c5e65ff07d 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -600,6 +600,7 @@ class OpenTelemetryV2(CustomLogger): trace=call.trace, session_id=call.session_id, metadata_keys=tuple(self.config.baggage_metadata_keys), + operation=call.operation, ) end_time_ns: Final = to_ns(end_time) if carrier is not None and carrier.span is not None: diff --git a/litellm/integrations/otel/model/metadata.py b/litellm/integrations/otel/model/metadata.py index b96ed55872d..a9a57de04ac 100644 --- a/litellm/integrations/otel/model/metadata.py +++ b/litellm/integrations/otel/model/metadata.py @@ -49,7 +49,7 @@ from litellm.constants import ( LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL, SESSION_ID_GENERATED_METADATA_KEY, ) -from litellm.integrations.otel.model.semconv import resolve_operation +from litellm.integrations.otel.model.semconv import PASSTHROUGH_OPERATION_KEY, GenAIOperation, resolve_operation from litellm.integrations.otel.model.trace_controls import TraceControls, caller_trace_controls from litellm.integrations.otel.model.utils import as_str, as_str_mapping, to_seconds from litellm.types.utils import InternalCallOrigin @@ -262,6 +262,7 @@ class LLMCallEvent: # never hands off to a provider, so ``upstream_started`` stays False even # though the call completed and is worth a span. cache_hit: bool + operation: GenAIOperation # A best-effort ``"{operation} {model}"`` name known at ``pre_call`` time. The # span is renamed from the typed payload at close (``finish_span``); this only # needs to be reasonable for a span that never gets closed (a leak). @@ -275,7 +276,7 @@ class LLMCallEvent: def from_dict(cls, kwargs: Mapping[str, object]) -> LLMCallEvent: raw_payload: Final = kwargs.get("standard_logging_object") payload: Final = cast("StandardLoggingPayload", raw_payload) if raw_payload else None - operation: Final = resolve_operation(as_str(kwargs.get("call_type"))) + operation: Final = resolve_operation(as_str(kwargs.get("call_type")), kwargs.get(PASSTHROUGH_OPERATION_KEY)) model: Final = as_str(kwargs.get("model")) or "" trace: Final = caller_trace_controls(kwargs) return cls( @@ -288,6 +289,7 @@ class LLMCallEvent: upstream_start_seconds=_epoch_seconds(kwargs.get("api_call_start_time")), purpose=internal_call_origin(payload, kwargs), cache_hit=bool(payload.get("cache_hit")) if payload else False, # pyright: ignore[reportUnknownMemberType] # TypedDict .get on a str key is partially unknown + operation=operation, 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")), diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index 21b2f44f077..7a62ed6aed9 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -464,6 +464,7 @@ class LLMCallSpanData: session_id: str | None = None, *, metadata_keys: tuple[str, ...] = (), + operation: GenAIOperation | None = None, ) -> LLMCallSpanData: params: Final = cast(Mapping[str, object], payload.get("model_parameters") or {}) cache_hit: Final = as_bool(payload.get("cache_hit")) is True # pyright: ignore[reportUnknownMemberType] # TypedDict .get on a str key is partially unknown @@ -486,13 +487,13 @@ class LLMCallSpanData: # no prompt/response text. finish_reasons: Final = _finish_reasons(choices_out) call_type: Final = as_str(payload.get("call_type")) - operation: Final = resolve_operation(call_type) + resolved_operation: Final = operation or resolve_operation(call_type) embedding_output: Final = ( - EmbeddingOutput.from_response(response) if operation is GenAIOperation.EMBEDDINGS else None + EmbeddingOutput.from_response(response) if resolved_operation is GenAIOperation.EMBEDDINGS else None ) cost_lookup_failed: Final = payload.get("response_cost_failure_debug_info") is not None and not cache_hit return cls( - operation=operation, + operation=resolved_operation, provider=resolve_provider(as_str(payload.get("custom_llm_provider"))), request_model=context.request_model, response_model=context.response_model, diff --git a/litellm/integrations/otel/model/semconv.py b/litellm/integrations/otel/model/semconv.py index b7f955a75ed..8d186a2715a 100644 --- a/litellm/integrations/otel/model/semconv.py +++ b/litellm/integrations/otel/model/semconv.py @@ -34,6 +34,7 @@ class GenAIOperation(str, Enum): LITELLM_VECTOR_STORE_FILE_MANAGEMENT = "litellm.vector_store_file_management" LITELLM_RESPONSES_MANAGEMENT = "litellm.responses_management" LITELLM_MODERATION = "litellm.moderation" + LITELLM_PASSTHROUGH = "litellm.passthrough" class GenAIOutputType(str, Enum): @@ -482,14 +483,21 @@ def resolve_provider(custom_llm_provider: str | None) -> str: return mapped.value if mapped is not None else custom_llm_provider -def resolve_operation(call_type: str | None) -> GenAIOperation: +PASSTHROUGH_OPERATION_KEY: Final = "passthrough_gen_ai_operation" + + +def resolve_operation(call_type: str | None, passthrough_operation: object = None) -> GenAIOperation: """Map a litellm ``call_type`` to a ``gen_ai.operation.name`` value. An unmapped call type still falls back to ``chat`` so every series keeps an operation label, but it logs at debug rather than falling through silently: a new call type mislabelled as ``chat`` mixes its latency and cost into everyone's chat charts, which is invisible until someone reads the numbers. + + A pass-through call's operation, stored under ``PASSTHROUGH_OPERATION_KEY``, wins over its call type. """ + if isinstance(passthrough_operation, GenAIOperation): + return passthrough_operation if not call_type: return GenAIOperation.CHAT mapped: Final = _OPERATION_BY_CALL_TYPE.get(call_type.lower()) diff --git a/litellm/integrations/otel/plumbing/metrics.py b/litellm/integrations/otel/plumbing/metrics.py index 72d51471ba2..1197ae055d0 100644 --- a/litellm/integrations/otel/plumbing/metrics.py +++ b/litellm/integrations/otel/plumbing/metrics.py @@ -26,6 +26,7 @@ from litellm.integrations.opentelemetry import ( ) from litellm.integrations.otel.model.metadata import time_to_first_chunk_seconds from litellm.integrations.otel.model.semconv import ( + PASSTHROUGH_OPERATION_KEY, Error, GenAI, Metric, @@ -280,7 +281,9 @@ class GenAIMetricRecorder: def _common_attributes(self, kwargs: Mapping[str, Any]) -> dict: params: Final = kwargs.get("litellm_params") or {} common_attrs: Final[dict] = { - GenAI.OPERATION_NAME: resolve_operation(kwargs.get("call_type")).value, + GenAI.OPERATION_NAME: resolve_operation( + kwargs.get("call_type"), kwargs.get(PASSTHROUGH_OPERATION_KEY) + ).value, **_provider_attributes(params.get("custom_llm_provider")), GenAI.REQUEST_MODEL: kwargs.get("model"), "gen_ai.framework": "litellm", diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index f2ea1752a86..650dcae3239 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -54,6 +54,7 @@ from litellm.integrations.custom_guardrail import ( without_server_streaming_classification, ) from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.otel.model.semconv import PASSTHROUGH_OPERATION_KEY, GenAIOperation from litellm.litellm_core_utils.core_helpers import ( bind_budget_reservation_to_callbacks, get_metadata_variable_name_from_kwargs, @@ -151,6 +152,7 @@ from .success_handler import PassThroughEndpointLogging from .upstream_usage_headers import ( UpstreamReportedUsage, apply_upstream_reported_model, + apply_upstream_reported_operation, apply_upstream_reported_usage, ) @@ -1134,6 +1136,7 @@ async def pass_through_request( guardrails_config: dict | None = None, timeout: float | None = None, upstream_headers_filter: Callable[[Mapping[str, str]], Mapping[str, str]] = _unfiltered_headers, + operation: GenAIOperation | None = None, ): """ Pass through endpoint handler, makes the httpx request for pass-through endpoints and ensures logging hooks are called @@ -1155,6 +1158,7 @@ async def pass_through_request( timeout: Optional per-endpoint timeout in seconds. Falls back to general_settings.pass_through_request_timeout, then 600s. upstream_headers_filter: Optional last rewrite of the headers sent upstream, after x-pass- expansion + operation: Optional GenAI operation the route reports, unless the upstream names its own """ from litellm.exceptions import ModifyResponseException from litellm.litellm_core_utils.litellm_logging import Logging @@ -1353,6 +1357,8 @@ async def pass_through_request( call_type="pass_through_endpoint", ) logging_obj.model_call_details["litellm_call_id"] = litellm_call_id + if operation is not None: + logging_obj.model_call_details[PASSTHROUGH_OPERATION_KEY] = operation ## PASSTHROUGH MANAGED ID RESOLUTION (INPUT) ## # Resolve managed IDs in path, query params, and body back to raw @@ -1553,6 +1559,7 @@ async def pass_through_request( logging_obj=logging_obj, headers=response.headers, ) + apply_upstream_reported_operation(logging_obj, response.headers) relay_response: Final = await _log_passthrough_upstream_failure( response=response, @@ -1646,6 +1653,7 @@ async def pass_through_request( logging_obj=logging_obj, headers=response.headers, ) + apply_upstream_reported_operation(logging_obj, response.headers) if _is_streaming_response(response) is True: logging_obj.stream = True @@ -2161,6 +2169,7 @@ def create_pass_through_route( config_file_path: str | None = None, timeout: float | None = None, upstream_headers_filter: Callable[[Mapping[str, str]], Mapping[str, str]] = _unfiltered_headers, + operation: GenAIOperation | None = None, ): # check if target is an adapter.py or a url from litellm._uuid import uuid @@ -2301,6 +2310,7 @@ def create_pass_through_route( guardrails_config=cast(dict | None, param_guardrails), timeout=cast(float | None, param_timeout), upstream_headers_filter=upstream_headers_filter, + operation=operation, ) finally: if hasattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY): @@ -3232,6 +3242,7 @@ class InitPassThroughEndpointHelpers: guardrails=guardrails, config_file_path=config_file_path, timeout=timeout, + operation=GenAIOperation.LITELLM_PASSTHROUGH, ), methods=methods, dependencies=dependencies, @@ -3316,6 +3327,7 @@ class InitPassThroughEndpointHelpers: guardrails=guardrails, config_file_path=config_file_path, timeout=timeout, + operation=GenAIOperation.LITELLM_PASSTHROUGH, ), methods=methods, dependencies=dependencies, diff --git a/litellm/proxy/pass_through_endpoints/upstream_usage_headers.py b/litellm/proxy/pass_through_endpoints/upstream_usage_headers.py index e9d93fb306d..223f6a5e760 100644 --- a/litellm/proxy/pass_through_endpoints/upstream_usage_headers.py +++ b/litellm/proxy/pass_through_endpoints/upstream_usage_headers.py @@ -5,18 +5,21 @@ price the request from the response body. Those targets report the totals for the whole HTTP request in ``x-litellm-*`` response headers instead; LiteLLM records the reported values without recomputing them. A target may also name the request in ``x-litellm-upstream-model`` so its spend can be told apart from -other requests that carry no ``model`` in the body. +other requests that carry no ``model`` in the body, and name the GenAI +operation it performed, which labels the request's trace. """ import math from collections.abc import Mapping from dataclasses import dataclass +from types import MappingProxyType from typing import Final import httpx from litellm._logging import verbose_proxy_logger from litellm.constants import MAX_SPEND_LOG_MODEL_NAME_LENGTH +from litellm.integrations.otel.model.semconv import PASSTHROUGH_OPERATION_KEY, GenAIOperation from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.utils import Usage @@ -25,6 +28,8 @@ UPSTREAM_TOTAL_TOKENS_HEADER: Final = "x-litellm-total-tokens" # Deliberately not x-litellm-model-name/-id/-group: the gateway already sends # those to clients, and pass-through relays upstream headers back to the client. UPSTREAM_MODEL_HEADER: Final = "x-litellm-upstream-model" +UPSTREAM_OPERATION_NAME_HEADER: Final = "x-litellm-operation-name" +_OPERATIONS_BY_NAME: Final = MappingProxyType({operation.value: operation for operation in GenAIOperation}) # model_call_details key holding what the upstream reported, so later stages of # the success path can tell an upstream-reported cost apart from one LiteLLM @@ -175,3 +180,19 @@ def has_upstream_reported_usage(logging_obj: LiteLLMLoggingObj) -> bool: means zero, never a fallback to someone else's estimate. """ return isinstance(logging_obj.model_call_details.get(UPSTREAM_REPORTED_USAGE_KEY), UpstreamReportedUsage) + + +def apply_upstream_reported_operation(logging_obj: LiteLLMLoggingObj, headers: Mapping[str, str]) -> None: + """Record the GenAI operation the upstream named, overriding the route's label. An unknown name is ignored.""" + raw_operation: Final = headers.get(UPSTREAM_OPERATION_NAME_HEADER) + if raw_operation is None: + return + operation: Final = _OPERATIONS_BY_NAME.get(raw_operation) + if operation is None: + verbose_proxy_logger.warning( + "pass_through_endpoint: upstream sent unknown %s=%r; ignoring it", + UPSTREAM_OPERATION_NAME_HEADER, + raw_operation, + ) + return + logging_obj.model_call_details[PASSTHROUGH_OPERATION_KEY] = operation # rebind-ok: annotates the request's record diff --git a/tests/integration/observability/test_passthrough_operation_name.py b/tests/integration/observability/test_passthrough_operation_name.py new file mode 100644 index 00000000000..7b01389a993 --- /dev/null +++ b/tests/integration/observability/test_passthrough_operation_name.py @@ -0,0 +1,163 @@ +import json +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, eventually, gateway_from_environment +from integration._support.otlp_sink import Span, SpanSinks, recorded_spans +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server + +pytestmark: Final = pytest.mark.timeout(180) + +_EXPORTERS: Final = { + "otel-v1-semconv": {"OTEL_SEMCONV_STABILITY_OPT_IN": "gen_ai_latest_experimental"}, + "otel-v2": {"LITELLM_OTEL_V2": "1"}, +} +_ANTHROPIC_MESSAGE: Final = { + "id": "msg_passthrough_operation", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 3, "output_tokens": 1}, +} +_STREAM_REQUESTS: Final = {"stream-requested": {"query": "q", "stream": True}, "stream-detected": {"query": "q"}} + + +@dataclass(frozen=True, slots=True) +class _Rig: + proxy: Gateway + sink: str + + +def _upstream(request: Request) -> Reply: + if request.target.endswith("/v1/messages"): + return Reply(body=json.dumps(_ANTHROPIC_MESSAGE).encode()) + _, _, reported = request.target.partition("/report/") + headers: Final = {"x-litellm-operation-name": reported} if reported else {} + if request.target.startswith("/app/stream/cut"): + return Reply( + content_type="text/event-stream", + chunks=(b'data: {"chunk": 1}\n\n', b'data: {"chunk": 2}\n\n'), + abort_after=1, + ) + if request.target.startswith("/app/stream"): + return Reply(content_type="text/event-stream", chunks=(b'data: {"chunk": 1}\n\n',), headers=headers) + return Reply(body=b'{"chunks": []}', headers=headers) + + +@pytest.fixture(scope="module", params=sorted(_EXPORTERS)) +def rig( + request: pytest.FixtureRequest, tmp_path_factory: pytest.TempPathFactory, audit_sinks: SpanSinks +) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("passthrough-operation") + exporter: Final = _EXPORTERS[request.param] + with gateway_from_environment() as base, wire_server(_upstream) as upstream: + shared: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config: Final = { + **shared, + "litellm_settings": {**shared["litellm_settings"], "callbacks": ["otel"]}, + "callback_settings": { + "otel": {"exporter": "http/json", "endpoint": audit_sinks.operator, "use_simple_processor": True} + }, + "general_settings": { + **shared["general_settings"], + "pass_through_endpoints": [ + {"path": "/doc-app", "target": upstream.url + "/app", "include_subpath": True} + ], + }, + "environment_variables": { + **shared["environment_variables"], + "ANTHROPIC_API_BASE": upstream.url, + "ANTHROPIC_API_KEY": "scripted", + }, + } + path: Final = directory / "passthrough-operation.yaml" + path.write_text(yaml.safe_dump(config)) + overrides: Final = {"OTEL_EXPORTER": "http/json", "OTEL_ENDPOINT": audit_sinks.operator, **exporter} + unset: Final = tuple( + name for name in ("LITELLM_OTEL_V2", "OTEL_SEMCONV_STABILITY_OPT_IN") if name not in exporter + ) + with owned_proxy(base, directory, overrides, config=path, remove_environment=unset) as proxy: + yield _Rig(proxy, audit_sinks.operator) + + +def _span_of(rig: _Rig, since: int, call_id: str) -> Span: + spans: Final = eventually( + lambda: tuple( + span + for span in recorded_spans(rig.sink, since)[1] + if span["attributes"].get("litellm.call_id") == call_id and "gen_ai.operation.name" in span["attributes"] + ), + lambda found: len(found) == 1, + seconds=60, + return_last_on_timeout=True, + ) + assert len(spans) == 1, spans + return spans[0] + + +def _operation_span(rig: _Rig, method: str, path: str, body: dict[str, object] | None = None) -> Span: + since, _ = recorded_spans(rig.sink) + response: Final = rig.proxy.request(method, path, body) + assert response.status_code == 200, response.text + return _span_of(rig, since, response.headers["x-litellm-call-id"]) + + +def _cut_stream_span(rig: _Rig, path: str) -> Span: + since, _ = recorded_spans(rig.sink) + headers: Final = {"Authorization": f"Bearer {rig.proxy.key}"} + with rig.proxy.client.stream("POST", path, json={"query": "q", "stream": True}, headers=headers) as response: + assert response.status_code == 200 + with pytest.raises(httpx.RemoteProtocolError): + response.read() + return _span_of(rig, since, response.headers["x-litellm-call-id"]) + + +def _assert_operation(span: Span, operation: str) -> None: + assert span["attributes"]["gen_ai.operation.name"] == operation, span + assert span["name"].split(" ")[0] == operation, span + + +def test_configured_route_without_the_header_is_labelled_passthrough(rig: _Rig) -> None: + _assert_operation(_operation_span(rig, "POST", "/doc-app/search", {"query": "q"}), "litellm.passthrough") + + +def test_operation_reported_by_the_target_labels_the_span(rig: _Rig) -> None: + _assert_operation(_operation_span(rig, "POST", "/doc-app/report/retrieval", {"query": "q"}), "retrieval") + + +def test_unknown_operation_reported_by_the_target_is_ignored(rig: _Rig) -> None: + _assert_operation(_operation_span(rig, "POST", "/doc-app/report/summarize", {"query": "q"}), "litellm.passthrough") + + +@pytest.mark.parametrize("body", _STREAM_REQUESTS.values(), ids=_STREAM_REQUESTS.keys()) +def test_streamed_configured_route_without_the_header_is_labelled_passthrough( + rig: _Rig, body: dict[str, object] +) -> None: + _assert_operation(_operation_span(rig, "POST", "/doc-app/stream", body), "litellm.passthrough") + + +@pytest.mark.parametrize("body", _STREAM_REQUESTS.values(), ids=_STREAM_REQUESTS.keys()) +def test_operation_reported_by_a_streamed_target_labels_the_span(rig: _Rig, body: dict[str, object]) -> None: + _assert_operation(_operation_span(rig, "POST", "/doc-app/stream/report/retrieval", body), "retrieval") + + +def test_configured_route_stream_cut_by_the_target_is_labelled_passthrough(rig: _Rig) -> None: + _assert_operation(_cut_stream_span(rig, "/doc-app/stream/cut"), "litellm.passthrough") + + +def test_provider_route_stays_chat(rig: _Rig) -> None: + body: Final = {"model": "claude-sonnet-5-5", "max_tokens": 8, "messages": [{"role": "user", "content": "hi"}]} + _assert_operation(_operation_span(rig, "POST", "/anthropic/v1/messages", body), "chat") + + +def test_provider_route_litellm_does_not_parse_stays_chat(rig: _Rig) -> None: + _assert_operation(_operation_span(rig, "GET", "/anthropic/v1/models"), "chat") diff --git a/tests/unit/integrations/otel/test_otel_v2_metrics.py b/tests/unit/integrations/otel/test_otel_v2_metrics.py index 016dbcd824b..74921512825 100644 --- a/tests/unit/integrations/otel/test_otel_v2_metrics.py +++ b/tests/unit/integrations/otel/test_otel_v2_metrics.py @@ -24,6 +24,7 @@ dashboard already querying that histogram. import asyncio import json from datetime import datetime, timedelta +from typing import Final import pytest @@ -40,6 +41,10 @@ from litellm.integrations.otel.logger import OpenTelemetryV2 # noqa: E402 from litellm.integrations.otel.model.config import ( # noqa: E402 OpenTelemetryV2Config, ) +from litellm.integrations.otel.model.semconv import ( # noqa: E402 + PASSTHROUGH_OPERATION_KEY, + GenAIOperation, +) from litellm.integrations.otel.plumbing.metrics import ( # noqa: E402 GenAIMetricRecorder, create_genai_metrics, @@ -500,6 +505,17 @@ def test_agent_message_is_not_labelled_as_chat(call_type): assert dp.attributes[OPERATION_KEY] == "invoke_agent" +def test_passthrough_operation_labels_the_metrics() -> None: + reader: Final = InMemoryMetricReader() + kwargs, response_obj, start, end = _build_call(call_type="pass_through_endpoint") + labelled: Final = {**kwargs, PASSTHROUGH_OPERATION_KEY: GenAIOperation.RETRIEVAL} + + asyncio.run(_logger(reader, enable_metrics=True).async_log_success_event(labelled, response_obj, start, end)) + + for dp in _metrics_by_name(reader)[OPERATION_DURATION]: + assert dp.attributes[OPERATION_KEY] == "retrieval" + + def test_provider_name_is_filterable(): """gen_ai.provider.name is a member of the metric-attribute allowlist, so an operator can include or exclude it; an unlisted name raises instead.""" diff --git a/tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py b/tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py index 318859cb3ed..2d0b4701c9f 100644 --- a/tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py +++ b/tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py @@ -38,6 +38,7 @@ from litellm.integrations.otel.model.payloads import ( RequestIdentity, _upstream_address_port, ) +from litellm.integrations.otel.model.semconv import PASSTHROUGH_OPERATION_KEY from litellm.integrations.otel.model.spans import ( SPAN_REGISTRY, LiteLLMSpanKind, @@ -223,6 +224,16 @@ def test_operation_resolution(): assert resolve_operation("call_mcp_tool") is GenAIOperation.EXECUTE_TOOL +def test_passthrough_operation_wins_over_the_call_type() -> None: + assert resolve_operation("pass_through_endpoint", GenAIOperation.RETRIEVAL) is GenAIOperation.RETRIEVAL + assert resolve_operation("pass_through_endpoint", "retrieval") is GenAIOperation.CHAT + event: Final = LLMCallEvent.from_dict( + {"call_type": "pass_through_endpoint", "model": "unknown", PASSTHROUGH_OPERATION_KEY: GenAIOperation.RETRIEVAL} + ) + assert event.operation is GenAIOperation.RETRIEVAL + assert event.provisional_span_name == "retrieval unknown" + + @pytest.mark.parametrize("call_type", ["vector_store_search", "avector_store_search"]) def test_vector_store_search_is_a_retrieval_operation(call_type): """A vector-store search is a retrieval, so its duration and cost must not diff --git a/tests/unit/integrations/test_opentelemetry.py b/tests/unit/integrations/test_opentelemetry.py index f908c685382..f0fca94c370 100644 --- a/tests/unit/integrations/test_opentelemetry.py +++ b/tests/unit/integrations/test_opentelemetry.py @@ -48,6 +48,7 @@ from litellm.integrations.opentelemetry import ( RAW_REQUEST_SPAN_NAME, _normalize_team_metadata_keys, ) +from litellm.integrations.otel.model.semconv import PASSTHROUGH_OPERATION_KEY, GenAIOperation from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy import proxy_server from litellm.proxy._types import ManagementEndpointLoggingPayload, ProxyException, SpanAttributes @@ -1048,6 +1049,15 @@ class TestOpenTelemetrySemconvStability(unittest.TestCase): kwargs = {"model": "claude-sonnet-4-5", "call_type": "unknown"} self.assertEqual(h._get_span_name(kwargs), "chat claude-sonnet-4-5") + def test_passthrough_operation_names_the_span(self) -> None: + h: Final = self._make(env="gen_ai_latest_experimental") + kwargs: Final = { + "model": "unknown", + "call_type": "pass_through_endpoint", + PASSTHROUGH_OPERATION_KEY: GenAIOperation.RETRIEVAL, + } + self.assertEqual(h._get_span_name(kwargs), "retrieval unknown") + def test_generation_name_metadata_overrides_semconv_name(self): h = self._make(env="gen_ai_latest_experimental") kwargs = { diff --git a/tests/unit/proxy/pass_through_endpoints/test_upstream_usage_headers.py b/tests/unit/proxy/pass_through_endpoints/test_upstream_usage_headers.py index 9dd0233ad8b..33328e418cd 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_upstream_usage_headers.py +++ b/tests/unit/proxy/pass_through_endpoints/test_upstream_usage_headers.py @@ -4,10 +4,12 @@ import httpx import pytest from litellm.constants import MAX_SPEND_LOG_MODEL_NAME_LENGTH +from litellm.integrations.otel.model.semconv import PASSTHROUGH_OPERATION_KEY, GenAIOperation from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.pass_through_endpoints.upstream_usage_headers import ( UpstreamReportedUsage, apply_upstream_reported_model, + apply_upstream_reported_operation, apply_upstream_reported_usage, has_upstream_reported_usage, parse_upstream_reported_model, @@ -211,3 +213,24 @@ def test_model_header_alone_does_not_make_the_upstream_own_the_cost(): assert has_upstream_reported_usage(logging_obj) is False assert "response_cost" not in logging_obj.model_call_details + + +@pytest.mark.parametrize("operation", [GenAIOperation.RETRIEVAL, GenAIOperation.EMBEDDINGS]) +def test_apply_operation_records_the_operation_the_upstream_named(operation: GenAIOperation) -> None: + logging_obj: Final = _logging_obj() + + apply_upstream_reported_operation(logging_obj, _headers(**{"x-litellm-operation-name": operation.value})) + + assert logging_obj.model_call_details[PASSTHROUGH_OPERATION_KEY] is operation + + +@pytest.mark.parametrize( + "headers", + [{}, {"x-litellm-operation-name": "summarize"}, {"x-litellm-operation-name": "Retrieval"}], +) +def test_apply_operation_ignores_a_missing_or_unknown_name(headers: dict[str, str]) -> None: + logging_obj: Final = _logging_obj() + + apply_upstream_reported_operation(logging_obj, _headers(**headers)) + + assert PASSTHROUGH_OPERATION_KEY not in logging_obj.model_call_details From d9aae590e344625a86cdccea2b853f8fda709fa9 Mon Sep 17 00:00:00 2001 From: "Onur C. Cakmak" Date: Wed, 28 Oct 2026 13:57:51 -0400 Subject: [PATCH 057/128] fix(chatgpt): derive the upstream session id from prompt_cache_key (#42014) * fix(chatgpt): derive Responses session id from prompt_cache_key ChatGPT partitions the Codex prompt cache by the Responses session-id header (the Codex CLI sends a stable per-conversation key derived from prompt_cache_key). The provider only consulted explicit session ids and otherwise stamped a per-request uuid4, so no turn ever reused the prompt cache -- 0% cache hits on repeat turns, an order-of-magnitude quota/cost multiplier on subscription accounts. Resolution order after this change: explicit session config (session_id / litellm_session_id / metadata.session_id) > prompt_cache_key (sanitized with _safe_header_value) > litellm-internal ids > uuid4. Two guards, mirroring get_fireworks_session_id: - proxy-generated session ids (missing_session_id: "generate") are per-request and are skipped, checked in both "metadata" and "litellm_metadata" (LITELLM_METADATA_ROUTES, responses included) - internal per-request ids (litellm_trace_id / litellm_call_id) lose to a stable conversation key Supersedes the body-key forwarding of #37630: a random session header overrides the body key on the provider side, so the header must carry the caller's key. Related: #29978, #29993. Verified against a real ChatGPT Codex account: stable session id -> 97-100% cached on repeat turns; random per-request -> 0%. * style: ruff format * fix(chatgpt): address review - hash unsafe cache keys, move tests, drop comments - hash prompt_cache_keys containing header-unsafe chars instead of mapping them through _safe_header_value, whose replacement char could collapse distinct keys onto one cache shard and serialize unrelated requests (greptile review) - strip the explanatory comment blocks per AGENTS.md comment policy - move the tests into the existing mapped responses transformation test file per AGENTS.md test-placement rule * fix(chatgpt): keep session typing and test diff scoped * fix(chatgpt): hash surrogate cache keys, type session params, flatten tests * test(chatgpt): record upstream requests with a mock instead of appending to a list * fix(chatgpt): keep proxy-mapped session headers ahead of body session_id and ignore non-string cache keys * fix(chatgpt): hash cache keys with leading or trailing whitespace so the session header stays valid * test(chatgpt): type the session id test parameters * fix(chatgpt): keep trace ids copied into litellm_session_id from outranking prompt_cache_key * fix(chatgpt): mark session ids copied from trace_id so explicit ones keep precedence * fix(chatgpt): keep the trace-id session marker out of provider params The marker now uses the _litellm_ internal prefix, so bridges that pass litellm_params back into completion() drop it instead of sending it as a provider param --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/constants.py | 1 + .../litellm_core_utils/get_litellm_params.py | 6 +- .../litellm_core_utils/provider_affinity.py | 55 +++++++++----- litellm/llms/chatgpt/common_utils.py | 39 ++++++---- .../test_get_litellm_params.py | 9 +++ .../test_chatgpt_responses_transformation.py | 52 +++++++++++++ .../llms/chatgpt/test_chatgpt_common_utils.py | 75 +++++++++++++++++++ 7 files changed, 202 insertions(+), 35 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 7c6d6deba66..2a20e812e68 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1737,6 +1737,7 @@ PASS_THROUGH_HEADER_PREFIX: Final = "x-pass-" TYPESAFE_DANGER_ZONE_HEADER_PREFIX: Final = "typesafe-danger-zone-" INTERNAL_KWARG_PREFIX: Final = "_litellm_" CONTROL_OPTIONS_KEY: Final = f"{INTERNAL_KWARG_PREFIX}control" +SESSION_ID_FROM_TRACE_ID_PARAM_KEY: Final = f"{INTERNAL_KWARG_PREFIX}session_id_from_trace_id" AZURE_SPEECH_CUSTOM_LLM_PROVIDER: Final = "azure_speech" AZURE_SPEECH_PASS_THROUGH_ROUTE_PREFIX: Final = "/azure_speech" diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index ca8e2af18e6..ba964d200ec 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -6,7 +6,7 @@ from typing import Final, cast from pydantic import TypeAdapter, ValidationError -from litellm.constants import CONTROL_OPTIONS_KEY +from litellm.constants import CONTROL_OPTIONS_KEY, SESSION_ID_FROM_TRACE_ID_PARAM_KEY from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.llms.openai.data_residency import infer_openai_data_residency from litellm.types.litellm_params import MAX_CONTROL_INT_DIGITS, ControlOptions @@ -212,6 +212,9 @@ def get_litellm_params( # Derive litellm_session_id / litellm_trace_id from metadata when not provided (call chaining) _meta: Final = resolved_metadata or {} + session_id_from_trace_id: Final = ( + litellm_session_id is None and not _meta.get("session_id") and bool(_meta.get("trace_id")) + ) if litellm_session_id is None: litellm_session_id = _meta.get("session_id") or _meta.get("trace_id") if litellm_trace_id is None: @@ -254,6 +257,7 @@ def get_litellm_params( or (get_base_model_from_litellm_call_metadata(metadata=metadata) if metadata else None), "litellm_trace_id": litellm_trace_id, "litellm_session_id": litellm_session_id, + **({SESSION_ID_FROM_TRACE_ID_PARAM_KEY: True} if session_id_from_trace_id else {}), "hf_model_name": hf_model_name, "custom_prompt_dict": custom_prompt_dict, "litellm_metadata": litellm_metadata, diff --git a/litellm/litellm_core_utils/provider_affinity.py b/litellm/litellm_core_utils/provider_affinity.py index e022839c5c7..c831e92c09d 100644 --- a/litellm/litellm_core_utils/provider_affinity.py +++ b/litellm/litellm_core_utils/provider_affinity.py @@ -2,7 +2,7 @@ import re from collections.abc import Mapping from typing import Final -from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY +from litellm.constants import SESSION_ID_FROM_TRACE_ID_PARAM_KEY, SESSION_ID_GENERATED_METADATA_KEY _HTTP_HEADER_NAME_PATTERN: Final = re.compile(r"^[!#$%&'*+\-.^_`|~0-9A-Za-z]+$") _FORBIDDEN_AFFINITY_HEADERS: Final = frozenset( @@ -53,36 +53,51 @@ def _get_provider_affinity_header_name(litellm_params: object | None, default_he return validate_provider_affinity_header_name(header) -def get_stable_session_id(litellm_params: object | None) -> str | None: - if litellm_params is None: - return None - - direct_session_id: Final = _get_value(litellm_params, "session_id") - if direct_session_id: - return str(direct_session_id) - - metadata_values: Final[tuple[object, ...]] = tuple( +def _metadata_values(litellm_params: object) -> tuple[object, ...]: + return tuple( value for key in ("metadata", "litellm_metadata") if (value := _get_value(litellm_params, key)) is not None ) + + +def _non_generated_litellm_session_id(litellm_params: object) -> str | None: has_generated_session_id: Final = any( - isinstance(metadata, Mapping) and metadata.get(SESSION_ID_GENERATED_METADATA_KEY) - for metadata in metadata_values + _get_value(metadata, SESSION_ID_GENERATED_METADATA_KEY) for metadata in _metadata_values(litellm_params) ) - litellm_session_id: Final = _get_value(litellm_params, "litellm_session_id") - if litellm_session_id and not has_generated_session_id: - return str(litellm_session_id) + return str(litellm_session_id) if litellm_session_id and not has_generated_session_id else None - for metadata in metadata_values: - if ( - isinstance(metadata, Mapping) - and not metadata.get(SESSION_ID_GENERATED_METADATA_KEY) - and (value := metadata.get("session_id")) + +def _metadata_session_id(litellm_params: object) -> str | None: + for metadata in _metadata_values(litellm_params): + if not _get_value(metadata, SESSION_ID_GENERATED_METADATA_KEY) and ( + value := _get_value(metadata, "session_id") ): return str(value) return None +def get_caller_litellm_session_id(litellm_params: object | None) -> str | None: + if litellm_params is None or _get_value(litellm_params, SESSION_ID_FROM_TRACE_ID_PARAM_KEY): + return None + return _non_generated_litellm_session_id(litellm_params) + + +def get_requested_session_id(litellm_params: object | None) -> str | None: + if litellm_params is None: + return None + direct_session_id: Final = _get_value(litellm_params, "session_id") + return str(direct_session_id) if direct_session_id else _metadata_session_id(litellm_params) + + +def get_stable_session_id(litellm_params: object | None) -> str | None: + if litellm_params is None: + return None + direct_session_id: Final = _get_value(litellm_params, "session_id") + if direct_session_id: + return str(direct_session_id) + return _non_generated_litellm_session_id(litellm_params) or _metadata_session_id(litellm_params) + + def add_provider_affinity_header( headers: Mapping[str, object], litellm_params: object | None, default_header: str | None = None ) -> dict[str, object]: # mutable-ok: downstream handlers add auth and signing headers diff --git a/litellm/llms/chatgpt/common_utils.py b/litellm/llms/chatgpt/common_utils.py index ec7fbfa92da..cd2ea6b9ccc 100644 --- a/litellm/llms/chatgpt/common_utils.py +++ b/litellm/llms/chatgpt/common_utils.py @@ -3,6 +3,7 @@ Constants and helpers for ChatGPT subscription OAuth. """ import base64 +import hashlib import json import os import platform @@ -17,6 +18,7 @@ import httpx from pydantic import JsonValue, TypeAdapter from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils +from litellm.litellm_core_utils.provider_affinity import get_caller_litellm_session_id, get_requested_session_id from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.utils import ProviderSpecificHeader @@ -395,22 +397,31 @@ def _normalize_litellm_params(litellm_params: object) -> Mapping[str, object]: return MappingProxyType({}) +def _prompt_cache_session_id(prompt_cache_key: object) -> str | None: + if not isinstance(prompt_cache_key, str) or not prompt_cache_key: + return None + return ( + prompt_cache_key + if _safe_header_value(prompt_cache_key.strip()) == prompt_cache_key + else hashlib.sha256(prompt_cache_key.encode("utf-8", "surrogatepass")).hexdigest() + ) + + def get_chatgpt_session_id(litellm_params: object) -> str | None: params: Final = _normalize_litellm_params(litellm_params) - for key in ("litellm_session_id", "session_id"): - value = params.get(key) - if value: - return str(value) - metadata: Final = _normalize_litellm_params(params.get("metadata")) - if metadata: - value = metadata.get("session_id") - if value: - return str(value) - for key in ("litellm_trace_id", "litellm_call_id"): - value = params.get(key) - if value: - return str(value) - return None + return ( + get_caller_litellm_session_id(params) + or get_requested_session_id(params) + or _prompt_cache_session_id(params.get("prompt_cache_key")) + or next( + ( + str(value) + for key in ("litellm_session_id", "litellm_trace_id", "litellm_call_id") + if (value := params.get(key)) + ), + None, + ) + ) def ensure_chatgpt_session_id(litellm_params: object) -> str: diff --git a/tests/unit/litellm_core_utils/test_get_litellm_params.py b/tests/unit/litellm_core_utils/test_get_litellm_params.py index 89173956dc6..738ab774b2f 100644 --- a/tests/unit/litellm_core_utils/test_get_litellm_params.py +++ b/tests/unit/litellm_core_utils/test_get_litellm_params.py @@ -22,6 +22,7 @@ from litellm.litellm_core_utils.get_litellm_params import ( stored_control_options, ) from litellm.types.litellm_params import ControlOptions +from litellm.utils import get_non_default_completion_params def _funnel_kwargs_completion_forwards(monkeypatch, wif_kwargs: dict[str, object]) -> dict[str, object]: @@ -411,6 +412,14 @@ class TestMetadataFallsBackToLitellmMetadata: result["metadata"].pop("trace_id") assert litellm_metadata == {"trace_id": "trace-1"} + def test_session_id_derived_from_trace_id_adds_no_provider_params(self): + from_trace_id: Final = get_litellm_params(metadata={"trace_id": "trace-1"}) + from_session_id: Final = get_litellm_params(metadata={"session_id": "session-1"}) + added_by_trace_id: Final = from_trace_id.keys() - from_session_id.keys() + + assert added_by_trace_id + assert not added_by_trace_id & get_non_default_completion_params(from_trace_id).keys() + @pytest.mark.parametrize( "value, expected", diff --git a/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py index 8985d0b7fd8..a6f7d613f28 100644 --- a/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py +++ b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py @@ -16,6 +16,7 @@ import pytest import litellm from litellm.llms.chatgpt.authenticator import Authenticator +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig from litellm.llms.openai.common_utils import OpenAIError from litellm.main import responses_api_bridge_check @@ -584,3 +585,54 @@ class TestChatGPTResponsesMultiAccount: assert headers_b["ChatGPT-Account-Id"] == "acct-b" assert headers_default["Authorization"] == "Bearer token-default" assert headers_default["ChatGPT-Account-Id"] == "acct-default" + + +@pytest.mark.parametrize("call", ["responses", "completion"]) +def test_prompt_cache_key_keeps_one_session_header_across_requests( + call: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + (tmp_path / "auth.json").write_text( + json.dumps({"access_token": "token", "account_id": "acct", "expires_at": 4102444800}) + ) + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + monkeypatch.delenv("CHATGPT_AUTH_FILE", raising=False) + completed: Final = { + "type": "response.completed", + "response": { + "id": "resp_1", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.5", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "OK", "annotations": []}], + } + ], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + }, + } + upstream: Final = MagicMock( + side_effect=lambda _request: httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=f"event: response.completed\ndata: {json.dumps(completed)}\n\n".encode(), + ) + ) + client: Final = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(upstream))) + for _ in range(2): + if call == "responses": + litellm.responses(model="chatgpt/gpt-5.5", input="hi", prompt_cache_key="conv-123", client=client) + else: + litellm.completion( + model="chatgpt/gpt-5.5", + messages=[{"role": "user", "content": "hi"}], + prompt_cache_key="conv-123", + client=client, + ) + + assert [call.args[0].headers["session_id"] for call in upstream.call_args_list] == ["conv-123", "conv-123"] diff --git a/tests/unit/llms/chatgpt/test_chatgpt_common_utils.py b/tests/unit/llms/chatgpt/test_chatgpt_common_utils.py index 6b280d97089..375c140187e 100644 --- a/tests/unit/llms/chatgpt/test_chatgpt_common_utils.py +++ b/tests/unit/llms/chatgpt/test_chatgpt_common_utils.py @@ -1,11 +1,17 @@ import base64 import json +from collections.abc import Mapping +from typing import Final +import h11 import pytest +from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY +from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.llms.chatgpt.common_utils import ( extract_chatgpt_account_id, get_chatgpt_client_credential, + get_chatgpt_session_id, is_chatgpt_oauth_key, select_chatgpt_client_credential_headers, ) @@ -88,3 +94,72 @@ def test_bearer_scheme_is_case_insensitive(scheme): def test_other_authentication_schemes_are_not_chatgpt_logins(): assert not is_chatgpt_oauth_key(f"Basic {CHATGPT_JWT}") + + +class TestGetChatGPTSessionId: + @pytest.mark.parametrize( + "explicit", + [ + {"litellm_session_id": "chosen"}, + {"session_id": "chosen"}, + {"metadata": {"session_id": "chosen"}}, + get_litellm_params(litellm_session_id="chosen", metadata={"trace_id": "chosen"}), + ], + ) + def test_explicit_session_id_wins_over_prompt_cache_key(self, explicit: Mapping[str, object]) -> None: + assert get_chatgpt_session_id({**explicit, "prompt_cache_key": "conv-123"}) == "chosen" + + @pytest.mark.parametrize( + "request_scoped", + [ + {"litellm_trace_id": "trace-1"}, + {"litellm_call_id": "call-1"}, + get_litellm_params(metadata={"trace_id": "trace-1"}), + get_litellm_params(litellm_metadata={"trace_id": "trace-1"}), + {"litellm_session_id": "gen-1", "metadata": {"session_id": "gen-1", SESSION_ID_GENERATED_METADATA_KEY: True}}, + { + "litellm_session_id": "gen-1", + "litellm_metadata": {"session_id": "gen-1", SESSION_ID_GENERATED_METADATA_KEY: True}, + }, + ], + ) + def test_prompt_cache_key_wins_over_request_scoped_ids(self, request_scoped: Mapping[str, object]) -> None: + assert get_chatgpt_session_id({**request_scoped, "prompt_cache_key": "conv-123"}) == "conv-123" + + @pytest.mark.parametrize( + "mapped_header", + [ + { + "litellm_session_id": "header-session", + "metadata": {"trace_id": "header-session", "session_id": "header-session"}, + }, + get_litellm_params(metadata={"trace_id": "header-session", "session_id": "header-session"}), + ], + ) + def test_proxy_mapped_session_header_wins_over_body_session_id(self, mapped_header: Mapping[str, object]) -> None: + assert ( + get_chatgpt_session_id({**mapped_header, "session_id": "body-session", "prompt_cache_key": "conv-123"}) + == "header-session" + ) + + @pytest.mark.parametrize("prompt_cache_key", [{"tenant": "a"}, 12345, True, ["x", "y"]]) + def test_non_string_prompt_cache_key_is_ignored(self, prompt_cache_key: object) -> None: + assert get_chatgpt_session_id({"prompt_cache_key": prompt_cache_key, "litellm_call_id": "call-1"}) == "call-1" + + @pytest.mark.parametrize( + "first, second", + [ + ("a\nb", "a\tb"), + (json.loads('"conv-\\ud800"'), json.loads('"conv-\\udc00"')), + ("conv-123 ", " conv-123"), + ], + ) + def test_header_unsafe_cache_keys_map_to_distinct_stable_header_values(self, first: str, second: str) -> None: + first_id: Final = get_chatgpt_session_id({"prompt_cache_key": first}) + second_id: Final = get_chatgpt_session_id({"prompt_cache_key": second}) + + assert first_id is not None and second_id is not None + assert first_id == get_chatgpt_session_id({"prompt_cache_key": first}) + assert first_id != second_id + for value in (first_id, second_id): + h11.Request(method="POST", target="/", headers=[("host", "chatgpt.com"), ("session_id", value)]) From 086daf465316c691fd28e6ba1a8285575e96f209 Mon Sep 17 00:00:00 2001 From: zoveki <349028329+zoveki@users.noreply.github.com> Date: Wed, 28 Oct 2026 11:12:55 -0700 Subject: [PATCH 058/128] fix(cost-map): price Mistral Large 4 at its list price now that the launch sale ended (#48825) mistral/mistral-large-4 and mistral/mistral-large-4-0 still carried the launch sale price (input 6.8e-07, cache read 6.8e-08, output 2.09e-06). Mistral's model card and pricing page now list only the list price (input 1.36e-06, cached input 1.4e-07, output 4.18e-06), so every call was billed at about half its real cost. The test derives expected cost from the row's own prices instead of pinning Mistral's numbers. Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 12 ++++++------ model_prices_and_context_window.json | 12 ++++++------ tests/unit/test_mistral_large_4_model_metadata.py | 13 ++++++++----- 3 files changed, 20 insertions(+), 17 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 43bf7a7d029..04ad3546b5b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -42058,13 +42058,13 @@ "mode": "embedding" }, "mistral/mistral-large-4": { - "cache_read_input_token_cost": 6.8e-08, - "input_cost_per_token": 6.8e-07, + "cache_read_input_token_cost": 1.4e-07, + "input_cost_per_token": 1.36e-06, "litellm_provider": "mistral", "max_input_tokens": 1048576, "max_output_tokens": 1048576, "mode": "chat", - "output_cost_per_token": 2.09e-06, + "output_cost_per_token": 4.18e-06, "reasoning_effort_levels": [ "none", "high" @@ -42079,13 +42079,13 @@ "supports_vision": true }, "mistral/mistral-large-4-0": { - "cache_read_input_token_cost": 6.8e-08, - "input_cost_per_token": 6.8e-07, + "cache_read_input_token_cost": 1.4e-07, + "input_cost_per_token": 1.36e-06, "litellm_provider": "mistral", "max_input_tokens": 1048576, "max_output_tokens": 1048576, "mode": "chat", - "output_cost_per_token": 2.09e-06, + "output_cost_per_token": 4.18e-06, "reasoning_effort_levels": [ "none", "high" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 43bf7a7d029..04ad3546b5b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -42058,13 +42058,13 @@ "mode": "embedding" }, "mistral/mistral-large-4": { - "cache_read_input_token_cost": 6.8e-08, - "input_cost_per_token": 6.8e-07, + "cache_read_input_token_cost": 1.4e-07, + "input_cost_per_token": 1.36e-06, "litellm_provider": "mistral", "max_input_tokens": 1048576, "max_output_tokens": 1048576, "mode": "chat", - "output_cost_per_token": 2.09e-06, + "output_cost_per_token": 4.18e-06, "reasoning_effort_levels": [ "none", "high" @@ -42079,13 +42079,13 @@ "supports_vision": true }, "mistral/mistral-large-4-0": { - "cache_read_input_token_cost": 6.8e-08, - "input_cost_per_token": 6.8e-07, + "cache_read_input_token_cost": 1.4e-07, + "input_cost_per_token": 1.36e-06, "litellm_provider": "mistral", "max_input_tokens": 1048576, "max_output_tokens": 1048576, "mode": "chat", - "output_cost_per_token": 2.09e-06, + "output_cost_per_token": 4.18e-06, "reasoning_effort_levels": [ "none", "high" diff --git a/tests/unit/test_mistral_large_4_model_metadata.py b/tests/unit/test_mistral_large_4_model_metadata.py index f45c53920dd..bfcd6690890 100644 --- a/tests/unit/test_mistral_large_4_model_metadata.py +++ b/tests/unit/test_mistral_large_4_model_metadata.py @@ -2,6 +2,7 @@ import json from pathlib import Path +from typing import Final import pytest @@ -32,9 +33,7 @@ def test_metadata_and_backup(model): assert info["litellm_provider"] == "mistral" assert info["mode"] == "chat" assert info["max_input_tokens"] == 1048576 - assert info["input_cost_per_token"] == 6.8e-07 - assert info["output_cost_per_token"] == 2.09e-06 - assert info["cache_read_input_token_cost"] == 6.8e-08 + assert 0 < info["cache_read_input_token_cost"] < info["input_cost_per_token"] < info["output_cost_per_token"] assert info["reasoning_effort_levels"] == ["none", "high"] for capability in ( "supports_vision", @@ -59,12 +58,16 @@ def test_reasoning_effort(model, effort): @pytest.mark.parametrize("model", MODELS) def test_completion_cost(model): + info: Final = litellm.get_model_info(model) + input_price: Final[float] = info["input_cost_per_token"] + output_price: Final[float] = info["output_cost_per_token"] + cache_read_price: Final[float] = info["cache_read_input_token_cost"] uncached_response = litellm.ModelResponse( model=model.split("/", 1)[1], usage={"prompt_tokens": 1000, "completion_tokens": 100, "total_tokens": 1100}, ) assert litellm.completion_cost(completion_response=uncached_response, model=model) == pytest.approx( - 1000 * 6.8e-07 + 100 * 2.09e-06 + 1000 * input_price + 100 * output_price ) response = litellm.ModelResponse( model=model.split("/", 1)[1], @@ -76,7 +79,7 @@ def test_completion_cost(model): }, ) assert litellm.completion_cost(completion_response=response, model=model) == pytest.approx( - 400 * 6.8e-07 + 600 * 6.8e-08 + 100 * 2.09e-06 + 400 * input_price + 600 * cache_read_price + 100 * output_price ) From 6e33d7e84b59e80b6ea5a491b758c24030da545d Mon Sep 17 00:00:00 2001 From: shrey Date: Wed, 28 Oct 2026 11:20:06 -0700 Subject: [PATCH 059/128] fix(anthropic-messages): merge deployment and client anthropic-beta instead of overwriting (#48594) --- litellm/llms/custom_httpx/llm_http_handler.py | 28 ++++-- .../custom_httpx/test_llm_http_handler.py | 91 ++++++++++++++++++- 2 files changed, 110 insertions(+), 9 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index ef009e70137..afa12b58b29 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -6,6 +6,7 @@ import ssl from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence from contextlib import asynccontextmanager from functools import lru_cache, partial, reduce +from itertools import chain from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -62,7 +63,7 @@ from litellm.litellm_core_utils.realtime_errors import ( ) from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base -from litellm.llms.anthropic.common_utils import AnthropicModelInfo +from litellm.llms.anthropic.common_utils import AnthropicModelInfo, merge_anthropic_beta_headers from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, ) @@ -572,6 +573,22 @@ def _unwrap_provider_error(error: Exception) -> Exception: return error +def _merge_messages_header_sources(*sources: Mapping[str, object] | None) -> Mapping[str, object]: + """Later sources win per header, except ``anthropic-beta``, which unions every source so client betas + never drop a deployment's configured ones.""" + entries: Final = tuple(chain.from_iterable((source or {}).items() for source in sources)) + beta_values: Final = tuple( + cast("str | Sequence[str]", value) # cast-ok: the merge helper skips non-str list entries + for name, value in entries + if name.lower() == "anthropic-beta" and isinstance(value, str | Sequence) + ) + merged_beta: Final = reduce(merge_anthropic_beta_headers, beta_values, "") + return MappingProxyType( + {name: value for name, value in entries if name.lower() != "anthropic-beta"} + | ({"anthropic-beta": merged_beta} if merged_beta else {}) + ) + + class BaseLLMHTTPHandler: async def _make_common_async_call( self, @@ -2481,13 +2498,8 @@ class BaseLLMHTTPHandler: forwarded_headers: Final = kwargs.get("headers", None) # Also check for extra_headers in kwargs (from config or direct calls) extra_headers_from_kwargs: Final = kwargs.get("extra_headers", None) - # Merge all header sources: forwarded < extra_headers < provider_specific - merged_headers: Final[Mapping[str, object]] = MappingProxyType( - { - **(forwarded_headers or {}), - **(extra_headers_from_kwargs or {}), - **(provider_specific_headers or {}), - } + merged_headers: Final = _merge_messages_header_sources( + forwarded_headers, extra_headers_from_kwargs, provider_specific_headers ) request_params: Final = reduce( delete_nested_value, diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 297e64629a8..f10f5a67b2f 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -66,7 +66,13 @@ from litellm.responses.streaming_iterator import ( ) from litellm.types.llms.openai import HttpxBinaryResponseContent, ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import ImageObject, ImageResponse, ModelResponse, TranscriptionResponse +from litellm.types.utils import ( + ImageObject, + ImageResponse, + ModelResponse, + ProviderSpecificHeader, + TranscriptionResponse, +) from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe from tests.unit.llms.bedrock.event_loop_probe import host_in_loop_probing_vault @@ -1607,6 +1613,89 @@ async def test_async_anthropic_messages_handler_header_priority(): assert captured_headers["X-Provider-Only"] == "keep-this-too" +class _AnthropicMessagesUpstream: + def __init__(self) -> None: + self.anthropic_beta: tuple[str, ...] = () + + def respond(self, request: httpx.Request) -> httpx.Response: + self.anthropic_beta = tuple(request.headers.get_list("anthropic-beta")) + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [{"type": "text", "text": "hi"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + +_BOTH_FOUNDRY_BETAS: Final = ("context-1m-2025-08-07,context-management-2025-06-27",) + + +@pytest.mark.parametrize( + ("forwarded_headers", "deployment_headers", "client_beta", "expected"), + [ + pytest.param( + None, + {"anthropic-beta": "context-1m-2025-08-07"}, + "context-management-2025-06-27", + _BOTH_FOUNDRY_BETAS, + id="deployment-and-client-betas-are-unioned", + ), + pytest.param( + None, + {"anthropic-beta": ["context-1m-2025-08-07"]}, + "context-management-2025-06-27", + _BOTH_FOUNDRY_BETAS, + id="deployment-list-and-client-string-are-unioned", + ), + pytest.param( + {"anthropic-beta": "interleaved-thinking-2025-05-14"}, + {"anthropic-beta": "context-1m-2025-08-07"}, + "context-management-2025-06-27", + ("context-1m-2025-08-07,context-management-2025-06-27,interleaved-thinking-2025-05-14",), + id="forwarded-deployment-and-client-betas-are-unioned", + ), + pytest.param( + None, + {"anthropic-beta": "context-1m-2025-08-07"}, + "claude-code-20250219", + ("context-1m-2025-08-07",), + id="unsupported-client-beta-is-dropped-and-deployment-beta-kept", + ), + ], +) +@pytest.mark.asyncio +async def test_messages_sends_deployment_and_client_anthropic_betas_upstream( + forwarded_headers: dict[str, str] | None, + deployment_headers: dict[str, str | list[str]], + client_beta: str, + expected: tuple[str, ...], +) -> None: + upstream: Final = _AnthropicMessagesUpstream() + + await litellm.anthropic.messages.acreate( + model="azure_ai/claude-sonnet-5-5", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=16, + api_base="https://test.services.ai.azure.com/anthropic", + api_key="test-key", + headers=forwarded_headers, + extra_headers=deployment_headers, + provider_specific_header=ProviderSpecificHeader( + custom_llm_provider="azure_ai", extra_headers={"anthropic-beta": client_beta} + ), + client=AsyncHTTPHandler(transport=httpx.MockTransport(upstream.respond)), + ) + + assert upstream.anthropic_beta == expected + + @pytest.mark.asyncio async def test_async_anthropic_messages_handler_drops_top_level_and_nested_params(): """ From ea6dbd210976b8b1a49cd586cea8caec714c0d8d Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 11:23:52 -0700 Subject: [PATCH 060/128] feat(proxy): log one line per token count call (#47657) * feat(proxy): log one line per token count call Each count call on /v1/messages/count_tokens, /utils/token_counter, /v1/responses/input_tokens and the Gemini countTokens routes now logs one INFO line with the litellm call id, the requested model, the model the count ran against, the count, the count source and the status code, so it joins the inference call that sends the same x-litellm-call-id * test(google_endpoints): give the count test's fake request a headers list * fix(proxy): keep token count log lines one per call and trusted in JSON logs, and log Bedrock native counts Quote log values that carry whitespace, quotes, equals, or control characters so a model name can't split or forge a line. Let JsonFormatter record extras win over a dict parsed from the message. Log the Bedrock native count-tokens route through the same line. * fix(proxy): escape braces in token count log values, log the Bedrock model AWS counts, and skip the line when there is nothing to count * fix(proxy): log the deployment id the count routed to on every token count line * fix(proxy): log count requests rejected before counting with the status they answer A count route that rejects a request before the shared count function runs now logs the same token count line with its 400. An unsupported content block on the Messages count route becomes its 400 inside the logged wrapper, so its line says 400 instead of 500. A count that ends in a BaseException other than cancellation logs a 500. * fix(proxy): compute the Bedrock counted model id only after a successful count * revert(logging): keep the JSON log formatter's field order as on main The count line already escapes braces in its values, so the formatter never parses a dict from it. Changing the order for every logger is a separate decision. --- .../bedrock/count_tokens/transformation.py | 11 +- litellm/proxy/_types.py | 4 +- .../proxy/anthropic_endpoints/endpoints.py | 20 +- litellm/proxy/common_utils/token_count_log.py | 133 +++++++++++++ litellm/proxy/google_endpoints/endpoints.py | 24 ++- .../llm_passthrough_endpoints.py | 55 ++++-- litellm/proxy/proxy_server.py | 97 +++++++-- .../proxy/response_api_endpoints/endpoints.py | 75 +++++-- .../anthropic_endpoints/test_endpoints.py | 180 +++++++++++++++++ .../common_utils/test_token_count_log.py | 45 +++++ tests/unit/proxy/conftest.py | 27 +++ .../proxy/google_endpoints/test_endpoints.py | 1 + .../test_google_api_endpoints.py | 78 ++++++++ .../test_llm_pass_through_endpoints.py | 185 +++++++++++++++++- .../response_api_endpoints/test_endpoints.py | 84 +++++++- tests/unit/proxy/test_proxy_token_counter.py | 154 ++++++++++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 +- 17 files changed, 1088 insertions(+), 90 deletions(-) create mode 100644 litellm/proxy/common_utils/token_count_log.py create mode 100644 tests/unit/proxy/common_utils/test_token_count_log.py diff --git a/litellm/llms/bedrock/count_tokens/transformation.py b/litellm/llms/bedrock/count_tokens/transformation.py index a9f95c1f9b9..bf9c119f6f2 100644 --- a/litellm/llms/bedrock/count_tokens/transformation.py +++ b/litellm/llms/bedrock/count_tokens/transformation.py @@ -206,6 +206,10 @@ class BedrockCountTokensConfig(BaseAWSLLM): encoded_body: Final = base64.b64encode(json.dumps(body_data).encode()).decode() return {"input": {"invokeModel": {"body": encoded_body}}} + @staticmethod + def count_tokens_model_id(model: str) -> str: + return get_bedrock_base_model(model).removeprefix("bedrock/") + def get_bedrock_count_tokens_endpoint( self, model: str, @@ -225,12 +229,7 @@ class BedrockCountTokensConfig(BaseAWSLLM): Returns: Complete endpoint URL for CountTokens API """ - # Use existing LiteLLM function to get the base model ID (removes region prefix) - model_id = get_bedrock_base_model(model) - - # Remove bedrock/ prefix if present - model_id = model_id.removeprefix("bedrock/") # Remove "bedrock/" prefix - encoded_model_id: Final = self.encode_model_id(model_id=model_id) + encoded_model_id: Final = self.encode_model_id(model_id=self.count_tokens_model_id(model)) base_url, _ = self.get_runtime_endpoint( api_base=api_base, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index dd91bb6eef3..5ce6cd93a22 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3461,12 +3461,12 @@ class LiteLLM_SpendLogs_ResponseObject(LiteLLMPydanticObjectBase): class TokenCountRequest(LiteLLMPydanticObjectBase): model: str prompt: str | None = None - messages: list[dict] | None = None + messages: list[dict[Any, Any]] | None = None """ Anthropic token counting endpoint uses /messages """ - contents: list[dict] | None = None + contents: list[dict[Any, Any]] | None = None """ Google /countTokens endpoint expects contents to be a list of dicts with the following structure: """ diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 1d2045eb743..1d88e133cf9 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -18,7 +18,6 @@ from litellm.anthropic_interface.exceptions import ( AnthropicExceptionMapping, ) from litellm.integrations.custom_guardrail import ModifyResponseException -from litellm.litellm_core_utils.token_counter import UnsupportedContentTypeError from litellm.litellm_core_utils.tool_fields import first_validation_error from litellm.llms.anthropic.messages.context_management import ( AnthropicContextManagementError, @@ -49,6 +48,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( openai_error_type, with_litellm_call_id, ) +from litellm.proxy.common_utils.token_count_log import rejected_token_count from litellm.proxy.common_utils.token_count_settings import count_source_fields from litellm.proxy.litellm_pre_call_utils import ( forwarded_anthropic_oauth_token, @@ -394,12 +394,17 @@ async def count_tokens( # Extract required fields model_name: Final = data.get("model") messages: Final = data.get("messages", []) + request_model: Final = model_name if isinstance(model_name, str) else None if not model_name: - raise HTTPException(status_code=400, detail="model parameter is required") + raise rejected_token_count( + HTTPException(status_code=400, detail="model parameter is required"), litellm_call_id, request_model + ) if not messages: - raise HTTPException(status_code=400, detail="messages parameter is required") + raise rejected_token_count( + HTTPException(status_code=400, detail="messages parameter is required"), litellm_call_id, request_model + ) auto_drop_params: Final = should_auto_drop_params_for_agentic_cli( request.headers.get("user-agent", ""), "drop_params" in data, proxy_config @@ -409,19 +414,18 @@ async def count_tokens( token_fields if isinstance(token_fields, ProxyException) else _token_count_request(token_fields) ) if isinstance(token_request, ProxyException): - return _anthropic_error_json_response(token_request, request) + return rejected_token_count( + _anthropic_error_json_response(token_request, request), litellm_call_id, request_model + ) token_response: Final = await count_request_tokens( request=token_request, call_endpoint=True, messages_format="anthropic", + litellm_call_id=litellm_call_id, client_oauth_token=forwarded_anthropic_oauth_token(request.headers), **_count_tokens_body_options(data, auto_drop_params), ) - except UnsupportedContentTypeError as e: - return _anthropic_error_json_response( - ProxyException(message=str(e), type="invalid_request_error", param="messages", code=400), request - ) except HTTPException as e: return _anthropic_error_json_response(proxy_exception_from_http_exception(e, {}), request) except ProxyException as e: diff --git a/litellm/proxy/common_utils/token_count_log.py b/litellm/proxy/common_utils/token_count_log.py new file mode 100644 index 00000000000..04c07a29117 --- /dev/null +++ b/litellm/proxy/common_utils/token_count_log.py @@ -0,0 +1,133 @@ +import json +import re +from asyncio import CancelledError +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, TypeVar + +from fastapi import status + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.common_utils.openai_error_payload import error_status_code +from litellm.types.token_count import TokenCountSource +from litellm.types.utils import TokenCountResponse + +_BARE_LOGFMT_VALUE: Final = re.compile(r"[\w.:/@+-]+", re.ASCII) +_JSON_ESCAPED_BRACES: Final = str.maketrans({"{": "\\u007b", "}": "\\u007d"}) + + +def _logfmt_value(value: str | None) -> str: + if value is None: + return "None" + if _BARE_LOGFMT_VALUE.fullmatch(value): + return value + return json.dumps(value).translate(_JSON_ESCAPED_BRACES) + + +_Counted: Final = TypeVar("_Counted") +_Rejection: Final = TypeVar("_Rejection") +_ProviderCount: Final = TypeVar("_ProviderCount", bound=Mapping[str, object]) + + +@dataclass(frozen=True, slots=True) +class _CountFields: + model_used: str | None + input_tokens: int | None + count_source: TokenCountSource | None + + +_FAILED_COUNT: Final = _CountFields(model_used=None, input_tokens=None, count_source=None) + + +@dataclass(frozen=True, slots=True) +class _CountRequest: + litellm_call_id: str + request_model: str | None + model_id: str | None + + +def deployment_model_id(deployment: Mapping[str, object] | None) -> str | None: + match deployment: + case {"model_info": {"id": str() as model_id}}: + return model_id + case _: + return None + + +def _log_token_count(request: _CountRequest, fields: _CountFields, status_code: int) -> None: + verbose_proxy_logger.info( + "token count - litellm_call_id=%s request_model=%s model_id=%s model_used=%s input_tokens=%s " + "count_source=%s status_code=%s", + _logfmt_value(request.litellm_call_id), + _logfmt_value(request.request_model), + _logfmt_value(request.model_id), + _logfmt_value(fields.model_used), + fields.input_tokens, + fields.count_source, + status_code, + extra=MappingProxyType( + { + "litellm_call_id": request.litellm_call_id, + "request_model": request.request_model, + "model_id": request.model_id, + "model_used": fields.model_used, + "input_tokens": fields.input_tokens, + "count_source": fields.count_source, + "status_code": status_code, + } + ), + ) + + +def rejected_token_count(rejection: _Rejection, litellm_call_id: str, request_model: str | None) -> _Rejection: + _log_token_count( + _CountRequest(litellm_call_id, request_model, model_id=None), + _FAILED_COUNT, + error_status_code(rejection, status.HTTP_500_INTERNAL_SERVER_ERROR), + ) + return rejection + + +async def _logged( + count: Awaitable[_Counted], request: _CountRequest, fields_of: Callable[[_Counted], _CountFields] +) -> _Counted: + try: + counted: Final = await count + except CancelledError: + raise + except BaseException as e: + _log_token_count(request, _FAILED_COUNT, error_status_code(e, status.HTTP_500_INTERNAL_SERVER_ERROR)) + raise + _log_token_count(request, fields_of(counted), status.HTTP_200_OK) + return counted + + +def _token_count_response_fields(counted: TokenCountResponse) -> _CountFields: + return _CountFields( + model_used=counted.model_used, input_tokens=counted.total_tokens, count_source=counted.count_source + ) + + +async def logged_token_count( + count: Awaitable[TokenCountResponse], litellm_call_id: str, request_model: str, model_id: str | None +) -> TokenCountResponse: + return await _logged(count, _CountRequest(litellm_call_id, request_model, model_id), _token_count_response_fields) + + +async def logged_provider_token_count( + count: Awaitable[_ProviderCount], + litellm_call_id: str, + request_model: str, + model_id: str | None, + model_used: Callable[[], str], +) -> _ProviderCount: + def provider_fields(counted: _ProviderCount) -> _CountFields: + input_tokens: Final = counted.get("input_tokens") + return _CountFields( + model_used=model_used(), + input_tokens=input_tokens if isinstance(input_tokens, int) and not isinstance(input_tokens, bool) else None, + count_source="provider", + ) + + return await _logged(count, _CountRequest(litellm_call_id, request_model, model_id), provider_fields) diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 3d1a9ab4cfe..54976857584 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -21,6 +21,8 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg _read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export read_request_body, ) +from litellm.proxy.common_utils.openai_error_payload import LITELLM_CALL_ID_HEADER +from litellm.proxy.common_utils.token_count_log import rejected_token_count from litellm.proxy.common_utils.token_count_settings import count_source_fields from litellm.types.llms.vertex_ai import TokenCountDetailsResponse @@ -286,25 +288,35 @@ async def google_count_tokens(request: Request, model_name: str): from litellm.proxy.common_utils.http_parsing_utils import read_request_body from litellm.proxy.proxy_server import token_counter as internal_token_counter + litellm_call_id: Final = resolve_litellm_call_id(request.headers.get(LITELLM_CALL_ID_HEADER)) data: Final = await read_request_body(request=request) try: token_request: Final = _google_count_tokens_request(model_name, data) except ValidationError: - return _google_error_response( - request, - ProxyException( - message="Invalid Google countTokens request body", type="invalid_request_error", param=None, code=400 + return rejected_token_count( + _google_error_response( + request, + ProxyException( + message="Invalid Google countTokens request body", + type="invalid_request_error", + param=None, + code=400, + ), ), + litellm_call_id, + model_name, ) try: - token_response: Final = await internal_token_counter(request=token_request, call_endpoint=True) + token_response: Final = await internal_token_counter( + request=token_request, call_endpoint=True, litellm_call_id=litellm_call_id + ) except ProxyException as error: return _google_error_response(request, error) except HTTPException as error: return _google_error_response(request, proxy_exception_from_http_exception(error, {})) except Exception as error: - log_llm_api_exception(error, resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))) + log_llm_api_exception(error, litellm_call_id) return _google_error_response( request, ProxyException(message="Internal server error", type="internal_server_error", param=None, code=500), diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 9eef4f0ef9d..e478096f7ad 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -87,6 +87,7 @@ from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, model_deployments_support_stream_options, open_sse_before_first_byte, + resolve_litellm_call_id, ) from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports _read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -99,10 +100,16 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg safe_get_request_headers, safe_set_request_parsed_body, ) +from litellm.proxy.common_utils.openai_error_payload import LITELLM_CALL_ID_HEADER from litellm.proxy.common_utils.resource_ownership import is_proxy_admin from litellm.proxy.common_utils.sse_keepalive import ( wrap_passthrough_sse_bytes_with_keepalive_pings, ) +from litellm.proxy.common_utils.token_count_log import ( + deployment_model_id, + logged_provider_token_count, + rejected_token_count, +) from litellm.proxy.pass_through_endpoints.common_utils import get_litellm_virtual_key from litellm.proxy.pass_through_endpoints.injected_stream_usage import ( stream_usage_relay, @@ -1259,6 +1266,7 @@ async def handle_bedrock_count_tokens( from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.proxy.proxy_server import llm_router + litellm_call_id: Final = resolve_litellm_call_id(request.headers.get(LITELLM_CALL_ID_HEADER)) try: # Initialize the handler handler: Final = BedrockCountTokensHandler() @@ -1266,40 +1274,45 @@ async def handle_bedrock_count_tokens( # Extract model from request body model: Final = _optional_str(request_body.get("model")) if not model: - raise HTTPException(status_code=400, detail={"error": "Model is required in request body"}) + raise rejected_token_count( + HTTPException(status_code=400, detail={"error": "Model is required in request body"}), + litellm_call_id, + None, + ) # Get model parameters from router litellm_params: Final = {"user_api_key_dict": user_api_key_dict} resolved_model = model # Default fallback - if llm_router: - deployments: Final = llm_router.get_model_list(model_name=model) - if deployments and len(deployments) > 0: - # Get the first matching deployment - deployment: Final = deployments[0] - model_litellm_params: Final = deployment.get("litellm_params", {}) + deployments: Final = llm_router.get_model_list(model_name=model) if llm_router else None + deployment: Final = deployments[0] if deployments else None + if deployment is not None: + model_litellm_params: Final = deployment.get("litellm_params", {}) - # Get the resolved model ID from the configuration - if "model" in model_litellm_params: - resolved_model = model_litellm_params["model"] + # Get the resolved model ID from the configuration + if "model" in model_litellm_params: + resolved_model = model_litellm_params["model"] - # Copy all litellm_params - BaseAWSLLM will handle AWS credential discovery - for key, value in model_litellm_params.items(): - if key != "user_api_key_dict": # Don't overwrite user_api_key_dict - litellm_params[key] = value + # Copy all litellm_params - BaseAWSLLM will handle AWS credential discovery + for key, value in model_litellm_params.items(): + if key != "user_api_key_dict": # Don't overwrite user_api_key_dict + litellm_params[key] = value verbose_proxy_logger.debug("Count tokens litellm_params: %s", litellm_params) verbose_proxy_logger.debug("Resolved model: %s", resolved_model) - # Handle the count tokens request - result: Final = await handler.handle_count_tokens_request( - request_data=request_body, - litellm_params=litellm_params, - resolved_model=resolved_model, + return await logged_provider_token_count( + handler.handle_count_tokens_request( + request_data=request_body, + litellm_params=litellm_params, + resolved_model=resolved_model, + ), + litellm_call_id=litellm_call_id, + request_model=model, + model_id=deployment_model_id(deployment), + model_used=partial(handler.count_tokens_model_id, resolved_model), ) - return result - except BedrockError as e: # Convert BedrockError to HTTPException for FastAPI verbose_proxy_logger.error("BedrockError in handle_bedrock_count_tokens: %s", e) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4e7c4601803..7fcb77e2bbf 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -100,7 +100,7 @@ from litellm.litellm_core_utils.litellm_logging import ( ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads -from litellm.litellm_core_utils.token_counter import offload_token_count +from litellm.litellm_core_utils.token_counter import UnsupportedContentTypeError, offload_token_count from litellm.llms.bedrock.count_tokens.kimi_k3 import ( HUGGING_FACE_SOURCE, KimiK3Encodings, @@ -552,6 +552,11 @@ from litellm.proxy.common_utils.timezone_utils import ( get_budget_reset_settings, get_budget_reset_time, ) +from litellm.proxy.common_utils.token_count_log import ( + deployment_model_id, + logged_token_count, + rejected_token_count, +) from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, end_user_cache_key, @@ -15253,16 +15258,21 @@ def _system_message(system: object) -> ChatCompletionSystemMessage | None: dependencies=[Depends(user_api_key_auth)], response_model=TokenCountResponse, ) -async def token_counter(request: TokenCountRequest, call_endpoint: bool = False): +async def token_counter( + request: TokenCountRequest, + call_endpoint: bool = False, + litellm_call_id: Annotated[str | None, Header(alias=LITELLM_CALL_ID_HEADER)] = None, +): """ Args: request: TokenCountRequest call_endpoint: bool - When set to "True" it will call the token counting endpoint - e.g Anthropic or Google AI Studio Token Counting APIs. + litellm_call_id: the x-litellm-call-id header, logged with the count so it joins the inference call's logs Returns: TokenCountResponse """ - return await count_request_tokens(request=request, call_endpoint=call_endpoint) + return await count_request_tokens(request=request, call_endpoint=call_endpoint, litellm_call_id=litellm_call_id) async def count_request_tokens( @@ -15273,7 +15283,69 @@ async def count_request_tokens( context_management: JsonValue = None, request_litellm_params: Mapping[str, JsonValue] = MappingProxyType({}), responses_request: Mapping[str, object] | None = None, + litellm_call_id: str | None = None, client_oauth_token: str | None = None, +) -> TokenCountResponse: + call_id: Final = resolve_litellm_call_id(litellm_call_id) + if request.prompt is None and request.messages is None and request.contents is None: + raise rejected_token_count( + HTTPException(status_code=400, detail="prompt or messages or contents must be provided"), + call_id, + request.model, + ) + routed: Final = await _routed_count_deployment(request.model, responses_request) + count: Final = _count_request_tokens( + request=request, + routed=routed, + call_endpoint=call_endpoint, + messages_format=messages_format, + output_format=output_format, + context_management=context_management, + request_litellm_params=request_litellm_params, + responses_request=responses_request, + client_oauth_token=client_oauth_token, + ) + return await logged_token_count( + _with_anthropic_content_errors(count) if messages_format == "anthropic" else count, + litellm_call_id=call_id, + request_model=request.model, + model_id=deployment_model_id(routed), + ) + + +async def _with_anthropic_content_errors(count: Awaitable[TokenCountResponse]) -> TokenCountResponse: + try: + return await count + except UnsupportedContentTypeError as e: + raise ProxyException(message=str(e), type="invalid_request_error", param="messages", code=400) from e + + +async def _routed_count_deployment( + model: str, responses_request: Mapping[str, object] | None +) -> Mapping[str, object] | None: + if llm_router is None: + return None + try: + return await llm_router.async_get_available_deployment( + model=model, request_kwargs=dict(_responses_routing_kwargs(responses_request)) + ) + except Exception: + verbose_proxy_logger.exception( + "litellm.proxy.proxy_server.token_counter(): Exception occured while getting deployment" + ) + return None + + +async def _count_request_tokens( + request: TokenCountRequest, + routed: Mapping[str, object] | None, + call_endpoint: bool, + messages_format: TokenCountMessagesFormat, + output_format: JsonValue, + context_management: JsonValue, + request_litellm_params: Mapping[str, JsonValue], + responses_request: Mapping[str, object] | None, + client_oauth_token: str | None, ) -> TokenCountResponse: prompt: Final = request.prompt messages: Final = request.messages @@ -15281,26 +15353,9 @@ async def count_request_tokens( tools: Final = request.tools system: Final = request.system - ######################################################### - # Validate request - ######################################################### - if prompt is None and messages is None and contents is None: - raise HTTPException(status_code=400, detail="prompt or messages or contents must be provided") - - deployment: dict[str, Any] | None = None + deployment: dict[str, Any] | None = None if routed is None else dict(routed) litellm_model_name = None model_info: ModelMapInfo | None = None - if llm_router is not None: - # get 1 deployment corresponding to the model - try: - deployment = await llm_router.async_get_available_deployment( - model=request.model, - request_kwargs=dict(_responses_routing_kwargs(responses_request)), - ) - except Exception: - verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.token_counter(): Exception occured while getting deployment" - ) if deployment is not None: litellm_model_name = deployment.get("litellm_params", {}).get("model") model_info = deployment.get("model_info", {}) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 97bbf893e3e..4f8b1697f7e 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -56,6 +56,7 @@ from litellm.proxy.common_request_processing import ( cost_headers, create_response, resolve_counted_tool_fields, + resolve_litellm_call_id, ) from litellm.proxy.common_utils.error_body_call_id import error_body_call_id, with_call_id from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports @@ -69,6 +70,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( error_status_code, openai_error_type_for_status, ) +from litellm.proxy.common_utils.token_count_log import rejected_token_count from litellm.proxy.common_utils.token_count_settings import count_source_fields from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity, decrypted_response_id @@ -1407,7 +1409,12 @@ class _TokenCountPayload(TypedDict): class _TokenCounter(Protocol): def __call__( - self, *, request: TokenCountRequest, call_endpoint: bool, responses_request: Mapping[str, object] | None + self, + *, + request: TokenCountRequest, + call_endpoint: bool, + responses_request: Mapping[str, object] | None, + litellm_call_id: str | None, ) -> Awaitable[TokenCountResponse]: ... @@ -1648,18 +1655,23 @@ async def responses_input_tokens( Returns: `{"object": "response.input_tokens", "input_tokens": }` """ + litellm_call_id: Final = resolve_litellm_call_id(request.headers.get(LITELLM_CALL_ID_HEADER)) data: Final = await read_request_body(request=request) model_name: Final = data.get("model") input_value: Final = data.get("input") if not isinstance(model_name, str) or not model_name: - return _missing_responses_param_response("model") + return rejected_token_count(_missing_responses_param_response("model"), litellm_call_id, None) if input_value is None: - return _missing_responses_param_response("input") + return rejected_token_count(_missing_responses_param_response("input"), litellm_call_id, model_name) if isinstance(input_value, (str, list)) and not input_value: - return _responses_invalid_request_response( - message="""One of "input" or "previous_response_id" or 'prompt' or 'conversation' must be provided.""", - param=None, - code="missing_required_parameter", + return rejected_token_count( + _responses_invalid_request_response( + message="""One of "input" or "previous_response_id" or 'prompt' or 'conversation' must be provided.""", + param=None, + code="missing_required_parameter", + ), + litellm_call_id, + model_name, ) counted_fields: Final = { @@ -1667,7 +1679,11 @@ async def responses_input_tokens( } token_fields: Final = resolve_counted_tool_fields(counted_fields, llm_router, user_api_key_dict.team_id) if isinstance(token_fields, ProxyException): - return _responses_invalid_request_response(message=token_fields.message, param=token_fields.param, code=None) + return rejected_token_count( + _responses_invalid_request_response(message=token_fields.message, param=token_fields.param, code=None), + litellm_call_id, + model_name, + ) try: payload: Final[_TokenCountPayload] = { @@ -1680,27 +1696,44 @@ async def responses_input_tokens( } token_request: Final = TokenCountRequest.model_validate(payload) except ValidationError as e: - return _responses_invalid_request_response( - message=first_validation_error(e), param=str(e.errors(include_url=False)[0]["loc"][0]), code=None + return rejected_token_count( + _responses_invalid_request_response( + message=first_validation_error(e), param=str(e.errors(include_url=False)[0]["loc"][0]), code=None + ), + litellm_call_id, + model_name, ) except Exception as e: - return _responses_invalid_request_response( - message=f"Invalid request for token counting: {e}", param=None, code=None + return rejected_token_count( + _responses_invalid_request_response( + message=f"Invalid request for token counting: {e}", param=None, code=None + ), + litellm_call_id, + model_name, ) 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, + try: + 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, - ), - managed_files, - user_api_key_dict, - ) + ) + except ProxyException as e: + raise rejected_token_count(e, litellm_call_id, model_name) try: token_response: Final = await token_counter( - request=token_request, call_endpoint=True, responses_request=responses_request + request=token_request, + call_endpoint=True, + responses_request=responses_request, + litellm_call_id=litellm_call_id, ) except ProxyException as e: if e.type != _PROVIDER_COUNT_ERROR_TYPE: diff --git a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py index c7bb05e5d30..037b0a6c29a 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py @@ -5,9 +5,11 @@ Test for anthropic_endpoints/endpoints.py, focusing on handling dictionary objec import json import logging import unittest +from collections.abc import Callable, Mapping from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from fastapi.testclient import TestClient @@ -615,6 +617,184 @@ async def test_count_tokens_drops_malformed_tools_the_counted_deployment_drops_e assert counter.await_args.kwargs["request"].tools is None +def _count_on_claude( + monkeypatch: pytest.MonkeyPatch, peer: Callable[[httpx.Request], httpx.Response], headers: Mapping[str, str] +) -> httpx.Response: + from fastapi import FastAPI + + import litellm.proxy.proxy_server as proxy_server + from litellm.llms.anthropic.count_tokens import handler as anthropic_count_tokens_handler + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.anthropic_endpoints.endpoints import router + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + claude: Final = {"model": "anthropic/claude-zeta-9", "api_key": "test-key", "api_base": "https://count-tokens.test"} + monkeypatch.setattr( + proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "count-log-claude", + "litellm_params": {**claude, "order": 2}, + "model_info": {"id": "count-log-fallback-deployment"}, + }, + { + "model_name": "count-log-claude", + "litellm_params": {**claude, "order": 1}, + "model_info": {"id": "count-log-routed-deployment"}, + }, + ] + ), + ) + monkeypatch.setattr( + anthropic_count_tokens_handler, + "get_async_httpx_client", + lambda *args, **kwargs: AsyncHTTPHandler(transport=httpx.MockTransport(peer)), + ) + app: Final = FastAPI() + app.include_router(router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth() + return TestClient(app).post( + "/v1/messages/count_tokens", + json={"model": "count-log-claude", "messages": [{"role": "user", "content": "Hello"}]}, + headers=dict(headers), + ) + + +def test_count_tokens_logs_one_line_with_the_callers_call_id_model_and_count( + monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] +) -> None: + response: Final = _count_on_claude( + monkeypatch, + lambda request: httpx.Response(200, json={"input_tokens": 77}, request=request), + {"x-litellm-call-id": "count-log-msg-claude"}, + ) + + assert response.json() == {"input_tokens": 77, "litellm_count_source": "provider"}, response.text + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-msg-claude", + "request_model": "count-log-claude", + "model_id": "count-log-routed-deployment", + "model_used": "claude-zeta-9", + "input_tokens": 77, + "count_source": "provider", + "status_code": 200, + }, + ) + + +def test_count_tokens_provider_rejection_logs_its_status_and_answers_the_same_error( + monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] +) -> None: + rejection: Final = { + "type": "error", + "error": {"type": "invalid_request_error", "message": "messages.0.content.0: invalid base64 data"}, + } + + response: Final = _count_on_claude( + monkeypatch, + lambda request: httpx.Response(400, json=rejection, request=request), + {"x-litellm-call-id": "count-log-msg-error"}, + ) + + assert response.status_code == 400, response.text + assert response.json()["error"] == rejection["error"] + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-msg-error", + "request_model": "count-log-claude", + "model_id": "count-log-routed-deployment", + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 400, + }, + ) + + +def _count_on_local_model(monkeypatch: pytest.MonkeyPatch, body: Mapping[str, object]) -> httpx.Response: + from fastapi import FastAPI + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.anthropic_endpoints.endpoints import router + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + monkeypatch.setattr( + proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "count-log-local", + "litellm_params": {"model": "hosted_vllm/count-log-local", "api_base": "https://local.test/v1"}, + "model_info": {"id": "count-log-local-deployment"}, + } + ] + ), + ) + app: Final = FastAPI() + app.include_router(router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth() + return TestClient(app).post( + "/v1/messages/count_tokens", json=dict(body), headers={"x-litellm-call-id": "count-log-local"} + ) + + +def test_count_tokens_unsupported_content_logs_the_400_it_answers( + monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] +) -> None: + response: Final = _count_on_local_model( + monkeypatch, + {"model": "count-log-local", "messages": [{"role": "user", "content": [{"type": "future_block"}]}]}, + ) + + assert (response.status_code, response.json()["error"]["type"]) == (400, "invalid_request_error"), response.text + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-local", + "request_model": "count-log-local", + "model_id": "count-log-local-deployment", + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 400, + }, + ) + + +@pytest.mark.parametrize( + ("body", "request_model"), + [ + ({"messages": [{"role": "user", "content": "Hello"}]}, None), + ({"model": "count-log-local", "messages": []}, "count-log-local"), + ], +) +def test_count_tokens_rejected_before_counting_logs_the_400_it_answers( + monkeypatch: pytest.MonkeyPatch, + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], + body: Mapping[str, object], + request_model: str | None, +) -> None: + response: Final = _count_on_local_model(monkeypatch, body) + + assert response.status_code == 400, response.text + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-local", + "request_model": request_model, + "model_id": None, + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 400, + }, + ) + + class TestEventLoggingBatchEndpoint: """Test the stubbed event logging batch endpoint""" diff --git a/tests/unit/proxy/common_utils/test_token_count_log.py b/tests/unit/proxy/common_utils/test_token_count_log.py new file mode 100644 index 00000000000..94b976d88cf --- /dev/null +++ b/tests/unit/proxy/common_utils/test_token_count_log.py @@ -0,0 +1,45 @@ +import asyncio +from collections.abc import Callable + +import pytest + +from litellm.proxy.common_utils.token_count_log import logged_token_count +from litellm.types.utils import TokenCountResponse + + +class _TokenizerPanic(BaseException): + pass + + +async def _count_ending_in(error: BaseException) -> TokenCountResponse: + raise error + + +@pytest.mark.asyncio +async def test_a_count_ending_in_a_base_exception_logs_a_500( + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], +) -> None: + with pytest.raises(_TokenizerPanic): + await logged_token_count(_count_ending_in(_TokenizerPanic()), "count-log-panic", "count-log-gpt", None) + + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-panic", + "request_model": "count-log-gpt", + "model_id": None, + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 500, + }, + ) + + +@pytest.mark.asyncio +async def test_a_cancelled_count_writes_no_line( + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], +) -> None: + with pytest.raises(asyncio.CancelledError): + await logged_token_count(_count_ending_in(asyncio.CancelledError()), "count-log-cancelled", "count-log-gpt", None) + + assert token_count_log_lines() == () diff --git a/tests/unit/proxy/conftest.py b/tests/unit/proxy/conftest.py index 58c9a5e4d76..e13eddef43b 100644 --- a/tests/unit/proxy/conftest.py +++ b/tests/unit/proxy/conftest.py @@ -3,6 +3,7 @@ import asyncio import copy import inspect +import logging import os import tempfile import warnings @@ -17,6 +18,7 @@ from prisma.errors import ClientNotConnectedError import litellm import litellm.proxy.proxy_server +from litellm._logging import verbose_proxy_logger from litellm._service_logger import ServiceTypes from litellm.integrations.otel.model.payloads import ServiceSpanData from litellm.integrations.otel.model.spans import service_span_name @@ -424,6 +426,31 @@ def fresh_agent_read_through(monkeypatch): return read_through +TOKEN_COUNT_LOG_FIELDS: Final = ( + "litellm_call_id", + "request_model", + "model_id", + "model_used", + "input_tokens", + "count_source", + "status_code", +) + + +@pytest.fixture +def token_count_log_lines( + caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch +) -> Iterator[Callable[[], tuple[dict[str, object], ...]]]: + """The structured fields of every INFO token count line the proxy logger writes while the test runs.""" + monkeypatch.setattr(verbose_proxy_logger, "propagate", True) + with caplog.at_level(logging.INFO, logger=verbose_proxy_logger.name): + yield lambda: tuple( + {field: record.__dict__[field] for field in TOKEN_COUNT_LOG_FIELDS} + for record in caplog.records + if record.levelno == logging.INFO and "count_source" in record.__dict__ + ) + + @pytest.fixture def postgres_span_names() -> Iterator[Callable[[], Awaitable[tuple[str, ...]]]]: """The ``postgres.{verb} {table}`` names OTel would render for every DB service event diff --git a/tests/unit/proxy/google_endpoints/test_endpoints.py b/tests/unit/proxy/google_endpoints/test_endpoints.py index 17af52b5579..80e45edf897 100644 --- a/tests/unit/proxy/google_endpoints/test_endpoints.py +++ b/tests/unit/proxy/google_endpoints/test_endpoints.py @@ -94,6 +94,7 @@ async def test_proxy_gemini_to_openai_like_model_token_counting(): request=Request( scope={ "type": "http", + "headers": [], "parsed_body": ( ["contents"], {"contents": [{"parts": [{"text": "Hello, how are you?"}]}]}, diff --git a/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py b/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py index 2777d18a399..4867ca9ab63 100644 --- a/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py +++ b/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py @@ -3,9 +3,11 @@ Test to verify the Google GenAI proxy API endpoints """ +from collections.abc import Callable from typing import Final from unittest.mock import AsyncMock, patch +import httpx import pytest from litellm.types.utils import TokenCountResponse @@ -240,3 +242,79 @@ def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs() assert body["promptTokensDetails"] == [] assert body["litellm_count_source"] == "token_count_model" assert body["litellm_count_model"] == "us.anthropic.claude-opus-5-5" + + +@pytest.mark.parametrize("path", ["/v1beta/models/count-log-claude:countTokens", "/models/count-log-claude:countTokens"]) +def test_google_count_tokens_logs_one_line_with_the_callers_call_id_model_and_count( + monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], path: str +) -> None: + import litellm.proxy.proxy_server as proxy_server + from litellm import Router + from litellm.llms.anthropic.count_tokens import handler as anthropic_count_tokens_handler + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + claude: Final = {"model": "anthropic/claude-zeta-9", "api_key": "test-key", "api_base": "https://count-tokens.test"} + monkeypatch.setattr( + proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "count-log-claude", + "litellm_params": claude, + "model_info": {"id": "count-log-gemini-deployment"}, + } + ] + ), + ) + monkeypatch.setattr( + anthropic_count_tokens_handler, + "get_async_httpx_client", + lambda *args, **kwargs: AsyncHTTPHandler( + transport=httpx.MockTransport( + lambda request: httpx.Response(200, json={"input_tokens": 77}, request=request) + ) + ), + ) + + response: Final = _build_test_client().post( + path, + json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, + headers={"x-litellm-call-id": "count-log-gemini"}, + ) + + assert (response.status_code, response.json()["totalTokens"]) == (200, 77), response.text + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-gemini", + "request_model": "count-log-claude", + "model_id": "count-log-gemini-deployment", + "model_used": "claude-zeta-9", + "input_tokens": 77, + "count_source": "provider", + "status_code": 200, + }, + ) + + +def test_google_count_tokens_rejected_body_logs_the_400_it_answers( + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], +) -> None: + response: Final = _build_test_client().post( + "/v1beta/models/count-log-claude:countTokens", + json={"contents": "not a list"}, + headers={"x-litellm-call-id": "count-log-gemini-bad-body"}, + ) + + assert response.status_code == 400, response.text + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-gemini-bad-body", + "request_model": "count-log-claude", + "model_id": None, + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 400, + }, + ) diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 60043daa402..43fe5d4e786 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -2142,7 +2142,6 @@ class TestBedrockPassthroughAwsErrors: assert exc_info.value.aws_error_type == "ValidationException" assert (exc_info.value.error.message, exc_info.value.error.code) == (message, status_code) - @pytest.mark.asyncio async def test_bedrock_request_check_is_raised_with_the_aws_error_type_it_carries( self, monkeypatch: pytest.MonkeyPatch @@ -2183,6 +2182,7 @@ class TestBedrockPassthroughAwsErrors: await self._answer(error) assert (exc_info.value.code, exc_info.value.headers.get("x-model-group")) == ("403", "modèle") + @pytest.mark.asyncio @pytest.mark.parametrize( ("error", "message", "code"), @@ -7250,6 +7250,189 @@ async def test_bedrock_count_tokens_error_forwards_provider_headers(): assert exc_info.value.headers["llm_provider-x-amzn-requestid"] == "req-count-tokens-500" +_BEDROCK_COUNT_ALIAS: Final = "count-log-bedrock" +_BEDROCK_COUNT_MODEL: Final = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" + + +async def _bedrock_native_count( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, aws_reply: httpx.Response +) -> dict[str, object]: + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + handle_bedrock_count_tokens, + ) + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.setattr( + "litellm.proxy.proxy_server.llm_router", + litellm.Router( + model_list=[ + { + "model_name": _BEDROCK_COUNT_ALIAS, + "litellm_params": { + "model": _BEDROCK_COUNT_MODEL, + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKIATESTACCESSKEY", + "aws_secret_access_key": "test-secret", + }, + "model_info": {"id": "count-log-bedrock-deployment"}, + } + ] + ), + ) + respx_mock.post(host="bedrock-runtime.us-east-1.amazonaws.com", path__regex=r"/count-tokens$").mock( + return_value=aws_reply + ) + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": f"/bedrock/model/{_BEDROCK_COUNT_ALIAS}/count-tokens", + "headers": [(b"x-litellm-call-id", b"bedrock-native-count")], + } + ) + return await handle_bedrock_count_tokens( + endpoint=f"model/{_BEDROCK_COUNT_ALIAS}/count-tokens", + request=request, + fastapi_response=Response(), + user_api_key_dict=MagicMock(), + request_body={"model": _BEDROCK_COUNT_ALIAS, "messages": [{"role": "user", "content": "hi"}]}, + ) + + +@pytest.mark.asyncio +async def test_bedrock_native_count_tokens_logs_the_count_against_the_router_model( + monkeypatch: pytest.MonkeyPatch, + httpx_transport: None, + respx_mock: respx.MockRouter, + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], +) -> None: + result: Final = await _bedrock_native_count( + monkeypatch, respx_mock, httpx.Response(200, json={"inputTokens": 12}) + ) + + assert result == {"input_tokens": 12} + assert len(respx_mock.calls) == 1 + assert respx_mock.calls.last.request.url.path == "/model/anthropic.claude-haiku-4-5-20251001-v1:0/count-tokens" + assert token_count_log_lines() == ( + { + "litellm_call_id": "bedrock-native-count", + "request_model": _BEDROCK_COUNT_ALIAS, + "model_id": "count-log-bedrock-deployment", + "model_used": "anthropic.claude-haiku-4-5-20251001-v1:0", + "input_tokens": 12, + "count_source": "provider", + "status_code": 200, + }, + ) + + +@pytest.mark.asyncio +async def test_bedrock_native_count_tokens_logs_the_provider_status_when_the_count_fails( + monkeypatch: pytest.MonkeyPatch, + httpx_transport: None, + respx_mock: respx.MockRouter, + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], +) -> None: + rejected: Final = httpx.Response(400, json={"message": "The provided model doesn't support counting tokens."}) + with pytest.raises(HTTPException) as exc_info: + await _bedrock_native_count(monkeypatch, respx_mock, rejected) + + assert exc_info.value.status_code == 400 + assert token_count_log_lines() == ( + { + "litellm_call_id": "bedrock-native-count", + "request_model": _BEDROCK_COUNT_ALIAS, + "model_id": "count-log-bedrock-deployment", + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 400, + }, + ) + + +@pytest.mark.asyncio +async def test_bedrock_native_count_tokens_without_a_model_logs_the_400_it_answers( + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], +) -> None: + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + handle_bedrock_count_tokens, + ) + + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": f"/bedrock/model/{_BEDROCK_COUNT_ALIAS}/count-tokens", + "headers": [(b"x-litellm-call-id", b"bedrock-native-no-model")], + } + ) + with pytest.raises(HTTPException) as exc_info: + await handle_bedrock_count_tokens( + endpoint=f"model/{_BEDROCK_COUNT_ALIAS}/count-tokens", + request=request, + fastapi_response=Response(), + user_api_key_dict=MagicMock(), + request_body={"messages": [{"role": "user", "content": "hi"}]}, + ) + + assert exc_info.value.status_code == 400 + assert token_count_log_lines() == ( + { + "litellm_call_id": "bedrock-native-no-model", + "request_model": None, + "model_id": None, + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 400, + }, + ) + + + +@pytest.mark.asyncio +async def test_bedrock_native_count_tokens_with_a_region_only_model_answers_the_validation_error_and_logs_it( + monkeypatch: pytest.MonkeyPatch, + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], +) -> None: + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + handle_bedrock_count_tokens, + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": "/bedrock/model/us/count-tokens", + "headers": [(b"x-litellm-call-id", b"bedrock-native-region-only")], + } + ) + with pytest.raises(HTTPException) as exc_info: + await handle_bedrock_count_tokens( + endpoint="model/us/count-tokens", + request=request, + fastapi_response=Response(), + user_api_key_dict=MagicMock(), + request_body={"model": "us"}, + ) + + assert exc_info.value.detail == { + "error": "CountTokens processing error: Request must contain content to count tokens" + } + assert token_count_log_lines() == ( + { + "litellm_call_id": "bedrock-native-region-only", + "request_model": "us", + "model_id": None, + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": exc_info.value.status_code, + }, + ) + class _AzureGroupRouter: def __init__(self, captured: list[dict]) -> None: self.captured = captured diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 985aee061bc..7760b6b7fcc 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -7,7 +7,7 @@ import contextlib import dataclasses import json import unittest -from collections.abc import Mapping +from collections.abc import Callable, Mapping from typing import Any, Final, Literal, cast from unittest.mock import AsyncMock, MagicMock, call, patch @@ -19,6 +19,7 @@ from fastapi.testclient import TestClient from httpx import Response import litellm +from litellm.constants import EMPTY_MAPPING from litellm.proxy._types import ProxyException from litellm.proxy.proxy_server import app from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse @@ -4275,6 +4276,7 @@ class TestResponsesInputTokens: llm_router: litellm.Router | None = None, response_id_security: object | None = None, managed_files: object | None = None, + headers: Mapping[str, str] = EMPTY_MAPPING, ) -> tuple[Response, AsyncMock]: from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -4307,7 +4309,7 @@ class TestResponsesInputTokens: 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"}) + response = client.post(path, json=body, headers={"Authorization": "Bearer sk-9876", **headers}) return response, token_counter_mock finally: app.dependency_overrides.pop(user_api_key_auth, None) @@ -4552,6 +4554,84 @@ class TestResponsesInputTokens: assert (response.status_code, response.json().get("input_tokens")) == (200, 7), response.text assert sent == [{"model": "gpt-4o-mini", "input": "hi"}] + @pytest.mark.parametrize( + "path", ["/v1/responses/input_tokens", "/responses/input_tokens", "/openai/v1/responses/input_tokens"] + ) + def test_count_logs_one_line_with_the_callers_call_id_model_and_count( + self, + monkeypatch: pytest.MonkeyPatch, + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], + path: str, + ) -> 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 + + gpt: Final = {"model": "openai/gpt-4o-mini", "api_key": "test-key", "api_base": "https://count.test/v1"} + router: Final = litellm.Router( + model_list=[ + { + "model_name": "count-log-gpt", + "litellm_params": gpt, + "model_info": {"id": "count-log-responses-deployment"}, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + monkeypatch.setattr( + openai_count_tokens_http, + "get_async_httpx_client", + lambda *args, **kwargs: AsyncHTTPHandler( + transport=httpx.MockTransport( + lambda request: httpx.Response( + 200, json={"object": "response.input_tokens", "input_tokens": 7}, request=request + ) + ) + ), + ) + + response, _ = self._post_input_tokens( + {"model": "count-log-gpt", "input": "hi"}, + path=path, + counter=AsyncMock(side_effect=count_request_tokens), + llm_router=router, + headers={"x-litellm-call-id": "count-log-responses"}, + ) + + assert (response.status_code, response.json().get("input_tokens")) == (200, 7), response.text + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-responses", + "request_model": "count-log-gpt", + "model_id": "count-log-responses-deployment", + "model_used": "gpt-4o-mini", + "input_tokens": 7, + "count_source": "provider", + "status_code": 200, + }, + ) + + def test_count_rejected_before_counting_logs_the_400_it_answers( + self, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] + ) -> None: + response, counter = self._post_input_tokens( + {"model": "count-log-gpt"}, headers={"x-litellm-call-id": "count-log-no-input"} + ) + + assert response.status_code == 400, response.text + counter.assert_not_awaited() + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-no-input", + "request_model": "count-log-gpt", + "model_id": None, + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 400, + }, + ) + def test_a_managed_file_the_key_cannot_read_answers_403_without_a_count(self) -> None: from litellm.integrations.custom_logger import CustomLogger diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index ebe3c7407d8..5484bc38317 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -10,6 +10,7 @@ import json import logging import threading import traceback +import uuid from collections.abc import Callable, Coroutine, Mapping from functools import partial from pathlib import Path @@ -40,7 +41,7 @@ from starlette.types import Message import litellm from litellm import Router -from litellm._logging import verbose_proxy_logger +from litellm._logging import JsonFormatter, verbose_proxy_logger from litellm.llms.anthropic.count_tokens import handler as anthropic_count_tokens_handler from litellm.llms.base_llm.base_utils import BaseTokenCounter from litellm.llms.bedrock.common_utils import BedrockError @@ -364,6 +365,7 @@ async def test_anthropic_messages_count_tokens_endpoint(): output_format=None, context_management=None, request_litellm_params=None, + litellm_call_id=None, client_oauth_token=None, ): assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" @@ -445,6 +447,7 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): output_format=None, context_management=None, request_litellm_params=None, + litellm_call_id=None, client_oauth_token=None, ): assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" @@ -1264,6 +1267,7 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): output_format=None, context_management=None, request_litellm_params=None, + litellm_call_id=None, client_oauth_token=None, ): raise ProxyException( @@ -1322,6 +1326,7 @@ async def test_anthropic_endpoint_403_permission_error_format(): output_format=None, context_management=None, request_litellm_params=None, + litellm_call_id=None, client_oauth_token=None, ): raise ProxyException( @@ -1380,6 +1385,7 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): output_format=None, context_management=None, request_litellm_params=None, + litellm_call_id=None, client_oauth_token=None, ): raise ProxyException( @@ -3897,3 +3903,149 @@ async def test_responses_count_on_claude_counts_the_turns_previous_response_id_r ] 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")) + + +_COUNTED_HELLO: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"model": "count-log-gpt", "messages": [{"role": "user", "content": "Hello, how many tokens is this?"}]} +) + + +def _utils_token_counter( + monkeypatch: pytest.MonkeyPatch, body: Mapping[str, JsonValue], headers: Mapping[str, str] +) -> httpx.Response: + from fastapi.testclient import TestClient + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.proxy_server import app + + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "count-log-gpt", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"}, + "model_info": {"id": "count-log-utils-deployment"}, + } + ] + ), + ) + monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, lambda: UserAPIKeyAuth(api_key="sk-test")) + return TestClient(app).post("/utils/token_counter", json=dict(body), headers=dict(headers)) + + +def test_utils_token_counter_logs_one_line_with_the_callers_call_id_model_and_count( + monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] +) -> None: + response: Final = _utils_token_counter(monkeypatch, _COUNTED_HELLO, {"x-litellm-call-id": "count-log-utils"}) + + assert response.status_code == 200, response.text + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-utils", + "request_model": "count-log-gpt", + "model_id": "count-log-utils-deployment", + "model_used": "gpt-4o-mini", + "input_tokens": response.json()["total_tokens"], + "count_source": "local_estimate", + "status_code": 200, + }, + ) + + +def test_utils_token_counter_logs_a_generated_call_id_when_the_client_sends_none( + monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] +) -> None: + response: Final = _utils_token_counter(monkeypatch, _COUNTED_HELLO, {}) + + assert response.status_code == 200, response.text + (line,) = token_count_log_lines() + assert str(uuid.UUID(str(line["litellm_call_id"]))) == line["litellm_call_id"] + + +def test_utils_token_counter_without_messages_logs_the_400_it_answers( + monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] +) -> None: + response: Final = _utils_token_counter( + monkeypatch, {"model": "count-log-gpt"}, {"x-litellm-call-id": "count-log-utils-empty"} + ) + + assert response.status_code == 400, response.text + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-utils-empty", + "request_model": "count-log-gpt", + "model_id": None, + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 400, + }, + ) + + +def test_utils_token_counter_failed_count_logs_its_status_in_place_of_a_count( + monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] +) -> None: + response: Final = _utils_token_counter( + monkeypatch, {**_COUNTED_HELLO, "prompt": "and a prompt"}, {"x-litellm-call-id": "count-log-utils-error"} + ) + + assert response.status_code == 400, response.text + assert response.json() == { + "error": { + "message": "prompt cannot be sent with messages or contents. Set drop_params to count them without the prompt", + "type": "invalid_request_error", + "param": "prompt", + "code": "400", + } + } + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-utils-error", + "request_model": "count-log-gpt", + "model_id": "count-log-utils-deployment", + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 400, + }, + ) + + +def test_utils_token_counter_json_line_keeps_its_own_fields_when_the_call_id_is_a_dict( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], +) -> None: + forged_call_id: Final = "{'component':'forged-comp','logger':'forged.py:1','injected_key':'injected'}" + + _utils_token_counter(monkeypatch, _COUNTED_HELLO, {"x-litellm-call-id": forged_call_id}) + + (line,) = token_count_log_lines() + (record,) = (record for record in caplog.records if "count_source" in record.__dict__) + rendered: Final = json.loads(JsonFormatter().format(record)) + assert line["litellm_call_id"] == forged_call_id + assert {key: rendered.get(key) for key in ("litellm_call_id", "component", "logger", "injected_key")} == { + "litellm_call_id": forged_call_id, + "component": verbose_proxy_logger.name, + "logger": f"{record.filename}:{record.lineno}", + "injected_key": None, + } + + +def test_utils_token_counter_line_stays_one_line_when_the_model_carries_a_newline( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], +) -> None: + forged_model: Final = "count-log-gpt\ntoken count - litellm_call_id=forged input_tokens=1" + + _utils_token_counter(monkeypatch, {**_COUNTED_HELLO, "model": forged_model}, {"x-litellm-call-id": "count log"}) + + (line,) = token_count_log_lines() + (message,) = (record.getMessage() for record in caplog.records if "count_source" in record.__dict__) + assert line["request_model"] == forged_model + assert "\n" not in message + assert message.startswith(f'token count - litellm_call_id="count log" request_model={json.dumps(forged_model)} ') diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 76e3c95fb75..f9117afdf3a 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -19058,6 +19058,7 @@ export interface paths { * @description Args: * request: TokenCountRequest * call_endpoint: bool - When set to "True" it will call the token counting endpoint - e.g Anthropic or Google AI Studio Token Counting APIs. + * litellm_call_id: the x-litellm-call-id header, logged with the count so it joins the inference call's logs * * Returns: * TokenCountResponse @@ -76773,7 +76774,9 @@ export interface operations { query?: { call_endpoint?: boolean; }; - header?: never; + header?: { + "x-litellm-call-id"?: string | null; + }; path?: never; cookie?: never; }; From 54a6a5fb3a86039b3f810c9d8083f8ae2d614dfd Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 28 Oct 2026 14:24:34 -0400 Subject: [PATCH 061/128] fix(otel): suppress spans for excluded proxy requests (#48720) --- litellm/integrations/otel/mount.py | 42 +++++- litellm/integrations/otel/plumbing/context.py | 15 ++ .../integrations/otel/plumbing/providers.py | 27 +++- .../otel/test_otel_v2_components.py | 22 +++ .../integrations/otel/test_otel_v2_mount.py | 141 ++++++++++++++++++ .../test_mcp_server_tool_calls_and_headers.py | 113 ++++++++++++++ 6 files changed, 354 insertions(+), 6 deletions(-) diff --git a/litellm/integrations/otel/mount.py b/litellm/integrations/otel/mount.py index c0d42eeb5c5..9d6c5fca75c 100644 --- a/litellm/integrations/otel/mount.py +++ b/litellm/integrations/otel/mount.py @@ -22,6 +22,7 @@ from litellm.integrations.otel.model.config import is_otel_v2_enabled if TYPE_CHECKING: from fastapi import FastAPI from opentelemetry.trace import Span, TracerProvider + from starlette.types import ASGIApp, Receive, Scope, Send # Routes excluded from server-span tracing by default: high-frequency pollers and # static UI/docs assets, none of which are LLM traffic. Entries are substring-matched @@ -44,6 +45,33 @@ _DEFAULT_EXCLUDED_ROUTES: Final = ( ) _DEFAULT_EXCLUDED_URLS: Final = ",".join(_DEFAULT_EXCLUDED_ROUTES) + +class _RequestTraceExclusionMiddleware: + def __init__(self, app: "ASGIApp", excluded_urls: str) -> None: + from opentelemetry.util.http import parse_excluded_urls + + self.app = app + self.excluded_urls = parse_excluded_urls(excluded_urls) + + async def __call__(self, scope: "Scope", receive: "Receive", send: "Send") -> None: + if scope["type"] not in ("http", "websocket"): + await self.app(scope, receive, send) + return + from opentelemetry.instrumentation.asgi import ( # pyright: ignore[reportMissingTypeStubs] # upstream ships no py.typed marker + get_host_port_url_tuple, # pyright: ignore[reportUnknownVariableType] # upstream URL helper has no annotations + ) + + from litellm.integrations.otel.plumbing.context import ( + request_trace_scope, + request_traces_excluded, + ) + + url: Final = str(get_host_port_url_tuple(scope)[2]) # pyright: ignore[reportUnknownArgumentType] # upstream constructs a str URL but leaves its return type unannotated + excluded: Final = request_traces_excluded() or self.excluded_urls.url_disabled(url) + with request_trace_scope(excluded): + await self.app(scope, receive, send) + + # Passthrough routes are catch-alls (e.g. "/openai/{endpoint:path}"), so the # default OTel server-span name "{method} {route}" collapses every upstream # endpoint into "POST /openai/{endpoint:path}". The hook below renames those spans @@ -129,6 +157,8 @@ def instrument_fastapi_app(app: "FastAPI", tracer_provider: "TracerProvider | No if app.middleware_stack is not None: verbose_logger.debug("Skipping OTel V2 FastAPI instrumentation: the app has already started") return + if getattr(app, "_is_instrumented_by_opentelemetry", False): + return # Lazy: only the V2-enabled path needs the optional # ``opentelemetry-instrumentation-fastapi`` package, which is not part of the @@ -151,11 +181,7 @@ def instrument_fastapi_app(app: "FastAPI", tracer_provider: "TracerProvider | No claim_server_span(span) _passthrough_span_name_hook(span, scope) - excluded_urls: Final = ( - os.environ.get("OTEL_PYTHON_FASTAPI_EXCLUDED_URLS") - if "OTEL_PYTHON_FASTAPI_EXCLUDED_URLS" in os.environ - else _DEFAULT_EXCLUDED_URLS - ) + excluded_urls: Final = os.environ.get("OTEL_PYTHON_FASTAPI_EXCLUDED_URLS", _DEFAULT_EXCLUDED_URLS) FastAPIInstrumentor.instrument_app( app, tracer_provider=tracer_provider, @@ -166,6 +192,12 @@ def instrument_fastapi_app(app: "FastAPI", tracer_provider: "TracerProvider | No # template in their name, which can't be rewritten from a hook. exclude_spans=["receive", "send"], ) + instrumented_stack: Final = app.build_middleware_stack + + def build_middleware_stack() -> "ASGIApp": + return _RequestTraceExclusionMiddleware(instrumented_stack(), excluded_urls) + + app.build_middleware_stack = build_middleware_stack # rebind-ok: wrap the complete ASGI stack if get_global_response_propagator() is None: set_global_response_propagator(UpstreamFirstTraceResponsePropagator()) except Exception as e: diff --git a/litellm/integrations/otel/plumbing/context.py b/litellm/integrations/otel/plumbing/context.py index d0c419abe87..b328b281dde 100644 --- a/litellm/integrations/otel/plumbing/context.py +++ b/litellm/integrations/otel/plumbing/context.py @@ -30,6 +30,21 @@ if TYPE_CHECKING: _PROPAGATOR: Final = TraceContextTextMapPropagator() _W3C_TRACE_HEADERS: Final = frozenset(("traceparent", "tracestate")) +_request_traces_excluded: Final[ContextVar[bool]] = ContextVar("litellm_otel_request_traces_excluded", default=False) + + +def request_traces_excluded() -> bool: + return _request_traces_excluded.get() + + +@contextmanager +def request_trace_scope(excluded: bool) -> Generator[None]: + token: Final = _request_traces_excluded.set(excluded) + try: + yield + finally: + _request_traces_excluded.reset(token) + # The request's root span — the FastAPI-owned SERVER span — captured ONCE when the # proxy first resolves it, so request-level spans (the LLM call, guardrails) can diff --git a/litellm/integrations/otel/plumbing/providers.py b/litellm/integrations/otel/plumbing/providers.py index 9736192a816..a369e352609 100644 --- a/litellm/integrations/otel/plumbing/providers.py +++ b/litellm/integrations/otel/plumbing/providers.py @@ -35,7 +35,8 @@ from opentelemetry.sdk.trace.export import ( from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( InMemorySpanExporter, ) -from opentelemetry.trace import Span, SpanContext, SpanKind, Status, StatusCode, Tracer +from opentelemetry.sdk.trace.sampling import Decision, Sampler, SamplingResult +from opentelemetry.trace import Link, Span, SpanContext, SpanKind, Status, StatusCode, Tracer, TraceState from opentelemetry.util.re import parse_env_headers from opentelemetry.util.types import Attributes, AttributeValue @@ -68,6 +69,7 @@ from litellm.integrations.otel.model.semconv import ( from litellm.integrations.otel.model.spans import LiteLLMSpanKind from litellm.integrations.otel.plumbing.context import ( request_destinations, + request_traces_excluded, suppressed_backends, ) from litellm.integrations.otel.plumbing.otlp_tls import resolve_otlp_http_tls @@ -1499,6 +1501,28 @@ def _spec_processor( return _MessageContentFilter(scoped) if strip_message_content else scoped +class _RequestExclusionSampler(Sampler): + def __init__(self, delegate: Sampler) -> None: + self.delegate = delegate + + def should_sample( + self, + parent_context: Context | None, + trace_id: int, + name: str, + kind: SpanKind | None = None, + attributes: Attributes = None, + links: Sequence[Link] | None = None, + trace_state: TraceState | None = None, + ) -> SamplingResult: + if request_traces_excluded(): + return SamplingResult(Decision.DROP) + return self.delegate.should_sample(parent_context, trace_id, name, kind, attributes, links, trace_state) + + def get_description(self) -> str: + return f"RequestExclusion({self.delegate.get_description()})" + + def build_tracer_provider( config: OpenTelemetryV2Config, exporter: SpanExporter | None = None, @@ -1531,6 +1555,7 @@ def build_tracer_provider( destination asked for. """ provider: Final = TracerProvider(resource=build_resource(config)) + provider.sampler = _RequestExclusionSampler(provider.sampler) if baggage_processor is None: baggage_processor = LiteLLMBaggageSpanProcessor(allowed_keys=config.baggage_promoted_keys) provider.add_span_processor(baggage_processor) diff --git a/tests/unit/integrations/otel/test_otel_v2_components.py b/tests/unit/integrations/otel/test_otel_v2_components.py index 8af13fd6e7e..aad06d33edb 100644 --- a/tests/unit/integrations/otel/test_otel_v2_components.py +++ b/tests/unit/integrations/otel/test_otel_v2_components.py @@ -1076,6 +1076,28 @@ def test_build_resource_includes_deployment_environment(): assert resource.attributes["deployment.environment"] == "prod" +@pytest.mark.parametrize("sampler", ["always_on", "always_off", "parentbased_always_on"]) +def test_request_exclusion_preserves_configured_sampling(monkeypatch: pytest.MonkeyPatch, sampler: str) -> None: + from opentelemetry.trace import NonRecordingSpan, SpanContext, set_span_in_context + + monkeypatch.setenv("OTEL_TRACES_SAMPLER", sampler) + provider: Final = providers.build_tracer_provider(OpenTelemetryV2Config(exporter="in_memory")) + tracer: Final = provider.get_tracer("request-exclusions") + sampled_parent: Final = set_span_in_context( + NonRecordingSpan(SpanContext(trace_id=1, span_id=2, is_remote=True, trace_flags=TraceFlags(1))) + ) + try: + with ctx_mod.request_trace_scope(True): + excluded: Final = tracer.start_span("excluded", context=sampled_parent) + assert not excluded.is_recording() + excluded.end() + included: Final = tracer.start_span("included", context=sampled_parent) + assert included.is_recording() == (sampler != "always_off") + included.end() + finally: + provider.shutdown() + + def test_build_tracer_provider_processor_selection(): cfg = OpenTelemetryV2Config(exporter="in_memory") simple = providers.build_tracer_provider(cfg, exporter=InMemorySpanExporter()) diff --git a/tests/unit/integrations/otel/test_otel_v2_mount.py b/tests/unit/integrations/otel/test_otel_v2_mount.py index 61f8b9ae0d9..32e846cc1b9 100644 --- a/tests/unit/integrations/otel/test_otel_v2_mount.py +++ b/tests/unit/integrations/otel/test_otel_v2_mount.py @@ -5,8 +5,11 @@ surface and the server-span + shared-provider behavior it produces. """ +import asyncio import contextvars +from contextlib import nullcontext from datetime import datetime, timezone +from typing import Final import pytest @@ -16,6 +19,7 @@ pytest.importorskip("opentelemetry.instrumentation.fastapi") fastapi = pytest.importorskip("fastapi") from fastapi.middleware.cors import CORSMiddleware # noqa: E402 +from fastapi import FastAPI, HTTPException, Request, Response # noqa: E402 # optional dependencies are skipped first from fastapi.responses import StreamingResponse # noqa: E402 from fastapi.testclient import TestClient # noqa: E402 from starlette.routing import Route # noqa: E402 @@ -36,6 +40,7 @@ from opentelemetry.sdk.trace.sampling import ALWAYS_OFF # noqa: E402 from opentelemetry.trace import SpanKind # noqa: E402 from litellm.constants import LITELLM_UI_ALLOW_HEADERS # noqa: E402 +from litellm.types.services import ServiceLoggerPayload, ServiceTypes # noqa: E402 # optional dependencies are skipped first from litellm.integrations.otel.model.config import ( # noqa: E402 OpenTelemetryV2Config, @@ -548,3 +553,139 @@ def test_the_unhandled_error_helper_names_the_span_it_is_given_in_any_context(mo headers = contextvars.Context().run(trace_response_headers, span) assert headers["traceresponse"] == f"00-{context.trace_id:032x}-{context.span_id:016x}-{context.trace_flags:02x}" + + +@pytest.mark.parametrize( + ("excluded_urls", "path", "root_path", "status", "excluded"), + [ + ("/health,/model/info", "/model/info", "", 200, True), + ("/model/info", "/model/info", "", 401, True), + ("/model/info", "/v1/models", "", 200, False), + (None, "/health/readiness", "", 200, True), + (None, "/model/info", "", 200, False), + ("", "/health/readiness", "", 200, False), + (r"^http://testserver/proxy/model/info$", "/model/info", "/proxy", 200, True), + (r"^http://other/model/info$", "/model/info", "", 200, False), + ], +) +@pytest.mark.parametrize("active_parent", [False, True]) +def test_url_exclusions_cover_auth_and_service_spans( + monkeypatch: pytest.MonkeyPatch, excluded_urls: str | None, path: str, + root_path: str, status: int, excluded: bool, active_parent: bool, +) -> None: + monkeypatch.setenv("LITELLM_OTEL_V2", "1") + if excluded_urls is None: + monkeypatch.delenv("OTEL_PYTHON_FASTAPI_EXCLUDED_URLS", raising=False) + else: + monkeypatch.setenv("OTEL_PYTHON_FASTAPI_EXCLUDED_URLS", excluded_urls) + app: Final = FastAPI(root_path=root_path) + logger: Final = OpenTelemetryV2(config=OpenTelemetryV2Config(exporter="in_memory")) + exporter: Final = InMemorySpanExporter() + logger.tracer_provider.add_span_processor(SimpleSpanProcessor(exporter)) + + @app.get(path) + async def endpoint() -> Response: + with logger.start_phase_span(f"auth {path}"), logger.tracer.start_as_current_span("direct SDK span"): + await logger.async_service_success_hook( + payload=ServiceLoggerPayload( + is_error=False, service=ServiceTypes.DB, duration=1, + call_type="get_data", event_metadata=None, + ), + start_time=1, end_time=2, + ) + if status == 401: + raise HTTPException(status_code=status) + return Response(status_code=status) + + instrument_fastapi_app(app, tracer_provider=logger.tracer_provider) + with ( + logger.tracer.start_as_current_span("caller") if active_parent else nullcontext(), + TestClient(app) as client, + ): + response: Final = client.get(f"{root_path}{path}?poll=1") + + names: Final = tuple(span.name for span in exporter.get_finished_spans()) + expected_request: Final = () if excluded else ( + "postgres get_data", "direct SDK span", f"auth {path}", f"GET {path}", + ) + assert response.status_code == status + assert names == expected_request + (("caller",) if active_parent else ()) + + +async def test_exclusion_survives_deferred_work_without_suppressing_other_requests( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("LITELLM_OTEL_V2", "1") + monkeypatch.setenv("OTEL_PYTHON_FASTAPI_EXCLUDED_URLS", "/model/info") + app: Final = FastAPI() + logger: Final = OpenTelemetryV2(config=OpenTelemetryV2Config(exporter="in_memory")) + exporter: Final = InMemorySpanExporter() + logger.tracer_provider.add_span_processor(SimpleSpanProcessor(exporter)) + entered: Final = asyncio.Event() + release: Final = asyncio.Event() + finish: Final = asyncio.Event() + deferred: Final[asyncio.Queue[asyncio.Task[None]]] = asyncio.Queue() + + async def after_response(path: str) -> None: + await finish.wait() + with logger.start_phase_span(f"deferred {path}"): + pass + + @app.get("/{path:path}") + async def endpoint(path: str) -> Response: + with logger.start_phase_span(f"auth {path}"): + if path == "model/info": + entered.set() + await release.wait() + await deferred.put(asyncio.create_task(after_response(path))) + return Response(status_code=200) + + instrument_fastapi_app(app, tracer_provider=logger.tracer_provider) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://testserver") as client: + excluded: Final = asyncio.create_task(client.get("/model/info")) + await entered.wait() + included: Final = await client.get("/v1/models") + release.set() + assert (await excluded).status_code == included.status_code == 200 + finish.set() + await asyncio.gather(await deferred.get(), await deferred.get()) + with logger.start_phase_span("independent SDK work"): + pass + + assert sorted(span.name for span in exporter.get_finished_spans()) == [ + "GET /{path:path}", "auth v1/models", "deferred v1/models", "independent SDK work", + ] + + +@pytest.mark.parametrize("path", ["/excluded", "/included"]) +def test_url_exclusions_cover_unhandled_error_handlers_after_repeated_mount( + monkeypatch: pytest.MonkeyPatch, path: str, +) -> None: + monkeypatch.setenv("LITELLM_OTEL_V2", "1") + monkeypatch.setenv("OTEL_PYTHON_FASTAPI_EXCLUDED_URLS", "/excluded") + app: Final = FastAPI() + logger: Final = OpenTelemetryV2(config=OpenTelemetryV2Config(exporter="in_memory")) + exporter: Final = InMemorySpanExporter() + logger.tracer_provider.add_span_processor(SimpleSpanProcessor(exporter)) + + async def error_handler(_request: Request, _exc: Exception) -> Response: + with logger.start_phase_span("error handler"), logger.tracer.start_as_current_span("error SDK span"): + return Response("handled", status_code=500) + + app.add_exception_handler(Exception, error_handler) + + @app.get(path) + async def endpoint() -> None: + with logger.start_phase_span("auth"): + raise ValueError("unhandled") + + instrument_fastapi_app(app, tracer_provider=logger.tracer_provider) + instrument_fastapi_app(app, tracer_provider=logger.tracer_provider) + with TestClient(app, raise_server_exceptions=False) as client: + response: Final = client.get(path) + + assert response.status_code == 500 + assert response.text == "handled" + assert tuple(span.name for span in exporter.get_finished_spans()) == ( + () if path == "/excluded" else ("auth", "error SDK span", "error handler", "GET /included") + ) diff --git a/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py b/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py index ccd5f7f2fea..78c37184d6c 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py +++ b/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py @@ -9209,6 +9209,119 @@ class TestMCPMetaTraceCarrier: assert _mcp_meta_trace_carrier(SimpleNamespace(meta=only_progress)) is None +@pytest.mark.asyncio +@pytest.mark.parametrize("initialized_excluded", [False, True]) +async def test_stateful_mcp_tracing_follows_each_message_exclusion( + monkeypatch: pytest.MonkeyPatch, initialized_excluded: bool +) -> None: + import httpx2 + from fastapi import FastAPI + from mcp.server.context import HandlerResult + from mcp.server.lowlevel import Server + from mcp.server.streamable_http_manager import StreamableHTTPSessionManager + from mcp.shared import _otel + from mcp.types import CallToolRequestParams, NotificationParams + from mcp_types import CLIENT_CAPABILITIES_META_KEY, PROTOCOL_VERSION_META_KEY + from opentelemetry.instrumentation.propagators import get_global_response_propagator, set_global_response_propagator + + from litellm.integrations.otel.model.config import is_otel_v2_enabled + from litellm.integrations.otel.mount import instrument_fastapi_app + from litellm.integrations.otel.plumbing.providers import in_memory_provider + from litellm.proxy import proxy_server + from litellm.proxy.mcp_server import server as gateway + + original_propagator: Final = get_global_response_propagator() + provider, exporter = in_memory_provider() + tracer: Final = provider.get_tracer("mcp-exclusion-test") + monkeypatch.setattr(_otel, "_tracer", tracer) + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + monkeypatch.setenv("OTEL_PYTHON_FASTAPI_EXCLUDED_URLS", "excluded.example") + monkeypatch.setitem(proxy_server.general_settings, "mcp_advertised_versions", ["2025-11-25", "2026-07-28"]) + is_otel_v2_enabled.cache_clear() + notified: Final = asyncio.Event() + + async def call_tool( + ctx: ServerRequestContext[object, object], params: CallToolRequestParams + ) -> CallToolResult: + with tracer.start_as_current_span("tool-handler"): + return CallToolResult(content=[TextContent(type="text", text="done")]) + + async def notification( + ctx: ServerRequestContext[object, object], params: NotificationParams + ) -> HandlerResult: + with tracer.start_as_current_span("notification-handler"): + notified.set() + return None + + protocol: Final = Server[object]("trace-exclusion", on_call_tool=call_tool) + protocol.middleware[:] = gateway.server.middleware + protocol.add_notification_handler("notifications/exclusion-probe", NotificationParams, notification) + manager: Final = StreamableHTTPSessionManager(protocol, stateless=False) + app: Final = FastAPI() + app.mount("/mcp", manager.handle_request) + instrument_fastapi_app(app, tracer_provider=provider) + initial_url: Final = f"http://{'excluded' if initialized_excluded else 'included'}.example/mcp/" + current_url: Final = f"http://{'included' if initialized_excluded else 'excluded'}.example/mcp/" + headers: Final = {"Accept": "application/json, text/event-stream", "MCP-Protocol-Version": "2025-11-25"} + try: + async with manager.run(): + async with httpx2.AsyncClient(transport=httpx2.ASGITransport(app=app)) as client: + initialized: Final = await client.post(initial_url, headers=headers, json={ + "jsonrpc": "2.0", "id": 0, "method": "initialize", "params": { + "protocolVersion": "2025-11-25", "capabilities": {}, + "clientInfo": {"name": "trace-exclusion", "version": "1"}, + }, + }) + assert initialized.status_code == 200, initialized.text + assert any( + span.name == "initialize" for span in exporter.get_finished_spans() + ) is not initialized_excluded + exporter.clear() + session_headers: Final = {**headers, "MCP-Session-Id": initialized.headers["mcp-session-id"]} + ready: Final = await client.post(initial_url, headers=session_headers, + json={"jsonrpc": "2.0", "method": "notifications/initialized"}) + assert ready.status_code == 202, ready.text + response: Final = await client.post(current_url, headers=session_headers, json={ + "jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": { + "name": "probe", "arguments": {}, "_meta": { + "traceparent": "00-11111111111111111111111111111111-2222222222222222-01", + }, + }, + }) + assert response.status_code == 200 and '"text":"done"' in response.text, response.text + names: Final = {span.name for span in exporter.get_finished_spans()} + assert ("tools/call probe" in names) is initialized_excluded + assert ("tool-handler" in names) is initialized_excluded + exporter.clear() + notification_response: Final = await client.post(current_url, headers=session_headers, + json={"jsonrpc": "2.0", "method": "notifications/exclusion-probe"}) + assert notification_response.status_code == 202, notification_response.text + await notified.wait() + assert any( + span.name == "notification-handler" for span in exporter.get_finished_spans() + ) is initialized_excluded + exporter.clear() + modern_response: Final = await client.post(current_url, headers={ + "Accept": "application/json, text/event-stream", "MCP-Protocol-Version": "2026-07-28", + "MCP-Method": "tools/call", "MCP-Name": "probe", + }, json={ + "jsonrpc": "2.0", "id": 2, "method": "tools/call", "params": { + "name": "probe", "arguments": {}, "_meta": { + PROTOCOL_VERSION_META_KEY: "2026-07-28", CLIENT_CAPABILITIES_META_KEY: {}, + }, + }, + }) + assert modern_response.status_code == 200, modern_response.text + assert '"text":"done"' in modern_response.text, modern_response.text + modern_names: Final = {span.name for span in exporter.get_finished_spans()} + assert ("tools/call probe" in modern_names) is initialized_excluded + assert ("tool-handler" in modern_names) is initialized_excluded + finally: + provider.shutdown() + set_global_response_propagator(original_propagator) + is_otel_v2_enabled.cache_clear() + + @pytest.mark.asyncio async def test_stateful_mcp_tool_call_uses_current_requests_otel_destinations(_mcp_request_ctx) -> None: from types import SimpleNamespace From 8c9a76859db5e6a7a2d752574f00cebfea90fa1a Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 28 Oct 2026 18:25:28 +0000 Subject: [PATCH 062/128] fix(baseten): price the two decision models at their published input rates (#48849) * fix(baseten): price the two decision models at their published input rates * chore: retrigger CI --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 4 ++-- model_prices_and_context_window.json | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 04ad3546b5b..cc7aaad4264 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -76938,7 +76938,7 @@ "supports_xhigh_reasoning_effort": true }, "baseten/decision-models/baseten-elephant-9b-latest": { - "input_cost_per_token": 0.0, + "input_cost_per_token": 3.5e-08, "litellm_provider": "baseten", "max_input_tokens": 16384, "mode": "decisions", @@ -76947,7 +76947,7 @@ "supports_vision": true }, "baseten/decision-models/baseten-mammoth-27b-latest": { - "input_cost_per_token": 0.0, + "input_cost_per_token": 1e-07, "litellm_provider": "baseten", "max_input_tokens": 16384, "mode": "decisions", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 04ad3546b5b..cc7aaad4264 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -76938,7 +76938,7 @@ "supports_xhigh_reasoning_effort": true }, "baseten/decision-models/baseten-elephant-9b-latest": { - "input_cost_per_token": 0.0, + "input_cost_per_token": 3.5e-08, "litellm_provider": "baseten", "max_input_tokens": 16384, "mode": "decisions", @@ -76947,7 +76947,7 @@ "supports_vision": true }, "baseten/decision-models/baseten-mammoth-27b-latest": { - "input_cost_per_token": 0.0, + "input_cost_per_token": 1e-07, "litellm_provider": "baseten", "max_input_tokens": 16384, "mode": "decisions", From c40bdc475f9adee7eae99a6394f0f6091ae437e9 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 28 Oct 2026 14:28:55 -0400 Subject: [PATCH 063/128] fix(router)!: resolve trusted JEV connection environment references (#48717) --- .../auto_router_endpoints.py | 12 +++ .../model_management_endpoints.py | 10 +++ .../auto_router_permissions.py | 35 +++++++- .../complexity_router/complexity_router.py | 33 ++++++-- .../test_auto_router_endpoints.py | 75 ++++++++++++++++-- .../test_model_management_endpoints.py | 70 ++++++++++++++++ .../complexity_router/test_jev_classifier.py | 79 +++++++++++++++++++ 7 files changed, 301 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 9c46e711cf6..7a843aecc0c 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -49,6 +49,8 @@ from litellm.proxy.management_endpoints.common_daily_activity import daily_activ from litellm.proxy.management_helpers.auto_router_permissions import ( authorize_member_auto_router_dependencies, authorize_member_auto_router_team, + auto_router_environment_reference_error, + raise_auto_router_environment_reference_error, validate_member_auto_router_config, ) from litellm.repositories.autorouter_session_repository import AutoRouterSessionRepository @@ -402,6 +404,10 @@ async def validate_complexity_router_config( """ member_team: Final = await _authorize_router_dry_run(user_api_key_dict=user_api_key_dict, team_id=data.team_id) + raise_auto_router_environment_reference_error( + auto_router_environment_reference_error(data.complexity_router_config, user_api_key_dict) + ) + from litellm.router_utils.auto_router_model_naming import ( validate_complexity_router_config_write, ) @@ -560,6 +566,12 @@ async def preview_auto_router_routing( detail={"error": CommonProxyErrors.no_llm_router.value}, ) resolved: Final = await _resolve_saved_routing_test(data, user_api_key_dict, llm_router) + if data.saved_model_id is None: + raise_auto_router_environment_reference_error( + auto_router_environment_reference_error( + data.complexity_router_config.model_dump(exclude_none=True), user_api_key_dict + ) + ) actor: Final = ( await _authorize_member_dry_run_config( config=resolved.complexity_router_config.model_dump(exclude_none=True), diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 1130fb73e9a..45006b14df6 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -102,6 +102,8 @@ from litellm.proxy.management_helpers.auto_router_permissions import ( authorize_member_auto_router_dependencies, authorize_member_auto_router_team, authorize_member_auto_router_write, + auto_router_environment_reference_error, + raise_auto_router_environment_reference_error, ) from litellm.proxy.management_helpers.model_allowlist_rename_sync import sync_model_allowlists_for_renamed_model from litellm.proxy.management_helpers.model_quota_policy import ( @@ -2627,6 +2629,14 @@ class ModelManagementAuthChecks: ModelManagementAuthChecks._reject_team_inline_api_key_secret_reference( model_params=model_params, incoming_params=incoming_params, incoming_model_params=incoming_model_params ) + raise_auto_router_environment_reference_error( + auto_router_environment_reference_error( + config=_OPTIONAL_JSON_OBJECT.validate_python( + incoming_params.get("complexity_router_config") if incoming_params is not None else None + ), + user_api_key_dict=user_api_key_dict, + ) + ) # Federation fields choose which server-side secret is read and where the org-scoped token # it buys is sent, so only a proxy admin may point a federated deployment somewhere else. # Evaluated on the RESULTING deployment: a patch attaching a credential by name inherits diff --git a/litellm/proxy/management_helpers/auto_router_permissions.py b/litellm/proxy/management_helpers/auto_router_permissions.py index a4428eb3a85..995446db7a6 100644 --- a/litellm/proxy/management_helpers/auto_router_permissions.py +++ b/litellm/proxy/management_helpers/auto_router_permissions.py @@ -5,8 +5,8 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal from fastapi import HTTPException -from pydantic import ConfigDict, Field, ValidationError -from typing_extensions import ReadOnly, TypedDict +from pydantic import ConfigDict, Field, TypeAdapter, ValidationError +from typing_extensions import ReadOnly, TypedDict, assert_never from litellm.models.organization import LiteLLM_OrganizationTable from litellm.models.project import LiteLLM_ProjectTable @@ -36,6 +36,7 @@ from litellm.repositories.table_repositories import TeamMembershipRepository from litellm.router import Router from litellm.router_strategy.complexity_router.config import ( ComplexityRouterConfigWrite, + normalize_classifier_config_aliases, resolve_complexity_router_config_write, ) from litellm.router_strategy.complexity_router.request_models import RequestComplexityRouterConfig @@ -47,6 +48,36 @@ if TYPE_CHECKING: from prisma import types as prisma_types +def auto_router_environment_reference_error( + config: Mapping[str, object] | None, user_api_key_dict: UserAPIKeyAuth +) -> Literal["environment_reference_denied"] | None: + if config is None or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: + return None + classifier: Final = normalize_classifier_config_aliases(config).get("opensource_classifier_config") + if not isinstance(classifier, Mapping): + return None + fields: Final = TypeAdapter(Mapping[str, object]).validate_python(classifier) + if any( + isinstance(value, str) and value.startswith("os.environ/") + for value in (fields.get("api_key"), fields.get("api_base")) + ): + return "environment_reference_denied" + return None + + +def raise_auto_router_environment_reference_error(error: Literal["environment_reference_denied"] | None) -> None: + match error: + case None: + return + case "environment_reference_denied": + raise HTTPException( + status_code=403, + detail="Only proxy admins can submit classifier connection environment references", + ) + case _: + assert_never(error) + + class _MemberRouterThinking(LiteLLMBaseModel): model_config = ConfigDict(extra="forbid", frozen=True) diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 2b0187fd667..43af0485a2d 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -23,6 +23,7 @@ import random import re import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from dataclasses import dataclass from functools import partial from itertools import accumulate, chain, islice, takewhile from threading import Lock @@ -1390,6 +1391,11 @@ def _session_affinity_cache_value(model: str, tier: ComplexityTier | str | None) return {"model": model, "tier": tier_value} +@dataclass(frozen=True, slots=True) +class _JevConnectionError: + message: str + + class ComplexityRouter(CustomLogger): """ Complexity router that classifies requests and routes to appropriate models. @@ -1405,7 +1411,7 @@ class ComplexityRouter(CustomLogger): - Question complexity (multiple questions) """ - def _build_jev_client(self, config: OpenSourceClassifierConfig) -> JevClassifierClient: + def _build_jev_client(self, config: OpenSourceClassifierConfig) -> JevClassifierClient | _JevConnectionError: if config.deployment_name is not None: return RouterJevClassifierClient(self.litellm_router_instance, config.deployment_name) if config.provider in ("laya", "bespoke"): @@ -1426,12 +1432,22 @@ class ComplexityRouter(CustomLogger): http_client=get_async_httpx_client(httpxSpecialProvider.PassThroughEndpoint), provider="databricks", ) - api_key: Final = config.api_key or get_secret_str("TYPESAFE_API_KEY") - if not api_key: - raise ValueError( + api_key: Final = ( + get_secret_str(config.api_key) + if config.api_key is not None and config.api_key.startswith("os.environ/") + else config.api_key or get_secret_str("TYPESAFE_API_KEY") + ) + if not api_key or not api_key.strip(): + return _JevConnectionError( "opensource_classifier_config.api_key or TYPESAFE_API_KEY is required for classifier_type 'oss_classifier'" ) - api_base: Final = config.api_base or get_secret_str("TYPESAFE_API_BASE") or "https://api.typesafe.ai" + api_base: Final = ( + get_secret_str(config.api_base) + if config.api_base is not None and config.api_base.startswith("os.environ/") + else config.api_base or get_secret_str("TYPESAFE_API_BASE") or "https://api.typesafe.ai" + ) + if not api_base or not api_base.strip(): + return _JevConnectionError("opensource_classifier_config.api_base must resolve to a non-empty value") return HttpJevClassifierClient( api_key=api_key, api_base=api_base, @@ -1474,13 +1490,18 @@ class ComplexityRouter(CustomLogger): self.config.default_model = default_model jev_config: Final = self.config.opensource_classifier_config - self._jev_client: JevClassifierClient | None = ( + jev_client_result: Final = ( jev_client if jev_client is not None else self._build_jev_client(jev_config) if self.config.classifier_type == "oss_classifier" and jev_config is not None else None ) + match jev_client_result: + case _JevConnectionError(message=message): + raise ValueError(message) + case _: + self._jev_client: JevClassifierClient | None = jev_client_result self._tier_affinity_config = hashlib.sha256( self.config.model_dump_json(include=MappingProxyType({"tiers": True, "tier_model_configs": True})).encode() diff --git a/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py index d5a8abaca9d..748b291d5ff 100644 --- a/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py @@ -2769,15 +2769,22 @@ async def test_jev_test_routing_authorizes_paid_evaluation_before_contacting_typ @pytest.mark.asyncio @pytest.mark.parametrize( - "case", ["allowed", "credential-free", "missing", "blocked", "key", "budget", "team", "not-router"] + "case", + ["allowed", "credential-free", "team-admin-environment", "missing", "blocked", "key", "budget", "team", "not-router"], ) async def test_saved_jev_probe_uses_authorized_server_configuration(monkeypatch: pytest.MonkeyPatch, case: str) -> None: router: Final = RecordingRouter("SIMPLE") stored_key: Final = "synthetic-server-jev-key" + monkeypatch.setenv("SAVED_JEV_KEY", stored_key) + monkeypatch.setenv("SAVED_JEV_BASE", "https://saved-jev.test") stored_config: Final = { "classifier_type": "jev", "tiers": TIERS, - "jev_classifier_config": {"api_key": stored_key, "api_base": "https://saved-jev.test"}, + "jev_classifier_config": ( + {"api_key": "os.environ/SAVED_JEV_KEY", "api_base": "os.environ/SAVED_JEV_BASE"} + if case == "team-admin-environment" + else {"api_key": stored_key, "api_base": "https://saved-jev.test"} + ), } router.add_deployment( Deployment.model_validate( @@ -2790,7 +2797,10 @@ async def test_saved_jev_probe_uses_authorized_server_configuration(monkeypatch: "model_info": { "id": "saved-jev-id", "blocked": case == "blocked", - "team_id": "owner-team" if case == "team" else None, + "team_id": ( + "owner-team" if case == "team" + else "admin-preview-team" if case == "team-admin-environment" else None + ), }, } ) @@ -2799,6 +2809,8 @@ async def test_saved_jev_probe_uses_authorized_server_configuration(monkeypatch: actor: Final = ( _configure_member_preview(monkeypatch) if case == "team" + else _configure_team_admin_preview(monkeypatch) + if case == "team-admin-environment" else UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-probe", @@ -2812,12 +2824,17 @@ async def test_saved_jev_probe_uses_authorized_server_configuration(monkeypatch: { "prompt": "what is 2+2", "saved_model_id": "missing-id" if case == "missing" else "saved-jev-id", - "team_id": "member-preview-team" if case == "team" else None, + "team_id": ( + "member-preview-team" if case == "team" + else "admin-preview-team" if case == "team-admin-environment" else None + ), }, classifier_type="jev", jev_classifier_config=( {"model": "jev-latest", "timeout_ms": 3000} if case == "credential-free" + else {"api_key": "os.environ/IGNORED_BROWSER_KEY", "api_base": "https://browser-override.test"} + if case == "team-admin-environment" else {"api_key": "masked-key", "api_base": "https://browser-override.test"} ), ) @@ -2856,7 +2873,7 @@ async def test_saved_jev_probe_uses_authorized_server_configuration(monkeypatch: assert result.routed_model == "cheap-model" assert evaluation.calls.last.request.headers["authorization"] == f"Bearer {stored_key}" assert stored_key not in result.model_dump_json() - assert evaluation.call_count == (1 if case in ("allowed", "credential-free") else 0) + assert evaluation.call_count == (1 if case in ("allowed", "credential-free", "team-admin-environment") else 0) assert router.recorded_calls == [] await handler.client.aclose() @@ -3378,6 +3395,54 @@ async def test_validate_config_gates_like_the_write_it_rehearses(monkeypatch: py assert not_their_team.value.status_code == 403 +def _configure_team_admin_preview(monkeypatch: pytest.MonkeyPatch) -> UserAPIKeyAuth: + from litellm.proxy._types import LiteLLM_TeamTable + + team: Final = LiteLLM_TeamTable( + team_id="admin-preview-team", + models=["saved-jev", "typesafe/jev-latest", *[TIERS[name][0] for name in TIERS]], + members_with_roles=[{"role": "admin", "user_id": "preview-admin"}], + ) + database: Final = MagicMock() + database.db.litellm_teamtable.find_unique = AsyncMock(return_value=team) + monkeypatch.setattr(proxy_server, "prisma_client", database) + monkeypatch.setattr(proxy_server, "premium_user", True) + return UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="preview-admin", api_key="sk-preview") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("legacy", [False, True]) +@pytest.mark.parametrize("field", ["api_key", "api_base"]) +async def test_team_admin_cannot_preview_or_validate_environment_references( + monkeypatch: pytest.MonkeyPatch, legacy: bool, field: str +) -> None: + from litellm.types.management_endpoints.auto_router_endpoints import ComplexityRouterConfigValidationRequest + + actor: Final = _configure_team_admin_preview(monkeypatch) + monkeypatch.setattr(proxy_server, "llm_router", _router()) + config: Final = { + "tiers": TIERS, + "classifier_type": "jev" if legacy else "oss_classifier", + "jev_classifier_config" if legacy else "opensource_classifier_config": { + "api_key": "caller-key", "api_base": "https://denied-jev.test", field: "os.environ/PRIVATE_JEV_VALUE", + }, + } + validation: Final = ComplexityRouterConfigValidationRequest( + team_id="admin-preview-team", complexity_router_config=config + ) + with pytest.raises(HTTPException) as denied_validation: + await auto_router_endpoints.validate_complexity_router_config(validation, actor) + assert denied_validation.value.status_code == 403 + with respx.mock(assert_all_called=False) as http, pytest.raises(HTTPException) as denied_preview: + await preview_auto_router_routing( + _request_from({"prompt": "what is 2+2", "team_id": "admin-preview-team"}, **config), + actor, + ROUTING_HTTP_REQUEST, + ) + assert denied_preview.value.status_code == 403 + assert len(http.calls) == 0 + + def _configure_member_preview(monkeypatch: pytest.MonkeyPatch, *, allowed: bool = True) -> UserAPIKeyAuth: from litellm.proxy import proxy_server from litellm.proxy._types import UI_TEAM_ID, LiteLLM_TeamTable diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index c8f601de280..fc6749aed0f 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -8475,6 +8475,76 @@ class TestTeamMemberAutoRouterWrites: "jev_classifier_config" if legacy else "opensource_classifier_config": classifier, } + @pytest.mark.asyncio + @pytest.mark.parametrize("endpoint", ["create", "patch", "legacy"]) + @pytest.mark.parametrize("legacy", [False, True]) + @pytest.mark.parametrize("field", ["api_key", "api_base"]) + async def test_team_admin_classifier_writes_reject_environment_references( + self, endpoint: str, legacy: bool, field: str + ) -> None: + from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model + + stored_config: Final = self._classifier_config( + {"api_key": "os.environ/PRIVATE_JEV_VALUE", "api_base": "https://trusted-jev.test"}, legacy + ) + row: Final = self._row().model_copy(update={"litellm_params": { + "model": "auto_router/complexity_router", "complexity_router_config": stored_config, + }}) + team: Final = self._team().model_copy(update={"members_with_roles": [Member(user_id="owner", role="admin")]}) + database: Final = self._database(team, row) + actor: Final = UserAPIKeyAuth(user_id="owner", user_role=LitellmUserRoles.INTERNAL_USER) + config: Final = self._classifier_config( + {"api_key": "caller-key", "api_base": "https://untrusted-jev.test", field: "os.environ/PRIVATE_JEV_VALUE"}, + legacy, + ) + request: Final = updateDeployment( + litellm_params=updateLiteLLMParams(complexity_router_config=config), + model_info=ModelInfo(id=row.model_id, team_id="member-team"), + ) + operation: Final = ( + add_new_model(Deployment( + model_name="new-jev-router", + litellm_params=LiteLLM_Params(model="auto_router/complexity_router", complexity_router_config=config), + model_info=ModelInfo(id=row.model_id, team_id="member-team"), + ), actor) + if endpoint == "create" + else patch_model(row.model_id, request, actor) if endpoint == "patch" else update_model(request, actor) + ) + with self._environment(database, row), pytest.raises((HTTPException, ProxyException)) as denied: + await operation + assert (denied.value.status_code if isinstance(denied.value, HTTPException) else int(denied.value.code)) == 403 + database.db.litellm_proxymodeltable.create.assert_not_awaited() + database.db.litellm_proxymodeltable.update.assert_not_awaited() + + @pytest.mark.asyncio + @pytest.mark.parametrize("endpoint", ["patch", "legacy"]) + @pytest.mark.parametrize("change_config", [False, True]) + async def test_team_admin_unrelated_updates_preserve_stored_environment_references( + self, endpoint: str, change_config: bool + ) -> None: + connection: Final = {"api_key": "os.environ/SAVED_JEV_KEY", "api_base": "os.environ/SAVED_JEV_BASE"} + stored_config: Final = self._classifier_config(connection, True) + row: Final = self._row().model_copy(update={"litellm_params": { + "model": "auto_router/complexity_router", "complexity_router_config": stored_config, + }}) + team: Final = self._team().model_copy(update={"members_with_roles": [Member(user_id="owner", role="admin")]}) + database: Final = self._database(team, row) + actor: Final = UserAPIKeyAuth(user_id="owner", user_role=LitellmUserRoles.INTERNAL_USER) + params: Final = ( + {"complexity_router_config": self._classifier_config({"timeout_ms": 8100}, False)} + if change_config else {"max_tokens": 100} + ) + request: Final = updateDeployment( + litellm_params=updateLiteLLMParams.model_validate(params), model_info=ModelInfo(id=row.model_id), + ) + with self._environment(database, row): + await (patch_model(row.model_id, request, actor) if endpoint == "patch" else update_model(request, actor)) + written: Final = database.db.litellm_proxymodeltable.update.await_args.kwargs["data"] + saved: Final = json.loads(written["litellm_params"])["complexity_router_config"] + assert saved == ( + self._classifier_config({**connection, "timeout_ms": 8100}, False) if change_config else stored_config + ) + @pytest.mark.asyncio @pytest.mark.parametrize( "stored_params,removed,denied_param", diff --git a/tests/unit/router_strategy/complexity_router/test_jev_classifier.py b/tests/unit/router_strategy/complexity_router/test_jev_classifier.py index 10a8f7929d9..77d465526a4 100644 --- a/tests/unit/router_strategy/complexity_router/test_jev_classifier.py +++ b/tests/unit/router_strategy/complexity_router/test_jev_classifier.py @@ -732,6 +732,85 @@ def test_jev_api_base_without_its_own_key_is_rejected_so_the_environment_key_sta assert JevClassifierConfig(api_key="sk-own").api_base is None +@pytest.mark.asyncio +@pytest.mark.parametrize("legacy", [False, True]) +@pytest.mark.parametrize("connection", ["literal", "key-reference", "base-reference", "references", "defaults"]) +async def test_jev_connection_resolves_only_explicit_environment_references( + monkeypatch: pytest.MonkeyPatch, legacy: bool, connection: str +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("JEV_CONNECTION_KEY", "resolved-key") + monkeypatch.setenv("JEV_CONNECTION_BASE", "https://jev.test") + monkeypatch.setenv("TYPESAFE_API_KEY", "default-key") + monkeypatch.setenv("TYPESAFE_API_BASE", "https://default-jev.test") + monkeypatch.setenv("literal-key", "must-not-replace-literal") + fields: Final = ( + {} + if connection == "defaults" + else { + "api_key": "os.environ/JEV_CONNECTION_KEY" + if connection in ("key-reference", "references") + else "literal-key", + "api_base": "os.environ/JEV_CONNECTION_BASE" + if connection in ("base-reference", "references") + else "https://jev.test", + } + ) + config: Final = { + "classifier_type": "jev" if legacy else "oss_classifier", + "jev_classifier_config" if legacy else "opensource_classifier_config": {"model": "openjev-latest", **fields}, + "tiers": {"SIMPLE": "cheap", "COMPLEX": "strong"}, + } + classifier: Final = ComplexityRouter("env-router", litellm.Router(), config, derive_savings_baseline=False) + with respx.mock(assert_all_called=True) as upstream: + route: Final = upstream.post( + "https://default-jev.test/v1/systemone" if connection == "defaults" else "https://jev.test/v1/systemone" + ).respond(json={"answers": {"tier": _answer("SIMPLE").model_dump()}}) + outcome: Final = await classifier.aclassify("What is two plus two?") + expected_key: Final = ( + "default-key" + if connection == "defaults" + else "resolved-key" if connection in ("key-reference", "references") else "literal-key" + ) + assert outcome.cause == "jev_classifier" + assert route.calls[0].request.headers["authorization"] == f"Bearer {expected_key}" + assert classifier.config.opensource_classifier_config is not None + assert classifier.config.opensource_classifier_config.api_key == fields.get("api_key") + assert classifier.config.opensource_classifier_config.api_base == fields.get("api_base") + + +@pytest.mark.parametrize("field", ["api_key", "api_base"]) +@pytest.mark.parametrize("value", [None, "", " "]) +def test_jev_missing_explicit_environment_reference_never_uses_provider_defaults( + monkeypatch: pytest.MonkeyPatch, field: str, value: str | None +) -> None: + monkeypatch.setenv("TYPESAFE_API_KEY", "never-send-default-key") + monkeypatch.setenv("TYPESAFE_API_BASE", "https://never-use-default.test") + if value is None: + monkeypatch.delenv("JEV_MISSING_CONNECTION", raising=False) + else: + monkeypatch.setenv("JEV_MISSING_CONNECTION", value) + with ( + respx.mock(assert_all_called=False) as upstream, + pytest.raises(ValueError, match=rf"opensource_classifier_config\.{field}"), + ): + ComplexityRouter( + "env-router", + litellm.Router(), + { + "classifier_type": "jev", + "jev_classifier_config": { + "api_key": "literal-key", + "api_base": "https://jev.test", + field: "os.environ/JEV_MISSING_CONNECTION", + }, + "tiers": {"SIMPLE": "cheap"}, + }, + derive_savings_baseline=False, + ) + assert len(upstream.calls) == 0 + + def test_databricks_api_base_without_its_own_key_is_rejected_so_the_workspace_token_stays_home() -> None: with pytest.raises(ValueError, match="DATABRICKS_API_KEY or DATABRICKS_TOKEN is only sent to DATABRICKS_API_BASE"): JevClassifierConfig.model_validate( From 20427f455536999e667134f2a2ab0fbd7bd479f8 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 28 Oct 2026 11:38:32 -0700 Subject: [PATCH 064/128] chore(cost-map): add fireworks spot serverless router rows from the prices api (#48858) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 146 +++++++++++------- model_prices_and_context_window.json | 146 +++++++++++------- 2 files changed, 172 insertions(+), 120 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index cc7aaad4264..50c23efca48 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -27480,18 +27480,6 @@ "supports_tool_choice": true, "supports_vision": false }, - "fireworks_ai/accounts/fireworks/routers/spot-deepseek-v4-flash": { - "cache_read_input_token_cost": 2.1e-08, - "input_cost_per_token": 1.05e-07, - "litellm_provider": "fireworks_ai", - "max_input_tokens": 1048576, - "mode": "chat", - "output_cost_per_token": 2.1e-07, - "source": "https://api.fireworks.ai/v1/serverless/models", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_vision": false - }, "fireworks_ai/accounts/fireworks/models/deepseek-v4-pro-0813": { "cache_read_input_token_cost": 4.4e-08, "cache_read_input_token_cost_priority": 5.5e-08, @@ -62352,18 +62340,6 @@ "supports_tool_choice": true, "supports_vision": true }, - "fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-spot": { - "cache_read_input_token_cost": 4.2e-09, - "input_cost_per_token": 2.1e-07, - "litellm_provider": "fireworks_ai", - "max_input_tokens": 1048576, - "mode": "chat", - "output_cost_per_token": 8.4e-07, - "source": "https://api.fireworks.ai/v1/serverless/models", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_vision": true - }, "fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-us": { "cache_read_input_token_cost": 9e-09, "cache_read_input_token_cost_priority": 1.125e-08, @@ -62424,18 +62400,6 @@ "supports_tool_choice": true, "supports_vision": true }, - "fireworks_ai/accounts/fireworks/routers/kimi-k3-spot": { - "cache_read_input_token_cost": 2.1e-07, - "input_cost_per_token": 2.1e-06, - "litellm_provider": "fireworks_ai", - "max_input_tokens": 1048576, - "mode": "chat", - "output_cost_per_token": 1.05e-05, - "source": "https://api.fireworks.ai/v1/serverless/models", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_vision": true - }, "fireworks_ai/deepseek-v4-flash-0731": { "cache_read_input_token_cost": 7e-09, "cache_read_input_token_cost_priority": 8.75e-09, @@ -65635,18 +65599,6 @@ "supports_tool_choice": true, "supports_vision": false }, - "fireworks_ai/accounts/fireworks/routers/glm-5p3-spot": { - "cache_read_input_token_cost": 1.82e-07, - "input_cost_per_token": 9.8e-07, - "litellm_provider": "fireworks_ai", - "max_input_tokens": 1048576, - "mode": "chat", - "output_cost_per_token": 3.08e-06, - "source": "https://api.fireworks.ai/v1/serverless/models", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_vision": false - }, "fireworks_ai/accounts/fireworks/routers/glm-5p3-us": { "cache_read_input_token_cost": 3.9e-07, "cache_read_input_token_cost_priority": 4.875e-07, @@ -75317,18 +75269,6 @@ "supports_tool_choice": true, "supports_vision": true }, - "fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-1017-spot": { - "cache_read_input_token_cost": 4.2e-09, - "input_cost_per_token": 2.1e-07, - "litellm_provider": "fireworks_ai", - "max_input_tokens": 1048576, - "mode": "chat", - "output_cost_per_token": 8.4e-07, - "source": "https://api.fireworks.ai/v1/serverless/models", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_vision": true - }, "github_copilot/claude-sonnet-4.6": { "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, @@ -77129,5 +77069,91 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "use_openai_responses_path": true + }, + "fireworks_ai/accounts/fireworks/routers/spot-deepseek-v4-flash": { + "cache_read_input_token_cost": 2.1e-08, + "deprecation_date": "2026-08-14", + "input_cost_per_token": 1.05e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 2.1e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/accounts/fireworks/routers/kimi-k3-spot": { + "cache_read_input_token_cost": 2.1e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.05e-05, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ], + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/accounts/fireworks/routers/glm-5p3-spot": { + "cache_read_input_token_cost": 1.82e-07, + "input_cost_per_token": 9.8e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3.08e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-spot": { + "cache_read_input_token_cost": 4.2e-09, + "input_cost_per_token": 2.1e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 8.4e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-1017-spot": { + "cache_read_input_token_cost": 4.2e-09, + "input_cost_per_token": 2.1e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 8.4e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index cc7aaad4264..50c23efca48 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -27480,18 +27480,6 @@ "supports_tool_choice": true, "supports_vision": false }, - "fireworks_ai/accounts/fireworks/routers/spot-deepseek-v4-flash": { - "cache_read_input_token_cost": 2.1e-08, - "input_cost_per_token": 1.05e-07, - "litellm_provider": "fireworks_ai", - "max_input_tokens": 1048576, - "mode": "chat", - "output_cost_per_token": 2.1e-07, - "source": "https://api.fireworks.ai/v1/serverless/models", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_vision": false - }, "fireworks_ai/accounts/fireworks/models/deepseek-v4-pro-0813": { "cache_read_input_token_cost": 4.4e-08, "cache_read_input_token_cost_priority": 5.5e-08, @@ -62352,18 +62340,6 @@ "supports_tool_choice": true, "supports_vision": true }, - "fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-spot": { - "cache_read_input_token_cost": 4.2e-09, - "input_cost_per_token": 2.1e-07, - "litellm_provider": "fireworks_ai", - "max_input_tokens": 1048576, - "mode": "chat", - "output_cost_per_token": 8.4e-07, - "source": "https://api.fireworks.ai/v1/serverless/models", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_vision": true - }, "fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-us": { "cache_read_input_token_cost": 9e-09, "cache_read_input_token_cost_priority": 1.125e-08, @@ -62424,18 +62400,6 @@ "supports_tool_choice": true, "supports_vision": true }, - "fireworks_ai/accounts/fireworks/routers/kimi-k3-spot": { - "cache_read_input_token_cost": 2.1e-07, - "input_cost_per_token": 2.1e-06, - "litellm_provider": "fireworks_ai", - "max_input_tokens": 1048576, - "mode": "chat", - "output_cost_per_token": 1.05e-05, - "source": "https://api.fireworks.ai/v1/serverless/models", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_vision": true - }, "fireworks_ai/deepseek-v4-flash-0731": { "cache_read_input_token_cost": 7e-09, "cache_read_input_token_cost_priority": 8.75e-09, @@ -65635,18 +65599,6 @@ "supports_tool_choice": true, "supports_vision": false }, - "fireworks_ai/accounts/fireworks/routers/glm-5p3-spot": { - "cache_read_input_token_cost": 1.82e-07, - "input_cost_per_token": 9.8e-07, - "litellm_provider": "fireworks_ai", - "max_input_tokens": 1048576, - "mode": "chat", - "output_cost_per_token": 3.08e-06, - "source": "https://api.fireworks.ai/v1/serverless/models", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_vision": false - }, "fireworks_ai/accounts/fireworks/routers/glm-5p3-us": { "cache_read_input_token_cost": 3.9e-07, "cache_read_input_token_cost_priority": 4.875e-07, @@ -75317,18 +75269,6 @@ "supports_tool_choice": true, "supports_vision": true }, - "fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-1017-spot": { - "cache_read_input_token_cost": 4.2e-09, - "input_cost_per_token": 2.1e-07, - "litellm_provider": "fireworks_ai", - "max_input_tokens": 1048576, - "mode": "chat", - "output_cost_per_token": 8.4e-07, - "source": "https://api.fireworks.ai/v1/serverless/models", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_vision": true - }, "github_copilot/claude-sonnet-4.6": { "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, @@ -77129,5 +77069,91 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "use_openai_responses_path": true + }, + "fireworks_ai/accounts/fireworks/routers/spot-deepseek-v4-flash": { + "cache_read_input_token_cost": 2.1e-08, + "deprecation_date": "2026-08-14", + "input_cost_per_token": 1.05e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 2.1e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/accounts/fireworks/routers/kimi-k3-spot": { + "cache_read_input_token_cost": 2.1e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.05e-05, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ], + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/accounts/fireworks/routers/glm-5p3-spot": { + "cache_read_input_token_cost": 1.82e-07, + "input_cost_per_token": 9.8e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3.08e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-spot": { + "cache_read_input_token_cost": 4.2e-09, + "input_cost_per_token": 2.1e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 8.4e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-1017-spot": { + "cache_read_input_token_cost": 4.2e-09, + "input_cost_per_token": 2.1e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 8.4e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true } } From 7abff3dd63893f372818bcc91d5e1e48fae6afa5 Mon Sep 17 00:00:00 2001 From: rj-cairn <349020452+rj-cairn@users.noreply.github.com> Date: Wed, 28 Oct 2026 11:41:33 -0700 Subject: [PATCH 065/128] fix(mistral): replay assistant reasoning as a Mistral thinking chunk (#48778) * fix(mistral): replay assistant reasoning as a Mistral thinking chunk When a client sends back an earlier Mistral assistant turn, its reasoning (reasoning_content, thinking_blocks, or Mistral's own thinking chunk in content) now reaches Mistral as a ThinkChunk inside content. The top-level LiteLLM fields stay stripped because Mistral rejects them with 422. * fix(mistral): replay reasoning only to Mistral models that accept it Mistral answers 400 to a thinking chunk on models without reasoning support, so a mid-chat model switch, a fallback, or Vertex AI Mistral broke once the replay was added. Gate the replay on the same check reasoning_effort uses and keep dropping the reasoning for the other models, as before. * fix(mistral): drop native thinking chunks for models without reasoning and mark the devstral, code agent, small and leanstral aliases as reasoning * fix(mistral): keep the leanstral rows as on main, since their reasoning support cannot be checked live * fix(mistral): send only the text items of a native thinking chunk, so a reference or image item no longer gets Mistral's 422 * fix(mistral): keep closed false only on the last chunk of a final prefix turn, so an earlier turn no longer gets Mistral's 400 * fix(mistral): drop closed false when merged text follows the thinking chunk, so a prefix turn with text first no longer gets Mistral's 400 --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/llms/mistral/chat/transformation.py | 207 ++++++- ...odel_prices_and_context_window_backup.json | 25 + litellm/types/llms/mistral.py | 3 +- model_prices_and_context_window.json | 25 + .../test_mistral_chat_transformation.py | 517 +++++++++++++++++- 5 files changed, 746 insertions(+), 31 deletions(-) diff --git a/litellm/llms/mistral/chat/transformation.py b/litellm/llms/mistral/chat/transformation.py index 2c0acacff23..2fb9402ed03 100644 --- a/litellm/llms/mistral/chat/transformation.py +++ b/litellm/llms/mistral/chat/transformation.py @@ -6,16 +6,17 @@ Why separate file? Make it easy to see how transformation works Docs - https://docs.mistral.ai/api/ """ -from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping +from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence from itertools import accumulate, count from typing import TYPE_CHECKING, Final, Literal, cast, get_type_hints, overload import httpx +from typing_extensions import TypeIs # noqa: TID251 # TypeIs reaches typing only on 3.13 from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.prompt_templates.common_utils import ( - handle_messages_with_content_list_to_str_conversion, + convert_content_list_to_str, strip_litellm_internal_message_fields, strip_none_values_from_message, ) @@ -29,8 +30,8 @@ from litellm.router_utils.reasoning_effort_capability import ( nearest_declared_reasoning_effort, ) from litellm.secret_managers.main import get_secret_str -from litellm.types.llms.mistral import MistralThinkingBlock, MistralToolCallMessage -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.mistral import MistralTextBlock, MistralThinkingBlock, MistralToolCallMessage +from litellm.types.llms.openai import AllMessageValues, ChatCompletionAssistantMessage, ChatCompletionTextObject from litellm.types.utils import ModelResponse, ModelResponseStream from litellm.utils import convert_to_model_response_object, supports_reasoning @@ -58,6 +59,10 @@ def _accepted_reasoning_effort(model: str, requested: str, custom_llm_provider: _NO_TOOL_CALL_TURN: Final[tuple[int, frozenset[str]]] = (-1, frozenset()) +def _is_json_array(value: object) -> TypeIs[Sequence[object]]: # guard-ok: a decoded JSON array is a list + return isinstance(value, list) + + class MistralConfig(OpenAIGPTConfig): """ Reference: https://docs.mistral.ai/api/ @@ -118,11 +123,14 @@ class MistralConfig(OpenAIGPTConfig): def custom_llm_provider(self) -> str: return "mistral" - def get_supported_openai_params(self, model: str) -> list[str]: - is_magistral: Final = "magistral" in model.lower() - accepts_reasoning_effort: Final = is_magistral or supports_reasoning( + def _accepts_reasoning(self, model: str) -> bool: + return "magistral" in model.lower() or supports_reasoning( model=model, custom_llm_provider=self.custom_llm_provider ) + + def get_supported_openai_params(self, model: str) -> list[str]: + is_magistral: Final = "magistral" in model.lower() + accepts_reasoning_effort: Final = self._accepts_reasoning(model) return [ "stream", "temperature", @@ -260,25 +268,32 @@ class MistralConfig(OpenAIGPTConfig): The above statement is not valid now. Need to plan to remove all the #1,2,3 Mistral API supports content as a list. """ - messages = [self._strip_output_only_fields(m) for m in messages] + keeps_reasoning: Final = self._accepts_reasoning(model) + final_position: Final = len(messages) - 1 + replayed_messages: Final = [ + MistralConfig._with_replayed_reasoning(m, keeps_reasoning, position == final_position) + for position, m in enumerate(messages) + ] ## 1. If 'image_url' or 'file' in content, then transform with base class and mistral-specific handling - for m in messages: + for m in replayed_messages: _content_block = m.get("content") if _content_block and isinstance(_content_block, list): if any(c.get("type") in ["image_url", "file"] for c in _content_block): if is_async: - return self._transform_messages_async(messages, model) + return self._transform_messages_async(replayed_messages, model) else: - messages = self._transform_messages_sync(messages, model) - return messages + return self._transform_messages_sync(replayed_messages, model) ## 2. If content is list, then convert to string - messages = handle_messages_with_content_list_to_str_conversion(messages) + flattened_messages: Final = tuple( + MistralConfig._with_flattened_content(m, keeps_reasoning) for m in replayed_messages + ) ## 3. Handle name in message prepared_messages: Final = tuple( - MistralConfig._handle_tool_call_message(MistralConfig._handle_name_in_message(m)) for m in messages + MistralConfig._handle_tool_call_message(MistralConfig._handle_name_in_message(m)) + for m in flattened_messages ) empty_flags: Final = tuple(map(MistralConfig._is_empty_assistant_message, prepared_messages)) tool_call_turns: Final = accumulate( @@ -442,17 +457,163 @@ class MistralConfig(OpenAIGPTConfig): return cleaned_tools - @classmethod - def _strip_output_only_fields(cls, message: AllMessageValues) -> AllMessageValues: - """ - LiteLLM attaches output-only fields (``reasoning_content``, ``thinking_blocks``, - ``provider_specific_fields``) to assistant responses. Mistral's input schema forbids - unknown fields, so replaying them verbatim in a follow-up turn triggers a - 422 ``extra_forbidden``. Drop them before the request is sent. - """ + @staticmethod + def _is_mistral_think_chunk(part: object) -> bool: + match part: + case {"type": "thinking", "thinking": list()}: + return True + case _: + return False + + @staticmethod + def _text_item_text(item: object) -> str | None: + match item: + case {"type": "text", "text": str(text)}: + return text + case _: + return None + + @staticmethod + def _sendable_think_chunk( + chunk: object, items: Sequence[object], may_stay_open: bool + ) -> MistralThinkingBlock | None: + text_items: Final = [ + MistralTextBlock(type="text", text=text) + for text in map(MistralConfig._text_item_text, items) + if text is not None + ] + if not text_items: + return None + match chunk: + case {"closed": False} if not may_stay_open: + return MistralThinkingBlock(type="thinking", thinking=text_items) + case {"closed": bool(closed)}: + return MistralThinkingBlock(type="thinking", thinking=text_items, closed=closed) + case _: + return MistralThinkingBlock(type="thinking", thinking=text_items) + + @staticmethod + def _sendable_part(part: object, may_stay_open: bool) -> object | None: + match part: + case {"type": "thinking", "thinking": object() as items} if _is_json_array(items): + return MistralConfig._sendable_think_chunk(part, items, may_stay_open) + case _: + return part + + @staticmethod + def _with_sendable_think_chunks( + message: AllMessageValues, content_parts: tuple[object, ...], may_end_open: bool + ) -> AllMessageValues: + last_position: Final = len(content_parts) - 1 + sendable_parts: Final = [ + sendable + for position, part in enumerate(content_parts) + if (sendable := MistralConfig._sendable_part(part, may_end_open and position == last_position)) is not None + ] + return cast( # cast-ok: Mistral's ThinkChunk is outside the OpenAI assistant content union + AllMessageValues, {**message, "content": sendable_parts or None} + ) + + @staticmethod + def _thinking_block_text(block: object) -> str: + match block: + case {"type": "thinking", "thinking": str(thinking)}: + return thinking + case _: + return "" + + @staticmethod + def _replayed_reasoning_text(message: ChatCompletionAssistantMessage) -> str | None: + reasoning_content: Final = message.get("reasoning_content") + if isinstance(reasoning_content, str) and reasoning_content.strip(): + return reasoning_content + thinking_blocks: Final = message.get("thinking_blocks") + thinking_text: Final = ( + "".join(map(MistralConfig._thinking_block_text, thinking_blocks)) + if isinstance(thinking_blocks, list) + else "" + ) + return thinking_text if thinking_text.strip() else None + + @staticmethod + def _assistant_content_parts(message: ChatCompletionAssistantMessage) -> tuple[object, ...] | None: + match message.get("content"): + case None | "": + return () + case str(text): + return (ChatCompletionTextObject(type="text", text=text),) + case list(parts) | tuple(parts): + return tuple(parts) + case _: + return None + + @staticmethod + def _is_prefix_turn(message: object) -> bool: + match message: + case {"prefix": True}: + return True + case _: + return False + + @staticmethod + def _with_replayed_reasoning( + message: AllMessageValues, keeps_reasoning: bool, is_final_turn: bool + ) -> AllMessageValues: if message["role"] != "assistant": return message - return strip_litellm_internal_message_fields(message) + stripped_message: Final = strip_litellm_internal_message_fields(message) + content_parts: Final = MistralConfig._assistant_content_parts(message) + if not keeps_reasoning: + return MistralConfig._without_think_chunks(stripped_message, content_parts) + if content_parts is not None and any(map(MistralConfig._is_mistral_think_chunk, content_parts)): + return MistralConfig._with_sendable_think_chunks( + stripped_message, content_parts, is_final_turn and MistralConfig._is_prefix_turn(message) + ) + reasoning_text: Final = MistralConfig._replayed_reasoning_text(message) + has_turn_output: Final = bool(content_parts) or bool(message.get("tool_calls")) + if reasoning_text is None or content_parts is None or not has_turn_output: + return stripped_message + think_chunk: Final = MistralThinkingBlock( + type="thinking", thinking=[MistralTextBlock(type="text", text=reasoning_text)], closed=True + ) + return cast( # cast-ok: Mistral's ThinkChunk is outside the OpenAI assistant content union + AllMessageValues, {**stripped_message, "content": [think_chunk, *content_parts]} + ) + + @staticmethod + def _without_think_chunks(message: AllMessageValues, content_parts: tuple[object, ...] | None) -> AllMessageValues: + if content_parts is None or not any(map(MistralConfig._is_mistral_think_chunk, content_parts)): + return message + kept_parts: Final = [part for part in content_parts if not MistralConfig._is_mistral_think_chunk(part)] + return cast( # cast-ok: an assistant turn takes a content list or no content + AllMessageValues, {**message, "content": kept_parts or None} + ) + + @staticmethod + def _with_flattened_content(message: AllMessageValues, keeps_reasoning: bool) -> AllMessageValues: + flattened_text: Final = convert_content_list_to_str(message) + match message: + case {"role": "assistant", "content": list(content)} if keeps_reasoning and any( + map(MistralConfig._is_mistral_think_chunk, content) + ): + text_parts: Final = ( + (ChatCompletionTextObject(type="text", text=flattened_text),) if flattened_text else () + ) + think_chunks: Final = tuple(filter(MistralConfig._is_mistral_think_chunk, content)) + sent_chunks: Final = tuple( + sendable + for chunk in think_chunks + if (sendable := MistralConfig._sendable_part(chunk, not text_parts)) is not None + ) + return cast( # cast-ok: Mistral's ThinkChunk is outside the OpenAI assistant content union + AllMessageValues, {**message, "content": [*sent_chunks, *text_parts]} + ) + case _ if flattened_text: + return cast( # cast-ok: every message role takes str content + AllMessageValues, {**message, "content": flattened_text} + ) + case _: + return message @classmethod def _handle_name_in_message(cls, message: AllMessageValues) -> AllMessageValues: diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 50c23efca48..1b65c87c315 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41495,9 +41495,14 @@ "max_output_tokens": 256000, "mode": "chat", "output_cost_per_token": 3e-07, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://docs.mistral.ai/models/devstral-small-2-25-12", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true }, @@ -41509,9 +41514,14 @@ "max_output_tokens": 256000, "mode": "chat", "output_cost_per_token": 2e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://mistral.ai/news/devstral-2-vibe-cli", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true }, @@ -41523,9 +41533,14 @@ "max_output_tokens": 256000, "mode": "chat", "output_cost_per_token": 2e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://mistral.ai/news/devstral-2-vibe-cli", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true }, @@ -42220,8 +42235,13 @@ "max_output_tokens": 8191, "mode": "chat", "output_cost_per_token": 3e-07, + "reasoning_effort_levels": [ + "none", + "high" + ], "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true }, @@ -65476,9 +65496,14 @@ "max_output_tokens": 256000, "mode": "chat", "output_cost_per_token": 2e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://mistral.ai/news/devstral-2-vibe-cli", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true }, diff --git a/litellm/types/llms/mistral.py b/litellm/types/llms/mistral.py index e64aa9c13ab..cd61f32f430 100644 --- a/litellm/types/llms/mistral.py +++ b/litellm/types/llms/mistral.py @@ -1,6 +1,6 @@ from typing import Literal -from typing_extensions import TypedDict +from typing_extensions import NotRequired, ReadOnly, TypedDict class FunctionCall(TypedDict): @@ -22,3 +22,4 @@ class MistralTextBlock(TypedDict): class MistralThinkingBlock(TypedDict): type: Literal["thinking"] thinking: list[MistralTextBlock] + closed: NotRequired[ReadOnly[bool]] diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 50c23efca48..1b65c87c315 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41495,9 +41495,14 @@ "max_output_tokens": 256000, "mode": "chat", "output_cost_per_token": 3e-07, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://docs.mistral.ai/models/devstral-small-2-25-12", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true }, @@ -41509,9 +41514,14 @@ "max_output_tokens": 256000, "mode": "chat", "output_cost_per_token": 2e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://mistral.ai/news/devstral-2-vibe-cli", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true }, @@ -41523,9 +41533,14 @@ "max_output_tokens": 256000, "mode": "chat", "output_cost_per_token": 2e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://mistral.ai/news/devstral-2-vibe-cli", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true }, @@ -42220,8 +42235,13 @@ "max_output_tokens": 8191, "mode": "chat", "output_cost_per_token": 3e-07, + "reasoning_effort_levels": [ + "none", + "high" + ], "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true }, @@ -65476,9 +65496,14 @@ "max_output_tokens": 256000, "mode": "chat", "output_cost_per_token": 2e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://mistral.ai/news/devstral-2-vibe-cli", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true }, diff --git a/tests/unit/llms/mistral/test_mistral_chat_transformation.py b/tests/unit/llms/mistral/test_mistral_chat_transformation.py index 10f57be33fe..b76e1707805 100644 --- a/tests/unit/llms/mistral/test_mistral_chat_transformation.py +++ b/tests/unit/llms/mistral/test_mistral_chat_transformation.py @@ -1,4 +1,4 @@ -from typing import List, cast +from typing import Final, List, cast from unittest.mock import MagicMock, patch import pytest @@ -8,8 +8,10 @@ from litellm.llms.mistral.chat.transformation import ( MistralChatResponseIterator, MistralConfig, ) +from litellm.llms.vertex_ai.vertex_ai_partner_models.mistral.transformation import VertexAIMistralConfig from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse +from litellm.utils import supports_reasoning @pytest.mark.asyncio @@ -805,7 +807,7 @@ class TestMistralStripsOutputOnlyFields: Regression for https://github.com/BerriAI/litellm/issues/30835. """ - def test_assistant_reasoning_content_is_dropped(self): + def test_assistant_reasoning_content_is_sent_as_a_thinking_chunk(self): messages = cast( List[AllMessageValues], [ @@ -824,11 +826,10 @@ class TestMistralStripsOutputOnlyFields: MistralConfig().transform_messages(messages=messages, model="mistral-medium-3-5"), ) - assistant_message = result[-1] - assert "reasoning_content" not in assistant_message - assert "thinking_blocks" not in assistant_message - assert assistant_message["content"] == "Follow-up" - assert assistant_message["role"] == "assistant" + assert result[-1] == { + "role": "assistant", + "content": [_think_chunk("Some internal reasoning text."), {"type": "text", "text": "Follow-up"}], + } def test_tool_search_reply_drops_litellm_response_fields(self): tool_search_reply = { @@ -920,6 +921,508 @@ class TestMistralStripsOutputOnlyFields: assert "reasoning_content" not in result[-1] +def _think_chunk(text: str) -> dict[str, object]: + return {"type": "thinking", "thinking": [{"type": "text", "text": text}], "closed": True} + + +def _mistral_request_messages( + messages: list[dict[str, object]], config: MistralConfig | None = None, model: str = "mistral-medium-3-5" +) -> list[object]: + request: Final = (config or MistralConfig()).transform_request( + model=model, + messages=cast(List[AllMessageValues], messages), + optional_params={}, + litellm_params={}, + headers={}, + ) + return request["messages"] + + +@pytest.mark.parametrize("content", [None, ""], ids=["null", "empty"]) +def test_mistral_sends_reasoning_of_a_tool_call_turn_as_a_thinking_chunk(content: str | None): + tool_call: Final = {"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": "{}"}} + + result: Final = _mistral_request_messages( + [ + {"role": "user", "content": "Weather in Paris?"}, + { + "role": "assistant", + "content": content, + "reasoning_content": "I should call get_weather.", + "tool_calls": [tool_call], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "sunny"}, + ] + ) + + assert result[1] == { + "role": "assistant", + "content": [_think_chunk("I should call get_weather.")], + "tool_calls": [tool_call], + } + assert result[2] == {"role": "tool", "tool_call_id": "call_1", "content": "sunny"} + + +def test_mistral_sends_reasoning_ahead_of_a_text_part_list(): + result: Final = _mistral_request_messages( + [ + {"role": "user", "content": "Question?"}, + { + "role": "assistant", + "content": [{"type": "text", "text": "First. "}, {"type": "text", "text": "Second."}], + "reasoning_content": "Two parts.", + }, + {"role": "user", "content": "And then?"}, + ] + ) + + assert result[1] == { + "role": "assistant", + "content": [_think_chunk("Two parts."), {"type": "text", "text": "First. Second."}], + } + + +def test_mistral_drops_a_reasoning_only_assistant_turn_between_a_tool_call_and_its_reply(): + tool_call: Final = {"id": "call_1", "type": "function", "function": {"name": "read", "arguments": "{}"}} + + result: Final = _mistral_request_messages( + [ + {"role": "user", "content": "read the log"}, + {"role": "assistant", "content": None, "tool_calls": [tool_call]}, + {"role": "assistant", "content": None, "reasoning_content": "The log is empty."}, + {"role": "tool", "tool_call_id": "call_1", "content": ""}, + ] + ) + + assert result == [ + {"role": "user", "content": "read the log"}, + {"role": "assistant", "tool_calls": [tool_call]}, + {"role": "tool", "tool_call_id": "call_1", "content": ""}, + ] + + +@pytest.mark.parametrize( + "assistant_content, reasoning_fields, expected_content", + [ + pytest.param( + [{"type": "thinking", "thinking": "R", "signature": "c2ln"}, {"type": "text", "text": "A"}], + {"reasoning_content": "R"}, + [_think_chunk("R"), {"type": "text", "text": "A"}], + id="anthropic_thinking_part_with_reasoning_content", + ), + pytest.param( + [{"type": "thinking", "thinking": "R", "signature": "c2ln"}, {"type": "text", "text": "A"}], + {}, + "A", + id="anthropic_thinking_part_alone", + ), + pytest.param( + {"type": "text", "text": "A"}, + {"reasoning_content": "R"}, + {"type": "text", "text": "A"}, + id="content_that_is_not_a_list", + ), + ], +) +def test_mistral_never_forwards_a_thinking_part_mistral_cannot_read( + assistant_content: object, reasoning_fields: dict[str, object], expected_content: object +): + result: Final = _mistral_request_messages( + [ + {"role": "user", "content": "Hi"}, + {"role": "assistant", "content": assistant_content, **reasoning_fields}, + {"role": "user", "content": "Bye"}, + ] + ) + + assert result[1] == {"role": "assistant", "content": expected_content} + + +@pytest.mark.parametrize( + "reasoning_fields, expected_reasoning", + [ + pytest.param( + { + "thinking_blocks": [ + {"type": "thinking", "thinking": "Let", "signature": "mistral"}, + {"type": "thinking", "thinking": " me", "signature": "mistral"}, + {"type": "thinking", "thinking": " think", "signature": "mistral"}, + ] + }, + "Let me think", + id="mistral_stream_fragments", + ), + pytest.param( + { + "thinking_blocks": [ + {"type": "thinking", "thinking": "Step one.", "signature": "c2lnLTE="}, + {"type": "redacted_thinking", "data": "ZW5jcnlwdGVk"}, + {"type": "thinking", "thinking": " Step two.", "signature": "c2lnLTI="}, + ] + }, + "Step one. Step two.", + id="anthropic_blocks_with_redacted", + ), + pytest.param( + {"reasoning_content": " ", "thinking_blocks": [{"type": "thinking", "thinking": "From blocks."}]}, + "From blocks.", + id="blank_reasoning_content_falls_back_to_blocks", + ), + ], +) +def test_mistral_sends_replayed_thinking_blocks_as_one_thinking_chunk( + reasoning_fields: dict[str, object], expected_reasoning: str +): + result: Final = _mistral_request_messages( + [ + {"role": "user", "content": "Question?"}, + {"role": "assistant", "content": "Answer.", **reasoning_fields}, + {"role": "user", "content": "And then?"}, + ] + ) + + assert result[1] == { + "role": "assistant", + "content": [_think_chunk(expected_reasoning), {"type": "text", "text": "Answer."}], + } + + +@pytest.mark.parametrize( + "native_content, reasoning_fields", + [ + pytest.param( + [ + {"type": "thinking", "thinking": [{"type": "text", "text": "I pick 518362."}]}, + {"type": "text", "text": "I have picked it."}, + ], + {}, + id="native_chunk", + ), + pytest.param( + [ + {"type": "thinking", "thinking": [{"type": "text", "text": "I pick 518362."}]}, + {"type": "text", "text": "I have picked it."}, + ], + {"reasoning_content": "I pick 518362.", "thinking_blocks": [{"type": "thinking", "thinking": "I pick"}]}, + id="native_chunk_and_reasoning_content", + ), + pytest.param( + [ + {"type": "text", "text": "I have "}, + {"type": "thinking", "thinking": [{"type": "text", "text": "I pick 518362."}]}, + {"type": "text", "text": "picked it."}, + ], + {}, + id="native_chunk_between_text_parts", + ), + ], +) +def test_mistral_keeps_a_native_thinking_chunk_once( + native_content: list[dict[str, object]], reasoning_fields: dict[str, object] +): + result: Final = _mistral_request_messages( + [ + {"role": "user", "content": "Pick a number."}, + {"role": "assistant", "content": native_content, **reasoning_fields}, + {"role": "user", "content": "Which one?"}, + ] + ) + + assert result[1] == { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": [{"type": "text", "text": "I pick 518362."}]}, + {"type": "text", "text": "I have picked it."}, + ], + } + + +def test_mistral_keeps_a_native_thinking_chunk_with_no_text_part(): + native_chunk: Final = {"type": "thinking", "thinking": [{"type": "text", "text": "I pick 518362."}], "closed": True} + + result: Final = _mistral_request_messages( + [ + {"role": "user", "content": "Pick a number."}, + {"role": "assistant", "content": [native_chunk]}, + {"role": "user", "content": "Which one?"}, + ] + ) + + assert result[1] == {"role": "assistant", "content": [native_chunk]} + + +@pytest.mark.parametrize( + ("inner_items", "expected_content"), + [ + pytest.param( + [{"type": "text", "text": "I pick 518362."}, {"type": "reference", "reference_ids": [1]}], + [ + {"type": "thinking", "thinking": [{"type": "text", "text": "I pick 518362."}], "closed": True}, + {"type": "text", "text": "I have picked it."}, + ], + id="text_and_reference", + ), + pytest.param( + [{"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}}], + "I have picked it.", + id="only_image_url", + ), + pytest.param([1, 2], "I have picked it.", id="not_items"), + ], +) +def test_mistral_sends_only_the_text_items_of_a_native_thinking_chunk( + inner_items: list[object], expected_content: object +): + result: Final = _mistral_request_messages( + [ + {"role": "user", "content": "Pick a number."}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": inner_items, "closed": True}, + {"type": "text", "text": "I have picked it."}, + ], + }, + {"role": "user", "content": "Which one?"}, + ] + ) + + assert result[1] == {"role": "assistant", "content": expected_content} + + +def test_mistral_sends_only_the_text_items_of_a_native_thinking_chunk_beside_an_image(): + image_part: Final = {"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}} + + result: Final = _mistral_request_messages( + [ + {"role": "user", "content": [{"type": "text", "text": "What is this?"}, image_part]}, + { + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": [{"type": "text", "text": "It has whiskers."}, {"type": "reference", "reference_ids": [1]}], + }, + {"type": "text", "text": "A cat."}, + ], + }, + {"role": "user", "content": "What color?"}, + ] + ) + + assert result[1] == { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": [{"type": "text", "text": "It has whiskers."}]}, + {"type": "text", "text": "A cat."}, + ], + } + + +_OPEN_THINK_CHUNK: Final = { + "type": "thinking", + "thinking": [{"type": "text", "text": "I pick 482913."}], + "closed": False, +} +_THINK_CHUNK_WITHOUT_CLOSED: Final = {"type": "thinking", "thinking": [{"type": "text", "text": "I pick 482913."}]} +_PICKED_TEXT: Final = {"type": "text", "text": "I have picked it."} + + +@pytest.mark.parametrize( + ("assistant_turn", "later_messages", "expected_turn"), + [ + pytest.param( + {"role": "assistant", "content": [_OPEN_THINK_CHUNK, _PICKED_TEXT]}, + [{"role": "user", "content": "Which one?"}], + {"role": "assistant", "content": [_THINK_CHUNK_WITHOUT_CLOSED, _PICKED_TEXT]}, + id="history_turn", + ), + pytest.param( + {"role": "assistant", "content": [_OPEN_THINK_CHUNK], "prefix": True}, + [], + {"role": "assistant", "content": [_OPEN_THINK_CHUNK], "prefix": True}, + id="last_chunk_of_final_prefix_turn", + ), + pytest.param( + {"role": "assistant", "content": [_OPEN_THINK_CHUNK, _PICKED_TEXT], "prefix": True}, + [], + {"role": "assistant", "content": [_THINK_CHUNK_WITHOUT_CLOSED, _PICKED_TEXT], "prefix": True}, + id="final_prefix_turn_with_text_after_the_chunk", + ), + pytest.param( + {"role": "assistant", "content": [_PICKED_TEXT, _OPEN_THINK_CHUNK], "prefix": True}, + [], + {"role": "assistant", "content": [_THINK_CHUNK_WITHOUT_CLOSED, _PICKED_TEXT], "prefix": True}, + id="final_prefix_turn_with_text_before_the_chunk", + ), + pytest.param( + {"role": "assistant", "content": [_OPEN_THINK_CHUNK]}, + [], + {"role": "assistant", "content": [_THINK_CHUNK_WITHOUT_CLOSED]}, + id="final_turn_without_prefix", + ), + ], +) +def test_mistral_keeps_an_open_thinking_chunk_only_as_the_last_chunk_of_a_final_prefix_turn( + assistant_turn: dict[str, object], later_messages: list[dict[str, object]], expected_turn: dict[str, object] +): + result: Final = _mistral_request_messages( + [{"role": "user", "content": "Pick a number."}, assistant_turn, *later_messages] + ) + + assert result[1] == expected_turn + + +@pytest.mark.parametrize( + "reasoning_fields", + [ + pytest.param({}, id="none"), + pytest.param({"reasoning_content": None, "thinking_blocks": None}, id="null"), + pytest.param({"reasoning_content": ""}, id="empty_reasoning_content"), + pytest.param({"reasoning_content": " \n "}, id="blank_reasoning_content"), + pytest.param({"thinking_blocks": []}, id="empty_thinking_blocks"), + pytest.param({"thinking_blocks": [{"type": "redacted_thinking", "data": "ZW5j"}]}, id="only_redacted"), + pytest.param({"thinking_blocks": [{"type": "thinking", "thinking": " "}]}, id="blank_thinking_block"), + pytest.param({"thinking_blocks": "not a list"}, id="thinking_blocks_not_a_list"), + ], +) +def test_mistral_sends_an_assistant_turn_without_usable_reasoning_as_plain_text(reasoning_fields: dict[str, object]): + result: Final = _mistral_request_messages( + [ + {"role": "user", "content": "Question?"}, + {"role": "assistant", "content": "Answer.", **reasoning_fields}, + {"role": "user", "content": "And then?"}, + ] + ) + + assert result[1] == {"role": "assistant", "content": "Answer."} + + +def test_mistral_still_flattens_text_part_lists_to_a_string(): + result: Final = _mistral_request_messages( + [ + {"role": "user", "content": [{"type": "text", "text": "Hello "}, {"type": "text", "text": "there"}]}, + {"role": "assistant", "content": [{"type": "text", "text": "Hi "}, {"type": "text", "text": "back"}]}, + {"role": "user", "content": "Bye"}, + ] + ) + + assert result == [ + {"role": "user", "content": "Hello there"}, + {"role": "assistant", "content": "Hi back"}, + {"role": "user", "content": "Bye"}, + ] + + +def test_mistral_sends_reasoning_as_a_thinking_chunk_when_the_request_carries_an_image(): + data_uri: Final = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUg==" + + result: Final = _mistral_request_messages( + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this"}, + {"type": "image_url", "image_url": {"url": data_uri}}, + ], + }, + {"role": "assistant", "content": "A cat.", "reasoning_content": "It has whiskers."}, + {"role": "user", "content": "What color?"}, + ] + ) + + assert result[1] == { + "role": "assistant", + "content": [_think_chunk("It has whiskers."), {"type": "text", "text": "A cat."}], + } + + +@pytest.mark.parametrize( + "config", + [pytest.param(MistralConfig(), id="mistral"), pytest.param(VertexAIMistralConfig(), id="vertex_ai")], +) +@pytest.mark.parametrize( + "assistant_message", + [ + pytest.param( + {"role": "assistant", "content": "I have picked it.", "reasoning_content": "I pick 518362."}, + id="reasoning_content", + ), + pytest.param( + { + "role": "assistant", + "content": "I have picked it.", + "thinking_blocks": [{"type": "thinking", "thinking": "I pick 518362.", "signature": "mistral"}], + }, + id="thinking_blocks", + ), + pytest.param( + { + "role": "assistant", + "content": [_think_chunk("I pick 518362."), {"type": "text", "text": "I have picked it."}], + }, + id="native_chunk", + ), + ], +) +def test_mistral_drops_replayed_reasoning_for_a_model_without_reasoning( + config: MistralConfig, assistant_message: dict[str, object] +): + model: Final = "mistral-chat-model-without-reasoning" + assert not supports_reasoning(model=model, custom_llm_provider=config.custom_llm_provider) + + result: Final = _mistral_request_messages( + [{"role": "user", "content": "Pick a number."}, assistant_message, {"role": "user", "content": "Which one?"}], + config=config, + model=model, + ) + + assert result[1] == {"role": "assistant", "content": "I have picked it."} + + +@pytest.mark.parametrize( + "config", + [pytest.param(MistralConfig(), id="mistral"), pytest.param(VertexAIMistralConfig(), id="vertex_ai")], +) +def test_mistral_drops_a_native_chunk_without_text_for_a_model_without_reasoning(config: MistralConfig): + tool_call: Final = {"id": "call_1", "type": "function", "function": {"name": "get_time", "arguments": "{}"}} + + result: Final = _mistral_request_messages( + [ + {"role": "user", "content": "What time is it?"}, + {"role": "assistant", "content": [_think_chunk("I should call get_time.")], "tool_calls": [tool_call]}, + {"role": "tool", "tool_call_id": "call_1", "content": "12:00"}, + ], + config=config, + model="mistral-chat-model-without-reasoning", + ) + + assert result[1] == {"role": "assistant", "tool_calls": [tool_call]} + + +@pytest.mark.parametrize( + "config", + [pytest.param(MistralConfig(), id="mistral"), pytest.param(VertexAIMistralConfig(), id="vertex_ai")], +) +def test_mistral_drops_a_native_chunk_from_an_image_history_for_a_model_without_reasoning(config: MistralConfig): + image_part: Final = {"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}} + + result: Final = _mistral_request_messages( + [ + {"role": "user", "content": [{"type": "text", "text": "What is this?"}, image_part]}, + { + "role": "assistant", + "content": [_think_chunk("It has whiskers."), {"type": "text", "text": "A cat."}], + }, + {"role": "user", "content": "What color?"}, + ], + config=config, + model="mistral-chat-model-without-reasoning", + ) + + assert result[1] == {"role": "assistant", "content": [{"type": "text", "text": "A cat."}]} + + def test_mistral_transform_request_keeps_an_empty_tool_message_after_its_tool_call(): messages: List[AllMessageValues] = cast( List[AllMessageValues], From 9b1d69e32c086dfed314beb11013949759f55b1a Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 28 Oct 2026 15:10:53 -0400 Subject: [PATCH 066/128] fix(images): normalize absent images in async edits (#48872) --- litellm/images/main.py | 4 ++-- tests/unit/images/test_image_edit.py | 16 +++++++++++++++- 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/litellm/images/main.py b/litellm/images/main.py index 3312ca417d5..84156919774 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -1062,7 +1062,7 @@ def image_edit( @client async def aimage_edit( - image: FileTypes | list[FileTypes], + image: FileTypes | list[FileTypes] | None, model: str, prompt: str, mask: str | None = None, @@ -1102,7 +1102,7 @@ async def aimage_edit( model=model, api_base=local_vars.get("base_url", None) ) - images: Final = image if isinstance(image, list) else [image] + images: Final = image if isinstance(image, list) else ([image] if image is not None else []) func: Final = partial( image_edit, diff --git a/tests/unit/images/test_image_edit.py b/tests/unit/images/test_image_edit.py index d8f1990af0e..fb7a2869b2a 100644 --- a/tests/unit/images/test_image_edit.py +++ b/tests/unit/images/test_image_edit.py @@ -12,7 +12,7 @@ from typing_extensions import ReadOnly, TypedDict, override import litellm from litellm.integrations.custom_logger import CustomLogger -from litellm.types.utils import ImageResponse +from litellm.types.utils import FileTypes, ImageResponse _PNG_SIGNATURE: Final = b"\x89PNG\r\n\x1a\n" _FIRST_IMAGE: Final = _PNG_SIGNATURE + b"first-reference-image" @@ -73,6 +73,20 @@ def _multipart_image_parts(request: httpx.Request) -> tuple[bytes, ...]: return tuple(part.split(b"\r\n\r\n", 1)[1].removesuffix(b"\r\n") for part in parts if b'name="image[]"' in part) +@pytest.mark.asyncio +@pytest.mark.parametrize("image", (None, []), ids=("none", "empty-list")) +async def test_bfl_image_edit_missing_image_returns_bad_request(image: list[FileTypes] | None) -> None: + with pytest.raises(litellm.BadRequestError, match="No image provided") as exc_info: + await litellm.aimage_edit( + model="black_forest_labs/flux-kontext-pro", + prompt="make it blue", + image=image, + api_key="fake-key", + ) + + assert exc_info.value.status_code == 400 + + def test_openai_image_edit_sync_accepts_bytesio_images(respx_mock: respx.MockRouter, httpx_transport: None) -> None: route: Final = respx_mock.post("https://api.openai.com/v1/images/edits").mock( return_value=httpx.Response(200, json=_EDIT_RESPONSE) From ff5bf4826dab9ee145de96966488cbd195a1063a Mon Sep 17 00:00:00 2001 From: rj-cairn <349020452+rj-cairn@users.noreply.github.com> Date: Wed, 28 Oct 2026 12:11:35 -0700 Subject: [PATCH 067/128] fix(count_tokens)!: audit #47441 and keep the framed estimate for Bedrock count models from the cost map (#48805) * fix(count_tokens): keep the framed estimate when a region-less token_count_model cannot count (audit #47441) * fix(count_tokens): let framing stand in for a shipped token_count_model, keep a configured one's error The framed local estimate now replaces a 400/404 from any count model the cost map ships, Mantle-pinned rows included, and an admin-set token_count_model keeps its error as the docs say. * fix(count_tokens): warn with the region when framing stands in for a pinned count model Test a Mantle row #47441 pinned in its own region, and keep the region hint check on an admin-set count model. * fix(count_tokens): name the deployment in the framing stand-in warning * fix(count_tokens): let framing stand in when a cost map count model's region denies the account A count model the cost map pins to a region can answer 401 or 403 when the account has not enabled that region or lacks Mantle count access there. With framing set, the estimate now stands in, as it does for 400 and 404. An admin set count model and a count model with no pinned region keep the error. The warning now names the model group, the count model and its region, and says the proxy falls back to the estimate. * test(count_tokens): cover framed 401 and 403 from a shipped count model region on the wire --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../common_utils/token_count_settings.py | 13 +- litellm/proxy/proxy_server.py | 33 +++- .../providers/test_count_tokens_model_wire.py | 185 +++++++++++++++++- .../common_utils/test_token_count_settings.py | 33 ++-- tests/unit/proxy/test_proxy_token_counter.py | 118 +++++++++-- 5 files changed, 334 insertions(+), 48 deletions(-) diff --git a/litellm/proxy/common_utils/token_count_settings.py b/litellm/proxy/common_utils/token_count_settings.py index 22b35bd1001..4e7dc012802 100644 --- a/litellm/proxy/common_utils/token_count_settings.py +++ b/litellm/proxy/common_utils/token_count_settings.py @@ -13,6 +13,7 @@ from litellm.types.token_count import TokenCountProvider, TokenCountSourceFields from litellm.types.utils import TokenCountResponse CANNOT_COUNT_MODEL_STATUSES: Final = frozenset({400, 404}) +REGION_DENIED_STATUSES: Final = frozenset({401, 403}) TOKENIZER_LOAD_WAIT_SECONDS: Final = 5.0 TOKENIZER_LOAD_REUSE_SECONDS: Final = 60.0 TOKENIZER_LOAD_MAX_RUNNING: Final = 4 @@ -47,6 +48,7 @@ class TokenCountModel: model: str region: str | None litellm_provider: TokenCountProvider | None = None + configured: bool = False @dataclass(frozen=True, slots=True) @@ -66,6 +68,10 @@ class ProviderCountFailure: def cannot_count_model(self) -> bool: return not self.rejected_before_provider_call and self.status_code in CANNOT_COUNT_MODEL_STATUSES + @property + def denied_in_region(self) -> bool: + return not self.rejected_before_provider_call and self.status_code in REGION_DENIED_STATUSES + def with_hint(self, hint: str) -> "ProviderCountFailure": return replace(self, message=_with_hint(self.message, hint)) @@ -111,7 +117,7 @@ def _count_provider(value: object) -> TokenCountProvider | None: return None -def _count_model(value: Mapping[str, object]) -> TokenCountModel | None: +def _count_model(value: Mapping[str, object], configured: bool) -> TokenCountModel | None: model: Final = _non_blank(value.get("model")) return ( None @@ -120,6 +126,7 @@ def _count_model(value: Mapping[str, object]) -> TokenCountModel | None: model=model, region=_non_blank(value.get("region")), litellm_provider=_count_provider(value.get("litellm_provider")), + configured=configured, ) ) @@ -146,12 +153,12 @@ def _resolved_count_model( ) -> TokenCountModel | None: configured_mapping: Final = as_str_mapping(configured) if configured_mapping is not None: - return _count_model(configured_mapping) + return _count_model(configured_mapping, configured=True) default_mapping: Final = as_str_mapping(default) return ( None if default_mapping is None - else _within_deployment_geography(_count_model(default_mapping), deployment_region) + else _within_deployment_geography(_count_model(default_mapping, configured=False), deployment_region) ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7fcb77e2bbf..64fb2791913 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -14853,6 +14853,7 @@ async def _try_native_responses_count( async def _provider_token_count( count_with: Callable[[str, str | None], Awaitable[Union["TokenCountResponse", ProviderCountFailure, None]]], model_to_use: str, + request_model: str, settings: TokenCountSettings, local_estimate_allowed: bool, ) -> Union["TokenCountResponse", ProviderCountFailure, None]: @@ -14875,17 +14876,28 @@ async def _provider_token_count( if count_model is None: return counted.model_copy(update={"count_source": "provider"}) return counted.model_copy(update={"count_source": "token_count_model", "model_used": count_model.model}) - failure: Final = ( - counted - if count_model is None or count_model.region is None - else counted.with_hint( - f"This model counts tokens with {count_model.model} in {count_model.region}. " - "If that region is not enabled for the AWS account, enable it or set token_count_model in model_info" - ) + pinned_count_model: Final = None if count_model is None or count_model.region is None else count_model + region_hint: Final = ( + None + if pinned_count_model is None + else f"This model counts tokens with {pinned_count_model.model} in {pinned_count_model.region}. " + "If that region is not enabled for the AWS account, enable it or set token_count_model in model_info" ) - if count_model is None and failure.cannot_count_model and settings.framing is not None and local_estimate_allowed: - return failure - raise failure.to_proxy_exception() + failure: Final = counted if region_hint is None else counted.with_hint(region_hint) + framing_can_stand_in: Final = (count_model is None or not count_model.configured) and ( + failure.cannot_count_model or (pinned_count_model is not None and failure.denied_in_region) + ) + if not (framing_can_stand_in and settings.framing is not None and local_estimate_allowed): + raise failure.to_proxy_exception() + if pinned_count_model is not None: + verbose_proxy_logger.warning( + "Token count for %s answered %s from %s in %s, so it falls back to the framed local estimate", + request_model, + failure.status_code, + pinned_count_model.model, + pinned_count_model.region, + ) + return failure @dataclass(frozen=True, slots=True) @@ -15455,6 +15467,7 @@ async def _count_request_tokens( chat_params=MappingProxyType({}) if chat_count is None else chat_count.params, ), model_to_use=model_to_use, + request_model=request.model, settings=settings, local_estimate_allowed=litellm.disable_token_counter is not True and contents_estimable, ) diff --git a/tests/integration/providers/test_count_tokens_model_wire.py b/tests/integration/providers/test_count_tokens_model_wire.py index b6ace7ddef3..f15545f4374 100644 --- a/tests/integration/providers/test_count_tokens_model_wire.py +++ b/tests/integration/providers/test_count_tokens_model_wire.py @@ -24,7 +24,10 @@ _SONNET: Final = "us.anthropic.claude-sonnet-5-5" _SONNET_BASE: Final = "anthropic.claude-sonnet-5-5" _OPUS: Final = "us.anthropic.claude-opus-5-5" _OPUS_BASE: Final = "anthropic.claude-opus-5-5" +_OPUS_4_7: Final = "us.anthropic.claude-opus-4-7" _NOVA: Final = "amazon.nova-lite-v1:0" +_RUNTIME_ROUTED: Final = "us.anthropic.claude-sonnet-4-6" +_CONFIGURED_COUNT_REGION: Final = "us-east-2" _KIMI: Final = "us.moonshotai.kimi-k3" _KIMI_BASE: Final = "moonshotai.kimi-k3" _PROFILE_ID: Final = "lit9563count" @@ -168,6 +171,7 @@ def _mantle(request: Request) -> Reply: return _json_reply(200, {"input_tokens": _COUNT_MODEL_COUNT}) error_type: Final = { 401: "authentication_error", + 403: "permission_error", 404: "not_found_error", 429: "rate_limit_error", }[status] @@ -206,9 +210,9 @@ def _mantle_count_requests(wire: Wire) -> tuple[tuple[tuple[str, str], dict[str, return tuple((_signer(request), _JSON_OBJECT.validate_json(request.body)) for request in requests) -def _mantle_body(prompt: str) -> dict[str, JsonValue]: +def _mantle_body(prompt: str, model: str = _OPUS_BASE) -> dict[str, JsonValue]: return { - "model": _OPUS_BASE, + "model": model, "messages": [{"role": "user", "content": prompt}], "system": _SYSTEM, "tools": _SDK_TOOLS, @@ -219,13 +223,13 @@ def _url(port: int) -> str: return f"http://127.0.0.1:{port}" -def _litellm_params(model: str, api_base: str) -> dict[str, JsonValue]: +def _litellm_params(model: str, api_base: str, region: str = _REGION) -> dict[str, JsonValue]: return { "model": f"bedrock/{model}", "api_base": api_base, "aws_access_key_id": _ACCESS_KEY, "aws_secret_access_key": _SECRET_KEY, - "aws_region_name": _REGION, + "aws_region_name": region, } @@ -259,9 +263,13 @@ def counting_proxy(tmp_path_factory: pytest.TempPathFactory, mantle_port: int) - def _deployment( - scenario: Scenario, runtime_port: int, model: str = _SONNET, model_info: Mapping[str, JsonValue] | None = None + scenario: Scenario, + runtime_port: int, + model: str = _SONNET, + model_info: Mapping[str, JsonValue] | None = None, + region: str = _REGION, ) -> str: - return scenario.model(model_info=model_info, api_key=None, **_litellm_params(model, _url(runtime_port))) + return scenario.model(model_info=model_info, api_key=None, **_litellm_params(model, _url(runtime_port), region)) def _count(gateway: Gateway, model: str, status: int, count_model_status: int | None = None) -> MessageTokensCount: @@ -382,6 +390,171 @@ def test_framing_on_one_deployment_leaves_a_sibling_of_the_same_model_raising_th assert _mantle_models(mantle) == () +def test_framing_is_added_to_the_local_estimate_when_a_shipped_runtime_count_model_cannot_count( + counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + count_model: Final = _shipped_token_count_model_field(_RUNTIME_ROUTED, "model") + assert _shipped_token_count_model(_RUNTIME_ROUTED).get("region") is None + with ( + wire_server(_runtime, port=runtime_port) as runtime, + wire_server(_mantle, port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + framed: Final = _count( + gateway, + _deployment(scenario, runtime_port, _RUNTIME_ROUTED, {"token_count_framing": _FRAMING}), + 400, + count_model_status=404, + ) + unframed: Final = _count( + gateway, + _deployment(scenario, runtime_port, _RUNTIME_ROUTED, {"token_count_framing": 0}), + 400, + count_model_status=404, + ) + assert (framed.model_extra, unframed.model_extra) == (_LOCAL_SOURCE, _LOCAL_SOURCE), (framed, unframed) + assert framed.input_tokens - unframed.input_tokens == _FRAMING, (framed, unframed) + assert _runtime_models(runtime) == (count_model, count_model) + assert _mantle_models(mantle) == (count_model, count_model) + + +_OPUS_4_7_PINNED: Final = ( + _OPUS_4_7, + _REGION, + _shipped_token_count_model_field(_OPUS_4_7, "region"), + "anthropic.claude-opus-4-7", +) + + +@pytest.mark.parametrize( + ("model", "deployment_region", "count_region", "count_model", "count_model_status"), + [ + (_SONNET, _REGION, _COUNT_REGION, _OPUS_BASE, 404), + (_SONNET_BASE, "ca-central-1", "ca-central-1", _OPUS_BASE, 404), + (*_OPUS_4_7_PINNED, 404), + (*_OPUS_4_7_PINNED, 401), + (*_OPUS_4_7_PINNED, 403), + ], + ids=[ + "sonnet_5_5_in_its_pinned_region", + "sonnet_5_5_outside_the_pinned_geography", + "opus_4_7_in_its_pinned_region", + "opus_4_7_pinned_region_denies_the_key", + "opus_4_7_pinned_region_denies_the_account", + ], +) +def test_framing_is_added_to_the_local_estimate_when_a_shipped_mantle_count_model_cannot_count( + counting_proxy: OwnedProxy, + runtime_port: int, + mantle_port: int, + model: str, + deployment_region: str, + count_region: str, + count_model: str, + count_model_status: int, +) -> None: + gateway: Final = counting_proxy.gateway + assert _shipped_token_count_model_field(model, "litellm_provider") == "bedrock_mantle" + with ( + wire_server(_runtime, port=runtime_port) as runtime, + wire_server(_mantle, port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + framed: Final = _count( + gateway, + _deployment(scenario, runtime_port, model, {"token_count_framing": _FRAMING}, deployment_region), + 200, + count_model_status=count_model_status, + ) + unframed: Final = _count( + gateway, + _deployment(scenario, runtime_port, model, {"token_count_framing": 0}, deployment_region), + 200, + count_model_status=count_model_status, + ) + assert (framed.model_extra, unframed.model_extra) == (_LOCAL_SOURCE, _LOCAL_SOURCE), (framed, unframed) + assert framed.input_tokens - unframed.input_tokens == _FRAMING, (framed, unframed) + assert _runtime_counts(runtime) == () + assert ( + _mantle_count_requests(mantle) + == (((_ACCESS_KEY, count_region), _mantle_body(_prompt(200, count_model_status), count_model)),) * 2 + ) + + +def test_framing_keeps_the_403_when_the_shipped_count_model_runs_in_the_deployment_region( + counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime, port=runtime_port) as runtime, + wire_server(_mantle, port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + with pytest.raises(anthropic.PermissionDeniedError, match=re.escape(f"scripted mantle 403 for {_OPUS_BASE}")): + _count( + gateway, + _deployment(scenario, runtime_port, _SONNET_BASE, {"token_count_framing": _FRAMING}, "ca-central-1"), + 200, + count_model_status=403, + ) + assert _runtime_counts(runtime) == () + assert _mantle_count_requests(mantle) == ( + ((_ACCESS_KEY, "ca-central-1"), _mantle_body(_prompt(200, 403))), + ) + + +def test_configured_token_count_model_error_stands_with_framing( + counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + count_model: Final = _shipped_token_count_model_field(_RUNTIME_ROUTED, "model") + model_info: Final[dict[str, JsonValue]] = { + "token_count_framing": _FRAMING, + "token_count_model": {"model": count_model}, + } + with ( + wire_server(_runtime, port=runtime_port) as runtime, + wire_server(_mantle, port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + with pytest.raises(anthropic.NotFoundError, match=re.escape(f"scripted mantle 404 for {count_model}")): + _count( + gateway, _deployment(scenario, runtime_port, _RUNTIME_ROUTED, model_info), 400, count_model_status=404 + ) + assert _runtime_models(runtime) == (count_model,) + assert _mantle_models(mantle) == (count_model,) + + +def test_configured_token_count_model_provider_counts_on_mantle_without_trying_the_runtime( + counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + count_model: Final = _shipped_token_count_model_field(_RUNTIME_ROUTED, "model") + assert _shipped_token_count_model_field(_RUNTIME_ROUTED, "litellm_provider") == "bedrock" + configured: Final[dict[str, JsonValue]] = { + "model": count_model, + "litellm_provider": "bedrock_mantle", + "region": _CONFIGURED_COUNT_REGION, + } + with ( + wire_server(_runtime, port=runtime_port) as runtime, + wire_server(_mantle, port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + counted: Final = _count( + gateway, _deployment(scenario, runtime_port, _RUNTIME_ROUTED, {"token_count_model": configured}), 400 + ) + assert (counted.input_tokens, counted.model_extra) == ( + _COUNT_MODEL_COUNT, + {"litellm_count_source": "token_count_model", "litellm_count_model": count_model}, + ), counted + assert _runtime_counts(runtime) == () + assert _mantle_count_requests(mantle) == ( + ((_ACCESS_KEY, _CONFIGURED_COUNT_REGION), {**_mantle_body(_prompt(400)), "model": count_model}), + ) + + def test_empty_configured_token_count_model_counts_the_model_itself_in_the_deployment_region( counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int ) -> None: diff --git a/tests/unit/proxy/common_utils/test_token_count_settings.py b/tests/unit/proxy/common_utils/test_token_count_settings.py index 366d8b010d0..7ff32f9d4f3 100644 --- a/tests/unit/proxy/common_utils/test_token_count_settings.py +++ b/tests/unit/proxy/common_utils/test_token_count_settings.py @@ -32,22 +32,31 @@ def _in_us_west_2() -> str: ({}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)), ( {"token_count_framing": 5, "token_count_model": {"model": "custom-model", "region": "ap-southeast-4"}}, - TokenCountSettings(framing=5, count_model=TokenCountModel(model="custom-model", region="ap-southeast-4")), + TokenCountSettings( + framing=5, + count_model=TokenCountModel(model="custom-model", region="ap-southeast-4", configured=True), + ), ), ( {"token_count_model": {"model": "custom-model", "litellm_provider": "bedrock_mantle"}}, TokenCountSettings( framing=97, - count_model=TokenCountModel(model="custom-model", region=None, litellm_provider="bedrock_mantle"), + count_model=TokenCountModel( + model="custom-model", region=None, litellm_provider="bedrock_mantle", configured=True + ), ), ), ( {"token_count_model": {"model": "custom-model"}}, - TokenCountSettings(framing=97, count_model=TokenCountModel(model="custom-model", region=None)), + TokenCountSettings( + framing=97, count_model=TokenCountModel(model="custom-model", region=None, configured=True) + ), ), ( {"token_count_model": {"model": "custom-model", "region": " "}}, - TokenCountSettings(framing=97, count_model=TokenCountModel(model="custom-model", region=None)), + TokenCountSettings( + framing=97, count_model=TokenCountModel(model="custom-model", region=None, configured=True) + ), ), ({"token_count_framing": 0}, TokenCountSettings(framing=0, count_model=_OPUS_US_EAST_1)), ( @@ -80,26 +89,26 @@ def test_config_overrides_the_cost_map_default_key_by_key( @pytest.mark.parametrize( - ("deployment_region", "configured", "expected_region"), + ("deployment_region", "configured", "expected"), [ - ("us-west-2", {}, "us-east-1"), - ("eu-central-1", {}, None), - ("ca-central-1", {}, None), - ("us-gov-west-1", {}, None), + ("us-west-2", {}, _OPUS_US_EAST_1), + ("eu-central-1", {}, TokenCountModel(model="anthropic.claude-opus-5-5", region=None)), + ("ca-central-1", {}, TokenCountModel(model="anthropic.claude-opus-5-5", region=None)), + ("us-gov-west-1", {}, TokenCountModel(model="anthropic.claude-opus-5-5", region=None)), ( "eu-central-1", {"token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-east-1"}}, - "us-east-1", + TokenCountModel(model="anthropic.claude-opus-5-5", region="us-east-1", configured=True), ), ], ) def test_shipped_count_region_never_takes_the_prompt_out_of_the_deployment_geography( - deployment_region: str, configured: dict[str, object], expected_region: str | None + deployment_region: str, configured: dict[str, object], expected: TokenCountModel ) -> None: settings: Final = resolve_token_count_settings( configured=configured, default=_DEFAULT, deployment_region=lambda: deployment_region ) - assert settings.count_model == TokenCountModel(model="anthropic.claude-opus-5-5", region=expected_region) + assert settings.count_model == expected @pytest.mark.parametrize("prefix", ("", "bedrock/", "bedrock_mantle/")) diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 5484bc38317..0e8fc6823bd 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -12,6 +12,7 @@ import threading import traceback import uuid from collections.abc import Callable, Coroutine, Mapping +from dataclasses import replace from functools import partial from pathlib import Path from types import MappingProxyType @@ -209,7 +210,6 @@ async def test_gpt_token_counting(): assert response.request_model == "gpt-4" - @pytest.mark.asyncio async def test_token_counter_leaves_the_routers_deployment_without_the_credentials_it_resolves(tmp_path, monkeypatch): """token_counter reads the router's own deployment, so loading the named credential into it would leave the @@ -1685,6 +1685,7 @@ async def _ladder( system=_COUNT_REQUEST.system, ), model_to_use="us.anthropic.claude-sonnet-5-5", + request_model=_COUNT_REQUEST.model, settings=settings, local_estimate_allowed=local_estimate_allowed, ) @@ -1748,6 +1749,7 @@ async def test_token_count_model_provider_selects_the_endpoint_and_preserves_the system=_COUNT_REQUEST.system, ), model_to_use="us.anthropic.claude-sonnet-5-5", + request_model=_COUNT_REQUEST.model, settings=settings, local_estimate_allowed=False, ) @@ -1786,15 +1788,22 @@ async def test_token_count_model_counts_the_same_request_in_its_own_region_witho assert _SONNET_DEPLOYMENT["litellm_params"]["aws_region_name"] == "us-west-2" +@pytest.mark.parametrize( + "count_model", + [None, TokenCountModel(model="us.anthropic.claude-sonnet-5-5", region=None)], + ids=["no_token_count_model", "shipped_token_count_model_in_the_deployment_region"], +) @pytest.mark.parametrize("status_code", [401, 403, 429, 500]) @pytest.mark.asyncio -async def test_errors_other_than_cannot_count_raise_even_with_framing(status_code: int): +async def test_errors_other_than_cannot_count_raise_even_with_framing( + count_model: TokenCountModel | None, status_code: int +): counter: Final = _ScriptedCounter( {"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=status_code)} ) with pytest.raises(ProxyException) as raised: - await _ladder(counter, TokenCountSettings(framing=10, count_model=None)) + await _ladder(counter, TokenCountSettings(framing=10, count_model=count_model)) assert (raised.value.code, raised.value.message) == (str(status_code), f"scripted {status_code} for sonnet") @@ -1840,28 +1849,97 @@ async def test_uncountable_model_without_framing_raises(): assert raised.value.code == "404" +@pytest.mark.parametrize( + ("count_model", "counted_as"), + [ + (None, "us.anthropic.claude-sonnet-5-5"), + ( + TokenCountModel(model="anthropic.claude-sonnet-5-5", region=None, litellm_provider="bedrock"), + "bedrock/anthropic.claude-sonnet-5-5", + ), + (_OPUS_IN_US_EAST_1, "us.anthropic.claude-opus-5-5"), + ], + ids=[ + "no_token_count_model", + "shipped_token_count_model_in_the_deployment_region", + "shipped_token_count_model_pinned_to_a_region", + ], +) @pytest.mark.asyncio -async def test_uncountable_model_with_framing_falls_through_to_the_local_estimate(): - counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=400)}) +async def test_uncountable_model_with_framing_falls_through_to_the_local_estimate( + count_model: TokenCountModel | None, counted_as: str +): + counter: Final = _ScriptedCounter({counted_as: _provider_reply("sonnet", status_code=400)}) - result: Final = await _ladder(counter, TokenCountSettings(framing=0, count_model=None)) + result: Final = await _ladder(counter, TokenCountSettings(framing=0, count_model=count_model)) assert isinstance(result, ProviderCountFailure) - assert len(counter.calls) == 1 + assert [call[0] for call in counter.calls] == [counted_as] -@pytest.mark.parametrize("status_code", [400, 404]) +@pytest.mark.parametrize("status_code", [401, 403]) @pytest.mark.asyncio -async def test_token_count_model_error_stands_even_with_framing(status_code: int): +async def test_framing_stands_in_when_a_shipped_token_count_model_region_denies_the_account(status_code: int): + counter: Final = _ScriptedCounter( + {"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=status_code)} + ) + + result: Final = await _ladder(counter, TokenCountSettings(framing=50, count_model=_OPUS_IN_US_EAST_1)) + + assert isinstance(result, ProviderCountFailure) + assert result.status_code == status_code + + +@pytest.mark.parametrize( + ("count_model", "region_hints"), + [ + (replace(_OPUS_IN_US_EAST_1, configured=True), 1), + (replace(_OPUS_IN_US_EAST_1, region=None, configured=True), 0), + ], + ids=["pinned_to_a_region", "in_the_deployment_region"], +) +@pytest.mark.parametrize("status_code", [400, 401, 403, 404]) +@pytest.mark.asyncio +async def test_configured_token_count_model_error_stands_even_with_framing( + count_model: TokenCountModel, region_hints: int, status_code: int +): counter: Final = _ScriptedCounter( {"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=status_code)} ) with pytest.raises(ProxyException) as raised: - await _ladder(counter, TokenCountSettings(framing=78, count_model=_OPUS_IN_US_EAST_1)) + await _ladder(counter, TokenCountSettings(framing=78, count_model=count_model)) assert raised.value.code == str(status_code) - assert "us.anthropic.claude-opus-5-5 in us-east-1" in raised.value.message + assert raised.value.message.startswith(f"scripted {status_code} for opus") + assert raised.value.message.count("counts tokens with us.anthropic.claude-opus-5-5 in us-east-1") == region_hints + + +@pytest.mark.parametrize( + ("count_model", "region_hints"), + [(_OPUS_IN_US_EAST_1, frozenset({True})), (replace(_OPUS_IN_US_EAST_1, region=None), frozenset())], + ids=["pinned_to_a_region", "in_the_deployment_region"], +) +@pytest.mark.asyncio +async def test_framing_standing_in_for_a_pinned_token_count_model_warns_with_its_model_and_region( + count_model: TokenCountModel, region_hints: frozenset[bool], caplog: pytest.LogCaptureFixture +): + counter: Final = _ScriptedCounter({"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=404)}) + verbose_proxy_logger.addHandler(caplog.handler) + try: + with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name): + result: Final = await _ladder(counter, TokenCountSettings(framing=50, count_model=count_model)) + finally: + verbose_proxy_logger.removeHandler(caplog.handler) + + assert isinstance(result, ProviderCountFailure) + stand_in_warnings: Final = [ + r.getMessage() for r in caplog.records if "falls back to the framed local estimate" in r.getMessage() + ] + assert { + f"for {_COUNT_REQUEST.model} answered 404 from us.anthropic.claude-opus-5-5 in us-east-1" in message + for message in stand_in_warnings + } == region_hints @pytest.mark.asyncio @@ -2837,7 +2915,9 @@ async def test_configured_token_count_model_overrides_the_shipped_one_and_an_emp async def test_shipped_token_count_model_keeps_the_prompt_in_the_deployment_geography( monkeypatch: pytest.MonkeyPatch, aws_region_name: str, counted_in: str ): - counter: Final = _ScriptedCounter({"bedrock_mantle/anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)}) + counter: Final = _ScriptedCounter( + {"bedrock_mantle/anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)} + ) monkeypatch.setattr( litellm.proxy.proxy_server, "llm_router", @@ -2888,7 +2968,9 @@ class _VaultWaitingForLoopProgress: async def test_region_lookup_for_a_shipped_count_model_leaves_the_event_loop_free(monkeypatch: pytest.MonkeyPatch): from litellm.proxy import proxy_server - counter: Final = _ScriptedCounter({"bedrock_mantle/anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)}) + counter: Final = _ScriptedCounter( + {"bedrock_mantle/anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)} + ) def bedrock_counter(deployment: object, model: str) -> tuple[_ScriptedCounter, str, str]: return counter, model, "bedrock" @@ -3653,7 +3735,11 @@ async def test_responses_count_sends_openai_the_provider_file_id_inference_sends model_list=[ { "model_name": "gpt-files", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key", "api_base": "https://count.test/v1"}, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "test-key", + "api_base": "https://count.test/v1", + }, "model_info": {"id": "deployment-a"}, } ] @@ -3833,9 +3919,7 @@ async def test_responses_count_on_claude_rejects_forced_tool_choice_without_drop 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"} - ) + 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) From 43ce3e22aa1a49bdbe1211ace8fdf362183cff71 Mon Sep 17 00:00:00 2001 From: tessellark <349031000+tessellark@users.noreply.github.com> Date: Wed, 28 Oct 2026 12:12:45 -0700 Subject: [PATCH 068/128] fix(responses)!: answer 400 before counting a lone Responses built-in tool on Claude (#48798) * test(responses): audit #47363 Adds integration cells for the input token count paths #47363 changed: the provider id of a managed file and the 403 for another team's file on OpenAI, file_search and additional_tools parity plus cache-mark stripping and untranslatable tools on Claude, and the Bedrock runtime count leaving out the options it refuses. * fix(count-tokens): reject a lone tool with no chat equivalent before the Claude count call A Claude input token count whose only tool was mcp or computer_use_preview skipped the pre-call check and went to the provider, so the caller got the provider error instead of the 400 naming the tool. The check now runs whenever a tool would get a placeholder name. * fix(responses): count a lone Anthropic toolset and name the uncountable tool The lone-tool 400 now fires only for OpenAI Responses built-in tool types, so a nameless Anthropic toolset still counts. The error says LiteLLM cannot count that tool type for the model. Adds a Bedrock unit cell for the lone-tool 400 and pins the runtime wire test's tool shape per request type. * fix(responses): take the uncountable tool types from the OpenAI SDK and keep the region hint off LiteLLM's 400 The Responses built-in tool types now come from openai's ToolParam union, so a type the SDK adds is refused before the count call too. The AWS region hint no longer lands on an error LiteLLM raised before calling the provider. Tests cover two more lone types and the hint. * fix(responses): keep LiteLLM's own list of uncountable tool types and pass a non-string tool type through * fix(responses): count tool_search and namespace tools the way Claude chat inference sends them The count refused a lone tool_search with LiteLLM's 400 and sent a lone namespace raw or squashed it into an empty tool, while Claude chat inference maps tool_search to Anthropic's tool search tool and unpacks a namespace into its function tools. Both now go through the deployment's own map_openai_params, so the count sends what inference sends, and tool_search leaves the list of uncountable types. * fix(bedrock): count a Claude request with its deployment's own chat translation Bedrock invoke and Mantle inference turn a chat tool_search tool into Bedrock's tool search tool, but the Bedrock counter always translated with the Anthropic config, so the count answered 400 for a request inference serves. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../anthropic/count_tokens/transformation.py | 59 ++++++++-- .../count_tokens/bedrock_token_counter.py | 23 +++- .../common_utils/token_count_settings.py | 2 + ...anthropic_count_tokens_system_lift_wire.py | 93 +++++++++++++++- .../test_bedrock_mantle_count_tokens_wire.py | 68 ++++++++++++ .../test_openai_count_tokens_tools_wire.py | 102 ++++++++++++++++++ ...t_anthropic_count_tokens_transformation.py | 95 ++++++++++++++++ .../test_bedrock_token_counter.py | 74 +++++++++++++ .../common_utils/test_token_count_settings.py | 7 ++ 9 files changed, 513 insertions(+), 10 deletions(-) diff --git a/litellm/llms/anthropic/count_tokens/transformation.py b/litellm/llms/anthropic/count_tokens/transformation.py index eb8e595c849..0967c22f93a 100644 --- a/litellm/llms/anthropic/count_tokens/transformation.py +++ b/litellm/llms/anthropic/count_tokens/transformation.py @@ -161,11 +161,58 @@ def _has_non_chat_user_block(message: Mapping[str, object]) -> bool: return any(isinstance(block, dict) and block.get("type") not in ValidUserMessageContentTypes for block in blocks) +_OPENAI_RESPONSES_BUILT_IN_TOOL_TYPES: Final = frozenset( + ( + "apply_patch", + "code_interpreter", + "computer", + "computer_use_preview", + "file_search", + "image_generation", + "local_shell", + "mcp", + "shell", + "web_search", + "web_search_2025_08_26", + "web_search_preview", + "web_search_preview_2025_03_11", + ) +) + + +_CHAT_INFERENCE_TRANSLATED_TOOL_TYPES: Final = frozenset(("namespace", "tool_search")) + + +def _is_responses_built_in_tool(tool: Mapping[str, object]) -> bool: + tool_type: Final = tool.get("type") + return isinstance(tool_type, str) and tool_type in _OPENAI_RESPONSES_BUILT_IN_TOOL_TYPES + + +def _is_chat_inference_translated_tool(tool: Mapping[str, object]) -> bool: + tool_type: Final = tool.get("type") + return isinstance(tool_type, str) and tool_type in _CHAT_INFERENCE_TRANSLATED_TOOL_TYPES + + +def _needs_tool_translation(tool: Mapping[str, object]) -> bool: + return _is_chat_tool(tool) or _is_chat_inference_translated_tool(tool) or _is_responses_built_in_tool(tool) + + def _gets_placeholder_name(tool: Mapping[str, object]) -> bool: from litellm.llms.anthropic.messages.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())) + return not ( + _is_chat_tool(tool) + or _is_chat_inference_translated_tool(tool) + or is_provider_native_tool_dict(tool) + or (isinstance(name, str) and name.strip()) + ) + + +def _chat_tools_for(tool: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + from litellm.litellm_core_utils.prompt_templates.factory import chat_completion_tools_for + + return (tool,) if _is_chat_inference_translated_tool(tool) else chat_completion_tools_for(tool) def _untranslatable_input( @@ -182,8 +229,8 @@ def _untranslatable_input( 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" + message=f"Invalid request for token counting: LiteLLM cannot count a tool of type " + f"{unnamed_tool.get('type')!r} for this model" ) return None @@ -223,9 +270,7 @@ def _translated_chat_request( 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))] + chat_tools: Final = [dict(chat_tool) for chat_tool in chain.from_iterable(map(_chat_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 @@ -255,7 +300,7 @@ def anthropic_fields_for_chat_tool_use( counted_tools: Final = tools or () if ( not chat_params - and not any(map(_is_chat_tool, counted_tools)) + and not any(map(_needs_tool_translation, counted_tools)) and not any(map(_needs_chat_translation, messages)) ): return None diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 21fbdb04896..26bc5904e06 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -10,6 +10,7 @@ from typing import Any, Final from pydantic import JsonValue from typing_extensions import assert_never +import litellm 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 @@ -53,6 +54,18 @@ class CountTokensFailure: CountTokensOutcome = CountedTokens | CountTokensFailure +def _claude_count_translator(litellm_params: Mapping[str, object]) -> AnthropicConfig: + deployment_model: Final = litellm_params.get("model") + if not isinstance(deployment_model, str): + return AnthropicConfig() + provider_hint: Final = litellm_params.get("custom_llm_provider") + model, provider, _, _ = litellm.get_llm_provider( + model=deployment_model, custom_llm_provider=provider_hint if isinstance(provider_hint, str) else None + ) + chat_config: Final = litellm.ProviderConfigManager.get_provider_chat_config(model, LlmProviders(provider)) + return chat_config if isinstance(chat_config, AnthropicConfig) else AnthropicConfig() + + def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model: str) -> bool: return ( isinstance(outcome, CountTokensFailure) @@ -138,16 +151,22 @@ class BedrockTokenCounter(BaseTokenCounter): 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 + litellm_params: Final = (deployment or {}).get("litellm_params", {}) translated: Final = ( await anthropic_count_fields( - counted_model, messages, system, tools, messages_format, AnthropicConfig(), chat_params + counted_model, + messages, + system, + tools, + messages_format, + _claude_count_translator(litellm_params), + 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 diff --git a/litellm/proxy/common_utils/token_count_settings.py b/litellm/proxy/common_utils/token_count_settings.py index 4e7dc012802..dd04e531f2d 100644 --- a/litellm/proxy/common_utils/token_count_settings.py +++ b/litellm/proxy/common_utils/token_count_settings.py @@ -73,6 +73,8 @@ class ProviderCountFailure: return not self.rejected_before_provider_call and self.status_code in REGION_DENIED_STATUSES def with_hint(self, hint: str) -> "ProviderCountFailure": + if self.rejected_before_provider_call: + return self return replace(self, message=_with_hint(self.message, hint)) def to_proxy_exception(self) -> ProxyException: 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 27cbbe81a13..380881a6796 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 @@ -37,6 +37,7 @@ from integration.providers._count_tokens_system_lift import ( REMINDER, RESPONSES_REQUEST_OPTIONS, RESPONSES_TOOL_USE, + RESPONSES_TOOLS, STRING_CASE, TOKEN_COUNTING_BETA, TOOLS, @@ -603,7 +604,24 @@ def _generating(provider: _Provider) -> Callable[[Request], Reply]: return respond -@pytest.mark.parametrize("request_body", RESPONSES_TOOL_USE.values(), ids=RESPONSES_TOOL_USE.keys()) +_REWRITTEN_TOOL_USE: Final[Mapping[str, dict[str, JsonValue]]] = MappingProxyType( + { + **RESPONSES_TOOL_USE, + "file-search": { + "input": WEATHER_QUESTION, + "tools": [{"type": "file_search", "vector_store_ids": ["vs_audit_weather"]}], + }, + "additional-tools": { + "input": [ + {"role": "user", "content": WEATHER_QUESTION}, + {"type": "additional_tools", "tools": RESPONSES_TOOLS}, + ], + }, + } +) + + +@pytest.mark.parametrize("request_body", _REWRITTEN_TOOL_USE.values(), ids=_REWRITTEN_TOOL_USE.keys()) def test_responses_input_tokens_counts_what_responses_inference_sends( deployment: _Deployment, gateway: Gateway, request_body: dict[str, JsonValue] ) -> None: @@ -642,6 +660,79 @@ def test_responses_input_tokens_counts_the_request_options_responses_inference_s assert count_request_fields(counted) == inference_count_fields(inference), (counted, inference) +def _without_cache_marks(value: JsonValue) -> JsonValue: + if isinstance(value, dict): + return {key: _without_cache_marks(item) for key, item in value.items() if key != "cache_control"} + if isinstance(value, list): + return [_without_cache_marks(item) for item in value] + return value + + +def test_responses_input_tokens_leaves_out_the_cache_marks_responses_inference_sends( + deployment: _Deployment, gateway: Gateway +) -> None: + marked: Final[dict[str, JsonValue]] = {"type": "ephemeral"} + request_body: Final[dict[str, JsonValue]] = { + "input": [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": f"Weather note {index}.", "cache_control": marked} + for index in range(2) + ], + }, + {"role": "user", "content": WEATHER_QUESTION}, + ], + "tools": RESPONSES_TOOLS, + } + 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(body) for body in _inference_and_count(peer)) + assert "cache_control" in json.dumps(inference), inference + assert "cache_control" not in json.dumps(counted), counted + assert counted == _without_cache_marks(inference), (counted, inference) + + +_UNCOUNTABLE_TOOLS: Final[Mapping[str, dict[str, JsonValue]]] = MappingProxyType( + { + "computer_use_preview": { + "type": "computer_use_preview", + "display_width": 1024, + "display_height": 768, + "environment": "browser", + }, + "mcp": { + "type": "mcp", + "server_label": "docs", + "server_url": "http://127.0.0.1:9/sse", + "require_approval": "never", + }, + } +) + + +@pytest.mark.parametrize("tool", _UNCOUNTABLE_TOOLS.values(), ids=_UNCOUNTABLE_TOOLS.keys()) +def test_responses_input_tokens_rejects_a_tool_it_cannot_count_without_calling_the_peer( + deployment: _Deployment, gateway: Gateway, tool: dict[str, JsonValue] +) -> None: + with _listening(deployment) as peer: + response: Final = gateway.request( + "POST", + "/v1/responses/input_tokens", + {"model": deployment.model, "input": WEATHER_QUESTION, "tools": [tool]}, + ) + assert response.status_code == 400, response.text + assert tool["type"] in response.text, response.text + assert peer.drain() == () + + def _text_blocks(message: JsonValue) -> JsonValue: if not isinstance(message, dict) or not isinstance(message.get("content"), str): return message 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 a4542f1f604..c392e355708 100644 --- a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py @@ -1,4 +1,5 @@ import asyncio +import base64 import json import re import threading @@ -32,6 +33,7 @@ from integration.providers._count_tokens_system_lift import ( LIFTED, MID_SYSTEM, REMINDER, + RESPONSES_REQUEST_OPTIONS, RESPONSES_TOOLS, STRING_CASE, USER, @@ -1702,6 +1704,72 @@ def test_responses_input_tokens_sends_mantle_the_tool_use_invoke_inference_sends assert _tool_use_fields(mantle_count) == expected, mantle_count +_RUNTIME_COUNTED_OPTIONS: Final[Mapping[str, dict[str, JsonValue]]] = MappingProxyType( + {**RESPONSES_REQUEST_OPTIONS, "forced-tool-choice": _FORCED_TOOL_CHOICE} +) +_CONVERSE_WEATHER_TOOLS: Final[dict[str, JsonValue]] = { + "tools": [ + { + "toolSpec": { + "name": "get_weather", + "description": "Look up the current weather for a city", + "inputSchema": {"json": WEATHER_SCHEMA}, + } + } + ] +} +_INVOKE_WEATHER_TOOLS: Final[list[JsonValue]] = [ + { + "type": "custom", + "name": "get_weather", + "description": "Look up the current weather for a city", + "input_schema": WEATHER_SCHEMA, + } +] +_RUNTIME_COUNTED_SHAPES: Final[Mapping[str, tuple[frozenset[str], str, JsonValue]]] = MappingProxyType( + { + "converse": (frozenset({"messages", "system", "toolConfig"}), "toolConfig", _CONVERSE_WEATHER_TOOLS), + "invokeModel": ( + frozenset({"anthropic_version", "max_tokens", "messages", "system", "tools"}), + "tools", + _INVOKE_WEATHER_TOOLS, + ), + } +) + + +def _runtime_counted_request(count_body: Mapping[str, JsonValue]) -> tuple[str, dict[str, JsonValue]]: + count_input: Final = _JSON_OBJECT.validate_python(count_body["input"]) + invoke: Final = count_input.get("invokeModel") + if invoke is None: + return "converse", _JSON_OBJECT.validate_python(count_input["converse"]) + encoded: Final = _JSON_OBJECT.validate_python(invoke)["body"] + assert isinstance(encoded, str), invoke + return "invokeModel", _JSON_OBJECT.validate_json(base64.b64decode(encoded)) + + +@pytest.mark.parametrize("request_body", _RUNTIME_COUNTED_OPTIONS.values(), ids=_RUNTIME_COUNTED_OPTIONS.keys()) +def test_responses_input_tokens_sends_the_runtime_the_tools_without_the_options_it_refuses( + counting_proxy: OwnedProxy, mantle_port: int, request_body: dict[str, JsonValue] +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _metadata_deployment(scenario, runtime.url, _SONNET) + response: Final = gateway.request("POST", "/v1/responses/input_tokens", {**request_body, "model": model}) + assert _payload(response)["input_tokens"] == _RUNTIME_COUNT, response.text + ((target, count_body),) = _runtime_count_requests(runtime) + assert target == f"/model/{_SONNET}/count-tokens", target + shape, counted = _runtime_counted_request(count_body) + allowed_fields, tools_field, weather_tools = _RUNTIME_COUNTED_SHAPES[shape] + assert set(counted) <= allowed_fields, counted + assert counted.get(tools_field) == (weather_tools if "tools" in request_body else None), counted + assert mantle.drain() == () + + def _counting_every_model(request: Request) -> Reply: if _COUNT_TARGET.match(target_of(request)): return _json_reply(200, {"inputTokens": _RUNTIME_COUNT}) 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 96410b03a92..9ecec3686a9 100644 --- a/tests/integration/providers/test_openai_count_tokens_tools_wire.py +++ b/tests/integration/providers/test_openai_count_tokens_tools_wire.py @@ -1,5 +1,7 @@ import json +import uuid from collections.abc import Iterator, Mapping +from dataclasses import dataclass from typing import Final import pytest @@ -488,3 +490,103 @@ def test_responses_input_tokens_sends_openai_the_owners_previous_response_id(cou 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"] + + +_PROVIDER_FILE_ID: Final = "file-audit-provider" +_UPLOADED_FILE: Final[Mapping[str, JsonValue]] = { + "id": _PROVIDER_FILE_ID, + "object": "file", + "bytes": 5, + "created_at": 1, + "filename": "notes.txt", + "purpose": "user_data", + "status": "processed", +} + + +def _files_peer(request: Request) -> Reply: + if request.target == "/v1/files": + return Reply(body=json.dumps(_UPLOADED_FILE).encode()) + return _count_response(request) + + +def _file_input(file_id: str) -> list[JsonValue]: + return [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "Summarize the notes."}, + {"type": "input_file", "file_id": file_id}, + ], + } + ] + + +@dataclass(frozen=True, slots=True) +class _FileTenants: + model_id: str + model: str + owner_key: str + outsider_key: str + file_id: str + + +def _file_tenants(scenario: Scenario, wire: Wire) -> _FileTenants: + model_id: Final = f"audit-openai-{uuid.uuid4().hex}" + model: Final = scenario.model( + model_info={"id": model_id}, + model="openai/gpt-4o-mini", + api_base=f"{wire.url}/v1", + api_key="sk-integration-openai", + ) + owner_key: Final = scenario.key( + user_id=scenario.user(user_role="internal_user"), team_id=scenario.team(models=[model]), models=[model] + ) + outsider_key: Final = scenario.key( + user_id=scenario.user(user_role="internal_user"), team_id=scenario.team(models=[model]), models=[model] + ) + uploaded: Final = scenario.gateway.request_multipart( + "/v1/files", + {"purpose": "user_data", "target_model_names": model}, + {"file": ("notes.txt", b"notes", "text/plain")}, + key=owner_key, + ) + assert uploaded.status_code == 200, uploaded.text + file_id: Final = _JSON_OBJECT.validate_json(uploaded.content)["id"] + assert isinstance(file_id, str) and file_id != _PROVIDER_FILE_ID, uploaded.text + assert [request.target for request in wire.drain()] == ["/v1/files"] + return _FileTenants(model_id, model, owner_key, outsider_key, file_id) + + +def test_responses_input_tokens_sends_openai_the_provider_id_of_a_managed_file(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_files_peer) as wire, gateway.scenario() as scenario: + tenants: Final = _file_tenants(scenario, wire) + response: Final = gateway.request( + "POST", + "/v1/responses/input_tokens", + { + "model": tenants.model, + "input": _file_input(tenants.file_id), + "model_file_id_mapping": {tenants.file_id: {tenants.model_id: "file-forged"}}, + }, + key=tenants.owner_key, + ) + 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": _file_input(_PROVIDER_FILE_ID)} + + +def test_responses_input_tokens_refuses_another_teams_managed_file_before_openai(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_files_peer) as wire, gateway.scenario() as scenario: + tenants: Final = _file_tenants(scenario, wire) + refused: Final = gateway.request( + "POST", + "/v1/responses/input_tokens", + {"model": tenants.model, "input": _file_input(tenants.file_id)}, + key=tenants.outsider_key, + ) + assert refused.status_code == 403, refused.text + assert _JSON_OBJECT.validate_json(refused.content)["error"]["type"] == "permission_error", refused.text + assert wire.drain() == () 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 63f149fabe4..89352ce8702 100644 --- a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py +++ b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py @@ -569,6 +569,7 @@ _ANTHROPIC_TOOL_USE: Final[list[dict[str, JsonValue]]] = [ ("messages", "tool", "messages_format"), ( (_ANTHROPIC_TOOL_USE, {"name": "get_weather", "input_schema": _WEATHER_SCHEMA}, "openai"), + ([{"role": "user", "content": "Open example.com"}], {"type": "computer_toolset_20260801"}, "openai"), ( [{"role": "user", "content": "Weather in SF?"}], {"type": "function", "name": "get_weather", "parameters": _WEATHER_SCHEMA}, @@ -685,6 +686,26 @@ _CHAT_WEATHER_TOOL: Final[dict[str, JsonValue]] = { [{"type": "mcp", "server_label": "wiki", "server_url": "https://mcp.example/mcp"}, _CHAT_WEATHER_TOOL], id="mcp-tool", ), + pytest.param( + [{"role": "user", "content": "Open the docs"}], + [{"type": "computer_use_preview", "display_width": 1024, "display_height": 768, "environment": "browser"}], + id="computer-use-preview-tool-alone", + ), + pytest.param( + [{"role": "user", "content": "Ask the wiki"}], + [{"type": "mcp", "server_label": "wiki", "server_url": "https://mcp.example/mcp"}], + id="mcp-tool-alone", + ), + pytest.param( + [{"role": "user", "content": "Plot the sales data"}], + [{"type": "code_interpreter", "container": {"type": "auto"}}], + id="code-interpreter-tool-alone", + ), + pytest.param( + [{"role": "user", "content": "Fix the failing test"}], + [{"type": "apply_patch"}], + id="apply-patch-tool-alone", + ), ), ) async def test_counter_rejects_input_chat_translation_would_leave_out_without_calling_anthropic( @@ -700,6 +721,80 @@ async def test_counter_rejects_input_chat_translation_would_leave_out_without_ca assert (result.error, result.status_code, result.rejected_before_provider_call) == (True, 400, True) +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport_clients") +@_COUNTERS +async def test_counter_sends_a_tool_with_a_non_string_type_to_anthropic_as_given( + counter_type: type[AnthropicTokenCounter] | type[AzureAIAnthropicTokenCounter], count_endpoint: str +) -> None: + tool: Final[dict[str, JsonValue]] = {"type": ["mcp"], "name": "lookup"} + + result, bodies = await _counted_body( + counter_type, count_endpoint, [{"role": "user", "content": "Ask the wiki"}], [tool], "openai" + ) + + assert result is not None + assert [body.get("tools") for body in bodies] == [[tool]] + + +_CRM_NAMESPACE_TOOL: Final[dict[str, JsonValue]] = { + "type": "namespace", + "name": "crm", + "description": "CRM lookups", + "tools": [ + { + "type": "function", + "name": "find_owner", + "description": "Find the owner of an account", + "parameters": {"type": "object", "properties": {"account": {"type": "string"}}, "required": ["account"]}, + } + ], +} + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport_clients") +@_COUNTERS +@pytest.mark.parametrize( + "tools", + ( + pytest.param([{"type": "tool_search"}], id="tool-search-alone"), + pytest.param( + [ + {"type": "tool_search"}, + { + "type": "function", + "function": {"name": "get_weather", "parameters": _WEATHER_SCHEMA}, + "defer_loading": True, + }, + ], + id="tool-search-with-deferred-function", + ), + pytest.param([_CRM_NAMESPACE_TOOL], id="namespace-alone"), + pytest.param([_CHAT_WEATHER_TOOL, _CRM_NAMESPACE_TOOL], id="namespace-with-function"), + ), +) +async def test_counter_sends_the_tools_chat_inference_makes_from_tool_search_and_namespace( + counter_type: type[AnthropicTokenCounter] | type[AzureAIAnthropicTokenCounter], + count_endpoint: str, + tools: list[dict[str, JsonValue]], +) -> None: + messages: Final[list[dict[str, JsonValue]]] = [{"role": "user", "content": "Who owns the north region account?"}] + config: Final = AnthropicConfig() + inference: Final = config.transform_request( + model="claude-opus-5-5", + messages=deepcopy(messages), + optional_params=config.map_openai_params({"tools": deepcopy(tools)}, {}, "claude-opus-5-5", drop_params=False), + litellm_params={}, + headers={}, + ) + + result, bodies = await _counted_body(counter_type, count_endpoint, messages, tools, "openai") + + assert result is not None + assert [body.get("tools") for body in bodies] == [inference["tools"]] + + _ANSWER_SCHEMA: Final[dict[str, JsonValue]] = { "type": "object", "properties": {"answer": {"type": "string"}}, 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 6fb47491323..cec7950b859 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 @@ -6,6 +6,7 @@ import httpx import pytest from pydantic import JsonValue, TypeAdapter +import litellm from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -264,6 +265,79 @@ async def test_claude_chat_tool_use_is_counted_in_the_shape_claude_inference_sen {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_1", "content": "sunny"}]}, ] + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tool", + ( + {"type": "mcp", "server_label": "wiki", "server_url": "https://mcp.example/mcp"}, + {"type": "computer_use_preview", "display_width": 1024, "display_height": 768, "environment": "browser"}, + ), + ids=("mcp", "computer_use_preview"), +) +async def test_claude_count_rejects_a_lone_responses_tool_without_calling_bedrock(tool: dict[str, JsonValue]) -> None: + bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 412}), mantle=(200, {"input_tokens": 2177})) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="global.anthropic.claude-opus-4-8", + messages=[{"role": "user", "content": "Ask the wiki"}], + contents=None, + deployment=DEPLOYMENT, + request_model="claude-opus-4-8", + tools=[tool], + messages_format="openai", + ) + + assert result is not None + assert (result.error, result.status_code, result.rejected_before_provider_call) == (True, 400, True) + assert bedrock.posted_hosts() == () + + +_CLAUDE_REPLY: Final = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 3, "output_tokens": 1}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "deployment_model", + ("bedrock/invoke/us.anthropic.claude-sonnet-5-5", "bedrock_mantle/us.anthropic.claude-sonnet-5-5"), +) +async def test_claude_count_sends_the_tool_search_tool_the_deployments_inference_sends(deployment_model: str) -> None: + credentials: Final = DEPLOYMENT["litellm_params"] + inference: Final = _Bedrock(runtime=(200, _CLAUDE_REPLY), mantle=(200, _CLAUDE_REPLY)) + await litellm.acompletion( + model=deployment_model, + messages=MESSAGES, + tools=[{"type": "tool_search"}], + max_tokens=8, + client=AsyncHTTPHandler(transport=httpx.MockTransport(inference)), + **credentials, + ) + bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 412}), mantle=(200, {"input_tokens": 2177})) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="bedrock_mantle/us.anthropic.claude-opus-5-5", + messages=MESSAGES, + contents=None, + deployment={"litellm_params": {**credentials, "model": deployment_model}}, + request_model="claude-sonnet-5-5", + tools=[{"type": "tool_search"}], + messages_format="openai", + ) + + assert result is not None + assert (result.error, result.total_tokens) == (False, 2177) + counted_tools: Final = _JSON_BODY.validate_json(bedrock.requests[0].content)["tools"] + assert counted_tools == _JSON_BODY.validate_json(inference.requests[0].content)["tools"] + + @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})) diff --git a/tests/unit/proxy/common_utils/test_token_count_settings.py b/tests/unit/proxy/common_utils/test_token_count_settings.py index 7ff32f9d4f3..43edc7e449f 100644 --- a/tests/unit/proxy/common_utils/test_token_count_settings.py +++ b/tests/unit/proxy/common_utils/test_token_count_settings.py @@ -229,3 +229,10 @@ def test_hint_lands_in_the_upstream_error_message_and_keeps_its_envelope(message assert (json.loads(hinted.message) if isinstance(expected, dict) else hinted.message) == expected assert hinted.status_code == 404 + + +def test_hint_stays_off_an_error_raised_before_any_provider_call() -> None: + message: Final = "Invalid request for token counting: LiteLLM cannot count a tool of type 'mcp' for this model" + failure: Final = ProviderCountFailure(status_code=400, message=message, rejected_before_provider_call=True) + + assert failure.with_hint("Enable eu-west-1").message == message From ca4a2ba006839a00187149c438207cc045c42923 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 28 Oct 2026 12:13:56 -0700 Subject: [PATCH 069/128] fix(router): drop the deployment rpm filter from enable_pre_call_checks (#48873) enable_pre_call_checks filtered deployments whose in-memory minute count reached litellm_params.rpm, undocumented since 0d1cca9aa0. Deployment rpm and tpm stay routing weights unless optional_pre_call_checks includes enforce_model_rate_limits, which is now the only deployment rate-limit path. Resolves LIT-10233 Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 41 +----- .../routing/test_end_user_region_routing.py | 29 +++++ tests/unit/test_router/test_router.py | 123 +++++++++++------- 3 files changed, 107 insertions(+), 86 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index f0bb6c3a7a9..ef39f5bd983 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -14439,7 +14439,6 @@ class Router: Filter out model in model group, if: - model context window < message length. For azure openai models, requires 'base_model' is set. - https://docs.litellm.ai/docs/proxy/cost_tracking#spend-tracking-for-azure-openai-models - - filter models above rpm limits - if region given, filter out models not in that region / unknown region - [TODO] function call and model doesn't support function calling """ @@ -14454,8 +14453,6 @@ class Router: _context_window_error = False _potential_error_str = "" - _rate_limit_error = False - parent_otel_span: Final = get_parent_otel_span_from_kwargs(request_kwargs) has_countable_input: Final = (messages is not None or input is not None) and not compaction_pending( request_kwargs @@ -14479,13 +14476,6 @@ class Router: else frozenset() ) - ## get model group RPM ## - dt: Final = get_utc_datetime() - current_minute: Final = dt.strftime("%H-%M") - rpm_key: Final = f"{model}:rpm:{current_minute}" - model_group_cache: Final = ( - self.cache.get_cache(key=rpm_key, local_only=True, parent_otel_span=parent_otel_span) or {} - ) # check the in-memory cache used by lowest_latency and usage-based routing. Only check the local cache. for idx, deployment in enumerate(_returned_deployments): # Cache nested dict access to avoid repeated temporary dict allocations _litellm_params = deployment.get("litellm_params", {}) @@ -14508,24 +14498,6 @@ class Router: ) continue - model_id = _model_info.get("id", "") - ## RPM CHECK ## - ### get local router cache ### - current_request_cache_local = ( - self.cache.get_cache(key=model_id, local_only=True, parent_otel_span=parent_otel_span) or 0 - ) - ### get usage based cache ### - if isinstance(model_group_cache, dict) and self.routing_strategy != "usage-based-routing-v2": - model_group_cache[model_id] = model_group_cache.get(model_id, 0) - - current_request = max(current_request_cache_local, model_group_cache[model_id]) - - if isinstance(_litellm_params, dict) and _litellm_params.get("rpm", None) is not None: - if isinstance(_litellm_params["rpm"], int) and _litellm_params["rpm"] <= current_request: - invalid_model_indices.add(idx) - _rate_limit_error = True - continue - ## REGION CHECK ## if request_kwargs is not None and request_kwargs.get("allowed_model_region") is not None: allowed_model_region = request_kwargs.get("allowed_model_region") @@ -14562,12 +14534,6 @@ class Router: ) if len(excluded_indices) == len(_returned_deployments): - """ - - no healthy deployments available b/c context window checks or rate limit error - - - First check for rate limit errors (if this is true, it means the model passed the context window check but failed the rate limit check) - """ - group_deployments: Final = ( tuple(request_deployments) if request_deployments is not None @@ -14589,12 +14555,7 @@ class Router: ), model=model, ) - if _rate_limit_error is True: # allow generic fallback logic to take place - raise RouterRateLimitErrorBasic( - model=model, - ) - - elif _context_window_error is True: + if _context_window_error is True: raise litellm.ContextWindowExceededError( message=f"litellm._pre_call_checks: Context Window exceeded for given call. No models have context window large enough for this call.\n{_potential_error_str}", model=model, diff --git a/tests/integration/routing/test_end_user_region_routing.py b/tests/integration/routing/test_end_user_region_routing.py index ce90280f173..884b61228cc 100644 --- a/tests/integration/routing/test_end_user_region_routing.py +++ b/tests/integration/routing/test_end_user_region_routing.py @@ -19,6 +19,9 @@ UNCOUNTABLE_MODEL: Final = "regional-uncountable-model" PINNED_REGION: Final = "eu" WEIGHT_BY_REGION: Final = {"eu": 1, "us": 99} SILENT_WAV: Final = "UklGRiQAAABXQVZFZm10IBAAAAABAAEAQB8AAIA+AAACABAAZGF0YQAAAAA=" +RPM_CAPPED_MODEL: Final = "regional-rpm-capped-model" +RPM_CAPPED_DEPLOYMENTS: Final = frozenset(f"{RPM_CAPPED_MODEL}-{region}" for region in UPSTREAM_BY_REGION) +RPM_CAPPED_CALLS: Final = 3 @pytest.fixture(scope="module") @@ -53,6 +56,19 @@ def candidate(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: } for region, upstream in UPSTREAM_BY_REGION.items() ), + *( + { + "model_name": RPM_CAPPED_MODEL, + "litellm_params": { + "model": f"openai/{upstream}", + "api_base": f"{base.upstream_url}/v1", + "api_key": "synthetic-region-key", + "rpm": 1, + }, + "model_info": {"id": f"{RPM_CAPPED_MODEL}-{region}"}, + } + for region, upstream in UPSTREAM_BY_REGION.items() + ), ] config["router_settings"] = {**config["router_settings"], "enable_pre_call_checks": True} path: Final = directory / "region-routing.yaml" @@ -143,3 +159,16 @@ def test_an_uncountable_prompt_still_pins_an_end_user_to_their_region( observed: Final = upstream.get("/__observations").json()["requests"] pinned_upstream: Final = UPSTREAM_BY_REGION[PINNED_REGION] assert [request["body"]["model"] for request in observed] == [pinned_upstream] * recorded_chat_calls + + +def test_pre_call_checks_keep_serving_deployments_past_their_rpm(candidate: Gateway) -> None: + with candidate.scenario() as scenario: + key: Final = scenario.key(models=[RPM_CAPPED_MODEL]) + body: Final = {"model": RPM_CAPPED_MODEL, "messages": [{"role": "user", "content": "integration control"}]} + responses: Final = tuple( + candidate.request("POST", "/v1/chat/completions", body, key=key) for _ in range(RPM_CAPPED_CALLS) + ) + assert [response.status_code for response in responses] == [200] * RPM_CAPPED_CALLS, [ + response.text for response in responses + ] + assert {response.headers.get("x-litellm-model-id") for response in responses} <= RPM_CAPPED_DEPLOYMENTS diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index c6d02dddae0..4cb994744f2 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -27633,38 +27633,89 @@ def test_router_get_model_group_info() -> None: assert model_group_info.output_cost_per_token > 0 -def test_router_rpm_pre_call_check(): - """ - - for a given model not in model cost map - - with rpm set - - check if rpm check is run - """ - try: - model_list = [ +def test_pre_call_checks_keep_deployment_whose_rpm_is_reached(): + model_list: Final = [ + { + "model_name": "fake-openai-endpoint", + "litellm_params": {"model": "openai/my-fake-model", "api_key": "my-fake-key", "rpm": 0}, + "model_info": {"id": "rpm-zero"}, + }, + { + "model_name": "fake-openai-endpoint", + "litellm_params": {"model": "openai/my-fake-model", "api_key": "my-fake-key"}, + "model_info": {"id": "uncapped"}, + }, + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + + kept: Final = router._pre_call_checks( + model="fake-openai-endpoint", + healthy_deployments=model_list, + messages=[{"role": "user", "content": "Hey, how's it going?"}], + ) + + assert [deployment["model_info"]["id"] for deployment in kept] == ["rpm-zero", "uncapped"] + + +@pytest.mark.asyncio +async def test_enable_pre_call_checks_alone_serves_past_deployment_rpm(): + router: Final = Router( + model_list=[ { - "model_name": "fake-openai-endpoint", # openai model name - "litellm_params": { # params for litellm completion/embedding call - "model": "openai/my-fake-model", - "api_key": "my-fake-key", - "api_base": "https://openai-function-calling-workers.tasslexyz.workers.dev/", - "rpm": 0, - }, - }, - ] + "model_name": "fake-openai-endpoint", + "litellm_params": {"model": "openai/my-fake-model", "api_key": "my-fake-key", "rpm": 1}, + "model_info": {"id": "rpm-one"}, + } + ], + enable_pre_call_checks=True, + num_retries=0, + ) - router = Router(model_list=model_list, set_verbose=True, enable_pre_call_checks=True, num_retries=0) # type: ignore + served: Final = [ + await router.acompletion( + model="fake-openai-endpoint", + messages=[{"role": "user", "content": "Hey, how's it going?"}], + mock_response="served", + ) + for _ in range(3) + ] - try: - router._pre_call_checks( + assert [response.choices[0].message.content for response in served] == ["served"] * 3 + + +@pytest.mark.asyncio +async def test_enforce_model_rate_limits_returns_429_at_deployment_rpm(): + frozen_minute: Final = datetime(2026, 1, 1, 12, 30, 15, tzinfo=timezone.utc) + router: Final = Router( + model_list=[ + { + "model_name": "fake-openai-endpoint", + "litellm_params": {"model": "openai/my-fake-model", "api_key": "my-fake-key", "rpm": 1}, + "model_info": {"id": "rpm-one"}, + } + ], + enable_pre_call_checks=True, + optional_pre_call_checks=["enforce_model_rate_limits"], + num_retries=0, + ) + + with patch( + "litellm.router_utils.pre_call_checks.model_rate_limit_check.get_utc_datetime", return_value=frozen_minute + ): + first: Final = await router.acompletion( + model="fake-openai-endpoint", + messages=[{"role": "user", "content": "Hey, how's it going?"}], + mock_response="served", + ) + with pytest.raises(litellm.RateLimitError) as rate_limited: + await router.acompletion( model="fake-openai-endpoint", - healthy_deployments=model_list, messages=[{"role": "user", "content": "Hey, how's it going?"}], + mock_response="served", ) - pytest.fail("Expected this to fail") - except Exception: - pass - except Exception as e: - pytest.fail(f"Got unexpected exception on router! - {str(e)}") + + assert first.choices[0].message.content == "served" + assert rate_limited.value.status_code == 429 @pytest.mark.parametrize( @@ -29494,26 +29545,6 @@ def test_pre_call_checks_filters_deployment_when_token_count_exceeds_limit(): assert [d["model_info"]["id"] for d in filtered] == ["large"] -def test_pre_call_checks_rpm_check_still_enforced_when_token_counting_fails(): - from litellm.types.router import RouterRateLimitErrorBasic - - model_list: Final = [ - { - "model_name": "fake-openai-endpoint", - "litellm_params": {"model": "openai/my-fake-model", "api_key": "my-fake-key", "rpm": 0}, - "model_info": {"max_input_tokens": 100}, - }, - ] - router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) - - with pytest.raises(RouterRateLimitErrorBasic): - router._pre_call_checks( - model="fake-openai-endpoint", - healthy_deployments=model_list, - messages=_PRE_CALL_CHECK_UNPARSEABLE_MESSAGES, - ) - - def test_pre_call_checks_context_window_check_is_skipped_but_not_fatal_when_token_counting_fails(): model_list: Final = [ { From 78d6c1aa823340f97a92619de6d42e90615a71b9 Mon Sep 17 00:00:00 2001 From: simonp9 <349053457+simonp9@users.noreply.github.com> Date: Wed, 28 Oct 2026 12:15:02 -0700 Subject: [PATCH 070/128] fix(router): pass lifecycle events through on background OpenAI responses streams that cannot fall back (#48817) * fix(router): pass lifecycle events through on background OpenAI responses streams that cannot fall back The Router held response.created, response.queued and response.in_progress until the first output event on every /v1/responses stream, so a background run showed the client no response id while it was queued or reasoning. A native OpenAI stream whose upstream echoes background: true now releases those events as they arrive when no fallback, order level or weighted failover could replace it, and once the id is out a later pre-output error is raised instead of starting a second run behind it. * test(router): type background stream test params and wait on the upstream pause instead of a clock * fix(router): keep the content-policy note on a released background stream error * fix(responses): return the routed id from cancel and add the router note only before output acancel_responses and cancel_responses now re-encode the reply id with the deployment, like get. The released stream's content-policy note follows the held path, which adds it only before output text. Renames _anthropic_messages_order_levels to _stream_order_levels. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/responses/main.py | 18 +- litellm/responses/utils.py | 8 +- litellm/router.py | 83 ++++- litellm/router_utils/common_utils.py | 37 ++ .../router_code_coverage.py | 4 + .../responses/test_responses_api_lifecycle.py | 44 ++- .../test_router_utils_common_utils.py | 9 + tests/unit/test_router/test_router.py | 317 +++++++++++++++++- 8 files changed, 497 insertions(+), 23 deletions(-) diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 8451834eab2..d474ed3991f 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -2341,6 +2341,7 @@ async def acancel_responses( timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, + litellm_metadata: Mapping[str, object] | None = None, **kwargs, ) -> ResponsesAPIResponse: """ @@ -2365,6 +2366,7 @@ async def acancel_responses( cancel_responses, response_id=response_id, custom_llm_provider=custom_llm_provider, + litellm_metadata=litellm_metadata, extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, @@ -2380,7 +2382,12 @@ async def acancel_responses( response = await init_response else: response = init_response - return response + + return ResponsesAPIRequestUtils.update_responses_api_response_id_with_model_id( + responses_api_response=response, + litellm_metadata=litellm_metadata, + custom_llm_provider=custom_llm_provider, + ) except Exception as e: raise litellm.exception_type( model=None, @@ -2402,6 +2409,7 @@ def cancel_responses( timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, + litellm_metadata: Mapping[str, object] | None = None, **kwargs, ) -> ResponsesAPIResponse | Coroutine[object, object, ResponsesAPIResponse]: """ @@ -2471,7 +2479,13 @@ def cancel_responses( shared_session=kwargs.get("shared_session"), ) - return response + if not isinstance(response, ResponsesAPIResponse): + return response + return ResponsesAPIRequestUtils.update_responses_api_response_id_with_model_id( + responses_api_response=response, + litellm_metadata=litellm_metadata, + custom_llm_provider=custom_llm_provider, + ) except Exception as e: raise litellm.exception_type( model=None, diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index a6b042a2234..02d07932533 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -399,7 +399,7 @@ class ResponsesAPIRequestUtils: def update_responses_api_response_id_with_model_id( responses_api_response: ResponsesAPIResponse, custom_llm_provider: str | None, - litellm_metadata: dict[str, object] | None = None, + litellm_metadata: Mapping[str, object] | None = None, ) -> ResponsesAPIResponse: ... @@ -408,7 +408,7 @@ class ResponsesAPIRequestUtils: def update_responses_api_response_id_with_model_id( responses_api_response: dict[str, object], custom_llm_provider: str | None, - litellm_metadata: dict[str, object] | None = None, + litellm_metadata: Mapping[str, object] | None = None, ) -> dict[str, object]: ... @@ -418,7 +418,7 @@ class ResponsesAPIRequestUtils: def update_responses_api_response_id_with_model_id( responses_api_response: ResponsesAPIResponse | dict[str, Any], custom_llm_provider: str | None, - litellm_metadata: dict[str, Any] | None = None, + litellm_metadata: Mapping[str, Any] | None = None, ) -> ResponsesAPIResponse | dict[str, object]: """Update the responses_api_response_id with model_id and custom_llm_provider. @@ -1110,7 +1110,7 @@ class ResponsesAPIRequestUtils: def _update_container_ids_in_response( responses_api_response: ResponsesAPIResponse | dict[str, object], custom_llm_provider: str | None, - litellm_metadata: dict[str, Any] | None = None, + litellm_metadata: Mapping[str, Any] | None = None, ) -> ResponsesAPIResponse | dict[str, object]: """Encode container IDs in the response output with provider/model info. diff --git a/litellm/router.py b/litellm/router.py index ef39f5bd983..d5cad013e8d 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -216,6 +216,8 @@ from litellm.router_utils.common_utils import ( deployments_the_caller_may_use, filter_team_based_models, filter_web_search_deployments, + format_content_policy_fallback_hint, + format_content_policy_no_fallback_note, format_fallback_outcome_message, format_no_fallback_group_message, get_request_team_id, @@ -832,6 +834,10 @@ def _request_model_group(kwargs: Mapping[str, object]) -> str | None: return model_group if isinstance(model_group, str) else None +def _mid_stream_error_before_output(error: "MidStreamFallbackError") -> bool: + return not error.generated_content or error.is_pre_first_chunk + + def _mid_stream_retry_trigger(error: "MidStreamFallbackError") -> Exception: """The provider's own error, which is what the retry policy and should_retry_this_error classify.""" return error.original_exception if error.original_exception is not None else error @@ -962,6 +968,10 @@ def _responses_stream_holds_event(item: object, held_event_count: int) -> bool: return getattr(item, "type", None) in PRE_OUTPUT_LIFECYCLE_EVENT_TYPES +def _responses_event_announces_background_run(item: object) -> bool: + return getattr(getattr(item, "response", None), "background", None) is True + + class FallbackAwareAnthropicMessagesStream: """ Bare async generators can't carry the `_hidden_params` attribute the @@ -4127,11 +4137,16 @@ class Router: async def stream_with_fallbacks(): held_lifecycle_events: tuple[object, ...] = () # rebind-ok: flushed at first output, dropped on fallback + response_id_exposed = False # rebind-ok: set once, before an early-released lifecycle event is yielded try: async for item in cast( # cast-ok: response iterators are async streams AsyncIterator[object], source_iterator ): - if _responses_stream_holds_event(item, len(held_lifecycle_events)): + if not response_id_exposed and self._aresponses_background_stream_can_release( + item, source_iterator, initial_kwargs + ): + response_id_exposed = True + if not response_id_exposed and _responses_stream_holds_event(item, len(held_lifecycle_events)): held_lifecycle_events = (*held_lifecycle_events, item) continue for held_event in held_lifecycle_events: @@ -4141,6 +4156,8 @@ class Router: for held_event in held_lifecycle_events: yield held_event except MidStreamFallbackError as e: + if response_id_exposed: + self._raise_released_stream_error(e, initial_kwargs) async with contextlib.aclosing( self._aresponses_fallback_attempt( e, source_iterator, initial_kwargs, wrapper.adopt_fallback_headers, held_lifecycle_events @@ -4166,6 +4183,54 @@ class Router: wrapper: Final = FallbackResponsesStreamWrapper(stream_with_fallbacks()) return wrapper + def _aresponses_background_stream_can_release( + self, + item: object, + source_iterator: "BaseResponsesAPIStreamingIterator", + kwargs: Mapping[str, object], + ) -> bool: + from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator + + if ( + type(source_iterator) is not ResponsesAPIStreamingIterator + or source_iterator.custom_llm_provider != "openai" + ): + return False + if not _responses_event_announces_background_run(item): + return False + model_group: Final = kwargs.get("model") + return isinstance(model_group, str) and not self._stream_can_fall_back(model_group, kwargs) + + def _raise_released_stream_error( + self, e: "MidStreamFallbackError", initial_kwargs: Mapping[str, object] + ) -> NoReturn: + original: Final = e.original_exception + model_group: Final = initial_kwargs.get("model") + if ( + litellm.expose_router_debug_in_errors + and isinstance(original, litellm.ContentPolicyViolationError) + and isinstance(model_group, str) + and not fallbacks_disabled_for_request(initial_kwargs) + and _mid_stream_error_before_output(e) + ): + hop_depth: Final = initial_kwargs.get("fallback_depth") + original.message += format_content_policy_no_fallback_note( + model_group, + fallback_lookup_groups(initial_kwargs, model_group), + _request_fallback_list( + initial_kwargs, + "content_policy_fallbacks", + cast("list[object] | None", self.content_policy_fallbacks), # cast-ok: untyped router attribute + ), + _request_fallback_list( + initial_kwargs, + "fallbacks", + cast("list[object] | None", self.fallbacks), # cast-ok: untyped router attribute + ), + isinstance(hop_depth, int) and hop_depth > 0, + ) + Router._raise_original_mid_stream_error(e) + def _responses_continuation( self, e: "MidStreamFallbackError", @@ -4175,7 +4240,7 @@ class Router: """What a fallback continues from once a Responses stream broke after text, raising the primary's error when the request constrains its output or the stream began a tool call, since a continued JSON document or tool call would not parse, and when a content filter cut it, so no other model finishes filtered text.""" - if not e.generated_content or e.is_pre_first_chunk: + if _mid_stream_error_before_output(e): return None if source_iterator.started_tool_call or isinstance(e.original_exception, litellm.ContentPolicyViolationError): Router._raise_original_mid_stream_error(e) @@ -6487,7 +6552,7 @@ class Router: model: Final = cast(str, initial_kwargs.get("model")) # cast-ok: kwargs always carries the model group has_generated_content = not ( # rebind-ok: set once real content is seen, the buffer cap is hit, or neither a retry nor a fallback can take over self._anthropic_messages_stream_can_retry(initial_kwargs) - or self._anthropic_messages_stream_can_fall_back(model, initial_kwargs) + or self._stream_can_fall_back(model, initial_kwargs) ) buffered_lifecycle_chunks: tuple[bytes, ...] = () # rebind-ok: flushed once committed or on decline try: @@ -6710,7 +6775,7 @@ class Router: budget, _ = self._stream_retry_budget(frame_error, kwargs) if budget > attempted_retries_for_request(kwargs): return frame_error - if self._anthropic_messages_stream_can_fall_back(model_group, kwargs): + if self._stream_can_fall_back(model_group, kwargs): return frame_error return None @@ -8725,7 +8790,9 @@ class Router: ) return response else: - error_message = f"model={model_group}. content_policy_fallback={mask_sensitive_structure(content_policy_fallbacks)}. fallbacks={mask_sensitive_structure(fallbacks)}.\n\nSet 'content_policy_fallback' - https://docs.litellm.ai/docs/routing#fallbacks" + error_message = format_content_policy_fallback_hint( + model_group, content_policy_fallbacks, fallbacks + ) verbose_router_logger.info( msg=f"Got 'ContentPolicyViolationError'. No content_policy_fallback set. Defaulting \ to fallbacks, if available.{error_message}" @@ -9873,7 +9940,7 @@ class Router: ) return has_unattempted_fallback_target(resolved, kwargs) - def _anthropic_messages_order_levels(self, model_group: str, kwargs: Mapping[str, Any]) -> tuple[int, ...]: + def _stream_order_levels(self, model_group: str, kwargs: Mapping[str, Any]) -> tuple[int, ...]: """ The distinct deployment order levels the fallback dispatcher would see for this request, computed the same way: the tier a pre-routing hook selected wins over the requested group. @@ -9891,7 +9958,7 @@ class Router: ) ) - def _anthropic_messages_stream_can_fall_back(self, model_group: str, kwargs: Mapping[str, Any]) -> bool: + def _stream_can_fall_back(self, model_group: str, kwargs: Mapping[str, Any]) -> bool: """ Whether async_function_with_fallbacks_common_utils could still route a MidStreamFallbackError somewhere for this request (order levels, weighted @@ -9903,7 +9970,7 @@ class Router: return False if self.enable_weighted_failover: return True - order_levels: Final = self._anthropic_messages_order_levels(model_group, kwargs) + order_levels: Final = self._stream_order_levels(model_group, kwargs) if len(order_levels) > 1: current_target: Final = kwargs.get("_target_order") skip_up_to: Final = current_target if current_target is not None else order_levels[0] diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 91b9e5cb114..2451422b5e6 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -39,6 +39,9 @@ _DEPLOYMENT_ROWS_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = tuple[Mapping[str, object], ...] ) _DEPLOYMENT_IDS_ADAPTER: Final = TypeAdapter(tuple[str, ...]) +_FALLBACK_ENTRIES_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter( + tuple[Mapping[str, object], ...] +) def with_admin_api_key_secret_reference(deployment: Deployment) -> Deployment: @@ -256,6 +259,40 @@ def format_fallback_outcome_message( return f"{lead} Fallback model group(s) configured: {targets}." +def format_content_policy_fallback_hint( + model_group: str | None, content_policy_fallbacks: object, fallbacks: object +) -> str: + return ( + f"model={model_group}. content_policy_fallback={mask_sensitive_structure(content_policy_fallbacks)}. " + f"fallbacks={mask_sensitive_structure(fallbacks)}.\n\nSet 'content_policy_fallback' - " + "https://docs.litellm.ai/docs/routing#fallbacks" + ) + + +def _as_fallback_entries(fallbacks: object) -> tuple[Mapping[str, object], ...] | None: + try: + return _FALLBACK_ENTRIES_ADAPTER.validate_python(fallbacks) + except ValidationError: + return None + + +def format_content_policy_no_fallback_note( + model_group: str, + lookup_groups: tuple[str, ...], + content_policy_fallbacks: object, + fallbacks: object, + nested_fallback_hop: bool, +) -> str: + outcome: Final = "" if nested_fallback_hop else format_fallback_outcome_message(model_group, None, "") + if content_policy_fallbacks is not None: + return outcome + hint: Final = f"\n{format_content_policy_fallback_hint(model_group, content_policy_fallbacks, fallbacks)}" + fallback_entries: Final = _as_fallback_entries(fallbacks) if fallbacks is not None else None + if fallback_entries is None or not lookup_groups or nested_fallback_hop: + return hint + outcome + return hint + format_no_fallback_group_message(lookup_groups, fallback_entries) + + def get_litellm_params_sensitive_credential_hash(litellm_params: dict) -> str: """ Hash of the credential params, used for mapping the file id to the right model diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index f4a295e697d..3183b9687ca 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -174,6 +174,10 @@ ignored_function_names = [ "_responses_continuation", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py "_responses_continuation_ceilings", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py "_team_deployments_named", # Tested through the team lookups in the concurrent delete tests in test_router.py + "_responses_event_announces_background_run", # Tested through the background Responses stream tests in test_router.py + "_aresponses_background_stream_can_release", # Tested through the background Responses stream tests in test_router.py + "_raise_released_stream_error", # Tested through the background Responses stream content-policy tests in test_router.py + "_mid_stream_error_before_output", # Tested through the Responses mid-stream fallback and background stream tests ] diff --git a/tests/unit/responses/test_responses_api_lifecycle.py b/tests/unit/responses/test_responses_api_lifecycle.py index 8805bc97d10..80da62f884a 100644 --- a/tests/unit/responses/test_responses_api_lifecycle.py +++ b/tests/unit/responses/test_responses_api_lifecycle.py @@ -6,6 +6,7 @@ import pytest import respx import litellm +from litellm.types.llms.openai import ResponsesAPIResponse from tests.unit.proxy.conftest import httpx_transport pytestmark: Final = pytest.mark.usefixtures(httpx_transport.__name__) @@ -42,6 +43,12 @@ async def _invoke_cancel_response(sync_mode: bool, call_kwargs: dict[str, object return await litellm.acancel_responses(**call_kwargs) +async def _invoke_get_response(sync_mode: bool, call_kwargs: dict[str, object]) -> object: + if sync_mode: + return litellm.get_responses(**call_kwargs) + return await litellm.aget_responses(**call_kwargs) + + @pytest.mark.parametrize( ("provider", "operation"), ( @@ -102,8 +109,43 @@ async def test_cancel_responses_404_surfaces_openai_api_error_with_exact_url( } with pytest.raises(openai.APIError) as exc_info: await _invoke_cancel_response(sync_mode, call_kwargs) - requests: Final = tuple(route.calls) + requests: Final[tuple[respx.models.Call, ...]] = tuple(route.calls) assert exc_info.value.status_code == 404 assert len(requests) == 1 assert str(requests[0].request.url) == expected_url + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_cancel_responses_reply_carries_the_id_retrieve_gives_for_the_same_run(sync_mode: bool) -> None: + upstream_id: Final = "resp_upstream123" + reply: Final = { + "id": upstream_id, + "object": "response", + "created_at": 0, + "status": "cancelled", + "model": "gpt-6.1-terra", + "output": [], + } + call_kwargs: Final = { + "custom_llm_provider": "openai", + "api_key": "sk-test", + "response_id": upstream_id, + "litellm_metadata": {"model_info": {"id": "deployment-1"}}, + } + + with respx.mock() as mock_router: + mock_router.post(f"https://api.openai.com/v1/responses/{upstream_id}/cancel").mock( + return_value=httpx.Response(status_code=200, json=reply) + ) + mock_router.get(f"https://api.openai.com/v1/responses/{upstream_id}").mock( + return_value=httpx.Response(status_code=200, json=reply) + ) + cancelled: Final = await _invoke_cancel_response(sync_mode, call_kwargs) + retrieved: Final = await _invoke_get_response(sync_mode, call_kwargs) + + assert isinstance(cancelled, ResponsesAPIResponse) + assert isinstance(retrieved, ResponsesAPIResponse) + assert cancelled.id == retrieved.id + assert cancelled.id != upstream_id diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index 3532870df1c..8011ca37cfb 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -15,6 +15,7 @@ from litellm.router_utils.common_utils import ( add_model_file_id_mappings, filter_team_based_models, filter_web_search_deployments, + format_content_policy_no_fallback_note, format_fallback_outcome_message, provider_for_generic_call, resolve_model_group_alias, @@ -679,6 +680,14 @@ def test_format_fallback_outcome_preserves_attempted_failure_wording_without_ski ) +def test_content_policy_note_explains_the_failure_when_fallbacks_are_a_plain_model_list() -> None: + note: Final = format_content_policy_no_fallback_note("gpt-bg", ("gpt-bg",), None, ["backup"], False) + + assert "content_policy_fallback=None" in note + assert "LiteLLM: model group 'gpt-bg' failed with the error above. No fallback was attempted." in note + assert "Fallbacks are configured for" not in note + + class TestWarnOnProviderCredentialMismatch: """A deployment that carries one provider's credentials while resolving to another is silently broken: litellm ignores the credentials and sends the diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 4cb994744f2..9875f1bd00a 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -7037,6 +7037,307 @@ async def test_aresponses_fallback_attempt_does_not_replay_held_lifecycle_events assert outcome == [fallback_created, fallback_error] +_BACKGROUND_PRIMARY_URL: Final = "https://api.openai.com/v1/responses" +_BACKGROUND_BACKUP_URL: Final = "https://backup.openai.test/v1/responses" + + +def _background_sse_frame(event: Mapping[str, object]) -> bytes: + return f"data: {json.dumps(event)}\n\n".encode() + + +def _background_response_body(status: str, background: bool) -> dict[str, object]: + return { + "id": "resp_upstream_run", + "object": "response", + "created_at": 0, + "status": status, + "model": "gpt-6.1-terra", + "background": background, + "output": [], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2} if status == "completed" else None, + } + + +def _paused_background_stream( + output_may_start: asyncio.Event, + upstream_paused: asyncio.Event, + background: bool, + after_pause: tuple[dict[str, object], ...], +) -> httpx.Response: + async def body() -> AsyncIterator[bytes]: + yield _background_sse_frame( + { + "type": "response.created", + "sequence_number": 0, + "response": _background_response_body("queued", background), + } + ) + yield _background_sse_frame( + { + "type": "response.queued", + "sequence_number": 1, + "response": _background_response_body("queued", background), + } + ) + upstream_paused.set() + await output_may_start.wait() + for event in after_pause: + yield _background_sse_frame(event) + + return httpx.Response(200, content=body(), headers={"content-type": "text/event-stream"}) + + +def _background_output_then_completed(background: bool) -> tuple[dict[str, object], ...]: + return ( + {"type": "response.output_text.delta", "sequence_number": 2, "item_id": "msg_1", "delta": "41"}, + { + "type": "response.completed", + "sequence_number": 3, + "response": _background_response_body("completed", background), + }, + ) + + +def _background_router( + model_name: str, + litellm_params: dict[str, object], + fallbacks: list[dict[str, list[str]]] | None, + content_policy_fallbacks: list[dict[str, list[str]]] | None = None, +) -> Router: + return litellm.Router( + model_list=[ + {"model_name": model_name, "litellm_params": litellm_params}, + { + "model_name": "backup", + "litellm_params": { + "model": "openai/gpt-6.1-terra", + "api_key": "backup-key", + "api_base": "https://backup.openai.test/v1", + }, + }, + ], + fallbacks=fallbacks, + content_policy_fallbacks=content_policy_fallbacks, + num_retries=0, + ) + + +async def _reached_client_before_upstream_paused( + first: "asyncio.Future[object]", upstream_paused: asyncio.Event +) -> bool: + paused: Final = asyncio.ensure_future(upstream_paused.wait()) + done, _ = await asyncio.wait({first, paused}, return_when=asyncio.FIRST_COMPLETED) + paused.cancel() + return first in done + + +@pytest.mark.parametrize( + ("fallbacks", "request_kwargs"), + [ + pytest.param(None, {}, id="no-fallback-configured"), + pytest.param([{"gpt-bg": ["backup"]}], {"disable_fallbacks": True}, id="fallbacks-disabled-for-the-request"), + ], +) +@pytest.mark.asyncio +async def test_aresponses_background_stream_hands_the_response_id_to_the_client_before_output( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + fallbacks: list[dict[str, list[str]]] | None, + request_kwargs: dict[str, bool], +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + output_may_start: Final = asyncio.Event() + upstream_paused: Final = asyncio.Event() + respx_mock.post(_BACKGROUND_PRIMARY_URL).mock( + return_value=_paused_background_stream( + output_may_start, upstream_paused, True, _background_output_then_completed(True) + ) + ) + router: Final = _background_router("gpt-bg", {"model": "openai/gpt-6.1-terra", "api_key": "k"}, fallbacks) + + stream: Final = await router.aresponses(model="gpt-bg", input="hi", background=True, stream=True, **request_kwargs) + first: Final = asyncio.ensure_future(anext(stream)) + arrived_before_output: Final = await _reached_client_before_upstream_paused(first, upstream_paused) + output_may_start.set() + created: Final = await first + rest: Final = [event async for event in stream] + + assert arrived_before_output is True + assert created.type == "response.created" + assert [event.type for event in rest] == [ + "response.queued", + "response.output_text.delta", + "response.completed", + ] + assert rest[-1].response.id == created.response.id + + +@pytest.mark.parametrize( + ("litellm_params", "background", "fallbacks"), + [ + pytest.param({"model": "openai/gpt-6.1-terra", "api_key": "k"}, False, None, id="foreground-request"), + pytest.param( + {"model": "openai/gpt-6.1-terra", "api_key": "k"}, True, [{"gpt-bg": ["backup"]}], id="fallback-configured" + ), + pytest.param( + { + "model": "azure/gpt-6.1-terra", + "api_key": "k", + "api_base": "https://bg-resource.openai.azure.com", + "api_version": "preview", + }, + True, + None, + id="azure-deployment", + ), + ], +) +@pytest.mark.asyncio +async def test_aresponses_stream_keeps_holding_lifecycle_events_until_output_off_the_eligible_background_route( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + litellm_params: dict[str, object], + background: bool, + fallbacks: list[dict[str, list[str]]] | None, +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + output_may_start: Final = asyncio.Event() + upstream_paused: Final = asyncio.Event() + respx_mock.route(method="POST", path__regex=r"/responses$").mock( + return_value=_paused_background_stream( + output_may_start, upstream_paused, background, _background_output_then_completed(background) + ) + ) + router: Final = _background_router("gpt-bg", litellm_params, fallbacks) + + stream: Final = await router.aresponses(model="gpt-bg", input="hi", background=background, stream=True) + first: Final = asyncio.ensure_future(anext(stream)) + arrived_before_output: Final = await _reached_client_before_upstream_paused(first, upstream_paused) + output_may_start.set() + events: Final = [await first, *[event async for event in stream]] + + assert arrived_before_output is False + assert upstream_paused.is_set() + assert [event.type for event in events] == [ + "response.created", + "response.queued", + "response.output_text.delta", + "response.completed", + ] + + +@pytest.mark.asyncio +async def test_aresponses_background_stream_surfaces_a_later_error_instead_of_a_fallback_once_its_id_is_out( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + output_may_start: Final = asyncio.Event() + upstream_paused: Final = asyncio.Event() + backup_output_may_start: Final = asyncio.Event() + backup_output_may_start.set() + server_error: Final = { + "type": "error", + "sequence_number": 2, + "error": {"type": "server_error", "code": "server_error", "message": "upstream died", "param": None}, + } + primary_route: Final = respx_mock.post(_BACKGROUND_PRIMARY_URL).mock( + return_value=_paused_background_stream(output_may_start, upstream_paused, True, (server_error,)) + ) + backup_route: Final = respx_mock.post(_BACKGROUND_BACKUP_URL).mock( + return_value=_paused_background_stream( + backup_output_may_start, asyncio.Event(), True, _background_output_then_completed(True) + ) + ) + router: Final = _background_router("gpt-bg", {"model": "openai/gpt-6.1-terra", "api_key": "k"}, None) + + stream: Final = await router.aresponses(model="gpt-bg", input="hi", background=True, stream=True) + first: Final = asyncio.ensure_future(anext(stream)) + arrived_before_output: Final = await _reached_client_before_upstream_paused(first, upstream_paused) + router.fallbacks = [{"gpt-bg": ["backup"]}] + output_may_start.set() + created: Final = await first + rest: Final = [item async for item in _events_until_error(stream)] + + assert arrived_before_output is True + assert created.type == "response.created" + assert [item.type for item in rest[:-1]] == ["response.queued"] + assert isinstance(rest[-1], litellm.InternalServerError) + assert "upstream died" in str(rest[-1]) + assert primary_route.call_count == 1 + assert backup_route.call_count == 0 + + + +_BACKGROUND_TEXT_DELTA: Final = { + "type": "response.output_text.delta", + "sequence_number": 2, + "item_id": "msg_1", + "delta": "41", +} + + +@pytest.mark.parametrize( + ("fallbacks", "content_policy_fallbacks", "output_before_refusal", "router_note_expected"), + [ + pytest.param(None, None, (), True, id="no-fallback-configured"), + pytest.param([{"other-group": ["backup"]}], None, (), True, id="fallbacks-for-another-group"), + pytest.param(None, [{"other-group": ["backup"]}], (), True, id="content-policy-fallbacks-for-another-group"), + pytest.param(None, None, (_BACKGROUND_TEXT_DELTA,), False, id="refused-after-output-text"), + ], +) +@pytest.mark.asyncio +async def test_aresponses_background_stream_content_policy_error_reads_the_same_once_its_id_is_out( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + fallbacks: list[dict[str, list[str]]] | None, + content_policy_fallbacks: list[dict[str, list[str]]] | None, + output_before_refusal: tuple[dict[str, object], ...], + router_note_expected: bool, +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + released_output_may_start: Final = asyncio.Event() + released_upstream_paused: Final = asyncio.Event() + held_output_may_start: Final = asyncio.Event() + held_output_may_start.set() + refusal: Final = { + "type": "error", + "sequence_number": 2 + len(output_before_refusal), + "error": { + "type": "invalid_request_error", + "code": "content_policy_violation", + "message": "Your request was rejected as a result of our safety system.", + "param": None, + }, + } + after_pause: Final = (*output_before_refusal, refusal) + primary_route: Final = respx_mock.post(_BACKGROUND_PRIMARY_URL).mock( + side_effect=[ + _paused_background_stream(released_output_may_start, released_upstream_paused, True, after_pause), + _paused_background_stream(held_output_may_start, asyncio.Event(), False, after_pause), + ] + ) + backup_route: Final = respx_mock.post(_BACKGROUND_BACKUP_URL) + router: Final = _background_router( + "gpt-bg", {"model": "openai/gpt-6.1-terra", "api_key": "k"}, fallbacks, content_policy_fallbacks + ) + + released: Final = await router.aresponses(model="gpt-bg", input="hi", background=True, stream=True) + first: Final = asyncio.ensure_future(anext(released)) + arrived_before_output: Final = await _reached_client_before_upstream_paused(first, released_upstream_paused) + released_output_may_start.set() + await first + released_error: Final = [item async for item in _events_until_error(released)][-1] + held: Final = await router.aresponses(model="gpt-bg", input="hi", stream=True) + held_error: Final = [item async for item in _events_until_error(held)][-1] + + assert arrived_before_output is True + assert isinstance(released_error, litellm.ContentPolicyViolationError) + assert isinstance(held_error, litellm.ContentPolicyViolationError) + assert ("LiteLLM: model group 'gpt-bg' failed with the error above" in held_error.message) is router_note_expected + assert released_error.message == held_error.message + assert primary_route.call_count == 2 + assert backup_route.call_count == 0 + @pytest.mark.asyncio async def test_aresponses_streaming_iterator_partial_content_injects_continuation(): """Mid-stream error: preserve partial assistant output and request continuation in a user turn.""" @@ -17296,9 +17597,9 @@ def _anthropic_messages_two_order_primary_model_list() -> list: pytest.param({"fallbacks": None, "enable_weighted_failover": True}, {"model": "primary"}, True, id="weighted-failover"), ], ) -def test_anthropic_messages_stream_can_fall_back_direct_call(router_kwargs, request_kwargs, expected): +def test_stream_can_fall_back_direct_call(router_kwargs, request_kwargs, expected): router = _anthropic_messages_make_router(**router_kwargs) - assert router._anthropic_messages_stream_can_fall_back("primary", request_kwargs) is expected + assert router._stream_can_fall_back("primary", request_kwargs) is expected @pytest.mark.parametrize( @@ -17308,7 +17609,7 @@ def test_anthropic_messages_stream_can_fall_back_direct_call(router_kwargs, requ pytest.param([1, 1], False, id="same-order-cannot-fall-back"), ], ) -def test_anthropic_messages_stream_can_fall_back_order_levels(orders, expected): +def test_stream_can_fall_back_order_levels(orders, expected): router = Router( model_list=[ { @@ -17319,7 +17620,7 @@ def test_anthropic_messages_stream_can_fall_back_order_levels(orders, expected): ], fallbacks=None, ) - assert router._anthropic_messages_stream_can_fall_back("primary", {"model": "primary"}) is expected + assert router._stream_can_fall_back("primary", {"model": "primary"}) is expected @pytest.mark.parametrize( @@ -17335,12 +17636,12 @@ def test_anthropic_messages_stream_can_fall_back_order_levels(orders, expected): ), ], ) -def test_anthropic_messages_stream_can_fall_back_order_target(request_kwargs, expected): +def test_stream_can_fall_back_order_target(request_kwargs, expected): router = Router(model_list=_anthropic_messages_two_order_primary_model_list(), fallbacks=None) - assert router._anthropic_messages_stream_can_fall_back("primary", request_kwargs) is expected + assert router._stream_can_fall_back("primary", request_kwargs) is expected -def test_anthropic_messages_order_levels_direct_call(): +def test_stream_order_levels_direct_call(): router = Router( model_list=[ { @@ -17351,7 +17652,7 @@ def test_anthropic_messages_order_levels_direct_call(): ], fallbacks=None, ) - assert router._anthropic_messages_order_levels("primary", {"model": "primary"}) == (1, 2) + assert router._stream_order_levels("primary", {"model": "primary"}) == (1, 2) @pytest.mark.asyncio From a0c9a5993b1e45479e85636305d581f015959433 Mon Sep 17 00:00:00 2001 From: tessellark <349031000+tessellark@users.noreply.github.com> Date: Wed, 28 Oct 2026 12:19:36 -0700 Subject: [PATCH 071/128] fix(batches): stop unbillable batches from blocking file deletes and new-batch polling (#48752) * test(batches): cover the batch cost poller gaps left after #47190 * fix(batches): keep unbillable batches from blocking file deletes and new-batch polling The file delete check now ignores batches past MANAGED_OBJECT_STALENESS_CUTOFF_DAYS, a provider 404 for the output object retires the batch, each poll page takes the arrival lane first, a batch expires only after the poller visits it past the cutoff, and the compare-and-swap matches batch_processed instead of updated_at * test(batches): cover the full arrival page and a failed expiry write in the batch cost poller Old batches still get polled every other cycle while new batches fill the page, old batches that fill a page never push a new batch out after boot, and a failed stale_expired write no longer stops the rest of the page from being polled * fix(batches): share the poll page between old and new batches and hold output files until billed Each lane now gets at least half of the poll page, a lane with nothing left gives its half to the other, and a lane read that resumed mid-lane wraps to the start, so neither old nor new batches can starve the other. The delete guard applies the 7 day cutoff to input files only. Output and error files stay locked until the batch is billed or retired, and stale_expired rows no longer lock anything. Only a lasting provider failure expires a batch past the cutoff, so a timeout or a 5xx no longer gives up on a batch that may still bill. Past the cutoff a lasting failure, a lasting cost-tracking failure, or an output file id that never appeared retires a finished row instead of polling it forever. * fix(batches): retry refused finished batches past 7 days, expire unfinished ones on a failed retrieve A finished batch the provider refuses after 7 days (rotated key, bucket 403, a 404 that does not name the batch) now keeps retrying so its spend is written once the config is fixed, instead of retiring unbilled. A batch the provider still reports running keeps polling. A batch that is not finished in the database becomes stale_expired once a retrieve fails after 7 days, so an unreachable provider is not polled forever. An S3 NoSuchKey for the output object now counts as gone. * fix(batches): expire an unfinished old batch only after 3 failed retrieves in a row * fix(batches): keep retrying an unfinished old batch through timeouts, 429 and 5xx * fix(batches): wait 3 polls before expiring an old batch whose output id lags * fix(batches): match a resource id in a 404 message followed by punctuation * fix(batches): keep refusal and pending-output expiry streaks apart, and void a pending-output retirement on a client write * fix(batches): give up on unsettled batches past twice the staleness cutoff Replace the per-reason refusal and pending-output streaks with one give-up age, and inject a clock into the poller and the delete check so tests pin exact age boundaries * fix(batches): keep polling finished batches past the give-up age so their spend stays recoverable * test(batches): expect a restarted poller's first page to split between old and new batches * test(batches): pin the delete-guard tests to a fixed clock and type a new journal list * test(batches): type the provider helper and the storage 404 cases --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../proxy/common_utils/check_batch_cost.py | 171 +++-- .../proxy/hooks/managed_files.py | 10 +- .../spend/test_batch_poll_starvation.py | 6 +- .../proxy/test_file_deletion_blocking.py | 7 +- .../common_utils/test_check_batch_cost.py | 598 ++++++++++++++++-- 5 files changed, 661 insertions(+), 131 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 f3b1240f733..47229466021 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -5,9 +5,9 @@ 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 collections.abc import Callable, Mapping, Sequence from dataclasses import replace as dataclasses_replace -from datetime import datetime, timedelta, timezone +from datetime import datetime, timedelta from types import MappingProxyType from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tuple, cast @@ -28,6 +28,7 @@ from litellm.constants import ( from litellm.repositories.table_repositories import ManagedObjectRepository from litellm.repositories.team_repository import TeamRepository from litellm.repositories.verification_token_repository import VerificationTokenRepository +from litellm.utils import get_utc_datetime if TYPE_CHECKING: from prisma import models as prisma_models @@ -42,11 +43,14 @@ if TYPE_CHECKING: def _error_names_resource(error: Exception, resource_id: str) -> bool: - return resource_id in re.findall(r"""[^\s"'`]+""", str(error)) + tokens: Final[list[str]] = re.findall(r"""[^\s"'`<>(]+""", str(error)) + return any(token.rstrip(".,:;)") == resource_id for token in tokens) CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost" RECENT_BATCH_POLL_WINDOW: Final = timedelta(hours=48) +GIVE_UP_POLLING_AFTER: Final = timedelta(days=2 * MANAGED_OBJECT_STALENESS_CUTOFF_DAYS) + PROVIDER_TERMINAL_BATCH_STATUSES: Final[Tuple[str, ...]] = ( "completed", @@ -91,6 +95,10 @@ def _team_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.Li return TeamRepository(prisma_client).table +def _unique_jobs(jobs: Sequence[ManagedObjectRow]) -> tuple[ManagedObjectRow, ...]: + return tuple({job.id: job for job in jobs}.values()) + + class CheckBatchCost: def __init__( self, @@ -98,6 +106,7 @@ class CheckBatchCost: prisma_client: "PrismaClient", llm_router: "Router", track_unmanaged_batch_cost: bool = False, + clock: Callable[[], datetime] = get_utc_datetime, ): from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router @@ -106,13 +115,14 @@ class CheckBatchCost: self.prisma_client: PrismaClient = prisma_client self.llm_router: Router = llm_router self._track_unmanaged_batch_cost = track_unmanaged_batch_cost + self._clock: Final = clock self._creator_attribution: Final = ManagedObjectCreatorAttribution(prisma_client, "CheckBatchCost") # Cached after the first poll cycle. Once we know the column is absent we skip # 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._poll_before: datetime = clock() self._arrival_cursor: tuple[datetime, str] | None = None self._arrival_before: datetime = self._poll_before self._poll_arrivals: bool = False @@ -147,27 +157,6 @@ class CheckBatchCost: return self.batch_processed_support_confirmed = True - async def _cleanup_stale_managed_objects(self) -> None: - """ - Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days - in non-terminal states as 'stale_expired'. These will never complete and - should not be polled. - """ - cutoff: Final = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS) - result: Final = await _managed_object_table(self.prisma_client).update_many( - where={ - "file_purpose": "batch", - "status": {"not_in": list(TERMINAL_MANAGED_OBJECT_STATUSES)}, - "created_at": {"lt": cutoff}, - }, - data={"status": "stale_expired"}, - ) - if result > 0: - verbose_proxy_logger.warning( - f"CheckBatchCost: marked {result} stale managed objects " - f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired" - ) - async def _fallback_find_jobs(self) -> "Sequence[ManagedObjectRow]": """Query batch jobs without the batch_processed filter (for older schemas).""" return await self._find_jobs( @@ -187,20 +176,43 @@ class CheckBatchCost: ) 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) + lead_arrivals: Final = self._poll_arrivals + self._poll_arrivals = not lead_arrivals + lead_share: Final = MAX_OBJECTS_PER_POLL_CYCLE - MAX_OBJECTS_PER_POLL_CYCLE // 2 + lead: Final = await self._find_job_page(where, arrivals=lead_arrivals, take=lead_share) + trail: Final = await self._find_job_page( + where, arrivals=not lead_arrivals, take=MAX_OBJECTS_PER_POLL_CYCLE - len(lead) + ) + spill: Final = ( + await self._find_job_page( + where, arrivals=lead_arrivals, take=MAX_OBJECTS_PER_POLL_CYCLE - len(lead) - len(trail) + ) + if len(lead) == lead_share + else () + ) + return _unique_jobs((*lead, *trail, *spill)) - async def _find_job_page(self, where: Mapping[str, object], *, arrivals: bool) -> Sequence[ManagedObjectRow]: + async def _find_job_page( + self, where: Mapping[str, object], *, arrivals: bool, take: int + ) -> Sequence[ManagedObjectRow]: + if take == 0: + return () + resumed_mid_lane: Final = (self._arrival_cursor if arrivals else self._poll_cursor) is not None + page: Final = await self._find_job_page_from_cursor(where, arrivals=arrivals, take=take) + if len(page) == take or not resumed_mid_lane: + return page + wrapped: Final = await self._find_job_page_from_cursor(where, arrivals=arrivals, take=take - len(page)) + return _unique_jobs((*page, *wrapped)) + + async def _find_job_page_from_cursor( + self, where: Mapping[str, object], *, arrivals: bool, take: int + ) -> 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) + self._arrival_before = self._clock() else: - self._poll_before = datetime.now(timezone.utc) - RECENT_BATCH_POLL_WINDOW + self._poll_before = self._clock() - 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={ @@ -220,10 +232,10 @@ class CheckBatchCost: else {} ), }, - take=MAX_OBJECTS_PER_POLL_CYCLE, + take=take, 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 + next_cursor: Final = (page[-1].created_at, page[-1].id) if len(page) == take else None if arrivals: self._arrival_cursor = next_cursor else: @@ -231,11 +243,47 @@ class CheckBatchCost: 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 {}), - } + if self._has_batch_processed_column: + return {"id": job.id, "batch_processed": False} + return {"id": job.id, "updated_at": job.updated_at} + + def _is_past_give_up_age(self, job: ManagedObjectRow) -> bool: + return job.created_at < self._clock() - GIVE_UP_POLLING_AFTER + + async def _expire_if_old_and_unfinished(self, job: ManagedObjectRow) -> None: + if job.status in TERMINAL_MANAGED_OBJECT_STATUSES or not self._is_past_give_up_age(job): + return + await self._expire_unfinished(job) + + @staticmethod + def _is_completed_with_output_id_pending(response: "LiteLLMBatch") -> bool: + from litellm.proxy.openai_files_endpoints.common_utils import ( + _completed_batch_safe_to_retire, + ) + + return ( + response.status in ("completed", "complete") + and response.output_file_id is None + and not _completed_batch_safe_to_retire(response) + ) + + async def _expire_unfinished(self, job: ManagedObjectRow) -> None: + try: + expired: Final = await _managed_object_table(self.prisma_client).update_many( + where={ + **self._unprocessed_snapshot(job), + "status": {"not_in": list(TERMINAL_MANAGED_OBJECT_STATUSES)}, + }, + data={"status": "stale_expired"}, + ) + except Exception as db_err: + verbose_proxy_logger.error(f"CheckBatchCost: failed to mark stale job {job.id} as stale_expired: {db_err}") + return + if expired > 0: + verbose_proxy_logger.warning( + f"CheckBatchCost: job {job.id} is older than {GIVE_UP_POLLING_AFTER.days} days " + "and still not finished, so it is marked stale_expired and no longer polled" + ) async def _retire_job(self, job: "ManagedObjectRow", reason: str) -> None: """ @@ -356,15 +404,21 @@ class CheckBatchCost: def _is_output_file_gone_at_provider(error: Exception, output_file_id: Optional[str]) -> bool: """A 404 naming the output file means there is nothing to fetch on this or any later poll: providers like Vertex AI advertise an output path for every batch, - including terminal ones that never wrote it. Any other failure may be - transient, so it keeps retrying until the staleness sweep bounds it.""" + including terminal ones that never wrote it. GCS names a gs:// object without its scheme, + and S3 names an s3:// object by its key alone. Any other failure keeps retrying.""" import openai from litellm.exceptions import NotFoundError + from litellm.llms.base_llm.chat.transformation import BaseLLMException if not output_file_id: return False - return isinstance(error, (NotFoundError, openai.NotFoundError)) and _error_names_resource(error, output_file_id) + is_not_found: Final = isinstance(error, (NotFoundError, openai.NotFoundError)) or ( + isinstance(error, BaseLLMException) and error.status_code == 404 + ) + object_path: Final = output_file_id.split("://", 1)[-1] + names: Final = (output_file_id, object_path, object_path.split("/", 1)[-1]) + return is_not_found and any(_error_names_resource(error, name) for name in names) async def _finalize_unbilled_terminal_job( self, job: "ManagedObjectRow", response: "LiteLLMBatch" @@ -1094,13 +1148,6 @@ class CheckBatchCost: processed_models: List[Tuple[Optional[str], Optional[str]]] = [] - try: - await self._cleanup_stale_managed_objects() - except Exception as cleanup_err: - verbose_proxy_logger.warning( - f"CheckBatchCost: stale cleanup failed (poll will continue): {cleanup_err}" - ) - # Look for all batches that have not yet been processed by CheckBatchCost. # self._has_batch_processed_column is cached after the first probe so that # older schemas don't pay a guaranteed-failing primary query + warning on @@ -1142,6 +1189,8 @@ class CheckBatchCost: if routing is None: if self._has_unified_id_without_model(job): await self._retire_job(job, "unified object id has no model id") + else: + await self._expire_if_old_and_unfinished(job) continue model_id, batch_id = routing @@ -1166,6 +1215,14 @@ class CheckBatchCost: prom_logger.record_check_batch_cost_error("provider_retrieval_error") if self._is_batch_gone_at_provider(e, batch_id) and self._batch_deployment_exists(model_id): await self._retire_job(job, f"batch {batch_id} no longer exists at the provider") + else: + await self._expire_if_old_and_unfinished(job) + continue + if self._is_completed_with_output_id_pending(response): + verbose_proxy_logger.info( + f"CheckBatchCost: batch {batch_id} is completed but its output file id " + f"has not appeared yet; leaving job {job.id} for the next poll cycle" + ) continue ## RETRIEVE THE BATCH JOB OUTPUT FILE @@ -1217,17 +1274,9 @@ class CheckBatchCost: self._store_batch_line_items_in_background(job, response, model_id, batch_id, billed=True) elif response.status in PROVIDER_TERMINAL_BATCH_STATUSES: - from litellm.proxy.openai_files_endpoints.common_utils import ( - _completed_batch_safe_to_retire, - ) - - if response.status in ("completed", "complete") and not _completed_batch_safe_to_retire(response): - verbose_proxy_logger.info( - f"CheckBatchCost: batch {batch_id} is completed but its output file id " - f"has not appeared yet; leaving job {job.id} for the next poll cycle" - ) - continue await self._retire_unbilled_terminal_job(job, response, model_id, batch_id) + else: + await self._expire_if_old_and_unfinished(job) # Record polling run metrics (always, even if nothing was processed) if prom_logger: diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 48cbc0ab8ff..7a07350e2d9 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -6,6 +6,7 @@ import base64 import json import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from datetime import datetime, timedelta from itertools import chain from types import MappingProxyType from typing import ( @@ -39,6 +40,7 @@ from litellm.constants import ( BATCH_OUTPUT_FILE_FALLBACK_RETRY_AFTER_SECONDS, BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, + MANAGED_OBJECT_STALENESS_CUTOFF_DAYS, MAX_FILE_LIST_LIMIT, PENDING_RESPONSE_STATUSES, ) @@ -113,6 +115,7 @@ from litellm.types.utils import ( LLMResponseTypes, SpecialEnums, ) +from litellm.utils import get_utc_datetime class _ManagedFileRetrieve(Protocol): @@ -507,10 +510,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): prisma_client: PrismaClient, *, sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + clock: Callable[[], datetime] = get_utc_datetime, ): self.internal_usage_cache = internal_usage_cache self.prisma_client = prisma_client self._sleep = sleep + self._clock = clock @staticmethod def _get_prometheus_logger(): @@ -2387,12 +2392,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): where={ "file_purpose": "batch", "batch_processed": False, - "status": {"not_in": ["failed", "expired", "cancelled"]}, + "status": {"not_in": ["failed", "expired", "cancelled", "stale_expired"]}, }, take=MAX_MATCHES_TO_RETURN, order={"created_at": "desc"}, ) + input_file_cutoff: Final = self._clock() - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS) referencing_batches: Final[list[dict[str, object]]] = [] for batch in batches: try: @@ -2405,7 +2411,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Extract file IDs from batch # Batches typically reference the unified file ID in input_file_id # Output and error files are generated by the provider - input_file_id = batch_data.get("input_file_id") + input_file_id = batch_data.get("input_file_id") if batch.created_at >= input_file_cutoff else None output_file_id = batch_data.get("output_file_id") error_file_id = batch_data.get("error_file_id") diff --git a/tests/integration/spend/test_batch_poll_starvation.py b/tests/integration/spend/test_batch_poll_starvation.py index 51f5f0871f5..f93cb807907 100644 --- a/tests/integration/spend/test_batch_poll_starvation.py +++ b/tests/integration/spend/test_batch_poll_starvation.py @@ -365,8 +365,10 @@ async def test_arrivals_progress_while_retry_scan_advances_and_wraps(count: int) 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] + restart_ids: Final = tuple(row.id for row in await restarted._find_jobs(where)) + restart_old: Final = tuple(identity for identity in restart_ids if identity in old_ids) + assert len(restart_ids) == MAX_OBJECTS_PER_POLL_CYCLE + assert restart_old == old_ids[: MAX_OBJECTS_PER_POLL_CYCLE - MAX_OBJECTS_PER_POLL_CYCLE // 2] finally: write_rows('DELETE FROM "LiteLLM_ManagedObjectTable" WHERE created_by = %s', (owner,)) diff --git a/tests/unit/enterprise/proxy/test_file_deletion_blocking.py b/tests/unit/enterprise/proxy/test_file_deletion_blocking.py index b7acdad35a1..4ca5e645f7f 100644 --- a/tests/unit/enterprise/proxy/test_file_deletion_blocking.py +++ b/tests/unit/enterprise/proxy/test_file_deletion_blocking.py @@ -10,6 +10,8 @@ This ensures cost tracking is not disrupted by premature file deletion. import base64 import json +from datetime import datetime, timezone +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -17,6 +19,8 @@ from fastapi import HTTPException from litellm.proxy._types import UserAPIKeyAuth +NOW: Final = datetime(2026, 10, 1, tzinfo=timezone.utc) + def _make_unified_file_id(file_id: str = "file-abc123") -> str: """Create a base64-encoded unified file ID.""" @@ -50,7 +54,7 @@ def _make_batch_db_record( mock_batch.status = status mock_batch.file_object = json.dumps(file_object) mock_batch.created_by = created_by - mock_batch.created_at = 1700000000 + mock_batch.created_at = NOW return mock_batch @@ -105,6 +109,7 @@ def _make_managed_files_instance_with_batches( instance = PROXY_LiteLLMManagedFiles( internal_usage_cache=mock_cache, prisma_client=mock_prisma, + clock=lambda: NOW, ) return instance 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 d7fc844171d..5ff159f80d3 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -9,22 +9,27 @@ ARN unified_object_id) batches with no managed unified id. import asyncio import json import logging +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from copy import copy -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from hashlib import sha256 from queue import SimpleQueue from types import SimpleNamespace from typing import TYPE_CHECKING, Final, Literal from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest import respx from fastapi import HTTPException +import litellm +from litellm.constants import MANAGED_OBJECT_STALENESS_CUTOFF_DAYS from litellm.types.router import LiteLLM_Params from litellm.types.utils import LiteLLMBatch +from litellm_enterprise.proxy.common_utils.check_batch_cost import GIVE_UP_POLLING_AFTER if TYPE_CHECKING: from litellm.batches.batch_utils import BatchCostUsageResult @@ -32,6 +37,11 @@ if TYPE_CHECKING: _IS_B64 = "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id" _CLAIM_UNIFIED_BATCH_ID = "dW5pZmllZF9iYXRjaF9pZA==" _CLAIM_OUTPUT_FILE_ID = "file-output-123" +_NOW: Final = datetime(2026, 1, 15, 12, tzinfo=timezone.utc) + + +def _fixed_clock() -> datetime: + return _NOW def test_batch_output_file_object_derives_metadata(): @@ -152,25 +162,9 @@ class TestCheckBatchCost: proxy_logging_obj=mock_proxy_logging_obj, prisma_client=mock_prisma_client, llm_router=mock_llm_router, + clock=_fixed_clock, ) - @pytest.mark.asyncio - async def test_cleanup_scoped_to_batch_file_purpose(self, check_batch_cost_instance, mock_prisma_client): - """_cleanup_stale_managed_objects scopes its update to file_purpose='batch' only.""" - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) - # Return empty so the main poll loop exits immediately - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[]) - - await check_batch_cost_instance.check_batch_cost() - - calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list - stale_call = calls[0] - assert stale_call[1]["data"] == {"status": "stale_expired"} - where = stale_call[1]["where"] - assert where["file_purpose"] == "batch" - assert "stale_expired" in where["status"]["not_in"] - assert "created_at" in where - @pytest.mark.asyncio async def test_startup_probe_confirms_batch_processed_support(self, check_batch_cost_instance, mock_prisma_client): mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None) @@ -251,7 +245,7 @@ class TestCheckBatchCost: fallback_where = calls[1][1]["where"] assert "batch_processed" not in fallback_where assert "stale_expired" in fallback_where["status"]["not_in"] - assert calls[1][1]["take"] == MAX_OBJECTS_PER_POLL_CYCLE + assert calls[2][1]["take"] == MAX_OBJECTS_PER_POLL_CYCLE # Column absence is now cached — next call should go straight to fallback assert check_batch_cost_instance._has_batch_processed_column is False assert check_batch_cost_instance.batch_processed_support_confirmed is False @@ -744,7 +738,7 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, ( + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 1, ( "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 +871,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 == 2, ( + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 1, ( "the job must still be marked processed once cost tracking succeeds" ) @@ -979,7 +973,7 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, ( + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 1, ( f"Expected update() to be called exactly once for a {terminal_status} job" ) update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"] @@ -1085,13 +1079,9 @@ class TestCheckBatchCost: assert owner.user_id == "user-1" assert owner.team_id == "team-1" - assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2 + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 1 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, - } + assert update_call.kwargs["where"] == {"id": "job-terminal-mint-1", "batch_processed": False} update_data = update_call.kwargs["data"] assert update_data["status"] == terminal_status assert update_data["batch_processed"] is True @@ -1156,7 +1146,7 @@ class TestCheckBatchCost: ) as mock_afile_content: await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, ( + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 1, ( "a completed batch with no output file must be marked processed exactly once" ) update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"] @@ -1205,6 +1195,7 @@ class TestCheckBatchCost: b"litellm_proxy;model_id:model-123;llm_batch_id:batch-456" ).decode() mock_job.created_by = "user-1" + mock_job.created_at = _NOW assert check_batch_cost_instance._has_batch_processed_column is True mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) @@ -1245,6 +1236,7 @@ class TestCheckBatchCost: mock_job.id = "job-in-progress-1" mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" mock_job.created_by = "user-1" + mock_job.created_at = _NOW mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) @@ -1375,7 +1367,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 == 2 + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 1 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, ( @@ -1548,7 +1540,7 @@ class TestCheckBatchCost: assert mock_afile_content.await_count == 1 mock_calculate.assert_not_awaited() - assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, ( + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 1, ( "a terminal batch with a 404ing output file must be retired, not retried forever" ) update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"] @@ -1946,7 +1938,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 == 2 + assert prisma.db.litellm_managedobjecttable.update_many.call_count == 1 update_data = prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["batch_processed"] is True assert update_data["status"] == "complete" @@ -2163,7 +2155,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 == 2 + assert prisma.db.litellm_managedobjecttable.update_many.call_count == 1 update_data = prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["batch_processed"] is True assert update_data["status"] == "complete" @@ -2710,6 +2702,7 @@ class TestPollPageStarvation: proxy_logging_obj=proxy_logging_obj, prisma_client=prisma, llm_router=llm_router, + clock=_fixed_clock, ) def _prisma(self, jobs): @@ -2725,6 +2718,7 @@ class TestPollPageStarvation: job.id = job_id job.unified_object_id = unified_object_id job.created_by = "user-1" + job.created_at = _NOW return job @staticmethod @@ -2910,6 +2904,173 @@ class TestPollPageStarvation: prisma.db.litellm_managedobjecttable.update.assert_not_awaited() + @pytest.mark.asyncio + async def test_a_new_batch_is_polled_every_cycle_while_an_older_batch_stays_unbilled(self) -> None: + """An unbilled batch older than the arrival window used to take every other cycle away from the arrival + lane, so a new batch waited two poll intervals to be billed.""" + now: Final = _NOW + rows: Final = tuple( + SimpleNamespace( + id=f"job-{name}", + unified_object_id=self._encode(f"litellm_proxy;model_id:model-123;llm_batch_id:batch_{name}"), + created_by="user-1", + created_at=now - age, + updated_at=now, + status="in_progress", + ) + for name, age in (("old", timedelta(days=3)), ("new", timedelta(hours=1))) + ) + prisma: Final = self._prisma([]) + prisma.db.litellm_managedobjecttable = _FakeManagedObjectPages(rows) + llm_router: Final = MagicMock() + llm_router.aretrieve_batch = AsyncMock(return_value=MagicMock(status="in_progress")) + poller: Final = self._instance(prisma, llm_router) + + async def _batches_polled_in_one_cycle() -> tuple[str, ...]: + llm_router.aretrieve_batch.reset_mock() + await poller.check_batch_cost() + return tuple(sorted(call.kwargs["batch_id"] for call in llm_router.aretrieve_batch.await_args_list)) + + assert [await _batches_polled_in_one_cycle() for _ in range(4)] == [("batch_new", "batch_old")] * 4 + + def _unbilled_rows(self, ages: dict[str, timedelta]) -> tuple[SimpleNamespace, ...]: + now: Final = _NOW + return tuple( + SimpleNamespace( + id=f"job-{name}", + unified_object_id=self._encode(f"litellm_proxy;model_id:model-123;llm_batch_id:batch_{name}"), + created_by="user-1", + created_at=now - age, + updated_at=now, + status="in_progress", + ) + for name, age in ages.items() + ) + + async def _poll_cycles( + self, monkeypatch: pytest.MonkeyPatch, ages: dict[str, timedelta], cycles: int, page_size: int = 2 + ) -> list[frozenset[str]]: + from litellm_enterprise.proxy.common_utils import check_batch_cost + + monkeypatch.setattr(check_batch_cost, "MAX_OBJECTS_PER_POLL_CYCLE", page_size) + prisma: Final = self._prisma([]) + prisma.db.litellm_managedobjecttable = _FakeManagedObjectPages(self._unbilled_rows(ages)) + llm_router: Final = MagicMock() + llm_router.aretrieve_batch = AsyncMock(return_value=MagicMock(status="in_progress")) + poller: Final = self._instance(prisma, llm_router) + + async def _batches_polled_in_one_cycle() -> frozenset[str]: + llm_router.aretrieve_batch.reset_mock() + await poller.check_batch_cost() + return frozenset(call.kwargs["batch_id"] for call in llm_router.aretrieve_batch.await_args_list) + + return [await _batches_polled_in_one_cycle() for _ in range(cycles)] + + @pytest.mark.asyncio + async def test_an_old_batch_is_polled_every_cycle_while_new_batches_fill_the_page( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + new_batches: Final = {f"new{n}": timedelta(hours=5 - n) for n in range(1, 5)} + polled: Final = await self._poll_cycles(monkeypatch, {"old": timedelta(days=3), **new_batches}, cycles=8) + + assert all("batch_old" in cycle for cycle in polled) + assert all(f"batch_{name}" in frozenset().union(*polled[:4]) for name in new_batches) + + @pytest.mark.asyncio + async def test_old_batches_that_fill_the_page_never_push_a_new_batch_out( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + polled: Final = await self._poll_cycles( + monkeypatch, + {"old1": timedelta(days=3, hours=1), "old2": timedelta(days=3), "new": timedelta(hours=1)}, + cycles=6, + ) + + assert all("batch_new" in cycle for cycle in polled) + assert frozenset({"batch_old1", "batch_old2"}) <= frozenset().union(*polled[:2]) + + @pytest.mark.asyncio + async def test_old_and_new_batches_each_get_half_the_page_when_both_overflow_it( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + ages: Final = { + **{f"old{n}": timedelta(days=3, hours=n) for n in range(3)}, + **{f"new{n}": timedelta(hours=n + 1) for n in range(3)}, + } + polled: Final = await self._poll_cycles(monkeypatch, ages, cycles=6) + + assert [sorted(batch_id[len("batch_") : -1] for batch_id in cycle) for cycle in polled] == [["new", "old"]] * 6 + assert frozenset().union(*polled) == frozenset(f"batch_{name}" for name in ages) + + @pytest.mark.asyncio + async def test_old_and_new_batches_take_turns_on_a_one_row_page(self, monkeypatch: pytest.MonkeyPatch) -> None: + polled: Final = await self._poll_cycles( + monkeypatch, {"old": timedelta(days=3), "new": timedelta(hours=1)}, cycles=4, page_size=1 + ) + + assert polled == [frozenset({"batch_old"}), frozenset({"batch_new"})] * 2 + + @pytest.mark.asyncio + async def test_new_batches_fill_the_page_when_no_old_batch_needs_polling( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + polled: Final = await self._poll_cycles( + monkeypatch, {f"new{n}": timedelta(hours=n + 1) for n in range(3)}, cycles=4 + ) + + assert [len(cycle) for cycle in polled] == [2] * 4 + + @pytest.mark.asyncio + async def test_a_failed_expiry_write_leaves_the_rest_of_the_page_polled(self) -> None: + rows: Final = self._unbilled_rows({"old": _PAST_THE_GIVE_UP_AGE, "new": timedelta(hours=1)}) + table: Final = _FakeManagedObjectPages(rows) + table.update_many = AsyncMock(side_effect=RuntimeError("database unavailable")) + prisma: Final = self._prisma([]) + prisma.db.litellm_managedobjecttable = table + llm_router: Final = MagicMock() + llm_router.aretrieve_batch = AsyncMock(side_effect=_A_REJECTED_KEY) + poller: Final = self._instance(prisma, llm_router) + + await poller.check_batch_cost() + + assert table.update_many.await_count == 1 + assert {call.kwargs["batch_id"] for call in llm_router.aretrieve_batch.await_args_list} == { + "batch_old", + "batch_new", + } + + +class _FakeManagedObjectPages: + """A LiteLLM_ManagedObjectTable double over several unbilled rows that pages them the way Postgres would: + inside the lane's created_at window, after the keyset cursor, in created_at then id order, take at most.""" + + def __init__(self, rows: tuple[SimpleNamespace, ...]) -> None: + self.rows = rows + self.find_many = AsyncMock(side_effect=self._find_many) + self.find_first = AsyncMock(return_value=None) + self.update_many = AsyncMock(return_value=0) + self.update = AsyncMock() + + @staticmethod + def _selected(row: SimpleNamespace, where: Mapping[str, object]) -> bool: + window: Final = where["created_at"] + cursor: Final = where.get("OR") + return ( + row.created_at <= window["lte"] + and ("gt" not in window or row.created_at > window["gt"]) + and ( + cursor is None + or row.created_at > cursor[0]["created_at"]["gt"] + or (row.created_at == cursor[1]["created_at"] and row.id > cursor[1]["id"]["gt"]) + ) + ) + + async def _find_many( + self, *, where: Mapping[str, object], take: int, order: Sequence[Mapping[str, str]] + ) -> tuple[SimpleNamespace, ...]: + selected: Final = (row for row in self.rows if self._selected(row, where)) + return tuple(sorted(selected, key=lambda row: (row.created_at, row.id))[:take]) + class _FakeManagedObjectRow: """One managed batch row the provider has finished but nothing has costed yet.""" @@ -2940,9 +3101,12 @@ class _FakeManagedObjectRow: class _FakeManagedObjectTable: """A LiteLLM_ManagedObjectTable double backed by one real, mutable row. - It honours the batch_processed and status filters, so the poller's compare-and-swap - and the managed-files deletion guard both read the same state a shared Postgres row - would give them. Staleness sweeps (the only queries scoped by created_at) never match. + It honours the batch_processed and status filters and the deletion guard's created_at cutoff, so the + poller's compare-and-swap and the managed-files deletion guard both read the state a shared Postgres row + would give them. The poll lanes' created_at windows are ignored, so every lane sees the row, and each + read hands back a snapshot, so a write that lands mid-cycle is invisible to the row the poller holds. + Writes journal "claim" or "release" when they flip batch_processed, "expire" when they mark the row + stale_expired, and "finalize" otherwise. """ def __init__(self, row: _FakeManagedObjectRow, journal: list): @@ -2956,7 +3120,7 @@ class _FakeManagedObjectTable: def _matches(self, where: dict) -> bool: for key, value in where.items(): if key == "created_at": - if "lt" in value and not self.row.created_at < value["lt"]: + if "gte" in value and not self.row.created_at >= value["gte"]: return False continue if key == "status": @@ -2971,10 +3135,12 @@ class _FakeManagedObjectTable: async def _update_many(self, *, where: dict, data: dict) -> int: if not self._matches(where): return 0 - if "updated_at" in where: + if "batch_processed" in data: + self.journal.append("claim" if data["batch_processed"] else "release") + elif data.get("status") == "stale_expired": + self.journal.append("expire") + else: 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) return 1 @@ -2985,7 +3151,7 @@ class _FakeManagedObjectTable: setattr(self.row, key, value) async def _find_many(self, *, where: dict, take=None, order=None) -> list: - return [self.row] if self._matches(where) else [] + return [copy(self.row)] if self._matches(where) else [] class TestMultiPodBatchCostClaim: @@ -3008,6 +3174,7 @@ class TestMultiPodBatchCostClaim: proxy_logging_obj=proxy_logging_obj, prisma_client=prisma, llm_router=llm_router, + clock=_fixed_clock, ) @staticmethod @@ -3465,7 +3632,7 @@ class TestMultiPodBatchCostClaim: cache = MagicMock() cache.async_get_cache = AsyncMock(return_value=None) cache.async_set_cache = AsyncMock() - guard = PROXY_LiteLLMManagedFiles(internal_usage_cache=cache, prisma_client=prisma) + guard = PROXY_LiteLLMManagedFiles(internal_usage_cache=cache, prisma_client=prisma, clock=_fixed_clock) scheduler = MagicMock() scheduler.get_job.return_value = MagicMock() @@ -3683,12 +3850,15 @@ class TestBatchLineItemsFromThePoller: @classmethod @contextmanager - def _provider(cls, journal: list): - """The provider's batch files by id, and a spend-log writer that journals each write.""" + def _provider(cls, journal: list[str], missing_file_error: Exception | None = None) -> Iterator[None]: + """The provider's batch files by id, and a spend-log writer that journals each write. Reading any other + file raises missing_file_error, or the NotFoundError OpenAI answers with.""" files = {"file-input-1": cls._INPUT, _CLAIM_OUTPUT_FILE_ID: cls._OUTPUT, "file-err-1": cls._ERROR} async def _afile_content(**kwargs): journal.append(f"read {kwargs['file_id']}") + if kwargs["file_id"] not in files and missing_file_error is not None: + raise missing_file_error if kwargs["file_id"] not in files: from litellm.exceptions import NotFoundError @@ -3861,7 +4031,7 @@ class TestBatchLineItemsFromThePoller: await instance.check_batch_cost() await self._stored(instance) - assert journal == ["finalize", "read file-err-1", "read file-input-1", "lines"] + assert journal == ["claim", "read file-err-1", "read file-input-1", "lines"] (payload,) = logger.payloads assert [(line["custom_id"], line["status"], line["response_cost"]) for line in payload["lines"]] == [ ("ok-1", "failure", 0.0) @@ -3892,7 +4062,7 @@ class TestBatchLineItemsFromThePoller: await instance.check_batch_cost() await self._stored(instance) - assert journal == ["finalize", "read file-err-1", "read file-input-1", "lines"] + assert journal == ["claim", "read file-err-1", "read file-input-1", "lines"] ((line,),) = [payload["lines"] for payload in logger.payloads] assert line["request"] == json.loads(self._INPUT) @@ -3910,7 +4080,7 @@ class TestBatchLineItemsFromThePoller: await instance.check_batch_cost() await self._stored(instance) - assert journal == ["finalize", "read file-err-1", "read file-input-1"] + assert journal == ["claim", "read file-err-1", "read file-input-1"] assert logger.payloads == [] @pytest.mark.asyncio @@ -3970,11 +4140,37 @@ class TestBatchLineItemsFromThePoller: await instance.check_batch_cost() await self._stored(instance) - assert journal == ["read file-out-deleted", "finalize", "read file-err-1", "read file-input-1", "lines"] + assert journal == ["read file-out-deleted", "claim", "read file-err-1", "read file-input-1", "lines"] (payload,) = logger.payloads assert [(line["custom_id"], line["status"]) for line in payload["lines"]] == [("ok-1", "failure")] assert row.batch_processed is True + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("status_code", "missing_object", "retired"), + ( + (404, "bucket/job-1/predictions.jsonl", True), + (404, "bucket/job-2/predictions.jsonl", False), + (500, "bucket/job-1/predictions.jsonl", False), + ), + ) + async def test_a_storage_404_for_the_output_object_retires_the_batch_unbilled( + self, status_code: int, missing_object: str, retired: bool + ) -> None: + """Vertex AI reports a deleted gs:// output as a VertexAIError 404 that names the object without its + scheme, and no later poll can read that object either.""" + from litellm.llms.vertex_ai.common_utils import VertexAIError + + journal: list[str] = [] + row, instance = self._setup(journal, output_file_id="gs://bucket/job-1/predictions.jsonl") + missing: Final = VertexAIError(status_code=status_code, message=f"No such object: {missing_object}") + + with self._provider(journal, missing_file_error=missing): + await instance.check_batch_cost() + + assert row.batch_processed is retired + assert "bill" not in journal + @pytest.mark.asyncio async def test_a_terminal_batch_with_no_lines_sends_nothing(self, monkeypatch): import litellm @@ -3988,7 +4184,7 @@ class TestBatchLineItemsFromThePoller: await instance.check_batch_cost() await self._stored(instance) - assert journal == ["finalize"] + assert journal == ["claim"] assert logger.payloads == [] @pytest.mark.asyncio @@ -4088,16 +4284,19 @@ class TestBatchLineItemsFromThePoller: _, instance = self._setup(journal) monkeypatch.setattr(litellm, "batch_line_item_callbacks", (self._recording_logger(journal),)) - async def _an_earlier_batch_waits_to_store_its_lines(): + find_jobs: Final = instance._find_jobs + + async def _an_earlier_batch_waits_to_store_its_lines(where: Mapping[str, object]) -> Sequence[object]: instance._store_batch_line_items_in_background( MagicMock(), MagicMock(), "model-123", "batch-earlier", billed=True ) await asyncio.sleep(0.01) + return await find_jobs(where) async def _store(job, response, model_id, batch_id, *, billed): journal.append(f"store {batch_id}") - monkeypatch.setattr(instance, "_cleanup_stale_managed_objects", _an_earlier_batch_waits_to_store_its_lines) + monkeypatch.setattr(instance, "_find_jobs", _an_earlier_batch_waits_to_store_its_lines) monkeypatch.setattr(instance, "_store_batch_line_items", _store) with self._provider(journal): await instance.check_batch_cost() @@ -4345,23 +4544,261 @@ class TestCreatorTurnedOffMessageLogging: instance.prisma_client.db.litellm_verificationtoken.find_unique.assert_not_called() +_PAST_THE_STALENESS_CUTOFF: Final = timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS, microseconds=1) +_INSIDE_THE_STALENESS_CUTOFF: Final = timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS) +_PAST_THE_GIVE_UP_AGE: Final = GIVE_UP_POLLING_AFTER + timedelta(microseconds=1) +_INSIDE_THE_GIVE_UP_AGE: Final = GIVE_UP_POLLING_AFTER + + +def _provider_batch(status: str) -> LiteLLMBatch: + return LiteLLMBatch( + id="batch-456", + status=status, + output_file_id=None, + error_file_id=None, + created_at=1, + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id="file-input-1", + object="batch", + ) + + @pytest.mark.asyncio -async def test_old_completed_batch_remains_available_for_costing() -> None: +async def test_an_unbilled_batch_past_the_staleness_cutoff_no_longer_blocks_deleting_its_input_file() -> None: row: Final = _FakeManagedObjectRow() row.status = "completed" - row.created_at = datetime(2000, 1, 1, tzinfo=timezone.utc) + row.created_at = _NOW - _PAST_THE_STALENESS_CUTOFF prisma: Final = TestMultiPodBatchCostClaim._prisma(row, []) - poller: Final = TestMultiPodBatchCostClaim._instance(prisma, TestMultiPodBatchCostClaim._router()) - await poller._cleanup_stale_managed_objects() + await TestMultiPodBatchCostClaim._run_deletion_guard(prisma, "file-input-1") assert row.batch_processed is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("age", "file_field", "file_id"), + ( + (_INSIDE_THE_STALENESS_CUTOFF, "input_file_id", "file-input-1"), + (_PAST_THE_STALENESS_CUTOFF, "output_file_id", _CLAIM_OUTPUT_FILE_ID), + (_PAST_THE_STALENESS_CUTOFF, "error_file_id", "file-error-1"), + ), + ids=("input_inside_the_cutoff", "output_past_the_cutoff", "error_past_the_cutoff"), +) +async def test_an_unbilled_batch_still_blocks_deleting_a_file_its_cost_may_need( + age: timedelta, file_field: str, file_id: str +) -> None: + row: Final = _FakeManagedObjectRow() + row.status = "completed" + row.created_at = _NOW - age + row.file_object = json.dumps({"id": "batch-456", "status": "completed", file_field: file_id}) + prisma: Final = TestMultiPodBatchCostClaim._prisma(row, []) + + with pytest.raises(HTTPException) as refused: + await TestMultiPodBatchCostClaim._run_deletion_guard(prisma, file_id) + + assert refused.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_a_stale_expired_batch_no_longer_blocks_deleting_its_output_file() -> None: + row: Final = _FakeManagedObjectRow() + row.status = "stale_expired" + prisma: Final = TestMultiPodBatchCostClaim._prisma(row, []) + + await TestMultiPodBatchCostClaim._run_deletion_guard(prisma, _CLAIM_OUTPUT_FILE_ID) + + assert row.status == "stale_expired" + + +@pytest.mark.asyncio +async def test_a_batch_past_the_give_up_age_that_the_provider_finished_is_still_billed() -> None: + row: Final = _FakeManagedObjectRow() + row.created_at = _NOW - _PAST_THE_GIVE_UP_AGE + journal: Final = [] + prisma: Final = TestMultiPodBatchCostClaim._prisma(row, journal) + poller: Final = TestMultiPodBatchCostClaim._instance(prisma, TestMultiPodBatchCostClaim._router()) + + with TestMultiPodBatchCostClaim._billing_patches(journal): + await poller.check_batch_cost() + + assert journal == ["fetch", "claim", "bill"] + assert row.batch_processed is True + + +_A_REJECTED_KEY: Final = litellm.AuthenticationError(message="invalid api key", llm_provider="openai", model="gpt-4") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("unified_object_id", "retrieve"), + ( + (_CLAIM_UNIFIED_BATCH_ID, _provider_batch("in_progress")), + (_CLAIM_UNIFIED_BATCH_ID, RuntimeError("timed out")), + ( + _CLAIM_UNIFIED_BATCH_ID, + litellm.InternalServerError(message="upstream error", llm_provider="openai", model="gpt-4"), + ), + (_CLAIM_UNIFIED_BATCH_ID, litellm.RateLimitError(message="rate limited", llm_provider="openai", model="gpt-4")), + (_CLAIM_UNIFIED_BATCH_ID, _A_REJECTED_KEY), + ("batch-raw-1", _provider_batch("in_progress")), + ), + ids=( + "still_running", + "retrieve_timed_out", + "provider_5xx", + "provider_429", + "key_rejected", + "unroutable", + ), +) +@pytest.mark.parametrize( + ("age", "db_status", "after_the_poll"), + ( + (_PAST_THE_GIVE_UP_AGE, "in_progress", ("stale_expired", False)), + (_PAST_THE_GIVE_UP_AGE, "completed", ("completed", False)), + (_INSIDE_THE_GIVE_UP_AGE, "in_progress", ("in_progress", False)), + (_INSIDE_THE_GIVE_UP_AGE, "completed", ("completed", False)), + ), + ids=("unfinished_past_the_age", "finished_past_the_age", "unfinished_at_the_age", "finished_at_the_age"), +) +async def test_only_an_unfinished_batch_a_poll_cannot_settle_is_given_up_once_past_the_give_up_age( + unified_object_id: str, + retrieve: LiteLLMBatch | Exception, + age: timedelta, + db_status: str, + after_the_poll: tuple[str, bool], +) -> None: + row: Final = _FakeManagedObjectRow() + row.created_at = _NOW - age + row.status = db_status + row.unified_object_id = unified_object_id + router: Final = TestMultiPodBatchCostClaim._router() + router.aretrieve_batch.side_effect = [retrieve] + poller: Final = TestMultiPodBatchCostClaim._instance(TestMultiPodBatchCostClaim._prisma(row, []), router) + + with TestMultiPodBatchCostClaim._billing_patches([]): + await poller.check_batch_cost() + + assert (row.status, row.batch_processed) == after_the_poll + + +async def _a_vertex_job_that_no_longer_exists( + *, model: str, batch_id: str, litellm_metadata: Mapping[str, object] +) -> LiteLLMBatch: + from litellm.llms.vertex_ai.common_utils import VertexAIError + + raise VertexAIError( + status_code=404, + message='Error: 404 {"error":{"code":404,"message":"The BatchPredictionJob does not exist.","status":"NOT_FOUND"}}', + ) + + +async def _a_finished_batch_whose_output_never_appeared( + *, model: str, batch_id: str, litellm_metadata: Mapping[str, object] +) -> LiteLLMBatch: + return _provider_batch("completed") + + +async def _an_output_file_the_key_cannot_read() -> None: + raise litellm.PermissionDeniedError( + message="forbidden", + llm_provider="openai", + model="gpt-4", + response=httpx.Response(403, request=httpx.Request("GET", "https://api.openai.com/v1/files")), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("db_status", ("in_progress", "completed")) +async def test_a_batch_past_the_give_up_age_whose_output_id_lags_is_billed_once_it_appears(db_status: str) -> None: + row: Final = _FakeManagedObjectRow() + row.created_at = _NOW - _PAST_THE_GIVE_UP_AGE + row.status = db_status + journal: Final = [] + router: Final = TestMultiPodBatchCostClaim._router() + router.aretrieve_batch.side_effect = [ + _provider_batch("completed"), + _provider_batch("completed"), + router.aretrieve_batch.return_value, + ] + poller: Final = TestMultiPodBatchCostClaim._instance(TestMultiPodBatchCostClaim._prisma(row, journal), router) + + with TestMultiPodBatchCostClaim._billing_patches(journal): + for _ in range(3): + await poller.check_batch_cost() + + assert journal == ["fetch", "claim", "bill"] + assert row.batch_processed is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("retrieve", "during_fetch"), + ( + (_A_REJECTED_KEY, None), + (_a_vertex_job_that_no_longer_exists, None), + (_a_finished_batch_whose_output_never_appeared, None), + (None, _an_output_file_the_key_cannot_read), + ), + ids=("key_rejected", "batch_404_without_its_id", "output_id_pending", "output_unreadable"), +) +async def test_a_finished_batch_a_poll_cannot_bill_past_the_give_up_age_bills_once_the_cause_is_fixed( + retrieve: Exception | Callable[..., Awaitable[LiteLLMBatch]] | None, + during_fetch: Callable[[], Awaitable[None]] | None, +) -> None: + row: Final = _FakeManagedObjectRow() + row.created_at = _NOW - _PAST_THE_GIVE_UP_AGE + row.status = "completed" + router: Final = TestMultiPodBatchCostClaim._router() + router.aretrieve_batch.side_effect = retrieve + journal: Final = [] + prisma: Final = TestMultiPodBatchCostClaim._prisma(row, journal) + poller: Final = TestMultiPodBatchCostClaim._instance(prisma, router) + + with TestMultiPodBatchCostClaim._billing_patches(journal, during_fetch=during_fetch): + await poller.check_batch_cost() + + assert (row.status, row.batch_processed) == ("completed", False) + with pytest.raises(HTTPException): + await TestMultiPodBatchCostClaim._run_deletion_guard(prisma, _CLAIM_OUTPUT_FILE_ID) + + router.aretrieve_batch.side_effect = None + with TestMultiPodBatchCostClaim._billing_patches(journal): + await poller.check_batch_cost() + + assert journal[-1] == "bill" + assert row.batch_processed is True + + +@pytest.mark.asyncio +async def test_giving_up_on_a_batch_keeps_a_finished_status_a_client_wrote_meanwhile() -> None: + row: Final = _FakeManagedObjectRow() + row.created_at = _NOW - _PAST_THE_GIVE_UP_AGE + router: Final = TestMultiPodBatchCostClaim._router() + + async def _a_client_reads_the_finished_batch_while_the_poller_retrieve_is_refused( + *, model: str, batch_id: str, litellm_metadata: Mapping[str, object] + ) -> LiteLLMBatch: + row.status = "completed" + raise _A_REJECTED_KEY + + router.aretrieve_batch.side_effect = _a_client_reads_the_finished_batch_while_the_poller_retrieve_is_refused + poller: Final = TestMultiPodBatchCostClaim._instance(TestMultiPodBatchCostClaim._prisma(row, []), router) + + with TestMultiPodBatchCostClaim._billing_patches([]): + await poller.check_batch_cost() + assert row.status == "completed" + assert row.batch_processed is False @pytest.mark.asyncio @pytest.mark.parametrize("progress", ["claimed", "refreshed"]) -async def test_stale_retirement_preserves_concurrent_progress(progress: str) -> None: +async def test_stale_retirement_lands_unless_another_pod_claimed_the_row(progress: str) -> None: + """A client GET that rewrote the row while the poller held it moves updated_at but leaves the batch unbilled, + so the retirement still lands.""" row: Final = _FakeManagedObjectRow() selected: Final = copy(row) if progress == "claimed": @@ -4372,13 +4809,16 @@ async def test_stale_retirement_preserves_concurrent_progress(progress: str) -> await poller._retire_job(selected, "provider batch disappeared") - assert row.batch_processed is (progress == "claimed") - assert row.updated_at > selected.updated_at + assert row.batch_processed is True @pytest.mark.asyncio -@pytest.mark.parametrize("progress", ["claimed", "refreshed"]) -async def test_stale_terminal_response_preserves_concurrent_progress(progress: str) -> None: +@pytest.mark.parametrize( + ("progress", "finalized", "expected_status"), [("claimed", False, "in_progress"), ("refreshed", True, "failed")] +) +async def test_stale_terminal_response_lands_unless_another_pod_claimed_the_row( + progress: str, finalized: bool, expected_status: str +) -> None: from litellm.types.utils import LiteLLMBatch row: Final = _FakeManagedObjectRow() @@ -4399,11 +4839,10 @@ async def test_stale_terminal_response_preserves_concurrent_progress(progress: s status="failed", ) - assert await poller._finalize_unbilled_terminal_job(selected, response) is False + assert await poller._finalize_unbilled_terminal_job(selected, response) is finalized - assert row.file_object == '{"revision":"fresh"}' - assert row.status == "in_progress" - assert row.batch_processed is (progress == "claimed") + assert row.status == expected_status + assert row.batch_processed is True @pytest.mark.asyncio @@ -4447,7 +4886,19 @@ def test_another_resource_in_a_404_is_not_permanent_loss(resource_id: str, other @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")) +@pytest.mark.parametrize( + "message", + ( + "No object found with id '{id}'.", + 'File "{id}" does not exist', + "{id} not found", + "No such object: {id}.", + "{id}: not found", + "Missing {id}, giving up", + "Not found (id {id})", + "Not found ({id})", + ), +) 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 @@ -4457,6 +4908,23 @@ def test_a_404_naming_the_exact_resource_is_permanent_loss(resource_id: str, mes assert CheckBatchCost._is_output_file_gone_at_provider(error, resource_id) is True +@pytest.mark.parametrize( + ("missing_key", "gone"), (("out/job-123/input.jsonl.out", True), ("out/job-123/other.jsonl.out", False)) +) +def test_an_s3_no_such_key_for_the_output_object_is_permanent_loss(missing_key: str, gone: bool) -> None: + from litellm.llms.bedrock.common_utils import BedrockError + from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + + # S3 GetObject 404 body shape, checked live against a missing key on 2026-10-20 + error: Final = BedrockError( + status_code=404, + message='\nNoSuchKey' + f"The specified key does not exist.{missing_key}" + "RH", + ) + assert CheckBatchCost._is_output_file_gone_at_provider(error, "s3://qa-bucket/out/job-123/input.jsonl.out") is gone + + @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 a45c2023ce67beaf18c99f1c8ddfce3f2b6d5c26 Mon Sep 17 00:00:00 2001 From: shrey Date: Wed, 28 Oct 2026 12:20:17 -0700 Subject: [PATCH 072/128] fix(audio): return 400 for unsupported transcription providers (#47971) --- litellm/__init__.py | 1 + litellm/exceptions.py | 5 + litellm/main.py | 12 +- litellm/router.py | 25 +++- litellm/router_utils/cooldown_handlers.py | 2 +- .../router_utils/fallback_event_handlers.py | 5 +- litellm/router_utils/get_retry_from_policy.py | 3 + .../router_code_coverage.py | 1 + .../router_utils/test_cooldown_handlers.py | 35 ++++++ .../test_get_retry_from_policy.py | 12 ++ tests/unit/test_main.py | 31 +++++ tests/unit/test_router/test_router.py | 110 ++++++++++++++++++ 12 files changed, 234 insertions(+), 8 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 2cff21bf29a..af771b2f0e6 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1449,6 +1449,7 @@ from .exceptions import ( Timeout, APIConnectionError, UnsupportedParamsError, + UnsupportedEndpointError as UnsupportedEndpointError, APIResponseValidationError, UnprocessableEntityError, InternalServerError, diff --git a/litellm/exceptions.py b/litellm/exceptions.py index c8eee9ad1f8..cd5139f5f79 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -317,6 +317,10 @@ class MalformedRequestFieldError(BadRequestError): pass +class UnsupportedEndpointError(BadRequestError): + pass + + class UnprocessableEntityError(openai.UnprocessableEntityError): def __init__( self, @@ -1056,6 +1060,7 @@ LITELLM_EXCEPTION_TYPES: Final = [ BadRequestError, UnprocessableEntityError, UnsupportedParamsError, + UnsupportedEndpointError, Timeout, PermissionDeniedError, RateLimitError, diff --git a/litellm/main.py b/litellm/main.py index 79d2d0370f7..75bd4431120 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8390,7 +8390,7 @@ def transcription( custom_llm_provider=custom_llm_provider, ) - response: TranscriptionResponse | Coroutine[object, object, TranscriptionResponse] | None = None + response: TranscriptionResponse | Coroutine[object, object, TranscriptionResponse] provider_config: Final = ProviderConfigManager.get_provider_audio_transcription_config( model=model, @@ -8552,10 +8552,16 @@ def transcription( provider_config=provider_config, shared_session=shared_session, ) + else: + raise litellm.UnsupportedEndpointError( + message=f"{custom_llm_provider}/{model} does not support audio transcription.", + model=model, + llm_provider=custom_llm_provider, + ) # Store duration in _hidden_params for cost calculation without # exposing it in the response body (see sync path comment above). - if response is not None and not isinstance(response, Coroutine): + if not isinstance(response, Coroutine): existing_duration: Final = getattr(response, "duration", None) if existing_duration is None: calculated_duration: Final = calculate_request_duration(file) @@ -8565,8 +8571,6 @@ def transcription( ) response_hidden_params["audio_transcription_duration"] = calculated_duration - if response is None: - raise ValueError("Unmapped provider passed in. Unable to get the response.") return response diff --git a/litellm/router.py b/litellm/router.py index d5cad013e8d..e8d83e1203e 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -633,6 +633,7 @@ _DEPLOYMENT_PARAMS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | GenericLite _OPAQUE_VALUE_ADAPTER: Final = TypeAdapter(object) _DEPLOYMENT_SELECTED_EVENT: Final = "litellm.request.deployment_selected" _CONTENT_BLOCKS_ADAPTER: Final = TypeAdapter(list[dict[str, object]]) +_DEPLOYMENTS_ADAPTER: Final = TypeAdapter(Sequence[Mapping[str, object]]) def _extend_trailing_text_block(blocks: Sequence[Mapping[str, object]], text: str) -> Sequence[Mapping[str, object]]: @@ -694,6 +695,12 @@ def _configured_model_info(deployment: Mapping[str, object]) -> Mapping[str, obj return _MODEL_INFO_ADAPTER.validate_python(deployment.get("model_info") or {}) +def _every_deployment_refused( + deployments: Sequence[Mapping[str, object]], refused_deployment_ids: tuple[str, ...] +) -> bool: + return all(_configured_model_info(deployment).get("id") in refused_deployment_ids for deployment in deployments) + + def _silent_experiment_targets(silent_model: object) -> tuple[str, ...]: if silent_model is None: return () @@ -4771,7 +4778,13 @@ class Router: @staticmethod def _stamp_retry_skip_deployment_id(exception: Exception, kwargs: Mapping[str, object]) -> None: effective_model_info: Final = kwargs.get("model_info") - deployment_id: Final = effective_model_info.get("id") if isinstance(effective_model_info, Mapping) else None + if not isinstance(effective_model_info, Mapping): + return + model_info: Final = _MODEL_INFO_ADAPTER.validate_python(effective_model_info) + configured_deployment_id: Final = ( + model_info.get("original_model_id") if isinstance(exception, litellm.UnsupportedEndpointError) else None + ) + deployment_id: Final = configured_deployment_id or model_info.get("id") if isinstance(deployment_id, str) and deployment_id: exception.retry_skip_deployment_id = deployment_id # pyright: ignore[reportAttributeAccessIssue] # dynamic stamp, read by _deployment_ids_to_skip_on_retry @@ -9115,6 +9128,10 @@ class Router: ) if first_skipped_ids: kwargs["_retry_skipped_deployment_ids"] = first_skipped_ids # rebind-ok: the next attempt reads it + if isinstance(original_exception, litellm.UnsupportedEndpointError) and _every_deployment_refused( + _DEPLOYMENTS_ADAPTER.validate_python(_healthy_deployments), first_skipped_ids + ): + raise else: raise @@ -9194,6 +9211,10 @@ class Router: ) if skipped_ids: kwargs["_retry_skipped_deployment_ids"] = skipped_ids # rebind-ok: the next attempt reads it + if isinstance(e, litellm.UnsupportedEndpointError) and _every_deployment_refused( + _DEPLOYMENTS_ADAPTER.validate_python(_healthy_deployments), skipped_ids + ): + raise _timeout = self._time_to_sleep_before_retry( e=e, remaining_retries=remaining_retries, @@ -9290,7 +9311,7 @@ class Router: status_code: Final = getattr(error, "status_code", None) if status_code is not None and not litellm.should_retry(status_code): # 401/403 are special cases - allow retry if multiple deployments exist (handled below) - if status_code not in (401, 403): + if status_code not in (401, 403) and not isinstance(error, litellm.UnsupportedEndpointError): raise error if isinstance(error, litellm.NotFoundError): diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index e1baf66b201..128bca6cca0 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -356,7 +356,7 @@ def _should_run_cooldown_logic( ) return False - if isinstance(original_exception, litellm.MalformedRequestFieldError): + if isinstance(original_exception, (litellm.MalformedRequestFieldError, litellm.UnsupportedEndpointError)): verbose_router_logger.debug( "Should Not Run Cooldown Logic: litellm rejected the request before calling the deployment" ) diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 6c93fe6d239..0a77ad5b9ea 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -261,7 +261,10 @@ def is_bad_request_without_fallback(exception: Exception, fallback_on_unsupporte return False if isinstance(exception, litellm.UnsupportedParamsError): return True - if isinstance(exception, (litellm.ContextWindowExceededError, litellm.ContentPolicyViolationError)): + if isinstance( + exception, + (litellm.ContextWindowExceededError, litellm.ContentPolicyViolationError, litellm.UnsupportedEndpointError), + ): return False return exception.status_code == 400 and bool(getattr(exception, "llm_provider", "")) diff --git a/litellm/router_utils/get_retry_from_policy.py b/litellm/router_utils/get_retry_from_policy.py index 7c412dc759b..835cfaacfd5 100644 --- a/litellm/router_utils/get_retry_from_policy.py +++ b/litellm/router_utils/get_retry_from_policy.py @@ -13,6 +13,7 @@ from litellm.exceptions import ( RateLimitError, ServiceUnavailableError, Timeout, + UnsupportedEndpointError, ) from litellm.types.router import RetryPolicy @@ -59,6 +60,8 @@ def get_num_retries_from_retry_policy( policy: Final = resolve_retry_policy(retry_policy, model_group, model_group_retry_policy) if policy is None: return None + if isinstance(exception, UnsupportedEndpointError): + return policy.DefaultRetries by_class: Final = ( _RETRIES_BY_EXCEPTION_TYPE[cls](policy) for cls in type(exception).__mro__ if cls in _RETRIES_BY_EXCEPTION_TYPE ) diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 3183b9687ca..eee02a2f6b6 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -95,6 +95,7 @@ ignored_function_names = [ "arm_routing_read_prefetch", # Tested in tests/unit/caching/test_request_redis_batch_pre_call.py (file lacks "router" in name) "_deployment_params_base_model", # Tested through log_retry in test_router_helper_utils.py "_configured_model_info", # Tested through get_configured_service_tiers in test_router.py + "_every_deployment_refused", # Tested through the unsupported transcription retry tests in test_router.py "_routable_deployments", # Tested through get_configured_service_tiers and get_routable_upstream_model in test_router.py "_async_get_available_deployment", # Body of the `route {model}` phase wrapper, exercised through async_get_available_deployment in test_router.py "_async_get_available_deployment_for_pass_through", # Same, through async_get_available_deployment_for_pass_through in test_router.py diff --git a/tests/unit/router_utils/test_cooldown_handlers.py b/tests/unit/router_utils/test_cooldown_handlers.py index 4299b61c703..0d39016ea79 100644 --- a/tests/unit/router_utils/test_cooldown_handlers.py +++ b/tests/unit/router_utils/test_cooldown_handlers.py @@ -35,6 +35,7 @@ from litellm.router_utils.cooldown_handlers import ( cast_exception_status_to_int, get_cooldown_deployments, mark_advisor_orchestration_failure, + set_cooldown_deployments, should_cooldown_based_on_allowed_fails_policy, ) from litellm.router_utils.fallback_event_handlers import ( @@ -55,6 +56,40 @@ from litellm.types.utils import CooldownDeploymentInfo, CredentialItem from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("exception_type", "expected_cooldown"), + ( + (litellm.APIConnectionError, False), + (litellm.UnsupportedEndpointError, False), + (litellm.BadRequestError, True), + ), +) +async def test_unsupported_endpoint_preserves_cooldown_exclusion_with_bad_request_policy( + exception_type: type[litellm.APIConnectionError | litellm.BadRequestError], expected_cooldown: bool +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "chat", + "litellm_params": {"model": "anthropic/claude-sonnet-5-5", "api_key": "not-needed"}, + "model_info": {"id": "deployment", "allowed_fails_policy": {"BadRequestErrorAllowedFails": 0}}, + } + ], + ) + exception: Final = exception_type(message="Request failed", model="claude-sonnet-5-5", llm_provider="anthropic") + set_cooldown_deployments( + litellm_router_instance=router, + original_exception=exception, + exception_status=exception.status_code, + deployment="deployment", + time_to_cooldown=60, + ) + assert await async_get_cooldown_deployments(router, parent_otel_span=None) == ( + ["deployment"] if expected_cooldown else [] + ) + + class TestGetDeploymentCooldownPolicy: def _make_router(self, deployment_id: str, model_info: dict | None = None): router = MagicMock() diff --git a/tests/unit/router_utils/test_get_retry_from_policy.py b/tests/unit/router_utils/test_get_retry_from_policy.py index 42580989a8c..4ff869849dd 100644 --- a/tests/unit/router_utils/test_get_retry_from_policy.py +++ b/tests/unit/router_utils/test_get_retry_from_policy.py @@ -89,6 +89,18 @@ def test_payment_required_error_uses_bad_request_retries(): ) +@pytest.mark.parametrize("default_retries", (None, 0, 3)) +@pytest.mark.parametrize("bad_request_retries", (0, 5)) +def test_unsupported_endpoint_preserves_the_connection_error_retry_budget( + default_retries: int | None, bad_request_retries: int +) -> None: + policy: Final = RetryPolicy(BadRequestErrorRetries=bad_request_retries, DefaultRetries=default_retries) + + assert get_num_retries_from_retry_policy( + exception=_error(litellm.UnsupportedEndpointError), retry_policy=policy + ) == get_num_retries_from_retry_policy(exception=_error(litellm.APIConnectionError), retry_policy=policy) + + @pytest.mark.parametrize("exception_type", (litellm.BadGatewayError,)) def test_default_retries_covers_exceptions_without_a_specific_field(exception_type: type[Exception]): exception: Final = _error(exception_type) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index f12129af1fc..ba39fc3f5ac 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -9780,6 +9780,37 @@ async def test_atranscription_answers_a_non_string_provider_hint_with_a_bad_requ assert not respx_mock.calls +def test_transcription_rejects_an_unsupported_provider_without_an_http_call(respx_mock: respx.MockRouter) -> None: + model: Final = "test-transcription-model" + provider: Final = "anthropic" + + with pytest.raises(litellm.UnsupportedEndpointError) as raised: + litellm.transcription(model=f"{provider}/{model}", file=("speech.wav", b"RIFF-audio", "audio/wav")) + + assert raised.value.status_code == 400 + assert raised.value.message == f"litellm.BadRequestError: {provider}/{model} does not support audio transcription." + assert raised.value.model == model + assert raised.value.llm_provider == provider + assert not respx_mock.calls + + +@pytest.mark.asyncio +async def test_atranscription_rejects_an_unsupported_provider_without_an_http_call( + respx_mock: respx.MockRouter, +) -> None: + model: Final = "test-transcription-model" + provider: Final = "anthropic" + + with pytest.raises(litellm.UnsupportedEndpointError) as raised: + await litellm.atranscription(model=f"{provider}/{model}", file=("speech.wav", b"RIFF-audio", "audio/wav")) + + assert raised.value.status_code == 400 + assert raised.value.message == f"litellm.BadRequestError: {provider}/{model} does not support audio transcription." + assert raised.value.model == model + assert raised.value.llm_provider == provider + assert not respx_mock.calls + + @pytest.mark.asyncio async def test_aspeech_infers_the_provider_from_base_url(respx_mock: respx.MockRouter) -> None: base_url: Final = "https://api.groq.com/openai/v1" diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 9875f1bd00a..63ee1c43edb 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -21917,6 +21917,116 @@ def test_router_retry_policy_400_moves_off_the_refusing_deployment_on_every_sync assert accepting.call_count == 1 +_TRANSCRIPTION_AUDIO: Final = ("speech.wav", b"RIFF", "audio/wav") +_TRANSCRIPTION_API_BASE: Final = "https://transcriber.local/v1" + + +def _transcription_deployment( + model_name: str, deployment_id: str, litellm_model: str, weight: int +) -> DeploymentTypedDict: + return { + "model_name": model_name, + "litellm_params": { + "model": litellm_model, + "api_key": "sk-fake", + "api_base": _TRANSCRIPTION_API_BASE, + "weight": weight, + }, + "model_info": {"id": deployment_id}, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("policy_scope", ["none", "router", "model_group", "request"]) +@pytest.mark.parametrize( + "caller_credentials", + ({}, {"api_key": "caller-key"}, {"api_base": _TRANSCRIPTION_API_BASE}, {"base_url": _TRANSCRIPTION_API_BASE}), + ids=("configured", "api_key", "api_base", "base_url"), +) +async def test_router_retries_a_transcription_off_a_deployment_whose_provider_cannot_transcribe( + monkeypatch: pytest.MonkeyPatch, + policy_scope: Literal["none", "router", "model_group", "request"], + caller_credentials: Mapping[str, str], +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + policy: Final = RetryPolicy(BadRequestErrorRetries=0) + router: Final = litellm.Router( + model_list=[ + _transcription_deployment("transcriber", "cannot-transcribe", "anthropic/claude-sonnet-5-5", weight=1), + _transcription_deployment("transcriber", "can-transcribe", "openai/gpt-transcribe", weight=0), + ], + num_retries=1, + retry_policy=policy if policy_scope == "router" else None, + model_group_retry_policy={"transcriber": policy} if policy_scope == "model_group" else None, + disable_cooldowns=True, + ) + request_kwargs: Final = {"model_group_retry_policy": {"transcriber": policy}} if policy_scope == "request" else {} + + with respx.mock as respx_mock: + upstream: Final = respx_mock.post(f"{_TRANSCRIPTION_API_BASE}/audio/transcriptions").mock( + return_value=httpx.Response(200, json={"text": "hi"}) + ) + response: Final = await router.atranscription( + model="transcriber", file=_TRANSCRIPTION_AUDIO, **request_kwargs, **caller_credentials + ) + + assert response.text == "hi" + assert upstream.call_count == 1 + + +@pytest.mark.parametrize("deployment_count", [1, 2, 3]) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "caller_credentials", + ({}, {"api_key": "caller-key"}, {"api_base": _TRANSCRIPTION_API_BASE}, {"base_url": _TRANSCRIPTION_API_BASE}), + ids=("configured", "api_key", "api_base", "base_url"), +) +async def test_router_tries_each_deployment_that_cannot_transcribe_once_then_answers_400( + deployment_count: int, caller_credentials: Mapping[str, str] +) -> None: + router: Final = litellm.Router( + model_list=[ + _transcription_deployment("transcriber", f"cannot-transcribe-{index}", "anthropic/claude-sonnet-5-5", 1) + for index in range(deployment_count) + ], + num_retries=5, + disable_cooldowns=True, + ) + + with respx.mock as respx_mock, pytest.raises(litellm.UnsupportedEndpointError) as raised: + await router.atranscription(model="transcriber", file=_TRANSCRIPTION_AUDIO, **caller_credentials) + + assert raised.value.status_code == 400 + assert router.fail_calls["transcriber"] == deployment_count + assert not respx_mock.calls + + +@pytest.mark.asyncio +async def test_router_falls_back_from_a_group_that_cannot_transcribe_when_unsupported_params_do_not( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = litellm.Router( + model_list=[ + _transcription_deployment("chat-only", "cannot-transcribe", "anthropic/claude-sonnet-5-5", weight=1), + _transcription_deployment("transcriber", "can-transcribe", "openai/gpt-transcribe", weight=1), + ], + fallbacks=[{"chat-only": ["transcriber"]}], + fallback_on_unsupported_params=False, + num_retries=0, + disable_cooldowns=True, + ) + + with respx.mock as respx_mock: + upstream: Final = respx_mock.post(f"{_TRANSCRIPTION_API_BASE}/audio/transcriptions").mock( + return_value=httpx.Response(200, json={"text": "hi"}) + ) + response: Final = await router.atranscription(model="chat-only", file=_TRANSCRIPTION_AUDIO) + + assert response.text == "hi" + assert upstream.call_count == 1 + + def _make_failure_logging_obj(): return LiteLLMLogging( model="gpt-5.6", From c9a5777739a38ef36281c161b69a98752cccdcc8 Mon Sep 17 00:00:00 2001 From: tessellark <349031000+tessellark@users.noreply.github.com> Date: Wed, 28 Oct 2026 12:24:20 -0700 Subject: [PATCH 073/128] fix(azure): share one Entra token request per worker and audit #47348 (#48769) * test(azure): audit #47348 Add proxy integration tests for #47348's cached Entra token providers and off-loop token refresh, and fix the token herd the audit found: once the fetch moved to a worker thread, concurrent cold requests each fetched their own token. Every cached bearer token provider now serializes its calls so the first fetch fills the cache and the rest reuse it * fix(azure): build each cached Entra token provider once under concurrent first calls lru_cache does not hold back concurrent misses, so worker threads that miss together each built their own credential and token provider, and each fetched its own token. A lock around each provider cache makes the first caller build it and the rest reuse it * fix(azure): share Entra token fetches only while no valid token is held Concurrent callers now share one result (or one error) from the cold fetch, while callers inside the refresh window keep the still-valid token. The oidc/azure provider build and the client-assertion exchange are single-flighted per audience, and the shared helpers no longer call subscripted generic aliases, which crashed on Python 3.10. * fix(azure): refresh Entra tokens once while other requests keep the valid token MSAL treats a token with less than five minutes left as expired, so every request in that window fetched a new one. Now one request refreshes while the others keep using the still-valid token. Requests that waited on a fetch that failed retry once through a fresh fetch instead of all failing with the first error. Credentials that only implement get_token keep working, and OIDC token exchanges are cached per scope. * fix(azure): bound Entra token waits during an outage Requests waiting on a failed token fetch now retry only when the fetch before it succeeded, so during an Entra outage later bursts fail after one fetch instead of two. Requests during a refresh only get the held token while it has more than 60 seconds left, and a failed refresh falls back to that token instead of failing its request. * fix(azure): share one retry per failed Entra token fetch and raise on a failed refresh Waiters of a failed fetch now retry the flight they joined, so a fast-failing Entra makes one retry instead of one per waiter. A failed refresh raises again instead of returning the held token, which Postgres token auth read as a fresh one. The DefaultAzureCredential fallback reuses one provider per scope. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/litellm_core_utils/single_flight.py | 92 +++ litellm/llms/azure/common_utils.py | 131 ++-- .../get_azure_ad_token_provider.py | 100 ++- litellm/secret_managers/main.py | 4 +- .../test_azure_entra_token_reuse_wire.py | 625 ++++++++++++++++++ tests/unit/llms/azure/test_azure.py | 8 +- .../llms/azure/test_azure_common_utils.py | 140 +++- .../test_get_azure_ad_token_provider.py | 270 +++++++- .../test_secret_managers_main.py | 40 ++ 9 files changed, 1335 insertions(+), 75 deletions(-) create mode 100644 litellm/litellm_core_utils/single_flight.py create mode 100644 tests/integration/providers/test_azure_entra_token_reuse_wire.py diff --git a/litellm/litellm_core_utils/single_flight.py b/litellm/litellm_core_utils/single_flight.py new file mode 100644 index 00000000000..44a84d05925 --- /dev/null +++ b/litellm/litellm_core_utils/single_flight.py @@ -0,0 +1,92 @@ +import threading +from collections.abc import Callable +from concurrent.futures import Future +from dataclasses import dataclass, field +from functools import lru_cache +from typing import Final, Generic, ParamSpec, TypeVar + +_P: Final = ParamSpec("_P") +_R: Final = TypeVar("_R") + + +@dataclass(slots=True) +class SingleFlight(Generic[_R]): + lock: threading.Lock = field(default_factory=threading.Lock) + in_flight: Future[_R] | None = None + flights_started: int = 0 + last_flight_failed: bool = False + + def run(self, fetch: Callable[[], _R], while_busy: _R | None = None) -> _R: + leader: Final[Future[_R]] = Future() + with self.lock: + in_flight: Final = self.in_flight + joined_flight: Final = self.flights_started + retry_if_it_fails: Final = not self.last_flight_failed + if in_flight is None: + self._take_off(leader) + if in_flight is None: + return self._lead(leader, fetch) + if while_busy is not None: + return while_busy + try: + return in_flight.result() + except Exception: + retry: Final = self._claim_retry(joined_flight) if retry_if_it_fails else None + if retry is None: + raise + retry_flight, leads_retry = retry + return self._lead(retry_flight, fetch) if leads_retry else retry_flight.result() + + def _take_off(self, leader: Future[_R]) -> None: + self.in_flight = leader + self.flights_started += 1 + + def _claim_retry(self, failed_flight: int) -> tuple[Future[_R], bool] | None: + leader: Final[Future[_R]] = Future() + with self.lock: + in_flight: Final = self.in_flight + if in_flight is not None: + return in_flight, False + if self.flights_started > failed_flight and self.last_flight_failed: + return None + self._take_off(leader) + return leader, True + + def _lead(self, flight: Future[_R], fetch: Callable[[], _R]) -> _R: + try: + result: Final = fetch() + except BaseException as error: + self._land(failed=True) + flight.set_exception(error) + raise + self._land(failed=False) + flight.set_result(result) + return result + + def _land(self, failed: bool) -> None: + with self.lock: + self.in_flight = None + self.last_flight_failed = failed + + +@dataclass(frozen=True, slots=True) +class SingleFlightCache(Generic[_P, _R]): + cached: Callable[..., _R] + clear: Callable[[], None] + lock: threading.Lock = field(default_factory=threading.Lock) + + def __call__( + self, + *args: _P.args, + **kwargs: _P.kwargs, # kwargs-ok: forwarded verbatim to the cached builder, typed by its own signature + ) -> _R: + with self.lock: + return self.cached(*args, **kwargs) + + def cache_clear(self) -> None: + self.clear() + + +def single_flight_lru_cache(build: Callable[_P, _R]) -> SingleFlightCache[_P, _R]: + cached: Final = lru_cache(maxsize=128)(build) + return SingleFlightCache(cached, cached.cache_clear) diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 23c5aa35f4b..4bd19889bcc 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -5,7 +5,6 @@ import os import re from collections.abc import Awaitable, Callable, Mapping from contextlib import closing -from functools import lru_cache from types import MappingProxyType from typing import Final, Literal, NamedTuple, cast @@ -18,15 +17,18 @@ import litellm from litellm._logging import verbose_logger from litellm.caching.caching import DualCache from litellm.constants import DEFAULT_MAX_RETRIES +from litellm.litellm_core_utils.single_flight import SingleFlight, single_flight_lru_cache from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.openai.common_utils import BaseOpenAILLM, OpenAIAsyncHTTPClient, OpenAIHTTPClient from litellm.secret_managers.get_azure_ad_token_provider import ( get_azure_ad_token_provider, + single_flight_bearer_token_provider, ) from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams +from litellm.types.secret_managers.get_azure_ad_token_provider import AzureCredentialType from litellm.utils import add_path_to_api_base azure_ad_cache: Final = DualCache() @@ -94,7 +96,7 @@ def process_azure_headers(headers: httpx.Headers | dict) -> dict: return {**llm_response_headers, **openai_headers} -@lru_cache(maxsize=128) +@single_flight_lru_cache def _cached_entra_id_token_provider( tenant_id: str, client_id: str, @@ -107,33 +109,38 @@ def _cached_entra_id_token_provider( to expiry, so reusing the provider keeps one AAD round trip per token lifetime instead of one per request. """ - from azure.identity import ClientSecretCredential, get_bearer_token_provider + from azure.identity import ClientSecretCredential - return get_bearer_token_provider(ClientSecretCredential(tenant_id, client_id, client_secret), scope) + return single_flight_bearer_token_provider(ClientSecretCredential(tenant_id, client_id, client_secret), scope) -@lru_cache(maxsize=128) +@single_flight_lru_cache def _cached_username_password_token_provider( client_id: str, azure_username: str, azure_password: str, scope: str, ) -> Callable[[], str]: - from azure.identity import UsernamePasswordCredential, get_bearer_token_provider + from azure.identity import UsernamePasswordCredential credential: Final = UsernamePasswordCredential( client_id=client_id, username=azure_username, password=azure_password, ) - return get_bearer_token_provider(credential, scope) + return single_flight_bearer_token_provider(credential, scope) -@lru_cache(maxsize=128) +@single_flight_lru_cache def _cached_azure_ad_token_refresh_provider(scope: str) -> Callable[[], str]: return get_azure_ad_token_provider(azure_scope=scope) +@single_flight_lru_cache +def _cached_default_azure_credential_provider(scope: str) -> Callable[[], str]: + return get_azure_ad_token_provider(azure_scope=scope, azure_credential=AzureCredentialType.DefaultAzureCredential) + + def cached_azure_ad_token_refresh_provider(scope: str) -> Callable[[], str]: return _cached_azure_ad_token_refresh_provider(scope) @@ -240,6 +247,58 @@ def get_azure_ad_token_from_username_password( return token_provider +@single_flight_lru_cache +def _oidc_token_exchange(cache_key: str) -> SingleFlight[str]: + return SingleFlight() + + +def _cached_oidc_token(cache_key: str) -> str | None: + cached: Final = azure_ad_cache.get_cache(cache_key) + return cached if isinstance(cached, str) else None + + +def _exchange_oidc_token(cache_key: str, token_url: str, client_id: str, scope: str, oidc_token: str) -> str: + cached_access_token: Final = _cached_oidc_token(cache_key) + if cached_access_token is not None: + return cached_access_token + + with closing(HTTPHandler(timeout=litellm.request_timeout)) as client: + req_token: Final = client.post( + token_url, + data={ + "client_id": client_id, + "grant_type": "client_credentials", + "scope": scope, + "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer", + "client_assertion": oidc_token, + }, + ) + + if req_token.status_code != 200: + raise AzureOpenAIError( + status_code=req_token.status_code, + message=req_token.text, + ) + + azure_ad_token_json: Final[_AzureAdTokenJson] = req_token.json() + azure_ad_token_access_token: Final = azure_ad_token_json.get("access_token", None) + azure_ad_token_expires_in: Final = azure_ad_token_json.get("expires_in", None) + + if azure_ad_token_access_token is None: + raise AzureOpenAIError(status_code=422, message="Azure AD Token access_token not returned") + + if azure_ad_token_expires_in is None: + raise AzureOpenAIError(status_code=422, message="Azure AD Token expires_in not returned") + + azure_ad_cache.set_cache( + key=cache_key, + value=azure_ad_token_access_token, + ttl=azure_ad_token_expires_in, + ) + + return azure_ad_token_access_token + + def get_azure_ad_token_from_oidc( azure_ad_token: str, azure_client_id: str | None = None, @@ -282,50 +341,25 @@ def get_azure_ad_token_from_oidc( "azure_client_id": azure_client_id, "azure_tenant_id": azure_tenant_id, "azure_authority_host": azure_authority_host, + "scope": scope, "oidc_token": oidc_token, } ) - azure_ad_token_access_token = azure_ad_cache.get_cache(azure_ad_token_cache_key) - if azure_ad_token_access_token is not None: - return azure_ad_token_access_token + cached_access_token: Final = _cached_oidc_token(azure_ad_token_cache_key) + if cached_access_token is not None: + return cached_access_token - with closing(HTTPHandler(timeout=litellm.request_timeout)) as client: - req_token: Final = client.post( - f"{azure_authority_host}/{azure_tenant_id}/oauth2/v2.0/token", - data={ - "client_id": azure_client_id, - "grant_type": "client_credentials", - "scope": scope, - "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer", - "client_assertion": oidc_token, - }, + return _oidc_token_exchange(azure_ad_token_cache_key).run( + lambda: _exchange_oidc_token( + azure_ad_token_cache_key, + token_url=f"{azure_authority_host}/{azure_tenant_id}/oauth2/v2.0/token", + client_id=azure_client_id, + scope=scope, + oidc_token=oidc_token, ) - - if req_token.status_code != 200: - raise AzureOpenAIError( - status_code=req_token.status_code, - message=req_token.text, - ) - - azure_ad_token_json: Final[_AzureAdTokenJson] = req_token.json() - azure_ad_token_access_token = azure_ad_token_json.get("access_token", None) - azure_ad_token_expires_in: Final = azure_ad_token_json.get("expires_in", None) - - if azure_ad_token_access_token is None: - raise AzureOpenAIError(status_code=422, message="Azure AD Token access_token not returned") - - if azure_ad_token_expires_in is None: - raise AzureOpenAIError(status_code=422, message="Azure AD Token expires_in not returned") - - azure_ad_cache.set_cache( - key=azure_ad_token_cache_key, - value=azure_ad_token_access_token, - ttl=azure_ad_token_expires_in, ) - return azure_ad_token_access_token - def select_azure_base_url_or_endpoint(azure_client_params: dict): azure_endpoint: Final = azure_client_params.get("azure_endpoint", None) @@ -503,17 +537,10 @@ class BaseAzureLLM(BaseOpenAILLM): Returns: Token provider callable if DefaultAzureCredential is enabled and available, None otherwise """ - from litellm.types.secret_managers.get_azure_ad_token_provider import ( - AzureCredentialType, - ) - verbose_logger.debug("Attempting to use DefaultAzureCredential for Azure Auth") try: - azure_ad_token_provider: Final = get_azure_ad_token_provider( - azure_scope=scope, - azure_credential=AzureCredentialType.DefaultAzureCredential, - ) + azure_ad_token_provider: Final = _cached_default_azure_credential_provider(scope) verbose_logger.debug("Successfully obtained Azure AD token provider using DefaultAzureCredential") return azure_ad_token_provider except Exception as e: diff --git a/litellm/secret_managers/get_azure_ad_token_provider.py b/litellm/secret_managers/get_azure_ad_token_provider.py index 2b707f179d0..4967e199218 100644 --- a/litellm/secret_managers/get_azure_ad_token_provider.py +++ b/litellm/secret_managers/get_azure_ad_token_provider.py @@ -1,14 +1,27 @@ import os +import time from collections.abc import Callable +from dataclasses import dataclass +from types import TracebackType from typing import TYPE_CHECKING, Final from litellm._logging import verbose_logger +from litellm.litellm_core_utils.single_flight import SingleFlight from litellm.types.secret_managers.get_azure_ad_token_provider import ( AzureCredentialType, ) +_HELD_TOKEN_EXPIRY_MARGIN_SECONDS: Final = 60 + if TYPE_CHECKING: - from azure.core.credentials import TokenProvider + from azure.core.credentials import ( + AccessToken, + AccessTokenInfo, + SupportsTokenInfo, + TokenCredential, + TokenRequestOptions, + ) + from typing_extensions import TypeIs # noqa: TID251 # tells the two azure-core credential protocols apart def infer_credential_type_from_environment() -> AzureCredentialType: @@ -37,6 +50,86 @@ def infer_credential_type_from_environment() -> AzureCredentialType: return AzureCredentialType.DefaultAzureCredential +@dataclass(slots=True) +class _HeldToken: + token: "AccessToken | AccessTokenInfo | None" = None + + def still_valid(self) -> str | None: + held: Final = self.token + if held is None or held.expires_on - time.time() <= _HELD_TOKEN_EXPIRY_MARGIN_SECONDS: + return None + return held.token + + +@dataclass(slots=True) +class _TokenInfoWatcher: + credential: "SupportsTokenInfo" + held: _HeldToken + + def get_token_info(self, *scopes: str, options: "TokenRequestOptions | None" = None) -> "AccessTokenInfo": + token: Final = self.credential.get_token_info(*scopes, options=options) + self.held.token = token + return token + + def close(self) -> None: + self.credential.close() + + def __enter__(self) -> "SupportsTokenInfo": + return self.credential.__enter__() + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + self.credential.__exit__(exc_type, exc_value, traceback) + + +@dataclass(slots=True) +class _TokenWatcher: + credential: "TokenCredential" + held: _HeldToken + + def get_token( + self, + *scopes: str, + claims: str | None = None, + tenant_id: str | None = None, + enable_cae: bool = False, + **kwargs: object, # kwargs-ok: azure-core forwards its request options verbatim to the wrapped credential + ) -> "AccessToken": + token: Final = self.credential.get_token( + *scopes, claims=claims, tenant_id=tenant_id, enable_cae=enable_cae, **kwargs + ) + self.held.token = token + return token + + +def _supports_token_info( + credential: "TokenCredential | SupportsTokenInfo", +) -> "TypeIs[SupportsTokenInfo]": + return hasattr(credential, "get_token_info") + + +def single_flight_bearer_token_provider( + credential: "TokenCredential | SupportsTokenInfo", scope: str +) -> Callable[[], str]: + from azure.identity import get_bearer_token_provider + + held: Final = _HeldToken() + watched: Final = ( + _TokenInfoWatcher(credential, held) if _supports_token_info(credential) else _TokenWatcher(credential, held) + ) + provider: Final = get_bearer_token_provider(watched, scope) + token_request: Final[SingleFlight[str]] = SingleFlight() + + def bearer_token() -> str: + return token_request.run(provider, while_busy=held.still_valid()) + + return bearer_token + + def get_azure_ad_token_provider( azure_scope: str | None = None, azure_credential: AzureCredentialType | None = None, @@ -65,7 +158,6 @@ def get_azure_ad_token_provider( DefaultAzureCredential, ManagedIdentityCredential, WorkloadIdentityCredential, - get_bearer_token_provider, ) if azure_scope is None: @@ -77,7 +169,7 @@ def get_azure_ad_token_provider( else None or os.environ.get("AZURE_CREDENTIAL") or infer_credential_type_from_environment() ) verbose_logger.info("For Azure AD Token Provider, choosing credential type: %s", cred) - credential: TokenProvider | None = None + credential: TokenCredential | SupportsTokenInfo | None = None if cred == AzureCredentialType.ClientSecretCredential: credential = ClientSecretCredential( client_id=os.environ["AZURE_CLIENT_ID"], @@ -133,4 +225,4 @@ def get_azure_ad_token_provider( if credential is None: raise ValueError("No credential provided") - return get_bearer_token_provider(credential, azure_scope) + return single_flight_bearer_token_provider(credential, azure_scope) diff --git a/litellm/secret_managers/main.py b/litellm/secret_managers/main.py index b33cf0e6b7e..588bdb1409c 100644 --- a/litellm/secret_managers/main.py +++ b/litellm/secret_managers/main.py @@ -5,7 +5,6 @@ import time import traceback from collections.abc import Callable from dataclasses import dataclass -from functools import lru_cache from typing import Final import httpx @@ -14,6 +13,7 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_logger from litellm.caching.caching import DualCache +from litellm.litellm_core_utils.single_flight import single_flight_lru_cache from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.secret_managers.get_azure_ad_token_provider import ( get_azure_ad_token_provider, @@ -27,7 +27,7 @@ _PARSED_LITERAL: Final = TypeAdapter(object) _SECRET_MANAGER_CLIENT: Final = TypeAdapter(object) -@lru_cache(maxsize=128) +@single_flight_lru_cache def _cached_azure_oidc_token_provider(audience: str) -> Callable[[], str]: return get_azure_ad_token_provider(azure_scope=audience) diff --git a/tests/integration/providers/test_azure_entra_token_reuse_wire.py b/tests/integration/providers/test_azure_entra_token_reuse_wire.py new file mode 100644 index 00000000000..85aef110124 --- /dev/null +++ b/tests/integration/providers/test_azure_entra_token_reuse_wire.py @@ -0,0 +1,625 @@ +import asyncio +import itertools +import json +import threading +import time +import uuid +from collections.abc import Callable, Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from enum import StrEnum +from pathlib import Path +from queue import SimpleQueue +from typing import Final +from urllib.parse import parse_qs, urlsplit + +import httpx +import pytest +import websockets +import yaml +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + eventually, + gateway_from_environment, + object_value, + string_value, +) +from integration._support.openai_wire import chat_reply, responses_reply +from integration._support.ports import allocate_port +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.responses_ws_peer import Peer, completed, responses_peer, session +from integration._support.tls import write_self_signed_cert +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import BaseModel, ConfigDict, JsonValue + +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state] + +_DEPLOYMENT: Final = "gpt-5.4-nano" +_DATED_API_VERSION: Final = "2024-10-21" +_REALTIME_API_VERSION: Final = "2024-10-01-preview" +_IDENTITY_PATH: Final = "/msi/token" +_EXCHANGE_SUFFIX: Final = "/oauth2/v2.0/token" +_FOUNDRY_PROJECT: Final = "/api/projects/audit" +_BING_TOOL: Final = "audit-bing-entra" +_BING_RESOURCE: Final = "https://ai.azure.com" +_BING_CITATION: Final = "https://result.invalid/bing" +_TENANT: Final = "audit-tenant" +_APP_CLIENT: Final = "audit-app-client" +_HOLD_SECONDS: Final = 60 +_LIVELINESS_SECONDS: Final = 3 +_HERD: Final = 6 +_HERD_WINDOW_SECONDS: Final = 3 +_TOKEN_LIFETIME_SECONDS: Final = 3600 +_INSIDE_REFRESH_WINDOW_SECONDS: Final = 120 +_REMOVED_ENVIRONMENT: Final = ( + "AZURE_API_KEY", + "AZURE_OPENAI_API_KEY", + "AZURE_AD_TOKEN", + "AZURE_TENANT_ID", + "AZURE_CLIENT_SECRET", + "AZURE_FEDERATED_TOKEN_FILE", + "AZURE_USERNAME", + "AZURE_PASSWORD", + "AZURE_CREDENTIAL", + "AZURE_SCOPE", + "MSI_ENDPOINT", + "MSI_SECRET", + "IMDS_ENDPOINT", + "IDENTITY_SERVER_THUMBPRINT", + "BING_GROUNDING_TOKEN", + "HTTP_PROXY", + "HTTPS_PROXY", + "http_proxy", + "https_proxy", +) + + +class _Surface(StrEnum): + LEGACY_CHAT = "legacy-chat" + V1_CHAT = "v1-chat" + RESPONSES = "responses" + RESPONSES_WEBSOCKET = "responses-websocket" + REALTIME = "realtime" + REALTIME_HEALTH_CHECK = "realtime-health-check" + + +_HTTP_SURFACES: Final = frozenset({_Surface.LEGACY_CHAT, _Surface.V1_CHAT, _Surface.RESPONSES}) + + +@dataclass(frozen=True, slots=True) +class _Rig: + gateway: Gateway + port: int + cert: tuple[Path, Path] + + @property + def edge(self) -> str: + return f"http://127.0.0.1:{self.port}" + + +@dataclass(frozen=True, slots=True) +class _Hold: + resource: str + arrivals: SimpleQueue[str] = field(default_factory=SimpleQueue) + release: threading.Event = field(default_factory=threading.Event) + + +def _config(directory: Path) -> Path: + stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + path: Final = directory / "azure-entra-token-reuse.yaml" + path.write_text( + yaml.safe_dump( + { + **stock, + "litellm_settings": {**object_value(stock["litellm_settings"]), "enable_azure_ad_token_refresh": True}, + "router_settings": {**object_value(stock["router_settings"]), "num_retries": 0}, + "search_tools": [ + {"search_tool_name": _BING_TOOL, "litellm_params": {"search_provider": "bing_grounding"}} + ], + } + ) + ) + return path + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("azure-entra-token-reuse") + port: Final = allocate_port() + cert: Final = write_self_signed_cert(directory) + edge: Final = f"http://127.0.0.1:{port}" + environment: Final = { + "AZURE_CLIENT_ID": "audit-managed-identity", + "IDENTITY_ENDPOINT": f"{edge}{_IDENTITY_PATH}", + "IDENTITY_HEADER": "audit-identity-header", + "AZURE_AUTHORITY_HOST": edge, + "BING_GROUNDING_PROJECT_ENDPOINT": f"{edge}{_FOUNDRY_PROJECT}", + "BING_GROUNDING_MODEL": _DEPLOYMENT, + "BING_GROUNDING_CONNECTION_ID": "audit-bing-connection", + "SSL_CERT_FILE": str(cert[0]), + } + with ( + gateway_from_environment() as shared, + owned_proxy( + shared, directory, environment, config=_config(directory), remove_environment=_REMOVED_ENVIRONMENT + ) as proxy, + ): + yield _Rig(proxy, port, cert) + + +def _json(payload: Mapping[str, JsonValue]) -> Reply: + return Reply(body=json.dumps(payload).encode()) + + +def _identity_token(resource: str) -> str: + return f"mi-token:{resource}" + + +def _identity_resource(request: Request) -> str | None: + target: Final = urlsplit(request.target) + if target.path != _IDENTITY_PATH: + return None + return parse_qs(target.query)["resource"][-1] + + +def _is_token_endpoint(request: Request) -> bool: + return _identity_resource(request) is not None or request.target.endswith(_EXCHANGE_SUFFIX) + + +def _managed_identity( + resource: str, + hold: _Hold | None, + lifetime_seconds: int = _TOKEN_LIFETIME_SECONDS, + access_token: str | None = None, +) -> Reply: + if hold is not None and resource == hold.resource: + hold.arrivals.put(resource) + hold.release.wait(_HOLD_SECONDS) + return _json( + { + "access_token": access_token or _identity_token(resource), + "expires_on": str(int(time.time()) + lifetime_seconds), + "resource": resource, + "token_type": "Bearer", + } + ) + + +def _entra_exchange(request: Request) -> Reply: + assertion: Final = parse_qs(request.body.decode())["client_assertion"][-1] + return _json({"access_token": f"aad:{assertion}", "expires_in": 3600, "token_type": "Bearer"}) + + +def _edge(upstream: Callable[[Request], Reply], hold: _Hold | None = None) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + resource: Final = _identity_resource(request) + if resource is not None: + return _managed_identity(resource, hold) + if request.target.endswith(_EXCHANGE_SUFFIX): + return _entra_exchange(request) + return upstream(request) + + return respond + + +def _azure_openai(request: Request) -> Reply: + if urlsplit(request.target).path.endswith("/chat/completions"): + return chat_reply(f"chatcmpl-{uuid.uuid4().hex}", _DEPLOYMENT, "azure answered", stream=False) + return responses_reply(f"resp_{uuid.uuid4().hex}", _DEPLOYMENT, "azure answered", stream=False) + + +def _identity_fetches(requests: tuple[Request, ...], resource: str) -> int: + return sum(1 for request in requests if _identity_resource(request) == resource) + + +def _provider_calls(requests: tuple[Request, ...]) -> tuple[Request, ...]: + return tuple(request for request in requests if not _is_token_endpoint(request)) + + +def _marker() -> str: + return f"entra audit {uuid.uuid4().hex}" + + +def _model(scenario: Scenario, surface: _Surface, api_base: str, scope: str) -> str: + match surface: + case _Surface.LEGACY_CHAT | _Surface.RESPONSES: + return scenario.model( + model=f"azure/{_DEPLOYMENT}", + api_base=api_base, + api_key=None, + api_version=_DATED_API_VERSION, + azure_scope=scope, + ) + case _Surface.V1_CHAT: + return scenario.model( + model=f"azure/{_DEPLOYMENT}", api_base=api_base, api_key=None, api_version="v1", azure_scope=scope + ) + case _Surface.RESPONSES_WEBSOCKET: + return scenario.model(model=f"azure/{_DEPLOYMENT}", api_base=api_base, api_key=None, azure_scope=scope) + case _Surface.REALTIME | _Surface.REALTIME_HEALTH_CHECK: + return scenario.model( + model_info={"mode": "realtime"}, + model=f"azure/{_DEPLOYMENT}", + api_base=api_base, + api_key=None, + api_version=_REALTIME_API_VERSION, + azure_scope=scope, + ) + + +def _post(surface: _Surface, gateway: Gateway, key: str, model: str) -> httpx.Response: + if surface is _Surface.RESPONSES: + return gateway.request("POST", "/v1/responses", {"model": model, "input": _marker()}, key=key) + return gateway.request( + "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": _marker()}]}, key=key + ) + + +async def _realtime_handshake(gateway: Gateway, key: str, model: str, peer: Peer) -> str | None: + proxy: Final = str(gateway.client.base_url).rstrip("/").replace("http://", "ws://") + async with websockets.connect( + f"{proxy}/v1/realtime?model={model}", + additional_headers={"Authorization": f"Bearer {key}"}, + open_timeout=_HOLD_SECONDS, + ): + credentials: Final = await asyncio.to_thread(peer.credentials.get, True, _HOLD_SECONDS) + return credentials.authorization + + +def _upstream_authorization(surface: _Surface, gateway: Gateway, key: str, model: str, wire: Wire, peer: Peer) -> str: + match surface: + case _Surface.LEGACY_CHAT | _Surface.V1_CHAT | _Surface.RESPONSES: + response: Final = _post(surface, gateway, key, model) + assert response.status_code == 200, response.text + (call,) = _provider_calls(wire.drain()) + return call.headers.get("authorization", "") + case _Surface.RESPONSES_WEBSOCKET: + (events,) = asyncio.run(session(str(gateway.client.base_url), key, model, (_marker(),))) + assert completed(events)["status"] == "completed", events + return peer.credentials.get(timeout=_HOLD_SECONDS).authorization or "" + case _Surface.REALTIME: + return asyncio.run(_realtime_handshake(gateway, key, model, peer)) or "" + case _Surface.REALTIME_HEALTH_CHECK: + gateway.request("GET", "/health", params={"model": model}) + return peer.credentials.get(timeout=_HOLD_SECONDS).authorization or "" + + +def _liveliness(gateway: Gateway) -> int | None: + try: + return httpx.get( + f"{str(gateway.client.base_url).rstrip('/')}/health/liveliness", timeout=_LIVELINESS_SECONDS + ).status_code + except httpx.TimeoutException: + return None + + +def _liveliness_while_held(gateway: Gateway, hold: _Hold) -> int | None: + try: + assert hold.arrivals.get(timeout=_HOLD_SECONDS) == hold.resource + return _liveliness(gateway) + finally: + hold.release.set() + + +@pytest.mark.parametrize("surface", tuple(_Surface)) +def test_a_slow_managed_identity_fetch_leaves_the_proxy_answering_other_requests(rig: _Rig, surface: _Surface) -> None: + resource: Final = f"api://audit-hold-{uuid.uuid4().hex}" + hold: Final = _Hold(resource) + with ( + wire_server(_edge(_azure_openai, hold), port=rig.port) as wire, + responses_peer(rig.cert) as peer, + rig.gateway.scenario() as scenario, + ThreadPoolExecutor(max_workers=1) as pool, + ): + api_base: Final = rig.edge if surface in _HTTP_SURFACES else peer.url.removesuffix("/v1") + model: Final = _model(scenario, surface, api_base, f"{resource}/.default") + key: Final = scenario.key() + call: Final = pool.submit(_upstream_authorization, surface, rig.gateway, key, model, wire, peer) + liveliness: Final = _liveliness_while_held(rig.gateway, hold) + authorization: Final = call.result(timeout=_HOLD_SECONDS) + + assert liveliness == 200, f"/health/liveliness did not answer within {_LIVELINESS_SECONDS}s during the fetch" + assert authorization == f"Bearer {_identity_token(resource)}" + + +@pytest.mark.parametrize("surface", (_Surface.LEGACY_CHAT, _Surface.RESPONSES)) +def test_concurrent_cold_requests_share_one_managed_identity_fetch(rig: _Rig, surface: _Surface) -> None: + resource: Final = f"api://audit-herd-{uuid.uuid4().hex}" + hold: Final = _Hold(resource) + with ( + wire_server(_edge(_azure_openai, hold), port=rig.port) as wire, + rig.gateway.scenario() as scenario, + ThreadPoolExecutor(max_workers=_HERD) as pool, + ): + model: Final = _model(scenario, surface, rig.edge, f"{resource}/.default") + key: Final = scenario.key() + calls: Final = tuple(pool.submit(_post, surface, rig.gateway, key, model) for _ in range(_HERD)) + try: + eventually( + hold.arrivals.qsize, + lambda arrived: arrived >= _HERD, + seconds=_HERD_WINDOW_SECONDS, + return_last_on_timeout=True, + ) + finally: + hold.release.set() + responses: Final = tuple(call.result(timeout=_HOLD_SECONDS) for call in calls) + requests: Final = wire.drain() + + assert [response.status_code for response in responses] == [200] * _HERD, [r.text for r in responses] + assert _identity_fetches(requests, resource) == 1 + assert [call.headers.get("authorization") for call in _provider_calls(requests)] == [ + f"Bearer {_identity_token(resource)}" + ] * _HERD + + +def _scripted_identity(resource: str, fetches: tuple[Callable[[], Reply], ...]) -> Callable[[Request], Reply]: + attempts: Final = itertools.count() + edge: Final = _edge(_azure_openai) + + def respond(request: Request) -> Reply: + if _identity_resource(request) != resource: + return edge(request) + return fetches[min(next(attempts), len(fetches) - 1)]() + + return respond + + +def _held_failure(hold: _Hold) -> Reply: + hold.arrivals.put(hold.resource) + hold.release.wait(_HOLD_SECONDS) + return Reply(body=b"not json") + + +def _herd( + rig: _Rig, edge: Callable[[Request], Reply], resource: str, hold: _Hold, warm_up: bool +) -> tuple[tuple[httpx.Response, ...], tuple[Request, ...]]: + with ( + wire_server(edge, port=rig.port) as wire, + rig.gateway.scenario() as scenario, + ThreadPoolExecutor(max_workers=_HERD) as pool, + ): + model: Final = _model(scenario, _Surface.LEGACY_CHAT, rig.edge, f"{resource}/.default") + key: Final = scenario.key() + warm_up_responses: Final = (_post(_Surface.LEGACY_CHAT, rig.gateway, key, model),) if warm_up else () + calls: Final = tuple(pool.submit(_post, _Surface.LEGACY_CHAT, rig.gateway, key, model) for _ in range(_HERD)) + try: + eventually( + hold.arrivals.qsize, + lambda arrived: arrived >= _HERD, + seconds=_HERD_WINDOW_SECONDS, + return_last_on_timeout=True, + ) + finally: + hold.release.set() + return warm_up_responses + tuple(call.result(timeout=_HOLD_SECONDS) for call in calls), wire.drain() + + +def test_concurrent_requests_inside_the_refresh_window_share_one_refresh_and_keep_the_valid_token(rig: _Rig) -> None: + resource: Final = f"api://audit-refresh-{uuid.uuid4().hex}" + refreshed: Final = f"mi-token-refreshed:{resource}" + hold: Final = _Hold(resource) + edge: Final = _scripted_identity( + resource, + ( + lambda: _managed_identity(resource, None, _INSIDE_REFRESH_WINDOW_SECONDS), + lambda: _managed_identity(resource, hold, access_token=refreshed), + ), + ) + + responses, requests = _herd(rig, edge, resource, hold, warm_up=True) + + assert [response.status_code for response in responses] == [200] * (_HERD + 1), [r.text for r in responses] + assert _identity_fetches(requests, resource) == 2 + assert sorted(call.headers.get("authorization", "") for call in _provider_calls(requests)) == sorted( + [f"Bearer {_identity_token(resource)}"] * _HERD + [f"Bearer {refreshed}"] + ) + + +def test_a_failed_cold_fetch_fails_only_the_request_that_made_it(rig: _Rig) -> None: + resource: Final = f"api://audit-failure-{uuid.uuid4().hex}" + hold: Final = _Hold(resource) + edge: Final = _scripted_identity(resource, (lambda: _held_failure(hold), lambda: _managed_identity(resource, None))) + + responses, _ = _herd(rig, edge, resource, hold, warm_up=False) + + assert sorted(response.status_code == 200 for response in responses) == [False] + [True] * (_HERD - 1), [ + r.text for r in responses + ] + + +def _oidc_model(scenario: Scenario, rig: _Rig, audience: str) -> str: + return scenario.model( + model=f"azure/{_DEPLOYMENT}", + api_base=rig.edge, + api_key=None, + api_version=_DATED_API_VERSION, + azure_ad_token=f"oidc/azure/{audience}", + tenant_id=_TENANT, + client_id=_APP_CLIENT, + ) + + +def _exchanges(requests: tuple[Request, ...]) -> tuple[Request, ...]: + return tuple(request for request in requests if request.target.endswith(_EXCHANGE_SUFFIX)) + + +def test_an_oidc_federated_deployment_fetches_its_assertion_once_across_responses_calls(rig: _Rig) -> None: + audience: Final = f"api://audit-oidc-{uuid.uuid4().hex}" + with wire_server(_edge(_azure_openai), port=rig.port) as wire, rig.gateway.scenario() as scenario: + model: Final = _oidc_model(scenario, rig, audience) + key: Final = scenario.key() + responses: Final = tuple(_post(_Surface.RESPONSES, rig.gateway, key, model) for _ in range(3)) + requests: Final = wire.drain() + + assert [response.status_code for response in responses] == [200] * 3, [r.text for r in responses] + assert _identity_fetches(requests, audience) == 1 + exchanges: Final = _exchanges(requests) + assert [request.target for request in exchanges] == [f"/{_TENANT}{_EXCHANGE_SUFFIX}"] + assert {name: values[-1] for name, values in parse_qs(exchanges[0].body.decode()).items()} == { + "client_id": _APP_CLIENT, + "grant_type": "client_credentials", + "scope": "https://cognitiveservices.azure.com/.default", + "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer", + "client_assertion": _identity_token(audience), + } + assert [call.headers.get("authorization") for call in _provider_calls(requests)] == [ + f"Bearer aad:{_identity_token(audience)}" + ] * 3 + + +def test_concurrent_cold_oidc_requests_share_one_assertion_fetch_and_one_exchange(rig: _Rig) -> None: + audience: Final = f"api://audit-oidc-herd-{uuid.uuid4().hex}" + hold: Final = _Hold(audience) + with ( + wire_server(_edge(_azure_openai, hold), port=rig.port) as wire, + rig.gateway.scenario() as scenario, + ThreadPoolExecutor(max_workers=_HERD) as pool, + ): + model: Final = _oidc_model(scenario, rig, audience) + key: Final = scenario.key() + calls: Final = tuple(pool.submit(_post, _Surface.RESPONSES, rig.gateway, key, model) for _ in range(_HERD)) + try: + eventually( + hold.arrivals.qsize, + lambda arrived: arrived >= _HERD, + seconds=_HERD_WINDOW_SECONDS, + return_last_on_timeout=True, + ) + finally: + hold.release.set() + responses: Final = tuple(call.result(timeout=_HOLD_SECONDS) for call in calls) + requests: Final = wire.drain() + + assert [response.status_code for response in responses] == [200] * _HERD, [r.text for r in responses] + assert _identity_fetches(requests, audience) == 1 + assert len(_exchanges(requests)) == 1 + assert [call.headers.get("authorization") for call in _provider_calls(requests)] == [ + f"Bearer aad:{_identity_token(audience)}" + ] * _HERD + + +def _foundry(request: Request) -> Reply: + assert urlsplit(request.target).path == f"{_FOUNDRY_PROJECT}/openai/v1/responses", request.target + text: Final = "Audit result." + return _json( + { + "status": "completed", + "error": None, + "incomplete_details": None, + "output": [ + { + "type": "message", + "content": [ + { + "type": "output_text", + "text": text, + "annotations": [ + { + "type": "url_citation", + "url": _BING_CITATION, + "title": "Bing result", + "start_index": len(text), + "end_index": len(text), + } + ], + } + ], + } + ], + } + ) + + +class _SearchResult(BaseModel): + model_config = ConfigDict(frozen=True) + + url: str + + +class _SearchResponse(BaseModel): + model_config = ConfigDict(frozen=True) + + results: tuple[_SearchResult, ...] + + +def test_bing_grounding_search_mints_its_entra_token_once_across_searches(rig: _Rig) -> None: + with wire_server(_edge(_foundry), port=rig.port) as wire: + responses: Final = tuple( + rig.gateway.request("POST", f"/v1/search/{_BING_TOOL}", {"query": _marker(), "max_results": 1}) + for _ in range(3) + ) + requests: Final = wire.drain() + + assert [response.status_code for response in responses] == [200] * 3, [r.text for r in responses] + assert [ + tuple(result.url for result in _SearchResponse.model_validate_json(response.content).results) + for response in responses + ] == [(_BING_CITATION,)] * 3 + assert _identity_fetches(requests, _BING_RESOURCE) == 1 + assert [call.headers.get("authorization") for call in _provider_calls(requests)] == [ + f"Bearer {_identity_token(_BING_RESOURCE)}" + ] * 3 + + +def _lifecycle(upstream_id: str) -> Callable[[Request], Reply]: + response: Final[dict[str, JsonValue]] = { + "id": upstream_id, + "object": "response", + "created_at": 1, + "model": _DEPLOYMENT, + "output": [], + } + + def respond(request: Request) -> Reply: + path: Final = urlsplit(request.target).path + match request.method: + case "DELETE": + return _json({"id": upstream_id, "object": "response", "deleted": True}) + case "GET" if path.endswith("/input_items"): + return _json({"object": "list", "data": [], "first_id": None, "last_id": None, "has_more": False}) + case "POST" if path.endswith("/cancel"): + return _json({**response, "status": "cancelled"}) + case _: + return _json({**response, "status": "completed"}) + + return respond + + +_LIFECYCLE_ROUTES: Final = ( + ("GET", "/v1/responses/{}", ""), + ("GET", "/v1/responses/{}/input_items", "/input_items"), + ("POST", "/v1/responses/{}/cancel", "/cancel"), + ("DELETE", "/v1/responses/{}", ""), +) + + +def test_every_responses_route_on_an_entra_deployment_reaches_azure_with_the_managed_identity_token( + rig: _Rig, +) -> None: + resource: Final = f"api://audit-routes-{uuid.uuid4().hex}" + upstream_id: Final = f"resp_{uuid.uuid4().hex}" + with wire_server(_edge(_lifecycle(upstream_id)), port=rig.port) as wire, rig.gateway.scenario() as scenario: + model: Final = _model(scenario, _Surface.RESPONSES, rig.edge, f"{resource}/.default") + key: Final = scenario.key() + created: Final = _post(_Surface.RESPONSES, rig.gateway, key, model) + assert created.status_code == 200, created.text + response_id: Final = string_value(JSON_OBJECT.validate_json(created.content)["id"]) + followups: Final = tuple( + rig.gateway.request(method, route.format(response_id), key=key) for method, route, _ in _LIFECYCLE_ROUTES + ) + requests: Final = wire.drain() + + assert [response.status_code for response in followups] == [200] * len(_LIFECYCLE_ROUTES), [ + response.text for response in followups + ] + calls: Final = _provider_calls(requests) + assert (calls[0].method, urlsplit(calls[0].target).path.endswith("/responses")) == ("POST", True) + assert [ + (call.method, urlsplit(call.target).path.endswith(f"/{upstream_id}{suffix}")) + for call, (_, _, suffix) in zip(calls[1:], _LIFECYCLE_ROUTES, strict=True) + ] == [(method, True) for method, _, _ in _LIFECYCLE_ROUTES] + assert {call.headers.get("authorization") for call in calls} == {f"Bearer {_identity_token(resource)}"} + assert _identity_fetches(requests, resource) == 1 diff --git a/tests/unit/llms/azure/test_azure.py b/tests/unit/llms/azure/test_azure.py index 7d5afde68ef..c52321baf02 100644 --- a/tests/unit/llms/azure/test_azure.py +++ b/tests/unit/llms/azure/test_azure.py @@ -6,7 +6,7 @@ import os import time from collections.abc import Mapping from typing import Final -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch import httpx import pytest @@ -291,11 +291,9 @@ def test_get_azure_ad_token_from_username_password(mock_get_bearer_token_provide mock_credential.assert_called_once_with(client_id=client_id, username=username, password=password) - mock_get_bearer_token_provider.assert_called_once_with( - mock_credential.return_value, "https://cognitiveservices.azure.com/.default" - ) + mock_get_bearer_token_provider.assert_called_once_with(ANY, "https://cognitiveservices.azure.com/.default") - assert result == mock_token_provider + assert result() == "mock-token" def test_azure_openai_gpt_4o_naming(monkeypatch): diff --git a/tests/unit/llms/azure/test_azure_common_utils.py b/tests/unit/llms/azure/test_azure_common_utils.py index 1dc7ff74323..f7e2e73d426 100644 --- a/tests/unit/llms/azure/test_azure_common_utils.py +++ b/tests/unit/llms/azure/test_azure_common_utils.py @@ -2,9 +2,12 @@ import json import os import threading import traceback +import uuid +from concurrent.futures import ThreadPoolExecutor +from queue import SimpleQueue from typing import Callable, Final, Optional from urllib.parse import parse_qs -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch import httpx import pytest @@ -13,6 +16,7 @@ import litellm from litellm.llms.azure.common_utils import ( BaseAzureLLM, _cached_azure_ad_token_refresh_provider, + _cached_default_azure_credential_provider, _cached_entra_id_token_provider, get_azure_ad_token, get_azure_ad_token_from_entra_id, @@ -41,6 +45,7 @@ def setup_mocks(monkeypatch): monkeypatch.delenv("AZURE_SCOPE", raising=False) monkeypatch.delenv("AZURE_AD_TOKEN", raising=False) _cached_azure_ad_token_refresh_provider.cache_clear() + _cached_default_azure_credential_provider.cache_clear() with ( patch( @@ -86,6 +91,7 @@ def setup_mocks(monkeypatch): "select_url": mock_select_url, } _cached_azure_ad_token_refresh_provider.cache_clear() + _cached_default_azure_credential_provider.cache_clear() def test_initialize_with_api_key(setup_mocks): @@ -1436,6 +1442,33 @@ def test_get_azure_ad_token_with_token_refresh_reuses_provider_per_scope(setup_m ] +def test_get_azure_ad_token_reuses_one_default_azure_credential_provider_per_scope(setup_mocks): + setup_mocks["litellm"].enable_azure_ad_token_refresh = True + default_scope: Final = "https://cognitiveservices.azure.com/.default" + other_scope: Final = "api://other-resource/.default" + + def provider_for(azure_scope: str, azure_credential: AzureCredentialType | None = None) -> Callable[[], str]: + if azure_credential != AzureCredentialType.DefaultAzureCredential: + raise ValueError("No credential provided") + return lambda: f"token-for-{azure_scope}" + + setup_mocks["token_provider"].side_effect = provider_for + + tokens: Final = ( + get_azure_ad_token(GenericLiteLLMParams()), + get_azure_ad_token(GenericLiteLLMParams()), + get_azure_ad_token(GenericLiteLLMParams(azure_scope=other_scope)), + ) + + default_credential_builds: Final = tuple( + call.kwargs["azure_scope"] + for call in setup_mocks["token_provider"].call_args_list + if call.kwargs.get("azure_credential") == AzureCredentialType.DefaultAzureCredential + ) + assert tokens == (f"token-for-{default_scope}", f"token-for-{default_scope}", f"token-for-{other_scope}") + assert default_credential_builds == (default_scope, other_scope) + + def test_get_azure_ad_token_with_token_refresh_error(setup_mocks): """Test get_azure_ad_token with token refresh enabled but raising an error.""" # Reset mocks to ensure clean state @@ -1536,9 +1569,7 @@ def test_get_azure_ad_token_provider_with_default_azure_credential(): mock_default_cred.assert_called_once_with() # Verify get_bearer_token_provider was called with the right parameters - mock_token_provider.assert_called_once_with( - mock_credential_instance, "https://cognitiveservices.azure.com/.default" - ) + mock_token_provider.assert_called_once_with(ANY, "https://cognitiveservices.azure.com/.default") # Verify the returned token provider works token = token_provider() @@ -2137,6 +2168,42 @@ class TestEntraIdTokenProviderCache: assert first is second assert mock_credential.call_count == 1 + def test_concurrent_first_callers_build_one_credential(self): + callers: Final = 8 + start: Final = threading.Barrier(callers) + built: Final = SimpleQueue[None]() + everyone_built: Final = threading.Event() + release: Final = threading.Event() + + def held_credential(*_: str) -> MagicMock: + built.put(None) + if built.qsize() >= callers: + everyone_built.set() + release.wait(5) + return MagicMock() + + def first_call() -> Callable[[], str]: + start.wait(5) + return get_azure_ad_token_from_entra_id( + tenant_id="tenant", + client_id="client", + client_secret="secret", + scope="https://cognitiveservices.azure.com/.default", + ) + + with ( + patch("azure.identity.ClientSecretCredential", side_effect=held_credential), + patch("azure.identity.get_bearer_token_provider", side_effect=lambda credential, scope: lambda: "token"), + ThreadPoolExecutor(max_workers=callers) as pool, + ): + calls: Final = tuple(pool.submit(first_call) for _ in range(callers)) + everyone_built.wait(0.5) + release.set() + providers: Final = tuple(call.result(timeout=5) for call in calls) + + assert built.qsize() == 1 + assert all(provider is providers[0] for provider in providers) + @pytest.mark.parametrize( "second_call_kwargs", [ @@ -2291,6 +2358,71 @@ def test_oidc_token_exchange_closes_its_own_http_client(monkeypatch: pytest.Monk assert token_endpoint.closed +def test_concurrent_oidc_callers_share_one_token_exchange(monkeypatch: pytest.MonkeyPatch) -> None: + callers: Final = 8 + start: Final = threading.Barrier(callers) + exchanges: Final = SimpleQueue[None]() + everyone_exchanging: Final = threading.Event() + release: Final = threading.Event() + + class HeldTokenEndpoint(httpx.BaseTransport): + def handle_request(self, request: httpx.Request) -> httpx.Response: + exchanges.put(None) + if exchanges.qsize() >= callers: + everyone_exchanging.set() + release.wait(5) + return httpx.Response(200, json={"access_token": "exchanged-token", "expires_in": 3600}) + + monkeypatch.setattr(litellm, "sync_transport", HeldTokenEndpoint(), raising=False) + monkeypatch.setenv("AZURE_TEST_OIDC_BURST_ASSERTION", "burst-jwt") + tenant: Final = f"tenant-{uuid.uuid4().hex}" + + def exchange() -> str: + start.wait(5) + return get_azure_ad_token_from_oidc( + azure_ad_token="os.environ/AZURE_TEST_OIDC_BURST_ASSERTION", + azure_client_id="client", + azure_tenant_id=tenant, + ) + + with ThreadPoolExecutor(max_workers=callers) as pool: + calls: Final = tuple(pool.submit(exchange) for _ in range(callers)) + exchanges.get(timeout=5) + everyone_exchanging.wait(0.5) + release.set() + tokens: Final = tuple(call.result(timeout=5) for call in calls) + + assert tokens == ("exchanged-token",) * callers + assert exchanges.qsize() == 0 + + +def test_oidc_tokens_for_different_scopes_are_exchanged_and_cached_separately(monkeypatch: pytest.MonkeyPatch) -> None: + exchanged_scopes: Final = SimpleQueue[str]() + + class ScopedTokenEndpoint(httpx.BaseTransport): + def handle_request(self, request: httpx.Request) -> httpx.Response: + scope: Final = parse_qs(request.read().decode())["scope"][-1] + exchanged_scopes.put(scope) + return httpx.Response(200, json={"access_token": f"token-for:{scope}", "expires_in": 3600}) + + monkeypatch.setattr(litellm, "sync_transport", ScopedTokenEndpoint(), raising=False) + monkeypatch.setenv("AZURE_TEST_OIDC_SCOPED_ASSERTION", "scoped-jwt") + tenant: Final = f"tenant-{uuid.uuid4().hex}" + + def token_for(scope: str) -> str: + return get_azure_ad_token_from_oidc( + azure_ad_token="os.environ/AZURE_TEST_OIDC_SCOPED_ASSERTION", + azure_client_id="client", + azure_tenant_id=tenant, + scope=scope, + ) + + tokens: Final = tuple(token_for(scope) for scope in ("api://first/.default", "api://second/.default") * 2) + + assert tokens == ("token-for:api://first/.default", "token-for:api://second/.default") * 2 + assert exchanged_scopes.qsize() == 2 + + def test_evicting_an_azure_client_built_on_the_callers_session_leaves_it_open(monkeypatch): """`initialize_azure_sdk_client` puts `litellm.aclient_session` on the SDK client. diff --git a/tests/unit/secret_managers/test_get_azure_ad_token_provider.py b/tests/unit/secret_managers/test_get_azure_ad_token_provider.py index 4bc7c21d8d2..141072635bc 100644 --- a/tests/unit/secret_managers/test_get_azure_ad_token_provider.py +++ b/tests/unit/secret_managers/test_get_azure_ad_token_provider.py @@ -1,16 +1,25 @@ +import itertools import json import os -from typing import Optional -from unittest.mock import MagicMock, patch +import threading +import time +from collections.abc import Callable, Iterator +from concurrent.futures import Future, ThreadPoolExecutor, wait +from dataclasses import dataclass, field +from queue import SimpleQueue +from typing import Final, Optional +from unittest.mock import ANY, MagicMock, patch # Adds the grandparent directory to sys.path to allow importing project modules import pytest +from azure.core.credentials import AccessToken, AccessTokenInfo, TokenRequestOptions from azure.core.exceptions import ClientAuthenticationError from litellm.secret_managers.get_azure_ad_token_provider import ( get_azure_ad_token_provider, infer_credential_type_from_environment, + single_flight_bearer_token_provider, ) from litellm.types.secret_managers.get_azure_ad_token_provider import ( AzureCredentialType, @@ -190,7 +199,7 @@ class TestGetAzureAdTokenProvider: tenant_id="test-tenant-id", ) mock_get_bearer_token_provider.assert_called_once_with( - mock_credential_instance, "https://cognitiveservices.azure.com/.default" + ANY, "https://cognitiveservices.azure.com/.default" ) # Test that the returned callable works @@ -228,7 +237,7 @@ class TestGetAzureAdTokenProvider: client_id="test-client-id" ) mock_get_bearer_token_provider.assert_called_once_with( - mock_credential_instance, "https://cognitiveservices.azure.com/.default" + ANY, "https://cognitiveservices.azure.com/.default" ) # Test that the returned callable works @@ -270,7 +279,7 @@ class TestGetAzureAdTokenProvider: certificate_path="/path/to/cert.pem", ) mock_get_bearer_token_provider.assert_called_once_with( - mock_credential_instance, "https://cognitiveservices.azure.com/.default" + ANY, "https://cognitiveservices.azure.com/.default" ) # Test that the returned callable works @@ -314,7 +323,7 @@ class TestGetAzureAdTokenProvider: password="pwd4cert.pem", ) mock_get_bearer_token_provider.assert_called_once_with( - mock_credential_instance, "https://cognitiveservices.azure.com/.default" + ANY, "https://cognitiveservices.azure.com/.default" ) # Test that the returned callable works @@ -348,7 +357,7 @@ class TestGetAzureAdTokenProvider: assert callable(result) mock_certificate_credential.assert_called_once_with() mock_get_bearer_token_provider.assert_called_once_with( - mock_credential_instance, "https://cognitiveservices.azure.com/.default" + ANY, "https://cognitiveservices.azure.com/.default" ) # Test that the returned callable works @@ -417,9 +426,254 @@ class TestGetAzureAdTokenProvider: assert callable(result) mock_default_azure_credential.assert_called_once_with() mock_get_bearer_token_provider.assert_called_once_with( - mock_credential_instance, "https://cognitiveservices.azure.com/.default" + ANY, "https://cognitiveservices.azure.com/.default" ) # Test that the returned callable works token = result() assert token == "mock-default-token" + + +_CONCURRENT_CALLERS: Final = 4 +_SCOPE: Final = "api://audit/.default" +_YEAR_2100: Final = 4_102_444_800 +_LONG_EXPIRED: Final = 1 +_JOIN_WINDOW_SECONDS: Final = 0.5 +_FAST_FAILING_CALLERS: Final = 16 + + +@dataclass(frozen=True, slots=True) +class _HeldCredential: + expires_on: int = _YEAR_2100 + failing_attempts: int = 0 + attempts: Iterator[int] = field(default_factory=itertools.count) + requests: SimpleQueue[tuple[str, ...]] = field(default_factory=SimpleQueue) + everyone_asked: threading.Event = field(default_factory=threading.Event) + release: threading.Event = field(default_factory=threading.Event) + + def get_token_info(self, *scopes: str, options: TokenRequestOptions | None = None) -> AccessTokenInfo: + self.requests.put(scopes) + if self.requests.qsize() >= _CONCURRENT_CALLERS: + self.everyone_asked.set() + self.release.wait(5) + attempt: Final = next(self.attempts) + if attempt < self.failing_attempts: + raise ClientAuthenticationError("Entra is down") + return AccessTokenInfo(f"token-{attempt}", self.expires_on) + + +def _call_together(provider: Callable[[], str], credential: _HeldCredential) -> tuple[Future[str], ...]: + start: Final = threading.Barrier(_CONCURRENT_CALLERS) + + def call() -> str: + start.wait(5) + return provider() + + with ThreadPoolExecutor(max_workers=_CONCURRENT_CALLERS) as pool: + calls: Final = tuple(pool.submit(call) for _ in range(_CONCURRENT_CALLERS)) + assert credential.requests.get(timeout=5) == (_SCOPE,) + credential.everyone_asked.wait(0.5) + credential.release.set() + return calls + + +def test_concurrent_cold_callers_share_one_token_request() -> None: + credential: Final = _HeldCredential() + provider: Final = single_flight_bearer_token_provider(credential, _SCOPE) + + calls: Final = _call_together(provider, credential) + + assert tuple(call.result() for call in calls) == ("token-0",) * _CONCURRENT_CALLERS + assert credential.requests.qsize() == 0 + + +def test_concurrent_callers_after_the_token_expired_share_one_token_request() -> None: + credential: Final = _HeldCredential(expires_on=_LONG_EXPIRED) + provider: Final = single_flight_bearer_token_provider(credential, _SCOPE) + credential.release.set() + provider() + credential.requests.get_nowait() + credential.release.clear() + + calls: Final = _call_together(provider, credential) + + assert tuple(call.result() for call in calls) == ("token-1",) * _CONCURRENT_CALLERS + assert credential.requests.qsize() == 0 + + +def test_a_failed_cold_token_request_fails_only_the_caller_that_made_it() -> None: + credential: Final = _HeldCredential(failing_attempts=1) + provider: Final = single_flight_bearer_token_provider(credential, _SCOPE) + + calls: Final = _call_together(provider, credential) + + failed: Final = sorted(isinstance(call.exception(), ClientAuthenticationError) for call in calls) + assert failed == [False] * (_CONCURRENT_CALLERS - 1) + [True] + + +@dataclass(frozen=True, slots=True) +class _ScriptedCredential: + expiries: Iterator[int | None] = field(default_factory=lambda: itertools.repeat(None)) + attempts: Iterator[int] = field(default_factory=itertools.count) + requests: SimpleQueue[None] = field(default_factory=SimpleQueue) + answers: threading.Semaphore = field(default_factory=lambda: threading.Semaphore(0)) + + def get_token_info(self, *_: str, options: TokenRequestOptions | None = None) -> AccessTokenInfo: + self.requests.put(None) + self.answers.acquire(timeout=5) + attempt: Final = next(self.attempts) + expires_on: Final = next(self.expiries) + if expires_on is None: + raise ClientAuthenticationError("Entra is down") + return AccessTokenInfo(f"token-{attempt}", expires_on) + + +def _answer_together( + provider: Callable[[], str], credential: _ScriptedCredential, token_requests: int +) -> tuple[Future[str], ...]: + start: Final = threading.Barrier(_CONCURRENT_CALLERS) + + def call() -> str: + start.wait(5) + return provider() + + with ThreadPoolExecutor(max_workers=_CONCURRENT_CALLERS) as pool: + calls: Final = tuple(pool.submit(call) for _ in range(_CONCURRENT_CALLERS)) + for _ in range(token_requests): + credential.requests.get(timeout=5) + time.sleep(_JOIN_WINDOW_SECONDS) + credential.answers.release() + return calls + + +def test_an_entra_outage_retries_only_the_first_failed_token_request() -> None: + credential: Final = _ScriptedCredential() + provider: Final = single_flight_bearer_token_provider(credential, _SCOPE) + + first_burst: Final = _answer_together(provider, credential, token_requests=2) + later_burst: Final = _answer_together(provider, credential, token_requests=1) + + assert all(isinstance(call.exception(), ClientAuthenticationError) for call in first_burst + later_burst) + assert credential.requests.qsize() == 0 + + +def test_waiting_callers_retry_a_failed_token_request_again_once_entra_recovers() -> None: + credential: Final = _ScriptedCredential(iter((None, None, _LONG_EXPIRED, None, _LONG_EXPIRED))) + provider: Final = single_flight_bearer_token_provider(credential, _SCOPE) + _answer_together(provider, credential, token_requests=2) + credential.answers.release() + assert provider() == "token-2" + credential.requests.get_nowait() + + recovered: Final = _answer_together(provider, credential, token_requests=2) + + assert sorted(call.exception() is None for call in recovered) == [False] + [True] * (_CONCURRENT_CALLERS - 1) + assert {call.result() for call in recovered if call.exception() is None} == {"token-4"} + + +def test_a_failed_refresh_fails_the_request_that_made_it_even_while_a_valid_token_is_held() -> None: + credential: Final = _ScriptedCredential(iter((int(time.time()) + 120, None))) + provider: Final = single_flight_bearer_token_provider(credential, _SCOPE) + credential.answers.release(2) + + assert provider() == "token-0" + with pytest.raises(ClientAuthenticationError): + provider() + + +def test_callers_waiting_on_a_failed_token_request_share_one_retry_when_entra_fails_fast() -> None: + credential: Final = _ScriptedCredential() + provider: Final = single_flight_bearer_token_provider(credential, _SCOPE) + start: Final = threading.Barrier(_FAST_FAILING_CALLERS) + + def call() -> str: + start.wait(5) + return provider() + + with ThreadPoolExecutor(max_workers=_FAST_FAILING_CALLERS) as pool: + calls: Final = tuple(pool.submit(call) for _ in range(_FAST_FAILING_CALLERS)) + credential.requests.get(timeout=5) + time.sleep(_JOIN_WINDOW_SECONDS) + credential.answers.release(_FAST_FAILING_CALLERS) + + assert all(isinstance(call.exception(), ClientAuthenticationError) for call in calls) + assert credential.requests.qsize() == 1 + + +@dataclass(frozen=True, slots=True) +class _ExpiringSoonTokens: + first_lifetime_seconds: int = 120 + requests: SimpleQueue[None] = field(default_factory=SimpleQueue) + refresh_started: threading.Event = field(default_factory=threading.Event) + release: threading.Event = field(default_factory=threading.Event) + + def next_token(self) -> tuple[str, int]: + self.requests.put(None) + if self.requests.qsize() == 1: + return "expiring-token", int(time.time()) + self.first_lifetime_seconds + self.refresh_started.set() + self.release.wait(5) + return "refreshed-token", _YEAR_2100 + + +@dataclass(frozen=True, slots=True) +class _ExpiringSoonTokenInfoCredential: + tokens: _ExpiringSoonTokens = field(default_factory=_ExpiringSoonTokens) + + def get_token_info(self, *_: str, options: TokenRequestOptions | None = None) -> AccessTokenInfo: + return AccessTokenInfo(*self.tokens.next_token()) + + +@dataclass(frozen=True, slots=True) +class _ExpiringSoonTokenOnlyCredential: + tokens: _ExpiringSoonTokens = field(default_factory=_ExpiringSoonTokens) + + def get_token( + self, + *_: str, + claims: str | None = None, + tenant_id: str | None = None, + enable_cae: bool = False, + **kwargs: object, + ) -> AccessToken: + return AccessToken(*self.tokens.next_token()) + + +@pytest.mark.parametrize("credential_type", (_ExpiringSoonTokenInfoCredential, _ExpiringSoonTokenOnlyCredential)) +def test_callers_during_a_refresh_keep_the_valid_token_while_one_of_them_refreshes( + credential_type: type[_ExpiringSoonTokenInfoCredential | _ExpiringSoonTokenOnlyCredential], +) -> None: + credential: Final = credential_type() + tokens: Final = credential.tokens + provider: Final = single_flight_bearer_token_provider(credential, _SCOPE) + assert provider() == "expiring-token" + + with ThreadPoolExecutor(max_workers=_CONCURRENT_CALLERS) as pool: + refresh: Final = pool.submit(provider) + assert tokens.refresh_started.wait(5) + others: Final = tuple(pool.submit(provider) for _ in range(_CONCURRENT_CALLERS - 1)) + try: + assert tuple(other.result(timeout=2) for other in others) == ("expiring-token",) * (_CONCURRENT_CALLERS - 1) + finally: + tokens.release.set() + + assert (refresh.result(), provider()) == ("refreshed-token", "refreshed-token") + assert tokens.requests.qsize() == 2 + + +def test_callers_during_a_refresh_wait_for_the_new_token_when_the_valid_one_is_about_to_expire() -> None: + tokens: Final = _ExpiringSoonTokens(first_lifetime_seconds=30) + provider: Final = single_flight_bearer_token_provider(_ExpiringSoonTokenInfoCredential(tokens), _SCOPE) + assert provider() == "expiring-token" + + with ThreadPoolExecutor(max_workers=_CONCURRENT_CALLERS) as pool: + refresh: Final = pool.submit(provider) + assert tokens.refresh_started.wait(5) + others: Final = tuple(pool.submit(provider) for _ in range(_CONCURRENT_CALLERS - 1)) + try: + assert not wait(others, timeout=_JOIN_WINDOW_SECONDS).done + finally: + tokens.release.set() + + assert tuple(call.result() for call in (refresh, *others)) == ("refreshed-token",) * _CONCURRENT_CALLERS + assert tokens.requests.qsize() == 2 diff --git a/tests/unit/secret_managers/test_secret_managers_main.py b/tests/unit/secret_managers/test_secret_managers_main.py index 7c6e1de6949..862658d9169 100644 --- a/tests/unit/secret_managers/test_secret_managers_main.py +++ b/tests/unit/secret_managers/test_secret_managers_main.py @@ -2,7 +2,11 @@ import base64 import json import logging import os +import threading import time +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from queue import SimpleQueue from typing import Final from unittest.mock import Mock, patch @@ -12,6 +16,7 @@ import litellm from litellm._logging import redact_internal_details_from_client_message from litellm.integrations.custom_secret_manager import CustomSecretManager from litellm.secret_managers.main import ( + _cached_azure_oidc_token_provider, get_secret, normalize_nonempty_secret_str, resolve_local_oidc_reference, @@ -313,6 +318,41 @@ def test_oidc_azure_ad_token_success(mock_get_azure_ad_token_provider, monkeypat mock_token_provider.assert_called_once_with() +def test_concurrent_azure_oidc_reads_build_one_token_provider(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("AZURE_FEDERATED_TOKEN_FILE", raising=False) + callers: Final = 8 + start: Final = threading.Barrier(callers) + builds: Final = SimpleQueue[str | None]() + everyone_building: Final = threading.Event() + release: Final = threading.Event() + + def held_build(azure_scope: str | None = None) -> Callable[[], str]: + builds.put(azure_scope) + if builds.qsize() >= callers: + everyone_building.set() + release.wait(5) + return lambda: "identity-token" + + def read() -> str | None: + start.wait(5) + return get_secret("oidc/azure/api://concurrent-audience") + + _cached_azure_oidc_token_provider.cache_clear() + with ( + patch("litellm.secret_managers.main.get_azure_ad_token_provider", side_effect=held_build), + ThreadPoolExecutor(max_workers=callers) as pool, + ): + calls: Final = tuple(pool.submit(read) for _ in range(callers)) + assert builds.get(timeout=5) == "api://concurrent-audience" + everyone_building.wait(0.5) + release.set() + tokens: Final = tuple(call.result(timeout=5) for call in calls) + _cached_azure_oidc_token_provider.cache_clear() + + assert tokens == ("identity-token",) * callers + assert builds.qsize() == 0 + + def test_oidc_file_success(tmp_path, monkeypatch): token_file = tmp_path / "token.txt" token_file.write_text("file_token") From 602cf333024ffcc21827ee29ad01e81e98679efa Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 28 Oct 2026 19:31:04 +0000 Subject: [PATCH 074/128] fix(cost-map): add mistral/deepseek-v4-1-flash-1017 (#48879) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../model_prices_and_context_window_backup.json | 16 ++++++++++++++++ model_prices_and_context_window.json | 16 ++++++++++++++++ 2 files changed, 32 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1b65c87c315..7199a075612 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41845,6 +41845,22 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "mistral/deepseek-v4-1-flash-1017": { + "cache_read_input_token_cost": 7e-09, + "input_cost_per_token": 3e-07, + "litellm_provider": "mistral", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://docs.mistral.ai/models/deepseek-v4-flash-1017", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "mistral/zai-glm-5-3": { "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.4e-06, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 1b65c87c315..7199a075612 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41845,6 +41845,22 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "mistral/deepseek-v4-1-flash-1017": { + "cache_read_input_token_cost": 7e-09, + "input_cost_per_token": 3e-07, + "litellm_provider": "mistral", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://docs.mistral.ai/models/deepseek-v4-flash-1017", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "mistral/zai-glm-5-3": { "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.4e-06, From 365d3d372971c8c77921a0688022007028b542a6 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Wed, 28 Oct 2026 12:35:26 -0700 Subject: [PATCH 075/128] feat(decisions): support Ollama decision models (#48871) * test: cover local Ollama Decisions wire contracts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(decisions): support ollama providers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(decisions): account for Ollama model-info requests 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> --- litellm/__init__.py | 3 + litellm/_lazy_imports_registry.py | 2 + .../llms/ollama/decisions/transformation.py | 17 ++ litellm/utils.py | 7 +- .../providers/test_decisions_wire.py | 222 +++++++++++++++++- tests/unit/decisions/test_main.py | 193 ++++++++++++++- .../test_health_check_helpers.py | 31 +++ .../decisions_endpoints/test_endpoints.py | 53 +++++ 8 files changed, 522 insertions(+), 6 deletions(-) create mode 100644 litellm/llms/ollama/decisions/transformation.py diff --git a/litellm/__init__.py b/litellm/__init__.py index af771b2f0e6..fc8f6b57254 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1739,6 +1739,9 @@ if TYPE_CHECKING: from .llms.lm_studio.decisions.transformation import ( LMStudioDecisionsConfig as LMStudioDecisionsConfig, ) + from .llms.ollama.decisions.transformation import ( + OllamaDecisionsConfig as OllamaDecisionsConfig, + ) from .llms.laya.decisions.transformation import ( LayaDecisionsConfig as LayaDecisionsConfig, ) diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index f43e369e092..599234bab1a 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -171,6 +171,7 @@ LLM_CONFIG_NAMES: Final = ( "DatabricksDecisionsConfig", "HostedVLLMDecisionsConfig", "LMStudioDecisionsConfig", + "OllamaDecisionsConfig", "LayaDecisionsConfig", "BespokeDecisionsConfig", "OpenAIDecisionsConfig", @@ -749,6 +750,7 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { "HostedVLLMDecisionsConfig", ), "LMStudioDecisionsConfig": (".llms.lm_studio.decisions.transformation", "LMStudioDecisionsConfig"), + "OllamaDecisionsConfig": (".llms.ollama.decisions.transformation", "OllamaDecisionsConfig"), "LayaDecisionsConfig": (".llms.laya.decisions.transformation", "LayaDecisionsConfig"), "BespokeDecisionsConfig": (".llms.bespoke.decisions.transformation", "BespokeDecisionsConfig"), "OpenAIDecisionsConfig": (".llms.openai.decisions.transformation", "OpenAIDecisionsConfig"), diff --git a/litellm/llms/ollama/decisions/transformation.py b/litellm/llms/ollama/decisions/transformation.py new file mode 100644 index 00000000000..3e69f5cfede --- /dev/null +++ b/litellm/llms/ollama/decisions/transformation.py @@ -0,0 +1,17 @@ +from typing import Final + +from litellm.constants import PLACEHOLDER_API_KEY +from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig + + +class OllamaDecisionsConfig(BaseDecisionsConfig): + api_key_env = ("OLLAMA_API_KEY",) + api_base_env = ("OLLAMA_API_BASE",) + api_key_required = False + + def get_default_api_base(self) -> str | None: + return "http://localhost:11434" + + def resolve_api_key(self, api_key: str | None) -> str | None: + resolved: Final = super().resolve_api_key(api_key) + return None if resolved == PLACEHOLDER_API_KEY else resolved diff --git a/litellm/utils.py b/litellm/utils.py index 1e095966a43..61b0e17e197 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9612,18 +9612,19 @@ class ProviderConfigManager: return litellm.HostedVLLMDecisionsConfig() if provider == LlmProviders.LM_STUDIO: return litellm.LMStudioDecisionsConfig() + if provider in (LlmProviders.OLLAMA, LlmProviders.OLLAMA_CHAT): + return litellm.OllamaDecisionsConfig() if provider == LlmProviders.LAYA: return litellm.LayaDecisionsConfig() if provider == LlmProviders.BESPOKE: return litellm.BespokeDecisionsConfig() - if provider in (LlmProviders.OPENAI, LlmProviders.AZURE): + if provider in (LlmProviders.OPENAI, LlmProviders.AZURE, LlmProviders.AZURE_AI): configs: Final = { LlmProviders.OPENAI: litellm.OpenAIDecisionsConfig, LlmProviders.AZURE: litellm.AzureDecisionsConfig, + LlmProviders.AZURE_AI: litellm.AzureAIDecisionsConfig, } return configs[provider]() - if provider == LlmProviders.AZURE_AI: - return litellm.AzureAIDecisionsConfig() return None @staticmethod diff --git a/tests/integration/providers/test_decisions_wire.py b/tests/integration/providers/test_decisions_wire.py index 04b78985b20..5af1532e9ed 100644 --- a/tests/integration/providers/test_decisions_wire.py +++ b/tests/integration/providers/test_decisions_wire.py @@ -1,7 +1,7 @@ import json import math import uuid -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from dataclasses import dataclass from pathlib import Path from typing import Final @@ -140,6 +140,38 @@ _PROVIDERS: Final = ( "databricks/databricks-openjev-qwen35-4b", ), ) +_LOCAL_PROVIDERS: Final = ( + pytest.param( + _Provider("ollama", "ollama/d3:4b", "/v1/systemone", "d3:4b", None, False, None), id="regression-ollama" + ), + pytest.param( + _Provider("ollama_chat", "ollama_chat/d3:4b", "/v1/systemone", "d3:4b", None, False, None), + id="regression-ollama_chat", + ), + pytest.param( + _Provider("hosted_vllm", "hosted_vllm/d3:4b", "/v1/systemone", "d3:4b", None, False, None), + id="guard-hosted_vllm", + ), + pytest.param( + _Provider("lm_studio", "lm_studio/d3:4b", "/v1/systemone", "d3:4b", None, False, None), + id="guard-lm_studio", + ), +) +_LOCAL_STATE: Final = "I was charged twice for order 4417." +_LOCAL_QUESTIONS: Final[Mapping[str, JsonValue]] = { + "billing": {"type": "noul", "instructions": "Is this a billing issue?"} +} +_LOCAL_ANSWERS: Final[Mapping[str, JsonValue]] = {"billing": {"type": "noul", "noul": 0.91}} +_LOCAL_OPENAI_QUESTIONS: Final[Sequence[JsonValue]] = ( + {"type": "predicate", "name": "billing", "instructions": "Is this a billing issue?"}, +) +_LOCAL_OPENAI_ANSWERS: Final[Sequence[JsonValue]] = ({"type": "predicate", "name": "billing", "probability": 0.91},) +_LOCAL_OPENAI_USAGE: Final[Mapping[str, JsonValue]] = { + **_USAGE, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 370, +} _PERPLEXITY: Final = _PROVIDERS[0] _OPENROUTER: Final = next(provider for provider in _PROVIDERS if provider.name == "openrouter") _OPENROUTER_CHAT_MODEL: Final = "openrouter/openai/gpt-5-mini" @@ -230,6 +262,26 @@ def _upstream_calls(gateway: Gateway, handle: ScenarioHandle) -> list[dict[str, return _calls_to(_observed_requests(gateway), handle) +def _sdk_inference_call(gateway: Gateway, handle: ScenarioHandle, provider: _Provider) -> dict[str, JsonValue]: + observed: Final = _upstream_calls(gateway, handle) + inference_path: Final = f"/{handle.scenario_id}/v1/systemone" + (call,) = tuple(request for request in observed if request["path"] == inference_path) + metadata_calls: Final = tuple(request for request in observed if request["path"] != inference_path) + assert metadata_calls in ( + (), + ( + { + "path": f"{inference_path}/api/show", + "authorization": "", + "body": {"name": provider.body_model}, + "method": "POST", + "api_key": "", + }, + ), + ), observed + return call + + def _spend_row(call_id: str) -> dict[str, JsonValue]: rows: Final = eventually(lambda: read_rows(_SPEND_QUERY, (call_id,)), lambda found: len(found) == 1, seconds=70) return rows[0] @@ -345,6 +397,174 @@ def test_each_provider_gets_its_own_path_key_and_body_and_is_billed(gateway: Gat assert math.isclose(_number(row["spend"]), expected_spend, rel_tol=1e-9), row +@pytest.fixture(autouse=True) +def _without_ambient_ollama_settings(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("OLLAMA_API_KEY", raising=False) + monkeypatch.delenv("OLLAMA_API_BASE", raising=False) + + +@pytest.mark.parametrize("provider", _LOCAL_PROVIDERS) +@pytest.mark.parametrize("route", ("/v1/systemone", "/v1/decisions"), ids=("systemone", "decisions")) +@pytest.mark.parametrize( + ("base_suffix", "api_key"), + ( + pytest.param("", None, id="root-no-key"), + pytest.param("/", _API_KEY, id="root-slash-key"), + pytest.param("/v1", None, id="v1-no-key"), + pytest.param("/v1/", _API_KEY, id="v1-slash-key"), + ), +) +def test_local_decisions_routes_preserve_wire_auth_usage_and_spend( + gateway: Gateway, provider: _Provider, route: str, base_suffix: str, api_key: str | None +) -> None: + native_reply: Final[Mapping[str, JsonValue]] = { + "model": provider.body_model, + "answers": dict(_LOCAL_ANSWERS), + "usage": _USAGE, + } + request_body: Final[Mapping[str, JsonValue]] = ( + {"state": _LOCAL_STATE, "questions": dict(_LOCAL_QUESTIONS)} + if route == "/v1/systemone" + else {"input": _LOCAL_STATE, "questions": list(_LOCAL_OPENAI_QUESTIONS)} + ) + expected_reply: Final = ( + native_reply + if route == "/v1/systemone" + else {"model": provider.body_model, "answers": list(_LOCAL_OPENAI_ANSWERS), "usage": dict(_LOCAL_OPENAI_USAGE)} + ) + with gateway.scenario() as scenario: + handle: Final = _register(scenario, dict(native_reply)) + model: Final = scenario.model( + model=provider.model, api_base=f"{handle.api_base()}{base_suffix}", api_key=api_key + ) + response: Final = gateway.request("POST", route, {"model": model, **request_body}) + assert response.status_code == 200, response.text + assert response.json() == expected_reply + assert response.headers["x-litellm-model-group"] == model + assert float(response.headers.get("x-litellm-response-cost", "0")) == 0.0 + (call,) = _upstream_calls(gateway, handle) + assert (call["path"], call["authorization"], call["body"]) == ( + f"/{handle.scenario_id}/v1/systemone", + f"Bearer {api_key}" if api_key else "", + {"model": provider.body_model, "state": _LOCAL_STATE, "questions": _LOCAL_QUESTIONS}, + ) + row: Final = _spend_row(response.headers["x-litellm-call-id"]) + assert {key: value for key, value in row.items() if key != "request_tags"} == { + "spend": 0.0, + "status": "success", + "call_type": "asystemone" if route == "/v1/systemone" else "adecisions", + "model_group": model, + "custom_llm_provider": provider.name, + "api_base": f"{handle.api_base()}/v1/systemone", + "prompt_tokens": _USAGE["input_tokens"], + "completion_tokens": _USAGE["output_tokens"], + } + + +@pytest.mark.parametrize("provider", _LOCAL_PROVIDERS) +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +async def test_local_systemone_sdk_preserves_native_wire( + gateway: Gateway, provider: _Provider, asynchronous: bool +) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _answer_body(provider)) + response: Final = ( + await litellm.asystemone( + model=provider.model, state=_STATE, questions=_SDK_QUESTIONS, api_base=f"{handle.api_base()}/v1/" + ) + if asynchronous + else litellm.systemone( + model=provider.model, state=_STATE, questions=_SDK_QUESTIONS, api_base=handle.api_base() + ) + ) + assert response.model_dump(mode="json") == _answer_body(provider) + call: Final = _sdk_inference_call(gateway, handle, provider) + assert (call["path"], call["authorization"], call["body"]) == ( + f"/{handle.scenario_id}/v1/systemone", + "", + {"model": provider.body_model, "state": _STATE, "questions": _QUESTIONS}, + ) + + +@pytest.mark.parametrize("provider", _LOCAL_PROVIDERS) +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +async def test_local_decisions_sdk_translates_openai_format( + gateway: Gateway, provider: _Provider, asynchronous: bool +) -> None: + questions: Final = TypeAdapter(list[dict[str, object]]).validate_python(_LOCAL_OPENAI_QUESTIONS) + native_reply: Final[Mapping[str, JsonValue]] = { + "model": provider.body_model, + "answers": dict(_LOCAL_ANSWERS), + "usage": _USAGE, + } + with gateway.scenario() as scenario: + handle: Final = _register(scenario, dict(native_reply)) + response: Final = ( + await litellm.adecisions( + model=provider.model, input=_LOCAL_STATE, questions=questions, api_base=f"{handle.api_base()}/v1/" + ) + if asynchronous + else litellm.decisions( + model=provider.model, input=_LOCAL_STATE, questions=questions, api_base=handle.api_base() + ) + ) + assert response.model_dump(mode="json") == { + "model": provider.body_model, + "answers": list(_LOCAL_OPENAI_ANSWERS), + "usage": dict(_LOCAL_OPENAI_USAGE), + } + call: Final = _sdk_inference_call(gateway, handle, provider) + assert (call["path"], call["authorization"], call["body"]) == ( + f"/{handle.scenario_id}/v1/systemone", + "", + {"model": provider.body_model, "state": _LOCAL_STATE, "questions": _LOCAL_QUESTIONS}, + ) + + +@pytest.mark.parametrize("provider", _LOCAL_PROVIDERS[:2]) +@pytest.mark.parametrize( + ("environment_base", "api_key", "environment_key", "authorization"), + ( + pytest.param(True, None, _ENV_KEY, f"Bearer {_ENV_KEY}", id="environment-base-and-key"), + pytest.param(False, _API_KEY, _ENV_KEY, f"Bearer {_API_KEY}", id="explicit-overrides-environment"), + pytest.param(False, "fake-api-key", _ENV_KEY, "", id="fake-key-suppresses-environment"), + pytest.param(False, None, None, "", id="unrelated-keys-not-forwarded"), + ), +) +def test_ollama_decisions_sdk_respects_environment_and_explicit_auth( + gateway: Gateway, + monkeypatch: pytest.MonkeyPatch, + provider: _Provider, + environment_base: bool, + api_key: str | None, + environment_key: str | None, + authorization: str, +) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "synthetic-unrelated-openai-key") + monkeypatch.setenv("LITELLM_API_KEY", "synthetic-unrelated-litellm-key") + if environment_key is not None: + monkeypatch.setenv("OLLAMA_API_KEY", environment_key) + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _answer_body(provider)) + unused: Final = _register(scenario, _answer_body(provider)) + monkeypatch.setenv("OLLAMA_API_BASE", handle.api_base() if environment_base else unused.api_base()) + response: Final = litellm.systemone( + model=provider.model, + state=_STATE, + questions=_SDK_QUESTIONS, + api_base=None if environment_base else f"{handle.api_base()}/v1/", + api_key=api_key, + ) + assert response.model_dump(mode="json") == _answer_body(provider) + call: Final = _sdk_inference_call(gateway, handle, provider) + assert (call["path"], call["authorization"], call["body"]) == ( + f"/{handle.scenario_id}/v1/systemone", + authorization, + {"model": provider.body_model, "state": _STATE, "questions": _QUESTIONS}, + ) + assert _upstream_calls(gateway, unused) == [] + + def test_test_connection_evaluation_mode_uses_typesafe_decisions_path(gateway: Gateway) -> None: provider: Final = _PROVIDERS[1] with gateway.scenario() as scenario: diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index b9813289749..e5501310556 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -578,8 +578,14 @@ async def test_unknown_provider_is_rejected_before_http(respx_mock: respx.MockRo @pytest.mark.asyncio -async def test_provider_without_decisions_support_is_rejected_before_http(respx_mock: respx.MockRouter) -> None: - with pytest.raises(litellm.BadRequestError, match=r"Unknown Decisions provider 'anthropic'\. Supported providers"): +async def test_provider_without_decisions_support_is_rejected_before_http( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.delenv("OLLAMA_API_KEY", raising=False) + monkeypatch.delenv("OLLAMA_API_BASE", raising=False) + with pytest.raises( + litellm.BadRequestError, match=r"Unknown Decisions provider 'anthropic'\. Supported providers" + ) as error: await litellm.asystemone( model="anthropic/claude-sonnet-4-5", state="review", @@ -587,6 +593,8 @@ async def test_provider_without_decisions_support_is_rejected_before_http(respx_ api_key="caller-key", ) + supported: Final = str(error.value).split("Supported providers: ", 1)[1].split(", ") + assert {"ollama", "ollama_chat"} <= set(supported) assert len(respx_mock.calls) == 0 @@ -1536,6 +1544,8 @@ _PROVIDERS_WHOSE_CHAT_MODELS_ALSO_DECIDE: Final = frozenset( LlmProviders.OPENROUTER.value, LlmProviders.HOSTED_VLLM.value, LlmProviders.LM_STUDIO.value, + LlmProviders.OLLAMA.value, + LlmProviders.OLLAMA_CHAT.value, } ) _DEDICATED_DECIDER_PROVIDERS: Final[tuple[str, ...]] = tuple( @@ -3911,6 +3921,185 @@ async def test_native_decisions_distinguishes_missing_usage_from_reported_zero( assert (expected == 0) == (usage["input_tokens"] == usage["output_tokens"] == 0) +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("ollama", "ollama_chat")) +@pytest.mark.parametrize("use_async", (True, False), ids=("asystemone", "systemone")) +@pytest.mark.parametrize("api_base", ("http://ollama.local:11434", "http://ollama.local:11434/v1")) +@pytest.mark.parametrize("api_key", (None, "fake-api-key")) +async def test_ollama_systemone_preserves_questions_without_unrelated_or_placeholder_auth( + provider: str, + use_async: bool, + api_base: str, + api_key: str | None, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.delenv("OLLAMA_API_KEY", raising=False) + monkeypatch.delenv("OLLAMA_API_BASE", raising=False) + monkeypatch.setenv("OPENAI_API_KEY", "unrelated-openai-key") + monkeypatch.setattr(litellm, "api_key", "unrelated-global-key") + monkeypatch.setattr(litellm, "openai_key", "unrelated-global-openai-key") + route: Final = respx_mock.post("http://ollama.local:11434/v1/systemone").respond(json=_RESPONSE) + kwargs: Final = { + "model": f"{provider}/d3:4b", + "state": "review", + "questions": _QUESTIONS, + "api_base": api_base, + "api_key": api_key, + } + + response: Final = await litellm.asystemone(**kwargs) if use_async else litellm.systemone(**kwargs) + + assert json.loads(route.calls.last.request.content) == { + "model": "d3:4b", + "state": "review", + "questions": dict(_QUESTIONS), + } + assert "authorization" not in route.calls.last.request.headers + assert response.answers["is_defect"] == NoulAnswer(type="noul", noul=0.9) + assert response.answers["sentiment"].choice == "positive" + assert response.answers["severity"].score == 1 + assert response.usage == DecisionsUsage(input_tokens=_INPUT_TOKENS, output_tokens=_OUTPUT_TOKENS) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("ollama", "ollama_chat")) +@pytest.mark.parametrize("use_async", (True, False), ids=("adecisions", "decisions")) +async def test_ollama_decisions_translates_openai_questions_and_answers( + provider: str, use_async: bool, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.delenv("OLLAMA_API_KEY", raising=False) + monkeypatch.delenv("OLLAMA_API_BASE", raising=False) + route: Final = respx_mock.post("http://ollama.local:11434/v1/systemone").respond(json=_RESPONSE) + kwargs: Final = { + "model": f"{provider}/d3:4b", + "input": "review", + "questions": [{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], + "api_base": "http://ollama.local:11434/v1", + } + + response: Final = await litellm.adecisions(**kwargs) if use_async else litellm.decisions(**kwargs) + + assert json.loads(route.calls.last.request.content) == { + "model": "d3:4b", + "state": "review", + "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + } + assert "authorization" not in route.calls.last.request.headers + assert response.answers[0].type == "predicate" + assert response.answers[0].name == "is_defect" + assert response.answers[0].probability == 0.9 + assert response.usage is not None + assert response.usage.input_tokens == _INPUT_TOKENS + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("ollama", "ollama_chat")) +@pytest.mark.parametrize("environment_base", (None, "http://ollama.local:11434/v1")) +async def test_ollama_decisions_resolves_environment_base_or_localhost_default( + provider: str, environment_base: str | None, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.delenv("OLLAMA_API_KEY", raising=False) + monkeypatch.delenv("OLLAMA_API_BASE", raising=False) + if environment_base is not None: + monkeypatch.setenv("OLLAMA_API_BASE", environment_base) + url: Final = f"{environment_base or 'http://localhost:11434/v1'}/systemone" + route: Final = respx_mock.post(url).respond(json=_RESPONSE) + + response: Final = await litellm.asystemone(model=f"{provider}/d3:4b", state="review", questions=_QUESTIONS) + + assert route.call_count == 1 + assert "authorization" not in route.calls.last.request.headers + assert response.answers["is_defect"] == NoulAnswer(type="noul", noul=0.9) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("ollama", "ollama_chat")) +@pytest.mark.parametrize("explicit_key", (None, "deployment-key")) +async def test_ollama_decisions_resolves_environment_key_and_explicit_connection_overrides( + provider: str, explicit_key: str | None, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.delenv("OLLAMA_API_KEY", raising=False) + monkeypatch.delenv("OLLAMA_API_BASE", raising=False) + monkeypatch.setenv("OLLAMA_API_BASE", "http://ignored.local:11434") + monkeypatch.setenv("OLLAMA_API_KEY", "environment-key") + route: Final = respx_mock.post("http://ollama.local:11434/v1/systemone").respond(json=_RESPONSE) + + response: Final = await litellm.asystemone( + model=f"{provider}/d3:4b", + state="review", + questions=_QUESTIONS, + api_base="http://ollama.local:11434", + api_key=explicit_key, + ) + + assert route.calls.last.request.headers["authorization"] == f"Bearer {explicit_key or 'environment-key'}" + assert response.answers["is_defect"] == NoulAnswer(type="noul", noul=0.9) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("ollama", "ollama_chat")) +async def test_ollama_decisions_maps_upstream_bad_requests( + provider: str, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.delenv("OLLAMA_API_KEY", raising=False) + monkeypatch.delenv("OLLAMA_API_BASE", raising=False) + route: Final = respx_mock.post("http://ollama.local:11434/v1/systemone").respond( + 400, json={"error": "decision rejected"} + ) + + with pytest.raises(litellm.BadRequestError, match="decision rejected") as error: + await litellm.asystemone( + model=f"{provider}/d3:4b", + state="review", + questions=_QUESTIONS, + api_base="http://ollama.local:11434", + num_retries=0, + ) + + assert error.value.status_code == 400 + assert route.call_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("ollama", "ollama_chat")) +@pytest.mark.parametrize("drop_params", (False, True)) +@pytest.mark.parametrize("image", ("data:image/png;base64,AA==", "https://images.example/input.png")) +async def test_ollama_decisions_rejects_input_images_before_http_or_drops_them( + provider: str, drop_params: bool, image: str, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.delenv("OLLAMA_API_KEY", raising=False) + monkeypatch.delenv("OLLAMA_API_BASE", raising=False) + monkeypatch.setattr(litellm, "drop_params", False) + kwargs: Final = { + "model": f"{provider}/d3:4b", + "input": [ + { + "role": "user", + "content": [{"type": "input_text", "text": "review"}, {"type": "input_image", "image_url": image}], + } + ], + "questions": [{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], + "api_base": "http://ollama.local:11434", + "drop_params": drop_params, + } + if not drop_params: + with pytest.raises(litellm.UnsupportedParamsError, match=r"\['input_image'\].*drop_params"): + await litellm.adecisions(**kwargs) + assert len(respx_mock.calls) == 0 + return + route: Final = respx_mock.post("http://ollama.local:11434/v1/systemone").respond(json=_RESPONSE) + + response: Final = await litellm.adecisions(**kwargs) + + assert json.loads(route.calls.last.request.content) == { + "model": "d3:4b", + "state": "review", + "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + } + assert response.answers[0].probability == 0.9 + + @pytest.mark.asyncio @pytest.mark.parametrize("use_async", (True, False), ids=("asystemone", "systemone")) @pytest.mark.parametrize("api_base", ("http://lmstudio.local:1234", "http://lmstudio.local:1234/v1")) diff --git a/tests/unit/litellm_core_utils/test_health_check_helpers.py b/tests/unit/litellm_core_utils/test_health_check_helpers.py index b584a72855a..88e2bc4b1fd 100644 --- a/tests/unit/litellm_core_utils/test_health_check_helpers.py +++ b/tests/unit/litellm_core_utils/test_health_check_helpers.py @@ -723,6 +723,37 @@ async def test_ahealth_check_probes_hosted_vllm_with_a_choice_question( } +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("ollama", "ollama_chat")) +async def test_ahealth_check_probes_ollama_decisions_with_one_noul_question_without_auth( + provider: str, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.delenv("OLLAMA_API_KEY", raising=False) + monkeypatch.delenv("OLLAMA_API_BASE", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + upstream: Final = respx_mock.post("http://ollama.local:11434/v1/systemone").respond( + json={ + "model": "d3:4b", + "answers": {"reachable": {"type": "noul", "noul": 1.0}}, + "usage": {"input_tokens": 12, "output_tokens": 1}, + } + ) + + result: Final = await ahealth_check( + {"model": f"{provider}/d3:4b", "api_base": "http://ollama.local:11434"}, mode="decisions" + ) + + assert "error" not in result, result + assert upstream.call_count == 1 + assert "authorization" not in upstream.calls.last.request.headers + assert json.loads(upstream.calls.last.request.content) == { + "model": "d3:4b", + "state": "health check", + "questions": {"reachable": {"type": "noul", "instructions": "Is the service reachable?"}}, + } + + @pytest.mark.asyncio @pytest.mark.parametrize("mode", ("evaluation", "decisions")) async def test_ahealth_check_keeps_the_noul_probe_for_other_decisions_providers( diff --git a/tests/unit/proxy/decisions_endpoints/test_endpoints.py b/tests/unit/proxy/decisions_endpoints/test_endpoints.py index a0a79f9b7cf..c67d1012475 100644 --- a/tests/unit/proxy/decisions_endpoints/test_endpoints.py +++ b/tests/unit/proxy/decisions_endpoints/test_endpoints.py @@ -1834,3 +1834,56 @@ def test_proxy_routes_lm_studio_decisions_with_optional_authentication( else [{"type": "predicate", "name": "is_defect", "probability": 0.9}] ) assert response.json()["answers"] == expected + + +@pytest.mark.parametrize("provider", ("ollama", "ollama_chat")) +@pytest.mark.parametrize("endpoint", ("/v1/systemone", "/v1/decisions")) +@pytest.mark.parametrize("api_key", (None, "deployment-key")) +def test_proxy_routes_ollama_decisions_with_optional_authentication( + client: TestClient, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + provider: str, + endpoint: str, + api_key: str | None, +) -> None: + monkeypatch.delenv("OLLAMA_API_KEY", raising=False) + monkeypatch.delenv("OLLAMA_API_BASE", raising=False) + deployment_router: Final = litellm.Router( + model_list=[ + { + "model_name": "local-decision", + "litellm_params": { + "model": f"{provider}/d3:4b", + "api_base": "http://ollama.local:11434/v1", + "api_key": api_key, + }, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", deployment_router) + upstream: Final = respx_mock.post("http://ollama.local:11434/v1/systemone").respond(json=_RESPONSE) + fields: Final = ( + {"state": "review", "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}} + if endpoint.endswith("systemone") + else { + "input": "review", + "questions": [{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], + } + ) + + response: Final = client.post(endpoint, json={"model": "local-decision", **fields}) + + assert response.status_code == 200, response.text + assert json.loads(upstream.calls.last.request.content) == { + "model": "d3:4b", + "state": "review", + "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + } + assert upstream.calls.last.request.headers.get("authorization") == (f"Bearer {api_key}" if api_key else None) + expected: Final = ( + _RESPONSE["answers"] + if endpoint.endswith("systemone") + else [{"type": "predicate", "name": "is_defect", "probability": 0.9}] + ) + assert response.json()["answers"] == expected From 41fcce62ffdb09430808f2a91a3434714d996ac6 Mon Sep 17 00:00:00 2001 From: zoveki <349028329+zoveki@users.noreply.github.com> Date: Wed, 28 Oct 2026 12:37:59 -0700 Subject: [PATCH 076/128] fix(gemini)!: accept Google's current Managed Agents fields on create and interactions (#48750) * fix(gemini): send Google's current Managed Agents fields on create and forward agent_config on interactions * fix(gemini): accept id on managed agent get, delete, and list_versions * fix(gemini): drop unsupported agent fields under drop_params and keep agent_config keyword-only A model interaction drops agent_config under drop_params, agent create drops wrongly typed optional fields under drop_params and otherwise answers 400 naming the field the caller sent, and agent_config on interactions create is keyword-only so positional response settings stay in place * fix(gemini): drop a non-object agent_config under drop_params and name the id alias the caller sent * fix(interactions): reject agent_config on bridged providers and drop a typeless one under drop_params Bridged providers answer 400 for agent_config without drop_params and drop it with drop_params. A Gemini agent run drops an agent_config without a string type under drop_params. The Gemini agents wire test expects Google's current upstream field names. * fix(interactions): honor additional_drop_params for agent_config on bridged providers * fix(agents): drop a clashing instructions alias under drop_params --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/interactions/__init__.py | 6 +- litellm/interactions/agents/__init__.py | 8 +- litellm/interactions/agents/http_handler.py | 29 +- litellm/interactions/agents/main.py | 270 ++++++++++++++---- litellm/interactions/main.py | 17 ++ .../llms/base_llm/agents/transformation.py | 8 +- litellm/llms/gemini/agents/transformation.py | 40 +-- .../gemini/interactions/transformation.py | 20 +- litellm/proxy/_lazy_openapi_snapshot.json | 2 +- .../google_endpoints/agents_endpoints.py | 7 +- litellm/proxy/route_llm_request.py | 3 +- litellm/types/agents.py | 14 +- .../test_missing_body_param_status.py | 15 +- .../providers/test_gemini_agents_wire.py | 4 +- .../interactions/test_agents_http_handler.py | 14 +- .../test_agents_main_and_utils.py | 256 ++++++++++++++++- tests/unit/interactions/test_main.py | 153 ++++++++++ .../test_gemini_agents_endpoints.py | 44 +++ .../test_interactions_agent_param.py | 41 +++ tests/unit/proxy/test_route_llm_request.py | 6 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 7 +- 21 files changed, 820 insertions(+), 144 deletions(-) diff --git a/litellm/interactions/__init__.py b/litellm/interactions/__init__.py index 6129cd87153..d5dd9b2cd00 100644 --- a/litellm/interactions/__init__.py +++ b/litellm/interactions/__init__.py @@ -32,11 +32,11 @@ Usage: # Create a managed agent on the provider side result = litellm.interactions.agents.create( - name="waverunner", + id="waverunner", custom_llm_provider="gemini", api_key="...", - base_agent="gemini-2.5-flash", - instructions="You are a helpful assistant.", + base_agent="antigravity-preview-09-2026", + system_instruction="You are a helpful assistant.", ) Methods: diff --git a/litellm/interactions/agents/__init__.py b/litellm/interactions/agents/__init__.py index 21fd2dc0742..a81b35c2fc7 100644 --- a/litellm/interactions/agents/__init__.py +++ b/litellm/interactions/agents/__init__.py @@ -3,11 +3,11 @@ litellm.interactions.agents Full CRUD SDK for provider-side managed agents (e.g. Gemini v1beta/agents). - litellm.interactions.agents.create(name=..., ...) + litellm.interactions.agents.create(id=..., ...) litellm.interactions.agents.list(api_key=...) - litellm.interactions.agents.get(name=..., ...) - litellm.interactions.agents.delete(name=..., ...) - litellm.interactions.agents.list_versions(name=..., ...) + litellm.interactions.agents.get(id=..., ...) + litellm.interactions.agents.delete(id=..., ...) + litellm.interactions.agents.list_versions(id=..., ...) Async counterparts: acreate, alist, aget, adelete, alist_versions """ diff --git a/litellm/interactions/agents/http_handler.py b/litellm/interactions/agents/http_handler.py index 2afedc34d36..fa1e52e6842 100644 --- a/litellm/interactions/agents/http_handler.py +++ b/litellm/interactions/agents/http_handler.py @@ -17,6 +17,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.llms.base_llm.agents.transformation import BaseAgentsAPIConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.types.agents import ( + AgentCreateRequest, AgentCreateResponse, AgentDeleteResult, AgentListResponse, @@ -35,7 +36,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): def create_agent( self, agents_api_config: BaseAgentsAPIConfig, - name: str, + create_request: AgentCreateRequest, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, extra_headers: dict[str, str] | None = None, @@ -47,7 +48,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): if _is_async: return self.async_create_agent( agents_api_config=agents_api_config, - name=name, + create_request=create_request, litellm_params=litellm_params, logging_obj=logging_obj, extra_headers=extra_headers, @@ -63,12 +64,13 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): api_base=litellm_params.get("api_base"), litellm_params=dict(litellm_params), ) - data: Final = agents_api_config.transform_create_request(name=name, litellm_params=dict(litellm_params)) - if extra_body: - data.update(extra_body) + data: Final = { + **agents_api_config.transform_create_request(create_request), + **(extra_body or {}), + } logging_obj.pre_call( - input=name, + input=create_request.id, api_key="", additional_args={ "complete_input_dict": data, @@ -85,12 +87,12 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): original_response=response.text, additional_args={"complete_input_dict": data}, ) - return agents_api_config.transform_create_response(raw_response=response, name=name) + return agents_api_config.transform_create_response(raw_response=response, name=create_request.id) async def async_create_agent( self, agents_api_config: BaseAgentsAPIConfig, - name: str, + create_request: AgentCreateRequest, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, extra_headers: dict[str, str] | None = None, @@ -106,12 +108,13 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): api_base=litellm_params.get("api_base"), litellm_params=dict(litellm_params), ) - data: Final = agents_api_config.transform_create_request(name=name, litellm_params=dict(litellm_params)) - if extra_body: - data.update(extra_body) + data: Final = { + **agents_api_config.transform_create_request(create_request), + **(extra_body or {}), + } logging_obj.pre_call( - input=name, + input=create_request.id, api_key="", additional_args={ "complete_input_dict": data, @@ -130,7 +133,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): original_response=response.text, additional_args={"complete_input_dict": data}, ) - return agents_api_config.transform_create_response(raw_response=response, name=name) + return agents_api_config.transform_create_response(raw_response=response, name=create_request.id) # ------------------------------------------------------------------ # # LIST # diff --git a/litellm/interactions/agents/main.py b/litellm/interactions/agents/main.py index 02b0ef4807f..5e2bfa9cea2 100644 --- a/litellm/interactions/agents/main.py +++ b/litellm/interactions/agents/main.py @@ -6,41 +6,47 @@ Usage: # Create response = litellm.interactions.agents.create( - name="waverunner", + id="waverunner", custom_llm_provider="gemini", api_key="...", - base_agent="gemini-2.5-flash", - instructions="You are a helpful assistant.", + base_agent="antigravity-preview-09-2026", + system_instruction="You are a helpful assistant.", + agent_config={"type": "antigravity", "model": "gemini-3.8-flash"}, ) # List response = litellm.interactions.agents.list(api_key="...", custom_llm_provider="gemini") # Get - response = litellm.interactions.agents.get(name="waverunner", api_key="...") + response = litellm.interactions.agents.get(id="waverunner", api_key="...") # Delete - result = litellm.interactions.agents.delete(name="waverunner", api_key="...") + result = litellm.interactions.agents.delete(id="waverunner", api_key="...") # List versions - result = litellm.interactions.agents.list_versions(name="waverunner", api_key="...") + result = litellm.interactions.agents.list_versions(id="waverunner", api_key="...") # Async versions: acreate, alist, aget, adelete, alist_versions """ import asyncio import contextvars -from collections.abc import Coroutine, Mapping +from collections.abc import Coroutine, Mapping, Sequence +from dataclasses import dataclass from functools import partial +from types import MappingProxyType from typing import Any, Final import httpx +from pydantic import ValidationError import litellm from litellm.interactions.agents.http_handler import agents_http_handler from litellm.interactions.agents.utils import get_provider_agents_api_config +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.agents import ( + AgentCreateRequest, AgentCreateResponse, AgentDeleteResult, AgentListResponse, @@ -89,14 +95,120 @@ def _make_logging_obj( return litellm_logging_obj +@dataclass(frozen=True, slots=True) +class _InvalidAgentRequest: + message: str + + +def _alias_conflict(field: str, value: str | None, alias: str, alias_value: str | None) -> str | None: + if value is None or alias_value is None or value == alias_value: + return None + return f"'{alias}' is an alias of '{field}', but the request gives them different values. Pass only '{field}'." + + +def _agent_id_or_error(id: str | None, name: str | None) -> str | _InvalidAgentRequest: + conflict: Final = _alias_conflict("id", id, "name", name) + if conflict is not None: + return _InvalidAgentRequest(conflict) + agent_id: Final = id if id is not None else name + if agent_id is None: + return _InvalidAgentRequest("Missing required parameter: 'id' (or its alias 'name').") + return agent_id + + +def _required_agent_id(id: str | None, name: str | None, llm_provider: str) -> str: + agent_id: Final = _agent_id_or_error(id, name) + if isinstance(agent_id, _InvalidAgentRequest): + raise litellm.BadRequestError(message=agent_id.message, model=id or name or "", llm_provider=llm_provider) + return agent_id + + # ================================================================== # # CREATE # # ================================================================== # +_CREATE_FIELD_SHAPES: Final = MappingProxyType( + { + "id": "a string", + "base_agent": "a string", + "agent_config": "an object", + "system_instruction": "a string", + "base_environment": "an environment id string or an object", + "description": "a string", + "tools": "a list of objects", + } +) + + +def _invalid_create_fields_message(invalid_fields: frozenset[str], caller_names: Mapping[str, str]) -> str: + problems: Final = "; ".join( + f"'{caller_names.get(field, field)}' must be {_CREATE_FIELD_SHAPES[field]}" for field in sorted(invalid_fields) + ) + drop_hint: Final = "" if "id" in invalid_fields else " Set drop_params to drop invalid optional fields." + return f"Invalid agent create request: {problems}.{drop_hint}" + + +def _validated_create_request( + fields: Mapping[str, object], drop_invalid_fields: bool, caller_names: Mapping[str, str] +) -> AgentCreateRequest | _InvalidAgentRequest: + try: + return AgentCreateRequest.model_validate(fields) + except ValidationError as e: + invalid_fields: Final = frozenset(str(error["loc"][0]) for error in e.errors()) + if drop_invalid_fields and "id" not in invalid_fields: + return AgentCreateRequest.model_validate( + {field: value for field, value in fields.items() if field not in invalid_fields} + ) + return _InvalidAgentRequest(_invalid_create_fields_message(invalid_fields, caller_names)) + + +def _create_request_or_error( + id: str | None, + name: str | None, + base_agent: str | None, + agent_config: Mapping[str, object] | None, + system_instruction: str | None, + instructions: str | None, + base_environment: InteractionEnvironment | None, + description: str | None, + tools: Sequence[Mapping[str, object]] | None, + drop_invalid_fields: bool, +) -> AgentCreateRequest | _InvalidAgentRequest: + agent_id: Final = _agent_id_or_error(id, name) + if isinstance(agent_id, _InvalidAgentRequest): + return agent_id + instruction_conflict: Final = _alias_conflict( + "system_instruction", system_instruction, "instructions", instructions + ) + if instruction_conflict is not None and not drop_invalid_fields: + return _InvalidAgentRequest(instruction_conflict) + sent_only_instructions: Final = system_instruction is None and instructions is not None + return _validated_create_request( + MappingProxyType( + { + "id": agent_id, + "base_agent": base_agent, + "agent_config": agent_config, + "system_instruction": instructions if sent_only_instructions else system_instruction, + "base_environment": base_environment, + "description": description, + "tools": tools, + } + ), + drop_invalid_fields, + MappingProxyType( + { + "id": "name" if id is None else "id", + "system_instruction": "instructions" if sent_only_instructions else "system_instruction", + } + ), + ) + + @client async def acreate( - name: str, + id: str | None = None, base_agent: str | None = None, instructions: str | None = None, base_environment: InteractionEnvironment | None = None, @@ -104,6 +216,12 @@ async def acreate( extra_headers: dict[str, str] | None = None, extra_body: Mapping[str, object] | None = None, timeout: float | httpx.Timeout | None = None, + *, + name: str | None = None, + agent_config: Mapping[str, object] | None = None, + system_instruction: str | None = None, + description: str | None = None, + tools: Sequence[Mapping[str, object]] | None = None, **kwargs, ) -> AgentCreateResponse: """Async: Create a managed agent on the provider side.""" @@ -113,7 +231,7 @@ async def acreate( kwargs["acreate_agent"] = True func: Final = partial( create, - name=name, + id=id, base_agent=base_agent, instructions=instructions, base_environment=base_environment, @@ -121,6 +239,11 @@ async def acreate( extra_headers=extra_headers, extra_body=extra_body, timeout=timeout, + name=name, + agent_config=agent_config, + system_instruction=system_instruction, + description=description, + tools=tools, **kwargs, ) ctx: Final = contextvars.copy_context() @@ -130,7 +253,7 @@ async def acreate( return init_response except Exception as e: raise litellm.exception_type( - model=name, + model=id or name or "", custom_llm_provider=custom_llm_provider or "gemini", original_exception=e, completion_kwargs=local_vars, @@ -140,7 +263,7 @@ async def acreate( @client def create( - name: str, + id: str | None = None, base_agent: str | None = None, instructions: str | None = None, base_environment: InteractionEnvironment | None = None, @@ -148,40 +271,61 @@ def create( extra_headers: dict[str, str] | None = None, extra_body: Mapping[str, object] | None = None, timeout: float | httpx.Timeout | None = None, + *, + name: str | None = None, + agent_config: Mapping[str, object] | None = None, + system_instruction: str | None = None, + description: str | None = None, + tools: Sequence[Mapping[str, object]] | None = None, **kwargs, ) -> AgentCreateResponse | Coroutine[object, object, AgentCreateResponse]: """ Sync: Create a managed agent on the provider side. Args: - name: Name for the agent (required). - base_agent: Base agent to derive from (e.g. "waverunner"). - instructions: System instructions for the agent. + id: Identifier for the agent. Exactly one of ``id`` or its alias ``name`` is required. + base_agent: Base agent to derive from (e.g. "antigravity-preview-09-2026"). + instructions: Alias of ``system_instruction``. base_environment: Environment to fork from — an env_id string or a dict like ``{"type": "remote", "sources": [...]}``. custom_llm_provider: Provider to use, e.g. "gemini". extra_headers: Additional HTTP headers. extra_body: Additional request body fields. timeout: Request timeout. + name: Alias of ``id``. + agent_config: Agent runtime config, e.g. ``{"type": "antigravity", "model": "gemini-3.8-flash"}``. + system_instruction: System instruction for the agent. + description: Description of the agent for developers. + tools: Tools available to the agent, e.g. ``[{"type": "code_execution"}]``. **kwargs: Forwarded to GenericLiteLLMParams (api_key, api_base, etc.). """ local_vars: Final = locals() - custom_llm_provider = custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini" + custom_llm_provider = custom_llm_provider or "gemini" try: _is_async: Final = kwargs.pop("acreate_agent", False) is True - if base_agent is not None: - kwargs["base_agent"] = base_agent - if instructions is not None: - kwargs["instructions"] = instructions - if base_environment is not None: - kwargs["base_environment"] = base_environment + create_request: Final = _create_request_or_error( + id=id, + name=name, + base_agent=base_agent, + agent_config=agent_config, + system_instruction=system_instruction, + instructions=instructions, + base_environment=base_environment, + description=description, + tools=tools, + drop_invalid_fields=drops_unsupported_params(kwargs.get("drop_params")), + ) + if isinstance(create_request, _InvalidAgentRequest): + raise litellm.BadRequestError( + message=create_request.message, model=id or name or "", llm_provider=custom_llm_provider + ) kwargs.setdefault("custom_llm_provider", custom_llm_provider) litellm_params: Final = GenericLiteLLMParams(**kwargs) - logging_obj: Final = _make_logging_obj(kwargs, name, custom_llm_provider, "create_agent", {}) + logging_obj: Final = _make_logging_obj(kwargs, create_request.id, custom_llm_provider, "create_agent", {}) config: Final = _get_agents_api_config(custom_llm_provider) return agents_http_handler.create_agent( agents_api_config=config, - name=name, + create_request=create_request, litellm_params=litellm_params, logging_obj=logging_obj, extra_headers=extra_headers, @@ -191,7 +335,7 @@ def create( ) except Exception as e: raise litellm.exception_type( - model=name, + model=id or name or "", custom_llm_provider=custom_llm_provider, original_exception=e, completion_kwargs=local_vars, @@ -247,7 +391,7 @@ def list( ) -> AgentListResponse | Coroutine[object, object, AgentListResponse]: """Sync: List all agents on the provider side.""" local_vars: Final = locals() - custom_llm_provider = custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini" + custom_llm_provider = custom_llm_provider or "gemini" try: _is_async: Final = kwargs.pop("alist_agents", False) is True kwargs.setdefault("custom_llm_provider", custom_llm_provider) @@ -279,19 +423,22 @@ def list( @client async def aget( - name: str, + id: str | None = None, custom_llm_provider: str | None = None, extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, + *, + name: str | None = None, **kwargs, ) -> AgentCreateResponse: - """Async: Get a specific agent by name.""" + """Async: Get a specific agent by id (``name`` is an alias of ``id``).""" local_vars: Final = locals() try: loop: Final = asyncio.get_event_loop() kwargs["aget_agent"] = True func: Final = partial( get, + id=id, name=name, custom_llm_provider=custom_llm_provider or "gemini", extra_headers=extra_headers, @@ -305,7 +452,7 @@ async def aget( return init_response except Exception as e: raise litellm.exception_type( - model=name, + model=id or name or "", custom_llm_provider=custom_llm_provider or "gemini", original_exception=e, completion_kwargs=local_vars, @@ -315,24 +462,27 @@ async def aget( @client def get( - name: str, + id: str | None = None, custom_llm_provider: str | None = None, extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, + *, + name: str | None = None, **kwargs, ) -> AgentCreateResponse | Coroutine[object, object, AgentCreateResponse]: - """Sync: Get a specific agent by name.""" + """Sync: Get a specific agent by id (``name`` is an alias of ``id``).""" local_vars: Final = locals() - custom_llm_provider = custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini" + custom_llm_provider = custom_llm_provider or "gemini" try: _is_async: Final = kwargs.pop("aget_agent", False) is True + agent_id: Final = _required_agent_id(id, name, custom_llm_provider) kwargs.setdefault("custom_llm_provider", custom_llm_provider) litellm_params: Final = GenericLiteLLMParams(**kwargs) - logging_obj: Final = _make_logging_obj(kwargs, name, custom_llm_provider, "get_agent", {"name": name}) + logging_obj: Final = _make_logging_obj(kwargs, agent_id, custom_llm_provider, "get_agent", {"name": agent_id}) config: Final = _get_agents_api_config(custom_llm_provider) return agents_http_handler.get_agent( agents_api_config=config, - name=name, + name=agent_id, litellm_params=litellm_params, logging_obj=logging_obj, extra_headers=extra_headers, @@ -341,7 +491,7 @@ def get( ) except Exception as e: raise litellm.exception_type( - model=name, + model=id or name or "", custom_llm_provider=custom_llm_provider, original_exception=e, completion_kwargs=local_vars, @@ -356,19 +506,22 @@ def get( @client async def adelete( - name: str, + id: str | None = None, custom_llm_provider: str | None = None, extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, + *, + name: str | None = None, **kwargs, ) -> AgentDeleteResult: - """Async: Delete a specific agent by name.""" + """Async: Delete a specific agent by id (``name`` is an alias of ``id``).""" local_vars: Final = locals() try: loop: Final = asyncio.get_event_loop() kwargs["adelete_agent"] = True func: Final = partial( delete, + id=id, name=name, custom_llm_provider=custom_llm_provider or "gemini", extra_headers=extra_headers, @@ -382,7 +535,7 @@ async def adelete( return init_response except Exception as e: raise litellm.exception_type( - model=name, + model=id or name or "", custom_llm_provider=custom_llm_provider or "gemini", original_exception=e, completion_kwargs=local_vars, @@ -392,24 +545,29 @@ async def adelete( @client def delete( - name: str, + id: str | None = None, custom_llm_provider: str | None = None, extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, + *, + name: str | None = None, **kwargs, ) -> AgentDeleteResult | Coroutine[object, object, AgentDeleteResult]: - """Sync: Delete a specific agent by name.""" + """Sync: Delete a specific agent by id (``name`` is an alias of ``id``).""" local_vars: Final = locals() - custom_llm_provider = custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini" + custom_llm_provider = custom_llm_provider or "gemini" try: _is_async: Final = kwargs.pop("adelete_agent", False) is True + agent_id: Final = _required_agent_id(id, name, custom_llm_provider) kwargs.setdefault("custom_llm_provider", custom_llm_provider) litellm_params: Final = GenericLiteLLMParams(**kwargs) - logging_obj: Final = _make_logging_obj(kwargs, name, custom_llm_provider, "delete_agent", {"name": name}) + logging_obj: Final = _make_logging_obj( + kwargs, agent_id, custom_llm_provider, "delete_agent", {"name": agent_id} + ) config: Final = _get_agents_api_config(custom_llm_provider) return agents_http_handler.delete_agent( agents_api_config=config, - name=name, + name=agent_id, litellm_params=litellm_params, logging_obj=logging_obj, extra_headers=extra_headers, @@ -418,7 +576,7 @@ def delete( ) except Exception as e: raise litellm.exception_type( - model=name, + model=id or name or "", custom_llm_provider=custom_llm_provider, original_exception=e, completion_kwargs=local_vars, @@ -433,19 +591,22 @@ def delete( @client async def alist_versions( - name: str, + id: str | None = None, custom_llm_provider: str | None = None, extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, + *, + name: str | None = None, **kwargs, ) -> AgentVersionsResponse: - """Async: List versions of a specific agent.""" + """Async: List versions of a specific agent by id (``name`` is an alias of ``id``).""" local_vars: Final = locals() try: loop: Final = asyncio.get_event_loop() kwargs["alist_agent_versions"] = True func: Final = partial( list_versions, + id=id, name=name, custom_llm_provider=custom_llm_provider or "gemini", extra_headers=extra_headers, @@ -459,7 +620,7 @@ async def alist_versions( return init_response except Exception as e: raise litellm.exception_type( - model=name, + model=id or name or "", custom_llm_provider=custom_llm_provider or "gemini", original_exception=e, completion_kwargs=local_vars, @@ -469,24 +630,29 @@ async def alist_versions( @client def list_versions( - name: str, + id: str | None = None, custom_llm_provider: str | None = None, extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, + *, + name: str | None = None, **kwargs, ) -> AgentVersionsResponse | Coroutine[object, object, AgentVersionsResponse]: - """Sync: List versions of a specific agent.""" + """Sync: List versions of a specific agent by id (``name`` is an alias of ``id``).""" local_vars: Final = locals() - custom_llm_provider = custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini" + custom_llm_provider = custom_llm_provider or "gemini" try: _is_async: Final = kwargs.pop("alist_agent_versions", False) is True + agent_id: Final = _required_agent_id(id, name, custom_llm_provider) kwargs.setdefault("custom_llm_provider", custom_llm_provider) litellm_params: Final = GenericLiteLLMParams(**kwargs) - logging_obj: Final = _make_logging_obj(kwargs, name, custom_llm_provider, "list_agent_versions", {"name": name}) + logging_obj: Final = _make_logging_obj( + kwargs, agent_id, custom_llm_provider, "list_agent_versions", {"name": agent_id} + ) config: Final = _get_agents_api_config(custom_llm_provider) return agents_http_handler.list_agent_versions( agents_api_config=config, - name=name, + name=agent_id, litellm_params=litellm_params, logging_obj=logging_obj, extra_headers=extra_headers, @@ -495,7 +661,7 @@ def list_versions( ) except Exception as e: raise litellm.exception_type( - model=name, + model=id or name or "", custom_llm_provider=custom_llm_provider, original_exception=e, completion_kwargs=local_vars, diff --git a/litellm/interactions/main.py b/litellm/interactions/main.py index 498e6e20aec..acd43c9b830 100644 --- a/litellm/interactions/main.py +++ b/litellm/interactions/main.py @@ -49,6 +49,7 @@ from litellm.interactions.utils import ( InteractionsAPIRequestUtils, get_provider_interactions_api_config, ) +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.interactions import ( CancelInteractionResult, @@ -100,6 +101,8 @@ async def acreate( timeout: float | httpx.Timeout | None = None, # LiteLLM params custom_llm_provider: str | None = None, + *, + agent_config: Mapping[str, object] | None = None, **kwargs, ) -> InteractionsAPIResponse | AsyncIterator[InteractionsAPIStreamingResponse]: """ @@ -121,6 +124,7 @@ async def acreate( string, or a config object such as ``{"type": "remote", "sources": [...]}`` / ``{"type": "remote", "network": {...}}`` + agent_config: Agent runtime config, e.g. ``{"type": "antigravity", "model": "gemini-3.8-flash"}`` response_modalities: Requested response modalities (TEXT, IMAGE, AUDIO) response_format: JSON schema for response format response_mime_type: MIME type of the response @@ -155,6 +159,7 @@ async def acreate( store=store, background=background, environment=environment, + agent_config=agent_config, response_modalities=response_modalities, response_format=response_format, response_mime_type=response_mime_type, @@ -225,6 +230,8 @@ def create( timeout: float | httpx.Timeout | None = None, # LiteLLM params custom_llm_provider: str | None = None, + *, + agent_config: Mapping[str, object] | None = None, **kwargs, ) -> ( InteractionsAPIResponse @@ -250,6 +257,7 @@ def create( string, or a config object such as ``{"type": "remote", "sources": [...]}`` / ``{"type": "remote", "network": {...}}`` + agent_config: Agent runtime config, e.g. ``{"type": "antigravity", "model": "gemini-3.8-flash"}`` response_modalities: Requested response modalities (TEXT, IMAGE, AUDIO) response_format: JSON schema for response format response_mime_type: MIME type of the response @@ -300,6 +308,15 @@ def create( # Check if this is a bridge provider (litellm_responses) - similar to responses API # Either provider is explicitly "litellm_responses" or no config found (bridge to responses) if custom_llm_provider == "litellm_responses" or interactions_api_config is None: + if optional_params.get("agent_config") is not None and not drops_unsupported_params( + kwargs.get("drop_params") + ): + raise litellm.UnsupportedParamsError( + message=f"'agent_config' is not supported for {custom_llm_provider} interactions. " + "Set drop_params to drop it.", + llm_provider=custom_llm_provider, + model=model, + ) # Bridge to litellm.responses() for non-native providers from litellm.interactions.litellm_responses_transformation.handler import ( LiteLLMResponsesInteractionsHandler, diff --git a/litellm/llms/base_llm/agents/transformation.py b/litellm/llms/base_llm/agents/transformation.py index 9d139b289c4..5ea6ea950d9 100644 --- a/litellm/llms/base_llm/agents/transformation.py +++ b/litellm/llms/base_llm/agents/transformation.py @@ -15,6 +15,7 @@ from collections.abc import Mapping import httpx from litellm.types.agents import ( + AgentCreateRequest, AgentCreateResponse, AgentDeleteResult, AgentListResponse, @@ -50,10 +51,9 @@ class BaseAgentsAPIConfig(ABC): @abstractmethod def transform_create_request( self, - name: str, - litellm_params: Mapping[str, object], - ) -> dict[str, object]: - """Map name + litellm_params to the provider's create-agent body.""" + create_request: AgentCreateRequest, + ) -> Mapping[str, object]: + """Map the create request to the provider's create-agent body.""" @abstractmethod def transform_create_response( diff --git a/litellm/llms/gemini/agents/transformation.py b/litellm/llms/gemini/agents/transformation.py index f7352da5fb6..45a6552a3fa 100644 --- a/litellm/llms/gemini/agents/transformation.py +++ b/litellm/llms/gemini/agents/transformation.py @@ -20,16 +20,13 @@ from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.agents.transformation import BaseAgentsAPIConfig from litellm.llms.gemini.common_utils import GeminiError, GeminiModelInfo from litellm.types.agents import ( + AgentCreateRequest, AgentCreateResponse, AgentDeleteResult, AgentListResponse, AgentVersionsResponse, ) -# Keys inside litellm_params that should be forwarded to the Gemini -# create-agent body verbatim. -_GEMINI_AGENT_BODY_KEYS: Final = ("base_agent", "instructions", "base_environment") - class _JsonResponse(Protocol): def json(self) -> object: ... @@ -39,30 +36,6 @@ def _json_body(response: _JsonResponse) -> object: return response.json() -# LiteLLM-internal keys that must never be forwarded to Gemini. -_LITELLM_INTERNAL_KEYS: Final = frozenset( - { - "custom_llm_provider", - "api_key", - "api_base", - "make_public", - "cost_per_query", - "input_cost_per_token", - "output_cost_per_token", - "require_trace_id_on_calls_to_agent", - "require_trace_id_on_calls_by_agent", - "max_iterations", - "max_budget_per_session", - "guardrails", - "is_public", - "agent_name", - "agent_id", - "agent_card_params", - "provider_agent_response", - } -) - - class GeminiAgentsConfig(BaseAgentsAPIConfig): """ Configuration for the Google AI Studio (Gemini) native Agents API. @@ -143,14 +116,9 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): def transform_create_request( self, - name: str, - litellm_params: Mapping[str, object], - ) -> dict[str, object]: - body: Final[dict[str, object]] = {"name": name} - for key in _GEMINI_AGENT_BODY_KEYS: - value = litellm_params.get(key) - if value is not None: - body[key] = value + create_request: AgentCreateRequest, + ) -> Mapping[str, object]: + body: Final[Mapping[str, object]] = create_request.model_dump(mode="json", exclude_none=True) verbose_logger.debug("GeminiAgentsConfig create body: %s", body) return body diff --git a/litellm/llms/gemini/interactions/transformation.py b/litellm/llms/gemini/interactions/transformation.py index 6f3e9a56b92..1895a97af7d 100644 --- a/litellm/llms/gemini/interactions/transformation.py +++ b/litellm/llms/gemini/interactions/transformation.py @@ -15,7 +15,7 @@ import httpx from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger -from litellm.litellm_core_utils.core_helpers import process_response_headers +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, process_response_headers from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.interactions.transformation import BaseInteractionsAPIConfig from litellm.llms.gemini.common_utils import GeminiError, GeminiModelInfo @@ -41,6 +41,14 @@ else: _JsonObject: TypeAlias = dict[str, object] +def _has_agent_config_type(agent_config: object) -> bool: + match agent_config: + case {"type": str()}: + return True + case _: + return False + + class _InteractionPayload(TypedDict, total=False): """JSON body of an Interactions API interaction, keyed as ``InteractionsAPIResponse`` fields.""" @@ -105,6 +113,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): return [ "model", "agent", + "agent_config", "input", "tools", "system_instruction", @@ -188,6 +197,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): request_body["input"] = input optional_keys: Final = [ + "agent_config", "tools", "system_instruction", "stream", @@ -197,8 +207,14 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): "response_modalities", "previous_interaction_id", ] + dropped_keys: Final = ( + frozenset({"agent_config"}) + if drops_unsupported_params(litellm_params.get("drop_params")) + and (model or not _has_agent_config_type(optional_params.get("agent_config"))) + else frozenset() + ) for key in optional_keys: - if optional_params.get(key) is not None: + if key not in dropped_keys and optional_params.get(key) is not None: request_body[key] = optional_params[key] # response_mime_type is removed — fold it into response_format. diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 2aaa987ba6a..b10b5ae98b5 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -15015,7 +15015,7 @@ ] }, "post": { - "description": "Create a named custom agent on the Gemini side.\n\nExample:\n```bash\ncurl -X POST \"http://localhost:4000/v1beta/agents\" \\\n -H \"Authorization: Bearer sk-...\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"name\": \"my-custom-slides-agent\",\n \"base_agent\": \"waverunner\",\n \"instructions\": \"You are a helpful assistant that creates slides.\",\n \"base_environment\": {\n \"type\": \"remote\",\n \"sources\": [\n {\"type\": \"gcs\", \"source\": \"gs://eap-templates/slides-skill\",\n \"target\": \"/.agents/skills/slides-skill\"}\n ]\n }\n }'\n```", + "description": "Create a named custom agent on the Gemini side.\n\nExample:\n```bash\ncurl -X POST \"http://localhost:4000/v1beta/agents\" \\\n -H \"Authorization: Bearer sk-...\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"id\": \"my-custom-slides-agent\",\n \"base_agent\": \"antigravity-preview-09-2026\",\n \"system_instruction\": \"You are a helpful assistant that creates slides.\",\n \"agent_config\": {\"type\": \"antigravity\", \"model\": \"gemini-3.8-flash\"},\n \"base_environment\": {\n \"type\": \"remote\",\n \"sources\": [\n {\"type\": \"gcs\", \"source\": \"gs://eap-templates/slides-skill\",\n \"target\": \"/.agents/skills/slides-skill\"}\n ]\n }\n }'\n```", "operationId": "create_gemini_agent_v1beta_agents_post", "responses": { "200": { diff --git a/litellm/proxy/google_endpoints/agents_endpoints.py b/litellm/proxy/google_endpoints/agents_endpoints.py index 00b23d846c2..7e3c2ddafe0 100644 --- a/litellm/proxy/google_endpoints/agents_endpoints.py +++ b/litellm/proxy/google_endpoints/agents_endpoints.py @@ -160,9 +160,10 @@ async def create_gemini_agent( -H "Authorization: Bearer sk-..." \\ -H "Content-Type: application/json" \\ -d '{ - "name": "my-custom-slides-agent", - "base_agent": "waverunner", - "instructions": "You are a helpful assistant that creates slides.", + "id": "my-custom-slides-agent", + "base_agent": "antigravity-preview-09-2026", + "system_instruction": "You are a helpful assistant that creates slides.", + "agent_config": {"type": "antigravity", "model": "gemini-3.8-flash"}, "base_environment": { "type": "remote", "sources": [ diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index cb226929167..db1951ec9e2 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -256,7 +256,6 @@ REQUIRED_PRESENT_BODY_PARAMS_BY_ROUTE: Final[Mapping[str, tuple[str, ...]]] = Ma "avideo_create_character": ("name", "video"), "acreate_container": ("name",), "aupload_container_file": ("file",), - "acreate_agent": ("name",), "acreate_interaction": ("input",), "acreate_eval": ("data_source_config", "testing_criteria"), "acreate_run": ("data_source",), @@ -264,7 +263,7 @@ REQUIRED_PRESENT_BODY_PARAMS_BY_ROUTE: Final[Mapping[str, tuple[str, ...]]] = Ma ) REQUIRED_ONE_OF_BODY_PARAMS_BY_ROUTE: Final[Mapping[str, tuple[str, str]]] = MappingProxyType( - {"acreate_interaction": ("model", "agent")} + {"acreate_interaction": ("model", "agent"), "acreate_agent": ("id", "name")} ) JSON_OBJECT_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object]) diff --git a/litellm/types/agents.py b/litellm/types/agents.py index 97dd67d3ff8..826aa5d9c98 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -3,7 +3,7 @@ from datetime import datetime from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, TypeAlias from urllib.parse import urlsplit -from pydantic import ConfigDict, Field, PrivateAttr, StrictInt, field_validator +from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, StrictInt, field_validator from typing_extensions import ReadOnly, Required, TypedDict from litellm.types.llms.base import LiteLLMBaseModel, LiteLLMPydanticObjectBase @@ -343,6 +343,18 @@ class ListAgentsResponse(LiteLLMBaseModel): agents: list[AgentResponse] +class AgentCreateRequest(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + id: str + base_agent: str | None = None + agent_config: Mapping[str, object] | None = None + system_instruction: str | None = None + base_environment: str | Mapping[str, object] | None = None + description: str | None = None + tools: Sequence[Mapping[str, object]] | None = None + + class AgentCreateResponse(LiteLLMPydanticObjectBase): """ Response from a provider-side agent creation or get call (e.g. Gemini v1beta/agents). diff --git a/tests/integration/compatibility/test_missing_body_param_status.py b/tests/integration/compatibility/test_missing_body_param_status.py index b2c5b1c74b6..005e581c643 100644 --- a/tests/integration/compatibility/test_missing_body_param_status.py +++ b/tests/integration/compatibility/test_missing_body_param_status.py @@ -254,8 +254,13 @@ _ROUTES: Final[dict[str, _Route]] = { "acreate_agent": ( "/v1beta/agents", "/v1beta/agents", - ("name",), - {"name": "agent", "base_agent": "waverunner", "instructions": "You are a helpful assistant."}, + ("id",), + { + "id": "agent", + "base_agent": "antigravity-preview-09-2026", + "system_instruction": "You are a helpful assistant.", + "agent_config": {"type": "antigravity", "model": "gemini-3.8-flash"}, + }, "gemini/gemini-2.5-flash", ), "acreate_interaction": ( @@ -1250,11 +1255,7 @@ def test_valid_agent_creation_reaches_upstream(gateway: Gateway) -> None: _assert_scripted_response("acreate_agent", caller) matches: Final = _observed(_Observations(gateway.upstream_url), identity) outbound: Final = object_value(matches[0]["body"]) - assert outbound == { - "name": "agent", - "base_agent": "waverunner", - "instructions": "You are a helpful assistant.", - }, matches + assert outbound == body, matches def test_valid_speech_returns_binary_audio_and_reaches_upstream(gateway: Gateway) -> None: diff --git a/tests/integration/providers/test_gemini_agents_wire.py b/tests/integration/providers/test_gemini_agents_wire.py index a619132d067..fd8f6b6ca63 100644 --- a/tests/integration/providers/test_gemini_agents_wire.py +++ b/tests/integration/providers/test_gemini_agents_wire.py @@ -75,9 +75,9 @@ def test_agents_forwards_caller_supplied_key_to_chosen_api_base(gateway: Gateway assert request.headers["x-goog-api-key"] == "synthetic-caller-key" assert _ENV_GEMINI_KEY not in request.headers.values() assert json.loads(request.body) == { - "name": _AGENT_NAME, + "id": _AGENT_NAME, "base_agent": "waverunner", - "instructions": "make slides", + "system_instruction": "make slides", } return Reply(body=json.dumps(_AGENT_REPLY).encode()) diff --git a/tests/unit/interactions/test_agents_http_handler.py b/tests/unit/interactions/test_agents_http_handler.py index 31b78d7a360..77e23e85e47 100644 --- a/tests/unit/interactions/test_agents_http_handler.py +++ b/tests/unit/interactions/test_agents_http_handler.py @@ -22,6 +22,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.gemini.agents.transformation import GeminiAgentsConfig from litellm.llms.gemini.common_utils import GeminiError from litellm.types.agents import ( + AgentCreateRequest, AgentCreateResponse, AgentDeleteResult, AgentListResponse, @@ -104,7 +105,7 @@ class TestCreateAgent: result = handler.create_agent( agents_api_config=config, - name="agent-x", + create_request=AgentCreateRequest(id="agent-x"), litellm_params=litellm_params, logging_obj=logging_obj, extra_headers={"X-Test": "1"}, @@ -117,8 +118,7 @@ class TestCreateAgent: client.post.assert_called_once() kwargs = client.post.call_args.kwargs assert kwargs["url"].endswith("/v1beta/agents") - assert kwargs["json"]["name"] == "agent-x" - assert kwargs["json"]["foo"] == "bar" + assert kwargs["json"] == {"id": "agent-x", "foo": "bar"} assert kwargs["headers"]["X-Test"] == "1" logging_obj.pre_call.assert_called_once() logging_obj.post_call.assert_called_once() @@ -130,7 +130,7 @@ class TestCreateAgent: result = handler.create_agent( agents_api_config=config, - name="agent-x", + create_request=AgentCreateRequest(id="agent-x"), litellm_params=litellm_params, logging_obj=_make_logging_obj(), client=client, @@ -152,7 +152,7 @@ class TestCreateAgent: with pytest.raises(GeminiError): handler.create_agent( agents_api_config=config, - name="agent-x", + create_request=AgentCreateRequest(id="agent-x"), litellm_params=litellm_params, logging_obj=_make_logging_obj(), client=client, @@ -169,7 +169,7 @@ class TestCreateAgent: result = await handler.async_create_agent( agents_api_config=config, - name="agent-y", + create_request=AgentCreateRequest(id="agent-y"), litellm_params=litellm_params, logging_obj=_make_logging_obj(), extra_body={"baz": "qux"}, @@ -193,7 +193,7 @@ class TestCreateAgent: with pytest.raises(GeminiError): await handler.async_create_agent( agents_api_config=config, - name="agent-y", + create_request=AgentCreateRequest(id="agent-y"), litellm_params=litellm_params, logging_obj=_make_logging_obj(), client=client, diff --git a/tests/unit/interactions/test_agents_main_and_utils.py b/tests/unit/interactions/test_agents_main_and_utils.py index 395801ff059..97d93563699 100644 --- a/tests/unit/interactions/test_agents_main_and_utils.py +++ b/tests/unit/interactions/test_agents_main_and_utils.py @@ -9,9 +9,14 @@ small helper utilities without touching the network. """ import asyncio +import json +from collections.abc import Callable, Mapping +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx import litellm @@ -137,7 +142,7 @@ class TestSyncEntryPoints: assert response is sentinel handler.create_agent.assert_called_once() kw = handler.create_agent.call_args.kwargs - assert kw["name"] == "waverunner" + assert kw["create_request"].id == "waverunner" assert kw["_is_async"] is False assert kw["extra_headers"] == {"X-Test": "1"} assert kw["extra_body"] == {"foo": "bar"} @@ -205,7 +210,7 @@ class TestAsyncEntryPoints: def fake_create_agent(**kwargs): assert kwargs["_is_async"] is True - assert kwargs["name"] == "waverunner" + assert kwargs["create_request"].id == "waverunner" return sentinel handler = MagicMock() @@ -349,3 +354,250 @@ class TestAsyncErrorWrapping: with patch(_HANDLER_PATH, handler): with pytest.raises(litellm.APIConnectionError): await alist_versions(name="waverunner", api_key="AIza") + + +_GEMINI_AGENTS_URL: Final = "https://generativelanguage.googleapis.com/v1beta/agents" + + +def _sent_json(route: respx.Route) -> object: + return json.loads(route.calls.last.request.content) + + +class TestCreateRequestBody: + def test_create_sends_every_current_field_verbatim(self, respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post(_GEMINI_AGENTS_URL).mock( + return_value=httpx.Response(200, json={"id": "sentinel", "base_agent": "antigravity-preview-09-2026"}) + ) + + response: Final = create( + id="sentinel", + base_agent="antigravity-preview-09-2026", + agent_config={"type": "antigravity", "model": "gemini-3.8-flash", "max_total_tokens": None}, + system_instruction="Always answer in exactly three words.", + base_environment={"type": "remote", "sources": [{"type": "gcs", "source": "gs://bucket/skill"}]}, + description="Watches the queue.", + tools=[{"type": "code_execution"}], + api_key="AIza-test", + ) + + assert _sent_json(route) == { + "id": "sentinel", + "base_agent": "antigravity-preview-09-2026", + "agent_config": {"type": "antigravity", "model": "gemini-3.8-flash", "max_total_tokens": None}, + "system_instruction": "Always answer in exactly three words.", + "base_environment": {"type": "remote", "sources": [{"type": "gcs", "source": "gs://bucket/skill"}]}, + "description": "Watches the queue.", + "tools": [{"type": "code_execution"}], + } + assert (response.id, response.name) == ("sentinel", "sentinel") + + @pytest.mark.parametrize( + ("identifier", "instruction", "expected"), + [ + pytest.param( + {"name": "slides"}, + {"instructions": "Make slides."}, + {"id": "slides", "system_instruction": "Make slides."}, + id="preview-aliases", + ), + pytest.param( + {"id": "slides", "name": "slides"}, + {"system_instruction": "Make slides.", "instructions": "Make slides."}, + {"id": "slides", "system_instruction": "Make slides."}, + id="same-value-under-both-names", + ), + ], + ) + def test_create_sends_preview_aliases_under_current_names( + self, + respx_mock: respx.MockRouter, + identifier: Mapping[str, str], + instruction: Mapping[str, str], + expected: Mapping[str, str], + ) -> None: + route: Final = respx_mock.post(_GEMINI_AGENTS_URL).mock( + return_value=httpx.Response(200, json={"id": "slides"}) + ) + + create(**identifier, **instruction, base_agent="antigravity-preview-09-2026", api_key="AIza-test") + + assert _sent_json(route) == {**expected, "base_agent": "antigravity-preview-09-2026"} + + @pytest.mark.parametrize( + "fields", + [ + pytest.param({"base_agent": "antigravity-preview-09-2026"}, id="no-identifier"), + pytest.param({"id": "slides", "name": "decks"}, id="different-id-and-name"), + pytest.param( + {"id": "slides", "system_instruction": "Make slides.", "instructions": "Make decks."}, + id="different-system-instruction-and-instructions", + ), + pytest.param({"id": "slides", "agent_config": "antigravity"}, id="agent-config-not-an-object"), + ], + ) + def test_create_rejects_a_malformed_identifier_or_instruction( + self, respx_mock: respx.MockRouter, fields: Mapping[str, str] + ) -> None: + with pytest.raises(litellm.BadRequestError) as exc_info: + create(**fields, api_key="AIza-test") + + assert exc_info.value.status_code == 400 + assert not respx_mock.calls + + def test_create_with_drop_params_drops_only_the_invalid_optional_fields( + self, respx_mock: respx.MockRouter + ) -> None: + route: Final = respx_mock.post(_GEMINI_AGENTS_URL).mock( + return_value=httpx.Response(200, json={"id": "slides"}) + ) + + create( + id="slides", + agent_config="antigravity", + tools="code_execution", + instructions="Make slides.", + drop_params=True, + api_key="AIza-test", + ) + + assert _sent_json(route) == {"id": "slides", "system_instruction": "Make slides."} + + def test_create_with_drop_params_drops_a_clashing_instructions_alias(self, respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post(_GEMINI_AGENTS_URL).mock( + return_value=httpx.Response(200, json={"id": "slides"}) + ) + + create( + id="slides", + system_instruction="Make slides.", + instructions="Make decks.", + drop_params=True, + api_key="AIza-test", + ) + + assert _sent_json(route) == {"id": "slides", "system_instruction": "Make slides."} + + def test_create_with_drop_params_still_rejects_a_clashing_identifier(self, respx_mock: respx.MockRouter) -> None: + with pytest.raises(litellm.BadRequestError) as exc_info: + create(id="slides", name="decks", drop_params=True, api_key="AIza-test") + + assert exc_info.value.status_code == 400 + assert not respx_mock.calls + + @pytest.mark.parametrize( + ("fields", "expected_message"), + [ + pytest.param( + {"id": "slides", "instructions": 5, "tools": "code_execution"}, + "Invalid agent create request: 'instructions' must be a string; 'tools' must be a list of objects." + " Set drop_params to drop invalid optional fields.", + id="names-the-alias-the-caller-sent", + ), + pytest.param( + {"name": 7, "base_agent": "antigravity-preview-09-2026"}, + "Invalid agent create request: 'name' must be a string.", + id="names-the-id-alias-the-caller-sent", + ), + pytest.param( + {"id": 7, "agent_config": "antigravity", "drop_params": True}, + "Invalid agent create request: 'agent_config' must be an object; 'id' must be a string.", + id="invalid-id-is-never-dropped", + ), + ], + ) + def test_create_rejects_invalid_fields_by_the_names_the_caller_sent( + self, respx_mock: respx.MockRouter, fields: Mapping[str, object], expected_message: str + ) -> None: + with pytest.raises(litellm.BadRequestError) as exc_info: + create(**fields, api_key="AIza-test") + + assert exc_info.value.status_code == 400 + assert exc_info.value.message.endswith(expected_message) + assert not respx_mock.calls + + @pytest.mark.asyncio + async def test_acreate_sends_every_current_field_verbatim( + self, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + route: Final = respx_mock.post(_GEMINI_AGENTS_URL).mock( + return_value=httpx.Response(200, json={"id": "sentinel"}) + ) + + response: Final = await acreate( + id="sentinel", + agent_config={"type": "antigravity", "model": "gemini-3.8-flash"}, + system_instruction="Always answer in exactly three words.", + api_key="AIza-test", + ) + + assert _sent_json(route) == { + "id": "sentinel", + "agent_config": {"type": "antigravity", "model": "gemini-3.8-flash"}, + "system_instruction": "Always answer in exactly three words.", + } + assert response.id == "sentinel" + + +class TestAgentIdentifier: + def test_get_by_id_fetches_that_agent(self, respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.get(f"{_GEMINI_AGENTS_URL}/slides").mock( + return_value=httpx.Response(200, json={"id": "slides", "base_agent": "antigravity-preview-09-2026"}) + ) + + response: Final = get(id="slides", api_key="AIza-test") + + assert route.call_count == 1 + assert (response.id, response.base_agent) == ("slides", "antigravity-preview-09-2026") + + def test_delete_by_id_deletes_that_agent(self, respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.delete(f"{_GEMINI_AGENTS_URL}/slides").mock( + return_value=httpx.Response(200, json={}) + ) + + result: Final = delete(id="slides", api_key="AIza-test") + + assert route.call_count == 1 + assert (result.name, result.deleted) == ("slides", True) + + def test_list_versions_by_id_lists_that_agents_versions(self, respx_mock: respx.MockRouter) -> None: + version: Final = {"agent": "slides", "name": "agents/slides/versions/v1"} + route: Final = respx_mock.get(f"{_GEMINI_AGENTS_URL}/slides/versions").mock( + return_value=httpx.Response(200, json={"agentVersions": [version]}) + ) + + response: Final = list_versions(id="slides", api_key="AIza-test") + + assert route.call_count == 1 + assert response.agent_versions == [version] + + @pytest.mark.parametrize( + "operation", + [pytest.param(get, id="get"), pytest.param(delete, id="delete"), pytest.param(list_versions, id="list_versions")], + ) + @pytest.mark.parametrize( + "identifier", + [ + pytest.param({}, id="no-identifier"), + pytest.param({"id": "slides", "name": "decks"}, id="different-id-and-name"), + ], + ) + def test_rejects_a_missing_or_conflicting_identifier( + self, + respx_mock: respx.MockRouter, + operation: Callable[..., object], + identifier: Mapping[str, str], + ) -> None: + with pytest.raises(litellm.BadRequestError) as exc_info: + operation(**identifier, api_key="AIza-test") + + assert exc_info.value.status_code == 400 + assert not respx_mock.calls + + @pytest.mark.asyncio + async def test_aget_rejects_a_missing_identifier(self, respx_mock: respx.MockRouter) -> None: + with pytest.raises(litellm.BadRequestError) as exc_info: + await aget(api_key="AIza-test") + + assert exc_info.value.status_code == 400 + assert not respx_mock.calls diff --git a/tests/unit/interactions/test_main.py b/tests/unit/interactions/test_main.py index 123f44ca779..5af5dd7d82a 100644 --- a/tests/unit/interactions/test_main.py +++ b/tests/unit/interactions/test_main.py @@ -236,6 +236,159 @@ class TestGoogleInteractionsCreate: assert json.loads(route.calls.last.request.content) == expected_body assert response.status == "completed" + @pytest.mark.usefixtures("fake_provider_credentials") + def test_create_with_agent_config_sends_it_to_gemini(self, respx_mock: MockRouter) -> None: + route: Final = respx_mock.post("https://generativelanguage.googleapis.com/v1beta/interactions").mock( + return_value=httpx.Response( + 200, + json={"id": "interaction-agent", "object": "interaction", "status": "completed"}, + ) + ) + + interactions.create( + agent="antigravity-preview-09-2026", + agent_config={"type": "antigravity", "model": "gemini-3.8-flash"}, + input="Say hi in one word.", + environment="remote", + api_key="gemini-offline", + ) + + assert json.loads(route.calls.last.request.content) == { + "agent": "antigravity-preview-09-2026", + "agent_config": {"type": "antigravity", "model": "gemini-3.8-flash"}, + "input": "Say hi in one word.", + "environment": "remote", + } + + @pytest.mark.usefixtures("fake_provider_credentials") + def test_create_keeps_positional_response_settings_in_place(self, respx_mock: MockRouter) -> None: + route: Final = respx_mock.post("https://generativelanguage.googleapis.com/v1beta/interactions").mock( + return_value=httpx.Response(200, json={"id": "interaction-model", "status": "completed"}) + ) + + interactions.create( + "gemini/gemini-2.5-flash", + None, + "Describe a cat.", + None, + None, + None, + None, + None, + None, + None, + ["TEXT"], + api_key="gemini-offline", + ) + + assert json.loads(route.calls.last.request.content) == { + "model": "gemini-2.5-flash", + "input": "Describe a cat.", + "response_modalities": ["TEXT"], + } + + @pytest.mark.parametrize( + ("target", "agent_config", "expected_target"), + [ + pytest.param( + {"model": "gemini/gemini-2.5-flash"}, + {"type": "antigravity", "model": "gemini-3.8-flash"}, + {"model": "gemini-2.5-flash"}, + id="model-drops-it", + ), + pytest.param( + {"agent": "antigravity-preview-09-2026"}, + {"type": "antigravity", "model": "gemini-3.8-flash"}, + { + "agent": "antigravity-preview-09-2026", + "agent_config": {"type": "antigravity", "model": "gemini-3.8-flash"}, + }, + id="agent-keeps-an-object", + ), + pytest.param( + {"agent": "antigravity-preview-09-2026"}, + "antigravity", + {"agent": "antigravity-preview-09-2026"}, + id="agent-drops-a-non-object", + ), + pytest.param( + {"agent": "antigravity-preview-09-2026"}, + {"model": "gemini-3.8-flash"}, + {"agent": "antigravity-preview-09-2026"}, + id="agent-drops-an-object-without-a-type", + ), + pytest.param( + {"agent": "antigravity-preview-09-2026"}, + {"type": 7, "model": "gemini-3.8-flash"}, + {"agent": "antigravity-preview-09-2026"}, + id="agent-drops-a-non-string-type", + ), + ], + ) + @pytest.mark.usefixtures("fake_provider_credentials") + def test_create_with_drop_params_drops_agent_config_gemini_cannot_take( + self, + respx_mock: MockRouter, + target: Mapping[str, str], + agent_config: object, + expected_target: Mapping[str, object], + ) -> None: + route: Final = respx_mock.post("https://generativelanguage.googleapis.com/v1beta/interactions").mock( + return_value=httpx.Response(200, json={"id": "interaction-drop", "status": "completed"}) + ) + + interactions.create( + **target, + agent_config=agent_config, + input="Say hi in one word.", + drop_params=True, + api_key="gemini-offline", + ) + + assert json.loads(route.calls.last.request.content) == {**expected_target, "input": "Say hi in one word."} + + @pytest.mark.usefixtures("fake_provider_credentials") + def test_create_rejects_agent_config_a_bridged_provider_cannot_take(self, respx_mock: MockRouter) -> None: + route: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, json=_openai_response_body("resp-agent-config", "Hi")) + ) + + with pytest.raises(litellm.UnsupportedParamsError, match="'agent_config' is not supported for openai") as e: + interactions.create( + model="gpt-5.6", + agent_config={"type": "antigravity", "model": "gemini-3.8-flash"}, + input="Say hi in one word.", + api_key="sk-offline", + ) + + assert e.value.status_code == 400 + assert not route.called + + @pytest.mark.parametrize( + "drop_kwargs", + [ + pytest.param({"drop_params": True}, id="drop-params"), + pytest.param({"additional_drop_params": ["agent_config"]}, id="additional-drop-params"), + ], + ) + @pytest.mark.usefixtures("fake_provider_credentials") + def test_create_with_drop_params_drops_agent_config_a_bridged_provider_cannot_take( + self, respx_mock: MockRouter, drop_kwargs: Mapping[str, object] + ) -> None: + route: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, json=_openai_response_body("resp-agent-config", "Hi")) + ) + + interactions.create( + model="gpt-5.6", + agent_config={"type": "antigravity", "model": "gemini-3.8-flash"}, + input="Say hi in one word.", + api_key="sk-offline", + **drop_kwargs, + ) + + assert json.loads(route.calls.last.request.content) == {"model": "gpt-5.6", "input": "Say hi in one word."} + @pytest.mark.parametrize( ("model", "api_key", "url", "expected_body", "sse_body", "expected_chunks"), [ diff --git a/tests/unit/proxy/google_endpoints/test_gemini_agents_endpoints.py b/tests/unit/proxy/google_endpoints/test_gemini_agents_endpoints.py index 79062f15eca..819e63dfd59 100644 --- a/tests/unit/proxy/google_endpoints/test_gemini_agents_endpoints.py +++ b/tests/unit/proxy/google_endpoints/test_gemini_agents_endpoints.py @@ -8,12 +8,18 @@ query parameter. Flat credential query params (e.g. ?api_key=…) are no longer accepted — they would appear in server logs. """ +import asyncio import json +from pathlib import Path +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx from fastapi import Request from fastapi.datastructures import Headers, QueryParams +from fastapi.testclient import TestClient from litellm.proxy.google_endpoints.agents_endpoints import ( @@ -470,3 +476,41 @@ def test_validate_environment_env_fallback_when_no_api_base_override(monkeypatch cfg = GeminiAgentsConfig() headers = cfg.validate_environment(headers={}, litellm_params={}) assert headers["x-goog-api-key"] == "AIzaFromEnv" + + +_PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "test_configs" / "test_bad_config.yaml" + + +@pytest.fixture +def proxy_client(monkeypatch: pytest.MonkeyPatch, httpx_transport: None) -> TestClient: + from litellm.proxy.proxy_server import app, cleanup_router_config_variables, initialize + + monkeypatch.setenv("LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY", "true") + cleanup_router_config_variables() + asyncio.run(initialize(config=str(_PROXY_CONFIG))) + return TestClient(app) + + +def test_create_agent_route_forwards_googles_create_body(proxy_client: TestClient, respx_mock: respx.MockRouter): + google_body: Final = { + "id": "sentinel", + "base_agent": "antigravity-preview-09-2026", + "agent_config": {"type": "antigravity", "model": "gemini-3.8-flash"}, + "system_instruction": "Always answer in exactly three words.", + "base_environment": {"type": "remote"}, + "description": "Watches the queue.", + "tools": [{"type": "code_execution"}], + } + upstream: Final = respx_mock.post("https://generativelanguage.googleapis.com/v1beta/agents").mock( + return_value=httpx.Response(200, json=google_body) + ) + + response: Final = proxy_client.post( + "/v1beta/agents", + json={**google_body, "litellm_params_template": {"api_key": "AIza-caller"}}, + ) + + assert response.status_code == 200, response.text + assert json.loads(upstream.calls.last.request.content) == google_body + assert upstream.calls.last.request.headers["x-goog-api-key"] == "AIza-caller" + assert response.json()["id"] == "sentinel" diff --git a/tests/unit/proxy/google_endpoints/test_interactions_agent_param.py b/tests/unit/proxy/google_endpoints/test_interactions_agent_param.py index f3cec320532..2a2a00d322e 100644 --- a/tests/unit/proxy/google_endpoints/test_interactions_agent_param.py +++ b/tests/unit/proxy/google_endpoints/test_interactions_agent_param.py @@ -6,9 +6,16 @@ The proxy must not pass the agent name as ``model`` or LiteLLM may route to openai/* wildcards instead of Gemini interactions. """ +import asyncio +import json +from pathlib import Path +from typing import Final from unittest.mock import MagicMock, patch +import httpx import pytest +import respx +from fastapi.testclient import TestClient class TestInteractionsAgentParameter: @@ -101,3 +108,37 @@ class TestInteractionsAgentOnlyProviderRouting: mock_get_llm_provider.assert_not_called() assert mock_handler.create_interaction.call_args.kwargs["model"] is None + + +_PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "test_configs" / "test_bad_config.yaml" + + +@pytest.fixture +def proxy_client(monkeypatch: pytest.MonkeyPatch, httpx_transport: None) -> TestClient: + from litellm.proxy.proxy_server import app, cleanup_router_config_variables, initialize + + monkeypatch.setenv("LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY", "true") + cleanup_router_config_variables() + asyncio.run(initialize(config=str(_PROXY_CONFIG))) + return TestClient(app) + + +@pytest.mark.usefixtures("fake_provider_credentials") +def test_create_interaction_route_forwards_agent_config(proxy_client: TestClient, respx_mock: respx.MockRouter): + request_body: Final = { + "agent": "antigravity-preview-09-2026", + "agent_config": {"type": "antigravity", "model": "gemini-3.8-flash"}, + "input": "Say hi in one word.", + "environment": "remote", + } + upstream: Final = respx_mock.post("https://generativelanguage.googleapis.com/v1beta/interactions").mock( + return_value=httpx.Response( + 200, + json={"id": "interaction-agent", "object": "interaction", "status": "completed"}, + ) + ) + + response: Final = proxy_client.post("/v1beta/interactions", json=request_body) + + assert response.status_code == 200, response.text + assert json.loads(upstream.calls.last.request.content) == request_body diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 6044475652c..9e31d3774f9 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -1494,7 +1494,7 @@ async def test_route_request_keeps_key_fallbacks_out_of_pass_through_provider_bo ("avideo_create_character", "name", "/videos/characters"), ("acreate_container", "name", "/containers"), ("aupload_container_file", "file", "/containers/{container_id}/files"), - ("acreate_agent", "name", "/v1beta/agents"), + ("acreate_agent", "id", "/v1beta/agents"), ("acreate_eval", "data_source_config", "/evals"), ("acreate_run", "data_source", "/evals/{eval_id}/runs"), ], @@ -1699,7 +1699,7 @@ def test_required_present_body_param_with_none_router_wide_default_still_raises( ("asearch", {"model": "search-model"}, "query"), ("acreate_eval", {"model": "eval-model"}, "data_source_config"), ("acreate_run", {"model": "eval-model"}, "data_source"), - ("acreate_agent", {"model": "agent-model"}, "name"), + ("acreate_agent", {"model": "agent-model"}, "id"), ("avector_store_search", {}, "query"), ], ) @@ -1842,6 +1842,8 @@ def test_raise_if_required_body_param_missing_names_each_missing_param(route_typ ("acreate_eval", {"data_source_config": {"type": "custom"}, "testing_criteria": []}), ("acreate_interaction", {"model": "gemini-3-pro-preview", "input": "hi"}), ("acreate_interaction", {"agent": "deep-research", "input": "hi"}), + ("acreate_agent", {"id": "sentinel"}), + ("acreate_agent", {"name": "sentinel"}), ("aimage_generation", {"model": "gpt-image-1", "prompt": "a cat"}), ("aspeech", {"model": "gpt-4o-mini-tts", "input": "hi", "voice": "alloy"}), ("amoderation", {"model": "omni-moderation-latest", "input": ""}), diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index f9117afdf3a..45f9630983f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -23366,9 +23366,10 @@ export interface paths { * -H "Authorization: Bearer sk-..." \ * -H "Content-Type: application/json" \ * -d '{ - * "name": "my-custom-slides-agent", - * "base_agent": "waverunner", - * "instructions": "You are a helpful assistant that creates slides.", + * "id": "my-custom-slides-agent", + * "base_agent": "antigravity-preview-09-2026", + * "system_instruction": "You are a helpful assistant that creates slides.", + * "agent_config": {"type": "antigravity", "model": "gemini-3.8-flash"}, * "base_environment": { * "type": "remote", * "sources": [ From 9675c13cb02b2e183deb04b0ae912ed85b20f1e2 Mon Sep 17 00:00:00 2001 From: m6atlas <349022726+m6atlas@users.noreply.github.com> Date: Wed, 28 Oct 2026 12:38:36 -0700 Subject: [PATCH 077/128] fix(openai)!: send the whole reasoning_effort dict when chat is bridged to Responses (#48807) * fix(openai): send the whole reasoning_effort dict when chat is bridged to Responses * fix(openai): drop only a refused effort and honor drop params on the bridged reasoning dict * fix(openai): drop generate_summary and forward a non-catalog model's reasoning dict whole on the bridge * fix(responses): drop a mode or context the model does not take under drop_params Read each OpenAI model's reasoning_modes and reasoning_contexts from the model map when drop_params filters the bridged reasoning object, send no reasoning object when validation dropped the effort, and ignore an additional_drop_params value that is not a list. * fix(responses): drop only a refused effort and forward an unlisted mode or context on the bridge --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../transformation.py | 19 +- litellm/main.py | 13 +- ...odel_prices_and_context_window_backup.json | 318 +++++++++++++++++ litellm/responses/utils.py | 101 +++++- litellm/types/utils.py | 2 + litellm/utils.py | 2 + model_prices_and_context_window.json | 318 +++++++++++++++++ tests/unit/test_main.py | 324 ++++++++++++++++++ tests/unit/test_utils.py | 2 + 9 files changed, 1083 insertions(+), 16 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 98322aecf6a..47a3a903d82 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -3,7 +3,6 @@ Handler for transforming /chat/completions api requests to litellm.responses req """ import json -import os from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence from itertools import accumulate, chain from types import MappingProxyType @@ -43,8 +42,8 @@ from litellm.responses.sse_output_recovery import ( ) from litellm.responses.utils import ( ResponsesAPIRequestUtils, - auto_reasoning_summary, normalize_responses_api_stream_options, + opt_in_auto_reasoning_summary, ) from litellm.types.llms.openai import ( ChatCompletionAnnotation, @@ -1745,19 +1744,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): **cast(Reasoning, reasoning_effort) # cast-ok: dict is forwarded verbatim to the provider ) - # Check if auto-summary is enabled via flag or environment variable - # Priority: litellm.reasoning_auto_summary flag > LITELLM_REASONING_AUTO_SUMMARY env var - auto_summary_enabled: Final = ( - litellm.reasoning_auto_summary or os.getenv("LITELLM_REASONING_AUTO_SUMMARY", "false").lower() == "true" - ) - + summary: Final = opt_in_auto_reasoning_summary(model, custom_llm_provider) return ( - Reasoning( - effort=reasoning_effort, - summary="detailed" if model is None else auto_reasoning_summary(model, custom_llm_provider), - ) - if auto_summary_enabled - else Reasoning(effort=reasoning_effort) + Reasoning(effort=reasoning_effort) + if summary is None + else Reasoning(effort=reasoning_effort, summary=summary) ) _map_reasoning_effort = map_reasoning_effort diff --git a/litellm/main.py b/litellm/main.py index 75bd4431120..41c4302499e 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6124,7 +6124,18 @@ def completion( optional_params, rs_val = strip_reasoning_summary_aliases_from_optional_params(optional_params) - if isinstance(resolved_reasoning_effort, dict) and "summary" in resolved_reasoning_effort: + if custom_llm_provider == "openai" and isinstance(resolved_reasoning_effort, dict): + from litellm.responses.utils import openai_bridge_reasoning + + optional_params["reasoning_effort"] = openai_bridge_reasoning( + requested=resolved_reasoning_effort, + effort_kept="reasoning_effort" in optional_params, + summary_alias=rs_val, + model=model, + drop_params=drop_unsupported_reasoning, + additional_drop_params=routing_kwargs.get("additional_drop_params"), + ) + elif isinstance(resolved_reasoning_effort, dict) and "summary" in resolved_reasoning_effort: optional_params["reasoning_effort"] = resolved_reasoning_effort elif rs_val is not None: eff: Final = optional_params.get("reasoning_effort", resolved_reasoning_effort) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 7199a075612..825de5ac3ee 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -35598,6 +35598,13 @@ "output_cost_per_token_batches": 5e-06, "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -35675,6 +35682,13 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_priority": 2e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -35851,6 +35865,13 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_priority": 2e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -35902,6 +35923,13 @@ "mode": "chat", "output_cost_per_token": 1.4e-05, "output_cost_per_token_priority": 2.8e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -36017,6 +36045,13 @@ "mode": "chat", "output_cost_per_token": 1.4e-05, "output_cost_per_token_priority": 2.8e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -36066,6 +36101,13 @@ "mode": "responses", "output_cost_per_token": 0.000168, "output_cost_per_token_batches": 8.4e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "pro" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -36107,6 +36149,13 @@ "mode": "responses", "output_cost_per_token": 0.000168, "output_cost_per_token_batches": 8.4e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "pro" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -36194,6 +36243,15 @@ "output_cost_per_token_flex": 2.5e-05, "output_cost_per_token_priority": 0.0001, "output_cost_per_token_ultrafast": 0.0003, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36278,6 +36336,15 @@ "output_cost_per_token_above_272k_tokens_batches": 7.5e-06, "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36370,6 +36437,15 @@ "output_cost_per_token_above_272k_tokens_batches": 3.75e-07, "output_cost_per_token_flex": 2.5e-07, "output_cost_per_token_priority": 1e-06, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36453,6 +36529,15 @@ "output_cost_per_token_above_272k_tokens_batches": 1.5e-05, "output_cost_per_token_flex": 1e-05, "output_cost_per_token_priority": 4e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36534,6 +36619,15 @@ "output_cost_per_token_above_272k_tokens_batches": 1.5e-05, "output_cost_per_token_flex": 1e-05, "output_cost_per_token_priority": 4e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36617,6 +36711,15 @@ "output_cost_per_token_above_272k_tokens_batches": 9e-06, "output_cost_per_token_flex": 6e-06, "output_cost_per_token_priority": 2.4e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36699,6 +36802,15 @@ "output_cost_per_token_above_272k_tokens_batches": 9e-07, "output_cost_per_token_flex": 6e-07, "output_cost_per_token_priority": 2.4e-06, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36936,6 +37048,14 @@ "reasoning_effort_levels": [ "medium" ], + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -36983,6 +37103,14 @@ "output_cost_per_token_batches": 1.5e-05, "output_cost_per_token_above_272k_tokens_batches": 2.25e-05, "output_cost_per_token_priority": 7.5e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37045,6 +37173,14 @@ "output_cost_per_token_batches": 1.5e-05, "output_cost_per_token_above_272k_tokens_batches": 2.25e-05, "output_cost_per_token_priority": 7.5e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37099,6 +37235,14 @@ "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, "output_cost_per_token_above_272k_tokens_batches": 0.000135, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37150,6 +37294,14 @@ "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, "output_cost_per_token_above_272k_tokens_batches": 0.000135, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37209,6 +37361,14 @@ "output_cost_per_token_batches": 7.5e-06, "output_cost_per_token_above_272k_tokens_batches": 1.125e-05, "output_cost_per_token_priority": 3e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ @@ -37266,6 +37426,14 @@ "output_cost_per_token_batches": 7.5e-06, "output_cost_per_token_above_272k_tokens_batches": 1.125e-05, "output_cost_per_token_priority": 3e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ @@ -37315,6 +37483,14 @@ "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, "output_cost_per_token_above_272k_tokens_batches": 0.000135, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37367,6 +37543,14 @@ "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, "output_cost_per_token_above_272k_tokens_batches": 0.000135, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37421,6 +37605,14 @@ "output_cost_per_token_flex": 2.25e-06, "output_cost_per_token_batches": 2.25e-06, "output_cost_per_token_priority": 9e-06, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37475,6 +37667,14 @@ "output_cost_per_token_flex": 2.25e-06, "output_cost_per_token_batches": 2.25e-06, "output_cost_per_token_priority": 9e-06, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37526,6 +37726,14 @@ "output_cost_per_token": 1.25e-06, "output_cost_per_token_flex": 6.25e-07, "output_cost_per_token_batches": 6.25e-07, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37578,6 +37786,14 @@ "output_cost_per_token": 1.25e-06, "output_cost_per_token_flex": 6.25e-07, "output_cost_per_token_batches": 6.25e-07, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37624,6 +37840,13 @@ "mode": "responses", "output_cost_per_token": 0.00012, "output_cost_per_token_batches": 6e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "pro" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -37666,6 +37889,13 @@ "mode": "responses", "output_cost_per_token": 0.00012, "output_cost_per_token_batches": 6e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "pro" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -37716,6 +37946,13 @@ "output_cost_per_token_batches": 5e-06, "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -37826,6 +38063,13 @@ "mode": "responses", "output_cost_per_token": 1.4e-05, "output_cost_per_token_priority": 2.8e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -37908,6 +38152,13 @@ "output_cost_per_token_batches": 1e-06, "output_cost_per_token_flex": 1e-06, "output_cost_per_token_priority": 3.6e-06, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -37959,6 +38210,13 @@ "output_cost_per_token_batches": 1e-06, "output_cost_per_token_flex": 1e-06, "output_cost_per_token_priority": 3.6e-06, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -38007,6 +38265,13 @@ "output_cost_per_token": 4e-07, "output_cost_per_token_batches": 2e-07, "output_cost_per_token_flex": 2e-07, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -38055,6 +38320,13 @@ "output_cost_per_token": 4e-07, "output_cost_per_token_batches": 2e-07, "output_cost_per_token_flex": 2e-07, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -43821,6 +44093,13 @@ "output_cost_per_token_batches": 4e-06, "output_cost_per_token_flex": 4e-06, "output_cost_per_token_priority": 1.4e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -43929,6 +44208,13 @@ "output_cost_per_token_batches": 4e-06, "output_cost_per_token_flex": 4e-06, "output_cost_per_token_priority": 1.4e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -44008,6 +44294,13 @@ "mode": "responses", "output_cost_per_token": 8e-05, "output_cost_per_token_batches": 4e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "pro" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -44045,6 +44338,13 @@ "mode": "responses", "output_cost_per_token": 8e-05, "output_cost_per_token_batches": 4e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "pro" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -73288,6 +73588,15 @@ "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, "output_cost_per_token_ultrafast": 6e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -76787,6 +77096,15 @@ "output_cost_per_token_batches": 2.5e-06, "output_cost_per_token_flex": 2.5e-06, "output_cost_per_token_priority": 1e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 02d07932533..cf993c69814 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1,8 +1,10 @@ import base64 +import os import re from collections.abc import Callable, Collection, Iterable, Mapping, Sequence from functools import reduce -from typing import Any, Final, Optional, TypeVar, Union, cast, get_type_hints, overload +from types import MappingProxyType +from typing import Any, Final, Optional, TypeVar, Union, cast, get_args, get_type_hints, overload from pydantic import BaseModel, TypeAdapter from typing_extensions import TypedDict as ExtraItemsTypedDict @@ -14,6 +16,7 @@ from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value, is_nested_path from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig, ReasoningSummaryLevel +from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.types.llms.openai import ( AllMessageValues, InputTokensDetails, @@ -29,6 +32,7 @@ from litellm.types.responses.main import DecodedResponseId from litellm.types.utils import ( CacheCreationTokenDetails, CompletionTokensDetailsWrapper, + ModelInfo, PromptTokensDetailsWrapper, ProviderSpecificHeader, SpecialEnums, @@ -1382,3 +1386,98 @@ def auto_reasoning_summary(model: str, custom_llm_provider: object) -> Reasoning provider, model.removeprefix(f"{provider}/") ) return "detailed" if responses_config is None else responses_config.auto_reasoning_summary() + + +def opt_in_auto_reasoning_summary(model: str | None, custom_llm_provider: object) -> ReasoningSummaryLevel | None: + if not (litellm.reasoning_auto_summary or os.getenv("LITELLM_REASONING_AUTO_SUMMARY", "false").lower() == "true"): + return None + return "detailed" if model is None else auto_reasoning_summary(model, custom_llm_provider) + + +def reasoning_with_summary_default( + reasoning: Mapping[str, object], summary_alias: object, model: str, custom_llm_provider: object +) -> Mapping[str, object]: + if "summary" in reasoning: + return {**reasoning} + summary: Final = ( + summary_alias if summary_alias is not None else opt_in_auto_reasoning_summary(model, custom_llm_provider) + ) + return {**reasoning} if summary is None else {**reasoning, "summary": summary} + + +_SUMMARY_VALUES: Final = frozenset((*get_args(ReasoningSummaryLevel), None)) + + +def _openai_model_info(model: str) -> ModelInfo | None: + from litellm.utils import ModelInfoLookupMiss, lookup_model_info_for_registration + + model_info: Final = lookup_model_info_for_registration(model.removeprefix("responses/"), "openai") + return None if isinstance(model_info, ModelInfoLookupMiss) else model_info + + +def _declared_values(values: Sequence[str] | None) -> frozenset[str | None] | None: + return None if values is None else frozenset(values) + + +def _openai_reasoning_field_values(model: str) -> Mapping[str, frozenset[str | None] | None]: + model_info: Final = _openai_model_info(model) + if model_info is None: + return MappingProxyType({"summary": _SUMMARY_VALUES, "mode": None, "context": None}) + if not model_info.get("supports_reasoning"): + return MappingProxyType({}) + return MappingProxyType( + { + "summary": _SUMMARY_VALUES, + "mode": _declared_values(model_info.get("reasoning_modes")), + "context": _declared_values(model_info.get("reasoning_contexts")), + } + ) + + +def _openai_accepts_reasoning_field( + key: str, value: object, values_by_field: Mapping[str, frozenset[str | None] | None] +) -> bool: + if key == "effort": + return True + if key not in values_by_field: + return False + allowed: Final = values_by_field[key] + return allowed is None or (isinstance(value, str | None) and value in allowed) + + +def _openai_accepted_reasoning(reasoning: Mapping[str, object], model: str) -> Mapping[str, object]: + values_by_field: Final = _openai_reasoning_field_values(model) + return { + key: value for key, value in reasoning.items() if _openai_accepts_reasoning_field(key, value, values_by_field) + } + + +def _additional_drop_paths(additional_drop_params: object) -> tuple[str, ...]: + if additional_drop_params.__class__ is not list: + return () + return tuple(path for path in iter_items(additional_drop_params) if isinstance(path, str)) + + +def openai_bridge_reasoning( + *, + requested: Mapping[str, object], + effort_kept: bool, + summary_alias: object, + model: str, + drop_params: bool, + additional_drop_params: object, +) -> Mapping[str, object] | None: + drop_paths: Final = _additional_drop_paths(additional_drop_params) + if "reasoning_effort" in drop_paths: + return None + kept: Final = requested if effort_kept else {key: value for key, value in requested.items() if key != "effort"} + with_summary: Final = reasoning_with_summary_default(kept, summary_alias, model, "openai") + accepted: Final = ( + _openai_accepted_reasoning(with_summary, model) + if drop_params and OpenAIGPTConfig.is_openai_catalog_model(model) + else {**with_summary} + ) + remaining: Final = as_str_mapping( + _apply_nested_drop_params({"reasoning_effort": accepted}, list(drop_paths)).get("reasoning_effort") + ) + return remaining or None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 055c18f2e36..4ce35bf86c1 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -225,6 +225,8 @@ class ProviderSpecificModelInfo(TypedDict, total=False): supports_xhigh_reasoning_effort: bool | None supports_max_reasoning_effort: bool | None reasoning_effort_levels: ReadOnly[Sequence[str] | None] + reasoning_modes: ReadOnly[Sequence[str] | None] + reasoning_contexts: ReadOnly[Sequence[str] | None] default_reasoning_effort: ReadOnly[Literal["none", "minimal", "low", "medium", "high", "xhigh"] | None] supports_output_config: bool | None supports_image_size: bool | None diff --git a/litellm/utils.py b/litellm/utils.py index 61b0e17e197..a231644addf 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7000,6 +7000,8 @@ def get_model_info_helper( supports_xhigh_reasoning_effort=_model_info.get("supports_xhigh_reasoning_effort", None), supports_max_reasoning_effort=_model_info.get("supports_max_reasoning_effort", None), reasoning_effort_levels=_model_info.get("reasoning_effort_levels", None), + reasoning_modes=_model_info.get("reasoning_modes", None), + reasoning_contexts=_model_info.get("reasoning_contexts", None), default_reasoning_effort=_model_info.get("default_reasoning_effort", None), bedrock_output_config_effort_ceiling=_model_info.get("bedrock_output_config_effort_ceiling", None), bedrock_converse_supports_strict_tools=_model_info.get("bedrock_converse_supports_strict_tools", None), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 7199a075612..825de5ac3ee 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -35598,6 +35598,13 @@ "output_cost_per_token_batches": 5e-06, "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -35675,6 +35682,13 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_priority": 2e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -35851,6 +35865,13 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_priority": 2e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -35902,6 +35923,13 @@ "mode": "chat", "output_cost_per_token": 1.4e-05, "output_cost_per_token_priority": 2.8e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -36017,6 +36045,13 @@ "mode": "chat", "output_cost_per_token": 1.4e-05, "output_cost_per_token_priority": 2.8e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -36066,6 +36101,13 @@ "mode": "responses", "output_cost_per_token": 0.000168, "output_cost_per_token_batches": 8.4e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "pro" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -36107,6 +36149,13 @@ "mode": "responses", "output_cost_per_token": 0.000168, "output_cost_per_token_batches": 8.4e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "pro" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -36194,6 +36243,15 @@ "output_cost_per_token_flex": 2.5e-05, "output_cost_per_token_priority": 0.0001, "output_cost_per_token_ultrafast": 0.0003, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36278,6 +36336,15 @@ "output_cost_per_token_above_272k_tokens_batches": 7.5e-06, "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36370,6 +36437,15 @@ "output_cost_per_token_above_272k_tokens_batches": 3.75e-07, "output_cost_per_token_flex": 2.5e-07, "output_cost_per_token_priority": 1e-06, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36453,6 +36529,15 @@ "output_cost_per_token_above_272k_tokens_batches": 1.5e-05, "output_cost_per_token_flex": 1e-05, "output_cost_per_token_priority": 4e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36534,6 +36619,15 @@ "output_cost_per_token_above_272k_tokens_batches": 1.5e-05, "output_cost_per_token_flex": 1e-05, "output_cost_per_token_priority": 4e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36617,6 +36711,15 @@ "output_cost_per_token_above_272k_tokens_batches": 9e-06, "output_cost_per_token_flex": 6e-06, "output_cost_per_token_priority": 2.4e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36699,6 +36802,15 @@ "output_cost_per_token_above_272k_tokens_batches": 9e-07, "output_cost_per_token_flex": 6e-07, "output_cost_per_token_priority": 2.4e-06, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36936,6 +37048,14 @@ "reasoning_effort_levels": [ "medium" ], + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -36983,6 +37103,14 @@ "output_cost_per_token_batches": 1.5e-05, "output_cost_per_token_above_272k_tokens_batches": 2.25e-05, "output_cost_per_token_priority": 7.5e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37045,6 +37173,14 @@ "output_cost_per_token_batches": 1.5e-05, "output_cost_per_token_above_272k_tokens_batches": 2.25e-05, "output_cost_per_token_priority": 7.5e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37099,6 +37235,14 @@ "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, "output_cost_per_token_above_272k_tokens_batches": 0.000135, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37150,6 +37294,14 @@ "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, "output_cost_per_token_above_272k_tokens_batches": 0.000135, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37209,6 +37361,14 @@ "output_cost_per_token_batches": 7.5e-06, "output_cost_per_token_above_272k_tokens_batches": 1.125e-05, "output_cost_per_token_priority": 3e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ @@ -37266,6 +37426,14 @@ "output_cost_per_token_batches": 7.5e-06, "output_cost_per_token_above_272k_tokens_batches": 1.125e-05, "output_cost_per_token_priority": 3e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ @@ -37315,6 +37483,14 @@ "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, "output_cost_per_token_above_272k_tokens_batches": 0.000135, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37367,6 +37543,14 @@ "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, "output_cost_per_token_above_272k_tokens_batches": 0.000135, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37421,6 +37605,14 @@ "output_cost_per_token_flex": 2.25e-06, "output_cost_per_token_batches": 2.25e-06, "output_cost_per_token_priority": 9e-06, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37475,6 +37667,14 @@ "output_cost_per_token_flex": 2.25e-06, "output_cost_per_token_batches": 2.25e-06, "output_cost_per_token_priority": 9e-06, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37526,6 +37726,14 @@ "output_cost_per_token": 1.25e-06, "output_cost_per_token_flex": 6.25e-07, "output_cost_per_token_batches": 6.25e-07, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37578,6 +37786,14 @@ "output_cost_per_token": 1.25e-06, "output_cost_per_token_flex": 6.25e-07, "output_cost_per_token_batches": 6.25e-07, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37624,6 +37840,13 @@ "mode": "responses", "output_cost_per_token": 0.00012, "output_cost_per_token_batches": 6e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "pro" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -37666,6 +37889,13 @@ "mode": "responses", "output_cost_per_token": 0.00012, "output_cost_per_token_batches": 6e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "pro" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -37716,6 +37946,13 @@ "output_cost_per_token_batches": 5e-06, "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -37826,6 +38063,13 @@ "mode": "responses", "output_cost_per_token": 1.4e-05, "output_cost_per_token_priority": 2.8e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -37908,6 +38152,13 @@ "output_cost_per_token_batches": 1e-06, "output_cost_per_token_flex": 1e-06, "output_cost_per_token_priority": 3.6e-06, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -37959,6 +38210,13 @@ "output_cost_per_token_batches": 1e-06, "output_cost_per_token_flex": 1e-06, "output_cost_per_token_priority": 3.6e-06, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -38007,6 +38265,13 @@ "output_cost_per_token": 4e-07, "output_cost_per_token_batches": 2e-07, "output_cost_per_token_flex": 2e-07, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -38055,6 +38320,13 @@ "output_cost_per_token": 4e-07, "output_cost_per_token_batches": 2e-07, "output_cost_per_token_flex": 2e-07, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -43821,6 +44093,13 @@ "output_cost_per_token_batches": 4e-06, "output_cost_per_token_flex": 4e-06, "output_cost_per_token_priority": 1.4e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -43929,6 +44208,13 @@ "output_cost_per_token_batches": 4e-06, "output_cost_per_token_flex": 4e-06, "output_cost_per_token_priority": 1.4e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -44008,6 +44294,13 @@ "mode": "responses", "output_cost_per_token": 8e-05, "output_cost_per_token_batches": 4e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "pro" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -44045,6 +44338,13 @@ "mode": "responses", "output_cost_per_token": 8e-05, "output_cost_per_token_batches": 4e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "pro" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -73288,6 +73588,15 @@ "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, "output_cost_per_token_ultrafast": 6e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -76787,6 +77096,15 @@ "output_cost_per_token_batches": 2.5e-06, "output_cost_per_token_flex": 2.5e-06, "output_cost_per_token_priority": 1e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index ba39fc3f5ac..f38bbcf4404 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -2127,6 +2127,330 @@ def test_responses_bridge_preserves_reasoning_effort_with_drop_params( assert request_body["reasoning"] == {"effort": "high"} +@pytest.fixture +def openai_responses_route( + restore_model_registry, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> respx.Route: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("LITELLM_REASONING_AUTO_SUMMARY", raising=False) + litellm.register_model( + { + "gpt-6.1-terra": { + "litellm_provider": "openai", + "mode": "chat", + "supports_reasoning": True, + "supports_minimal_reasoning_effort": False, + "reasoning_modes": ["standard", "pro"], + "reasoning_contexts": ["auto", "current_turn", "all_turns"], + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + } + }, + persist_across_reloads=False, + ) + return respx_mock.post("https://api.openai.com/v1/responses").respond( + json={ + "id": "resp_test", + "object": "response", + "created_at": 1734366691, + "status": "completed", + "model": "gpt-6.1-terra", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Done.", "annotations": []}], + } + ], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + } + ) + + +@pytest.mark.parametrize( + ("reasoning_effort", "summary_alias", "expected_reasoning"), + [ + ( + {"effort": "low", "mode": "pro", "context": "current_turn"}, + {}, + {"effort": "low", "mode": "pro", "context": "current_turn"}, + ), + ( + {"effort": "low", "mode": "pro", "context": "current_turn"}, + {"reasoning_summary": "auto"}, + {"effort": "low", "mode": "pro", "context": "current_turn", "summary": "auto"}, + ), + ( + {"effort": "low", "mode": "pro", "summary": "concise"}, + {"reasoningSummary": "auto"}, + {"effort": "low", "mode": "pro", "summary": "concise"}, + ), + ( + {"effort": "low", "context": "current_turn", "summary": None}, + {"reasoning_summary": "auto"}, + {"effort": "low", "context": "current_turn", "summary": None}, + ), + ({"mode": "pro", "context": "all_turns"}, {}, {"mode": "pro", "context": "all_turns"}), + ("low", {}, {"effort": "low"}), + ], +) +def test_openai_responses_bridge_sends_the_whole_reasoning_effort_dict( + reasoning_effort: str | dict[str, str | None], + summary_alias: dict[str, str], + expected_reasoning: dict[str, str | None], + openai_responses_route: respx.Route, +): + caller_reasoning_effort: Final = copy.deepcopy(reasoning_effort) + + litellm.completion( + model="openai/responses/gpt-6.1-terra", + messages=[{"role": "user", "content": "hello"}], + reasoning_effort=reasoning_effort, + api_key="fake-key", + **summary_alias, + ) + + request_body: Final = json.loads(openai_responses_route.calls[0].request.content) + assert request_body["reasoning"] == expected_reasoning + assert reasoning_effort == caller_reasoning_effort + + +@pytest.mark.parametrize( + "reasoning_effort, summary_alias, expected_reasoning", + [ + ( + {"effort": "minimal", "mode": "pro", "context": "current_turn"}, + {}, + {"mode": "pro", "context": "current_turn"}, + ), + ({"effort": "minimal", "mode": "pro", "summary": "auto"}, {}, {"mode": "pro", "summary": "auto"}), + ({"effort": "minimal", "mode": "pro"}, {"reasoning_summary": "auto"}, {"mode": "pro", "summary": "auto"}), + ], +) +def test_openai_responses_bridge_drops_only_a_refused_effort_under_drop_params( + reasoning_effort: dict[str, str], + summary_alias: dict[str, str], + expected_reasoning: dict[str, str], + openai_responses_route: respx.Route, +): + litellm.completion( + model="openai/responses/gpt-6.1-terra", + messages=[{"role": "user", "content": "hello"}], + reasoning_effort=reasoning_effort, + drop_params=True, + api_key="fake-key", + **summary_alias, + ) + + request_body: Final = json.loads(openai_responses_route.calls[0].request.content) + assert request_body["reasoning"] == expected_reasoning + + +def test_openai_responses_bridge_drops_a_reasoning_effort_dict_for_a_model_without_reasoning( + openai_responses_route: respx.Route, +): + litellm.completion( + model="openai/responses/gpt-4.1", + messages=[{"role": "user", "content": "hello"}], + reasoning_effort={"effort": "low", "mode": "pro", "context": "current_turn"}, + drop_params=True, + api_key="fake-key", + ) + + request_body: Final = json.loads(openai_responses_route.calls[0].request.content) + assert "reasoning" not in request_body + + +@pytest.mark.parametrize( + "drop_params, expected_reasoning", + [ + (True, {"effort": "low", "context": "current_turn"}), + ( + False, + { + "effort": "low", + "exclude": False, + "mode": "turbo", + "context": "current_turn", + "generate_summary": "concise", + }, + ), + ], +) +def test_openai_responses_bridge_drops_reasoning_fields_openai_refuses_only_under_drop_params( + drop_params: bool, + expected_reasoning: dict[str, object], + openai_responses_route: respx.Route, +): + litellm.completion( + model="openai/responses/gpt-6.1-terra", + messages=[{"role": "user", "content": "hello"}], + reasoning_effort={ + "effort": "low", + "exclude": False, + "mode": "turbo", + "context": "current_turn", + "generate_summary": "concise", + }, + drop_params=drop_params, + api_key="fake-key", + ) + + request_body: Final = json.loads(openai_responses_route.calls[0].request.content) + assert request_body["reasoning"] == expected_reasoning + + +@pytest.mark.parametrize( + "declared, drop_params, expected_reasoning", + [ + ( + {"reasoning_modes": ["standard"], "reasoning_contexts": ["auto", "current_turn"]}, + True, + {"effort": "low"}, + ), + ( + {"reasoning_modes": ["pro"], "reasoning_contexts": ["all_turns"]}, + True, + {"effort": "low", "mode": "pro", "context": "all_turns"}, + ), + ({}, True, {"effort": "low", "mode": "pro", "context": "all_turns"}), + ({"reasoning_modes": ["standard"]}, True, {"effort": "low", "context": "all_turns"}), + ( + {"reasoning_modes": ["standard"], "reasoning_contexts": ["auto", "current_turn"]}, + False, + {"effort": "low", "mode": "pro", "context": "all_turns"}, + ), + ], +) +def test_openai_responses_bridge_drops_a_mode_or_context_the_model_does_not_list_only_under_drop_params( + declared: dict[str, list[str]], + drop_params: bool, + expected_reasoning: dict[str, str], + openai_responses_route: respx.Route, +): + litellm.register_model( + { + "gpt-6.1-reasoning-test": { + "litellm_provider": "openai", + "mode": "chat", + "supports_reasoning": True, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + **declared, + } + }, + persist_across_reloads=False, + ) + + litellm.completion( + model="openai/responses/gpt-6.1-reasoning-test", + messages=[{"role": "user", "content": "hello"}], + reasoning_effort={"effort": "low", "mode": "pro", "context": "all_turns"}, + drop_params=drop_params, + api_key="fake-key", + ) + + request_body: Final = json.loads(openai_responses_route.calls[0].request.content) + assert request_body["reasoning"] == expected_reasoning + + +def test_openai_responses_bridge_forwards_mode_and_context_under_drop_params_for_a_catalog_model_with_no_cost_entry( + openai_responses_route: respx.Route, +): + litellm.open_ai_chat_completion_models.add("unpriced-reasoning-model") + + litellm.completion( + model="openai/responses/unpriced-reasoning-model", + messages=[{"role": "user", "content": "hello"}], + reasoning_effort={"effort": "low", "mode": "pro", "context": "all_turns", "generate_summary": "auto"}, + drop_params=True, + api_key="fake-key", + ) + + request_body: Final = json.loads(openai_responses_route.calls[0].request.content) + assert request_body["reasoning"] == {"effort": "low", "mode": "pro", "context": "all_turns"} + + +def test_openai_responses_bridge_forwards_a_reasoning_effort_dict_whole_for_a_model_outside_the_openai_catalog( + openai_responses_route: respx.Route, +): + reasoning_effort: Final = {"effort": "high", "exclude": True, "budget_tokens": 2000} + + litellm.completion( + model="openai/responses/my-proxy-alias", + messages=[{"role": "user", "content": "hello"}], + reasoning_effort=reasoning_effort, + drop_params=True, + api_key="fake-key", + ) + + request_body: Final = json.loads(openai_responses_route.calls[0].request.content) + assert request_body["reasoning"] == reasoning_effort + + +@pytest.mark.parametrize( + "additional_drop_params, summary_alias, expected_reasoning", + [ + (["reasoning_effort.mode"], {}, {"effort": "low", "context": "current_turn"}), + ( + ["reasoning_effort.summary"], + {"reasoning_summary": "auto"}, + {"effort": "low", "mode": "pro", "context": "current_turn"}, + ), + (["reasoning_effort"], {"reasoning_summary": "auto"}, None), + ("reasoning_effort", {}, {"effort": "low", "mode": "pro", "context": "current_turn"}), + ], +) +def test_openai_responses_bridge_applies_additional_drop_params_to_a_reasoning_effort_dict( + additional_drop_params: list[str] | str, + summary_alias: dict[str, str], + expected_reasoning: dict[str, str] | None, + openai_responses_route: respx.Route, +): + litellm.completion( + model="openai/responses/gpt-6.1-terra", + messages=[{"role": "user", "content": "hello"}], + reasoning_effort={"effort": "low", "mode": "pro", "context": "current_turn"}, + additional_drop_params=additional_drop_params, + api_key="fake-key", + **summary_alias, + ) + + request_body: Final = json.loads(openai_responses_route.calls[0].request.content) + assert request_body.get("reasoning") == expected_reasoning + + +def test_openai_responses_bridge_gives_a_reasoning_effort_dict_the_auto_summary_a_string_gets( + openai_responses_route: respx.Route, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setenv("LITELLM_REASONING_AUTO_SUMMARY", "true") + litellm.completion( + model="openai/responses/gpt-6.1-terra", + messages=[{"role": "user", "content": "hello"}], + reasoning_effort="low", + api_key="fake-key", + ) + litellm.completion( + model="openai/responses/gpt-6.1-terra", + messages=[{"role": "user", "content": "hello"}], + reasoning_effort={"effort": "low", "mode": "pro", "context": "current_turn"}, + api_key="fake-key", + ) + + string_reasoning, dict_reasoning = ( + json.loads(call.request.content)["reasoning"] for call in openai_responses_route.calls + ) + assert string_reasoning["summary"] is not None + assert dict_reasoning == { + "effort": "low", + "mode": "pro", + "context": "current_turn", + "summary": string_reasoning["summary"], + } + + _BRIDGED_RESPONSE_ANSWER: Final = json.dumps( { "id": "resp_slow", diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 38287b25169..146e2df6413 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -1195,6 +1195,8 @@ _COST_MAP_ROW_SCHEMA: Final = { "type": "array", "items": {"type": "string", "enum": ["none", "minimal", "low", "medium", "high", "xhigh", "max"]}, }, + "reasoning_modes": {"type": "array", "items": {"type": "string"}}, + "reasoning_contexts": {"type": "array", "items": {"type": "string"}}, "default_reasoning_effort": { "type": "string", "enum": ["none", "minimal", "low", "medium", "high", "xhigh"], From ca6c83022193c563b42c347d7e54ed8f66edc009 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 28 Oct 2026 15:39:41 -0400 Subject: [PATCH 078/128] fix(proxy): account for media in combined TPM reservations (#48865) --- .../litellm_core_utils/audio_utils/utils.py | 43 +++ .../prompt_templates/common_utils.py | 155 +++++++++ litellm/litellm_core_utils/token_counter.py | 232 +++++++++++++- .../hooks/parallel_request_limiter_v3.py | 91 +++++- litellm/router.py | 38 ++- ...ore_utils_prompt_templates_common_utils.py | 116 +++++++ .../litellm_core_utils/test_audio_utils.py | 64 ++++ .../litellm_core_utils/test_token_counter.py | 131 ++++++++ tests/unit/proxy/hooks/test_tpm_concurrent.py | 299 +++++++++++++++++- tests/unit/test_router/test_router.py | 31 ++ 10 files changed, 1169 insertions(+), 31 deletions(-) diff --git a/litellm/litellm_core_utils/audio_utils/utils.py b/litellm/litellm_core_utils/audio_utils/utils.py index dab3e48f91a..640559ac7e2 100644 --- a/litellm/litellm_core_utils/audio_utils/utils.py +++ b/litellm/litellm_core_utils/audio_utils/utils.py @@ -2,8 +2,12 @@ Utils used for litellm.transcription() and litellm.atranscription() """ +import base64 +import binascii import hashlib +import math import os +from collections.abc import Sequence from dataclasses import dataclass from typing import Final @@ -331,6 +335,45 @@ def calculate_request_duration(file: FileTypes) -> float | None: return None +DEFAULT_AUDIO_TOKEN_ESTIMATE: Final = 300 +AUDIO_TOKENS_PER_SECOND: Final = 10 +AUDIO_BYTES_PER_TOKEN: Final = 1600 +_MAX_AUDIO_ESTIMATE_BYTES: Final = 16 * 1024 * 1024 + + +def _audio_token_rate(model: str) -> int | None: + if model.startswith( + ("gemini-", "gemini/gemini-", "vertex_ai/gemini-", "vertex_ai_beta/gemini-", "openrouter/google/gemini-") + ): + family: Final = model.rsplit("/", 1)[-1] + return 25 if family.startswith(("gemini-3-", "gemini-3.")) else 32 + if model.startswith(("gpt-", "openai/gpt-", "azure/gpt-", "openrouter/openai/gpt-")): + return AUDIO_TOKENS_PER_SECOND + return None + + +def estimate_audio_tokens(data: str | None, models: Sequence[str] = (), fallback_model: str | None = None) -> int: + """Estimate audio tokens from container duration within a bounded decoding budget.""" + if not data: + return DEFAULT_AUDIO_TOKEN_ESTIMATE + fallback_estimate: Final = max(len(data) * 3 // 4 // AUDIO_BYTES_PER_TOKEN, DEFAULT_AUDIO_TOKEN_ESTIMATE) + padding: Final = 2 if data.endswith("==") else 1 if data.endswith("=") else 0 + if len(data) * 3 // 4 - padding > _MAX_AUDIO_ESTIMATE_BYTES: + return fallback_estimate + try: + duration: Final = calculate_request_duration(base64.b64decode(data, validate=True)) + except (ValueError, binascii.Error): + return fallback_estimate + if duration is None or not math.isfinite(duration) or duration <= 0: + return fallback_estimate + fallback_rate: Final = _audio_token_rate(fallback_model) if fallback_model is not None else None + rate: Final = max( + (_audio_token_rate(model) or fallback_rate or AUDIO_TOKENS_PER_SECOND for model in models), + default=fallback_rate or AUDIO_TOKENS_PER_SECOND, + ) + return max(1, math.ceil(duration * rate)) + + DEFAULT_SPEECH_MEDIA_TYPE: Final = "audio/mpeg" diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 3c06c522dd6..cf6bdc69556 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -9,6 +9,7 @@ import json import mimetypes import re from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence +from dataclasses import dataclass from graphlib import CycleError, TopologicalSorter from itertools import chain, groupby, islice from os import PathLike @@ -236,6 +237,160 @@ _IMAGE_CONTENT_PART_TYPES: Final = frozenset({"image_url", "input_image", "image _IMAGE_SCAN_MAX_DEPTH: Final = 4 +@dataclass(frozen=True, slots=True) +class _TPMMediaProjection: + content: object + media: tuple[Mapping[str, object], ...] = () + + +def _tpm_image_part(url: str, detail: object = "auto") -> Mapping[str, object]: + return {"type": "image_url", "image_url": {"url": url, "detail": detail}} + + +def _tpm_blank_fields(fields: Mapping[str, object], names: tuple[str, ...]) -> Mapping[str, object]: + return {key: "" if key in names and isinstance(value, str) else value for key, value in fields.items()} + + +def _tpm_file_media(part: Mapping[str, object]) -> _TPMMediaProjection: + from litellm.litellm_core_utils.prompt_templates.factory import parse_mime_type + + nested: Final = part.get("type") == "file" + fields: Final = as_str_mapping(part.get("file")) if nested else part + if fields is None: + return _TPMMediaProjection(part) + raw_data: Final = fields.get("file_data") + data: Final = raw_data if isinstance(raw_data, str) else "" + declared: Final = next( + ( + value + for key in ("format", "mime_type", "content_type") + if isinstance(value := fields.get(key), str) and value + ), + "", + ) + mime: Final = (declared or parse_mime_type(data[:256]) or "").split(";", 1)[0].strip().lower() + if not mime.startswith(("image/", "audio/")): + return _TPMMediaProjection(part) + raw_url: Final = fields.get("file_url") or fields.get("file_id") + url: Final = data or (raw_url if isinstance(raw_url, str) else "") + header_end: Final = data.find(",", 0, 256) + encoded: Final = ( + data[header_end + 1 :] + if header_end >= 0 and data[:header_end].endswith(";base64") + else "" + if data.startswith(("data:", "https://", "http://", "gs://")) + else data + ) + media: Final = ( + _tpm_image_part( + url + if not data or url.startswith(("data:", "https://", "http://", "gs://")) + else f"data:{mime};base64,{url}", + fields.get("detail", "auto"), + ) + if mime.startswith("image/") + else {"type": "input_audio", "input_audio": {"data": encoded, "format": mime.removeprefix("audio/")}} + ) + projected: Final = _tpm_blank_fields(fields, ("file_data", "file_url")) + return _TPMMediaProjection({**part, "file": projected} if nested else projected, (media,)) + + +def _tpm_media_leaf(value: object) -> _TPMMediaProjection: + part: Final = as_str_mapping(value) + if part is None: + return _TPMMediaProjection(value) + kind: Final = part.get("type") + if kind in ("file", "input_file"): + return _tpm_file_media(part) + if kind == "input_audio": + audio: Final = as_str_mapping(part.get("input_audio")) + return _TPMMediaProjection( + {**part, "input_audio": _tpm_blank_fields(audio, ("data",))} if audio is not None else value, + ({"type": "input_audio", "input_audio": part.get("input_audio")},), + ) + if kind in ("image_url", "input_image", "image", "computer_screenshot"): + image: Final = as_str_mapping(part.get("image_url")) + source: Final = as_str_mapping(part.get("source")) + canonical: Final = _tpm_image_part( + _image_part_url({**part, "type": "input_image"} if kind == "computer_screenshot" else part) or "", + (image or part).get("detail", part.get("detail", "auto")), + ) + projected: Final = ( + {**part, "image_url": _tpm_blank_fields(image, ("url",))} + if image is not None and kind in ("image_url", "input_image", "computer_screenshot") + else {**part, "source": _tpm_blank_fields(source, ("data", "url"))} + if source is not None and kind == "image" and source.get("type") in ("base64", "url") + else _tpm_blank_fields(part, ("image_url", "url")) + ) + return _TPMMediaProjection(projected, (canonical,)) + return _TPMMediaProjection(value) + + +def _tpm_media_list(value: object, split: Callable[[object], _TPMMediaProjection]) -> _TPMMediaProjection: + if not isinstance(value, (list, tuple)): + return _TPMMediaProjection(value) + parts: Final = cast(Sequence[object], value) # cast-ok: finite JSON content arrays + projected: Final = tuple(map(split, parts)) + media: Final = tuple(chain.from_iterable(part.media for part in projected)) + content: Final = tuple(part.content for part in projected) + return _TPMMediaProjection(parts if not media else content if isinstance(value, tuple) else list(content), media) + + +def _tpm_replace_field(part: Mapping[str, object], key: str, projection: _TPMMediaProjection) -> _TPMMediaProjection: + return _TPMMediaProjection({**part, key: projection.content} if projection.media else part, projection.media) + + +def _tpm_document_media(value: object) -> _TPMMediaProjection: + part: Final = as_str_mapping(value) + source: Final = as_str_mapping(part.get("source")) if part is not None else None + if part is None or part.get("type") != "document" or source is None or source.get("type") != "content": + return _tpm_media_leaf(value) + content: Final = _tpm_media_list(source.get("content"), _tpm_media_leaf) + return _tpm_replace_field(part, "source", _tpm_replace_field(source, "content", content)) + + +def _tpm_tool_result_media(value: object) -> _TPMMediaProjection: + part: Final = as_str_mapping(value) + if part is None or part.get("type") != "tool_result": + return _tpm_document_media(value) + return _tpm_replace_field(part, "content", _tpm_media_list(part.get("content"), _tpm_document_media)) + + +def _tpm_code_output_media(value: object) -> _TPMMediaProjection: + part: Final = as_str_mapping(value) + url: Final = part.get("url") if part is not None else None + if part is None or part.get("type") != "image" or not isinstance(url, str) or not url: + return _TPMMediaProjection(value) + return _TPMMediaProjection(_tpm_blank_fields(part, ("url",)), (_tpm_image_part(url),)) + + +def _tpm_input_media(value: object) -> _TPMMediaProjection: + part: Final = as_str_mapping(value) + if part is None: + return _TPMMediaProjection(value) + kind: Final = part.get("type") + if kind in ("function_call_output", "custom_tool_call_output"): + return _tpm_replace_field(part, "output", _tpm_media_list(part.get("output"), _tpm_tool_result_media)) + if kind == "computer_call_output": + return _tpm_replace_field(part, "output", _tpm_media_leaf(part.get("output"))) + if kind == "code_interpreter_call": + return _tpm_replace_field(part, "outputs", _tpm_media_list(part.get("outputs"), _tpm_code_output_media)) + result: Final = part.get("result") + if kind == "image_generation_call" and isinstance(result, str) and result: + return _TPMMediaProjection( + _tpm_blank_fields(part, ("result",)), + (_tpm_image_part("data:application/octet-stream;base64," + result),), + ) + if kind == "message" or (kind is None and isinstance(part.get("role"), str)): + return _tpm_replace_field(part, "content", _tpm_media_list(part.get("content"), _tpm_tool_result_media)) + return _tpm_tool_result_media(value) + + +def split_media_for_tpm(value: object) -> tuple[object, tuple[Mapping[str, object], ...]]: + projection: Final = _tpm_media_list(value, _tpm_input_media) + return projection.content, projection.media + + def _content_parts_contain_image(parts: Sequence[object]) -> bool: """Depth-bounded frontier walk over nested content lists, iterative because the repo bans recursion; an Anthropic tool_result nests its image parts exactly one level down.""" diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 6d50e4b065e..f42e4eae08f 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -3,9 +3,12 @@ import base64 import io import math +import re import struct +from bisect import bisect_right from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from itertools import accumulate +from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, cast import anyio @@ -35,6 +38,7 @@ from litellm.constants import ( TOKEN_COUNTER_MAX_EXACT_CHARS, ) from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.prompt_templates.chat_template_tools import ( ChatTemplateToolSection, chat_template_tool_section, @@ -188,22 +192,23 @@ def resize_image_high_res( # Test the function with the given example def calculate_tiles_needed( - resized_width, - resized_height, - tile_width=MAX_TILE_WIDTH, - tile_height=MAX_TILE_HEIGHT, -): + resized_width: int, + resized_height: int, + tile_width: int = MAX_TILE_WIDTH, + tile_height: int = MAX_TILE_HEIGHT, +) -> int: tiles_across: Final = (resized_width + tile_width - 1) // tile_width tiles_down: Final = (resized_height + tile_height - 1) // tile_height total_tiles: Final = tiles_across * tiles_down return total_tiles -def high_detail_image_token_upper_bound(base_tokens: int = 85) -> int: +def high_detail_image_token_upper_bound(base_tokens: int = 85, tile_tokens: int | None = None) -> int: largest_tile_count: Final = calculate_tiles_needed( MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES, MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES ) - return base_tokens + (base_tokens * 2) * largest_tile_count + per_tile: Final = base_tokens * 2 if tile_tokens is None else tile_tokens + return base_tokens + per_tile * largest_tile_count def _unpack_ints(fmt: str, buffer: bytes) -> tuple[int, ...]: @@ -381,6 +386,219 @@ def calculate_img_tokens( return total_tokens +_IMAGE_HEADER_BASE64_CHARS: Final = 65_536 +_MAX_IMAGE_PATCHES: Final = 30_000 +_ANTHROPIC_IMAGE_LIMITS: Final = MappingProxyType( + {"claude-standard": (ANTHROPIC_IMAGE_MAX_LONG_EDGE_PX, 1568), "claude-high": (2576, 4784)} +) +_TILE_IMAGE_TOKEN_COSTS: Final = MappingProxyType( + { + "gpt-4o": (85, 170), + "gpt-4.1": (85, 170), + "gpt-4o-mini": (2833, 5667), + "gpt-5": (70, 140), + "gpt-5.1": (70, 140), + "o1": (75, 150), + "o1-pro": (75, 150), + "o3": (75, 150), + } +) +_PATCH_IMAGE_MULTIPLIERS: Final = MappingProxyType( + { + "gpt-6-astra": 1.2, + "gpt-5.6-sol": 1.2, + "gpt-5.6-terra": 1.2, + "gpt-5.6-luna": 1.2, + "gpt-5.5": 1.2, + "gpt-5.4": 1.2, + "gpt-5.4-mini": 1.2, + "gpt-5.4-nano": 1.2, + "gpt-5.2": 1.2, + "gpt-4.1-mini": 1.62, + } +) + + +def _inline_image_dimensions(data: str) -> tuple[int, int] | None: + separator: Final = data.find(",", 0, 256) + if not data.startswith("data:") or separator < 0 or not data[:separator].endswith(";base64"): + return None + encoded_header: Final = data[separator + 1 : separator + 1 + _IMAGE_HEADER_BASE64_CHARS] + try: + header: Final = base64.b64decode(encoded_header, validate=True) + image_type: Final = get_image_type(header) + if image_type not in ("png", "gif", "jpeg", "webp"): + return None + if image_type == "webp" and header[12:16] not in (b"VP8X", b"VP8 ", b"VP8L"): + return None + dimensions: Final = _image_dimensions_from_bytes(header) + except (ValueError, TypeError, struct.error): + return None + width, height = dimensions + return (width, height) if width > 0 and height > 0 else None + + +def _patch_image_limits(model: str, detail: str) -> tuple[int, int | None] | None: + if model in ("gpt-5.2", "gpt-4.1-mini"): + return None if detail == "original" else (2048, 6144) + original_default: Final = model in ("gpt-6-astra", "gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna", "gpt-5.5") + effective_detail: Final = ("original" if original_default else "high") if detail == "auto" else detail + if effective_detail == "low": + return (2048, 6144) if model in ("gpt-5.4", "gpt-5.4-mini", "gpt-5.4-nano") else (512, None) + if effective_detail == "high": + return (65_535 if model == "gpt-6-astra" else 2048, 2500) + return ( + (65_535, None) if model in ("gpt-6-astra", "gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna") else (6000, 10_000) + ) + + +def _image_patch_count(width: int, height: int, max_dimension: int, budget: int | None) -> int: + dimension_scale: Final = min(1.0, max_dimension / max(width, height)) + fitted_width: Final = max(1, round(width * dimension_scale)) + fitted_height: Final = max(1, round(height * dimension_scale)) + patches: Final = math.ceil(fitted_width / 32) * math.ceil(fitted_height / 32) + if budget is None or patches <= budget: + return patches + shrink: Final = math.sqrt(32**2 * budget / (fitted_width * fitted_height)) + width_patches: Final = fitted_width * shrink / 32 + height_patches: Final = fitted_height * shrink / 32 + aligned_shrink: Final = shrink * min( + math.floor(width_patches) / width_patches, math.floor(height_patches) / height_patches + ) + resized_width: Final = max(1, math.floor(fitted_width * aligned_shrink)) + resized_height: Final = max(1, math.floor(fitted_height * aligned_shrink)) + return math.ceil(resized_width / 32) * math.ceil(resized_height / 32) + + +def _image_policy_model_name(model: str) -> str | None: + provider, separator, provider_model = model.partition("/") + if not separator: + return model + if provider == "openai": + return provider_model + if provider == "openrouter": + return provider_model.removeprefix("openai/") if provider_model.startswith("openai/") else None + if provider not in ("azure", "azure_ai"): + return None + catalog_entry: Final = as_str_mapping( + cast( # cast-ok: read-only view of the legacy string-keyed registry; as_str_mapping checks the entry + Mapping[str, object], litellm.model_cost + ).get(model) + ) + if catalog_entry is not None and catalog_entry.get("litellm_provider") in ("azure", "azure_ai"): + return provider_model.rsplit("/", 1)[-1] + return provider_model.removeprefix("responses/").removeprefix("o_series/").removeprefix("openai/") + + +def _anthropic_image_family(model: str) -> str | None: + from litellm.utils import get_potential_model_names + + provider, separator, _ = model.partition("/") + if separator and provider not in ( + "anthropic", + "bedrock", + "bedrock_converse", + "bedrock_mantle", + "vertex_ai", + "vertex_ai_beta", + "azure_ai", + "openrouter", + ): + return None + candidates: Final = get_potential_model_names(model, provider if separator else "anthropic") + catalog: Final = cast(Mapping[str, object], litellm.model_cost) # cast-ok: read-only view of the legacy registry + for candidate in dict.fromkeys(candidates.values()): + if not isinstance(candidate, str): + continue + if (entry := as_str_mapping(catalog.get(candidate))) is None or entry.get("supports_vision") is not True: + continue + if ( + version := re.search( + r"(?:^|[./])claude-(?:(?:haiku|sonnet|opus)-)?(\d+)(?:[-.](\d{1,2})(?!\d))?", candidate + ) + ) is None: + continue + major, minor = int(version[1]), int(version[2] or 0) + return "claude-high" if (major, minor) >= (4, 7) else "claude-standard" + return None + + +def _anthropic_image_tokens(dimensions: tuple[int, int] | None, limits: tuple[int, int]) -> int: + max_edge, token_budget = limits + if dimensions is None: + return token_budget + long_edge, short_edge = max(dimensions), min(dimensions) + aspect_ratio: Final = long_edge / short_edge + + def tokens_at_edge(edge: int) -> int: + return math.ceil(edge / 28) * math.ceil(max(round(edge / aspect_ratio), 1) / 28) + + fitted_edge: Final = bisect_right(range(1, min(long_edge, max_edge) + 1), token_budget, key=tokens_at_edge) + return tokens_at_edge(max(fitted_edge, 1)) + + +def _image_model_family(model: str) -> str | None: + provider_model: Final = _image_policy_model_name(model) + if provider_model is None: + return _anthropic_image_family(model) + family: Final = re.sub(r"-\d{4}-\d{2}-\d{2}$", "", provider_model) + canonical: Final = litellm.azure_llms.get(f"azure/{family}", family) if model.startswith("azure/") else family + return ( + canonical + if canonical in _TILE_IMAGE_TOKEN_COSTS or canonical in _PATCH_IMAGE_MULTIPLIERS + else _anthropic_image_family(model) + ) + + +def _image_policy_tokens(family: str | None, dimensions: tuple[int, int] | None, detail: str) -> int: + if family is None: + return DEFAULT_IMAGE_TOKEN_COUNT + anthropic_limits: Final = _ANTHROPIC_IMAGE_LIMITS.get(family) + if anthropic_limits is not None: + return _anthropic_image_tokens(dimensions, anthropic_limits) + tile_costs: Final = _TILE_IMAGE_TOKEN_COSTS.get(family) + multiplier: Final = _PATCH_IMAGE_MULTIPLIERS.get(family) + if tile_costs is not None: + base_tokens, tile_tokens = tile_costs + if detail == "original": + return DEFAULT_IMAGE_TOKEN_COUNT + if detail == "low": + return base_tokens + if dimensions is None: + return high_detail_image_token_upper_bound(base_tokens, tile_tokens) + width, height = dimensions + scale: Final = min(1.0, 2048 / max(width, height), 768 / min(width, height)) + tiles: Final = calculate_tiles_needed(max(1, math.floor(width * scale)), max(1, math.floor(height * scale))) + return base_tokens + tiles * tile_tokens + limits: Final = _patch_image_limits(family, detail) + if limits is None or multiplier is None: + return DEFAULT_IMAGE_TOKEN_COUNT + max_dimension, patch_budget = limits + patches: Final = ( + min(math.ceil(max_dimension / 32) ** 2, patch_budget if patch_budget is not None else _MAX_IMAGE_PATCHES) + if dimensions is None + else _image_patch_count(*dimensions, max_dimension, patch_budget) + ) + return math.ceil(patches * multiplier) + + +def estimate_image_tokens( + model: str | Sequence[str], data: str, detail: str = "auto", *, fallback_model: str | None = None +) -> int: + """Estimate image policies from at most 48 KiB of inline headers, without fetching media.""" + if detail not in ("low", "high", "auto", "original"): + return DEFAULT_IMAGE_TOKEN_COUNT + models: Final = (model,) if isinstance(model, str) else model + fallback_family: Final = _image_model_family(fallback_model) if fallback_model is not None else None + families: Final = tuple( + dict.fromkeys(_image_model_family(candidate) or fallback_family for candidate in models or ("",)) + ) + if all(family is None for family in families): + return DEFAULT_IMAGE_TOKEN_COUNT + dimensions: Final = _inline_image_dimensions(data) + return max(_image_policy_tokens(family, dimensions, detail) for family in families) + + TokenCounterFunction = Callable[[str], int] """ Type for a function that counts tokens in a string. diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 567b8452ae2..e44d9304ac1 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -45,11 +45,14 @@ from litellm.caching.redis_cache import log_redis_failure from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.cost_calculator import get_usage_object from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.audio_utils.utils import estimate_audio_tokens +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, + split_media_for_tpm, ) from litellm.litellm_core_utils.ptu_pricing import is_ptu_cost_attribution_enabled -from litellm.litellm_core_utils.token_counter import offload_token_count +from litellm.litellm_core_utils.token_counter import estimate_image_tokens, offload_token_count from litellm.llms.azure.ptu_capacity import normalized_tokens from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( @@ -98,6 +101,7 @@ if TYPE_CHECKING: from opentelemetry.trace import Span as _Span from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache + from litellm.router import Router from litellm.types.agents import AgentResponse from litellm.types.caching import RedisPipelineIncrementOperation @@ -135,6 +139,18 @@ class RateLimitedModel: return requested_limit if requested_limit is not None else limits.get(self.group) +def _resolve_media_models(router: "Router | None", model: str, team_id: str | None) -> tuple[str, ...]: + if router is None: + return (model,) + return tuple(dict.fromkeys(router.resolved_litellm_models(model, team_id, use_base_model=True))) or (model,) + + +def _resolve_media_models_via_proxy_router(model: str, team_id: str | None) -> tuple[str, ...]: + from litellm.proxy.proxy_server import llm_router + + return _resolve_media_models(llm_router, model, team_id) + + def _resolve_model_group_alias_via_proxy_router(model: str) -> str | None: from litellm.proxy.proxy_server import llm_router @@ -862,6 +878,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): time_provider: Callable[[], datetime] | None = None, tag_rate_limit_resolver: TagRateLimitResolver = resolve_tag_rate_limits_from_db, model_group_resolver: Callable[[str], str | None] = _resolve_model_group_alias_via_proxy_router, + media_model_resolver: Callable[[str, str | None], tuple[str, ...]] = _resolve_media_models_via_proxy_router, ptu_team_ceiling_resolver: Callable[ [str, str], PTUTeamCeiling | None ] = _resolve_ptu_team_ceiling_via_proxy_router, @@ -872,6 +889,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self._time_provider = time_provider or datetime.now self._tag_rate_limit_resolver = tag_rate_limit_resolver self._model_group_resolver = model_group_resolver + self._media_model_resolver = media_model_resolver self._ptu_team_ceiling_resolver = ptu_team_ceiling_resolver self._fail_closed_resolver = fail_closed_resolver self._force_hash_tag_grouping_resolver = force_hash_tag_grouping_resolver @@ -1133,11 +1151,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def _estimate_tokens_for_request( self, - data: dict, + data: object, model: str | None = None, min_configured_tpm_limit: int | None = None, call_type: str | None = None, configured_output_tokens: int | None = None, + team_id: str | None = None, ) -> int: """ Estimate total tokens this request will consume so we can reserve them @@ -1164,6 +1183,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): min_configured_tpm_limit=min_configured_tpm_limit, call_type=call_type, configured_output_tokens=configured_output_tokens, + team_id=team_id, + media_model_resolver=self._media_model_resolver, ) total_estimated: Final = estimated_input_tokens + max_tokens_estimate @@ -1184,6 +1205,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): call_type: str | None, configured_output_tokens: int | None, raw_estimate: int, + team_id: str | None = None, ) -> int: """The team PTU ceiling counts Azure normalized tokens, so its reservation weighs the output budget the way the ceiling does instead of the raw sum the other scopes reserve.""" @@ -1194,6 +1216,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): min_configured_tpm_limit=min_configured_tpm_limit, call_type=call_type, configured_output_tokens=configured_output_tokens, + team_id=team_id, + media_model_resolver=self._media_model_resolver, ) normalized: Final = normalized_tokens( ceiling, prompt_tokens=estimated_input_tokens, completion_tokens=max_tokens_estimate @@ -1206,6 +1230,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): min_configured_tpm_limit: int | None = None, call_type: str | None = None, configured_output_tokens: int | None = None, + include_media: bool = True, + team_id: str | None = None, + media_model_resolver: Callable[[str, str | None], tuple[str, ...]] | None = None, ) -> tuple[int, int]: """ Estimate input tokens and output (max_tokens) budget separately, so @@ -1229,10 +1256,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): floor entirely, so the reservation reflects what this tenant's model actually emits rather than one constant shared by every tenant. """ - if not isinstance(data, dict): + request_data: Final = as_str_mapping(data) + if not isinstance(data, dict) or request_data is None: return 0, 0 + request_model: Final = request_data.get("model") translated_data: Final = self._translate_google_genai_native_request(data, call_type) - estimable_data: Final = translated_data if translated_data is not None else data + estimable_data: Final = translated_data if translated_data is not None else request_data selected_fields: Final[tuple[object | None, object | None, object | None]] = ( (None, None, estimable_data.get("input")) if call_type in RESPONSES_API_CALL_TYPES or call_type in EMBEDDING_API_CALL_TYPES @@ -1249,10 +1278,27 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) messages, prompt, raw_input = selected_fields - input_text: Final = ( + original_input: Final = ( billable_input_since_newest_compaction(raw_input) if call_type in RESPONSES_API_CALL_TYPES else raw_input ) + media_source: Final = messages if messages is not None else original_input + projected, media = ( + split_media_for_tpm(media_source) + if include_media and (isinstance(messages, list) or call_type in RESPONSES_API_CALL_TYPES) + else (original_input, ()) + ) + input_text: Final = projected if call_type in RESPONSES_API_CALL_TYPES else original_input + fallback_model: Final = request_model if isinstance(request_model, str) else None + resolved_models: Final = ( + media_model_resolver(fallback_model or "", team_id) + if media and media_model_resolver is not None + else (fallback_model or "",) + ) + media_tokens: Final = sum( + self._estimate_media_block_tokens(block, resolved_models, fallback_model) for block in media + ) + total_chars: Final = ( len(get_str_from_messages(messages)) if isinstance(messages, list) and messages @@ -1267,7 +1313,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): else 0 ) - estimated_input_tokens: Final = max(1, total_chars // DEFAULT_CHARS_PER_TOKEN) if total_chars > 0 else 0 + estimated_input_tokens: Final = ( + max(1, total_chars // DEFAULT_CHARS_PER_TOKEN) if total_chars > 0 else 0 + ) + media_tokens explicit_max_tokens: Final = self._get_explicit_output_cap(data, call_type) generates_no_output: Final = ( @@ -1283,7 +1331,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): max_tokens_estimate: Final = ( 0 if generates_no_output - or (explicit_max_tokens is None and total_chars == 0 and configured_output_tokens is None) + or (explicit_max_tokens is None and estimated_input_tokens == 0 and configured_output_tokens is None) else explicit_max_tokens if explicit_max_tokens is not None else configured_output_tokens @@ -1293,6 +1341,24 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return estimated_input_tokens, max_tokens_estimate * self.get_output_candidate_count(data, call_type) + @staticmethod + def _estimate_media_block_tokens( + block: Mapping[str, object], models: tuple[str, ...], fallback_model: str | None + ) -> int: + if block.get("type") == "input_audio": + audio: Final = _callback_kwargs_mapping(block.get("input_audio")) + data: Final = audio.get("data") + return estimate_audio_tokens(data if isinstance(data, str) else None, models, fallback_model=fallback_model) + image: Final = _callback_kwargs_mapping(block.get("image_url")) + url: Final = image.get("url") + detail: Final = image.get("detail") + return estimate_image_tokens( + models, + url if isinstance(url, str) else "", + detail if isinstance(detail, str) else "auto", + fallback_model=fallback_model, + ) + def _is_redis_cluster(self) -> bool: """ Check if the dual cache is using Redis cluster. @@ -3958,7 +4024,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): except Exception: # noqa: BLE001 # tokenizer failures degrade to the cheap estimate if call_type in RERANK_API_CALL_TYPES and isinstance(selected_text, str): return max(0, len(selected_text) // DEFAULT_CHARS_PER_TOKEN) - estimated_input_tokens, _ = self._estimate_input_and_output_tokens(data=data, call_type=call_type) + estimated_input_tokens, _ = self._estimate_input_and_output_tokens( + data=data, call_type=call_type, include_media=False + ) return estimated_input_tokens + audio_token_estimate async def _reserve_project_io_tokens_or_raise( @@ -4003,6 +4071,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): data=data, min_configured_tpm_limit=min_configured_otpm_limit, call_type=call_type, + include_media=False, ) raw_estimated_input_tokens: Final = await offload_token_count(self._estimate_precise_input_tokens)( data=data, model=requested_model, call_type=call_type @@ -4360,12 +4429,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Post-call reconciliation refunds the over-reservation # delta when actual usage comes in below the floor. estimated_tokens: Final = max( - self._estimate_tokens_for_request( + await offload_token_count(self._estimate_tokens_for_request)( data=data, model=requested_model, min_configured_tpm_limit=min_configured_tpm_limit, call_type=call_type, configured_output_tokens=configured_output_tokens, + team_id=user_api_key_dict.team_id, ), 1, ) @@ -4383,13 +4453,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) capped_request: Final = _REQUEST_RATE_LIMIT_DATA.validate_python(data) - ptu_estimated_tokens: Final = self._estimate_ptu_tokens_for_request( + ptu_estimated_tokens: Final = await offload_token_count(self._estimate_ptu_tokens_for_request)( ceiling=stash.ptu_ceiling, data=capped_request, min_configured_tpm_limit=min_configured_tpm_limit, call_type=call_type, configured_output_tokens=configured_output_tokens, raw_estimate=estimated_tokens, + team_id=user_api_key_dict.team_id, ) if stash.ptu_ceiling is not None and ptu_estimated_tokens > stash.ptu_ceiling.tpm_limit: if not reuses_request_slot: diff --git a/litellm/router.py b/litellm/router.py index e8d83e1203e..d712ed4ac52 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -10364,28 +10364,38 @@ class Router: model_info[field] = backend_value @staticmethod - def _cost_map_backend_model(deployment: Deployment) -> str: - model_info_base_model: Final = deployment.model_info.base_model + def _cost_map_backend_model(deployment: Deployment | DeploymentTypedDict) -> str: + model_info_base_model: Final = ( + deployment.model_info.base_model + if isinstance(deployment, Deployment) + else _configured_model_info(deployment).get("base_model") + ) if isinstance(model_info_base_model, str) and model_info_base_model: return model_info_base_model - params_base_model: Final = deployment.litellm_params.get("base_model") + params: Final = _LITELLM_PARAMS_ADAPTER.validate_python( + deployment.litellm_params.model_dump(exclude_none=True) + if isinstance(deployment, Deployment) + else deployment["litellm_params"] + ) + params_base_model: Final = params.get("base_model") if isinstance(params_base_model, str) and params_base_model: return params_base_model - params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(deployment.litellm_params.model_dump(exclude_none=True)) + model: Final = ( + deployment.litellm_params.model + if isinstance(deployment, Deployment) + else deployment["litellm_params"].get("model", "") + ) invoked_model: Final = params.get("model_id") if ( isinstance(invoked_model, str) and invoked_model - and ( - deployment.litellm_params.custom_llm_provider == "bedrock" - or deployment.litellm_params.model.startswith("bedrock/") - ) + and (params.get("custom_llm_provider") == "bedrock" or model.startswith("bedrock/")) ): profile_model: Final = invoked_model if invoked_model.startswith("bedrock/") else f"bedrock/{invoked_model}" profile_info: Final = lookup_model_info_for_registration(model=profile_model, custom_llm_provider="bedrock") if not isinstance(profile_info, ModelInfoLookupMiss): return profile_model - return deployment.litellm_params.model + return model @staticmethod def _inherit_builtin_service_tier_pricing( @@ -13945,7 +13955,9 @@ class Router: ] return named or self._get_wildcard_deployments(model_name=model_group) - def resolved_litellm_models(self, model_name: str, team_id: str | None = None) -> tuple[str, ...]: + def resolved_litellm_models( + self, model_name: str, team_id: str | None = None, *, use_base_model: bool = False + ) -> tuple[str, ...]: """The provider model strings `model_name` can actually be served by on this proxy. `get_model_list` composes every channel the request path itself uses (exact name, @@ -13958,6 +13970,12 @@ class Router: deployment behind it - so the fallback for an empty result is the caller's policy, never this function's. """ + if use_base_model: + return tuple( + backend_model + for deployment in self.get_model_list(model_name=model_name, team_id=team_id) or () + if (backend_model := self._cost_map_backend_model(deployment)) + ) return tuple( litellm_model for deployment in self.get_model_list(model_name=model_name, team_id=team_id) or () 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 d50ab908b43..b9675f6fd5c 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 @@ -37,6 +37,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( schema_has_unresolvable_ref, schema_property_count, split_concatenated_json_objects, + split_media_for_tpm, strip_encrypted_reasoning_from_messages, strip_litellm_response_only_fields, system_messages_first, @@ -2512,3 +2513,118 @@ def test_shared_image_url_agrees_with_responses_normalization( expected: Final = {"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE"}} assert as_openai_image_part(part) == expected assert as_openai_image_part(normalized) == expected + + +@pytest.mark.parametrize( + "part,projected,media", + [ + ( + {"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "high"}}, + {"type": "image_url", "image_url": {"url": "", "detail": "high"}}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "high"}}, + ), + ( + {"type": "input_image", "image_url": "https://example.com/image.png", "detail": "low"}, + {"type": "input_image", "image_url": "", "detail": "low"}, + {"type": "image_url", "image_url": {"url": "https://example.com/image.png", "detail": "low"}}, + ), + ( + {"type": "computer_screenshot", "image_url": "data:image/png;base64,IMAGE"}, + {"type": "computer_screenshot", "image_url": ""}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "auto"}}, + ), + ( + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "IMAGE"}}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": ""}}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "auto"}}, + ), + ( + {"type": "input_audio", "input_audio": {"data": "AUDIO", "format": "wav"}}, + {"type": "input_audio", "input_audio": {"data": "", "format": "wav"}}, + {"type": "input_audio", "input_audio": {"data": "AUDIO", "format": "wav"}}, + ), + ( + {"type": "file", "file": {"file_data": "IMAGE", "format": "image/png", "filename": "image.png"}}, + {"type": "file", "file": {"file_data": "", "format": "image/png", "filename": "image.png"}}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "auto"}}, + ), + ( + {"type": "input_file", "file_data": "data:audio/wav;base64,AUDIO", "filename": "audio.wav"}, + {"type": "input_file", "file_data": "", "filename": "audio.wav"}, + {"type": "input_audio", "input_audio": {"data": "AUDIO", "format": "wav"}}, + ), + ], +) +def test_tpm_projection_blanks_only_declared_media_transport( + part: Mapping[str, object], projected: Mapping[str, object], media: Mapping[str, object], +) -> None: + text: Final = {"type": "text", "text": "data:image/png;base64,IMAGE"} + request: Final = [{"role": "user", "content": [text, {**part, "cache_control": {"type": "ephemeral"}}]}] + original: Final = copy.deepcopy(request) + assert split_media_for_tpm(request) == ( + [{"role": "user", "content": [text, {**projected, "cache_control": {"type": "ephemeral"}}]}], + (media,), + ) + assert request == original + + +def test_tpm_projection_retains_nested_and_native_history_around_media() -> None: + image: Final = {"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "high"}} + empty_image: Final = {"type": "image_url", "image_url": {"url": "", "detail": "high"}} + document: Final = {"type": "document", "title": "keep", "source": {"type": "content", "content": [image]}} + empty_document: Final = {**document, "source": {"type": "content", "content": [empty_image]}} + tool: Final = {"type": "tool_result", "tool_use_id": "call", "content": [document, {"type": "text", "text": "keep"}]} + empty_tool: Final = {**tool, "content": [empty_document, {"type": "text", "text": "keep"}]} + request: Final = [ + {"role": "user", "content": [tool]}, + {"type": "function_call_output", "call_id": "call", "output": [image]}, + {"type": "custom_tool_call_output", "call_id": "custom", "output": [image]}, + {"type": "computer_call_output", "call_id": "computer", "output": { + "type": "computer_screenshot", "image_url": "data:image/png;base64,IMAGE", + }}, + {"type": "image_generation_call", "id": "ig", "result": "IMAGE", "status": "completed"}, + {"type": "code_interpreter_call", "id": "ci", "code": "keep()", "outputs": [ + {"type": "logs", "logs": "keep"}, {"type": "image", "url": "https://example.com/image.png"}, + ]}, + ] + original: Final = copy.deepcopy(request) + projected, media = split_media_for_tpm(request) + assert projected == [ + {"role": "user", "content": [empty_tool]}, + {"type": "function_call_output", "call_id": "call", "output": [empty_image]}, + {"type": "custom_tool_call_output", "call_id": "custom", "output": [empty_image]}, + {"type": "computer_call_output", "call_id": "computer", "output": { + "type": "computer_screenshot", "image_url": "", + }}, + {"type": "image_generation_call", "id": "ig", "result": "", "status": "completed"}, + {"type": "code_interpreter_call", "id": "ci", "code": "keep()", "outputs": [ + {"type": "logs", "logs": "keep"}, {"type": "image", "url": ""}, + ]}, + ] + assert media == ( + image, image, image, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "auto"}}, + {"type": "image_url", "image_url": {"url": "data:application/octet-stream;base64,IMAGE", "detail": "auto"}}, + {"type": "image_url", "image_url": {"url": "https://example.com/image.png", "detail": "auto"}}, + ) + assert request == original + + +def test_tpm_projection_preserves_non_media_identity_and_opaque_values() -> None: + image: Final = {"type": "input_image", "image_url": "data:image/png;base64,OPAQUE"} + request: Final = [ + {"role": "user", "content": json.dumps(image)}, + {"type": "custom_tool_call_output", "call_id": "call", "output": json.dumps(image)}, + {"type": "function_call", "call_id": "call", "name": "keep", "arguments": {"content": [image]}}, + {"type": "file_search_call", "results": [{"text": "keep", "content": [image]}]}, + {"type": "image_generation_call", "id": "ig", "status": "in_progress", "result": None}, + {"role": "user", "content": [ + {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "PDF"}}, + {"type": "document", "source": {"type": "text", "data": json.dumps(image)}}, + {"type": "file", "file": {"format": "application/pdf", "file_data": "PDF"}}, + {"type": "text", "text": "keep", "metadata": {"content": [image]}}, + ]}, + ] + projected, media = split_media_for_tpm(request) + assert projected is request + assert media == () diff --git a/tests/unit/litellm_core_utils/test_audio_utils.py b/tests/unit/litellm_core_utils/test_audio_utils.py index 155f6680416..122445e8e30 100644 --- a/tests/unit/litellm_core_utils/test_audio_utils.py +++ b/tests/unit/litellm_core_utils/test_audio_utils.py @@ -2,16 +2,23 @@ Test the audio utils functionality in litellm_core_utils/audio_utils/utils.py """ +import base64 import io import os import tempfile +from typing import Final from unittest.mock import mock_open, patch import pytest +import soundfile as sf from litellm.litellm_core_utils.audio_utils.utils import ( + AUDIO_BYTES_PER_TOKEN, + AUDIO_TOKENS_PER_SECOND, + DEFAULT_AUDIO_TOKEN_ESTIMATE, ProcessedAudioFile, calculate_request_duration, + estimate_audio_tokens, get_audio_file_content_hash, get_audio_file_for_health_check, get_audio_file_name, @@ -19,6 +26,63 @@ from litellm.litellm_core_utils.audio_utils.utils import ( ) +def _encoded_audio(seconds: int, sample_rate: int, channels: int, container: str) -> str: + buffer: Final = io.BytesIO() + with sf.SoundFile(buffer, "w", samplerate=sample_rate, channels=channels, format=container) as audio: + audio.buffer_write(b"\0" * (seconds * sample_rate * channels * 2), dtype="int16") + return base64.b64encode(buffer.getvalue()).decode() + + +@pytest.mark.parametrize( + ("sample_rate", "channels", "container"), [(8000, 1, "WAV"), (48000, 2, "WAV"), (24000, 1, "FLAC")] +) +def test_audio_estimate_uses_duration_independent_of_sample_rate_channels_and_compression( + sample_rate: int, channels: int, container: str +) -> None: + assert estimate_audio_tokens(_encoded_audio(1, sample_rate, channels, container)) == AUDIO_TOKENS_PER_SECOND + assert estimate_audio_tokens(_encoded_audio(31, sample_rate, channels, container)) == 31 * AUDIO_TOKENS_PER_SECOND + + +# Audio rates verified 2026-10-28: https://ai.google.dev/gemini-api/docs/generate-content/media-resolution and +# https://developers.openai.com/api/docs/guides/voice-latency-cost +@pytest.mark.parametrize( + ("models", "fallback", "expected"), + [ + (("openai/gpt-audio-1.5",), "gemini-3.8-flash", 10), + (("gpt-audio-1.5", "gemini/gemini-3.8-flash"), None, 25), + (("vertex_ai/gemini-3.8-flash",), None, 25), + (("openrouter/google/gemini-3.8-flash",), None, 25), + (("azure/audio-deployment",), "gpt-audio-1.5", 10), + (("opaque-deployment",), "gemini-3.8-flash", 25), + ((), "gemini-3.8-flash", 25), + (("gemini/gemini-2.5-flash", "gemini/gemini-3-pro-preview"), None, 32), + (("other/gemini-3.8-flash",), None, 10), + ], +) +def test_audio_estimate_prefers_each_backend_rate_over_the_requested_model( + models: tuple[str, ...], fallback: str | None, expected: int +) -> None: + assert estimate_audio_tokens(_encoded_audio(1, 8000, 1, "WAV"), models, fallback) == expected + + +@pytest.mark.parametrize("data", [None, "", "invalid-base64!", base64.b64encode(b"unreadable audio").decode()]) +def test_audio_estimate_retains_nonzero_fallback_for_unreadable_audio(data: str | None) -> None: + assert estimate_audio_tokens(data) == DEFAULT_AUDIO_TOKEN_ESTIMATE + + +def test_audio_estimate_retains_size_based_fallback_for_large_unreadable_audio() -> None: + data: Final = base64.b64encode(b"\0" * (AUDIO_BYTES_PER_TOKEN * 400)).decode() + assert estimate_audio_tokens(data) == 400 + + +def test_audio_estimate_bounds_base64_allocation_before_decoding() -> None: + data: Final = base64.b64encode(b"\0" * (16 * 1024 * 1024 + 1)).decode() + with patch("base64.b64decode", side_effect=AssertionError("oversized audio must not be decoded")) as decode: + estimate: Final = estimate_audio_tokens(data) + assert estimate == len(data) * 3 // 4 // AUDIO_BYTES_PER_TOKEN + decode.assert_not_called() + + class TestProcessAudioFile: """Test the process_audio_file function with various input types""" diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 92f971f64e9..b02e0732817 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -37,6 +37,7 @@ from litellm.litellm_core_utils.token_counter import ( _get_extrapolating_count_function, _get_tiktoken_count_function, calculate_img_tokens, + estimate_image_tokens, high_detail_image_token_upper_bound, offload_token_count, ) @@ -2152,6 +2153,136 @@ def test_high_detail_image_token_upper_bound_is_reached_by_the_largest_high_res_ assert calculate_img_tokens(_png_data_url(1, 1), mode="high") < high_detail_image_token_upper_bound() +# Image sizing verified 2026-10-28: https://developers.openai.com/api/docs/guides/images-vision +# and https://platform.claude.com/docs/en/build-with-claude/vision-coordinates +@pytest.mark.parametrize( + ("model", "width", "height", "detail", "expected"), + [ + ("gpt-4o-mini", 1024, 1024, "low", 2833), + ("gpt-4o-mini", 1024, 1024, "high", 25501), + ("gpt-5.1", 512, 512, "auto", 210), + ("gpt-6-astra", 2048, 2048, "high", 3000), + ("gpt-6-astra", 4096, 512, "high", 2458), + ("gpt-6-astra", 2048, 2048, "auto", 4916), + ("gpt-5.6-sol", 2048, 2048, "low", 308), + ("gpt-5.5", 6400, 6400, "original", 12000), + ("gpt-5.4-mini", 2048, 2048, "low", 4916), + ("gpt-5.4-nano", 2048, 2048, "auto", 3000), + ("gpt-5.2", 4096, 512, "auto", 615), + ("gpt-4.1-mini-2025-04-14", 1024, 1024, "low", 1659), + ("claude-haiku-5-5", 200, 200, "auto", 64), + ("anthropic/claude-haiku-5-5", 1024, 1024, "auto", 1369), + ("claude-haiku-5-5", 1920, 1080, "low", 2691), + ("claude-haiku-4-5", 1920, 1080, "auto", 1560), + ("claude-haiku-4-5", 1075, 1520, "auto", 1551), + ("claude-haiku-5-5", 3840, 2160, "auto", 4784), + ("claude-haiku-5-5", 8000, 100, "auto", 184), + ], +) +def test_image_estimate_uses_model_resolution_and_detail( + model: str, width: int, height: int, detail: str, expected: int +) -> None: + assert estimate_image_tokens(model, _png_data_url(width, height), detail) == expected + + +@pytest.mark.parametrize( + ("models", "fallback", "expected_model"), + [ + (("azure/deployment-prod",), "gpt-4o", "gpt-4o"), + (("openai/gpt-4o",), "gpt-4o-mini", "gpt-4o"), + (("azure/deployment-prod", "openai/gpt-4o"), "gpt-4o-mini", "gpt-4o-mini"), + (("openai/gpt-4o", "openrouter/openai/gpt-4o-mini"), None, "gpt-4o-mini"), + (("azure_ai/openai/gpt-6-astra",), None, "gpt-6-astra"), + (("azure/gpt-41",), None, "gpt-4.1"), + (("azure/eu/gpt-4o-mini",), None, "gpt-4o-mini"), + (("unknown", "gpt-4o-mini"), None, "gpt-4o-mini"), + ((), "gpt-6-astra", "gpt-6-astra"), + (("other/gpt-4o-mini",), None, "unregistered-model"), + (("azure/opaque/gpt-4o-mini",), None, "unregistered-model"), + (("bedrock/anthropic.claude-haiku-5-5",), None, "claude-haiku-5-5"), + (("bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0",), None, "claude-haiku-4-5"), + (("vertex_ai/claude-haiku-5-5@default",), None, "claude-haiku-5-5"), + (("azure_ai/claude-haiku-5-5",), None, "claude-haiku-5-5"), + (("anthropic/claude-haiku-5-5", "openai/gpt-4o"), "gpt-4o-mini", "claude-haiku-5-5"), + (("other/claude-haiku-5-5",), None, "unregistered-model"), + ], +) +def test_image_estimate_prefers_each_backend_policy_over_the_requested_model( + models: tuple[str, ...], fallback: str | None, expected_model: str +) -> None: + data: Final = _png_data_url(1024, 1024) + assert estimate_image_tokens(models, data, "high", fallback_model=fallback) == estimate_image_tokens( + expected_model, data, "high" + ) + + +@pytest.mark.parametrize( + ("model", "detail"), + [ + ("unknown", "high"), + ("claude-unregistered", "auto"), + ("gpt-4o", "original"), + ("gpt-5.2", "original"), + ("gpt-6-astra", "invalid"), + ], +) +def test_image_estimate_preserves_default_when_policy_is_unavailable(model: str, detail: str) -> None: + assert ( + estimate_image_tokens(model, _png_data_url(1024, 1024), detail) == litellm.constants.DEFAULT_IMAGE_TOKEN_COUNT + ) + + +@pytest.mark.parametrize( + "data", + [ + "https://image.invalid/photo.png", + "data:image/png;base64,invalid!", + "data:application/octet-stream;base64,bm90LWFuLWltYWdl", + "data:image/jpeg;base64,/9j/4AAQ", + "data:image/webp;base64,UklGRgAAAABXRUJQ", + _png_data_url(0, 1024), + ], +) +# Accepted patch limits verified 2026-10-28: https://developers.openai.com/api/docs/guides/images-vision +# and https://platform.claude.com/docs/en/build-with-claude/vision +@pytest.mark.parametrize( + ("model", "detail", "bound"), + [ + ("gpt-4o", "high", high_detail_image_token_upper_bound()), + ("gpt-6-astra", "high", 3000), + ("gpt-6-astra", "auto", 36000), + ("claude-haiku-5-5", "auto", 4784), + ("claude-haiku-4-5", "auto", 1568), + ], +) +def test_image_estimate_uses_bounded_fallback_without_fetching_media( + data: str, model: str, detail: str, bound: int +) -> None: + with patch("httpx.Client.send") as send: + estimate: Final = estimate_image_tokens(model, data, detail) + send.assert_not_called() + assert estimate == bound + assert estimate > estimate_image_tokens(model, _png_data_url(1024, 1024), detail) + + +@pytest.mark.parametrize("container", ["PNG", "JPEG", "GIF", "WEBP"]) +def test_image_estimate_depends_on_dimensions_and_not_compressed_payload_size(container: str) -> None: + output: Final = io.BytesIO() + Image.new("RGB", (1024, 1024)).save(output, format=container) + data: Final = "data:application/octet-stream;base64," + base64.b64encode(output.getvalue()).decode() + assert estimate_image_tokens("gpt-6-astra", data) == estimate_image_tokens("gpt-6-astra", _png_data_url(1024, 1024)) + + +def test_image_estimate_bounds_header_decoding_for_large_inline_payloads() -> None: + header: Final = base64.b64decode(_png_data_url(1024, 1024).split(",", 1)[1]) + data: Final = "data:image/png;base64," + base64.b64encode(header + b"\0" * 4_000_000).decode() + with patch("base64.b64decode", wraps=base64.b64decode) as decode_header: + estimate: Final = estimate_image_tokens("gpt-6-astra", data) + assert estimate == estimate_image_tokens("gpt-6-astra", _png_data_url(1024, 1024)) + assert decode_header.call_count == 1 + assert len(decode_header.call_args.args[0]) <= 65_536 + + HUB_SETUP_SCRIPT: Final = """ import json import sys diff --git a/tests/unit/proxy/hooks/test_tpm_concurrent.py b/tests/unit/proxy/hooks/test_tpm_concurrent.py index 3d1e494e36c..8245c261a12 100644 --- a/tests/unit/proxy/hooks/test_tpm_concurrent.py +++ b/tests/unit/proxy/hooks/test_tpm_concurrent.py @@ -15,28 +15,38 @@ Redis. """ import asyncio +import base64 +import copy +import io import json from datetime import datetime, timedelta from typing import Any, Dict, Final import pytest +import soundfile as sf +from fastapi import HTTPException +from PIL import Image +from litellm import token_counter from litellm.caching.caching import DualCache +from litellm.litellm_core_utils.token_counter import UnsupportedContentTypeError, estimate_image_tokens from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _AUDIO_BYTES_PER_TOKEN, + DEFAULT_AUDIO_TOKEN_ESTIMATE, PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY, RateLimitedModel, - _AUDIO_BYTES_PER_TOKEN, - PROXY_MaxParallelRequestsHandler_v3 as RateLimitHandler, -) -from litellm.proxy.hooks.parallel_request_limiter_v3 import ( _call_id_from_callback_kwargs, _request_stash, get_or_create_request_stash, get_request_stash, ) +from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + PROXY_MaxParallelRequestsHandler_v3 as RateLimitHandler, +) from litellm.proxy.utils import InternalUsageCache, hash_token +from litellm.router_utils.ptu_shares import PTUTeamCeiling from litellm.types.llms.openai import ( InputTokensDetails, ResponseAPIUsage, @@ -3775,3 +3785,284 @@ def test_reservation_counts_a_compaction_item_by_its_summary(rate_limiter): ) == handler._estimate_input_and_output_tokens( data={"model": "gpt-4o", "input": _COMPACTED_INPUT[2:], "max_output_tokens": 1}, call_type="aresponses" ) + + +_QUOTA_TEXT: Final = "a" * 400 +_QUOTA_JSON_TEXT: Final = '{"type":"input_image","image_url":"data:image/png;base64,AAAA"}' + + +@pytest.mark.parametrize( + ("call_type", "request_data", "expected"), + [ + ("acompletion", {"messages": [{"role": "user", "content": _QUOTA_TEXT}], "max_tokens": 16}, (100, 16)), + ("acompletion", {"messages": [{"role": "user", "content": _QUOTA_JSON_TEXT}], "max_tokens": 16}, (15, 16)), + ( + "anthropic_messages", + { + "messages": [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "t", + "content": [{"type": "text", "text": _QUOTA_TEXT}], + } + ], + } + ], + "max_tokens": 16, + }, + (0, 16), + ), + ( + "acompletion", + { + "messages": [ + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "t", "type": "function", "function": {"name": "f", "arguments": _QUOTA_TEXT}} + ], + } + ], + "max_tokens": 16, + }, + (0, 16), + ), + ("aresponses", {"input": _QUOTA_TEXT, "max_output_tokens": 16}, (100, 16)), + ("aresponses", {"input": _QUOTA_TEXT, "instructions": "b" * 400, "max_output_tokens": 16}, (100, 16)), + ( + "aresponses", + {"input": [{"role": "user", "content": _QUOTA_TEXT}], "instructions": "b" * 400, "max_output_tokens": 16}, + (107, 16), + ), + ( + "aresponses", + { + "input": [ + { + "type": "file_search_call", + "id": "f", + "status": "completed", + "queries": ["q"], + "results": [ + { + "file_id": "file_1", + "filename": "result.txt", + "attributes": {}, + "score": 1, + "text": _QUOTA_TEXT, + } + ], + } + ], + "max_output_tokens": 16, + }, + (146, 16), + ), + ( + "aresponses", + { + "input": [{"type": "custom_tool_call_output", "call_id": "t", "output": _QUOTA_JSON_TEXT}], + "max_output_tokens": 16, + }, + (32, 16), + ), + ("aresponses", {"input": "", "instructions": "b" * 400}, (0, 0)), + ("aembedding", {"input": [_QUOTA_TEXT]}, (100, 0)), + ("atext_completion", {"prompt": [_QUOTA_TEXT], "max_tokens": 16}, (100, 16)), + ], +) +def test_media_correction_preserves_exact_non_media_reservations( + rate_limiter: tuple[RateLimitHandler, DualCache], + call_type: str, + request_data: dict[str, object], + expected: tuple[int, int], +) -> None: + handler, _ = rate_limiter + assert handler._estimate_input_and_output_tokens(request_data, call_type=call_type) == expected + + +def _quota_media_payload(kind: str, compressed: bool) -> str: + output: Final = io.BytesIO() + if kind == "image": + Image.new("RGB", (1024, 1024)).save(output, format="PNG", compress_level=9 if compressed else 0) + else: + sample_rate: Final = 8000 if compressed else 48000 + with sf.SoundFile(output, "w", samplerate=sample_rate, channels=1, format="WAV") as audio: + audio.buffer_write(b"\0" * (sample_rate * 2), dtype="int16") + return base64.b64encode(output.getvalue()).decode() + + +def _quota_media_request(kind: str, payload: str, call_type: str) -> dict[str, object]: + block: Final = ( + {"type": "input_audio", "input_audio": {"data": payload, "format": "wav"}} + if kind == "audio" + else {"type": "input_image", "image_url": "data:image/png;base64," + payload} + if call_type == "aresponses" + else {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": payload}} + if call_type == "anthropic_messages" + else {"type": "image_url", "image_url": {"url": "data:image/png;base64," + payload}} + ) + content: Final = ( + [{"type": "tool_result", "tool_use_id": "t", "content": [block]}] + if call_type == "anthropic_messages" + else [block] + ) + field: Final = "input" if call_type == "aresponses" else "messages" + return { + "model": "gpt-6-astra", + field: [{"role": "user", "content": content}], + "max_tokens": 16, + "max_output_tokens": 16, + } + + +@pytest.mark.parametrize( + ("kind", "call_type"), + [ + ("image", "acompletion"), + ("image", "aresponses"), + ("image", "anthropic_messages"), + ("audio", "acompletion"), + ("audio", "aresponses"), + ], +) +def test_media_reservation_is_independent_of_encoding_and_preserves_request( + rate_limiter: tuple[RateLimitHandler, DualCache], kind: str, call_type: str +) -> None: + handler, _ = rate_limiter + compressed: Final = _quota_media_request(kind, _quota_media_payload(kind, True), call_type) + uncompressed: Final = _quota_media_request(kind, _quota_media_payload(kind, False), call_type) + original: Final = copy.deepcopy(uncompressed) + estimate: Final = handler._estimate_input_and_output_tokens(uncompressed, call_type=call_type) + assert estimate == handler._estimate_input_and_output_tokens(compressed, call_type=call_type) + assert estimate[0] > 0 and estimate[1] == 16 + assert uncompressed == original + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", ["acompletion", "aresponses"]) +@pytest.mark.parametrize("limit", [1000, 20000]) +async def test_media_pre_call_rejects_tight_quota_and_reserves_sufficient_quota( + rate_limiter: tuple[RateLimitHandler, DualCache], call_type: str, limit: int +) -> None: + handler, cache = rate_limiter + request: Final = { + key: value + for key, value in _quota_media_request("image", _quota_media_payload("image", True), call_type).items() + if key not in ("max_tokens", "max_output_tokens") + } + auth: Final = UserAPIKeyAuth(api_key=hash_token("sk-media-quota"), tpm_limit=limit) + if limit == 1000: + with pytest.raises(HTTPException) as error: + await handler.async_pre_call_hook(user_api_key_dict=auth, cache=cache, data=request, call_type=call_type) + assert error.value.status_code == 429 + return + await handler.async_pre_call_hook(user_api_key_dict=auth, cache=cache, data=request, call_type=call_type) + stash: Final = get_request_stash() + assert stash is not None + expected: Final = handler._estimate_input_and_output_tokens( + request, call_type=call_type, min_configured_tpm_limit=limit + ) + assert stash.reserved_tokens == sum(expected) + assert expected[0] > 1000 and expected[1] >= expected[0] + + +@pytest.mark.asyncio +async def test_ptu_pre_call_weights_media_input_and_output() -> None: + cache: Final = DualCache() + ceiling: Final = PTUTeamCeiling( + model_group="gpt-6-astra", tpm_limit=10000, output_to_input_ratio=4, cached_input_ratio=0 + ) + + def resolve_ceiling(team: str, model: str) -> PTUTeamCeiling | None: + return ceiling if (team, model) == ("test-team", "gpt-6-astra") else None + + handler: Final = RateLimitHandler( + internal_usage_cache=InternalUsageCache(cache), ptu_team_ceiling_resolver=resolve_ceiling + ) + request: Final = _quota_media_request("image", _quota_media_payload("image", True), "acompletion") + await handler.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-media-ptu"), team_id="test-team"), + cache=cache, + data=request, + call_type="acompletion", + ) + stash: Final = get_request_stash() + assert stash is not None + assert ( + stash.ptu_reserved_tokens + == estimate_image_tokens("gpt-6-astra", "data:image/png;base64," + _quota_media_payload("image", True)) + 4 * 16 + ) + + +@pytest.mark.asyncio +async def test_project_fallback_keeps_legacy_audio_estimate_and_output_policy( + rate_limiter: tuple[RateLimitHandler, DualCache], +) -> None: + handler, cache = rate_limiter + unsupported: Final = {"type": "unsupported_content"} + with pytest.raises(UnsupportedContentTypeError): + token_counter(model="gpt-6-astra", messages=[{"role": "user", "content": [unsupported]}]) + request: Final = { + "model": "gpt-6-astra", + "messages": [ + { + "role": "user", + "content": [ + unsupported, + { + "type": "input_audio", + "input_audio": {"data": _quota_media_payload("audio", True), "format": "wav"}, + }, + ], + } + ], + } + auth: Final = UserAPIKeyAuth( + api_key=hash_token("sk-media-project"), + project_id="media-project", + project_metadata={"model_itpm_limit": {"gpt-6-astra": 1000}, "model_otpm_limit": {"gpt-6-astra": 1000}}, + ) + await handler.async_pre_call_hook(user_api_key_dict=auth, cache=cache, data=request, call_type="acompletion") + stash: Final = get_request_stash() + assert stash is not None + assert stash.itpm_reserved_tokens == DEFAULT_AUDIO_TOKEN_ESTIMATE + assert stash.otpm_reserved_tokens == 1 + + +@pytest.mark.asyncio +async def test_single_anthropic_image_exceeds_tight_quota_with_explicit_output_cap( + rate_limiter: tuple[RateLimitHandler, DualCache], +) -> None: + handler, cache = rate_limiter + request: Final = { + "model": "anthropic/claude-haiku-5-5", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": _quota_media_payload("image", True), + }, + } + ], + } + ], + "max_tokens": 64, + } + with pytest.raises(HTTPException) as error: + await handler.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-claude-media-quota"), tpm_limit=500), + cache=cache, + data=request, + call_type="anthropic_messages", + ) + assert error.value.status_code == 429 diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 63ee1c43edb..566024b5d94 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -20630,6 +20630,37 @@ def test_resolved_litellm_models_answers_through_every_channel_a_request_uses( assert set(_resolution_router().resolved_litellm_models(model_name)) == set(expected) +def test_resolved_litellm_models_base_model_precedence_is_opt_in() -> None: + router: Final = Router( + model_list=[ + { + "model_name": "pooled", + "litellm_params": {"model": "openai/deployment-a", "api_key": "sk-test", "base_model": "gpt-5.5"}, + "model_info": {"base_model": "gpt-6-astra"}, + }, + { + "model_name": "pooled", + "litellm_params": {"model": "openai/deployment-b", "api_key": "sk-test", "base_model": "gpt-5.6-sol"}, + }, + { + "model_name": "team-deployment", + "litellm_params": {"model": "openai/deployment-c", "api_key": "sk-test"}, + "model_info": { + "base_model": "gpt-6-astra", + "team_id": "test-team", + "team_public_model_name": "team-model", + }, + }, + ], + model_group_alias={"nickname": "pooled"}, + ) + assert router.resolved_litellm_models("nickname") == ("openai/deployment-a", "openai/deployment-b") + assert router.resolved_litellm_models("nickname", use_base_model=True) == ("gpt-6-astra", "gpt-5.6-sol") + assert router.resolved_litellm_models("team-model", "test-team") == ("openai/deployment-c",) + assert router.resolved_litellm_models("team-model", "test-team", use_base_model=True) == ("gpt-6-astra",) + assert router.resolved_litellm_models("never-configured", use_base_model=True) == () + + class TestTierParamsTheTargetAccepts: """A tier's litellm_params are applied to every request that tier routes, so one the target cannot take raised UnsupportedParamsError before the request left the proxy, turning the whole From 7c7304c22b4a855242a101303a33215ca7f1c466 Mon Sep 17 00:00:00 2001 From: Amaan Date: Thu, 29 Oct 2026 01:10:32 +0530 Subject: [PATCH 079/128] feat(providers): add ZeroGPU as an OpenAI-compatible provider (#43874) * feat(providers): add ZeroGPU as an OpenAI-compatible provider Registers zerogpu through the JSON provider registry so zerogpu/ routes to https://api.zerogpu.ai/v1/chat/completions with ZEROGPU_API_KEY and an optional ZEROGPU_API_BASE. Adds LlmProviders.ZEROGPU, the constants.py provider and endpoint lists, cost map entries for the 10 ZeroGPU text generation models, the endpoint matrix, the Add Model form fields, the README row and the dashboard provider entry * feat(zerogpu): add capability flags verified against the live API Live calls to https://api.zerogpu.ai/v1/chat/completions confirmed tools, tool_choice and json_schema response_format on 8 of the 10 catalog text models, cached token reporting on deepseek-v4.1-flash, gpt-4.1-mini, gpt-5.4-nano and gpt-5.6-luna, and reasoning_effort on gpt-5.6-luna. LFM2.5-1.2B-Instruct answers stream requests with plain JSON instead of SSE and LFM2.5-1.2B-Thinking returns 500 on every request, so both stay out of the cost map until the API serves them correctly * feat(zerogpu): price LFM2.5-1.2B-Instruct Adds zerogpu/LFM2.5-1.2B-Instruct to the cost map with the catalog prices and context window. It gets no capability flags because the live API ignores tools, response_format and max_completion_tokens for this model, so LiteLLM drops tool params for it instead of sending them. LFM2.5-1.2B-Thinking stays out of the cost map * fix(zerogpu): reject transcription requests before calling the upstream API ZeroGPU only serves chat completions, so drop it from the OpenAI transcription transport set. A zerogpu transcription request now fails locally instead of hitting an /audio/transcriptions endpoint that does not exist * test(zerogpu): cover the responses and messages bridges Both endpoints reach ZeroGPU through chat completions, so assert the request each one sends upstream and the reply it returns * refactor(zerogpu): move zerogpu out of the provider list tails New upstream providers append to the end of openai_compatible_endpoints, openai_compatible_providers and providers.json, so an entry at the tail conflicts with every one of them. Placing zerogpu next to long-stable entries keeps those merges clean. No behavior change * fix(zerogpu): price cached input for qwen3-30b-a3b-fp8 The ZeroGPU model catalog (https://docs.zerogpu.ai/docs/model-catalog, checked 2026-10-09) lists cached input for qwen3-30b-a3b-fp8 at $0.05 per 1M tokens. Without cache_read_input_token_cost, cached tokens were billed at the full $0.10 input rate * revert(zerogpu): leave the audio transcription provider set unchanged Per review, keep OPENAI_AUDIO_TRANSCRIPTION_PROVIDERS as upstream defines it. This PR only covers chat completions, and ZeroGPU now serves /audio/transcriptions anyway, so the local refusal and its test are removed * fix(zerogpu): add max_tokens so trim_messages can size ZeroGPU history trim_messages reads max_tokens as the fallback for max_input_tokens and evaluates it eagerly, so without it the lookup raised, got swallowed, and the full history came back untrimmed. max_tokens matches each model's maxTokens in ZeroGPU's models API (https://api-dashboard.zerogpu.ai/api/models, checked 2026-10-10) * test(zerogpu): record mocked requests on a recorder object The mock handler appended to a list the caller passed in, which AGENTS.md forbids. A recorder now owns the captured requests as a tuple * fix(zerogpu): cap max_tokens at each model's output limit Each cap is the max_output_tokens ZeroGPU stores for the model in its production models table (read 2026-10-20), the same column its /v1/models serves as max_length. That gives 65536 for deepseek-v4.1-flash, 32768 for gpt-4.1-mini and gpt-oss-120b, 128000 for gpt-5.4-nano and gpt-5.6-luna, 8192 for LFM2.5-1.2B-Instruct and llama-3.1-8b-instruct-fast, and 16384 for qwen3-30b-a3b-fp8. glm-5.3-flash is stored as 32786, which reads as a typo for 32768, so that row uses 32768 trim_messages sizes history from max_input_tokens, so lowering max_tokens leaves it unaffected * fix(zerogpu): flag prompt caching on glm-5.3-flash Repeating the same ~4.3K-token prompt against ZeroGPU on 2026-10-20 returned cached_tokens of 4320 on some glm-5.3-flash calls. gpt-oss-120b and llama-3.1-8b-instruct-fast returned 0 on all nine calls, and qwen3-30b-a3b-fp8 never returned cached_tokens, so they stay unflagged. Their cached prices stay as ZeroGPU publishes them, so any cache hit is still billed at the right rate * fix(zerogpu): set gpt-5.4-nano context window to 272000 ZeroGPU accepts a 271001-token gpt-5.4-nano prompt and rejects 273000 tokens as over the context window (checked 2026-10-21), so max_input_tokens is 272000 rather than the 400000 its catalog lists * fix(zerogpu): list zerogpu models for zerogpu/* wildcards zerogpu had no models_by_provider entry, so get_provider_models("zerogpu") returned None and a zerogpu/* wildcard listed no models. zerogpu_models is now filled from the cost map rows like openinfer_models and registered in models_by_provider * fix(zerogpu): flag gpt-5.4-nano as a reasoning model ZeroGPU honors reasoning_effort on gpt-5.4-nano (checked 2026-10-21). Without it the model spends 0 reasoning tokens, and with low or high it spends 32 or more. LiteLLM only treated it as a reasoning model by falling back to OpenAI's bare gpt-5.4-nano row, so get_model_info and /model/info reported supports_reasoning as null. The row now carries the flag itself * fix(zerogpu): flag deepseek, glm, gpt-oss and qwen3 as reasoning models ZeroGPU's /v1/models marks these four as reasoning models and accepts reasoning_effort on all of them (checked 2026-10-21), but LiteLLM rejected the param with UnsupportedParamsError because their rows had no supports_reasoning flag. glm reports reasoning_tokens under reasoning_effort, and gpt-oss and qwen3 spend hidden reasoning tokens on every reply. The gpt-oss test case now expects reasoning_effort to reach ZeroGPU * fix(zerogpu): set gpt-oss-120b context window to 128000 ZeroGPU routes gpt-oss-120b across several upstream backends, and at least one of them returns 413 past a 128000-token window, so max_input_tokens uses that as the safe floor. The backends reached on 2026-10-21 accepted up to about 131K prompt tokens, so this only gives up a little headroom on the larger ones * feat(zerogpu): call ZeroGPU's Responses API natively ZeroGPU serves /v1/responses for all nine catalog models, so LiteLLM no longer bridges Responses requests through chat completions. Native calls return ZeroGPU's reasoning output items for glm, gpt-oss and qwen3, which the bridge dropped. Checked live on 2026-10-22 through litellm.responses for every model, non-streaming and streaming, plus a gpt-oss tool call. LFM2.5-1.2B-Instruct streams correctly now that ZeroGPU's text_generation lane emits Responses SSE * fix(zerogpu): stop inferring the zerogpu provider from api_base alone Listing https://api.zerogpu.ai/v1 in openai_compatible_endpoints rerouted bare model names on that api_base from openai to zerogpu. A bare gpt-5.6-luna call with verbosity used to return 200 and started raising UnsupportedParamsError. Drop the entry and its api_base-alone test, so zerogpu is only selected through the zerogpu/ prefix * fix(zerogpu): price gpt-6-luna, llama-3.2-3b, LFM2.5-1.2B-Thinking and qwen3.6 A zerogpu/* wildcard served these four chat models with no cost-map row, so their calls logged $0 spend. Add rows at ZeroGPU's catalog prices from https://api-dashboard.zerogpu.ai/api/models, with context windows and output caps from ZeroGPU's models table. LFM2.5-1.2B-Thinking has no cached price, like LFM2.5-1.2B-Instruct. gpt-6-luna is flagged supports_reasoning because LiteLLM already routes gpt-6 names as reasoning models. Checked live on 2026-10-28, all four now report nonzero response_cost * fix(zerogpu): drop the retired max_tokens field from the zerogpu rows Upstream #48471 retired max_tokens from the cost map in favor of max_input_tokens and max_output_tokens, and test_no_row_carries_the_retired_max_tokens_field failed on all 13 zerogpu rows. Every row already carries the same value in max_output_tokens, so nothing is lost. The zerogpu limits test now checks that each row's output limit fits its context window * chore(zerogpu): regenerate provider endpoint support for zerogpu Upstream #48619 generates provider_endpoints_support.json from the cost map and checks it in CI, and the hand-written zerogpu entry no longer matched. Regenerated with scripts/generate_provider_endpoints_support.py. Only the zerogpu entry changed, and it now matches what the generator emits for other JSON-registry providers such as openinfer * fix(zerogpu): flag tools, json_schema and prompt caching on gpt-6-luna ZeroGPU serves function calling, tool_choice, json_schema response_format and prompt caching on gpt-6-luna, so set supports_function_calling, supports_tool_choice, supports_response_schema and supports_prompt_caching on its row in both cost maps. Checked live on 2026-10-28, a weather tool call came back with the right arguments, json_schema output validated, and repeated prompts reported 3260 of 3263 tokens cached * fix(zerogpu): set qwen3.6-35b-a3b context to 262144 and flag its tools and reasoning ZeroGPU accepts up to 262,144 input tokens on qwen3.6-35b-a3b and serves tools, tool_choice, json_schema response_format and reasoning_effort on it. Raise max_input_tokens from 16384 and set supports_function_calling, supports_tool_choice, supports_response_schema and supports_reasoning in both cost maps. Before this, LiteLLM rejected tools and reasoning_effort for the model. Checked live on 2026-10-28, a 209,905-token prompt succeeded, an oversized one was refused with a 262144-token limit, and tool calls, json_schema output and reasoning_effort all worked * fix(zerogpu): flag vision on gpt-6-luna and qwen3.6-35b-a3b ZeroGPU reads image inputs on gpt-6-luna and qwen3.6-35b-a3b, so set supports_vision on both rows in both cost maps. Checked live on 2026-10-28, both models named the animal in an image_url photo, and the image showed up in their prompt token counts --- README.md | 1 + .../src/exception_mapping_utils/original.rs | 1 + litellm/__init__.py | 5 + litellm/constants.py | 1 + litellm/llms/openai_like/providers.json | 6 + ...odel_prices_and_context_window_backup.json | 174 +++++++++++++ .../provider_endpoints_support_backup.json | 26 ++ .../provider_create_fields.json | 28 ++ litellm/types/utils.py | 1 + model_prices_and_context_window.json | 174 +++++++++++++ provider_endpoints_support.json | 26 ++ .../llms/openai_like/test_zerogpu_provider.py | 239 ++++++++++++++++++ .../public_endpoints/test_public_endpoints.py | 20 ++ .../test_model_prices_and_context_window.py | 17 ++ .../components/provider_info_helpers.test.tsx | 5 + .../src/components/provider_info_helpers.tsx | 3 + 16 files changed, 727 insertions(+) create mode 100644 tests/unit/llms/openai_like/test_zerogpu_provider.py diff --git a/README.md b/README.md index 66b6acb7a71..a7311882bd0 100644 --- a/README.md +++ b/README.md @@ -427,6 +427,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th | [Watsonx Text (`watsonx_text`)](https://docs.litellm.ai/docs/providers/watsonx) | ✅ | ✅ | ✅ | | | | | | | | | [xAI (`xai`)](https://docs.litellm.ai/docs/providers/xai) | ✅ | ✅ | ✅ | | | | | | | | | [Xinference (`xinference`)](https://docs.litellm.ai/docs/providers/xinference) | | | | ✅ | | | | | | | +| [ZeroGPU (`zerogpu`)](https://docs.litellm.ai/docs/providers/zerogpu) | ✅ | ✅ | ✅ | | | | | | | | [**Read the Docs**](https://docs.litellm.ai/docs/) diff --git a/litellm-rust/crates/core-utils/src/exception_mapping_utils/original.rs b/litellm-rust/crates/core-utils/src/exception_mapping_utils/original.rs index 6517b093bbd..615d530af19 100644 --- a/litellm-rust/crates/core-utils/src/exception_mapping_utils/original.rs +++ b/litellm-rust/crates/core-utils/src/exception_mapping_utils/original.rs @@ -113,6 +113,7 @@ const OPENAI_COMPATIBLE_PROVIDERS: &[&str] = &[ "meta", "cognition", "scx-ai", + "zerogpu", ]; impl ExceptionFamily { diff --git a/litellm/__init__.py b/litellm/__init__.py index fc8f6b57254..8dfe5adca1c 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -681,6 +681,7 @@ moonshot_models: Set = set() publicai_models: Set = set() darkbloom_models: Set = set() openinfer_models: Final[Set[str]] = set() +zerogpu_models: Final[Set[str]] = set() v0_models: Set = set() morph_models: Set = set() opencode_go_models: Final[set[str]] = set() # mutable-ok: populated when the shared model catalog loads @@ -969,6 +970,8 @@ def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, obje darkbloom_models.add(key) elif value.get("litellm_provider") == "openinfer": openinfer_models.add(key) + elif value.get("litellm_provider") == "zerogpu": + zerogpu_models.add(key) elif value.get("litellm_provider") == "v0": v0_models.add(key) elif value.get("litellm_provider") == "morph": @@ -1142,6 +1145,7 @@ model_list = list( | publicai_models | darkbloom_models | openinfer_models + | zerogpu_models | v0_models | morph_models | opencode_go_models @@ -1261,6 +1265,7 @@ def _build_models_by_provider() -> dict: "publicai": publicai_models, "darkbloom": darkbloom_models, "openinfer": openinfer_models, + "zerogpu": zerogpu_models, "v0": v0_models, "morph": morph_models, "opencode_go": opencode_go_models, diff --git a/litellm/constants.py b/litellm/constants.py index 2a20e812e68..2b083b03766 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1096,6 +1096,7 @@ openai_compatible_providers: Final[list[str]] = [ "llamafile", "lm_studio", "galadriel", + "zerogpu", "github_copilot", # GitHub Copilot Chat API "chatgpt", # ChatGPT subscription API "novita", diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index 0034ed2b9d4..2b71707239b 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -58,6 +58,12 @@ "max_completion_tokens": "max_tokens" } }, + "zerogpu": { + "base_url": "https://api.zerogpu.ai/v1", + "api_key_env": "ZEROGPU_API_KEY", + "api_base_env": "ZEROGPU_API_BASE", + "supported_endpoints": ["/v1/chat/completions", "/v1/responses"] + }, "chutes": { "base_url": "https://llm.chutes.ai/v1/", "api_key_env": "CHUTES_API_KEY", diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 825de5ac3ee..4e022e1a4b2 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -66163,6 +66163,180 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "zerogpu/deepseek-v4.1-flash": { + "cache_read_input_token_cost": 3e-09, + "input_cost_per_token": 3e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/glm-5.3-flash": { + "cache_read_input_token_cost": 1.5e-08, + "input_cost_per_token": 1e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 1048576, + "max_output_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 3.5e-07, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/gpt-4.1-mini": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 4e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 1047576, + "max_output_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.6e-06, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/gpt-5.4-nano": { + "cache_read_input_token_cost": 2e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.25e-06, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/gpt-5.6-luna": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/gpt-6-luna": { + "cache_read_input_token_cost": 1e-08, + "input_cost_per_token": 1e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "zerogpu/gpt-oss-120b": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 128000, + "max_output_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/LFM2.5-1.2B-Instruct": { + "input_cost_per_token": 2e-08, + "litellm_provider": "zerogpu", + "max_input_tokens": 32768, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 5e-08, + "source": "https://docs.zerogpu.ai/docs/model-catalog" + }, + "zerogpu/LFM2.5-1.2B-Thinking": { + "input_cost_per_token": 2e-08, + "litellm_provider": "zerogpu", + "max_input_tokens": 32768, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 5e-08, + "source": "https://docs.zerogpu.ai/docs/model-catalog" + }, + "zerogpu/llama-3.1-8b-instruct-fast": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 131072, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.8e-07, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/llama-3.2-3b-instruct-fast": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_token": 5e-08, + "litellm_provider": "zerogpu", + "max_input_tokens": 131072, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 3e-07, + "source": "https://docs.zerogpu.ai/docs/model-catalog" + }, + "zerogpu/qwen3-30b-a3b-fp8": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 1e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 32768, + "max_output_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 4.5e-07, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/qwen3.6-35b-a3b": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 262144, + "max_output_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "together_ai/Qwen/Qwen3.8-Flash": { "input_cost_per_token": 1.5e-07, "litellm_provider": "together_ai", diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index 1ab03fb3fdc..90395421109 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -4437,6 +4437,32 @@ "generateContent": true } }, + "zerogpu": { + "display_name": "ZeroGPU (`zerogpu`)", + "url": "https://docs.litellm.ai/docs/providers/zerogpu", + "endpoints": { + "chat_completions": true, + "responses": true, + "text_completion": true, + "embeddings": false, + "image_generations": false, + "image_edits": false, + "audio_transcriptions": false, + "audio_speech": false, + "rerank": false, + "ocr": false, + "search": false, + "moderations": false, + "realtime": false, + "video_generations": false, + "batches": false, + "messages": true, + "interactions": true, + "systemone": false, + "a2a": true, + "generateContent": true + } + }, "ragflow": { "display_name": "RAGFlow (`ragflow`)", "url": "https://docs.litellm.ai/docs/providers/ragflow", diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index 0dd96c856f8..f2d88702470 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -4145,6 +4145,34 @@ ], "default_model_placeholder": "gpt-3.5-turbo" }, + { + "provider": "ZeroGPU", + "provider_display_name": "ZeroGPU", + "litellm_provider": "zerogpu", + "credential_fields": [ + { + "key": "api_base", + "label": "API Base", + "placeholder": "https://api.zerogpu.ai/v1", + "tooltip": null, + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "api_key", + "label": "API Key", + "placeholder": null, + "tooltip": null, + "required": true, + "field_type": "password", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "zerogpu/gpt-oss-120b" + }, { "provider": "CURSOR", "provider_display_name": "Cursor", diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 4ce35bf86c1..b21bf0e14d4 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4588,6 +4588,7 @@ class LlmProviders(str, Enum): PRISM = "prism" OPENINFER = "openinfer" INFERENCEAPIS = "inferenceapis" + ZEROGPU = "zerogpu" DARKBLOOM = "darkbloom" META = "meta" SAIL = "sail" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 825de5ac3ee..4e022e1a4b2 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -66163,6 +66163,180 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "zerogpu/deepseek-v4.1-flash": { + "cache_read_input_token_cost": 3e-09, + "input_cost_per_token": 3e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/glm-5.3-flash": { + "cache_read_input_token_cost": 1.5e-08, + "input_cost_per_token": 1e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 1048576, + "max_output_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 3.5e-07, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/gpt-4.1-mini": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 4e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 1047576, + "max_output_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.6e-06, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/gpt-5.4-nano": { + "cache_read_input_token_cost": 2e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.25e-06, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/gpt-5.6-luna": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/gpt-6-luna": { + "cache_read_input_token_cost": 1e-08, + "input_cost_per_token": 1e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "zerogpu/gpt-oss-120b": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 128000, + "max_output_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/LFM2.5-1.2B-Instruct": { + "input_cost_per_token": 2e-08, + "litellm_provider": "zerogpu", + "max_input_tokens": 32768, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 5e-08, + "source": "https://docs.zerogpu.ai/docs/model-catalog" + }, + "zerogpu/LFM2.5-1.2B-Thinking": { + "input_cost_per_token": 2e-08, + "litellm_provider": "zerogpu", + "max_input_tokens": 32768, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 5e-08, + "source": "https://docs.zerogpu.ai/docs/model-catalog" + }, + "zerogpu/llama-3.1-8b-instruct-fast": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 131072, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.8e-07, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/llama-3.2-3b-instruct-fast": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_token": 5e-08, + "litellm_provider": "zerogpu", + "max_input_tokens": 131072, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 3e-07, + "source": "https://docs.zerogpu.ai/docs/model-catalog" + }, + "zerogpu/qwen3-30b-a3b-fp8": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 1e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 32768, + "max_output_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 4.5e-07, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/qwen3.6-35b-a3b": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 262144, + "max_output_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "together_ai/Qwen/Qwen3.8-Flash": { "input_cost_per_token": 1.5e-07, "litellm_provider": "together_ai", diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 1ab03fb3fdc..90395421109 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -4437,6 +4437,32 @@ "generateContent": true } }, + "zerogpu": { + "display_name": "ZeroGPU (`zerogpu`)", + "url": "https://docs.litellm.ai/docs/providers/zerogpu", + "endpoints": { + "chat_completions": true, + "responses": true, + "text_completion": true, + "embeddings": false, + "image_generations": false, + "image_edits": false, + "audio_transcriptions": false, + "audio_speech": false, + "rerank": false, + "ocr": false, + "search": false, + "moderations": false, + "realtime": false, + "video_generations": false, + "batches": false, + "messages": true, + "interactions": true, + "systemone": false, + "a2a": true, + "generateContent": true + } + }, "ragflow": { "display_name": "RAGFlow (`ragflow`)", "url": "https://docs.litellm.ai/docs/providers/ragflow", diff --git a/tests/unit/llms/openai_like/test_zerogpu_provider.py b/tests/unit/llms/openai_like/test_zerogpu_provider.py new file mode 100644 index 00000000000..1021237c335 --- /dev/null +++ b/tests/unit/llms/openai_like/test_zerogpu_provider.py @@ -0,0 +1,239 @@ +import json +from collections.abc import Mapping +from typing import Final + +import httpx +import pytest +import respx +from openai import OpenAI + +import litellm +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.proxy.auth.model_checks import get_provider_models + +ZEROGPU_API_BASE: Final = "https://api.zerogpu.ai/v1" + +ZEROGPU_MODELS: Final = ( + "deepseek-v4.1-flash", + "glm-5.3-flash", + "gpt-4.1-mini", + "gpt-5.4-nano", + "gpt-5.6-luna", + "gpt-6-luna", + "gpt-oss-120b", + "LFM2.5-1.2B-Instruct", + "LFM2.5-1.2B-Thinking", + "llama-3.1-8b-instruct-fast", + "llama-3.2-3b-instruct-fast", + "qwen3-30b-a3b-fp8", + "qwen3.6-35b-a3b", +) + + +WEATHER_TOOL: Final = { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + }, +} + + +class _RecordingZeroGPU: + def __init__(self, usage: Mapping[str, object]) -> None: + self.usage: Final = usage + self.requests: tuple[httpx.Request, ...] = () + + def respond(self, request: httpx.Request) -> httpx.Response: + self.requests = (*self.requests, request) + return httpx.Response( + 200, + json={ + "id": "chatcmpl-zerogpu", + "object": "chat.completion", + "created": 1234567890, + "model": json.loads(request.content)["model"], + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": dict(self.usage), + }, + ) + + def client(self) -> OpenAI: + return OpenAI( + api_key="zerogpu-test-key", + base_url=ZEROGPU_API_BASE, + http_client=httpx.Client(transport=httpx.MockTransport(self.respond)), + ) + + +def test_prefixed_model_reaches_zerogpu_chat_completions_without_the_prefix(): + zerogpu: Final = _RecordingZeroGPU({"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}) + + with zerogpu.client() as client: + litellm.completion( + model="zerogpu/gpt-oss-120b", + messages=[{"role": "user", "content": "hello"}], + client=client, + ) + + assert len(zerogpu.requests) == 1 + assert str(zerogpu.requests[0].url) == f"{ZEROGPU_API_BASE}/chat/completions" + assert json.loads(zerogpu.requests[0].content) == { + "model": "gpt-oss-120b", + "messages": [{"role": "user", "content": "hello"}], + } + + +def test_responses_request_reaches_zerogpu_responses_natively(): + with respx.mock() as upstream: + route: Final = upstream.post(f"{ZEROGPU_API_BASE}/responses").respond( + 200, + json={ + "id": "resp_zerogpu", + "object": "response", + "created_at": 1_792_000_000, + "model": "gpt-oss-120b", + "status": "completed", + "output": [ + { + "id": "msg_zerogpu", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + "usage": {"input_tokens": 3, "output_tokens": 1, "total_tokens": 4}, + }, + ) + response: Final = litellm.responses(model="zerogpu/gpt-oss-120b", input="hello", api_key="zerogpu-test-key") + + body: Final = json.loads(route.calls.last.request.content) + assert body["model"] == "gpt-oss-120b" + assert body["input"] == "hello" + assert response.output[0].content[0].text == "hi" + + +def test_anthropic_messages_request_is_bridged_to_zerogpu_chat_completions(): + zerogpu: Final = _RecordingZeroGPU({"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}) + + with zerogpu.client() as client: + response: Final = litellm.anthropic.messages.create( + model="zerogpu/gpt-oss-120b", + messages=[{"role": "user", "content": "hello"}], + max_tokens=10, + client=client, + ) + + assert [str(request.url) for request in zerogpu.requests] == [f"{ZEROGPU_API_BASE}/chat/completions"] + assert json.loads(zerogpu.requests[0].content) == { + "model": "gpt-oss-120b", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 10, + } + assert response["content"] == [{"type": "text", "text": "hi"}] + + +@pytest.mark.parametrize("model", ZEROGPU_MODELS) +def test_zerogpu_completion_is_charged_at_the_cost_map_rates(model: str): + usage: Final = { + "prompt_tokens": 1000, + "completion_tokens": 200, + "total_tokens": 1200, + "prompt_tokens_details": {"cached_tokens": 400}, + } + zerogpu: Final = _RecordingZeroGPU(usage) + + with zerogpu.client() as client: + response: Final = litellm.completion( + model=f"zerogpu/{model}", + messages=[{"role": "user", "content": "hello"}], + client=client, + ) + + rates: Final = litellm.model_cost[f"zerogpu/{model}"] + cached_rate: Final = rates.get("cache_read_input_token_cost", rates["input_cost_per_token"]) + expected_cost: Final = ( + 600 * rates["input_cost_per_token"] + 400 * cached_rate + 200 * rates["output_cost_per_token"] + ) + assert response._hidden_params["response_cost"] == pytest.approx(expected_cost) + assert expected_cost > 0 + + +@pytest.mark.parametrize("model", ZEROGPU_MODELS) +def test_trim_messages_fits_zerogpu_history_into_the_model_context_window(model: str): + budget: Final = 20 + trim_ratio: Final = budget / litellm.model_cost[f"zerogpu/{model}"]["max_input_tokens"] + messages: Final = [{"role": "user", "content": "lorem ipsum " * 300}] + + trimmed: Final = litellm.utils.trim_messages(messages, model=f"zerogpu/{model}", trim_ratio=trim_ratio) + + trimmed_tokens: Final = litellm.token_counter(model=f"zerogpu/{model}", messages=trimmed) + assert trimmed_tokens <= budget, trimmed + + +@pytest.mark.parametrize( + ("model", "forwards_tools", "forwards_reasoning_effort"), + [ + ("gpt-5.6-luna", True, True), + ("gpt-oss-120b", True, True), + ("qwen3.6-35b-a3b", True, True), + ("LFM2.5-1.2B-Instruct", False, False), + ], +) +def test_tools_and_reasoning_effort_reach_zerogpu_only_for_models_that_support_them( + model: str, forwards_tools: bool, forwards_reasoning_effort: bool +): + zerogpu: Final = _RecordingZeroGPU({"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}) + + with zerogpu.client() as client: + litellm.completion( + model=f"zerogpu/{model}", + messages=[{"role": "user", "content": "weather in Paris?"}], + tools=[WEATHER_TOOL], + tool_choice="auto", + reasoning_effort="high", + drop_params=True, + client=client, + ) + + tools: Final = {"tools": [WEATHER_TOOL], "tool_choice": "auto"} if forwards_tools else {} + reasoning: Final = {"reasoning_effort": "high"} if forwards_reasoning_effort else {} + assert json.loads(zerogpu.requests[0].content) == { + "model": model, + "messages": [{"role": "user", "content": "weather in Paris?"}], + **tools, + **reasoning, + } + + +@pytest.mark.parametrize("model", ZEROGPU_MODELS) +def test_model_info_reports_the_reasoning_support_litellm_routes_with(model: str): + model_info: Final = litellm.get_model_info(f"zerogpu/{model}") + + assert bool(model_info.get("supports_reasoning")) == litellm.utils.supports_reasoning(f"zerogpu/{model}") + + +def test_zerogpu_key_and_base_come_from_zerogpu_env_vars(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("ZEROGPU_API_KEY", "zerogpu-env-key") + monkeypatch.delenv("ZEROGPU_API_BASE", raising=False) + + assert get_llm_provider(model="zerogpu/gpt-oss-120b") == ( + "gpt-oss-120b", + "zerogpu", + "zerogpu-env-key", + ZEROGPU_API_BASE, + ) + + monkeypatch.setenv("ZEROGPU_API_BASE", "https://zerogpu.internal.example/v1") + + assert get_llm_provider(model="zerogpu/gpt-oss-120b")[3] == "https://zerogpu.internal.example/v1" + + +def test_zerogpu_wildcard_lists_every_zerogpu_catalog_model(): + catalog_models: Final = sorted( + key for key, row in litellm.model_cost.items() if row.get("litellm_provider") == "zerogpu" + ) + + assert catalog_models + assert sorted(get_provider_models("zerogpu") or []) == catalog_models diff --git a/tests/unit/proxy/public_endpoints/test_public_endpoints.py b/tests/unit/proxy/public_endpoints/test_public_endpoints.py index f60c2254b41..d55ce25d878 100644 --- a/tests/unit/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/unit/proxy/public_endpoints/test_public_endpoints.py @@ -363,6 +363,26 @@ def test_cognition_provider_fields(): assert fields_by_key["api_base"]["required"] is False +def test_zerogpu_add_model_form_defaults_resolve_to_a_priced_zerogpu_model(): + import litellm + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + app_instance: Final = FastAPI() + app_instance.include_router(router) + providers: Final = TestClient(app_instance).get("/public/providers/fields").json() + + zerogpu: Final = next(p for p in providers if p["litellm_provider"] == "zerogpu") + fields_by_key: Final = {f["key"]: f for f in zerogpu["credential_fields"]} + placeholder: Final = zerogpu["default_model_placeholder"] + _, provider, _, api_base = get_llm_provider(model=placeholder) + + assert provider == "zerogpu" + assert api_base == fields_by_key["api_base"]["placeholder"] + assert litellm.model_cost[placeholder]["litellm_provider"] == "zerogpu" + assert fields_by_key["api_key"]["required"] is True + assert fields_by_key["api_key"]["field_type"] == "password" + + def test_qwen_mainland_provider_fields_carry_the_qianwen_brand(): app_instance = FastAPI() app_instance.include_router(router) diff --git a/tests/unit/test_model_prices_and_context_window.py b/tests/unit/test_model_prices_and_context_window.py index bbfe304429d..2aceed493a7 100644 --- a/tests/unit/test_model_prices_and_context_window.py +++ b/tests/unit/test_model_prices_and_context_window.py @@ -115,6 +115,23 @@ def test_github_copilot_backup_rows_match_the_main_catalog() -> None: assert _github_copilot_catalog_rows(PRICES_PATH) == _github_copilot_catalog_rows(BACKUP_PRICES_PATH) +class _TokenLimits(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + max_input_tokens: int + max_output_tokens: int + + +@pytest.mark.parametrize("path", (PRICES_PATH, BACKUP_PRICES_PATH), ids=("main", "backup")) +def test_zerogpu_rows_have_an_output_limit_that_fits_the_context_window(path: Path) -> None: + prices: Final[dict[str, JsonValue]] = TypeAdapter(dict[str, JsonValue]).validate_json(path.read_bytes()) + rows: Final = MappingProxyType( + {key: _TokenLimits.model_validate(value) for key, value in prices.items() if key.startswith("zerogpu/")} + ) + assert rows, f"No ZeroGPU rows found in {path}" + assert tuple(key for key, row in rows.items() if row.max_output_tokens > row.max_input_tokens) == () + + def test_github_copilot_rows_resolve_through_get_model_info() -> None: rows: Final[Mapping[str, _CopilotEndpointRow]] = _copilot_endpoint_rows(PRICES_PATH) assert { diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx index 0ebd2cdff17..83e80911cff 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx @@ -224,6 +224,7 @@ describe("provider_info_helpers", () => { Providers.Sail, Providers.WANDB, Providers.ZAI, + Providers.ZeroGPU, ]; const logolessProviders = Object.values(Providers).filter((provider) => !providerLogoMap[provider]); expect([...logolessProviders].sort()).toEqual([...knownLogolessProviders].sort()); @@ -357,6 +358,10 @@ describe("provider_info_helpers", () => { expect(getPlaceholder(Providers.Cognition)).toBe("cognition/swe-1.7"); }); + it("should return zerogpu/gpt-oss-120b placeholder for ZeroGPU provider", () => { + expect(getPlaceholder("ZeroGPU")).toBe("zerogpu/gpt-oss-120b"); + }); + it("should return a chatgpt/ placeholder for the CHATGPT dropdown key", () => { expect(getPlaceholder("CHATGPT")).toBe("chatgpt/gpt-5.4"); }); diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index 2541aa45050..bfb578b218a 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -197,6 +197,7 @@ export enum Providers { xAI = "xAI", XINFERENCE = "Xinference", ZAI = "Z.AI (Zhipu AI)", + ZeroGPU = "ZeroGPU", } export const provider_map: Record = { @@ -325,6 +326,7 @@ export const provider_map: Record = { xAI: "xai", XINFERENCE: "xinference", ZAI: "zai", + ZeroGPU: "zerogpu", }; const standaloneSubproviderSlugs = new Set(["bedrock_mantle"]); @@ -497,6 +499,7 @@ const providerPlaceholderMap: Partial> = { [Providers.Voyage]: "voyage/", [Providers.WATSONX]: "watsonx/ibm/granite-3-3-8b-instruct", [Providers.ZAI]: "zai/glm-4.5", + [Providers.ZeroGPU]: "zerogpu/gpt-oss-120b", }; export const getPlaceholder = (selectedProvider: string): string => { From 1bd9a0780ab710ad266ebae27a819172bb489100 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 28 Oct 2026 12:51:24 -0700 Subject: [PATCH 080/128] refactor(ui): use the shadcn Kbd primitive for keyboard hints (#48885) Replaces the hand-rolled kbd and code key chips across the dashboard with the shared Kbd component from components/ui/kbd, so every shortcut hint renders the same way --- .../guardrails/_components/GuardrailTestPanel.tsx | 5 ++--- .../mcp-servers/_components/mcp_connect.tsx | 9 ++++----- .../components/complianceUI/ComplianceUI.tsx | 6 ++---- .../policies/_components/ai_suggestion_modal.tsx | 4 ++-- .../components/CommandPalette/CommandPalette.tsx | 10 ++++------ .../CommandPalette/CommandPaletteHint.tsx | 5 ++--- .../CommandPalette/CommandPaletteTrigger.tsx | 5 ++--- .../src/components/liteadmin/LiteAdmin.tsx | 3 ++- .../src/components/shared/ShortcutHints.tsx | 11 ++--------- .../src/components/shared/search/SearchBox.tsx | 13 +++---------- 10 files changed, 25 insertions(+), 46 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx index d19374f48ca..183d68da261 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx @@ -2,6 +2,7 @@ import React, { useState } from "react"; import { Copy, Info } from "lucide-react"; import { toast } from "@/lib/toast"; import { Button } from "@/components/ui/button"; +import { Kbd } from "@/components/ui/kbd"; import { Textarea } from "@/components/ui/textarea"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; @@ -165,9 +166,7 @@ export function GuardrailTestPanel({ />

- Press Enter to - submit • Shift+Enter{" "} - for new line + Press Enter to submit • Shift+Enter for new line Characters: {inputText.length}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_connect.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_connect.tsx index 81a3994070e..d609c9f2527 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_connect.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_connect.tsx @@ -5,6 +5,7 @@ import React, { useId, useState } from "react"; import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert"; import { Button } from "@/components/ui/button"; import { Card, CardContent } from "@/components/ui/card"; +import { Kbd } from "@/components/ui/kbd"; import { Label } from "@/components/ui/label"; import { Switch } from "@/components/ui/switch"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; @@ -369,8 +370,7 @@ const MCPConnect: React.FC = ({ currentServerAccessGroups = []
- Use the keyboard shortcut ⇧+⌘+J (Mac) or{" "} - Ctrl+Shift+J (Windows/Linux) + Use the keyboard shortcut ⇧+⌘+J (Mac) or Ctrl+Shift+J (Windows/Linux) @@ -380,9 +380,8 @@ const MCPConnect: React.FC = ({ currentServerAccessGroups = [] - Copy the JSON configuration below and paste it into Cursor, then save with{" "} - Cmd+S or{" "} - Ctrl+S + Copy the JSON configuration below and paste it into Cursor, then save with Cmd+S or{" "} + Ctrl+S } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx index db576ed03a5..ec4ac9a8b88 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx @@ -12,6 +12,7 @@ import { getGuardrailsList, testPoliciesAndGuardrails } from "@/components/netwo import PolicySelector, { getPolicyOptionEntries } from "@/components/policies/PolicySelector"; import { Policy } from "@/components/policies/types"; import { makeOpenAIChatCompletionRequest } from "@/components/llm_calls/chat_completion"; +import { Kbd } from "@/components/ui/kbd"; import { AlertTriangle, BarChart3, @@ -1306,10 +1307,7 @@ export default function ComplianceUI({ />
- Press Enter to - submit ·{" "} - Shift+Enter for new - line + Press Enter to submit · Shift+Enter for new line {quickTestInput.length}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/ai_suggestion_modal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/ai_suggestion_modal.tsx index 50ffb89ae83..c7eb8aaec48 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/ai_suggestion_modal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/ai_suggestion_modal.tsx @@ -5,6 +5,7 @@ import { Card, CardContent } from "@/components/ui/card"; import { Checkbox } from "@/components/ui/checkbox"; import { Dialog, DialogContent, DialogTitle } from "@/components/ui/dialog"; import { Input } from "@/components/ui/input"; +import { Kbd } from "@/components/ui/kbd"; import { Textarea } from "@/components/ui/textarea"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; @@ -654,8 +655,7 @@ const AiSuggestionModal: React.FC = ({ />
- Press Enter to - submit + Press Enter to submit
diff --git a/ui/litellm-dashboard/src/components/CommandPalette/CommandPalette.tsx b/ui/litellm-dashboard/src/components/CommandPalette/CommandPalette.tsx index 9b23dbebd6c..38bb987b7c4 100644 --- a/ui/litellm-dashboard/src/components/CommandPalette/CommandPalette.tsx +++ b/ui/litellm-dashboard/src/components/CommandPalette/CommandPalette.tsx @@ -7,6 +7,7 @@ import { usePathname, useRouter } from "next/navigation"; import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import { Dialog, DialogContent, DialogTitle } from "@/components/ui/dialog"; +import { Kbd } from "@/components/ui/kbd"; import { useTerminology } from "@/contexts/TerminologyContext"; import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { keyDetailHref } from "@/utils/entityLinks"; @@ -31,9 +32,6 @@ interface PaletteGroup { isLoading?: boolean; } -const shortcutChipClass = - "rounded-sm border border-border border-b-2 bg-muted px-[3px] font-mono text-muted-foreground"; - export function CommandPalette({ open, setOpen }: { open: boolean; setOpen: (open: boolean) => void }) { if (!open) return null; return ; @@ -305,13 +303,13 @@ function CommandPaletteDialog({ setOpen }: { setOpen: (open: boolean) => void })
- ↵ Open + ↵ Open - ↑↓ Navigate + ↑↓ Navigate - esc Close + esc Close
diff --git a/ui/litellm-dashboard/src/components/CommandPalette/CommandPaletteHint.tsx b/ui/litellm-dashboard/src/components/CommandPalette/CommandPaletteHint.tsx index 0a4c3fcda18..6fa362f4509 100644 --- a/ui/litellm-dashboard/src/components/CommandPalette/CommandPaletteHint.tsx +++ b/ui/litellm-dashboard/src/components/CommandPalette/CommandPaletteHint.tsx @@ -2,6 +2,7 @@ import { Search, X } from "lucide-react"; import { z } from "zod"; +import { Kbd } from "@/components/ui/kbd"; import { storageKey } from "@/lib/storage"; import type { PaletteScope } from "./utils"; import { useCommandPalette } from "./CommandPaletteProvider"; @@ -26,9 +27,7 @@ export function CommandPaletteHint({ scope, onDismiss }: CommandPaletteHintProps ); } diff --git a/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.tsx b/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.tsx index c512c3bd0af..198b90e9adf 100644 --- a/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.tsx +++ b/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.tsx @@ -29,6 +29,7 @@ import { Button } from "@/components/ui/button"; import { Card, CardContent, CardFooter, CardHeader, CardTitle } from "@/components/ui/card"; import { Dialog, DialogClose, DialogContent, DialogDescription, DialogTitle } from "@/components/ui/dialog"; import { FieldError } from "@/components/ui/field"; +import { Kbd } from "@/components/ui/kbd"; import { Skeleton } from "@/components/ui/skeleton"; import { useTerminology } from "@/contexts/TerminologyContext"; import { cn } from "@/lib/cva.config"; @@ -185,7 +186,7 @@ export default function LiteAdminTrigger() { > {ARIA_AGENT_LABEL} - ⌘J + ⌘J ); } diff --git a/ui/litellm-dashboard/src/components/shared/ShortcutHints.tsx b/ui/litellm-dashboard/src/components/shared/ShortcutHints.tsx index 6ed5b374cbe..934ba5aa9e0 100644 --- a/ui/litellm-dashboard/src/components/shared/ShortcutHints.tsx +++ b/ui/litellm-dashboard/src/components/shared/ShortcutHints.tsx @@ -1,8 +1,9 @@ "use client"; -import { Fragment, type ReactNode } from "react"; +import { Fragment } from "react"; import { useHotkeysContext } from "react-hotkeys-hook"; +import { Kbd } from "@/components/ui/kbd"; import { cn } from "@/lib/cva.config"; type BoundShortcut = ReturnType["hotkeys"][number]; @@ -59,14 +60,6 @@ export function shortcutHints(shortcuts: readonly BoundShortcut[]): readonly Sho >((hints, [key, shortcut]) => mergeAdjacent(hints, key, shortcut.description ?? ""), []); } -function Kbd({ children }: { children: ReactNode }) { - return ( - - {children} - - ); -} - /** The shortcuts active right now, read from what is bound rather than typed by hand. */ export function ShortcutHints({ className }: { className?: string }) { const { hotkeys } = useHotkeysContext(); diff --git a/ui/litellm-dashboard/src/components/shared/search/SearchBox.tsx b/ui/litellm-dashboard/src/components/shared/search/SearchBox.tsx index 4e45f8902f0..ae6d5a6ec09 100644 --- a/ui/litellm-dashboard/src/components/shared/search/SearchBox.tsx +++ b/ui/litellm-dashboard/src/components/shared/search/SearchBox.tsx @@ -8,9 +8,10 @@ import { Check, Copy, CornerDownLeft, Loader2, type LucideIcon, Search } from "l import { Schema } from "prosemirror-model"; import { EditorState, Plugin, TextSelection, type Transaction } from "prosemirror-state"; import { Decoration, DecorationSet, type EditorView } from "prosemirror-view"; -import { type ComponentProps, createContext, type ReactNode, useContext, useId, useMemo, useState } from "react"; +import { type ComponentProps, createContext, useContext, useId, useMemo, useState } from "react"; import { Button } from "@/components/ui/button"; +import { Kbd } from "@/components/ui/kbd"; import { cn } from "@/lib/cva.config"; import { type FilterOp, languageOps, parseQuery, type QueryClause, type QueryLanguage } from "./language"; @@ -326,7 +327,7 @@ function Suggestions({ className, children, ...props }: SearchBoxSuggestionsProp - + Select @@ -392,12 +393,4 @@ function OperatorHints({ ops }: { ops: readonly FilterOp[] }) { ); } -function Kbd({ children }: { children: ReactNode }) { - return ( - - {children} - - ); -} - export const SearchBox = { Root, Input, Suggestions, CopyCommand } as const; From 83ca14b3927d5bdf92f0dc6e9939558886380217 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 28 Oct 2026 19:58:04 +0000 Subject: [PATCH 081/128] fix(pricing): sync databricks model prices (#48884) * fix(pricing): sync databricks model prices * chore: regenerate provider endpoint support for databricks image rows --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 714 ++++++++++++++++-- .../provider_endpoints_support_backup.json | 2 +- model_prices_and_context_window.json | 714 ++++++++++++++++-- provider_endpoints_support.json | 2 +- .../test_databricks_cost_calculator.py | 22 +- 5 files changed, 1330 insertions(+), 124 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 4e022e1a4b2..3e0e86012a4 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -21009,10 +21009,11 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-claude-fable-5": { - "cache_creation_input_token_cost": 1.250004e-05, + "cache_creation_input_token_cost": 1.249997e-05, + "cache_creation_input_token_cost_above_1hr": 1.999998e-05, "cache_read_input_token_cost": 1.00002e-06, - "input_cost_per_token": 1.000006e-05, - "input_dbu_cost_per_token": 0.000142858, + "input_cost_per_token": 9.99999e-06, + "input_dbu_cost_per_token": 0.00014286, "litellm_provider": "databricks", "max_input_tokens": 1000000, "max_output_tokens": 128000, @@ -21021,7 +21022,7 @@ }, "mode": "chat", "output_cost_per_token": 5.000002e-05, - "output_dbu_cost_per_token": 0.000714286, + "output_dbu_cost_per_token": 0.00071429, "prompt_cache_min_tokens": 512, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, @@ -21037,10 +21038,11 @@ "thinking_always_on": true }, "databricks/databricks-claude-fable-5-1": { - "cache_creation_input_token_cost": 1.250004e-05, - "cache_read_input_token_cost": 2.5004e-07, - "input_cost_per_token": 1.000006e-05, - "input_dbu_cost_per_token": 0.000142858, + "cache_creation_input_token_cost": 1.249997e-05, + "cache_creation_input_token_cost_above_1hr": 1.999998e-05, + "cache_read_input_token_cost": 2.4997e-07, + "input_cost_per_token": 9.99999e-06, + "input_dbu_cost_per_token": 0.00014286, "litellm_provider": "databricks", "max_input_tokens": 1000000, "max_output_tokens": 128000, @@ -21049,7 +21051,7 @@ }, "mode": "chat", "output_cost_per_token": 5.000002e-05, - "output_dbu_cost_per_token": 0.000714286, + "output_dbu_cost_per_token": 0.00071429, "prompt_cache_min_tokens": 512, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, @@ -21067,6 +21069,7 @@ }, "databricks/databricks-claude-haiku-4-5": { "cache_creation_input_token_cost": 1.24999e-06, + "cache_creation_input_token_cost_above_1hr": 1.99997e-06, "cache_read_input_token_cost": 1.0003e-07, "input_cost_per_token": 1.00002e-06, "input_dbu_cost_per_token": 1.4286e-05, @@ -21090,9 +21093,10 @@ }, "databricks/databricks-claude-opus-4": { "cache_creation_input_token_cost": 1.874999e-05, + "cache_creation_input_token_cost_above_1hr": 2.999997e-05, "cache_read_input_token_cost": 1.50003e-06, - "input_cost_per_token": 1.5000020000000002e-05, - "input_dbu_cost_per_token": 0.000214286, + "input_cost_per_token": 1.500002e-05, + "input_dbu_cost_per_token": 0.00021429, "litellm_provider": "databricks", "max_input_tokens": 200000, "max_output_tokens": 32000, @@ -21100,8 +21104,8 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 7.500003000000001e-05, - "output_dbu_cost_per_token": 0.001071429, + "output_cost_per_token": 7.50001e-05, + "output_dbu_cost_per_token": 0.0010714, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -21113,9 +21117,10 @@ }, "databricks/databricks-claude-opus-4-1": { "cache_creation_input_token_cost": 1.874999e-05, + "cache_creation_input_token_cost_above_1hr": 2.999997e-05, "cache_read_input_token_cost": 1.50003e-06, - "input_cost_per_token": 1.5000020000000002e-05, - "input_dbu_cost_per_token": 0.000214286, + "input_cost_per_token": 1.500002e-05, + "input_dbu_cost_per_token": 0.00021429, "litellm_provider": "databricks", "max_input_tokens": 200000, "max_output_tokens": 32000, @@ -21123,8 +21128,8 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 7.500003000000001e-05, - "output_dbu_cost_per_token": 0.001071429, + "output_cost_per_token": 7.50001e-05, + "output_dbu_cost_per_token": 0.0010714, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -21136,6 +21141,7 @@ }, "databricks/databricks-claude-opus-4-5": { "cache_creation_input_token_cost": 6.25002e-06, + "cache_creation_input_token_cost_above_1hr": 9.99999e-06, "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, @@ -21146,8 +21152,8 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 2.5000010000000002e-05, - "output_dbu_cost_per_token": 0.000357143, + "output_cost_per_token": 2.500001e-05, + "output_dbu_cost_per_token": 0.00035714, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -21160,6 +21166,7 @@ }, "databricks/databricks-claude-opus-4-6": { "cache_creation_input_token_cost": 6.25002e-06, + "cache_creation_input_token_cost_above_1hr": 9.99999e-06, "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, @@ -21170,8 +21177,8 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 2.5000010000000002e-05, - "output_dbu_cost_per_token": 0.000357143, + "output_cost_per_token": 2.500001e-05, + "output_dbu_cost_per_token": 0.00035714, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": false, "supports_function_calling": true, @@ -21184,6 +21191,7 @@ }, "databricks/databricks-claude-opus-4-7": { "cache_creation_input_token_cost": 6.25002e-06, + "cache_creation_input_token_cost_above_1hr": 9.99999e-06, "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, @@ -21195,7 +21203,7 @@ }, "mode": "chat", "output_cost_per_token": 2.500001e-05, - "output_dbu_cost_per_token": 0.000357143, + "output_dbu_cost_per_token": 0.00035714, "prompt_cache_min_tokens": 2048, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, @@ -21210,6 +21218,7 @@ }, "databricks/databricks-claude-opus-4-8": { "cache_creation_input_token_cost": 6.25002e-06, + "cache_creation_input_token_cost_above_1hr": 9.99999e-06, "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, @@ -21221,7 +21230,7 @@ }, "mode": "chat", "output_cost_per_token": 2.500001e-05, - "output_dbu_cost_per_token": 0.000357143, + "output_dbu_cost_per_token": 0.00035714, "prompt_cache_min_tokens": 1024, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, @@ -21237,6 +21246,7 @@ }, "databricks/databricks-claude-opus-5": { "cache_creation_input_token_cost": 6.25002e-06, + "cache_creation_input_token_cost_above_1hr": 9.99999e-06, "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, @@ -21248,7 +21258,7 @@ }, "mode": "chat", "output_cost_per_token": 2.500001e-05, - "output_dbu_cost_per_token": 0.000357143, + "output_dbu_cost_per_token": 0.00035714, "prompt_cache_min_tokens": 512, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, @@ -21344,8 +21354,9 @@ }, "databricks/databricks-claude-sonnet-4-5": { "cache_creation_input_token_cost": 3.74997e-06, + "cache_creation_input_token_cost_above_1hr": 5.99998e-06, "cache_read_input_token_cost": 3.0002e-07, - "input_cost_per_token": 2.9999900000000002e-06, + "input_cost_per_token": 2.99999e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", "max_input_tokens": 200000, @@ -21354,8 +21365,8 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 1.5000020000000002e-05, - "output_dbu_cost_per_token": 0.000214286, + "output_cost_per_token": 1.500002e-05, + "output_dbu_cost_per_token": 0.00021429, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -21367,8 +21378,9 @@ }, "databricks/databricks-claude-sonnet-4-6": { "cache_creation_input_token_cost": 3.74997e-06, + "cache_creation_input_token_cost_above_1hr": 5.99998e-06, "cache_read_input_token_cost": 3.0002e-07, - "input_cost_per_token": 2.9999900000000002e-06, + "input_cost_per_token": 2.99999e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", "max_input_tokens": 1000000, @@ -21377,8 +21389,8 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 1.5000020000000002e-05, - "output_dbu_cost_per_token": 0.000214286, + "output_cost_per_token": 1.500002e-05, + "output_dbu_cost_per_token": 0.00021429, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": false, "supports_function_calling": true, @@ -21390,19 +21402,20 @@ "prompt_cache_min_tokens": 1024 }, "databricks/databricks-claude-sonnet-5": { - "cache_creation_input_token_cost": 3.74997e-06, - "cache_read_input_token_cost": 3.0002e-07, - "input_cost_per_token": 2.99999e-06, - "input_dbu_cost_per_token": 4.2857e-05, + "cache_creation_input_token_cost": 2.49998e-06, + "cache_creation_input_token_cost_above_1hr": 4.00001e-06, + "cache_read_input_token_cost": 1.9999e-07, + "input_cost_per_token": 1.99997e-06, + "input_dbu_cost_per_token": 2.8571e-05, "litellm_provider": "databricks", "max_input_tokens": 1000000, "max_output_tokens": 128000, "metadata": { - "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields. Introductory launch rates of 28.571 input / 142.857 output / 35.714 cache write / 2.857 cache read DBU run through 2026-08-31; the standard rates are listed here because entries carry no expiry date." + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", - "output_cost_per_token": 1.500002e-05, - "output_dbu_cost_per_token": 0.000214286, + "output_cost_per_token": 9.99999e-06, + "output_dbu_cost_per_token": 0.00014286, "prompt_cache_min_tokens": 1024, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, @@ -21546,7 +21559,10 @@ }, "databricks/databricks-gemini-3-1-flash-lite": { "cache_creation_input_token_cost": 3.1248e-07, + "cache_read_input_audio_token_cost": 6.251e-08, "cache_read_input_token_cost": 3.122e-08, + "input_cost_per_audio_token": 6.2503e-07, + "input_cost_per_image_token": 3.1248e-07, "input_cost_per_token": 3.1248e-07, "input_dbu_cost_per_token": 4.464e-06, "litellm_provider": "databricks", @@ -21576,13 +21592,19 @@ "supports_video_input": true }, "databricks/databricks-gemini-3-1-flash-image": { + "cache_creation_input_token_cost": 6.2503e-07, + "cache_read_input_token_cost": 6.2503e-07, + "input_cost_per_image_token": 6.2503e-07, + "input_cost_per_token": 6.2503e-07, "litellm_provider": "databricks", "max_input_tokens": 131072, "max_output_tokens": 32768, "metadata": { - "notes": "Databricks DBU rates not yet published for this model; endpoint metadata only." + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", + "output_cost_per_image_token": 7.50001e-05, + "output_cost_per_token": 3.74997e-06, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -21596,13 +21618,19 @@ "supports_vision": true }, "databricks/databricks-gemini-3-pro-image": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_read_input_token_cost": 2.49998e-06, + "input_cost_per_image_token": 2.49998e-06, + "input_cost_per_token": 2.49998e-06, "litellm_provider": "databricks", "max_input_tokens": 65536, "max_output_tokens": 32768, "metadata": { - "notes": "Databricks DBU rates not yet published for this model; endpoint metadata only." + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", + "output_cost_per_image_token": 0.0001500002, + "output_cost_per_token": 1.500002e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -21618,7 +21646,9 @@ "databricks/databricks-gemini-3-1-pro": { "cache_creation_input_token_cost": 2.49998e-06, "cache_read_input_token_cost": 2.4997e-07, + "cache_read_input_token_cost_above_200k_tokens": 5.0001e-07, "input_cost_per_token": 2.49998e-06, + "input_cost_per_token_above_200k_tokens": 5.00003e-06, "input_dbu_cost_per_token": 3.5714e-05, "litellm_provider": "databricks", "max_input_tokens": 1048576, @@ -21628,12 +21658,9 @@ }, "mode": "chat", "output_cost_per_token": 1.5000020000000002e-05, + "output_cost_per_token_above_200k_tokens": 2.25e-05, "output_dbu_cost_per_token": 0.000214286, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_tool_choice": true, - "supports_reasoning": true, "supported_modalities": [ "text", "image", @@ -21643,13 +21670,20 @@ "supported_output_modalities": [ "text" ], - "supports_vision": true, "supports_audio_input": true, - "supports_video_input": true + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true }, "databricks/databricks-gemini-3-flash": { "cache_creation_input_token_cost": 6.2503e-07, + "cache_read_input_audio_token_cost": 1.2502e-07, "cache_read_input_token_cost": 6.251e-08, + "input_cost_per_audio_token": 1.24999e-06, + "input_cost_per_image_token": 6.2503e-07, "input_cost_per_token": 6.2503e-07, "input_dbu_cost_per_token": 8.929e-06, "litellm_provider": "databricks", @@ -21698,13 +21732,17 @@ "supports_tool_choice": true }, "databricks/databricks-gemini-3-8-flash": { + "cache_creation_input_token_cost": 7.4998e-07, + "cache_read_input_token_cost": 7.497e-08, + "input_cost_per_token": 7.4998e-07, "litellm_provider": "databricks", "max_input_tokens": 1048576, "max_output_tokens": 65536, "metadata": { - "notes": "Databricks DBU rates not yet published for this model; endpoint metadata only." + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", + "output_cost_per_token": 3.74997e-06, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -21722,13 +21760,17 @@ "supports_vision": true }, "databricks/databricks-gemini-3-7-flash": { + "cache_creation_input_token_cost": 7.4998e-07, + "cache_read_input_token_cost": 7.497e-08, + "input_cost_per_token": 7.4998e-07, "litellm_provider": "databricks", "max_input_tokens": 1048576, "max_output_tokens": 65536, "metadata": { - "notes": "Databricks DBU rates not yet published for this model; endpoint metadata only." + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", + "output_cost_per_token": 3.74997e-06, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -21747,9 +21789,9 @@ }, "databricks/databricks-gemini-3-6-flash": { "cache_creation_input_token_cost": 1.87502e-06, - "cache_read_input_token_cost": 1.8753e-07, - "input_cost_per_token": 1.87502e-06, - "input_dbu_cost_per_token": 2.6786e-05, + "cache_read_input_token_cost": 7.497e-08, + "input_cost_per_token": 7.4998e-07, + "input_dbu_cost_per_token": 1.0714e-05, "litellm_provider": "databricks", "max_input_tokens": 1048576, "max_output_tokens": 65536, @@ -21757,8 +21799,8 @@ "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", - "output_cost_per_token": 9.37503e-06, - "output_dbu_cost_per_token": 0.000133929, + "output_cost_per_token": 3.74997e-06, + "output_dbu_cost_per_token": 5.3571e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -22019,7 +22061,9 @@ "databricks/databricks-gpt-5-4": { "cache_creation_input_token_cost": 2.49998e-06, "cache_read_input_token_cost": 2.4997e-07, + "cache_read_input_token_cost_above_272k_tokens": 5.0001e-07, "input_cost_per_token": 2.49998e-06, + "input_cost_per_token_above_272k_tokens": 5.00003e-06, "input_dbu_cost_per_token": 3.5714e-05, "litellm_provider": "databricks", "max_input_tokens": 922000, @@ -22029,6 +22073,7 @@ }, "mode": "chat", "output_cost_per_token": 1.5000020000000002e-05, + "output_cost_per_token_above_272k_tokens": 2.25e-05, "output_dbu_cost_per_token": 0.000214286, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ @@ -22102,8 +22147,11 @@ }, "databricks/databricks-gpt-5-6-sol": { "cache_creation_input_token_cost": 5.00003e-06, + "cache_creation_input_token_cost_above_272k_tokens": 9.99999e-06, "cache_read_input_token_cost": 3.9998e-07, + "cache_read_input_token_cost_above_272k_tokens": 8.0003e-07, "input_cost_per_token": 4.00001e-06, + "input_cost_per_token_above_272k_tokens": 8.00002e-06, "input_dbu_cost_per_token": 5.7143e-05, "litellm_provider": "databricks", "max_input_tokens": 922000, @@ -22113,6 +22161,7 @@ }, "mode": "chat", "output_cost_per_token": 1.999998e-05, + "output_cost_per_token_above_272k_tokens": 3e-05, "output_dbu_cost_per_token": 0.000285714, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ @@ -22130,8 +22179,11 @@ }, "databricks/databricks-gpt-5-6-terra": { "cache_creation_input_token_cost": 2.49998e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5.00003e-06, "cache_read_input_token_cost": 1.9999e-07, + "cache_read_input_token_cost_above_272k_tokens": 3.9998e-07, "input_cost_per_token": 1.99997e-06, + "input_cost_per_token_above_272k_tokens": 4.00001e-06, "input_dbu_cost_per_token": 2.8571e-05, "litellm_provider": "databricks", "max_input_tokens": 922000, @@ -22141,6 +22193,7 @@ }, "mode": "chat", "output_cost_per_token": 1.200003e-05, + "output_cost_per_token_above_272k_tokens": 1.8e-05, "output_dbu_cost_per_token": 0.000171429, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ @@ -22158,8 +22211,11 @@ }, "databricks/databricks-gpt-5-6-luna": { "cache_creation_input_token_cost": 2.4997e-07, + "cache_creation_input_token_cost_above_272k_tokens": 5.0001e-07, "cache_read_input_token_cost": 2.002e-08, + "cache_read_input_token_cost_above_272k_tokens": 3.997e-08, "input_cost_per_token": 1.9999e-07, + "input_cost_per_token_above_272k_tokens": 3.9998e-07, "input_dbu_cost_per_token": 2.857e-06, "litellm_provider": "databricks", "max_input_tokens": 922000, @@ -22169,6 +22225,7 @@ }, "mode": "chat", "output_cost_per_token": 1.20001e-06, + "output_cost_per_token_above_272k_tokens": 1.79998e-06, "output_dbu_cost_per_token": 1.7143e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ @@ -22187,7 +22244,9 @@ "databricks/databricks-gpt-5-5": { "cache_creation_input_token_cost": 5.00003e-06, "cache_read_input_token_cost": 5.0001e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.00002e-06, "input_cost_per_token": 5.00003e-06, + "input_cost_per_token_above_272k_tokens": 9.99999e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", "max_input_tokens": 922000, @@ -22197,6 +22256,7 @@ }, "mode": "responses", "output_cost_per_token": 2.999997e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, "output_dbu_cost_per_token": 0.000428571, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ @@ -22216,7 +22276,8 @@ "cache_creation_input_token_cost": 2.999997e-05, "cache_read_input_token_cost": 2.999997e-05, "input_cost_per_token": 2.999997e-05, - "input_dbu_cost_per_token": 0.000428571, + "input_cost_per_token_above_272k_tokens": 6e-05, + "input_dbu_cost_per_token": 0.00042857, "litellm_provider": "databricks", "max_input_tokens": 922000, "max_output_tokens": 128000, @@ -22224,8 +22285,9 @@ "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "responses", - "output_cost_per_token": 0.00018000003, - "output_dbu_cost_per_token": 0.002571429, + "output_cost_per_token": 0.0001800001, + "output_cost_per_token_above_272k_tokens": 0.00027, + "output_dbu_cost_per_token": 0.0025714, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -22544,6 +22606,25 @@ "output_dbu_cost_per_token": 0.0, "source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models" }, + "databricks/databricks-typesafeai-jev-1-14-0": { + "cache_creation_input_token_cost": 4.2e-08, + "cache_read_input_token_cost": 4.2e-08, + "input_cost_per_token": 4.2e-08, + "input_dbu_cost_per_token": 6e-07, + "litellm_provider": "databricks", + "metadata": { + "notes": "TypeSafe AI Jev 1.14.0, provider-hosted by TypeSafe AI and served through Unity Gateway, queried through the System One API. The proprietary pricing page lists 'Jev 1.14' at 0.600 DBU per 1M input tokens; input_cost_per_token is that times $0.070 per DBU. Output is free. No cache rates are published, so cache tokens bill at the input rate." + }, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "output_dbu_cost_per_token": 0.0, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supports_vision": true + }, "databricks/ai_decide": { "cache_creation_input_token_cost": 0.0, "cache_read_input_token_cost": 0.0, @@ -77603,6 +77684,529 @@ "supports_xhigh_reasoning_effort": true, "use_openai_responses_path": true }, + "databricks/databricks-claude-haiku-5-5": { + "cache_creation_input_token_cost": 1.2502e-07, + "cache_creation_input_token_cost_above_100k_tokens": 6.2503e-07, + "cache_creation_input_token_cost_above_1hr": 1.9999e-07, + "cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1.00002e-06, + "cache_read_input_token_cost": 1.001e-08, + "cache_read_input_token_cost_above_100k_tokens": 4.998e-08, + "input_cost_per_token": 1.0003e-07, + "input_cost_per_token_above_100k_tokens": 5.0001e-07, + "input_dbu_cost_per_token": 1.429e-06, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 5.0001e-07, + "output_cost_per_token_above_100k_tokens": 2.49998e-06, + "output_dbu_cost_per_token": 7.143e-06, + "prompt_cache_min_tokens": 512.0, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_anthropic_thinking_payload": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_mid_conversation_system": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "thinking_always_on": false + }, + "databricks/databricks-claude-sonnet-5-5": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_creation_input_token_cost_above_1hr": 4.00001e-06, + "cache_read_input_token_cost": 1.0003e-07, + "input_cost_per_token": 1.99997e-06, + "input_dbu_cost_per_token": 2.8571e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 9.99999e-06, + "output_dbu_cost_per_token": 0.00014286, + "prompt_cache_min_tokens": 512.0, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_anthropic_thinking_payload": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_mid_conversation_system": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": false, + "thinking_always_on": true + }, + "databricks/databricks-deepseek-v4-1-flash": { + "cache_creation_input_token_cost": 3.0002e-07, + "cache_read_input_token_cost": 3.003e-08, + "deprecation_date": "2026-11-19", + "input_cost_per_token": 3.0002e-07, + "input_dbu_cost_per_token": 4.286e-06, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 384000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields. The docs list only databricks-deepseek-v4-1-flash (retiring 2026-11-19) and databricks-deepseek-v4-1-flash-1017, so the pricing page's 'DeepSeek V4.1 Flash (0910)' row prices the unsuffixed endpoint." + }, + "mode": "chat", + "output_cost_per_token": 1.20001e-06, + "output_dbu_cost_per_token": 1.7143e-05, + "source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-deepseek-v4-1-flash-1017": { + "cache_creation_input_token_cost": 3.0002e-07, + "cache_read_input_token_cost": 7e-09, + "input_cost_per_token": 3.0002e-07, + "input_dbu_cost_per_token": 4.286e-06, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 384000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 1.20001e-06, + "output_dbu_cost_per_token": 1.7143e-05, + "source": "https://www.databricks.com/product/pricing/foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-gemini-2-5-flash-lite": { + "cache_creation_input_token_cost": 1.2502e-07, + "cache_read_input_audio_token_cost": 3.752e-08, + "cache_read_input_token_cost": 1.253e-08, + "input_cost_per_audio_token": 3.7499e-07, + "input_cost_per_image_token": 1.2502e-07, + "input_cost_per_token": 1.2502e-07, + "input_dbu_cost_per_token": 1.786e-06, + "litellm_provider": "databricks", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 5.0001e-07, + "output_dbu_cost_per_token": 7.143e-06, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true + }, + "databricks/databricks-gemini-3-1-flash-lite-image": { + "cache_creation_input_token_cost": 3.1248e-07, + "cache_read_input_token_cost": 3.1248e-07, + "input_cost_per_image_token": 3.1248e-07, + "input_cost_per_token": 3.1248e-07, + "input_dbu_cost_per_token": 4.464e-06, + "litellm_provider": "databricks", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_image_token": 3.749998e-05, + "output_cost_per_token": 1.87502e-06, + "output_dbu_cost_per_token": 2.6786e-05, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_vision": true + }, + "databricks/databricks-gemini-4-argon": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_read_input_token_cost": 1.2502e-07, + "input_cost_per_token": 2.49998e-06, + "input_dbu_cost_per_token": 3.5714e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 1.249997e-05, + "output_dbu_cost_per_token": 0.00017857, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true + }, + "databricks/databricks-gpt-5-4-pro": { + "cache_creation_input_token_cost": 2.999997e-05, + "cache_read_input_token_cost": 2.999997e-05, + "input_cost_per_token": 2.999997e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, + "input_dbu_cost_per_token": 0.00042857, + "litellm_provider": "databricks", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "responses", + "output_cost_per_token": 0.0001800001, + "output_cost_per_token_above_272k_tokens": 0.00027, + "output_dbu_cost_per_token": 0.0025714, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-gpt-6-1-sol": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5.00003e-06, + "cache_read_input_token_cost": 1.0003e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.9999e-07, + "input_cost_per_token": 1.99997e-06, + "input_cost_per_token_above_272k_tokens": 4.00001e-06, + "input_dbu_cost_per_token": 2.8571e-05, + "litellm_provider": "databricks", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "responses", + "output_cost_per_token": 9.99999e-06, + "output_cost_per_token_above_272k_tokens": 1.5e-05, + "output_dbu_cost_per_token": 0.00014286, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-gpt-6-1-terra": { + "cache_creation_input_token_cost": 1.24999e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.49998e-06, + "cache_read_input_token_cost": 4.998e-08, + "cache_read_input_token_cost_above_272k_tokens": 1.0003e-07, + "input_cost_per_token": 1.00002e-06, + "input_cost_per_token_above_272k_tokens": 1.99997e-06, + "input_dbu_cost_per_token": 1.4286e-05, + "litellm_provider": "databricks", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 5.00003e-06, + "output_cost_per_token_above_272k_tokens": 7.50001e-06, + "output_dbu_cost_per_token": 7.1429e-05, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-gpt-6-astra": { + "cache_creation_input_token_cost": 1.249997e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, + "cache_read_input_token_cost": 1.00002e-06, + "cache_read_input_token_cost_above_272k_tokens": 1.99997e-06, + "input_cost_per_token": 9.99999e-06, + "input_cost_per_token_above_272k_tokens": 2e-05, + "input_dbu_cost_per_token": 0.00014286, + "litellm_provider": "databricks", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 5.000002e-05, + "output_cost_per_token_above_272k_tokens": 7.50001e-05, + "output_dbu_cost_per_token": 0.00071429, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-gpt-6-luna": { + "cache_creation_input_token_cost": 1.2502e-07, + "cache_creation_input_token_cost_above_272k_tokens": 2.4997e-07, + "cache_read_input_token_cost": 1.001e-08, + "cache_read_input_token_cost_above_272k_tokens": 2.002e-08, + "input_cost_per_token": 1.0003e-07, + "input_cost_per_token_above_272k_tokens": 1.9999e-07, + "input_dbu_cost_per_token": 1.429e-06, + "litellm_provider": "databricks", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "responses", + "output_cost_per_token": 5.0001e-07, + "output_cost_per_token_above_272k_tokens": 7.497e-07, + "output_dbu_cost_per_token": 7.143e-06, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-gpt-6-sol": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5.00003e-06, + "cache_read_input_token_cost": 1.9999e-07, + "cache_read_input_token_cost_above_272k_tokens": 3.9998e-07, + "input_cost_per_token": 1.99997e-06, + "input_cost_per_token_above_272k_tokens": 4.00001e-06, + "input_dbu_cost_per_token": 2.8571e-05, + "litellm_provider": "databricks", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "responses", + "output_cost_per_token": 9.99999e-06, + "output_cost_per_token_above_272k_tokens": 1.50003e-05, + "output_dbu_cost_per_token": 0.00014286, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-gpt-image-1": { + "cache_creation_input_token_cost": 5.00003e-06, + "cache_read_input_token_cost": 1.24999e-06, + "input_cost_per_image_token": 9.99999e-06, + "input_cost_per_token": 5.00003e-06, + "input_dbu_cost_per_token": 7.1429e-05, + "litellm_provider": "databricks", + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "image_generation", + "output_cost_per_image_token": 4.000003e-05, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + }, + "databricks/databricks-gpt-image-1-5": { + "cache_creation_input_token_cost": 5.00003e-06, + "cache_read_input_token_cost": 1.24999e-06, + "input_cost_per_image_token": 8.00002e-06, + "input_cost_per_token": 5.00003e-06, + "input_dbu_cost_per_token": 7.1429e-05, + "litellm_provider": "databricks", + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "image_generation", + "output_cost_per_image_token": 3.200001e-05, + "output_cost_per_token": 9.99999e-06, + "output_dbu_cost_per_token": 0.00014286, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_pdf_input": true, + "supports_vision": true + }, + "databricks/databricks-gpt-image-1-mini": { + "cache_creation_input_token_cost": 1.99997e-06, + "cache_read_input_token_cost": 1.9999e-07, + "input_cost_per_image_token": 2.49998e-06, + "input_cost_per_token": 1.99997e-06, + "input_dbu_cost_per_token": 2.8571e-05, + "litellm_provider": "databricks", + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "image_generation", + "output_cost_per_image_token": 8.00002e-06, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + }, + "databricks/databricks-gpt-image-2": { + "cache_creation_input_token_cost": 5.00003e-06, + "cache_read_input_token_cost": 1.24999e-06, + "input_cost_per_image_token": 8.00002e-06, + "input_cost_per_token": 5.00003e-06, + "input_dbu_cost_per_token": 7.1429e-05, + "litellm_provider": "databricks", + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "image_generation", + "output_cost_per_image_token": 2.999997e-05, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_pdf_input": true, + "supports_vision": true + }, + "databricks/databricks-grok-4-7": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_read_input_token_cost": 6.2503e-07, + "input_cost_per_token": 2.49998e-06, + "input_dbu_cost_per_token": 3.5714e-05, + "litellm_provider": "databricks", + "max_input_tokens": 500000, + "max_output_tokens": 500000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 7.50001e-06, + "output_dbu_cost_per_token": 0.00010714, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-kimi-k2-7": { + "cache_creation_input_token_cost": 9.4997e-07, + "cache_read_input_token_cost": 1.8998e-07, + "input_cost_per_token": 9.4997e-07, + "input_dbu_cost_per_token": 1.3571e-05, + "litellm_provider": "databricks", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 4.00001e-06, + "output_dbu_cost_per_token": 5.7143e-05, + "source": "https://www.databricks.com/product/pricing/foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "fireworks_ai/accounts/fireworks/routers/spot-deepseek-v4-flash": { "cache_read_input_token_cost": 2.1e-08, "deprecation_date": "2026-08-14", diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index 90395421109..d73ac19bae1 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -1225,7 +1225,7 @@ "responses": true, "text_completion": true, "embeddings": true, - "image_generations": false, + "image_generations": true, "image_edits": false, "audio_transcriptions": false, "audio_speech": false, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4e022e1a4b2..3e0e86012a4 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -21009,10 +21009,11 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-claude-fable-5": { - "cache_creation_input_token_cost": 1.250004e-05, + "cache_creation_input_token_cost": 1.249997e-05, + "cache_creation_input_token_cost_above_1hr": 1.999998e-05, "cache_read_input_token_cost": 1.00002e-06, - "input_cost_per_token": 1.000006e-05, - "input_dbu_cost_per_token": 0.000142858, + "input_cost_per_token": 9.99999e-06, + "input_dbu_cost_per_token": 0.00014286, "litellm_provider": "databricks", "max_input_tokens": 1000000, "max_output_tokens": 128000, @@ -21021,7 +21022,7 @@ }, "mode": "chat", "output_cost_per_token": 5.000002e-05, - "output_dbu_cost_per_token": 0.000714286, + "output_dbu_cost_per_token": 0.00071429, "prompt_cache_min_tokens": 512, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, @@ -21037,10 +21038,11 @@ "thinking_always_on": true }, "databricks/databricks-claude-fable-5-1": { - "cache_creation_input_token_cost": 1.250004e-05, - "cache_read_input_token_cost": 2.5004e-07, - "input_cost_per_token": 1.000006e-05, - "input_dbu_cost_per_token": 0.000142858, + "cache_creation_input_token_cost": 1.249997e-05, + "cache_creation_input_token_cost_above_1hr": 1.999998e-05, + "cache_read_input_token_cost": 2.4997e-07, + "input_cost_per_token": 9.99999e-06, + "input_dbu_cost_per_token": 0.00014286, "litellm_provider": "databricks", "max_input_tokens": 1000000, "max_output_tokens": 128000, @@ -21049,7 +21051,7 @@ }, "mode": "chat", "output_cost_per_token": 5.000002e-05, - "output_dbu_cost_per_token": 0.000714286, + "output_dbu_cost_per_token": 0.00071429, "prompt_cache_min_tokens": 512, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, @@ -21067,6 +21069,7 @@ }, "databricks/databricks-claude-haiku-4-5": { "cache_creation_input_token_cost": 1.24999e-06, + "cache_creation_input_token_cost_above_1hr": 1.99997e-06, "cache_read_input_token_cost": 1.0003e-07, "input_cost_per_token": 1.00002e-06, "input_dbu_cost_per_token": 1.4286e-05, @@ -21090,9 +21093,10 @@ }, "databricks/databricks-claude-opus-4": { "cache_creation_input_token_cost": 1.874999e-05, + "cache_creation_input_token_cost_above_1hr": 2.999997e-05, "cache_read_input_token_cost": 1.50003e-06, - "input_cost_per_token": 1.5000020000000002e-05, - "input_dbu_cost_per_token": 0.000214286, + "input_cost_per_token": 1.500002e-05, + "input_dbu_cost_per_token": 0.00021429, "litellm_provider": "databricks", "max_input_tokens": 200000, "max_output_tokens": 32000, @@ -21100,8 +21104,8 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 7.500003000000001e-05, - "output_dbu_cost_per_token": 0.001071429, + "output_cost_per_token": 7.50001e-05, + "output_dbu_cost_per_token": 0.0010714, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -21113,9 +21117,10 @@ }, "databricks/databricks-claude-opus-4-1": { "cache_creation_input_token_cost": 1.874999e-05, + "cache_creation_input_token_cost_above_1hr": 2.999997e-05, "cache_read_input_token_cost": 1.50003e-06, - "input_cost_per_token": 1.5000020000000002e-05, - "input_dbu_cost_per_token": 0.000214286, + "input_cost_per_token": 1.500002e-05, + "input_dbu_cost_per_token": 0.00021429, "litellm_provider": "databricks", "max_input_tokens": 200000, "max_output_tokens": 32000, @@ -21123,8 +21128,8 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 7.500003000000001e-05, - "output_dbu_cost_per_token": 0.001071429, + "output_cost_per_token": 7.50001e-05, + "output_dbu_cost_per_token": 0.0010714, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -21136,6 +21141,7 @@ }, "databricks/databricks-claude-opus-4-5": { "cache_creation_input_token_cost": 6.25002e-06, + "cache_creation_input_token_cost_above_1hr": 9.99999e-06, "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, @@ -21146,8 +21152,8 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 2.5000010000000002e-05, - "output_dbu_cost_per_token": 0.000357143, + "output_cost_per_token": 2.500001e-05, + "output_dbu_cost_per_token": 0.00035714, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -21160,6 +21166,7 @@ }, "databricks/databricks-claude-opus-4-6": { "cache_creation_input_token_cost": 6.25002e-06, + "cache_creation_input_token_cost_above_1hr": 9.99999e-06, "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, @@ -21170,8 +21177,8 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 2.5000010000000002e-05, - "output_dbu_cost_per_token": 0.000357143, + "output_cost_per_token": 2.500001e-05, + "output_dbu_cost_per_token": 0.00035714, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": false, "supports_function_calling": true, @@ -21184,6 +21191,7 @@ }, "databricks/databricks-claude-opus-4-7": { "cache_creation_input_token_cost": 6.25002e-06, + "cache_creation_input_token_cost_above_1hr": 9.99999e-06, "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, @@ -21195,7 +21203,7 @@ }, "mode": "chat", "output_cost_per_token": 2.500001e-05, - "output_dbu_cost_per_token": 0.000357143, + "output_dbu_cost_per_token": 0.00035714, "prompt_cache_min_tokens": 2048, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, @@ -21210,6 +21218,7 @@ }, "databricks/databricks-claude-opus-4-8": { "cache_creation_input_token_cost": 6.25002e-06, + "cache_creation_input_token_cost_above_1hr": 9.99999e-06, "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, @@ -21221,7 +21230,7 @@ }, "mode": "chat", "output_cost_per_token": 2.500001e-05, - "output_dbu_cost_per_token": 0.000357143, + "output_dbu_cost_per_token": 0.00035714, "prompt_cache_min_tokens": 1024, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, @@ -21237,6 +21246,7 @@ }, "databricks/databricks-claude-opus-5": { "cache_creation_input_token_cost": 6.25002e-06, + "cache_creation_input_token_cost_above_1hr": 9.99999e-06, "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, @@ -21248,7 +21258,7 @@ }, "mode": "chat", "output_cost_per_token": 2.500001e-05, - "output_dbu_cost_per_token": 0.000357143, + "output_dbu_cost_per_token": 0.00035714, "prompt_cache_min_tokens": 512, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, @@ -21344,8 +21354,9 @@ }, "databricks/databricks-claude-sonnet-4-5": { "cache_creation_input_token_cost": 3.74997e-06, + "cache_creation_input_token_cost_above_1hr": 5.99998e-06, "cache_read_input_token_cost": 3.0002e-07, - "input_cost_per_token": 2.9999900000000002e-06, + "input_cost_per_token": 2.99999e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", "max_input_tokens": 200000, @@ -21354,8 +21365,8 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 1.5000020000000002e-05, - "output_dbu_cost_per_token": 0.000214286, + "output_cost_per_token": 1.500002e-05, + "output_dbu_cost_per_token": 0.00021429, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -21367,8 +21378,9 @@ }, "databricks/databricks-claude-sonnet-4-6": { "cache_creation_input_token_cost": 3.74997e-06, + "cache_creation_input_token_cost_above_1hr": 5.99998e-06, "cache_read_input_token_cost": 3.0002e-07, - "input_cost_per_token": 2.9999900000000002e-06, + "input_cost_per_token": 2.99999e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", "max_input_tokens": 1000000, @@ -21377,8 +21389,8 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 1.5000020000000002e-05, - "output_dbu_cost_per_token": 0.000214286, + "output_cost_per_token": 1.500002e-05, + "output_dbu_cost_per_token": 0.00021429, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": false, "supports_function_calling": true, @@ -21390,19 +21402,20 @@ "prompt_cache_min_tokens": 1024 }, "databricks/databricks-claude-sonnet-5": { - "cache_creation_input_token_cost": 3.74997e-06, - "cache_read_input_token_cost": 3.0002e-07, - "input_cost_per_token": 2.99999e-06, - "input_dbu_cost_per_token": 4.2857e-05, + "cache_creation_input_token_cost": 2.49998e-06, + "cache_creation_input_token_cost_above_1hr": 4.00001e-06, + "cache_read_input_token_cost": 1.9999e-07, + "input_cost_per_token": 1.99997e-06, + "input_dbu_cost_per_token": 2.8571e-05, "litellm_provider": "databricks", "max_input_tokens": 1000000, "max_output_tokens": 128000, "metadata": { - "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields. Introductory launch rates of 28.571 input / 142.857 output / 35.714 cache write / 2.857 cache read DBU run through 2026-08-31; the standard rates are listed here because entries carry no expiry date." + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", - "output_cost_per_token": 1.500002e-05, - "output_dbu_cost_per_token": 0.000214286, + "output_cost_per_token": 9.99999e-06, + "output_dbu_cost_per_token": 0.00014286, "prompt_cache_min_tokens": 1024, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, @@ -21546,7 +21559,10 @@ }, "databricks/databricks-gemini-3-1-flash-lite": { "cache_creation_input_token_cost": 3.1248e-07, + "cache_read_input_audio_token_cost": 6.251e-08, "cache_read_input_token_cost": 3.122e-08, + "input_cost_per_audio_token": 6.2503e-07, + "input_cost_per_image_token": 3.1248e-07, "input_cost_per_token": 3.1248e-07, "input_dbu_cost_per_token": 4.464e-06, "litellm_provider": "databricks", @@ -21576,13 +21592,19 @@ "supports_video_input": true }, "databricks/databricks-gemini-3-1-flash-image": { + "cache_creation_input_token_cost": 6.2503e-07, + "cache_read_input_token_cost": 6.2503e-07, + "input_cost_per_image_token": 6.2503e-07, + "input_cost_per_token": 6.2503e-07, "litellm_provider": "databricks", "max_input_tokens": 131072, "max_output_tokens": 32768, "metadata": { - "notes": "Databricks DBU rates not yet published for this model; endpoint metadata only." + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", + "output_cost_per_image_token": 7.50001e-05, + "output_cost_per_token": 3.74997e-06, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -21596,13 +21618,19 @@ "supports_vision": true }, "databricks/databricks-gemini-3-pro-image": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_read_input_token_cost": 2.49998e-06, + "input_cost_per_image_token": 2.49998e-06, + "input_cost_per_token": 2.49998e-06, "litellm_provider": "databricks", "max_input_tokens": 65536, "max_output_tokens": 32768, "metadata": { - "notes": "Databricks DBU rates not yet published for this model; endpoint metadata only." + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", + "output_cost_per_image_token": 0.0001500002, + "output_cost_per_token": 1.500002e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -21618,7 +21646,9 @@ "databricks/databricks-gemini-3-1-pro": { "cache_creation_input_token_cost": 2.49998e-06, "cache_read_input_token_cost": 2.4997e-07, + "cache_read_input_token_cost_above_200k_tokens": 5.0001e-07, "input_cost_per_token": 2.49998e-06, + "input_cost_per_token_above_200k_tokens": 5.00003e-06, "input_dbu_cost_per_token": 3.5714e-05, "litellm_provider": "databricks", "max_input_tokens": 1048576, @@ -21628,12 +21658,9 @@ }, "mode": "chat", "output_cost_per_token": 1.5000020000000002e-05, + "output_cost_per_token_above_200k_tokens": 2.25e-05, "output_dbu_cost_per_token": 0.000214286, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_tool_choice": true, - "supports_reasoning": true, "supported_modalities": [ "text", "image", @@ -21643,13 +21670,20 @@ "supported_output_modalities": [ "text" ], - "supports_vision": true, "supports_audio_input": true, - "supports_video_input": true + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true }, "databricks/databricks-gemini-3-flash": { "cache_creation_input_token_cost": 6.2503e-07, + "cache_read_input_audio_token_cost": 1.2502e-07, "cache_read_input_token_cost": 6.251e-08, + "input_cost_per_audio_token": 1.24999e-06, + "input_cost_per_image_token": 6.2503e-07, "input_cost_per_token": 6.2503e-07, "input_dbu_cost_per_token": 8.929e-06, "litellm_provider": "databricks", @@ -21698,13 +21732,17 @@ "supports_tool_choice": true }, "databricks/databricks-gemini-3-8-flash": { + "cache_creation_input_token_cost": 7.4998e-07, + "cache_read_input_token_cost": 7.497e-08, + "input_cost_per_token": 7.4998e-07, "litellm_provider": "databricks", "max_input_tokens": 1048576, "max_output_tokens": 65536, "metadata": { - "notes": "Databricks DBU rates not yet published for this model; endpoint metadata only." + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", + "output_cost_per_token": 3.74997e-06, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -21722,13 +21760,17 @@ "supports_vision": true }, "databricks/databricks-gemini-3-7-flash": { + "cache_creation_input_token_cost": 7.4998e-07, + "cache_read_input_token_cost": 7.497e-08, + "input_cost_per_token": 7.4998e-07, "litellm_provider": "databricks", "max_input_tokens": 1048576, "max_output_tokens": 65536, "metadata": { - "notes": "Databricks DBU rates not yet published for this model; endpoint metadata only." + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", + "output_cost_per_token": 3.74997e-06, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -21747,9 +21789,9 @@ }, "databricks/databricks-gemini-3-6-flash": { "cache_creation_input_token_cost": 1.87502e-06, - "cache_read_input_token_cost": 1.8753e-07, - "input_cost_per_token": 1.87502e-06, - "input_dbu_cost_per_token": 2.6786e-05, + "cache_read_input_token_cost": 7.497e-08, + "input_cost_per_token": 7.4998e-07, + "input_dbu_cost_per_token": 1.0714e-05, "litellm_provider": "databricks", "max_input_tokens": 1048576, "max_output_tokens": 65536, @@ -21757,8 +21799,8 @@ "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", - "output_cost_per_token": 9.37503e-06, - "output_dbu_cost_per_token": 0.000133929, + "output_cost_per_token": 3.74997e-06, + "output_dbu_cost_per_token": 5.3571e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -22019,7 +22061,9 @@ "databricks/databricks-gpt-5-4": { "cache_creation_input_token_cost": 2.49998e-06, "cache_read_input_token_cost": 2.4997e-07, + "cache_read_input_token_cost_above_272k_tokens": 5.0001e-07, "input_cost_per_token": 2.49998e-06, + "input_cost_per_token_above_272k_tokens": 5.00003e-06, "input_dbu_cost_per_token": 3.5714e-05, "litellm_provider": "databricks", "max_input_tokens": 922000, @@ -22029,6 +22073,7 @@ }, "mode": "chat", "output_cost_per_token": 1.5000020000000002e-05, + "output_cost_per_token_above_272k_tokens": 2.25e-05, "output_dbu_cost_per_token": 0.000214286, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ @@ -22102,8 +22147,11 @@ }, "databricks/databricks-gpt-5-6-sol": { "cache_creation_input_token_cost": 5.00003e-06, + "cache_creation_input_token_cost_above_272k_tokens": 9.99999e-06, "cache_read_input_token_cost": 3.9998e-07, + "cache_read_input_token_cost_above_272k_tokens": 8.0003e-07, "input_cost_per_token": 4.00001e-06, + "input_cost_per_token_above_272k_tokens": 8.00002e-06, "input_dbu_cost_per_token": 5.7143e-05, "litellm_provider": "databricks", "max_input_tokens": 922000, @@ -22113,6 +22161,7 @@ }, "mode": "chat", "output_cost_per_token": 1.999998e-05, + "output_cost_per_token_above_272k_tokens": 3e-05, "output_dbu_cost_per_token": 0.000285714, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ @@ -22130,8 +22179,11 @@ }, "databricks/databricks-gpt-5-6-terra": { "cache_creation_input_token_cost": 2.49998e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5.00003e-06, "cache_read_input_token_cost": 1.9999e-07, + "cache_read_input_token_cost_above_272k_tokens": 3.9998e-07, "input_cost_per_token": 1.99997e-06, + "input_cost_per_token_above_272k_tokens": 4.00001e-06, "input_dbu_cost_per_token": 2.8571e-05, "litellm_provider": "databricks", "max_input_tokens": 922000, @@ -22141,6 +22193,7 @@ }, "mode": "chat", "output_cost_per_token": 1.200003e-05, + "output_cost_per_token_above_272k_tokens": 1.8e-05, "output_dbu_cost_per_token": 0.000171429, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ @@ -22158,8 +22211,11 @@ }, "databricks/databricks-gpt-5-6-luna": { "cache_creation_input_token_cost": 2.4997e-07, + "cache_creation_input_token_cost_above_272k_tokens": 5.0001e-07, "cache_read_input_token_cost": 2.002e-08, + "cache_read_input_token_cost_above_272k_tokens": 3.997e-08, "input_cost_per_token": 1.9999e-07, + "input_cost_per_token_above_272k_tokens": 3.9998e-07, "input_dbu_cost_per_token": 2.857e-06, "litellm_provider": "databricks", "max_input_tokens": 922000, @@ -22169,6 +22225,7 @@ }, "mode": "chat", "output_cost_per_token": 1.20001e-06, + "output_cost_per_token_above_272k_tokens": 1.79998e-06, "output_dbu_cost_per_token": 1.7143e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ @@ -22187,7 +22244,9 @@ "databricks/databricks-gpt-5-5": { "cache_creation_input_token_cost": 5.00003e-06, "cache_read_input_token_cost": 5.0001e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.00002e-06, "input_cost_per_token": 5.00003e-06, + "input_cost_per_token_above_272k_tokens": 9.99999e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", "max_input_tokens": 922000, @@ -22197,6 +22256,7 @@ }, "mode": "responses", "output_cost_per_token": 2.999997e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, "output_dbu_cost_per_token": 0.000428571, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ @@ -22216,7 +22276,8 @@ "cache_creation_input_token_cost": 2.999997e-05, "cache_read_input_token_cost": 2.999997e-05, "input_cost_per_token": 2.999997e-05, - "input_dbu_cost_per_token": 0.000428571, + "input_cost_per_token_above_272k_tokens": 6e-05, + "input_dbu_cost_per_token": 0.00042857, "litellm_provider": "databricks", "max_input_tokens": 922000, "max_output_tokens": 128000, @@ -22224,8 +22285,9 @@ "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "responses", - "output_cost_per_token": 0.00018000003, - "output_dbu_cost_per_token": 0.002571429, + "output_cost_per_token": 0.0001800001, + "output_cost_per_token_above_272k_tokens": 0.00027, + "output_dbu_cost_per_token": 0.0025714, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -22544,6 +22606,25 @@ "output_dbu_cost_per_token": 0.0, "source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models" }, + "databricks/databricks-typesafeai-jev-1-14-0": { + "cache_creation_input_token_cost": 4.2e-08, + "cache_read_input_token_cost": 4.2e-08, + "input_cost_per_token": 4.2e-08, + "input_dbu_cost_per_token": 6e-07, + "litellm_provider": "databricks", + "metadata": { + "notes": "TypeSafe AI Jev 1.14.0, provider-hosted by TypeSafe AI and served through Unity Gateway, queried through the System One API. The proprietary pricing page lists 'Jev 1.14' at 0.600 DBU per 1M input tokens; input_cost_per_token is that times $0.070 per DBU. Output is free. No cache rates are published, so cache tokens bill at the input rate." + }, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "output_dbu_cost_per_token": 0.0, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supports_vision": true + }, "databricks/ai_decide": { "cache_creation_input_token_cost": 0.0, "cache_read_input_token_cost": 0.0, @@ -77603,6 +77684,529 @@ "supports_xhigh_reasoning_effort": true, "use_openai_responses_path": true }, + "databricks/databricks-claude-haiku-5-5": { + "cache_creation_input_token_cost": 1.2502e-07, + "cache_creation_input_token_cost_above_100k_tokens": 6.2503e-07, + "cache_creation_input_token_cost_above_1hr": 1.9999e-07, + "cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1.00002e-06, + "cache_read_input_token_cost": 1.001e-08, + "cache_read_input_token_cost_above_100k_tokens": 4.998e-08, + "input_cost_per_token": 1.0003e-07, + "input_cost_per_token_above_100k_tokens": 5.0001e-07, + "input_dbu_cost_per_token": 1.429e-06, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 5.0001e-07, + "output_cost_per_token_above_100k_tokens": 2.49998e-06, + "output_dbu_cost_per_token": 7.143e-06, + "prompt_cache_min_tokens": 512.0, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_anthropic_thinking_payload": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_mid_conversation_system": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "thinking_always_on": false + }, + "databricks/databricks-claude-sonnet-5-5": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_creation_input_token_cost_above_1hr": 4.00001e-06, + "cache_read_input_token_cost": 1.0003e-07, + "input_cost_per_token": 1.99997e-06, + "input_dbu_cost_per_token": 2.8571e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 9.99999e-06, + "output_dbu_cost_per_token": 0.00014286, + "prompt_cache_min_tokens": 512.0, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_anthropic_thinking_payload": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_mid_conversation_system": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": false, + "thinking_always_on": true + }, + "databricks/databricks-deepseek-v4-1-flash": { + "cache_creation_input_token_cost": 3.0002e-07, + "cache_read_input_token_cost": 3.003e-08, + "deprecation_date": "2026-11-19", + "input_cost_per_token": 3.0002e-07, + "input_dbu_cost_per_token": 4.286e-06, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 384000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields. The docs list only databricks-deepseek-v4-1-flash (retiring 2026-11-19) and databricks-deepseek-v4-1-flash-1017, so the pricing page's 'DeepSeek V4.1 Flash (0910)' row prices the unsuffixed endpoint." + }, + "mode": "chat", + "output_cost_per_token": 1.20001e-06, + "output_dbu_cost_per_token": 1.7143e-05, + "source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-deepseek-v4-1-flash-1017": { + "cache_creation_input_token_cost": 3.0002e-07, + "cache_read_input_token_cost": 7e-09, + "input_cost_per_token": 3.0002e-07, + "input_dbu_cost_per_token": 4.286e-06, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 384000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 1.20001e-06, + "output_dbu_cost_per_token": 1.7143e-05, + "source": "https://www.databricks.com/product/pricing/foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-gemini-2-5-flash-lite": { + "cache_creation_input_token_cost": 1.2502e-07, + "cache_read_input_audio_token_cost": 3.752e-08, + "cache_read_input_token_cost": 1.253e-08, + "input_cost_per_audio_token": 3.7499e-07, + "input_cost_per_image_token": 1.2502e-07, + "input_cost_per_token": 1.2502e-07, + "input_dbu_cost_per_token": 1.786e-06, + "litellm_provider": "databricks", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 5.0001e-07, + "output_dbu_cost_per_token": 7.143e-06, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true + }, + "databricks/databricks-gemini-3-1-flash-lite-image": { + "cache_creation_input_token_cost": 3.1248e-07, + "cache_read_input_token_cost": 3.1248e-07, + "input_cost_per_image_token": 3.1248e-07, + "input_cost_per_token": 3.1248e-07, + "input_dbu_cost_per_token": 4.464e-06, + "litellm_provider": "databricks", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_image_token": 3.749998e-05, + "output_cost_per_token": 1.87502e-06, + "output_dbu_cost_per_token": 2.6786e-05, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_vision": true + }, + "databricks/databricks-gemini-4-argon": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_read_input_token_cost": 1.2502e-07, + "input_cost_per_token": 2.49998e-06, + "input_dbu_cost_per_token": 3.5714e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 1.249997e-05, + "output_dbu_cost_per_token": 0.00017857, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true + }, + "databricks/databricks-gpt-5-4-pro": { + "cache_creation_input_token_cost": 2.999997e-05, + "cache_read_input_token_cost": 2.999997e-05, + "input_cost_per_token": 2.999997e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, + "input_dbu_cost_per_token": 0.00042857, + "litellm_provider": "databricks", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "responses", + "output_cost_per_token": 0.0001800001, + "output_cost_per_token_above_272k_tokens": 0.00027, + "output_dbu_cost_per_token": 0.0025714, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-gpt-6-1-sol": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5.00003e-06, + "cache_read_input_token_cost": 1.0003e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.9999e-07, + "input_cost_per_token": 1.99997e-06, + "input_cost_per_token_above_272k_tokens": 4.00001e-06, + "input_dbu_cost_per_token": 2.8571e-05, + "litellm_provider": "databricks", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "responses", + "output_cost_per_token": 9.99999e-06, + "output_cost_per_token_above_272k_tokens": 1.5e-05, + "output_dbu_cost_per_token": 0.00014286, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-gpt-6-1-terra": { + "cache_creation_input_token_cost": 1.24999e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.49998e-06, + "cache_read_input_token_cost": 4.998e-08, + "cache_read_input_token_cost_above_272k_tokens": 1.0003e-07, + "input_cost_per_token": 1.00002e-06, + "input_cost_per_token_above_272k_tokens": 1.99997e-06, + "input_dbu_cost_per_token": 1.4286e-05, + "litellm_provider": "databricks", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 5.00003e-06, + "output_cost_per_token_above_272k_tokens": 7.50001e-06, + "output_dbu_cost_per_token": 7.1429e-05, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-gpt-6-astra": { + "cache_creation_input_token_cost": 1.249997e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, + "cache_read_input_token_cost": 1.00002e-06, + "cache_read_input_token_cost_above_272k_tokens": 1.99997e-06, + "input_cost_per_token": 9.99999e-06, + "input_cost_per_token_above_272k_tokens": 2e-05, + "input_dbu_cost_per_token": 0.00014286, + "litellm_provider": "databricks", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 5.000002e-05, + "output_cost_per_token_above_272k_tokens": 7.50001e-05, + "output_dbu_cost_per_token": 0.00071429, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-gpt-6-luna": { + "cache_creation_input_token_cost": 1.2502e-07, + "cache_creation_input_token_cost_above_272k_tokens": 2.4997e-07, + "cache_read_input_token_cost": 1.001e-08, + "cache_read_input_token_cost_above_272k_tokens": 2.002e-08, + "input_cost_per_token": 1.0003e-07, + "input_cost_per_token_above_272k_tokens": 1.9999e-07, + "input_dbu_cost_per_token": 1.429e-06, + "litellm_provider": "databricks", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "responses", + "output_cost_per_token": 5.0001e-07, + "output_cost_per_token_above_272k_tokens": 7.497e-07, + "output_dbu_cost_per_token": 7.143e-06, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-gpt-6-sol": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5.00003e-06, + "cache_read_input_token_cost": 1.9999e-07, + "cache_read_input_token_cost_above_272k_tokens": 3.9998e-07, + "input_cost_per_token": 1.99997e-06, + "input_cost_per_token_above_272k_tokens": 4.00001e-06, + "input_dbu_cost_per_token": 2.8571e-05, + "litellm_provider": "databricks", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "responses", + "output_cost_per_token": 9.99999e-06, + "output_cost_per_token_above_272k_tokens": 1.50003e-05, + "output_dbu_cost_per_token": 0.00014286, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-gpt-image-1": { + "cache_creation_input_token_cost": 5.00003e-06, + "cache_read_input_token_cost": 1.24999e-06, + "input_cost_per_image_token": 9.99999e-06, + "input_cost_per_token": 5.00003e-06, + "input_dbu_cost_per_token": 7.1429e-05, + "litellm_provider": "databricks", + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "image_generation", + "output_cost_per_image_token": 4.000003e-05, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + }, + "databricks/databricks-gpt-image-1-5": { + "cache_creation_input_token_cost": 5.00003e-06, + "cache_read_input_token_cost": 1.24999e-06, + "input_cost_per_image_token": 8.00002e-06, + "input_cost_per_token": 5.00003e-06, + "input_dbu_cost_per_token": 7.1429e-05, + "litellm_provider": "databricks", + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "image_generation", + "output_cost_per_image_token": 3.200001e-05, + "output_cost_per_token": 9.99999e-06, + "output_dbu_cost_per_token": 0.00014286, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_pdf_input": true, + "supports_vision": true + }, + "databricks/databricks-gpt-image-1-mini": { + "cache_creation_input_token_cost": 1.99997e-06, + "cache_read_input_token_cost": 1.9999e-07, + "input_cost_per_image_token": 2.49998e-06, + "input_cost_per_token": 1.99997e-06, + "input_dbu_cost_per_token": 2.8571e-05, + "litellm_provider": "databricks", + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "image_generation", + "output_cost_per_image_token": 8.00002e-06, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + }, + "databricks/databricks-gpt-image-2": { + "cache_creation_input_token_cost": 5.00003e-06, + "cache_read_input_token_cost": 1.24999e-06, + "input_cost_per_image_token": 8.00002e-06, + "input_cost_per_token": 5.00003e-06, + "input_dbu_cost_per_token": 7.1429e-05, + "litellm_provider": "databricks", + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "image_generation", + "output_cost_per_image_token": 2.999997e-05, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_pdf_input": true, + "supports_vision": true + }, + "databricks/databricks-grok-4-7": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_read_input_token_cost": 6.2503e-07, + "input_cost_per_token": 2.49998e-06, + "input_dbu_cost_per_token": 3.5714e-05, + "litellm_provider": "databricks", + "max_input_tokens": 500000, + "max_output_tokens": 500000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 7.50001e-06, + "output_dbu_cost_per_token": 0.00010714, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-kimi-k2-7": { + "cache_creation_input_token_cost": 9.4997e-07, + "cache_read_input_token_cost": 1.8998e-07, + "input_cost_per_token": 9.4997e-07, + "input_dbu_cost_per_token": 1.3571e-05, + "litellm_provider": "databricks", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 4.00001e-06, + "output_dbu_cost_per_token": 5.7143e-05, + "source": "https://www.databricks.com/product/pricing/foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "fireworks_ai/accounts/fireworks/routers/spot-deepseek-v4-flash": { "cache_read_input_token_cost": 2.1e-08, "deprecation_date": "2026-08-14", diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 90395421109..d73ac19bae1 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -1225,7 +1225,7 @@ "responses": true, "text_completion": true, "embeddings": true, - "image_generations": false, + "image_generations": true, "image_edits": false, "audio_transcriptions": false, "audio_speech": false, diff --git a/tests/unit/llms/databricks/test_databricks_cost_calculator.py b/tests/unit/llms/databricks/test_databricks_cost_calculator.py index 7120a130462..34d3806d6eb 100644 --- a/tests/unit/llms/databricks/test_databricks_cost_calculator.py +++ b/tests/unit/llms/databricks/test_databricks_cost_calculator.py @@ -33,6 +33,13 @@ PRICE_FIELDS: Final = ( "cache_read_input_token_cost", ) PUBLISHED_DBU_PER_MILLION: Final = { + "databricks/databricks-claude-haiku-5-5": ("1.429", "7.143", "1.786", "0.143"), + "databricks/databricks-claude-sonnet-5-5": ("28.571", "142.857", "35.714", "1.429"), + "databricks/databricks-gpt-6-1-sol": ("28.571", "142.857", "35.714", "1.429"), + "databricks/databricks-gpt-6-1-terra": ("14.286", "71.429", "17.857", "0.714"), + "databricks/databricks-gpt-6-astra": ("142.857", "714.286", "178.571", "14.286"), + "databricks/databricks-gpt-6-luna": ("1.429", "7.143", "1.786", "0.143"), + "databricks/databricks-gpt-6-sol": ("28.571", "142.857", "35.714", "2.857"), "databricks/databricks-claude-opus-5-5": ("57.143", "285.714", "71.429", "2.857"), "databricks/databricks-claude-fable-5-1": ("142.858", "714.286", "178.572", "3.572"), "databricks/databricks-claude-fable-5": ("142.858", "714.286", "178.572", "14.286"), @@ -43,7 +50,7 @@ PUBLISHED_DBU_PER_MILLION: Final = { "databricks/databricks-claude-opus-4-5": ("71.429", "357.143", "89.286", "7.143"), "databricks/databricks-claude-opus-4-1": ("214.286", "1071.429", "267.857", "21.429"), "databricks/databricks-claude-opus-4": ("214.286", "1071.429", "267.857", "21.429"), - "databricks/databricks-claude-sonnet-5": ("42.857", "214.286", "53.571", "4.286"), + "databricks/databricks-claude-sonnet-5": ("28.571", "142.857", "35.714", "2.857"), "databricks/databricks-claude-sonnet-4-6": ("42.857", "214.286", "53.571", "4.286"), "databricks/databricks-claude-sonnet-4-5": ("42.857", "214.286", "53.571", "4.286"), "databricks/databricks-claude-sonnet-4-1": ("42.857", "214.286", "53.571", "4.286"), @@ -197,7 +204,7 @@ def test_models_without_a_cache_discount_bill_cache_tokens_at_the_input_rate( assert prompt_cost > 8000 * info["input_cost_per_token"] -def test_every_model_without_published_cache_dbu_bills_cache_at_its_own_input_rate( +def test_every_model_without_published_cache_dbu_bills_cache_writes_at_its_own_input_rate( local_model_cost_map: None, ) -> None: without_published_rates: Final = [ @@ -210,8 +217,7 @@ def test_every_model_without_published_cache_dbu_bills_cache_at_its_own_input_ra for model in without_published_rates: info = _model_info(model) - for field in CACHE_FIELDS: - assert info[field] == pytest.approx(info["input_cost_per_token"]), (model, field) + assert info["cache_creation_input_token_cost"] == pytest.approx(info["input_cost_per_token"]), model @pytest.mark.parametrize("model", NEW_MODELS) @@ -224,14 +230,6 @@ def test_backup_price_map_matches_main(model: str) -> None: assert backup_cost[model] == main_cost[model] -def test_sonnet_5_ships_standard_rates_not_introductory(local_model_cost_map: None) -> None: - sonnet_5: Final = _model_info("databricks/databricks-claude-sonnet-5") - sonnet_4_6: Final = _model_info("databricks/databricks-claude-sonnet-4-6") - - for field in PRICE_FIELDS: - assert sonnet_5[field] == pytest.approx(sonnet_4_6[field]), field - - def test_cost_per_token_bills_the_served_priority_tier( local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch ) -> None: From 26e53cb7386bff7c953c6c5d7a1ac3fd6de1db0d Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 13:07:17 -0700 Subject: [PATCH 082/128] fix(cli): let the proxy decide when a classic lite login token has expired (#47712) * test(cli): print-token and whoami let the proxy judge a classic token past the client window * fix(cli): let the proxy decide when a classic lite login token has expired * fix(cli): treat 403 and 404 as a rejected token and judge lite up logins the same way * fix(cli): keep lite up on the client window and bound the proxy check to 3 seconds * test(cli): prove the proxy check gives up on a silent proxy and pin stale fixtures Point the timeout test at a local socket that never answers, use fixed timestamps for stale CLI token fixtures, and correct the print-token docs for tokens younger than the client window. * fix(cli): expire a CLI token with no recorded proxy and stop treating 404 as rejected A token file without base_url cannot be checked, so it keeps the client window's verdict instead of printing forever. A 404 also comes from a proxy that does not serve /v1/models, so only 401 and 403 count as a rejected key, matching the CLI's existing key check. The timeout test now runs in process. * test(cli): judge the whoami proxy check on a fixed timestamp and name the client window for lite up --- litellm/litellm_core_utils/cli_token_utils.py | 6 +- litellm/proxy/client/README.md | 4 +- litellm/proxy/client/cli/README.md | 2 +- litellm/proxy/client/cli/commands/auth.py | 44 ++++- .../proxy/client/cli/test_auth_commands.py | 169 +++++++++++++----- 5 files changed, 165 insertions(+), 60 deletions(-) diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index ffe2c4433a0..37ed4efa0ca 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -353,9 +353,9 @@ def is_cli_token_fresh( buffer_hours: float = CLI_TOKEN_FRESHNESS_BUFFER_SECONDS / 3600, ) -> bool: """Check whether a cached CLI token is still within its expiration window. - Used by `lite auth print-token` to fail fast, without a network round trip, - once the cached token is past `LITELLM_CLI_JWT_EXPIRATION_HOURS`. A `--pkce` - credential carries its own `expires_at`, which is authoritative when present.""" + Used by `lite auth print-token` to skip asking the proxy while the cached token + is younger than `LITELLM_CLI_JWT_EXPIRATION_HOURS`. A `--pkce` credential + carries its own `expires_at`, which is authoritative when present.""" from litellm.constants import CLI_JWT_EXPIRATION_HOURS expires_at: Final = ( diff --git a/litellm/proxy/client/README.md b/litellm/proxy/client/README.md index 1fff68677cc..ceec636567f 100644 --- a/litellm/proxy/client/README.md +++ b/litellm/proxy/client/README.md @@ -342,7 +342,7 @@ The CLI provides these authentication commands: - **`lite login --pkce`** - Sign in through the system browser with OAuth authorization code + PKCE; the key renews itself with a refresh token - **`lite logout`** - Clear stored authentication token (and revoke a `--pkce` refresh token on the proxy) - **`lite whoami`** - Show current authentication status -- **`lite auth print-token`** - Print the cached token (used as Claude Code's `apiKeyHelper`); renews a `--pkce` key first and fails once a classic token has expired +- **`lite auth print-token`** - Print the cached token (used as Claude Code's `apiKeyHelper`); renews a `--pkce` key first, and fails once a classic token is past 24 hours and the proxy rejects it ### Authentication Flow Steps @@ -382,7 +382,7 @@ Keychain storage needs the `keyring` package, which ships with `pip install 'lit `lite logout` clears both stores. If the keychain is locked at that moment it says so, and re-running it once the keychain is unlocked finishes the job. -The stored credential is a short-lived, per-session agent token, not a managed virtual key. It is scoped to the user and team you logged in as and inherits their models and budgets; spend is tracked against the shared team and user budgets rather than a separate per-session cap, so multiple logins or several concurrent agents all draw down the same allowance. It is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); re-run `lite login` to refresh it and pick up your latest team and user settings. `lite auth print-token` (usable as Claude Code's `apiKeyHelper`) prints it while fresh and fails once it expires -- there is no silent renewal. It is accepted on a default deployment without `EXPERIMENTAL_UI_LOGIN`, does not appear in the Keys UI, and cannot be rotated or revoked mid-session. A credential from `lite login --pkce` is the exception: it carries a refresh token, so the CLI renews the key shortly before it expires and `lite logout` revokes the refresh token on the proxy (see [Browser sign-in with PKCE](https://docs.litellm.ai/docs/proxy/cli_sso#browser-sign-in-with-pkce)). Only the holder can end a `--pkce` session early, with `lite logout`; an admin has no button for it, but every renewal re-reads the user on the proxy, so deactivating the user or removing them from the team makes the next renewal fail and the key runs out within `LITELLM_CLI_JWT_EXPIRATION_HOURS`. On a proxy with more than one worker or replica, configure Redis (`litellm_settings.cache` with Redis `cache_params`, or `general_settings.coordination_redis`) so a refresh token stays single-use and `lite logout` holds on every worker; without Redis each worker keeps its own record. For a long-lived, rotatable, Keys-UI-visible credential, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY`. +The stored credential is a short-lived, per-session agent token, not a managed virtual key. It is scoped to the user and team you logged in as and inherits their models and budgets; spend is tracked against the shared team and user budgets rather than a separate per-session cap, so multiple logins or several concurrent agents all draw down the same allowance. It is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); re-run `lite login` to refresh it and pick up your latest team and user settings. `lite auth print-token` (usable as Claude Code's `apiKeyHelper`) prints it, and once it is past 24 hours (or the client's `LITELLM_CLI_JWT_EXPIRATION_HOURS`) asks the proxy first and fails if the proxy rejects it. There is no silent renewal. It is accepted on a default deployment without `EXPERIMENTAL_UI_LOGIN`, does not appear in the Keys UI, and cannot be rotated or revoked mid-session. A credential from `lite login --pkce` is the exception: it carries a refresh token, so the CLI renews the key shortly before it expires and `lite logout` revokes the refresh token on the proxy (see [Browser sign-in with PKCE](https://docs.litellm.ai/docs/proxy/cli_sso#browser-sign-in-with-pkce)). Only the holder can end a `--pkce` session early, with `lite logout`; an admin has no button for it, but every renewal re-reads the user on the proxy, so deactivating the user or removing them from the team makes the next renewal fail and the key runs out within `LITELLM_CLI_JWT_EXPIRATION_HOURS`. On a proxy with more than one worker or replica, configure Redis (`litellm_settings.cache` with Redis `cache_params`, or `general_settings.coordination_redis`) so a refresh token stays single-use and `lite logout` holds on every worker; without Redis each worker keeps its own record. For a long-lived, rotatable, Keys-UI-visible credential, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY`. ### Usage diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index 4038ff6f582..819908f9e70 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -506,7 +506,7 @@ To pin the model, pass the agent's own model flag (for example `lite claude --mo The token minted by `lite login` is a short-lived, per-session agent credential, not a managed virtual key. It is scoped to the user and team you authenticated as, inherits that user's and team's models and budgets, and is enforced on the proxy exactly like a virtual key on the same team (guardrails, routing, logging, spend). Spend is tracked against the shared team and user budgets, so running several agents (or logging in more than once) does not hand each session its own separate budget; they all draw down the same team/user allowance. There is no separate per-session cap, so sustained agent use is not capped at a small chat-session limit. -The credential is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); run `lite login` again to refresh it, which also re-reads your latest team and user settings. It does not appear in the Keys UI and cannot be rotated or revoked mid-session. `lite auth print-token` (usable as Claude Code's `apiKeyHelper`) prints it while it's still fresh and fails once it expires -- there is no silent renewal, so a long-running session needs a fresh `lite login` once a day. `lite claude`, `lite codex`, `lite opencode`, and `lite pi` work with it on a default deployment; `EXPERIMENTAL_UI_LOGIN` is not required. `lite login --pkce` is the exception to the daily re-login: it signs in through your system browser with OAuth authorization code and PKCE and stores a refresh token next to the key, so every `lite` command and `lite auth print-token` renew the key on their own shortly before it expires, `lite whoami` shows when the current key expires, and `lite logout` revokes the refresh token on the proxy (it needs a proxy that serves `/.well-known/litellm-cli-auth`; see [Browser sign-in with PKCE](https://docs.litellm.ai/docs/proxy/cli_sso#browser-sign-in-with-pkce)). When a renewal is refused, for example after a `lite logout` run from another copy of the credential, the command prints why on stderr and, once the key has run out, tells you to run `lite login --pkce` again. Only the holder can end a `--pkce` session early, with `lite logout`; an admin has no button for it, but every renewal re-reads the user on the proxy, so deactivating the user or removing them from the team makes the next renewal fail and the key runs out within `LITELLM_CLI_JWT_EXPIRATION_HOURS`. On a proxy with more than one worker or replica, configure Redis (`litellm_settings.cache` with Redis `cache_params`, or `general_settings.coordination_redis`) so a refresh token stays single-use and `lite logout` holds on every worker; without Redis each worker keeps its own record. If you need a long-lived, rotatable key that shows up in the Keys UI, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY` instead. +The credential is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); run `lite login` again to refresh it, which also re-reads your latest team and user settings. It does not appear in the Keys UI and cannot be rotated or revoked mid-session. `lite auth print-token` (usable as Claude Code's `apiKeyHelper`) prints it without asking the proxy for its first 24 hours, or `LITELLM_CLI_JWT_EXPIRATION_HOURS` if you set it on the client. After that, `print-token` and `lite whoami` ask the proxy whether it still accepts the token, and `print-token` fails once the proxy rejects it. There is no silent renewal, so a long-running session needs a fresh `lite login` each time the proxy's token lifetime runs out. If the proxy's tokens last less than 24 hours, set the same `LITELLM_CLI_JWT_EXPIRATION_HOURS` on the client so `print-token` stops printing the token once it runs out. `lite up` still starts a new login once the token is older than that same client window. It copies the token into Claude Code's settings as a static key that it never renews, so a fresh login gives Claude Code the proxy's full token lifetime instead of whatever an old token has left. `lite claude`, `lite codex`, `lite opencode`, and `lite pi` work with it on a default deployment; `EXPERIMENTAL_UI_LOGIN` is not required. `lite login --pkce` is the exception to the re-login: it signs in through your system browser with OAuth authorization code and PKCE and stores a refresh token next to the key, so every `lite` command and `lite auth print-token` renew the key on their own shortly before it expires, `lite whoami` shows when the current key expires, and `lite logout` revokes the refresh token on the proxy (it needs a proxy that serves `/.well-known/litellm-cli-auth`; see [Browser sign-in with PKCE](https://docs.litellm.ai/docs/proxy/cli_sso#browser-sign-in-with-pkce)). When a renewal is refused, for example after a `lite logout` run from another copy of the credential, the command prints why on stderr and, once the key has run out, tells you to run `lite login --pkce` again. Only the holder can end a `--pkce` session early, with `lite logout`; an admin has no button for it, but every renewal re-reads the user on the proxy, so deactivating the user or removing them from the team makes the next renewal fail and the key runs out within `LITELLM_CLI_JWT_EXPIRATION_HOURS`. On a proxy with more than one worker or replica, configure Redis (`litellm_settings.cache` with Redis `cache_params`, or `general_settings.coordination_redis`) so a refresh token stays single-use and `lite logout` holds on every worker; without Redis each worker keeps its own record. If you need a long-lived, rotatable key that shows up in the Keys UI, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY` instead. ### Route Every Claude Code Session Through the Proxy diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index b5302d1f561..e38f54adc4f 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -13,7 +13,6 @@ from rich.console import Console from rich.table import Table from typing_extensions import NotRequired, ReadOnly, TypedDict, assert_never -from litellm.constants import CLI_JWT_EXPIRATION_HOURS from litellm.litellm_core_utils.cli_keyring import ( DISABLE_KEYRING_ENV_VAR, SYSTEM_KEYRING, @@ -143,6 +142,9 @@ _TeamMapping: Final = TypeVar("_TeamMapping", bound=Mapping[str, object]) _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_PROXY_KEY_CHECK_TIMEOUT_SECONDS: Final = 3 +_PROXY_REJECTED_KEY_STATUSES: Final = frozenset({401, 403}) + KEYRING_INSTALL_HINT: Final = "pip install 'litellm[cli]'" KEYRING_ENABLE_HINT: Final = "keyring --enable (or unset PYTHON_KEYRING_BACKEND)" @@ -1033,10 +1035,12 @@ def print_token(ctx: click.Context): Designed to be used as Claude Code's `apiKeyHelper` (https://docs.claude.com/en/docs/claude-code/settings): stdout must - contain only the token, so all diagnostics go to stderr. The token - expires after `LITELLM_CLI_JWT_EXPIRATION_HOURS` (default 24h); a - `lite login --pkce` token renews itself here first, and once a token - has expired for good, run the same `lite login` command again. + contain only the token, so all diagnostics go to stderr. A token younger + than `LITELLM_CLI_JWT_EXPIRATION_HOURS` (default 24h) is printed without + asking the proxy. An older one is checked against the proxy and refused + only if the proxy rejects it. A `lite login --pkce` + token renews itself here first, and once a token has expired for good, + run the same `lite login` command again. """ vault: Final = context_secret_vault(ctx) token_data: Final = load_token(vault=vault) @@ -1055,7 +1059,7 @@ def print_token(ctx: click.Context): sys.exit(1) renews: Final = "refresh_token" in token_data - if not is_cli_token_fresh(token_data) and not renews: + if not renews and token_expired(token_data): click.echo("Token expired. Run 'lite login' again.", err=True) sys.exit(1) @@ -1111,8 +1115,32 @@ def whoami(ctx: click.Context): expires_at: Final = token_data.get("expires_at") if isinstance(expires_at, (int, float)): click.echo(_key_expiry_line(expires_at, renews="refresh_token" in token_data)) - elif age_hours > CLI_JWT_EXPIRATION_HOURS: - click.echo(f"Warning: Token is more than {CLI_JWT_EXPIRATION_HOURS} hours old and may have expired.") + elif token_expired(token_data): + click.echo("Token expired. Run 'lite login' again.") + + +def token_expired(token_data: Mapping[str, object]) -> bool: + if is_cli_token_fresh(token_data): + return False + if isinstance(token_data.get("expires_at"), (int, float)): + return True + return _proxy_rejects_key(token_data) + + +def _proxy_rejects_key(token_data: Mapping[str, object]) -> bool: + key: Final = token_data.get("key") + base_url: Final = token_data.get("base_url") + if not isinstance(key, str) or not isinstance(base_url, str) or not base_url: + return True + try: + response: Final = requests.get( + f"{base_url}/v1/models", + headers={"Authorization": f"Bearer {key}"}, + timeout=_PROXY_KEY_CHECK_TIMEOUT_SECONDS, + ) + except requests.RequestException: + return False + return response.status_code in _PROXY_REJECTED_KEY_STATUSES def _key_expiry_line(expires_at: float, renews: bool) -> str: diff --git a/tests/unit/proxy/client/cli/test_auth_commands.py b/tests/unit/proxy/client/cli/test_auth_commands.py index cb1fcba9df7..5d5ebcf5dfd 100644 --- a/tests/unit/proxy/client/cli/test_auth_commands.py +++ b/tests/unit/proxy/client/cli/test_auth_commands.py @@ -3,13 +3,15 @@ import os import stat import time from pathlib import Path +from typing import Final from unittest.mock import Mock, patch import pytest +import requests import responses from click.testing import CliRunner +from responses import matchers -from litellm.constants import CLI_JWT_EXPIRATION_HOURS from litellm.litellm_core_utils.cli_keyring import ( DISABLE_KEYRING_ENV_VAR, KeyringDisabled, @@ -56,6 +58,20 @@ def _write_token_file(home: Path, *, key: str | None) -> None: _write_home_json(home, "token.json", payload if key is None else {**payload, "key": key}) +STALE_TOKEN_BASE_URL: Final = "https://stale-token.example.com" +STALE_TOKEN_TIMESTAMP: Final = 1_700_000_000.0 + + +def _register_models_answer(key: str, answer: int | Exception) -> None: + models_url: Final = f"{STALE_TOKEN_BASE_URL}/v1/models" + authorized: Final = [matchers.header_matcher({"Authorization": f"Bearer {key}"})] + if isinstance(answer, Exception): + responses.get(models_url, body=answer, match=authorized) + else: + responses.get(models_url, json={"data": []}, status=answer, match=authorized) + responses.get(models_url, status=401) + + def _secret_blob(base_url: str, key: str) -> str: return json.dumps({"base_url": base_url, "key": key, "jwt_token": ""}) @@ -696,22 +712,32 @@ class TestWhoamiCommand: assert "Not authenticated" in result.output assert "Run 'lite login'" in result.output - def test_whoami_old_token(self): - """Test whoami with old token showing warning""" - token_data = CliTokenRecord( + @pytest.mark.parametrize( + ("proxy_status", "expected_last_line_start"), + [ + (200, "Token age: "), + (401, "Token expired. Run 'lite login' again."), + ], + ) + @responses.activate + def test_whoami_lets_the_proxy_judge_a_token_older_than_the_client_window( + self, proxy_status: int, expected_last_line_start: str + ): + _register_models_answer("sk-live", proxy_status) + token_data: Final = CliTokenRecord( + base_url=STALE_TOKEN_BASE_URL, user_email="test@example.com", - user_id="test-user-123", - user_role="admin", key="sk-live", - timestamp=time.time() - (25 * 3600), + timestamp=STALE_TOKEN_TIMESTAMP, ) with patch("litellm.proxy.client.cli.commands.auth.load_cli_token", return_value=token_data): - result = self.runner.invoke(whoami) + result: Final = self.runner.invoke(whoami) - assert result.exit_code == 0 - assert "Authenticated" in result.output - assert "Warning: Token is more than 24 hours old" in result.output + assert result.exit_code == 0 + assert result.output.splitlines()[-1].startswith(expected_last_line_start) + assert "Warning" not in result.output + assert len(responses.calls) == 1 def test_whoami_missing_fields(self): """Test whoami with token missing some fields""" @@ -1024,51 +1050,93 @@ class TestPrintTokenCommand: assert result.exit_code == 0 assert result.output.strip() == "sk-matches" + @responses.activate def test_fresh_cached_key_printed_without_network_call(self): """A recently-issued key should be printed straight from cache -- no refresh call on every single invocation (apiKeyHelper gets called frequently).""" - with ( - patch( - "litellm.proxy.client.cli.commands.auth.load_cli_token", - return_value=CliTokenRecord( - base_url="http://localhost:4000", - key="sk-cached-fresh", - timestamp=time.time(), - ), + with patch( + "litellm.proxy.client.cli.commands.auth.load_cli_token", + return_value=CliTokenRecord( + base_url="http://localhost:4000", + key="sk-cached-fresh", + timestamp=time.time(), ), - patch("requests.post") as mock_post, ): - result = self.runner.invoke(print_token, obj={}) + result: Final = self.runner.invoke(print_token, obj={}) assert result.exit_code == 0 assert result.output.strip() == "sk-cached-fresh" - mock_post.assert_not_called() + assert len(responses.calls) == 0 - def test_stale_key_fails_fast_without_network_call(self): - """There is no silent refresh: an expired cached key must fail - loudly (stderr, nonzero exit) telling the user to `lite login` - again, rather than making a network call or printing a dead key - that will just 401 Claude Code.""" - old_timestamp = time.time() - (CLI_JWT_EXPIRATION_HOURS + 1) * 3600 + @pytest.mark.parametrize( + ("proxy_answer", "expected_exit_code", "expected_stdout", "expected_stderr"), + [ + (200, 0, "sk-stale-key\n", ""), + (401, 1, "", "Token expired. Run 'lite login' again.\n"), + (403, 1, "", "Token expired. Run 'lite login' again.\n"), + (404, 0, "sk-stale-key\n", ""), + (500, 0, "sk-stale-key\n", ""), + (requests.ConnectionError("proxy unreachable"), 0, "sk-stale-key\n", ""), + ], + ) + @responses.activate + def test_key_older_than_the_client_window_is_judged_by_the_proxy( + self, + proxy_answer: int | Exception, + expected_exit_code: int, + expected_stdout: str, + expected_stderr: str, + ): + _register_models_answer("sk-stale-key", proxy_answer) - with ( - patch( - "litellm.proxy.client.cli.commands.auth.load_cli_token", - return_value=CliTokenRecord( - base_url="http://localhost:4000", - key="sk-stale-key", - timestamp=old_timestamp, - ), + with patch( + "litellm.proxy.client.cli.commands.auth.load_cli_token", + return_value=CliTokenRecord( + base_url=STALE_TOKEN_BASE_URL, + key="sk-stale-key", + timestamp=STALE_TOKEN_TIMESTAMP, ), - patch("requests.post") as mock_post, ): - result = self.runner.invoke(print_token, obj={}) + result: Final = self.runner.invoke(print_token, obj={}) - assert result.exit_code != 0 - assert "sk-stale-key" not in result.output - assert "lite login" in result.output - mock_post.assert_not_called() + assert result.exit_code == expected_exit_code + assert result.stdout == expected_stdout + assert result.stderr == expected_stderr + assert len(responses.calls) == 1 + + @responses.activate + def test_proxy_check_waits_only_a_few_seconds_for_a_silent_proxy(self): + _register_models_answer("sk-stale-key", requests.ReadTimeout("proxy sent nothing")) + + with patch( + "litellm.proxy.client.cli.commands.auth.load_cli_token", + return_value=CliTokenRecord( + base_url=STALE_TOKEN_BASE_URL, + key="sk-stale-key", + timestamp=STALE_TOKEN_TIMESTAMP, + ), + ): + result: Final = self.runner.invoke(print_token, obj={}) + + timeout: Final = responses.calls[0].request.req_kwargs["timeout"] + assert result.exit_code == 0 + assert result.stdout == "sk-stale-key\n" + assert isinstance(timeout, int | float) + assert 0 < timeout <= 5 + + @responses.activate + def test_key_without_a_recorded_proxy_still_expires_after_the_client_window(self): + with patch( + "litellm.proxy.client.cli.commands.auth.load_cli_token", + return_value=CliTokenRecord(key="sk-old-format", timestamp=STALE_TOKEN_TIMESTAMP), + ): + result: Final = self.runner.invoke(print_token, obj={}) + + assert result.exit_code == 1 + assert result.stdout == "" + assert result.stderr == "Token expired. Run 'lite login' again.\n" + assert len(responses.calls) == 0 class TestPrintTokenWithConfigFile: @@ -2019,18 +2087,27 @@ class TestPkcePrintToken: assert sum(len(session.posts) for session in _FakeSession.instances) == 1 assert save.call_count == 1 - def test_print_token_for_an_expired_classic_token_makes_no_request(self): + @responses.activate + def test_print_token_trusts_a_recorded_expiry_without_asking_the_proxy(self): + _register_models_answer("sk-classic", 200) with ( patch( "litellm.proxy.client.cli.commands.auth.load_token", - return_value={"key": "sk-classic", "timestamp": time.time() - (CLI_JWT_EXPIRATION_HOURS + 1) * 3600}, + return_value={ + "key": "sk-classic", + "base_url": STALE_TOKEN_BASE_URL, + "timestamp": STALE_TOKEN_TIMESTAMP, + "expires_at": STALE_TOKEN_TIMESTAMP + 3600, + }, ), patch("litellm.proxy.client.cli.commands.auth.requests.Session", _FakeSession), ): - result = self.runner.invoke(print_token, obj={}) + result: Final = self.runner.invoke(print_token, obj={}) assert result.exit_code == 1 - assert "Token expired" in result.output + assert result.stdout == "" + assert result.stderr == "Token expired. Run 'lite login' again.\n" + assert len(responses.calls) == 0 assert _FakeSession.instances == [] From 99ac0d3e39e19276947f58176b382cc2daea3fed Mon Sep 17 00:00:00 2001 From: Temo <72650579+TemoSulava@users.noreply.github.com> Date: Thu, 29 Oct 2026 00:14:44 +0400 Subject: [PATCH 083/128] fix(aiohttp_openai): forward service_tier in transform_response (#47263) * fix(aiohttp_openai): forward service_tier in transform_response * chore(aiohttp_openai): mark service_tier assignment rebind-ok for the type gate * fix(aiohttp_openai): set service_tier with setattr like other providers ModelResponse has no declared service_tier field, so attribute assignment added a reportAttributeAccessIssue error to the type gate. * test(aiohttp_openai): verify serialized served service tier * fix(aiohttp_openai): copy response when preserving served tier --------- Co-authored-by: tin-berri --- .../aiohttp_openai/chat/transformation.py | 2 ++ .../chat/test_transformation.py | 26 ++++++++++++++++--- 2 files changed, 24 insertions(+), 4 deletions(-) diff --git a/litellm/llms/aiohttp_openai/chat/transformation.py b/litellm/llms/aiohttp_openai/chat/transformation.py index 74f782b56e7..9253ca389ee 100644 --- a/litellm/llms/aiohttp_openai/chat/transformation.py +++ b/litellm/llms/aiohttp_openai/chat/transformation.py @@ -82,4 +82,6 @@ class AiohttpOpenAIChatConfig(OpenAILikeChatConfig): model_response.model = _json_response.get("model") model_response.object = _json_response.get("object") model_response.system_fingerprint = _json_response.get("system_fingerprint") + if "service_tier" in _json_response: + return model_response.model_copy(update={"service_tier": _json_response.get("service_tier")}) return model_response diff --git a/tests/unit/llms/aiohttp_openai/chat/test_transformation.py b/tests/unit/llms/aiohttp_openai/chat/test_transformation.py index 253f9dbe746..ae242e69a9a 100644 --- a/tests/unit/llms/aiohttp_openai/chat/test_transformation.py +++ b/tests/unit/llms/aiohttp_openai/chat/test_transformation.py @@ -1,15 +1,16 @@ +from typing import Final from unittest.mock import AsyncMock, Mock import pytest from aiohttp import ClientResponse -from pydantic import ValidationError +from pydantic import TypeAdapter, ValidationError from litellm.llms.aiohttp_openai.chat.transformation import AiohttpOpenAIChatConfig from litellm.types.utils import ModelResponse -async def _transform(body: object) -> ModelResponse: - raw_response = Mock(spec=ClientResponse) +async def _transform(body: object, *, requested_tier: str | None = None) -> ModelResponse: + raw_response: Final = Mock(spec=ClientResponse) raw_response.json = AsyncMock(return_value=body) return await AiohttpOpenAIChatConfig().transform_response( model="gpt-4o", @@ -18,7 +19,7 @@ async def _transform(body: object) -> ModelResponse: logging_obj=Mock(), request_data={}, messages=[{"role": "user", "content": "Hello"}], - optional_params={}, + optional_params={"service_tier": requested_tier} if requested_tier is not None else {}, litellm_params={}, encoding=None, ) @@ -99,3 +100,20 @@ async def test_transform_response_rejects_choices_that_are_not_a_list_of_objects async def test_transform_response_rejects_a_choice_whose_message_is_not_an_object(): with pytest.raises(ValidationError, match="validation error for Choices"): await _transform({"choices": [{"message": "text"}]}) + + +@pytest.mark.parametrize("tier", ("flex", "priority", "default", "custom-tier", None)) +async def test_transform_response_serializes_served_tier_instead_of_requested_tier(tier: str | None) -> None: + response: Final = await _transform( + {"id": "chatcmpl-1", "choices": [{}], "service_tier": tier}, requested_tier="auto" + ) + serialized: Final = TypeAdapter(dict[str, object]).validate_json(response.model_dump_json()) + + assert serialized["service_tier"] == tier + + +async def test_transform_response_does_not_invent_a_served_tier_when_omitted() -> None: + response: Final = await _transform({"id": "chatcmpl-1", "choices": [{}]}, requested_tier="flex") + serialized: Final = TypeAdapter(dict[str, object]).validate_json(response.model_dump_json()) + + assert "service_tier" not in serialized From 1f68cc5649483aba0515e7e2e9e416a2633a2ad4 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Wed, 28 Oct 2026 13:16:56 -0700 Subject: [PATCH 084/128] refactor(rust-bridge): require the native extension (#48852) * refactor(rust-bridge): require the native extension _native ships in every wheel and the sdist will not install without it, yet the bridge loaded it optionally and carried None paths that could fail open for fork safety, exception mapping and OTLP errors. Delete loader.py and import _native directly. Tests inject a fake module through NativeBinding(source=...) or with_source() instead of patching the loader The install smoke checks now verify the compiled module, and tests fail instead of skipping when _native is missing. The interactions spec canary installs the project so it has the extension Co-Authored-By: Claude Opus 5.5 * fix(ci): build Rust bridge in Windows smoke tests --------- Co-authored-by: Claude Opus 5.5 Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../workflows/interactions-spec-canary.yml | 7 ++- .github/workflows/test-merge-smoke.yml | 7 ++- .github/workflows/test-unit.yml | 10 ++++ gateway/Dockerfile | 2 +- .../crates/python-bridge/src/logger/tests.rs | 2 + litellm/litellm_core_utils/README.md | 2 +- litellm/rust_bridge/AGENTS.md | 6 ++ litellm/rust_bridge/__init__.py | 7 +-- litellm/rust_bridge/bindings.py | 26 +++------ litellm/rust_bridge/clickhouse.py | 37 ++---------- litellm/rust_bridge/fork_guard.py | 20 ++----- litellm/rust_bridge/loader.py | 35 ------------ litellm/rust_bridge/responses/websocket.py | 25 +++----- litellm/rust_bridge/runtime.py | 19 ++----- litellm/tracing/otlp_http.py | 18 +----- pyproject.toml | 2 +- .../base_sdk_tests/check_base_sdk_install.py | 20 ++++++- tests/test_litellm_rust/conftest.py | 12 ---- .../unit/litellm_core_utils/test_tokenizer.py | 8 +-- .../proxy_server/test_exception_handlers.py | 11 +--- tests/unit/proxy/test_tracing_endpoints.py | 1 - .../responses/test_rust_bridge_websocket.py | 15 ----- .../rust_bridge/host/test_secret_manager.py | 46 +-------------- .../unit/rust_bridge/messages/test_secrets.py | 2 - tests/unit/rust_bridge/ocr/test_secrets.py | 3 +- tests/unit/rust_bridge/test_bindings.py | 57 ++++++++++--------- tests/unit/rust_bridge/test_fork_guard.py | 23 ++------ tests/unit/rust_bridge/test_logger.py | 3 - tests/unit/rust_bridge/test_runtime.py | 17 +----- tests/unit/rust_bridge/test_token_counter.py | 4 +- .../test_audio_transcription_rust_bridge.py | 14 +---- tests/unit/test_logging.py | 14 ----- tests/unit/test_secret_redaction.py | 2 - .../check_windows_wheel_install.py | 12 +++- 34 files changed, 144 insertions(+), 345 deletions(-) delete mode 100644 litellm/rust_bridge/loader.py diff --git a/.github/workflows/interactions-spec-canary.yml b/.github/workflows/interactions-spec-canary.yml index 7ec3defdcb7..a2e4646dce5 100644 --- a/.github/workflows/interactions-spec-canary.yml +++ b/.github/workflows/interactions-spec-canary.yml @@ -19,7 +19,7 @@ jobs: name: interactions-spec-canary if: github.repository == 'BerriAI/litellm' runs-on: ubuntu-latest - timeout-minutes: 15 + timeout-minutes: 30 permissions: contents: read outputs: @@ -39,8 +39,11 @@ jobs: with: version: "0.10.9" + - name: Set up the Rust bridge + uses: ./.github/actions/rust-bridge + - name: Install dependencies - run: .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy --no-install-project + run: .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy - name: Run the compliance tests against Google's live spec id: compliance diff --git a/.github/workflows/test-merge-smoke.yml b/.github/workflows/test-merge-smoke.yml index f3f09523491..99dbd1c101b 100644 --- a/.github/workflows/test-merge-smoke.yml +++ b/.github/workflows/test-merge-smoke.yml @@ -132,7 +132,7 @@ jobs: windows-sdk: name: windows-sdk runs-on: windows-2025 - timeout-minutes: 10 + timeout-minutes: 30 env: LITELLM_LOCAL_MODEL_COST_MAP: "True" defaults: @@ -154,8 +154,11 @@ jobs: with: version: "0.10.9" + - name: Set up the Rust bridge + uses: ./.github/actions/rust-bridge + - name: Install dependencies - run: .github/scripts/uv_sync_with_retries.sh --frozen --no-install-workspace --group dev --python 3.11 + run: .github/scripts/uv_sync_with_retries.sh --frozen --group dev --python 3.11 - name: Run the Windows tests run: uv run --no-sync python -m pytest --tb=short -v tests/windows_tests diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index b4b7afb9a09..e3980bd8f3b 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -1306,6 +1306,8 @@ jobs: done openai-sdk-compat: name: OpenAI ${{ matrix.openai-version }} compatibility + needs: rust-bridge + if: ${{ !cancelled() }} permissions: contents: read pull-requests: read @@ -1342,10 +1344,18 @@ jobs: with: version: "0.10.9" + - name: Set up the Rust bridge + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 5 + uses: ./.github/actions/rust-bridge + with: + artifact: ${{ needs.rust-bridge.outputs.artifact }} + - name: Install dependencies if: steps.changes.outputs.decision != 'skip' run: | .github/scripts/uv_sync_with_retries.sh --frozen --group dev --group proxy-dev --extra proxy --no-install-workspace + cp rust-bridge-dist/litellm/rust_bridge/_native.* litellm/rust_bridge/ uv pip install "openai==${OPENAI_VERSION:?}" --exclude-newer-package openai=2026-10-07T00:00:00Z uv run --no-sync python -c 'import os, openai; assert openai.__version__ == os.environ["OPENAI_VERSION"], openai.__version__' diff --git a/gateway/Dockerfile b/gateway/Dockerfile index 69068cfc8a5..6dc8bab45ce 100644 --- a/gateway/Dockerfile +++ b/gateway/Dockerfile @@ -153,7 +153,7 @@ RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \ find /app/.venv -type d -path "*/tornado/test" -delete && \ chmod -R a+rX /opt/prisma && \ python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths" && \ - python -c "import litellm; from litellm.rust_bridge.loader import native_bridge_available; assert litellm.__file__ == '/app/litellm/__init__.py', litellm.__file__; assert native_bridge_available()" + python -c "import litellm, litellm.rust_bridge._native; assert litellm.__file__ == '/app/litellm/__init__.py', litellm.__file__" USER nonroot diff --git a/litellm-rust/crates/python-bridge/src/logger/tests.rs b/litellm-rust/crates/python-bridge/src/logger/tests.rs index 1fca3be720e..9a07428808a 100644 --- a/litellm-rust/crates/python-bridge/src/logger/tests.rs +++ b/litellm-rust/crates/python-bridge/src/logger/tests.rs @@ -187,6 +187,7 @@ fn native_events_reach_python_with_levels_context_reentry_and_http_deduplication locals .set_item("http_warning", wrap_pyfunction!(http_warning, py).unwrap()) .unwrap(); + locals.set_item("native", crate::native_module(py)).unwrap(); let importable = py .eval( c"__import__('importlib.util', fromlist=['util']).find_spec('dotenv') is not None", @@ -205,6 +206,7 @@ import asyncio import logging import sys sys.path.insert(0, repo_root) +sys.modules['litellm.rust_bridge._native'] = native import litellm from litellm._logging import verbose_logger, session_id_var, trace_id_var diff --git a/litellm/litellm_core_utils/README.md b/litellm/litellm_core_utils/README.md index a5f5e8326b9..331f4adccbb 100644 --- a/litellm/litellm_core_utils/README.md +++ b/litellm/litellm_core_utils/README.md @@ -11,4 +11,4 @@ Core files: - `duration_parser.py`: code for parsing durations - e.g. "1d", "1mo", "10s" - `api_route_to_call_types.py`: mapping of API routes to their corresponding CallTypes (e.g., `/chat/completions` -> [acompletion, completion]) -Tokenizer factories return Python tokenizer objects by default. Set `LITELLM_RUST=1` or call `litellm.rust(True)` before constructing tokenizers to select the Rust backend through `Route.TOKENIZER` in the Rust catalog. Missing native bindings or unsupported native features fall back to Python. Existing tokenizer objects keep their selected backend. Rust-backed tokenizer objects carry the read-only `tiktoken.Encoding` / `tokenizers.Tokenizer` surface and are immutable: `enable_padding`, `enable_truncation` and `add_tokens` stay on the Python tokenizer. +Tokenizer factories return Python tokenizer objects by default. Set `LITELLM_RUST=1` or call `litellm.rust(True)` before constructing tokenizers to select the Rust backend through `Route.TOKENIZER` in the Rust catalog. Unsupported native features fall back to Python. Existing tokenizer objects keep their selected backend. Rust-backed tokenizer objects carry the read-only `tiktoken.Encoding` / `tokenizers.Tokenizer` surface and are immutable: `enable_padding`, `enable_truncation` and `add_tokens` stay on the Python tokenizer. diff --git a/litellm/rust_bridge/AGENTS.md b/litellm/rust_bridge/AGENTS.md index c197b21bad0..c2a030ca4d9 100644 --- a/litellm/rust_bridge/AGENTS.md +++ b/litellm/rust_bridge/AGENTS.md @@ -46,6 +46,12 @@ - A failure after execution starts is terminal and never triggers fallback - Creating an unstarted native coroutine acquires no clients or credentials and captures no execution context +# Native extension + +- `_native` ships in every wheel and the sdist builds it, so it is always importable +- Import it with `from litellm.rust_bridge import _native` and use the attributes `_native.pyi` declares +- Never add a `None` branch, `getattr` default or stand-in class for a missing extension + # Verification - Check cross-language behavior against a freshly installed extension, covering unstarted-call release, setup ordering, exception identity, sync and async streams, cancellation and GC diff --git a/litellm/rust_bridge/__init__.py b/litellm/rust_bridge/__init__.py index 8f6f4390b8a..52cc3c8cff1 100644 --- a/litellm/rust_bridge/__init__.py +++ b/litellm/rust_bridge/__init__.py @@ -1,10 +1,5 @@ """LiteLLM Rust bridge package.""" from litellm.rust_bridge.configuration import rust -from litellm.rust_bridge.loader import ( - get_native_bridge, - native_bridge_available, - reset_native_bridge_cache, -) -__all__ = ["get_native_bridge", "native_bridge_available", "reset_native_bridge_cache", "rust"] +__all__ = ["rust"] diff --git a/litellm/rust_bridge/bindings.py b/litellm/rust_bridge/bindings.py index d16f150a2aa..91570c6a8bc 100644 --- a/litellm/rust_bridge/bindings.py +++ b/litellm/rust_bridge/bindings.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import Callable from typing import Final, Generic, TypeVar -from litellm.rust_bridge.loader import get_native_bridge +from litellm.rust_bridge import _native BindingT = TypeVar("BindingT") @@ -18,32 +18,24 @@ _UNSET: Final = _Unset() class NativeBinding(Generic[BindingT]): """Resolve one native attribute with an explicit, resettable test override.""" - def __init__(self, attribute: str, *, validate: Callable[[object], BindingT | None]) -> None: + def __init__( + self, attribute: str, *, validate: Callable[[object], BindingT | None], source: object = _native + ) -> None: self._attribute: Final = attribute self._validate: Final = validate + self._source: Final = source self._override: BindingT | None | _Unset = _UNSET def load(self) -> BindingT | None: if not isinstance(self._override, _Unset): return self._override - native: Final = get_native_bridge() - if native is None: - return None - return self._validate(getattr(native, self._attribute, None)) + return self._validate(getattr(self._source, self._attribute, None)) + + def with_source(self, source: object) -> NativeBinding[BindingT]: + return NativeBinding(self._attribute, validate=self._validate, source=source) def override(self, value: BindingT | None) -> None: self._override = value def reset(self) -> None: self._override = _UNSET - - -def native_exception_types() -> tuple[type[BaseException], type[BaseException]] | None: - native: Final = get_native_bridge() - if native is None: - return None - declined: Final = getattr(native, "RustBridgeDeclined", None) - upstream: Final = getattr(native, "RustUpstreamError", None) - if not isinstance(declined, type) or not isinstance(upstream, type): - return None - return declined, upstream diff --git a/litellm/rust_bridge/clickhouse.py b/litellm/rust_bridge/clickhouse.py index cf2c8fdd3e2..fe2809ef856 100644 --- a/litellm/rust_bridge/clickhouse.py +++ b/litellm/rust_bridge/clickhouse.py @@ -1,12 +1,10 @@ import os -from collections.abc import Awaitable, Mapping, Sequence +from collections.abc import Mapping, Sequence from dataclasses import dataclass -from typing import Final, Protocol, runtime_checkable - -from pydantic import ConfigDict, TypeAdapter +from typing import Final from litellm.constants import DEFAULT_AGENT_TRACING_RETENTION_DAYS, DEFAULT_CLICKHOUSE_DATABASE -from litellm.rust_bridge.loader import get_native_bridge +from litellm.rust_bridge import _native @dataclass(frozen=True, slots=True, repr=False) @@ -16,27 +14,6 @@ class SpendStorageConfig: retention_days: int = DEFAULT_AGENT_TRACING_RETENTION_DAYS -class NativeSpendConfig(Protocol): - def __init__(self, database: str, url: str, retention_days: int) -> None: ... - - -class NativeSpendStorage(Protocol): - def __init__(self, config: NativeSpendConfig) -> None: ... - def ensure_schema(self) -> Awaitable[None]: ... - def insert_rows(self, rows: Sequence[Mapping[str, object]]) -> Awaitable[None]: ... - - -@runtime_checkable -class NativeSpendBridge(Protocol): - NativeClickHouseSpendConfig: type[NativeSpendConfig] - NativeClickHouseSpendStorage: type[NativeSpendStorage] - - -_NATIVE: Final[TypeAdapter[NativeSpendBridge]] = TypeAdapter( - NativeSpendBridge, config=ConfigDict(arbitrary_types_allowed=True) -) - - def spend_storage_config(environ: Mapping[str, str] = os.environ) -> SpendStorageConfig: url: Final = environ.get("CLICKHOUSE_URL") if not url: @@ -53,12 +30,8 @@ def spend_storage_config(environ: Mapping[str, str] = os.environ) -> SpendStorag class ClickHouseSpendStorage: def __init__(self, config: SpendStorageConfig) -> None: - bridge: Final = get_native_bridge() - if bridge is None: - raise RuntimeError("ClickHouse spend logging requires the Rust extension") - native: Final = _NATIVE.validate_python(bridge) - self._native: Final = native.NativeClickHouseSpendStorage( - native.NativeClickHouseSpendConfig(config.database, config.url, config.retention_days) + self._native: Final = _native.NativeClickHouseSpendStorage( + _native.NativeClickHouseSpendConfig(config.database, config.url, config.retention_days) ) async def ensure_schema(self) -> None: diff --git a/litellm/rust_bridge/fork_guard.py b/litellm/rust_bridge/fork_guard.py index c94665fb8db..82eb5c96295 100644 --- a/litellm/rust_bridge/fork_guard.py +++ b/litellm/rust_bridge/fork_guard.py @@ -11,32 +11,22 @@ site, so the runtime can never start there. Workers forked from it are unaffecte from __future__ import annotations +from collections.abc import Callable from typing import Final -from litellm.rust_bridge.loader import get_native_bridge +from litellm.rust_bridge import _native class NativeStateStartedBeforeFork(RuntimeError): pass -class _NeverRaised(RuntimeError): - """Stands in for a native exception when the extension is unavailable or predates it.""" +ForkedAfterNativeRuntimeStarted: Final = _native.ForkedAfterNativeRuntimeStarted +ProcessReservedForForking: Final = _native.ProcessReservedForForking -_native: Final = get_native_bridge() -ForkedAfterNativeRuntimeStarted: Final[type[RuntimeError]] = getattr( - _native, "ForkedAfterNativeRuntimeStarted", _NeverRaised -) -ProcessReservedForForking: Final[type[RuntimeError]] = getattr(_native, "ProcessReservedForForking", _NeverRaised) - - -def reserve_process_for_forking(where: str) -> None: +def reserve_process_for_forking(where: str, reserve: Callable[[], None] = _native.reserve_process_for_forking) -> None: """Forbid native routes in this process. Raises if one already ran here.""" - native: Final = get_native_bridge() - reserve: Final = getattr(native, "reserve_process_for_forking", None) - if not callable(reserve): - return try: reserve() except RuntimeError as error: diff --git a/litellm/rust_bridge/loader.py b/litellm/rust_bridge/loader.py deleted file mode 100644 index 022c38f5a85..00000000000 --- a/litellm/rust_bridge/loader.py +++ /dev/null @@ -1,35 +0,0 @@ -"""Loader for the packaged LiteLLM Rust extension.""" - -from __future__ import annotations - -from types import ModuleType -from typing import Final - -_BRIDGE_SENTINEL: Final = object() -_cached_bridge: ModuleType | None | object = _BRIDGE_SENTINEL - - -def get_native_bridge() -> ModuleType | None: - """Return the packaged Rust extension, or ``None`` when unavailable.""" - global _cached_bridge - if _cached_bridge is not _BRIDGE_SENTINEL: - return _cached_bridge if isinstance(_cached_bridge, ModuleType) else None - - try: - from litellm.rust_bridge import _native - except ImportError: - _cached_bridge = None - return None - _cached_bridge = _native - return _native - - -def reset_native_bridge_cache() -> None: - """Forget the cached extension so the next lookup reimports it from disk.""" - global _cached_bridge - _cached_bridge = _BRIDGE_SENTINEL - - -def native_bridge_available() -> bool: - """Whether the packaged Rust extension is importable.""" - return get_native_bridge() is not None diff --git a/litellm/rust_bridge/responses/websocket.py b/litellm/rust_bridge/responses/websocket.py index f10631958bb..83a1fa53376 100644 --- a/litellm/rust_bridge/responses/websocket.py +++ b/litellm/rust_bridge/responses/websocket.py @@ -2,32 +2,33 @@ from __future__ import annotations +from collections.abc import Awaitable from dataclasses import dataclass from typing import Final, Protocol import httpx from websockets.exceptions import ConnectionClosedOK +from litellm.rust_bridge import _native from litellm.rust_bridge.host.timeouts import timeout_to_seconds -from litellm.rust_bridge.loader import get_native_bridge class RustResponsesWebSocket(Protocol): - async def send_text(self, text: str) -> None: ... + def send_text(self, text: str) -> Awaitable[None]: ... - async def recv_text(self) -> str | None: ... + def recv_text(self) -> Awaitable[str | None]: ... - async def close(self) -> None: ... + def close(self) -> Awaitable[None]: ... class RustResponsesWebSocketConnection(Protocol): @classmethod - async def connect( + def connect( cls, url: str, headers: dict[str, str], timeout_seconds: float | None, - ) -> RustResponsesWebSocket: ... + ) -> Awaitable[RustResponsesWebSocket]: ... class _Unset: @@ -53,16 +54,10 @@ def set_rust_responses_websocket( _STATE.connection = connection -def load_rust_responses_websocket() -> RustResponsesWebSocketConnection | None: +def load_rust_responses_websocket() -> RustResponsesWebSocketConnection: if _STATE.connection is not None: return _STATE.connection - native_bridge: Final = get_native_bridge() - if native_bridge is None: - return None - connection_type: Final[RustResponsesWebSocketConnection | None] = getattr( - native_bridge, "ResponsesWebSocketConnection", None - ) - return connection_type + return _native.ResponsesWebSocketConnection class _ConnectionAdapter: @@ -89,8 +84,6 @@ async def connect( timeout: float | httpx.Timeout | None, ) -> _ConnectionAdapter | None: connection_type: Final = load_rust_responses_websocket() - if connection_type is None: - return None try: connection: Final = await connection_type.connect( url=url, diff --git a/litellm/rust_bridge/runtime.py b/litellm/rust_bridge/runtime.py index cf632c09074..e22ebe73b08 100644 --- a/litellm/rust_bridge/runtime.py +++ b/litellm/rust_bridge/runtime.py @@ -7,7 +7,8 @@ from typing import Final, Generic, NoReturn, TypeAlias, TypeVar from typing_extensions import assert_never from litellm.exceptions import APIError -from litellm.rust_bridge.bindings import NativeBinding, native_exception_types +from litellm.rust_bridge import _native +from litellm.rust_bridge.bindings import NativeBinding from litellm.rust_bridge.catalog import RouteContext, Rules, decision from litellm.rust_bridge.configuration import Decision from litellm.rust_bridge.response_metadata import mark_rust_response @@ -153,15 +154,11 @@ def attempt( ) -> RustAttempt[ResultT]: if native_call is None: return RustUnavailable() - exceptions: Final = native_exception_types() - if exceptions is None: - return RustHandled(adapt(native_call())) - declined, upstream = exceptions try: value: Final = native_call() - except declined as error: + except _native.RustBridgeDeclined as error: return RustDeclined(reason=_decline_reason(error)) - except upstream as error: + except _native.RustUpstreamError as error: _raise_upstream(error, context) return RustHandled(adapt(value)) @@ -174,15 +171,11 @@ async def aattempt( ) -> RustAttempt[ResultT]: if native_call is None: return RustUnavailable() - exceptions: Final = native_exception_types() - if exceptions is None: - return RustHandled(adapt(await native_call())) - declined, upstream = exceptions try: value: Final = await native_call() - except declined as error: + except _native.RustBridgeDeclined as error: return RustDeclined(reason=_decline_reason(error)) - except upstream as error: + except _native.RustUpstreamError as error: _raise_upstream(error, context) return RustHandled(adapt(value)) diff --git a/litellm/tracing/otlp_http.py b/litellm/tracing/otlp_http.py index 493cc7cf984..81dbab6c023 100644 --- a/litellm/tracing/otlp_http.py +++ b/litellm/tracing/otlp_http.py @@ -1,20 +1,9 @@ import json -from typing import Final, Protocol, runtime_checkable +from typing import Final -from pydantic import ConfigDict, TypeAdapter from typing_extensions import ReadOnly, TypedDict -from litellm.rust_bridge.loader import get_native_bridge - - -@runtime_checkable -class NativeOtlpError(Protocol): - def trace_encode_error(self, message: str) -> bytes: ... - - -_NATIVE: Final[TypeAdapter[NativeOtlpError]] = TypeAdapter( - NativeOtlpError, config=ConfigDict(arbitrary_types_allowed=True) -) +from litellm.rust_bridge import _native class OTLPError(TypedDict): @@ -22,8 +11,7 @@ class OTLPError(TypedDict): def encode_error(message: str) -> bytes: - native: Final = get_native_bridge() - return b"" if native is None else _NATIVE.validate_python(native).trace_encode_error(message) + return _native.trace_encode_error(message) def encode_otlp_response(content_type: str | None, error: str | None = None) -> tuple[bytes, str]: diff --git a/pyproject.toml b/pyproject.toml index 2d25c1ec737..00cdb1587aa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -401,7 +401,7 @@ markers = [ "asyncio: mark test as an asyncio test", "limit_leaks: mark test with memory limit for leak detection (e.g., '40 MB')", "no_parallel: mark test to run sequentially (not in parallel) - typically for memory measurement tests", - "requires_rust_extension: public Python contract requiring an enabled, compiled Rust extension", + "requires_rust_extension: public Python contract tests that run against the release-built wheel", ] filterwarnings = [ # Suppress Pydantic serializer warnings from mock server responses (non-critical for memory tests) diff --git a/tests/base_sdk_tests/check_base_sdk_install.py b/tests/base_sdk_tests/check_base_sdk_install.py index e546549d25f..e37593adb4b 100644 --- a/tests/base_sdk_tests/check_base_sdk_install.py +++ b/tests/base_sdk_tests/check_base_sdk_install.py @@ -8,6 +8,8 @@ the very class of undeclared-dependency bug this guards against. import argparse import importlib.util +from importlib.machinery import EXTENSION_SUFFIXES +from pathlib import Path import re import sys import traceback @@ -89,6 +91,18 @@ def check_import() -> str: return f"imported {distribution} {sdk_version}" +def check_native_extension() -> str: + import litellm + from litellm.rust_bridge import _native + + native_path: Final = Path(_native.__file__).resolve() + package_path: Final = Path(litellm.__file__).resolve().parent + _require(native_path.parent == package_path / "rust_bridge", "native extension is outside the SDK package") + _require(any(str(native_path).endswith(suffix) for suffix in EXTENSION_SUFFIXES), "native module is not compiled") + _require(_native.trace_encode_error("ping") == b"\x12\x04ping", "native OTLP error encoding failed") + return f"loaded and exercised {native_path}" + + def check_completion() -> str: import litellm @@ -230,6 +244,7 @@ def check_bedrock_credential_resolution() -> str: CHECKS: tuple[tuple[str, Callable[[], str]], ...] = ( ("environment is base-only", check_environment_is_base_only), ("import litellm", check_import), + ("native extension", check_native_extension), ("optional MCP installation guidance", check_mcp_install_guidance), ("chat completion", check_completion), ("chat completion with a function tool", check_completion_with_function_tool), @@ -250,8 +265,11 @@ def _run(check: Callable[[], str]) -> tuple[bool, str]: def main() -> int: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--profile", choices=("legacy", "core", "dependencies"), default="legacy") + parser.add_argument("--profile", choices=("legacy", "core", "dependencies", "native"), default="legacy") profile = parser.parse_args().profile + if profile == "native": + print(check_native_extension()) + return 0 for module in ("boto3", "botocore", "tokenizers", "huggingface_hub"): present = importlib.util.find_spec(module) is not None _require(present == (profile != "core"), f"{profile}: unexpected presence of {module}: {present}") diff --git a/tests/test_litellm_rust/conftest.py b/tests/test_litellm_rust/conftest.py index 06fbcf58591..8bcb48a09b5 100644 --- a/tests/test_litellm_rust/conftest.py +++ b/tests/test_litellm_rust/conftest.py @@ -1,6 +1,4 @@ from collections.abc import AsyncIterator, Generator -from importlib import import_module -from typing import Final import pytest import pytest_asyncio @@ -23,16 +21,6 @@ def recording_server() -> Generator[RecordingServer]: yield server -def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: - try: - import_module("litellm.rust_bridge._native") - except ImportError: - skip: Final = pytest.mark.skip(reason="requires a compiled litellm.rust_bridge._native extension") - for item in items: - if "test_litellm_rust" in item.path.parts: - item.add_marker(skip) - - @pytest.fixture def isolated_azure_auth(monkeypatch: pytest.MonkeyPatch) -> None: for name in ( diff --git a/tests/unit/litellm_core_utils/test_tokenizer.py b/tests/unit/litellm_core_utils/test_tokenizer.py index 61e797d7166..54185c08c10 100644 --- a/tests/unit/litellm_core_utils/test_tokenizer.py +++ b/tests/unit/litellm_core_utils/test_tokenizer.py @@ -241,20 +241,14 @@ print("compatible") @pytest.mark.parametrize("rust", (None, "0", "1")) -def test_tokenization_without_native_extension_stays_offline(tmp_path: Path, rust: str | None) -> None: +def test_tokenization_stays_offline(tmp_path: Path, rust: str | None) -> None: script: Final = """ -import importlib.abc import sys sys.path.insert(0, sys.argv[1]) def reject_network(event, args): if event == "socket.connect": raise AssertionError("tokenizer attempted a network connection") sys.addaudithook(reject_network) -class Block(importlib.abc.MetaPathFinder): - def find_spec(self, fullname, path=None, target=None): - if fullname == "litellm.rust_bridge._native": - raise ImportError("native extension is unavailable") -sys.meta_path.insert(0, Block()) import litellm from litellm.rust_bridge.tokenizer import get_encoding import tiktoken diff --git a/tests/unit/proxy/proxy_server/test_exception_handlers.py b/tests/unit/proxy/proxy_server/test_exception_handlers.py index f6efdecc884..b34c2d2fb08 100644 --- a/tests/unit/proxy/proxy_server/test_exception_handlers.py +++ b/tests/unit/proxy/proxy_server/test_exception_handlers.py @@ -488,22 +488,17 @@ async def test_otel_unhandled_exception_handler_reraises_http_exception_invalid( @pytest.mark.asyncio @pytest.mark.parametrize("media_type", ["application/json", "application/x-protobuf"]) @pytest.mark.parametrize("root_path", ["", "/tenant-a"]) -@pytest.mark.parametrize("native_available", [True, False]) @pytest.mark.parametrize("error", [ ProxyException("database credentials: secret", "auth_error", None, 401), HTTPException(403, "database credentials: secret"), ]) -async def test_otlp_auth_errors_hide_internal_details_and_survive_missing_native( - media_type: str, root_path: str, native_available: bool, - error: ProxyException | HTTPException, monkeypatch: pytest.MonkeyPatch, +async def test_otlp_auth_errors_hide_internal_details( + media_type: str, root_path: str, error: ProxyException | HTTPException, ) -> None: from google.rpc.status_pb2 import Status from litellm.proxy.proxy_server import otlp_http_exception_handler - from litellm.rust_bridge import loader - if not native_available: - monkeypatch.setattr(loader, "_cached_bridge", None) request: Final = Request({ "type": "http", "method": "POST", "path": root_path + "/v1/traces", "root_path": root_path, "headers": [(b"content-type", media_type.encode())], @@ -522,7 +517,7 @@ async def test_otlp_auth_errors_hide_internal_details_and_survive_missing_native else Status.FromString(response.body).message ) expected: Final = "Unauthorized" if isinstance(error, ProxyException) else "Forbidden" - assert message == (expected if native_available or media_type == "application/json" else "") + assert message == expected def _websocket(sent: list[Message]) -> WebSocket: diff --git a/tests/unit/proxy/test_tracing_endpoints.py b/tests/unit/proxy/test_tracing_endpoints.py index 62ca71b6cf6..7cff44da985 100644 --- a/tests/unit/proxy/test_tracing_endpoints.py +++ b/tests/unit/proxy/test_tracing_endpoints.py @@ -30,7 +30,6 @@ from litellm.proxy.auth.authorization import OwnedRows, ReadScope from litellm.proxy.auth.authorization_dependencies import get_log_team_lookup from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.tracing_runtime import manage_tracing, provide_storage -from litellm.rust_bridge import loader from litellm.tracing.errors import TraceChanged from litellm.tracing.generated.models import TraceQueryHelp from litellm.tracing.generated.responses import TraceSQLResponse diff --git a/tests/unit/responses/test_rust_bridge_websocket.py b/tests/unit/responses/test_rust_bridge_websocket.py index 00ae5eb970f..57b161d2147 100644 --- a/tests/unit/responses/test_rust_bridge_websocket.py +++ b/tests/unit/responses/test_rust_bridge_websocket.py @@ -55,21 +55,6 @@ async def test_adapter_raises_clean_close_when_rust_connection_ends() -> None: await adapter.recv() -@pytest.mark.asyncio -async def test_bridge_unavailable_returns_none(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(responses_websocket, "_STATE", responses_websocket._RustResponsesWebSocketState()) - monkeypatch.setattr(responses_websocket, "get_native_bridge", lambda: None) - - assert ( - await responses_websocket.connect( - url="wss://example.test/responses", - headers={}, - timeout=None, - ) - is None - ) - - @pytest.mark.asyncio async def test_enabled_bridge_connects_and_adapts_socket( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/rust_bridge/host/test_secret_manager.py b/tests/unit/rust_bridge/host/test_secret_manager.py index 9d876124ffd..1f8acc5792e 100644 --- a/tests/unit/rust_bridge/host/test_secret_manager.py +++ b/tests/unit/rust_bridge/host/test_secret_manager.py @@ -17,6 +17,7 @@ from botocore.credentials import Credentials from pydantic import JsonValue import litellm +from litellm.rust_bridge import _native as native from litellm.rust_bridge.host.secret_manager import capture_secret_manager, native_secret_manager_config from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2 from litellm.secret_managers.cyberark_secret_manager import CyberArkSecretManager @@ -67,7 +68,6 @@ def _vault_body(value: str) -> dict[str, object]: def _native_handle(client: object) -> "_SecretManagerRuntime": - native: Final = pytest.importorskip("litellm.rust_bridge._native") handle: Final = native._SecretManagerRuntime.from_client(client) assert handle is not None return handle @@ -75,7 +75,6 @@ def _native_handle(client: object) -> "_SecretManagerRuntime": @pytest.mark.parametrize("system", ("aws_secret_manager", "hashicorp_vault", "cyberark")) async def test_python_native_handle_reuses_backend_across_sync_and_async_reads(system: str) -> None: - native: Final = pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: environment: Final = { "AWS_REGION_NAME": "us-east-1", @@ -109,7 +108,6 @@ async def test_python_native_handle_reuses_backend_across_sync_and_async_reads(s def test_shared_initializer_captures_credentials_and_tracks_instance_settings(monkeypatch: pytest.MonkeyPatch) -> None: - native: Final = pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.default_response = ResponseSpec(body={"SecretString": "native-value"}) server.expected_requests = 2 @@ -141,7 +139,6 @@ def test_shared_initializer_captures_credentials_and_tracks_instance_settings(mo def test_configuration_replacement_rebuilds_without_invalidating_existing_handles( monkeypatch: pytest.MonkeyPatch, ) -> None: - native: Final = pytest.importorskip("litellm.rust_bridge._native") with recording_service() as first_server, recording_service() as second_server: first_server.default_response = ResponseSpec(body=_vault_body("first")) second_server.default_response = ResponseSpec(body=_vault_body("second")) @@ -158,7 +155,6 @@ def test_configuration_replacement_rebuilds_without_invalidating_existing_handle def test_custom_subclass_is_not_captured_for_a_native_handle() -> None: - native: Final = pytest.importorskip("litellm.rust_bridge._native") class CustomManager(AWSSecretsManagerV2): def sync_read_secret(self, secret_name: str, primary_secret_name: str | None = None) -> str: @@ -171,19 +167,16 @@ def test_custom_subclass_is_not_captured_for_a_native_handle() -> None: @pytest.mark.parametrize("system", ("google_secret_manager", "hashicorp_vault", "cyberark")) def test_enterprise_backends_cannot_initialize_without_entitlement(system: str) -> None: - native: Final = pytest.importorskip("litellm.rust_bridge._native") with pytest.raises(ValueError, match=r"[Ee]nterprise|[Pp]remium"): native._SecretManagerRuntime.from_config(system, {"CYBERARK_API_KEY": "key"}) def test_azure_factory_rejects_unencrypted_vault_endpoints() -> None: - native: Final = pytest.importorskip("litellm.rust_bridge._native") with pytest.raises(ValueError, match="https"): native._SecretManagerRuntime.from_config("azure_key_vault", {"AZURE_KEY_VAULT_URI": "http://127.0.0.1:1"}) def test_direct_builtin_constructor_can_use_native_without_registration(monkeypatch: pytest.MonkeyPatch) -> None: - native: Final = pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.default_response = ResponseSpec(body=_vault_body("native-value")) manager: Final = _vault(monkeypatch, server.base_url) @@ -193,12 +186,10 @@ def test_direct_builtin_constructor_can_use_native_without_registration(monkeypa def test_uncaptured_clients_have_no_native_handle() -> None: - native: Final = pytest.importorskip("litellm.rust_bridge._native") assert native._SecretManagerRuntime.from_client(SimpleNamespace()) is None def test_native_failure_is_not_replayed_in_python(monkeypatch: pytest.MonkeyPatch) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.default_response = ResponseSpec(status=403, body={"errors": ["denied"]}) server.expected_requests = 1 @@ -267,7 +258,6 @@ def test_same_named_custom_client_is_not_captured() -> None: @pytest.mark.parametrize("value", ("text", "", True, False, 42, 2**100, [1, "two"], {"nested": True}, None)) async def test_aws_primary_values_match_python_handler(monkeypatch: pytest.MonkeyPatch, value: JsonValue) -> None: - native: Final = pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.default_response = ResponseSpec(body={"SecretString": json.dumps({"KEY": value})}) server.expected_requests = 3 @@ -290,7 +280,6 @@ async def test_aws_primary_values_match_python_handler(monkeypatch: pytest.Monke @pytest.mark.parametrize("primary", (None, "primary")) async def test_aws_primary_values_match_python_reads(monkeypatch: pytest.MonkeyPatch, primary: str | None) -> None: - native: Final = pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.default_response = ResponseSpec(body={"SecretString": '{"KEY":"value"}'}) server.expected_requests = 3 @@ -343,7 +332,6 @@ def test_aws_absence_and_failed_reads_match_python_without_environment_fallback( async def test_aws_primary_json_errors_preserve_python_exception_details( monkeypatch: pytest.MonkeyPatch, document: str ) -> None: - native: Final = pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.default_response = ResponseSpec(body={"SecretString": document}) server.expected_requests = 3 @@ -370,7 +358,6 @@ async def test_aws_primary_json_errors_preserve_python_exception_details( async def test_public_aws_reads_preserve_coroutines_and_per_call_credentials( monkeypatch: pytest.MonkeyPatch, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as initial, recording_service() as selected: initial.expected_requests = 0 selected.expected_requests = 2 @@ -421,7 +408,6 @@ async def test_public_aws_reads_preserve_coroutines_and_per_call_credentials( async def test_public_aws_primary_reads_ignore_operation_overrides(monkeypatch: pytest.MonkeyPatch) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server, recording_service() as unused: server.default_response = ResponseSpec(body={"SecretString": '{"KEY":true}'}) server.expected_requests = 2 @@ -440,7 +426,6 @@ async def test_public_aws_primary_reads_ignore_operation_overrides(monkeypatch: async def test_public_aws_bootstrap_names_only_bypass_sync_reads(monkeypatch: pytest.MonkeyPatch) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.default_response = ResponseSpec(body={"SecretString": "remote-access"}) server.expected_requests = 1 @@ -457,7 +442,6 @@ async def test_public_aws_bootstrap_names_only_bypass_sync_reads(monkeypatch: py async def test_public_aws_read_timeouts_follow_the_python_http_handler( monkeypatch: pytest.MonkeyPatch, timeout: float | httpx.Timeout ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.default_response = ResponseSpec(body={"SecretString": "too-late"}, delay=0.25) server.expected_requests = 2 @@ -470,7 +454,6 @@ async def test_public_aws_read_timeouts_follow_the_python_http_handler( async def test_public_vault_reads_keep_overrides_cache_and_coroutines(monkeypatch: pytest.MonkeyPatch) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.default_response = ResponseSpec(body=_vault_body("value")) server.expected_requests = 2 @@ -490,7 +473,6 @@ async def test_public_vault_reads_keep_overrides_cache_and_coroutines(monkeypatc async def test_public_vault_failed_reads_return_none_without_replay( monkeypatch: pytest.MonkeyPatch, status: int ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.default_response = ResponseSpec(status=status, body={"errors": ["unavailable"]}) server.expected_requests = 2 @@ -502,7 +484,6 @@ async def test_public_vault_failed_reads_return_none_without_replay( async def test_public_cyberark_reads_reuse_authentication_and_cached_values( monkeypatch: pytest.MonkeyPatch, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") from litellm.proxy import proxy_server from litellm.secret_managers.cyberark_secret_manager import CyberArkSecretManager @@ -536,7 +517,6 @@ def test_public_aws_bootstrap_read_does_not_hit_the_provider(monkeypatch: pytest def test_public_vault_prefix_overrides_match_python_string_conversion( monkeypatch: pytest.MonkeyPatch, override: str | int ) -> None: - pytest.importorskip("litellm.rust_bridge._native") monkeypatch.setenv("HCP_VAULT_PATH_PREFIX", "default-prefix") with recording_service() as server: server.default_response = ResponseSpec(body=_vault_body("value")) @@ -562,7 +542,6 @@ def _cyberark(monkeypatch: pytest.MonkeyPatch, address: str) -> CyberArkSecretMa async def test_public_cyberark_writes_and_deletes_share_the_read_cache( monkeypatch: pytest.MonkeyPatch, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: for body in (b"token", b"old", {}, {}, b"provider-after-delete"): server.enqueue(ResponseSpec(body=body)) @@ -609,7 +588,6 @@ async def test_public_cyberark_write_errors_match_python_without_http_retries( status: int, authentication: bool, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: responses: Final = ( (ResponseSpec(status=status, body={}),) * 2 @@ -635,7 +613,6 @@ async def test_public_cyberark_write_errors_match_python_without_http_retries( async def test_public_cyberark_write_recovers_from_initial_policy_authentication_failure( monkeypatch: pytest.MonkeyPatch, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.enqueue(ResponseSpec(status=401, body={})) server.enqueue(ResponseSpec(body=b"token")) @@ -658,7 +635,6 @@ async def test_public_cyberark_write_rejects_unsafe_names_before_authentication( monkeypatch: pytest.MonkeyPatch, name: str, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.expected_requests = 0 manager: Final = _cyberark(monkeypatch, server.base_url) @@ -673,7 +649,6 @@ async def test_public_cyberark_rotation_returns_the_write_response_and_retains_o monkeypatch: pytest.MonkeyPatch, same_name: bool, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: for body in (b"token", b"old-value", {}, {}): server.enqueue(ResponseSpec(body=body)) @@ -696,7 +671,6 @@ async def test_public_native_cyberark_rotation_requires_a_fresh_matching_replace monkeypatch: pytest.MonkeyPatch, replacement: bytes | None, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: for body in (b"token", b"old-value", {}, {}): server.enqueue(ResponseSpec(body=body)) @@ -720,7 +694,6 @@ async def test_public_native_cyberark_rotation_requires_a_fresh_matching_replace async def test_public_cyberark_cached_authentication_does_not_retry_denied_reads( monkeypatch: pytest.MonkeyPatch, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.enqueue(ResponseSpec(body=b"token")) server.enqueue(ResponseSpec(body=b"value")) @@ -734,7 +707,6 @@ async def test_public_cyberark_cached_authentication_does_not_retry_denied_reads async def test_cyberark_handler_errors_match_python_after_cached_authentication_is_denied( monkeypatch: pytest.MonkeyPatch, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: for response in ( ResponseSpec(body=b"token"), @@ -757,7 +729,6 @@ async def test_cyberark_handler_errors_match_python_after_cached_authentication_ async def test_public_cyberark_connection_errors_match_python( monkeypatch: pytest.MonkeyPatch, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.expected_requests = 0 address: Final = server.base_url @@ -769,7 +740,6 @@ async def test_public_cyberark_connection_errors_match_python( async def test_public_native_cyberark_rotation_stops_after_a_failed_write(monkeypatch: pytest.MonkeyPatch) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: for body in (b"token", b"old-value", {}): server.enqueue(ResponseSpec(body=body)) @@ -788,7 +758,6 @@ async def test_public_vault_writes_preserve_complete_responses_and_request_field monkeypatch: pytest.MonkeyPatch, description: str | None, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: response: Final = { "request_id": "test-request", @@ -829,7 +798,6 @@ async def test_public_vault_mutation_http_errors_match_python_without_retry( operation: str, status: int, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.default_response = ResponseSpec(status=status, body={"errors": ["denied"]}) server.expected_requests = 2 @@ -856,7 +824,6 @@ async def test_public_vault_write_response_conversion_matches_python( monkeypatch: pytest.MonkeyPatch, body: bytes, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.default_response = ResponseSpec(body=body) server.expected_requests = 2 @@ -870,7 +837,6 @@ async def test_public_vault_write_response_conversion_matches_python( async def test_public_vault_deletion_invalidates_cached_fields( monkeypatch: pytest.MonkeyPatch, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.enqueue(ResponseSpec(body=_vault_body("old"))) server.enqueue(ResponseSpec(status=204, body=b"")) @@ -893,7 +859,6 @@ async def test_public_vault_rotation_preserves_response_and_best_effort_deletion same_name: bool, delete_status: int, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: response: Final = {"request_id": "write-id", "data": {"version": 3}, "extra": [1, 2]} server.enqueue(ResponseSpec(body=b"current-existence-is-status-only")) @@ -924,7 +889,6 @@ async def test_public_vault_rotation_failure_messages_and_request_counts_match_p stage: str, status: int, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: before: Final = ( () @@ -960,7 +924,6 @@ async def test_public_vault_rotation_mismatches_do_not_delete_the_old_alias( monkeypatch: pytest.MonkeyPatch, value: JsonValue, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: responses: Final = ( ResponseSpec(body=_vault_body("old")), @@ -983,7 +946,6 @@ async def test_public_vault_mutation_timeouts_match_python( monkeypatch: pytest.MonkeyPatch, operation: str, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.default_response = ResponseSpec(body=_vault_body("value"), delay=0.25) server.expected_requests = 2 @@ -1021,7 +983,6 @@ async def test_public_vault_unsafe_names_fail_before_authentication( monkeypatch: pytest.MonkeyPatch, operation: str, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.expected_requests = 0 manager: Final = _vault(monkeypatch, server.base_url) @@ -1038,7 +999,6 @@ async def test_public_vault_authentication_errors_match_python( monkeypatch: pytest.MonkeyPatch, operation: str, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") monkeypatch.setenv("HCP_VAULT_APPROLE_ROLE_ID", "role") monkeypatch.setenv("HCP_VAULT_APPROLE_SECRET_ID", "secret-id") with recording_service() as server: @@ -1064,7 +1024,6 @@ async def test_public_vault_authentication_errors_match_python( async def test_public_vault_rotation_stops_on_a_success_response_containing_an_error( monkeypatch: pytest.MonkeyPatch, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: result: Final = {"status": "error", "message": "write rejected", "extra": 2**100} for response in (ResponseSpec(body=_vault_body("old")), ResponseSpec(body=result)) * 2: @@ -1078,7 +1037,6 @@ async def test_public_vault_rotation_stops_on_a_success_response_containing_an_e async def test_public_native_vault_write_invalidates_stale_cached_values(monkeypatch: pytest.MonkeyPatch) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.enqueue(ResponseSpec(body=_vault_body("old"))) server.enqueue(ResponseSpec(body={"data": {"version": 2}})) @@ -1094,7 +1052,6 @@ async def test_public_native_vault_write_invalidates_stale_cached_values(monkeyp async def test_public_native_vault_write_rejects_description_overwriting_the_secret( monkeypatch: pytest.MonkeyPatch, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.expected_requests = 0 handle: Final = _native_handle(_vault(monkeypatch, server.base_url)) @@ -1111,7 +1068,6 @@ async def test_public_vault_rotation_preserves_malformed_verification_errors( monkeypatch: pytest.MonkeyPatch, body: JsonValue, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: responses: Final = ( ResponseSpec(body=_vault_body("old")), diff --git a/tests/unit/rust_bridge/messages/test_secrets.py b/tests/unit/rust_bridge/messages/test_secrets.py index f6439753809..64559d63193 100644 --- a/tests/unit/rust_bridge/messages/test_secrets.py +++ b/tests/unit/rust_bridge/messages/test_secrets.py @@ -17,8 +17,6 @@ from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagem from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_MODEL, MESSAGES_RESPONSE from tests.test_litellm_rust.support.recording_server import ResponseSpec, recording_service -pytest.importorskip("litellm.rust_bridge._native") - pytestmark = pytest.mark.usefixtures("local_model_cost_map") diff --git a/tests/unit/rust_bridge/ocr/test_secrets.py b/tests/unit/rust_bridge/ocr/test_secrets.py index 33ded916888..4ea6e09f885 100644 --- a/tests/unit/rust_bridge/ocr/test_secrets.py +++ b/tests/unit/rust_bridge/ocr/test_secrets.py @@ -12,6 +12,7 @@ import pytest import litellm from litellm.integrations.custom_secret_manager import CustomSecretManager from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.rust_bridge import _native as native from litellm.rust_bridge.host import settings from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR from litellm.rust_bridge.public_call import NativeCall @@ -19,8 +20,6 @@ from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagem from tests.test_litellm_rust.ocr.data import OCR_DOCUMENT, OCR_MODEL, OCR_RESPONSE from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, recording_service -native: Final = pytest.importorskip("litellm.rust_bridge._native") - AccessMode: TypeAlias = Literal["read_only", "write_only", "read_and_write"] diff --git a/tests/unit/rust_bridge/test_bindings.py b/tests/unit/rust_bridge/test_bindings.py index 39a157220dc..086be65ee26 100644 --- a/tests/unit/rust_bridge/test_bindings.py +++ b/tests/unit/rust_bridge/test_bindings.py @@ -6,7 +6,7 @@ from typing import Final, cast # noqa: TID251 # narrows the async binding's un import pytest -from litellm.rust_bridge import bindings +from litellm.rust_bridge import _native, bindings from litellm.rust_bridge.chat_completions import entrypoints as chat_completions from litellm.rust_bridge.embeddings import entrypoints as embeddings from litellm.rust_bridge.execution import Complete, Step @@ -17,12 +17,15 @@ from litellm.rust_bridge.responses import entrypoints as responses from litellm.rust_bridge.transcription import native as transcription -def test_binding_distinguishes_disable_from_reset(monkeypatch) -> None: - native = SimpleNamespace(route=lambda: "native") - monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) - binding: bindings.NativeBinding[object] = bindings.NativeBinding("route", validate=lambda value: value) +def test_binding_distinguishes_disable_from_reset() -> None: + def route() -> str: + return "native" - assert binding.load() is native.route + binding: bindings.NativeBinding[object] = bindings.NativeBinding( + "route", validate=lambda value: value, source=SimpleNamespace(route=route) + ) + + assert binding.load() is route binding.override(None) assert binding.load() is None @@ -32,16 +35,15 @@ def test_binding_distinguishes_disable_from_reset(monkeypatch) -> None: assert binding.load() is replacement binding.reset() - assert binding.load() is native.route + assert binding.load() is route @pytest.mark.parametrize(("value", "expected"), ((3, 3), ("invalid", None), (None, None))) -def test_binding_validates_native_attribute( - monkeypatch: pytest.MonkeyPatch, value: object, expected: int | None -) -> None: +def test_binding_validates_native_attribute(value: object, expected: int | None) -> None: native: Final = SimpleNamespace(route=value) - monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) - binding: Final = bindings.NativeBinding("route", validate=lambda item: item if isinstance(item, int) else None) + binding: Final = bindings.NativeBinding( + "route", validate=lambda item: item if isinstance(item, int) else None, source=native + ) assert binding.load() == expected @@ -72,28 +74,33 @@ ROUTE_BINDINGS: Final = tuple((attribute, binding) for attribute, binding, _ in @pytest.mark.parametrize( ("attribute", "route_binding"), ROUTE_BINDINGS, ids=[attribute for attribute, _ in ROUTE_BINDINGS] ) -def test_route_bindings_reject_non_callable_native_attributes( - monkeypatch: pytest.MonkeyPatch, attribute: str, route_binding: bindings.NativeBinding[object] +def test_route_bindings_resolve_the_packaged_native_route( + attribute: str, route_binding: bindings.NativeBinding[object] ) -> None: - monkeypatch.setattr(bindings, "get_native_bridge", lambda: SimpleNamespace(**{attribute: "not callable"})) - route_binding.reset() - assert route_binding.load() is None route_binding.reset() + assert attribute in dir(_native) + assert callable(route_binding.load()) + + +@pytest.mark.parametrize( + ("attribute", "route_binding"), ROUTE_BINDINGS, ids=[attribute for attribute, _ in ROUTE_BINDINGS] +) +def test_route_bindings_reject_non_callable_native_attributes( + attribute: str, route_binding: bindings.NativeBinding[object] +) -> None: + assert route_binding.with_source(SimpleNamespace(**{attribute: "not callable"})).load() is None @pytest.mark.parametrize( ("attribute", "route_binding"), TRANSCRIPTION_BINDINGS, ids=[attribute for attribute, _ in TRANSCRIPTION_BINDINGS] ) def test_transcription_bindings_expose_the_native_callable_itself( - monkeypatch: pytest.MonkeyPatch, attribute: str, route_binding: bindings.NativeBinding[object] + attribute: str, route_binding: bindings.NativeBinding[object] ) -> None: def native_route() -> None: pass - monkeypatch.setattr(bindings, "get_native_bridge", lambda: SimpleNamespace(**{attribute: native_route})) - route_binding.reset() - assert route_binding.load() is native_route - route_binding.reset() + assert route_binding.with_source(SimpleNamespace(**{attribute: native_route})).load() is native_route @dataclass(frozen=True, slots=True) @@ -116,7 +123,6 @@ class _EchoingExecution: ids=[attribute for attribute, _, _ in EXECUTION_BINDINGS], ) def test_execution_bindings_start_the_native_execution_and_return_its_result( - monkeypatch: pytest.MonkeyPatch, attribute: str, route_binding: bindings.NativeBinding[object], asynchronous: bool, @@ -128,9 +134,7 @@ def test_execution_bindings_start_the_native_execution_and_return_its_result( received.append(incoming) return _EchoingExecution() - monkeypatch.setattr(bindings, "get_native_bridge", lambda: SimpleNamespace(**{attribute: native_route})) - route_binding.reset() - bound: Final = route_binding.load() + bound: Final = route_binding.with_source(SimpleNamespace(**{attribute: native_route})).load() assert callable(bound) outcome: Final = bound(call) @@ -139,7 +143,6 @@ def test_execution_bindings_start_the_native_execution_and_return_its_result( if asynchronous else outcome ) - route_binding.reset() assert received == [call] assert result == ("completed with", None) diff --git a/tests/unit/rust_bridge/test_fork_guard.py b/tests/unit/rust_bridge/test_fork_guard.py index 88ae017ec39..aad50437f03 100644 --- a/tests/unit/rust_bridge/test_fork_guard.py +++ b/tests/unit/rust_bridge/test_fork_guard.py @@ -1,36 +1,21 @@ -from types import SimpleNamespace - import pytest from litellm.rust_bridge import fork_guard -def _reserve_with(monkeypatch: pytest.MonkeyPatch, native: object) -> None: - monkeypatch.setattr(fork_guard, "get_native_bridge", lambda: native) - fork_guard.reserve_process_for_forking("the gunicorn master") - - -def test_missing_extension_has_nothing_to_reserve(monkeypatch: pytest.MonkeyPatch) -> None: - assert _reserve_with(monkeypatch, None) is None - - -def test_extension_built_before_reservation_existed_passes(monkeypatch: pytest.MonkeyPatch) -> None: - assert _reserve_with(monkeypatch, SimpleNamespace()) is None - - -def test_unused_extension_is_reserved(monkeypatch: pytest.MonkeyPatch) -> None: +def test_unused_extension_is_reserved() -> None: calls: list[None] = [] - _reserve_with(monkeypatch, SimpleNamespace(reserve_process_for_forking=lambda: calls.append(None))) + fork_guard.reserve_process_for_forking("the gunicorn master", reserve=lambda: calls.append(None)) assert calls == [None] -def test_used_extension_refuses_and_names_the_place(monkeypatch: pytest.MonkeyPatch) -> None: +def test_used_extension_refuses_and_names_the_place() -> None: def reserve() -> None: raise RuntimeError("the native runtime already started in this process") with pytest.raises(fork_guard.NativeStateStartedBeforeFork, match="the gunicorn master") as raised: - _reserve_with(monkeypatch, SimpleNamespace(reserve_process_for_forking=reserve)) + fork_guard.reserve_process_for_forking("the gunicorn master", reserve=reserve) assert isinstance(raised.value.__cause__, RuntimeError) diff --git a/tests/unit/rust_bridge/test_logger.py b/tests/unit/rust_bridge/test_logger.py index 5a7b6c270a1..4745aca9d50 100644 --- a/tests/unit/rust_bridge/test_logger.py +++ b/tests/unit/rust_bridge/test_logger.py @@ -80,7 +80,6 @@ def test_native_logging_observes_level_changes(caplog: pytest.LogCaptureFixture) ), ) def test_native_credential_patterns_match_python(text: str) -> None: - pytest.importorskip("litellm.rust_bridge._native") from litellm.rust_bridge._native import NativeDiagnosticProcessor processor: Final = NativeDiagnosticProcessor(MINIMUM_CUSTOM_KEY_LENGTH) @@ -91,7 +90,6 @@ def test_native_credential_patterns_match_python(text: str) -> None: def test_native_client_redaction_matches_python() -> None: - pytest.importorskip("litellm.rust_bridge._native") from litellm.rust_bridge._native import NativeDiagnosticProcessor text: Final = "error at /etc/secrets/config on db.internal\nTraceback (most recent call last):\nsecret" @@ -100,7 +98,6 @@ def test_native_client_redaction_matches_python() -> None: def test_native_diagnostic_batch_matches_python() -> None: - pytest.importorskip("litellm.rust_bridge._native") from litellm.rust_bridge._native import NativeDiagnosticProcessor message: Final = "é" * 110 + "sk-" + "q" * 48 + "界" * 1000 diff --git a/tests/unit/rust_bridge/test_runtime.py b/tests/unit/rust_bridge/test_runtime.py index 95fa3e2cc7b..c0c416a62c9 100644 --- a/tests/unit/rust_bridge/test_runtime.py +++ b/tests/unit/rust_bridge/test_runtime.py @@ -1,7 +1,6 @@ from __future__ import annotations from collections.abc import Callable, Generator -from types import SimpleNamespace from typing import Final, Protocol import pytest @@ -10,26 +9,14 @@ from litellm.exceptions import APIError from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.rust_bridge import bindings, configuration, runtime +from litellm.rust_bridge._native import RustBridgeDeclined, RustUpstreamError from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule from litellm.rust_bridge.configuration import Rollout from litellm.rust_bridge.execution import AsyncStream, Complete, StreamChunk, SyncStream -class RustBridgeDeclined(Exception): - pass - - -class RustUpstreamError(Exception): - pass - - @pytest.fixture(autouse=True) -def native_exceptions(monkeypatch: pytest.MonkeyPatch) -> Generator[None]: - native: Final = SimpleNamespace( - RustBridgeDeclined=RustBridgeDeclined, - RustUpstreamError=RustUpstreamError, - ) - monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) +def rust_configuration(monkeypatch: pytest.MonkeyPatch) -> Generator[None]: monkeypatch.delenv("LITELLM_RUST", raising=False) configuration.reset_rust_configuration() yield diff --git a/tests/unit/rust_bridge/test_token_counter.py b/tests/unit/rust_bridge/test_token_counter.py index 3af7127a9b3..3513a0ef44c 100644 --- a/tests/unit/rust_bridge/test_token_counter.py +++ b/tests/unit/rust_bridge/test_token_counter.py @@ -16,6 +16,7 @@ import litellm from litellm.constants import TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS from litellm.litellm_core_utils.token_counter import openai_tokenizer_encoding from litellm.proxy.spend_tracking.input_tokens import count_input_tokens_for_model +from litellm.rust_bridge import _native as native from litellm.rust_bridge import token_counter as bridge from litellm.rust_bridge import tokenizer as tokenizer_dispatch from litellm.rust_bridge._native import Tokenizer @@ -313,7 +314,6 @@ PARITY_MODELS: Final[tuple[tuple[str, bridge.RustTokenizer], ...]] = ( async def test_native_count_matches_python_budget_counter( request_body: dict[str, object], model: str, tokenizer: bridge.RustTokenizer ) -> None: - native: Final = pytest.importorskip("litellm.rust_bridge._native") body: Final = json.dumps(request_body).replace(MODEL, model) parsed: Final = json.loads(body) @@ -328,7 +328,6 @@ async def test_tiktoken_counts_long_text_exactly_where_python_chunks( model: str, tokenizer: bridge.RustTokenizer ) -> None: """Python encodes tiktoken text in fixed-size chunks (drift of up to one token per chunk boundary); Rust does not.""" - native: Final = pytest.importorskip("litellm.rust_bridge._native") text: Final = "x " * 20_000 body: Final = {"model": model, "messages": [{"role": "user", "content": text}]} encoding: Final = Tokenizer.from_tiktoken(tokenizer) @@ -362,7 +361,6 @@ DECLINED_REQUESTS: Final[tuple[dict[str, object], ...]] = ( async def test_native_declines_shapes_python_prices_differently( request_body: dict[str, object], tokenizer: bridge.RustTokenizer ) -> None: - native: Final = pytest.importorskip("litellm.rust_bridge._native") raw, _ = _counted(request_body, MODEL_BY_TOKENIZER[tokenizer]) with pytest.raises(native.RustBridgeDeclined): diff --git a/tests/unit/test_audio_transcription_rust_bridge.py b/tests/unit/test_audio_transcription_rust_bridge.py index e3a6781dee7..715423c59d9 100644 --- a/tests/unit/test_audio_transcription_rust_bridge.py +++ b/tests/unit/test_audio_transcription_rust_bridge.py @@ -1,14 +1,14 @@ from __future__ import annotations from collections.abc import Generator -from types import SimpleNamespace from typing import Final import pytest import litellm from litellm.llms.bedrock.audio_transcription import BedrockAudioTranscriptionRustDispatch -from litellm.rust_bridge import bindings, configuration +from litellm.rust_bridge import configuration +from litellm.rust_bridge._native import RustBridgeDeclined, RustUpstreamError from litellm.rust_bridge.public_call import NativeCall from litellm.rust_bridge.transcription.native import NATIVE_ATRANSCRIPTION, NATIVE_TRANSCRIPTION @@ -28,18 +28,8 @@ TRANSCRIPTION_FIELDS: Final = frozenset( ) -class RustBridgeDeclined(Exception): - pass - - -class RustUpstreamError(Exception): - pass - - @pytest.fixture(autouse=True) def isolated_bridge(monkeypatch: pytest.MonkeyPatch) -> Generator[None]: - native: Final = SimpleNamespace(RustBridgeDeclined=RustBridgeDeclined, RustUpstreamError=RustUpstreamError) - monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) monkeypatch.delenv("LITELLM_RUST", raising=False) configuration.reset_rust_configuration() yield diff --git a/tests/unit/test_logging.py b/tests/unit/test_logging.py index 36be864a6f7..195d5afe000 100644 --- a/tests/unit/test_logging.py +++ b/tests/unit/test_logging.py @@ -862,8 +862,6 @@ def test_secret_filter_keeps_truncated_traceback(monkeypatch): @pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) def test_diagnostic_redaction_precedes_a_credential_cut(monkeypatch, native): - if native: - pytest.importorskip("litellm.rust_bridge._native") monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") monkeypatch.setenv("MAX_STRING_LENGTH_STDOUT_LOG", "500") monkeypatch.setenv("MAX_BASE64_LENGTH_STDOUT_LOG", "0") @@ -890,8 +888,6 @@ def test_correlation_id_redacts_before_its_length_bound(monkeypatch): @pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) def test_malformed_interpolation_still_scrubs_a_record(monkeypatch, native): - if native: - pytest.importorskip("litellm.rust_bridge._native") monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True) record = _make_record(logging.WARNING, "bad % api_key=secret123", ("value",)) @@ -905,8 +901,6 @@ def test_malformed_interpolation_still_scrubs_a_record(monkeypatch, native): @pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) def test_key_pattern_template_keeps_the_rendered_redacted_line(monkeypatch, native): - if native: - pytest.importorskip("litellm.rust_bridge._native") monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True) record = _make_record(logging.INFO, "password=%s ok", ("hunter2",)) @@ -1733,8 +1727,6 @@ def test_access_log_path_filter_keeps_a_record_without_a_string_path_arg(monkeyp @pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) def test_diagnostic_filter_scrubs_exc_stack_and_nested_extras(monkeypatch, native): - if native: - pytest.importorskip("litellm.rust_bridge._native") monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True) secret = "sk-" + "q" * 48 @@ -1761,8 +1753,6 @@ def test_diagnostic_filter_scrubs_exc_stack_and_nested_extras(monkeypatch, nativ @pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) def test_diagnostic_filter_stamps_records_so_a_second_pass_is_free(monkeypatch, native): - if native: - pytest.importorskip("litellm.rust_bridge._native") monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True) record = _make_record(logging.WARNING, "api_key=secret123") @@ -1775,8 +1765,6 @@ def test_diagnostic_filter_stamps_records_so_a_second_pass_is_free(monkeypatch, @pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) def test_json_formatter_scrubs_unfiltered_extras(monkeypatch, native): - if native: - pytest.importorskip("litellm.rust_bridge._native") monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True) secret = "sk-" + "q" * 48 @@ -1791,8 +1779,6 @@ def test_json_formatter_scrubs_unfiltered_extras(monkeypatch, native): @pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) def test_diagnostic_filter_redacts_a_non_string_message_object(monkeypatch, native): - if native: - pytest.importorskip("litellm.rust_bridge._native") monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True) secret = "sk-" + "q" * 48 diff --git a/tests/unit/test_secret_redaction.py b/tests/unit/test_secret_redaction.py index 3b0e9390a61..cce9c6c558d 100644 --- a/tests/unit/test_secret_redaction.py +++ b/tests/unit/test_secret_redaction.py @@ -76,8 +76,6 @@ def test_redact_string_catches_secret_patterns(): @pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) def test_diagnostic_redaction_policy_matches_across_backends(monkeypatch: pytest.MonkeyPatch, native: bool) -> None: - if native: - pytest.importorskip("litellm.rust_bridge._native") monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") assert redact_string("GET /v1?api_key=abcdefgh12345&page=2") == "GET /v1?REDACTED&page=2" diff --git a/tests/windows_tests/check_windows_wheel_install.py b/tests/windows_tests/check_windows_wheel_install.py index a6c2e7f2984..b554221030d 100644 --- a/tests/windows_tests/check_windows_wheel_install.py +++ b/tests/windows_tests/check_windows_wheel_install.py @@ -87,10 +87,20 @@ def main(argv): f"::error::installing {os.path.basename(wheel)} into a deep prefix failed" ) return 1 - if _run([python, "-c", "import litellm; import litellm.types.utils"]) != 0: + if _run([python, "-I", "-c", "import litellm; import litellm.types.utils"]) != 0: print("::error::litellm did not import after install (half-unpacked package)") return 1 + checker = os.path.abspath(os.path.join("tests", "base_sdk_tests", "check_base_sdk_install.py")) + for rust in ("0", "1"): + result = subprocess.run( + [python, "-I", checker, "--profile", "native"], + env={**os.environ, "LITELLM_RUST": rust, "LITELLM_LOCAL_MODEL_COST_MAP": "True"}, + check=False, + ) + if result.returncode != 0: + return result.returncode + print( f"ok: {os.path.basename(wheel)} installs into a worst-case prefix and imports" ) From aa769fbc6992c9b82f3e46f3cadbaddb5ba7a194 Mon Sep 17 00:00:00 2001 From: Simon <80467011+sorgfresser@users.noreply.github.com> Date: Wed, 28 Oct 2026 21:18:05 +0100 Subject: [PATCH 085/128] fix(streaming): surface upstream abort as server errors (LIT-10159) (#48328) * fix(streaming): surface upstream abort finish reasons as errors * fix(streaming): keep SDK chunks with null choices streaming past the abort check * fix(streaming): treat null choices as empty in the raw SSE abort check * refactor(streaming): bind raw SSE choices as Final instead of rebinding --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../litellm_core_utils/streaming_handler.py | 6 + .../llms/openai/chat/gpt_transformation.py | 10 +- .../test_streaming_handler.py | 145 ++++++++++++++++++ 3 files changed, 159 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index b64e146ca08..3fb25de43cc 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1557,6 +1557,12 @@ class CustomStreamWrapper: else: # openai / azure chat model if isinstance(chunk, ChatCompletionChunk): self._openai_handler_receives_chat_chunks = True + if any(str(choice.finish_reason) == "abort" for choice in chunk.choices or ()): + raise litellm.InternalServerError( + message="Upstream generation aborted (finish_reason=abort)", + model=chunk.model, + llm_provider=self.custom_llm_provider, + ) if self.custom_llm_provider in [ LlmProviders.AZURE.value, LlmProviders.AZURE_AI.value, diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 97e988ef2b1..b7396e8fb03 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -76,6 +76,7 @@ else: _NO_TOOLS_UPDATE: Final[Mapping[str, object]] = MappingProxyType({}) _OBJECT_MAPPING: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) +_CHOICE_MAPPINGS: Final = TypeAdapter(tuple[Mapping[str, object], ...]) _DECODED_JSON: Final = TypeAdapter(object) @@ -958,8 +959,13 @@ class OpenAIChatCompletionStreamingHandler(BaseModelResponseIterator): status_code=error_status_code, message=error_message, ) - choices = chunk.get("choices", []) - choices = self._map_reasoning_to_reasoning_content(choices) + raw_choices: Final = chunk.get("choices") or [] + if any(choice.get("finish_reason") == "abort" for choice in _CHOICE_MAPPINGS.validate_python(raw_choices)): + raise OpenAIError( + status_code=500, + message="Upstream generation aborted (finish_reason=abort)", + ) + choices: Final = self._map_reasoning_to_reasoning_content(raw_choices) fingerprint: Final = _OBJECT_MAPPING.validate_python(chunk).get("system_fingerprint") kwargs: Final[dict[str, Any]] = { diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 6d700bae514..a1b58521d7e 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -1,5 +1,6 @@ import asyncio import datetime +import itertools import json import queue import sys @@ -13,6 +14,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx import pytest +from openai import AsyncOpenAI, AsyncStream, OpenAI, Stream from openai.types.chat import ChatCompletionChunk from pydantic import BaseModel, TypeAdapter import litellm @@ -28,6 +30,7 @@ from litellm.litellm_core_utils.streaming_handler import ( _ProviderChunkParsed, ) from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler from litellm.responses.mcp.chat_completions_handler import _add_mcp_metadata_to_response from litellm.types.utils import ( CompletionTokensDetailsWrapper, @@ -1367,6 +1370,148 @@ async def test_bedrock_validation_error_raises_directly(logging_obj: Logging): assert getattr(excinfo.value, "status_code", None) == 400 +@pytest.mark.parametrize("async_stream", [False, True]) +@pytest.mark.parametrize("transport", ["sdk", "http_handler"]) +@pytest.mark.parametrize("prefix", ["", "Partial answer"]) +@pytest.mark.parametrize("abort_index", [0, 1]) +@pytest.mark.parametrize("abort_content", ["", "unfinished"]) +async def test_aborted_stream_raises_instead_of_finishing_successfully( + async_stream: bool, transport: str, prefix: str, abort_index: int, abort_content: str, logging_obj: Logging +) -> None: + initial: Final = ( + (({"index": 0, "delta": {"role": "assistant", "content": prefix}, "finish_reason": None},),) if prefix else () + ) + final: Final = tuple( + { + "index": index, + "delta": {"content": abort_content}, + "finish_reason": "abort" if index == abort_index else "stop", + } + for index in range(abort_index + 1) + ) + body: Final = ( + "".join( + "data: " + + json.dumps( + { + "id": "chatcmpl-aborted", + "object": "chat.completion.chunk", + "created": 1, + "model": "test-model", + "choices": choices, + } + ) + + "\n\n" + for choices in (*initial, final) + ) + + "data: [DONE]\n\n" + ) + raw_response: Final = httpx.Response( + 200, + content=body, + headers={"content-type": "text/event-stream"}, + request=httpx.Request("POST", "https://provider.example/v1/chat/completions"), + ) + client: Final = AsyncOpenAI(api_key="unused") if async_stream else OpenAI(api_key="unused") + stream: Final = ( + ( + AsyncStream(cast_to=ChatCompletionChunk, response=raw_response, client=client) + if isinstance(client, AsyncOpenAI) + else Stream(cast_to=ChatCompletionChunk, response=raw_response, client=client) + ) + if transport == "sdk" + else OpenAIChatCompletionStreamingHandler( + streaming_response=raw_response.aiter_lines() if async_stream else raw_response.iter_lines(), + sync_stream=not async_stream, + ) + ) + response: Final = CustomStreamWrapper( + completion_stream=stream, + model="test-model", + logging_obj=logging_obj, + custom_llm_provider="openai" if transport == "sdk" else "hosted_vllm", + ) + + try: + with pytest.raises(MidStreamFallbackError, match="abort") as excinfo: + _chunks: Final = [chunk async for chunk in response] if async_stream else list(response) + + assert excinfo.value.is_pre_first_chunk is (not prefix) + assert excinfo.value.generated_content == prefix + assert getattr(excinfo.value.original_exception, "status_code", None) == 500 + finally: + if isinstance(client, AsyncOpenAI): + await client.close() + else: + client.close() + + +@pytest.mark.parametrize("async_stream", [False, True]) +@pytest.mark.parametrize("transport", ["sdk", "http_handler"]) +async def test_stream_with_null_choices_chunk_finishes_normally( + async_stream: bool, transport: str, logging_obj: Logging +) -> None: + body: Final = ( + "".join( + "data: " + + json.dumps( + { + "id": "chatcmpl-null-choices", + "object": "chat.completion.chunk", + "created": 1, + "model": "test-model", + "choices": choices, + **extra, + } + ) + + "\n\n" + for choices, extra in ( + ([{"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": None}], {}), + ([{"index": 0, "delta": {}, "finish_reason": "stop"}], {}), + (None, {"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}}), + ) + ) + + "data: [DONE]\n\n" + ) + raw_response: Final = httpx.Response( + 200, + content=body, + headers={"content-type": "text/event-stream"}, + request=httpx.Request("POST", "https://provider.example/v1/chat/completions"), + ) + client: Final = AsyncOpenAI(api_key="unused") if async_stream else OpenAI(api_key="unused") + stream: Final = ( + ( + AsyncStream(cast_to=ChatCompletionChunk, response=raw_response, client=client) + if isinstance(client, AsyncOpenAI) + else Stream(cast_to=ChatCompletionChunk, response=raw_response, client=client) + ) + if transport == "sdk" + else OpenAIChatCompletionStreamingHandler( + streaming_response=raw_response.aiter_lines() if async_stream else raw_response.iter_lines(), + sync_stream=not async_stream, + ) + ) + response: Final = CustomStreamWrapper( + completion_stream=stream, + model="test-model", + logging_obj=logging_obj, + custom_llm_provider="openai" if transport == "sdk" else "hosted_vllm", + ) + + try: + chunks: Final = [chunk async for chunk in response] if async_stream else list(response) + finally: + if isinstance(client, AsyncOpenAI): + await client.close() + else: + client.close() + + streamed_choices: Final = tuple(itertools.chain.from_iterable(chunk.choices for chunk in chunks)) + assert "".join(choice.delta.content or "" for choice in streamed_choices) == "hi" + assert [choice.finish_reason for choice in streamed_choices if choice.finish_reason] == ["stop"] + + def _hosted_vllm_stream_wrapper(logging_obj: Logging, error_payload: dict) -> CustomStreamWrapper: """A CustomStreamWrapper over the real OpenAI-compatible line iterator, fed an HTTP 200 SSE body that carries an in-body error payload the way From a604c890590f88729b98e02df1d57d6518e00003 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 28 Oct 2026 13:25:53 -0700 Subject: [PATCH 086/128] fix(ui): bundle the Aria avatar so it loads in every deployment (#48890) AriaAvatar pointed at /ui/assets/aria-agent.png, a path that only exists when the proxy serves the static export under /ui. The dashboard dev server and the split-chart nginx image 404 it, so the header Aria button and every other AriaAvatar showed a broken image. Import the PNG statically like the provider logos (#34125) so the bundle always serves it under /_next/static/media --- .../components/liteadmin/AriaAvatar.test.tsx | 21 +++++++++++++++++++ .../src/components/liteadmin/AriaAvatar.tsx | 4 ++-- 2 files changed, 23 insertions(+), 2 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/liteadmin/AriaAvatar.test.tsx diff --git a/ui/litellm-dashboard/src/components/liteadmin/AriaAvatar.test.tsx b/ui/litellm-dashboard/src/components/liteadmin/AriaAvatar.test.tsx new file mode 100644 index 00000000000..5c9cf27470e --- /dev/null +++ b/ui/litellm-dashboard/src/components/liteadmin/AriaAvatar.test.tsx @@ -0,0 +1,21 @@ +import { render, screen } from "@testing-library/react"; +import { afterEach, describe, expect, it } from "vitest"; +import { setServerRootPath } from "@/lib/serverRootPath"; +import { AriaAvatar } from "./AriaAvatar"; + +describe("AriaAvatar", () => { + afterEach(() => { + setServerRootPath("/"); + }); + + it.each(["/", "/litellm"])("loads the bundled image when the server root is %s", (root) => { + setServerRootPath(root); + + render(); + + expect(screen.getByRole("presentation", { hidden: true })).toHaveAttribute( + "src", + "/_next/static/media/aria-agent.png", + ); + }); +}); diff --git a/ui/litellm-dashboard/src/components/liteadmin/AriaAvatar.tsx b/ui/litellm-dashboard/src/components/liteadmin/AriaAvatar.tsx index 8c78908d3a9..e0f014624fe 100644 --- a/ui/litellm-dashboard/src/components/liteadmin/AriaAvatar.tsx +++ b/ui/litellm-dashboard/src/components/liteadmin/AriaAvatar.tsx @@ -1,10 +1,10 @@ -import { resolveLogoSrc } from "@/lib/assetPaths"; +import ariaAgentImage from "../../../public/assets/aria-agent.png"; import { cn } from "@/lib/cva.config"; export function AriaAvatar({ size = 28, className }: { size?: number; className?: string }) { return ( Date: Wed, 28 Oct 2026 13:32:09 -0700 Subject: [PATCH 087/128] fix(responses)!: bill partial usage on mid-stream failures and report bridged /v1/messages failures to the proxy (#47537) A Responses API stream that fails partway now bills the usage it already consumed, ends with an error instead of a clean stop, and honors stream_timeout between chunks. Bridged /v1/messages calls report those failures to the proxy's failure hooks, so spend, budgets and logs see them --- litellm/litellm_core_utils/litellm_logging.py | 7 + .../responses_adapters/streaming_iterator.py | 65 ++- litellm/proxy/common_request_processing.py | 89 +++- litellm/responses/streaming_iterator.py | 78 +++- litellm/router.py | 3 + .../test_exception_mapping_utils.py | 4 +- ...t_responses_adapters_streaming_iterator.py | 196 +++++++++ .../proxy/test_common_request_processing.py | 174 +++++++- .../mcp/test_aresponses_api_with_mcp.py | 59 +++ .../test_responses_streaming_iterator.py | 391 +++++++++++++++++- .../unit/responses/test_streaming_iterator.py | 10 +- tests/unit/test_router/test_router.py | 126 ++++++ 12 files changed, 1155 insertions(+), 47 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 48e8a44b9a4..3b8157b2249 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -842,6 +842,7 @@ class Logging(LiteLLMLoggingBaseClass): self._async_success_enqueued: bool = False self._sync_success_scheduled: bool = False self.on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None + self._stream_failure_claimed: bool = False self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None def copy_for_new_call(self, litellm_call_id: str, start_time: datetime.datetime) -> "Logging": @@ -2561,6 +2562,12 @@ class Logging(LiteLLMLoggingBaseClass): }, ) + def claim_stream_failure_report(self) -> bool: + if self._stream_failure_claimed: + return False + self._stream_failure_claimed = True + return True + def record_assembled_response_for_failure(self, assembled: ModelResponse) -> None: """Bill a fully streamed response on the failure log when a post-call hook rejects it.""" usage: Final = getattr(assembled, "usage", None) diff --git a/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py index 614a7c96848..44f233f2254 100644 --- a/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py @@ -2,15 +2,16 @@ ## Translates OpenAI call to Anthropic `/v1/messages` format import asyncio from collections import deque -from collections.abc import AsyncGenerator, Iterator, Mapping +from collections.abc import AsyncGenerator, Awaitable, Callable, Iterator, Mapping from typing import TYPE_CHECKING, Any, Final +import anyio from pydantic import ConfigDict, field_validator from litellm import verbose_logger from litellm._logging import redact_internal_details_from_client_message from litellm._uuid import uuid -from litellm.exceptions import MidStreamFallbackError +from litellm.exceptions import APIError, MidStreamFallbackError from litellm.litellm_core_utils.prompt_templates.common_utils import ( encrypted_reasoning_signature, ) @@ -22,6 +23,7 @@ from litellm.llms.anthropic.messages.native.utils import ( from litellm.llms.anthropic.messages.safeguards import StreamedSafeguardResults from litellm.llms.anthropic.messages.utils import anthropic_sse_frame from litellm.responses.streaming_iterator import ( + map_stream_error_to_exception, partial_responses_incomplete_event_on_disconnect, stream_error_status_and_message, ) @@ -130,10 +132,12 @@ class AnthropicResponsesStreamWrapper: ) -> None: self.responses_stream = responses_stream self.model = model + self.litellm_logging_obj = litellm_logging_obj self.safeguards_evaluator: Final = safeguards_evaluator self._message_id: str = f"msg_{uuid.uuid4()}" if litellm_logging_obj is not None: litellm_logging_obj.record_streamed_anthropic_message_id(self._message_id) + self._unreported_failure: Exception | None = None self._current_block_index: int = -1 # Map item_id -> content_block_index so we can stop the right block later self._item_id_to_block_index: dict[str, int] = {} @@ -377,16 +381,24 @@ class AnthropicResponsesStreamWrapper: if event_type == "response.failed": failed: Final = _FailedResponseEvent.model_validate(event) - status_code, message = stream_error_status_and_message( - failed.response.error if failed.response is not None else None - ) + upstream_error: Final = failed.response.error if failed.response is not None else None + status_code, message = stream_error_status_and_message(upstream_error) verbose_logger.error( "AnthropicResponsesStreamWrapper: upstream Responses stream for %s failed (%s): %s", self.model, status_code, message, ) - self._fail_stream(status_code, message) + self._fail_stream( + status_code, + message, + failure=map_stream_error_to_exception( + upstream_error, + self.model, + getattr(self.responses_stream, "custom_llm_provider", None) or "", + logging_obj=self.litellm_logging_obj, + ), + ) return # ---- response completed -> message_delta + message_stop ---- @@ -439,10 +451,36 @@ class AnthropicResponsesStreamWrapper: self._sent_message_stop = True return - def _fail_stream(self, status_code: int, message: str) -> None: + def _fail_stream(self, status_code: int, message: str, failure: Exception | None = None) -> None: self._stream_failed = True + self._unreported_failure = failure or APIError( + status_code=status_code, + message=message, + llm_provider=getattr(self.responses_stream, "custom_llm_provider", None), + model=self.model, + ) self._chunk_queue.append(_anthropic_error_chunk(status_code, message)) + async def _report_unreported_failure(self) -> None: + failure: Final = self._unreported_failure + if failure is None: + return + self._unreported_failure = None + on_failure: Final[Callable[[Exception], Awaitable[None]] | None] = getattr( + self.litellm_logging_obj, "_on_detached_stream_failure", None + ) + if on_failure is None: + return + try: + with anyio.CancelScope(shield=True): + await on_failure(failure) + except Exception as hook_failure: # noqa: BLE001 # a failing proxy hook must not hide the error event from the client + verbose_logger.warning( + "AnthropicResponsesStreamWrapper: stream failure hook raised: %s(%s)", + type(hook_failure).__name__, + hook_failure, + ) + def __aiter__(self) -> "AnthropicResponsesStreamWrapper": return self @@ -450,6 +488,7 @@ class AnthropicResponsesStreamWrapper: if self._chunk_queue: return self._chunk_queue.popleft() if self._stream_failed: + await self._report_unreported_failure() raise StopAsyncIteration if not self._sent_message_start: @@ -478,7 +517,7 @@ class AnthropicResponsesStreamWrapper: verbose_logger.exception( "AnthropicResponsesStreamWrapper: upstream Responses stream for %s failed", self.model ) - self._fail_stream(*_failure_status_and_message(e)) + self._fail_stream(*_failure_status_and_message(e), failure=_original_failure(e)) if not self._chunk_queue and not self._sent_message_stop and not self._stream_failed: verbose_logger.error( @@ -498,8 +537,14 @@ class AnthropicResponsesStreamWrapper: async def _sse_frames(self) -> AsyncGenerator[bytes, None]: safeguard_results: Final = StreamedSafeguardResults(self.safeguards_evaluator) - async for chunk in self: - yield anthropic_sse_frame(await safeguard_results.observe(chunk)) + try: + async for chunk in self: + yield anthropic_sse_frame(await safeguard_results.observe(chunk)) + finally: + from litellm.llms.anthropic.messages.native.streaming_iterator import closing_superseded_attempt + + if not closing_superseded_attempt(): + await self._report_unreported_failure() def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None": return partial_responses_incomplete_event_on_disconnect(self.responses_stream) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 24a36a7f2dc..f5d6417fb79 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -4,6 +4,7 @@ import json import logging import math from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime from functools import lru_cache from itertools import chain @@ -136,7 +137,7 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports _check_and_merge_model_level_guardrails, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export check_and_merge_model_level_guardrails, ) -from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect +from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect, spawn_logging_task from litellm.router import Router from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX @@ -2006,6 +2007,49 @@ def _content_text(fields: Mapping[str, object]) -> str: return str(fields.get("content", "")) +async def _no_stream_failure_to_release() -> None: + return None + + +@dataclass(slots=True) +class _ClientStreamFailureReport: + logging_obj: LiteLLMLoggingObj + report: Callable[[Exception], Awaitable[None]] + held: Exception | None = None + + @classmethod + def hold_while_streaming(cls, logging_obj: LiteLLMLoggingObj | None) -> Callable[[], Awaitable[None]]: + if logging_obj is None or logging_obj.on_detached_stream_failure is None: + return _no_stream_failure_to_release + return cls(logging_obj, logging_obj.on_detached_stream_failure)._start_holding() + + def _start_holding(self) -> Callable[[], Awaitable[None]]: + self.logging_obj.on_detached_stream_failure = self.hold + return self.release + + async def hold(self, failure: Exception) -> None: + if self.held is None: + self.held = failure + + async def release(self) -> None: + self.logging_obj.on_detached_stream_failure = self.report + if self.held is None: + return + await asyncio.shield( + spawn_logging_task( + asyncio.get_running_loop(), self._report_held(self.held), task_name="stream failure report" + ) + ) + + async def _report_held(self, failure: Exception) -> None: + try: + await self.report(failure) + except Exception as hook_failure: # noqa: BLE001 # a failing failure hook must not skip the stream cleanup + verbose_proxy_logger.warning( + "Stream failure hook raised: %s(%s)", type(hook_failure).__name__, sanitize_for_log(hook_failure) + ) + + class ProxyBaseLLMRequestProcessing: def __init__(self, data: dict, served_model: str | None = None): self.data = data @@ -3498,6 +3542,8 @@ class ProxyBaseLLMRequestProcessing: request_data: Final = self.data async def _on_detached_stream_failure(exc: Exception) -> None: + if not logging_obj.claim_stream_failure_report(): + return try: await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, @@ -4448,11 +4494,15 @@ class ProxyBaseLLMRequestProcessing: client_disconnected = False delivered_chunk = False recent_tail = SSE_STREAM_START_TAIL # rebind-ok: rolling window over the yielded bytes + stream_logging_obj: Final[object] = request_data.get("litellm_logging_obj") # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # legacy request data is narrowed to Logging below guarded_stream: Final[AsyncGenerator[object, None]] = proxy_logging_obj.async_post_call_streaming_iterator_hook( user_api_key_dict=user_api_key_dict, response=response, request_data=request_data, ) + release_stream_failure: Final = _ClientStreamFailureReport.hold_while_streaming( + stream_logging_obj if isinstance(stream_logging_obj, LiteLLMLoggingObj) else None + ) try: str_so_far = "" async for chunk in guarded_stream: @@ -4523,10 +4573,18 @@ class ProxyBaseLLMRequestProcessing: verbose_proxy_logger.exception( "litellm.proxy.proxy_server.async_data_generator(): Exception occured - %s", e ) - transformed_exception: Final = await proxy_logging_obj.post_call_failure_hook( - user_api_key_dict=user_api_key_dict, - original_exception=e, - request_data=request_data, + already_reported: Final = ( + isinstance(stream_logging_obj, LiteLLMLoggingObj) + and not stream_logging_obj.claim_stream_failure_report() + ) + transformed_exception: Final = ( + None + if already_reported + else await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=e, + request_data=request_data, + ) ) if transformed_exception is not None: e = transformed_exception @@ -4549,15 +4607,18 @@ class ProxyBaseLLMRequestProcessing: seal: Final = "" if seal_open_frame is None else seal_open_frame(recent_tail) yield seal + error_frame if seal else error_frame finally: - await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup( - request=request, - request_data=request_data, - response=response, - stream_completed=stream_completed, - client_disconnected=client_disconnected, - user_api_key_dict=user_api_key_dict, - proxy_logging_obj=proxy_logging_obj, - ) + try: + await release_stream_failure() + finally: + await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup( + request=request, + request_data=request_data, + response=response, + stream_completed=stream_completed, + client_disconnected=client_disconnected, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + ) @staticmethod def async_sse_data_generator( diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index c1f422c7d7b..4e418a91156 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -41,7 +41,11 @@ from litellm.exceptions import MidStreamFallbackError, RateLimitError from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.exception_mapping_utils import with_deployment_api_base -from litellm.litellm_core_utils.internal_call_metadata import without_budget_reservation +from litellm.litellm_core_utils.internal_call_metadata import ( + is_background_response, + requests_background, + without_budget_reservation, +) from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base @@ -65,7 +69,7 @@ from litellm.types.llms.openai import ( ResponsesAPIStreamingResponse, ) from litellm.types.responses.main import GenericResponseOutputItem, OutputText -from litellm.types.utils import CallTypes +from litellm.types.utils import CallTypes, ModelResponse from litellm.utils import async_post_call_success_deployment_hook if TYPE_CHECKING: @@ -288,7 +292,7 @@ def stream_error_status_and_message(error_obj: object) -> tuple[int, str]: return _status_code_for_error_fields(error_type, error_code), message -def _map_stream_error_to_exception( +def map_stream_error_to_exception( error_obj: object, model: str, custom_llm_provider: str, @@ -578,6 +582,7 @@ class BaseResponsesAPIStreamingIterator: self.request_data.get("input"), self.request_data, self._generated_content + self._generated_tool_arguments, + use_default_image_token_count=False, ) if _estimate_wanted else None @@ -731,7 +736,7 @@ class BaseResponsesAPIStreamingIterator: self._earlier_rounds_spend = earlier def _map_error_event_exception(self, error_obj: object) -> Exception: - return _map_stream_error_to_exception( + return map_stream_error_to_exception( error_obj, self.model or "", self.custom_llm_provider or "", @@ -981,6 +986,39 @@ class BaseResponsesAPIStreamingIterator: except Exception: pass + def _record_partial_usage_for_failure(self) -> None: + generated_text: Final = self._generated_content + self._generated_tool_arguments + if ( + not generated_text + or requests_background(self.request_data) + or is_background_response(self._created_response) + or self._completed_response_logged + or self.logging_obj.model_call_details.get("combined_usage_object") is not None + ): + return + estimate: Final = _estimate_usage_safely( + self.model or "", + self.request_data.get("input"), + self.request_data, + generated_text, + use_default_image_token_count=True, + ) + if estimate is None: + return + usage: Final = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # same shared transform the success path above uses + estimate + ) + served_tier: Final[object] = getattr(self._created_response, "service_tier", None) + assembled: Final = ModelResponse( + model=self.model, usage=usage, service_tier=served_tier if isinstance(served_tier, str) else None + ) + if self._earlier_rounds_spend is None: + self.logging_obj.record_assembled_response_for_failure(assembled) + return + self.logging_obj.record_failure_after_earlier_rounds( + self._earlier_rounds_spend, usage, self.logging_obj.response_cost_calculator(result=assembled) or 0.0 + ) + def _handle_failure(self, exception: Exception): """ Trigger failure handlers before bubbling the exception. @@ -990,6 +1028,7 @@ class BaseResponsesAPIStreamingIterator: if self._failure_handled: return self._failure_handled = True + self._record_partial_usage_for_failure() traceback_exception: Final = traceback.format_exc() end_time: Final = datetime.now() @@ -1060,15 +1099,20 @@ class BaseResponsesAPIStreamingIterator: the same shape the success handler assembles for a natural end of stream. Its response is the response.created payload (it carries the served service_tier) plus the text and tool arguments relayed so far, with counted - usage since the provider never sent any. None when nothing was relayed yet - or the stream already reached its terminal event and billed itself. + usage since the provider never sent any. None when nothing was relayed yet, + the stream already reached its terminal event and billed itself, or its failure + log already billed the relayed output. """ created: Final = self._created_response generated: Final = self._generated_content + self._generated_tool_arguments - if self.completed_response is not None or created is None or not generated: + if created is None or not generated or self._billed_without_disconnect(): return None usage: Final = _estimate_usage_safely( - self.model or "", self.request_data.get("input"), self.request_data, generated + self.model or "", + self.request_data.get("input"), + self.request_data, + generated, + use_default_image_token_count=False, ) if usage is None: return None @@ -1082,6 +1126,9 @@ class BaseResponsesAPIStreamingIterator: stamp_responses_usage_cost(partial, self.logging_obj) return ResponseIncompleteEvent(type=ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, response=partial) + def _billed_without_disconnect(self) -> bool: + return self.completed_response is not None or self._failure_handled + def _note_yielded_event(self, event: ResponsesAPIStreamingResponse) -> None: self._yielded_first_chunk = True if event.type not in PRE_OUTPUT_LIFECYCLE_EVENT_TYPES: @@ -1218,6 +1265,9 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): raise StopAsyncIteration from e self._raise_for_transport_error(e) except httpx.TimeoutException as e: + if self.completed_response is not None: + self.finished = True + raise StopAsyncIteration from e self._raise_for_transport_timeout(e) except httpx.HTTPError as e: # Handle HTTP errors @@ -1303,6 +1353,9 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): raise StopIteration from e self._raise_for_transport_error(e) except httpx.TimeoutException as e: + if self.completed_response is not None: + self.finished = True + raise StopIteration from e self._raise_for_transport_timeout(e) except httpx.HTTPError as e: # Handle HTTP errors @@ -1727,6 +1780,8 @@ def _estimate_usage_from_text( request_input: object, responses_api_request: Mapping[str, object], generated_text: str, + *, + use_default_image_token_count: bool = False, ) -> ResponseAPIUsage: messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # the transformer's signature is partially untyped input=request_input, # pyright: ignore[reportArgumentType] # the raw Responses API input is a str or ResponseInputParam list, matching the helper's declared union @@ -1734,7 +1789,7 @@ def _estimate_usage_from_text( only_since_compaction=True, ) input_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped - model=model, messages=messages + model=model, messages=messages, use_default_image_token_count=use_default_image_token_count ) output_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped model=model, text=generated_text, count_response_tokens=True @@ -1751,6 +1806,8 @@ def _estimate_usage_safely( request_input: object, responses_api_request: Mapping[str, object], generated_text: str, + *, + use_default_image_token_count: bool, ) -> ResponseAPIUsage | None: try: return _estimate_usage_from_text( @@ -1758,6 +1815,7 @@ def _estimate_usage_safely( request_input=request_input, responses_api_request=responses_api_request, generated_text=generated_text, + use_default_image_token_count=use_default_image_token_count, ) except Exception as e: verbose_logger.debug("Could not estimate usage from stream text, billing $0: %s", e) @@ -2276,7 +2334,7 @@ class _ResponsesWsTurn: ) if failed_event is None: return None - return _map_stream_error_to_exception( + return map_stream_error_to_exception( _ws_event_error(failed_event), model, custom_llm_provider, logging_obj=connection_logging_obj ) diff --git a/litellm/router.py b/litellm/router.py index d712ed4ac52..5b18d36316f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6456,6 +6456,9 @@ class Router: hop_kwargs: Final = mid_stream_fallback_hop_kwargs( model=model, original_generic_function=original_generic_function, controls=controls, kwargs=kwargs ) + shared_logging: Final = kwargs.get("litellm_logging_obj") + if isinstance(shared_logging, LiteLLMLogging): + shared_logging.carry_failed_attempt_usage() response: Final = await self._ageneric_api_call_with_fallbacks_helper( model=model, original_generic_function=original_generic_function, **kwargs ) diff --git a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py index 9e8e9cf2a80..55e7c47e2c9 100644 --- a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py @@ -1516,9 +1516,9 @@ def test_bedrock_mantle_openai_envelope_context_overflow_maps_to_context_window_ def test_bedrock_mantle_streamed_context_overflow_event_maps_to_context_window_exceeded(): - from litellm.responses.streaming_iterator import _map_stream_error_to_exception + from litellm.responses.streaming_iterator import map_stream_error_to_exception - mapped_exception = _map_stream_error_to_exception( + mapped_exception = map_stream_error_to_exception( { "code": "context_length_exceeded", "message": "Your input exceeds the context window of this model. Please adjust your input and try again.", diff --git a/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_streaming_iterator.py b/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_streaming_iterator.py index 5af8e69a0ea..4626f76c8fc 100644 --- a/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_streaming_iterator.py +++ b/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_streaming_iterator.py @@ -4,20 +4,25 @@ Tests for AnthropicResponsesStreamWrapper """ import asyncio +import datetime import json import os import sys +from collections.abc import AsyncIterator, Awaitable, Callable from types import SimpleNamespace +import anyio import pytest sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../../.."))) import litellm from litellm.exceptions import MidStreamFallbackError +from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.litellm_core_utils.prompt_templates.common_utils import ( encrypted_reasoning_signature, ) +from litellm.llms.anthropic.messages.native.streaming_iterator import aclose_superseded_attempt from litellm.llms.anthropic.messages.native.utils import INCOMPLETE_STREAM_ERROR_MESSAGE from litellm.llms.anthropic.messages.responses_adapters.streaming_iterator import ( AnthropicResponsesStreamWrapper, @@ -622,3 +627,194 @@ class TestUpstreamFailureEndsStreamWithErrorEvent: chunks = _collect(iter(())) assert [chunk["type"] for chunk in chunks] == ["message_start", "error"] assert chunks[1]["error"] == {"type": "api_error", "message": INCOMPLETE_STREAM_ERROR_MESSAGE} + + +class _FailureHookRecorder: + def __init__(self) -> None: + self.received: tuple[Exception, ...] = () + + async def __call__(self, failure: Exception) -> None: + self.received = (*self.received, failure) + + +def _logging_obj_with_failure_hook(on_failure: Callable[[Exception], Awaitable[None]]) -> LitellmLogging: + logging_obj = LitellmLogging( + model="m", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.datetime(2026, 1, 1), + litellm_call_id="lit-6871-test", + function_id="lit-6871-test", + ) + logging_obj._on_detached_stream_failure = on_failure + return logging_obj + + +def _collect_with_logging( + stream: AsyncIterator[dict[str, object]], logging_obj: LitellmLogging +) -> list[dict[str, object]]: + async def _run() -> list[dict[str, object]]: + wrapper = AnthropicResponsesStreamWrapper(responses_stream=stream, model="m", litellm_logging_obj=logging_obj) + return [chunk async for chunk in wrapper] + + return asyncio.run(_run()) + + +class _DeliveryCountingConsumer: + def __init__(self) -> None: + self.delivered: tuple[dict[str, object], ...] = () + self.delivered_when_hook_fired: tuple[int, ...] = () + + async def hook(self, failure: Exception) -> None: + self.delivered_when_hook_fired = (*self.delivered_when_hook_fired, len(self.delivered)) + + async def drain(self, wrapper: AnthropicResponsesStreamWrapper) -> None: + async for chunk in wrapper: + self.delivered = (*self.delivered, chunk) + + +def _failing_after_one_delta() -> AsyncIterator[dict[str, object]]: + async def _gen() -> AsyncIterator[dict[str, object]]: + yield {"type": "response.created"} + yield {"type": "response.output_item.added", "item": {"type": "message", "id": "msg_1"}} + yield {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "Hi"} + raise litellm.RateLimitError(message="You have no credits remaining.", llm_provider="openai", model="m") + + return _gen() + + +class TestUpstreamFailureReachesTheProxyFailureHook: + def test_raised_upstream_failure_reaches_the_hook_unwrapped_once_the_error_event_was_consumed(self): + rate_limit = litellm.RateLimitError(message="You have no credits remaining.", llm_provider="openai", model="m") + wrapped = MidStreamFallbackError( + message=str(rate_limit), + model="m", + llm_provider="openai", + original_exception=rate_limit, + is_pre_first_chunk=False, + ) + recorder = _FailureHookRecorder() + + async def _gen() -> AsyncIterator[dict[str, object]]: + yield {"type": "response.created"} + yield {"type": "response.output_item.added", "item": {"type": "message", "id": "msg_1"}} + yield {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "Hi"} + raise wrapped + + chunks = _collect_with_logging(_gen(), _logging_obj_with_failure_hook(recorder)) + assert [chunk["type"] for chunk in chunks] == [ + "message_start", + "content_block_start", + "content_block_delta", + "error", + ] + assert recorder.received == (rate_limit,) + + def test_response_failed_event_reaches_the_hook_as_an_exception_carrying_the_upstream_status(self): + recorder = _FailureHookRecorder() + failed = SimpleNamespace( + status="failed", + output=[], + usage=None, + error={"code": "rate_limit_exceeded", "message": "Rate limit reached for gpt-5.5, try again in 20s."}, + ) + + async def _gen() -> AsyncIterator[dict[str, object]]: + yield {"type": "response.created"} + yield {"type": "response.failed", "response": failed} + + chunks = _collect_with_logging(_gen(), _logging_obj_with_failure_hook(recorder)) + assert chunks[-1]["type"] == "error" + (failure,) = recorder.received + assert isinstance(failure, litellm.RateLimitError) + assert failure.status_code == 429 + assert "Rate limit reached for gpt-5.5, try again in 20s." in failure.message + + def test_a_raising_hook_never_hides_the_error_event_from_the_client(self): + async def _exploding_hook(failure: Exception) -> None: + raise RuntimeError("spend writer is down") + + logging_obj = _logging_obj_with_failure_hook(_exploding_hook) + + async def _gen() -> AsyncIterator[dict[str, object]]: + yield {"type": "response.created"} + raise ConnectionResetError("Response payload is not completed") + + chunks = _collect_with_logging(_gen(), logging_obj) + assert [chunk["type"] for chunk in chunks] == ["message_start", "error"] + assert chunks[-1]["error"] == {"type": "api_error", "message": "Response payload is not completed"} + + def test_the_hook_fires_only_after_the_consumer_took_the_error_event(self): + consumer = _DeliveryCountingConsumer() + + async def _run() -> None: + wrapper = AnthropicResponsesStreamWrapper( + responses_stream=_failing_after_one_delta(), + model="m", + litellm_logging_obj=_logging_obj_with_failure_hook(consumer.hook), + ) + await consumer.drain(wrapper) + with pytest.raises(StopAsyncIteration): + await wrapper.__anext__() + + asyncio.run(_run()) + assert [chunk["type"] for chunk in consumer.delivered] == [ + "message_start", + "content_block_start", + "content_block_delta", + "error", + ] + assert consumer.delivered_when_hook_fired == (len(consumer.delivered),) + + @pytest.mark.parametrize( + ("close", "expected_reports"), + [ + pytest.param(lambda sse: sse.aclose(), 1, id="client_closes"), + pytest.param(aclose_superseded_attempt, 0, id="fallback_takes_over"), + ], + ) + def test_closing_the_stream_on_the_error_event_reports_unless_a_fallback_takes_over(self, close, expected_reports): + recorder = _FailureHookRecorder() + + async def _close_on_the_error_event() -> None: + wrapper = AnthropicResponsesStreamWrapper( + responses_stream=_failing_after_one_delta(), model="m", litellm_logging_obj=_logging_obj_with_failure_hook(recorder) + ) + sse = wrapper.async_anthropic_sse_wrapper() + async for frame in sse: + if frame.startswith(b"event: error\n"): + await close(sse) + return + pytest.fail("the stream ended without an error event") + + asyncio.run(_close_on_the_error_event()) + assert len(recorder.received) == expected_reports + + def test_a_client_disconnect_during_the_failure_hook_lets_the_hook_finish(self): + hook_started = anyio.Event() + hook_released = anyio.Event() + finished: list[Exception] = [] + + async def _held_hook(failure: Exception) -> None: + hook_started.set() + await hook_released.wait() + finished.append(failure) + + async def _disconnect_mid_hook() -> None: + wrapper = AnthropicResponsesStreamWrapper( + responses_stream=_failing_after_one_delta(), model="m", litellm_logging_obj=_logging_obj_with_failure_hook(_held_hook) + ) + + async def _consume() -> None: + async for _ in wrapper.async_anthropic_sse_wrapper(): + pass + + async with anyio.create_task_group() as task_group: + task_group.start_soon(_consume) + await hook_started.wait() + task_group.cancel_scope.cancel() + hook_released.set() + + asyncio.run(_disconnect_mid_hook()) + assert len(finished) == 1 diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 52816dc50a8..790790459c3 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -10826,12 +10826,19 @@ class TestResolveLitellmCallId: class _FailureHookRecorder: """Stands in for ProxyLogging.post_call_failure_hook, recording what the detached-failure closure hands it.""" - def __init__(self, raises: Optional[Exception] = None): + def __init__(self, raises: Optional[Exception] = None, gate: Optional[asyncio.Event] = None): self.calls = [] + self.entered = asyncio.Event() + self.reported = asyncio.Event() self._raises = raises + self._gate = gate async def post_call_failure_hook(self, **kwargs): + self.entered.set() + if self._gate is not None: + await self._gate.wait() self.calls.append(kwargs) + self.reported.set() if self._raises is not None: raise self._raises @@ -10932,6 +10939,171 @@ class TestDetachedStreamFailureHook: assert [call["original_exception"] for call in recorder.calls] == [failure] + @staticmethod + def _armed_stream(recorder: _FailureHookRecorder, iterator_hook_for, response: Optional[MagicMock] = None): + from litellm.proxy._types import UserAPIKeyAuth as RealUserAPIKeyAuth + + logging_obj = LiteLLMLoggingObj( + model="gpt-6.1-sol", + messages=[{"role": "user", "content": "Write a story."}], + stream=True, + call_type="anthropic_messages", + start_time=None, + litellm_call_id="call-lit6871", + function_id="function-lit6871", + ) + user_api_key_dict = RealUserAPIKeyAuth(api_key="sk-test") + processing_obj = ProxyBaseLLMRequestProcessing(data={"litellm_logging_obj": logging_obj}) + ProxyLogging._callback_capabilities_cache.clear() + proxy_logging_obj = ProxyLogging(user_api_key_cache=MagicMock()) + proxy_logging_obj.post_call_failure_hook = recorder.post_call_failure_hook + proxy_logging_obj.async_post_call_streaming_iterator_hook = iterator_hook_for(logging_obj) + processing_obj._arm_detached_stream_failure_hook( + logging_obj=logging_obj, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + ) + stream = ProxyBaseLLMRequestProcessing.async_streaming_data_generator( + response=MagicMock() if response is None else response, + user_api_key_dict=user_api_key_dict, + request_data=processing_obj.data, + proxy_logging_obj=proxy_logging_obj, + serialize_chunk=lambda chunk: "data: chunk\n\n", + serialize_error=lambda error: "data: error\n\n", + ) + return logging_obj, stream + + @staticmethod + def _hook_reporting(upstream_failure: Exception, after_report=None): + def iterator_hook_for(logging_obj): + async def iterator_hook(**_kwargs): + yield b"event: message_start\n\n" + await logging_obj.on_detached_stream_failure(upstream_failure) + yield b"event: error\n\n" + yield b"event: message_stop\n\n" + if after_report is not None: + raise after_report + + return iterator_hook + + return iterator_hook_for + + @pytest.mark.asyncio + async def test_iterator_hook_failure_after_a_bridged_failure_is_the_one_the_failure_hook_gets(self): + recorder = _FailureHookRecorder() + upstream_failure = RuntimeError("upstream Responses stream failed") + hook_failure = RuntimeError("iterator hook rejected the stream") + + _, stream = self._armed_stream(recorder, self._hook_reporting(upstream_failure, after_report=hook_failure)) + + frames = [frame async for frame in stream] + + assert frames == ["data: chunk\n\n"] * 3 + ["data: error\n\n"] + assert [call["original_exception"] for call in recorder.calls] == [hook_failure] + + @pytest.mark.asyncio + async def test_failure_hook_replacing_the_iterator_hook_failure_aborts_the_stream(self): + replacement = HTTPException(status_code=418, detail="rejected by a custom failure hook") + recorder = _FailureHookRecorder(raises=replacement) + upstream_failure = RuntimeError("upstream Responses stream failed") + hook_failure = RuntimeError("iterator hook rejected the stream") + + _, stream = self._armed_stream(recorder, self._hook_reporting(upstream_failure, after_report=hook_failure)) + + frames = [await stream.__anext__() for _ in range(3)] + with pytest.raises(HTTPException) as raised: + await stream.__anext__() + + assert raised.value is replacement + assert frames == ["data: chunk\n\n"] * 3 + assert [call["original_exception"] for call in recorder.calls] == [hook_failure] + + @pytest.mark.asyncio + async def test_failure_reported_by_the_stream_reaches_the_failure_hook_once_after_the_stream_ends(self): + recorder = _FailureHookRecorder() + upstream_failure = RuntimeError("upstream Responses stream failed") + + _, stream = self._armed_stream(recorder, self._hook_reporting(upstream_failure)) + + frames = [] + async for frame in stream: + frames.append(frame) + assert recorder.calls == [] + + assert frames == ["data: chunk\n\n"] * 3 + assert [call["original_exception"] for call in recorder.calls] == [upstream_failure] + + @pytest.mark.asyncio + async def test_failure_report_and_stream_cleanup_survive_the_client_task_being_cancelled_mid_report(self): + gate = asyncio.Event() + recorder = _FailureHookRecorder(gate=gate) + upstream_failure = RuntimeError("upstream Responses stream failed") + response = MagicMock() + response.aclose = AsyncMock() + + _, stream = self._armed_stream(recorder, self._hook_reporting(upstream_failure), response=response) + + async def client_reads_the_whole_stream(): + return [frame async for frame in stream] + + client = asyncio.create_task(client_reads_the_whole_stream()) + await asyncio.wait_for(recorder.entered.wait(), timeout=5) + client.cancel() + with pytest.raises(asyncio.CancelledError): + await client + + response.aclose.assert_awaited_once() + gate.set() + await asyncio.wait_for(recorder.reported.wait(), timeout=5) + assert [call["original_exception"] for call in recorder.calls] == [upstream_failure] + + @pytest.mark.asyncio + async def test_failure_reported_by_the_stream_reaches_the_failure_hook_once_when_the_client_leaves_early(self): + recorder = _FailureHookRecorder() + upstream_failure = RuntimeError("upstream Responses stream failed") + + _, stream = self._armed_stream(recorder, self._hook_reporting(upstream_failure)) + + await stream.__anext__() + await stream.__anext__() + await stream.aclose() + + assert [call["original_exception"] for call in recorder.calls] == [upstream_failure] + + @pytest.mark.asyncio + async def test_failure_reported_after_the_stream_ends_still_reaches_the_failure_hook(self): + recorder = _FailureHookRecorder() + upstream_failure = RuntimeError("upstream Responses stream failed after the client stream closed") + + async def iterator_hook(**_kwargs): + yield b"event: message_start\n\n" + + logging_obj, stream = self._armed_stream(recorder, lambda _logging_obj: iterator_hook) + + frames = [frame async for frame in stream] + await logging_obj.on_detached_stream_failure(upstream_failure) + + assert frames == ["data: chunk\n\n"] + assert [call["original_exception"] for call in recorder.calls] == [upstream_failure] + + @pytest.mark.asyncio + async def test_failure_reported_by_an_iterator_hook_is_not_reported_again_when_the_stream_closes(self): + recorder = _FailureHookRecorder() + upstream_failure = RuntimeError("upstream Responses stream failed") + hook_failure = RuntimeError("iterator hook rejected the error frame") + + async def iterator_hook(**_kwargs): + yield b"event: message_start\n\n" + raise hook_failure + + logging_obj, stream = self._armed_stream(recorder, lambda _logging_obj: iterator_hook) + + frames = [frame async for frame in stream] + await logging_obj.on_detached_stream_failure(upstream_failure) + + assert frames == ["data: chunk\n\n", "data: error\n\n"] + assert [call["original_exception"] for call in recorder.calls] == [hook_failure] + class TestPostCallMaskedOutputReachesDeferredLogging: @pytest.mark.asyncio diff --git a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py index 4cbcc6c1519..ec0be896095 100644 --- a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py +++ b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py @@ -2668,6 +2668,65 @@ async def test_stream_auto_execute_bills_the_paid_first_round_on_its_success_row assert _request_failure_row(logging_obj, recorder.logged[0]).get("response_cost", 0.0) == 0.0 +class _CutAfterStreamedText(httpx.AsyncByteStream): + async def __aiter__(self): + opened: Final = {**json.loads(json.dumps(_ROUND_TWO_BODY, default=dict)), "status": "in_progress"} + yield _sse( + ( + {"type": "response.created", "response": opened, "sequence_number": 0}, + { + "type": "response.output_text.delta", + "item_id": "msg_round_two", + "output_index": 0, + "content_index": 0, + "delta": "PING " * 40, + "sequence_number": 1, + }, + ) + ) + raise httpx.ReadError("connection reset mid-stream") + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_a_follow_up_cut_after_streamed_text_with_no_second_fixed_margin( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, + monkeypatch, + (_streamed(_ROUND_ONE_BODY), httpx.Response(200, stream=_CutAfterStreamedText(), headers=SSE_HEADERS)), + ) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-stream-follow-up-cut-after-text") + + with pytest.raises(httpx.ReadError): + await _stream_auto_execute(logging_obj) + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + + _assert_one_row_bills(recorder.logged, _ROUND_ONE_BODY) + assert len(recorder.failed) == 1 + failed: Final = recorder.failed[0] + assert failed["completion_tokens"] > 0 + priced_with_its_own_margin: Final = litellm.completion_cost( + completion_response=ModelResponse( + model="gpt-5.6", + usage=Usage( + prompt_tokens=failed["prompt_tokens"], + completion_tokens=failed["completion_tokens"], + total_tokens=failed["prompt_tokens"] + failed["completion_tokens"], + ), + ), + model="gpt-5.6", + custom_llm_provider="openai", + ) + assert failed["response_cost"] == pytest.approx(priced_with_its_own_margin - _FIXED_MARGIN) + _assert_breakdown_agrees_with_the_billed_cost(failed) + + @pytest.mark.asyncio async def test_stream_auto_execute_blocked_by_a_post_call_guardrail_keeps_both_rounds_on_the_success_row( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch diff --git a/tests/unit/responses/test_responses_streaming_iterator.py b/tests/unit/responses/test_responses_streaming_iterator.py index 37a58c32f38..ae688ff9173 100644 --- a/tests/unit/responses/test_responses_streaming_iterator.py +++ b/tests/unit/responses/test_responses_streaming_iterator.py @@ -9,20 +9,23 @@ crashed with exit 139 whenever any CustomLogger was registered). import asyncio import json import time -from collections.abc import Mapping -from typing import Final +from collections.abc import AsyncIterator, Mapping +from datetime import datetime +from typing import Final, cast import httpx import pytest import litellm +from litellm.exceptions import MidStreamFallbackError from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils import thread_pool_executor as thread_pool_executor_module -from litellm.responses import streaming_iterator as responses_streaming_iterator_module from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig +from litellm.responses import streaming_iterator as responses_streaming_iterator_module from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator -from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesAPIStreamingResponse +from litellm.types.utils import StandardLoggingPayload class RecordingCustomLogger(CustomLogger): @@ -169,6 +172,240 @@ async def test_sync_callbacks_run_only_after_async_handler_completes(recording_e assert submit_times[0] >= recorder.async_hook_finished +class FailureRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.failure_payloads: tuple[StandardLoggingPayload, ...] = () + self.success_payloads: tuple[StandardLoggingPayload, ...] = () + + async def async_log_failure_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.failure_payloads = ( + *self.failure_payloads, + cast(StandardLoggingPayload, kwargs["standard_logging_object"]), + ) + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.success_payloads = ( + *self.success_payloads, + cast(StandardLoggingPayload, kwargs["standard_logging_object"]), + ) + + +class _ChunkRecorder: + def __init__(self) -> None: + self.chunks: tuple[ResponsesAPIStreamingResponse, ...] = () + + async def drain(self, iterator: ResponsesAPIStreamingIterator) -> None: + async for chunk in iterator: + self.chunks = (*self.chunks, chunk) + + +_SEA_DELTAS = ("The sea ", "is wide ", "and deep.") +_SEA_PROMPT = "Write a 300 word story about the sea." +_UNREACHABLE_IMAGE_INPUT: list[dict[str, object]] = [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "Describe this picture."}, + {"type": "input_image", "image_url": "http://images.invalid/photo.png", "detail": "high"}, + ], + } +] + + +_TOOL_ARGUMENT_DELTAS = ('{"city": ', '"Lisbon", ', '"unit": "celsius"}') + + +class _UpstreamThatTimesOutAfterThreeDeltas: + headers: dict[str, str] = {} + + def __init__(self, event_type: str = "response.output_text.delta", deltas: tuple[str, ...] = _SEA_DELTAS) -> None: + self.event_type = event_type + self.deltas = deltas + + async def aiter_bytes(self) -> AsyncIterator[bytes]: + for sequence_number, delta in enumerate(self.deltas): + event = { + "type": self.event_type, + "item_id": "item_1", + "output_index": 0, + "content_index": 0, + "delta": delta, + "sequence_number": sequence_number, + } + yield f"data: {json.dumps(event)}\n\n".encode() + raise httpx.ReadTimeout("Timeout on reading data from socket") + + +class _UpstreamThatFailsWithUsageAfterThreeDeltas: + headers: dict[str, str] = {} + + async def aiter_bytes(self) -> AsyncIterator[bytes]: + for sequence_number, delta in enumerate(_SEA_DELTAS): + event = { + "type": "response.output_text.delta", + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": delta, + "sequence_number": sequence_number, + } + yield f"data: {json.dumps(event)}\n\n".encode() + failed = { + "type": "response.failed", + "sequence_number": len(_SEA_DELTAS), + "response": { + "id": "resp_1", + "object": "response", + "created_at": 1700000000, + "status": "failed", + "model": "gpt-5.4-nano", + "output": [], + "error": {"code": "server_error", "message": "The model had an internal error"}, + "usage": {"input_tokens": 400, "output_tokens": 9, "total_tokens": 409}, + }, + } + yield f"data: {json.dumps(failed)}\n\n".encode() + + +async def _failure_payload_after_mid_stream_timeout( + monkeypatch: pytest.MonkeyPatch, + request_input: str | list[dict[str, object]], + upstream: _UpstreamThatTimesOutAfterThreeDeltas | None = None, + background_fields: Mapping[str, object] | None = None, +) -> tuple[FailureRecorder, StandardLoggingPayload]: + upstream_response: Final = upstream or _UpstreamThatTimesOutAfterThreeDeltas() + recorder = FailureRecorder() + monkeypatch.setattr(litellm, "failure_callback", [recorder]) + monkeypatch.setattr(litellm, "_async_failure_callback", [recorder]) + monkeypatch.setattr(litellm, "success_callback", [recorder]) + monkeypatch.setattr(litellm, "_async_success_callback", [recorder]) + logging_obj = _make_logging_obj() + logging_obj.update_environment_variables( + model="gpt-5.4-nano", + optional_params={"stream": True}, + litellm_params={"custom_llm_provider": "openai", "aresponses": True}, + ) + iterator = ResponsesAPIStreamingIterator( + response=cast(httpx.Response, upstream_response), + model="gpt-5.4-nano", + responses_api_provider_config=OpenAIResponsesAPIConfig(), + logging_obj=logging_obj, + custom_llm_provider="openai", + request_data={"model": "gpt-5.4-nano", "input": request_input, "stream": True, **(background_fields or {})}, + ) + received = _ChunkRecorder() + + with pytest.raises(litellm.Timeout, match="Timeout on reading data from socket"): + await received.drain(iterator) + assert [getattr(chunk, "delta", None) for chunk in received.chunks] == list(upstream_response.deltas) + + await iterator._await_pending_logging() + (payload,) = recorder.failure_payloads + return recorder, payload + + +@pytest.mark.asyncio +async def test_mid_stream_read_timeout_logs_failure_with_partial_usage_and_no_success( + monkeypatch: pytest.MonkeyPatch, +) -> None: + recorder, payload = await _failure_payload_after_mid_stream_timeout(monkeypatch, _SEA_PROMPT) + + assert payload["status"] == "failure" + assert payload["prompt_tokens"] == litellm.token_counter( + model="gpt-5.4-nano", messages=[{"role": "user", "content": _SEA_PROMPT}] + ) + assert payload["completion_tokens"] == litellm.token_counter( + model="gpt-5.4-nano", text="".join(_SEA_DELTAS), count_response_tokens=True + ) + assert payload["response_cost"] > 0 + assert recorder.success_payloads == () + + +@pytest.mark.asyncio +async def test_mid_stream_read_timeout_bills_the_tool_call_arguments_already_streamed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + upstream: Final = _UpstreamThatTimesOutAfterThreeDeltas( + event_type="response.function_call_arguments.delta", deltas=_TOOL_ARGUMENT_DELTAS + ) + _, payload = await _failure_payload_after_mid_stream_timeout(monkeypatch, _SEA_PROMPT, upstream) + + assert payload["status"] == "failure" + assert payload["completion_tokens"] == litellm.token_counter( + model="gpt-5.4-nano", text="".join(_TOOL_ARGUMENT_DELTAS), count_response_tokens=True + ) + assert payload["response_cost"] > 0 + + +@pytest.mark.asyncio +async def test_mid_stream_failure_of_a_background_response_bills_nothing_since_retrieving_it_bills_the_whole_run( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _, payload = await _failure_payload_after_mid_stream_timeout( + monkeypatch, _SEA_PROMPT, background_fields={"background": True} + ) + + assert payload["status"] == "failure" + assert payload["completion_tokens"] == 0 + assert payload["response_cost"] == 0 + + +@pytest.mark.asyncio +async def test_mid_stream_failure_usage_estimate_never_fetches_image_dimensions( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _, payload = await _failure_payload_after_mid_stream_timeout(monkeypatch, _UNREACHABLE_IMAGE_INPUT) + + assert payload["status"] == "failure" + assert payload["prompt_tokens"] == litellm.token_counter( + model="gpt-5.4-nano", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this picture."}, + {"type": "image_url", "image_url": {"url": "http://images.invalid/photo.png", "detail": "high"}}, + ], + } + ], + use_default_image_token_count=True, + ) + assert payload["completion_tokens"] > 0 + + +@pytest.mark.asyncio +async def test_mid_stream_failure_keeps_the_usage_the_provider_reported(monkeypatch: pytest.MonkeyPatch) -> None: + recorder = FailureRecorder() + monkeypatch.setattr(litellm, "failure_callback", [recorder]) + monkeypatch.setattr(litellm, "_async_failure_callback", [recorder]) + logging_obj = _make_logging_obj() + logging_obj.update_environment_variables( + model="gpt-5.4-nano", + optional_params={"stream": True}, + litellm_params={"custom_llm_provider": "openai", "aresponses": True}, + ) + iterator = ResponsesAPIStreamingIterator( + response=cast(httpx.Response, _UpstreamThatFailsWithUsageAfterThreeDeltas()), + model="gpt-5.4-nano", + responses_api_provider_config=OpenAIResponsesAPIConfig(), + logging_obj=logging_obj, + custom_llm_provider="openai", + request_data={"model": "gpt-5.4-nano", "input": _SEA_PROMPT, "stream": True}, + ) + + with pytest.raises(MidStreamFallbackError, match="The model had an internal error"): + await _ChunkRecorder().drain(iterator) + await iterator._await_pending_logging() + + (payload,) = recorder.failure_payloads + assert (payload["prompt_tokens"], payload["completion_tokens"]) == (400, 9) + + _CREATED_RESPONSE_WITH_PRIORITY_TIER: Final[Mapping[str, object]] = { "id": "resp_partial", "object": "response", @@ -211,17 +448,135 @@ def _completed_after(text: str) -> Mapping[str, object]: } -def _live_iterator(logging_obj: LitellmLogging) -> ResponsesAPIStreamingIterator: +def _live_iterator( + logging_obj: LitellmLogging, upstream: object = None, background: bool = False +) -> ResponsesAPIStreamingIterator: return ResponsesAPIStreamingIterator( - response=httpx.Response(200), + response=cast(httpx.Response, upstream or httpx.Response(200)), model="gpt-5.4-nano", responses_api_provider_config=OpenAIResponsesAPIConfig(), logging_obj=logging_obj, - request_data={"model": "gpt-5.4-nano", "input": "tell me a story"}, + request_data={"model": "gpt-5.4-nano", "input": "tell me a story", "background": background}, custom_llm_provider="openai", ) +class _UpstreamThatTimesOutAfterRelayingText: + headers: dict[str, str] = {} + + async def aiter_bytes(self) -> AsyncIterator[bytes]: + for event in (_CREATED_WITH_PRIORITY_TIER, _text_delta("Once upon a time", 1)): + yield f"data: {json.dumps(event)}\n\n".encode() + raise httpx.ReadTimeout("Timeout on reading data from socket") + + +class _UpstreamThatCompletesThenTimesOut: + headers: dict[str, str] = {} + + async def aiter_bytes(self) -> AsyncIterator[bytes]: + events: Final = ( + _CREATED_WITH_PRIORITY_TIER, + *(_text_delta(delta, sequence_number) for sequence_number, delta in enumerate(_SEA_DELTAS, start=1)), + _completed_after("".join(_SEA_DELTAS)), + ) + for event in events: + yield f"data: {json.dumps(event)}\n\n".encode() + raise httpx.ReadTimeout("Timeout on reading data from socket") + + +@pytest.mark.asyncio +async def test_timeout_after_completed_response_ends_the_stream_and_bills_it_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + recorder: Final = FailureRecorder() + monkeypatch.setattr(litellm, "failure_callback", [recorder]) + monkeypatch.setattr(litellm, "_async_failure_callback", [recorder]) + monkeypatch.setattr(litellm, "success_callback", [recorder]) + monkeypatch.setattr(litellm, "_async_success_callback", [recorder]) + logging_obj: Final = _make_logging_obj() + logging_obj.update_environment_variables( + model="gpt-5.4-nano", + optional_params={"stream": True}, + litellm_params={"custom_llm_provider": "openai", "aresponses": True}, + ) + iterator: Final = _live_iterator(logging_obj, upstream=_UpstreamThatCompletesThenTimesOut()) + + chunks: Final = _ChunkRecorder() + await chunks.drain(iterator) + await iterator._await_pending_logging() + + assert chunks.chunks[-1].type == "response.completed" + assert len(recorder.success_payloads) == 1 + success_payload: Final = recorder.success_payloads[0] + response_payload: Final = success_payload["response"] + assert isinstance(response_payload, dict) + usage: Final = cast(dict[str, object], response_payload["usage"]) + assert (usage["prompt_tokens"], usage["completion_tokens"], usage["total_tokens"]) == (5, 7, 12) + assert recorder.failure_payloads == () + assert logging_obj.model_call_details.get("combined_usage_object") is None + + +class _UpstreamThatTimesOutAfterCreated: + headers: dict[str, str] = {} + + def __init__(self, **created_fields: object) -> None: + self.created_fields = created_fields + + async def aiter_bytes(self) -> AsyncIterator[bytes]: + created: Final = { + **_CREATED_WITH_PRIORITY_TIER, + "response": {**_CREATED_RESPONSE_WITH_PRIORITY_TIER, **self.created_fields}, + } + for event in (created, _text_delta("Once upon a time", 1)): + yield f"data: {json.dumps(event)}\n\n".encode() + raise httpx.ReadTimeout("Timeout on reading data from socket") + + +async def _mid_stream_failure_payload( + monkeypatch: pytest.MonkeyPatch, requested_tier: str | None, **created_fields: object +) -> StandardLoggingPayload: + recorder: Final = FailureRecorder() + monkeypatch.setattr(litellm, "failure_callback", [recorder]) + monkeypatch.setattr(litellm, "_async_failure_callback", [recorder]) + logging_obj: Final = _make_logging_obj() + logging_obj.update_environment_variables( + model="gpt-5.4-nano", + optional_params={"stream": True} + if requested_tier is None + else {"stream": True, "service_tier": requested_tier}, + litellm_params={"custom_llm_provider": "openai", "aresponses": True}, + ) + iterator: Final = _live_iterator(logging_obj, upstream=_UpstreamThatTimesOutAfterCreated(**created_fields)) + + with pytest.raises(litellm.Timeout): + await _ChunkRecorder().drain(iterator) + await iterator._await_pending_logging() + + (payload,) = recorder.failure_payloads + return payload + + +@pytest.mark.asyncio +async def test_mid_stream_failure_bills_the_relayed_text_at_the_tier_the_provider_served( + monkeypatch: pytest.MonkeyPatch, +) -> None: + flex_requested: Final = await _mid_stream_failure_payload(monkeypatch, "flex", service_tier="default") + nothing_requested: Final = await _mid_stream_failure_payload(monkeypatch, None, service_tier="default") + + assert nothing_requested["response_cost"] > 0 + assert flex_requested["response_cost"] == nothing_requested["response_cost"] + + +@pytest.mark.asyncio +async def test_mid_stream_failure_bills_nothing_when_the_provider_runs_the_response_in_background( + monkeypatch: pytest.MonkeyPatch, +) -> None: + payload: Final = await _mid_stream_failure_payload(monkeypatch, None, background=True) + + assert payload["status"] == "failure" + assert (payload["completion_tokens"], payload["response_cost"]) == (0, 0) + + def test_partial_response_on_disconnect_bills_the_relayed_text_at_the_served_service_tier(): iterator: Final = _live_iterator(_make_logging_obj()) for event in (_CREATED_WITH_PRIORITY_TIER, _text_delta("Once upon a time ", 1), _text_delta("a fox slept.", 2)): @@ -263,3 +618,25 @@ async def test_partial_response_on_disconnect_is_none_when_there_is_nothing_to_b assert iterator.partial_response_on_disconnect() is None await iterator._await_pending_logging() + + +@pytest.mark.asyncio +async def test_partial_response_on_disconnect_is_none_once_the_failure_log_billed_the_relayed_text(): + iterator: Final = _live_iterator(_make_logging_obj(), upstream=_UpstreamThatTimesOutAfterRelayingText()) + + with pytest.raises(litellm.Timeout): + await _ChunkRecorder().drain(iterator) + + assert iterator.partial_response_on_disconnect() is None + await iterator._await_pending_logging() + + +def test_partial_response_on_disconnect_reports_the_relayed_text_of_a_background_response(): + iterator: Final = _live_iterator(_make_logging_obj(), background=True) + for event in (_CREATED_WITH_PRIORITY_TIER, _text_delta("Once upon a time", 1)): + iterator._process_chunk(json.dumps(event)) + + incomplete_event: Final = iterator.partial_response_on_disconnect() + + assert incomplete_event is not None + assert incomplete_event.response.output[0].content[0].text == "Once upon a time" diff --git a/tests/unit/responses/test_streaming_iterator.py b/tests/unit/responses/test_streaming_iterator.py index 6c3e6177503..664f4dba163 100644 --- a/tests/unit/responses/test_streaming_iterator.py +++ b/tests/unit/responses/test_streaming_iterator.py @@ -210,9 +210,11 @@ _TRAILING_ERRORS = [ httpx.RemoteProtocolError("peer closed connection without sending complete message body"), ] +_TRAILING_ERRORS_AFTER_COMPLETED = [*_TRAILING_ERRORS, httpx.ReadTimeout("Timeout on reading data from socket")] + @pytest.mark.asyncio -@pytest.mark.parametrize("trailing_error", _TRAILING_ERRORS, ids=type) +@pytest.mark.parametrize("trailing_error", _TRAILING_ERRORS_AFTER_COMPLETED, ids=type) async def test_transport_error_after_completed_event_ends_stream_cleanly(trailing_error): """A sloppy connection close after `response.completed` must not turn a complete stream into an error (regression guard for the transport no longer @@ -243,7 +245,7 @@ async def test_transport_error_before_completed_event_raises(): pass -@pytest.mark.parametrize("trailing_error", _TRAILING_ERRORS, ids=type) +@pytest.mark.parametrize("trailing_error", _TRAILING_ERRORS_AFTER_COMPLETED, ids=type) def test_sync_transport_error_after_completed_event_ends_stream_cleanly(trailing_error): iterator = _make_sync_iterator( sse_events=_COMPLETE_STREAM_EVENTS, @@ -1205,7 +1207,9 @@ async def test_completed_event_survives_a_failing_usage_estimate(): falls back to usage None.""" malformed_input: Final = [{"type": "message", "role": "user", "content": 42}] with pytest.raises(ValueError, match="Invalid content type"): - _estimate_usage_from_text("gpt-4o-mini", malformed_input, {"input": malformed_input}, "hello world") + _estimate_usage_from_text( + "gpt-4o-mini", malformed_input, {"input": malformed_input}, "hello world" + ) response = _responses_api_response_without_usage() iterator = _make_iterator( diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 566024b5d94..b4a3d207bd4 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -3774,6 +3774,132 @@ async def test_continued_stream_sharing_the_proxy_logging_object_bills_the_cut_t ) +def _responses_sse(*events: Mapping[str, object]) -> str: + return "".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events) + + +def _responses_envelope(response_id: str, model: str, status: str, **fields: object) -> dict[str, object]: + return { + "id": response_id, + "object": "response", + "created_at": 1700000000, + "status": status, + "model": model, + "output": [], + **fields, + } + + +@pytest.mark.asyncio +async def test_responses_fallback_sharing_the_proxy_logging_object_logs_the_fallback_usage( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + recorder: Final = _SpendRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + respx_mock.post("https://primary.example/v1/responses").mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text=_responses_sse( + {"type": "response.created", "response": _responses_envelope("resp_primary", "gpt-4o", "in_progress")}, + *( + { + "type": "response.output_text.delta", + "item_id": "msg_primary", + "output_index": 0, + "content_index": 0, + "delta": word, + } + for word in ("Once ", "upon ", "a ", "time ", "there") + ), + { + "type": "response.failed", + "response": _responses_envelope( + "resp_primary", + "gpt-4o", + "failed", + error={"code": "server_error", "message": "The model had an internal error"}, + ), + }, + ), + ) + ) + respx_mock.post("https://backup.example/v1/responses").mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text=_responses_sse( + {"type": "response.created", "response": _responses_envelope("resp_backup", "gpt-4o-mini", "in_progress")}, + { + "type": "response.completed", + "response": _responses_envelope( + "resp_backup", + "gpt-4o-mini", + "completed", + output=[ + { + "type": "message", + "id": "msg_backup", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Once upon a time", "annotations": []}], + } + ], + usage={"input_tokens": 1000, "output_tokens": 700, "total_tokens": 1700}, + ), + }, + ), + ) + ) + router: Final = Router( + model_list=[ + { + "model_name": "primary", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "k1", "api_base": "https://primary.example/v1"}, + }, + { + "model_name": "backup", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "k2", + "api_base": "https://backup.example/v1", + }, + }, + ], + fallbacks=[{"primary": ["backup"]}], + num_retries=0, + ) + shared_logging, _ = litellm.utils.function_setup( + original_function="aresponses", + rules_obj=litellm.utils.Rules(), + start_time=datetime.now(), + model="primary", + input="Tell me a story", + stream=True, + litellm_call_id="shared-call", + ) + + response: Final = await router.aresponses( + model="primary", + input="Tell me a story", + stream=True, + litellm_call_id="shared-call", + litellm_logging_obj=shared_logging, + ) + _ = [event async for event in response] + await asyncio.wait_for(recorder.succeeded.wait(), timeout=10) + + success: Final = recorder.successes[-1] + assert (success["prompt_tokens"], success["completion_tokens"]) == (1000, 700) + spend_row: Final = recorder.spend_rows[-1] + assert (spend_row["prompt_tokens"], spend_row["completion_tokens"], spend_row["spend"]) == ( + 1000, + 700, + success["response_cost"], + ) + + def test_carry_partial_spend_to_fallback_adds_the_partial_the_cut_stream_inherited(): from litellm.types.utils import Usage From 9007e0fa70ca38b6b1ad78f6fd082d03d32d886c Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 13:42:22 -0700 Subject: [PATCH 088/128] fix(files)!: scope provider and model file lists to the caller's managed files under require_managed_files (#47682) * test(files): cover owner-scoped file listing under require_managed_files * fix(files): list only the caller's managed files on every list route under require_managed_files * fix(files): narrow owner-scoped managed file lists to the named provider or model and page with after_id and before_id * fix(files): narrow scoped lists the way uploads route, and refuse lists that cannot be scoped A model name now matches every deployment that serves it, wildcards and aliases included, and a provider matches deployments by the provider uploads resolve, so anthropic/* and unprefixed Claude models count. The OpenAI and Azure lists also return the caller's pass-through uploads. An unknown provider and different after and after_id cursors answer 400, and a list with no matching deployment skips the database read. * fix(files): match stored files by model name and keep flag-off lists as on main Under require_managed_files, a narrowed list now matches a file's model mappings by deployment id or model name, so files LiteLLM stores in its own database show on their model and provider lists. The Anthropic list defaults to 20 rows when no limit is sent, after_id, before_id and the cursor 400s apply only under the flag, and an admin keeps the provider list when the managed files hook is not loaded. List errors are returned as values and mapped to HTTP errors in one place. * fix(files): keep admin provider lists and index the owner-scoped walk Under require_managed_files, an admin keeps the provider list as on main, and only a non-admin list reads the caller's own and team managed files. A stored file now also shows on the provider its model name routes to through a wildcard, with exact model names first, and a list routed to a model takes the default page size of the provider serving it. A new index on created_by and created_at keeps each chunk of the walk an index read, and the model check runs before the file object is parsed. * fix(files): filter provider and model lists in the query A non-admin list scoped to a provider or model now asks Postgres only for the caller's rows that map to one of that scope's deployments or models, or that live in file storage, instead of reading every owned row and filtering them in Python. A caller with many files on other providers no longer waits for the whole set to be read. Team members' files are matched in the same query. * test(files): type the fake raw query in the files list tests * refactor(proxy): list managed files through one typed database handle * fix(files): list alias and routing-group storage uploads on their provider --- .../proxy/hooks/managed_files.py | 151 ++++- .../migration.sql | 2 + .../litellm_proxy_extras/schema.prisma | 1 + litellm/constants.py | 1 + litellm/llms/base_llm/files/transformation.py | 23 +- .../openai_files_endpoints/files_endpoints.py | 255 +++++++- .../managed_id_rewriter.py | 6 + litellm/proxy/schema.prisma | 1 + schema.prisma | 1 + .../security/test_managed_file_listing.py | 263 ++++++++ .../proxy/test_managed_files_hook.py | 17 + .../test_files_endpoint.py | 569 +++++++++++++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 6 + 13 files changed, 1245 insertions(+), 51 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20261021000000_managed_file_created_by_created_at_index/migration.sql create mode 100644 tests/integration/security/test_managed_file_listing.py diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 7a07350e2d9..416444664a3 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -28,7 +28,7 @@ import httpx from fastapi import HTTPException from openai import APIConnectionError, AsyncOpenAI from openai.types.file_deleted import FileDeleted -from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError from typing_extensions import ReadOnly, Unpack import litellm @@ -53,7 +53,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_file_metadata, ) from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend -from litellm.llms.base_llm.files.transformation import BaseFileEndpoints +from litellm.llms.base_llm.files.transformation import BaseFileEndpoints, FileListModelScope from litellm.llms.base_llm.managed_resources.isolation import ( build_list_page, build_owner_filter, @@ -309,6 +309,7 @@ def _is_xai_batch_output( class _ManagedFileRow(Protocol): unified_file_id: str file_object: OpenAIFileObject + model_mappings: Mapping[str, str] flat_model_file_ids: Sequence[str] storage_backend: Optional[str] storage_url: Optional[str] @@ -362,6 +363,8 @@ class _ManagedResourceDatabase(Protocol): @property def litellm_managedfiletable(self) -> _ManagedFileTableActions: ... + async def query_raw(self, query: str, *args: object) -> object: ... + @property def litellm_managedobjecttable(self) -> _ManagedObjectTableActions: ... @@ -385,10 +388,83 @@ class _RouterFileCallKwargs(TypedDict, total=False): custom_llm_provider: ReadOnly[str | None] +def _managed_database(prisma_client: _ManagedResourcePrismaClient) -> _ManagedResourceDatabase: + return prisma_client.db + + def _managed_file_table(prisma_client: _ManagedResourcePrismaClient) -> _ManagedFileTableActions: return prisma_client.db.litellm_managedfiletable +class _ListedFileRow(Protocol): + @property + def unified_file_id(self) -> str: ... + + @property + def file_object(self) -> object: ... + + @property + def model_mappings(self) -> Mapping[str, object]: ... + + @property + def storage_backend(self) -> str | None: ... + + +class _ScopedFileRow(BaseModel): + model_config = ConfigDict(frozen=True) + + unified_file_id: str + file_object: JsonValue = None + model_mappings: Mapping[str, JsonValue] + storage_backend: str | None = None + + +_SCOPED_FILE_ROWS: Final = TypeAdapter(tuple[_ScopedFileRow, ...]) + + +def _owner_condition(owner_filter: Mapping[str, object]) -> tuple[str, tuple[str, ...]]: + match owner_filter: + case {"OR": [{"created_by": str() as user_id}, {"team_id": str() as team_id}]}: + return '("created_by" = $1 OR "team_id" = $2)', (user_id, team_id) + case {"created_by": str() as user_id}: + return '"created_by" = $1', (user_id,) + case {"team_id": str() as team_id}: + return '"team_id" = $1', (team_id,) + case _: + return ("FALSE", ()) if owner_filter else ("TRUE", ()) + + +async def _scoped_file_rows( + database: _ManagedResourceDatabase, + owner_filter: Mapping[str, object], + model_mapping_keys: frozenset[str], + cursor_id: str | None, + direction: Literal["asc", "desc"], + take: int, +) -> tuple[_ScopedFileRow, ...]: + owner_condition, owner_params = _owner_condition(owner_filter) + keys_param: Final = len(owner_params) + 1 + comparison: Final = "<" if direction == "desc" else ">" + cursor_condition: Final = ( + f'AND ("created_at", "unified_file_id") {comparison} (SELECT "created_at", "unified_file_id" ' + f'FROM "LiteLLM_ManagedFileTable" WHERE "unified_file_id" = ${keys_param + 1}) ' + if cursor_id + else "" + ) + order: Final = direction.upper() + rows: Final = await database.query_raw( + 'SELECT "unified_file_id", "file_object", "model_mappings", "storage_backend" FROM "LiteLLM_ManagedFileTable" ' + f'WHERE {owner_condition} AND ("model_mappings" ?| ${keys_param}::text[] OR "storage_backend" IS NOT NULL) ' + f'{cursor_condition}ORDER BY "created_at" {order}, "unified_file_id" {order} ' + f"LIMIT ${keys_param + (2 if cursor_id else 1)}", + *owner_params, + sorted(model_mapping_keys), + *((cursor_id,) if cursor_id else ()), + take, + ) + return _SCOPED_FILE_ROWS.validate_python(rows) + + def _iter_provider_file_id_pairs( rows: Sequence[_ManagedFileRow], requested_provider_file_ids: frozenset[str], @@ -2265,14 +2341,23 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): user_api_key_dict: UserAPIKeyAuth, limit: Optional[int] = None, after: Optional[str] = None, + model_scope: Optional[FileListModelScope] = None, + before: Optional[str] = None, **data: Dict, ) -> FileListPage: """List the managed files the caller owns, newest first. Pagination is keyset based on ``unified_file_id`` so a key that owns every file on the proxy still reads one bounded page at a time. - ``purpose`` is applied after parsing, because the managed file table - keeps it inside the ``file_object`` blob instead of a column, and rows + ``after`` pages toward older files and ``before`` toward newer ones, + each page still ordered newest first. + ``model_scope`` is applied in the query, so only rows whose + ``model_mappings`` name one of its deployment ids or model names, and + storage-backed rows, leave the database; a storage-backed row then + stays only when its model name is one the router sends to the listed + provider. ``purpose`` is applied after parsing, because the managed + file table keeps it inside the ``file_object`` blob instead of a + column. Rows whose blob will not parse drop out there too, so a chunk of rows can yield fewer matches than the page holds. Successive chunks are read until the page is full or the caller's rows run out, which keeps @@ -2281,9 +2366,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): scan has to continue past it, the chunk widens to ``FILE_LIST_CONTINUATION_CHUNK_SIZE``, so the walk costs one query per that many rows instead of one per page. That bound is per query, not - per request: the work is still linear in the rows the caller owns, and - a filter matching nothing reads every one of them, with no index - covering either the owner filter or the sort. + per request: a ``purpose`` filter matching nothing still reads every + row that passes the query. """ validate_file_list_limit(limit) validate_file_list_purpose(purpose) @@ -2292,36 +2376,60 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if owner_filter is None: return FileListPage.model_validate(build_list_page([])) - if after: - cursor_row = await _managed_file_table(self.prisma_client).find_first( - where={**owner_filter, "unified_file_id": after} + if after and before: + raise ProxyException( + message="Pass either 'after' (or 'after_id') or 'before_id' when listing files, not both.", + type="invalid_request_error", + param="before_id", + code=400, + openai_code="invalid_value", + ) + cursor_param: Final = "before_id" if before else "after" + cursor: Final = before or after + database: Final = _managed_database(self.prisma_client) + if cursor: + cursor_row = await database.litellm_managedfiletable.find_first( + where={**owner_filter, "unified_file_id": cursor} ) if cursor_row is None: raise ProxyException( - message=f"Invalid 'after' cursor: no file found with id '{after}'.", + message=f"Invalid '{cursor_param}' cursor: no file found with id '{cursor}'.", type="invalid_request_error", - param="after", + param=cursor_param, code=400, openai_code="invalid_value", ) + if model_scope is not None and not model_scope.model_mapping_keys: + return FileListPage.model_validate(build_list_page([])) + page_size: Final = min(limit or MAX_FILE_LIST_LIMIT, MAX_FILE_LIST_LIMIT) matches: Final[List[OpenAIFileObject]] = [] - cursor_id = after + cursor_id = cursor chunk_size = page_size + 1 + direction: Final[Literal["asc", "desc"]] = "asc" if before else "desc" while len(matches) <= page_size: cursor_args: _CursorPageArgs = {"cursor": {"unified_file_id": cursor_id}, "skip": 1} if cursor_id else {} - chunk = await _managed_file_table(self.prisma_client).find_many( - where=owner_filter, - take=chunk_size, - order=[{"created_at": "desc"}, {"unified_file_id": "desc"}], - **cursor_args, + chunk: Sequence[_ListedFileRow] = ( + await _scoped_file_rows( + database, owner_filter, model_scope.model_mapping_keys, cursor_id, direction, chunk_size + ) + if model_scope is not None + else await database.litellm_managedfiletable.find_many( + where=owner_filter, + take=chunk_size, + order=[{"created_at": direction}, {"unified_file_id": direction}], + **cursor_args, + ) ) matches.extend( _public_file_object(parsed_file_object, row.unified_file_id) for row in chunk - if (parsed_file_object := _parse_managed_file_object(row.file_object, row.unified_file_id)) is not None + if ( + model_scope is None or model_scope.holds(row.model_mappings, stored=row.storage_backend is not None) + ) + and (parsed_file_object := _parse_managed_file_object(row.file_object, row.unified_file_id)) is not None and (purpose is None or parsed_file_object.purpose == purpose) ) if len(chunk) < chunk_size: @@ -2329,7 +2437,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): cursor_id = chunk[-1].unified_file_id chunk_size = max(chunk_size, FILE_LIST_CONTINUATION_CHUNK_SIZE) - return FileListPage.model_validate(build_list_page(matches[:page_size], has_more=len(matches) > page_size)) + page: Final = matches[:page_size] + return FileListPage.model_validate( + build_list_page(page[::-1] if before else page, has_more=len(matches) > page_size) + ) def _is_batch_polling_enabled(self) -> bool: """ diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261021000000_managed_file_created_by_created_at_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261021000000_managed_file_created_by_created_at_index/migration.sql new file mode 100644 index 00000000000..8bbcb67bddf --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261021000000_managed_file_created_by_created_at_index/migration.sql @@ -0,0 +1,2 @@ +CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_ManagedFileTable_created_by_created_at_idx" +ON "LiteLLM_ManagedFileTable" ("created_by", "created_at" DESC); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 83a0e6f9e05..4187f76c97c 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1163,6 +1163,7 @@ model LiteLLM_ManagedFileTable { @@index([unified_file_id]) @@index([flat_model_file_ids], type: Gin) @@index([team_id, created_at(sort: Desc)]) + @@index([created_by, created_at(sort: Desc)]) } model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use the diff --git a/litellm/constants.py b/litellm/constants.py index 2b083b03766..1d5c6404bfd 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -129,6 +129,7 @@ MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES: Final = 1024 S3_LOG_PROMPTS_ONLY_ENV_VAR: Final = "S3_LOG_PROMPTS_ONLY" S3_PARTITION_GRANULARITY_ENV_VAR: Final = "S3_PARTITION_GRANULARITY" MAX_FILE_LIST_LIMIT: Final = 10000 +ANTHROPIC_DEFAULT_FILE_LIST_LIMIT: Final = 20 DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = get_env_int("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10) DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = get_env_int("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1) GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS: Final = 10 diff --git a/litellm/llms/base_llm/files/transformation.py b/litellm/llms/base_llm/files/transformation.py index 4346e214dc8..491fb1919b6 100644 --- a/litellm/llms/base_llm/files/transformation.py +++ b/litellm/llms/base_llm/files/transformation.py @@ -1,6 +1,7 @@ from abc import ABC, abstractmethod -from collections.abc import AsyncGenerator, Iterator, Mapping -from typing import TYPE_CHECKING, Any, Union +from collections.abc import AsyncGenerator, Callable, Iterable, Iterator, Mapping +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Final, Union import httpx from openai.types.file_deleted import FileDeleted @@ -38,6 +39,22 @@ else: UserAPIKeyAuth = Any +def _routes_nowhere(_model_name: str) -> bool: + return False + + +@dataclass(frozen=True, slots=True) +class FileListModelScope: + model_mapping_keys: frozenset[str] + routes_stored_model_name: Callable[[str], bool] = _routes_nowhere + + def holds(self, model_mapping_keys: Iterable[str], stored: bool) -> bool: + mapped_keys: Final = frozenset(model_mapping_keys) + return not self.model_mapping_keys.isdisjoint(mapped_keys) or ( + stored and any(map(self.routes_stored_model_name, mapped_keys)) + ) + + class BaseFileUploadStream(ABC): """Re-iterable request body that yields an upload's bytes lazily. @@ -274,6 +291,8 @@ class BaseFileEndpoints(ABC): user_api_key_dict: UserAPIKeyAuth, limit: int | None = None, after: str | None = None, + model_scope: FileListModelScope | None = None, + before: str | None = None, **data: dict, ) -> FileListPage: pass diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index d2ceefffe47..5d63c3db222 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -6,9 +6,12 @@ ###################################################################### import asyncio +import functools +import itertools import traceback -from collections.abc import Mapping, Sequence -from typing import Any, BinaryIO, Final, TypedDict, cast, get_args +from collections.abc import Iterable, Iterator, Mapping, Sequence +from dataclasses import dataclass +from typing import Any, BinaryIO, Final, NoReturn, TypedDict, cast, get_args import httpx from fastapi import ( @@ -28,12 +31,13 @@ from typing_extensions import ReadOnly import litellm from litellm import CreateFileRequest, get_secret_str from litellm._logging import verbose_proxy_logger +from litellm.constants import ANTHROPIC_DEFAULT_FILE_LIST_LIMIT from litellm.litellm_core_utils.cloud_storage_security import ( is_managed_cloud_storage_uri, ) from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket from litellm.llms.base_llm.files.litellm_db_storage_backend import LITELLM_DB_STORAGE_BACKEND_NAME -from litellm.llms.base_llm.files.transformation import BaseFileEndpoints +from litellm.llms.base_llm.files.transformation import BaseFileEndpoints, FileListModelScope from litellm.llms.base_llm.managed_resources.isolation import build_list_page from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -112,6 +116,9 @@ from litellm.types.llms.openai import ( OpenAIFileObject, OpenAIFilesPurpose, ) +from litellm.types.proxy.auth.user_api_key_auth import user_api_key_has_admin_view +from litellm.types.router import Deployment +from litellm.types.utils import LlmProviders router: Final = APIRouter() @@ -1673,6 +1680,185 @@ def _as_file_list_page(response: object) -> object: return FileListPage.model_validate(page) +async def _list_owned_managed_files( + managed_files_obj: BaseFileEndpoints, + user_api_key_dict: UserAPIKeyAuth, + purpose: str | None, + limit: int | None, + after: str | None, + before: str | None, + model_scope: FileListModelScope | None, +) -> FileListPage: + return await managed_files_obj.afile_list( + purpose=purpose, + litellm_parent_otel_span=user_api_key_dict.parent_otel_span, + user_api_key_dict=user_api_key_dict, + limit=limit, + after=after, + before=before, + model_scope=model_scope, + ) + + +@dataclass(frozen=True, slots=True) +class _UnscopedFileList: + status_code: int + detail: str + + +@dataclass(frozen=True, slots=True) +class _ListedFiles: + model_scope: FileListModelScope | None + custom_llm_provider: str | None + + +def _raise_unscoped_file_list(failure: _UnscopedFileList) -> NoReturn: + raise HTTPException(status_code=failure.status_code, detail=failure.detail) + + +def _single_target_model_name(target_model_names: str) -> str | _UnscopedFileList: + target_model_names_list: Final = target_model_names.split(",") + if len(target_model_names_list) != 1: + return _UnscopedFileList( + status_code=400, + detail="target_model_names on list files must be a list of one model name. Example: ['gpt-4o']", + ) + return target_model_names_list[0] + + +async def _requested_custom_llm_provider(request: Request, provider: str | None) -> str | None: + return ( + provider + or get_custom_llm_provider_from_request_headers(request=request) + or get_custom_llm_provider_from_request_query(request=request) + or await get_custom_llm_provider_from_request_body(request=request) + ) + + +def _default_file_list_limit(custom_llm_provider: str | None) -> int | None: + return ANTHROPIC_DEFAULT_FILE_LIST_LIMIT if custom_llm_provider == LlmProviders.ANTHROPIC.value else None + + +def _deployment_model_mapping_keys(deployment: Deployment) -> frozenset[str]: + return frozenset(key for key in (deployment.model_name, deployment.model_info.id) if key is not None) + + +def _model_mapping_keys(deployments: Iterable[Deployment]) -> frozenset[str]: + return frozenset(itertools.chain.from_iterable(map(_deployment_model_mapping_keys, deployments))) + + +def _deployments_serving_model(llm_router: Router, model: str) -> tuple[Deployment, ...]: + deployment: Final = llm_router.get_deployment(model_id=model) if llm_router.has_model_id(model) else None + if deployment is not None: + return (deployment,) + return tuple(Deployment.model_validate(listed) for listed in llm_router.get_model_list(model_name=model) or ()) + + +def _resolved_deployment_provider(deployment: Deployment) -> str | None: + try: + return litellm.get_llm_provider( + model=deployment.litellm_params.model, + custom_llm_provider=deployment.litellm_params.custom_llm_provider, + )[1] + except litellm.BadRequestError: + return None + + +def _deployments_on_provider(llm_router: Router | None, custom_llm_provider: str) -> Iterator[Deployment]: + if llm_router is None: + return iter(()) + return ( + deployment + for deployment_id in llm_router.get_model_ids() + if (deployment := llm_router.get_deployment(model_id=deployment_id)) is not None + and _resolved_deployment_provider(deployment) == custom_llm_provider + ) + + +def _routed_model_name_on_provider( + llm_router: Router, custom_llm_provider: str, provider_model_names: frozenset[str], model_name: str +) -> bool: + if model_name in llm_router.model_names: + return model_name in provider_model_names + return llm_router.pattern_router.routes_to_provider(model_name, custom_llm_provider) + + +def _stored_model_name_routes_to_provider( + llm_router: Router, custom_llm_provider: str, provider_model_names: frozenset[str], model_name: str +) -> bool: + routing_group: Final = llm_router.get_routing_group(model_name) + routed_model_names: Final = ( + tuple(routing_group.models) if routing_group is not None else (llm_router.routable_model_group(model_name),) + ) + return any( + _routed_model_name_on_provider(llm_router, custom_llm_provider, provider_model_names, routed_model_name) + for routed_model_name in routed_model_names + ) + + +def _files_on_provider(llm_router: Router | None, custom_llm_provider: str) -> _ListedFiles | _UnscopedFileList: + from litellm.proxy.pass_through_endpoints.managed_id_rewriter import passthrough_file_model_id + + if custom_llm_provider not in litellm.provider_list: + return _UnscopedFileList(status_code=400, detail=f"Unknown provider '{custom_llm_provider}' for listing files.") + passthrough_model_id: Final = passthrough_file_model_id(custom_llm_provider) + deployments: Final = tuple(_deployments_on_provider(llm_router, custom_llm_provider)) + model_mapping_keys: Final = _model_mapping_keys(deployments) | frozenset( + (passthrough_model_id,) if passthrough_model_id is not None else () + ) + if llm_router is None: + return _ListedFiles(FileListModelScope(model_mapping_keys), custom_llm_provider) + return _ListedFiles( + FileListModelScope( + model_mapping_keys, + routes_stored_model_name=functools.partial( + _stored_model_name_routes_to_provider, + llm_router, + custom_llm_provider, + frozenset(deployment.model_name for deployment in deployments), + ), + ), + custom_llm_provider, + ) + + +def _files_serving_model(llm_router: Router, model: str) -> _ListedFiles: + deployments: Final = _deployments_serving_model(llm_router, model) + providers: Final = frozenset(map(_resolved_deployment_provider, deployments)) + return _ListedFiles( + FileListModelScope(_model_mapping_keys(deployments)), next(iter(providers)) if len(providers) == 1 else None + ) + + +async def _listed_files( + llm_router: Router | None, + user_api_key_dict: UserAPIKeyAuth, + custom_llm_provider: str | None, + target_model_names: str | None, + routed_model: str | None, +) -> _ListedFiles | _UnscopedFileList: + if routed_model is not None and llm_router is not None: + return _files_serving_model(llm_router, routed_model) + if target_model_names: + target_model_name: Final = _single_target_model_name(target_model_names) + if isinstance(target_model_name, _UnscopedFileList): + return target_model_name + if llm_router is None: + return _UnscopedFileList( + status_code=500, detail="LLM Router not initialized. Ensure models added to proxy." + ) + await get_authorized_credentials_for_model( + llm_router=llm_router, + model_id=target_model_name, + user_api_key_dict=user_api_key_dict, + operation_context="file list", + ) + return _files_serving_model(llm_router, target_model_name) + if custom_llm_provider is None: + return _ListedFiles(None, None) + return _files_on_provider(llm_router, custom_llm_provider) + + @router.get( "/{provider}/v1/files", dependencies=[Depends(user_api_key_auth)], @@ -1697,6 +1883,8 @@ async def list_files( purpose: str | None = None, limit: int | None = None, after: str | None = None, + after_id: str | None = None, + before_id: str | None = None, ): """ Returns information about a specific file. that can be used across - Assistants API, Batch API @@ -1739,6 +1927,10 @@ async def list_files( ) response: Any | None = None + owned_files_hook: Final = proxy_logging_obj.get_proxy_hook("managed_files") + list_owned_files_only: Final = litellm.require_managed_files is True and not user_api_key_has_admin_view( + user_api_key_dict + ) # Check for model-based credential routing (no file_id encoding check for list) should_route, model_used, _, credentials = await handle_model_based_routing( @@ -1750,7 +1942,36 @@ async def list_files( check_file_id_encoding=False, ) - if should_route and credentials is not None: + if list_owned_files_only: + if after and after_id and after != after_id: + raise HTTPException( + status_code=400, + detail="Pass the list cursor as either 'after' or 'after_id', not both with different ids.", + ) + if not isinstance(owned_files_hook, BaseFileEndpoints): + raise HTTPException( + status_code=500, + detail="Managed file listing is unavailable, so require_managed_files cannot scope the list to the caller.", + ) + listed_files: Final = await _listed_files( + llm_router=llm_router, + user_api_key_dict=user_api_key_dict, + custom_llm_provider=await _requested_custom_llm_provider(request=request, provider=provider), + target_model_names=target_model_names, + routed_model=model_used if should_route else None, + ) + if isinstance(listed_files, _UnscopedFileList): + _raise_unscoped_file_list(listed_files) + response = await _list_owned_managed_files( # rebind-ok: each dispatch branch binds response + managed_files_obj=owned_files_hook, + user_api_key_dict=user_api_key_dict, + purpose=purpose, + limit=limit if limit is not None else _default_file_list_limit(listed_files.custom_llm_provider), + after=after or after_id, + before=before_id, + model_scope=listed_files.model_scope, + ) + elif should_route and credentials is not None: # Use model-based routing with credentials from config prepare_data_with_credentials(data=data, credentials=credentials, include_internal_credentials=True) response = await litellm.afile_list( @@ -1762,12 +1983,9 @@ async def list_files( verbose_proxy_logger.debug("Listed files using model: %s", model_used) elif target_model_names and isinstance(target_model_names, str): - target_model_names_list: Final = target_model_names.split(",") - if len(target_model_names_list) != 1: - raise HTTPException( - status_code=400, - detail="target_model_names on list files must be a list of one model name. Example: ['gpt-4o']", - ) + target_model_name: Final = _single_target_model_name(target_model_names) + if isinstance(target_model_name, _UnscopedFileList): + _raise_unscoped_file_list(target_model_name) if llm_router is None: raise HTTPException( status_code=500, @@ -1775,7 +1993,7 @@ async def list_files( ) credentials = await get_authorized_credentials_for_model( llm_router=llm_router, - model_id=target_model_names_list[0], + model_id=target_model_name, user_api_key_dict=user_api_key_dict, operation_context="file list", ) @@ -1786,20 +2004,17 @@ async def list_files( **data, ) else: - custom_llm_provider: Final = ( - provider - or get_custom_llm_provider_from_request_headers(request=request) - or get_custom_llm_provider_from_request_query(request=request) - or await get_custom_llm_provider_from_request_body(request=request) - ) + custom_llm_provider: Final = await _requested_custom_llm_provider(request=request, provider=provider) managed_files_obj: Final = proxy_logging_obj.get_proxy_hook("managed_files") if custom_llm_provider is None and isinstance(managed_files_obj, BaseFileEndpoints): - response = await managed_files_obj.afile_list( - purpose=purpose, - litellm_parent_otel_span=user_api_key_dict.parent_otel_span, + response = await _list_owned_managed_files( # rebind-ok: each dispatch branch binds response + managed_files_obj=managed_files_obj, user_api_key_dict=user_api_key_dict, + purpose=purpose, limit=limit, after=after, + before=None, + model_scope=None, ) else: resolved_custom_llm_provider: Final = custom_llm_provider or "openai" diff --git a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py index a5a04fb68d2..9790c150ae1 100644 --- a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py +++ b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py @@ -242,6 +242,12 @@ def _passthrough_sentinel_model_id(provider: str) -> str: return f"_passthrough_{provider}" +def passthrough_file_model_id(provider: str) -> str | None: + if (provider, "/v1/files") not in _LIST_ROUTE_TABLE: + return None + return _passthrough_sentinel_model_id(provider) + + # Key under which the provider marker is stored in a file row's model_mappings. # Its value lands in flat_model_file_ids (built from model_mappings.values()), # giving the file table a DB-queryable provider scope it otherwise lacks. diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 83a0e6f9e05..4187f76c97c 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1163,6 +1163,7 @@ model LiteLLM_ManagedFileTable { @@index([unified_file_id]) @@index([flat_model_file_ids], type: Gin) @@index([team_id, created_at(sort: Desc)]) + @@index([created_by, created_at(sort: Desc)]) } model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use the diff --git a/schema.prisma b/schema.prisma index 83a0e6f9e05..4187f76c97c 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1163,6 +1163,7 @@ model LiteLLM_ManagedFileTable { @@index([unified_file_id]) @@index([flat_model_file_ids], type: Gin) @@index([team_id, created_at(sort: Desc)]) + @@index([created_by, created_at(sort: Desc)]) } model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use the diff --git a/tests/integration/security/test_managed_file_listing.py b/tests/integration/security/test_managed_file_listing.py new file mode 100644 index 00000000000..8c736b2b2bf --- /dev/null +++ b/tests/integration/security/test_managed_file_listing.py @@ -0,0 +1,263 @@ +import json +import uuid +from collections.abc import Callable, Iterator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass, field +from itertools import count, islice +from pathlib import Path +from types import MappingProxyType +from typing import Final +from urllib.parse import parse_qs, urlsplit + +import pytest +import yaml +from integration._support.client import Gateway, gateway_from_environment +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, Wire, wire_server +from anthropic import Anthropic +from openai import OpenAI +from pydantic import BaseModel, ConfigDict, JsonValue + +pytestmark: Final = pytest.mark.shares_worker_state + +UNMANAGED_FILE: Final = "file-uploaded-straight-to-the-provider" +UPLOAD: Final = b"isolation check\n" + + +@dataclass(frozen=True, slots=True) +class Deployment: + provider: str + model: str + litellm_model: str + api_key: str + api_path: str + + +OPENAI: Final = Deployment("openai", "gpt-4.1-managed-listing", "openai/gpt-4.1", "sk-managed-listing", "/v1") +ANTHROPIC: Final = Deployment( + "anthropic", "claude-managed-listing", "anthropic/claude-haiku-5-5", "sk-ant-managed-listing", "" +) +DEPLOYMENTS: Final = (pytest.param(OPENAI, id="openai"), pytest.param(ANTHROPIC, id="anthropic")) + + +@dataclass(frozen=True, slots=True) +class ListCall: + path: str + params: Mapping[str, str] = field(default_factory=lambda: MappingProxyType({})) + headers: Mapping[str, str] = field(default_factory=lambda: MappingProxyType({})) + + +LIST_FORMS: Final = ( + pytest.param(lambda d: ListCall(f"/{d.provider}/v1/files"), id="provider-path"), + pytest.param(lambda d: ListCall("/v1/files", params={"provider": d.provider}), id="provider-query"), + pytest.param(lambda d: ListCall("/v1/files", headers={"custom-llm-provider": d.provider}), id="provider-header"), + pytest.param(lambda d: ListCall("/v1/files", params={"target_model_names": d.model}), id="target-model-names"), + pytest.param(lambda d: ListCall("/v1/files", headers={"x-litellm-model": d.model}), id="model-header"), +) + + +class FilePage(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + data: list[dict[str, JsonValue]] + has_more: bool + + @property + def ids(self) -> list[str]: + return [str(entry["id"]) for entry in self.data] + + +def _openai_file(file_id: str) -> dict[str, JsonValue]: + return { + "id": file_id, + "object": "file", + "bytes": len(UPLOAD), + "created_at": 1700000000, + "filename": "notes.txt", + "purpose": "user_data", + "status": "processed", + } + + +def _anthropic_file(file_id: str) -> dict[str, JsonValue]: + return { + "id": file_id, + "type": "file", + "filename": "notes.txt", + "mime_type": "text/plain", + "size_bytes": len(UPLOAD), + "created_at": "2026-01-01T00:00:00Z", + "downloadable": False, + } + + +def _reply(body: Mapping[str, JsonValue]) -> Reply: + return Reply(body=json.dumps(body).encode()) + + +def _list_reply(provider: str, file_ids: tuple[str, ...]) -> Reply: + if provider == ANTHROPIC.provider: + return _reply( + { + "data": [_anthropic_file(file_id) for file_id in file_ids], + "has_more": False, + "first_id": file_ids[0], + "last_id": file_ids[-1], + } + ) + return _reply({"object": "list", "data": [_openai_file(file_id) for file_id in file_ids]}) + + +def _responder(provider_files: Mapping[str, tuple[str, ...]]) -> Callable[[Request], Reply]: + uploads: Final = {provider: count() for provider in provider_files} + + def respond(request: Request) -> Reply: + target: Final = urlsplit(request.target) + path: Final = target.path + provider: Final = ANTHROPIC.provider if "x-api-key" in request.headers else OPENAI.provider + files: Final = provider_files[provider] + if (request.method, path) == ("GET", "/v1/models"): + return _reply({"object": "list", "data": []}) + if (request.method, path) == ("POST", "/v1/files"): + uploaded: Final = files[next(uploads[provider])] + return _reply(_anthropic_file(uploaded) if provider == ANTHROPIC.provider else _openai_file(uploaded)) + if (request.method, path) == ("GET", "/v1/files"): + limit: Final = int(parse_qs(target.query).get("limit", ["10000"])[0]) + return _list_reply(provider, (UNMANAGED_FILE, *reversed(files))[:limit]) + return Reply(status=404, body=json.dumps({"error": {"message": f"unscripted {request.target}"}}).encode()) + + return respond + + +def _config(wire: Wire) -> dict[str, JsonValue]: + return { + "model_list": [ + { + "model_name": deployment.model, + "litellm_params": { + "model": deployment.litellm_model, + "api_key": deployment.api_key, + "api_base": f"{wire.url}{deployment.api_path}", + }, + } + for deployment in (OPENAI, ANTHROPIC) + ], + "litellm_settings": {"require_managed_files": True}, + "general_settings": {"master_key": "os.environ/LITELLM_MASTER_KEY", "database_url": "os.environ/DATABASE_URL"}, + } + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + owner: str + intruder: str + files_by_provider: Mapping[str, tuple[str, ...]] + every_file: tuple[str, ...] + + @property + def base_url(self) -> str: + return str(self.proxy.client.base_url).rstrip("/") + + def list(self, call: ListCall, key: str | None = None) -> FilePage: + response: Final = self.proxy.request("GET", call.path, key=key, params=call.params, headers=call.headers) + assert response.status_code == 200, response.text + return FilePage.model_validate_json(response.text) + + +def _upload(proxy: Gateway, key: str, deployments: tuple[Deployment, ...]) -> str: + with OpenAI(base_url=f"{str(proxy.client.base_url).rstrip('/')}/openai/v1", api_key=key, max_retries=0) as client: + uploaded: Final = client.files.with_raw_response.create( + file=("notes.txt", UPLOAD, "text/plain"), + purpose="user_data", + extra_body={"target_model_names": ",".join(deployment.model for deployment in deployments)}, + ) + assert uploaded.status_code == 200, uploaded.text + return uploaded.parse().id + + +@contextmanager +def _rig(directory: Path) -> Iterator[Rig]: + build: Final = uuid.uuid4().hex + provider_files: Final = { + OPENAI.provider: (f"file-older-{build}", f"file-newer-{build}"), + ANTHROPIC.provider: (f"file_older_{build}", f"file_newer_{build}"), + } + with gateway_from_environment() as shared, wire_server(_responder(provider_files)) as wire: + config: Final = directory / "config.yaml" + config.write_text(yaml.safe_dump(_config(wire))) + with owned_proxy(shared, directory, {}, config=config) as proxy, proxy.scenario() as scenario: + owner: Final = scenario.key(user_id=scenario.user()) + intruder: Final = scenario.key(user_id=scenario.user()) + on_both: Final = _upload(proxy, owner, (OPENAI, ANTHROPIC)) + on_openai: Final = _upload(proxy, owner, (OPENAI,)) + on_anthropic: Final = _upload(proxy, owner, (ANTHROPIC,)) + yield Rig( + proxy, + owner, + intruder, + MappingProxyType( + {OPENAI.provider: (on_openai, on_both), ANTHROPIC.provider: (on_anthropic, on_both)} + ), + (on_anthropic, on_openai, on_both), + ) + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with _rig(tmp_path_factory.mktemp("managed-file-listing")) as built: + yield built + + +@pytest.mark.parametrize("deployment", DEPLOYMENTS) +@pytest.mark.parametrize("form", LIST_FORMS) +def test_another_key_lists_none_of_the_files_on_the_shared_key( + rig: Rig, deployment: Deployment, form: Callable[[Deployment], ListCall] +) -> None: + page: Final = rig.list(form(deployment), key=rig.intruder) + assert (page.ids, page.has_more) == ([], False) + + +@pytest.mark.parametrize("deployment", DEPLOYMENTS) +@pytest.mark.parametrize("form", LIST_FORMS) +def test_owner_lists_their_own_managed_files_on_that_provider_or_model_newest_first( + rig: Rig, deployment: Deployment, form: Callable[[Deployment], ListCall] +) -> None: + page: Final = rig.list(form(deployment), key=rig.owner) + assert (page.ids, page.has_more) == (list(rig.files_by_provider[deployment.provider]), False) + + +def test_owner_lists_every_managed_file_when_no_provider_or_model_is_named(rig: Rig) -> None: + page: Final = rig.list(ListCall("/v1/files"), key=rig.owner) + assert (page.ids, page.has_more) == (list(rig.every_file), False) + + +@pytest.mark.parametrize("cursor_param", ["after", "after_id"]) +def test_owner_pages_through_a_provider_list_with_limit_and_a_cursor(rig: Rig, cursor_param: str) -> None: + first: Final = rig.list(ListCall("/openai/v1/files", params={"limit": "1"}), key=rig.owner) + second: Final = rig.list( + ListCall("/openai/v1/files", params={"limit": "1", cursor_param: first.ids[0]}), key=rig.owner + ) + expected: Final = rig.files_by_provider[OPENAI.provider] + assert [(page.ids, page.has_more) for page in (first, second)] == [ + ([expected[0]], True), + ([expected[1]], False), + ] + + +def test_anthropic_sdk_auto_pages_the_owners_files_to_the_end(rig: Rig) -> None: + with Anthropic(base_url=f"{rig.base_url}/anthropic", api_key=rig.owner, max_retries=0) as client: + listed: Final = [file.id for file in islice(client.beta.files.list(limit=1), 5)] + assert listed == list(rig.files_by_provider[ANTHROPIC.provider]) + + +def test_anthropic_sdk_pages_backwards_to_the_newest_file_and_stops(rig: Rig) -> None: + newest_first: Final = rig.files_by_provider[ANTHROPIC.provider] + with Anthropic(base_url=f"{rig.base_url}/anthropic", api_key=rig.owner, max_retries=0) as client: + listed: Final = [file.id for file in islice(client.beta.files.list(limit=1, before_id=newest_first[-1]), 5)] + assert listed == list(reversed(newest_first[:-1])) + + +def test_master_key_still_lists_the_files_on_the_shared_provider_key(rig: Rig) -> None: + listed: Final = rig.list(ListCall(f"/{ANTHROPIC.provider}/v1/files")).ids + assert UNMANAGED_FILE in listed diff --git a/tests/unit/enterprise/proxy/test_managed_files_hook.py b/tests/unit/enterprise/proxy/test_managed_files_hook.py index 01f20689e89..f9f4e17c760 100644 --- a/tests/unit/enterprise/proxy/test_managed_files_hook.py +++ b/tests/unit/enterprise/proxy/test_managed_files_hook.py @@ -15,6 +15,7 @@ import pytest from typing import Final, Optional from unittest.mock import AsyncMock, MagicMock, patch +from litellm.llms.base_llm.files.transformation import FileListModelScope from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth from litellm.proxy.openai_files_endpoints.common_utils import BATCH_CREATE_HIDDEN_PARAM from litellm.types.llms.openai import FileListPage, OpenAIFileObject @@ -649,6 +650,22 @@ async def test_afile_list_denies_a_caller_with_no_identity_at_all(): assert table.find_many_calls == [] +@pytest.mark.asyncio +async def test_afile_list_reads_no_rows_when_no_deployment_can_hold_a_file() -> None: + managed_files, table = _make_managed_files_over_rows([_make_managed_file_row("unified-mine")]) + + response: Final = await managed_files.afile_list( + purpose=None, + litellm_parent_otel_span=None, + user_api_key_dict=_make_user_api_key_dict(), + model_scope=FileListModelScope(frozenset()), + ) + + assert response.data == [] + assert response.has_more is False + assert table.find_many_calls == [] + + @pytest.mark.asyncio async def test_afile_list_filters_by_purpose(): managed_files, _ = _make_managed_files_over_rows( 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 d47db69257b..ff63e37bcd8 100644 --- a/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py @@ -1,3 +1,4 @@ +import itertools import json from typing import Final, List, Optional from urllib.parse import quote @@ -12,6 +13,7 @@ from pytest_mock import MockerFixture import litellm from litellm import Router +from litellm.constants import ANTHROPIC_DEFAULT_FILE_LIST_LIMIT from litellm.files.types import FileContentStreamingResult from litellm.proxy._types import LiteLLM_UserTableFiltered, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.hooks import get_proxy_hook @@ -4969,6 +4971,8 @@ def _managed_file_row( stored_file_id: Optional[str] = None, created_by: str = "test-user", team_id: Optional[str] = None, + model_mappings: Optional[dict[str, str]] = None, + storage_backend: Optional[str] = None, ) -> MagicMock: file_object = OpenAIFileObject( id=stored_file_id or unified_file_id, @@ -4984,6 +4988,8 @@ def _managed_file_row( file_object=file_object.model_dump(), created_by=created_by, team_id=team_id, + model_mappings=model_mappings or {}, + storage_backend=storage_backend, ) @@ -5010,16 +5016,49 @@ class _ManagedFileTableOverRows: return next(iter(self._owned_rows(where)), None) async def find_many(self, where, take=None, order=None, cursor=None, skip=0): - rows = self._owned_rows(where) - if cursor is not None: - start = next( - index - for index, row in enumerate(rows) - if row.unified_file_id == cursor["unified_file_id"] - ) - rows = rows[start + skip :] + owned_rows: Final = self._owned_rows(where) + oldest_first: Final = bool(order) and order[0].get("created_at") == "asc" + ordered_rows: Final = owned_rows[::-1] if oldest_first else owned_rows + start: Final = ( + 0 + if cursor is None + else skip + + next(index for index, row in enumerate(ordered_rows) if row.unified_file_id == cursor["unified_file_id"]) + ) + rows: Final = ordered_rows[start:] return rows if take is None else rows[:take] + async def query_raw(self, query: str, *params: object) -> list[dict[str, object]]: + keys_index: Final = next(index for index, param in enumerate(params) if isinstance(param, list)) + owner_ids: Final = frozenset(params[:keys_index]) + model_mapping_keys: Final = frozenset(params[keys_index]) + cursor_id: Final = params[keys_index + 1] if len(params) == keys_index + 3 else None + ordered_rows: Final = [ + row + for row in (self.rows[::-1] if '"created_at" ASC' in query else self.rows) + if not owner_ids or {row.created_by, row.team_id} & owner_ids + ] + start: Final = ( + 0 + if cursor_id is None + else 1 + next(index for index, row in enumerate(ordered_rows) if row.unified_file_id == cursor_id) + ) + listed_rows: Final = [ + row + for row in ordered_rows[start:] + if not model_mapping_keys.isdisjoint(row.model_mappings) + or ('"storage_backend" IS NOT NULL' in query and row.storage_backend is not None) + ] + return [ + { + "unified_file_id": row.unified_file_id, + "file_object": row.file_object, + "model_mappings": row.model_mappings, + "storage_backend": row.storage_backend, + } + for row in listed_rows[: params[-1]] + ] + def _setup_unscoped_list_files_route_over_real_hook( mocker, monkeypatch, llm_router: Router, rows @@ -5041,7 +5080,9 @@ def _setup_unscoped_list_files_route_over_real_hook( managed_files = PROXY_LiteLLMManagedFiles( internal_usage_cache=MagicMock(), prisma_client=MagicMock() ) - managed_files.prisma_client.db.litellm_managedfiletable = _ManagedFileTableOverRows(rows) + managed_file_table = _ManagedFileTableOverRows(rows) + managed_files.prisma_client.db.litellm_managedfiletable = managed_file_table + managed_files.prisma_client.db.query_raw = managed_file_table.query_raw proxy_logging_obj = setup_proxy_logging_object(monkeypatch, llm_router) proxy_logging_obj.proxy_hook_mapping["managed_files"] = managed_files @@ -6263,3 +6304,513 @@ def test_create_file_passthrough_fails_closed_when_guardrails_would_scan_the_bat assert error["param"] == "passthrough" assert "guardrails" in error["message"] assert forwarded_calls == [] + + +def _callers_files_on_every_deployment() -> list[MagicMock]: + return [ + _managed_file_row("unified-gpt-newer", model_mappings={"gpt-3.5-turbo-id": "file-gpt-newer"}), + _managed_file_row("unified-azure", model_mappings={"azure-gpt-3-5-turbo-id": "file-azure"}), + _managed_file_row("unified-gemini", model_mappings={"gemini-2.0-flash-id": "file-gemini"}), + _managed_file_row( + "unified-theirs", created_by="other-user", model_mappings={"gpt-3.5-turbo-id": "file-theirs"} + ), + _managed_file_row("unified-gpt-older", model_mappings={"gpt-3.5-turbo-id": "file-gpt-older"}), + ] + + +def _list_files_as_test_user(path: str, params: dict[str, str], headers: dict[str, str]) -> httpx.Response: + import litellm.proxy.proxy_server as ps + + try: + return client.get(path, params=params, headers={"Authorization": "Bearer test-key", **headers}) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.parametrize( + ("path", "params", "headers", "expected_ids"), + [ + ("/openai/v1/files", {}, {}, ["unified-gpt-newer", "unified-gpt-older"]), + ("/v1/files", {"provider": "azure"}, {}, ["unified-azure"]), + ("/v1/files", {}, {"custom-llm-provider": "gemini"}, ["unified-gemini"]), + ("/v1/files", {"target_model_names": "gemini-2.0-flash"}, {}, ["unified-gemini"]), + ("/v1/files", {}, {"x-litellm-model": "gpt-3.5-turbo"}, ["unified-gpt-newer", "unified-gpt-older"]), + ("/v1/files", {}, {}, ["unified-gpt-newer", "unified-azure", "unified-gemini", "unified-gpt-older"]), + ], + ids=["provider-path", "provider-query", "provider-header", "target-model-names", "model-header", "no-provider"], +) +def test_require_managed_files_lists_the_callers_files_on_the_named_provider_or_model( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + llm_router: Router, + path: str, + params: dict[str, str], + headers: dict[str, str], + expected_ids: list[str], +) -> None: + _, provider_list = _setup_unscoped_list_files_route_over_real_hook( + mocker, monkeypatch, llm_router, _callers_files_on_every_deployment() + ) + monkeypatch.setattr("litellm.require_managed_files", True) + + response: Final = _list_files_as_test_user(path, params, headers) + + assert response.status_code == 200, response.text + assert [file["id"] for file in response.json()["data"]] == expected_ids + provider_list.assert_not_awaited() + + +def test_require_managed_files_lists_a_teammates_files_on_the_provider_and_no_other_teams( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, llm_router: Router +) -> None: + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _setup_unscoped_list_files_route_over_real_hook( + mocker, + monkeypatch, + llm_router, + [ + _managed_file_row( + "unified-teammate", created_by="teammate", team_id="team-a", model_mappings={"gpt-3.5-turbo-id": "f-1"} + ), + _managed_file_row("unified-own", model_mappings={"gpt-3.5-turbo-id": "f-2"}), + _managed_file_row( + "unified-other-team", created_by="outsider", team_id="team-b", model_mappings={"gpt-3.5-turbo-id": "f-3"} + ), + ], + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="test-key", user_role=LitellmUserRoles.INTERNAL_USER, user_id="test-user", team_id="team-a" + ) + monkeypatch.setattr("litellm.require_managed_files", True) + + response: Final = _list_files_as_test_user("/openai/v1/files", {}, {}) + + assert response.status_code == 200, response.text + assert [file["id"] for file in response.json()["data"]] == ["unified-teammate", "unified-own"] + + +def _router_with_wildcard_and_unprefixed_anthropic_deployments() -> Router: + unprefixed_anthropic_model: Final = next(model for model in sorted(litellm.anthropic_models) if "/" not in model) + return Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "openai/gpt-3.5-turbo", "api_key": "openai_api_key"}, + "model_info": {"id": "gpt-3.5-turbo-id"}, + }, + { + "model_name": "anthropic/*", + "litellm_params": {"model": "anthropic/*", "api_key": "anthropic_api_key"}, + "model_info": {"id": "anthropic-wildcard-id"}, + }, + { + "model_name": "claude-unprefixed", + "litellm_params": {"model": unprefixed_anthropic_model, "api_key": "anthropic_api_key"}, + "model_info": {"id": "claude-unprefixed-id"}, + }, + ], + model_group_alias={"claude-alias": "claude-unprefixed"}, + routing_groups=[ + {"group_name": "claude-group", "models": ["claude-unprefixed"], "routing_strategy": "simple-shuffle"} + ], + ) + + +def _callers_files_on_wildcard_unprefixed_and_passthrough_deployments() -> list[MagicMock]: + return [ + _managed_file_row("unified-wildcard", model_mappings={"anthropic-wildcard-id": "file-wildcard"}), + _managed_file_row("unified-unprefixed", model_mappings={"claude-unprefixed-id": "file-unprefixed"}), + _managed_file_row("unified-gpt", model_mappings={"gpt-3.5-turbo-id": "file-gpt"}), + _managed_file_row( + "unified-passthrough", + model_mappings={ + "_passthrough_openai": "file-passthrough", + "_passthrough_provider_marker": "_passthrough_provider:openai", + }, + ), + _managed_file_row( + "unified-storage", + model_mappings={"claude-unprefixed": "litellm_db://unified-storage"}, + storage_backend="litellm_db", + ), + _managed_file_row( + "unified-wildcard-storage", + model_mappings={"anthropic/claude-haiku-5-5": "litellm_db://unified-wildcard-storage"}, + storage_backend="litellm_db", + ), + _managed_file_row( + "unified-alias-storage", + model_mappings={"claude-alias": "s3://bucket/unified-alias-storage"}, + storage_backend="s3", + ), + _managed_file_row( + "unified-group-storage", + model_mappings={"claude-group": "s3://bucket/unified-group-storage"}, + storage_backend="s3", + ), + ] + + +@pytest.mark.parametrize( + ("path", "params", "headers", "expected_ids"), + [ + ( + "/v1/files", + {"target_model_names": "anthropic/claude-haiku-5-5"}, + {}, + ["unified-wildcard", "unified-wildcard-storage"], + ), + ( + "/v1/files", + {}, + {"x-litellm-model": "anthropic/claude-haiku-5-5"}, + ["unified-wildcard", "unified-wildcard-storage"], + ), + ( + "/anthropic/v1/files", + {}, + {}, + [ + "unified-wildcard", + "unified-unprefixed", + "unified-storage", + "unified-wildcard-storage", + "unified-alias-storage", + "unified-group-storage", + ], + ), + ("/openai/v1/files", {}, {}, ["unified-gpt", "unified-passthrough"]), + ("/v1/files", {"target_model_names": "claude-unprefixed"}, {}, ["unified-unprefixed", "unified-storage"]), + ], + ids=[ + "wildcard-target-model-names", + "wildcard-model-header", + "unprefixed-provider", + "passthrough-provider", + "storage-backed-model", + ], +) +def test_require_managed_files_lists_the_callers_files_wherever_the_router_put_them( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + path: str, + params: dict[str, str], + headers: dict[str, str], + expected_ids: list[str], +) -> None: + _, provider_list = _setup_unscoped_list_files_route_over_real_hook( + mocker, + monkeypatch, + _router_with_wildcard_and_unprefixed_anthropic_deployments(), + _callers_files_on_wildcard_unprefixed_and_passthrough_deployments(), + ) + monkeypatch.setattr("litellm.require_managed_files", True) + + response: Final = _list_files_as_test_user(path, params, headers) + + assert response.status_code == 200, response.text + assert [file["id"] for file in response.json()["data"]] == expected_ids + provider_list.assert_not_awaited() + + +def _router_with_a_catch_all_openai_deployment() -> Router: + unprefixed_anthropic_model: Final = next(model for model in sorted(litellm.anthropic_models) if "/" not in model) + return Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "openai/gpt-3.5-turbo", "api_key": "openai_api_key"}, + "model_info": {"id": "gpt-3.5-turbo-id"}, + }, + { + "model_name": "*", + "litellm_params": {"model": "openai/*", "api_key": "openai_api_key"}, + "model_info": {"id": "catch-all-id"}, + }, + { + "model_name": "claude-unprefixed", + "litellm_params": {"model": unprefixed_anthropic_model, "api_key": "anthropic_api_key"}, + "model_info": {"id": "claude-unprefixed-id"}, + }, + ] + ) + + +@pytest.mark.parametrize( + ("path", "expected_ids"), + [ + ("/openai/v1/files", ["unified-gpt", "unified-catch-all-storage"]), + ("/anthropic/v1/files", ["unified-unprefixed", "unified-storage"]), + ], + ids=["catch-all-provider", "explicit-provider"], +) +def test_require_managed_files_keeps_a_catch_all_route_from_claiming_files_on_other_providers( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + path: str, + expected_ids: list[str], +) -> None: + _, provider_list = _setup_unscoped_list_files_route_over_real_hook( + mocker, + monkeypatch, + _router_with_a_catch_all_openai_deployment(), + [ + _managed_file_row("unified-unprefixed", model_mappings={"claude-unprefixed-id": "file-unprefixed"}), + _managed_file_row("unified-gpt", model_mappings={"gpt-3.5-turbo-id": "file-gpt"}), + _managed_file_row( + "unified-storage", + model_mappings={"claude-unprefixed": "litellm_db://unified-storage"}, + storage_backend="litellm_db", + ), + _managed_file_row( + "unified-catch-all-storage", + model_mappings={"model-only-the-catch-all-serves": "litellm_db://unified-catch-all-storage"}, + storage_backend="litellm_db", + ), + ], + ) + monkeypatch.setattr("litellm.require_managed_files", True) + + response: Final = _list_files_as_test_user(path, {}, {}) + + assert response.status_code == 200, response.text + assert [file["id"] for file in response.json()["data"]] == expected_ids + provider_list.assert_not_awaited() + + +@pytest.mark.parametrize( + ("path", "params"), + [ + ("/antropic/v1/files", {}), + ("/openai/v1/files", {"after": "unified-gpt", "after_id": "unified-wildcard"}), + ("/v1/files", {"target_model_names": "gpt-3.5-turbo,claude-unprefixed"}), + ], + ids=["unknown-provider", "conflicting-after-cursors", "two-target-model-names"], +) +def test_require_managed_files_rejects_a_list_it_cannot_scope( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, path: str, params: dict[str, str] +) -> None: + _, provider_list = _setup_unscoped_list_files_route_over_real_hook( + mocker, + monkeypatch, + _router_with_wildcard_and_unprefixed_anthropic_deployments(), + _callers_files_on_wildcard_unprefixed_and_passthrough_deployments(), + ) + monkeypatch.setattr("litellm.require_managed_files", True) + + response: Final = _list_files_as_test_user(path, params, {}) + + assert response.status_code == 400, response.text + provider_list.assert_not_awaited() + + +@pytest.mark.parametrize("cursor_param", ["after", "after_id"]) +def test_require_managed_files_pages_a_provider_list_with_either_cursor_param( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, llm_router: Router, cursor_param: str +) -> None: + import litellm.proxy.proxy_server as ps + + _setup_unscoped_list_files_route_over_real_hook( + mocker, monkeypatch, llm_router, _callers_files_on_every_deployment() + ) + monkeypatch.setattr("litellm.require_managed_files", True) + + try: + first_page: Final = client.get( + "/openai/v1/files", params={"limit": "1"}, headers={"Authorization": "Bearer test-key"} + ) + second_page: Final = client.get( + "/openai/v1/files", + params={"limit": "1", cursor_param: first_page.json()["last_id"]}, + headers={"Authorization": "Bearer test-key"}, + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert (first_page.status_code, second_page.status_code) == (200, 200), (first_page.text, second_page.text) + assert [file["id"] for file in first_page.json()["data"]] == ["unified-gpt-newer"] + assert first_page.json()["has_more"] is True + assert [file["id"] for file in second_page.json()["data"]] == ["unified-gpt-older"] + assert second_page.json()["has_more"] is False + + +def test_require_managed_files_pages_toward_newer_files_with_before_id( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, llm_router: Router +) -> None: + import litellm.proxy.proxy_server as ps + + _setup_unscoped_list_files_route_over_real_hook( + mocker, monkeypatch, llm_router, _callers_files_on_every_deployment() + ) + monkeypatch.setattr("litellm.require_managed_files", True) + + try: + page_before_oldest: Final = client.get( + "/v1/files", + params={"limit": "2", "before_id": "unified-gpt-older"}, + headers={"Authorization": "Bearer test-key"}, + ) + newest_page: Final = client.get( + "/v1/files", + params={"limit": "2", "before_id": page_before_oldest.json()["first_id"]}, + headers={"Authorization": "Bearer test-key"}, + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert (page_before_oldest.status_code, newest_page.status_code) == (200, 200), ( + page_before_oldest.text, + newest_page.text, + ) + assert [file["id"] for file in page_before_oldest.json()["data"]] == ["unified-azure", "unified-gemini"] + assert page_before_oldest.json()["has_more"] is True + assert [file["id"] for file in newest_page.json()["data"]] == ["unified-gpt-newer"] + assert newest_page.json()["has_more"] is False + + +def test_managed_file_list_rejects_both_cursor_directions( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, llm_router: Router +) -> None: + _setup_unscoped_list_files_route_over_real_hook( + mocker, monkeypatch, llm_router, _callers_files_on_every_deployment() + ) + monkeypatch.setattr("litellm.require_managed_files", True) + + response: Final = _list_files_as_test_user( + "/v1/files", {"after": "unified-gpt-newer", "before_id": "unified-gpt-older"}, {} + ) + + assert response.status_code == 400, response.text + assert "not both" in response.json()["error"]["message"] + + +@pytest.mark.parametrize( + ("params", "headers"), + [({"target_model_names": "gemini-2.0-flash"}, {}), ({}, {"x-litellm-model": "gemini-2.0-flash"})], + ids=["target-model-names", "model-header"], +) +def test_require_managed_files_list_rejects_a_model_the_key_cannot_use( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + llm_router: Router, + params: dict[str, str], + headers: dict[str, str], +) -> None: + import litellm.proxy.proxy_server as ps + + _setup_unscoped_list_files_route_over_real_hook( + mocker, monkeypatch, llm_router, _callers_files_on_every_deployment() + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="test-key", + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="test-user", + models=["gpt-3.5-turbo"], + ) + monkeypatch.setattr("litellm.require_managed_files", True) + + response: Final = _list_files_as_test_user("/v1/files", params, headers) + + assert response.status_code == 403, response.text + assert "unified-gemini" not in response.text + + +def test_require_managed_files_list_fails_closed_without_the_managed_files_hook( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, llm_router: Router +) -> None: + proxy_logging_obj, provider_list = _setup_unscoped_list_files_route_over_real_hook( + mocker, monkeypatch, llm_router, _callers_files_on_every_deployment() + ) + proxy_logging_obj.proxy_hook_mapping.pop("managed_files") + provider_list.return_value = [_managed_file("file-on-the-shared-key")] + monkeypatch.setattr("litellm.require_managed_files", True) + + response: Final = _list_files_as_test_user("/openai/v1/files", {}, {}) + + assert response.status_code == 500, response.text + assert "file-on-the-shared-key" not in response.text + assert "require_managed_files cannot scope the list to the caller" in response.json()["error"]["message"] + provider_list.assert_not_awaited() + + +@pytest.mark.parametrize("hook_loaded", [True, False], ids=["hook-loaded", "no-hook"]) +def test_require_managed_files_keeps_the_provider_list_for_an_admin( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + llm_router: Router, + hook_loaded: bool, +) -> None: + import litellm.proxy.proxy_server as ps + + proxy_logging_obj, provider_list = _setup_unscoped_list_files_route_over_real_hook( + mocker, monkeypatch, llm_router, _callers_files_on_every_deployment() + ) + if not hook_loaded: + proxy_logging_obj.proxy_hook_mapping.pop("managed_files") + provider_list.return_value = [_managed_file("file-on-the-shared-key")] + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="test-key", user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user" + ) + monkeypatch.setattr("litellm.require_managed_files", True) + + response: Final = _list_files_as_test_user("/anthropic/v1/files", {}, {}) + + assert response.status_code == 200, response.text + assert [file["id"] for file in response.json()["data"]] == ["file-on-the-shared-key"] + + +def test_managed_file_list_without_require_managed_files_ignores_cursor_params_openai_does_not_define( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, llm_router: Router +) -> None: + _setup_unscoped_list_files_route_over_real_hook( + mocker, monkeypatch, llm_router, _callers_files_on_every_deployment() + ) + monkeypatch.setattr("litellm.require_managed_files", False) + + response: Final = _list_files_as_test_user( + "/v1/files", {"after": "unified-gpt-newer", "after_id": "unified-gemini", "before_id": "file-raw"}, {} + ) + + assert response.status_code == 200, response.text + assert [file["id"] for file in response.json()["data"]] == ["unified-azure", "unified-gemini", "unified-gpt-older"] + + +@pytest.mark.parametrize( + ("path", "headers", "expected_count", "has_more"), + [ + ("/anthropic/v1/files", {}, ANTHROPIC_DEFAULT_FILE_LIST_LIMIT, True), + ("/v1/files", {"x-litellm-model": "claude-unprefixed"}, ANTHROPIC_DEFAULT_FILE_LIST_LIMIT, True), + ("/anthropic/v1/files", {"x-litellm-model": "gpt-3.5-turbo"}, ANTHROPIC_DEFAULT_FILE_LIST_LIMIT + 1, False), + ("/v1/files", {}, 2 * (ANTHROPIC_DEFAULT_FILE_LIST_LIMIT + 1), False), + ], + ids=["anthropic-provider", "anthropic-model-header", "openai-model-header", "every-provider"], +) +def test_require_managed_files_pages_an_anthropic_list_at_the_anthropic_default_size( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + path: str, + headers: dict[str, str], + expected_count: int, + has_more: bool, +) -> None: + _setup_unscoped_list_files_route_over_real_hook( + mocker, + monkeypatch, + _router_with_wildcard_and_unprefixed_anthropic_deployments(), + [ + _managed_file_row( + f"unified-{deployment_id}-{index:03d}", model_mappings={deployment_id: f"file-{index:03d}"} + ) + for deployment_id, index in itertools.product( + ("claude-unprefixed-id", "gpt-3.5-turbo-id"), range(ANTHROPIC_DEFAULT_FILE_LIST_LIMIT + 1) + ) + ], + ) + monkeypatch.setattr("litellm.require_managed_files", True) + + response: Final = _list_files_as_test_user(path, {}, headers) + + assert response.status_code == 200, response.text + assert (len(response.json()["data"]), response.json()["has_more"]) == (expected_count, has_more) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 45f9630983f..70406105193 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -59857,6 +59857,8 @@ export interface operations { purpose?: string | null; limit?: number | null; after?: string | null; + after_id?: string | null; + before_id?: string | null; }; header?: never; path?: never; @@ -79053,6 +79055,8 @@ export interface operations { purpose?: string | null; limit?: number | null; after?: string | null; + after_id?: string | null; + before_id?: string | null; }; header?: never; path?: never; @@ -86388,6 +86392,8 @@ export interface operations { purpose?: string | null; limit?: number | null; after?: string | null; + after_id?: string | null; + before_id?: string | null; }; header?: never; path: { From cf7c35a1174eb9864cf35c1513c2d6d50b5df26a Mon Sep 17 00:00:00 2001 From: moe-berri Date: Wed, 28 Oct 2026 13:48:37 -0700 Subject: [PATCH 089/128] fix(anthropic): map model_context_window_exceeded stop reason to length (#48874) * test(anthropic): cover model_context_window_exceeded finish reason on chat and responses Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(anthropic): map model_context_window_exceeded stop reason to length Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(rust): mirror model_context_window_exceeded finish reason mapping 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> --- .../crates/core-utils/src/core_helpers.rs | 2 + litellm/litellm_core_utils/core_helpers.py | 1 + ...test_anthropic_context_window_stop_wire.py | 122 ++++++++++++++++++ .../litellm_core_utils/test_core_helpers.py | 1 + 4 files changed, 126 insertions(+) create mode 100644 tests/integration/providers/test_anthropic_context_window_stop_wire.py diff --git a/litellm-rust/crates/core-utils/src/core_helpers.rs b/litellm-rust/crates/core-utils/src/core_helpers.rs index ada1c3ceb1a..58ee88d301c 100644 --- a/litellm-rust/crates/core-utils/src/core_helpers.rs +++ b/litellm-rust/crates/core-utils/src/core_helpers.rs @@ -13,6 +13,7 @@ const FINISH_REASONS: &[(&str, &str)] = &[ ("max_tokens", "length"), ("refusal", "content_filter"), ("compaction", "length"), + ("model_context_window_exceeded", "length"), ("guardrail_intervened", "content_filter"), ("content_filtered", "content_filter"), ("content_filter", "content_filter"), @@ -74,6 +75,7 @@ mod tests { #[case::end_turn("end_turn", "stop")] #[case::stop_sequence("stop_sequence", "stop")] #[case::max_tokens("max_tokens", "length")] + #[case::model_context_window_exceeded("model_context_window_exceeded", "length")] #[case::refusal("refusal", "content_filter")] #[case::guardrail_intervened("guardrail_intervened", "content_filter")] #[case::content_filtered("content_filtered", "content_filter")] diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 00c6ce97e78..55025508e40 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -266,6 +266,7 @@ _FINISH_REASON_MAP: Final[dict[str, OpenAIChatCompletionFinishReason]] = { "tool_use": "tool_calls", "refusal": "content_filter", "compaction": "length", + "model_context_window_exceeded": "length", # Cohere "COMPLETE": "stop", "ERROR_TOXIC": "content_filter", diff --git a/tests/integration/providers/test_anthropic_context_window_stop_wire.py b/tests/integration/providers/test_anthropic_context_window_stop_wire.py new file mode 100644 index 00000000000..38e3e7dffcd --- /dev/null +++ b/tests/integration/providers/test_anthropic_context_window_stop_wire.py @@ -0,0 +1,122 @@ +import json +import uuid +from collections.abc import Mapping +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +MODEL: Final = "claude-haiku-4-5" + + +def _sse(event: Mapping[str, JsonValue]) -> bytes: + return f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() + + +def _message(identity: str, stop_reason: str | None) -> dict[str, JsonValue]: + return { + "id": "msg_" + identity, + "type": "message", + "role": "assistant", + "model": MODEL, + "content": [{"type": "text", "text": "alpha bravo"}], + "stop_reason": stop_reason, + "stop_sequence": None, + "usage": {"input_tokens": 199878, "output_tokens": 12799}, + } + + +def _stream(identity: str, stop_reason: str) -> tuple[bytes, ...]: + return ( + _sse({"type": "message_start", "message": {**_message(identity, None), "content": []}}), + _sse({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + _sse({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "alpha bravo"}}), + _sse({"type": "content_block_stop", "index": 0}), + _sse( + { + "type": "message_delta", + "delta": {"stop_reason": stop_reason, "stop_sequence": None}, + "usage": {"output_tokens": 12799}, + } + ), + _sse({"type": "message_stop"}), + ) + + +def _anthropic(identity: str, stop_reason: str): + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/messages", request.target + if json.loads(request.body).get("stream"): + return Reply(content_type="text/event-stream", chunks=_stream(identity, stop_reason)) + return Reply(body=json.dumps(_message(identity, stop_reason)).encode()) + + return respond + + +def _events(gateway: Gateway, path: str, body: Mapping[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]: + with gateway.client.stream( + "POST", path, json={**body, "stream": True}, headers={"Authorization": f"Bearer {gateway.key}"} + ) as response: + assert response.status_code == 200, response.read().decode() + return tuple( + json.loads(line.removeprefix("data: ")) + for line in response.iter_lines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + + +def _chat_finish_reason(gateway: Gateway, model: str, *, stream: bool) -> JsonValue: + body: Final = {"model": model, "max_tokens": 64000, "messages": [{"role": "user", "content": "copy the list"}]} + if stream: + reasons: Final = tuple( + choice["finish_reason"] + for chunk in _events(gateway, "/v1/chat/completions", body) + for choice in chunk["choices"] # comprehension-ok: chunk -> choices flatten + if choice.get("finish_reason") + ) + assert len(reasons) == 1, reasons + return reasons[0] + response: Final = gateway.request("POST", "/v1/chat/completions", body) + assert response.status_code == 200, response.text + return response.json()["choices"][0]["finish_reason"] + + +def _responses_status(gateway: Gateway, model: str, *, stream: bool) -> JsonValue: + body: Final = {"model": model, "max_output_tokens": 64000, "input": "copy the list"} + if stream: + terminal: Final = tuple( + event["response"] + for event in _events(gateway, "/v1/responses", body) + if event["type"] in ("response.completed", "response.incomplete", "response.failed") + ) + assert len(terminal) == 1, terminal + return terminal[0]["status"] + response: Final = gateway.request("POST", "/v1/responses", body) + assert response.status_code == 200, response.text + return response.json()["status"] + + +@pytest.mark.parametrize("stream", [False, True], ids=["json", "stream"]) +@pytest.mark.parametrize( + ("stop_reason", "finish_reason"), + [("model_context_window_exceeded", "length"), ("max_tokens", "length"), ("end_turn", "stop")], +) +def test_chat_maps_anthropic_stop_reason(gateway: Gateway, stop_reason: str, finish_reason: str, stream: bool) -> None: + identity: Final = uuid.uuid4().hex + with wire_server(_anthropic(identity, stop_reason)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{MODEL}", api_base=wire.url, api_key="scripted-anthropic-key") + assert _chat_finish_reason(gateway, model, stream=stream) == finish_reason + + +@pytest.mark.parametrize("stream", [False, True], ids=["json", "stream"]) +@pytest.mark.parametrize( + ("stop_reason", "status"), + [("model_context_window_exceeded", "incomplete"), ("max_tokens", "incomplete"), ("end_turn", "completed")], +) +def test_responses_maps_anthropic_stop_reason(gateway: Gateway, stop_reason: str, status: str, stream: bool) -> None: + identity: Final = uuid.uuid4().hex + with wire_server(_anthropic(identity, stop_reason)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{MODEL}", api_base=wire.url, api_key="scripted-anthropic-key") + assert _responses_status(gateway, model, stream=stream) == status diff --git a/tests/unit/litellm_core_utils/test_core_helpers.py b/tests/unit/litellm_core_utils/test_core_helpers.py index d66c37ba8c2..66a8e6bbda6 100644 --- a/tests/unit/litellm_core_utils/test_core_helpers.py +++ b/tests/unit/litellm_core_utils/test_core_helpers.py @@ -256,6 +256,7 @@ class TestMapFinishReasonAnthropic: ("max_tokens", "length"), ("tool_use", "tool_calls"), ("compaction", "length"), + ("model_context_window_exceeded", "length"), ("content_filtered", "content_filter"), ], ) From 4516a251c940cc6ff7977b460a26a0582a0a18ff Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 28 Oct 2026 14:39:25 -0700 Subject: [PATCH 090/128] feat(accounting): add accounting contracts, PricingSnapshot and ledger schema (LIT-10193) (#48862) * feat(accounting): add typed contracts and spend log adapters * refactor(accounting): reuse spend identity constants and register tests * refactor(accounting): remove unused spend log conversion API * feat(accounting): add PricingSnapshot, delivery enums and projection dimensions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(db): add LiteLLM_AccountingEvent and LiteLLM_AccountingDelivery ledger tables Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(accounting): export JSON schema and semantic fixtures for the Rust cost crate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(accounting): drop the schema and fixture export, the Python package is the contract Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(accounting): make tests/unit/accounting a package Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(accounting): validate pricing snapshots and ledger rows * fix(otel): register the accounting ledger tables as span tables Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(accounting): address review findings on dimensions, token subsets, multipliers and agent target Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Co-authored-by: joshua Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-unit.yml | 1 + .../migration.sql | 47 ++ .../litellm_proxy_extras/schema.prisma | 70 +++ litellm/accounting/__init__.py | 62 +++ litellm/accounting/measurements.py | 13 + litellm/accounting/protocols.py | 85 +++ litellm/accounting/types.py | 428 ++++++++++++++++ litellm/integrations/otel/model/spans.py | 2 + litellm/proxy/schema.prisma | 70 +++ schema.prisma | 70 +++ .../database/test_accounting_ledger.py | 128 +++++ tests/unit/accounting/__init__.py | 0 tests/unit/accounting/test_measurements.py | 235 +++++++++ tests/unit/accounting/test_types.py | 484 ++++++++++++++++++ 14 files changed, 1695 insertions(+) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20261027000000_add_accounting_ledger_tables/migration.sql create mode 100644 litellm/accounting/__init__.py create mode 100644 litellm/accounting/measurements.py create mode 100644 litellm/accounting/protocols.py create mode 100644 litellm/accounting/types.py create mode 100644 tests/integration/database/test_accounting_ledger.py create mode 100644 tests/unit/accounting/__init__.py create mode 100644 tests/unit/accounting/test_measurements.py create mode 100644 tests/unit/accounting/test_types.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index e3980bd8f3b..2ddf21bca17 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -137,6 +137,7 @@ jobs: include: - shard: core-utils test-path: |- + tests/unit/accounting tests/unit/batch_completion tests/unit/decisions tests/unit/extraction diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261027000000_add_accounting_ledger_tables/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261027000000_add_accounting_ledger_tables/migration.sql new file mode 100644 index 00000000000..c0ad8692d36 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261027000000_add_accounting_ledger_tables/migration.sql @@ -0,0 +1,47 @@ +SET LOCAL lock_timeout = '10s'; + +CREATE TYPE "AccountingEventKind" AS ENUM ('inbound_usage', 'outbound_usage', 'cost'); + +CREATE TYPE "AccountingDeliveryState" AS ENUM ('pending', 'applied', 'parked'); + +CREATE TYPE "AccountingDeliveryTarget" AS ENUM ('LiteLLM_VerificationToken', 'LiteLLM_UserTable', 'LiteLLM_TeamTable', 'LiteLLM_TeamMembership', 'LiteLLM_OrganizationTable', 'LiteLLM_OrganizationMembership', 'LiteLLM_EndUserTable', 'LiteLLM_AgentsTable', 'LiteLLM_TagTable', 'LiteLLM_ProjectTable', 'LiteLLM_ModelAccessGroupBudgetTable', 'LiteLLM_BudgetTable', 'LiteLLM_BudgetWindowSpend', 'LiteLLM_DailyUserSpend', 'LiteLLM_DailyTeamSpend', 'LiteLLM_DailyOrganizationSpend', 'LiteLLM_DailyEndUserSpend', 'LiteLLM_DailyAgentSpend', 'LiteLLM_DailyTagSpend', 'LiteLLM_DailyModelUsage', 'LiteLLM_DailyToolSpend', 'LiteLLM_AutoRouterDailySpend'); + +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_AccountingEvent" ( + "event_id" TEXT NOT NULL, + "kind" "AccountingEventKind" NOT NULL, + "source_event_id" TEXT, + "schema_version" INTEGER NOT NULL, + "payload_hash" TEXT NOT NULL, + "occurred_at" TIMESTAMP(3) NOT NULL, + "recorded_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "payload" JSONB NOT NULL, + + CONSTRAINT "LiteLLM_AccountingEvent_schema_version_check" CHECK ("schema_version" > 0), + CONSTRAINT "LiteLLM_AccountingEvent_pkey" PRIMARY KEY ("event_id") +); + +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_AccountingDelivery" ( + "event_id" TEXT NOT NULL, + "target" "AccountingDeliveryTarget" NOT NULL, + "version" INTEGER NOT NULL, + "state" "AccountingDeliveryState" NOT NULL DEFAULT 'pending', + "attempts" INTEGER NOT NULL DEFAULT 0, + "next_attempt_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "applied_at" TIMESTAMP(3), + "last_error" TEXT, + + CONSTRAINT "LiteLLM_AccountingDelivery_version_check" CHECK ("version" > 0), + CONSTRAINT "LiteLLM_AccountingDelivery_attempts_check" CHECK ("attempts" >= 0), + CONSTRAINT "LiteLLM_AccountingDelivery_pkey" PRIMARY KEY ("event_id","target","version") +); + +-- CreateIndex +CREATE INDEX IF NOT EXISTS "LiteLLM_AccountingEvent_recorded_at_idx" ON "LiteLLM_AccountingEvent"("recorded_at"); + +-- CreateIndex +CREATE INDEX IF NOT EXISTS "LiteLLM_AccountingEvent_kind_source_event_id_idx" ON "LiteLLM_AccountingEvent"("kind", "source_event_id"); + +-- CreateIndex +CREATE INDEX IF NOT EXISTS "LiteLLM_AccountingDelivery_target_state_next_attempt_at_idx" ON "LiteLLM_AccountingDelivery"("target", "state", "next_attempt_at"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 4187f76c97c..2965174114b 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -2061,3 +2061,73 @@ model LiteLLM_TelemetryReport { @@index([window_end]) @@index([created_at]) } + +enum AccountingEventKind { + inbound_usage + outbound_usage + cost +} + +enum AccountingDeliveryState { + pending + applied + parked +} + +enum AccountingDeliveryTarget { + LiteLLM_VerificationToken + LiteLLM_UserTable + LiteLLM_TeamTable + LiteLLM_TeamMembership + LiteLLM_OrganizationTable + LiteLLM_OrganizationMembership + LiteLLM_EndUserTable + LiteLLM_AgentsTable + LiteLLM_TagTable + LiteLLM_ProjectTable + LiteLLM_ModelAccessGroupBudgetTable + LiteLLM_BudgetTable + LiteLLM_BudgetWindowSpend + LiteLLM_DailyUserSpend + LiteLLM_DailyTeamSpend + LiteLLM_DailyOrganizationSpend + LiteLLM_DailyEndUserSpend + LiteLLM_DailyAgentSpend + LiteLLM_DailyTagSpend + LiteLLM_DailyModelUsage + LiteLLM_DailyToolSpend + LiteLLM_AutoRouterDailySpend +} + +// Durable accounting ledger. One row per usage or cost event; payload holds the +// private event envelope versioned by schema_version. Usage rows retain request-scoped +// pricing context; cost rows include resolved rates and the finalized pricing snapshot. +model LiteLLM_AccountingEvent { + event_id String @id + kind AccountingEventKind + source_event_id String? + schema_version Int + payload_hash String + occurred_at DateTime + recorded_at DateTime @default(now()) + payload Json + + @@index([recorded_at]) + @@index([kind, source_event_id]) +} + +// One row per (event, target table, version): pending work for the projector and, +// once applied in the same transaction as the target write, the receipt that stops a retry. +model LiteLLM_AccountingDelivery { + event_id String + target AccountingDeliveryTarget + version Int + state AccountingDeliveryState @default(pending) + attempts Int @default(0) + next_attempt_at DateTime @default(now()) + applied_at DateTime? + last_error String? + + @@id([event_id, target, version]) + @@index([target, state, next_attempt_at]) +} diff --git a/litellm/accounting/__init__.py b/litellm/accounting/__init__.py new file mode 100644 index 00000000000..d4e56f61279 --- /dev/null +++ b/litellm/accounting/__init__.py @@ -0,0 +1,62 @@ +"""Public accounting contracts. Importing these values does not enable accounting.""" + +from litellm.accounting.measurements import measurements_for +from litellm.accounting.protocols import BudgetEnforcer, CostAggregator, CostCompiler, UsageMeter +from litellm.accounting.types import ( + AdmissionDecision, + AggregateQuery, + AggregateRow, + CostEvent, + DeliveryState, + DeliveryTarget, + Dimension, + DimensionFilter, + DimensionName, + EventKind, + Granularity, + InboundUsageEvent, + Measurement, + Metric, + MetricName, + Multiplier, + NamedCost, + OutboundUsageEvent, + Outcome, + PricedQuantity, + PricingSnapshot, + Rate, + Unpriced, + validate_aggregate_rows, +) + +__all__ = [ + "AdmissionDecision", + "AggregateQuery", + "AggregateRow", + "BudgetEnforcer", + "CostAggregator", + "CostCompiler", + "CostEvent", + "DeliveryState", + "DeliveryTarget", + "Dimension", + "DimensionFilter", + "DimensionName", + "EventKind", + "Granularity", + "InboundUsageEvent", + "Measurement", + "Metric", + "MetricName", + "Multiplier", + "NamedCost", + "OutboundUsageEvent", + "Outcome", + "PricedQuantity", + "PricingSnapshot", + "Rate", + "Unpriced", + "UsageMeter", + "measurements_for", + "validate_aggregate_rows", +] diff --git a/litellm/accounting/measurements.py b/litellm/accounting/measurements.py new file mode 100644 index 00000000000..83011a280da --- /dev/null +++ b/litellm/accounting/measurements.py @@ -0,0 +1,13 @@ +from typing import Final + +from litellm.accounting.types import CostEvent, InboundUsageEvent, Measurement, OutboundUsageEvent + + +def measurements_for(event: InboundUsageEvent | OutboundUsageEvent | CostEvent) -> tuple[Measurement, ...]: + """Return the complete contribution set for one event, ready for one apply call.""" + metrics: Final = event.charges if isinstance(event, CostEvent) else event.metrics + occurred_at: Final = event.occurred_at if isinstance(event, CostEvent) else event.started_at + return tuple( + Measurement(source_event_id=event.event_id, occurred_at=occurred_at, dimensions=event.dimensions, metric=metric) + for metric in metrics or () + ) diff --git a/litellm/accounting/protocols.py b/litellm/accounting/protocols.py new file mode 100644 index 00000000000..8955f56e7a0 --- /dev/null +++ b/litellm/accounting/protocols.py @@ -0,0 +1,85 @@ +"""Contracts only. Implementations own persistence, pricing, projections and live budgets. + +The event row's ``payload`` is a private envelope versioned by the row's ``schema_version``. +Usage envelopes retain the public event and request-scoped pricing context: quantities with units, +model/provider, requested and billed tier/region when known, and per-request/deployment overrides. +These inputs must survive request teardown. Cost-map rates are resolved at compilation, not capture. +Cost envelopes retain the public CostEvent, including its required finalized pricing snapshot. +T5/T6 own envelope serialization; this public API does not introduce a second persisted schema. +""" + +from typing import Protocol + +from litellm.accounting.types import ( + AdmissionDecision, + AggregateQuery, + AggregateRow, + CostEvent, + Dimension, + InboundUsageEvent, + Measurement, + Metric, + OutboundUsageEvent, + Unpriced, +) + + +class UsageMeter(Protocol): + async def record(self, event: InboundUsageEvent | OutboundUsageEvent) -> None: + """Acknowledge commit, no-op identical IDs and reject conflicting content. + + Persist the event and its request-scoped pricing context before acknowledging. + HTTP acknowledgement is a separate policy. Unknown usage remains None. + """ + ... + + +class CostCompiler(Protocol): + async def compile(self, event: OutboundUsageEvent) -> CostEvent | Unpriced: + """Reuse the committed charge or finalize one charge and its delivery intent. + + Resolve captured request-scoped context by event identity and cost-map rates + at compilation. Persist the resolved inputs in the required snapshot. Replay + reuses the committed result and never consumes free allowance again. + Missing pricing is Unpriced, never an inferred zero. + """ + ... + + +class CostAggregator(Protocol): + async def apply(self, measurements: tuple[Measurement, ...]) -> None: + """Apply one source event's complete contribution set per target/version. + + Each target commits all contributions and its receipt atomically. A receipt + for the first metric must not suppress the rest. Usage need not wait for cost. + """ + ... + + async def query(self, query: AggregateQuery) -> tuple[AggregateRow, ...]: + """Return disjoint [start, end) rows per group using existing daily projections. + + Filters combine with AND, values within a filter with IN. DAY emits daily + buckets; None emits whole-interval totals. Reject unsupported metric/scope + combinations. Different groups may overlap, including non-additive tags. + """ + ... + + +class BudgetEnforcer(Protocol): + async def reserve( + self, reservation_id: str, dimensions: tuple[Dimension, ...], estimated: tuple[Metric, ...] + ) -> AdmissionDecision: + """Apply existing admission policies and estimates; duplicate IDs add no hold.""" + ... + + async def settle(self, reservation_id: str | None, cost_event: CostEvent) -> None: + """Settle the committed charge with atomic deduplication and hold mutation. + + Preserve scope/window attribution and retry unfinished settlement after + restart. SQL projection receipts do not provide live-counter atomicity. + """ + ... + + async def release(self, reservation_id: str) -> None: + """Idempotently release unused exposure at the request boundary, preserving charges.""" + ... diff --git a/litellm/accounting/types.py b/litellm/accounting/types.py new file mode 100644 index 00000000000..d0be29bf5b9 --- /dev/null +++ b/litellm/accounting/types.py @@ -0,0 +1,428 @@ +"""Accounting values use UTC instants, decimal strings on the wire and USD charges.""" + +from collections.abc import Mapping +from datetime import datetime, timezone +from decimal import Decimal +from enum import Enum +from typing import Annotated, Final, Literal, TypeAlias + +from pydantic import AfterValidator, AwareDatetime, BaseModel, ConfigDict, Field, model_validator +from typing_extensions import Self + +from litellm.types.utils import CallTypes + +_CONFIG: Final = ConfigDict(frozen=True, extra="forbid", validate_default=True) +Identifier: TypeAlias = Annotated[str, Field(strict=True, min_length=1, pattern=r"\S")] +Quantity: TypeAlias = Annotated[Decimal, Field(strict=True, ge=0, allow_inf_nan=False)] + + +def _utc(value: datetime) -> datetime: + return value.astimezone(timezone.utc) + + +Instant: TypeAlias = Annotated[AwareDatetime, AfterValidator(_utc)] + + +class MetricName(str, Enum): + """Reasoning tokens are included in completion tokens; above-1hr writes are included in cache creation tokens.""" + + REQUESTS = "requests" + SUCCESSFUL_REQUESTS = "successful_requests" + FAILED_REQUESTS = "failed_requests" + PROMPT_TOKENS = "prompt_tokens" + COMPLETION_TOKENS = "completion_tokens" + REASONING_TOKENS = "reasoning_tokens" + TOTAL_TOKENS = "total_tokens" + CACHE_READ_INPUT_TOKENS = "cache_read_input_tokens" + CACHE_CREATION_INPUT_TOKENS = "cache_creation_input_tokens" + CACHE_CREATION_INPUT_TOKENS_ABOVE_1HR = "cache_creation_input_tokens_above_1hr" + SPEND = "spend" + + +class DimensionName(str, Enum): + KEY = "key" + USER = "user" + TEAM = "team" + ORGANIZATION = "organization" + END_USER = "end_user" + MODEL = "model" + MODEL_GROUP = "model_group" + DEPLOYMENT = "deployment" + PROVIDER = "provider" + TOOL = "tool" + AGENT = "agent" + PROJECT = "project" + MODEL_ACCESS_GROUP = "model_access_group" + API_BASE = "api_base" + SERVICE_TIER = "service_tier" + TAG = "tag" + + +class Outcome(str, Enum): + SUCCEEDED = "succeeded" + FAILED = "failed" + CANCELLED = "cancelled" + + +class Granularity(str, Enum): + DAY = "day" + + +class EventKind(str, Enum): + INBOUND_USAGE = "inbound_usage" + OUTBOUND_USAGE = "outbound_usage" + COST = "cost" + + +class DeliveryState(str, Enum): + PENDING = "pending" + APPLIED = "applied" + PARKED = "parked" + + +class DeliveryTarget(str, Enum): + """One member per existing table the projector writes; each is applied and receipted independently.""" + + KEY_SPEND = "LiteLLM_VerificationToken" + USER_SPEND = "LiteLLM_UserTable" + TEAM_SPEND = "LiteLLM_TeamTable" + TEAM_MEMBER_SPEND = "LiteLLM_TeamMembership" + ORGANIZATION_SPEND = "LiteLLM_OrganizationTable" + ORGANIZATION_MEMBER_SPEND = "LiteLLM_OrganizationMembership" + END_USER_SPEND = "LiteLLM_EndUserTable" + AGENT_SPEND = "LiteLLM_AgentsTable" + TAG_SPEND = "LiteLLM_TagTable" + PROJECT_SPEND = "LiteLLM_ProjectTable" + MODEL_ACCESS_GROUP_SPEND = "LiteLLM_ModelAccessGroupBudgetTable" + BUDGET_SPEND = "LiteLLM_BudgetTable" + BUDGET_WINDOW_SPEND = "LiteLLM_BudgetWindowSpend" + DAILY_USER_SPEND = "LiteLLM_DailyUserSpend" + DAILY_TEAM_SPEND = "LiteLLM_DailyTeamSpend" + DAILY_ORGANIZATION_SPEND = "LiteLLM_DailyOrganizationSpend" + DAILY_END_USER_SPEND = "LiteLLM_DailyEndUserSpend" + DAILY_AGENT_SPEND = "LiteLLM_DailyAgentSpend" + DAILY_TAG_SPEND = "LiteLLM_DailyTagSpend" + DAILY_MODEL_USAGE = "LiteLLM_DailyModelUsage" + DAILY_TOOL_SPEND = "LiteLLM_DailyToolSpend" + AUTOROUTER_DAILY_SPEND = "LiteLLM_AutoRouterDailySpend" + + +_SUBSET_PAIRS: Final = ( + (MetricName.REASONING_TOKENS, MetricName.COMPLETION_TOKENS), + (MetricName.CACHE_CREATION_INPUT_TOKENS_ABOVE_1HR, MetricName.CACHE_CREATION_INPUT_TOKENS), +) + + +def _check_subsets(values: Mapping[str, Decimal]) -> None: + for part, whole in _SUBSET_PAIRS: + if part.value in values and whole.value in values and values[part.value] > values[whole.value]: + raise ValueError(f"{part.value} cannot exceed {whole.value}") + + +class Metric(BaseModel): + model_config = _CONFIG + + name: MetricName + value: Quantity + unit: Literal["count", "USD"] = "count" + + @model_validator(mode="after") + def validate_unit(self) -> Self: + if (self.name is MetricName.SPEND) != (self.unit == "USD"): + raise ValueError("SPEND requires USD; usage metrics require count") + if self.unit == "count" and self.value != self.value.to_integral_value(): + raise ValueError("Count metrics require integral values") + return self + + +class Dimension(BaseModel): + model_config = _CONFIG + + name: DimensionName + value: Identifier + + +_REPEATABLE_DIMENSIONS: Final = frozenset({DimensionName.TAG, DimensionName.MODEL_ACCESS_GROUP}) + + +def _dimensions(values: tuple[Dimension, ...]) -> tuple[Dimension, ...]: + if len(frozenset(values)) != len(values): + raise ValueError("Duplicate dimension values") + single_names: Final = tuple(value.name for value in values if value.name not in _REPEATABLE_DIMENSIONS) + if len(frozenset(single_names)) != len(single_names): + raise ValueError("Only tag and model access group dimensions may have multiple values") + return values + + +Dimensions: TypeAlias = Annotated[tuple[Dimension, ...], AfterValidator(_dimensions)] +_CALLER_METRICS: Final = frozenset((MetricName.REQUESTS, MetricName.SUCCESSFUL_REQUESTS, MetricName.FAILED_REQUESTS)) + + +def _metrics(values: tuple[Metric, ...]) -> tuple[Metric, ...]: + if len(frozenset(value.name for value in values)) != len(values): + raise ValueError("Duplicate metric names") + return values + + +Metrics: TypeAlias = Annotated[tuple[Metric, ...], AfterValidator(_metrics)] + + +def _interval(start: datetime, end: datetime, *, allow_empty: bool = False) -> None: + if end < start or (end == start and not allow_empty): + raise ValueError("End must follow start") + + +class InboundUsageEvent(BaseModel): + model_config = _CONFIG + + event_id: Identifier + request_id: Identifier + call_type: CallTypes + started_at: Instant + ended_at: Instant + outcome: Outcome + dimensions: Dimensions + metrics: Metrics + + @model_validator(mode="after") + def validate_caller_activity(self) -> Self: + _interval(self.started_at, self.ended_at, allow_empty=True) + if any(metric.name not in _CALLER_METRICS for metric in self.metrics): + raise ValueError("Inbound events own only caller request metrics") + if not any(metric.name is MetricName.REQUESTS and metric.value == 1 for metric in self.metrics): + raise ValueError("Inbound events count exactly one caller request") + for metric in self.metrics: + if metric.name is MetricName.SUCCESSFUL_REQUESTS and metric.value != int(self.outcome is Outcome.SUCCEEDED): + raise ValueError("Successful request count must match outcome") + if metric.name is MetricName.FAILED_REQUESTS and metric.value != int(self.outcome is not Outcome.SUCCEEDED): + raise ValueError("Failed request count must match outcome") + return self + + +class OutboundUsageEvent(BaseModel): + model_config = _CONFIG + + event_id: Identifier + attempt_id: Identifier + request_id: Identifier | None + retry_of_attempt_id: Identifier | None + call_type: CallTypes + started_at: Instant + ended_at: Instant + outcome: Outcome + dimensions: Dimensions + metrics: Metrics | None + + @model_validator(mode="after") + def validate_attempt_usage(self) -> Self: + _interval(self.started_at, self.ended_at, allow_empty=True) + if self.retry_of_attempt_id == self.attempt_id: + raise ValueError("An attempt cannot retry itself") + if self.metrics is not None and any( + metric.name in _CALLER_METRICS or metric.name is MetricName.SPEND for metric in self.metrics + ): + raise ValueError("Outbound events own only attempt usage metrics") + if self.metrics is not None: + _check_subsets({metric.name.value: metric.value for metric in self.metrics}) + return self + + +PricingUnit: TypeAlias = Literal["token", "second", "character", "image", "request"] + + +class Rate(BaseModel): + model_config = _CONFIG + + measure: Identifier + unit: PricingUnit + usd_per_unit: Quantity + + +class PricedQuantity(BaseModel): + """A pricing input in its billing unit. Only seconds may be fractional.""" + + model_config = _CONFIG + + measure: Identifier + unit: PricingUnit + value: Quantity + + @model_validator(mode="after") + def validate_discrete_quantity(self) -> Self: + if self.unit != "second" and self.value != self.value.to_integral_value(): + raise ValueError("Discrete priced quantities require integral values") + return self + + +class Multiplier(BaseModel): + """A regional or tier factor applied on top of the priced key, e.g. data_residency=eu or vertex_location=us-east5.""" + + model_config = _CONFIG + + name: Identifier + value: Identifier + factor: Annotated[Decimal, Field(strict=True, gt=0, allow_inf_nan=False)] + + +class NamedCost(BaseModel): + """A free-form extra charge by name, e.g. azure_model_router_flat_cost.""" + + model_config = _CONFIG + + name: Identifier + amount: Quantity + + +class PricingSnapshot(BaseModel): + """Retained inputs and result of authoritative compilation, reproducible after cost-map rates change. + + Request-scoped inputs and overrides come from the usage row's versioned payload. Cost-map rates are + resolved at compilation, not frozen at capture. Retries reuse the committed charge and snapshot. + Component names match ``CostBreakdown``: reasoning_cost is part of output_cost, not an additional charge. + Component sums and float tolerances belong to reconciliation, not this structural contract. + """ + + model_config = _CONFIG + + schema_version: Literal[1] = 1 + calculator_version: Identifier + model: Identifier + provider: Identifier | None + pricing_source: Literal["cost_map", "deployment", "per_request", "upstream_header"] + priced_model_key: Identifier + multipliers: tuple[Multiplier, ...] = () + rates: tuple[Rate, ...] = () + usage: tuple[PricedQuantity, ...] = () + input_cost: Quantity | None = None + output_cost: Quantity | None = None + cache_read_cost: Quantity | None = None + cache_creation_cost: Quantity | None = None + reasoning_cost: Quantity | None = None + tool_usage_cost: Quantity | None = None + guardrail_cost: Quantity | None = None + additional_costs: tuple[NamedCost, ...] = () + original_cost: Quantity | None = None + discount_percent: Quantity | None = None + discount_amount: Quantity | None = None + margin_percent: Quantity | None = None + margin_fixed_amount: Quantity | None = None + margin_total_amount: Quantity | None = None + upstream_reported_cost: Quantity | None = None + total_cost: Quantity + + @model_validator(mode="after") + def validate_snapshot(self) -> Self: + rate_keys: Final = frozenset((rate.measure, rate.unit) for rate in self.rates) + quantity_keys: Final = frozenset((quantity.measure, quantity.unit) for quantity in self.usage) + if len(rate_keys) != len(self.rates): + raise ValueError("Duplicate rate measures and units") + if len(quantity_keys) != len(self.usage): + raise ValueError("Duplicate priced quantity measures and units") + if not rate_keys.issubset(quantity_keys): + raise ValueError("Every rate requires one priced quantity with the same measure and unit") + if len(frozenset(cost.name for cost in self.additional_costs)) != len(self.additional_costs): + raise ValueError("Duplicate additional cost names") + if len(frozenset(multiplier.name for multiplier in self.multipliers)) != len(self.multipliers): + raise ValueError("Duplicate multiplier names") + _check_subsets({quantity.measure: quantity.value for quantity in self.usage if quantity.unit == "token"}) + return self + + +class CostEvent(BaseModel): + model_config = _CONFIG + + event_id: Identifier + usage_event_id: Identifier + occurred_at: Instant + dimensions: Dimensions + charges: Metrics + snapshot: PricingSnapshot + + @model_validator(mode="after") + def validate_charge(self) -> Self: + if len(self.charges) != 1 or self.charges[0].name is not MetricName.SPEND: + raise ValueError("A cost event requires one authoritative USD spend metric") + if self.event_id == self.usage_event_id: + raise ValueError("Cost and usage events require distinct identities") + if self.snapshot.total_cost != self.charges[0].value: + raise ValueError("Snapshot total must equal the authoritative spend") + return self + + +class Unpriced(BaseModel): + model_config = _CONFIG + + usage_event_id: Identifier + reason: Identifier + + +class Measurement(BaseModel): + model_config = _CONFIG + + source_event_id: Identifier + occurred_at: Instant + dimensions: Dimensions + metric: Metric + + +class DimensionFilter(BaseModel): + model_config = _CONFIG + + name: DimensionName + values: Annotated[tuple[Identifier, ...], Field(min_length=1)] + + +class AggregateQuery(BaseModel): + model_config = _CONFIG + + metric: MetricName + filters: tuple[DimensionFilter, ...] + group_by: tuple[DimensionName, ...] + start: Instant + end: Instant + granularity: Granularity | None + + @model_validator(mode="after") + def validate_query(self) -> Self: + _interval(self.start, self.end) + if len(frozenset(item.name for item in self.filters)) != len(self.filters): + raise ValueError("Use one IN filter per dimension; filters combine with AND") + if len(frozenset(self.group_by)) != len(self.group_by): + raise ValueError("Duplicate group-by dimensions") + if any( + (instant.hour, instant.minute, instant.second, instant.microsecond) != (0, 0, 0, 0) + for instant in (self.start, self.end) + ): + raise ValueError("Initial reporting queries require UTC day boundaries") + return self + + +class AggregateRow(BaseModel): + model_config = _CONFIG + + start: Instant + end: Instant + dimensions: Dimensions + value: Quantity + + @model_validator(mode="after") + def validate_bounds(self) -> Self: + _interval(self.start, self.end) + return self + + +def validate_aggregate_rows(rows: tuple[AggregateRow, ...]) -> tuple[AggregateRow, ...]: + """Validate [start, end) buckets per group; overlapping tag groups are not additive.""" + groups: Final = frozenset(frozenset(row.dimensions) for row in rows) + for group in groups: + ordered = sorted((row for row in rows if frozenset(row.dimensions) == group), key=lambda row: row.start) + if any(left.end > right.start for left, right in zip(ordered, ordered[1:])): + raise ValueError("Aggregate buckets overlap within a dimension group") + return rows + + +class AdmissionDecision(BaseModel): + model_config = _CONFIG + + allowed: bool + reason: Identifier | None = None diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index d8be607150c..e423b6698e6 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -338,6 +338,8 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset( "LiteLLM_ToolTable", "LiteLLM_TelemetryReport", "LiteLLM_AccessGroupTable", + "LiteLLM_AccountingDelivery", + "LiteLLM_AccountingEvent", "LiteLLM_ClaudeCodePluginTable", "LiteLLM_MemoryTable", "LiteLLM_AdaptiveRouterState", diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 4187f76c97c..2965174114b 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -2061,3 +2061,73 @@ model LiteLLM_TelemetryReport { @@index([window_end]) @@index([created_at]) } + +enum AccountingEventKind { + inbound_usage + outbound_usage + cost +} + +enum AccountingDeliveryState { + pending + applied + parked +} + +enum AccountingDeliveryTarget { + LiteLLM_VerificationToken + LiteLLM_UserTable + LiteLLM_TeamTable + LiteLLM_TeamMembership + LiteLLM_OrganizationTable + LiteLLM_OrganizationMembership + LiteLLM_EndUserTable + LiteLLM_AgentsTable + LiteLLM_TagTable + LiteLLM_ProjectTable + LiteLLM_ModelAccessGroupBudgetTable + LiteLLM_BudgetTable + LiteLLM_BudgetWindowSpend + LiteLLM_DailyUserSpend + LiteLLM_DailyTeamSpend + LiteLLM_DailyOrganizationSpend + LiteLLM_DailyEndUserSpend + LiteLLM_DailyAgentSpend + LiteLLM_DailyTagSpend + LiteLLM_DailyModelUsage + LiteLLM_DailyToolSpend + LiteLLM_AutoRouterDailySpend +} + +// Durable accounting ledger. One row per usage or cost event; payload holds the +// private event envelope versioned by schema_version. Usage rows retain request-scoped +// pricing context; cost rows include resolved rates and the finalized pricing snapshot. +model LiteLLM_AccountingEvent { + event_id String @id + kind AccountingEventKind + source_event_id String? + schema_version Int + payload_hash String + occurred_at DateTime + recorded_at DateTime @default(now()) + payload Json + + @@index([recorded_at]) + @@index([kind, source_event_id]) +} + +// One row per (event, target table, version): pending work for the projector and, +// once applied in the same transaction as the target write, the receipt that stops a retry. +model LiteLLM_AccountingDelivery { + event_id String + target AccountingDeliveryTarget + version Int + state AccountingDeliveryState @default(pending) + attempts Int @default(0) + next_attempt_at DateTime @default(now()) + applied_at DateTime? + last_error String? + + @@id([event_id, target, version]) + @@index([target, state, next_attempt_at]) +} diff --git a/schema.prisma b/schema.prisma index 4187f76c97c..2965174114b 100644 --- a/schema.prisma +++ b/schema.prisma @@ -2061,3 +2061,73 @@ model LiteLLM_TelemetryReport { @@index([window_end]) @@index([created_at]) } + +enum AccountingEventKind { + inbound_usage + outbound_usage + cost +} + +enum AccountingDeliveryState { + pending + applied + parked +} + +enum AccountingDeliveryTarget { + LiteLLM_VerificationToken + LiteLLM_UserTable + LiteLLM_TeamTable + LiteLLM_TeamMembership + LiteLLM_OrganizationTable + LiteLLM_OrganizationMembership + LiteLLM_EndUserTable + LiteLLM_AgentsTable + LiteLLM_TagTable + LiteLLM_ProjectTable + LiteLLM_ModelAccessGroupBudgetTable + LiteLLM_BudgetTable + LiteLLM_BudgetWindowSpend + LiteLLM_DailyUserSpend + LiteLLM_DailyTeamSpend + LiteLLM_DailyOrganizationSpend + LiteLLM_DailyEndUserSpend + LiteLLM_DailyAgentSpend + LiteLLM_DailyTagSpend + LiteLLM_DailyModelUsage + LiteLLM_DailyToolSpend + LiteLLM_AutoRouterDailySpend +} + +// Durable accounting ledger. One row per usage or cost event; payload holds the +// private event envelope versioned by schema_version. Usage rows retain request-scoped +// pricing context; cost rows include resolved rates and the finalized pricing snapshot. +model LiteLLM_AccountingEvent { + event_id String @id + kind AccountingEventKind + source_event_id String? + schema_version Int + payload_hash String + occurred_at DateTime + recorded_at DateTime @default(now()) + payload Json + + @@index([recorded_at]) + @@index([kind, source_event_id]) +} + +// One row per (event, target table, version): pending work for the projector and, +// once applied in the same transaction as the target write, the receipt that stops a retry. +model LiteLLM_AccountingDelivery { + event_id String + target AccountingDeliveryTarget + version Int + state AccountingDeliveryState @default(pending) + attempts Int @default(0) + next_attempt_at DateTime @default(now()) + applied_at DateTime? + last_error String? + + @@id([event_id, target, version]) + @@index([target, state, next_attempt_at]) +} diff --git a/tests/integration/database/test_accounting_ledger.py b/tests/integration/database/test_accounting_ledger.py new file mode 100644 index 00000000000..089ed872dee --- /dev/null +++ b/tests/integration/database/test_accounting_ledger.py @@ -0,0 +1,128 @@ +import os +import subprocess +import sys +from collections.abc import Iterator +from pathlib import Path +from typing import Final + +import psycopg +import pytest +from psycopg import sql + +from litellm.accounting import DeliveryState, DeliveryTarget, EventKind +from tests.integration._support.database import scratch_database + +_SCHEMA: Final = Path(__file__).resolve().parents[3] / "litellm-proxy-extras/litellm_proxy_extras/schema.prisma" +_EVENT_INSERT: Final = ( + 'INSERT INTO "LiteLLM_AccountingEvent" ' + "(event_id, kind, schema_version, payload_hash, occurred_at, payload) " + "VALUES ('event', %s, %s, 'hash', '2026-01-01', '{}') RETURNING kind::text, schema_version" +) +_DELIVERY_INSERT: Final = ( + 'INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version, state, attempts) ' + "VALUES ('event', %s, %s, %s, %s) RETURNING target::text, version, state::text, attempts" +) + + +@pytest.fixture(scope="module") +def ledger_database() -> Iterator[str]: + with scratch_database("accounting") as database_url: + result: Final = subprocess.run( + [sys.executable, "-I", "-m", "prisma", "migrate", "deploy", "--schema", str(_SCHEMA)], + env={**os.environ, "DATABASE_URL": database_url}, + capture_output=True, + text=True, + timeout=300, + ) + assert result.returncode == 0, result.stdout + result.stderr + yield database_url + + +@pytest.fixture +def ledger(ledger_database: str) -> Iterator[psycopg.Connection[tuple[object, ...]]]: + with psycopg.connect(ledger_database) as connection: + try: + yield connection + finally: + connection.rollback() + + +@pytest.mark.parametrize("kind", tuple(EventKind)) +def test_ledger_accepts_every_public_event_kind( + ledger: psycopg.Connection[tuple[object, ...]], kind: EventKind +) -> None: + row: Final = ledger.execute(_EVENT_INSERT, (kind.value, 1)).fetchone() + assert row == (kind.value, 1) + + +@pytest.mark.parametrize("target", tuple(DeliveryTarget)) +@pytest.mark.parametrize("state", tuple(DeliveryState)) +def test_delivery_accepts_every_public_target_and_state( + ledger: psycopg.Connection[tuple[object, ...]], target: DeliveryTarget, state: DeliveryState +) -> None: + row: Final = ledger.execute(_DELIVERY_INSERT, (target.value, 1, state.value, 0)).fetchone() + assert row == (target.value, 1, state.value, 0) + + +@pytest.mark.parametrize( + ("enum_name", "values"), + [ + ("AccountingEventKind", tuple(value.value for value in EventKind)), + ("AccountingDeliveryState", tuple(value.value for value in DeliveryState)), + ("AccountingDeliveryTarget", tuple(value.value for value in DeliveryTarget)), + ], +) +def test_migrated_enums_match_the_public_contract( + ledger: psycopg.Connection[tuple[object, ...]], enum_name: str, values: tuple[str, ...] +) -> None: + rows: Final = ledger.execute( + "SELECT e.enumlabel FROM pg_enum e JOIN pg_type t ON t.oid = e.enumtypid WHERE t.typname = %s", + (enum_name,), + ).fetchall() + assert frozenset(row[0] for row in rows) == frozenset(values) + + +@pytest.mark.parametrize("version", [0, -1]) +def test_event_rejects_nonpositive_schema_versions( + ledger: psycopg.Connection[tuple[object, ...]], version: int +) -> None: + with pytest.raises(psycopg.errors.CheckViolation, match="AccountingEvent_schema_version_check"): + ledger.execute(_EVENT_INSERT, (EventKind.OUTBOUND_USAGE.value, version)) + + +@pytest.mark.parametrize( + ("version", "attempts", "constraint"), + [ + (0, 0, "AccountingDelivery_version_check"), + (-1, 0, "AccountingDelivery_version_check"), + (1, -1, "AccountingDelivery_attempts_check"), + ], +) +def test_delivery_rejects_invalid_versions_and_attempts( + ledger: psycopg.Connection[tuple[object, ...]], version: int, attempts: int, constraint: str +) -> None: + with pytest.raises(psycopg.errors.CheckViolation, match=constraint): + ledger.execute(_DELIVERY_INSERT, (DeliveryTarget.KEY_SPEND.value, version, "pending", attempts)) + + +@pytest.mark.parametrize("column", ["state", "target"]) +def test_delivery_rejects_unknown_enum_values(ledger: psycopg.Connection[tuple[object, ...]], column: str) -> None: + ledger.execute(_DELIVERY_INSERT, (DeliveryTarget.KEY_SPEND.value, 1, "pending", 0)) + with pytest.raises(psycopg.errors.InvalidTextRepresentation, match="invalid input value for enum"): + ledger.execute( + sql.SQL('UPDATE "LiteLLM_AccountingDelivery" SET {} = %s').format(sql.Identifier(column)), ("invalid",) + ) + + +def test_event_rejects_unknown_kind(ledger: psycopg.Connection[tuple[object, ...]]) -> None: + with pytest.raises(psycopg.errors.InvalidTextRepresentation, match="invalid input value for enum"): + ledger.execute(_EVENT_INSERT, ("invalid", 1)) + + +def test_new_delivery_defaults_to_pending_with_no_attempts(ledger: psycopg.Connection[tuple[object, ...]]) -> None: + row: Final = ledger.execute( + 'INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version) ' + "VALUES ('event', %s, 1) RETURNING state::text, attempts", + (DeliveryTarget.KEY_SPEND.value,), + ).fetchone() + assert row == ("pending", 0) diff --git a/tests/unit/accounting/__init__.py b/tests/unit/accounting/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/accounting/test_measurements.py b/tests/unit/accounting/test_measurements.py new file mode 100644 index 00000000000..b65c3076c5d --- /dev/null +++ b/tests/unit/accounting/test_measurements.py @@ -0,0 +1,235 @@ +from datetime import datetime, timezone +from decimal import Decimal +from typing import Final + +import pytest + +from litellm.accounting import ( + CostEvent, + Dimension, + DimensionName, + InboundUsageEvent, + Measurement, + Metric, + MetricName, + Outcome, + OutboundUsageEvent, + PricingSnapshot, + measurements_for, +) +from litellm.types.utils import CallTypes + +_START: Final = datetime(2026, 1, 1, 23, 59, 59, tzinfo=timezone.utc) +_END: Final = datetime(2026, 1, 2, tzinfo=timezone.utc) +_DIMENSIONS: Final = ( + Dimension(name=DimensionName.TEAM, value="team"), + Dimension(name=DimensionName.TAG, value="alpha"), + Dimension(name=DimensionName.TAG, value="beta"), +) + + +def _snapshot(amount: Decimal) -> PricingSnapshot: + return PricingSnapshot( + calculator_version="test", + model="test-model", + provider=None, + pricing_source="upstream_header", + priced_model_key="test-model", + upstream_reported_cost=amount, + total_cost=amount, + ) + + +@pytest.mark.parametrize( + ("outcome", "outcome_metric"), + [ + (Outcome.SUCCEEDED, MetricName.SUCCESSFUL_REQUESTS), + (Outcome.FAILED, MetricName.FAILED_REQUESTS), + (Outcome.CANCELLED, MetricName.FAILED_REQUESTS), + ], +) +def test_caller_measurements_keep_both_request_and_outcome_counts(outcome, outcome_metric): + event = InboundUsageEvent( + event_id="caller-event", + request_id="caller", + call_type=CallTypes.acompletion, + started_at=_START, + ended_at=_END, + outcome=outcome, + dimensions=_DIMENSIONS, + metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)), Metric(name=outcome_metric, value=Decimal(1))), + ) + assert measurements_for(event) == ( + Measurement( + source_event_id="caller-event", + occurred_at=_START, + dimensions=_DIMENSIONS, + metric=Metric(name=MetricName.REQUESTS, value=Decimal(1)), + ), + Measurement( + source_event_id="caller-event", + occurred_at=_START, + dimensions=_DIMENSIONS, + metric=Metric(name=outcome_metric, value=Decimal(1)), + ), + ) + + +def test_usage_measurements_keep_every_metric_in_the_start_time_bucket(): + event = OutboundUsageEvent( + event_id="usage-event", + attempt_id="attempt", + request_id="caller", + retry_of_attempt_id=None, + call_type=CallTypes.acompletion, + started_at=_START, + ended_at=_END, + outcome=Outcome.SUCCEEDED, + dimensions=_DIMENSIONS, + metrics=( + Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(10)), + Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(2)), + Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(12)), + Metric(name=MetricName.CACHE_READ_INPUT_TOKENS, value=Decimal(3)), + Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS, value=Decimal(1)), + ), + ) + assert measurements_for(event) == ( + Measurement( + source_event_id="usage-event", + occurred_at=_START, + dimensions=_DIMENSIONS, + metric=Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(10)), + ), + Measurement( + source_event_id="usage-event", + occurred_at=_START, + dimensions=_DIMENSIONS, + metric=Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(2)), + ), + Measurement( + source_event_id="usage-event", + occurred_at=_START, + dimensions=_DIMENSIONS, + metric=Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(12)), + ), + Measurement( + source_event_id="usage-event", + occurred_at=_START, + dimensions=_DIMENSIONS, + metric=Metric(name=MetricName.CACHE_READ_INPUT_TOKENS, value=Decimal(3)), + ), + Measurement( + source_event_id="usage-event", + occurred_at=_START, + dimensions=_DIMENSIONS, + metric=Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS, value=Decimal(1)), + ), + ) + + +@pytest.mark.parametrize("metrics", [None, ()]) +def test_unknown_or_empty_usage_does_not_fabricate_zero_measurements(metrics): + event = OutboundUsageEvent( + event_id="usage-event", + attempt_id="attempt", + request_id=None, + retry_of_attempt_id=None, + call_type=CallTypes.acompletion, + started_at=_START, + ended_at=_END, + outcome=Outcome.SUCCEEDED, + dimensions=(), + metrics=metrics, + ) + assert measurements_for(event) == () + + +@pytest.mark.parametrize("amount", [Decimal(0), Decimal("0.1234567890123456789")]) +def test_cost_measurements_preserve_authoritative_amount_identity_and_occurrence(amount): + event = CostEvent( + event_id="cost-event", + usage_event_id="usage-event", + occurred_at=_END, + dimensions=_DIMENSIONS, + charges=(Metric(name=MetricName.SPEND, value=amount, unit="USD"),), + snapshot=_snapshot(amount), + ) + assert measurements_for(event) == ( + Measurement( + source_event_id="cost-event", + occurred_at=_END, + dimensions=_DIMENSIONS, + metric=Metric(name=MetricName.SPEND, value=amount, unit="USD"), + ), + ) + + +def test_billed_retry_and_success_contribute_one_caller_request_and_both_charges(): + caller = InboundUsageEvent( + event_id="caller-event", + request_id="caller", + call_type=CallTypes.acompletion, + started_at=_START, + ended_at=_END, + outcome=Outcome.SUCCEEDED, + dimensions=_DIMENSIONS, + metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),), + ) + retry = OutboundUsageEvent( + event_id="retry-usage", + attempt_id="retry", + request_id="caller", + retry_of_attempt_id=None, + call_type=CallTypes.acompletion, + started_at=_START, + ended_at=_END, + outcome=Outcome.FAILED, + dimensions=_DIMENSIONS, + metrics=(Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(4)),), + ) + success = OutboundUsageEvent( + event_id="success-usage", + attempt_id="success", + request_id="caller", + retry_of_attempt_id="retry", + call_type=CallTypes.acompletion, + started_at=_START, + ended_at=_END, + outcome=Outcome.SUCCEEDED, + dimensions=_DIMENSIONS, + metrics=(Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(12)),), + ) + retry_cost = CostEvent( + event_id="retry-cost", + usage_event_id="retry-usage", + occurred_at=_START, + dimensions=_DIMENSIONS, + charges=(Metric(name=MetricName.SPEND, value=Decimal("0.025"), unit="USD"),), + snapshot=_snapshot(Decimal("0.025")), + ) + success_cost = CostEvent( + event_id="success-cost", + usage_event_id="success-usage", + occurred_at=_START, + dimensions=_DIMENSIONS, + charges=(Metric(name=MetricName.SPEND, value=Decimal("0.125"), unit="USD"),), + snapshot=_snapshot(Decimal("0.125")), + ) + measurements = ( + *measurements_for(caller), + *measurements_for(retry), + *measurements_for(success), + *measurements_for(retry_cost), + *measurements_for(success_cost), + ) + assert sum(m.metric.value for m in measurements if m.metric.name is MetricName.REQUESTS) == Decimal(1) + assert sum(m.metric.value for m in measurements if m.metric.name is MetricName.TOTAL_TOKENS) == Decimal(16) + assert sum(m.metric.value for m in measurements if m.metric.name is MetricName.SPEND) == Decimal("0.150") + assert {m.source_event_id for m in measurements} == { + "caller-event", + "retry-usage", + "success-usage", + "retry-cost", + "success-cost", + } diff --git a/tests/unit/accounting/test_types.py b/tests/unit/accounting/test_types.py new file mode 100644 index 00000000000..f645ee6c2cc --- /dev/null +++ b/tests/unit/accounting/test_types.py @@ -0,0 +1,484 @@ +import re +from collections.abc import Mapping +from datetime import datetime, timedelta, timezone +from decimal import Decimal +from enum import Enum +from pathlib import Path +from typing import Final + +import pytest +from pydantic import ValidationError + +from litellm.accounting import ( + AdmissionDecision, + AggregateQuery, + AggregateRow, + CostEvent, + DeliveryState, + DeliveryTarget, + Dimension, + DimensionFilter, + DimensionName, + EventKind, + Granularity, + InboundUsageEvent, + Metric, + MetricName, + Multiplier, + Outcome, + OutboundUsageEvent, + PricedQuantity, + PricingSnapshot, + Rate, + validate_aggregate_rows, +) +from litellm.accounting.types import PricingUnit +from litellm.types.utils import CallTypes + +_START: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) +_END: Final = datetime(2026, 1, 2, tzinfo=timezone.utc) +_REQUEST: Final = Metric(name=MetricName.REQUESTS, value=Decimal(1)) +_TOKEN: Final = Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(10)) +_SPEND: Final = Metric(name=MetricName.SPEND, value=Decimal("0.1234567890123456789"), unit="USD") +_TEAM: Final = Dimension(name=DimensionName.TEAM, value="team") + + +def _inbound(metrics=(_REQUEST,), outcome=Outcome.SUCCEEDED): + return InboundUsageEvent( + event_id="inbound:request", + request_id="request", + call_type=CallTypes.acompletion, + started_at=_START, + ended_at=_START, + outcome=outcome, + dimensions=(_TEAM,), + metrics=metrics, + ) + + +def _outbound(metrics=(_TOKEN,), **changes): + return OutboundUsageEvent.model_validate( + { + "event_id": "outbound:attempt", + "attempt_id": "attempt", + "request_id": "request", + "retry_of_attempt_id": None, + "call_type": CallTypes.acompletion, + "started_at": _START, + "ended_at": _END, + "outcome": Outcome.SUCCEEDED, + "dimensions": (_TEAM,), + "metrics": metrics, + **changes, + } + ) + + +def _query(**changes): + return AggregateQuery.model_validate( + { + "metric": MetricName.SPEND, + "filters": (DimensionFilter(name=DimensionName.TEAM, values=("one", "two")),), + "group_by": (DimensionName.MODEL,), + "start": _START, + "end": _END, + "granularity": Granularity.DAY, + **changes, + } + ) + + +def test_wire_values_preserve_decimal_precision_and_normalize_aware_times(): + offset = timezone(timedelta(hours=-8)) + event = CostEvent( + event_id="cost:attempt", + usage_event_id="outbound:attempt", + occurred_at=datetime(2025, 12, 31, 16, tzinfo=offset), + dimensions=(_TEAM,), + charges=(_SPEND,), + snapshot=_snapshot(), + ) + wire = event.model_dump(mode="json") + assert wire["occurred_at"] == "2026-01-01T00:00:00Z" + assert wire["charges"] == [{"name": "spend", "value": "0.1234567890123456789", "unit": "USD"}] + assert event.occurred_at == _START + assert CostEvent.model_validate_json(event.model_dump_json()).charges[0].value == _SPEND.value + + +@pytest.mark.parametrize( + ("name", "value", "unit", "message"), + [ + (MetricName.SPEND, Decimal(1), "count", "SPEND requires USD"), + (MetricName.REQUESTS, Decimal(1), "USD", "SPEND requires USD"), + (MetricName.PROMPT_TOKENS, Decimal("1.5"), "count", "integral"), + (MetricName.SPEND, Decimal("NaN"), "USD", "finite"), + (MetricName.SPEND, Decimal("Infinity"), "USD", "finite"), + (MetricName.SPEND, Decimal(-1), "USD", "greater than or equal"), + (MetricName.SPEND, 0.1, "USD", "instance of Decimal"), + ], +) +def test_metrics_reject_wrong_units_nonfinite_values_and_binary_float_inputs(name, value, unit, message): + with pytest.raises(ValidationError, match=message): + Metric(name=name, value=value, unit=unit) + + +@pytest.mark.parametrize("metrics", [(_TOKEN,), (_SPEND,), ()]) +def test_inbound_activity_cannot_own_usage_or_cost_or_omit_request_count(metrics): + with pytest.raises(ValidationError, match="Inbound events"): + _inbound(metrics) + + +@pytest.mark.parametrize( + ("outcome", "metric"), + [ + (Outcome.FAILED, Metric(name=MetricName.SUCCESSFUL_REQUESTS, value=Decimal(1))), + (Outcome.SUCCEEDED, Metric(name=MetricName.FAILED_REQUESTS, value=Decimal(1))), + ], +) +def test_inbound_outcome_counts_cannot_contradict_outcome(outcome, metric): + with pytest.raises(ValidationError, match="count must match outcome"): + _inbound((_REQUEST, metric), outcome) + + +@pytest.mark.parametrize("metrics", [(_REQUEST,), (_SPEND,)]) +def test_outbound_attempts_cannot_own_caller_counts_or_charges(metrics): + with pytest.raises(ValidationError, match="Outbound events own only"): + _outbound(metrics) + + +@pytest.mark.parametrize( + ("changes", "message"), + [ + ({"started_at": datetime(2026, 1, 1)}, "timezone"), + ({"ended_at": _START - timedelta(seconds=1)}, "End must follow start"), + ({"retry_of_attempt_id": "attempt"}, "cannot retry itself"), + ({"dimensions": (_TEAM, _TEAM)}, "Duplicate dimension values"), + ({"dimensions": (_TEAM, Dimension(name=DimensionName.TEAM, value="other"))}, "Only tag and model access group"), + ( + { + "metrics": ( + Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(5)), + Metric(name=MetricName.REASONING_TOKENS, value=Decimal(6)), + ) + }, + "reasoning_tokens cannot exceed completion_tokens", + ), + ( + { + "metrics": ( + Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS, value=Decimal(5)), + Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS_ABOVE_1HR, value=Decimal(6)), + ) + }, + "cache_creation_input_tokens_above_1hr cannot exceed cache_creation_input_tokens", + ), + ({"event_id": " "}, "pattern"), + ({"metrics": (_TOKEN, _TOKEN)}, "Duplicate metric names"), + ], +) +def test_outbound_events_reject_invalid_identity_time_and_duplicate_attribution(changes, message): + with pytest.raises(ValidationError, match=message): + _outbound(**changes) + + +def test_unknown_and_known_zero_usage_are_distinct_and_values_are_immutable(): + unknown = _outbound(None) + zero = _outbound((Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(0)),)) + assert unknown.metrics is None + assert zero.metrics is not None and zero.metrics[0].value == 0 + assert _inbound().started_at == _inbound().ended_at + with pytest.raises(ValidationError, match="frozen"): + zero.attempt_id = "different" + with pytest.raises(ValidationError, match="frozen"): + zero.metrics[0].value = Decimal(10) + + +@pytest.mark.parametrize("charges", [(), (_TOKEN,), (_SPEND, _SPEND)]) +def test_cost_events_require_one_authoritative_spend(charges): + with pytest.raises(ValidationError, match=r"authoritative USD spend|Duplicate metric"): + CostEvent( + event_id="cost", + usage_event_id="usage", + occurred_at=_START, + dimensions=(), + charges=charges, + snapshot=_snapshot(), + ) + + +def test_cost_event_cannot_reuse_its_usage_identity(): + with pytest.raises(ValidationError, match="distinct identities"): + CostEvent( + event_id="usage", + usage_event_id="usage", + occurred_at=_START, + dimensions=(), + charges=(_SPEND,), + snapshot=_snapshot(), + ) + + +@pytest.mark.parametrize("granularity", [Granularity.DAY, None]) +def test_daily_and_whole_interval_queries_preserve_filter_and_group_semantics(granularity): + query = _query(granularity=granularity) + assert query.granularity is granularity + assert query.filters[0].values == ("one", "two") + assert query.group_by == (DimensionName.MODEL,) + assert query.start == _START and query.end == _END + + +@pytest.mark.parametrize( + ("changes", "message"), + [ + ({"end": _START}, "End must follow start"), + ({"start": _START + timedelta(hours=1)}, "UTC day boundaries"), + ({"end": _END + timedelta(microseconds=1)}, "UTC day boundaries"), + ({"granularity": "hour"}, "day"), + ({"group_by": (DimensionName.TEAM, DimensionName.TEAM)}, "Duplicate group-by"), + ({"filters": (DimensionFilter(name=DimensionName.TEAM, values=("one",)),) * 2}, "one IN filter"), + ({"filters": ({"name": "team", "values": ()},)}, "at least 1"), + ], +) +def test_queries_reject_unsupported_or_ambiguous_ranges_and_filters(changes, message): + with pytest.raises(ValidationError, match=message): + _query(**changes) + + +def test_rows_allow_adjacent_buckets_and_overlapping_distinct_groups(): + first = AggregateRow(start=_START, end=_END, dimensions=(_TEAM,), value=Decimal(1)) + second = AggregateRow(start=_END, end=_END + timedelta(days=1), dimensions=(_TEAM,), value=Decimal(2)) + other = AggregateRow(start=_START, end=_END, dimensions=(), value=Decimal(3)) + rows = (second, other, first) + assert validate_aggregate_rows(rows) is rows + assert validate_aggregate_rows(()) == () + + +def test_rows_reject_overlap_within_a_group_regardless_of_dimension_order(): + model = Dimension(name=DimensionName.MODEL, value="model") + first = AggregateRow(start=_START, end=_END, dimensions=(_TEAM, model), value=Decimal(1)) + overlapping = AggregateRow(start=_START, end=_END, dimensions=(model, _TEAM), value=Decimal(2)) + with pytest.raises(ValueError, match="overlap within a dimension group"): + validate_aggregate_rows((first, overlapping)) + + +def test_rows_reject_empty_intervals(): + with pytest.raises(ValidationError, match="End must follow start"): + AggregateRow(start=_START, end=_START, dimensions=(), value=Decimal(0)) + + +def test_admission_decision_preserves_a_denial_reason(): + decision = AdmissionDecision(allowed=False, reason="Team budget exhausted") + assert not decision.allowed + assert decision.reason == "Team budget exhausted" + + +def _snapshot(**changes: object) -> PricingSnapshot: + return PricingSnapshot.model_validate( + { + "calculator_version": "1.80.0", + "model": "gpt-4o", + "provider": "openai", + "pricing_source": "cost_map", + "priced_model_key": "openai/gpt-4o", + "multipliers": (Multiplier(name="data_residency", value="eu", factor=Decimal("1.1")),), + "rates": (Rate(measure="input", unit="token", usd_per_unit=Decimal("0.0000025")),), + "usage": (PricedQuantity(measure="input", unit="token", value=_TOKEN.value),), + "input_cost": Decimal("0.000025"), + "total_cost": _SPEND.value, + **changes, + } + ) + + +def test_cost_event_carries_a_snapshot_whose_total_matches_the_spend(): + event = CostEvent( + event_id="cost:attempt", + usage_event_id="outbound:attempt", + occurred_at=_START, + dimensions=(_TEAM,), + charges=(_SPEND,), + snapshot=_snapshot(), + ) + assert event.snapshot is not None + assert event.snapshot.total_cost == _SPEND.value + assert event.snapshot.multipliers[0].factor == Decimal("1.1") + with pytest.raises(ValidationError, match="Snapshot total must equal the authoritative spend"): + CostEvent( + event_id="cost:attempt", + usage_event_id="outbound:attempt", + occurred_at=_START, + dimensions=(_TEAM,), + charges=(_SPEND,), + snapshot=_snapshot(total_cost=Decimal("0.5")), + ) + + +@pytest.mark.parametrize( + ("changes", "message"), + ( + ( + { + "rates": ( + Rate(measure="input", unit="token", usd_per_unit=Decimal(1)), + Rate(measure="input", unit="token", usd_per_unit=Decimal(2)), + ) + }, + "Duplicate rate measures and units", + ), + ({"usage": (_SPEND,)}, "PricedQuantity"), + ({"usage": (_REQUEST,)}, "PricedQuantity"), + ({"pricing_source": "guess"}, "pricing_source"), + ({"multipliers": ({"name": "x", "value": "y", "factor": Decimal(0)},)}, "factor"), + ({"schema_version": 2}, "schema_version"), + ( + {"additional_costs": ({"name": "fee", "amount": Decimal(1)}, {"name": "fee", "amount": Decimal(2)})}, + "Duplicate additional cost names", + ), + ), +) +def test_snapshots_reject_ambiguous_rates_wrong_usage_and_unknown_sources(changes: Mapping[str, object], message: str): + with pytest.raises(ValidationError, match=message): + _snapshot(**changes) + + +@pytest.mark.parametrize( + ("unit", "value"), + [("token", "10"), ("second", "1.25"), ("character", "25"), ("image", "2"), ("request", "1")], +) +def test_snapshot_preserves_matching_rates_and_quantities_for_each_billing_unit(unit: PricingUnit, value: str) -> None: + quantity: Final = PricedQuantity(measure="input", unit=unit, value=Decimal(value)) + rate: Final = Rate(measure="input", unit=unit, usd_per_unit=Decimal("0.003")) + amount: Final = quantity.value * rate.usd_per_unit + snapshot: Final = _snapshot(rates=(rate,), usage=(quantity,), total_cost=amount) + restored: Final = PricingSnapshot.model_validate_json(snapshot.model_dump_json()) + assert restored == snapshot + assert restored.usage[0].value * restored.rates[0].usd_per_unit == amount + + +@pytest.mark.parametrize( + ("usage", "message"), + [ + ((), "Every rate requires"), + (({"measure": "output", "unit": "token", "value": Decimal(10)},), "Every rate requires"), + (({"measure": "input", "unit": "second", "value": Decimal(10)},), "Every rate requires"), + ( + ({"measure": "input", "unit": "token", "value": Decimal(10)},) * 2, + "Duplicate priced quantity", + ), + ], +) +def test_snapshot_rejects_missing_mismatched_or_duplicate_priced_quantities( + usage: tuple[Mapping[str, object], ...], message: str +) -> None: + with pytest.raises(ValidationError, match=message): + _snapshot(usage=usage) + + +@pytest.mark.parametrize("unit", ["token", "character", "image", "request"]) +def test_priced_counts_remain_integral(unit: PricingUnit) -> None: + with pytest.raises(ValidationError, match="integral"): + PricedQuantity(measure="input", unit=unit, value=Decimal("1.25")) + + +@pytest.mark.parametrize( + ("value", "message"), + [ + (Decimal(-1), "greater than or equal"), + (Decimal("NaN"), "finite"), + (Decimal("Infinity"), "finite"), + (1.25, "instance of Decimal"), + ], +) +def test_pricing_values_reject_invalid_amounts(value: Decimal | float, message: str) -> None: + with pytest.raises(ValidationError, match=message): + PricedQuantity.model_validate({"measure": "input", "unit": "second", "value": value}) + with pytest.raises(ValidationError, match=message): + Rate.model_validate({"measure": "input", "unit": "second", "usd_per_unit": value}) + + +@pytest.mark.parametrize("snapshot_fields", [{}, {"snapshot": None}]) +def test_cost_event_cannot_omit_its_derivation(snapshot_fields: Mapping[str, object]) -> None: + with pytest.raises(ValidationError, match="snapshot"): + CostEvent.model_validate( + { + "event_id": "cost", + "usage_event_id": "usage", + "occurred_at": _START, + "dimensions": (), + "charges": (_SPEND,), + **snapshot_fields, + } + ) + + +def test_upstream_charge_snapshot_preserves_reported_cost_without_inventing_rates() -> None: + snapshot: Final = _snapshot( + pricing_source="upstream_header", rates=(), usage=(), upstream_reported_cost=_SPEND.value + ) + restored: Final = PricingSnapshot.model_validate_json(snapshot.model_dump_json()) + assert restored.upstream_reported_cost == restored.total_cost == _SPEND.value + assert restored.rates == restored.usage == () + + +def test_subset_token_counts_round_trip_without_changing_totals() -> None: + metrics: Final = ( + Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(10)), + Metric(name=MetricName.REASONING_TOKENS, value=Decimal(3)), + Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS, value=Decimal(8)), + Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS_ABOVE_1HR, value=Decimal(2)), + ) + event: Final = _outbound(metrics=metrics) + assert OutboundUsageEvent.model_validate_json(event.model_dump_json()).metrics == metrics + + +@pytest.mark.parametrize( + "schema_path", + [ + "schema.prisma", + "litellm/proxy/schema.prisma", + "litellm-proxy-extras/litellm_proxy_extras/schema.prisma", + ], +) +@pytest.mark.parametrize( + ("enum_name", "python_enum"), + [ + ("AccountingEventKind", EventKind), + ("AccountingDeliveryState", DeliveryState), + ("AccountingDeliveryTarget", DeliveryTarget), + ], +) +def test_prisma_accepts_exactly_the_public_accounting_enum_values( + schema_path: str, enum_name: str, python_enum: type[Enum] +) -> None: + schema: Final = (Path(__file__).resolve().parents[3] / schema_path).read_text() + declaration: Final = re.search(r"enum " + enum_name + r"\s*\{([^}]+)\}", schema) + assert declaration is not None, f"Missing persisted enum {enum_name}" + assert frozenset(declaration[1].split()) == frozenset(member.value for member in python_enum) + + +def test_one_event_carries_every_matched_model_access_group(): + groups = ( + Dimension(name=DimensionName.MODEL_ACCESS_GROUP, value="gpt-group"), + Dimension(name=DimensionName.MODEL_ACCESS_GROUP, value="claude-group"), + ) + assert _outbound(dimensions=(_TEAM, *groups)).dimensions == (_TEAM, *groups) + with pytest.raises(ValidationError, match="Duplicate dimension values"): + _outbound(dimensions=(groups[0], groups[0])) + + +def test_snapshots_reject_duplicate_multipliers_and_contradictory_token_counts(): + multiplier = Multiplier(name="data_residency", value="eu", factor=Decimal("1.1")) + with pytest.raises(ValidationError, match="Duplicate multiplier names"): + _snapshot(multipliers=(multiplier, multiplier)) + with pytest.raises(ValidationError, match="reasoning_tokens cannot exceed completion_tokens"): + _snapshot( + usage=( + PricedQuantity(measure="input", unit="token", value=_TOKEN.value), + PricedQuantity(measure="completion_tokens", unit="token", value=Decimal(5)), + PricedQuantity(measure="reasoning_tokens", unit="token", value=Decimal(6)), + ) + ) + + +def test_agent_live_spend_is_a_delivery_target(): + assert DeliveryTarget.AGENT_SPEND.value == "LiteLLM_AgentsTable" From 66ddbaab8624703e0868a6366e994fce73b5cc90 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 28 Oct 2026 14:48:40 -0700 Subject: [PATCH 091/128] fix(ci): spend logs guard requires the approval to cover the head's spend-log lines (#48904) * fix(ci): spend logs guard requires the approval to cover the head's spend-log lines An approving review survives later pushes (dismiss_stale_reviews_on_push is off on main), so a PR approved once could add more LiteLLM_SpendLogs schema or migration changes and the guard kept passing. The guard now reads the review's commit_id and passes only when the approval is on the head, or the head adds no spend-log line the approved commit did not already carry. * fix(ci): keep approved_by untouched, add latest_approval beside it for the stale-approval check --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ci_cd/spend_logs_guard.py | 47 +++++++++++++++++++++++--- tests/unit/test_spend_logs_guard.py | 52 ++++++++++++++++++++++++++++- 2 files changed, 93 insertions(+), 6 deletions(-) diff --git a/ci_cd/spend_logs_guard.py b/ci_cd/spend_logs_guard.py index a04999c47c0..64e33c4fb70 100644 --- a/ci_cd/spend_logs_guard.py +++ b/ci_cd/spend_logs_guard.py @@ -4,7 +4,9 @@ needs an approving review from the spend-log owner before it can merge. CODEOWNERS can only match file paths, and every schema change edits the same schema.prisma copies, so a path rule would put every migration on the owner. This gate reads the diff instead: when no added line in the watched files mentions the table it passes at once, otherwise it passes only once -the owner's latest review on the pull request is an approval. The workflow re-runs on every review. +the owner's latest review on the pull request is an approval that still covers every spend-log line +in the head (approvals survive pushes, so a later commit may not add spend-log lines the owner never +saw). The workflow re-runs on every review. """ from __future__ import annotations @@ -13,7 +15,7 @@ import argparse import re import subprocess import sys -from collections.abc import Sequence +from collections.abc import Callable, Sequence from pathlib import Path from typing import Final @@ -37,6 +39,7 @@ class Review(BaseModel): model_config = ConfigDict(frozen=True, extra="ignore") state: str user: Reviewer + commit_id: str | None = None _REVIEWS: Final = TypeAdapter(list[Review]) @@ -75,6 +78,24 @@ def approved_by(reviews: Sequence[Review], login: str) -> bool: return bool(states) and states[-1] == "APPROVED" +def latest_approval(reviews: Sequence[Review], login: str) -> Review | None: + """`login`'s most recent APPROVED review on the pull request, the one GitHub keeps in force after pushes.""" + own: Final = tuple(review for review in reviews if review.user.login == login and review.state == "APPROVED") + return own[-1] if own else None + + +def approval_covers( + approval: Review, head: str, hits: Sequence[str], hits_at: Callable[[str], Sequence[str] | None] +) -> bool: + """True when the approval is on the head, or the head adds no spend-log line the approved commit did not.""" + if approval.commit_id is None: + return False + if approval.commit_id == head: + return True + reviewed: Final = hits_at(approval.commit_id) + return reviewed is not None and set(hits) <= set(reviewed) + + def added_diff(base: str, head: str) -> str: result: Final = subprocess.run( ("git", "diff", "--unified=0", "--no-color", f"{base}..{head}", "--", *WATCHED, MIGRATIONS_PREFIX), @@ -85,6 +106,15 @@ def added_diff(base: str, head: str) -> str: return result.stdout +def fetched(commit: str) -> bool: + result: Final = subprocess.run(("git", "fetch", "--no-tags", "--depth=1", "origin", commit), check=False) + return result.returncode == 0 + + +def reviewed_hits(base: str, commit: str) -> tuple[str, ...] | None: + return spend_log_hits(added_diff(base, commit)) if fetched(commit) else None + + class Arguments(BaseModel): model_config = ConfigDict(frozen=True) base: str @@ -104,14 +134,21 @@ def main(argv: Sequence[str]) -> int: if not hits: sys.stdout.write("No added line in schema.prisma or a migration names LiteLLM_SpendLogs; nothing to review.\n") return 0 - if approved_by(_REVIEWS.validate_json(args.reviews.read_bytes()), args.approver): + reviews: Final = _REVIEWS.validate_json(args.reviews.read_bytes()) + approval: Final = latest_approval(reviews, args.approver) if approved_by(reviews, args.approver) else None + if approval is not None and approval_covers( + approval, args.head, hits, lambda commit: reviewed_hits(args.base, commit) + ): sys.stdout.write(f"LiteLLM_SpendLogs change approved by {args.approver}: {', '.join(hits)}\n") return 0 listed: Final = "".join(f" {hit}\n" for hit in hits) - sys.stdout.write( + reason: Final = ( f"This pull request changes LiteLLM_SpendLogs and needs an approving review from @{args.approver}:\n" - f"{listed}The check re-runs when a review is submitted.\n" + if approval is None + else f"LiteLLM_SpendLogs lines changed after @{args.approver} approved {approval.commit_id}, " + f"a new approving review is needed:\n" ) + sys.stdout.write(f"{reason}{listed}The check re-runs when a review is submitted.\n") return 1 diff --git a/tests/unit/test_spend_logs_guard.py b/tests/unit/test_spend_logs_guard.py index 9a033045dd0..6a3355e1c52 100644 --- a/tests/unit/test_spend_logs_guard.py +++ b/tests/unit/test_spend_logs_guard.py @@ -42,6 +42,10 @@ class TestSpendLogHits: assert guard.spend_log_hits(diff) == ("litellm/proxy/schema.prisma: a LiteLLM_SpendLogs",) +def _review(login: str, state: str, commit_id: str | None = "head") -> "guard.Review": + return guard.Review(state=state, user=guard.Reviewer(login=login), commit_id=commit_id) + + class TestApprovedBy: def test_the_latest_review_of_the_approver_decides(self): reviews: Final = tuple( @@ -61,6 +65,42 @@ class TestApprovedBy: assert guard.approved_by((), "kerry-berri") is False +class TestLatestApproval: + def test_the_approvers_last_approved_review_is_returned_with_its_commit(self): + reviews: Final = ( + _review("kerry-berri", "APPROVED", "old"), + _review("someone-else", "APPROVED", "head"), + _review("kerry-berri", "APPROVED", "new"), + _review("kerry-berri", "COMMENTED", "head"), + ) + assert guard.latest_approval(reviews, "kerry-berri") is reviews[2] + assert guard.latest_approval(reviews[:2], "kerry-berri") is reviews[0] + assert guard.latest_approval(reviews[1:2], "kerry-berri") is None + assert guard.latest_approval((), "kerry-berri") is None + + +class TestApprovalCovers: + _HITS: Final = ("schema.prisma: a LiteLLM_SpendLogs", "schema.prisma: b LiteLLM_SpendLogs") + + def test_an_approval_on_the_head_covers_it_without_reading_the_approved_commit(self): + def unreadable(commit: str) -> None: + raise AssertionError(f"read {commit}") + + assert guard.approval_covers(_review("kerry-berri", "APPROVED", "head"), "head", self._HITS, unreadable) is True + + def test_an_older_approval_covers_the_head_only_when_it_saw_every_spend_log_line(self): + approval: Final = _review("kerry-berri", "APPROVED", "old") + assert guard.approval_covers(approval, "head", self._HITS, lambda commit: self._HITS) is True + assert guard.approval_covers(approval, "head", self._HITS[:1], lambda commit: self._HITS) is True + assert guard.approval_covers(approval, "head", self._HITS, lambda commit: self._HITS[:1]) is False + + def test_an_approval_whose_commit_cannot_be_read_does_not_cover_the_head(self): + gone: Final = _review("kerry-berri", "APPROVED", "gone") + assert guard.approval_covers(gone, "head", self._HITS, lambda commit: None) is False + without_commit: Final = _review("kerry-berri", "APPROVED", None) + assert guard.approval_covers(without_commit, "head", self._HITS, lambda commit: self._HITS) is False + + class TestExitCode: def test_a_pull_request_without_spend_log_lines_passes_without_reviews(self, tmp_path, monkeypatch, capsys): monkeypatch.setattr(guard, "added_diff", lambda base, head: _diff("schema.prisma", " foo String", 1)) @@ -75,5 +115,15 @@ class TestExitCode: reviews.write_text('[{"user": {"login": "someone-else"}, "state": "APPROVED"}]') assert guard.main(["--base", "a", "--head", "b", "--approver", "kerry-berri", "--reviews", str(reviews)]) == 1 assert "schema.prisma: x LiteLLM_SpendLogs" in capsys.readouterr().out - reviews.write_text('[{"user": {"login": "kerry-berri"}, "state": "APPROVED"}]') + reviews.write_text('[{"user": {"login": "kerry-berri"}, "state": "APPROVED", "commit_id": "b"}]') assert guard.main(["--base", "a", "--head", "b", "--approver", "kerry-berri", "--reviews", str(reviews)]) == 0 + + def test_an_approval_before_a_new_spend_log_line_was_pushed_fails(self, tmp_path, monkeypatch, capsys): + monkeypatch.setattr(guard, "added_diff", lambda base, head: _diff("schema.prisma", " x LiteLLM_SpendLogs", 7)) + monkeypatch.setattr(guard, "reviewed_hits", lambda base, commit: () if commit == "old" else None, raising=False) + reviews: Final = tmp_path / "r.json" + reviews.write_text('[{"user": {"login": "kerry-berri"}, "state": "APPROVED", "commit_id": "old"}]') + assert guard.main(["--base", "a", "--head", "b", "--approver", "kerry-berri", "--reviews", str(reviews)]) == 1 + assert "changed after @kerry-berri approved old" in capsys.readouterr().out + reviews.write_text('[{"user": {"login": "kerry-berri"}, "state": "APPROVED", "commit_id": "gone"}]') + assert guard.main(["--base", "a", "--head", "b", "--approver", "kerry-berri", "--reviews", str(reviews)]) == 1 From c6ca22ebc58baa4398728bdcbba5b5156105ee8b Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Wed, 28 Oct 2026 14:49:48 -0700 Subject: [PATCH 092/128] fix(guardrails)!: make hide-secrets honor skip_tool_message_in_guardrail (#48689) * fix(guardrails): make hide-secrets honor skip_tool_message_in_guardrail With the flag on, per guardrail or in litellm_settings, hide-secrets leaves role: tool messages and Anthropic tool_result blocks as sent, as the unified guardrails do. Tool call arguments are still masked. * test(guardrails): type the skip-tool walker test and assert the whole request --- .../enterprise_callbacks/secret_detection.py | 8 ++- litellm/proxy/guardrails/_content_utils.py | 35 ++++++---- .../test_hide_secrets_tool_traffic.py | 68 +++++++++++++++++++ .../proxy/guardrails/test_content_utils.py | 29 ++++++++ 4 files changed, 126 insertions(+), 14 deletions(-) diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py b/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py index 65e4d50587f..c190e680347 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py @@ -32,6 +32,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.llms.base_llm.guardrail_translation.utils import effective_skip_tool_message_for_guardrail from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails._content_utils import iter_tool_json_fields, map_tool_json_text, walk_user_text from litellm.types.guardrails import GuardrailEventHooks @@ -753,7 +754,12 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail): _masked_entity_count.set({}) # Covers multimodal list content + Responses-API input. - inspected = walk_user_text(data, self.redact_text, self.redact_tool_payload) + self._redact_prompt(data) + inspected = walk_user_text( + data, + self.redact_text, + self.redact_tool_payload, + skip_tool_message=effective_skip_tool_message_for_guardrail(self), + ) + self._redact_prompt(data) if inspected == 0: # Image-only, empty-text, and unsupported payloads inspected diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index b4830b7c771..9bea22acd71 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -305,6 +305,11 @@ def iter_message_text(data: Mapping[str, object]) -> Iterator[str]: yield from _iter_text_parts_in_content(message.get("content")) +def _is_tool_message(message: object) -> bool: + fields: Final = _text_object(message) + return fields is not None and str(fields.get("role") or "").lower() == "tool" + + def _call_item(item: object) -> tuple[str, object] | None: """Return ``(field, value)`` for a Responses tool-call input item, else None.""" fields: Final = _text_object(item) @@ -316,11 +321,14 @@ def walk_user_text( data: dict[str, object], visit: Callable[[str], str], visit_tool: Callable[[JsonValue], JsonValue] | None = None, + *, + skip_tool_message: bool = False, ) -> int: """Rewrite every text fragment in place via ``visit``. Mutates ``data["messages"]`` and ``data["input"]`` (including tool payloads), ``data["system"]`` and - ``data["instructions"]``. Returns the number + ``data["instructions"]``. ``skip_tool_message`` leaves ``role: tool`` messages and ``tool_result`` + blocks as sent. Returns the number of fragments visited so callers can short-circuit when nothing was inspected. """ @@ -353,6 +361,8 @@ def walk_user_text( def _rewrite_part(part: Mapping[str, object], native: bool) -> Mapping[str, object]: nonlocal visited + if skip_tool_message and part.get("type") == "tool_result": + return part text: Final = _part_text(part) rewritten: Final = {**part, **({"text": visit(text)} if text is not None else {})} if text is not None: @@ -384,18 +394,17 @@ def walk_user_text( return content messages: Final = data.get("messages") - if isinstance(messages, list): - for message in messages: - if not isinstance(message, dict): - continue - if "content" in message: - message["content"] = _rewrite_content(message["content"]) - calls = message.get("tool_calls") - if isinstance(calls, list): - for call in calls: - function = call.get("function") if isinstance(call, dict) else None - if isinstance(function, dict) and "arguments" in function: - function["arguments"] = _rewrite_arguments(function["arguments"]) + for message in _object_list(messages) or (): + if (skip_tool_message and _is_tool_message(message)) or not isinstance(message, dict): + continue + if "content" in message: + message["content"] = _rewrite_content(message["content"]) + calls = message.get("tool_calls") + if isinstance(calls, list): + for call in calls: + function = call.get("function") if isinstance(call, dict) else None + if isinstance(function, dict) and "arguments" in function: + function["arguments"] = _rewrite_arguments(function["arguments"]) if "system" in data: data["system"] = _rewrite_content(data["system"]) # rebind-ok: preserve the native hook's in-place contract diff --git a/tests/integration/security/test_hide_secrets_tool_traffic.py b/tests/integration/security/test_hide_secrets_tool_traffic.py index 42fdc0f5a44..60b49520b1f 100644 --- a/tests/integration/security/test_hide_secrets_tool_traffic.py +++ b/tests/integration/security/test_hide_secrets_tool_traffic.py @@ -413,3 +413,71 @@ def test_resent_tool_history_is_scanned_once(single_worker: Rig, shape: str) -> second: Final = timed(lambda: single_worker.proxy.request("POST", path, following)) assert first.value.status_code == 200 and second.value.status_code == 200, (first.value.text, second.value.text) assert_within_factor_of_control(second.seconds, control=first.seconds, factor=0.5) + + +@pytest.mark.parametrize("placement", ("guardrail", "global")) +def test_skip_tool_message_in_guardrail_forwards_tool_results(tmp_path: Path, placement: str) -> None: + def configure(configuration: dict[str, JsonValue]) -> None: + guardrails: Final = configuration["guardrails"] + settings: Final = configuration["litellm_settings"] + assert isinstance(guardrails, list) and isinstance(guardrails[0], dict) and isinstance(settings, dict) + params: Final = guardrails[0]["litellm_params"] + assert isinstance(params, dict) + (params if placement == "guardrail" else settings)["skip_tool_message_in_guardrail"] = True + + with tool_rig(tmp_path, configure=configure) as rig: + rig.wire.drain() + replies: Final = ( + rig.proxy.request( + "POST", + "/v1/messages", + { + "model": "claude", + "max_tokens": 32, + "cache": {"no-cache": True}, + "messages": [ + {"role": "user", "content": TEXT}, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "tool_1", "name": "read_file", "input": {"env": TEXT}} + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "tool_1", "content": TEXT}], + }, + ], + }, + ), + rig.proxy.request( + "POST", + "/v1/chat/completions", + { + "model": "gpt", + "cache": {"no-cache": True}, + "messages": [ + {"role": "user", "content": TEXT}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "read_file", "arguments": json.dumps({"env": TEXT})}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": TEXT}, + ], + }, + ), + ) + assert all(reply.status_code == 200 for reply in replies), [reply.text for reply in replies] + sent: Final = rig.wire.drain() + assert len(sent) == 2 + for request in sent: + messages = json.loads(request.body)["messages"] + assert SECRET not in json.dumps(messages[:2]) + assert SECRET in json.dumps(messages[2]) diff --git a/tests/unit/proxy/guardrails/test_content_utils.py b/tests/unit/proxy/guardrails/test_content_utils.py index a668d5f609d..21d96a3ea97 100644 --- a/tests/unit/proxy/guardrails/test_content_utils.py +++ b/tests/unit/proxy/guardrails/test_content_utils.py @@ -2,6 +2,7 @@ import json from functools import reduce +from typing import Final import pytest from pydantic import JsonValue @@ -924,6 +925,34 @@ def test_native_writer_visits_each_tool_and_system_slot(data): assert "payload" not in json.dumps(data) +def test_skip_tool_message_leaves_tool_results_and_masks_the_rest() -> None: + data: Final[dict[str, object]] = { + "messages": [ + {"role": "user", "content": "secret"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "tool", "name": "read", "input": {"env": "secret"}}], + "tool_calls": [{"id": "call", "function": {"name": "read", "arguments": '{"env": "secret"}'}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "tool", "content": "secret"}]}, + {"role": "tool", "tool_call_id": "call", "content": "secret"}, + ] + } + assert walk_user_text(data, lambda text: text.replace("secret", "[MASKED]"), skip_tool_message=True) == 3 + assert data == { + "messages": [ + {"role": "user", "content": "[MASKED]"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "tool", "name": "read", "input": {"env": "[MASKED]"}}], + "tool_calls": [{"id": "call", "function": {"name": "read", "arguments": '{"env": "[MASKED]"}'}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "tool", "content": "secret"}]}, + {"role": "tool", "tool_call_id": "call", "content": "secret"}, + ] + } + + @pytest.mark.parametrize( "part", ( From 5215c2d3f77cb7cad6db861f323a4e892f5c16bf Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 28 Oct 2026 14:59:23 -0700 Subject: [PATCH 093/128] fix(pricing): add databricks priority pay per token prices (#48891) * fix(pricing): add databricks priority pay per token prices * fix(pricing): bill gemini 3.6 flash cache writes at its databricks input rate * fix(pricing): drop image token priority rate the cost map schema does not carry --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 233 +++++++++++++++--- model_prices_and_context_window.json | 233 +++++++++++++++--- .../test_databricks_cost_calculator.py | 26 ++ 3 files changed, 428 insertions(+), 64 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 3e0e86012a4..1f2bd7d7837 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -21589,7 +21589,13 @@ ], "supports_vision": true, "supports_audio_input": true, - "supports_video_input": true + "supports_video_input": true, + "input_cost_per_token_priority": 5.6252e-07, + "output_cost_per_token_priority": 3.37498e-06, + "cache_read_input_token_cost_priority": 5.628e-08, + "cache_creation_input_token_cost_priority": 5.6252e-07, + "input_cost_per_audio_token_priority": 1.12497e-06, + "cache_read_input_audio_token_cost_priority": 1.1249e-07 }, "databricks/databricks-gemini-3-1-flash-image": { "cache_creation_input_token_cost": 6.2503e-07, @@ -21676,7 +21682,14 @@ "supports_reasoning": true, "supports_tool_choice": true, "supports_video_input": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 4.50002e-06, + "output_cost_per_token_priority": 2.699998e-05, + "cache_read_input_token_cost_priority": 4.5003e-07, + "cache_creation_input_token_cost_priority": 4.50002e-06, + "input_cost_per_token_above_200k_tokens_priority": 8.99997e-06, + "output_cost_per_token_above_200k_tokens_priority": 4.049997e-05, + "cache_read_input_token_cost_above_200k_tokens_priority": 8.9999e-07 }, "databricks/databricks-gemini-3-flash": { "cache_creation_input_token_cost": 6.2503e-07, @@ -21710,7 +21723,13 @@ ], "supports_vision": true, "supports_audio_input": true, - "supports_video_input": true + "supports_video_input": true, + "input_cost_per_token_priority": 1.12497e-06, + "output_cost_per_token_priority": 6.75003e-06, + "cache_read_input_token_cost_priority": 1.1249e-07, + "cache_creation_input_token_cost_priority": 1.12497e-06, + "input_cost_per_audio_token_priority": 2.25001e-06, + "cache_read_input_audio_token_cost_priority": 2.2498e-07 }, "databricks/databricks-gemini-3-pro": { "cache_creation_input_token_cost": 2.49998e-06, @@ -21757,7 +21776,11 @@ "supports_reasoning": true, "supports_sampling_params": false, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 1.35002e-06, + "output_cost_per_token_priority": 6.75003e-06, + "cache_read_input_token_cost_priority": 1.3503e-07, + "cache_creation_input_token_cost_priority": 1.35002e-06 }, "databricks/databricks-gemini-3-7-flash": { "cache_creation_input_token_cost": 7.4998e-07, @@ -21785,10 +21808,14 @@ "supports_reasoning": true, "supports_sampling_params": false, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 1.35002e-06, + "output_cost_per_token_priority": 6.75003e-06, + "cache_read_input_token_cost_priority": 1.3503e-07, + "cache_creation_input_token_cost_priority": 1.35002e-06 }, "databricks/databricks-gemini-3-6-flash": { - "cache_creation_input_token_cost": 1.87502e-06, + "cache_creation_input_token_cost": 7.4998e-07, "cache_read_input_token_cost": 7.497e-08, "input_cost_per_token": 7.4998e-07, "input_dbu_cost_per_token": 1.0714e-05, @@ -21816,7 +21843,11 @@ "supports_reasoning": true, "supports_sampling_params": false, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 1.35002e-06, + "output_cost_per_token_priority": 6.75003e-06, + "cache_read_input_token_cost_priority": 1.3503e-07, + "cache_creation_input_token_cost_priority": 1.35002e-06 }, "databricks/databricks-gemini-3-5-flash": { "cache_creation_input_token_cost": 1.87502e-06, @@ -21846,7 +21877,11 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 3.37498e-06, + "output_cost_per_token_priority": 2.025002e-05, + "cache_read_input_token_cost_priority": 3.3747e-07, + "cache_creation_input_token_cost_priority": 3.37498e-06 }, "databricks/databricks-gemini-3-5-flash-lite": { "cache_creation_input_token_cost": 3.7499e-07, @@ -21878,7 +21913,11 @@ "supports_response_schema": true, "supports_sampling_params": false, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 6.7501e-07, + "output_cost_per_token_priority": 5.62499e-06, + "cache_read_input_token_cost_priority": 6.748e-08, + "cache_creation_input_token_cost_priority": 6.7501e-07 }, "databricks/databricks-gemma-3-12b": { "cache_creation_input_token_cost": 1.5001e-07, @@ -21921,7 +21960,11 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": false, + "input_cost_per_token_priority": 2.45e-06, + "output_cost_per_token_priority": 7.7e-06, + "cache_read_input_token_cost_priority": 4.55e-07, + "cache_creation_input_token_cost_priority": 2.45e-06 }, "databricks/databricks-glm-5-3": { "cache_creation_input_token_cost": 1.4e-06, @@ -21951,7 +21994,11 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": false, - "thinking_always_on": true + "thinking_always_on": true, + "input_cost_per_token_priority": 2.45e-06, + "output_cost_per_token_priority": 7.7e-06, + "cache_read_input_token_cost_priority": 4.55e-07, + "cache_creation_input_token_cost_priority": 2.45e-06 }, "databricks/databricks-glm-5-3-flash": { "cache_creation_input_token_cost": 1.5001e-07, @@ -21980,7 +22027,11 @@ "supports_reasoning": true, "supports_tool_choice": true, "supports_vision": true, - "thinking_always_on": true + "thinking_always_on": true, + "input_cost_per_token_priority": 2.625e-07, + "output_cost_per_token_priority": 8.75e-07, + "cache_read_input_token_cost_priority": 5.25e-08, + "cache_creation_input_token_cost_priority": 2.625e-07 }, "databricks/databricks-gpt-5": { "cache_creation_input_token_cost": 1.24999e-06, @@ -21999,7 +22050,11 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token_priority": 2.49998e-06, + "output_cost_per_token_priority": 1.999998e-05, + "cache_read_input_token_cost_priority": 2.4997e-07, + "cache_creation_input_token_cost_priority": 2.49998e-06 }, "databricks/databricks-gpt-5-1": { "cache_creation_input_token_cost": 1.24999e-06, @@ -22018,7 +22073,11 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token_priority": 2.49998e-06, + "output_cost_per_token_priority": 1.999998e-05, + "cache_read_input_token_cost_priority": 2.4997e-07, + "cache_creation_input_token_cost_priority": 2.49998e-06 }, "databricks/databricks-gpt-5-2": { "cache_creation_input_token_cost": 1.75e-06, @@ -22037,7 +22096,11 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token_priority": 3.5e-06, + "output_cost_per_token_priority": 2.8e-05, + "cache_read_input_token_cost_priority": 3.5e-07, + "cache_creation_input_token_cost_priority": 3.5e-06 }, "databricks/databricks-gpt-5-3-codex": { "cache_creation_input_token_cost": 1.75e-06, @@ -22056,7 +22119,11 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token_priority": 3.5e-06, + "output_cost_per_token_priority": 2.8e-05, + "cache_read_input_token_cost_priority": 3.5e-07, + "cache_creation_input_token_cost_priority": 3.5e-06 }, "databricks/databricks-gpt-5-4": { "cache_creation_input_token_cost": 2.49998e-06, @@ -22087,7 +22154,14 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 5.00003e-06, + "output_cost_per_token_priority": 2.999997e-05, + "cache_read_input_token_cost_priority": 5.0001e-07, + "cache_creation_input_token_cost_priority": 5.00003e-06, + "input_cost_per_token_above_272k_tokens_priority": 5.00003e-06, + "output_cost_per_token_above_272k_tokens_priority": 2.999997e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 5.0001e-07 }, "databricks/databricks-gpt-5-4-mini": { "cache_creation_input_token_cost": 7.4998e-07, @@ -22115,7 +22189,11 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 1.50003e-06, + "output_cost_per_token_priority": 8.99997e-06, + "cache_read_input_token_cost_priority": 1.5001e-07, + "cache_creation_input_token_cost_priority": 1.50003e-06 }, "databricks/databricks-gpt-5-4-nano": { "cache_creation_input_token_cost": 1.9999e-07, @@ -22175,7 +22253,15 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 8.00002e-06, + "output_cost_per_token_priority": 4.000003e-05, + "cache_read_input_token_cost_priority": 8.0003e-07, + "cache_creation_input_token_cost_priority": 9.99999e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.599997e-05, + "output_cost_per_token_above_272k_tokens_priority": 6.000001e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.59999e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1.999998e-05 }, "databricks/databricks-gpt-5-6-terra": { "cache_creation_input_token_cost": 2.49998e-06, @@ -22207,7 +22293,15 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 4.00001e-06, + "output_cost_per_token_priority": 2.399999e-05, + "cache_read_input_token_cost_priority": 3.9998e-07, + "cache_creation_input_token_cost_priority": 5.00003e-06, + "input_cost_per_token_above_272k_tokens_priority": 8.00002e-06, + "output_cost_per_token_above_272k_tokens_priority": 3.600002e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 8.0003e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06 }, "databricks/databricks-gpt-5-6-luna": { "cache_creation_input_token_cost": 2.4997e-07, @@ -22239,7 +22333,15 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 3.9998e-07, + "output_cost_per_token_priority": 2.40002e-06, + "cache_read_input_token_cost_priority": 3.997e-08, + "cache_creation_input_token_cost_priority": 5.0001e-07, + "input_cost_per_token_above_272k_tokens_priority": 8.0003e-07, + "output_cost_per_token_above_272k_tokens_priority": 3.60003e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 8.001e-08, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1.00002e-06 }, "databricks/databricks-gpt-5-5": { "cache_creation_input_token_cost": 5.00003e-06, @@ -22270,7 +22372,14 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 1.249997e-05, + "output_cost_per_token_priority": 7.50001e-05, + "cache_read_input_token_cost_priority": 1.24999e-06, + "cache_creation_input_token_cost_priority": 1.249997e-05, + "input_cost_per_token_above_272k_tokens_priority": 1.249997e-05, + "output_cost_per_token_above_272k_tokens_priority": 7.50001e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.24999e-06 }, "databricks/databricks-gpt-5-5-pro": { "cache_creation_input_token_cost": 2.999997e-05, @@ -22319,7 +22428,11 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token_priority": 4.5003e-07, + "output_cost_per_token_priority": 3.60003e-06, + "cache_read_input_token_cost_priority": 4.501e-08, + "cache_creation_input_token_cost_priority": 4.5003e-07 }, "databricks/databricks-gpt-5-nano": { "cache_creation_input_token_cost": 4.998e-08, @@ -22469,7 +22582,11 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 5.25e-06, + "output_cost_per_token_priority": 2.625e-05, + "cache_read_input_token_cost_priority": 5.25e-07, + "cache_creation_input_token_cost_priority": 5.25e-06 }, "databricks/system.ai.kimi-k3": { "litellm_provider": "databricks", @@ -22551,7 +22668,11 @@ "supports_reasoning": true, "supports_tool_choice": true, "supports_vision": false, - "thinking_always_on": true + "thinking_always_on": true, + "input_cost_per_token_priority": 3.85e-07, + "output_cost_per_token_priority": 3.85e-06, + "cache_read_input_token_cost_priority": 3.85e-07, + "cache_creation_input_token_cost_priority": 3.85e-07 }, "databricks/databricks-qwen3-next-80b-a3b-instruct": { "cache_creation_input_token_cost": 1.5001e-07, @@ -77774,7 +77895,11 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 5.25e-07, + "output_cost_per_token_priority": 2.1e-06, + "cache_read_input_token_cost_priority": 5.25e-08, + "cache_creation_input_token_cost_priority": 5.25e-07 }, "databricks/databricks-deepseek-v4-1-flash-1017": { "cache_creation_input_token_cost": 3.0002e-07, @@ -77897,7 +78022,11 @@ "supports_reasoning": true, "supports_tool_choice": true, "supports_video_input": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 4.50002e-06, + "output_cost_per_token_priority": 2.250003e-05, + "cache_read_input_token_cost_priority": 2.2498e-07, + "cache_creation_input_token_cost_priority": 4.50002e-06 }, "databricks/databricks-gpt-5-4-pro": { "cache_creation_input_token_cost": 2.999997e-05, @@ -77959,7 +78088,15 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 4.00001e-06, + "output_cost_per_token_priority": 1.99997e-05, + "cache_read_input_token_cost_priority": 1.9999e-07, + "cache_creation_input_token_cost_priority": 5.00003e-06, + "input_cost_per_token_above_272k_tokens_priority": 8.00002e-06, + "output_cost_per_token_above_272k_tokens_priority": 2.99999e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 3.9998e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06 }, "databricks/databricks-gpt-6-1-terra": { "cache_creation_input_token_cost": 1.24999e-06, @@ -77991,7 +78128,15 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 1.99997e-06, + "output_cost_per_token_priority": 9.99999e-06, + "cache_read_input_token_cost_priority": 1.0003e-07, + "cache_creation_input_token_cost_priority": 2.49998e-06, + "input_cost_per_token_above_272k_tokens_priority": 4.00001e-06, + "output_cost_per_token_above_272k_tokens_priority": 1.500002e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.9999e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 5.00003e-06 }, "databricks/databricks-gpt-6-astra": { "cache_creation_input_token_cost": 1.249997e-05, @@ -78023,7 +78168,15 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 1.999998e-05, + "output_cost_per_token_priority": 9.99999e-05, + "cache_read_input_token_cost_priority": 1.99997e-06, + "cache_creation_input_token_cost_priority": 2.500001e-05, + "input_cost_per_token_above_272k_tokens_priority": 4.000003e-05, + "output_cost_per_token_above_272k_tokens_priority": 0.0001500002, + "cache_read_input_token_cost_above_272k_tokens_priority": 4.00001e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 5.000002e-05 }, "databricks/databricks-gpt-6-luna": { "cache_creation_input_token_cost": 1.2502e-07, @@ -78055,7 +78208,15 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 1.9999e-07, + "output_cost_per_token_priority": 1.0003e-06, + "cache_read_input_token_cost_priority": 2.002e-08, + "cache_creation_input_token_cost_priority": 2.4997e-07, + "input_cost_per_token_above_272k_tokens_priority": 3.9998e-07, + "output_cost_per_token_above_272k_tokens_priority": 1.5001e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 3.997e-08, + "cache_creation_input_token_cost_above_272k_tokens_priority": 5.0001e-07 }, "databricks/databricks-gpt-6-sol": { "cache_creation_input_token_cost": 2.49998e-06, @@ -78087,7 +78248,15 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 4.00001e-06, + "output_cost_per_token_priority": 1.99997e-05, + "cache_read_input_token_cost_priority": 3.9998e-07, + "cache_creation_input_token_cost_priority": 5.00003e-06, + "input_cost_per_token_above_272k_tokens_priority": 8.00002e-06, + "output_cost_per_token_above_272k_tokens_priority": 2.99999e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 8.0003e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06 }, "databricks/databricks-gpt-image-1": { "cache_creation_input_token_cost": 5.00003e-06, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3e0e86012a4..1f2bd7d7837 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -21589,7 +21589,13 @@ ], "supports_vision": true, "supports_audio_input": true, - "supports_video_input": true + "supports_video_input": true, + "input_cost_per_token_priority": 5.6252e-07, + "output_cost_per_token_priority": 3.37498e-06, + "cache_read_input_token_cost_priority": 5.628e-08, + "cache_creation_input_token_cost_priority": 5.6252e-07, + "input_cost_per_audio_token_priority": 1.12497e-06, + "cache_read_input_audio_token_cost_priority": 1.1249e-07 }, "databricks/databricks-gemini-3-1-flash-image": { "cache_creation_input_token_cost": 6.2503e-07, @@ -21676,7 +21682,14 @@ "supports_reasoning": true, "supports_tool_choice": true, "supports_video_input": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 4.50002e-06, + "output_cost_per_token_priority": 2.699998e-05, + "cache_read_input_token_cost_priority": 4.5003e-07, + "cache_creation_input_token_cost_priority": 4.50002e-06, + "input_cost_per_token_above_200k_tokens_priority": 8.99997e-06, + "output_cost_per_token_above_200k_tokens_priority": 4.049997e-05, + "cache_read_input_token_cost_above_200k_tokens_priority": 8.9999e-07 }, "databricks/databricks-gemini-3-flash": { "cache_creation_input_token_cost": 6.2503e-07, @@ -21710,7 +21723,13 @@ ], "supports_vision": true, "supports_audio_input": true, - "supports_video_input": true + "supports_video_input": true, + "input_cost_per_token_priority": 1.12497e-06, + "output_cost_per_token_priority": 6.75003e-06, + "cache_read_input_token_cost_priority": 1.1249e-07, + "cache_creation_input_token_cost_priority": 1.12497e-06, + "input_cost_per_audio_token_priority": 2.25001e-06, + "cache_read_input_audio_token_cost_priority": 2.2498e-07 }, "databricks/databricks-gemini-3-pro": { "cache_creation_input_token_cost": 2.49998e-06, @@ -21757,7 +21776,11 @@ "supports_reasoning": true, "supports_sampling_params": false, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 1.35002e-06, + "output_cost_per_token_priority": 6.75003e-06, + "cache_read_input_token_cost_priority": 1.3503e-07, + "cache_creation_input_token_cost_priority": 1.35002e-06 }, "databricks/databricks-gemini-3-7-flash": { "cache_creation_input_token_cost": 7.4998e-07, @@ -21785,10 +21808,14 @@ "supports_reasoning": true, "supports_sampling_params": false, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 1.35002e-06, + "output_cost_per_token_priority": 6.75003e-06, + "cache_read_input_token_cost_priority": 1.3503e-07, + "cache_creation_input_token_cost_priority": 1.35002e-06 }, "databricks/databricks-gemini-3-6-flash": { - "cache_creation_input_token_cost": 1.87502e-06, + "cache_creation_input_token_cost": 7.4998e-07, "cache_read_input_token_cost": 7.497e-08, "input_cost_per_token": 7.4998e-07, "input_dbu_cost_per_token": 1.0714e-05, @@ -21816,7 +21843,11 @@ "supports_reasoning": true, "supports_sampling_params": false, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 1.35002e-06, + "output_cost_per_token_priority": 6.75003e-06, + "cache_read_input_token_cost_priority": 1.3503e-07, + "cache_creation_input_token_cost_priority": 1.35002e-06 }, "databricks/databricks-gemini-3-5-flash": { "cache_creation_input_token_cost": 1.87502e-06, @@ -21846,7 +21877,11 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 3.37498e-06, + "output_cost_per_token_priority": 2.025002e-05, + "cache_read_input_token_cost_priority": 3.3747e-07, + "cache_creation_input_token_cost_priority": 3.37498e-06 }, "databricks/databricks-gemini-3-5-flash-lite": { "cache_creation_input_token_cost": 3.7499e-07, @@ -21878,7 +21913,11 @@ "supports_response_schema": true, "supports_sampling_params": false, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 6.7501e-07, + "output_cost_per_token_priority": 5.62499e-06, + "cache_read_input_token_cost_priority": 6.748e-08, + "cache_creation_input_token_cost_priority": 6.7501e-07 }, "databricks/databricks-gemma-3-12b": { "cache_creation_input_token_cost": 1.5001e-07, @@ -21921,7 +21960,11 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": false, + "input_cost_per_token_priority": 2.45e-06, + "output_cost_per_token_priority": 7.7e-06, + "cache_read_input_token_cost_priority": 4.55e-07, + "cache_creation_input_token_cost_priority": 2.45e-06 }, "databricks/databricks-glm-5-3": { "cache_creation_input_token_cost": 1.4e-06, @@ -21951,7 +21994,11 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": false, - "thinking_always_on": true + "thinking_always_on": true, + "input_cost_per_token_priority": 2.45e-06, + "output_cost_per_token_priority": 7.7e-06, + "cache_read_input_token_cost_priority": 4.55e-07, + "cache_creation_input_token_cost_priority": 2.45e-06 }, "databricks/databricks-glm-5-3-flash": { "cache_creation_input_token_cost": 1.5001e-07, @@ -21980,7 +22027,11 @@ "supports_reasoning": true, "supports_tool_choice": true, "supports_vision": true, - "thinking_always_on": true + "thinking_always_on": true, + "input_cost_per_token_priority": 2.625e-07, + "output_cost_per_token_priority": 8.75e-07, + "cache_read_input_token_cost_priority": 5.25e-08, + "cache_creation_input_token_cost_priority": 2.625e-07 }, "databricks/databricks-gpt-5": { "cache_creation_input_token_cost": 1.24999e-06, @@ -21999,7 +22050,11 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token_priority": 2.49998e-06, + "output_cost_per_token_priority": 1.999998e-05, + "cache_read_input_token_cost_priority": 2.4997e-07, + "cache_creation_input_token_cost_priority": 2.49998e-06 }, "databricks/databricks-gpt-5-1": { "cache_creation_input_token_cost": 1.24999e-06, @@ -22018,7 +22073,11 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token_priority": 2.49998e-06, + "output_cost_per_token_priority": 1.999998e-05, + "cache_read_input_token_cost_priority": 2.4997e-07, + "cache_creation_input_token_cost_priority": 2.49998e-06 }, "databricks/databricks-gpt-5-2": { "cache_creation_input_token_cost": 1.75e-06, @@ -22037,7 +22096,11 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token_priority": 3.5e-06, + "output_cost_per_token_priority": 2.8e-05, + "cache_read_input_token_cost_priority": 3.5e-07, + "cache_creation_input_token_cost_priority": 3.5e-06 }, "databricks/databricks-gpt-5-3-codex": { "cache_creation_input_token_cost": 1.75e-06, @@ -22056,7 +22119,11 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token_priority": 3.5e-06, + "output_cost_per_token_priority": 2.8e-05, + "cache_read_input_token_cost_priority": 3.5e-07, + "cache_creation_input_token_cost_priority": 3.5e-06 }, "databricks/databricks-gpt-5-4": { "cache_creation_input_token_cost": 2.49998e-06, @@ -22087,7 +22154,14 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 5.00003e-06, + "output_cost_per_token_priority": 2.999997e-05, + "cache_read_input_token_cost_priority": 5.0001e-07, + "cache_creation_input_token_cost_priority": 5.00003e-06, + "input_cost_per_token_above_272k_tokens_priority": 5.00003e-06, + "output_cost_per_token_above_272k_tokens_priority": 2.999997e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 5.0001e-07 }, "databricks/databricks-gpt-5-4-mini": { "cache_creation_input_token_cost": 7.4998e-07, @@ -22115,7 +22189,11 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 1.50003e-06, + "output_cost_per_token_priority": 8.99997e-06, + "cache_read_input_token_cost_priority": 1.5001e-07, + "cache_creation_input_token_cost_priority": 1.50003e-06 }, "databricks/databricks-gpt-5-4-nano": { "cache_creation_input_token_cost": 1.9999e-07, @@ -22175,7 +22253,15 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 8.00002e-06, + "output_cost_per_token_priority": 4.000003e-05, + "cache_read_input_token_cost_priority": 8.0003e-07, + "cache_creation_input_token_cost_priority": 9.99999e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.599997e-05, + "output_cost_per_token_above_272k_tokens_priority": 6.000001e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.59999e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1.999998e-05 }, "databricks/databricks-gpt-5-6-terra": { "cache_creation_input_token_cost": 2.49998e-06, @@ -22207,7 +22293,15 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 4.00001e-06, + "output_cost_per_token_priority": 2.399999e-05, + "cache_read_input_token_cost_priority": 3.9998e-07, + "cache_creation_input_token_cost_priority": 5.00003e-06, + "input_cost_per_token_above_272k_tokens_priority": 8.00002e-06, + "output_cost_per_token_above_272k_tokens_priority": 3.600002e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 8.0003e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06 }, "databricks/databricks-gpt-5-6-luna": { "cache_creation_input_token_cost": 2.4997e-07, @@ -22239,7 +22333,15 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 3.9998e-07, + "output_cost_per_token_priority": 2.40002e-06, + "cache_read_input_token_cost_priority": 3.997e-08, + "cache_creation_input_token_cost_priority": 5.0001e-07, + "input_cost_per_token_above_272k_tokens_priority": 8.0003e-07, + "output_cost_per_token_above_272k_tokens_priority": 3.60003e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 8.001e-08, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1.00002e-06 }, "databricks/databricks-gpt-5-5": { "cache_creation_input_token_cost": 5.00003e-06, @@ -22270,7 +22372,14 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 1.249997e-05, + "output_cost_per_token_priority": 7.50001e-05, + "cache_read_input_token_cost_priority": 1.24999e-06, + "cache_creation_input_token_cost_priority": 1.249997e-05, + "input_cost_per_token_above_272k_tokens_priority": 1.249997e-05, + "output_cost_per_token_above_272k_tokens_priority": 7.50001e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.24999e-06 }, "databricks/databricks-gpt-5-5-pro": { "cache_creation_input_token_cost": 2.999997e-05, @@ -22319,7 +22428,11 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token_priority": 4.5003e-07, + "output_cost_per_token_priority": 3.60003e-06, + "cache_read_input_token_cost_priority": 4.501e-08, + "cache_creation_input_token_cost_priority": 4.5003e-07 }, "databricks/databricks-gpt-5-nano": { "cache_creation_input_token_cost": 4.998e-08, @@ -22469,7 +22582,11 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 5.25e-06, + "output_cost_per_token_priority": 2.625e-05, + "cache_read_input_token_cost_priority": 5.25e-07, + "cache_creation_input_token_cost_priority": 5.25e-06 }, "databricks/system.ai.kimi-k3": { "litellm_provider": "databricks", @@ -22551,7 +22668,11 @@ "supports_reasoning": true, "supports_tool_choice": true, "supports_vision": false, - "thinking_always_on": true + "thinking_always_on": true, + "input_cost_per_token_priority": 3.85e-07, + "output_cost_per_token_priority": 3.85e-06, + "cache_read_input_token_cost_priority": 3.85e-07, + "cache_creation_input_token_cost_priority": 3.85e-07 }, "databricks/databricks-qwen3-next-80b-a3b-instruct": { "cache_creation_input_token_cost": 1.5001e-07, @@ -77774,7 +77895,11 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 5.25e-07, + "output_cost_per_token_priority": 2.1e-06, + "cache_read_input_token_cost_priority": 5.25e-08, + "cache_creation_input_token_cost_priority": 5.25e-07 }, "databricks/databricks-deepseek-v4-1-flash-1017": { "cache_creation_input_token_cost": 3.0002e-07, @@ -77897,7 +78022,11 @@ "supports_reasoning": true, "supports_tool_choice": true, "supports_video_input": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 4.50002e-06, + "output_cost_per_token_priority": 2.250003e-05, + "cache_read_input_token_cost_priority": 2.2498e-07, + "cache_creation_input_token_cost_priority": 4.50002e-06 }, "databricks/databricks-gpt-5-4-pro": { "cache_creation_input_token_cost": 2.999997e-05, @@ -77959,7 +78088,15 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 4.00001e-06, + "output_cost_per_token_priority": 1.99997e-05, + "cache_read_input_token_cost_priority": 1.9999e-07, + "cache_creation_input_token_cost_priority": 5.00003e-06, + "input_cost_per_token_above_272k_tokens_priority": 8.00002e-06, + "output_cost_per_token_above_272k_tokens_priority": 2.99999e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 3.9998e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06 }, "databricks/databricks-gpt-6-1-terra": { "cache_creation_input_token_cost": 1.24999e-06, @@ -77991,7 +78128,15 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 1.99997e-06, + "output_cost_per_token_priority": 9.99999e-06, + "cache_read_input_token_cost_priority": 1.0003e-07, + "cache_creation_input_token_cost_priority": 2.49998e-06, + "input_cost_per_token_above_272k_tokens_priority": 4.00001e-06, + "output_cost_per_token_above_272k_tokens_priority": 1.500002e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.9999e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 5.00003e-06 }, "databricks/databricks-gpt-6-astra": { "cache_creation_input_token_cost": 1.249997e-05, @@ -78023,7 +78168,15 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 1.999998e-05, + "output_cost_per_token_priority": 9.99999e-05, + "cache_read_input_token_cost_priority": 1.99997e-06, + "cache_creation_input_token_cost_priority": 2.500001e-05, + "input_cost_per_token_above_272k_tokens_priority": 4.000003e-05, + "output_cost_per_token_above_272k_tokens_priority": 0.0001500002, + "cache_read_input_token_cost_above_272k_tokens_priority": 4.00001e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 5.000002e-05 }, "databricks/databricks-gpt-6-luna": { "cache_creation_input_token_cost": 1.2502e-07, @@ -78055,7 +78208,15 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 1.9999e-07, + "output_cost_per_token_priority": 1.0003e-06, + "cache_read_input_token_cost_priority": 2.002e-08, + "cache_creation_input_token_cost_priority": 2.4997e-07, + "input_cost_per_token_above_272k_tokens_priority": 3.9998e-07, + "output_cost_per_token_above_272k_tokens_priority": 1.5001e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 3.997e-08, + "cache_creation_input_token_cost_above_272k_tokens_priority": 5.0001e-07 }, "databricks/databricks-gpt-6-sol": { "cache_creation_input_token_cost": 2.49998e-06, @@ -78087,7 +78248,15 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 4.00001e-06, + "output_cost_per_token_priority": 1.99997e-05, + "cache_read_input_token_cost_priority": 3.9998e-07, + "cache_creation_input_token_cost_priority": 5.00003e-06, + "input_cost_per_token_above_272k_tokens_priority": 8.00002e-06, + "output_cost_per_token_above_272k_tokens_priority": 2.99999e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 8.0003e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06 }, "databricks/databricks-gpt-image-1": { "cache_creation_input_token_cost": 5.00003e-06, diff --git a/tests/unit/llms/databricks/test_databricks_cost_calculator.py b/tests/unit/llms/databricks/test_databricks_cost_calculator.py index 34d3806d6eb..c35aa6e4cf0 100644 --- a/tests/unit/llms/databricks/test_databricks_cost_calculator.py +++ b/tests/unit/llms/databricks/test_databricks_cost_calculator.py @@ -253,3 +253,29 @@ def test_cost_per_token_bills_the_served_priority_tier( prompt_cost, completion_cost = cost_per_token(model="databricks/dbrx-tiered-test", usage=usage) assert prompt_cost == pytest.approx(30 * 0.001) assert completion_cost == pytest.approx(40 * 0.002) + + +RATES_WITHOUT_A_PRIORITY_FIELD: Final = frozenset({"input_cost_per_image_token"}) + + +def _is_standard_rate_key(key: str) -> bool: + return ( + ("_cost_per_" in key or key.endswith("_token_cost")) + and "dbu" not in key + and key not in RATES_WITHOUT_A_PRIORITY_FIELD + and not key.endswith(("_priority", "_flex", "_batches")) + ) + + +def test_every_databricks_priority_row_prices_every_standard_rate_at_the_priority_tier() -> None: + cost_map: Final = json.loads(MAIN_PRICES.read_text()) + priority_rows: Final = { + key: row + for key, row in cost_map.items() + if key.startswith("databricks/") and "input_cost_per_token_priority" in row + } + + assert priority_rows + for key, row in priority_rows.items(): + missing = sorted(f for f in row if _is_standard_rate_key(f) and f"{f}_priority" not in row) + assert not missing, f"{key} bills these at the standard rate on priority requests: {missing}" From b8fb50f70e82f5eed8a8d56cdb5b6b1ca4e1e7b7 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 28 Oct 2026 15:03:40 -0700 Subject: [PATCH 094/128] fix(ci): keep the spend logs guard approval through a later COMMENTED review (LIT-10242) (#48901) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ci_cd/spend_logs_guard.py | 11 ++++++++--- tests/unit/test_spend_logs_guard.py | 23 +++++++++++++++++++++++ 2 files changed, 31 insertions(+), 3 deletions(-) diff --git a/ci_cd/spend_logs_guard.py b/ci_cd/spend_logs_guard.py index 64e33c4fb70..fac9d4ba5d2 100644 --- a/ci_cd/spend_logs_guard.py +++ b/ci_cd/spend_logs_guard.py @@ -27,7 +27,8 @@ WATCHED: Final = ( "litellm-proxy-extras/litellm_proxy_extras/schema.prisma", ) MIGRATIONS_PREFIX: Final = "litellm-proxy-extras/litellm_proxy_extras/migrations/" -TABLE: Final = re.compile(r"(? str: def approved_by(reviews: Sequence[Review], login: str) -> bool: - """True when `login`'s most recent review on the pull request is an approval.""" - states: Final = tuple(review.state for review in reviews if review.user.login == login) + """True when `login`'s most recent review that changes GitHub's approval state is an approval. + + COMMENTED and PENDING reviews leave an earlier approval in force on GitHub, so they are skipped here too. + """ + own: Final = tuple(review.state for review in reviews if review.user.login == login) + states: Final = tuple(state for state in own if state in DECIDING_STATES) return bool(states) and states[-1] == "APPROVED" diff --git a/tests/unit/test_spend_logs_guard.py b/tests/unit/test_spend_logs_guard.py index 6a3355e1c52..9ac2c923632 100644 --- a/tests/unit/test_spend_logs_guard.py +++ b/tests/unit/test_spend_logs_guard.py @@ -3,6 +3,8 @@ import sys from pathlib import Path from typing import Final +import pytest + _SCRIPT: Final = Path(__file__).resolve().parents[2] / "ci_cd" / "spend_logs_guard.py" _SPEC: Final = importlib.util.spec_from_file_location("spend_logs_guard", _SCRIPT) assert _SPEC is not None and _SPEC.loader is not None @@ -35,6 +37,13 @@ class TestSpendLogHits: elsewhere: Final = _diff("litellm/proxy/utils.py", 'table = "LiteLLM_SpendLogs"', 5) assert guard.spend_log_hits(removed + other + elsewhere) == () + def test_a_table_whose_name_only_starts_with_the_watched_one_is_not_a_hit(self): + archive: Final = _diff("schema.prisma", "model LiteLLM_SpendLogsArchive {", 900) + index: Final = _diff("schema.prisma", " @@index([team_id], map: \"LiteLLM_SpendLogs_team_id_idx\")", 760) + assert guard.spend_log_hits(archive + index) == ( + 'schema.prisma: @@index([team_id], map: "LiteLLM_SpendLogs_team_id_idx")', + ) + def test_two_files_in_one_diff_are_attributed_to_their_own_path(self): diff: Final = _diff("litellm/proxy/schema.prisma", " a LiteLLM_SpendLogs", 1) + _diff( "schema.prisma", " b String", 1 @@ -59,6 +68,20 @@ class TestApprovedBy: assert guard.approved_by(reviews, "kerry-berri") is True assert guard.approved_by(reviews[::-1], "kerry-berri") is False + @pytest.mark.parametrize( + ("states", "approved"), + ( + (("APPROVED", "COMMENTED"), True), + (("APPROVED", "PENDING"), True), + (("CHANGES_REQUESTED", "COMMENTED"), False), + (("APPROVED", "DISMISSED"), False), + (("COMMENTED",), False), + ), + ) + def test_a_comment_or_pending_review_does_not_change_the_answer(self, states, approved): + reviews: Final = tuple(guard.Review(state=state, user=guard.Reviewer(login="kerry-berri")) for state in states) + assert guard.approved_by(reviews, "kerry-berri") is approved + def test_no_review_from_the_approver_is_not_an_approval(self): other: Final = guard.Review(state="APPROVED", user=guard.Reviewer(login="someone-else")) assert guard.approved_by((other,), "kerry-berri") is False From e21de0e418c076be488eca83cf97d2ba3c1998cd Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 28 Oct 2026 15:11:26 -0700 Subject: [PATCH 095/128] fix(logging): bill each SDK call its own native Bedrock guardrail trace when calls share one metadata dict (LIT-10239) (#48923) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 2 +- ...ck_native_guardrail_shared_metadata_sdk.py | 67 +++++++++++++++++++ .../test_litellm_logging.py | 32 +++++++++ 3 files changed, 100 insertions(+), 1 deletion(-) create mode 100644 tests/integration/sdk/test_bedrock_native_guardrail_shared_metadata_sdk.py diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 3b8157b2249..80f1f8f5eca 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2758,7 +2758,7 @@ class Logging(LiteLLMLoggingBaseClass): if entry is None or not is_str_keyed_dict(litellm_params): return metadata: Final[object] = litellm_params.get("metadata") - bucket: Final = metadata if is_str_keyed_dict(metadata) else {} + bucket: Final = {**metadata} if is_str_keyed_dict(metadata) else {} litellm_params["metadata"] = bucket with_guardrail_entry(bucket, entry) diff --git a/tests/integration/sdk/test_bedrock_native_guardrail_shared_metadata_sdk.py b/tests/integration/sdk/test_bedrock_native_guardrail_shared_metadata_sdk.py new file mode 100644 index 00000000000..d0235310aaa --- /dev/null +++ b/tests/integration/sdk/test_bedrock_native_guardrail_shared_metadata_sdk.py @@ -0,0 +1,67 @@ +"""An SDK caller that reuses one ``metadata=`` dict across ``litellm.completion`` calls to a Bedrock deployment whose +``guardrailConfig`` has ``trace: enabled`` is billed, on each call, for that call's guardrail trace only, and the +caller's dict is left as the caller wrote it (LIT-10239)""" + +import json +from collections.abc import Mapping +from typing import Final + +import pytest +from integration._support.client import eventually +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +import litellm + +_MODEL: Final = "bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0" +_PROMPT_TOKENS: Final = 5 +_COMPLETION_TOKENS: Final = 1 +_USAGE: Final[Mapping[str, int]] = {"contentPolicyUnits": 1, "topicPolicyUnits": 1} +_GUARDRAIL_IDS: Final = ("guardrail-one", "guardrail-two") + + +def _bedrock(request: Request) -> Reply: + """Bedrock's Converse runtime: a one-token answer and the trace of the guardrail the request named.""" + guardrail_id: Final = json.loads(request.body)["guardrailConfig"]["guardrailIdentifier"] + body: Final[dict[str, JsonValue]] = { + "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": _PROMPT_TOKENS, "outputTokens": _COMPLETION_TOKENS, "totalTokens": 6}, + "metrics": {"latencyMs": 3}, + "trace": {"guardrail": {"inputAssessment": {guardrail_id: {"invocationMetrics": {"usage": dict(_USAGE)}}}}}, + } + return Reply(body=json.dumps(body).encode()) + + +def _guardrail_cost() -> float: + rates: Final = litellm.model_cost["bedrock/guardrails"]["guardrail_cost_per_unit"] + return sum(units * rates[counter] for counter, units in _USAGE.items()) + + +def _llm_cost() -> float: + prices: Final = litellm.model_cost[_MODEL.removeprefix("bedrock/")] + return _PROMPT_TOKENS * prices["input_cost_per_token"] + _COMPLETION_TOKENS * prices["output_cost_per_token"] + + +def test_each_call_sharing_one_metadata_dict_is_billed_its_own_trace_only() -> None: + billed: list[float] = [] # mutable-ok: the success callback appends what each call was billed + litellm.success_callback.append( + lambda kwargs, response, start, end: billed.append(kwargs["standard_logging_object"]["response_cost"]) + ) + shared_metadata: Final = {"tag": "x"} + with wire_server(_bedrock) as wire: + for guardrail_id in _GUARDRAIL_IDS: + litellm.completion( + model=_MODEL, + messages=[{"role": "user", "content": "hello"}], + api_base=wire.url, + aws_access_key_id="AKIASCRIPTED", + aws_secret_access_key="scripted", + aws_region_name="us-east-1", + guardrailConfig={"guardrailIdentifier": guardrail_id, "guardrailVersion": "1", "trace": "enabled"}, + metadata=shared_metadata, + ) + costs: Final = eventually(lambda: tuple(billed), lambda seen: len(seen) == len(_GUARDRAIL_IDS), seconds=30) + + assert costs == pytest.approx((_llm_cost() + _guardrail_cost(),) * len(_GUARDRAIL_IDS)) + assert shared_metadata == {"tag": "x"} diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index e786753e1a8..299088a6c4f 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -13775,3 +13775,35 @@ def test_record_native_bedrock_guardrail_bills_nothing_for_a_response_served_fro logging_obj.record_native_bedrock_guardrail({"trace": trace}) assert "standard_logging_guardrail_information" not in logging_obj.model_call_details["litellm_params"]["metadata"] + + +def test_record_native_bedrock_guardrail_bills_each_call_its_own_trace_when_calls_share_one_metadata_dict() -> None: + import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost + + shared_metadata = {"tag": "x"} + traces = tuple( + {"guardrail": {"inputAssessment": {"gid": {"invocationMetrics": {"usage": {"contentPolicyUnits": units}}}}}} + for units in (1, 2) + ) + costs = [] + for call_index, trace in enumerate(traces): + logging_obj = Logging( + model="bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": "hello"}], + stream=False, + call_type="completion", + start_time=datetime.datetime.now(), + litellm_call_id=f"call-shared-metadata-{call_index}", + function_id="fn", + ) + logging_obj.model_call_details["litellm_params"] = {"metadata": shared_metadata} + logging_obj.record_native_bedrock_guardrail({"trace": trace}) + recorded = logging_obj.model_call_details["litellm_params"]["metadata"]["standard_logging_guardrail_information"] + costs.append(guardrail_information_cost(recorded)) + + (first_cost, second_cost) = costs + assert second_cost == pytest.approx(2 * first_cost) + assert shared_metadata == {"tag": "x"} From 4a70d80aef0c600e4c493a77ab0a69d001d3d3bc Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 28 Oct 2026 15:23:52 -0700 Subject: [PATCH 096/128] fix(aiohttp_transport): keep the owned session open while a response stream is still being read (LIT-10237) (#48920) * fix(aiohttp_transport): keep the owned session open while a response stream is still being read (LIT-10237) * test(aiohttp_transport): annotate the stream regression test locals with Final and drop its docstring --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/custom_httpx/aiohttp_transport.py | 23 +++++++- .../custom_httpx/test_aiohttp_transport.py | 55 +++++++++++++++++++ 2 files changed, 77 insertions(+), 1 deletion(-) diff --git a/litellm/llms/custom_httpx/aiohttp_transport.py b/litellm/llms/custom_httpx/aiohttp_transport.py index 0e4d97fd770..9297e76624d 100644 --- a/litellm/llms/custom_httpx/aiohttp_transport.py +++ b/litellm/llms/custom_httpx/aiohttp_transport.py @@ -6,6 +6,7 @@ import ssl import sys import typing import urllib.request +import weakref from collections.abc import Callable, Generator from typing import ClassVar, Final @@ -109,6 +110,24 @@ def map_aiohttp_exceptions() -> Generator[None, None, None]: raise mapped_exc(message) from exc +def _release_transport_anchor(_transport: object) -> None: + """Release a transport anchored to a response stream. See ``_anchor_transport_to``.""" + + +def _anchor_transport_to(stream: "AiohttpResponseStream", transport: object) -> None: + """Keep the transport alive for as long as the stream can still read. + + The transport finalizer closes the session it owns, and nothing on the path from a + live response back to that session references the transport, so a client dropped + mid-stream would close the session under the body. ``weakref.finalize`` holds the + transport in its own registry rather than on the stream, which keeps the transport + out of the response's reference cycle (an exception carrying the response, for one) + so it is finalized by refcount once the anchor drops and can still schedule an async + close, instead of being collected in the same cyclic pass as the session it closes. + """ + weakref.finalize(stream, _release_transport_anchor, transport) + + class AiohttpResponseStream(httpx.AsyncByteStream): CHUNK_SIZE = 1024 * 16 @@ -463,10 +482,12 @@ class LiteLLMAiohttpTransport(AiohttpTransport): # Re-raise if it's a different RuntimeError raise + stream: Final = AiohttpResponseStream(response) + _anchor_transport_to(stream, self) return httpx.Response( status_code=response.status, headers=response.headers, - stream=AiohttpResponseStream(response), + stream=stream, request=request, ) diff --git a/tests/unit/llms/custom_httpx/test_aiohttp_transport.py b/tests/unit/llms/custom_httpx/test_aiohttp_transport.py index c7d5518ef5b..aed808223bd 100644 --- a/tests/unit/llms/custom_httpx/test_aiohttp_transport.py +++ b/tests/unit/llms/custom_httpx/test_aiohttp_transport.py @@ -5,6 +5,7 @@ import socket import sys import struct import threading +import weakref from typing import Final import aiohttp @@ -1383,3 +1384,57 @@ def test_a_transport_collected_while_its_loop_runs_elsewhere_closes_its_session_ loop.call_soon_threadsafe(loop.stop) thread.join(timeout=10) loop.close() + + +@pytest.mark.asyncio +async def test_a_stream_keeps_reading_after_the_client_that_opened_it_is_collected(): + from aiohttp import web + + release: Final = asyncio.Event() + body: Final = tuple(f"chunk{i}\n".encode() for i in range(4)) + + async def handler(request: web.Request) -> web.StreamResponse: + response: Final = web.StreamResponse() + await response.prepare(request) + await response.write(body[0]) + await release.wait() + for chunk in body[1:]: + await response.write(chunk) + await response.write_eof() + return response + + app: Final = web.Application() + app.router.add_get("/stream", handler) + runner: Final = web.AppRunner(app) + await runner.setup() + sock: Final = socket.socket() + sock.bind(("127.0.0.1", 0)) + port: Final[int] = sock.getsockname()[1] + await web.SockSite(runner, sock).start() + + transport: Final = LiteLLMAiohttpTransport(client=aiohttp.ClientSession) + client: Final = httpx.AsyncClient(transport=transport) + transport_ref: Final = weakref.ref(transport) + try: + response: Final = await client.send(client.build_request("GET", f"http://127.0.0.1:{port}/stream"), stream=True) + session: Final = transport.client + assert isinstance(session, aiohttp.ClientSession) + chunks: Final = response.aiter_bytes() + first: Final = await chunks.__anext__() + + del client, transport + gc.collect() + await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks) + release.set() + + rest: Final = [chunk async for chunk in chunks] + assert first + b"".join(rest) == b"".join(body), (first, rest) + assert not session.closed, "the session was closed under a stream that was still being read" + + del response, chunks + gc.collect() + await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks) + assert transport_ref() is None, "the finished stream kept the transport alive" + assert session.closed, "the owned session stayed open after its last stream finished" + finally: + await runner.cleanup() From 2791822f921fab9469b86ed37d8a2deadad0fc58 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 28 Oct 2026 15:28:35 -0700 Subject: [PATCH 097/128] ci(spend-logs-guard): run only on pull_request_target, re-run by hand after approving (#48941) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/spend-logs-guard.yml | 4 ---- ci_cd/spend_logs_guard.py | 4 ++-- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/.github/workflows/spend-logs-guard.yml b/.github/workflows/spend-logs-guard.yml index 354a1b7c947..29e9f236202 100644 --- a/.github/workflows/spend-logs-guard.yml +++ b/.github/workflows/spend-logs-guard.yml @@ -4,10 +4,6 @@ on: # zizmor: ignore[dangerous-triggers] runs the base branch's code only; the P pull_request_target: branches: - main - pull_request_review: - types: - - submitted - - dismissed permissions: contents: read diff --git a/ci_cd/spend_logs_guard.py b/ci_cd/spend_logs_guard.py index fac9d4ba5d2..a4d9a2e0a0d 100644 --- a/ci_cd/spend_logs_guard.py +++ b/ci_cd/spend_logs_guard.py @@ -6,7 +6,7 @@ so a path rule would put every migration on the owner. This gate reads the diff added line in the watched files mentions the table it passes at once, otherwise it passes only once the owner's latest review on the pull request is an approval that still covers every spend-log line in the head (approvals survive pushes, so a later commit may not add spend-log lines the owner never -saw). The workflow re-runs on every review. +saw). Reviews do not trigger the workflow, so the check is re-run by hand after approving. """ from __future__ import annotations @@ -153,7 +153,7 @@ def main(argv: Sequence[str]) -> int: else f"LiteLLM_SpendLogs lines changed after @{args.approver} approved {approval.commit_id}, " f"a new approving review is needed:\n" ) - sys.stdout.write(f"{reason}{listed}The check re-runs when a review is submitted.\n") + sys.stdout.write(f"{reason}{listed}After approving, re-run this job.\n") return 1 From 84c4c5aa5c530610863afe104b32dc50ebca37da Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 28 Oct 2026 15:32:27 -0700 Subject: [PATCH 098/128] fix(router): cool down failing deployments when fallbacks use string targets (LIT-10238) (#48913) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router_utils/cooldown_handlers.py | 32 ++++++++++++++----- .../sdk/test_router_cooldown_fallback_sdk.py | 28 ++++++++++++++++ .../router_utils/test_cooldown_handlers.py | 32 +++++++++++++++++++ 3 files changed, 84 insertions(+), 8 deletions(-) diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index 128bca6cca0..a1b1e327802 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -13,7 +13,7 @@ from datetime import datetime from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError import litellm from litellm._internal_context import service_target @@ -417,17 +417,35 @@ def _should_run_cooldown_logic( return True -_FALLBACK_ENTRIES_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, list[str | Mapping[str, object]]] | str, ...]) +_FALLBACK_ENTRIES_ADAPTER: Final = TypeAdapter(tuple[object, ...]) +_FALLBACK_ENTRY_ADAPTER: Final = TypeAdapter(Mapping[str, list[str | Mapping[str, object]] | str]) def _fallback_target_name(target: str | Mapping[str, object]) -> str: return target if isinstance(target, str) else str(target.get("model", "")) +def _fallback_target_names( + targets: Sequence[str | Mapping[str, object]] | str, +) -> list[str]: # mutable-ok: mirrors get_fallback_model_group's fallbacks shape + return [targets] if isinstance(targets, str) else [_fallback_target_name(target) for target in targets] + + def _fallback_chain_names( - entry: Mapping[str, Sequence[str | Mapping[str, object]]], -) -> Mapping[str, list[str]]: # mutable-ok: mirrors get_fallback_model_group's fallbacks shape - return {key: [_fallback_target_name(target) for target in targets] for key, targets in entry.items()} + entry: object, +) -> Mapping[str, list[str]] | None: # mutable-ok: mirrors get_fallback_model_group's fallbacks shape + try: + validated: Final = _FALLBACK_ENTRY_ADAPTER.validate_python(entry) + except ValidationError: + return None + return {key: _fallback_target_names(targets) for key, targets in validated.items()} + + +def _fallback_entries(configured: object) -> tuple[object, ...]: + try: + return _FALLBACK_ENTRIES_ADAPTER.validate_python(configured) + except ValidationError: + return () def _has_fallback_group(litellm_router_instance: LitellmRouter, requested_model_group: str | None) -> bool: @@ -437,9 +455,7 @@ def _has_fallback_group(litellm_router_instance: LitellmRouter, requested_model_ if requested_model_group is None or not configured: return False chains: Final = tuple( - _fallback_chain_names(entry) - for entry in _FALLBACK_ENTRIES_ADAPTER.validate_python(configured) - if isinstance(entry, Mapping) + chain for chain in map(_fallback_chain_names, _fallback_entries(configured)) if chain is not None ) fallback_group, _ = get_fallback_model_group(fallbacks=chains, model_group=requested_model_group) return bool(fallback_group) diff --git a/tests/integration/sdk/test_router_cooldown_fallback_sdk.py b/tests/integration/sdk/test_router_cooldown_fallback_sdk.py index 747adeee41d..e7508a2670e 100644 --- a/tests/integration/sdk/test_router_cooldown_fallback_sdk.py +++ b/tests/integration/sdk/test_router_cooldown_fallback_sdk.py @@ -12,6 +12,7 @@ from pydantic import JsonValue, TypeAdapter import litellm from litellm import Router +from litellm.router_utils.cooldown_handlers import async_get_cooldown_deployments from litellm.types.router import DeploymentTypedDict _MODEL: Final = "gpt-5.4" @@ -118,3 +119,30 @@ async def test_cooldown_skips_only_the_failed_duplicate_model_deployment() -> No "/primary/chat/completions", "/backup/chat/completions", ) + + +@pytest.mark.parametrize( + "router_kwargs", + [{"fallbacks": [{"primary": "backup"}]}, {"default_fallbacks": "backup"}], + ids=["string_fallback_target", "string_default_fallbacks"], +) +@pytest.mark.asyncio +async def test_string_valued_fallbacks_still_cool_down_the_failing_primary(router_kwargs: dict[str, object]) -> None: + with wire_server(_peer) as wire: + router: Final = Router( + model_list=[ + _deployment("primary", f"{wire.url}/primary", "primary-deployment"), + _deployment("backup", f"{wire.url}/backup", "backup-deployment"), + ], + allowed_fails=0, + num_retries=0, + **router_kwargs, + ) + with pytest.raises(litellm.AuthenticationError): + await router.acompletion( + model="primary", + messages=[{"role": "user", "content": f"cooldown {uuid.uuid4().hex}"}], + ) + cooled: Final = await async_get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None) + + assert tuple(cooled) == ("primary-deployment",) diff --git a/tests/unit/router_utils/test_cooldown_handlers.py b/tests/unit/router_utils/test_cooldown_handlers.py index 0d39016ea79..24459a5a250 100644 --- a/tests/unit/router_utils/test_cooldown_handlers.py +++ b/tests/unit/router_utils/test_cooldown_handlers.py @@ -25,6 +25,7 @@ from litellm.router_utils.cooldown_callbacks import router_cooldown_event_callba from litellm.router_utils.cooldown_handlers import ( _get_deployment_cooldown_policy, _has_explicit_allowed_fails_policy_for_exception, + _has_fallback_group, _increment_allowed_fails, _is_cooldown_required, _resolve_allowed_fails_from_policy, @@ -3149,6 +3150,8 @@ BENCHED: Final = (litellm.Timeout, litellm.Timeout, RouterRateLimitError, Router {}, BENCHED, ), + ({"allowed_fails": 1, "fallbacks": [{"test-model": "backup-model"}]}, {}, BENCHED), + ({"allowed_fails": 1, "default_fallbacks": "backup-model"}, {}, BENCHED), ], ids=[ "allowed_fails", @@ -3157,6 +3160,8 @@ BENCHED: Final = (litellm.Timeout, litellm.Timeout, RouterRateLimitError, Router "policy_names_timeout", "allowed_fails_with_a_fallback_group", "allowed_fails_with_a_dict_fallback_target", + "allowed_fails_with_a_string_fallback_target", + "allowed_fails_with_a_string_default_fallbacks", ], ) @pytest.mark.asyncio @@ -3180,6 +3185,33 @@ async def test_router_allowed_fails_benches_a_lone_deployment_only_when_the_poli assert len(route.calls) == outcomes.count(litellm.Timeout) +@pytest.mark.parametrize( + ("fallbacks", "expected"), + [ + ([{"test-model": "backup-model"}], True), + ([{"*": "backup-model"}], True), + ([{"test-model": ["backup-model"]}], True), + ([{"test-model": 5}, {"*": "backup-model"}], True), + ([{"test-model": 5}], False), + ([{"other-model": ["backup-model"]}], False), + ], + ids=[ + "string_target", + "string_catch_all_from_default_fallbacks", + "list_target", + "unparseable_entry_does_not_hide_the_catch_all", + "unparseable_entry_alone", + "another_groups_chain", + ], +) +def test_has_fallback_group_reads_string_targets_and_skips_unparseable_entries( + fallbacks: Sequence[Mapping[str, object]], expected: bool +) -> None: + router: Final = Router(model_list=[TIMED_OUT_DEPLOYMENT], fallbacks=list(fallbacks)) + + assert _has_fallback_group(router, "test-model") is expected + + async def _group_call_outcome(router: Router) -> str: try: response: Final = await router.acompletion( From e157c5e520acbd83ed75bb26ec3633c4b2cc09b1 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 28 Oct 2026 15:34:15 -0700 Subject: [PATCH 099/128] fix(migration_data_checks): size btree entries by the bytes Postgres stores (LIT-10244) (#48911) * fix(migration_data_checks): size btree entries by the bytes Postgres stores (LIT-10244) * test(migration_data_checks): flag exactly the compressed keys Postgres rejects at the entry limit (LIT-10244) --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../migration_data_checks.py | 11 +++- .../migrations/test_request_log_indexes.py | 64 +++++++++++++++++++ .../test_migration_data_checks.py | 15 +++++ 3 files changed, 87 insertions(+), 3 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py b/litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py index 7b01e64dd73..e22639daaef 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py +++ b/litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py @@ -343,11 +343,16 @@ def btree_max_entry_bytes(block_size: int) -> int: def _value_bytes(column: ColumnType, identifier: "sql.Identifier") -> "sql.Composed": + """The bytes a value takes in the index entry: its length, or its stored size when + Postgres compressed it, since `index_form_tuple` keeps the compressed bytes.""" from psycopg import sql - if column.type_name in _OCTET_LENGTH_TYPES: - return sql.SQL("octet_length({})").format(identifier) - return sql.SQL("octet_length({}::text)").format(identifier) + length: Final = ( + sql.SQL("octet_length({})").format(identifier) + if column.type_name in _OCTET_LENGTH_TYPES + else sql.SQL("octet_length({}::text)").format(identifier) + ) + return sql.SQL("least({}, pg_column_size({}))").format(length, identifier) def _aligned(offset: "sql.Composable", align: str) -> "sql.Composed": diff --git a/tests/integration/database/migrations/test_request_log_indexes.py b/tests/integration/database/migrations/test_request_log_indexes.py index 1de2090b487..f558c355c84 100644 --- a/tests/integration/database/migrations/test_request_log_indexes.py +++ b/tests/integration/database/migrations/test_request_log_indexes.py @@ -1,3 +1,4 @@ +import hashlib import os import queue import secrets @@ -973,6 +974,7 @@ def test_the_migration_job_refuses_before_the_first_migration_when_spend_log_key _insert_spend_log_with_key(conn, "legacy-1", _legacy_bearer_token()) _insert_spend_log_with_key(conn, "legacy-2", _legacy_bearer_token()) _insert_spend_log_with_key(conn, "hashed-1", "a" * 64) + _insert_spend_log_with_key(conn, "repetitive-1", "z" * 3000) ledger_before: Final = _ledger(scratch_database) with pytest.raises(RuntimeError) as refused: @@ -1109,6 +1111,68 @@ def test_the_entry_size_check_draws_the_line_exactly_where_postgres_does( assert finding.detail == f"with an index entry over 2704 bytes, largest 2712 bytes ({described_as})" +@requires_db +def test_a_value_postgres_compresses_under_the_limit_is_not_counted_while_random_text_of_that_length_is( + scratch_database: str, +) -> None: + """Postgres compresses a long value before it checks the entry size, so 3000 repeated + bytes index fine while 3000 random hex bytes do not. The check must draw the same line.""" + create_index: Final = sql.SQL('CREATE INDEX "compress_probe_idx" ON "compress_probe" ("k")') + with psycopg.connect(scratch_database, autocommit=True) as conn: + conn.execute('CREATE TABLE "compress_probe" ("k" TEXT)') + conn.execute('INSERT INTO "compress_probe" ("k") VALUES (%s), (%s)', ("a" * 3000, secrets.token_hex(1500))) + findings: Final = migration_data_checks.find_blockers( + conn, "public", migration_data_checks.checks_for_statement(create_index.as_string(conn), "probe") + ) + with pytest.raises(psycopg.errors.ProgramLimitExceeded, match="index row size 3016 exceeds"): + conn.execute(create_index) + conn.execute('DELETE FROM "compress_probe" WHERE "k" <> %s', ("a" * 3000,)) + conn.execute(create_index) + stored: Final = conn.execute('SELECT octet_length("k"), pg_column_size("k") FROM "compress_probe"').fetchone() + + assert stored is not None and stored[0] == 3000 and stored[1] < 2704, stored + (finding,) = findings + assert isinstance(finding, migration_data_checks.BlockedChange) + assert finding.rows == 1, "only the random key is too long for the index once Postgres compresses the other" + assert finding.detail == "with an index entry over 2704 bytes, largest 3016 bytes (k up to 3000 bytes)" + + +def _postgres_rejects_an_index_on(conn: "psycopg.Connection[tuple[object, ...]]", key: str) -> bool: + conn.execute('DROP TABLE IF EXISTS "one_key"') + conn.execute('CREATE TABLE "one_key" ("k" TEXT)') + conn.execute('INSERT INTO "one_key" ("k") VALUES (%s)', (key,)) + try: + conn.execute('CREATE INDEX "one_key_idx" ON "one_key" ("k")') + except psycopg.errors.ProgramLimitExceeded: + return True + return False + + +@requires_db +def test_the_check_flags_exactly_the_compressed_keys_postgres_rejects_at_the_entry_limit(scratch_database: str) -> None: + """Keys that compress to a stored size just under and just over the entry limit go through the + check and through Postgres one by one. The check must flag a key exactly when Postgres rejects it.""" + noise: Final = "".join(hashlib.sha256(str(i).encode()).hexdigest() for i in range(40)) + keys: Final = tuple(noise[:n] + "a" * 1500 for n in range(2400, 2500)) + create_index: Final = sql.SQL('CREATE INDEX "stored_probe_idx" ON "stored_probe" ("k")') + with psycopg.connect(scratch_database, autocommit=True) as conn: + conn.execute('CREATE TABLE "stored_probe" ("k" TEXT)') + conn.execute('INSERT INTO "stored_probe" ("k") SELECT unnest(%s::text[])', (list(keys),)) + compressed: Final = conn.execute( + 'SELECT count(*) FROM "stored_probe" WHERE pg_column_size("k") < octet_length("k") AND octet_length("k") > 2704' + ).fetchone() + findings: Final = migration_data_checks.find_blockers( + conn, "public", migration_data_checks.checks_for_statement(create_index.as_string(conn), "probe") + ) + rejected: Final = tuple(key for key in keys if _postgres_rejects_an_index_on(conn, key)) + + assert compressed == (len(keys),), "every probe key must be longer than the limit yet stored compressed" + assert 0 < len(rejected) < len(keys), f"the probe must straddle the limit, Postgres rejected {len(rejected)}" + (finding,) = findings + assert isinstance(finding, migration_data_checks.BlockedChange) + assert finding.rows == len(rejected), f"the check must flag the {len(rejected)} keys Postgres rejects and no other" + + @requires_db def test_a_scan_that_cannot_finish_in_time_is_reported_as_unverified_rather_than_passed(scratch_database: str) -> None: with psycopg.connect(scratch_database, autocommit=True) as conn: diff --git a/tests/unit/litellm_proxy_extras/test_migration_data_checks.py b/tests/unit/litellm_proxy_extras/test_migration_data_checks.py index f3b787beda1..d4490927e29 100644 --- a/tests/unit/litellm_proxy_extras/test_migration_data_checks.py +++ b/tests/unit/litellm_proxy_extras/test_migration_data_checks.py @@ -7,7 +7,9 @@ from litellm_proxy_extras.migration_data_checks import ( MIGRATION_DATA_CHECK_TIMEOUT_ENV_VAR, BlockedChange, Check, + ColumnType, CouldNotVerify, + _count_query, btree_max_entry_bytes, checks_for_index, checks_for_migration, @@ -113,6 +115,19 @@ def test_the_btree_entry_limit_on_the_default_page_is_the_number_in_postgres_err assert btree_max_entry_bytes(8192) == 2704 +def test_the_entry_size_query_sizes_a_text_column_by_its_stored_bytes_when_postgres_compressed_it() -> None: + """Postgres compresses a long value before it checks the entry size, so a repetitive + 3000-byte key indexes fine. The query has to size it as stored, never by its length alone.""" + (check,) = checks_for_index( + "LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", ("api_key", "startTime") + ) + types: Final = (ColumnType("text", -1, "i", -1), ColumnType("timestamp", 8, "d", 3)) + query: Final = _count_query(check, "public", types, 2704).as_string(None) + stored: Final = 'least(octet_length("api_key"), pg_column_size("api_key"))' + assert stored in query + assert 'octet_length("api_key")' not in query.replace(stored, ""), query + + def test_the_report_names_the_rows_the_statement_and_the_fix_and_says_nothing_was_applied() -> None: (check,) = checks_for_index( "LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", ("api_key", "startTime") From 0b57b9976a609ab74ba5c5cf21f75079527527ff Mon Sep 17 00:00:00 2001 From: mrinal-berri Date: Wed, 28 Oct 2026 15:35:08 -0700 Subject: [PATCH 100/128] feat(daily-activity): per-guardrail spend breakdown in user, team and organization daily activity (#48702) * feat(daily-activity): read guardrail spend per key from the daily guardrail usage table Adds DailyActivityRepository.guardrail_spend, which sums LiteLLM_DailyGuardrailUsageUnits.cost by date, guardrail and key. A guardrail row counts only when its key has a spend row on that day under the same scope filters as the rest of the daily activity response, so user, key and date filters apply to guardrail cost too * feat(daily-activity): add a per-guardrail spend breakdown to /user/daily/activity/aggregated Each day now has breakdown.guardrails[] with metrics.spend and an api_key_breakdown carrying each key's spend and metadata. Guardrail entries only report spend, since that is the one metric stored per key. The breakdown is on the user endpoint only, is empty when a model filter is set because guardrail cost is not stored per model, and lists only keys within api_key_limit while the guardrail total keeps every key. spend and budgets are unchanged Resolves LIT-9785 * feat(daily-activity): add the guardrail breakdown to team and organization activity, with guardrail names /team/daily/activity/aggregated and /organization/daily/activity/aggregated now return breakdown.guardrails the same way the user endpoint does, since each key belongs to one team and one organization. Each guardrail entry also carries metadata.guardrail_name, read from LiteLLM_GuardrailsTable for guardrails created through the UI or API and falling back to the breakdown key for config.yaml guardrails, whose key is already their name * fix(daily-activity): keep daily activity up when the guardrail breakdown query fails, and treat an empty model filter as none A database error from the guardrail spend query now logs a warning and returns the response without guardrails, where it used to fail the whole aggregated response with a 500. An empty model= no longer drops the guardrail breakdown, matching how the rest of the endpoint already ignores it. The daily activity golden files and fixtures gain the empty guardrails key every breakdown now carries * refactor(daily-activity): fold guardrail breakdown in a worker thread and drop docstrings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(daily-activity): report untracked guardrail units next to spend 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> --- litellm/proxy/_lazy_openapi_snapshot.json | 75 ++++++++ .../common_daily_activity.py | 68 ++++++- .../daily_activity_routes.py | 1 + .../daily_activity_scopes.py | 4 + .../repositories/daily_activity_repository.py | 6 + litellm/repositories/daily_activity_sql.py | 19 ++ .../common_daily_activity.py | 29 +++ litellm/types/repositories/daily_activity.py | 10 + .../spend/fixtures/daily_activity_team.json | 2 +- .../spend/fixtures/daily_activity_user.json | 2 +- .../daily_activity_team_aggregated.json | 2 +- .../golden/daily_activity_team_paginated.json | 2 +- .../daily_activity_user_aggregated.json | 2 +- .../golden/daily_activity_user_paginated.json | 2 +- ...st_daily_activity_aggregated_breakdowns.py | 180 ++++++++++++++++++ .../test_daily_activity_routes.py | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 36 ++++ 17 files changed, 432 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index b10b5ae98b5..39adb99a0d0 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -3487,6 +3487,13 @@ "title": "Entities", "type": "object" }, + "guardrails": { + "additionalProperties": { + "$ref": "#/components/schemas/GuardrailMetric" + }, + "title": "Guardrails", + "type": "object" + }, "mcp_servers": { "additionalProperties": { "$ref": "#/components/schemas/MetricWithMetadata" @@ -3802,6 +3809,74 @@ "title": "ExportType", "type": "string" }, + "GuardrailKeyMetric": { + "properties": { + "metadata": { + "$ref": "#/components/schemas/KeyMetadata" + }, + "metrics": { + "$ref": "#/components/schemas/GuardrailSpendMetrics" + } + }, + "required": [ + "metrics", + "metadata" + ], + "title": "GuardrailKeyMetric", + "type": "object" + }, + "GuardrailMetadata": { + "properties": { + "guardrail_name": { + "title": "Guardrail Name", + "type": "string" + } + }, + "required": [ + "guardrail_name" + ], + "title": "GuardrailMetadata", + "type": "object" + }, + "GuardrailMetric": { + "properties": { + "api_key_breakdown": { + "additionalProperties": { + "$ref": "#/components/schemas/GuardrailKeyMetric" + }, + "title": "Api Key Breakdown", + "type": "object" + }, + "metadata": { + "$ref": "#/components/schemas/GuardrailMetadata" + }, + "metrics": { + "$ref": "#/components/schemas/GuardrailSpendMetrics" + } + }, + "required": [ + "metrics", + "metadata" + ], + "title": "GuardrailMetric", + "type": "object" + }, + "GuardrailSpendMetrics": { + "properties": { + "spend": { + "default": 0.0, + "title": "Spend", + "type": "number" + }, + "untracked_units": { + "default": 0, + "title": "Untracked Units", + "type": "integer" + } + }, + "title": "GuardrailSpendMetrics", + "type": "object" + }, "HTTPAuthSecurityScheme": { "description": "Defines a security scheme using HTTP authentication.", "properties": { diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index f948a49441a..66f386b792d 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -3,6 +3,7 @@ from collections.abc import Awaitable, Callable, Mapping, Sequence from collections.abc import Set as AbstractSet from dataclasses import dataclass, replace from datetime import date, datetime, timedelta +from itertools import groupby from types import MappingProxyType from typing import Final, Literal, NoReturn, Protocol @@ -28,6 +29,10 @@ from litellm.types.proxy.management_endpoints.common_daily_activity import ( DailySpendData, DailySpendMetadata, GroupedData, + GuardrailKeyMetric, + GuardrailMetadata, + GuardrailMetric, + GuardrailSpendMetrics, KeyMetadata, KeyMetricWithMetadata, MetricWithMetadata, @@ -39,6 +44,7 @@ from litellm.types.repositories.daily_activity import ( DailyActivityTable, EntityRollupRow, GroupingSetsRow, + GuardrailSpendRow, KeyMetadataRow, RollupMetricsRow, SpendLogsWindow, @@ -1013,17 +1019,68 @@ def _fold_entity_rollups_sync( ) +def _guardrail_metric(day: DailySpendData, rows: Sequence[GuardrailSpendRow]) -> GuardrailMetric: + return GuardrailMetric( + metrics=GuardrailSpendMetrics( + spend=sum(row.spend for row in rows), + untracked_units=sum(row.untracked_units for row in rows), + ), + metadata=GuardrailMetadata(guardrail_name=rows[0].guardrail_name), + api_key_breakdown={ + row.api_key: GuardrailKeyMetric( + metrics=GuardrailSpendMetrics(spend=row.spend, untracked_units=row.untracked_units), + metadata=day.breakdown.api_keys[row.api_key].metadata, + ) + for row in rows + if row.api_key in day.breakdown.api_keys + }, + ) + + +def _day_with_guardrails(day: DailySpendData, rows: Sequence[GuardrailSpendRow]) -> DailySpendData: + by_guardrail: Final = groupby(sorted(rows, key=lambda row: row.guardrail_id), key=lambda row: row.guardrail_id) + guardrails: Final = {guardrail_id: _guardrail_metric(day, tuple(group)) for guardrail_id, group in by_guardrail} + return day.model_copy(update={"breakdown": day.breakdown.model_copy(update={"guardrails": guardrails})}) + + +def _with_guardrail_breakdown( + results: Sequence[DailySpendData], rows: Sequence[GuardrailSpendRow] +) -> tuple[DailySpendData, ...]: + rows_by_date: Final = MappingProxyType( + {day: tuple(group) for day, group in groupby(sorted(rows, key=lambda row: row.date), key=lambda row: row.date)} + ) + return tuple(_day_with_guardrails(day, rows_by_date.get(day.date.strftime("%Y-%m-%d"), ())) for day in results) + + +async def _guardrail_spend_or_empty( + repository: DailyActivityRepository, scope: DailyActivityScope +) -> tuple[GuardrailSpendRow, ...]: + from prisma.errors import PrismaError + + try: + return await repository.guardrail_spend(scope) + except PrismaError as exc: + verbose_proxy_logger.warning("Returning daily activity without its guardrail breakdown: %s", exc) + return () + + async def get_daily_activity_aggregated( repository: DailyActivityRepository, scope: DailyActivityScope, *, entity_metadata_field: Mapping[str, dict[str, object]] | None = None, include_entity_breakdown: bool = False, + include_guardrail_breakdown: bool = False, api_key_limit: int = constants.USAGE_TOP_API_KEYS_DEFAULT, ) -> SpendAnalyticsPaginatedResponse: try: - aggregated_rows: Final = await repository.aggregated( - scope, include_entity_breakdown=include_entity_breakdown, api_key_limit=api_key_limit + aggregated_rows, guardrail_rows = await asyncio.gather( + repository.aggregated( + scope, include_entity_breakdown=include_entity_breakdown, api_key_limit=api_key_limit + ), + _guardrail_spend_or_empty(repository, scope) + if include_guardrail_breakdown and not scope.model + else asyncio.sleep(0, result=()), ) records: Final = aggregated_rows.grouping_rows aggregated: Final = await _aggregate_grouping_sets_records( @@ -1058,8 +1115,13 @@ async def get_daily_activity_aggregated( api_key_metadata=entity_key_metadata, entity_metadata_field=entity_metadata_field, ) + results: Final = ( + list(await asyncio.to_thread(_with_guardrail_breakdown, aggregated["results"], guardrail_rows)) + if guardrail_rows + else aggregated["results"] + ) return SpendAnalyticsPaginatedResponse( - results=aggregated["results"], + results=results, metadata=DailySpendMetadata( total_spend=aggregated["totals"].spend, total_flat_cost=aggregated["totals"].flat_cost, diff --git a/litellm/proxy/management_endpoints/daily_activity_routes.py b/litellm/proxy/management_endpoints/daily_activity_routes.py index 842873129e4..79e515944ce 100644 --- a/litellm/proxy/management_endpoints/daily_activity_routes.py +++ b/litellm/proxy/management_endpoints/daily_activity_routes.py @@ -300,6 +300,7 @@ def _register_aggregated_route(router: APIRouter, resolver: EntityScopeResolver, resolved.scope, entity_metadata_field=resolved.entity_metadata, include_entity_breakdown=resolver.include_entity_breakdown, + include_guardrail_breakdown=resolver.include_guardrail_breakdown, api_key_limit=api_key_limit, ) if resolver.entity != "team": diff --git a/litellm/proxy/management_endpoints/daily_activity_scopes.py b/litellm/proxy/management_endpoints/daily_activity_scopes.py index bc373388a4e..77246f5dda6 100644 --- a/litellm/proxy/management_endpoints/daily_activity_scopes.py +++ b/litellm/proxy/management_endpoints/daily_activity_scopes.py @@ -56,6 +56,7 @@ class EntityScopeResolver: resolve: EntityScopeResolve include_entity_breakdown: bool operation_names: OperationNames + include_guardrail_breakdown: bool = False def _query_ids(value: str | None) -> tuple[str, ...] | None: @@ -364,6 +365,7 @@ USER_RESOLVER = EntityScopeResolver( query=_user_query, resolve=_resolve_user, include_entity_breakdown=False, + include_guardrail_breakdown=True, operation_names=MappingProxyType( { "aggregated": "get_user_daily_activity_aggregated", @@ -385,6 +387,7 @@ TEAM_RESOLVER = EntityScopeResolver( query=_team_query, resolve=_resolve_team, include_entity_breakdown=True, + include_guardrail_breakdown=True, operation_names=MappingProxyType( { "aggregated": "get_team_daily_activity_aggregated", @@ -423,6 +426,7 @@ ORGANIZATION_RESOLVER = EntityScopeResolver( query=_organization_query, resolve=_resolve_organization, include_entity_breakdown=True, + include_guardrail_breakdown=True, operation_names=MappingProxyType( { "aggregated": "get_organization_daily_activity_aggregated", diff --git a/litellm/repositories/daily_activity_repository.py b/litellm/repositories/daily_activity_repository.py index 10c603fe56f..4ca23e4e89e 100644 --- a/litellm/repositories/daily_activity_repository.py +++ b/litellm/repositories/daily_activity_repository.py @@ -19,6 +19,7 @@ from litellm.repositories.daily_activity_sql import ( build_cache_leakage_keys_sql, build_entity_rollup_sql, build_export_sql, + build_guardrail_spend_sql, build_key_page_sql, build_key_search_sql, build_model_top_keys_sql, @@ -36,6 +37,7 @@ from litellm.types.repositories.daily_activity import ( ExportRow, ExportType, GroupingSetsRow, + GuardrailSpendRow, KeyMetadataRow, KeyPage, KeySpendRow, @@ -100,6 +102,7 @@ class DailyActivityDatabase(Protocol): _GROUPING_ADAPTER: Final = TypeAdapter(tuple[GroupingSetsRow, ...]) _ENTITY_ADAPTER: Final = TypeAdapter(tuple[EntityRollupRow, ...]) +_GUARDRAIL_SPEND_ADAPTER: Final = TypeAdapter(tuple[GuardrailSpendRow, ...]) _KEY_SPEND_ADAPTER: Final = TypeAdapter(tuple[KeySpendRow, ...]) _KEY_PAGE_TOTAL_ADAPTER: Final[TypeAdapter[int]] = TypeAdapter(int) _USER_SPEND_ADAPTER: Final = TypeAdapter(tuple[UserSpendRow, ...]) @@ -191,6 +194,9 @@ class DailyActivityRepository: distinct_api_keys=distinct_api_keys, ) + async def guardrail_spend(self, scope: DailyActivityScope) -> tuple[GuardrailSpendRow, ...]: + return _GUARDRAIL_SPEND_ADAPTER.validate_python(await self._query(build_guardrail_spend_sql(scope))) + async def search_keys(self, scope: DailyActivityScope, *, search: str, limit: int) -> tuple[str, ...]: if not 1 <= limit <= constants.USAGE_KEY_SEARCH_MAX: raise ValueError(f"limit must be between 1 and {constants.USAGE_KEY_SEARCH_MAX}") diff --git a/litellm/repositories/daily_activity_sql.py b/litellm/repositories/daily_activity_sql.py index 3d27e7d937e..7ec9fa67a67 100644 --- a/litellm/repositories/daily_activity_sql.py +++ b/litellm/repositories/daily_activity_sql.py @@ -261,6 +261,25 @@ def build_entity_rollup_sql(scope: DailyActivityScope, *, api_key_limit: int) -> return SqlQuery(sql=sql, params=(*where_params, PTU_SENTINEL_API_KEY, api_key_limit)) +def build_guardrail_spend_sql(scope: DailyActivityScope) -> SqlQuery: + pg_table: Final = PRISMA_TO_PG_TABLE[scope.table] + where_clause, where_params = build_where_clause(scope) + sql: Final = f""" + SELECT g.date, g.guardrail_id, COALESCE(t.guardrail_name, g.guardrail_id) AS guardrail_name, g.api_key, + COALESCE(SUM(g.cost), 0)::float AS spend, + COALESCE(SUM(CASE WHEN g.cost IS NULL THEN g.units ELSE g.untracked_units END), 0)::bigint AS untracked_units + FROM "LiteLLM_DailyGuardrailUsageUnits" g + LEFT JOIN "LiteLLM_GuardrailsTable" t ON t.guardrail_id = g.guardrail_id + WHERE g.date >= $1 AND g.date <= $2 + AND EXISTS ( + SELECT 1 FROM "{pg_table}" + WHERE {where_clause} AND "{pg_table}".date = g.date AND "{pg_table}".api_key = g.api_key + ) + GROUP BY g.date, g.guardrail_id, t.guardrail_name, g.api_key + """ + return SqlQuery(sql=sql, params=where_params) + + def _key_spend_select() -> str: return """ COALESCE(SUM(spend), 0)::float AS spend, diff --git a/litellm/types/proxy/management_endpoints/common_daily_activity.py b/litellm/types/proxy/management_endpoints/common_daily_activity.py index 213c911cfdc..3f1a7980b27 100644 --- a/litellm/types/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/types/proxy/management_endpoints/common_daily_activity.py @@ -65,6 +65,34 @@ class MetricWithMetadata(MetricBase): api_key_breakdown: dict[str, KeyMetricWithMetadata] = Field(default_factory=dict) # api_key -> {metrics, metadata} +class GuardrailSpendMetrics(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + spend: float = 0.0 + untracked_units: int = 0 + + +class GuardrailKeyMetric(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + metrics: GuardrailSpendMetrics + metadata: KeyMetadata + + +class GuardrailMetadata(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + guardrail_name: str + + +class GuardrailMetric(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + metrics: GuardrailSpendMetrics + metadata: GuardrailMetadata + api_key_breakdown: Mapping[str, GuardrailKeyMetric] = Field(default_factory=dict) + + class BreakdownMetrics(LiteLLMBaseModel): """Breakdown of spend by different dimensions""" @@ -75,6 +103,7 @@ class BreakdownMetrics(LiteLLMBaseModel): endpoints: dict[str, MetricWithMetadata] = Field(default_factory=dict) # endpoint -> {metrics, metadata} api_keys: dict[str, KeyMetricWithMetadata] = Field(default_factory=dict) # api_key -> {metrics, metadata} entities: dict[str, MetricWithMetadata] = Field(default_factory=dict) # entity -> {metrics, metadata} + guardrails: Mapping[str, GuardrailMetric] = Field(default_factory=dict) # guardrail_id -> spend class DailySpendData(LiteLLMBaseModel): diff --git a/litellm/types/repositories/daily_activity.py b/litellm/types/repositories/daily_activity.py index 4111b83a891..b013ce43113 100644 --- a/litellm/types/repositories/daily_activity.py +++ b/litellm/types/repositories/daily_activity.py @@ -168,6 +168,16 @@ class EntityRollupRow(RollupMetricsRow): distinct_api_keys: int | None +@dataclass(frozen=True, slots=True) +class GuardrailSpendRow: + date: str + guardrail_id: str + guardrail_name: str + api_key: str + spend: float + untracked_units: int + + @dataclass(frozen=True, slots=True) class AggregatedRows: grouping_rows: tuple[GroupingSetsRow, ...] diff --git a/tests/integration/spend/fixtures/daily_activity_team.json b/tests/integration/spend/fixtures/daily_activity_team.json index 62377c07432..8f4956dd7fe 100644 --- a/tests/integration/spend/fixtures/daily_activity_team.json +++ b/tests/integration/spend/fixtures/daily_activity_team.json @@ -1 +1 @@ -{"results":[{"date":"2026-06-01","metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"breakdown":{"mcp_servers":{},"models":{"model-ptu":{"metrics":{"spend":1000.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":0,"completion_tokens":0,"cache_read_input_tokens":0,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":0,"successful_requests":0,"failed_requests":0,"api_requests":0,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{}},"model-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}},"model-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}},"model-popular":{"metrics":{"spend":270.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":9,"completion_tokens":6,"cache_read_input_tokens":3,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":15,"successful_requests":3,"failed_requests":0,"api_requests":3,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"model_groups":{"model-ptu":{"metrics":{"spend":1000.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":0,"completion_tokens":0,"cache_read_input_tokens":0,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":0,"successful_requests":0,"failed_requests":0,"api_requests":0,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{}},"model-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}},"model-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}},"model-popular":{"metrics":{"spend":270.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":9,"completion_tokens":6,"cache_read_input_tokens":3,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":15,"successful_requests":3,"failed_requests":0,"api_requests":3,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"providers":{"provider-a":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"endpoints":{"/v1/chat/completions":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"api_keys":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}},"entities":{"team-1":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{"team_alias":"Usage Team"},"api_key_breakdown":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}}}}}],"metadata":{"total_spend":1273.0,"total_flat_cost":0.0,"total_ptu_hours":0.0,"total_prompt_tokens":1014,"total_completion_tokens":10,"total_tokens":1024,"total_api_requests":5,"total_successful_requests":5,"total_failed_requests":0,"total_cache_read_input_tokens":6,"total_cache_creation_input_tokens":0,"total_compression_saved_tokens":0,"total_compression_savings_spend":0.0,"total_prompt_caching_savings_spend":0.0,"total_gateway_injected_caching_savings_spend":0.0,"total_autorouter_savings_spend":0.0,"total_response_time_ms":0,"total_timed_requests":0,"page":1,"total_pages":1,"has_more":false,"api_key_limit":100,"total_api_keys":5,"entity_total_api_keys":{"team-1":5}}} +{"results":[{"date":"2026-06-01","metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"breakdown":{"mcp_servers":{},"models":{"model-ptu":{"metrics":{"spend":1000.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":0,"completion_tokens":0,"cache_read_input_tokens":0,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":0,"successful_requests":0,"failed_requests":0,"api_requests":0,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{}},"model-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}},"model-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}},"model-popular":{"metrics":{"spend":270.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":9,"completion_tokens":6,"cache_read_input_tokens":3,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":15,"successful_requests":3,"failed_requests":0,"api_requests":3,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"model_groups":{"model-ptu":{"metrics":{"spend":1000.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":0,"completion_tokens":0,"cache_read_input_tokens":0,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":0,"successful_requests":0,"failed_requests":0,"api_requests":0,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{}},"model-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}},"model-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}},"model-popular":{"metrics":{"spend":270.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":9,"completion_tokens":6,"cache_read_input_tokens":3,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":15,"successful_requests":3,"failed_requests":0,"api_requests":3,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"providers":{"provider-a":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"endpoints":{"/v1/chat/completions":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"api_keys":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}},"entities":{"team-1":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{"team_alias":"Usage Team"},"api_key_breakdown":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}}},"guardrails":{}}}],"metadata":{"total_spend":1273.0,"total_flat_cost":0.0,"total_ptu_hours":0.0,"total_prompt_tokens":1014,"total_completion_tokens":10,"total_tokens":1024,"total_api_requests":5,"total_successful_requests":5,"total_failed_requests":0,"total_cache_read_input_tokens":6,"total_cache_creation_input_tokens":0,"total_compression_saved_tokens":0,"total_compression_savings_spend":0.0,"total_prompt_caching_savings_spend":0.0,"total_gateway_injected_caching_savings_spend":0.0,"total_autorouter_savings_spend":0.0,"total_response_time_ms":0,"total_timed_requests":0,"page":1,"total_pages":1,"has_more":false,"api_key_limit":100,"total_api_keys":5,"entity_total_api_keys":{"team-1":5}}} diff --git a/tests/integration/spend/fixtures/daily_activity_user.json b/tests/integration/spend/fixtures/daily_activity_user.json index c4302ee7045..ae05e661071 100644 --- a/tests/integration/spend/fixtures/daily_activity_user.json +++ b/tests/integration/spend/fixtures/daily_activity_user.json @@ -1 +1 @@ -{"results":[{"date":"2026-06-01","metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"breakdown":{"mcp_servers":{},"models":{"model-ptu":{"metrics":{"spend":1000.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":0,"completion_tokens":0,"cache_read_input_tokens":0,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":0,"successful_requests":0,"failed_requests":0,"api_requests":0,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{}},"model-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}},"model-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}},"model-popular":{"metrics":{"spend":270.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":9,"completion_tokens":6,"cache_read_input_tokens":3,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":15,"successful_requests":3,"failed_requests":0,"api_requests":3,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"model_groups":{"model-ptu":{"metrics":{"spend":1000.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":0,"completion_tokens":0,"cache_read_input_tokens":0,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":0,"successful_requests":0,"failed_requests":0,"api_requests":0,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{}},"model-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}},"model-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}},"model-popular":{"metrics":{"spend":270.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":9,"completion_tokens":6,"cache_read_input_tokens":3,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":15,"successful_requests":3,"failed_requests":0,"api_requests":3,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"providers":{"provider-a":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"endpoints":{"/v1/chat/completions":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"api_keys":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}},"entities":{"user-1":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}}}}}],"metadata":{"total_spend":1273.0,"total_flat_cost":0.0,"total_ptu_hours":0.0,"total_prompt_tokens":1014,"total_completion_tokens":10,"total_tokens":1024,"total_api_requests":5,"total_successful_requests":5,"total_failed_requests":0,"total_cache_read_input_tokens":6,"total_cache_creation_input_tokens":0,"total_compression_saved_tokens":0,"total_compression_savings_spend":0.0,"total_prompt_caching_savings_spend":0.0,"total_gateway_injected_caching_savings_spend":0.0,"total_autorouter_savings_spend":0.0,"total_response_time_ms":0,"total_timed_requests":0,"page":1,"total_pages":1,"has_more":false,"api_key_limit":100,"total_api_keys":5,"entity_total_api_keys":{"user-1":5}}} +{"results":[{"date":"2026-06-01","metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"breakdown":{"mcp_servers":{},"models":{"model-ptu":{"metrics":{"spend":1000.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":0,"completion_tokens":0,"cache_read_input_tokens":0,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":0,"successful_requests":0,"failed_requests":0,"api_requests":0,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{}},"model-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}},"model-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}},"model-popular":{"metrics":{"spend":270.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":9,"completion_tokens":6,"cache_read_input_tokens":3,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":15,"successful_requests":3,"failed_requests":0,"api_requests":3,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"model_groups":{"model-ptu":{"metrics":{"spend":1000.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":0,"completion_tokens":0,"cache_read_input_tokens":0,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":0,"successful_requests":0,"failed_requests":0,"api_requests":0,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{}},"model-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}},"model-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}},"model-popular":{"metrics":{"spend":270.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":9,"completion_tokens":6,"cache_read_input_tokens":3,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":15,"successful_requests":3,"failed_requests":0,"api_requests":3,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"providers":{"provider-a":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"endpoints":{"/v1/chat/completions":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"api_keys":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}},"entities":{"user-1":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}}},"guardrails":{}}}],"metadata":{"total_spend":1273.0,"total_flat_cost":0.0,"total_ptu_hours":0.0,"total_prompt_tokens":1014,"total_completion_tokens":10,"total_tokens":1024,"total_api_requests":5,"total_successful_requests":5,"total_failed_requests":0,"total_cache_read_input_tokens":6,"total_cache_creation_input_tokens":0,"total_compression_saved_tokens":0,"total_compression_savings_spend":0.0,"total_prompt_caching_savings_spend":0.0,"total_gateway_injected_caching_savings_spend":0.0,"total_autorouter_savings_spend":0.0,"total_response_time_ms":0,"total_timed_requests":0,"page":1,"total_pages":1,"has_more":false,"api_key_limit":100,"total_api_keys":5,"entity_total_api_keys":{"user-1":5}}} diff --git a/tests/integration/spend/golden/daily_activity_team_aggregated.json b/tests/integration/spend/golden/daily_activity_team_aggregated.json index 2956adda17a..91a5d33391b 100644 --- a/tests/integration/spend/golden/daily_activity_team_aggregated.json +++ b/tests/integration/spend/golden/daily_activity_team_aggregated.json @@ -26,7 +26,7 @@ }, "results": [ { - "breakdown": { + "breakdown": {"guardrails": {}, "api_keys": { "key-a": { "metadata": { diff --git a/tests/integration/spend/golden/daily_activity_team_paginated.json b/tests/integration/spend/golden/daily_activity_team_paginated.json index ca7239f8fa9..4d91b0da209 100644 --- a/tests/integration/spend/golden/daily_activity_team_paginated.json +++ b/tests/integration/spend/golden/daily_activity_team_paginated.json @@ -26,7 +26,7 @@ }, "results": [ { - "breakdown": { + "breakdown": {"guardrails": {}, "api_keys": { "key-a": { "metadata": { diff --git a/tests/integration/spend/golden/daily_activity_user_aggregated.json b/tests/integration/spend/golden/daily_activity_user_aggregated.json index 261ac22b2cb..185d333965a 100644 --- a/tests/integration/spend/golden/daily_activity_user_aggregated.json +++ b/tests/integration/spend/golden/daily_activity_user_aggregated.json @@ -26,7 +26,7 @@ }, "results": [ { - "breakdown": { + "breakdown": {"guardrails": {}, "api_keys": { "key-a": { "metadata": { diff --git a/tests/integration/spend/golden/daily_activity_user_paginated.json b/tests/integration/spend/golden/daily_activity_user_paginated.json index f1fb7ba50a4..00389289ebe 100644 --- a/tests/integration/spend/golden/daily_activity_user_paginated.json +++ b/tests/integration/spend/golden/daily_activity_user_paginated.json @@ -26,7 +26,7 @@ }, "results": [ { - "breakdown": { + "breakdown": {"guardrails": {}, "api_keys": { "key-a": { "metadata": { diff --git a/tests/integration/spend/test_daily_activity_aggregated_breakdowns.py b/tests/integration/spend/test_daily_activity_aggregated_breakdowns.py index 33baf9d0e2e..1c315a6e694 100644 --- a/tests/integration/spend/test_daily_activity_aggregated_breakdowns.py +++ b/tests/integration/spend/test_daily_activity_aggregated_breakdowns.py @@ -243,3 +243,183 @@ async def test_get_daily_activity_aggregated_model_group_rollups_fall_back_to_mo assert object_value(object_value(models["gpt-5"])["metrics"])["spend"] == 10.0 finally: _clean(day) + + +def _seed_guardrail_units(day: str, rows: Sequence[tuple[str, str, str, float | None, int, int]]) -> None: + for guardrail_id, api_key, usage_unit, cost, units, untracked_units in rows: + write_rows( + 'INSERT INTO "LiteLLM_DailyGuardrailUsageUnits" (guardrail_id, date, team_id, api_key, usage_unit,' + " units, cost, untracked_units, updated_at) VALUES (%s, %s, '', %s, %s, %s, %s, %s, now())", + ( + guardrail_id, + day, + api_key, + usage_unit, + str(units), + None if cost is None else str(cost), + str(untracked_units), + ), + ) + + +def _guardrail_spend(breakdown: dict[str, JsonValue]) -> dict[str, tuple[float, dict[str, float]]]: + return { + guardrail_id: ( + float(object_value(object_value(entry)["metrics"])["spend"]), + { + api_key: float(object_value(object_value(key_entry)["metrics"])["spend"]) + for api_key, key_entry in object_value(object_value(entry)["api_key_breakdown"]).items() + }, + ) + for guardrail_id, entry in object_value(breakdown["guardrails"]).items() + } + + +def _day_breakdown(gateway: Gateway, day: str, **params: str) -> dict[str, JsonValue]: + results: Final = _RESULTS.validate_python(_activity(gateway, day, **params)["results"]) + assert len(results) == 1 + return object_value(object_value(results[0])["breakdown"]) + + +@pytest.mark.asyncio +async def test_get_daily_activity_aggregated_breaks_down_guardrail_spend(gateway: Gateway) -> None: + day: Final = _unique_day() + _seed( + day, + [ + (_row_id(), "user-a", day, "key-a", "gpt-5", "", "bedrock", None, "/v1/chat/completions", 10, 1.0, 1, 1), + (_row_id(), "user-b", day, "key-b", "gpt-5", "", "bedrock", None, "/v1/chat/completions", 10, 2.0, 1, 1), + ], + ) + _seed_guardrail_units( + day, + [ + ("pii-filter", "key-a", "contentPolicyUnits", 0.10, 5, 2), + ("pii-filter", "key-a", "sensitiveInformationPolicyUnits", 0.05, 1, 0), + ("pii-filter", "key-b", "contentPolicyUnits", 0.20, 1, 0), + ("pii-filter", "key-b", "legacyUnits", None, 7, 0), + ("topic-filter", "key-a", "topicPolicyUnits", 0.01, 1, 0), + ("topic-filter", "key-b", "topicPolicyUnits", None, 1, 0), + ("pii-filter", "orphan-key", "contentPolicyUnits", 9.0, 1, 0), + ], + ) + try: + everyone: Final = _day_breakdown(gateway, day) + assert _guardrail_spend(everyone) == { + "pii-filter": (pytest.approx(0.35), {"key-a": pytest.approx(0.15), "key-b": pytest.approx(0.20)}), + "topic-filter": (pytest.approx(0.01), {"key-a": pytest.approx(0.01), "key-b": 0.0}), + } + api_keys: Final = object_value(everyone["api_keys"]) + pii_filter: Final = object_value(object_value(everyone["guardrails"])["pii-filter"]) + assert object_value(pii_filter["metrics"]) == {"spend": pytest.approx(0.35), "untracked_units": 9} + pii_api_keys: Final = object_value(pii_filter["api_key_breakdown"]) + assert object_value(object_value(pii_api_keys["key-a"])["metrics"]) == { + "spend": pytest.approx(0.15), + "untracked_units": 2, + } + assert object_value(object_value(pii_api_keys["key-b"])["metrics"]) == { + "spend": pytest.approx(0.20), + "untracked_units": 7, + } + + assert object_value(everyone["guardrails"])["topic-filter"] == { + "metrics": {"spend": pytest.approx(0.01), "untracked_units": 1}, + "metadata": {"guardrail_name": "topic-filter"}, + "api_key_breakdown": { + "key-a": { + "metrics": {"spend": pytest.approx(0.01), "untracked_units": 0}, + "metadata": object_value(api_keys["key-a"])["metadata"], + }, + "key-b": { + "metrics": {"spend": 0.0, "untracked_units": 1}, + "metadata": object_value(api_keys["key-b"])["metadata"], + }, + }, + } + assert object_value(object_value(object_value(everyone["models"])["gpt-5"])["metrics"])["spend"] == 3.0 + + assert _guardrail_spend(_day_breakdown(gateway, day, user_id="user-a")) == { + "pii-filter": (pytest.approx(0.15), {"key-a": pytest.approx(0.15)}), + "topic-filter": (pytest.approx(0.01), {"key-a": pytest.approx(0.01)}), + } + assert _guardrail_spend(_day_breakdown(gateway, day, api_key_limit="1")) == { + "pii-filter": (pytest.approx(0.35), {"key-b": pytest.approx(0.20)}), + "topic-filter": (pytest.approx(0.01), {"key-b": 0.0}), + } + limited_pii_filter: Final = object_value( + object_value(_day_breakdown(gateway, day, api_key_limit="1")["guardrails"])["pii-filter"] + ) + assert object_value(limited_pii_filter["metrics"])["untracked_units"] == 9 + assert set(object_value(limited_pii_filter["api_key_breakdown"])) == {"key-b"} + assert _guardrail_spend(_day_breakdown(gateway, day, model="gpt-5")) == {} + assert _guardrail_spend(_day_breakdown(gateway, day, model="")) == _guardrail_spend(everyone) + finally: + _clean(day) + write_rows('DELETE FROM "LiteLLM_DailyGuardrailUsageUnits" WHERE date = %s', (day,)) + + +def _entity_breakdown(gateway: Gateway, url: str, day: str, **params: str) -> dict[str, JsonValue]: + response: Final = gateway.request("GET", url, params={"start_date": day, "end_date": day, **params}) + assert response.status_code == 200, response.text + results: Final = _RESULTS.validate_python(object_value(response.json())["results"]) + assert len(results) == 1 + return object_value(object_value(results[0])["breakdown"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("url", "table", "entity_column", "filter_param"), + [ + ("/team/daily/activity/aggregated", "LiteLLM_DailyTeamSpend", "team_id", "team_ids"), + ( + "/organization/daily/activity/aggregated", + "LiteLLM_DailyOrganizationSpend", + "organization_id", + "organization_ids", + ), + ], +) +async def test_team_and_organization_activity_break_down_guardrail_spend_by_name( + gateway: Gateway, url: str, table: str, entity_column: str, filter_param: str +) -> None: + day: Final = _unique_day() + db_guardrail_id: Final = str(uuid.uuid4()) + db_guardrail_name: Final = f"pii-filter-{db_guardrail_id[:8]}" + for entity_id, api_key, spend in (("entity-a", "key-a", 1.0), ("entity-b", "key-b", 2.0)): + write_rows( + f'INSERT INTO "{table}" (id, {entity_column}, date, api_key, model, spend, api_requests,' + " successful_requests, updated_at) VALUES (%s, %s, %s, %s, 'gpt-5', %s, 1, 1, now())", + (_row_id(), entity_id, day, api_key, str(spend)), + ) + write_rows( + 'INSERT INTO "LiteLLM_GuardrailsTable" (guardrail_id, guardrail_name, litellm_params, updated_at)' + " VALUES (%s, %s, '{}', now())", + (db_guardrail_id, db_guardrail_name), + ) + _seed_guardrail_units( + day, + [ + (db_guardrail_id, "key-a", "contentPolicyUnits", 0.10, 1, 0), + (db_guardrail_id, "key-b", "contentPolicyUnits", 0.20, 1, 0), + ("config-topic-filter", "key-a", "topicPolicyUnits", 0.05, 1, 0), + (db_guardrail_id, "orphan-key", "contentPolicyUnits", 9.0, 1, 0), + ], + ) + try: + everyone: Final = _entity_breakdown(gateway, url, day) + assert _guardrail_spend(everyone) == { + db_guardrail_id: (pytest.approx(0.30), {"key-a": pytest.approx(0.10), "key-b": pytest.approx(0.20)}), + "config-topic-filter": (pytest.approx(0.05), {"key-a": pytest.approx(0.05)}), + } + guardrails: Final = object_value(everyone["guardrails"]) + assert object_value(guardrails[db_guardrail_id])["metadata"] == {"guardrail_name": db_guardrail_name} + assert object_value(guardrails["config-topic-filter"])["metadata"] == {"guardrail_name": "config-topic-filter"} + + assert _guardrail_spend(_entity_breakdown(gateway, url, day, **{filter_param: "entity-a"})) == { + db_guardrail_id: (pytest.approx(0.10), {"key-a": pytest.approx(0.10)}), + "config-topic-filter": (pytest.approx(0.05), {"key-a": pytest.approx(0.05)}), + } + finally: + write_rows(f'DELETE FROM "{table}" WHERE date = %s', (day,)) + write_rows('DELETE FROM "LiteLLM_DailyGuardrailUsageUnits" WHERE date = %s', (day,)) + write_rows('DELETE FROM "LiteLLM_GuardrailsTable" WHERE guardrail_id = %s', (db_guardrail_id,)) diff --git a/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py b/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py index 8fbead8c01c..2a273034a97 100644 --- a/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py +++ b/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py @@ -248,6 +248,7 @@ class _FakeRepository: def __init__(self, rows: tuple[_Activity, ...]) -> None: self._rows: Final = rows self.aggregated = AsyncMock(side_effect=self._aggregated) + self.guardrail_spend = AsyncMock(return_value=()) self.key_page = AsyncMock(side_effect=self._key_page) self.key_page_call: tuple[DailyActivityScope, int, int] | None = None self.search_keys = AsyncMock(side_effect=self._search_keys) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 70406105193..72cc8cdc94e 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -27721,6 +27721,10 @@ export interface components { entities?: { [key: string]: components["schemas"]["MetricWithMetadata"]; }; + /** Guardrails */ + guardrails?: { + [key: string]: components["schemas"]["GuardrailMetric"]; + }; /** Mcp Servers */ mcp_servers?: { [key: string]: components["schemas"]["MetricWithMetadata"]; @@ -33236,6 +33240,38 @@ export interface components { /** Updated At */ updated_at?: string | null; }; + /** GuardrailKeyMetric */ + GuardrailKeyMetric: { + metadata: components["schemas"]["KeyMetadata"]; + metrics: components["schemas"]["GuardrailSpendMetrics"]; + }; + /** GuardrailMetadata */ + GuardrailMetadata: { + /** Guardrail Name */ + guardrail_name: string; + }; + /** GuardrailMetric */ + GuardrailMetric: { + /** Api Key Breakdown */ + api_key_breakdown?: { + [key: string]: components["schemas"]["GuardrailKeyMetric"]; + }; + metadata: components["schemas"]["GuardrailMetadata"]; + metrics: components["schemas"]["GuardrailSpendMetrics"]; + }; + /** GuardrailSpendMetrics */ + GuardrailSpendMetrics: { + /** + * Spend + * @default 0 + */ + spend: number; + /** + * Untracked Units + * @default 0 + */ + untracked_units: number; + }; /** GuardrailSubmissionItem */ GuardrailSubmissionItem: { /** Created At */ From cf1006a3687de717f247e846cefedf823f2d8203 Mon Sep 17 00:00:00 2001 From: mrinal-berri Date: Wed, 28 Oct 2026 15:35:33 -0700 Subject: [PATCH 101/128] feat(ui): search and multi-select tools in the MCP toolset editor (#48601) * feat(ui): add per-server MCP tool catalog query hook One query per server keyed by server_id, so the toolset editor can load tools for many servers in parallel and a failing server surfaces as an error on its own row instead of an empty list. (LIT-10012) * feat(ui): add toolset tool search and bulk selection helpers Pure functions for the toolset editor: filter a server's tools by a search query, select/deselect many tools without duplicates or dropping other selections, and compute tri-state selection. (LIT-10012) * feat(ui): search tools and bulk select in the toolset editor The toolset editor search now matches tool names and descriptions as well as server names. Typing fetches every server's tools in parallel through the per-server query hook, auto-expands matching servers and hides servers that loaded with no match. Tool rows become checkboxes, each server header gets a tri-state select-all over its visible tools, and a bar offers select/deselect all matching tools. Bulk actions only touch visible tools, so selections on other servers are kept on save. (LIT-10012) * feat(ui): clear all and flag unconfirmed tools in the toolset panel Adds a Clear all action to the Your Toolset panel and labels selected tools whose server is gone, unreachable, or no longer lists the tool. Those tools stay selected and are still saved, so an outage never silently drops part of a toolset. (LIT-10012) * test(ui): cover toolset editor tool search and bulk selection Integration tests that drive the real toolset modal against stubbed networking and assert the saved tools payload: cross-server search by tool name and description, select/deselect all matching, server select-all, single-row toggles, Clear all, and that tools on unreachable or deleted servers survive a save. (LIT-10012) * fix(ui): pointer cursor on the toolset server select-all checkbox Base UI renders the checkbox as a span, which the global button cursor rule does not cover, so the server header select-all showed the default arrow. Use a pointer, and not-allowed while it is disabled. (LIT-10012) * fix(ui): stop toolset editor background fan-out and keep tools on failed refetch The toolset modal stays mounted when closed, so a finished search kept every server's tool query enabled and refocusing the window refetched all of them once stale. Fetch only while the editor is open, require both the live and debounced search before fanning out, and disable refetch on window focus for these queries. A failed refetch now keeps the tools that already loaded instead of flagging the server unreachable, and tools without a name are dropped. Trim comments that restated the code. (LIT-10012) * fix(ui): clear toolset search on close and only keep non-empty tool lists on failed refetch Reopening the editor after a search rendered once with the old query, which fetched every stale server; clear the search whenever the editor closes. A failed refetch now keeps the previous list only when it has tools, so a server that loaded none and then fails reports the error. (LIT-10012) * fix(ui): reset toolset editor state on close and avoid mutating selection sets Closing the editor now also clears expanded servers, so reopening it cannot refetch them for one render. Build the selection and expanded-server sets without in-place mutation, drop a comment that restated the code, and use screen queries in the new tests to stay within the prefer-screen-queries lint budget. (LIT-10012) --- ui/litellm-dashboard/eslint-suppressions.json | 3 - .../mcpServers/useMCPServerTools.test.ts | 118 ++++ .../hooks/mcpServers/useMCPServerTools.ts | 70 +++ .../MCPToolsetsTab.integration.test.tsx | 219 +++++++- .../_components/MCPToolsetsTab.tsx | 503 ++++++++++++------ .../_components/toolsetSelection.test.ts | 90 ++++ .../_components/toolsetSelection.ts | 63 +++ 7 files changed, 890 insertions(+), 176 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerTools.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerTools.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/toolsetSelection.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/toolsetSelection.ts diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index e67c1105d86..9fce93782b2 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -498,9 +498,6 @@ } }, "src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx": { - "no-nested-ternary": { - "count": 1 - }, "react-hooks/set-state-in-effect": { "count": 1 } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerTools.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerTools.test.ts new file mode 100644 index 00000000000..3ef8e88b2c9 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerTools.test.ts @@ -0,0 +1,118 @@ +/* @vitest-environment jsdom */ +import React from "react"; +import { act, renderHook, waitFor } from "@testing-library/react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { mcpServerToolKeys, useMCPServerTools } from "./useMCPServerTools"; +import * as networking from "@/components/networking"; + +vi.mock("@/components/networking", () => ({ + listMCPTools: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: vi.fn(() => ({ accessToken: "test-token-123" })), +})); + +const wrapper = ({ children }: { children: React.ReactNode }) => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: 0 } } }); + return React.createElement(QueryClientProvider, { client: queryClient }, children); +}; + +const toolsFor = (serverId: string) => ({ + tools: [{ name: `${serverId}_tool`, description: `from ${serverId}` }], + error: null, +}); + +describe("useMCPServerTools", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("fetches only the enabled servers and leaves the rest idle", async () => { + vi.mocked(networking.listMCPTools).mockImplementation(async (_token, serverId) => toolsFor(serverId)); + + const { result } = renderHook(() => useMCPServerTools(["a", "b"], new Set(["a"])), { wrapper }); + + await waitFor(() => expect(result.current.get("a")?.tools).toEqual([{ name: "a_tool", description: "from a" }])); + expect(result.current.get("b")).toEqual({ tools: [], status: "idle" }); + expect(networking.listMCPTools).toHaveBeenCalledTimes(1); + expect(networking.listMCPTools).toHaveBeenCalledWith("test-token-123", "a"); + }); + + it("isolates a failing server from the others", async () => { + vi.mocked(networking.listMCPTools).mockImplementation(async (_token, serverId) => + serverId === "bad" ? { tools: [], error: "http_500", message: "upstream down" } : toolsFor(serverId), + ); + + const { result } = renderHook(() => useMCPServerTools(["good", "bad"], new Set(["good", "bad"])), { wrapper }); + + await waitFor(() => expect(result.current.get("bad")).toEqual({ tools: [], status: "error" })); + expect(result.current.get("good")).toEqual({ + tools: [{ name: "good_tool", description: "from good" }], + status: "success", + }); + }); + + it("reports a server that is being fetched as loading", () => { + vi.mocked(networking.listMCPTools).mockReturnValue(new Promise(() => {})); + + const { result } = renderHook(() => useMCPServerTools(["a"], new Set(["a"])), { wrapper }); + + expect(result.current.get("a")).toEqual({ tools: [], status: "loading" }); + }); + + it("normalizes a bare array response, tools without descriptions, and drops tools without a name", async () => { + vi.mocked(networking.listMCPTools).mockResolvedValue([ + { tool_name: "legacy" }, + "plain", + { description: "nameless" }, + ]); + + const { result } = renderHook(() => useMCPServerTools(["a"], new Set(["a"])), { wrapper }); + + await waitFor(() => + expect(result.current.get("a")?.tools).toEqual([ + { name: "legacy", description: "" }, + { name: "plain", description: "" }, + ]), + ); + }); + + it("reports an error when a refetch fails for a server that had loaded no tools", async () => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + const sharedWrapper = ({ children }: { children: React.ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + vi.mocked(networking.listMCPTools) + .mockResolvedValueOnce({ tools: [], error: null }) + .mockResolvedValueOnce({ tools: [], error: "http_502", message: "upstream down" }); + + const { result } = renderHook(() => useMCPServerTools(["a"], new Set(["a"])), { wrapper: sharedWrapper }); + await waitFor(() => expect(result.current.get("a")?.status).toBe("success")); + + await act(() => queryClient.refetchQueries({ queryKey: mcpServerToolKeys.detail("a") })); + + await waitFor(() => expect(result.current.get("a")).toEqual({ tools: [], status: "error" })); + }); + + it("keeps showing the loaded tools when a later refetch fails", async () => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + const sharedWrapper = ({ children }: { children: React.ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + vi.mocked(networking.listMCPTools) + .mockResolvedValueOnce(toolsFor("a")) + .mockResolvedValueOnce({ tools: [], error: "http_502", message: "upstream down" }); + + const { result } = renderHook(() => useMCPServerTools(["a"], new Set(["a"])), { wrapper: sharedWrapper }); + await waitFor(() => expect(result.current.get("a")?.status).toBe("success")); + + await act(() => queryClient.refetchQueries({ queryKey: mcpServerToolKeys.detail("a") })); + + await waitFor(() => expect(queryClient.getQueryState(mcpServerToolKeys.detail("a"))?.status).toBe("error")); + expect(networking.listMCPTools).toHaveBeenCalledTimes(2); + expect(result.current.get("a")).toEqual({ + tools: [{ name: "a_tool", description: "from a" }], + status: "success", + }); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerTools.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerTools.ts new file mode 100644 index 00000000000..bc4a6696a6e --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerTools.ts @@ -0,0 +1,70 @@ +import { useQueries, type QueryObserverResult } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; +import { listMCPTools } from "@/components/networking"; +import useAuthorized from "../useAuthorized"; + +export const mcpServerToolKeys = createQueryKeys("mcpServerTools"); + +export interface MCPServerToolEntry { + name: string; + description: string; +} + +export type MCPServerToolsStatus = "idle" | "loading" | "error" | "success"; + +export interface MCPServerToolsState { + tools: MCPServerToolEntry[]; + status: MCPServerToolsStatus; +} + +type RawTool = string | { name?: string; tool_name?: string; description?: string }; + +const toEntry = (tool: RawTool): MCPServerToolEntry => + typeof tool === "string" + ? { name: tool, description: "" } + : { name: tool.name ?? tool.tool_name ?? "", description: tool.description ?? "" }; + +// listMCPTools never throws: it reports failures as `{ tools: [], error }`. +// Throw here so a broken server surfaces as a query error instead of an empty tool list. +export const fetchServerTools = async (accessToken: string, serverId: string): Promise => { + const result = await listMCPTools(accessToken, serverId); + if (!Array.isArray(result) && result?.error) { + throw new Error(result.message || result.error); + } + const toolList: RawTool[] = Array.isArray(result) ? result : result?.tools ?? []; + return toolList.map(toEntry).filter((tool) => tool.name !== ""); +}; + +const toStatus = ({ + status, + fetchStatus, + data, +}: Pick, "status" | "fetchStatus" | "data">): MCPServerToolsStatus => { + if (status === "error" && data !== undefined && data.length > 0) return "success"; + if (status !== "pending") return status; + return fetchStatus === "idle" ? "idle" : "loading"; +}; + +export const useMCPServerTools = ( + serverIds: readonly string[], + enabledServerIds: ReadonlySet, +): ReadonlyMap => { + const { accessToken } = useAuthorized(); + return useQueries({ + queries: serverIds.map((serverId) => ({ + queryKey: mcpServerToolKeys.detail(serverId), + queryFn: () => fetchServerTools(accessToken!, serverId), + enabled: !!accessToken && enabledServerIds.has(serverId), + staleTime: 5 * 60 * 1000, + retry: false, + refetchOnWindowFocus: false, + })), + combine: (results) => + new Map( + results.map((result, idx): [string, MCPServerToolsState] => [ + serverIds[idx], + { tools: result.data ?? [], status: toStatus(result) }, + ]), + ), + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.integration.test.tsx index 81462c6bf9b..ba9f2bfa839 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.integration.test.tsx @@ -1,5 +1,5 @@ import React from "react"; -import { fireEvent, render, screen, waitFor, within } from "@testing-library/react"; +import { act, fireEvent, render, screen, waitFor, within } from "@testing-library/react"; import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event"; import { describe, it, expect, vi, beforeEach } from "vitest"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; @@ -7,6 +7,7 @@ import { MCPToolsetsTab } from "./MCPToolsetsTab"; import * as networking from "@/components/networking"; import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; +import { mcpServerToolKeys } from "@/app/(dashboard)/hooks/mcpServers/useMCPServerTools"; import { MCPToolset } from "@/components/mcp_tools/types"; vi.mock("@/components/networking", () => ({ @@ -20,6 +21,7 @@ vi.mock("@/components/networking", () => ({ vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPToolsets", () => ({ useMCPToolsets: vi.fn() })); vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPServers", () => ({ useMCPServers: vi.fn() })); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: () => ({ accessToken: "sk-test" }) })); const setup = () => userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); @@ -197,6 +199,221 @@ describe("MCPToolsetsTab create/edit toolset form", () => { }); }); +describe("MCPToolsetsTab tool search and bulk selection", () => { + const servers = [ + { server_id: "gh", alias: "github", server_name: "GitHub" }, + { server_id: "lin", alias: "linear", server_name: "Linear" }, + { server_id: "broken", alias: "flaky", server_name: "Flaky" }, + ]; + const catalogs: Record = { + gh: [ + { name: "create_issue", description: "Open a GitHub issue" }, + { name: "list_pulls", description: "List pull requests" }, + ], + lin: [ + { name: "create_ticket", description: "Create a Linear issue" }, + { name: "list_projects", description: "List projects" }, + ], + }; + + const renderWithServers = (toolsets: MCPToolset[] = []) => { + vi.mocked(useMCPToolsets).mockReturnValue({ + data: toolsets, + isLoading: false, + } as unknown as ReturnType); + vi.mocked(useMCPServers).mockReturnValue({ data: servers, isSuccess: true } as unknown as ReturnType< + typeof useMCPServers + >); + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: 0 } } }); + render( + + + , + ); + return queryClient; + }; + + const editToolset = (tools: MCPToolset["tools"]) => + ({ toolset_id: "ts-1", toolset_name: "triage", description: "", tools }) as MCPToolset; + + const expectedUpdate = (tools: MCPToolset["tools"]) => ({ + toolset_id: "ts-1", + toolset_name: "triage", + description: "", + tools, + }); + + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(networking.listMCPTools).mockImplementation(async (_token, serverId) => + serverId in catalogs + ? { tools: catalogs[serverId], error: null } + : { tools: [], error: "http_502", message: "upstream down" }, + ); + vi.mocked(networking.createMCPToolset).mockResolvedValue( + {} as Awaited>, + ); + vi.mocked(networking.updateMCPToolset).mockResolvedValue( + {} as Awaited>, + ); + }); + + it("searches tool names and descriptions across collapsed servers and selects every match", async () => { + const user = setup(); + renderWithServers(); + + await openCreate(user); + fireEvent.change(screen.getByPlaceholderText("e.g. github-linear-tools"), { target: { value: "issues" } }); + await user.type(screen.getByRole("textbox", { name: "Search tools or servers" }), "issue"); + + expect(await screen.findByRole("checkbox", { name: /^create_ticket/ })).toBeInTheDocument(); + expect(screen.getByRole("checkbox", { name: /^create_issue/ })).toBeInTheDocument(); + expect(screen.queryByRole("checkbox", { name: /^list_pulls/ })).not.toBeInTheDocument(); + expect(await screen.findByText("2 matching tools")).toBeInTheDocument(); + + await user.click(screen.getByRole("button", { name: "Select all matching" })); + await user.click(screen.getByRole("button", { name: "Create Toolset" })); + + await waitFor(() => { + expect(networking.createMCPToolset).toHaveBeenCalledWith("sk-test", { + toolset_name: "issues", + description: "", + tools: [ + { server_id: "gh", tool_name: "create_issue" }, + { server_id: "lin", tool_name: "create_ticket" }, + ], + }); + }); + }); + + it("deselecting all matches keeps selected tools the search hides", async () => { + const user = setup(); + renderWithServers([ + editToolset([ + { server_id: "gh", tool_name: "list_pulls" }, + { server_id: "gh", tool_name: "create_issue" }, + { server_id: "lin", tool_name: "create_ticket" }, + ]), + ]); + + await openEditFor(user); + await user.type(screen.getByRole("textbox", { name: "Search tools or servers" }), "create"); + expect(await screen.findByRole("checkbox", { name: /^create_ticket/ })).toBeChecked(); + + await user.click(screen.getByRole("button", { name: "Deselect all matching" })); + await user.click(screen.getByRole("button", { name: "Save Changes" })); + + await waitFor(() => { + expect(networking.updateMCPToolset).toHaveBeenCalledWith( + "sk-test", + expectedUpdate([{ server_id: "gh", tool_name: "list_pulls" }]), + ); + }); + }); + + it("selects all of one server's tools from its header and toggles single rows, fetching only that server", async () => { + const user = setup(); + renderWithServers(); + + await openCreate(user); + fireEvent.change(screen.getByPlaceholderText("e.g. github-linear-tools"), { target: { value: "gh-only" } }); + await user.click(screen.getByRole("button", { name: /^github/ })); + expect(await screen.findByRole("checkbox", { name: /^list_pulls/ })).toBeInTheDocument(); + + await user.click(screen.getByRole("checkbox", { name: "Select all tools on github" })); + expect(screen.getByRole("checkbox", { name: /^list_pulls/ })).toBeChecked(); + // Clicking anywhere on a tool row toggles it exactly once. + await user.click(screen.getByText("List pull requests")); + expect(screen.getByRole("checkbox", { name: /^list_pulls/ })).not.toBeChecked(); + await user.click(screen.getByRole("button", { name: "Create Toolset" })); + + await waitFor(() => { + expect(networking.createMCPToolset).toHaveBeenCalledWith("sk-test", { + toolset_name: "gh-only", + description: "", + tools: [{ server_id: "gh", tool_name: "create_issue" }], + }); + }); + expect(networking.listMCPTools).toHaveBeenCalledTimes(1); + }); + + it("keeps and flags tools whose server is unreachable or gone when saving", async () => { + const user = setup(); + const tools = [ + { server_id: "broken", tool_name: "flaky_tool" }, + { server_id: "deleted", tool_name: "old_tool" }, + { server_id: "gh", tool_name: "create_issue" }, + ]; + renderWithServers([editToolset(tools)]); + + await openEditFor(user); + await user.type(screen.getByRole("textbox", { name: "Search tools or servers" }), "create"); + + expect(await screen.findByText("Couldn't load tools for this server.")).toBeInTheDocument(); + expect(await screen.findByText("server unreachable")).toBeInTheDocument(); + expect(screen.getByText("server not found")).toBeInTheDocument(); + + await user.click(screen.getByRole("button", { name: "Save Changes" })); + + await waitFor(() => { + expect(networking.updateMCPToolset).toHaveBeenCalledWith("sk-test", expectedUpdate(tools)); + }); + }); + + it("stops fetching server tools once the editor is closed after a search", async () => { + const user = setup(); + const queryClient = renderWithServers(); + + await openCreate(user); + await user.type(screen.getByRole("textbox", { name: "Search tools or servers" }), "issue"); + expect(await screen.findByRole("checkbox", { name: /^create_ticket/ })).toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Cancel" })); + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + const callsWhileOpen = vi.mocked(networking.listMCPTools).mock.calls.length; + + await act(() => queryClient.invalidateQueries({ queryKey: mcpServerToolKeys.all })); + await new Promise((resolve) => setTimeout(resolve, 50)); + expect(networking.listMCPTools).toHaveBeenCalledTimes(callsWhileOpen); + + await openCreate(user); + expect(screen.getByRole("textbox", { name: "Search tools or servers" })).toHaveValue(""); + await new Promise((resolve) => setTimeout(resolve, 50)); + expect(networking.listMCPTools).toHaveBeenCalledTimes(callsWhileOpen); + }); + + it("does not refetch a previously expanded server when the editor is reopened", async () => { + const user = setup(); + const queryClient = renderWithServers(); + + await openCreate(user); + await user.click(screen.getByRole("button", { name: /^github/ })); + expect(await screen.findByRole("checkbox", { name: /^list_pulls/ })).toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Cancel" })); + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + await act(() => queryClient.invalidateQueries({ queryKey: mcpServerToolKeys.all })); + + await openCreate(user); + await new Promise((resolve) => setTimeout(resolve, 50)); + + expect(networking.listMCPTools).toHaveBeenCalledTimes(1); + expect(screen.queryByRole("checkbox", { name: /^list_pulls/ })).not.toBeInTheDocument(); + }); + + it("clears the whole toolset with Clear all", async () => { + const user = setup(); + renderWithServers([editToolset([{ server_id: "gh", tool_name: "create_issue" }])]); + + await openEditFor(user); + await user.click(screen.getByRole("button", { name: "Clear all" })); + expect(screen.getByText("No tools added yet")).toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Save Changes" })); + + await waitFor(() => { + expect(networking.updateMCPToolset).toHaveBeenCalledWith("sk-test", expectedUpdate([])); + }); + }); +}); + describe("MCPToolsetsTab usage guide", () => { it("shows the toolset URL on the gateway host", () => { renderTab(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx index ff0e8f7b6a9..dc1477b79d3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx @@ -1,28 +1,39 @@ -import React, { useState, useCallback } from "react"; +import React, { useState } from "react"; import { z } from "zod"; import { SortingState } from "@tanstack/react-table"; import { Inbox, Plus, X } from "lucide-react"; import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; import { useQueryClient } from "@tanstack/react-query"; -import { DataTable } from "@/components/shared/DataTable"; +import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; import { - createMCPToolset, - updateMCPToolset, - deleteMCPToolset, - listMCPTools, - getGatewayBaseUrl, -} from "@/components/networking"; -import { MCPToolset, MCPToolsetTool } from "@/components/mcp_tools/types"; + useMCPServerTools, + type MCPServerToolEntry, + type MCPServerToolsState, +} from "@/app/(dashboard)/hooks/mcpServers/useMCPServerTools"; +import { DataTable } from "@/components/shared/DataTable"; +import { createMCPToolset, updateMCPToolset, deleteMCPToolset, getGatewayBaseUrl } from "@/components/networking"; +import { MCPServer, MCPToolset, MCPToolsetTool } from "@/components/mcp_tools/types"; import { FieldGroup } from "@/components/ui/field"; import { FormField } from "@/components/shared/form/FormField"; import { Button } from "@/components/ui/button"; +import { Checkbox } from "@/components/ui/checkbox"; import { Input } from "@/components/ui/input"; import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { useZodForm } from "@/lib/forms/useZodForm"; import { toast } from "@/lib/toast"; import { displayToolName, getMCPToolsetTableColumns } from "./MCPToolsetTableColumns"; +import { + deselectTools, + filterServerTools, + isToolSelected, + matchesServerName, + selectTools, + selectionState, + toggleTool, + type SelectionState, +} from "./toolsetSelection"; import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; interface MCPToolsetsTabProps { @@ -37,110 +48,191 @@ const toolsetSchema = z.object({ type ToolsetFormValues = z.infer; -interface MCPToolListProps { +const SEARCH_DEBOUNCE_WAIT_MS = 300; + +const IDLE_SERVER_TOOLS: MCPServerToolsState = { tools: [], status: "idle" }; + +const NO_SERVERS: ReadonlySet = new Set(); + +// The modal stays mounted while closed, so fetch only while it is open; otherwise a finished +// search keeps every server's query live and refetches all of them in the background. +const serversToFetch = ( + open: boolean, + searching: boolean, + serverIds: readonly string[], + expandedServerIds: ReadonlySet, +): ReadonlySet => { + if (!open) return NO_SERVERS; + return searching ? new Set(serverIds) : expandedServerIds; +}; + +const checkboxStateProps = (state: SelectionState) => ({ + checked: state === "all", + indeterminate: state === "some", +}); + +interface ServerView { + server: MCPServer; + serverName: string; + state: MCPServerToolsState; + visibleTools: readonly MCPServerToolEntry[]; +} + +const isPending = (state: MCPServerToolsState) => state.status === "idle" || state.status === "loading"; + +// While searching, a server stays listed until its tools load and turn out not to match. +const isHiddenBySearch = ( + search: string, + serverNames: readonly string[], + state: MCPServerToolsState, + visibleTools: readonly MCPServerToolEntry[], +) => { + if (!search.trim() || state.status !== "success") return false; + return visibleTools.length === 0 && !matchesServerName(search, serverNames); +}; + +const buildServerViews = ( + servers: readonly MCPServer[], + toolsByServer: ReadonlyMap, + search: string, +): ServerView[] => + servers.flatMap((server) => { + const state = toolsByServer.get(server.server_id) ?? IDLE_SERVER_TOOLS; + const names = [server.alias ?? "", server.server_name ?? ""]; + const visibleTools = filterServerTools(search, names, state.tools); + if (isHiddenBySearch(search, names, state, visibleTools)) return []; + return [{ server, serverName: server.alias || server.server_name || server.server_id, state, visibleTools }]; + }); + +const summarizeSearch = (serverViews: readonly ServerView[]) => ({ + matchingTools: serverViews.flatMap(({ server, visibleTools }) => + visibleTools.map((tool) => ({ server_id: server.server_id, tool_name: tool.name })), + ), + pendingServerCount: serverViews.filter(({ state }) => isPending(state)).length, +}); + +// Why a selected tool cannot be confirmed against its server. The tool stays selected and is +// still saved; the note only keeps it from looking lost. +const selectedToolIssue = ( + tool: MCPToolsetTool, + serverNameById: ReadonlyMap, + toolsByServer: ReadonlyMap, +): string | undefined => { + if (!serverNameById.has(tool.server_id)) return "server not found"; + const state = toolsByServer.get(tool.server_id); + if (state?.status === "error") return "server unreachable"; + if (state?.status === "success" && !state.tools.some((t) => t.name === tool.tool_name)) return "tool not found"; + return undefined; +}; + +interface ServerToolGroupProps { serverId: string; serverName: string; - accessToken: string | null; + state: MCPServerToolsState; + visibleTools: readonly MCPServerToolEntry[]; + expanded: boolean; + searchActive: boolean; selectedTools: MCPToolsetTool[]; - onToggle: (tool: MCPToolsetTool) => void; + onToggleExpanded: () => void; + onToggleTool: (tool: MCPToolsetTool) => void; + onSetToolsSelected: (tools: MCPToolsetTool[], selected: boolean) => void; } -interface ToolEntry { - name: string; - description?: string; -} +function ServerToolGroup({ + serverId, + serverName, + state, + visibleTools, + expanded, + searchActive, + selectedTools, + onToggleExpanded, + onToggleTool, + onSetToolsSelected, +}: ServerToolGroupProps) { + const selectedCount = selectedTools.filter((t) => t.server_id === serverId).length; + const visibleRefs = visibleTools.map((tool) => ({ server_id: serverId, tool_name: tool.name })); + const groupState = selectionState(selectedTools, visibleRefs); -function MCPToolList({ serverId, serverName, accessToken, selectedTools, onToggle }: MCPToolListProps) { - const [tools, setTools] = useState([]); - const [loading, setLoading] = useState(false); - const [expanded, setExpanded] = useState(false); - - const selectedSet = new Set(selectedTools.filter((t) => t.server_id === serverId).map((t) => t.tool_name)); - - const fetchTools = useCallback(async () => { - if (!accessToken || tools.length > 0) return; - setLoading(true); - try { - const result = await listMCPTools(accessToken, serverId); - const toolList = Array.isArray(result) ? result : result?.tools ?? []; - setTools(toolList.map((t: any) => ({ name: t.name ?? t.tool_name ?? t, description: t.description ?? "" }))); - } catch { - setTools([]); - } finally { - setLoading(false); + const renderBody = () => { + if (isPending(state)) { + return ( +
+ +
+ ); } - }, [accessToken, serverId, tools.length]); - - const handleToggle = () => { - if (!expanded) fetchTools(); - setExpanded(!expanded); + if (state.status === "error") { + return

Couldn't load tools for this server.

; + } + if (visibleTools.length === 0) { + return ( +

+ {searchActive ? "No matching tools on this server." : "No tools found for this server."} +

+ ); + } + return ( +
+ {visibleRefs.map((ref, idx) => { + const tool = visibleTools[idx]; + const selected = isToolSelected(selectedTools, ref); + return ( + + ); + })} +
+ ); }; return (
- - {expanded && ( -
- {loading ? ( -
- -
- ) : tools.length === 0 ? ( -

No tools found for this server.

- ) : ( -
- {tools.map((tool) => { - const selected = selectedSet.has(tool.name); - return ( - - ); - })} -
- )} -
- )} +
+ onSetToolsSelected(visibleRefs, checked)} + /> + +
+ {expanded &&
{renderBody()}
}
); } @@ -149,21 +241,29 @@ interface CreateToolsetModalProps { open: boolean; onClose: () => void; onSave: (name: string, description: string | undefined, tools: MCPToolsetTool[]) => Promise; - accessToken: string | null; initialToolset?: MCPToolset; } -function CreateToolsetModal({ open, onClose, onSave, accessToken, initialToolset }: CreateToolsetModalProps) { - const form = useZodForm(toolsetSchema, { - defaultValues: { - toolset_name: initialToolset?.toolset_name || "", - description: initialToolset?.description || "", - }, - }); +const toolsetFormDefaults = (toolset?: MCPToolset): ToolsetFormValues => ({ + toolset_name: toolset?.toolset_name || "", + description: toolset?.description || "", +}); + +function CreateToolsetModal({ open, onClose, onSave, initialToolset }: CreateToolsetModalProps) { + const form = useZodForm(toolsetSchema, { defaultValues: toolsetFormDefaults(initialToolset) }); const [selectedTools, setSelectedTools] = useState(initialToolset?.tools || []); const [saving, setSaving] = useState(false); - const [serverSearch, setServerSearch] = useState(""); - const { data: mcpServers = [] } = useMCPServers(); + const [search, setSearch] = useState(""); + const [debouncedSearch] = useDebouncedValue(search, { wait: SEARCH_DEBOUNCE_WAIT_MS }); + const [expandedServerIds, setExpandedServerIds] = useState>(() => new Set()); + const { data: mcpServers = [], isSuccess: serversLoaded } = useMCPServers(); + const serverIds = React.useMemo(() => mcpServers.map((s) => s.server_id), [mcpServers]); + const searching = search.trim() !== "" && debouncedSearch.trim() !== ""; + const fetchServerIds = React.useMemo( + () => serversToFetch(open, searching, serverIds, expandedServerIds), + [open, searching, serverIds, expandedServerIds], + ); + const toolsByServer = useMCPServerTools(serverIds, fetchServerIds); const serverPrefixById = React.useMemo( () => new Map(mcpServers.map((s) => [s.server_id, s.alias || s.server_name || s.server_id])), [mcpServers], @@ -171,41 +271,54 @@ function CreateToolsetModal({ open, onClose, onSave, accessToken, initialToolset React.useEffect(() => { if (open) { - form.reset({ - toolset_name: initialToolset?.toolset_name || "", - description: initialToolset?.description || "", - }); + form.reset(toolsetFormDefaults(initialToolset)); setSelectedTools(initialToolset?.tools || []); - setServerSearch(""); + setSearch(""); + setExpandedServerIds(new Set()); } }, [open, initialToolset, form]); const handleToggleTool = (tool: MCPToolsetTool) => { - setSelectedTools((prev) => { - const exists = prev.some((t) => t.server_id === tool.server_id && t.tool_name === tool.tool_name); - return exists - ? prev.filter((t) => !(t.server_id === tool.server_id && t.tool_name === tool.tool_name)) - : [...prev, tool]; - }); + setSelectedTools((prev) => toggleTool(prev, tool)); + }; + + const handleSetToolsSelected = (tools: MCPToolsetTool[], selected: boolean) => { + setSelectedTools((prev) => (selected ? selectTools(prev, tools) : deselectTools(prev, tools))); + }; + + const handleToggleExpanded = (serverId: string) => { + setExpandedServerIds((prev) => + prev.has(serverId) ? new Set([...prev].filter((id) => id !== serverId)) : new Set([...prev, serverId]), + ); + }; + + // The create modal stays mounted, so state left from the last session would fetch servers + // on the first render after reopening, before the open effect resets it. + const handleClose = () => { + setSearch(""); + setExpandedServerIds(new Set()); + onClose(); }; const handleSubmit = async (values: ToolsetFormValues) => { setSaving(true); try { await onSave(values.toolset_name, values.description, selectedTools); - onClose(); + handleClose(); } finally { setSaving(false); } }; - const filteredServers = mcpServers.filter((s) => { - const q = serverSearch.toLowerCase(); - return !q || (s.alias || "").toLowerCase().includes(q) || (s.server_name || "").toLowerCase().includes(q); - }); + const searchActive = search.trim() !== ""; + const serverViews = buildServerViews(mcpServers, toolsByServer, search); + const { matchingTools, pendingServerCount } = searchActive + ? summarizeSearch(serverViews) + : { matchingTools: [], pendingServerCount: 0 }; + const matchingState = selectionState(selectedTools, matchingTools); return ( - !open && onClose()}> + !open && handleClose()}> {initialToolset ? "Edit Toolset" : "New Toolset"} @@ -229,32 +342,65 @@ function CreateToolsetModal({ open, onClose, onSave, accessToken, initialToolset setServerSearch(e.target.value)} + placeholder="Search tools or servers..." + aria-label="Search tools or servers" + value={search} + onChange={(e) => setSearch(e.target.value)} /> - {serverSearch && ( + {search && ( - setServerSearch("")}> + setSearch("")}> )} + {searchActive && ( +
+

+ {matchingTools.length} matching {matchingTools.length === 1 ? "tool" : "tools"} + {pendingServerCount > 0 && + ` · loading ${pendingServerCount} ${pendingServerCount === 1 ? "server" : "servers"}…`} +

+
+ + +
+
+ )}
- {filteredServers.length === 0 ? ( + {serverViews.length === 0 ? (

- {mcpServers.length === 0 ? "No MCP servers configured" : "No servers match your search"} + {mcpServers.length === 0 ? "No MCP servers configured" : "No tools or servers match your search"}

) : ( - filteredServers.map((server) => ( - ( + handleToggleExpanded(server.server_id)} + onToggleTool={handleToggleTool} + onSetToolsSelected={handleSetToolsSelected} /> )) )} @@ -266,41 +412,60 @@ function CreateToolsetModal({ open, onClose, onSave, accessToken, initialToolset {/* Right panel: Your Toolset */}
-

- Your Toolset{" "} - ({selectedTools.length} tools) -

+
+

+ Your Toolset{" "} + ({selectedTools.length} tools) +

+ +
{selectedTools.length === 0 ? (

No tools added yet

) : ( - selectedTools.map((tool, idx) => ( -
- - ✕ - - - )) + + ); + }) )}
-
diff --git a/ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx b/ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx index 9f320049b09..ecb8dd161f3 100644 --- a/ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx @@ -63,6 +63,24 @@ describe("SearchSelect", () => { expect(onValueChange).toHaveBeenLastCalledWith("team-3"); }); + it("offers the typed text as a value only when the caller accepts it", async () => { + const user = userEvent.setup(); + const onValueChange = vi.fn(); + const customValue = { hint: "Custom value", accepts: (typed: string) => typed.endsWith("*") }; + const { rerender } = render(); + const input = screen.getByRole("combobox"); + await user.click(input); + fireEvent.change(input, { target: { value: "team-*" } }); + expect(screen.queryByRole("option", { name: /team-\*/ })).not.toBeInTheDocument(); + + rerender(); + fireEvent.change(input, { target: { value: "team-4" } }); + expect(screen.queryByRole("option", { name: /team-4/ })).not.toBeInTheDocument(); + fireEvent.change(input, { target: { value: "team-*" } }); + await user.click(await screen.findByRole("option", { name: /team-\*/ })); + expect(onValueChange).toHaveBeenLastCalledWith("team-*"); + }); + it("filters the options client-side as you type", async () => { const user = userEvent.setup(); render(); diff --git a/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx b/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx index 3649114826d..e139720a2ed 100644 --- a/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx +++ b/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx @@ -1,6 +1,6 @@ "use client"; -import React from "react"; +import React, { useState } from "react"; import { Combobox, ComboboxContent, @@ -19,6 +19,11 @@ export interface SearchSelectOption { icon?: React.ReactNode; } +interface CustomValue { + hint: string; + accepts: (typed: string) => boolean; +} + interface SearchSelectProps { options: SearchSelectOption[]; value?: string | null; @@ -29,6 +34,8 @@ interface SearchSelectProps { className?: string; inputId?: string; allowClear?: boolean; + customValue?: CustomValue; + reservedValues?: readonly string[]; "aria-label"?: string; id?: string; } @@ -49,13 +56,22 @@ export function SearchSelect({ className, inputId, allowClear = true, + customValue, + reservedValues = [], "aria-label": ariaLabel, id, }: SearchSelectProps) { + const [query, setQuery] = useState(""); const selected = value == null || value === "" ? null : options.find((option) => option.value === value) ?? { label: value, value }; - const items = + const known = selected !== null && !options.some((option) => option.value === selected.value) ? [selected, ...options] : options; + const typed = query.trim(); + const isNewValue = typed !== "" && !reservedValues.includes(typed) && !known.some((option) => option.value === typed); + const items = + customValue !== undefined && isNewValue && customValue.accepts(typed) + ? [...known, { label: typed, value: typed, sublabel: customValue.hint }] + : known; return ( a.value === b.value} itemToStringLabel={(item: SearchSelectOption) => item.label} filter={matchesQuery} + onInputValueChange={setQuery} disabled={disabled} > x. - * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) + * - model_max_budget: Optional[dict] - Model-specific max budget for user, keyed by model name, `*` pattern, or model access group. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) * - mcp_rpm_limit: Optional[dict] - Per-MCP-server rpm limit, keyed by MCP server name {"github": 100, "slack": 200}. Enforced for keys and teams only; values set on a user are stored but not enforced per user. @@ -18917,7 +18917,7 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) + * - model_max_budget: Optional[dict] - Model-specific max budget for user, keyed by model name, `*` pattern, or model access group. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) * - mcp_rpm_limit: Optional[dict] - Per-MCP-server rpm limit, keyed by MCP server name {"github": 100, "slack": 200}. Enforced for keys and teams only; values set on a user are stored but not enforced per user. From 41c4c3788e0e39c972f6bf4ed71451e61cbb24f8 Mon Sep 17 00:00:00 2001 From: tessellark <349031000+tessellark@users.noreply.github.com> Date: Wed, 28 Oct 2026 15:42:51 -0700 Subject: [PATCH 104/128] test(routing): audit #46601 (#48910) * test(routing): audit #46601 Adds integration cells for the falsy budget 400, the empty and none-only declarations under drop_params, Vertex AI Claude, Databricks Claude, and the Bedrock native output retry * test(routing): expect the Anthropic default max_tokens on Bedrock Invoke chat Reads the default from litellm instead of pinning 4096, which #48502 raised to 128000 * test(routing): assert whole sent bodies in the Vertex AI and Bedrock native output cells --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../_reasoning_effort_levels_support.py | 5 + .../test_reasoning_effort_levels_wire.py | 320 +++++++++++++++++- 2 files changed, 316 insertions(+), 9 deletions(-) diff --git a/tests/integration/routing/_reasoning_effort_levels_support.py b/tests/integration/routing/_reasoning_effort_levels_support.py index 8d43695e419..d295f6e6d56 100644 --- a/tests/integration/routing/_reasoning_effort_levels_support.py +++ b/tests/integration/routing/_reasoning_effort_levels_support.py @@ -32,6 +32,7 @@ GEMINI: Final = "gemini-3-pro-preview" ZAI_GLM: Final = "glm-5" ECHOED_KEY_MODEL: Final = f"{KIMI}-echoed-key" BEDROCK_CLAUDE: Final = f"us.anthropic.{CLAUDE}" +DATABRICKS_CLAUDE: Final = f"databricks-{CLAUDE}" AZURE_API_VERSION: Final = "2025-04-01-preview" SYNTHETIC_KEY: Final = "synthetic-provider-key" _MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") @@ -215,6 +216,10 @@ def backend(kind: str, wire_url: str) -> Backend: f"bedrock/invoke/{BEDROCK_CLAUDE}", MappingProxyType({"api_base": wire_url, "api_key": SYNTHETIC_KEY, "aws_region_name": "us-east-1"}), ) + case "databricks_claude": + return Backend( + f"databricks/{DATABRICKS_CLAUDE}", MappingProxyType({"api_base": wire_url, "api_key": SYNTHETIC_KEY}) + ) case _: raise AssertionError(f"unknown backend kind {kind}") diff --git a/tests/integration/routing/test_reasoning_effort_levels_wire.py b/tests/integration/routing/test_reasoning_effort_levels_wire.py index 3d88aad547c..d15867ac572 100644 --- a/tests/integration/routing/test_reasoning_effort_levels_wire.py +++ b/tests/integration/routing/test_reasoning_effort_levels_wire.py @@ -1,4 +1,5 @@ import asyncio +import json from collections.abc import Callable, Mapping from dataclasses import dataclass from types import MappingProxyType @@ -8,9 +9,11 @@ import anthropic import httpx import openai import pytest +from integration._support.anthropic_sse import message_json from integration._support.client import Gateway, eventually, list_value, object_value, string_value from integration._support.openai_wire import openai_error -from integration._support.wire import Request, wire_server +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, wire_server from integration.routing._reasoning_effort_levels_support import ( AZURE_API_VERSION, CLAUDE, @@ -34,6 +37,7 @@ from integration.routing._reasoning_effort_levels_support import ( marker_of, messages_request, new_marker, + path_of, prompt, provider_reply, responses_request, @@ -41,6 +45,8 @@ from integration.routing._reasoning_effort_levels_support import ( ) from pydantic import JsonValue +from litellm.constants import DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS + CHAT: Final = "/v1/chat/completions" MESSAGES: Final = "/v1/messages" RESPONSES: Final = "/v1/responses" @@ -49,6 +55,7 @@ _ADAPTIVE: Final[dict[str, JsonValue]] = {"type": "adaptive"} _STRING_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": "2048"} _XHIGH_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": 9000} _BOOL_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": True} +_FALSE_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": False} _NULL_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": None} _LIST_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": [2048]} _EMPTY_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": ""} @@ -67,6 +74,8 @@ _KIMI_INFO: Final[Mapping[str, JsonValue]] = MappingProxyType({"reasoning_effort _CLAUDE_INFO: Final[Mapping[str, JsonValue]] = MappingProxyType({"reasoning_effort_levels": list(CLAUDE_LEVELS)}) _GPT_INFO: Final[Mapping[str, JsonValue]] = MappingProxyType({"reasoning_effort_levels": list(GPT_LEVELS)}) _UNDECLARED: Final[Mapping[str, JsonValue]] = MappingProxyType({}) +_EMPTY_DECLARATION: Final[Mapping[str, JsonValue]] = MappingProxyType({"reasoning_effort_levels": []}) +_NONE_ONLY_DECLARATION: Final[Mapping[str, JsonValue]] = MappingProxyType({"reasoning_effort_levels": ["none"]}) _ECHOED_MAP_KEY: Final[Mapping[str, JsonValue]] = MappingProxyType( {"key": "azure/chat-latest", "litellm_provider": "azure"} ) @@ -134,7 +143,7 @@ def _invoke_chat(effort: str, thinking: Mapping[str, JsonValue], betas: tuple[st "messages": [{"role": "user", "content": [{"type": "text", "text": prompt(marker)}]}], "thinking": dict(thinking), "output_config": {"effort": effort}, - "max_tokens": 4096, + "max_tokens": DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, "anthropic_version": "bedrock-2023-05-31", "anthropic_beta": list(betas), } @@ -144,6 +153,14 @@ def _with_output_config(expected: Expected, output_config: JsonValue) -> Expecte return lambda marker: {**expected(marker), "output_config": output_config} +def _with_fields(expected: Expected, fields: Mapping[str, JsonValue]) -> Expected: + return lambda marker: {**expected(marker), **fields} + + +def _without(expected: Expected, field: str) -> Expected: + return lambda marker: {key: value for key, value in expected(marker).items() if key != field} + + def _converse_chat(effort: str, thinking: Mapping[str, JsonValue]) -> Expected: return lambda marker: { "messages": [{"role": "user", "content": [{"text": prompt(marker)}]}], @@ -320,6 +337,7 @@ _OPENAI_RESPONSES: Final = "/v1/responses" _AZURE_RESPONSES: Final = f"/openai/responses?api-version={AZURE_API_VERSION}" _UNSUPPORTED: Final = "does not support parameters" _GEMINI: Final = f"/models/{GEMINI}:generateContent" +_DATABRICKS: Final = "/chat/completions" CELLS: Final[Mapping[str, Cell]] = MappingProxyType( { @@ -1750,6 +1768,139 @@ CELLS: Final[Mapping[str, Cell]] = MappingProxyType( _anthropic_chat_without_effort(_ADAPTIVE), _ANTHROPIC, ), + "chat-anthropic-false-budget-400": Cell( + CHAT, "anthropic", _CLAUDE_INFO, {"thinking": _FALSE_BUDGET}, None, _ANTHROPIC, 400, _NON_NUMERIC_BUDGET + ), + "chat-anthropic-undeclared-null-budget-400": Cell( + CHAT, "anthropic", _UNDECLARED, {"thinking": _NULL_BUDGET}, None, _ANTHROPIC, 400, _NON_NUMERIC_BUDGET + ), + "chat-anthropic-undeclared-false-budget-drop-params-sends-adaptive-without-effort": Cell( + CHAT, + "anthropic", + _UNDECLARED, + {"thinking": _FALSE_BUDGET, "drop_params": True}, + _anthropic_chat_without_effort(_ADAPTIVE), + _ANTHROPIC, + ), + "chat-anthropic-undeclared-missing-budget-sends-low": Cell( + CHAT, + "anthropic", + _UNDECLARED, + {"thinking": _MISSING_BUDGET}, + _anthropic_chat("low", _ADAPTIVE), + _ANTHROPIC, + ), + "messages-anthropic-undeclared-false-budget-400": Cell( + MESSAGES, "anthropic", _UNDECLARED, {"thinking": _FALSE_BUDGET}, None, _ANTHROPIC, 400, _NON_NUMERIC_BUDGET + ), + "messages-anthropic-empty-string-budget-400": Cell( + MESSAGES, "anthropic", _CLAUDE_INFO, {"thinking": _EMPTY_BUDGET}, None, _ANTHROPIC, 400, _NON_NUMERIC_BUDGET + ), + "messages-anthropic-undeclared-null-budget-drop-params-sends-adaptive-without-effort": Cell( + MESSAGES, + "anthropic", + _UNDECLARED, + {"thinking": _NULL_BUDGET, "drop_params": True}, + _native_messages(None, _ADAPTIVE), + _ANTHROPIC, + ), + "responses-anthropic-false-budget-400": Cell( + RESPONSES, + "anthropic", + _CLAUDE_INFO, + {"thinking": _FALSE_BUDGET}, + None, + _ANTHROPIC, + 400, + _NON_NUMERIC_BUDGET, + ), + "chat-converse-undeclared-false-budget-400": Cell( + CHAT, + "bedrock_converse", + _UNDECLARED, + {"thinking": _FALSE_BUDGET}, + None, + _CONVERSE, + 400, + _NON_NUMERIC_BUDGET, + ), + "messages-invoke-undeclared-null-budget-400": Cell( + MESSAGES, "bedrock_invoke", _UNDECLARED, {"thinking": _NULL_BUDGET}, None, _INVOKE, 400, _NON_NUMERIC_BUDGET + ), + "chat-databricks-claude-undeclared-false-budget-400": Cell( + CHAT, + "databricks_claude", + _UNDECLARED, + {"thinking": _FALSE_BUDGET}, + None, + _DATABRICKS, + 400, + _NON_NUMERIC_BUDGET, + ), + "chat-anthropic-empty-declaration-budget-drop-params-sends-adaptive-without-effort": Cell( + CHAT, + "anthropic", + _EMPTY_DECLARATION, + {"thinking": _BUDGET, "drop_params": True}, + _anthropic_chat_without_effort(_ADAPTIVE), + _ANTHROPIC, + ), + "chat-anthropic-none-only-declaration-budget-drop-params-sends-adaptive-without-effort": Cell( + CHAT, + "anthropic", + _NONE_ONLY_DECLARATION, + {"thinking": _BUDGET, "drop_params": True}, + _anthropic_chat_without_effort(_ADAPTIVE), + _ANTHROPIC, + ), + "chat-anthropic-empty-declaration-budget-keeps-its-effort": Cell( + CHAT, + "anthropic", + _EMPTY_DECLARATION, + {"thinking": _BUDGET}, + _anthropic_chat("medium", _ADAPTIVE), + _ANTHROPIC, + ), + "chat-converse-empty-declaration-budget-drop-params-sends-adaptive-without-effort": Cell( + CHAT, + "bedrock_converse", + _EMPTY_DECLARATION, + {"thinking": _BUDGET, "drop_params": True}, + _converse_chat_without_effort(_ADAPTIVE), + _CONVERSE, + ), + "messages-anthropic-empty-declaration-budget-drop-params-sends-adaptive-without-effort": Cell( + MESSAGES, + "anthropic", + _EMPTY_DECLARATION, + {"thinking": _BUDGET, "drop_params": True}, + _native_messages(None, _ADAPTIVE), + _ANTHROPIC, + ), + "messages-anthropic-none-only-declaration-budget-drop-params-sends-adaptive-without-effort": Cell( + MESSAGES, + "anthropic", + _NONE_ONLY_DECLARATION, + {"thinking": _BUDGET, "drop_params": True}, + _native_messages(None, _ADAPTIVE), + _ANTHROPIC, + ), + "messages-anthropic-empty-declaration-budget-keeps-its-effort": Cell( + MESSAGES, + "anthropic", + _EMPTY_DECLARATION, + {"thinking": _BUDGET}, + _native_messages({"effort": "medium"}, _ADAPTIVE), + _ANTHROPIC, + ), + "messages-converse-empty-declaration-budget-drop-params-sends-adaptive-without-effort": Cell( + MESSAGES, + "bedrock_converse", + _EMPTY_DECLARATION, + {"thinking": _BUDGET, "drop_params": True}, + _converse_messages(None), + _CONVERSE, + ), } ) @@ -1815,20 +1966,65 @@ def _bedrock_native_schema(body: Mapping[str, JsonValue], *, invoke: bool) -> Js return JSON_OBJECT.validate_json(string_value(json_schema["schema"])) +_CONVERSE_NATIVE_OUTPUT: Final[dict[str, JsonValue]] = { + "textFormat": { + "type": "json_schema", + "structure": {"jsonSchema": {"name": "answer", "schema": json.dumps(_NATIVE_OUTPUT_SCHEMA)}}, + } +} +_INVOKE_NATIVE_OUTPUT: Final[dict[str, JsonValue]] = { + "effort": "high", + "format": {"type": "json_schema", "schema": _NATIVE_OUTPUT_SCHEMA}, +} +_JSON_TOOL: Final = "json_tool_call" +_CONVERSE_TOOL_SCHEMA: Final[dict[str, JsonValue]] = { + key: value for key, value in _NATIVE_OUTPUT_SCHEMA.items() if key != "additionalProperties" +} +_CONVERSE_WITH_NATIVE_OUTPUT: Final = _with_fields( + _converse_chat("high", _ADAPTIVE), {"outputConfig": _CONVERSE_NATIVE_OUTPUT} +) +_INVOKE_WITH_NATIVE_OUTPUT: Final = _with_output_config( + _invoke_chat("high", _ADAPTIVE, ("structured-outputs-2025-11-13",)), _INVOKE_NATIVE_OUTPUT +) +_CONVERSE_WITH_JSON_TOOL: Final = _with_fields( + _converse_chat("high", _ADAPTIVE), + { + "toolConfig": { + "tools": [ + { + "toolSpec": { + "name": _JSON_TOOL, + "description": _JSON_TOOL, + "inputSchema": {"json": _CONVERSE_TOOL_SCHEMA}, + } + } + ] + } + }, +) +_INVOKE_WITH_JSON_TOOL: Final = _with_fields( + _without(_invoke_chat("high", _ADAPTIVE, ()), "anthropic_beta"), + {"tools": [{"name": _JSON_TOOL, "input_schema": _NATIVE_OUTPUT_SCHEMA}]}, +) + + @pytest.mark.parametrize( - ("kind", "target", "fields"), + ("kind", "target", "fields", "expected"), [ - pytest.param("bedrock_converse", _CONVERSE, {"thinking": _BUDGET}, id="converse-budget"), + pytest.param( + "bedrock_converse", _CONVERSE, {"thinking": _BUDGET}, _CONVERSE_WITH_NATIVE_OUTPUT, id="converse-budget" + ), pytest.param( "bedrock_invoke", _INVOKE, {"thinking": _BUDGET, "output_config": {"effort": "medium"}}, + _INVOKE_WITH_NATIVE_OUTPUT, id="invoke-output-config-beside-budget", ), ], ) def test_a_deployment_opted_into_bedrock_native_output_keeps_it_beside_its_declared_level( - gateway: Gateway, kind: str, target: str, fields: Mapping[str, JsonValue] + gateway: Gateway, kind: str, target: str, fields: Mapping[str, JsonValue], expected: Expected ) -> None: marker: Final = new_marker() response_format: Final[dict[str, JsonValue]] = { @@ -1844,10 +2040,116 @@ def test_a_deployment_opted_into_bedrock_native_output_keeps_it_beside_its_decla assert response.status_code == 200, response.text[:2000] assert answer(marker) in response.text, response.text[:2000] assert [(request.method, request.target) for request in received] == [("POST", target)] - body: Final = body_of(received[0]) - assert _bedrock_native_schema(body, invoke=target == _INVOKE) == _NATIVE_OUTPUT_SCHEMA, body - assert {"toolConfig", "tools"}.isdisjoint(body), body - assert sent_effort(received[0]) == "high", body + assert body_of(received[0]) == expected(marker) + + +_NATIVE_OUTPUT_REJECTION: Final = ( + "The model returned the following errors: output_config.format: Extra inputs are not permitted" +) + + +def _rejecting_native_output(request: Request) -> Reply: + native: Final = _bedrock_native_schema(body_of(request), invoke=path_of(request).endswith("/invoke")) + if native is None: + return provider_reply(request) + return Reply(status=400, body=json.dumps({"message": _NATIVE_OUTPUT_REJECTION}).encode()) + + +@pytest.mark.parametrize( + ("kind", "target", "rejected", "retried"), + [ + pytest.param( + "bedrock_converse", _CONVERSE, _CONVERSE_WITH_NATIVE_OUTPUT, _CONVERSE_WITH_JSON_TOOL, id="converse" + ), + pytest.param("bedrock_invoke", _INVOKE, _INVOKE_WITH_NATIVE_OUTPUT, _INVOKE_WITH_JSON_TOOL, id="invoke"), + ], +) +def test_the_retry_after_bedrock_rejects_native_output_keeps_the_declared_level( + gateway: Gateway, kind: str, target: str, rejected: Expected, retried: Expected +) -> None: + marker: Final = new_marker() + response_format: Final[dict[str, JsonValue]] = { + "type": "json_schema", + "json_schema": {"name": "answer", "schema": _NATIVE_OUTPUT_SCHEMA}, + } + with wire_server(_rejecting_native_output) as wire, gateway.scenario() as scenario: + model: Final = deployment(scenario, kind, wire, _CLAUDE_INFO_WITH_NATIVE_OUTPUT) + response: Final = gateway.request( + "POST", CHAT, chat_request(model, marker, response_format=response_format, thinking=_BUDGET) + ) + received: Final = wire.drain() + assert response.status_code == 200, response.text[:2000] + assert [(request.method, request.target) for request in received] == [("POST", target), ("POST", target)] + assert [body_of(request) for request in received] == [rejected(marker), retried(marker)] + + +_VERTEX_PROJECT: Final = "scripted-project" +_VERTEX_LOCATION: Final = "us-east5" +_VERTEX_TOKEN: Final = "/_oauth/token" +_VERTEX_RAW_PREDICT: Final = ( + f"/v1/projects/{_VERTEX_PROJECT}/locations/{_VERTEX_LOCATION}/publishers/anthropic/models/{CLAUDE}:rawPredict" +) +_VERTEX_ACCESS_TOKEN: Final = json.dumps({"access_token": "scripted", "token_type": "Bearer", "expires_in": 3600}) + + +def _vertex(expected: Expected) -> Expected: + return _with_fields(_without(expected, "model"), {"anthropic_version": "vertex-2023-10-16"}) + + +def _vertex_reply(request: Request) -> Reply: + if request.target == _VERTEX_TOKEN: + return Reply(body=_VERTEX_ACCESS_TOKEN.encode()) + marker: Final = marker_of(request) + return Reply(body=message_json(f"msg_{marker}", CLAUDE, answer(marker))) + + +@pytest.mark.parametrize( + ("route", "fields", "expected", "status"), + [ + pytest.param( + CHAT, + {"reasoning_effort": "medium"}, + _vertex(_anthropic_chat("high", {"type": "adaptive", "display": "summarized"})), + 200, + id="chat-medium-to-high", + ), + pytest.param( + MESSAGES, + {"thinking": _BUDGET}, + _vertex(_native_messages({"effort": "high"}, _ADAPTIVE)), + 200, + id="messages-budget-to-high", + ), + pytest.param(CHAT, {"thinking": _FALSE_BUDGET}, None, 400, id="chat-false-budget-400"), + ], +) +def test_a_vertex_ai_claude_deployment_gets_one_of_its_declared_levels( + gateway: Gateway, route: str, fields: dict[str, JsonValue], expected: Expected | None, status: int +) -> None: + marker: Final = new_marker() + with wire_server(_vertex_reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"vertex_ai/{CLAUDE}", + api_base=wire.url, + api_key=None, + vertex_project=_VERTEX_PROJECT, + vertex_location=_VERTEX_LOCATION, + vertex_credentials=service_account_json(_VERTEX_PROJECT, wire.url), + model_info=dict(_CLAUDE_INFO), + ) + body: Final = ( + chat_request(model, marker, **fields) if route == CHAT else messages_request(model, marker, **fields) + ) + response: Final = gateway.request("POST", route, body) + received: Final = tuple(request for request in wire.drain() if request.target != _VERTEX_TOKEN) + assert response.status_code == status, response.text[:2000] + if expected is None: + assert _NON_NUMERIC_BUDGET in response.text, response.text[:2000] + assert received == (), [request.body[:300] for request in received] + return + assert answer(marker) in response.text, response.text[:2000] + assert [(request.method, request.target) for request in received] == [("POST", _VERTEX_RAW_PREDICT)] + assert body_of(received[0]) == expected(marker) def _single(received: tuple[Request, ...], marker: str) -> Request: From ddb8a14ee97e3abb6c251edfd3a4ba53467688db Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 15:44:43 -0700 Subject: [PATCH 105/128] feat(skills): add a litellm-setup agent skill (#47731) * feat(skills): add a litellm-setup agent skill * fix(skills): run the quickstart for Docker with Postgres in litellm-setup Use scripts/quickstart.sh for the Docker plus Postgres path instead of a hand-rolled Postgres container, keep uv with the user's own Postgres, load .env as data instead of sourcing it, pin the uv Python range, chmod an existing .env, list bedrock_converse and vertex_ai rows in the cost map lookup, ask before editing the shell profile, and stop the health wait when the proxy exits * fix(skills): guide empty Postgres, docker sudo, and key fixes in litellm-setup Ask for a new empty database on the uv path and DIRECT_URL behind a pooler, add the baseline error to troubleshooting, and drop inherited database settings in the loader line. Check python3 and lsof, explain Docker's permission denied, give quickstart commands by the folder's path, and fix provider keys by recreating the container or deleting the quickstart model. Drop the structure-only test. * fix(skills): keep shell profiles untouched when litellm-setup installs uv The uv installer edits shell profiles unless UV_NO_MODIFY_PATH is set, so the skill now sets it and puts ~/.local/bin on PATH for later commands. It also says litellm --setup runs in a separate terminal and how to retire the shared master-key UI login * fix(skills): offer litellm --setup only on the uv path The Docker paths install no host litellm CLI, and the old wording also passed the skill's own step number on to the user * fix(skills): ask before using an env provider key, print model names config.yaml accepts, probe the port with curl * fix(skills): state the env key default in the question, chmod the wizard config, keep uv's bin dir on PATH * fix(skills): keep litellm-setup on its own .env and install when the shell has others * fix(skills): read the setup wizard's config without printing keys and hand off the exact restart command * fix(skills): hide the database-access dev skill from end-user skill installs * fix(skills): give the litellm --setup path a restart command that keeps its config * docs(skills): point litellm --setup users to the master key in litellm_config.yaml * fix(skills): hide every non-end-user skill, use Azure deployment names, keep the restart block out of the user's shell * test(skills): find non-end-user skills in-process instead of through git * docs(skills): price Azure deployments with base_model and put uv on PATH in the start block --- .agents/skills/database-access/SKILL.md | 2 + .../.agents/skills/rust-error/SKILL.md | 2 + .../.agents/skills/rust-gateway/SKILL.md | 2 + .../.agents/skills/rust-string-enums/SKILL.md | 2 + .../.agents/skills/rust-tracing/SKILL.md | 2 + skills/litellm-setup/SKILL.md | 205 ++++++++++++++++++ .../slack-gif-creator/SKILL.md | 2 + .../test-delete-skill/SKILL.md | 2 + .../test-skill-litellm/SKILL.md | 2 + tests/unit/test_repo_agent_skills.py | 68 ++++++ .../.agents/skills/url-state/SKILL.md | 2 + 11 files changed, 291 insertions(+) create mode 100644 skills/litellm-setup/SKILL.md create mode 100644 tests/unit/test_repo_agent_skills.py diff --git a/.agents/skills/database-access/SKILL.md b/.agents/skills/database-access/SKILL.md index 796dfafe4f6..13fb7b7e286 100644 --- a/.agents/skills/database-access/SKILL.md +++ b/.agents/skills/database-access/SKILL.md @@ -1,6 +1,8 @@ --- name: database-access description: Add database queries, extract existing Prisma calls into repositories, or explicitly port queries to Rust and SQLx without expanding migration scope +metadata: + internal: true --- # Database access diff --git a/litellm-rust/.agents/skills/rust-error/SKILL.md b/litellm-rust/.agents/skills/rust-error/SKILL.md index fd5d6222287..74da1ecbcf4 100644 --- a/litellm-rust/.agents/skills/rust-error/SKILL.md +++ b/litellm-rust/.agents/skills/rust-error/SKILL.md @@ -1,6 +1,8 @@ --- name: rust-error description: Define, wrap, convert, classify, report or test Rust error types in litellm-rust with thiserror. Use when adding an error type or variant, mapping one crate's error into another, turning an error into a gateway response or Python exception, deciding retry or fault from an error, or replacing call sites that build errors from strings +metadata: + internal: true --- # Rust errors diff --git a/litellm-rust/.agents/skills/rust-gateway/SKILL.md b/litellm-rust/.agents/skills/rust-gateway/SKILL.md index 66b93ee128d..4fee0b54a68 100644 --- a/litellm-rust/.agents/skills/rust-gateway/SKILL.md +++ b/litellm-rust/.agents/skills/rust-gateway/SKILL.md @@ -1,6 +1,8 @@ --- name: rust-gateway description: Add or change an HTTP route, extractor, response envelope, stream or router mount in the litellm-rust gateway crates (crates/gateway and crates/gateway-*). Use when adding an endpoint, parsing a request body, choosing an error envelope for an API surface, relaying SSE, or writing a router-level test +metadata: + internal: true --- # Rust gateway diff --git a/litellm-rust/.agents/skills/rust-string-enums/SKILL.md b/litellm-rust/.agents/skills/rust-string-enums/SKILL.md index 5fc067790a0..352f6427dd9 100644 --- a/litellm-rust/.agents/skills/rust-string-enums/SKILL.md +++ b/litellm-rust/.agents/skills/rust-string-enums/SKILL.md @@ -1,6 +1,8 @@ --- name: rust-string-enums description: Define or refactor Rust string-valued enums and their Serde adapters in litellm-rust, using Strum and serde_with while preserving parsing, wire values, and schemas +metadata: + internal: true --- # Rust string enums diff --git a/litellm-rust/.agents/skills/rust-tracing/SKILL.md b/litellm-rust/.agents/skills/rust-tracing/SKILL.md index c1c7f43c456..a8da3f3bb14 100644 --- a/litellm-rust/.agents/skills/rust-tracing/SKILL.md +++ b/litellm-rust/.agents/skills/rust-tracing/SKILL.md @@ -1,6 +1,8 @@ --- name: rust-tracing description: Add or change Rust diagnostic tracing in litellm-rust, including route spans, subscriber layers, and Python logger delivery +metadata: + internal: true --- # Rust tracing diff --git a/skills/litellm-setup/SKILL.md b/skills/litellm-setup/SKILL.md new file mode 100644 index 00000000000..7c824751d6f --- /dev/null +++ b/skills/litellm-setup/SKILL.md @@ -0,0 +1,205 @@ +--- +name: litellm-setup +description: Set up a self-hosted LiteLLM proxy, the OpenAI-compatible LLM gateway, on this machine, from install through provider keys, config.yaml, a master key, optional Postgres and Admin UI, and a verified chat completion. Use when the user asks to set up LiteLLM, install the LiteLLM proxy, self-host an LLM gateway, or put one endpoint in front of OpenAI, Anthropic, Gemini, Azure, Bedrock, Ollama, or other providers. Not for using the LiteLLM Python SDK inside an app +--- + +# Set up a LiteLLM proxy + +Take the user from nothing to a running LiteLLM proxy that has answered one chat completion. Run every command yourself unless a step says the user does it + +Never print a secret. Provider keys, the master key, the salt key, the database URL, and virtual keys live only in `.env`, and `config.yaml` reads them through `os.environ/NAME`. Check that a value is there with `grep -c '^NAME=.' .env`, never with `cat .env` or `echo "$NAME"`. Load `.env` into a shell with this line, which exports each line as it is, so a `&` or `$` in a value stays literal, and first drops any database settings the shell already had and sets `LITELLM_MODE=PRODUCTION`, which stops `litellm` from also loading a `.env` it finds above its install folder, such as `~/.env`, so only this `.env` picks the database. The steps that need it start with it: + +```bash +unset DATABASE_URL DIRECT_URL DATABASE_HOST; export LITELLM_MODE=PRODUCTION; while IFS= read -r line || [ -n "$line" ]; do case "$line" in [A-Za-z_]*=*) export "$line" ;; esac; done < .env +``` + +## 1. Ask once + +Ask everything in one message, with defaults the user can accept by saying "go": + +- Providers and models. Default is OpenAI with its newest flagship model, using the `OPENAI_API_KEY` in your environment when `[ -n "$OPENAI_API_KEY" ] && echo set` says it is there. Run that check before asking, and say in the question that the default uses that variable, naming it but never its value, since a key in your environment can belong to another account +- Whether they want virtual keys, teams, spend tracking, or the Admin UI. Those need Postgres, plain routing does not. Default is no +- uv (default) or Docker. With Postgres, Docker runs the repo's quickstart, which brings its own database, and uv needs a Postgres server they already have, with a new, empty database for LiteLLM, so someone with their own Postgres server picks uv +- Folder and port. Defaults are `./litellm-proxy` (`./litellm-gateway` for the quickstart) and 4000 + +Then follow the steps for their answers. Plain routing skips steps 3 and 6. The quickstart is steps 2, 3, 8, and 9. uv with their own Postgres skips step 3 + +In the same message, say that someone who picks uv and would rather answer prompts can run `litellm --setup` in a separate terminal window once LiteLLM is installed, not through Claude Code. Never run it yourself, since it needs a real terminal and without one prints "Setup cancelled." and writes nothing. Give them the command as `LITELLM_MODE=PRODUCTION /litellm --setup`, with `` replaced by the folder `uv tool dir --bin` prints, to run from the folder they want the config in. It writes `litellm_config.yaml` there with the provider keys and master key in plain text and offers to start the proxy, which then listens on every network interface and stops with Ctrl+C in that window. When the user is done, run this in that folder. It makes the file readable only by them, keeps it out of git (the wizard does neither), and prints only the model names. Never print that file or run `grep`, `sed`, or `cat` on it: + +```bash +chmod 600 litellm_config.yaml; grep -qx 'litellm_config.yaml' .gitignore 2>/dev/null || echo 'litellm_config.yaml' >> .gitignore +"$(uv tool dir)/litellm/bin/python" -I -c 'import yaml; print(*(m["model_name"] for m in yaml.safe_load(open("litellm_config.yaml"))["model_list"]), sep="\n")' +``` + +Then skip to step 8 and run it against their proxy, starting that shell command with this line instead of the loader line, which reads the master key without printing it: + +```bash +export LITELLM_MASTER_KEY="$("$(uv tool dir)/litellm/bin/python" -I -c 'import yaml; print(yaml.safe_load(open("litellm_config.yaml"))["general_settings"]["master_key"])')" +``` + +## 2. Check prerequisites + +```bash +uv --version; docker info --format '{{.ServerVersion}}'; docker compose version; curl --version | head -1; openssl version; command -v python3 +``` + +On native Windows, ask the user to run Claude Code inside WSL, since every command here is POSIX shell. Every path needs curl, openssl, and python3, so ask the user to install any that is missing. The uv path needs uv, which downloads a supported Python (LiteLLM supports 3.10 to 3.14) when the `--python` range in step 4 rules out the one on the machine. If uv is missing, offer to install it, and only after the user says yes run `curl -LsSf https://astral.sh/uv/install.sh | env UV_NO_MODIFY_PATH=1 sh`, which leaves their shell profile alone (https://docs.astral.sh/uv/getting-started/installation/). It puts uv in `~/.local/bin`, so start every later command that runs `uv` or `litellm` with `export PATH="$HOME/.local/bin:$PATH";`. The Docker path needs a running Docker daemon, and the quickstart also needs Docker Compose v2. If `docker info` says `permission denied`, the user's account cannot use Docker without sudo. Offer the uv path, or ask them to add themselves to the `docker` group the way https://docs.docker.com/engine/install/linux-postinstall/ shows and start a new login session before you check again + +Check the port is free with `curl -s -o /dev/null --max-time 2 http://127.0.0.1:4000; echo $?`, which prints `7` (could not connect) only when nothing listens there. Don't use `lsof` for this, since it misses ports that root or Docker hold. If it is taken, pick another port and use it everywhere below, because a proxy asked for a busy 4000 silently moves to a random port + +## 3. Postgres with Docker: run the quickstart + +The repo's quickstart starts the gateway, Postgres, and the Admin UI in Docker, listening on this machine only, and puts a `litellm-gateway` command in that folder that stops, starts, updates, or removes them. It asks nothing when it has no terminal. It writes the master key, salt key, and database password to its own `.env` and never prints them. Download it, then run it with the folder and port from step 1. It returns once the gateway is ready: + +```bash +curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/quickstart.sh -o quickstart.sh +LITELLM_DIR="$PWD/litellm-gateway" LITELLM_PORT=4000 sh quickstart.sh --yes +cd litellm-gateway +``` + +Models live in the database here, so add each one through the API instead of a `config.yaml`, naming it the way step 5 says. Pipe the request body in, so the provider key never shows up in output or on a command line. Unless the user agreed to the key in your environment, as step 5 says, ask them to add the line `OPENAI_API_KEY=` to this `.env` themselves. Run the loader line first: + +```bash +python3 -c 'import json, os, sys; json.dump({"model_name": "gpt-6.1-sol", "litellm_params": {"model": "openai/gpt-6.1-sol", "api_key": os.environ["OPENAI_API_KEY"]}}, sys.stdout)' \ + | curl -sS -o /dev/null -w '%{http_code}\n' http://127.0.0.1:4000/model/new \ + -H "Authorization: Bearer $LITELLM_MASTER_KEY" -H "Content-Type: application/json" -d @- +``` + +Run that in one shell command that starts with the loader line, since the master key is in `.env`. `200` means the model was added. Then go to step 8 + +## 4. Install + +uv: + +```bash +uv tool install --python '>=3.10,<3.15' 'litellm[proxy]' +export PATH="$(uv tool dir --bin):$PATH"; litellm --version +``` + +With their own Postgres, install `'litellm[proxy,extra_proxy]'` instead, with the same `--python` range, which adds the Prisma client the proxy uses for the database. Running `uv tool install` again with the new extras replaces the old install. Start every later command that runs `litellm` with `export PATH="$(uv tool dir --bin):$PATH";`, so this install runs even when the shell has no `litellm` or an older one from a venv or `pip install litellm`. `uv tool update-shell` puts it on the PATH of new shells by editing the user's shell profile, so run it only after the user says yes + +Docker: + +```bash +docker pull docker.litellm.ai/berriai/litellm:main-stable +``` + +## 5. Write `.env` and `config.yaml` + +Run every command from here on in the folder from step 1. Create `.env` readable only by the user, keep it out of git, and add a master key unless it already has one. The proxy refuses to start without one, and `sk-` is the format of every key it issues: + +```bash +mkdir -p litellm-proxy && cd litellm-proxy +touch .env && chmod 600 .env; grep -qx '.env' .gitignore 2>/dev/null || echo '.env' >> .gitignore +grep -q '^LITELLM_MASTER_KEY=.' .env || printf 'LITELLM_MASTER_KEY=sk-%s\n' "$(openssl rand -hex 32)" >> .env +``` + +Copy a provider key that is already in your environment only after the user agreed to that variable, in their step 1 answer or when you ask now, and copy it without printing it, as in `printf 'OPENAI_API_KEY=%s\n' "$OPENAI_API_KEY" >> .env`. Otherwise ask the user to add the line `OPENAI_API_KEY=` to `.env` themselves and tell you when it is there. If they paste a key into the chat anyway, write it to `.env` and never repeat it. Keep one `NAME=value` per line with no quotes, so `docker --env-file` and the loader line read the same value + +Confirm each model's current LiteLLM name in the cost map instead of trusting memory, since model names go stale fast. This prints the chat models of the provider named last (`openai`, `anthropic`, `gemini`, `azure`, `bedrock`, `vertex_ai`, and so on) the way `config.yaml` takes them, with the provider in front. It matches by prefix, so `bedrock` also lists the newer Bedrock models and `vertex_ai` lists every Vertex family. For Ollama, use a model from `ollama list` instead: + +```bash +curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json | python3 -c ' +import json, sys +for name, row in json.load(sys.stdin).items(): + provider = str(row.get("litellm_provider", "")) if isinstance(row, dict) else "" + if provider.startswith(sys.argv[1]) and row.get("mode") == "chat": + prefix = "bedrock" if provider == "bedrock_converse" else "vertex_ai" if provider.startswith("vertex_ai") else provider + print(name if name.startswith(prefix + "/") else prefix + "/" + name) +' openai +``` + +In `config.yaml`, `model` is a name from that list, as in `openai/gpt-6.1-sol` or `bedrock/us.anthropic.`, and `model_name` is what clients will send. If the folder already has a `config.yaml`, show the user its `model_name` lines and ask before replacing it. This example uses the newest GPT model in the cost map when this skill was written, so replace it with what the user picked: + +```yaml +model_list: + - model_name: gpt-6.1-sol + litellm_params: + model: openai/gpt-6.1-sol + api_key: os.environ/OPENAI_API_KEY + +general_settings: + master_key: os.environ/LITELLM_MASTER_KEY +``` + +Add one entry per model. On Azure, ask the user for each deployment's name and the model it runs. Use `azure/` as `model`, since Azure routes by deployment and a name from the lookup works only when the deployment has that same name, and add `model_info:` with `base_model: azure/` to the entry, since LiteLLM prices each call by that name. Providers that need more than a key, like Azure or Bedrock, list their `litellm_params` on their page under https://docs.litellm.ai/docs/providers. Point every secret at `.env` through `os.environ/NAME` the same way + +## 6. Their own Postgres, on the uv path + +Ask the user to create a new, empty database for LiteLLM on their server (`CREATE DATABASE litellm;`) and add `DATABASE_URL=postgresql://...` for it to `.env` themselves. The proxy creates its tables there on first start and stops on a database that already holds other tables. If that URL goes through a transaction pooler such as PgBouncer, they also add `DIRECT_URL` for the same database without the pooler, which the migrations use. Then add a salt key. It encrypts the provider credentials the proxy stores in the database and falls back to the master key when unset, so set it once and never change it: + +```bash +grep -q '^LITELLM_SALT_KEY=.' .env || printf 'LITELLM_SALT_KEY=sk-%s\n' "$(openssl rand -hex 32)" >> .env +``` + +Generate the Prisma client once inside the tool's environment. A pip or uv install doesn't ship it generated (Docker images do): + +```bash +tool="$(uv tool dir)/litellm" +schema=$("$tool/bin/python" -I -c 'import importlib.util, os; print(os.path.join(importlib.util.find_spec("litellm").submodule_search_locations[0], "proxy", "schema.prisma"))') +PATH="$tool/bin:$PATH" prisma generate --schema "$schema" +``` + +## 7. Start the proxy + +Bind to `127.0.0.1` unless the user wants other machines to reach it + +uv, in one shell command that starts with the loader line, since the installed `litellm` does not read `.env` by itself: + +```bash +export PATH="$HOME/.local/bin:$PATH"; export PATH="$(uv tool dir --bin):$PATH" +nohup litellm --config config.yaml --host 127.0.0.1 --port 4000 > litellm.log 2>&1 & +echo $! > litellm.pid +``` + +Docker: + +```bash +docker run -d --name litellm --env-file .env -v "$PWD/config.yaml:/app/config.yaml:ro" -p 127.0.0.1:4000:4000 \ + docker.litellm.ai/berriai/litellm:main-stable --config /app/config.yaml --port 4000 +``` + +Wait on the health route, stopping early if the proxy exits. The first start with Postgres applies the database migrations and takes longer. Define `alive` for the path you took, `alive() { kill -0 "$(cat litellm.pid)" 2>/dev/null; }` for uv or `alive() { [ "$(docker inspect -f '{{.State.Running}}' litellm 2>/dev/null)" = true ]; }` for Docker, in the same shell command as the loop: + +```bash +for i in $(seq 1 180); do curl -fsS http://127.0.0.1:4000/health/liveliness >/dev/null 2>&1 && break; alive || break; sleep 1; done +curl -sS http://127.0.0.1:4000/health/readiness +``` + +Readiness reports `"db":"connected"` with Postgres and `"db":"Not connected"` without it. If the proxy never comes up, read the end of `litellm.log` (or `docker logs litellm`) and check the troubleshooting list + +## 8. Verify with a chat completion + +Use a `model_name` from `config.yaml` (or step 3), in one shell command that starts with the loader line: + +```bash +curl -sS http://127.0.0.1:4000/v1/chat/completions \ + -H "Authorization: Bearer $LITELLM_MASTER_KEY" -H "Content-Type: application/json" \ + -d '{"model": "gpt-6.1-sol", "messages": [{"role": "user", "content": "Say hello in five words."}]}' +``` + +Show the user the reply in `choices[0].message.content`. An error here means the setup is not done, so fix it before handing off + +## 9. Hand off + +Tell the user, without printing any secret: + +- The base URL is `http://127.0.0.1:4000`. Any OpenAI-compatible client works with it as the base URL and the master key or a virtual key as the API key +- The master key is `LITELLM_MASTER_KEY` in the folder's `.env`, or `general_settings.master_key` in `litellm_config.yaml` when they used `litellm --setup` +- To stop it, `kill "$(cat litellm.pid)"` (uv), `docker stop litellm` (Docker), `/litellm-gateway stop` (quickstart), or Ctrl+C in its window (`litellm --setup`). To start it again, `docker start litellm`, `/litellm-gateway start`, on uv the whole step 7 command as you ran it, loader line and PATH export included, wrapped in `( ... )` so their shell keeps its own settings, shown as a code block and never paraphrased, since a new shell has neither, or for `litellm --setup` `LITELLM_MODE=PRODUCTION /litellm --config litellm_config.yaml --host 127.0.0.1 --port ` in a terminal window in that folder, with `` and `` filled in, since running the wizard again writes a new config with a new master key. `update` and `uninstall` update or remove a quickstart install. Give the quickstart commands with the folder's full path, since a bare `litellm-gateway` is on PATH only when `~/.local/bin` is, and can belong to an earlier install +- With Postgres, the Admin UI is at `http://127.0.0.1:4000/ui` (user `admin`, password the master key), and `POST /key/generate` with the master key creates virtual keys with their own budgets and model lists. Without Postgres, both need step 3 or 6 first +- The Admin UI warns that the master-key login is a shared secret, which is fine while only this machine reaches the proxy. Before anyone else can reach it, invite an admin with their own password on the Internal Users page, then set `general_settings.disable_env_credential_login: true` in the proxy config +- To manage the running proxy from a coding agent (keys, teams, users, models, MCP servers), install `npx skills add BerriAI/litellm-skills -g` + +## Troubleshooting + +- `refused to start: no master key is set` means the loader line did not run in the shell command that started `litellm`, or `--env-file .env` is missing from `docker run` +- `a database URL is set but the prisma CLI is neither on PATH nor importable` means the install lacks `extra_proxy`, so rerun step 4 with it, or, on plain routing, that the command that started `litellm` did not start with the loader line +- `The Client hasn't been generated yet` or `Unable to find Prisma binaries` means the Prisma generate command in step 6 has not run +- `Incorrect API key provided` or another provider `AuthenticationError` means a provider key is wrong. Fix it in `.env` and restart, since the proxy reads the environment only at startup. On uv, stop it and rerun step 7. On Docker, `docker rm -f litellm` and rerun step 7, since `docker restart` keeps the environment the container was created with. On the quickstart the key is saved with the model, so delete the model and add it again: `GET /model/info` lists each model's `model_info.id`, and `POST /model/delete` with `{"id": ""}` deletes it +- `Cannot automatically baseline this database` means `DATABASE_URL` points at a database that already holds tables LiteLLM did not create. Point it at a new, empty database +- `LLM Provider NOT provided` means a `model` in `config.yaml` (or step 3) lacks its provider in front. Use the name as the step 5 lookup prints it +- `Invalid model name passed in model=...` means the request's `model` is not a `model_name` in `config.yaml` or step 3. `GET /v1/models` lists the ones the key can use +- `No connected db.` on a proxy without Postgres means the bearer token is not the master key +- `Conflict. The container name "/litellm" is already in use` means an earlier run left a container. `docker start litellm` reuses it, or with the user's OK `docker rm -f litellm` and rerun step 7 +- Inside Docker, `localhost` in `config.yaml` (an Ollama `api_base`, say) is the container itself. Use `host.docker.internal`, and on Linux also pass `--add-host=host.docker.internal:host-gateway` diff --git a/tests/unit/skills/test_skills_data/slack-gif-creator/SKILL.md b/tests/unit/skills/test_skills_data/slack-gif-creator/SKILL.md index 3cae971b731..d2cdfd17bde 100644 --- a/tests/unit/skills/test_skills_data/slack-gif-creator/SKILL.md +++ b/tests/unit/skills/test_skills_data/slack-gif-creator/SKILL.md @@ -1,6 +1,8 @@ --- name: slack-gif-creator description: Knowledge and utilities for creating animated GIFs optimized for Slack. Provides constraints, validation tools, and animation concepts. Use when users request animated GIFs for Slack like "make me a GIF of X doing Y for Slack." +metadata: + internal: true --- # Slack GIF Creator diff --git a/tests/unit/skills/test_skills_data/test-delete-skill/SKILL.md b/tests/unit/skills/test_skills_data/test-delete-skill/SKILL.md index ba3fda9ab7b..ef1fd1a7b24 100644 --- a/tests/unit/skills/test_skills_data/test-delete-skill/SKILL.md +++ b/tests/unit/skills/test_skills_data/test-delete-skill/SKILL.md @@ -1,6 +1,8 @@ --- name: test-delete-skill description: A test skill created specifically for deletion testing +metadata: + internal: true --- # Test Delete Skill diff --git a/tests/unit/skills/test_skills_data/test-skill-litellm/SKILL.md b/tests/unit/skills/test_skills_data/test-skill-litellm/SKILL.md index ec83b580f0e..b3d54ccdf3c 100644 --- a/tests/unit/skills/test_skills_data/test-skill-litellm/SKILL.md +++ b/tests/unit/skills/test_skills_data/test-skill-litellm/SKILL.md @@ -1,6 +1,8 @@ --- name: test-skill-litellm description: A test skill created by LiteLLM automated tests +metadata: + internal: true --- # Test Skill diff --git a/tests/unit/test_repo_agent_skills.py b/tests/unit/test_repo_agent_skills.py new file mode 100644 index 00000000000..591241f0cf9 --- /dev/null +++ b/tests/unit/test_repo_agent_skills.py @@ -0,0 +1,68 @@ +from collections.abc import Iterator +from fnmatch import fnmatch +from pathlib import Path +from typing import Final + +import pytest +import yaml + +_REPO_ROOT: Final = Path(__file__).resolve().parents[2] +_SKILLS_DIR: Final = _REPO_ROOT / "skills" +_SKILL_DIRECTORIES: Final = tuple(sorted(skill_file.parent for skill_file in _SKILLS_DIR.glob("*/SKILL.md"))) +_FULL_DEPTH_MAX_DEPTH: Final = 5 +_FULL_DEPTH_SKIPPED_NAMES: Final = frozenset(("node_modules", ".git", "dist", "build", "__pycache__")) +_GITIGNORED_NAME_PATTERNS: Final = tuple( + line.strip().rstrip("/") + for line in (_REPO_ROOT / ".gitignore").read_text(encoding="utf-8").splitlines() + if line.strip() and not line.startswith(("#", "!")) and "/" not in line.strip().rstrip("/") +) + + +def _absent_from_a_clone_or_skipped_by_the_cli(name: str) -> bool: + return name in _FULL_DEPTH_SKIPPED_NAMES or any(fnmatch(name, pattern) for pattern in _GITIGNORED_NAME_PATTERNS) + + +def _full_depth_skill_directories(directory: Path, depth: int) -> Iterator[Path]: + if (directory / "SKILL.md").is_file(): + yield directory + if depth == _FULL_DEPTH_MAX_DEPTH: + return + for child in sorted(directory.iterdir()): + if child.is_dir() and not _absent_from_a_clone_or_skipped_by_the_cli(child.name): + yield from _full_depth_skill_directories(child, depth + 1) + + +_NON_END_USER_SKILL_DIRECTORIES: Final = tuple( + directory for directory in _full_depth_skill_directories(_REPO_ROOT, 0) if directory.parent != _SKILLS_DIR +) + + +def _frontmatter(skill_directory: Path) -> object: + opening, frontmatter, _body = (skill_directory / "SKILL.md").read_text(encoding="utf-8").split("---\n", 2) + assert opening == "", f"{skill_directory.name}/SKILL.md must open with a --- frontmatter block" + return yaml.safe_load(frontmatter) + + +@pytest.mark.parametrize("skill_directory", _SKILL_DIRECTORIES, ids=lambda path: path.name) +def test_skill_frontmatter_has_string_name_matching_its_directory_and_a_description(skill_directory: Path) -> None: + frontmatter: Final = _frontmatter(skill_directory) + assert isinstance(frontmatter, dict), frontmatter + assert frontmatter.get("name") == skill_directory.name, frontmatter + description: Final = frontmatter.get("description") + assert isinstance(description, str) and description.strip(), frontmatter + + +@pytest.mark.parametrize( + "skill_directory", + _NON_END_USER_SKILL_DIRECTORIES, + ids=lambda path: str(path.relative_to(_REPO_ROOT)), +) +def test_skill_outside_skills_folder_is_hidden_from_end_user_skill_installs(skill_directory: Path) -> None: + frontmatter: Final = _frontmatter(skill_directory) + assert isinstance(frontmatter, dict), frontmatter + metadata: Final = frontmatter.get("metadata") + assert isinstance(metadata, dict) and metadata.get("internal") is True, ( + f"{skill_directory} needs `metadata: {{internal: true}}`, or `npx skills add BerriAI/litellm` installs it " + "for end users, from a root dot-folder always and from anywhere with --full-depth (skills CLI 1.7.4, " + "checked 2026-10-28)" + ) diff --git a/ui/litellm-dashboard/.agents/skills/url-state/SKILL.md b/ui/litellm-dashboard/.agents/skills/url-state/SKILL.md index 986118d81a7..8bf8fadd19e 100644 --- a/ui/litellm-dashboard/.agents/skills/url-state/SKILL.md +++ b/ui/litellm-dashboard/.agents/skills/url-state/SKILL.md @@ -1,6 +1,8 @@ --- name: url-state description: Read or write dashboard URL query state (tabs, filters, pagination, deep links, demo flags) with nuqs instead of raw history or URLSearchParams +metadata: + internal: true --- # URL state From 752b5c053a50f1cb56d1d099f2660a4459dafc1c Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 28 Oct 2026 15:52:15 -0700 Subject: [PATCH 106/128] fix(ci): flag any schema.prisma change inside model LiteLLM_SpendLogs in the spend logs guard (LIT-10240) (#48905) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ci_cd/spend_logs_guard.py | 107 ++++++++++++++---- tests/unit/test_spend_logs_guard.py | 163 +++++++++++++++++++++++----- 2 files changed, 220 insertions(+), 50 deletions(-) diff --git a/ci_cd/spend_logs_guard.py b/ci_cd/spend_logs_guard.py index a4d9a2e0a0d..5e6d38fba57 100644 --- a/ci_cd/spend_logs_guard.py +++ b/ci_cd/spend_logs_guard.py @@ -2,11 +2,12 @@ needs an approving review from the spend-log owner before it can merge. CODEOWNERS can only match file paths, and every schema change edits the same schema.prisma copies, -so a path rule would put every migration on the owner. This gate reads the diff instead: when no -added line in the watched files mentions the table it passes at once, otherwise it passes only once -the owner's latest review on the pull request is an approval that still covers every spend-log line -in the head (approvals survive pushes, so a later commit may not add spend-log lines the owner never -saw). Reviews do not trigger the workflow, so the check is re-run by hand after approving. +so a path rule would put every migration on the owner. This gate reads the diff instead: a schema.prisma +hunk counts when it adds or removes a line inside the `model LiteLLM_SpendLogs { ... }` block, a migration +hunk counts when an added line names the table. Without such a hunk it passes at once, otherwise it passes +only once the owner's latest review on the pull request is an approval that still covers every spend-log line +in the head (approvals survive pushes, so a later commit may not add spend-log lines the owner never saw). +Reviews do not trigger the workflow, so the check is re-run by hand after approving. """ from __future__ import annotations @@ -15,8 +16,10 @@ import argparse import re import subprocess import sys -from collections.abc import Callable, Sequence +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass from pathlib import Path +from types import MappingProxyType from typing import Final from pydantic import BaseModel, ConfigDict, TypeAdapter @@ -29,6 +32,8 @@ WATCHED: Final = ( MIGRATIONS_PREFIX: Final = "litellm-proxy-extras/litellm_proxy_extras/migrations/" TABLE: Final = re.compile(r"(? bool: return path in WATCHED or (path.startswith(MIGRATIONS_PREFIX) and path.endswith("/migration.sql")) -def _added_lines(section: str) -> tuple[str, ...]: - body: Final = section.split("\n@@", 1) - return tuple(line[1:] for line in body[1].splitlines() if line.startswith("+")) if len(body) == 2 else () +def model_block(schema: str) -> range: + """1-based line range of `model LiteLLM_SpendLogs { ... }` including both braces, empty when the schema has none.""" + lines: Final = schema.splitlines() + starts: Final = tuple(number for number, line in enumerate(lines, 1) if MODEL_START.match(line)) + if not starts: + return range(0) + ends: Final = tuple(number for number, line in enumerate(lines, 1) if number > starts[0] and line.startswith("}")) + return range(starts[0], (ends[0] if ends else len(lines)) + 1) -def spend_log_hits(diff: str) -> tuple[str, ...]: - """`path: added line` for every added line in a watched file that names the table.""" +def _hunks(section: str) -> tuple[Hunk, ...]: + headers: Final = tuple(HUNK.finditer(section)) + ends: Final = tuple(header.start() for header in headers[1:]) + (len(section),) + return tuple(_hunk(header, section[header.end() : end]) for header, end in zip(headers, ends)) + + +def _hunk(header: re.Match[str], body: str) -> Hunk: + lines: Final = body.splitlines() + return Hunk( + removed=tuple(enumerate((line[1:] for line in lines if line.startswith("-")), int(header.group(1)))), + added=tuple(enumerate((line[1:] for line in lines if line.startswith("+")), int(header.group(2)))), + ) + + +def spend_log_hits(diff: str, blocks: Mapping[str, Block]) -> tuple[str, ...]: + """`path: +line` or `path: -line` for every changed line that touches the table. + + A schema.prisma line counts when its number falls inside that copy's `blocks` entry (head for added lines, + merge base for removed ones), a migration line when it is added and names the table.""" sections: Final = re.split(r"^diff --git [^\n]*\n", diff, flags=re.MULTILINE) - return tuple(f"{path}: {line.strip()}" for path, line in _watched_added_lines(sections) if TABLE.search(line)) - - -def _watched_added_lines(sections: Sequence[str]) -> tuple[tuple[str, str], ...]: - pairs: Final = tuple((_path_of(section), section) for section in sections) + paths: Final = tuple((_path_of(section), section) for section in sections) return tuple( - (path, line) for path, section in pairs if watched(path) for line in _added_lines(section) - ) # comprehension-ok: path then its lines + hit for path, section in paths if watched(path) for hit in _section_hits(path, _hunks(section), blocks) + ) # comprehension-ok: path then its hits + + +def _section_hits(path: str, hunks: Sequence[Hunk], blocks: Mapping[str, Block]) -> tuple[str, ...]: + if path not in blocks: + return tuple(f"{path}: +{line.strip()}" for hunk in hunks for _, line in hunk.added if TABLE.search(line)) + block: Final = blocks[path] + removed: Final = tuple( + f"{path}: -{line.strip()}" for hunk in hunks for number, line in hunk.removed if number in block.base + ) # comprehension-ok: hunk then its lines + added: Final = tuple( + f"{path}: +{line.strip()}" for hunk in hunks for number, line in hunk.added if number in block.head + ) # comprehension-ok: hunk then its lines + return removed + added def _path_of(section: str) -> str: @@ -111,13 +161,25 @@ def added_diff(base: str, head: str) -> str: return result.stdout +def file_at(revision: str, path: str) -> str: + """Contents of `path` at `revision`, empty when the file does not exist there.""" + result: Final = subprocess.run(("git", "show", f"{revision}:{path}"), check=False, capture_output=True, text=True) + return result.stdout if result.returncode == 0 else "" + + +def schema_blocks(base: str, head: str) -> Mapping[str, Block]: + return MappingProxyType( + {path: Block(base=model_block(file_at(base, path)), head=model_block(file_at(head, path))) for path in WATCHED} + ) + + def fetched(commit: str) -> bool: result: Final = subprocess.run(("git", "fetch", "--no-tags", "--depth=1", "origin", commit), check=False) return result.returncode == 0 def reviewed_hits(base: str, commit: str) -> tuple[str, ...] | None: - return spend_log_hits(added_diff(base, commit)) if fetched(commit) else None + return spend_log_hits(added_diff(base, commit), schema_blocks(base, commit)) if fetched(commit) else None class Arguments(BaseModel): @@ -135,9 +197,12 @@ def main(argv: Sequence[str]) -> int: parser.add_argument("--approver", required=True, help="GitHub login whose approval unlocks the merge") parser.add_argument("--reviews", required=True, help="JSON file with the pull request's reviews") args: Final = Arguments.model_validate(vars(parser.parse_args(argv))) - hits: Final = spend_log_hits(added_diff(args.base, args.head)) + hits: Final = spend_log_hits(added_diff(args.base, args.head), schema_blocks(args.base, args.head)) if not hits: - sys.stdout.write("No added line in schema.prisma or a migration names LiteLLM_SpendLogs; nothing to review.\n") + sys.stdout.write( + "No change inside model LiteLLM_SpendLogs in schema.prisma and no added migration line names it; " + "nothing to review.\n" + ) return 0 reviews: Final = _REVIEWS.validate_json(args.reviews.read_bytes()) approval: Final = latest_approval(reviews, args.approver) if approved_by(reviews, args.approver) else None diff --git a/tests/unit/test_spend_logs_guard.py b/tests/unit/test_spend_logs_guard.py index 9ac2c923632..eb54099cf8b 100644 --- a/tests/unit/test_spend_logs_guard.py +++ b/tests/unit/test_spend_logs_guard.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import importlib.util import sys from pathlib import Path @@ -12,46 +14,144 @@ guard: Final = importlib.util.module_from_spec(_SPEC) sys.modules["spend_logs_guard"] = guard _SPEC.loader.exec_module(guard) +SCHEMA: Final = ( + "generator client {\n" + ' provider = "prisma-client-py"\n' + "}\n" + "\n" + "model LiteLLM_TeamTable {\n" + " team_id String @id\n" + "}\n" + "\n" + "model LiteLLM_SpendLogs {\n" + " request_id String @id\n" + " api_key String\n" + ' @@index([api_key], map: "LiteLLM_SpendLogs_api_key_idx")\n' + "}\n" + "\n" + "model LiteLLM_UserTable {\n" + " user_id String @id\n" + "}\n" +) +_NO_BLOCKS: Final = {} -def _diff(path: str, added: str, start: int = 10) -> str: - return f"diff --git a/{path} b/{path}\n--- a/{path}\n+++ b/{path}\n@@ -{start},0 +{start},1 @@\n+{added}\n" + +def _hunk(old_start: int, old_count: int, new_start: int, new_count: int, *lines: str) -> str: + body: Final = "".join(f"{line}\n" for line in lines) + return f"@@ -{old_start},{old_count} +{new_start},{new_count} @@\n{body}" + + +def _diff(path: str, *hunks: str) -> str: + return f"diff --git a/{path} b/{path}\n--- a/{path}\n+++ b/{path}\n" + "".join(hunks) + + +def _added(path: str, added: str, start: int = 10) -> str: + return _diff(path, _hunk(start, 0, start, 1, f"+{added}")) + + +def _blocks(path: str, base: str = SCHEMA, head: str = SCHEMA) -> dict[str, guard.Block]: + return {path: guard.Block(base=guard.model_block(base), head=guard.model_block(head))} + + +def _with_line(schema: str, after: int, line: str) -> str: + lines: Final = schema.splitlines(keepends=True) + return "".join(lines[:after]) + f"{line}\n" + "".join(lines[after:]) + + +class TestModelBlock: + def test_the_block_runs_from_the_model_line_through_its_closing_brace(self): + assert guard.model_block(SCHEMA) == range(9, 14) + + def test_a_schema_without_the_model_has_an_empty_block(self): + assert guard.model_block(SCHEMA.replace("LiteLLM_SpendLogs", "LiteLLM_ErrorLogs")) == range(0) class TestSpendLogHits: - def test_an_added_schema_line_naming_the_table_is_a_hit(self): + def test_a_field_added_inside_the_block_is_a_hit_without_naming_the_table(self): + head: Final = _with_line(SCHEMA, 9, " gauntlet_probe String?") + diff: Final = _added("schema.prisma", " gauntlet_probe String?") + assert guard.spend_log_hits(diff, _blocks("schema.prisma", head=head)) == ( + "schema.prisma: +gauntlet_probe String?", + ) + + def test_a_field_removed_from_the_block_is_a_hit(self): + head: Final = SCHEMA.replace(" api_key String\n", "") + diff: Final = _diff("schema.prisma", _hunk(11, 1, 10, 0, "- api_key String")) + assert guard.spend_log_hits(diff, _blocks("schema.prisma", head=head)) == ("schema.prisma: -api_key String",) + + def test_the_whole_model_appearing_or_disappearing_is_a_hit(self): + without: Final = SCHEMA.replace( + SCHEMA[SCHEMA.index("model LiteLLM_SpendLogs") : SCHEMA.index("model LiteLLM_UserTable")], "" + ) + body: Final = tuple(SCHEMA.splitlines()[8:14]) + added: Final = _diff("schema.prisma", _hunk(8, 0, 9, 6, *(f"+{line}" for line in body))) + removed: Final = _diff("schema.prisma", _hunk(9, 6, 8, 0, *(f"-{line}" for line in body))) + assert len(guard.spend_log_hits(added, _blocks("schema.prisma", base=without))) == 5 + assert len(guard.spend_log_hits(removed, _blocks("schema.prisma", head=without))) == 5 + + def test_a_line_naming_the_table_in_another_model_is_not_a_hit(self): + head: Final = _with_line(SCHEMA, 6, " spend_logs LiteLLM_SpendLogs[]") diff: Final = _diff( - "schema.prisma", ' @@index([api_key, startTime], map: "LiteLLM_SpendLogs_api_key_startTime_idx")', 740 - ) - assert guard.spend_log_hits(diff) == ( - 'schema.prisma: @@index([api_key, startTime], map: "LiteLLM_SpendLogs_api_key_startTime_idx")', + "schema.prisma", + _hunk(6, 0, 7, 1, "+ spend_logs LiteLLM_SpendLogs[]"), + _hunk(16, 1, 17, 0, "- user_id String @id"), ) + assert guard.spend_log_hits(diff, _blocks("schema.prisma", head=head)) == () - def test_an_added_migration_line_is_a_hit_in_any_case(self): + def test_each_hunk_is_placed_by_its_own_header(self): + head: Final = _with_line(_with_line(SCHEMA, 2, ' output = "./client"'), 10, " new_col String?") + diff: Final = _diff( + "schema.prisma", + _hunk(2, 0, 3, 1, '+ output = "./client"'), + _hunk(9, 0, 11, 1, "+ new_col String?"), + ) + assert guard.spend_log_hits(diff, _blocks("schema.prisma", head=head)) == ("schema.prisma: +new_col String?",) + + def test_an_added_migration_line_is_a_hit_in_any_case_and_a_removed_one_is_not(self): path: Final = "litellm-proxy-extras/litellm_proxy_extras/migrations/20260101000000_x/migration.sql" - hits: Final = guard.spend_log_hits(_diff(path, 'CREATE INDEX "i" ON "litellm_spendlogs"("api_key");', 3)) - assert hits == (f'{path}: CREATE INDEX "i" ON "litellm_spendlogs"("api_key");',) + diff: Final = _diff( + path, + _hunk( + 3, + 1, + 3, + 1, + '-DROP INDEX "LiteLLM_SpendLogs_old_idx";', + '+CREATE INDEX "i" ON "litellm_spendlogs"("api_key");', + ), + ) + assert guard.spend_log_hits(diff, _NO_BLOCKS) == ( + f'{path}: +CREATE INDEX "i" ON "litellm_spendlogs"("api_key");', + ) - def test_a_removed_line_or_another_table_or_an_unwatched_file_is_not_a_hit(self): - removed: Final = "--- a/schema.prisma\n+++ b/schema.prisma\n@@ -1,1 +1,0 @@\n-model LiteLLM_SpendLogs {\n" - other: Final = _diff("schema.prisma", " team_id String // LiteLLM_TeamTable, not MyLiteLLM_SpendLogs", 5) - elsewhere: Final = _diff("litellm/proxy/utils.py", 'table = "LiteLLM_SpendLogs"', 5) - assert guard.spend_log_hits(removed + other + elsewhere) == () + def test_another_table_in_a_migration_or_an_unwatched_file_is_not_a_hit(self): + path: Final = "litellm-proxy-extras/litellm_proxy_extras/migrations/20260101000000_x/migration.sql" + other: Final = _added( + path, 'ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "x" TEXT; -- not MyLiteLLM_SpendLogs', 5 + ) + elsewhere: Final = _added("litellm/proxy/utils.py", 'table = "LiteLLM_SpendLogs"', 5) + assert guard.spend_log_hits(other + elsewhere, _NO_BLOCKS) == () def test_a_table_whose_name_only_starts_with_the_watched_one_is_not_a_hit(self): - archive: Final = _diff("schema.prisma", "model LiteLLM_SpendLogsArchive {", 900) - index: Final = _diff("schema.prisma", " @@index([team_id], map: \"LiteLLM_SpendLogs_team_id_idx\")", 760) - assert guard.spend_log_hits(archive + index) == ( - 'schema.prisma: @@index([team_id], map: "LiteLLM_SpendLogs_team_id_idx")', + index: Final = ' @@index([team_id], map: "LiteLLM_SpendLogs_team_id_idx")' + head: Final = _with_line(_with_line(SCHEMA, 18, "model LiteLLM_SpendLogsArchive {"), 12, index) + diff: Final = _diff( + "schema.prisma", + _hunk(12, 0, 13, 1, f"+{index}"), + _hunk(18, 0, 20, 1, "+model LiteLLM_SpendLogsArchive {"), + ) + assert guard.spend_log_hits(diff, _blocks("schema.prisma", head=head)) == ( + 'schema.prisma: +@@index([team_id], map: "LiteLLM_SpendLogs_team_id_idx")', ) def test_two_files_in_one_diff_are_attributed_to_their_own_path(self): - diff: Final = _diff("litellm/proxy/schema.prisma", " a LiteLLM_SpendLogs", 1) + _diff( - "schema.prisma", " b String", 1 - ) - assert guard.spend_log_hits(diff) == ("litellm/proxy/schema.prisma: a LiteLLM_SpendLogs",) + head: Final = _with_line(SCHEMA, 9, " a String") + diff: Final = _added("litellm/proxy/schema.prisma", " a String") + _added("schema.prisma", " b String", 2) + blocks: Final = _blocks("litellm/proxy/schema.prisma", head=head) | _blocks("schema.prisma") + assert guard.spend_log_hits(diff, blocks) == ("litellm/proxy/schema.prisma: +a String",) -def _review(login: str, state: str, commit_id: str | None = "head") -> "guard.Review": +def _review(login: str, state: str, commit_id: str | None = "head") -> guard.Review: return guard.Review(state=state, user=guard.Reviewer(login=login), commit_id=commit_id) @@ -126,24 +226,29 @@ class TestApprovalCovers: class TestExitCode: def test_a_pull_request_without_spend_log_lines_passes_without_reviews(self, tmp_path, monkeypatch, capsys): - monkeypatch.setattr(guard, "added_diff", lambda base, head: _diff("schema.prisma", " foo String", 1)) + monkeypatch.setattr(guard, "added_diff", lambda base, head: _added("schema.prisma", " foo String", 6)) + monkeypatch.setattr(guard, "schema_blocks", lambda base, head: _blocks("schema.prisma")) reviews: Final = tmp_path / "r.json" reviews.write_text("[]") assert guard.main(["--base", "a", "--head", "b", "--approver", "kerry-berri", "--reviews", str(reviews)]) == 0 assert "nothing to review" in capsys.readouterr().out def test_a_spend_log_change_fails_until_the_approver_approves(self, tmp_path, monkeypatch, capsys): - monkeypatch.setattr(guard, "added_diff", lambda base, head: _diff("schema.prisma", " x LiteLLM_SpendLogs", 7)) + head: Final = _with_line(SCHEMA, 9, " x String?") + monkeypatch.setattr(guard, "added_diff", lambda base_sha, head_sha: _added("schema.prisma", " x String?")) + monkeypatch.setattr(guard, "schema_blocks", lambda base_sha, head_sha: _blocks("schema.prisma", head=head)) reviews: Final = tmp_path / "r.json" reviews.write_text('[{"user": {"login": "someone-else"}, "state": "APPROVED"}]') assert guard.main(["--base", "a", "--head", "b", "--approver", "kerry-berri", "--reviews", str(reviews)]) == 1 - assert "schema.prisma: x LiteLLM_SpendLogs" in capsys.readouterr().out + assert "schema.prisma: +x String?" in capsys.readouterr().out reviews.write_text('[{"user": {"login": "kerry-berri"}, "state": "APPROVED", "commit_id": "b"}]') assert guard.main(["--base", "a", "--head", "b", "--approver", "kerry-berri", "--reviews", str(reviews)]) == 0 def test_an_approval_before_a_new_spend_log_line_was_pushed_fails(self, tmp_path, monkeypatch, capsys): - monkeypatch.setattr(guard, "added_diff", lambda base, head: _diff("schema.prisma", " x LiteLLM_SpendLogs", 7)) - monkeypatch.setattr(guard, "reviewed_hits", lambda base, commit: () if commit == "old" else None, raising=False) + head: Final = _with_line(SCHEMA, 9, " x String?") + monkeypatch.setattr(guard, "added_diff", lambda base_sha, head_sha: _added("schema.prisma", " x String?")) + monkeypatch.setattr(guard, "schema_blocks", lambda base_sha, head_sha: _blocks("schema.prisma", head=head)) + monkeypatch.setattr(guard, "reviewed_hits", lambda base_sha, commit: () if commit == "old" else None) reviews: Final = tmp_path / "r.json" reviews.write_text('[{"user": {"login": "kerry-berri"}, "state": "APPROVED", "commit_id": "old"}]') assert guard.main(["--base", "a", "--head", "b", "--approver", "kerry-berri", "--reviews", str(reviews)]) == 1 From 0f786fa1c7b99ee860643597675ed7daf0d289d0 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 28 Oct 2026 15:52:44 -0700 Subject: [PATCH 107/128] feat(router): add deployment rpd and tpd daily limits to enforce_model_rate_limits (#48875) * feat(router): add deployment rpd and tpd daily limits to enforce_model_rate_limits Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): route rate limit cache increments through typed helpers * chore(ui): regenerate schema.d.ts for deployment rpd and tpd fields * test(types): add rpd and tpd to the owned litellm_params inventory * fix(router): run daily rpd/tpd checks after the minute checks and increment rpm and rpd in one cache trip * fix(router): flatten the async rpm reject branch for the ruff strict gate * fix(router): hand rpm and rpd back on any rate limit reject A request the rpm, rpd or tpd check rejects no longer keeps its request counts, so retries against a deployment whose rpd is spent cannot fill its rpm and block the in-group fallover with the minute 429. * test(router): a minute-rejected request no longer keeps its rpm count --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/caching/in_memory_cache.py | 4 +- .../litellm_core_utils/get_litellm_params.py | 2 + .../pre_call_checks/model_rate_limit_check.py | 228 +++++++++++--- litellm/types/litellm_params.py | 2 + litellm/types/router.py | 4 + ..._enforce_model_rate_limits_daily_limits.py | 96 ++++++ tests/unit/caching/test_in_memory_cache.py | 25 +- .../test_enforce_model_rate_limits.py | 286 +++++++++++++++++- .../test_router/test_io_token_rate_limits.py | 2 +- tests/unit/types/test_litellm_params.py | 2 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 8 + 11 files changed, 589 insertions(+), 70 deletions(-) create mode 100644 tests/integration/routing/test_enforce_model_rate_limits_daily_limits.py diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index 1cc3e106aa0..6b40c48c4ef 100644 --- a/litellm/caching/in_memory_cache.py +++ b/litellm/caching/in_memory_cache.py @@ -274,7 +274,9 @@ class InMemoryCache(BaseCache): ) -> list[float] | None: results: Final = [] for increment in increment_list: - result = await self.async_increment(increment["key"], increment["increment_value"], **kwargs) + result = await self.async_increment( + increment["key"], increment["increment_value"], ttl=increment["ttl"], **kwargs + ) results.append(result) return results diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index ba964d200ec..8178b7e98b8 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -75,6 +75,8 @@ OPTIONAL_KWARGS_KEYS: Final = ( "rpm", "itpm", "otpm", + "tpd", + "rpd", "use_xai_oauth", "chatgpt_auth_file", "github_copilot_auth_type", diff --git a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py index ea3fdd601dc..9eebcf39515 100644 --- a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py @@ -5,15 +5,19 @@ When enabled via router_settings.optional_pre_call_checks: ["enforce_model_rate_ - tpm/rpm: combined TPM + optional RPM (legacy) - itpm/otpm: separate input/output tokens per minute +- tpd/rpd: tokens and requests per UTC calendar day When a deployment sets both itpm/otpm and tpm/rpm, both are enforced. A warning is logged the first time such a deployment is seen. """ import contextlib -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Callable +from datetime import datetime, timedelta +from typing import TYPE_CHECKING, Any, Final, Literal import httpx +from pydantic import BaseModel, ConfigDict import litellm from litellm._internal_context import with_service_target @@ -33,6 +37,7 @@ from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( io_token_refund_failure, ) from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET +from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.router import RouterErrors from litellm.types.utils import StandardLoggingPayload from litellm.utils import get_utc_datetime @@ -47,6 +52,62 @@ else: class RoutingArgs: ttl: int = 60 # 1min (RPM/TPM expire key) + daily_ttl: int = 86400 + + +DAY_FORMAT: Final = "%Y-%m-%d" + + +def seconds_until_utc_midnight(now: datetime) -> int: + next_midnight: Final = (now + timedelta(days=1)).replace(hour=0, minute=0, second=0, microsecond=0) + return max(1, int((next_midnight - now).total_seconds())) + + +def _first_set(*values: int | None) -> int | None: + return next((value for value in values if value is not None), None) + + +class _DailyLimitFields(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore", from_attributes=True) + tpd: int | None = None + rpd: int | None = None + + +class _DailyLimitParams(_DailyLimitFields): + model: str | None = None + + +class _DailyLimitModelInfo(_DailyLimitFields): + id: str | None = None + + +class _DailyLimitDeployment(_DailyLimitFields): + model_name: str = "" + litellm_params: _DailyLimitParams = _DailyLimitParams() + model_info: _DailyLimitModelInfo = _DailyLimitModelInfo() + + def limits(self) -> tuple[int | None, int | None]: + return ( + _first_set(self.tpd, self.litellm_params.tpd, self.model_info.tpd), + _first_set(self.rpd, self.litellm_params.rpd, self.model_info.rpd), + ) + + def cache_keys(self, current_day: str) -> tuple[str, str]: + prefix: Final = f"{self.model_info.id}:{self.litellm_params.model}" + return f"{prefix}:tpd:{current_day}", f"{prefix}:rpd:{current_day}" + + def limit_error(self, kind: Literal["tpd", "rpd"], limit: int, usage: int, now: datetime) -> litellm.RateLimitError: + return litellm.RateLimitError( + message=f"Model rate limit exceeded. {kind.upper()} limit={limit}, current usage={usage}", + llm_provider="", + model=self.litellm_params.model, + response=httpx.Response( + status_code=429, + content=f"{RouterErrors.user_defined_ratelimit_error.value} {kind} limit={limit}. current usage={usage}. id={self.model_info.id}, model_group={self.model_name}", + headers={"retry-after": str(seconds_until_utc_midnight(now))}, + request=httpx.Request(method="model_rate_limit_check", url="https://github.com/BerriAI/litellm"), + ), + ) class ModelRateLimitingCheck(CustomLogger): @@ -60,12 +121,60 @@ class ModelRateLimitingCheck(CustomLogger): this check actively enforces those limits across ALL routing strategies. """ - def __init__(self, dual_cache: DualCache): + def __init__(self, dual_cache: DualCache, clock: Callable[[], datetime] = get_utc_datetime): self.dual_cache = dual_cache + self.clock = clock # model_ids already warned about itpm/otpm + tpm/rpm on the same deployment, # so the warning is logged once per deployment rather than per request. self._io_token_conflict_warned_ids: set[str] = set() + def _increment(self, key: str, value: int, ttl: int) -> int: + return self.dual_cache.increment_cache(key=key, value=value, ttl=ttl) + + async def _async_increment( + self, key: str, value: int, ttl: int, parent_otel_span: Span | None = None + ) -> int | None: + count: Final = await self.dual_cache.async_increment_cache( + key=key, value=value, ttl=ttl, parent_otel_span=parent_otel_span + ) + return None if count is None else int(count) + + @staticmethod + def _request_count_operations( + rpm_key: str | None, rpd_key: str | None, increment_value: int + ) -> tuple[RedisPipelineIncrementOperation, ...]: + return tuple( + RedisPipelineIncrementOperation(key=key, increment_value=increment_value, ttl=ttl) + for key, ttl in ((rpm_key, RoutingArgs.ttl), (rpd_key, RoutingArgs.daily_ttl)) + if key is not None + ) + + async def _async_refund_request_counts( + self, rpm_key: str | None, rpd_key: str | None, parent_otel_span: Span | None + ) -> None: + operations: Final = self._request_count_operations(rpm_key, rpd_key, -1) + if operations: + await self.dual_cache.async_increment_cache_pipeline( + increment_list=list(operations), parent_otel_span=parent_otel_span + ) + + async def _async_increment_request_counts( + self, rpm_key: str | None, rpd_key: str | None, parent_otel_span: Span | None + ) -> tuple[int | None, int | None]: + operations: Final = self._request_count_operations(rpm_key, rpd_key, 1) + if not operations: + return None, None + counts: Final = await self.dual_cache.async_increment_cache_pipeline( + increment_list=list(operations), parent_otel_span=parent_otel_span + ) + if counts is None or len(counts) != len(operations): + return None, None + count_by_key: Final = {op["key"]: int(count) for op, count in zip(operations, counts)} + return ( + None if rpm_key is None else count_by_key.get(rpm_key), + None if rpd_key is None else count_by_key.get(rpd_key), + ) + def _warn_io_token_and_tpm_rpm_coexist_once(self, deployment: dict) -> None: tpm_limit, rpm_limit = self._get_deployment_limits(deployment) if tpm_limit is None and rpm_limit is None: @@ -129,6 +238,21 @@ class ModelRateLimitingCheck(CustomLogger): return tpm, rpm + def _check_daily_limits(self, deployment: _DailyLimitDeployment, now: datetime) -> None: + tpd_limit, rpd_limit = deployment.limits() + if tpd_limit is None and rpd_limit is None: + return + tpd_key, rpd_key = deployment.cache_keys(now.strftime(DAY_FORMAT)) + if tpd_limit is not None: + current_tpd: Final = self._get_current_tpm(tpd_key, tpd_limit) + if current_tpd is not None and current_tpd >= tpd_limit: + raise deployment.limit_error("tpd", tpd_limit, current_tpd, now) + if rpd_limit is not None: + current_rpd: Final = self._increment(rpd_key, 1, RoutingArgs.daily_ttl) + if current_rpd > rpd_limit: + self._increment(rpd_key, -1, RoutingArgs.daily_ttl) + raise deployment.limit_error("rpd", rpd_limit, current_rpd, now) + def _get_cache_keys(self, deployment: dict, current_minute: str) -> tuple[str, str]: """Get the cache keys for TPM and RPM tracking.""" model_id: Final = deployment.get("model_info", {}).get("id") @@ -178,13 +302,10 @@ class ModelRateLimitingCheck(CustomLogger): ) io_reservation_made = True + dt: Final = self.clock() + daily: Final = _DailyLimitDeployment.model_validate(deployment) tpm_limit, rpm_limit = self._get_deployment_limits(deployment) - # If no limits are set, allow the request - if tpm_limit is None and rpm_limit is None: - return deployment - - dt: Final = get_utc_datetime() current_minute: Final = dt.strftime("%H-%M") tpm_key, rpm_key = self._get_cache_keys(deployment, current_minute) @@ -213,8 +334,9 @@ class ModelRateLimitingCheck(CustomLogger): # Check RPM limit (atomic increment-first to avoid race conditions) if rpm_limit is not None: - current_rpm: Final = self.dual_cache.increment_cache(key=rpm_key, value=1, ttl=RoutingArgs.ttl) + current_rpm: Final = self._increment(rpm_key, 1, RoutingArgs.ttl) if current_rpm is not None and current_rpm > rpm_limit: + self._increment(rpm_key, -1, RoutingArgs.ttl) raise litellm.RateLimitError( message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}", llm_provider="", @@ -230,6 +352,12 @@ class ModelRateLimitingCheck(CustomLogger): ), ) + try: + self._check_daily_limits(daily, dt) + except litellm.RateLimitError: + if rpm_limit is not None: + self._increment(rpm_key, -1, RoutingArgs.ttl) + raise return deployment except litellm.RateLimitError: @@ -259,15 +387,14 @@ class ModelRateLimitingCheck(CustomLogger): ) io_reservation_made = True + dt: Final = self.clock() + daily: Final = _DailyLimitDeployment.model_validate(deployment) + tpd_limit, rpd_limit = daily.limits() tpm_limit, rpm_limit = self._get_deployment_limits(deployment) - # If no limits are set, allow the request - if tpm_limit is None and rpm_limit is None: - return deployment - - dt: Final = get_utc_datetime() current_minute: Final = dt.strftime("%H-%M") tpm_key, rpm_key = self._get_cache_keys(deployment, current_minute) + tpd_key, rpd_key = daily.cache_keys(dt.strftime(DAY_FORMAT)) model_id: Final = deployment.get("model_info", {}).get("id") model_name: Final = deployment.get("litellm_params", {}).get("model") @@ -293,30 +420,37 @@ class ModelRateLimitingCheck(CustomLogger): num_retries=0, # Don't retry - return 429 immediately ) - # Check RPM limit (atomic increment-first to avoid race conditions) - if rpm_limit is not None: - current_rpm: Final = await self.dual_cache.async_increment_cache( - key=rpm_key, - value=1, - ttl=RoutingArgs.ttl, - parent_otel_span=parent_otel_span, - ) - if current_rpm is not None and current_rpm > rpm_limit: - raise litellm.RateLimitError( - message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}", - llm_provider="", - model=model_name, - response=httpx.Response( - status_code=429, - content=f"{RouterErrors.user_defined_ratelimit_error.value} rpm limit={rpm_limit}. current usage={current_rpm}. id={model_id}, model_group={model_group}", - headers={"retry-after": str(60)}, - request=httpx.Request( - method="model_rate_limit_check", - url="https://github.com/BerriAI/litellm", - ), + if tpd_limit is not None: + current_tpd: Final = await self._async_get_current_tpm(tpd_key, tpd_limit, parent_otel_span) + if current_tpd is not None and current_tpd >= tpd_limit: + raise daily.limit_error("tpd", tpd_limit, current_tpd, dt) + + charged_rpm_key: Final = rpm_key if rpm_limit is not None else None + charged_rpd_key: Final = rpd_key if rpd_limit is not None else None + current_rpm, current_rpd = await self._async_increment_request_counts( + charged_rpm_key, charged_rpd_key, parent_otel_span + ) + if rpm_limit is not None and current_rpm is not None and current_rpm > rpm_limit: + await self._async_refund_request_counts(charged_rpm_key, charged_rpd_key, parent_otel_span) + raise litellm.RateLimitError( + message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}", + llm_provider="", + model=model_name, + response=httpx.Response( + status_code=429, + content=f"{RouterErrors.user_defined_ratelimit_error.value} rpm limit={rpm_limit}. current usage={current_rpm}. id={model_id}, model_group={model_group}", + headers={"retry-after": str(60)}, + request=httpx.Request( + method="model_rate_limit_check", + url="https://github.com/BerriAI/litellm", ), - num_retries=0, # Don't retry - return 429 immediately - ) + ), + num_retries=0, # Don't retry - return 429 immediately + ) + + if rpd_limit is not None and current_rpd is not None and current_rpd > rpd_limit: + await self._async_refund_request_counts(charged_rpm_key, charged_rpd_key, parent_otel_span) + raise daily.limit_error("rpd", rpd_limit, current_rpd, dt) return deployment @@ -373,17 +507,16 @@ class ModelRateLimitingCheck(CustomLogger): if not model or not total_tokens: return - dt: Final = get_utc_datetime() + dt: Final = self.clock() current_minute: Final = dt.strftime("%H-%M") tpm_key: Final = f"{model_id}:{model}:tpm:{current_minute}" + tpd_key: Final = f"{model_id}:{model}:tpd:{dt.strftime(DAY_FORMAT)}" verbose_router_logger.debug("[TPM TRACKING] Incrementing %s by %s", tpm_key, total_tokens) - await self.dual_cache.async_increment_cache( - key=tpm_key, - value=total_tokens, - ttl=RoutingArgs.ttl, - ) + tokens: Final = int(total_tokens) + await self._async_increment(tpm_key, tokens, RoutingArgs.ttl) + await self._async_increment(tpd_key, tokens, RoutingArgs.daily_ttl) except Exception as e: verbose_router_logger.debug("Error in ModelRateLimitingCheck.async_log_success_event: %s", e) @@ -434,15 +567,14 @@ class ModelRateLimitingCheck(CustomLogger): if not model or not total_tokens: return - dt: Final = get_utc_datetime() + dt: Final = self.clock() current_minute: Final = dt.strftime("%H-%M") tpm_key: Final = f"{model_id}:{model}:tpm:{current_minute}" + tpd_key: Final = f"{model_id}:{model}:tpd:{dt.strftime(DAY_FORMAT)}" - self.dual_cache.increment_cache( - key=tpm_key, - value=total_tokens, - ttl=RoutingArgs.ttl, - ) + tokens: Final = int(total_tokens) + self._increment(tpm_key, tokens, RoutingArgs.ttl) + self._increment(tpd_key, tokens, RoutingArgs.daily_ttl) except Exception as e: verbose_router_logger.debug("Error in ModelRateLimitingCheck.log_success_event: %s", e) diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 717b74e5741..9196b7a226f 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -189,6 +189,8 @@ class DeploymentOptions: tpm: int | None = None itpm: int | None = None otpm: int | None = None + rpd: int | None = None + tpd: int | None = None kubernetes_pod_discovery: bool | None = None default_api_key_rpm_limit: int | None = None default_api_key_tpm_limit: int | None = None diff --git a/litellm/types/router.py b/litellm/types/router.py index ba9b82e2a57..f8de2128a80 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -533,6 +533,8 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): rpm: int | None = None itpm: int | None = None otpm: int | None = None + tpd: int | None = None + rpd: int | None = None timeout: float | str | httpx.Timeout | SDKTimeout | None = None # if str, pass in as os.environ/ # read timeout of a live stream, a faked stream takes the longer of this and timeout. If str, pass in as os.environ/ stream_timeout: float | str | None = None @@ -738,6 +740,8 @@ class LiteLLMParamsTypedDict(TypedDict, total=False): rpm: int | None itpm: int | None otpm: int | None + tpd: ReadOnly[int | None] + rpd: ReadOnly[int | None] order: int | None weight: int | None max_parallel_requests: int | None diff --git a/tests/integration/routing/test_enforce_model_rate_limits_daily_limits.py b/tests/integration/routing/test_enforce_model_rate_limits_daily_limits.py new file mode 100644 index 00000000000..c670171a149 --- /dev/null +++ b/tests/integration/routing/test_enforce_model_rate_limits_daily_limits.py @@ -0,0 +1,96 @@ +from __future__ import annotations + +import uuid +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import httpx +import pytest +from integration._support.client import gateway_from_environment +from integration._support.openai_wire import chat_reply +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._cache_control_marks_support import owned_config +from integration.routing.test_payment_required_mapping import OPENAI_MODEL, is_model_info_probe, model_list_reply +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +GROUP: Final = "daily-limited-group" +LIMITED_ID: Final = "free-tier-two-a-day" +UNLIMITED_ID: Final = "paid-tier" +RPD: Final = 2 +EXTRA_REQUESTS: Final = 2 + + +@dataclass(frozen=True, slots=True) +class Rig: + owned: OwnedProxy + limited: Wire + unlimited: Wire + + +def _reply(request: Request) -> Reply: + if is_model_info_probe(request): + return model_list_reply() + return chat_reply("chatcmpl-daily", OPENAI_MODEL, "served", stream=False) + + +def _deployment(model_id: str, wire: Wire, **params: JsonValue) -> dict[str, JsonValue]: + return { + "model_name": GROUP, + "litellm_params": { + "model": OPENAI_MODEL, + "api_base": wire.url + "/v1", + "api_key": "synthetic-openai-key", + **params, + }, + "model_info": {"id": model_id}, + } + + +def _config(directory: Path, limited: Wire, unlimited: Wire) -> Path: + return owned_config( + directory, + [_deployment(LIMITED_ID, limited, rpd=RPD, order=1), _deployment(UNLIMITED_ID, unlimited, order=2)], + litellm_settings={"cache": False}, + router_settings={"optional_pre_call_checks": ["enforce_model_rate_limits"], "num_retries": 1}, + ) + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + directory: Final = tmp_path_factory.mktemp("daily-limits") + with ( + gateway_from_environment() as gateway, + wire_server(_reply) as limited, + wire_server(_reply) as unlimited, + owned_proxy_process(gateway, directory, {}, config=_config(directory, limited, unlimited)) as owned, + ): + yield Rig(owned, limited, unlimited) + + +def _chat(rig: Rig) -> httpx.Response: + return rig.owned.gateway.request( + "POST", + "/v1/chat/completions", + {"model": GROUP, "messages": [{"role": "user", "content": f"daily {uuid.uuid4().hex}"}]}, + ) + + +def _forwarded(wire: Wire) -> int: + return sum(1 for request in wire.drain() if not is_model_info_probe(request)) + + +def test_requests_past_the_deployment_rpd_are_served_by_the_other_deployment_in_the_group(rig: Rig) -> None: + responses: Final = tuple(_chat(rig) for _ in range(RPD + EXTRA_REQUESTS)) + + assert [response.status_code for response in responses] == [200] * (RPD + EXTRA_REQUESTS), [ + response.text for response in responses + ] + assert [response.headers["x-litellm-model-id"] for response in responses] == [LIMITED_ID] * RPD + [ + UNLIMITED_ID + ] * EXTRA_REQUESTS + assert (_forwarded(rig.limited), _forwarded(rig.unlimited)) == (RPD, EXTRA_REQUESTS) diff --git a/tests/unit/caching/test_in_memory_cache.py b/tests/unit/caching/test_in_memory_cache.py index 75232476ecd..901832937b5 100644 --- a/tests/unit/caching/test_in_memory_cache.py +++ b/tests/unit/caching/test_in_memory_cache.py @@ -235,13 +235,9 @@ def test_in_memory_cache_eviction_order(): # Add items with different TTLs now = time.time() - in_memory_cache.set_cache( - key="early_expire", value="value_1", ttl=100 - ) # expires in 100 seconds + in_memory_cache.set_cache(key="early_expire", value="value_1", ttl=100) # expires in 100 seconds time.sleep(0.01) - in_memory_cache.set_cache( - key="late_expire", value="value_2", ttl=200 - ) # expires in 200 seconds + in_memory_cache.set_cache(key="late_expire", value="value_2", ttl=200) # expires in 200 seconds # Verify TTL order early_ttl = in_memory_cache.ttl_dict["early_expire"] @@ -379,3 +375,20 @@ def test_in_memory_cache_injected_clock_controls_expiry_and_eviction() -> None: assert cache.get_cache("second") is None assert cache.get_cache("third") == "replacement" assert cache.get_cache("fourth") == "new" + + +@pytest.mark.asyncio +async def test_async_increment_pipeline_keeps_each_key_for_its_own_ttl() -> None: + now: Final = {"t": 1_000.0} + cache: Final = InMemoryCache(default_ttl=600, clock=lambda: now["t"]) + + await cache.async_increment_pipeline( + increment_list=[ + {"key": "minute", "increment_value": 1, "ttl": 60}, + {"key": "day", "increment_value": 1, "ttl": 86400}, + ] + ) + now["t"] += 601 + + assert cache.get_cache("minute") is None + assert cache.get_cache("day") == 1 diff --git a/tests/unit/test_router/test_enforce_model_rate_limits.py b/tests/unit/test_router/test_enforce_model_rate_limits.py index 7577064b7f9..2a260fdf0b8 100644 --- a/tests/unit/test_router/test_enforce_model_rate_limits.py +++ b/tests/unit/test_router/test_enforce_model_rate_limits.py @@ -6,7 +6,7 @@ regardless of the routing strategy being used. """ import asyncio -from datetime import timedelta +from datetime import datetime, timedelta, timezone from unittest.mock import AsyncMock, MagicMock import pytest @@ -222,11 +222,9 @@ class TestModelRateLimitingCheck: check.log_success_event(kwargs, None, None, None) - # Verify increment_cache was called - mock_cache.increment_cache.assert_called_once() - _, kwarg_params = mock_cache.increment_cache.call_args - assert "test-id:gpt-4:tpm:" in kwarg_params["key"] - assert kwarg_params["value"] == 50 + keys = tuple(call.kwargs["key"] for call in mock_cache.increment_cache.call_args_list) + assert any(key.startswith("test-id:gpt-4:tpm:") for key in keys), keys + assert all(call.kwargs["value"] == 50 for call in mock_cache.increment_cache.call_args_list) class TestModelRateLimitingCheckAsync: @@ -253,7 +251,7 @@ class TestModelRateLimitingCheckAsync: """Test that RateLimitError is raised when RPM limit is exceeded (async).""" mock_cache = MagicMock() mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.async_increment_cache = AsyncMock(return_value=11) # Over limit + mock_cache.async_increment_cache_pipeline = AsyncMock(return_value=[11.0]) # Over limit check = ModelRateLimitingCheck(dual_cache=mock_cache) @@ -274,7 +272,7 @@ class TestModelRateLimitingCheckAsync: """Test that requests are allowed when under the limit (async).""" mock_cache = MagicMock() mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.async_increment_cache = AsyncMock(return_value=6) + mock_cache.async_increment_cache_pipeline = AsyncMock(return_value=[6.0]) check = ModelRateLimitingCheck(dual_cache=mock_cache) @@ -339,7 +337,7 @@ class TestModelRateLimitingCheckAsync: ): redis_cache = MagicMock() redis_cache.async_get_cache = AsyncMock(side_effect=RedisCircuitBreakerOpenError()) - redis_cache.async_increment = AsyncMock(return_value=2) + redis_cache.async_increment_pipeline = AsyncMock(return_value=[2.0]) check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) with pytest.raises(litellm.RateLimitError) as exc_info: @@ -375,11 +373,9 @@ class TestModelRateLimitingCheckAsync: await check.async_log_success_event(kwargs, None, None, None) - # Verify async_increment_cache was called - mock_cache.async_increment_cache.assert_called_once() - _, kwarg_params = mock_cache.async_increment_cache.call_args - assert "test-id:gpt-4:tpm:" in kwarg_params["key"] - assert kwarg_params["value"] == 50 + keys = tuple(call.kwargs["key"] for call in mock_cache.async_increment_cache.call_args_list) + assert any(key.startswith("test-id:gpt-4:tpm:") for key in keys), keys + assert all(call.kwargs["value"] == 50 for call in mock_cache.async_increment_cache.call_args_list) class TestRouterWithEnforceModelRateLimits: @@ -466,3 +462,265 @@ class TestModelRateLimitConcurrency: assert len(successes) == 2, f"Expected 2 successes, got {len(successes)}" assert len(failures) == 2, f"Expected 2 rate limit errors, got {len(failures)}" + + +DAY_ONE = datetime(2026, 10, 28, 23, 59, 30, tzinfo=timezone.utc) +RPD_DEPLOYMENT = { + "litellm_params": {"model": "gpt-4.1-mini", "rpd": 2}, + "model_info": {"id": "daily-test-id"}, + "model_name": "daily-model", +} +RPM_AND_RPD_DEPLOYMENT = {**RPD_DEPLOYMENT, "litellm_params": {"model": "gpt-4.1-mini", "rpm": 1, "rpd": 2}} + + +class Clock: + def __init__(self, now: datetime): + self.now = now + + def __call__(self) -> datetime: + return self.now + + +def _pre_call(check: ModelRateLimitingCheck, deployment: dict[str, object]) -> object: + return check.pre_call_check(deployment) + + +async def _async_pre_call(check: ModelRateLimitingCheck, deployment: dict[str, object]) -> object: + return await check.async_pre_call_check(deployment) + + +async def _async_log_success(check: ModelRateLimitingCheck, kwargs: dict[str, object]) -> None: + await check.async_log_success_event(kwargs, None, None, None) + + +class TestDailyLimits: + def test_pre_call_check_rejects_the_request_over_rpd(self): + check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(DAY_ONE)) + + assert _pre_call(check, RPD_DEPLOYMENT) == RPD_DEPLOYMENT + assert _pre_call(check, RPD_DEPLOYMENT) == RPD_DEPLOYMENT + with pytest.raises(litellm.RateLimitError) as exc_info: + _pre_call(check, RPD_DEPLOYMENT) + + assert "RPD limit=2, current usage=3" in str(exc_info.value) + assert "rpd limit=2. current usage=3. id=daily-test-id, model_group=daily-model" in exc_info.value.response.text + + @pytest.mark.asyncio + async def test_async_pre_call_check_rejects_the_request_over_rpd_and_lets_the_router_retry(self): + check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(DAY_ONE)) + + assert await _async_pre_call(check, RPD_DEPLOYMENT) == RPD_DEPLOYMENT + assert await _async_pre_call(check, RPD_DEPLOYMENT) == RPD_DEPLOYMENT + with pytest.raises(litellm.RateLimitError) as exc_info: + await _async_pre_call(check, RPD_DEPLOYMENT) + + assert "RPD limit=2" in str(exc_info.value) + assert exc_info.value.num_retries is None + + @pytest.mark.asyncio + async def test_async_pre_call_check_rejects_once_the_day_tokens_reach_tpd(self): + deployment = {**RPD_DEPLOYMENT, "litellm_params": {"model": "gpt-4.1-mini", "tpd": 100}} + check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(DAY_ONE)) + kwargs = { + "standard_logging_object": { + "model_id": "daily-test-id", + "total_tokens": 60, + "hidden_params": {"litellm_model_name": "gpt-4.1-mini"}, + } + } + + await _async_log_success(check, kwargs) + assert await _async_pre_call(check, deployment) == deployment + await _async_log_success(check, kwargs) + with pytest.raises(litellm.RateLimitError) as exc_info: + await _async_pre_call(check, deployment) + + assert "TPD limit=100, current usage=120" in str(exc_info.value) + + def test_pre_call_check_minute_rejected_request_does_not_spend_rpd(self): + clock = Clock(DAY_ONE) + check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) + + assert _pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT + with pytest.raises(litellm.RateLimitError, match="RPM limit=1"): + _pre_call(check, RPM_AND_RPD_DEPLOYMENT) + clock.now = DAY_ONE - timedelta(minutes=1) + + assert _pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT + clock.now = DAY_ONE - timedelta(minutes=2) + with pytest.raises(litellm.RateLimitError, match="RPD limit=2, current usage=3"): + _pre_call(check, RPM_AND_RPD_DEPLOYMENT) + + @pytest.mark.asyncio + async def test_async_pre_call_check_minute_rejected_request_does_not_spend_rpd(self): + clock = Clock(DAY_ONE) + check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) + + assert await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT + with pytest.raises(litellm.RateLimitError, match="RPM limit=1"): + await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) + clock.now = DAY_ONE - timedelta(minutes=1) + + assert await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT + clock.now = DAY_ONE - timedelta(minutes=2) + with pytest.raises(litellm.RateLimitError, match="RPD limit=2, current usage=3"): + await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) + + def test_pre_call_check_rejected_requests_hand_back_rpm_and_rpd(self): + clock = Clock(DAY_ONE) + check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) + + assert _pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT + with pytest.raises(litellm.RateLimitError, match="RPM limit=1, current usage=2"): + _pre_call(check, RPM_AND_RPD_DEPLOYMENT) + with pytest.raises(litellm.RateLimitError, match="RPM limit=1, current usage=2"): + _pre_call(check, RPM_AND_RPD_DEPLOYMENT) + clock.now = DAY_ONE - timedelta(minutes=1) + assert _pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT + + clock.now = DAY_ONE - timedelta(minutes=2) + with pytest.raises(litellm.RateLimitError, match="RPD limit=2, current usage=3"): + _pre_call(check, RPM_AND_RPD_DEPLOYMENT) + with pytest.raises(litellm.RateLimitError, match="RPD limit=2, current usage=3"): + _pre_call(check, RPM_AND_RPD_DEPLOYMENT) + + @pytest.mark.asyncio + async def test_async_pre_call_check_rejected_requests_hand_back_rpm_and_rpd(self): + clock = Clock(DAY_ONE) + check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) + + assert await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT + with pytest.raises(litellm.RateLimitError, match="RPM limit=1, current usage=2"): + await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) + with pytest.raises(litellm.RateLimitError, match="RPM limit=1, current usage=2"): + await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) + clock.now = DAY_ONE - timedelta(minutes=1) + assert await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT + + clock.now = DAY_ONE - timedelta(minutes=2) + with pytest.raises(litellm.RateLimitError, match="RPD limit=2, current usage=3"): + await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) + with pytest.raises(litellm.RateLimitError, match="RPD limit=2, current usage=3"): + await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) + + @pytest.mark.asyncio + async def test_async_pre_call_check_refunds_rpm_and_rpd_in_one_cache_trip(self): + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_increment_cache_pipeline = AsyncMock(side_effect=[[1.0, 3.0], [0.0, 2.0]]) + check = ModelRateLimitingCheck(dual_cache=mock_cache, clock=Clock(DAY_ONE)) + + with pytest.raises(litellm.RateLimitError, match="RPD limit=2, current usage=3"): + await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) + + assert mock_cache.async_increment_cache_pipeline.await_count == 2 + refund = mock_cache.async_increment_cache_pipeline.call_args_list[1].kwargs["increment_list"] + assert [(op["key"], op["increment_value"], op["ttl"]) for op in refund] == [ + ("daily-test-id:gpt-4.1-mini:rpm:23-59", -1, 60), + ("daily-test-id:gpt-4.1-mini:rpd:2026-10-28", -1, 86400), + ] + + @pytest.mark.asyncio + async def test_async_pre_call_check_increments_rpm_and_rpd_in_one_cache_trip(self): + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_increment_cache_pipeline = AsyncMock(return_value=[1.0, 1.0]) + check = ModelRateLimitingCheck(dual_cache=mock_cache, clock=Clock(DAY_ONE)) + + assert await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT + + assert mock_cache.async_increment_cache_pipeline.await_count == 1 + operations = mock_cache.async_increment_cache_pipeline.call_args.kwargs["increment_list"] + assert [(op["key"], op["ttl"]) for op in operations] == [ + ("daily-test-id:gpt-4.1-mini:rpm:23-59", 60), + ("daily-test-id:gpt-4.1-mini:rpd:2026-10-28", 86400), + ] + assert not mock_cache.async_increment_cache.called + + def test_pre_call_check_reads_the_shared_tpd_count_from_redis(self): + redis_cache = MagicMock() + redis_cache.get_cache.return_value = 100 + deployment = {**RPD_DEPLOYMENT, "litellm_params": {"model": "gpt-4.1-mini", "tpd": 100}} + check = ModelRateLimitingCheck(dual_cache=DualCache(redis_cache=redis_cache), clock=Clock(DAY_ONE)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + _pre_call(check, deployment) + + assert "TPD limit=100, current usage=100" in str(exc_info.value) + assert redis_cache.get_cache.call_args.kwargs["key"] == "daily-test-id:gpt-4.1-mini:tpd:2026-10-28" + + def test_rpd_count_starts_fresh_after_utc_midnight(self): + clock = Clock(DAY_ONE) + check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) + _pre_call(check, RPD_DEPLOYMENT) + _pre_call(check, RPD_DEPLOYMENT) + with pytest.raises(litellm.RateLimitError): + _pre_call(check, RPD_DEPLOYMENT) + + clock.now = DAY_ONE + timedelta(seconds=30) + + assert _pre_call(check, RPD_DEPLOYMENT) == RPD_DEPLOYMENT + assert _pre_call(check, RPD_DEPLOYMENT) == RPD_DEPLOYMENT + with pytest.raises(litellm.RateLimitError) as exc_info: + _pre_call(check, RPD_DEPLOYMENT) + assert "current usage=3" in str(exc_info.value) + + @pytest.mark.parametrize( + ("now", "seconds_to_midnight"), + [ + (DAY_ONE, 30), + (datetime(2026, 10, 28, 0, 0, 0, tzinfo=timezone.utc), 86400), + (datetime(2026, 10, 28, 12, 30, 15, 500000, tzinfo=timezone.utc), 41384), + ], + ) + def test_daily_429_retry_after_is_seconds_until_utc_midnight(self, now: datetime, seconds_to_midnight: int): + deployment = {**RPD_DEPLOYMENT, "litellm_params": {"model": "gpt-4.1-mini", "rpd": 0}} + check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(now)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + _pre_call(check, deployment) + + assert exc_info.value.response.headers["retry-after"] == str(seconds_to_midnight) + + @pytest.mark.asyncio + async def test_router_serves_the_request_from_the_other_deployment_once_rpd_is_spent(self): + router = Router( + model_list=[ + { + "model_name": "daily-group", + "litellm_params": { + "model": "openai/gpt-4.1-mini", + "api_key": "k", + "mock_response": "a", + "rpd": 1, + "order": 1, + }, + "model_info": {"id": "spent-first"}, + }, + { + "model_name": "daily-group", + "litellm_params": { + "model": "openai/gpt-4.1-mini", + "api_key": "k", + "mock_response": "b", + "order": 2, + }, + "model_info": {"id": "always-free"}, + }, + ], + num_retries=1, + ) + check = ModelRateLimitingCheck(dual_cache=router.cache, clock=Clock(DAY_ONE)) + router.optional_callbacks = [check] + litellm.logging_callback_manager.add_litellm_callback(check) + try: + served = [ + (await router.acompletion(model="daily-group", messages=[{"role": "user", "content": "hi"}])) + .choices[0] + .message.content + for _ in range(4) + ] + finally: + router.discard() + + assert served == ["a", "b", "b", "b"] diff --git a/tests/unit/test_router/test_io_token_rate_limits.py b/tests/unit/test_router/test_io_token_rate_limits.py index a5a68271111..0ae22534278 100644 --- a/tests/unit/test_router/test_io_token_rate_limits.py +++ b/tests/unit/test_router/test_io_token_rate_limits.py @@ -859,7 +859,7 @@ class TestModelRateLimitingCheckIOTokens: with pytest.raises(litellm.RateLimitError): await check.async_pre_call_check(deployment) - assert await dual_cache.async_get_cache(key=rpm_key) == 6 + assert await dual_cache.async_get_cache(key=rpm_key) == 5, "the rejected request hands its rpm count back" assert (await dual_cache.async_get_cache(key=itpm_key) or 0) == 0 assert ITPM_RESERVED_KEY not in request_kwargs["metadata"] assert any("both limit types are enforced" in record.message for record in caplog.records) diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index 740342f54d9..cc1ea645012 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/tests/unit/types/test_litellm_params.py @@ -154,6 +154,8 @@ OPTION_NAMES: Final = ( "model_info", "rpm", "tpm", + "rpd", + "tpd", "itpm", "otpm", "default_api_key_rpm_limit", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index a09672d7d8f..a49f1099121 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35647,6 +35647,8 @@ export interface components { regional_processing_uplift_multiplier_us?: number | null; /** Responses Api Version */ responses_api_version?: string | null; + /** Rpd */ + rpd?: number | null; /** Rpm */ rpm?: number | null; /** S3 Access Key Id */ @@ -35699,6 +35701,8 @@ export interface components { token_exchange_profile?: string | null; /** Token Exchange Scope */ token_exchange_scope?: string | null; + /** Tpd */ + tpd?: number | null; /** Tpm */ tpm?: number | null; /** Use Chat Completions Api */ @@ -51039,6 +51043,8 @@ export interface components { regional_processing_uplift_multiplier_us?: number | null; /** Responses Api Version */ responses_api_version?: string | null; + /** Rpd */ + rpd?: number | null; /** Rpm */ rpm?: number | null; /** S3 Access Key Id */ @@ -51091,6 +51097,8 @@ export interface components { token_exchange_profile?: string | null; /** Token Exchange Scope */ token_exchange_scope?: string | null; + /** Tpd */ + tpd?: number | null; /** Tpm */ tpm?: number | null; /** Use Chat Completions Api */ From 8e8f6dc66ed693f21f60eddce1939fd179621832 Mon Sep 17 00:00:00 2001 From: tessellark <349031000+tessellark@users.noreply.github.com> Date: Wed, 28 Oct 2026 15:53:00 -0700 Subject: [PATCH 108/128] feat(tencent): add Seedream image generation and editing through TokenHub (#48735) * feat(tencent): add Seedream image generation and editing through TokenHub * feat(tencent): price Seedream image generation and edits from TokenHub usage * test(tencent): cover Seedream image generation, edits, and pricing through TokenHub * fix(tencent): keep non-Seedream image models on the OpenAI-compatible route Only seedream- models use TokenHub's wand route, so Hy-Image and other tencent image models keep working as on main. The response now carries TokenHub's image size and format, Pro's flat fallback bills its default 2K size, and the new tests are typed. * fix(tencent): bill other Tencent image models as before and forward user and TokenHub params on edits * fix(tencent): keep dropped edit params out and parse form-string watermark and seed * fix(tencent): send extra_body edit fields to TokenHub as given * fix(tencent): keep the routed model on Seedream edits and bill a deployment's per-token price * fix(tencent): accept string drop params on edits, rebuild nested form fields, and honor a per-image deployment price An image edit with a non-list additional_drop_params answered 500 for every provider after this branch validated it as a list, so a non-list value now drops nothing again, as chat treats it. Seedream edits sent through the OpenAI SDK arrive as multipart form fields, so TokenHub-only booleans and integers are parsed and bracketed keys like optimize_prompt_options[mode] are rebuilt into objects. A deployment that sets only output_cost_per_image is billed that price, and the TokenHub URL uses the shared api_base helpers. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/images/main.py | 57 ++++- litellm/images/utils.py | 40 ++- .../litellm_core_utils/llm_cost_calc/utils.py | 83 +++---- .../flux_lora_depth_transformation.py | 5 +- .../llms/fal_ai/image_edit/transformation.py | 44 +--- litellm/llms/tencent/cost_calculator.py | 47 +++- litellm/llms/tencent/image_edit/__init__.py | 0 .../llms/tencent/image_edit/transformation.py | 147 +++++++++++ .../llms/tencent/image_generation/__init__.py | 0 .../image_generation/transformation.py | 200 +++++++++++++++ ...odel_prices_and_context_window_backup.json | 33 +++ .../provider_endpoints_support_backup.json | 4 +- litellm/utils.py | 16 ++ model_prices_and_context_window.json | 33 +++ provider_endpoints_support.json | 4 +- tests/unit/images/test_main.py | 18 ++ .../unit/llms/tencent/image_edit/__init__.py | 0 .../test_tencent_image_edit_transformation.py | 232 ++++++++++++++++++ .../llms/tencent/image_generation/__init__.py | 0 ...tencent_image_generation_transformation.py | 215 ++++++++++++++++ .../tencent/test_tencent_cost_calculator.py | 128 ++++++++++ 21 files changed, 1204 insertions(+), 102 deletions(-) create mode 100644 litellm/llms/tencent/image_edit/__init__.py create mode 100644 litellm/llms/tencent/image_edit/transformation.py create mode 100644 litellm/llms/tencent/image_generation/__init__.py create mode 100644 litellm/llms/tencent/image_generation/transformation.py create mode 100644 tests/unit/llms/tencent/image_edit/__init__.py create mode 100644 tests/unit/llms/tencent/image_edit/test_tencent_image_edit_transformation.py create mode 100644 tests/unit/llms/tencent/image_generation/__init__.py create mode 100644 tests/unit/llms/tencent/image_generation/test_tencent_image_generation_transformation.py create mode 100644 tests/unit/llms/tencent/test_tencent_cost_calculator.py diff --git a/litellm/images/main.py b/litellm/images/main.py index 84156919774..d0618e5e804 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -1,9 +1,9 @@ import asyncio import contextvars import importlib -from collections.abc import Coroutine, Mapping +from collections.abc import Collection, Coroutine, Mapping from functools import partial -from typing import TYPE_CHECKING, Final, Literal, cast, overload +from typing import TYPE_CHECKING, Final, Literal, cast, get_type_hints, overload from pydantic import TypeAdapter, ValidationError @@ -61,6 +61,7 @@ from litellm.types.utils import ( ProviderRoutingHints, ) from litellm.utils import ( + EXTRA_BODY_ROUTING_KEYS, ImageResponse, ProviderConfigManager, get_llm_provider, @@ -466,7 +467,7 @@ def image_generation( litellm.LlmProviders.QWENCLOUD, litellm.LlmProviders.QWEN_AI_PLATFORM, litellm.LlmProviders.EDENAI, - ): + ) or (custom_llm_provider == litellm.LlmProviders.TENCENT and image_generation_config is not None): if image_generation_config is None: raise ValueError(f"image generation config is not supported for {custom_llm_provider}") @@ -807,6 +808,32 @@ def image_variation( return response +_IMAGE_EDIT_OPENAI_PARAMS: Final = frozenset(get_type_hints(ImageEditOptionalRequestParams)) +_DROP_PARAM_KWARGS: Final = frozenset({"drop_params", "additional_drop_params"}) +_ADDITIONAL_DROP_PARAMS: Final = TypeAdapter[tuple[str, ...]](tuple[str, ...]) + + +def _listed_drop_params(additional_drop_params: object) -> tuple[str, ...]: + try: + return _ADDITIONAL_DROP_PARAMS.validate_python(additional_drop_params) + except ValidationError: + return () + + +def _tokenhub_passthrough_edit_params( + non_default_params: Mapping[str, object], + extra_body: Mapping[str, object] | None, + mapped_params: Collection[str], + additional_drop_params: Collection[str] | None, +) -> Mapping[str, object]: + dropped: Final = frozenset(additional_drop_params or ()) + unmapped_excluded: Final = _DROP_PARAM_KWARGS | _IMAGE_EDIT_OPENAI_PARAMS | frozenset(mapped_params) | dropped + return { + **{key: value for key, value in non_default_params.items() if key not in unmapped_excluded}, + **{key: value for key, value in (extra_body or {}).items() if key not in dropped | EXTRA_BODY_ROUTING_KEYS}, + } + + @client def image_edit( image: FileTypes | list[FileTypes] | None = None, @@ -935,22 +962,24 @@ def image_edit( raise ValueError(f"image edit is not supported for {custom_llm_provider}") local_vars.update(kwargs) + provider_supported_params: Final[frozenset[str]] = frozenset( + image_edit_provider_config.get_supported_openai_params(model) + ) # Get ImageEditOptionalRequestParams with only valid parameters image_edit_optional_params: Final[ImageEditOptionalRequestParams] = ( _get_ImageEditRequestUtils().get_requested_image_edit_optional_param( local_vars, - provider_supported_params=frozenset( - image_edit_provider_config.get_supported_openai_params(model) - ).intersection(non_default_params), + provider_supported_params=provider_supported_params.intersection(non_default_params), ) ) # Get optional parameters for the responses API + additional_drop_params: Final = _listed_drop_params(kwargs.get("additional_drop_params")) image_edit_request_params: Final[dict] = _get_ImageEditRequestUtils().get_optional_params_image_edit( model=model, image_edit_provider_config=image_edit_provider_config, image_edit_optional_params=image_edit_optional_params, drop_params=kwargs.get("drop_params"), - additional_drop_params=kwargs.get("additional_drop_params"), + additional_drop_params=additional_drop_params, ) if image_edit_provider_config.use_multipart_form_data() and ( @@ -1039,7 +1068,19 @@ def image_edit( image=images, prompt=prompt, image_edit_provider_config=image_edit_provider_config, - image_edit_optional_request_params=image_edit_request_params, + image_edit_optional_request_params={ + **image_edit_request_params, + **( + _tokenhub_passthrough_edit_params( + non_default_params, + extra_body if isinstance(extra_body, dict) else None, + provider_supported_params, + additional_drop_params, + ) + if custom_llm_provider == litellm.LlmProviders.TENCENT.value + else {} + ), + }, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params, logging_obj=litellm_logging_obj, diff --git a/litellm/images/utils.py b/litellm/images/utils.py index 24454954714..4e0cde24a68 100644 --- a/litellm/images/utils.py +++ b/litellm/images/utils.py @@ -1,6 +1,9 @@ +import base64 +import os from collections.abc import Collection, Mapping from io import BufferedReader, BytesIO -from typing import Any, Final, cast, get_type_hints +from pathlib import Path +from typing import Any, Final, Protocol, cast, get_type_hints, runtime_checkable import litellm from litellm.litellm_core_utils.token_counter import get_image_type @@ -16,7 +19,7 @@ class ImageEditRequestUtils: image_edit_provider_config: BaseImageEditConfig, image_edit_optional_params: ImageEditOptionalRequestParams, drop_params: bool | None = None, - additional_drop_params: list[str] | None = None, + additional_drop_params: Collection[str] | None = None, ) -> dict: """ Get optional parameters for the image edit API. @@ -145,3 +148,36 @@ class ImageEditRequestUtils: except Exception: # If anything goes wrong, default to PNG return FILE_MIME_TYPES[FileType.PNG] + + +@runtime_checkable +class _SeekableBinaryReader(Protocol): + def tell(self) -> int: ... + + def seek(self, offset: int) -> int: ... + + def read(self) -> bytes: ... + + +def _read_image_bytes(image: object) -> bytes: + if isinstance(image, bytes): + return image + if isinstance(image, tuple): + return _read_image_bytes(image[1]) + if isinstance(image, os.PathLike): + return Path(image).read_bytes() + if isinstance(image, _SeekableBinaryReader): + position: Final = image.tell() + image.seek(0) + data: Final = image.read() + image.seek(position) + return data + raise ValueError(f"Unsupported image type for image edit: {type(image).__name__}") + + +def image_to_data_url(image: object) -> str: + if isinstance(image, str): + return image + image_bytes: Final = _read_image_bytes(image) + mime_type: Final = ImageEditRequestUtils.get_image_content_type(image_bytes) + return f"data:{mime_type};base64,{base64.b64encode(image_bytes).decode('utf-8')}" diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 3b49f239afe..4b0b875af6d 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -6,7 +6,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass, field, replace from datetime import datetime, timezone, tzinfo from types import MappingProxyType -from typing import Final, Literal, TypedDict, cast +from typing import Final, Literal, Protocol, TypedDict, cast from zoneinfo import ZoneInfo, ZoneInfoNotFoundError from typing_extensions import ReadOnly @@ -2085,6 +2085,41 @@ def calculate_image_response_web_search_cost( ) +class _ModelPricedImageCostCalculator(Protocol): + def __call__(self, model: str, image_response: ImageResponse, model_info: ModelInfo | None = None) -> float: ... + + +def _model_priced_image_cost_calculator( + custom_llm_provider: str | None, model: str +) -> _ModelPricedImageCostCalculator | None: + match custom_llm_provider: + case litellm.LlmProviders.RECRAFT.value: + from litellm.llms.recraft.cost_calculator import cost_calculator as recraft_image_cost_calculator + + return recraft_image_cost_calculator + case litellm.LlmProviders.AIML.value: + from litellm.llms.aiml.image_generation.cost_calculator import cost_calculator as aiml_image_cost_calculator + + return aiml_image_cost_calculator + case litellm.LlmProviders.COMETAPI.value: + from litellm.llms.cometapi.image_generation.cost_calculator import ( + cost_calculator as cometapi_image_cost_calculator, + ) + + return cometapi_image_cost_calculator + case litellm.LlmProviders.RUNWAYML.value: + from litellm.llms.runwayml.cost_calculator import cost_calculator as runwayml_image_cost_calculator + + return runwayml_image_cost_calculator + case litellm.LlmProviders.TENCENT.value: + from litellm.llms.tencent.cost_calculator import image_cost_calculator as tencent_image_cost_calculator + from litellm.llms.tencent.image_generation.transformation import is_tokenhub_image_model + + return tencent_image_cost_calculator if is_tokenhub_image_model(model) else None + case _: + return None + + class CostCalculatorUtils: @staticmethod def call_type_has_image_response(call_type: str) -> bool: @@ -2131,9 +2166,6 @@ class CostCalculatorUtils: from litellm.llms.gemini.image_generation.cost_calculator import ( cost_calculator as gemini_image_cost_calculator, ) - from litellm.llms.recraft.cost_calculator import ( - cost_calculator as recraft_image_cost_calculator, - ) from litellm.llms.vertex_ai.image_generation.cost_calculator import ( cost_calculator as vertex_ai_image_cost_calculator, ) @@ -2146,7 +2178,10 @@ class CostCalculatorUtils: ) resolved_n: Final = n if n is not None else (len(completion_response.data) if completion_response.data else 0) pricing: Final = deployment_pricing(model_info) + model_priced_calculator: Final = _model_priced_image_cost_calculator(custom_llm_provider, model) + if model_priced_calculator is not None: + return model_priced_calculator(model=model, image_response=completion_response, model_info=pricing) if custom_llm_provider == litellm.LlmProviders.VERTEX_AI.value: if isinstance(completion_response, ImageResponse): return vertex_ai_image_cost_calculator( @@ -2164,36 +2199,6 @@ class CostCalculatorUtils: optional_params=optional_params, ) raise TypeError("completion_response must be of type ImageResponse for bedrock image cost calculation") - elif custom_llm_provider == litellm.LlmProviders.RECRAFT.value: - from litellm.llms.recraft.cost_calculator import ( - cost_calculator as recraft_image_cost_calculator, - ) - - return recraft_image_cost_calculator( - model=model, - image_response=completion_response, - model_info=pricing, - ) - elif custom_llm_provider == litellm.LlmProviders.AIML.value: - from litellm.llms.aiml.image_generation.cost_calculator import ( - cost_calculator as aiml_image_cost_calculator, - ) - - return aiml_image_cost_calculator( - model=model, - image_response=completion_response, - model_info=pricing, - ) - elif custom_llm_provider == litellm.LlmProviders.COMETAPI.value: - from litellm.llms.cometapi.image_generation.cost_calculator import ( - cost_calculator as cometapi_image_cost_calculator, - ) - - return cometapi_image_cost_calculator( - model=model, - image_response=completion_response, - model_info=pricing, - ) elif custom_llm_provider == litellm.LlmProviders.GEMINI.value: if call_type in ( CallTypes.image_edit.value, @@ -2237,16 +2242,6 @@ class CostCalculatorUtils: optional_params=optional_params, model_info=pricing, ) - elif custom_llm_provider == litellm.LlmProviders.RUNWAYML.value: - from litellm.llms.runwayml.cost_calculator import ( - cost_calculator as runwayml_image_cost_calculator, - ) - - return runwayml_image_cost_calculator( - model=model, - image_response=completion_response, - model_info=pricing, - ) elif ( custom_llm_provider == litellm.LlmProviders.OPENAI.value or custom_llm_provider == litellm.LlmProviders.AZURE.value diff --git a/litellm/llms/fal_ai/image_edit/flux_lora_depth_transformation.py b/litellm/llms/fal_ai/image_edit/flux_lora_depth_transformation.py index 344d9d3c75b..6d0891295bd 100644 --- a/litellm/llms/fal_ai/image_edit/flux_lora_depth_transformation.py +++ b/litellm/llms/fal_ai/image_edit/flux_lora_depth_transformation.py @@ -4,13 +4,14 @@ from typing import Final from httpx._types import RequestFiles +from litellm.images.utils import image_to_data_url from litellm.litellm_core_utils.url_utils import join_api_base from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import FileTypes -from .transformation import DEFAULT_BASE_URL, FalAIImageEditConfig, to_data_url +from .transformation import DEFAULT_BASE_URL, FalAIImageEditConfig FLUX_LORA_DEPTH_ENDPOINT: Final[str] = "fal-ai/flux-lora-depth" SUPPORTED_OPENAI_PARAMS: Final[tuple[str, ...]] = ("n", "size") @@ -68,7 +69,7 @@ class FalAIFluxLoraDepthEditConfig(FalAIImageEditConfig): ) request_body: Final[dict[str, object]] = { # mutable-ok: base class contract returns a dict "prompt": prompt, - "image_url": to_data_url(next(iter(images))), + "image_url": image_to_data_url(next(iter(images))), **provider_params, } return request_body, () diff --git a/litellm/llms/fal_ai/image_edit/transformation.py b/litellm/llms/fal_ai/image_edit/transformation.py index ea4e8cac5bf..5cbf635fd47 100644 --- a/litellm/llms/fal_ai/image_edit/transformation.py +++ b/litellm/llms/fal_ai/image_edit/transformation.py @@ -1,15 +1,12 @@ -import base64 -import os from collections.abc import Mapping -from pathlib import Path from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Final import httpx from httpx._types import RequestFiles from pydantic import ConfigDict, TypeAdapter -from litellm.images.utils import ImageEditRequestUtils +from litellm.images.utils import image_to_data_url from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.fal_ai.image_generation.gpt_image_2_transformation import ( map_gpt_image_quality, @@ -39,39 +36,6 @@ PARAM_TRANSLATION: Final[Mapping[str, str]] = MappingProxyType( ) -@runtime_checkable -class _SeekableBinaryReader(Protocol): - def tell(self) -> int: ... - - def seek(self, offset: int) -> int: ... - - def read(self) -> bytes: ... - - -def _read_image_bytes(image: object) -> bytes: - if isinstance(image, bytes): - return image - if isinstance(image, tuple): - return _read_image_bytes(image[1]) - if isinstance(image, os.PathLike): - return Path(image).read_bytes() - if isinstance(image, _SeekableBinaryReader): - position: Final = image.tell() - image.seek(0) - data: Final = image.read() - image.seek(position) - return data - raise ValueError(f"Unsupported image type for Fal AI image edit: {type(image).__name__}") - - -def to_data_url(image: object) -> str: - if isinstance(image, str): - return image - image_bytes: Final = _read_image_bytes(image) - mime_type: Final = ImageEditRequestUtils.get_image_content_type(image_bytes) - return f"data:{mime_type};base64,{base64.b64encode(image_bytes).decode('utf-8')}" - - def _first(value: object) -> object: return value[0] if isinstance(value, list) and value else value @@ -146,14 +110,14 @@ class FalAIImageEditConfig(BaseImageEditConfig): raise ValueError("Fal AI image edit requires at least one input image") mask: Final = _first(image_edit_optional_request_params.get("mask")) mask_field: Final[Mapping[str, str]] = ( - MappingProxyType({"mask_url": to_data_url(mask)}) if mask is not None else MappingProxyType({}) + MappingProxyType({"mask_url": image_to_data_url(mask)}) if mask is not None else MappingProxyType({}) ) provider_params: Final[Mapping[str, object]] = MappingProxyType( {key: value for key, value in image_edit_optional_request_params.items() if key != "mask"} ) request_body: Final[dict[str, object]] = { # mutable-ok: base class contract returns a dict "prompt": prompt, - "image_urls": tuple(to_data_url(img) for img in images), + "image_urls": tuple(image_to_data_url(img) for img in images), **mask_field, **provider_params, } diff --git a/litellm/llms/tencent/cost_calculator.py b/litellm/llms/tencent/cost_calculator.py index d9aebdc3284..5d2abe3acf4 100644 --- a/litellm/llms/tencent/cost_calculator.py +++ b/litellm/llms/tencent/cost_calculator.py @@ -1,6 +1,49 @@ -from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token -from litellm.types.utils import Usage +from typing import Final + +from litellm.litellm_core_utils.llm_cost_calc.utils import ( + calculate_image_response_cost_from_usage, + flat_image_cost, + generic_cost_per_token, + resolve_image_model_info, +) +from litellm.types.utils import ImageResponse, ModelInfo, Usage def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: return generic_cost_per_token(model=model, usage=usage, custom_llm_provider="tencent") + + +def _deployment_image_token_pricing(model_info: ModelInfo | None) -> ModelInfo | None: + if model_info is None or model_info.get("output_cost_per_image_token") is not None: + return model_info + output_cost_per_token: Final = model_info.get("output_cost_per_token") + if output_cost_per_token is None: + return model_info + priced: Final[ModelInfo] = {**model_info, "output_cost_per_image_token": output_cost_per_token} + return priced + + +def _deployment_prices_per_image_only(model_info: ModelInfo | None) -> bool: + return ( + model_info is not None + and model_info.get("output_cost_per_image") is not None + and model_info.get("output_cost_per_image_token") is None + and model_info.get("output_cost_per_token") is None + ) + + +def image_cost_calculator(model: str, image_response: ImageResponse, model_info: ModelInfo | None = None) -> float: + resolved_model_info: Final = resolve_image_model_info( + model=model, custom_llm_provider="tencent", model_info=_deployment_image_token_pricing(model_info) + ) + if _deployment_prices_per_image_only(model_info): + return flat_image_cost(resolved_model_info, image_response) + token_based_cost: Final = calculate_image_response_cost_from_usage( + model=model, + image_response=image_response, + custom_llm_provider="tencent", + model_info=resolved_model_info, + ) + if token_based_cost is not None: + return token_based_cost + return flat_image_cost(resolved_model_info, image_response) diff --git a/litellm/llms/tencent/image_edit/__init__.py b/litellm/llms/tencent/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/tencent/image_edit/transformation.py b/litellm/llms/tencent/image_edit/transformation.py new file mode 100644 index 00000000000..3033d161ce7 --- /dev/null +++ b/litellm/llms/tencent/image_edit/transformation.py @@ -0,0 +1,147 @@ +from collections.abc import Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, TypeAlias + +import httpx +from httpx._types import RequestFiles +from pydantic import TypeAdapter +from typing_extensions import override + +from litellm.images.utils import image_to_data_url +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig +from litellm.llms.tencent.image_generation.transformation import ( + TOKENHUB_IMAGE_PARAMS, + tokenhub_image_headers, + tokenhub_image_params, + tokenhub_image_response, + tokenhub_image_url, +) +from litellm.types.images.main import ImageEditOptionalRequestParams +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import ImageResponse + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +_IMAGES: Final = TypeAdapter(tuple[object, ...]) +_FORM_BOOLEANS: Final = MappingProxyType({"true": True, "false": False}) +_OPENAI_PARAMS: Final = frozenset(TOKENHUB_IMAGE_PARAMS) + +_FormEntries: TypeAlias = tuple[tuple[tuple[str, ...], object], ...] + + +def _form_scalar(value: object) -> object: + match value: + case str() if value.lower() in _FORM_BOOLEANS: + return _FORM_BOOLEANS[value.lower()] + case str() if value.removeprefix("-").isdigit(): + return int(value) + case _: + return value + + +def _form_path(key: str) -> tuple[str, ...]: + root, bracket, rest = key.partition("[") + if not bracket or not rest.endswith("]"): + return (key,) + return (root, *rest.removesuffix("]").split("][")) + + +def _nested_form_value(entries: _FormEntries) -> object: + match entries: + case (((), value),): + return _form_scalar(value) + case _: + return _nested_form_fields(tuple(entry for entry in entries if entry[0])) + + +def _nested_form_fields(entries: _FormEntries) -> Mapping[str, object]: + roots: Final = tuple(dict.fromkeys(path[0] for path, _ in entries)) + return { + root: _nested_form_value(tuple((path[1:], value) for path, value in entries if path[0] == root)) + for root in roots + } + + +def _tokenhub_form_params(params: Mapping[str, object]) -> Mapping[str, object]: + tokenhub_entries: Final = tuple( + (_form_path(key), value) for key, value in params.items() if key not in _OPENAI_PARAMS + ) + return { + **{key: value for key, value in params.items() if key in _OPENAI_PARAMS}, + **_nested_form_fields(tokenhub_entries), + } + + +class TencentImageEditConfig(BaseImageEditConfig): + @override + def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: mirrors override contract + return list(TOKENHUB_IMAGE_PARAMS) + + @override + def map_openai_params( + self, + image_edit_optional_params: ImageEditOptionalRequestParams, + model: str, + drop_params: bool, + ) -> dict[str, object]: # mutable-ok: mirrors override contract + return dict(tokenhub_image_params(image_edit_optional_params, model, drop_params)) + + @override + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + api_key: str | None = None, + litellm_params: Mapping[str, object] | None = None, + api_base: str | None = None, + ) -> dict[str, str]: # mutable-ok: mirrors override contract + return dict(tokenhub_image_headers(headers, api_key, model)) + + @override + def get_complete_url( + self, + model: str, + api_base: str | None, + litellm_params: Mapping[str, object], + ) -> str: + return tokenhub_image_url(api_base) + + @override + def use_multipart_form_data(self) -> bool: + return False + + @override + def transform_image_edit_request( + self, + model: str, + prompt: str | None, + image: object, + image_edit_optional_request_params: Mapping[str, object], + litellm_params: GenericLiteLLMParams, + headers: Mapping[str, str], + ) -> tuple[dict[str, object], RequestFiles]: # mutable-ok: mirrors override contract + images: Final = _IMAGES.validate_python(image) if isinstance(image, list) else (image,) + return { + "model": model, + "prompt": prompt, + "images": tuple(image_to_data_url(img) for img in images if img is not None), + **_tokenhub_form_params(image_edit_optional_request_params), + }, () + + @override + def transform_image_edit_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: "LiteLLMLoggingObj", + ) -> ImageResponse: + image_response: Final = tokenhub_image_response(raw_response) + if image_response is None: + raise BaseLLMException( + status_code=raw_response.status_code, + message=raw_response.text, + headers=raw_response.headers, + ) + return image_response diff --git a/litellm/llms/tencent/image_generation/__init__.py b/litellm/llms/tencent/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/tencent/image_generation/transformation.py b/litellm/llms/tencent/image_generation/transformation.py new file mode 100644 index 00000000000..fb4b102e8be --- /dev/null +++ b/litellm/llms/tencent/image_generation/transformation.py @@ -0,0 +1,200 @@ +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Final + +import httpx +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from typing_extensions import override + +from litellm.exceptions import AuthenticationError, UnsupportedParamsError +from litellm.litellm_core_utils.url_utils import ensure_api_route, join_api_base +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.image_generation.transformation import ( + BaseImageGenerationConfig, +) +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import ( + AllMessageValues, + OpenAIImageGenerationOptionalParams, +) +from litellm.types.utils import ( + ImageObject, + ImageResponse, + ImageUsage, + ImageUsageInputTokensDetails, + LlmProviders, +) + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +TOKENHUB_API_BASE: Final = "https://tokenhub-intl.tencentcloudmaas.com/v1" +TOKENHUB_IMAGE_PATH: Final = "/wand/si-image/generation" +TOKENHUB_IMAGE_MODEL_PREFIX: Final = "seedream-" +TOKENHUB_IMAGE_PARAMS: Final[tuple[OpenAIImageGenerationOptionalParams, ...]] = ( + "n", + "size", + "response_format", + "output_format", + "background", + "user", +) + +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + + +class _TokenHubImage(BaseModel): + model_config = ConfigDict(frozen=True) + + url: str | None = None + b64_json: str | None = None + size: str | None = None + output_format: str | None = None + + +class _TokenHubUsage(BaseModel): + model_config = ConfigDict(frozen=True) + + total_tokens: int + + +class _TokenHubImageResponse(BaseModel): + model_config = ConfigDict(frozen=True, hide_input_in_errors=True) + + created: int + data: tuple[_TokenHubImage, ...] + tokenhub_usage: _TokenHubUsage | None = None + + +def is_tokenhub_image_model(model: str) -> bool: + return model.removeprefix(f"{LlmProviders.TENCENT.value}/").startswith(TOKENHUB_IMAGE_MODEL_PREFIX) + + +def tokenhub_image_url(api_base: str | None) -> str: + base: Final = api_base or get_secret_str("TENCENT_API_BASE") or TOKENHUB_API_BASE + return join_api_base(ensure_api_route(base, "v1"), TOKENHUB_IMAGE_PATH) + + +def tokenhub_image_headers(headers: Mapping[str, str], api_key: str | None, model: str) -> Mapping[str, str]: + resolved_api_key: Final = api_key or get_secret_str("TENCENT_API_KEY") + if not resolved_api_key: + raise AuthenticationError( + message="TENCENT_API_KEY is not set. Set it in the environment or pass api_key.", + llm_provider="tencent", + model=model, + ) + return {**headers, "Authorization": f"Bearer {resolved_api_key}", "Content-Type": "application/json"} + + +def tokenhub_image_params(params: Mapping[str, object], model: str, drop_params: bool) -> Mapping[str, object]: + n: Final = params.get("n") + if n is not None and n != 1 and not drop_params: + raise UnsupportedParamsError( + status_code=400, + message=f"Tencent TokenHub returns one image per request, so n={n} is not supported for {model}. " + + "To send n=1 instead, set `litellm.drop_params = True`.", + ) + return { + key: value + for key, value in params.items() + if key != "n" and value is not None and not (key == "background" and value == "auto") + } + + +def tokenhub_image_response(raw_response: httpx.Response) -> ImageResponse | None: + try: + parsed: Final = _TokenHubImageResponse.model_validate_json(raw_response.content) + except ValidationError: + return None + total_tokens: Final = parsed.tokenhub_usage.total_tokens if parsed.tokenhub_usage is not None else 0 + first_image: Final = parsed.data[0] if parsed.data else _TokenHubImage() + return ImageResponse( + created=parsed.created, + data=[ImageObject(url=image.url, b64_json=image.b64_json) for image in parsed.data], + size=first_image.size, + output_format=first_image.output_format, + usage=ImageUsage( + input_tokens=0, + input_tokens_details=ImageUsageInputTokensDetails(image_tokens=0, text_tokens=0), + output_tokens=total_tokens, + total_tokens=total_tokens, + ), + ) + + +class TencentImageGenerationConfig(BaseImageGenerationConfig): + @override + def get_supported_openai_params( + self, model: str + ) -> list[OpenAIImageGenerationOptionalParams]: # mutable-ok: mirrors override contract + return list(TOKENHUB_IMAGE_PARAMS) + + @override + def map_openai_params( + self, + non_default_params: Mapping[str, object], + optional_params: Mapping[str, object], + model: str, + drop_params: bool, + ) -> dict[str, object]: # mutable-ok: mirrors override contract + return {**optional_params, **tokenhub_image_params(non_default_params, model, drop_params)} + + @override + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + stream: bool | None = None, + ) -> str: + return tokenhub_image_url(api_base) + + @override + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, str]: # mutable-ok: mirrors override contract + return dict(tokenhub_image_headers(headers, api_key, model)) + + @override + def transform_image_generation_request( + self, + model: str, + prompt: str, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + headers: Mapping[str, str], + ) -> dict[str, object]: # mutable-ok: mirrors override contract + extra_body: Final = _JSON_OBJECT.validate_python(optional_params.get("extra_body") or {}) + params: Final = {key: value for key, value in optional_params.items() if key != "extra_body"} + return {"model": model, "prompt": prompt, **params, **extra_body} + + @override + def transform_image_generation_response( + self, + model: str, + raw_response: httpx.Response, + model_response: ImageResponse, + logging_obj: "LiteLLMLoggingObj", + request_data: Mapping[str, object], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + encoding: object, + api_key: str | None = None, + json_mode: bool | None = None, + ) -> ImageResponse: + image_response: Final = tokenhub_image_response(raw_response) + if image_response is None: + raise BaseLLMException( + status_code=raw_response.status_code, + message=raw_response.text, + headers=raw_response.headers, + ) + return image_response diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1f2bd7d7837..dc043dc4d26 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -61986,6 +61986,39 @@ "supports_reasoning": true, "supports_vision": false }, + "tencent/seedream-image-v5.0-flash": { + "litellm_provider": "tencent", + "mode": "image_generation", + "output_cost_per_image": 0.018, + "output_cost_per_image_token": 1.6e-06, + "source": "https://cloud.tencent.com/document/product/1823/130055", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ] + }, + "tencent/seedream-image-v5.0-lite": { + "litellm_provider": "tencent", + "mode": "image_generation", + "output_cost_per_image": 0.035, + "output_cost_per_image_token": 1.6e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ] + }, + "tencent/seedream-image-v5.0-pro": { + "litellm_provider": "tencent", + "mode": "image_generation", + "output_cost_per_image": 0.09, + "output_cost_per_image_token": 1.6e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ] + }, "cognition/swe-1.6": { "input_cost_per_token": 5e-07, "output_cost_per_token": 2.5e-06, diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index d73ac19bae1..0a6c2bcb2a9 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -3793,8 +3793,8 @@ "responses": true, "text_completion": true, "embeddings": false, - "image_generations": false, - "image_edits": false, + "image_generations": true, + "image_edits": true, "audio_transcriptions": false, "audio_speech": false, "rerank": false, diff --git a/litellm/utils.py b/litellm/utils.py index a231644addf..1ef39059e1b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -10418,6 +10418,13 @@ class ProviderConfigManager: return get_modelscope_image_generation_config(model) elif LlmProviders.EDENAI == provider: return litellm.EdenAIImageGenerationConfig() + elif LlmProviders.TENCENT == provider: + from litellm.llms.tencent.image_generation.transformation import ( + TencentImageGenerationConfig, + is_tokenhub_image_model, + ) + + return TencentImageGenerationConfig() if is_tokenhub_image_model(model) else None return None @staticmethod @@ -10606,6 +10613,15 @@ class ProviderConfigManager: ) return get_openrouter_image_edit_config(model) + elif LlmProviders.TENCENT == provider: + from litellm.llms.tencent.image_edit.transformation import ( + TencentImageEditConfig, + ) + from litellm.llms.tencent.image_generation.transformation import ( + is_tokenhub_image_model, + ) + + return TencentImageEditConfig() if is_tokenhub_image_model(model) else None return None @staticmethod diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 1f2bd7d7837..dc043dc4d26 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -61986,6 +61986,39 @@ "supports_reasoning": true, "supports_vision": false }, + "tencent/seedream-image-v5.0-flash": { + "litellm_provider": "tencent", + "mode": "image_generation", + "output_cost_per_image": 0.018, + "output_cost_per_image_token": 1.6e-06, + "source": "https://cloud.tencent.com/document/product/1823/130055", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ] + }, + "tencent/seedream-image-v5.0-lite": { + "litellm_provider": "tencent", + "mode": "image_generation", + "output_cost_per_image": 0.035, + "output_cost_per_image_token": 1.6e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ] + }, + "tencent/seedream-image-v5.0-pro": { + "litellm_provider": "tencent", + "mode": "image_generation", + "output_cost_per_image": 0.09, + "output_cost_per_image_token": 1.6e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ] + }, "cognition/swe-1.6": { "input_cost_per_token": 5e-07, "output_cost_per_token": 2.5e-06, diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index d73ac19bae1..0a6c2bcb2a9 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -3793,8 +3793,8 @@ "responses": true, "text_completion": true, "embeddings": false, - "image_generations": false, - "image_edits": false, + "image_generations": true, + "image_edits": true, "audio_transcriptions": false, "audio_speech": false, "rerank": false, diff --git a/tests/unit/images/test_main.py b/tests/unit/images/test_main.py index 88d5ae1fc2e..29424e50079 100644 --- a/tests/unit/images/test_main.py +++ b/tests/unit/images/test_main.py @@ -357,6 +357,24 @@ def test_image_edit_prices_a_vertex_deployment_at_its_configured_location( assert cost_at("us-central1") == pytest.approx(0.044) +def test_image_edit_accepts_a_string_additional_drop_params(respx_mock: respx.MockRouter) -> None: + edits_route: Final = respx_mock.post("https://api.openai.com/v1/images/edits").mock( + return_value=httpx.Response(status_code=200, json={"created": 1792551838, "data": [{"b64_json": "aW1n"}]}) + ) + + response: Final = litellm.image_edit( + model="openai/gpt-image-2.5-sunburst", + image=b"\x89PNG\r\n\x1a\nfakepng", + prompt="make the circle blue", + api_key="sk-test", + quality="high", + additional_drop_params="quality", + ) + + assert edits_route.called + assert response.data[0].b64_json == "aW1n" + + class TestCustomLogger(CustomLogger): __test__ = False diff --git a/tests/unit/llms/tencent/image_edit/__init__.py b/tests/unit/llms/tencent/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/tencent/image_edit/test_tencent_image_edit_transformation.py b/tests/unit/llms/tencent/image_edit/test_tencent_image_edit_transformation.py new file mode 100644 index 00000000000..506c43f6cba --- /dev/null +++ b/tests/unit/llms/tencent/image_edit/test_tencent_image_edit_transformation.py @@ -0,0 +1,232 @@ +import base64 +import io +import json +from typing import Final + +import httpx +import pytest + +import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler + +PNG_BYTES: Final = b"\x89PNG\r\n\x1a\n" + b"\x00" * 32 +TOKENHUB_RESPONSE: Final[dict[str, object]] = { + "created": 1792551838, + "data": [{"output_format": "jpeg", "size": "2048x2048", "url": "https://cos.example/edited.jpg"}], + "model": "doubao-seedream-5-0-lite", + "tokenhub_usage": {"total_tokens": 21875}, + "request_id": "4f795d48-b227-4c35-af2c-1a0ba288b6b6", +} + + +def _client(captured: dict[str, object]) -> HTTPHandler: + def respond(request: httpx.Request) -> httpx.Response: + captured["url"] = str(request.url) + captured["content_type"] = request.headers.get("content-type") + captured["authorization"] = request.headers.get("authorization") + captured["body"] = json.loads(request.content) + return httpx.Response(200, json=TOKENHUB_RESPONSE) + + return HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))) + + +def test_edit_sends_uploaded_images_as_data_urls_in_a_json_body() -> None: + captured: Final[dict[str, object]] = {} + png_data_url: Final = f"data:image/png;base64,{base64.b64encode(PNG_BYTES).decode()}" + + response: Final = litellm.image_edit( + model="tencent/seedream-image-v5.0-lite", + image=[PNG_BYTES, io.BytesIO(PNG_BYTES), "https://ref.example/style.png"], + prompt="turn the square into a paper lantern", + size="2K", + api_key="sk-tokenhub", + client=_client(captured), + ) + + assert captured["url"] == "https://tokenhub-intl.tencentcloudmaas.com/v1/wand/si-image/generation" + assert captured["content_type"] == "application/json" + assert captured["authorization"] == "Bearer sk-tokenhub" + assert captured["body"] == { + "model": "seedream-image-v5.0-lite", + "prompt": "turn the square into a paper lantern", + "images": [png_data_url, png_data_url, "https://ref.example/style.png"], + "size": "2K", + } + assert response.data[0].url == "https://cos.example/edited.jpg" + assert response.usage.total_tokens == 21875 + + +def test_edit_rejects_a_mask_without_drop_params() -> None: + captured: Final[dict[str, object]] = {} + + with pytest.raises(litellm.UnsupportedParamsError): + litellm.image_edit( + model="tencent/seedream-image-v5.0-pro", + image=PNG_BYTES, + mask=PNG_BYTES, + prompt="turn the square into a paper lantern", + api_key="sk-tokenhub", + client=_client(captured), + ) + + assert captured == {} + + +def test_edit_drops_a_mask_under_drop_params() -> None: + captured: Final[dict[str, object]] = {} + + litellm.image_edit( + model="tencent/seedream-image-v5.0-pro", + image=PNG_BYTES, + mask=PNG_BYTES, + prompt="turn the square into a paper lantern", + api_key="sk-tokenhub", + client=_client(captured), + drop_params=True, + ) + + assert captured["body"] == { + "model": "seedream-image-v5.0-pro", + "prompt": "turn the square into a paper lantern", + "images": [f"data:image/png;base64,{base64.b64encode(PNG_BYTES).decode()}"], + } + + +def test_edit_is_not_offered_for_tencent_models_outside_seedream() -> None: + captured: Final[dict[str, object]] = {} + + with pytest.raises(Exception, match="image edit is not supported for tencent"): + litellm.image_edit( + model="tencent/hy-image-v3.5-preview", + image=PNG_BYTES, + prompt="turn the square into a paper lantern", + api_key="sk-tokenhub", + client=_client(captured), + ) + + assert captured == {} + + +@pytest.mark.parametrize( + ("form_params", "expected_body_params"), + [ + ({"watermark": "false", "seed": "7"}, {"watermark": False, "seed": 7}), + ({"watermark": "True", "seed": "-1"}, {"watermark": True, "seed": -1}), + ({"watermark": "maybe"}, {"watermark": "maybe"}), + ( + {"layer_decomposition": "true", "optimize_prompt_options[mode]": "fast"}, + {"layer_decomposition": True, "optimize_prompt_options": {"mode": "fast"}}, + ), + ( + {"sequential_image_generation_options[max_images]": "3", "user": "12345"}, + {"sequential_image_generation_options": {"max_images": 3}, "user": "12345"}, + ), + ], +) +def test_edit_parses_tokenhub_values_sent_as_form_strings( + form_params: dict[str, object], expected_body_params: dict[str, object] +) -> None: + captured: Final[dict[str, object]] = {} + + litellm.image_edit( + model="tencent/seedream-image-v5.0-flash", + image=PNG_BYTES, + prompt="turn the square into a paper lantern", + api_key="sk-tokenhub", + client=_client(captured), + **form_params, + ) + + assert captured["body"] == { + "model": "seedream-image-v5.0-flash", + "prompt": "turn the square into a paper lantern", + "images": [f"data:image/png;base64,{base64.b64encode(PNG_BYTES).decode()}"], + **expected_body_params, + } + + +@pytest.mark.parametrize( + ("params", "expected_body_params"), + [ + ({"background": "auto", "watermark": False}, {"watermark": False}), + ({"watermark": False, "seed": 7, "additional_drop_params": ["watermark"]}, {"seed": 7}), + ({"extra_body": {"watermark": False}, "additional_drop_params": ["watermark"]}, {}), + ], +) +def test_edit_keeps_dropped_params_out_of_the_tokenhub_body( + params: dict[str, object], expected_body_params: dict[str, object] +) -> None: + captured: Final[dict[str, object]] = {} + + litellm.image_edit( + model="tencent/seedream-image-v5.0-flash", + image=PNG_BYTES, + prompt="turn the square into a paper lantern", + api_key="sk-tokenhub", + client=_client(captured), + **params, + ) + + assert captured["body"] == { + "model": "seedream-image-v5.0-flash", + "prompt": "turn the square into a paper lantern", + "images": [f"data:image/png;base64,{base64.b64encode(PNG_BYTES).decode()}"], + **expected_body_params, + } + + +@pytest.mark.parametrize( + ("params", "expected_body_params"), + [ + ({"extra_body": {"response_format": "b64_json"}}, {"response_format": "b64_json"}), + ({"size": "1K", "extra_body": {"size": "2K"}}, {"size": "2K"}), + ({"extra_body": {"model": "seedream-image-v5.0-pro", "seed": 7}}, {"seed": 7}), + ({"extra_body": '{"seed": 7}'}, {}), + ], +) +def test_edit_sends_extra_body_fields_as_given( + params: dict[str, object], expected_body_params: dict[str, object] +) -> None: + captured: Final[dict[str, object]] = {} + + litellm.image_edit( + model="tencent/seedream-image-v5.0-flash", + image=PNG_BYTES, + prompt="turn the square into a paper lantern", + api_key="sk-tokenhub", + client=_client(captured), + **params, + ) + + assert captured["body"] == { + "model": "seedream-image-v5.0-flash", + "prompt": "turn the square into a paper lantern", + "images": [f"data:image/png;base64,{base64.b64encode(PNG_BYTES).decode()}"], + **expected_body_params, + } + + +def test_edit_forwards_tokenhub_only_params_and_extra_body() -> None: + captured: Final[dict[str, object]] = {} + + litellm.image_edit( + model="tencent/seedream-image-v5.0-flash", + image=PNG_BYTES, + prompt="turn the square into a paper lantern", + api_key="sk-tokenhub", + client=_client(captured), + user="end-user-42", + watermark=False, + seed=7, + extra_body={"sequential_image_generation": "disabled"}, + ) + + assert captured["body"] == { + "model": "seedream-image-v5.0-flash", + "prompt": "turn the square into a paper lantern", + "images": [f"data:image/png;base64,{base64.b64encode(PNG_BYTES).decode()}"], + "user": "end-user-42", + "watermark": False, + "seed": 7, + "sequential_image_generation": "disabled", + } diff --git a/tests/unit/llms/tencent/image_generation/__init__.py b/tests/unit/llms/tencent/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/tencent/image_generation/test_tencent_image_generation_transformation.py b/tests/unit/llms/tencent/image_generation/test_tencent_image_generation_transformation.py new file mode 100644 index 00000000000..ab16382cdae --- /dev/null +++ b/tests/unit/llms/tencent/image_generation/test_tencent_image_generation_transformation.py @@ -0,0 +1,215 @@ +import json +from typing import Final + +import httpx +import openai +import pytest + +import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler + +TOKENHUB_RESPONSE: Final[dict[str, object]] = { + "created": 1792551825, + "data": [{"output_format": "jpeg", "size": "1248x832", "url": "https://cos.example/lantern.jpg"}], + "model": "doubao-seedream-5-0-pro-260628", + "tokenhub_usage": {"total_tokens": 28125}, + "request_id": "c57b0374-b5fe-47db-b4b2-751315f920aa", +} + + +def _client(captured: dict[str, object], status_code: int = 200, body: dict[str, object] | None = None) -> HTTPHandler: + def respond(request: httpx.Request) -> httpx.Response: + captured["url"] = str(request.url) + captured["authorization"] = request.headers.get("authorization") + captured["body"] = json.loads(request.content) + return httpx.Response(status_code, json=TOKENHUB_RESPONSE if body is None else body) + + return HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))) + + +def test_generation_posts_openai_and_tokenhub_params_to_the_wand_route() -> None: + captured: Final[dict[str, object]] = {} + + response: Final = litellm.image_generation( + model="tencent/seedream-image-v5.0-pro", + prompt="a red paper lantern", + size="2K", + response_format="url", + api_key="sk-tokenhub", + client=_client(captured), + images=["https://ref.example/lantern.png"], + watermark=False, + ) + + assert captured["url"] == "https://tokenhub-intl.tencentcloudmaas.com/v1/wand/si-image/generation" + assert captured["authorization"] == "Bearer sk-tokenhub" + assert captured["body"] == { + "model": "seedream-image-v5.0-pro", + "prompt": "a red paper lantern", + "size": "2K", + "response_format": "url", + "images": ["https://ref.example/lantern.png"], + "watermark": False, + } + assert response.data[0].url == "https://cos.example/lantern.jpg" + assert response.size == "1248x832" + assert response.output_format == "jpeg" + assert response.created == 1792551825 + assert response.usage.total_tokens == 28125 + assert response.usage.output_tokens == 28125 + + +def test_generation_returns_base64_images() -> None: + captured: Final[dict[str, object]] = {} + body: Final[dict[str, object]] = { + **TOKENHUB_RESPONSE, + "data": [{"output_format": "png", "size": "1024x1024", "b64_json": "aW1n"}], + } + + response: Final = litellm.image_generation( + model="tencent/seedream-image-v5.0-flash", + prompt="a red paper lantern", + response_format="b64_json", + api_key="sk-tokenhub", + client=_client(captured, body=body), + ) + + assert response.data[0].b64_json == "aW1n" + assert response.data[0].url is None + + +@pytest.mark.parametrize( + ("api_base", "expected_url"), + [ + ("https://tokenhub.example", "https://tokenhub.example/v1/wand/si-image/generation"), + ("https://tokenhub.example/v1", "https://tokenhub.example/v1/wand/si-image/generation"), + ("https://tokenhub.example/v1/", "https://tokenhub.example/v1/wand/si-image/generation"), + ], +) +def test_generation_resolves_a_configured_api_base_to_the_wand_route(api_base: str, expected_url: str) -> None: + captured: Final[dict[str, object]] = {} + + litellm.image_generation( + model="tencent/seedream-image-v5.0-lite", + prompt="a red paper lantern", + api_key="sk-tokenhub", + api_base=api_base, + client=_client(captured), + ) + + assert captured["url"] == expected_url + + +@pytest.mark.parametrize( + ("params", "expected_body_params"), + [ + ({"n": 1}, {}), + ({"n": 3, "drop_params": True}, {}), + ({"background": "auto"}, {}), + ({"background": "transparent", "output_format": "png"}, {"background": "transparent", "output_format": "png"}), + ({"user": "end-user-42"}, {"user": "end-user-42"}), + ], +) +def test_generation_sends_only_values_tokenhub_accepts( + params: dict[str, object], expected_body_params: dict[str, object] +) -> None: + captured: Final[dict[str, object]] = {} + + litellm.image_generation( + model="tencent/seedream-image-v5.0-pro", + prompt="a red paper lantern", + api_key="sk-tokenhub", + client=_client(captured), + **params, + ) + + assert captured["body"] == { + "model": "seedream-image-v5.0-pro", + "prompt": "a red paper lantern", + **expected_body_params, + } + + +def test_generation_rejects_more_than_one_image_without_drop_params() -> None: + captured: Final[dict[str, object]] = {} + + with pytest.raises(litellm.UnsupportedParamsError) as exc_info: + litellm.image_generation( + model="tencent/seedream-image-v5.0-pro", + prompt="a red paper lantern", + n=2, + api_key="sk-tokenhub", + client=_client(captured), + ) + + assert exc_info.value.status_code == 400 + assert captured == {} + + +@pytest.mark.parametrize( + ("error", "expected_message"), + [ + ( + { + "error": { + "message": "The model or service ID 'seedream-image-v5.0-pro' does not exist.", + "code": "400004", + "type": "gateway_error", + } + }, + "does not exist", + ), + ({"Type": 2, "Code": 1001, "Msg": "Invalid param", "Desc": ""}, "Invalid param"), + ], +) +def test_generation_surfaces_the_tokenhub_error_message(error: dict[str, object], expected_message: str) -> None: + with pytest.raises(litellm.BadRequestError) as exc_info: + litellm.image_generation( + model="tencent/seedream-image-v5.0-pro", + prompt="a red paper lantern", + api_key="sk-tokenhub", + client=_client({}, status_code=400, body=error), + ) + + assert expected_message in str(exc_info.value) + + +def test_generation_without_a_tokenhub_key_is_an_authentication_error(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("TENCENT_API_KEY", raising=False) + captured: Final[dict[str, object]] = {} + + with pytest.raises(litellm.AuthenticationError) as exc_info: + litellm.image_generation( + model="tencent/seedream-image-v5.0-pro", + prompt="a red paper lantern", + client=_client(captured), + ) + + assert exc_info.value.status_code == 401 + assert "TENCENT_API_KEY" in str(exc_info.value) + assert captured == {} + + +def test_generation_keeps_other_tencent_image_models_on_the_openai_compatible_route() -> None: + captured: Final[dict[str, object]] = {} + + def respond(request: httpx.Request) -> httpx.Response: + captured["url"] = str(request.url) + captured["body"] = json.loads(request.content) + return httpx.Response(200, json={"created": 1792551825, "data": [{"url": "https://cos.example/hy.png"}]}) + + response: Final = litellm.image_generation( + model="tencent/hy-image-v3.5-preview", + prompt="a red paper lantern", + quality="high", + api_key="sk-tokenhub", + client=openai.OpenAI( + api_key="sk-tokenhub", + base_url="https://tokenhub-intl.tencentcloudmaas.com/v1", + http_client=httpx.Client(transport=httpx.MockTransport(respond)), + ), + ) + + assert captured["url"] == "https://tokenhub-intl.tencentcloudmaas.com/v1/images/generations" + assert captured["body"] == {"model": "hy-image-v3.5-preview", "prompt": "a red paper lantern", "quality": "high"} + assert response.data[0].url == "https://cos.example/hy.png" diff --git a/tests/unit/llms/tencent/test_tencent_cost_calculator.py b/tests/unit/llms/tencent/test_tencent_cost_calculator.py new file mode 100644 index 00000000000..7f98ee02810 --- /dev/null +++ b/tests/unit/llms/tencent/test_tencent_cost_calculator.py @@ -0,0 +1,128 @@ +from typing import Final + +import httpx +import openai +import pytest + +import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.types.utils import ImageObject, ImageResponse, ImageUsage, ImageUsageInputTokensDetails + +SEEDREAM_MODELS: Final = ( + "tencent/seedream-image-v5.0-pro", + "tencent/seedream-image-v5.0-flash", + "tencent/seedream-image-v5.0-lite", +) + + +def _image_response(total_tokens: int, images: int = 1) -> ImageResponse: + return ImageResponse( + created=1792551825, + data=[ImageObject(url=f"https://cos.example/{index}.jpg") for index in range(images)], + usage=ImageUsage( + input_tokens=0, + input_tokens_details=ImageUsageInputTokensDetails(image_tokens=0, text_tokens=0), + output_tokens=total_tokens, + total_tokens=total_tokens, + ), + ) + + +@pytest.mark.parametrize("model", SEEDREAM_MODELS) +@pytest.mark.parametrize("call_type", ["image_generation", "image_edit"]) +def test_seedream_bills_tokenhub_usage_tokens_for_a_tier_size( + local_model_cost_map: None, model: str, call_type: str +) -> None: + token_rate: Final = litellm.model_cost[model]["output_cost_per_image_token"] + + cost: Final = litellm.completion_cost( + completion_response=_image_response(total_tokens=56250), + model=model, + call_type=call_type, + size="2K", + ) + + assert cost == pytest.approx(56250 * token_rate) + + +@pytest.mark.parametrize("model", SEEDREAM_MODELS) +def test_seedream_falls_back_to_the_per_image_price_without_usage(local_model_cost_map: None, model: str) -> None: + per_image_price: Final = litellm.model_cost[model]["output_cost_per_image"] + + cost: Final = litellm.completion_cost( + completion_response=_image_response(total_tokens=0, images=2), + model=model, + call_type="image_generation", + size="2K", + ) + + assert cost == pytest.approx(2 * per_image_price) + + +async def test_a_deployment_priced_tencent_image_model_outside_seedream_bills_its_deployment_price() -> None: + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"created": 1792551825, "data": [{"url": "https://cos.example/hy.png"}]}) + + router: Final = litellm.Router( + model_list=[ + { + "model_name": "hy-image", + "litellm_params": {"model": "tencent/hy-image-v3.5-preview", "api_key": "sk-tokenhub"}, + "model_info": {"input_cost_per_image": 0.04}, + } + ] + ) + + response: Final = await router.aimage_generation( + model="hy-image", + prompt="a red paper lantern", + client=openai.AsyncOpenAI( + api_key="sk-tokenhub", + base_url="https://tokenhub-intl.tencentcloudmaas.com/v1", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(respond)), + ), + ) + + assert response._hidden_params["response_cost"] == pytest.approx(0.04) + + +@pytest.mark.parametrize( + ("deployment_price", "expected_cost"), + [ + ({"output_cost_per_token": 0.000001}, 11250 * 0.000001), + ({"output_cost_per_image": 0.05}, 0.05), + ], +) +def test_a_seedream_deployment_priced_in_the_ui_bills_its_own_rate( + deployment_price: dict[str, float], expected_cost: float +) -> None: + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "created": 1792551825, + "data": [{"output_format": "jpeg", "size": "1024x1024", "url": "https://cos.example/0.jpg"}], + "tokenhub_usage": {"total_tokens": 11250}, + }, + ) + + router: Final = litellm.Router( + model_list=[ + { + "model_name": "seedream-flash", + "litellm_params": { + "model": "tencent/seedream-image-v5.0-flash", + "api_key": "sk-tokenhub", + **deployment_price, + }, + } + ] + ) + + response: Final = router.image_generation( + model="seedream-flash", + prompt="a red paper lantern", + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))), + ) + + assert response._hidden_params["response_cost"] == pytest.approx(expected_cost) From 6063187fea3715087bddb7d23e7eaefaa5dea82d Mon Sep 17 00:00:00 2001 From: quietjules <349048734+quietjules@users.noreply.github.com> Date: Wed, 28 Oct 2026 15:53:06 -0700 Subject: [PATCH 109/128] test(spend): audit #47506 billed background cost reaching the creator's loggers (#48919) Adds integration cells for a different key reading the response, message redaction, a rotated, deleted, or team-deleted creator, and a creator whose Langfuse keys resolve to nothing Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../test_background_response_billed_once.py | 183 ++++++++++++++++-- 1 file changed, 172 insertions(+), 11 deletions(-) diff --git a/tests/integration/spend/test_background_response_billed_once.py b/tests/integration/spend/test_background_response_billed_once.py index b207e46204a..5f948fac3ab 100644 --- a/tests/integration/spend/test_background_response_billed_once.py +++ b/tests/integration/spend/test_background_response_billed_once.py @@ -1,6 +1,7 @@ import asyncio import base64 import json +import math import uuid from collections.abc import Callable, Mapping from concurrent.futures import ThreadPoolExecutor @@ -10,9 +11,18 @@ 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 +from integration._support.client import ( + Gateway, + Scenario, + delete_key_if_present, + eventually, + held, + list_value, + object_value, + string_value, +) from integration._support.database import read_rows, write_rows -from integration._support.otlp_sink import decode_spans +from integration._support.otlp_sink import Span, decode_spans from integration._support.upstream import delete_scenario, register_scenario from integration._support.wire import Request, wire_server from integration.caching.response_cache_case import entry_key, redis_store @@ -859,7 +869,7 @@ LANGFUSE_COST: Final = "langfuse.observation.cost_details" _LANGFUSE_COST_DETAILS: Final = TypeAdapter(dict[str, float]) -def _langfuse_callback(host: str) -> dict[str, JsonValue]: +def _langfuse_callback(host: str, **callback_vars: JsonValue) -> dict[str, JsonValue]: return { "callback_name": "langfuse", "callback_type": "success", @@ -867,12 +877,17 @@ def _langfuse_callback(host: str) -> dict[str, JsonValue]: "langfuse_public_key": LANGFUSE_PUBLIC_KEY, "langfuse_secret_key": LANGFUSE_SECRET_KEY, "langfuse_host": host, + **callback_vars, }, } +def _logging_to(langfuse_url: str, **callback_vars: JsonValue) -> dict[str, JsonValue]: + return {"logging": [_langfuse_callback(langfuse_url, **callback_vars)]} + + def _key_with_its_own_logging(scenario: Scenario, langfuse_url: str) -> str: - return scenario.key(metadata={"logging": [_langfuse_callback(langfuse_url)]}) + return scenario.key(metadata=_logging_to(langfuse_url)) def _key_in_a_team_with_logging(scenario: Scenario, langfuse_url: str) -> str: @@ -892,17 +907,33 @@ def _trace_exports(langfuse: Recorder) -> tuple[Request, ...]: return tuple(request for request in langfuse.requests() if request.target.startswith(LANGFUSE_TRACES_PATH)) -def _exported_costs(langfuse: Recorder) -> tuple[float, ...]: - spans: Final = tuple( +def _exported_spans(langfuse: Recorder) -> tuple[Span, ...]: + return tuple( chain.from_iterable( decode_spans(export.body, export.headers.get("content-type", "")) for export in _trace_exports(langfuse) ) ) - return tuple( - _LANGFUSE_COST_DETAILS.validate_json(string_value(span["attributes"][LANGFUSE_COST]))["total"] - for span in spans - if LANGFUSE_COST in span["attributes"] - ) + + +def _span_cost(span: Span) -> float | None: + cost_details: Final = span["attributes"].get(LANGFUSE_COST) + return None if cost_details is None else _LANGFUSE_COST_DETAILS.validate_json(string_value(cost_details))["total"] + + +def _exported_costs(langfuse: Recorder) -> tuple[float, ...]: + return tuple(cost for cost in map(_span_cost, _exported_spans(langfuse)) if cost is not None) + + +def _priced_spans(langfuse: Recorder) -> tuple[Span, ...]: + return tuple(span for span in _exported_spans(langfuse) if (_span_cost(span) or 0.0) > 0) + + +def _billed_once(langfuse: Recorder, seconds: float) -> Span: + priced: Final = eventually(lambda: _priced_spans(langfuse), lambda spans: len(spans) > 0, seconds=seconds) + costs: Final = tuple(_span_cost(span) or 0.0 for span in priced) + assert len(costs) == 1, costs + assert math.isclose(costs[0], PRICE), costs + return priced[0] @pytest.mark.timeout(240) @@ -925,3 +956,133 @@ def test_the_billed_cost_reaches_the_langfuse_project_set_on_the_creating_key_or assert [cost for cost in exported if cost > 0] == [pytest.approx(PRICE)] credentials: Final = base64.b64encode(f"{LANGFUSE_PUBLIC_KEY}:{LANGFUSE_SECRET_KEY}".encode()).decode() assert {export.headers.get("authorization") for export in _trace_exports(langfuse)} == {f"Basic {credentials}"} + + +@pytest.mark.timeout(240) +def test_the_billed_cost_reaches_the_creating_keys_langfuse_project_and_never_the_readers(gateway: Gateway) -> None: + with ( + wire_server(langfuse_sink) as creator_wire, + wire_server(langfuse_sink) as reader_wire, + gateway.scenario() as scenario, + ): + creator_langfuse: Final = Recorder(creator_wire) + reader_langfuse: Final = Recorder(reader_wire) + team_id: Final = scenario.team() + user_id: Final = scenario.member(team_id) + creator: Final = scenario.key(team_id=team_id, user_id=user_id, metadata=_logging_to(creator_wire.url)) + reader: Final = scenario.key(team_id=team_id, user_id=user_id, metadata=_logging_to(reader_wire.url)) + upstream: Final = scripted_upstream(scenario, background_routes("completed")) + response_id: Final = create_background_response(scenario, creator, upstream.model) + reads: Final = tuple(read_response(gateway, reader, response_id) for _ in range(READS)) + assert reads in (settling_reads("completed", READS), free_reads("completed", READS)) + _billed_once(creator_langfuse, SPEND_LANDING_SECONDS) + eventually(lambda: _trace_exports(reader_langfuse), lambda exports: len(exports) > 0) + assert _priced_spans(reader_langfuse) == () + + +MESSAGE_LOGGING: Final[Mapping[str, Mapping[str, JsonValue]]] = { + "logged": {}, + "redacted": {"turn_off_message_logging": "true"}, +} + + +@pytest.mark.timeout(240) +@pytest.mark.parametrize("message_logging", MESSAGE_LOGGING.keys()) +def test_the_billed_observation_carries_the_response_text_only_when_the_creating_key_keeps_message_logging( + gateway: Gateway, message_logging: str +) -> None: + with wire_server(langfuse_sink) as wire, gateway.scenario() as scenario: + langfuse: Final = Recorder(wire) + key: Final = scenario.key(metadata=_logging_to(wire.url, **MESSAGE_LOGGING[message_logging])) + upstream: Final = scripted_upstream(scenario, background_routes("completed")) + response_id: Final = create_background_response(scenario, key, upstream.model) + reads: Final = tuple(read_response(gateway, key, response_id) for _ in range(READS)) + assert reads in (settling_reads("completed", READS), free_reads("completed", READS)) + billed: Final = _billed_once(langfuse, SPEND_LANDING_SECONDS) + assert ("pong" in json.dumps(billed["attributes"])) is (message_logging == "logged"), billed["attributes"] + + +def _generated_key(scenario: Scenario, **fields: JsonValue) -> str: + key: Final = string_value(scenario.gateway.post("/key/generate", fields)["key"]) + scenario.cleanups.callback(delete_key_if_present, scenario.gateway, key) + return key + + +def _delete_team_if_present(gateway: Gateway, team_id: str) -> None: + if gateway.rows('SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team_id,)): + gateway.post("/team/delete", {"team_ids": [team_id]}) + + +Retirement = Callable[[Scenario, str, str], tuple[str, Callable[[], object]]] + + +def _rotated_with_grace_then_relogged( + scenario: Scenario, billed_url: str, stale_url: str +) -> tuple[str, Callable[[], object]]: + key: Final = _generated_key(scenario, metadata=_logging_to(stale_url)) + + def rotate() -> None: + rotated: Final = string_value( + scenario.gateway.post("/key/regenerate", {"key": key, "grace_period": "1h"})["key"] + ) + scenario.cleanups.callback(delete_key_if_present, scenario.gateway, rotated) + scenario.gateway.post("/key/update", {"key": rotated, "metadata": _logging_to(billed_url)}) + + return key, rotate + + +def _deleted_key(scenario: Scenario, billed_url: str, _stale_url: str) -> tuple[str, Callable[[], object]]: + key: Final = _generated_key(scenario, metadata=_logging_to(billed_url)) + return key, lambda: scenario.gateway.post("/key/delete", {"keys": [key]}) + + +def _deleted_team(scenario: Scenario, billed_url: str, _stale_url: str) -> tuple[str, Callable[[], object]]: + team_id: Final = string_value( + scenario.gateway.post("/team/new", {"team_alias": f"integration-{uuid.uuid4().hex}"})["team_id"] + ) + scenario.cleanups.callback(_delete_team_if_present, scenario.gateway, team_id) + added: Final = scenario.gateway.request("POST", f"/team/{team_id}/callback", _langfuse_callback(billed_url)) + assert added.status_code == 200, added.text + key: Final = _generated_key(scenario, team_id=team_id) + return key, lambda: scenario.gateway.post("/team/delete", {"team_ids": [team_id]}) + + +RETIREMENTS: Final[Mapping[str, Retirement]] = { + "key_rotated_with_grace": _rotated_with_grace_then_relogged, + "key_deleted": _deleted_key, + "team_deleted": _deleted_team, +} + + +@pytest.mark.timeout(240) +@pytest.mark.parametrize("retired", RETIREMENTS.keys()) +def test_a_retired_creators_billed_cost_reaches_its_own_langfuse_project(gateway: Gateway, retired: str) -> None: + with ( + wire_server(langfuse_sink) as billed_wire, + wire_server(langfuse_sink) as stale_wire, + gateway.scenario() as scenario, + ): + billed_langfuse: Final = Recorder(billed_wire) + stale_langfuse: Final = Recorder(stale_wire) + key, retire = RETIREMENTS[retired](scenario, billed_wire.url, stale_wire.url) + upstream: Final = scripted_upstream(scenario, background_routes("queued")) + response_id: Final = create_background_response(scenario, key, upstream.model) + retire() + upstream.script(background_routes("completed")) + _billed_once(billed_langfuse, ONE_POLL_SECONDS + SPEND_LANDING_SECONDS) + assert [row["status"] for row in managed_rows(response_id)] == ["completed"] + assert _priced_spans(stale_langfuse) == () + + +@pytest.mark.timeout(240) +def test_a_creator_whose_langfuse_keys_resolve_to_nothing_is_still_billed_once(gateway: Gateway) -> None: + unset: Final = f"os.environ/INTEGRATION_UNSET_LANGFUSE_{uuid.uuid4().hex.upper()}" + with wire_server(langfuse_sink) as wire, gateway.scenario() as scenario: + key: Final = scenario.key(metadata=_logging_to(wire.url, langfuse_public_key=unset, langfuse_secret_key=unset)) + upstream: Final = scripted_upstream(scenario, background_routes("completed")) + response_id: Final = create_background_response(scenario, key, upstream.model) + reads: Final = tuple(read_response(gateway, key, response_id) for _ in range(READS)) + assert reads in (settling_reads("completed", READS), free_reads("completed", READS)) + assert_charged_exactly(lambda: key_spend(key), PRICE) + assert math.isclose(number(one_charged_row(upstream.model_id)["spend"]), PRICE) + assert [row["status"] for row in managed_rows(response_id)] == ["completed"] From 8bbce867d05b239593890118d31797c77060f8a4 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 28 Oct 2026 18:58:52 -0400 Subject: [PATCH 110/128] feat(proxy): accept request priority from the x-litellm-priority header (#47618) * feat(proxy): accept request priority from the x-litellm-priority header Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): only apply x-litellm-priority on scheduler-queued routes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): apply x-litellm-priority only where the router schedules the request The header now maps onto priority only for chat and text completions that a router serves, so the Cursor route's responses call and proxies started without a config never forward it to the provider. /queue/chat/completions now honors the header too, with the header winning over a body priority. * test(proxy): use a real ProxyLogging in the priority header routing test * fix(proxy): ignore a malformed priority header under drop_params instead of dropping the body priority * fix(proxy): treat an empty priority header as unset and pass a valid one on as an integer --------- Co-authored-by: Mubashir Osmani 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> --- litellm/proxy/common_request_processing.py | 12 +- litellm/proxy/litellm_pre_call_utils.py | 3 + litellm/proxy/proxy_server.py | 11 +- litellm/router.py | 15 +- litellm/router_utils/request_priority.py | 9 + .../router_code_coverage.py | 2 + .../proxy/proxy_server/test_routes_queue.py | 48 ++++++ .../proxy/test_common_request_processing.py | 163 ++++++++++++++++++ 8 files changed, 257 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index f5d6417fb79..70e305f46fb 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -266,6 +266,8 @@ from litellm.proxy.anthropic_endpoints.streaming_model_restamp import ( AnthropicStreamModelRestamper, ) from litellm.proxy.litellm_pre_call_utils import ( + ROUTER_SCHEDULED_ROUTE_TYPES, + SCHEDULER_PRIORITY_HEADER, LiteLLMProxyRequestSetup, add_litellm_data_to_request, get_metadata_variable_name, @@ -2478,14 +2480,20 @@ class ProxyBaseLLMRequestProcessing: self.data["model"] = alias_target aliased_model: Final = self.data.get("model") + if llm_router is not None and route_type in ROUTER_SCHEDULED_ROUTE_TYPES and isinstance(aliased_model, str): + self.data.update( + llm_router.priority_header_update( + model=aliased_model, kwargs=self.data, header=request.headers.get(SCHEDULER_PRIORITY_HEADER) + ) + ) litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) self.data["litellm_call_id"] = litellm_call_id bind_telemetry_call_id(litellm_call_id) - DDSpanTagger.tag_call_id(self.data.get("litellm_call_id")) + DDSpanTagger.tag_call_id(litellm_call_id) DDSpanTagger.tag_request( user_api_key_dict=user_api_key_dict, - requested_model=self.data.get("model"), + requested_model=aliased_model, ) ### AUTO STREAM USAGE TRACKING ### diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 24a31723faa..601d741c189 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -118,6 +118,9 @@ _SESSION_ID_VALUE_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]{8,}$") _SHA256_HEX_RE: Final = re.compile(r"^[0-9a-f]{64}$") +SCHEDULER_PRIORITY_HEADER: Final = "x-litellm-priority" +ROUTER_SCHEDULED_ROUTE_TYPES: Final = frozenset({"acompletion", "atext_completion"}) + # W3C Trace Context traceparent header: https://www.w3.org/TR/trace-context/ # e.g. "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01" _TRACEPARENT_RE: Final = re.compile(r"^[0-9a-f]{2}-([0-9a-f]{32})-[0-9a-f]{16}-[0-9a-f]{2}$", re.IGNORECASE) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 64fb2791913..70e1c25a70a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -666,7 +666,11 @@ from litellm.proxy.list_api.common import ( problem_response, request_validation_problem, ) -from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, add_litellm_data_to_request +from litellm.proxy.litellm_pre_call_utils import ( + SCHEDULER_PRIORITY_HEADER, + LiteLLMProxyRequestSetup, + add_litellm_data_to_request, +) from litellm.proxy.logging_endpoints.callback_logs_endpoints import ( rust_control_plane_router, ) @@ -18481,6 +18485,11 @@ async def async_queue_request( raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value}) router: Final = llm_router + data.update( + router.priority_header_update( + model=data["model"], kwargs=data, header=request.headers.get(SCHEDULER_PRIORITY_HEADER) + ) + ) queued_priority, _ = router.split_request_priority(model=data["model"], kwargs=data) if "stream" in data and data["stream"] is True: # use generate_responses to stream responses diff --git a/litellm/router.py b/litellm/router.py index 3e6d82d9b07..5ca09499c3b 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -359,6 +359,7 @@ from litellm.router_utils.request_priority import ( InvalidPriority, parse_default_priority, priority_error, + priority_header_fields, request_drops_params, resolve_request_priority, unqueueable_priority_message, @@ -5411,14 +5412,22 @@ class Router: request_priority: Final = resolve_request_priority( requested=kwargs.get("priority"), default_priority=self.default_priority, - drops_params=lambda: request_drops_params( - kwargs, self.default_litellm_params, self._request_deployment_params(model=model, kwargs=kwargs) - ), + drops_params=lambda: self._drops_params(model=model, kwargs=kwargs), ) if isinstance(request_priority, InvalidPriority): raise priority_error(model=model, message=request_priority.message) return request_priority, {key: value for key, value in kwargs.items() if key != "priority"} + def priority_header_update( + self, model: str, kwargs: Mapping[str, object], header: str | None + ) -> Mapping[str, int | str]: + return priority_header_fields(header, drops_params=lambda: self._drops_params(model=model, kwargs=kwargs)) + + def _drops_params(self, model: str, kwargs: Mapping[str, object]) -> bool: + return request_drops_params( + kwargs, self.default_litellm_params, self._request_deployment_params(model=model, kwargs=kwargs) + ) + async def _schedule_factory( self, model: str, diff --git a/litellm/router_utils/request_priority.py b/litellm/router_utils/request_priority.py index 14608cc1bcf..61627d4b0cf 100644 --- a/litellm/router_utils/request_priority.py +++ b/litellm/router_utils/request_priority.py @@ -47,6 +47,15 @@ def resolve_request_priority( return priority +def priority_header_fields(header: str | None, drops_params: Callable[[], bool]) -> Mapping[str, int | str]: + if not header: + return {} + priority: Final = _as_priority(header) + if isinstance(priority, InvalidPriority): + return {} if drops_params() else {"priority": header} + return {"priority": priority} + + def parse_default_priority(value: object) -> int | None | InvalidPriority: if value is None: return None diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index eee02a2f6b6..4c9afd26dad 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -175,6 +175,8 @@ ignored_function_names = [ "_responses_continuation", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py "_responses_continuation_ceilings", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py "_team_deployments_named", # Tested through the team lookups in the concurrent delete tests in test_router.py + "priority_header_update", # Tested through the priority header tests in test_common_request_processing.py + "_drops_params", # Read by split_request_priority and priority_header_update, tested through both "_responses_event_announces_background_run", # Tested through the background Responses stream tests in test_router.py "_aresponses_background_stream_can_release", # Tested through the background Responses stream tests in test_router.py "_raise_released_stream_error", # Tested through the background Responses stream content-policy tests in test_router.py diff --git a/tests/unit/proxy/proxy_server/test_routes_queue.py b/tests/unit/proxy/proxy_server/test_routes_queue.py index a86314e3063..5cbf89f5e68 100644 --- a/tests/unit/proxy/proxy_server/test_routes_queue.py +++ b/tests/unit/proxy/proxy_server/test_routes_queue.py @@ -191,6 +191,54 @@ def test_queue_chat_completions_rejects_a_request_it_cannot_queue_as_an_invalid_ assert add_request.await_count == 0 +@pytest.mark.parametrize( + ("default_priority", "priority_header", "priority_fields", "expected_priority"), + [ + (5, "0", {"priority": 5}, 0), + (5, "0", {}, 0), + (None, "high", {"priority": 1, "drop_params": True}, 1), + (None, "", {"priority": 1}, 1), + ], + ids=("header-beats-body", "header-only", "drop-params-keeps-body", "empty-header-keeps-body"), +) +def test_queue_chat_completions_queues_at_the_priority_header( + client, + auth_as, + monkeypatch: pytest.MonkeyPatch, + default_priority: int | None, + priority_header: str, + priority_fields: dict[str, object], + expected_priority: int, +) -> None: + router: Final = _queue_router(monkeypatch, default_priority=default_priority) + payload: Final = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}], **priority_fields} + with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request: + with auth_as(): + response: Final = client.post( + "/queue/chat/completions", json=payload, headers={"x-litellm-priority": priority_header} + ) + assert response.status_code == 200, response.text + assert [call.kwargs["request"].priority for call in add_request.await_args_list] == [expected_priority] + assert response.headers.get("x-litellm-priority") == str(expected_priority) + + +def test_queue_chat_completions_rejects_a_malformed_priority_header_like_a_malformed_body_value( + client, auth_as, monkeypatch: pytest.MonkeyPatch +) -> None: + router: Final = _queue_router(monkeypatch, default_priority=5) + payload: Final = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}], "priority": 1} + with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request: + with auth_as(): + response: Final = client.post( + "/queue/chat/completions", json=payload, headers={"x-litellm-priority": "high"} + ) + assert response.status_code == 400, response.text + error: Final = response.json()["error"] + assert (error["type"], error["param"]) == ("invalid_request_error", "priority") + assert "priority must be an integer" in error["message"] + assert add_request.await_count == 0 + + def test_streamed_queue_responses_report_the_priority_the_router_queued_at( client, auth_as, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 790790459c3..fdcbcbfcd05 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -35,6 +35,7 @@ from litellm.constants import ( ) from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.opentelemetry import UserAPIKeyAuth +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, ProxyConfig, @@ -513,6 +514,168 @@ class TestProxyBaseLLMRequestProcessing: assert exc_info.value.param == "model" add_litellm_data_to_request.assert_not_awaited() + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("path", "route_type", "has_router", "priority_header", "body_priority", "expected_priority"), + [ + ("/v1/chat/completions", "acompletion", True, "0", 5, 0), + ("/v1/completions", "atext_completion", True, "0", 5, 0), + ("/v1/chat/completions", "acompletion", True, "1", None, 1), + ("/v1/completions", "atext_completion", True, "1", None, 1), + ("/v1/chat/completions", "acompletion", True, None, 5, 5), + ("/v1/chat/completions", "acompletion", True, "", 5, 5), + ("/cursor/chat/completions", "aresponses", True, "0", None, None), + ("/v1/responses", "aresponses", True, "0", 5, 5), + ("/v1/embeddings", "aembedding", True, "0", None, None), + ("/v1/chat/completions", "acompletion", False, "0", None, None), + ("/v1/chat/completions", "acompletion", False, "0", 5, 5), + ], + ids=( + "chat-header-beats-body", + "text-completion-header-beats-body", + "chat-header-only", + "text-completion-header-only", + "chat-body-only", + "chat-empty-header-keeps-body", + "cursor-responses-route-ignores-header", + "responses-keeps-body", + "embeddings-ignores-header", + "routerless-ignores-header", + "routerless-keeps-body", + ), + ) + async def test_common_processing_pre_call_logic_applies_priority_header_only_where_the_router_schedules( + self, + path: str, + route_type: Literal["acompletion", "atext_completion", "aresponses", "aembedding"], + has_router: bool, + priority_header: str | None, + body_priority: int | None, + expected_priority: str | int | None, + ) -> None: + from litellm.proxy.proxy_server import ProxyConfig as RealProxyConfig + + header_entries: Final[tuple[tuple[bytes, bytes], ...]] = ( + () if priority_header is None else ((b"x-litellm-priority", priority_header.encode()),) + ) + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": path, + "scheme": "http", + "server": ("localhost", 4000), + "client": ("127.0.0.1", 12345), + "query_string": b"", + "headers": ((b"content-type", b"application/json"), *header_entries), + } + ) + router: Final = ( + Router( + model_list=[ + { + "model_name": "qa-model", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-test", "mock_response": "hi"}, + } + ] + ) + if has_router + else None + ) + processing_obj: Final = ProxyBaseLLMRequestProcessing( + data={ + "model": "qa-model", + "messages": [{"role": "user", "content": "hi"}], + **({"priority": body_priority} if body_priority is not None else {}), + } + ) + returned_data, _ = await processing_obj.common_processing_pre_call_logic( + request=request, + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + proxy_logging_obj=ProxyLogging(user_api_key_cache=UserApiKeyCache()), + proxy_config=RealProxyConfig(), + route_type=route_type, + llm_router=router, + ) + + assert ("priority" in returned_data, returned_data.get("priority")) == ( + expected_priority is not None, + expected_priority, + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("requested_model", "deployment_drops_params", "body_priority", "expected_priority"), + [ + ("qa-model", True, 5, 5), + ("qa-model", True, None, None), + ("key-alias", True, 5, 5), + ("qa-model", False, 5, "high"), + ], + ids=( + "drop-params-keeps-body", + "drop-params-without-body-leaves-default", + "drop-params-read-from-the-aliased-deployment", + "no-drop-params-passes-header-on-to-400", + ), + ) + async def test_common_processing_pre_call_logic_drops_only_a_malformed_priority_header_under_drop_params( + self, + requested_model: str, + deployment_drops_params: bool, + body_priority: int | None, + expected_priority: str | int | None, + ) -> None: + from litellm.proxy.proxy_server import ProxyConfig as RealProxyConfig + + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/chat/completions", + "scheme": "http", + "server": ("localhost", 4000), + "client": ("127.0.0.1", 12345), + "query_string": b"", + "headers": ((b"content-type", b"application/json"), (b"x-litellm-priority", b"high")), + } + ) + router: Final = Router( + model_list=[ + { + "model_name": "qa-model", + "litellm_params": { + "model": "openai/gpt-4o", + "api_key": "sk-test", + "mock_response": "hi", + "drop_params": deployment_drops_params, + }, + } + ] + ) + processing_obj: Final = ProxyBaseLLMRequestProcessing( + data={ + "model": requested_model, + "messages": [{"role": "user", "content": "hi"}], + **({"priority": body_priority} if body_priority is not None else {}), + } + ) + returned_data, _ = await processing_obj.common_processing_pre_call_logic( + request=request, + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test", aliases={"key-alias": "qa-model"}), + proxy_logging_obj=ProxyLogging(user_api_key_cache=UserApiKeyCache()), + proxy_config=RealProxyConfig(), + route_type="acompletion", + llm_router=router, + ) + + assert ("priority" in returned_data, returned_data.get("priority")) == ( + expected_priority is not None, + expected_priority, + ) + @pytest.mark.asyncio @pytest.mark.parametrize("safe_memory_mode", [False, True]) @pytest.mark.parametrize( From 3f526a5017219f477fb24aa752511b9667ce51c0 Mon Sep 17 00:00:00 2001 From: agustin18 Date: Wed, 28 Oct 2026 20:03:03 -0300 Subject: [PATCH 111/128] fix(websearch): strip Claude Code instructional prefix in short-circuit search (#43458) * fix(websearch): strip Claude Code instructional prefix in short-circuit search Signed-off-by: agustin18 * fix(websearch): preserve quotes for unwrapped search queries Signed-off-by: agustin18 * test(websearch): replace mock with pure unit tests and cover empty queries Signed-off-by: agustin18 * style(websearch): fix ruff formatting blank line in handler Signed-off-by: agustin18 * refactor(websearch): enforce Final annotations, prevent rebinding, and retain e2e coverage Signed-off-by: agustin18 * test(websearch): assert server_tool_use block unconditionally Signed-off-by: agustin18 * fix(websearch): search only the query inside Claude Code's web search prompt Match Claude Code's exact "Perform a web search for the query:" prefix instead of a broad regex family, keep any quotes the user typed so exact-phrase searches still work, and test the search request at the HTTP boundary. --------- Signed-off-by: agustin18 Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../websearch_interception/handler.py | 5 +- .../test_websearch_short_circuit.py | 64 +++++++++++++++++++ 2 files changed, 68 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index a38b8bed595..6e1ba3eea08 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -97,6 +97,8 @@ WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY: Final = "websearch_native_blocks" _RESPONSE_CONTENT_FIELD: Final = "content" +_CLAUDE_CODE_WEB_SEARCH_PROMPT_PREFIX: Final = "Perform a web search for the query:" + _ResponseT: Final = TypeVar("_ResponseT") @@ -363,7 +365,8 @@ class WebSearchInterceptionLogger(CustomLogger): get_last_user_message, ) - query: Final = get_last_user_message(cast(list[AllMessageValues], messages)) + last_user_message: Final = get_last_user_message(cast(list[AllMessageValues], messages)) or "" + query: Final = last_user_message.removeprefix(_CLAUDE_CODE_WEB_SEARCH_PROMPT_PREFIX).strip() if not query: return None diff --git a/tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py b/tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py index fdc3771e903..414bcc11602 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py @@ -6,10 +6,13 @@ and executes the search directly without routing through the backend LLM. """ import json +from typing import Final from unittest.mock import AsyncMock, patch import pytest +import respx +import litellm from litellm.integrations.websearch_interception.handler import ( WebSearchInterceptionLogger, ) @@ -169,6 +172,27 @@ class TestTryShortCircuitSearch: assert result is None + @pytest.mark.asyncio + @pytest.mark.parametrize( + "empty_prompt", + [ + " ", + "Perform a web search for the query: ", + ], + ) + async def test_does_not_short_circuit_empty_or_whitespace_user_message(self, empty_prompt: str): + """User message that is whitespace or empty query → returns None without executing search.""" + logger = WebSearchInterceptionLogger(enabled_providers=["github_copilot"]) + + result = await logger.try_short_circuit_search( + model="github_copilot/claude-sonnet-4", + messages=[{"role": "user", "content": empty_prompt}], + tools=[{"type": "web_search_20250305", "name": "web_search", "max_uses": 8}], + custom_llm_provider="github_copilot", + ) + + assert result is None + @pytest.mark.asyncio async def test_search_failure_returns_error_text(self): """Search failure → response with error message, not exception""" @@ -226,6 +250,46 @@ class TestTryShortCircuitSearch: # Query extraction tests # --------------------------------------------------------------------------- +PERPLEXITY_SEARCH_URL: Final = "https://api.perplexity.ai/search" + + +@pytest.fixture +def perplexity_search(monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter) -> respx.Route: + monkeypatch.setenv("PERPLEXITYAI_API_KEY", "test-perplexity-key") + monkeypatch.delenv("PERPLEXITY_API_BASE", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + return respx_mock.post(PERPLEXITY_SEARCH_URL).respond( + json={"results": [{"title": "Releases", "url": "https://example.com/releases", "snippet": "Release notes"}]} + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("user_message", "searched_query"), + [ + ("Perform a web search for the query: litellm latest release", "litellm latest release"), + ('Perform a web search for the query: "litellm proxy"', '"litellm proxy"'), + ], +) +async def test_short_circuit_searches_the_query_inside_claude_codes_prompt( + perplexity_search: respx.Route, user_message: str, searched_query: str +): + logger: Final = WebSearchInterceptionLogger(enabled_providers=["bedrock"]) + + result: Final = await logger.try_short_circuit_search( + model="bedrock/converse/us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": user_message}], + tools=[{"type": "web_search_20250305", "name": "web_search", "max_uses": 8}], + custom_llm_provider="bedrock", + ) + + assert perplexity_search.call_count == 1 + assert json.loads(perplexity_search.calls.last.request.content)["query"] == searched_query + assert result is not None + server_tool_use: Final = next(block for block in result["content"] if block["type"] == "server_tool_use") + assert server_tool_use["input"]["query"] == searched_query + # --------------------------------------------------------------------------- # Integration with entry point From d30c702c5cfe98d0f2c6832ce5ef5bb0ca8019de Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Wed, 28 Oct 2026 16:11:51 -0700 Subject: [PATCH 112/128] feat(accounting): add durable event ledger writer (#48883) * feat(accounting): add durable event ledger writer * fix(accounting): validate delegated transaction managers locally * docs(accounting): clarify ledger integration contracts * fix(accounting): return ledger conflicts as values and keep cleanup on the kind index Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(accounting): import assert_never from typing_extensions for Python 3.10 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(accounting): annotate ledger integration test locals as Final Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(accounting): finish ledger test conventions --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/db/accounting_ledger.py | 269 ++++++++++++ .../database/test_accounting_ledger.py | 412 +++++++++++++++++- tests/unit/proxy/db/test_accounting_ledger.py | 103 +++++ 3 files changed, 781 insertions(+), 3 deletions(-) create mode 100644 litellm/proxy/db/accounting_ledger.py create mode 100644 tests/unit/proxy/db/test_accounting_ledger.py diff --git a/litellm/proxy/db/accounting_ledger.py b/litellm/proxy/db/accounting_ledger.py new file mode 100644 index 00000000000..91f8307283e --- /dev/null +++ b/litellm/proxy/db/accounting_ledger.py @@ -0,0 +1,269 @@ +from __future__ import annotations + +import hashlib +import json +from collections.abc import AsyncGenerator, Callable +from contextlib import AbstractAsyncContextManager, asynccontextmanager +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from typing import TYPE_CHECKING, Final, Literal, TypeAlias + +from pydantic import AwareDatetime, ConfigDict, JsonValue, TypeAdapter +from typing_extensions import assert_never + +from litellm.accounting import CostEvent, DeliveryTarget, EventKind, InboundUsageEvent, OutboundUsageEvent +from litellm.proxy.db.db_span import db_span +from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper + +if TYPE_CHECKING: + from prisma import Prisma + + from litellm.proxy.utils import PrismaClient + +AccountingEvent: TypeAlias = InboundUsageEvent | OutboundUsageEvent | CostEvent +_JSON: Final = TypeAdapter(dict[str, JsonValue]) +_INSTANT: Final[TypeAdapter[datetime]] = TypeAdapter(AwareDatetime) +_STRINGS: Final = TypeAdapter(tuple[dict[str, str], ...]) +_TIMEOUT: Final = timedelta(seconds=10) + + +class AccountingConflictError(ValueError): + pass + + +_Transaction: TypeAlias = AbstractAsyncContextManager["Prisma"] + + +@dataclass(frozen=True, slots=True) +class _AcceptedEvent: + event_id: str + kind: EventKind + source_event_id: str | None + occurred_at: datetime + payload: str + payload_hash: str + targets: tuple[str, ...] + schema_version: int + version: int + + +def _json(value: object) -> str: + return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) + + +def _prepare( + event: AccountingEvent, + targets: tuple[DeliveryTarget, ...], + envelope: str, + schema_version: int, + version: int, +) -> _AcceptedEvent: + if type(schema_version) is not int or type(version) is not int or min(schema_version, version) < 1: + raise ValueError("Accounting versions must be positive integers") + validated: Final = type(event).model_validate_json(event.model_dump_json()) + payload: Final = _JSON.validate_json(envelope) + if payload.get("event") != validated.model_dump(mode="json"): + raise ValueError("Accounting envelope must contain the matching public event") + destinations: Final = tuple(sorted(frozenset(DeliveryTarget(target).value for target in targets))) + serialized: Final = _json(payload) + identity: Final = _json((schema_version, version, destinations, payload)) + kind: Final = ( + EventKind.COST + if isinstance(validated, CostEvent) + else EventKind.INBOUND_USAGE + if isinstance(validated, InboundUsageEvent) + else EventKind.OUTBOUND_USAGE + ) + return _AcceptedEvent( + event_id=validated.event_id, + kind=kind, + source_event_id=validated.usage_event_id if isinstance(validated, CostEvent) else None, + occurred_at=validated.occurred_at if isinstance(validated, CostEvent) else validated.ended_at, + payload=serialized, + payload_hash=hashlib.sha256(identity.encode()).hexdigest(), + targets=destinations, + schema_version=schema_version, + version=version, + ) + + +_INSERT_EVENT: Final = """ +INSERT INTO "LiteLLM_AccountingEvent" + (event_id, kind, source_event_id, schema_version, payload_hash, occurred_at, payload) +VALUES ($1, $2::"AccountingEventKind", $3, $4::int, $5, $6::timestamp, $7::jsonb) +ON CONFLICT (event_id) DO NOTHING RETURNING event_id +""" +_INSERT_DELIVERIES: Final = """ +INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version) +SELECT $1, value::"AccountingDeliveryTarget", $3::int FROM jsonb_array_elements_text($2::jsonb) +""" +_LOCK_SOURCE: Final = """ +SELECT event_id FROM "LiteLLM_AccountingEvent" +WHERE event_id = $1 AND kind = 'outbound_usage' FOR UPDATE +""" +_EXISTING_COST: Final = """ +SELECT event_id FROM "LiteLLM_AccountingEvent" WHERE kind = 'cost' AND source_event_id = $1 +""" +_LOCK_EVENT: Final = """ +SELECT payload_hash FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 FOR UPDATE +""" +_COMPLETED_FAMILY: Final = """ +root.source_event_id IS NULL AND root.kind IN ('inbound_usage', 'outbound_usage') +AND root.recorded_at < $1::timestamp +AND (root.kind = 'inbound_usage' OR EXISTS ( + SELECT 1 FROM "LiteLLM_AccountingEvent" AS cost WHERE cost.source_event_id = root.event_id AND cost.kind = 'cost' +)) +AND NOT EXISTS ( + SELECT 1 FROM "LiteLLM_AccountingEvent" AS member + WHERE (member.event_id = root.event_id OR (member.kind = 'cost' AND member.source_event_id = root.event_id)) + AND (member.recorded_at >= $1::timestamp + OR EXISTS (SELECT 1 FROM "LiteLLM_AccountingDelivery" AS delivery + WHERE delivery.event_id = member.event_id AND delivery.state != 'applied')) +) +""" +_RETENTION_CANDIDATES: Final = ( + 'SELECT root.event_id FROM "LiteLLM_AccountingEvent" AS root WHERE ' + + _COMPLETED_FAMILY + + " ORDER BY root.recorded_at, root.event_id LIMIT $2::int FOR UPDATE OF root SKIP LOCKED" +) +_ELIGIBLE_GROUP: Final = ( + 'SELECT root.event_id FROM "LiteLLM_AccountingEvent" AS root WHERE root.event_id = $2 AND ' + _COMPLETED_FAMILY +) +_DELETE_DELIVERIES: Final = """ +DELETE FROM "LiteLLM_AccountingDelivery" WHERE event_id IN ( + SELECT event_id FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 OR (kind = 'cost' AND source_event_id = $1) +) +""" +_DELETE_EVENTS: Final = """ +DELETE FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 OR (kind = 'cost' AND source_event_id = $1) +""" + + +_Conflict: TypeAlias = Literal["missing_source", "competing_cost", "identity_conflict"] + + +async def _accept(db: Prisma, record: _AcceptedEvent) -> Literal["inserted", "duplicate"] | _Conflict: + if record.source_event_id is not None: + source: Final = _STRINGS.validate_python(await db.query_raw(_LOCK_SOURCE, record.source_event_id)) + if not source: + return "missing_source" + costs: Final = _STRINGS.validate_python(tuple(await db.query_raw(_EXISTING_COST, record.source_event_id))) + if any(cost["event_id"] != record.event_id for cost in costs): + return "competing_cost" + inserted: Final = _STRINGS.validate_python( + await db.query_raw( + _INSERT_EVENT, + record.event_id, + record.kind.value, + record.source_event_id, + record.schema_version, + record.payload_hash, + record.occurred_at.astimezone(timezone.utc).isoformat(), + record.payload, + ) + ) + if not inserted: + existing: Final = _STRINGS.validate_python(tuple(await db.query_raw(_LOCK_EVENT, record.event_id))) + if not existing or existing[0]["payload_hash"] != record.payload_hash: + return "identity_conflict" + return "duplicate" + await db.execute_raw(_INSERT_DELIVERIES, record.event_id, _json(record.targets), record.version) + return "inserted" + + +async def _delete_completed(db: Prisma, event_id: str, cutoff: str) -> int: + if not _STRINGS.validate_python(await db.query_raw(_ELIGIBLE_GROUP, cutoff, event_id)): + return 0 + await db.execute_raw(_DELETE_DELIVERIES, event_id) + return await db.execute_raw(_DELETE_EVENTS, event_id) + + +class AccountingLedger: + def __init__(self, transaction: Callable[[], _Transaction]) -> None: + self._transaction: Final = transaction + + @classmethod + def for_client(cls, client: PrismaClient) -> AccountingLedger: + from prisma.client import TransactionManager + + managers: Final = TypeAdapter(TransactionManager, config=ConfigDict(arbitrary_types_allowed=True)) + + @asynccontextmanager + async def transaction() -> AsyncGenerator[Prisma]: + primary: Final = writer_wrapper(client.db) + manager: Final = managers.validate_python(primary.tx(timeout=_TIMEOUT)) + async with db_span("accounting_ledger", "LiteLLM_AccountingEvent"), manager as db: + await db.execute_raw("SET LOCAL statement_timeout = 5000") + await db.execute_raw("SET LOCAL lock_timeout = 1000") + yield db + + return cls(transaction) + + async def insert_event( + self, + event: AccountingEvent, + targets: tuple[DeliveryTarget, ...], + *, + envelope: str, + schema_version: int = 1, + version: int = 1, + ) -> Literal["inserted", "duplicate"]: + """Accept T5/T6's serialized envelope, containing its public event under `event`. + + The immutable acceptance hash includes the target set and delivery version. + Changing either for an accepted event conflicts. T7 must provide a separate + add-target or new-version workflow that preserves receipts and prevents + duplicate contributions; redelivery cannot silently create or reset work. + + The Prisma adapter opens one interactive transaction per event, including + two SET LOCAL statements and a source-row lock for costs. I/O is async, + but awaiting this method inline adds commit latency per event, potentially + several times per request. T5 owns bounded asynchronous capture and an + explicit batching strategy; looping over this method does not batch commits. + Success acknowledges commit, independently of HTTP response acknowledgement. + """ + record: Final = _prepare(event, targets, envelope, schema_version, version) + async with self._transaction() as db: + result: Final = await _accept(db, record) + match result: + case "inserted" | "duplicate": + return result + case "missing_source": + raise AccountingConflictError("Cost source must be a durably accepted outbound usage event") + case "competing_cost": + raise AccountingConflictError("Usage event already has an authoritative cost") + case "identity_conflict": + raise AccountingConflictError("Event identity conflicts with its accepted content or delivery set") + case _: + assert_never(result) + + async def purge_completed( + self, + *, + redelivery_before: datetime | None = None, + recovery_before: datetime | None = None, + limit: int = 100, + ) -> int: + """Delete at most `limit` completed families once both retention owners permit it. + + T8 must supply a coherent recovery watermark including outstanding holds. + Omitting either watermark disables deletion. Each usage row and its cost + remain together so a retained charge never loses its source or vice versa. + Applied receipts are deleted with the family, ending its retained dedup + protection. The caller must guarantee events recorded before + redelivery_before cannot be redelivered through supported paths. + T8's recovery_before must also protect checkpoints, + outstanding holds and unfinished settlements; age alone is insufficient. + """ + if type(limit) is not int or not 1 <= limit <= 1000: + raise ValueError("Accounting retention batch size must be between 1 and 1000") + if redelivery_before is None or recovery_before is None: + return 0 + cutoff: Final = min( + _INSTANT.validate_python(redelivery_before).astimezone(timezone.utc), + _INSTANT.validate_python(recovery_before).astimezone(timezone.utc), + ).isoformat() + async with self._transaction() as db: + roots: Final = _STRINGS.validate_python(tuple(await db.query_raw(_RETENTION_CANDIDATES, cutoff, limit))) + deleted: Final = tuple([await _delete_completed(db, root["event_id"], cutoff) for root in roots]) + return sum(deleted) diff --git a/tests/integration/database/test_accounting_ledger.py b/tests/integration/database/test_accounting_ledger.py index 089ed872dee..0c2ceb62eaf 100644 --- a/tests/integration/database/test_accounting_ledger.py +++ b/tests/integration/database/test_accounting_ledger.py @@ -1,16 +1,48 @@ +from __future__ import annotations + +import asyncio +import json import os import subprocess import sys -from collections.abc import Iterator +from collections.abc import AsyncGenerator, AsyncIterator, Iterator, Mapping, Sequence +from contextlib import asynccontextmanager +from datetime import datetime, timezone +from decimal import Decimal +from itertools import chain from pathlib import Path -from typing import Final +from typing import TYPE_CHECKING, Final, TypeAlias import psycopg import pytest +import pytest_asyncio from psycopg import sql +from pydantic import JsonValue -from litellm.accounting import DeliveryState, DeliveryTarget, EventKind +from litellm.accounting import ( + CostEvent, + DeliveryState, + DeliveryTarget, + EventKind, + InboundUsageEvent, + Metric, + MetricName, + OutboundUsageEvent, + Outcome, + PricingSnapshot, +) +from litellm.proxy.db import accounting_ledger +from litellm.proxy.db.accounting_ledger import AccountingConflictError, AccountingEvent, AccountingLedger +from litellm.types.utils import CallTypes from tests.integration._support.database import scratch_database +from tests.integration._support.prisma_clients import connected_proxy_prisma + +if TYPE_CHECKING: + from prisma import Prisma + + from litellm.proxy.utils import PrismaClient + +_Writer: TypeAlias = tuple["PrismaClient", AccountingLedger] _SCHEMA: Final = Path(__file__).resolve().parents[3] / "litellm-proxy-extras/litellm_proxy_extras/schema.prisma" _EVENT_INSERT: Final = ( @@ -126,3 +158,377 @@ def test_new_delivery_defaults_to_pending_with_no_attempts(ledger: psycopg.Conne (DeliveryTarget.KEY_SPEND.value,), ).fetchone() assert row == ("pending", 0) + + +_AT: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) +_CUTOFF: Final = datetime(2026, 2, 1, tzinfo=timezone.utc) +_TARGETS: Final = (DeliveryTarget.USER_SPEND, DeliveryTarget.TEAM_SPEND) + + +def _usage(event_id: str = "usage", *, inbound: bool = False) -> InboundUsageEvent | OutboundUsageEvent: + fields: Final = dict( + event_id=event_id, + request_id="request", + call_type=CallTypes.completion, + started_at=_AT, + ended_at=_AT, + outcome=Outcome.SUCCEEDED, + dimensions=(), + ) + if inbound: + return InboundUsageEvent(**fields, metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),)) + return OutboundUsageEvent(**fields, attempt_id=event_id, retry_of_attempt_id=None, metrics=None) + + +def _cost(source: str = "usage", event_id: str = "cost", amount: str = "0.125") -> CostEvent: + return CostEvent( + event_id=event_id, + usage_event_id=source, + occurred_at=_AT, + dimensions=(), + charges=(Metric(name=MetricName.SPEND, value=Decimal(amount), unit="USD"),), + snapshot=PricingSnapshot( + calculator_version="fixture", + model="fixture", + provider=None, + pricing_source="cost_map", + priced_model_key="fixture", + total_cost=Decimal(amount), + ), + ) + + +def _envelope(event: AccountingEvent, *, context: Mapping[str, JsonValue] | None = None) -> str: + return json.dumps({"event": event.model_dump(mode="json"), "context": context or {}}) + + +async def _insert( + store: AccountingLedger, + event: AccountingEvent, + targets: tuple[DeliveryTarget, ...] = _TARGETS, + *, + context: Mapping[str, JsonValue] | None = None, +) -> str: + return await store.insert_event(event, targets, envelope=_envelope(event, context=context)) + + +def _rows(database_url: str, statement: str, parameters: tuple[object, ...] = ()) -> Sequence[tuple[object, ...]]: + with psycopg.connect(database_url) as connection: + cursor: Final = connection.execute(statement, parameters) + return cursor.fetchall() if cursor.description is not None else [] + + +def _complete_and_age(database_url: str) -> None: + with psycopg.connect(database_url) as connection: + connection.execute('UPDATE "LiteLLM_AccountingEvent" SET recorded_at = %s', (_AT,)) + connection.execute("UPDATE \"LiteLLM_AccountingDelivery\" SET state = 'applied', applied_at = %s", (_AT,)) + + +@pytest_asyncio.fixture(loop_scope="function") +async def writer(ledger_database: str, monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[_Writer]: + monkeypatch.setenv("DATABASE_URL", ledger_database) + monkeypatch.delenv("DATABASE_URL_READ_REPLICA", raising=False) + _rows(ledger_database, 'TRUNCATE "LiteLLM_AccountingDelivery", "LiteLLM_AccountingEvent"') + async with connected_proxy_prisma(ledger_database) as client: + yield client, AccountingLedger.for_client(client) + + +@pytest.mark.asyncio +async def test_event_and_complete_delivery_set_commit_together(writer: _Writer, ledger_database: str) -> None: + _, store = writer + event: Final = _usage() + assert await _insert(store, event, context={"requested_region": "eu", "quantity": "1.25"}) == "inserted" + rows: Final = _rows(ledger_database, 'SELECT kind::text, source_event_id, payload FROM "LiteLLM_AccountingEvent"') + assert rows == [ + ("outbound_usage", None, json.loads(_envelope(event, context={"requested_region": "eu", "quantity": "1.25"}))) + ] + deliveries: Final = _rows( + ledger_database, 'SELECT target::text, state::text, attempts FROM "LiteLLM_AccountingDelivery"' + ) + assert set(deliveries) == {(target.value, "pending", 0) for target in _TARGETS} + + +@pytest.mark.asyncio +async def test_concurrent_duplicates_preserve_completed_deliveries(writer: _Writer, ledger_database: str) -> None: + _, store = writer + event: Final = _usage(inbound=True) + results: Final = await asyncio.gather(*(_insert(store, event) for _ in range(8))) + assert results.count("inserted") == 1 and results.count("duplicate") == 7 + _complete_and_age(ledger_database) + before: Final = _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY target') + reordered: Final = json.dumps(json.loads(_envelope(event)), sort_keys=True, indent=2) + assert await store.insert_event(event, tuple(reversed(_TARGETS)) + _TARGETS, envelope=reordered) == "duplicate" + assert _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY target') == before + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(1,)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("change", ["payload", "context", "targets", "schema_version", "version"]) +async def test_conflicting_acceptance_cannot_overwrite_or_add_work( + writer: _Writer, ledger_database: str, change: str +) -> None: + _, store = writer + event: Final = _usage() + await _insert(store, event) + altered: Final = event.model_copy(update={"outcome": Outcome.FAILED}) if change == "payload" else event + targets: Final = (DeliveryTarget.KEY_SPEND,) if change == "targets" else _TARGETS + context: Final = {"override": "0.5"} if change == "context" else {} + options: Final = {change: 2} if change in ("schema_version", "version") else {} + with pytest.raises(AccountingConflictError, match="conflicts"): + await store.insert_event(altered, targets, envelope=_envelope(altered, context=context), **options) + assert _rows(ledger_database, 'SELECT payload FROM "LiteLLM_AccountingEvent"') == [(json.loads(_envelope(event)),)] + assert set(_rows(ledger_database, 'SELECT target::text FROM "LiteLLM_AccountingDelivery"')) == { + (t.value,) for t in _TARGETS + } + + +@pytest.mark.asyncio +async def test_competing_cost_identities_accept_only_one_authoritative_charge( + writer: _Writer, ledger_database: str +) -> None: + _, store = writer + await _insert(store, _usage()) + results: Final = await asyncio.gather( + _insert(store, _cost()), _insert(store, _cost(event_id="other-cost")), return_exceptions=True + ) + assert sum(result == "inserted" for result in results) == 1 + assert sum(isinstance(result, AccountingConflictError) for result in results) == 1 + accepted: Final = _rows( + ledger_database, "SELECT event_id, payload FROM \"LiteLLM_AccountingEvent\" WHERE kind = 'cost'" + ) + assert len(accepted) == 1 + event: Final = CostEvent.model_validate_json(json.dumps(accepted[0][1]["event"])) + assert await _insert(store, event) == "duplicate" + assert event.snapshot.total_cost == Decimal("0.125") + with pytest.raises(AccountingConflictError): + await _insert(store, _cost(event_id=event.event_id, amount="0.5")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("source", ["missing", "inbound", "cost"]) +async def test_cost_requires_an_existing_outbound_source(writer: _Writer, ledger_database: str, source: str) -> None: + _, store = writer + if source == "inbound": + await _insert(store, _usage("source", inbound=True)) + if source == "cost": + await _insert(store, _usage()) + await _insert(store, _cost(event_id="source")) + with pytest.raises(AccountingConflictError, match="outbound usage"): + await _insert(store, _cost(source="source", event_id="invalid")) + assert _rows(ledger_database, "SELECT count(*) FROM \"LiteLLM_AccountingEvent\" WHERE event_id = 'invalid'") == [ + (0,) + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["before_commit", "after_commit", "cancellation"]) +async def test_interruption_rolls_back_or_retry_recognizes_commit( + writer: _Writer, ledger_database: str, failure: str +) -> None: + client, store = writer + + @asynccontextmanager + async def interrupted() -> AsyncGenerator[Prisma]: + async with client.writer_db.tx() as tx: + yield tx + if failure == "before_commit": + raise RuntimeError("injected pre-commit failure") + if failure == "cancellation": + raise asyncio.CancelledError + raise RuntimeError("injected lost acknowledgement") + + error: Final = asyncio.CancelledError if failure == "cancellation" else RuntimeError + with pytest.raises(error): + await _insert(AccountingLedger(interrupted), _usage()) + expected: Final = 1 if failure == "after_commit" else 0 + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(expected,)] + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(expected * len(_TARGETS),)] + assert await _insert(store, _usage()) == ("duplicate" if expected else "inserted") + + +@pytest.mark.asyncio +async def test_delivery_insert_failure_rolls_back_event(writer: _Writer, ledger_database: str) -> None: + from prisma.errors import RawQueryError + + _, store = writer + with psycopg.connect(ledger_database) as connection: + connection.execute( + 'ALTER TABLE "LiteLLM_AccountingDelivery" ADD CONSTRAINT reject_delivery CHECK (version < 1)' + ) + try: + with pytest.raises(RawQueryError, match="reject_delivery"): + await _insert(store, _usage()) + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(0,)] + finally: + with psycopg.connect(ledger_database) as connection: + connection.execute('ALTER TABLE "LiteLLM_AccountingDelivery" DROP CONSTRAINT reject_delivery') + + +@pytest.mark.asyncio +@pytest.mark.parametrize("blocked", ["pending", "parked", "unpriced", "recent_cost", "recovery"]) +async def test_retention_preserves_incomplete_or_recovery_needed_families( + writer: _Writer, ledger_database: str, blocked: str +) -> None: + _, store = writer + await _insert(store, _usage()) + if blocked != "unpriced": + await _insert(store, _cost()) + _complete_and_age(ledger_database) + with psycopg.connect(ledger_database) as connection: + if blocked in ("pending", "parked"): + connection.execute( + "UPDATE \"LiteLLM_AccountingDelivery\" SET state = %s WHERE event_id = 'cost'", (blocked,) + ) + if blocked == "recent_cost": + connection.execute( + "UPDATE \"LiteLLM_AccountingEvent\" SET recorded_at = %s WHERE event_id = 'cost'", (_CUTOFF,) + ) + before: Final = _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') + assert ( + await store.purge_completed( + redelivery_before=_CUTOFF, recovery_before=_AT if blocked == "recovery" else _CUTOFF + ) + == 0 + ) + assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == before + + +@pytest.mark.asyncio +async def test_retention_removes_completed_family_and_receipts_together(writer: _Writer, ledger_database: str) -> None: + _, store = writer + await _insert(store, _usage()) + await _insert(store, _cost()) + _complete_and_age(ledger_database) + assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2 + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(0,)] + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(0,)] + + +@pytest.mark.asyncio +async def test_retention_respects_batch_bound_and_skips_locked_families(writer: _Writer, ledger_database: str) -> None: + client, store = writer + for name in ("a", "b", "c"): + await _insert(store, _usage(name, inbound=True)) + _complete_and_age(ledger_database) + async with client.writer_db.tx() as tx: + await tx.query_raw("SELECT event_id FROM \"LiteLLM_AccountingEvent\" WHERE event_id = 'a' FOR UPDATE") + assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF, limit=1) == 1 + assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == [ + ("a",), + ("c",), + ] + assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2 + + +@pytest.mark.asyncio +async def test_retention_makes_progress_past_pending_and_unpriced_events(writer: _Writer, ledger_database: str) -> None: + _, store = writer + await _insert(store, _usage("a", inbound=True)) + await _insert(store, _usage("b")) + await _insert(store, _usage("c", inbound=True)) + _complete_and_age(ledger_database) + _rows(ledger_database, "UPDATE \"LiteLLM_AccountingDelivery\" SET state = 'pending' WHERE event_id = 'a'") + assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF, limit=1) == 1 + assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == [ + ("a",), + ("b",), + ] + + +@pytest.mark.asyncio +async def test_zero_cost_and_empty_delivery_sets_are_durable(writer: _Writer, ledger_database: str) -> None: + _, store = writer + await _insert(store, _usage(), targets=()) + event: Final = _cost(amount="0") + assert await _insert(store, event, targets=()) == "inserted" + assert await _insert(store, event, targets=()) == "duplicate" + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(0,)] + _complete_and_age(ledger_database) + assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2 + + +@pytest.mark.asyncio +async def test_retention_cannot_delete_source_while_cost_acceptance_is_uncommitted( + writer: _Writer, ledger_database: str +) -> None: + client, store = writer + await _insert(store, _usage()) + _complete_and_age(ledger_database) + inserted: Final = asyncio.Event() + release: Final = asyncio.Event() + + @asynccontextmanager + async def paused_commit() -> AsyncGenerator[Prisma]: + async with client.writer_db.tx() as tx: + yield tx + inserted.set() + await release.wait() + + task: Final = asyncio.create_task(_insert(AccountingLedger(paused_commit), _cost())) + try: + await asyncio.wait_for(inserted.wait(), timeout=5) + assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 0 + finally: + release.set() + await task + assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == [ + ("cost",), + ("usage",), + ] + + +@pytest.mark.asyncio +async def test_retention_failure_rolls_back_event_and_receipt_deletion(writer: _Writer, ledger_database: str) -> None: + client, store = writer + await _insert(store, _usage(inbound=True)) + _complete_and_age(ledger_database) + + @asynccontextmanager + async def interrupted() -> AsyncGenerator[Prisma]: + async with client.writer_db.tx() as tx: + yield tx + raise RuntimeError("injected cleanup failure") + + with pytest.raises(RuntimeError, match="cleanup failure"): + await AccountingLedger(interrupted).purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(1,)] + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(len(_TARGETS),)] + assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 1 + + +@pytest.mark.asyncio +async def test_ledger_reads_and_writes_use_primary_when_replica_is_configured( + writer: _Writer, ledger_database: str, monkeypatch: pytest.MonkeyPatch +) -> None: + with scratch_database("replica") as replica_url: + monkeypatch.setenv("DATABASE_URL_READ_REPLICA", replica_url) + async with connected_proxy_prisma(ledger_database) as client: + store: Final = AccountingLedger.for_client(client) + await _insert(store, _usage()) + assert await _insert(store, _usage()) == "duplicate" + assert await _insert(store, _cost()) == "inserted" + _complete_and_age(ledger_database) + assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2 + assert _rows(replica_url, "SELECT to_regclass('public.\"LiteLLM_AccountingEvent\"')") == [(None,)] + + +def _index_conditions(plan: object) -> tuple[tuple[object, object], ...]: + if isinstance(plan, dict): + own: Final = ((plan.get("Index Name"), plan.get("Index Cond")),) if "Index Name" in plan else () + children: Final = (plan.get("Plan"), *plan.get("Plans", ())) + return own + tuple(chain.from_iterable(_index_conditions(child) for child in children if child is not None)) + return tuple(chain.from_iterable(_index_conditions(item) for item in plan)) if isinstance(plan, list) else () + + +@pytest.mark.parametrize("statement", [accounting_ledger._DELETE_EVENTS, accounting_ledger._DELETE_DELIVERIES]) +def test_family_cleanup_lookups_use_the_kind_source_index(ledger_database: str, statement: str) -> None: + explained: Final = "EXPLAIN (FORMAT JSON) " + statement.replace("$1", "%s") + with psycopg.connect(ledger_database) as connection: + connection.execute("SET enable_seqscan = off") + plan: Final = connection.execute(explained, ("usage",) * statement.count("$1")).fetchone() + assert plan is not None + conditions: Final = _index_conditions(plan[0]) + assert any( + name == "LiteLLM_AccountingEvent_kind_source_event_id_idx" and "source_event_id" in str(condition) + for name, condition in conditions + ), conditions diff --git a/tests/unit/proxy/db/test_accounting_ledger.py b/tests/unit/proxy/db/test_accounting_ledger.py new file mode 100644 index 00000000000..3916fe869e6 --- /dev/null +++ b/tests/unit/proxy/db/test_accounting_ledger.py @@ -0,0 +1,103 @@ +import json +from datetime import datetime, timezone +from decimal import Decimal +from typing import Final, Literal, NoReturn + +import pytest + +from litellm.accounting import InboundUsageEvent, Metric, MetricName, Outcome +from litellm.proxy.db.accounting_ledger import AccountingLedger +from litellm.types.utils import CallTypes + + +def _no_transaction() -> NoReturn: + raise AssertionError("Invalid accounting input must not open a database transaction") + + +@pytest.fixture +def event() -> InboundUsageEvent: + return InboundUsageEvent( + event_id="caller", + request_id="caller", + call_type=CallTypes.completion, + started_at=datetime(2026, 1, 1, tzinfo=timezone.utc), + ended_at=datetime(2026, 1, 1, tzinfo=timezone.utc), + outcome=Outcome.SUCCEEDED, + dimensions=(), + metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("envelope", ["{}", "[]", '{"event":null}', '{"event":{}}', "not json"]) +async def test_invalid_envelopes_fail_before_database_access(event: InboundUsageEvent, envelope: str) -> None: + with pytest.raises(ValueError, match=r"Accounting envelope|validation error"): + await AccountingLedger(_no_transaction).insert_event(event, (), envelope=envelope) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("version", [0, -1, True, 1.5]) +@pytest.mark.parametrize("field", ["version", "schema_version"]) +async def test_invalid_versions_fail_before_database_access( + event: InboundUsageEvent, version: int | float, field: Literal["version", "schema_version"] +) -> None: + with pytest.raises(ValueError, match="positive integers"): + await AccountingLedger(_no_transaction).insert_event( + event, + (), + envelope="{}", + **{field: version}, # pyright: ignore[reportArgumentType] # Exercise invalid runtime versions + ) + + +@pytest.mark.asyncio +async def test_constructed_invalid_events_are_revalidated(event: InboundUsageEvent) -> None: + invalid: Final = event.model_copy(update={"event_id": ""}) + with pytest.raises(ValueError, match="event_id"): + await AccountingLedger(_no_transaction).insert_event(invalid, (), envelope="{}") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("limit", [0, -1, 1001, True, 1.5]) +async def test_retention_rejects_invalid_batch_bounds(limit: int | float) -> None: + with pytest.raises(ValueError, match="batch size"): + await AccountingLedger(_no_transaction).purge_completed( + limit=limit # pyright: ignore[reportArgumentType] # Exercise invalid runtime batch sizes + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("missing", ["both", "redelivery", "recovery"]) +async def test_retention_requires_both_safety_watermarks(missing: Literal["both", "redelivery", "recovery"]) -> None: + cutoff: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + deleted: Final = await AccountingLedger(_no_transaction).purge_completed( + redelivery_before=cutoff if missing == "recovery" else None, + recovery_before=cutoff if missing == "redelivery" else None, + ) + assert deleted == 0 + + +@pytest.mark.asyncio +async def test_retention_rejects_ambiguous_naive_watermarks() -> None: + with pytest.raises(ValueError, match="timezone"): + await AccountingLedger(_no_transaction).purge_completed( + redelivery_before=datetime(2026, 1, 1), recovery_before=datetime(2026, 1, 1, tzinfo=timezone.utc) + ) + + +@pytest.mark.asyncio +async def test_nonfinite_context_is_rejected_before_database_access(event: InboundUsageEvent) -> None: + envelope: Final = json.dumps({"event": event.model_dump(mode="json"), "context": {"price": float("nan")}}) + with pytest.raises(ValueError, match="Out of range float"): + await AccountingLedger(_no_transaction).insert_event(event, (), envelope=envelope) + + +@pytest.mark.asyncio +async def test_unknown_delivery_target_is_rejected_before_database_access(event: InboundUsageEvent) -> None: + envelope: Final = json.dumps({"event": event.model_dump(mode="json")}) + with pytest.raises(ValueError, match="not a valid DeliveryTarget"): + await AccountingLedger(_no_transaction).insert_event( + event, + ("unknown",), # pyright: ignore[reportArgumentType] # Exercise an unknown runtime target + envelope=envelope, + ) From f89b9c61504801fcc99e40e37af532bb38bbfa30 Mon Sep 17 00:00:00 2001 From: rj-cairn <349020452+rj-cairn@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:12:10 -0700 Subject: [PATCH 113/128] refactor(types): replace Any with proven types in 28 files (#48903) * refactor(types): replace Any with proven types in 52 files Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(types): use typing_extensions.TypedDict for validated seams Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(types): revert strict Anthropic response validation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(types): accept a null finish reason from Bedrock Cohere and TwelveLabs Pegasus * test(bedrock): annotate the Pegasus null finish reason test --------- Co-authored-by: mateo 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> --- litellm/litellm_core_utils/litellm_logging.py | 4 +- litellm/litellm_core_utils/safe_json_dumps.py | 32 +++++++-------- ...mazon_twelvelabs_pegasus_transformation.py | 15 ++++++- .../base_invoke_transformation.py | 39 +++++++++++++++---- .../llms/bedrock/realtime/transformation.py | 3 +- .../llms/dashscope/embed/transformation.py | 30 +++++++++++--- litellm/llms/databricks/streaming_utils.py | 9 ++++- litellm/llms/gigachat/authenticator.py | 15 ++++++- litellm/llms/openai/realtime/handler.py | 7 +++- .../perplexity/embedding/transformation.py | 33 ++++++++++++---- .../vertex_gemma_models/transformation.py | 22 ++++++++--- .../volcengine/responses/transformation.py | 5 ++- litellm/proxy/client/models.py | 9 +++-- .../guardrail_hooks/openai/moderations.py | 2 +- litellm/proxy/litellm_pre_call_utils.py | 12 +++--- .../auto_router_endpoints.py | 5 ++- .../router_settings_endpoints.py | 4 +- .../management_helpers/team_roster_sync.py | 3 +- .../mcp_server/discoverable_endpoints.py | 3 +- .../pass_through_endpoints.py | 24 ++++++++---- .../spend_tracking/key_metadata_recovery.py | 5 ++- .../secret_managers/aws_secret_manager_v2.py | 23 ++++++++--- .../integrations/rag/bedrock_knowledgebase.py | 10 ++--- litellm/types/llms/bedrock_invoke_agents.py | 10 ++--- litellm/types/llms/gemini.py | 11 +++--- litellm/types/llms/langgraph.py | 11 +++--- litellm/types/mcp_server/tool_registry.py | 10 ++--- .../guardrails/guardrail_hooks/enkryptai.py | 10 ++--- .../test_base_invoke_transformation.py | 27 +++++++++++++ .../test_twelvelabs_pegasus_transformation.py | 28 +++++++++++++ 30 files changed, 308 insertions(+), 113 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 80f1f8f5eca..8b66c27b930 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3481,7 +3481,7 @@ class Logging(LiteLLMLoggingBaseClass): start_time=start_time, end_time=end_time, print_verbose=print_verbose, - level=LogfireLevel.INFO.value, + level=LogfireLevel.INFO, ) if callback == "lunary" and lunaryLogger is not None: @@ -4442,7 +4442,7 @@ class Logging(LiteLLMLoggingBaseClass): response_obj=result, start_time=start_time, end_time=end_time, - level=LogfireLevel.ERROR.value, + level=LogfireLevel.ERROR, print_verbose=print_verbose, ) diff --git a/litellm/litellm_core_utils/safe_json_dumps.py b/litellm/litellm_core_utils/safe_json_dumps.py index 63242a580e7..24c401281ef 100644 --- a/litellm/litellm_core_utils/safe_json_dumps.py +++ b/litellm/litellm_core_utils/safe_json_dumps.py @@ -34,7 +34,7 @@ def safe_json_structure( def _transform(key: str | None, value: str) -> str: return value if value_transform is None else value_transform(key, value) - def _serialize(obj: object, seen: set[int], depth: int, key: str | None = None) -> Any: + def _serialize(obj: object, seen: set[int], depth: int, key: str | None = None) -> object: # Check for maximum depth. if depth > max_depth: return "MaxDepthExceeded" @@ -48,32 +48,32 @@ def safe_json_structure( if id(obj) in seen: return "CircularReference Detected" seen.add(id(obj)) - result: dict | list | tuple | set | str if isinstance(obj, dict): - result = {} - for k, v in obj.items(): - if isinstance(k, (str)): - clean_k = k.replace("\x00", "") if "\x00" in k else k - result[clean_k] = _serialize(v, seen, depth + 1, clean_k) + mapping_result: Final = { + (clean_k := k.replace("\x00", "") if "\x00" in k else k): _serialize(v, seen, depth + 1, clean_k) + for k, v in obj.items() + if isinstance(k, str) + } seen.remove(id(obj)) - return result + return mapping_result elif isinstance(obj, list): - result = [_serialize(item, seen, depth + 1, key) for item in obj] + list_result: Final = [_serialize(item, seen, depth + 1, key) for item in obj] seen.remove(id(obj)) - return result + return list_result elif isinstance(obj, tuple): - result = tuple(_serialize(item, seen, depth + 1, key) for item in obj) + tuple_result: Final = tuple(_serialize(item, seen, depth + 1, key) for item in obj) seen.remove(id(obj)) - return result + return tuple_result elif isinstance(obj, set): - result = sorted([_serialize(item, seen, depth + 1, key) for item in obj]) + sortable_values: Final[list[Any]] = [_serialize(item, seen, depth + 1, key) for item in obj] + set_result: Final = sorted(sortable_values) seen.remove(id(obj)) - return result + return set_result elif isinstance(obj, BaseModel): dumped: Final = obj.model_dump() - result = _serialize(dumped, seen, depth + 1, key) + model_result: Final = _serialize(dumped, seen, depth + 1, key) seen.remove(id(obj)) - return result + return model_result else: # Fall back to string conversion for non-serializable objects. try: diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py index f7566ac4923..dc14ef005b5 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py @@ -10,6 +10,8 @@ import time from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -34,6 +36,15 @@ else: LiteLLMLoggingObj = Any +@with_config(ConfigDict(extra="allow", strict=True)) +class _PegasusResponse(TypedDict): + message: ReadOnly[NotRequired[str | None]] + finishReason: ReadOnly[NotRequired[str | None]] + + +_PEGASUS_RESPONSE: Final = TypeAdapter(_PegasusResponse) + + class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): """ Handles transforming OpenAI-style requests into Bedrock InvokeModel requests for @@ -207,7 +218,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): ModelResponse with choices[0].message.content and finish_reason """ try: - completion_response: Final = raw_response.json() + completion_response: Final = _PEGASUS_RESPONSE.validate_python(raw_response.json()) except Exception as e: raise BedrockError( message=f"Error parsing response: {raw_response.text}, error: {e}", @@ -224,7 +235,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): message_content: Final = completion_response.get("message", "") # Extract finish reason and map to LiteLLM format - finish_reason_raw: Final = completion_response.get("finishReason", "stop") + finish_reason_raw: Final = completion_response.get("finishReason") or "stop" finish_reason: Final = map_finish_reason(finish_reason_raw) # Set the response content diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index 51b5a99272e..852179e8caa 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -1,11 +1,12 @@ import copy import json import time -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, cast, get_args import httpx -from pydantic import TypeAdapter, ValidationError +from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -48,6 +49,27 @@ from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM _GUARDRAIL_CONFIG_VALIDATOR: Final["TypeAdapter[GuardrailConfigBlock]"] = TypeAdapter(GuardrailConfigBlock) + +@with_config(ConfigDict(extra="allow", strict=True)) +class _CohereGeneration(TypedDict): + text: ReadOnly[str | None] + finish_reason: ReadOnly[str | None] + + +@with_config(ConfigDict(extra="allow", strict=True)) +class _CohereInvokeResponse(TypedDict): + text: NotRequired[ReadOnly[str | None]] + generations: NotRequired[ReadOnly[Sequence[_CohereGeneration]]] + + +@with_config(ConfigDict(extra="allow", strict=True)) +class _LlamaInvokeResponse(TypedDict): + generation: ReadOnly[str | None] + + +_COHERE_INVOKE_RESPONSE: Final = TypeAdapter(_CohereInvokeResponse) +_LLAMA_INVOKE_RESPONSE: Final = TypeAdapter(_LlamaInvokeResponse) + _GUARDRAIL_CONFIG_EXPECTED_FORMAT: Final = ( "{'guardrailIdentifier': str, 'guardrailVersion': str, 'trace': 'enabled'|'disabled'|'enabled_full'}" ) @@ -310,12 +332,13 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): outputText: str | None = None try: if provider == "cohere": - if "text" in completion_response: - outputText = completion_response["text"] - elif "generations" in completion_response: - outputText = completion_response["generations"][0]["text"] + cohere_response: Final = _COHERE_INVOKE_RESPONSE.validate_python(completion_response) + if "text" in cohere_response: + outputText = cohere_response["text"] + elif "generations" in cohere_response: + outputText = cohere_response["generations"][0]["text"] model_response.choices[0].finish_reason = map_finish_reason( - completion_response["generations"][0]["finish_reason"] + cohere_response["generations"][0]["finish_reason"] or "stop" ) elif provider == "anthropic": return litellm.AmazonAnthropicClaudeConfig().transform_response( @@ -361,7 +384,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): elif provider == "ai21": outputText = completion_response.get("completions")[0].get("data").get("text") elif provider == "meta" or provider == "llama" or provider == "deepseek_r1": - outputText = completion_response["generation"] + outputText = _LLAMA_INVOKE_RESPONSE.validate_python(completion_response)["generation"] elif provider == "mistral": outputText = litellm.AmazonMistralConfig.get_outputText(completion_response, model_response) else: # amazon titan diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index 9cb38822881..952d8858f78 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -53,6 +53,7 @@ from litellm.types.realtime import ( RealtimeResponseTypedDict, ) +_JSON_OBJECT_MESSAGE: Final = TypeAdapter(dict[str, JsonValue], config=ConfigDict(strict=True)) _JSON_VALUE: Final = TypeAdapter(object) _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) @@ -856,7 +857,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): List of Bedrock format messages (JSON strings) """ try: - json_message: Final[dict[str, JsonValue]] = json.loads(message) + json_message: Final = _JSON_OBJECT_MESSAGE.validate_python(json.loads(message)) except json.JSONDecodeError: verbose_logger.warning("Invalid JSON message: %s", message[:200]) return [] diff --git a/litellm/llms/dashscope/embed/transformation.py b/litellm/llms/dashscope/embed/transformation.py index d0c2b4dff0c..35132346261 100644 --- a/litellm/llms/dashscope/embed/transformation.py +++ b/litellm/llms/dashscope/embed/transformation.py @@ -14,6 +14,8 @@ Docs - https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.url_utils import ensure_api_route @@ -28,6 +30,23 @@ from ..common_utils import DashScopeError DEFAULT_API_BASE: Final = "https://dashscope.aliyuncs.com/compatible-mode/v1" +@with_config(ConfigDict(extra="allow", strict=True)) +class _DashScopeUsage(TypedDict, total=False): + prompt_tokens: NotRequired[ReadOnly[int]] + total_tokens: NotRequired[ReadOnly[int]] + + +@with_config(ConfigDict(extra="allow", strict=True)) +class _DashScopeEmbeddingResponse(TypedDict, total=False): + data: NotRequired[ReadOnly[list[object]]] + model: NotRequired[ReadOnly[str | None]] + usage: NotRequired[ReadOnly[_DashScopeUsage | None]] + id: NotRequired[ReadOnly[object]] + + +_DASHSCOPE_EMBEDDING_RESPONSE: Final = TypeAdapter(_DashScopeEmbeddingResponse) + + class DashScopeEmbeddingConfig(BaseEmbeddingConfig): """ Reference: https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api @@ -152,11 +171,12 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig): message=message, ) + parsed_response: Final = _DASHSCOPE_EMBEDDING_RESPONSE.validate_python(response_json) model_response.object = "list" - model_response.data = response_json.get("data", []) - model_response.model = response_json.get("model", model) + model_response.data = parsed_response.get("data", []) + model_response.model = parsed_response.get("model", model) - usage: Final = response_json.get("usage") or {} + usage: Final = parsed_response.get("usage") or {} prompt_tokens: Final = usage.get("prompt_tokens", 0) total_tokens: Final = usage.get("total_tokens", prompt_tokens) setattr( @@ -169,8 +189,8 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig): ), ) - if "id" in response_json: - setattr(model_response, "id", response_json["id"]) + if "id" in parsed_response: + setattr(model_response, "id", parsed_response["id"]) return model_response diff --git a/litellm/llms/databricks/streaming_utils.py b/litellm/llms/databricks/streaming_utils.py index eca87f9dd5f..cc85012b803 100644 --- a/litellm/llms/databricks/streaming_utils.py +++ b/litellm/llms/databricks/streaming_utils.py @@ -1,6 +1,8 @@ import json from typing import Final +from pydantic import ConfigDict, TypeAdapter + import litellm from litellm import verbose_logger from litellm.types.llms.openai import ( @@ -23,6 +25,9 @@ def _usage_block(usage: Usage) -> ChatCompletionUsageBlock: ) +_JSON_CHUNK: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) + + class ModelResponseIterator: def __init__(self, streaming_response, sync_stream: bool): self.streaming_response = streaming_response @@ -99,7 +104,7 @@ class ModelResponseIterator: self.response_iterator = self.streaming_response chunk: Final = self._next_sse_data() try: - json_chunk: Final = json.loads(chunk) + json_chunk: Final = _JSON_CHUNK.validate_python(json.loads(chunk)) return self.chunk_parser(chunk=json_chunk) except ValueError as e: verbose_logger.debug( @@ -126,7 +131,7 @@ class ModelResponseIterator: if chunk == "[DONE]": raise StopAsyncIteration try: - json_chunk: Final = json.loads(chunk) + json_chunk: Final = _JSON_CHUNK.validate_python(json.loads(chunk)) return self.chunk_parser(chunk=json_chunk) except ValueError as e: verbose_logger.debug( diff --git a/litellm/llms/gigachat/authenticator.py b/litellm/llms/gigachat/authenticator.py index d43e203d7a1..f37c6e88896 100644 --- a/litellm/llms/gigachat/authenticator.py +++ b/litellm/llms/gigachat/authenticator.py @@ -12,6 +12,8 @@ from types import MappingProxyType from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.caching.caching import InMemoryCache @@ -37,6 +39,17 @@ _token_cache: Final = InMemoryCache() _NO_LITELLM_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) +@with_config(ConfigDict(extra="allow", strict=True)) +class _GigaChatTokenResponse(TypedDict): + tok: NotRequired[ReadOnly[str | None]] + access_token: NotRequired[ReadOnly[str | None]] + exp: NotRequired[ReadOnly[int | float | str | None]] + expires_at: NotRequired[ReadOnly[int | float | str | None]] + + +_TOKEN_RESPONSE: Final = TypeAdapter(_GigaChatTokenResponse) + + class GigaChatAuthError(BaseLLMException): """GigaChat authentication error.""" @@ -230,7 +243,7 @@ async def _request_token_async( def _parse_token_response(response: httpx.Response) -> tuple[str, int]: """Parse OAuth token response.""" - data: Final = response.json() + data: Final = _TOKEN_RESPONSE.validate_python(response.json()) # GigaChat returns either 'tok'/'exp' or 'access_token'/'expires_at' access_token: Final = data.get("tok") or data.get("access_token") diff --git a/litellm/llms/openai/realtime/handler.py b/litellm/llms/openai/realtime/handler.py index 080fc3d24c9..b7e4cc6868f 100644 --- a/litellm/llms/openai/realtime/handler.py +++ b/litellm/llms/openai/realtime/handler.py @@ -5,7 +5,7 @@ This requires websockets, and is currently only supported on LiteLLM Proxy. """ import ssl -from typing import Any, Final, cast +from typing import TYPE_CHECKING, Final, cast from litellm._logging import redact_string, verbose_logger from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES @@ -23,6 +23,9 @@ from ....llms.custom_httpx.http_handler import get_shared_realtime_ssl_context from ..common_utils import is_openai_backed_api_base from ..openai import OpenAIChatCompletion +if TYPE_CHECKING: + from fastapi import WebSocket + class OpenAIRealtime(OpenAIChatCompletion): """ @@ -131,7 +134,7 @@ class OpenAIRealtime(OpenAIChatCompletion): async def async_realtime( self, model: str, - websocket: Any, + websocket: "WebSocket", logging_obj: LiteLLMLogging, api_base: str | None = None, api_key: str | None = None, diff --git a/litellm/llms/perplexity/embedding/transformation.py b/litellm/llms/perplexity/embedding/transformation.py index 7b4da933976..147d344795d 100644 --- a/litellm/llms/perplexity/embedding/transformation.py +++ b/litellm/llms/perplexity/embedding/transformation.py @@ -12,12 +12,14 @@ This module decodes them into float arrays for OpenAI-compatible responses. """ import base64 +import builtins import struct -from collections.abc import Mapping -from typing import Final +from collections.abc import Mapping, Sequence +from typing import Final, Literal import httpx -from pydantic import ConfigDict, TypeAdapter +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import ReadOnly, TypedDict from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.url_utils import join_api_base @@ -31,6 +33,17 @@ _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_i _INT: Final = TypeAdapter(int) +@with_config(ConfigDict(extra="allow", strict=True)) +class _PerplexityEmbeddingResponse(TypedDict, total=False): + model: ReadOnly[str | None] + object: ReadOnly[Literal["list"]] + data: ReadOnly[Sequence[builtins.object] | Mapping[str, builtins.object] | Literal[""]] + usage: ReadOnly[builtins.object] + + +_EMBEDDING_RESPONSE_ADAPTER: Final = TypeAdapter(_PerplexityEmbeddingResponse) + + class PerplexityEmbeddingError(BaseLLMException): def __init__( self, @@ -153,14 +166,18 @@ class PerplexityEmbeddingConfig(BaseEmbeddingConfig): litellm_params: dict = {}, ) -> EmbeddingResponse: try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final[object] = raw_response.json() except Exception: raise PerplexityEmbeddingError(message=raw_response.text, status_code=raw_response.status_code) - model_response.model = raw_response_json.get("model", model) - model_response.object = raw_response_json.get("object", "list") + response_data: Final[_PerplexityEmbeddingResponse] = _EMBEDDING_RESPONSE_ADAPTER.validate_python( + raw_response_json + ) - raw_data: Final = raw_response_json.get("data", []) + model_response.model = response_data.get("model", model) + model_response.object = response_data.get("object", "list") + + raw_data: Final = response_data.get("data", []) decoded_data: Final[list[dict[str, object]]] = [] for item in raw_data: decoded_item = dict(_JSON_OBJECT.validate_python(item)) @@ -168,7 +185,7 @@ class PerplexityEmbeddingConfig(BaseEmbeddingConfig): decoded_data.append(decoded_item) model_response.data = decoded_data - usage_data: Final = _JSON_OBJECT.validate_python(raw_response_json.get("usage", {})) + usage_data: Final = _JSON_OBJECT.validate_python(response_data.get("usage", {})) usage: Final = Usage( prompt_tokens=_INT.validate_python( usage_data.get("prompt_tokens", 0) or usage_data.get("total_tokens", 0) or 0 diff --git a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py index 9fa86ec30f5..ec2fd16f93e 100644 --- a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py +++ b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py @@ -9,10 +9,11 @@ The actual message transformation reuses OpenAIGPTConfig since Gemma uses OpenAI """ from collections.abc import Callable -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Final, cast import httpx -from pydantic import ValidationError +from pydantic import ConfigDict, TypeAdapter, ValidationError +from typing_extensions import NotRequired, ReadOnly, TypeAliasType, TypedDict from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.custom_httpx.http_handler import ( @@ -32,6 +33,17 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer +class _PredictionsEnvelope(TypedDict): + predictions: NotRequired[ReadOnly[dict[str, object]]] + + +_PredictionsEnvelopeType = TypeAliasType("_PredictionsEnvelopeType", _PredictionsEnvelope) +_PREDICTIONS_ENVELOPE: Final[TypeAdapter[_PredictionsEnvelope]] = TypeAdapter( + _PredictionsEnvelopeType, + config=ConfigDict(strict=True, extra="allow"), +) + + def parse_vertex_gemma_container_error(predictions: object) -> VertexGemmaContainerError | None: try: return VertexGemmaContainerError.model_validate(predictions) @@ -142,7 +154,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): def _unwrap_predictions_response( self, - response_json: dict[str, Any], + response_json: _PredictionsEnvelope, ) -> dict[str, object]: """ Unwrap the Vertex Gemma predictions format to OpenAI format. @@ -354,7 +366,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): message=f"Request failed: {response.text}", ) - response_json: Final = response.json() + response_json: Final = _PREDICTIONS_ENVELOPE.validate_python(response.json()) # Unwrap predictions to get OpenAI-compatible response openai_response: Final = self._unwrap_predictions_response(response_json) @@ -449,7 +461,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): message=f"Request failed: {response.text}", ) - response_json: Final = response.json() + response_json: Final = _PREDICTIONS_ENVELOPE.validate_python(response.json()) # Unwrap predictions to get OpenAI-compatible response openai_response: Final = self._unwrap_predictions_response(response_json) diff --git a/litellm/llms/volcengine/responses/transformation.py b/litellm/llms/volcengine/responses/transformation.py index 747ea21e4fb..4076ed06fc2 100644 --- a/litellm/llms/volcengine/responses/transformation.py +++ b/litellm/llms/volcengine/responses/transformation.py @@ -3,6 +3,7 @@ from types import UnionType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, Union, get_args, get_origin import httpx +from pydantic import ConfigDict, TypeAdapter from pydantic import fields as pyd_fields import litellm @@ -37,6 +38,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_RESPONSE_BODY: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) + class _EventModelClass(Protocol): @property @@ -420,7 +423,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): @staticmethod def _parsed_response_body(raw_response: httpx.Response) -> dict[str, object]: - return raw_response.json() + return _RESPONSE_BODY.validate_python(raw_response.json()) @staticmethod def _annotation_origin(annotation: object) -> object: diff --git a/litellm/proxy/client/models.py b/litellm/proxy/client/models.py index 1361a3e6364..859322f061d 100644 --- a/litellm/proxy/client/models.py +++ b/litellm/proxy/client/models.py @@ -3,9 +3,12 @@ from collections.abc import Mapping from typing import Any, Final import requests +from pydantic import ConfigDict, TypeAdapter from .exceptions import NotFoundError, UnauthorizedError +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) + class ModelsManagementClient: def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30): @@ -109,7 +112,7 @@ class ModelsManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _JSON_OBJECT.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) @@ -145,7 +148,7 @@ class ModelsManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _JSON_OBJECT.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) @@ -287,7 +290,7 @@ class ModelsManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _JSON_OBJECT.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) diff --git a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py index 98be35079c3..16b3c45f6f8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py @@ -169,7 +169,7 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): from litellm.types.llms.openai import OpenAIModerationResponse - return OpenAIModerationResponse(**response.json()) + return OpenAIModerationResponse.model_validate(response.json()) def _check_moderation_result(self, moderation_response: "OpenAIModerationResponse") -> None: """ diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 601d741c189..d4528bc85e8 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -691,7 +691,7 @@ def _metadata_string_tags(metadata_objects: Iterable[Mapping[str, object]]) -> t def _strip_untrusted_request_header_controls( - headers: Any, + headers: object, *, allow_client_message_redaction_opt_out: bool = False, ) -> None: @@ -1695,7 +1695,7 @@ class LiteLLMProxyRequestSetup: @staticmethod def add_internal_user_from_user_mapping( - general_settings: dict | None, + general_settings: Mapping[str, object] | None, user_api_key_dict: UserAPIKeyAuth, headers: dict, ) -> UserAPIKeyAuth: @@ -1714,7 +1714,7 @@ class LiteLLMProxyRequestSetup: return user_api_key_dict @staticmethod - def get_user_from_headers(headers: dict, general_settings: dict | None = None) -> str | None: + def get_user_from_headers(headers: dict, general_settings: Mapping[str, object] | None = None) -> str | None: """ Get the user from the specified header if `general_settings.user_header_name` is set. """ @@ -1735,7 +1735,9 @@ class LiteLLMProxyRequestSetup: return user @staticmethod - def get_openai_org_id_from_headers(headers: dict, general_settings: dict | None = None) -> str | None: + def get_openai_org_id_from_headers( + headers: dict, general_settings: Mapping[str, object] | None = None + ) -> str | None: """ Get the OpenAI Org ID from the headers. """ @@ -1826,7 +1828,7 @@ class LiteLLMProxyRequestSetup: headers: dict, request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, - general_settings: dict[str, object] | None = None, + general_settings: Mapping[str, object] | None = None, ) -> LitellmDataForBackendLLMCall: """ - Adds user from headers diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 7a843aecc0c..4f48f3c57b9 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -231,9 +231,12 @@ def _shadow_eval_attempts(prisma_client: "PrismaClient") -> _ShadowEvalAttemptTa return _ShadowEvalAttemptRepository(prisma_client).table +_RAW_ROWS: Final = TypeAdapter(Sequence[dict[str, object]], config=ConfigDict(strict=True)) + + async def _query_raw(prisma_client: "PrismaClient", query: str, *args: object) -> Sequence[Mapping[str, object]]: async with db_span("auto_router_report_query", sql_relation(query)): - return await prisma_client.db.query_raw(query, *args) + return _RAW_ROWS.validate_python(await prisma_client.db.query_raw(query, *args)) async def _execute_raw(prisma_client: "PrismaClient", query: str, *args: object) -> int: diff --git a/litellm/proxy/management_endpoints/router_settings_endpoints.py b/litellm/proxy/management_endpoints/router_settings_endpoints.py index 1f4ca442d20..41c4b5cfaf7 100644 --- a/litellm/proxy/management_endpoints/router_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/router_settings_endpoints.py @@ -10,7 +10,7 @@ GET /router/fields - Get router settings field definitions without values (for U import inspect from collections.abc import Mapping from types import MappingProxyType -from typing import Any, Final, get_args +from typing import Final, get_args from fastapi import APIRouter, Depends from pydantic import Field @@ -32,7 +32,7 @@ router: Final = APIRouter() class RouterSettingsResponse(LiteLLMBaseModel): fields: list[RouterSettingsField] = Field(description="List of all configurable router settings with metadata") - current_values: dict[str, Any] = Field(description="Current values of router settings") + current_values: dict[str, object] = Field(description="Current values of router settings") routing_strategy_descriptions: dict[str, str] = Field(description="Descriptions for each routing strategy option") routing_group_strategies: tuple[str, ...] = Field( description="Strategies supported when constructing a routing group" diff --git a/litellm/proxy/management_helpers/team_roster_sync.py b/litellm/proxy/management_helpers/team_roster_sync.py index dc0ba2e47d4..08be455a55d 100644 --- a/litellm/proxy/management_helpers/team_roster_sync.py +++ b/litellm/proxy/management_helpers/team_roster_sync.py @@ -40,6 +40,7 @@ from litellm.proxy.management_helpers.access_group_team_sync import ( ) from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.chunked_in import delete_many_in, find_many_in, update_many_in +from litellm.repositories.team_repository import TeamRepository if TYPE_CHECKING: from prisma import Prisma @@ -144,7 +145,7 @@ def _changes_nothing(plan: RosterPlan) -> bool: async def _unchanged_roster(prisma_client: PrismaClient, team_id: str) -> RosterSync | TeamGone: - team_row: Final = await prisma_client.writer_db.litellm_teamtable.find_unique(where={"team_id": team_id}) + team_row: Final = await TeamRepository(prisma_client, use_writer=True).table.find_unique(where={"team_id": team_id}) if team_row is None: return TeamGone(team_id=team_id) team: Final = LiteLLM_TeamTable.model_validate(team_row.model_dump()) diff --git a/litellm/proxy/mcp_server/discoverable_endpoints.py b/litellm/proxy/mcp_server/discoverable_endpoints.py index eb12b15cdae..07a267d6caa 100644 --- a/litellm/proxy/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/mcp_server/discoverable_endpoints.py @@ -300,12 +300,13 @@ def decode_state_hash(encrypted_state: str) -> dict: if decrypted_json is None: raise ValueError("Failed to decrypt state parameter") - state_data: Final = json.loads(decrypted_json) + state_data: Final = _DECODED_STATE.validate_python(json.loads(decrypted_json)) return state_data _MANAGED_AUTH_CODE_PREFIX: Final = "llm_scope_" _DECODED_JSON: Final = TypeAdapter(object) +_DECODED_STATE: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) class _ManagedAuthorizationCode(LiteLLMBaseModel): diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 650dcae3239..e017eee3d27 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -30,6 +30,7 @@ from pydantic import TypeAdapter from starlette.datastructures import UploadFile as StarletteUploadFile from starlette.routing import BaseRoute, Route from starlette.websockets import WebSocketState +from typing_extensions import ReadOnly from websockets.asyncio.client import connect from websockets.exceptions import ( ConnectionClosedError, @@ -173,8 +174,17 @@ def _is_mergeable_client_metadata(key: str, value: object) -> bool: return not key.startswith("user_api_key_") and (key != "model_info" or isinstance(value, Mapping)) +class _RegisteredPassThroughRoute(TypedDict): + endpoint_id: ReadOnly[str] + path: ReadOnly[str | None] + type: ReadOnly[str] + methods: ReadOnly[str | bool | list[str] | tuple[str, ...] | Mapping[str, object]] + auth: ReadOnly[bool] + passthrough_params: ReadOnly[Mapping[str, object]] + + # Global registry to track registered pass-through routes and prevent memory leaks -_registered_pass_through_routes: Final[dict[str, dict[str, str | bool | list[str] | Mapping[str, object]]]] = {} +_registered_pass_through_routes: Final[dict[str, _RegisteredPassThroughRoute]] = {} def get_response_body(response: httpx.Response) -> dict | None: @@ -2239,7 +2249,7 @@ def create_pass_through_route( status_code=status.HTTP_405_METHOD_NOT_ALLOWED, detail=f"Method {request.method} is not allowed for pass-through endpoint {path}.", ) - target_params: Final = { + target_params: Final[dict[str, object]] = { "target": target, "custom_headers": custom_headers, "forward_headers": _forward_headers, @@ -2271,7 +2281,9 @@ def create_pass_through_route( # Ensure custom_headers is a dict. Botocore returns a HeadersDict # for SigV4-prepared requests, which is a Mapping but not a dict. - headers_dict: Final = dict(param_custom_headers) if isinstance(param_custom_headers, Mapping) else {} + headers_dict: Final[dict[str, object]] = ( + {**param_custom_headers} if isinstance(param_custom_headers, Mapping) else {} + ) # Ensure query_params and custom_body are dicts or None final_query_params: Final = query_params_data if isinstance(query_params_data, dict) else {} @@ -3104,9 +3116,7 @@ class SafeRouteAdder: return f"{path}/{{subpath:path}}" return None - def registered_route_methods( - route_info: Mapping[str, str | bool | Sequence[str] | Mapping[str, object]], - ) -> tuple[tuple[str, str], ...]: + def registered_route_methods(route_info: _RegisteredPassThroughRoute) -> tuple[tuple[str, str], ...]: starlette_path: Final = registered_route_path(route_info) methods: Final = route_info.get("methods") if starlette_path is None or not isinstance(methods, (list, tuple)): @@ -3462,7 +3472,7 @@ class InitPassThroughEndpointHelpers: return (comparison_route, *forwarded) @staticmethod - def get_registered_pass_through_route(route: str, method: str | None = None) -> dict[str, Any] | None: + def get_registered_pass_through_route(route: str, method: str | None = None) -> _RegisteredPassThroughRoute | None: """Get passthrough params for a given route and optionally filter by HTTP method""" comparison_route: Final = InitPassThroughEndpointHelpers._route_for_registry_lookup(route) for key in _registered_pass_through_routes: diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index 5442c0f3f81..8fd4fbfcb8b 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -6,7 +6,7 @@ from datetime import datetime, timedelta from types import MappingProxyType from typing import Final, Literal, TypeVar -from pydantic import TypeAdapter +from pydantic import ConfigDict, TypeAdapter from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger @@ -155,6 +155,7 @@ class _DailyUserSpendOwnerRow(LiteLLMBaseModel): last_owner: str | None = None +_RAW_ROWS: Final = TypeAdapter(Sequence[dict[str, object]], config=ConfigDict(strict=True)) _TOKEN_DIGEST_ROWS: Final = TypeAdapter(tuple[_TokenDigestRow, ...]) _SPEND_LOG_DIGEST_ROWS: Final = TypeAdapter(tuple[_SpendLogDigestRow, ...]) _DAILY_USER_SPEND_OWNER_ROWS: Final = TypeAdapter(tuple[_DailyUserSpendOwnerRow, ...]) @@ -194,7 +195,7 @@ async def _rows_within_the_statement_timeout( await transaction.execute_raw(_SPEND_LOG_STATEMENT_TIMEOUT_SQL) for setting in planner_settings: await transaction.execute_raw(setting) - return await transaction.query_raw(sql, *params) + return _RAW_ROWS.validate_python(await transaction.query_raw(sql, *params)) async def _reverse_hash_key_metadata( diff --git a/litellm/secret_managers/aws_secret_manager_v2.py b/litellm/secret_managers/aws_secret_manager_v2.py index 98d12e497da..de8566e6fdc 100644 --- a/litellm/secret_managers/aws_secret_manager_v2.py +++ b/litellm/secret_managers/aws_secret_manager_v2.py @@ -21,6 +21,8 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Final import httpx +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -41,6 +43,15 @@ if TYPE_CHECKING: from botocore.awsrequest import HTTPHeaders +@with_config(ConfigDict(extra="allow", strict=True)) +class _SecretStringResponse(TypedDict): + SecretString: ReadOnly[str] + + +_SECRET_STRING_RESPONSE: Final = TypeAdapter(_SecretStringResponse) +_AWS_SECRET_MANAGER_RESPONSE: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) + + class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): def __init__( self, @@ -154,7 +165,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): try: response: Final = await async_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8")) response.raise_for_status() - return response.json()["SecretString"] + return _SECRET_STRING_RESPONSE.validate_python(response.json())["SecretString"] except httpx.TimeoutException: raise ValueError("Timeout error occurred") except Exception as e: @@ -204,7 +215,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): try: response: Final = sync_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8")) - return response.json()["SecretString"] + return _SECRET_STRING_RESPONSE.validate_python(response.json())["SecretString"] except httpx.TimeoutException: raise ValueError("Timeout error occurred") except httpx.HTTPStatusError as e: @@ -424,7 +435,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): try: response: Final = await async_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8")) response.raise_for_status() - return response.json() + return _AWS_SECRET_MANAGER_RESPONSE.validate_python(response.json()) except httpx.HTTPStatusError as err: raise ValueError(f"HTTP error occurred: {err.response.text}") except httpx.TimeoutException: @@ -481,7 +492,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): try: response: Final = await async_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8")) response.raise_for_status() - return response.json() + return _AWS_SECRET_MANAGER_RESPONSE.validate_python(response.json()) except httpx.HTTPStatusError as err: raise ValueError(f"HTTP error occurred: {err.response.text}") except httpx.TimeoutException: @@ -533,7 +544,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): try: response: Final = await async_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8")) response.raise_for_status() - return response.json() + return _AWS_SECRET_MANAGER_RESPONSE.validate_python(response.json()) except httpx.HTTPStatusError as err: raise ValueError(f"HTTP error occurred: {err.response.text}") except httpx.TimeoutException: @@ -613,7 +624,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): try: response: Final = await async_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8")) response.raise_for_status() - return response.json() + return _AWS_SECRET_MANAGER_RESPONSE.validate_python(response.json()) except httpx.HTTPStatusError as err: raise ValueError(f"HTTP error occurred: {err.response.text}") except httpx.TimeoutException: diff --git a/litellm/types/integrations/rag/bedrock_knowledgebase.py b/litellm/types/integrations/rag/bedrock_knowledgebase.py index 7156d8101e1..7288eca03b4 100644 --- a/litellm/types/integrations/rag/bedrock_knowledgebase.py +++ b/litellm/types/integrations/rag/bedrock_knowledgebase.py @@ -1,4 +1,4 @@ -from typing import Any, Literal +from typing import Literal from typing_extensions import ReadOnly, TypedDict @@ -40,7 +40,7 @@ class BedrockKBRetrievalResult(TypedDict, total=False): content: BedrockKBContent | None location: BedrockKBLocation | None score: float | None - metadata: dict[str, Any] | None + metadata: dict[str, object] | None class BedrockKBResponse(TypedDict, total=False): @@ -87,21 +87,21 @@ class BedrockKBMetadataConfiguration(TypedDict, total=False): class BedrockKBModelConfiguration(TypedDict, total=False): """Model configuration for reranking.""" - additionalModelRequestFields: dict[str, Any] | None + additionalModelRequestFields: dict[str, object] | None modelArn: str | None class BedrockKBRerankingConfiguration(TypedDict, total=False): """Configuration for reranking in vector search.""" - bedrockRerankingConfiguration: dict[str, Any] | None # This could be further typed if needed + bedrockRerankingConfiguration: dict[str, object] | None # This could be further typed if needed type: str | None class BedrockKBVectorSearchConfiguration(TypedDict, total=False): """Configuration for vector search.""" - filter: dict[str, Any] | None + filter: dict[str, object] | None implicitFilterConfiguration: BedrockKBImplicitFilterConfiguration | None numberOfResults: int | None overrideSearchType: str | None diff --git a/litellm/types/llms/bedrock_invoke_agents.py b/litellm/types/llms/bedrock_invoke_agents.py index 62f3c18ca7a..b334c986d6b 100644 --- a/litellm/types/llms/bedrock_invoke_agents.py +++ b/litellm/types/llms/bedrock_invoke_agents.py @@ -4,7 +4,7 @@ Type definitions for AWS Bedrock Invoke Agent API responses. https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_InvokeAgent.html """ -from typing import Any, Final +from typing import Final from typing_extensions import TypedDict @@ -39,7 +39,7 @@ class InvokeAgentModelInvocationInput(TypedDict, total=False): """Model invocation input details.""" foundationModel: str | None - inferenceConfiguration: dict[str, Any] | None + inferenceConfiguration: dict[str, object] | None text: str | None traceId: str | None type: str | None @@ -49,9 +49,9 @@ class InvokeAgentModelInvocationOutput(TypedDict, total=False): """Model invocation output details.""" metadata: InvokeAgentMetadata | None - parsedResponse: dict[str, Any] | None - rawResponse: dict[str, Any] | None - reasoningContent: dict[str, Any] | None + parsedResponse: dict[str, object] | None + rawResponse: dict[str, object] | None + reasoningContent: dict[str, object] | None traceId: str | None diff --git a/litellm/types/llms/gemini.py b/litellm/types/llms/gemini.py index c300c26bcd3..2ad81e64e05 100644 --- a/litellm/types/llms/gemini.py +++ b/litellm/types/llms/gemini.py @@ -1,5 +1,6 @@ +from collections.abc import Mapping from enum import Enum -from typing import Any, Literal +from typing import Literal from typing_extensions import Required, TypedDict @@ -235,7 +236,7 @@ class GeminiVideoGenerationInstance(TypedDict, total=False): """Instance data for Gemini video generation request""" prompt: Required[str] - image: dict[str, Any] + image: Mapping[str, object] class GeminiVideoGenerationParameters(LiteLLMBaseModel): @@ -263,7 +264,7 @@ class GeminiVideoGenerationParameters(LiteLLMBaseModel): negativePrompt: str | None = None """Text describing what not to include in the video.""" - lastFrame: Any | None = None + lastFrame: object | None = None """ The final image for interpolation video to transition. Should be used with the 'image' parameter. @@ -275,7 +276,7 @@ class GeminiVideoGenerationParameters(LiteLLMBaseModel): Only supported in Veo 3.1 (list of VideoGenerationReferenceImage objects). """ - video: Any | None = None + video: object | None = None """ Video to be used for video extension (Video object). Only supported in Veo 3.1 & Veo 3 Fast. @@ -354,5 +355,5 @@ class GeminiLongRunningOperationResponse(LiteLLMBaseModel): response: GeminiOperationResponse | None = None """Response object when operation is complete""" - error: dict[str, Any] | None = None + error: dict[str, object] | None = None """Error details if operation failed""" diff --git a/litellm/types/llms/langgraph.py b/litellm/types/llms/langgraph.py index 7e34329992d..83f7bb090d1 100644 --- a/litellm/types/llms/langgraph.py +++ b/litellm/types/llms/langgraph.py @@ -4,7 +4,8 @@ Type definitions for LangGraph API. LangGraph provides a streaming and non-streaming API for running agents. """ -from typing import Any, Literal +from collections.abc import Mapping +from typing import Literal from typing_extensions import TypedDict @@ -29,8 +30,8 @@ class LangGraphRequest(TypedDict, total=False): assistant_id: str input: LangGraphInput stream_mode: str | None - config: dict[str, Any] | None - metadata: dict[str, Any] | None + config: Mapping[str, object] | None + metadata: Mapping[str, object] | None # Response Types - Streaming @@ -38,7 +39,7 @@ class LangGraphStreamEvent(TypedDict, total=False): """Single event in a LangGraph stream response.""" event: str - data: Any + data: object # Response Types - Non-streaming @@ -55,7 +56,7 @@ class LangGraphResponse(TypedDict, total=False): """Non-streaming response structure from LangGraph.""" messages: list[LangGraphResponseMessage] - values: dict[str, Any] + values: dict[str, object] # Parsed response for internal use diff --git a/litellm/types/mcp_server/tool_registry.py b/litellm/types/mcp_server/tool_registry.py index 738d37c9503..080e27b6ed3 100644 --- a/litellm/types/mcp_server/tool_registry.py +++ b/litellm/types/mcp_server/tool_registry.py @@ -1,5 +1,5 @@ from collections.abc import Callable -from typing import Any, ClassVar +from typing import ClassVar from pydantic import ConfigDict, Field @@ -10,7 +10,7 @@ class MCPTool(LiteLLMBaseModel): model_config: ClassVar[ConfigDict] = ConfigDict(arbitrary_types_allowed=True) name: str description: str - input_schema: dict[str, Any] + input_schema: dict[str, object] handler: Callable server_id: str | None = Field(default=None, frozen=True) @@ -18,18 +18,18 @@ class MCPTool(LiteLLMBaseModel): class ToolSchema(LiteLLMBaseModel): name: str description: str - inputSchema: dict[str, Any] + inputSchema: dict[str, object] class ListToolsResponse(LiteLLMBaseModel): tools: list[ToolSchema] nextCursor: str | None = None - _meta: dict[str, Any] | None = None + _meta: dict[str, object] | None = None class CallToolRequest(LiteLLMBaseModel): method: str = "tools/call" - params: dict[str, Any] + params: dict[str, object] class ContentItem(LiteLLMBaseModel): diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/enkryptai.py b/litellm/types/proxy/guardrails/guardrail_hooks/enkryptai.py index d1eba1dd4a9..f936c9d85cf 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/enkryptai.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/enkryptai.py @@ -1,5 +1,3 @@ -from typing import Any - from pydantic import Field from typing_extensions import TypedDict @@ -20,7 +18,7 @@ class EnkryptAIPolicyViolationDetail(TypedDict, total=False): class EnkryptAIPIIDetail(TypedDict, total=False): """Details for PII detection.""" - pii: dict[str, Any] + pii: dict[str, object] class EnkryptAIToxicityDetail(TypedDict, total=False): @@ -85,8 +83,8 @@ class EnkryptAIResponseDetails(TypedDict, total=False): toxicity: EnkryptAIToxicityDetail | None keyword_detected: EnkryptAIKeywordDetail | None bias: EnkryptAIBiasDetail | None - prompt_injection: dict[str, Any] | None - jailbreak: dict[str, Any] | None + prompt_injection: dict[str, object] | None + jailbreak: dict[str, object] | None class EnkryptAIResponse(TypedDict, total=False): @@ -107,7 +105,7 @@ class EnkryptAIProcessedResult(TypedDict): | EnkryptAIToxicityDetail | EnkryptAIKeywordDetail | EnkryptAIBiasDetail - | dict[str, Any], + | dict[str, object], ] diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py index 1890c68b30d..c36644ba536 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py @@ -242,6 +242,33 @@ def test_transform_response_hands_json_mode_to_nova(): assert json.loads(result.choices[0].message.content) == {"city": "Paris", "temperature": 21} +@pytest.mark.parametrize(("finish_reason", "expected"), [("MAX_TOKENS", "length"), (None, "stop")]) +def test_transform_response_accepts_null_cohere_finish_reason(finish_reason: str | None, expected: str) -> None: + from litellm.types.utils import ModelResponse + + raw_response: Final = httpx.Response( + 200, + json={"generations": [{"text": "Hello there", "finish_reason": finish_reason}]}, + request=httpx.Request("POST", "https://bedrock"), + ) + + result: Final = AmazonInvokeConfig().transform_response( + model="invoke/cohere.command-text-v14", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + api_key=None, + ) + + assert result.choices[0].message.content == "Hello there" + assert result.choices[0].finish_reason == expected + + def _stream_invoke_completion_with_spied_client(**kwargs: object) -> tuple[MagicMock, MagicMock]: mock_response = MagicMock() mock_response.status_code = 200 diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_twelvelabs_pegasus_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_twelvelabs_pegasus_transformation.py index d90e9933b88..5b3cb2e0c92 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_twelvelabs_pegasus_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_twelvelabs_pegasus_transformation.py @@ -1,6 +1,13 @@ +from typing import Final +from unittest.mock import Mock + +import httpx +import pytest + from litellm.llms.bedrock.chat.invoke_transformations.amazon_twelvelabs_pegasus_transformation import ( AmazonTwelveLabsPegasusConfig, ) +from litellm.types.utils import ModelResponse def _make_messages() -> list[dict]: @@ -84,3 +91,24 @@ def test_transform_request_includes_s3_media(): s3_location = request["mediaSource"]["s3Location"] assert s3_location["uri"] == "s3://test-bucket/video.mp4" assert s3_location["bucketOwner"] == "123456789012" + + +@pytest.mark.parametrize( + ("finish_reason", "expected"), + [("length", "length"), (None, "stop")], +) +def test_transform_response_accepts_null_finish_reason(finish_reason: str | None, expected: str) -> None: + result: Final = AmazonTwelveLabsPegasusConfig().transform_response( + model="twelvelabs.pegasus-1-2-v1:0", + raw_response=httpx.Response(200, json={"message": "A cat jumps.", "finishReason": finish_reason}), + model_response=ModelResponse(), + logging_obj=Mock(), + request_data={}, + messages=_make_messages(), + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert result.choices[0].message.content == "A cat jumps." + assert result.choices[0].finish_reason == expected From 9e878cfae2edb7a66ecb3d977b2f3320c5955e6d Mon Sep 17 00:00:00 2001 From: moe-berri Date: Wed, 28 Oct 2026 16:26:28 -0700 Subject: [PATCH 114/128] test(gemini): cover every web source for a grounded passage (#48861) * fix(gemini): preserve every grounding web source * test(gemini): cover all grounding sources per passage Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Zabir Raihan Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_gemini_grounding_annotations_wire.py | 65 +++++++++++++++++++ ...test_vertex_and_google_ai_studio_gemini.py | 33 ++++++++++ 2 files changed, 98 insertions(+) create mode 100644 tests/integration/providers/test_gemini_grounding_annotations_wire.py diff --git a/tests/integration/providers/test_gemini_grounding_annotations_wire.py b/tests/integration/providers/test_gemini_grounding_annotations_wire.py new file mode 100644 index 00000000000..11df7c214f3 --- /dev/null +++ b/tests/integration/providers/test_gemini_grounding_annotations_wire.py @@ -0,0 +1,65 @@ +import json +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server + + +def _grounded_reply(request: Request) -> Reply: + body: Final = { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "The passage"}]}, + "finishReason": "STOP", + "groundingMetadata": { + "groundingChunks": [ + {"retrievedContext": {"uri": "gs://documents/source"}}, + {"web": {"uri": "https://example.com/first", "title": "First"}}, + {"web": {"uri": "https://example.com/second", "title": "Second"}}, + ], + "groundingSupports": [ + {"segment": {"startIndex": 4, "endIndex": 11}, "groundingChunkIndices": [0, 99, 2, 1]} + ], + }, + } + ], + "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 2, "totalTokenCount": 7}, + } + assert json.loads(request.body)["contents"] == [{"role": "user", "parts": [{"text": "Find sources"}]}] + if ":streamGenerateContent" in request.target: + return Reply(content_type="text/event-stream", chunks=(f"data: {json.dumps(body)}\n\n".encode(),)) + return Reply(body=json.dumps(body).encode()) + + +@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) +def test_chat_grounding_annotations_include_every_web_source(gateway: Gateway, stream: bool) -> None: + with wire_server(_grounded_reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model="gemini/gemini-3.8-flash", api_base=wire.url, api_key="synthetic-key") + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "Find sources"}], "stream": stream}, + ) + assert response.status_code == 200, response.text + payloads: Final = ( + [json.loads(line[6:]) for line in response.text.splitlines() if line.startswith("data: {")] + if stream + else [response.json()] + ) + messages: Final = tuple(payload["choices"][0]["delta" if stream else "message"] for payload in payloads) + annotations: Final = tuple(message["annotations"] for message in messages if message.get("annotations")) + assert annotations == ( + [ + { + "type": "url_citation", + "url_citation": { + "start_index": 4, + "end_index": 11, + "url": f"https://example.com/{name.lower()}", + "title": name, + }, + } + for name in ("Second", "First") + ], + ), response.text diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 1546d1e9fe0..3e1bca18f08 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -3514,6 +3514,39 @@ def test_vertex_ai_annotation_conversion(): assert "Current time" in time_annotation["url_citation"]["title"] +@pytest.mark.parametrize( + "indices, expected_indices", + [([2, 0], [2, 0]), ([1, 99, -1, 2, 0], [2, 0]), ([0], [0]), ([], []), ([1, 99], [])], + ids=["multiple-sources", "non-web-first", "single-source", "empty", "no-web"], +) +def test_grounding_annotations_preserve_all_web_sources(indices: list[int], expected_indices: list[int]) -> None: + chunks: Final = [ + {"web": {"uri": "https://example.com/first", "title": "First"}}, + {"retrievedContext": {"uri": "gs://documents/source"}}, + {"web": {"uri": "https://example.com/second", "title": "Second"}}, + ] + metadata: Final = { + "groundingChunks": chunks, + "groundingSupports": [ + {"segment": {"startIndex": 4, "endIndex": 11}, "groundingChunkIndices": indices} + ], + } + annotations: Final = VertexGeminiConfig._convert_grounding_metadata_to_annotations([metadata], "The passage") + expected: Final = [ + { + "type": "url_citation", + "url_citation": { + "start_index": 4, + "end_index": 11, + "url": chunks[index]["web"]["uri"], + "title": chunks[index]["web"]["title"], + }, + } + for index in expected_indices + ] + assert json.dumps(annotations) == json.dumps(expected) + + def test_vertex_ai_annotation_empty_grounding_metadata(): """ Test handling of empty or missing grounding metadata. From 96fe78f96ecb6d0a5f92a6b7e68a3b770bb25213 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 28 Oct 2026 16:32:45 -0700 Subject: [PATCH 115/128] fix(router): reject string fallback targets at Router init (LIT-10238) (#48943) * fix(router): reject string fallback targets at Router init (LIT-10238) * fix(router): keep direct target fallback entries when rejecting string targets * test(router): set string fallbacks after init in cooldown tolerance tests --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 15 +++++- .../router_utils/test_cooldown_handlers.py | 7 +-- tests/unit/test_router/test_router.py | 48 +++++++++++++++++++ 3 files changed, 63 insertions(+), 7 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 5ca09499c3b..408cb20945f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1370,7 +1370,7 @@ class Router: }, ] - router = Router(model_list=model_list, fallbacks=[{"azure-gpt-3.5-turbo": "openai-gpt-3.5-turbo"}]) + router = Router(model_list=model_list, fallbacks=[{"azure-gpt-3.5-turbo": ["openai-gpt-3.5-turbo"]}]) ``` """ @@ -1596,6 +1596,10 @@ class Router: if default_fallbacks is not None or litellm.default_fallbacks is not None: _fallbacks = default_fallbacks or litellm.default_fallbacks + if isinstance(_fallbacks, str): + raise TypeError( + f"default_fallbacks must be a list of model names, got '{_fallbacks}'. Use ['{_fallbacks}']" + ) if self.fallbacks is not None: self.fallbacks.append({"*": _fallbacks}) else: @@ -2839,7 +2843,7 @@ class Router: self._initialize_core_endpoints() self._initialize_specialized_endpoints() - def validate_fallbacks(self, fallback_param: list | None): + def validate_fallbacks(self, fallback_param: Sequence[Mapping[str, object]] | None): """ Validate the fallbacks parameter. """ @@ -2852,6 +2856,13 @@ class Router: raise ValueError( f"Dictionary '{fallback_dict}' must have exactly one key, but has {len(fallback_dict)} keys." ) + ((model_group, targets),) = fallback_dict.items() + if isinstance(targets, list) or check_non_standard_fallback_format(fallbacks=[fallback_dict]): + continue + raise TypeError( + f"Fallbacks for '{model_group}' must be a list of model names, got '{targets}'. " + f"Use {{'{model_group}': ['{targets}']}}" + ) def _add_encrypted_content_affinity_check(self, enable_global_affinity: bool) -> None: def _move_before_deployment_affinity( diff --git a/tests/unit/router_utils/test_cooldown_handlers.py b/tests/unit/router_utils/test_cooldown_handlers.py index 24459a5a250..2e1e3acaabb 100644 --- a/tests/unit/router_utils/test_cooldown_handlers.py +++ b/tests/unit/router_utils/test_cooldown_handlers.py @@ -3150,8 +3150,6 @@ BENCHED: Final = (litellm.Timeout, litellm.Timeout, RouterRateLimitError, Router {}, BENCHED, ), - ({"allowed_fails": 1, "fallbacks": [{"test-model": "backup-model"}]}, {}, BENCHED), - ({"allowed_fails": 1, "default_fallbacks": "backup-model"}, {}, BENCHED), ], ids=[ "allowed_fails", @@ -3160,8 +3158,6 @@ BENCHED: Final = (litellm.Timeout, litellm.Timeout, RouterRateLimitError, Router "policy_names_timeout", "allowed_fails_with_a_fallback_group", "allowed_fails_with_a_dict_fallback_target", - "allowed_fails_with_a_string_fallback_target", - "allowed_fails_with_a_string_default_fallbacks", ], ) @pytest.mark.asyncio @@ -3207,7 +3203,8 @@ async def test_router_allowed_fails_benches_a_lone_deployment_only_when_the_poli def test_has_fallback_group_reads_string_targets_and_skips_unparseable_entries( fallbacks: Sequence[Mapping[str, object]], expected: bool ) -> None: - router: Final = Router(model_list=[TIMED_OUT_DEPLOYMENT], fallbacks=list(fallbacks)) + router: Final = Router(model_list=[TIMED_OUT_DEPLOYMENT]) + router.fallbacks = list(fallbacks) assert _has_fallback_group(router, "test-model") is expected diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 3c7a222d9a1..300c9cff6ff 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -31228,3 +31228,51 @@ async def test_anthropic_caller_rate_limit_does_not_block_other_credentials( finally: await GLOBAL_LOGGING_WORKER.flush() await client.close() + + +_FALLBACK_VALIDATION_MODEL_LIST: Final = ( + {"model_name": "gpt-4", "litellm_params": {"model": "openai/gpt-4", "api_key": "fake-key"}}, + {"model_name": "gpt-3.5", "litellm_params": {"model": "openai/gpt-3.5-turbo", "api_key": "fake-key"}}, +) + + +@pytest.mark.parametrize( + ("router_kwargs", "expected_message"), + [ + ({"fallbacks": [{"gpt-4": "gpt-3.5"}]}, "Fallbacks for 'gpt-4' must be a list of model names, got 'gpt-3.5'"), + ({"default_fallbacks": "gpt-3.5"}, "default_fallbacks must be a list of model names, got 'gpt-3.5'"), + ( + {"content_policy_fallbacks": [{"gpt-4": "gpt-3.5"}]}, + "Fallbacks for 'gpt-4' must be a list of model names, got 'gpt-3.5'", + ), + ], + ids=["string_fallback_target", "string_default_fallbacks", "string_content_policy_fallback_target"], +) +def test_router_rejects_a_string_fallback_target_at_init( + router_kwargs: Mapping[str, object], expected_message: str +) -> None: + with pytest.raises(TypeError, match="must be a list of model names") as excinfo: + Router(model_list=list(_FALLBACK_VALIDATION_MODEL_LIST), **router_kwargs) + + assert expected_message in str(excinfo.value) + + +@pytest.mark.parametrize( + ("router_kwargs", "expected_fallbacks"), + [ + ({"fallbacks": [{"gpt-4": ["gpt-3.5"]}]}, [{"gpt-4": ["gpt-3.5"]}]), + ({"default_fallbacks": ["gpt-3.5"]}, [{"*": ["gpt-3.5"]}]), + ( + {"fallbacks": [{"gpt-4": [{"model": "gpt-3.5", "api_key": "other-key"}]}]}, + [{"gpt-4": [{"model": "gpt-3.5", "api_key": "other-key"}]}], + ), + ({"fallbacks": [{"model": "gpt-3.5"}]}, [{"model": "gpt-3.5"}]), + ], + ids=["list_fallback_target", "list_default_fallbacks", "dict_override_fallback_target", "direct_target_entry"], +) +def test_router_keeps_supported_fallback_shapes( + router_kwargs: Mapping[str, object], expected_fallbacks: list[dict[str, object]] +) -> None: + router: Final = Router(model_list=list(_FALLBACK_VALIDATION_MODEL_LIST), **router_kwargs) + + assert router.fallbacks == expected_fallbacks From ab07b7207e37e52f1ac8126817d101c90f5d01a0 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:59:57 -0700 Subject: [PATCH 116/128] fix(vertex-ai): add batch above 200k and image output batch prices for gemini models (#48945) * fix(vertex-ai): add batch and flex above 200k and image output tier prices for gemini models * fix(vertex-ai): drop flex above 200k and image priority fields the cost map schema rejects Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 28 ++++++++++++++++--- model_prices_and_context_window.json | 28 ++++++++++++++++--- 2 files changed, 48 insertions(+), 8 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index dc043dc4d26..8719565a96f 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -28853,6 +28853,7 @@ "mode": "image_generation", "output_cost_per_image": 0.039, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_reasoning_token": 2.5e-06, "output_cost_per_token": 2.5e-06, "output_cost_per_token_batches": 1.25e-06, @@ -28910,6 +28911,7 @@ "mode": "image_generation", "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, + "output_cost_per_image_token_batches": 6e-05, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, @@ -28963,12 +28965,13 @@ "mode": "image_generation", "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, + "output_cost_per_image_token_batches": 6e-05, "output_cost_per_token": 1.2e-05, "output_cost_per_token_priority": 2.16e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/images/generations", @@ -29014,6 +29017,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, "output_cost_per_token": 3e-06, "output_cost_per_token_batches": 1.5e-06, "output_cost_per_token_flex": 1.5e-06, @@ -29059,9 +29063,10 @@ "mode": "image_generation", "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, "output_cost_per_token": 3e-06, "output_cost_per_token_batches": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/images/generations", @@ -29106,6 +29111,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 1.5e-06, "output_cost_per_token_batches": 7.5e-07, "output_cost_per_token_flex": 7.5e-07, @@ -29672,10 +29678,12 @@ }, "search_free_requests": {"count": 10000, "period": "day", "shared_group": "gemini-2.5-pro"}, "google_maps_grounding_cost_per_query": 0.025, + "input_cost_per_token_above_200k_tokens_batches": 1.25e-06, "input_cost_per_token_above_200k_tokens_priority": 4.5e-06, "input_cost_per_token_batches": 6.25e-07, "input_cost_per_token_flex": 6.25e-07, "input_cost_per_token_priority": 2.25e-06, + "output_cost_per_token_above_200k_tokens_batches": 7.5e-06, "output_cost_per_token_above_200k_tokens_priority": 2.7e-05, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_flex": 5e-06, @@ -29690,6 +29698,7 @@ "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, @@ -29701,6 +29710,7 @@ "output_cost_per_image": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_batches": 9e-06, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "output_cost_per_token_flex": 6e-06, @@ -30229,6 +30239,7 @@ "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, @@ -30240,6 +30251,7 @@ "output_cost_per_image": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_batches": 9e-06, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "output_cost_per_token_flex": 6e-06, @@ -51910,6 +51922,7 @@ "mode": "image_generation", "output_cost_per_image": 0.039, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_reasoning_token": 2.5e-06, "output_cost_per_token": 2.5e-06, "output_cost_per_token_batches": 1.25e-06, @@ -51968,6 +51981,7 @@ "mode": "image_generation", "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, + "output_cost_per_image_token_batches": 6e-05, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, @@ -52002,13 +52016,14 @@ "mode": "image_generation", "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, + "output_cost_per_image_token_batches": 6e-05, "output_cost_per_token": 1.2e-05, "output_cost_per_token_priority": 2.16e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "supports_reasoning": false, - "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/images/generations", @@ -52034,6 +52049,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, "output_cost_per_token": 3e-06, "output_cost_per_token_batches": 1.5e-06, "output_cost_per_token_flex": 1.5e-06, @@ -52061,10 +52077,11 @@ "mode": "image_generation", "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, "output_cost_per_token": 3e-06, "output_cost_per_token_batches": 1.5e-06, "supports_reasoning": false, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/images/generations", @@ -52089,6 +52106,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 1.5e-06, "output_cost_per_token_batches": 7.5e-07, "output_cost_per_token_flex": 7.5e-07, @@ -74400,6 +74418,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, @@ -74445,6 +74464,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index dc043dc4d26..8719565a96f 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -28853,6 +28853,7 @@ "mode": "image_generation", "output_cost_per_image": 0.039, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_reasoning_token": 2.5e-06, "output_cost_per_token": 2.5e-06, "output_cost_per_token_batches": 1.25e-06, @@ -28910,6 +28911,7 @@ "mode": "image_generation", "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, + "output_cost_per_image_token_batches": 6e-05, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, @@ -28963,12 +28965,13 @@ "mode": "image_generation", "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, + "output_cost_per_image_token_batches": 6e-05, "output_cost_per_token": 1.2e-05, "output_cost_per_token_priority": 2.16e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/images/generations", @@ -29014,6 +29017,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, "output_cost_per_token": 3e-06, "output_cost_per_token_batches": 1.5e-06, "output_cost_per_token_flex": 1.5e-06, @@ -29059,9 +29063,10 @@ "mode": "image_generation", "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, "output_cost_per_token": 3e-06, "output_cost_per_token_batches": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/images/generations", @@ -29106,6 +29111,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 1.5e-06, "output_cost_per_token_batches": 7.5e-07, "output_cost_per_token_flex": 7.5e-07, @@ -29672,10 +29678,12 @@ }, "search_free_requests": {"count": 10000, "period": "day", "shared_group": "gemini-2.5-pro"}, "google_maps_grounding_cost_per_query": 0.025, + "input_cost_per_token_above_200k_tokens_batches": 1.25e-06, "input_cost_per_token_above_200k_tokens_priority": 4.5e-06, "input_cost_per_token_batches": 6.25e-07, "input_cost_per_token_flex": 6.25e-07, "input_cost_per_token_priority": 2.25e-06, + "output_cost_per_token_above_200k_tokens_batches": 7.5e-06, "output_cost_per_token_above_200k_tokens_priority": 2.7e-05, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_flex": 5e-06, @@ -29690,6 +29698,7 @@ "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, @@ -29701,6 +29710,7 @@ "output_cost_per_image": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_batches": 9e-06, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "output_cost_per_token_flex": 6e-06, @@ -30229,6 +30239,7 @@ "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, @@ -30240,6 +30251,7 @@ "output_cost_per_image": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_batches": 9e-06, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "output_cost_per_token_flex": 6e-06, @@ -51910,6 +51922,7 @@ "mode": "image_generation", "output_cost_per_image": 0.039, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_reasoning_token": 2.5e-06, "output_cost_per_token": 2.5e-06, "output_cost_per_token_batches": 1.25e-06, @@ -51968,6 +51981,7 @@ "mode": "image_generation", "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, + "output_cost_per_image_token_batches": 6e-05, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, @@ -52002,13 +52016,14 @@ "mode": "image_generation", "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, + "output_cost_per_image_token_batches": 6e-05, "output_cost_per_token": 1.2e-05, "output_cost_per_token_priority": 2.16e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "supports_reasoning": false, - "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/images/generations", @@ -52034,6 +52049,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, "output_cost_per_token": 3e-06, "output_cost_per_token_batches": 1.5e-06, "output_cost_per_token_flex": 1.5e-06, @@ -52061,10 +52077,11 @@ "mode": "image_generation", "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, "output_cost_per_token": 3e-06, "output_cost_per_token_batches": 1.5e-06, "supports_reasoning": false, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/images/generations", @@ -52089,6 +52106,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 1.5e-06, "output_cost_per_token_batches": 7.5e-07, "output_cost_per_token_flex": 7.5e-07, @@ -74400,6 +74418,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, @@ -74445,6 +74464,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, From d06b889d413237881e3dc27486e4345dd7aa18f9 Mon Sep 17 00:00:00 2001 From: rj-cairn <349020452+rj-cairn@users.noreply.github.com> Date: Wed, 28 Oct 2026 17:00:29 -0700 Subject: [PATCH 117/128] fix(streaming): keep merged think tags balanced and split SSE lines only on CR/LF (#48755) * fix(vertex_ai): split Gemini SSE lines on CR/LF only so U+2028 in text keeps the stream The Gemini and Vertex AI streaming path read the SSE body with httpx iter_lines()/aiter_lines(). Those follow str.splitlines(), which also breaks on U+0085, U+2028, U+2029, \x0b, \x0c and \x1c-\x1e. JSON leaves U+0085, U+2028 and U+2029 unescaped inside strings, so a data: line whose text holds one of them was cut in two. The tail is not a data: line and the head never parses, and the accumulate-and-retry fallback then carries the unparsed prefix forward, so every later event is appended to it and the stream ends with no more content and no usage. SSE lines end only at CRLF, LF or CR. The new iter_sse_lines/aiter_sse_lines split decoded text on exactly those, treating a CR at the end of one chunk and an LF at the start of the next as one break. The first request and every long-decode hop now read through them. Co-Authored-By: Claude Opus 5.5 * fix(streaming): keep merged think tags balanced and split shared SSE lines only on CR/LF Close an open block on the first content or tool-call delta and at stream end, reopen it when reasoning returns, and route the shared chat completions stream through the SSE line splitter so U+2028 and U+2029 inside JSON no longer drop chunks. * fix(streaming): exact-match the SSE [DONE] marker in BaseModelResponseIterator * fix(streaming): split every SSE reader on CR/LF only and keep the logged think block closed Switch the remaining httpx stream readers (Anthropic chat, Google GenAI, Interactions, Codestral, openai_like, Predibase, LangGraph, Vertex Agent Engine, A2A providers, Azure AI agents, sandboxes, the A2A endpoint, and others) from iter_lines/aiter_lines to iter_sse_lines/aiter_sse_lines, so a raw U+2028 or U+2029 inside a JSON string no longer splits an event. Keep the chunk that closes a think block in the logged and cached response when a stream ends without a finish_reason, close the block before a legacy function_call, and make sync Codestral FIM streaming send stream=true and use the text-completion-codestral stream parser. Test fakes that only mocked iter_lines now stream real httpx bytes and assert the content they feed in. * fix(streaming): split buffered SSE parsers only on CR and LF, shard the codestral tests Buffered SSE and JSONL parsers used str.splitlines(), which also breaks on U+2028, U+2029 and U+0085 that JSON allows raw inside strings. They now use split_sse_lines from sse_lines.py, which also gains a str.split fast path. Adds tests/unit/llms/codestral to the unit workflow shard. * fix(streaming): keep U+2028 in restamped Anthropic frames and fold CRLF before splitting SSE lines * test: type the new stream tests, drop a comment, and benchmark the SSE line splitter * test: type the iterator class parameters in the line-separator stream tests * fix(responses): stream tool call arguments opened in a text chunk * fix(responses): hold a tool call that shares a chunk with text until the text ends * test: type the a2a SSE handler helper and the sandbox fake responses --------- Co-authored-by: Tim Vink Co-authored-by: Claude Opus 5.5 Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: michelligabriele --- .github/workflows/test-unit.yml | 1 + .../bedrock_agentcore/transformation.py | 5 +- .../providers/watsonx_orchestrate/handler.py | 5 +- .../transformation.py | 3 +- litellm/google_genai/streaming_iterator.py | 9 +- litellm/interactions/streaming_iterator.py | 5 +- .../litellm_core_utils/streaming_handler.py | 57 +++--- litellm/llms/anthropic/chat/handler.py | 5 +- .../anthropic/messages/stream_assembly.py | 3 +- litellm/llms/azure_ai/agents/handler.py | 3 +- litellm/llms/base_llm/base_model_iterator.py | 2 +- .../base_llm/guardrail_translation/utils.py | 3 +- .../llms/base_llm/sandbox/transformation.py | 3 +- litellm/llms/bedrock/files/transformation.py | 3 +- .../guardrail_translation/handler.py | 3 +- litellm/llms/bedrock/search/transformation.py | 5 +- .../llms/chatgpt/responses/transformation.py | 3 +- litellm/llms/codestral/completion/handler.py | 9 +- litellm/llms/custom_httpx/aiohttp_handler.py | 3 +- litellm/llms/custom_httpx/llm_http_handler.py | 5 +- litellm/llms/custom_httpx/sse_lines.py | 74 ++++++++ .../guardrail_translation/handler.py | 3 +- .../huggingface/embedding/transformation.py | 3 +- litellm/llms/langgraph/chat/sse_iterator.py | 5 +- litellm/llms/oobabooga/chat/oobabooga.py | 3 +- litellm/llms/openai_like/chat/handler.py | 9 +- litellm/llms/predibase/chat/handler.py | 5 +- .../vertex_ai/agent_engine/transformation.py | 5 +- .../vertex_and_google_ai_studio_gemini.py | 9 +- .../proxy/agent_endpoints/a2a_endpoints.py | 3 +- .../streaming_model_restamp.py | 4 +- .../anthropic_passthrough_logging_handler.py | 7 +- .../tinyfish_passthrough_logging_handler.py | 3 +- .../managed_id_rewriter.py | 3 +- .../response_polling/background_streaming.py | 3 +- .../streaming_iterator.py | 13 +- tests/benchmarks/test_sse_lines_benchmarks.py | 27 +++ .../test_bedrock_agentcore_a2a.py | 32 ++++ ...test_watsonx_orchestrate_transformation.py | 20 ++- ...responses_transformation_transformation.py | 18 ++ .../test_google_genai_streaming_iterator.py | 83 +++++---- .../test_interactions_streaming_iterator.py | 45 ++++- .../test_streaming_handler.py | 114 +++++++++++- .../chat/test_anthropic_chat_handler.py | 59 ++++++- .../messages/test_stream_assembly.py | 24 +++ .../unit/llms/azure_ai/agents/test_handler.py | 23 ++- .../azure_ai/agents/test_transformation.py | 20 +-- .../guardrail_translation/test_utils.py | 17 ++ .../test_bedrock_files_transformation.py | 14 ++ .../guardrail_translation/test_handler.py | 14 +- .../test_agentcore_search_transformation.py | 15 ++ .../test_chatgpt_responses_transformation.py | 31 ++++ .../test_cloudflare_transformation.py | 47 ++++- tests/unit/llms/codestral/__init__.py | 0 .../llms/codestral/completion/__init__.py | 0 .../test_codestral_completion_handler.py | 50 ++++++ .../llms/custom_httpx/test_aiohttp_handler.py | 37 ++++ .../unit/llms/custom_httpx/test_sse_lines.py | 64 +++++++ .../test_databricks_chat_transformation.py | 14 +- .../test_fireworks_ai_chat_transformation.py | 6 +- .../test_google_genai_guardrail_handler.py | 16 ++ .../chat/test_langgraph_sse_iterator.py | 32 ++++ .../llms/oobabooga/chat/test_oobabooga.py | 29 ++++ .../chat/test_openai_like_chat_handler.py | 37 +++- .../test_snowflake_chat_transformation.py | 30 +--- .../llms/test_predibase_transformation.py | 47 +++++ tests/unit/llms/triton/test_triton.py | 11 +- .../agent_engine/test_transformation.py | 47 ++++- ...test_vertex_and_google_ai_studio_gemini.py | 64 ++++++- .../llms/vertex_ai/test_http_status_201.py | 8 +- .../agent_endpoints/test_a2a_endpoints.py | 138 +++++++-------- .../test_streaming_model_restamp.py | 23 +++ ...t_anthropic_passthrough_logging_handler.py | 48 +++++ ...st_tinyfish_passthrough_logging_handler.py | 7 + .../test_managed_id_rewriter.py | 33 +++- .../test_response_polling_handler.py | 31 +++- .../test_streaming_iterator_transformation.py | 100 +++++++++++ tests/unit/sandbox/test_e2b_sandbox.py | 28 ++- .../unit/sandbox/test_opensandbox_sandbox.py | 13 +- tests/unit/test_groq_streaming_encoding.py | 164 ++++-------------- 80 files changed, 1554 insertions(+), 413 deletions(-) create mode 100644 litellm/llms/custom_httpx/sse_lines.py create mode 100644 tests/benchmarks/test_sse_lines_benchmarks.py create mode 100644 tests/unit/llms/codestral/__init__.py create mode 100644 tests/unit/llms/codestral/completion/__init__.py create mode 100644 tests/unit/llms/codestral/completion/test_codestral_completion_handler.py create mode 100644 tests/unit/llms/custom_httpx/test_sse_lines.py create mode 100644 tests/unit/llms/langgraph/chat/test_langgraph_sse_iterator.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 2ddf21bca17..80d5018fe9a 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -245,6 +245,7 @@ jobs: tests/unit/llms/clarifai tests/unit/llms/claude_code tests/unit/llms/cloudflare + tests/unit/llms/codestral tests/unit/llms/codex tests/unit/llms/cohere tests/unit/llms/cometapi diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py index 1cea0194715..d4400fbb4cc 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py @@ -19,6 +19,7 @@ from litellm.a2a_protocol.utils import ( ) from litellm.exceptions import BadRequestError from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines RUNTIME_SESSION_ID_MIN_LENGTH: Final = 33 RUNTIME_SESSION_ID_MAX_LENGTH: Final = 256 @@ -50,7 +51,7 @@ _RESERVED_PREFIX_HEADERS: Final[tuple[str, ...]] = ( class _SSELineSource(Protocol): """Minimal streaming-response surface used to read SSE lines.""" - def aiter_lines(self) -> AsyncIterator[str]: ... + def aiter_text(self) -> AsyncIterator[str]: ... def _filter_reserved_headers( @@ -244,7 +245,7 @@ class BedrockAgentCoreA2ATransformation: Yields: Parsed JSON dicts from SSE data lines """ - async for line in response.aiter_lines(): + async for line in aiter_sse_lines(response.aiter_text()): line = line.strip() if not line: continue diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py index 0a2871d25cc..3937c47ea3c 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py @@ -20,6 +20,7 @@ from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, get_async_httpx_client, ) +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines from litellm.types.llms.custom_http import httpxSpecialProvider _IBM_CLOUD_IAM_URL: Final = "https://iam.cloud.ibm.com/identity/token" @@ -74,7 +75,7 @@ class _WXORun(TypedDict, total=False): class _SSELineSource(Protocol): - def aiter_lines(self) -> AsyncIterator[str]: ... + def aiter_text(self) -> AsyncIterator[str]: ... class _WXOView(TypedDict, total=False): @@ -231,7 +232,7 @@ class WatsonxOrchestrateHandler: async def _accumulate_wxo_sse_text(response: _SSELineSource) -> str: source: Final[_WXOView] = {"sse_source": response} accumulated_text = "" - async for line in source["sse_source"].aiter_lines(): + async for line in aiter_sse_lines(source["sse_source"].aiter_text()): if not line.startswith("data:"): continue data_str = line[5:].strip() diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 47a3a903d82..ee3639ac387 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -34,6 +34,7 @@ from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.bridges.completion_transformation import ( CompletionTransformationBridge, ) +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.responses.litellm_completion_transformation.tool_search import is_anthropic_tool_search_chat_tool from litellm.responses.sse_output_recovery import ( parse_sse_json_chunk, @@ -1351,7 +1352,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): recovered_output_items: Final[dict[int, dict[str, object]]] = {} recovered_text_only_items: Final[dict[int, dict[str, object]]] = {} - for chunk in raw_sse.splitlines(): + for chunk in split_sse_lines(raw_sse): parsed_chunk = parse_sse_json_chunk(chunk) if parsed_chunk is None: continue diff --git a/litellm/google_genai/streaming_iterator.py b/litellm/google_genai/streaming_iterator.py index a22c47cb1e7..08ac7d4f9dc 100644 --- a/litellm/google_genai/streaming_iterator.py +++ b/litellm/google_genai/streaming_iterator.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Final import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) @@ -131,9 +132,9 @@ class GoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContent self.model = model self.generate_content_provider_config = generate_content_provider_config self.litellm_metadata = litellm_metadata - # Gemini streamGenerateContent uses SSE line framing; iter_lines keeps + # Gemini streamGenerateContent uses SSE line framing; iter_sse_lines keeps # large inlineData payloads (e.g. image/jpeg) intact within one event. - self.stream_iterator = response.iter_lines() + self.stream_iterator = iter_sse_lines(response.iter_text()) def __iter__(self): return self @@ -182,9 +183,9 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateCo self.model = model self.generate_content_provider_config = generate_content_provider_config self.litellm_metadata = litellm_metadata - # Gemini streamGenerateContent uses SSE line framing; aiter_lines keeps + # Gemini streamGenerateContent uses SSE line framing; aiter_sse_lines keeps # large inlineData payloads (e.g. image/jpeg) intact within one event. - self.stream_iterator = response.aiter_lines() + self.stream_iterator = aiter_sse_lines(response.aiter_text()) def __aiter__(self): return self diff --git a/litellm/interactions/streaming_iterator.py b/litellm/interactions/streaming_iterator.py index f4ad57bf7b4..c7faa3fc510 100644 --- a/litellm/interactions/streaming_iterator.py +++ b/litellm/interactions/streaming_iterator.py @@ -21,6 +21,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.llms.base_llm.interactions.transformation import BaseInteractionsAPIConfig +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.types.interactions import ( InteractionsAPIStreamingResponse, ) @@ -145,7 +146,7 @@ class InteractionsAPIStreamingIterator(BaseInteractionsAPIStreamingIterator): litellm_metadata=litellm_metadata, custom_llm_provider=custom_llm_provider, ) - self.stream_iterator = response.aiter_lines() + self.stream_iterator = aiter_sse_lines(response.aiter_text()) def __aiter__(self): return self @@ -212,7 +213,7 @@ class SyncInteractionsAPIStreamingIterator(BaseInteractionsAPIStreamingIterator) litellm_metadata=litellm_metadata, custom_llm_provider=custom_llm_provider, ) - self.stream_iterator = response.iter_lines() + self.stream_iterator = iter_sse_lines(response.iter_text()) def __iter__(self): return self diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 3fb25de43cc..2348d8e57d6 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1263,7 +1263,7 @@ class CustomStreamWrapper: ) if _is_delta_empty: - model_response.choices[0].delta = Delta(content=None) # ensure empty delta chunk returned + model_response.choices[0].delta = Delta(content=self._close_open_think_block()) # get any function call arguments model_response.choices[0].finish_reason = map_finish_reason( finish_reason=self.received_finish_reason @@ -1292,21 +1292,26 @@ class CustomStreamWrapper: if self.merge_reasoning_content_in_choices is not True: return delta: Final = model_response.choices[0].delta - reasoning_content: Final = getattr(delta, "reasoning_content", None) + reasoning_content: Final = getattr(delta, "reasoning_content", None) or "" content: Final = delta.content or "" - if reasoning_content: - self._emitted_disqualifying_content = True - closes_thinking: Final = bool(content) and not self.sent_last_thinking_block - opening: Final = "" if self.sent_first_thinking_block else "" - delta.content = opening + reasoning_content + ("" + content if closes_thinking else "") - self.sent_first_thinking_block = True - self.sent_last_thinking_block = self.sent_last_thinking_block or closes_thinking - elif self.sent_first_thinking_block and not self.sent_last_thinking_block and content: - delta.content = "" + content - self.sent_last_thinking_block = True - + ends_thinking: Final = bool(content or delta.tool_calls or delta.function_call) + inside_think_block: Final = self.sent_first_thinking_block and not self.sent_last_thinking_block + opening: Final = "" if reasoning_content and not inside_think_block else "" + closing: Final = "" if ends_thinking and (reasoning_content or inside_think_block) else "" if hasattr(delta, "reasoning_content"): del delta.reasoning_content + if not reasoning_content and not closing: + return + delta.content = opening + reasoning_content + closing + content + self.sent_first_thinking_block = True + self.sent_last_thinking_block = bool(closing) + self._emitted_disqualifying_content = self._emitted_disqualifying_content or bool(reasoning_content) + + def _close_open_think_block(self) -> str | None: + if not self.sent_first_thinking_block or self.sent_last_thinking_block: + return None + self.sent_last_thinking_block = True + return "" def _dispatch_provider_chunk( self, @@ -1932,8 +1937,20 @@ class CustomStreamWrapper: model_response.choices[0].finish_reason == "stop" and self.tool_call ): # don't overwrite for other - potential error finish reasons model_response.choices[0].finish_reason = "tool_calls" + close_tag: Final = self._close_open_think_block() + if close_tag is not None: + model_response.choices[0].delta.content = close_tag return model_response + def _final_chunk_for_logging(self, processed_chunk: "ModelResponseStream") -> "ModelResponseStream | None": + if self.received_finish_reason is not None or self.intermittent_finish_reason is not None: + return processed_chunk + if processed_chunk.choices[0].delta.content is None: + return None + return processed_chunk.model_copy( + update={"choices": [processed_chunk.choices[0].model_copy(update={"finish_reason": None})]} + ) + def _record_usage_only_chunk(self, model_response: "ModelResponseStream") -> None: """ Keep provider usage-only chunks (e.g. OpenRouter's post-finish chunk, which carries a @@ -2173,10 +2190,9 @@ class CustomStreamWrapper: else: self.sent_last_chunk = True processed_chunk: Final = self.finish_reason_handler() - # The logged response is built from self.chunks; keep a finish_reason the provider sent on its - # last content chunk (stripped there), but never add the synthetic "stop" used when it sent none. - if self.received_finish_reason is not None or self.intermittent_finish_reason is not None: - self.chunks.append(processed_chunk) + final_chunk_for_logging: Final = self._final_chunk_for_logging(processed_chunk) + if final_chunk_for_logging is not None: + self.chunks.append(final_chunk_for_logging) if self.stream_options is None: self._set_final_chunk_usage(processed_chunk) ## LOGGING @@ -2406,10 +2422,9 @@ class CustomStreamWrapper: else: self.sent_last_chunk = True processed_chunk: Final = self.finish_reason_handler() - # The logged response is built from self.chunks; keep a finish_reason the provider sent on its - # last content chunk (stripped there), but never add the synthetic "stop" used when it sent none. - if self.received_finish_reason is not None or self.intermittent_finish_reason is not None: - self.chunks.append(processed_chunk) + final_chunk_for_logging: Final = self._final_chunk_for_logging(processed_chunk) + if final_chunk_for_logging is not None: + self.chunks.append(final_chunk_for_logging) if self.stream_options is None: self._set_final_chunk_usage(processed_chunk) # see sync __next__'s sibling branch: deliberately do NOT restore diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 13ceea3e31b..8a82ed98991 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -38,6 +38,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, get_httpx_client, ) +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.types.llms.anthropic import ( ContentBlockDelta, ContentBlockStart, @@ -233,7 +234,7 @@ async def make_call( raise AnthropicError(status_code=500, message=str(e)) completion_stream: Final = ModelResponseIterator( - streaming_response=response.aiter_lines(), + streaming_response=aiter_sse_lines(response.aiter_text()), sync_stream=False, json_mode=json_mode, speed=speed, @@ -304,7 +305,7 @@ def make_sync_call( ) completion_stream: Final = ModelResponseIterator( - streaming_response=response.iter_lines(), + streaming_response=iter_sse_lines(response.iter_text()), sync_stream=True, json_mode=json_mode, speed=speed, diff --git a/litellm/llms/anthropic/messages/stream_assembly.py b/litellm/llms/anthropic/messages/stream_assembly.py index 90bdad8f4bc..c103eb0f975 100644 --- a/litellm/llms/anthropic/messages/stream_assembly.py +++ b/litellm/llms/anthropic/messages/stream_assembly.py @@ -13,6 +13,7 @@ from litellm.llms.anthropic.chat.handler import ( ) from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.llms.anthropic.messages.native.utils import MESSAGE_DELTA_OPTIONAL_KEYS +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.types.utils import ( Choices, Message, @@ -269,7 +270,7 @@ def build_complete_streaming_response_legacy( def extract_sse_data(event_str: str) -> dict | None: """Parse the JSON object from the ``data:`` line of an Anthropic SSE event.""" - for line in event_str.splitlines(): + for line in split_sse_lines(event_str): stripped = line.strip() if stripped.startswith("data:"): payload = stripped[len("data:") :].strip() diff --git a/litellm/llms/azure_ai/agents/handler.py b/litellm/llms/azure_ai/agents/handler.py index 94f1399f01b..423a09c5c92 100644 --- a/litellm/llms/azure_ai/agents/handler.py +++ b/litellm/llms/azure_ai/agents/handler.py @@ -35,6 +35,7 @@ from litellm.llms.azure_ai.agents.transformation import ( AzureAIAgentsConfig, AzureAIAgentsError, ) +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines from litellm.types.llms.openai import ( ChatCompletionAnnotation, ChatCompletionAnnotationURLCitation, @@ -632,7 +633,7 @@ class AzureAIAgentsHandler: current_event = None - async for line in response.aiter_lines(): + async for line in aiter_sse_lines(response.aiter_text()): line = line.strip() if line.startswith("event:"): diff --git a/litellm/llms/base_llm/base_model_iterator.py b/litellm/llms/base_llm/base_model_iterator.py index 032951be74a..3afc685b3dd 100644 --- a/litellm/llms/base_llm/base_model_iterator.py +++ b/litellm/llms/base_llm/base_model_iterator.py @@ -96,7 +96,7 @@ class BaseModelResponseIterator: released (and a backend like vLLM aborts generation) when the stream is abandoned before its natural end. - ``streaming_response`` is usually a bare ``aiter_lines()`` generator + ``streaming_response`` is usually a bare ``aiter_sse_lines()`` generator that holds no reference to the response, so the handler that owns the response attaches it here after construction.""" if self.http_response is not None: diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 725146fe894..fb886e42780 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -6,6 +6,7 @@ from typing import Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row ha from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage from litellm.types.llms.openai import AllMessageValues, ResponseAPIUsage @@ -24,7 +25,7 @@ def _anthropic_stream_chunk_events(item: object) -> list[dict]: events: Final[list[dict]] = [] for block in chunk.split("\n\n"): - for line in block.splitlines(): + for line in split_sse_lines(block): stripped = line.strip() if not stripped.startswith("data:"): continue diff --git a/litellm/llms/base_llm/sandbox/transformation.py b/litellm/llms/base_llm/sandbox/transformation.py index fe0e3abb87e..9b9b4e0a947 100644 --- a/litellm/llms/base_llm/sandbox/transformation.py +++ b/litellm/llms/base_llm/sandbox/transformation.py @@ -12,6 +12,7 @@ from typing import Final import httpx from pydantic import Field, PrivateAttr +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines from litellm.types.llms.base import LiteLLMPydanticObjectBase SANDBOX_MAX_OUTPUT_BYTES: Final = 10 * 1024 * 1024 @@ -99,7 +100,7 @@ class BaseSandboxConfig: async def _read_capped_lines(self, response: httpx.Response) -> list[str]: lines: Final[list[str]] = [] total = 0 - async for line in response.aiter_lines(): + async for line in aiter_sse_lines(response.aiter_text()): total += len(line.encode("utf-8")) if total > SANDBOX_MAX_OUTPUT_BYTES: raise ValueError( diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 8c4702bcf99..0624c2c65db 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -49,6 +49,7 @@ from litellm.llms.base_llm.files.transformation import ( BaseFilesConfig, LiteLLMLoggingObj, ) +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.bedrock import AwsAuthParams, BedrockBatchRecordKind from litellm.types.llms.openai import ( @@ -400,7 +401,7 @@ def _batch_input_line(line_number: int, line: str) -> _BatchInputLine: def _non_blank_lines(file_content: str) -> Iterator[tuple[int, str]]: - numbered_lines: Final = enumerate(file_content.splitlines(), start=1) + numbered_lines: Final = enumerate(split_sse_lines(file_content), start=1) return ((line_number, line) for line_number, line in numbered_lines if line.strip()) diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index 1b01b0ec6d0..3a65a8928bf 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -45,6 +45,7 @@ from litellm.llms.bedrock.passthrough.guardrail_translation.structured_edits imp structured_edits, tool_edits, ) +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.proxy.guardrails.auto_router_compression import is_compression_guardrail from litellm.types.utils import GenericGuardrailAPIInputs @@ -947,7 +948,7 @@ def _invoke_frames(items: Sequence[object]) -> tuple[InvokeStreamFrame, ...]: def _sse_events(sse: bytes) -> tuple[Mapping[str, object], ...]: return tuple( _JSON_OBJECT.validate_python(json.loads(line[len("data:") :])) - for line in sse.decode().splitlines() + for line in split_sse_lines(sse.decode()) if line.startswith("data:") ) diff --git a/litellm/llms/bedrock/search/transformation.py b/litellm/llms/bedrock/search/transformation.py index 565fc12aa1e..6f72cadad67 100644 --- a/litellm/llms/bedrock/search/transformation.py +++ b/litellm/llms/bedrock/search/transformation.py @@ -47,6 +47,7 @@ from litellm.llms.base_llm.search.transformation import ( ) from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import BedrockError +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.secret_managers.main import get_secret_str # AgentCore web-search rejects queries longer than 200 characters @@ -146,7 +147,9 @@ def _iter_sse_events(text: str) -> Iterator[Mapping[str, object]]: progress notifications before the JSON-RPC response. """ for chunk in _SSE_EVENT_SEPARATOR.split(text): - payload = "\n".join(line[len("data:") :].lstrip() for line in chunk.splitlines() if line.startswith("data:")) + payload = "\n".join( + line[len("data:") :].lstrip() for line in split_sse_lines(chunk) if line.startswith("data:") + ) if not payload: continue try: diff --git a/litellm/llms/chatgpt/responses/transformation.py b/litellm/llms/chatgpt/responses/transformation.py index 040b7b8c1d8..72b5cde21d1 100644 --- a/litellm/llms/chatgpt/responses/transformation.py +++ b/litellm/llms/chatgpt/responses/transformation.py @@ -11,6 +11,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo safe_convert_created_field, ) from litellm.litellm_core_utils.url_utils import join_api_base +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.llms.openai.common_utils import OpenAIError from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.responses.sse_output_recovery import ( @@ -199,7 +200,7 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): error_message = None streamed_output_items: Final[dict[int, dict[str, object]]] = {} text_only_output_items: Final[dict[int, dict[str, object]]] = {} - for chunk in body_text.splitlines(): + for chunk in split_sse_lines(body_text): parsed_chunk = parse_sse_json_chunk(chunk) if parsed_chunk is None: continue diff --git a/litellm/llms/codestral/completion/handler.py b/litellm/llms/codestral/completion/handler.py index 6bebfe2d444..6b172e77295 100644 --- a/litellm/llms/codestral/completion/handler.py +++ b/litellm/llms/codestral/completion/handler.py @@ -20,6 +20,7 @@ from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, get_async_httpx_client, ) +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.types.utils import TextChoices from litellm.utils import CustomStreamWrapper, TextCompletionResponse @@ -109,7 +110,7 @@ async def make_call( if response.status_code != 200: raise TextCompletionCodestralError(status_code=response.status_code, message=response.text) - completion_stream: Final = response.aiter_lines() + completion_stream: Final = aiter_sse_lines(response.aiter_text()) # LOGGING logging_obj.post_call( input=messages, @@ -339,13 +340,13 @@ class CodestralTextCompletion: response = litellm.module_level_client.post( completion_url, headers=headers, - data=json.dumps(data), + data=json.dumps({**data, "stream": True}), stream=stream, ) _response: Final = CustomStreamWrapper( - response.iter_lines(), + iter_sse_lines(response.iter_text()), model, - custom_llm_provider="codestral", + custom_llm_provider="text-completion-codestral", logging_obj=logging_obj, ) return _response diff --git a/litellm/llms/custom_httpx/aiohttp_handler.py b/litellm/llms/custom_httpx/aiohttp_handler.py index 447d7b685f7..4060d84f7f4 100644 --- a/litellm/llms/custom_httpx/aiohttp_handler.py +++ b/litellm/llms/custom_httpx/aiohttp_handler.py @@ -21,6 +21,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_httpx_client, get_ssl_configuration, ) +from litellm.llms.custom_httpx.sse_lines import iter_sse_lines from litellm.types.llms.openai import FileTypes from litellm.types.utils import HttpHandlerRequestFields, ImageResponse, LlmProviders from litellm.utils import CustomStreamWrapper, ModelResponse, ProviderConfigManager @@ -466,7 +467,7 @@ class BaseLLMAIOHTTPHandler: ) else: completion_stream = provider_config.get_model_response_iterator( - streaming_response=response.iter_lines(), sync_stream=True + streaming_response=iter_sse_lines(response.iter_text()), sync_stream=True ) # LOGGING diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index afa12b58b29..e7da6586b97 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -118,6 +118,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_httpx_client, header_value, ) +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.responses.streaming_iterator import ( BaseResponsesAPIStreamingIterator, MockResponsesAPIStreamingIterator, @@ -1086,7 +1087,7 @@ class BaseLLMHTTPHandler: completion_stream: object = MockResponseIterator(model_response=model_response, json_mode=json_mode) else: completion_stream = provider_config.get_model_response_iterator_for_request( - streaming_response=response.iter_lines(), + streaming_response=iter_sse_lines(response.iter_text()), sync_stream=True, json_mode=json_mode, optional_params=_OPTIONAL_PARAMS.validate_python(optional_params), @@ -1229,7 +1230,7 @@ class BaseLLMHTTPHandler: completion_stream: object = MockResponseIterator(model_response=model_response, json_mode=json_mode) else: completion_stream = provider_config.get_model_response_iterator_for_request( - streaming_response=response.aiter_lines(), + streaming_response=aiter_sse_lines(response.aiter_text()), sync_stream=False, json_mode=json_mode, optional_params=_OPTIONAL_PARAMS.validate_python(optional_params), diff --git a/litellm/llms/custom_httpx/sse_lines.py b/litellm/llms/custom_httpx/sse_lines.py new file mode 100644 index 00000000000..71ea7d621eb --- /dev/null +++ b/litellm/llms/custom_httpx/sse_lines.py @@ -0,0 +1,74 @@ +""" +Split a decoded Server-Sent Events body into lines the way the SSE spec does. + +httpx's `iter_lines()` / `aiter_lines()` follow `str.splitlines()`, which also breaks on U+0085, U+2028, +U+2029, \\x0b, \\x0c and \\x1c-\\x1e. JSON leaves U+0085, U+2028 and U+2029 unescaped inside strings, so a +`data:` line carrying one of them gets cut in half. SSE lines end only at CRLF, LF or CR +(https://html.spec.whatwg.org/multipage/server-sent-events.html#parsing-an-event-stream) +""" + +import io +import re +from collections.abc import AsyncIterable, AsyncIterator, Iterable, Iterator, Sequence +from typing import Final + +_SSE_LINE_BREAK: Final = re.compile(r"\r\n|\r|\n") +_SSE_LINE_WITH_BREAK: Final = re.compile(r"[^\r\n]*(?:\r\n|\r|\n)|[^\r\n]+") + + +def _split_lines(text: str) -> Sequence[str]: + if "\r" not in text: + return text.split("\n") + crlf_folded: Final = text.replace("\r\n", "\n") + if "\r" not in crlf_folded: + return crlf_folded.split("\n") + return _SSE_LINE_BREAK.split(text) + + +class _SSELineSplitter: + def __init__(self) -> None: + self._partial_line: io.StringIO = io.StringIO() + self._after_cr: bool = False + + def feed(self, text: str) -> tuple[str, ...]: + if not text: + return () + fresh: Final = text[1:] if self._after_cr and text.startswith("\n") else text + self._after_cr = fresh.endswith("\r") + *ended, tail = _split_lines(fresh) + if not ended: + _ = self._partial_line.write(tail) + return () + first: Final = self._partial_line.getvalue() + ended[0] + self._partial_line = io.StringIO() + _ = self._partial_line.write(tail) + return (first, *ended[1:]) + + def flush(self) -> tuple[str, ...]: + rest: Final = self._partial_line.getvalue() + self._partial_line = io.StringIO() + return (rest,) if rest else () + + +def iter_sse_lines(text_chunks: Iterable[str]) -> Iterator[str]: + splitter: Final = _SSELineSplitter() + for text in text_chunks: + yield from splitter.feed(text) + yield from splitter.flush() + + +def split_sse_lines(text: str) -> tuple[str, ...]: + return tuple(iter_sse_lines((text,))) + + +def split_sse_lines_keepends(text: str) -> tuple[str, ...]: + return tuple(_SSE_LINE_WITH_BREAK.findall(text)) + + +async def aiter_sse_lines(text_chunks: AsyncIterable[str]) -> AsyncIterator[str]: + splitter: Final = _SSELineSplitter() + async for text in text_chunks: + for line in splitter.feed(text): + yield line + for line in splitter.flush(): + yield line diff --git a/litellm/llms/gemini/google_genai/guardrail_translation/handler.py b/litellm/llms/gemini/google_genai/guardrail_translation/handler.py index 34007bc8fd4..b6b2aa6c28f 100644 --- a/litellm/llms/gemini/google_genai/guardrail_translation/handler.py +++ b/litellm/llms/gemini/google_genai/guardrail_translation/handler.py @@ -23,6 +23,7 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, StreamTransformSink, ) +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -121,7 +122,7 @@ def _parse_json_dict_or_none(payload: str) -> Mapping[str, object] | None: def _sse_payload_texts(sse_text: str) -> tuple[str, ...]: return tuple( text - for line in sse_text.splitlines() + for line in split_sse_lines(sse_text) if line.startswith("data:") for payload in (line[len("data:") :].strip(),) if payload and payload != "[DONE]" diff --git a/litellm/llms/huggingface/embedding/transformation.py b/litellm/llms/huggingface/embedding/transformation.py index 1b7d6483739..6e6a117227e 100644 --- a/litellm/llms/huggingface/embedding/transformation.py +++ b/litellm/llms/huggingface/embedding/transformation.py @@ -17,6 +17,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( ) from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.llms.custom_httpx.sse_lines import iter_sse_lines from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import Choices, Message, ModelResponse, Usage @@ -357,7 +358,7 @@ class HuggingFaceEmbeddingConfig(BaseConfig): api_key: str | None = None, ) -> list[dict[str, str]]: streamed_response: Final = CustomStreamWrapper( - completion_stream=response.iter_lines(), + completion_stream=iter_sse_lines(response.iter_text()), model=model, custom_llm_provider="huggingface", logging_obj=logging_obj, diff --git a/litellm/llms/langgraph/chat/sse_iterator.py b/litellm/llms/langgraph/chat/sse_iterator.py index 575820bceee..6fe355ad98c 100644 --- a/litellm/llms/langgraph/chat/sse_iterator.py +++ b/litellm/llms/langgraph/chat/sse_iterator.py @@ -12,6 +12,7 @@ import httpx from pydantic import TypeAdapter from litellm._logging import verbose_logger +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices _DECODED_JSON: Final = TypeAdapter(object) @@ -36,12 +37,12 @@ class LangGraphSSEStreamIterator: def __iter__(self): """Initialize sync iteration.""" - self.line_iterator = self.response.iter_lines() + self.line_iterator = iter_sse_lines(self.response.iter_text()) return self def __aiter__(self): """Initialize async iteration.""" - self.async_line_iterator = self.response.aiter_lines() + self.async_line_iterator = aiter_sse_lines(self.response.aiter_text()) return self def _parse_sse_line(self, line: str) -> ModelResponseStream | None: diff --git a/litellm/llms/oobabooga/chat/oobabooga.py b/litellm/llms/oobabooga/chat/oobabooga.py index 8879c20c5eb..b4dc57ff5ed 100644 --- a/litellm/llms/oobabooga/chat/oobabooga.py +++ b/litellm/llms/oobabooga/chat/oobabooga.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Final import litellm from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.custom_httpx.http_handler import get_httpx_client +from litellm.llms.custom_httpx.sse_lines import iter_sse_lines from litellm.utils import EmbeddingResponse, ModelResponse, Usage from ..common_utils import OobaboogaError @@ -74,7 +75,7 @@ def completion( stream=optional_params["stream"] if "stream" in optional_params else False, ) if "stream" in optional_params and optional_params["stream"] is True: - return response.iter_lines() + return iter_sse_lines(response.iter_text()) else: return oobabooga_config.transform_response( model=model, diff --git a/litellm/llms/openai_like/chat/handler.py b/litellm/llms/openai_like/chat/handler.py index 2c74ef19e7d..a450947afdf 100644 --- a/litellm/llms/openai_like/chat/handler.py +++ b/litellm/llms/openai_like/chat/handler.py @@ -16,6 +16,7 @@ from litellm import LlmProviders from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.base_model_iterator import MockResponseIterator from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.llms.databricks.streaming_utils import ModelResponseIterator from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.llms.openai.openai import OpenAIConfig @@ -65,7 +66,9 @@ async def make_call( model_response: Final = _fake_streamed_model_response(response.json()) completion_stream = MockResponseIterator(model_response=model_response) else: - completion_stream = ModelResponseIterator(streaming_response=response.aiter_lines(), sync_stream=False) + completion_stream = ModelResponseIterator( + streaming_response=aiter_sse_lines(response.aiter_text()), sync_stream=False + ) # LOGGING logging_obj.post_call( input=messages, @@ -103,7 +106,9 @@ def make_sync_call( model_response: Final = _fake_streamed_model_response(response.json()) completion_stream = MockResponseIterator(model_response=model_response) else: - completion_stream = ModelResponseIterator(streaming_response=response.iter_lines(), sync_stream=True) + completion_stream = ModelResponseIterator( + streaming_response=iter_sse_lines(response.iter_text()), sync_stream=True + ) # LOGGING logging_obj.post_call( diff --git a/litellm/llms/predibase/chat/handler.py b/litellm/llms/predibase/chat/handler.py index d0a61ea6e00..2ba4a63e08b 100644 --- a/litellm/llms/predibase/chat/handler.py +++ b/litellm/llms/predibase/chat/handler.py @@ -14,6 +14,7 @@ from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, get_async_httpx_client, ) +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.utils import CustomStreamWrapper, ModelResponse from ..common_utils import PredibaseError @@ -34,7 +35,7 @@ async def make_call( if response.status_code != 200: raise PredibaseError(status_code=response.status_code, message=response.text) - completion_stream: Final = response.aiter_lines() + completion_stream: Final = aiter_sse_lines(response.aiter_text()) # LOGGING logging_obj.post_call( input=messages, @@ -163,7 +164,7 @@ class PredibaseChatCompletion: timeout=timeout, ) _response: Final = CustomStreamWrapper( - response.iter_lines(), + iter_sse_lines(response.iter_text()), model, custom_llm_provider="predibase", logging_obj=logging_obj, diff --git a/litellm/llms/vertex_ai/agent_engine/transformation.py b/litellm/llms/vertex_ai/agent_engine/transformation.py index 9f592398f74..0339533e1f5 100644 --- a/litellm/llms/vertex_ai/agent_engine/transformation.py +++ b/litellm/llms/vertex_ai/agent_engine/transformation.py @@ -22,6 +22,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, ) from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.llms.vertex_ai.agent_engine.sse_iterator import ( VertexAgentEngineResponseIterator, ) @@ -357,7 +358,7 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase): ) -> VertexAgentEngineResponseIterator: """Return a streaming iterator for SSE responses.""" return VertexAgentEngineResponseIterator( - streaming_response=raw_response.iter_lines(), + streaming_response=iter_sse_lines(raw_response.iter_text()), sync_stream=True, ) @@ -465,7 +466,7 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase): # Create iterator for SSE stream (async) completion_stream: Final = VertexAgentEngineResponseIterator( - streaming_response=response.aiter_lines(), + streaming_response=aiter_sse_lines(response.aiter_text()), sync_stream=False, ) 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 f27593dc94e..c506192f253 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 @@ -41,6 +41,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, get_httpx_client, ) +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.types.llms.anthropic import AnthropicThinkingParam from litellm.types.llms.base import CachedTokensDetails from litellm.types.llms.gemini import BidiGenerateContentServerMessage @@ -3241,7 +3242,7 @@ async def make_call( raise completion_stream: Final = ModelResponseIterator( - streaming_response=response.aiter_lines(), + streaming_response=aiter_sse_lines(response.aiter_text()), sync_stream=False, logging_obj=logging_obj, response_headers=response.headers, @@ -3312,7 +3313,7 @@ def make_sync_call( raise completion_stream: Final = ModelResponseIterator( - streaming_response=response.iter_lines(), + streaming_response=iter_sse_lines(response.iter_text()), sync_stream=True, logging_obj=logging_obj, response_headers=response.headers, @@ -4190,7 +4191,7 @@ class ModelResponseIterator: log_failed_hop(e) return self._truncated_stream_end(token) self._start_hop(hop_response, token) - self.streaming_response = hop_response.aiter_lines() + self.streaming_response = aiter_sse_lines(hop_response.aiter_text()) self.async_response_iterator = self.streaming_response return None @@ -4204,7 +4205,7 @@ class ModelResponseIterator: log_failed_hop(e) return self._truncated_stream_end(token) self._start_hop(hop_response, token) - self.streaming_response = hop_response.iter_lines() + self.streaming_response = iter_sse_lines(hop_response.iter_text()) self.response_iterator = self.streaming_response return None diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 4691c7b0bc5..219ce906340 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -26,6 +26,7 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.url_utils import SSRFError, validate_url from litellm.llms.a2a.common_utils import resolve_a2a_hop_auth_header +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.a2a.version_convert import ( A2AVersion, @@ -320,7 +321,7 @@ async def _a2a_sse_event_source( "error": {"code": -32603, "message": resp.reason_phrase}, } return - async for line in resp.aiter_lines(): + async for line in aiter_sse_lines(resp.aiter_text()): stripped = line.strip() if not stripped.startswith("data:"): continue diff --git a/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py b/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py index f748a754fe0..32e0f6e8682 100644 --- a/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py +++ b/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py @@ -14,6 +14,8 @@ from typing import Final from pydantic import TypeAdapter, ValidationError +from litellm.llms.custom_httpx.sse_lines import split_sse_lines_keepends + _MESSAGE_START_EVENT: Final = "message_start" _MESSAGE_START_MARKER: Final = b"message_start" _SSE_DATA_FIELD: Final = "data:" @@ -52,7 +54,7 @@ def _restamped_data_line(line: str, requested_model: str) -> str | None: def _restamped_frame(frame: str, requested_model: str) -> str | None: - lines: Final = frame.splitlines(keepends=True) + lines: Final = split_sse_lines_keepends(frame) restamped: Final = tuple(_restamped_data_line(line, requested_model) for line in lines) if all(line is None for line in restamped): return None diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index 7c7eb51bc5d..5415314f29b 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -29,6 +29,7 @@ from litellm.llms.anthropic.messages.stream_assembly import ( split_sse_chunk_into_events, ) from litellm.llms.bedrock.common_utils import BEDROCK_INVOKE_TRACE_FIELD +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import ( @@ -61,7 +62,7 @@ _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_i def _event_payload(event: str) -> Mapping[str, object] | None: data_lines: Final = tuple( - line.strip()[len("data:") :].strip() for line in event.splitlines() if line.strip().startswith("data:") + line.strip()[len("data:") :].strip() for line in split_sse_lines(event) if line.strip().startswith("data:") ) if not data_lines or data_lines[0] in ("", "[DONE]"): return None @@ -362,7 +363,7 @@ class AnthropicPassthroughLoggingHandler: ) -> str | None: for raw in all_chunks: text = raw.decode("utf-8") if isinstance(raw, bytes) else raw - for line in text.splitlines(): + for line in split_sse_lines(text): if not line.startswith("data:"): continue try: @@ -402,7 +403,7 @@ class AnthropicPassthroughLoggingHandler: """ for raw in reversed(all_chunks): text = raw.decode("utf-8") if isinstance(raw, bytes) else raw - for line in reversed(text.splitlines()): + for line in reversed(split_sse_lines(text)): if not line.startswith("data:"): continue try: diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py index 2afc1d17bc2..0c4a3f64852 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py @@ -21,6 +21,7 @@ from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, get_async_httpx_client, ) +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.passthrough_endpoints.tinyfish import ( @@ -81,7 +82,7 @@ def sse_poller_spawned(logging_obj: LiteLLMLoggingObj) -> bool: def run_id_from_sse_frames(frames: bytes) -> str | None: - return _run_id_from_sse_chunks(frames.decode("utf-8", errors="replace").splitlines()) + return _run_id_from_sse_chunks(split_sse_lines(frames.decode("utf-8", errors="replace"))) def resolve_tinyfish_agent_api_base() -> str: diff --git a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py index 9790c150ae1..cae8ccc93ee 100644 --- a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py +++ b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py @@ -51,6 +51,7 @@ from litellm.llms.base_llm.managed_resources.isolation import ( can_access_resource, resolve_resource_owner_id, ) +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames @@ -830,7 +831,7 @@ _SSE_EVENT_ADAPTER: Final = TypeAdapter(Mapping[str, JsonValue]) def _first_streamed_response(frames: bytes) -> tuple[str, Mapping[str, JsonValue]] | None: - for line in frames.decode("utf-8", errors="replace").splitlines(): + for line in split_sse_lines(frames.decode("utf-8", errors="replace")): if not line.startswith(_SSE_DATA_PREFIX): continue try: diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index 5f6a7432afb..1eb76f8f558 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -18,6 +18,7 @@ from starlette.types import Message from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler @@ -81,7 +82,7 @@ class _StreamsBody(Protocol): def _sse_frame_data(frame: str) -> str | None: - return next((line[6:].strip() for line in frame.splitlines() if line.startswith("data: ")), None) + return next((line[6:].strip() for line in split_sse_lines(frame) if line.startswith("data: ")), None) async def _never_receive() -> Message: diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 219944de521..982158040c4 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -165,6 +165,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._upstream_exhausted: bool = False self._response_id_primed: bool = False self._pending_tool_events: list[BaseLiteLLMOpenAIResponseObject] = [] + self._tool_calls_held_behind_text: tuple[object, ...] = () self._tool_output_index_by_call_id: dict[str, int] = {} self._tool_args_by_call_id: dict[str, str] = {} self._tool_item_id_by_call_id: dict[str, str] = {} # mutable-ok: filled per call id as tool call events stream @@ -388,6 +389,15 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return "" return self._tool_call_id_by_index.get(index, "") + def _hold_tool_calls_behind_text(self, tool_calls: object) -> None: + if isinstance(tool_calls, list): + self._tool_calls_held_behind_text = (*self._tool_calls_held_behind_text, *tool_calls) + + def _release_held_tool_calls(self) -> tuple[object, ...]: + held: Final = self._tool_calls_held_behind_text + self._tool_calls_held_behind_text = () + return held + def _queue_tool_call_delta_events(self, tool_calls: object) -> None: """ Convert chat-completions streaming `tool_calls` deltas into Responses API streaming events. @@ -1365,6 +1375,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if delta_content: if not self.sent_message_item_added_event: self._queue_message_item_added_events() + self._hold_tool_calls_behind_text(getattr(chunk.choices[0].delta, "tool_calls", None)) self._sequence_number += 1 text_delta_event: Final = OutputTextDeltaEvent( type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, @@ -1379,7 +1390,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): # Priority 3: Handle tool call deltas (if any) -> queue events and emit them # For each tool call delta, we emit events one at a time to match OpenAI's streaming behavior if chunk.choices and hasattr(chunk.choices[0].delta, "tool_calls") and chunk.choices[0].delta.tool_calls: - self._queue_tool_call_delta_events(chunk.choices[0].delta.tool_calls) + self._queue_tool_call_delta_events([*self._release_held_tool_calls(), *chunk.choices[0].delta.tool_calls]) # Return one pending tool event at a time if self._pending_tool_events: return self._pending_tool_events.pop(0) diff --git a/tests/benchmarks/test_sse_lines_benchmarks.py b/tests/benchmarks/test_sse_lines_benchmarks.py new file mode 100644 index 00000000000..1320fae49a6 --- /dev/null +++ b/tests/benchmarks/test_sse_lines_benchmarks.py @@ -0,0 +1,27 @@ +from collections.abc import Sequence +from typing import Final + +import pytest + +from litellm.llms.custom_httpx.sse_lines import iter_sse_lines + +EVENT: Final = 'data: {"choices":[{"delta":{"content":"hello world"}}]}' + + +def _chunks(line_break: str) -> Sequence[str]: + body: Final = "".join(f"{EVENT}{line_break}{line_break}" for _ in range(2000)) + return tuple(body[i : i + 512] for i in range(0, len(body), 512)) + + +LF_CHUNKS: Final = _chunks("\n") +CRLF_CHUNKS: Final = _chunks("\r\n") + + +@pytest.mark.benchmark +def test_iter_sse_lines_lf_stream() -> None: + assert sum(1 for _ in iter_sse_lines(LF_CHUNKS)) == 4000 + + +@pytest.mark.benchmark +def test_iter_sse_lines_crlf_stream() -> None: + assert sum(1 for _ in iter_sse_lines(CRLF_CHUNKS)) == 4000 diff --git a/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py b/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py index 09f09170241..cd3f7622fa2 100644 --- a/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py +++ b/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py @@ -10,6 +10,7 @@ Verifies that: """ import json +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -560,6 +561,37 @@ class TestStreaming: assert sent_body["params"]["message"]["messageId"] == "msg-001" assert [event["result"]["kind"] for event in events] == ["task", "status-update"] + @pytest.mark.asyncio + @pytest.mark.usefixtures("httpx_transport") + async def test_streaming_keeps_unicode_line_separators_inside_an_sse_event(self) -> None: + from litellm.a2a_protocol.providers.bedrock_agentcore.config import ( + BedrockAgentCoreA2AConfig, + ) + + event: Final = { + "jsonrpc": "2.0", + "id": "req-001", + "result": {"kind": "artifact-update", "artifact": {"parts": [{"kind": "text", "text": "a\u2028b\u2029c"}]}}, + } + with respx.mock(assert_all_called=True) as router: + router.post(url__regex=r".*/invocations.*").mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text=f"data: {json.dumps(event, ensure_ascii=False)}\n\n", + ) + ) + events: Final = [ + streamed + async for streamed in BedrockAgentCoreA2AConfig().handle_streaming( + request_id="req-001", + params=SAMPLE_PARAMS, + litellm_params=SAMPLE_LITELLM_PARAMS, + ) + ] + + assert events == [event] + class TestConfigManager: """Test that config manager routes 'bedrock' correctly.""" diff --git a/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py b/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py index a300560ae9d..ac3f0390f84 100644 --- a/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py +++ b/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py @@ -1,6 +1,7 @@ import asyncio import json import time +from typing import Final import httpx import pytest @@ -36,13 +37,8 @@ class _ShortTtlTokenClient: return _JsonResponse({"access_token": f"token-{self.calls}", "expires_in": 30}) -class _SSELines: - def __init__(self, lines): - self.lines = lines - - async def aiter_lines(self): - for line in self.lines: - yield line +def _sse_response(lines: list[str]) -> httpx.Response: + return httpx.Response(200, stream=httpx.ByteStream("".join(f"{line}\n" for line in lines).encode())) class _InvalidJsonStreamResponse: @@ -218,7 +214,7 @@ def test_cp4d_token_ttl_from_absolute_expiration(): @pytest.mark.asyncio async def test_accumulate_wxo_sse_text_ignores_non_dict_json_events(): - response = _SSELines( + response = _sse_response( [ "data: null", "data: true", @@ -230,6 +226,14 @@ async def test_accumulate_wxo_sse_text_ignores_non_dict_json_events(): ) +@pytest.mark.asyncio +async def test_accumulate_wxo_sse_text_keeps_unicode_line_separators_inside_an_event() -> None: + separated_text: Final = "before\u2028between\u2029after" + response: Final = _sse_response([f"data: {json.dumps({'results': separated_text}, ensure_ascii=False)}", ""]) + + assert await WatsonxOrchestrateHandler._accumulate_wxo_sse_text(response) == separated_text + + @pytest.mark.asyncio async def test_short_lived_tokens_are_not_served_from_cache(): client = _ShortTtlTokenClient() diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 9aa494c1aab..46f1dbf2e7a 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -1102,6 +1102,24 @@ def test_recover_output_items_merges_text_only_items_at_distinct_indices(): assert recovered[1]["content"][0]["text"] == "From OUTPUT_TEXT_DONE only" +def test_recover_output_items_keeps_an_item_whose_text_holds_unicode_line_separators() -> None: + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + + item: Final = { + "type": "message", + "id": "msg_item_0", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "From\u2028OUTPUT_ITEM\u2029DONE", "annotations": []}], + } + event: Final = {"type": "response.output_item.done", "output_index": 0, "item": item} + raw_sse: Final = f"data: {json.dumps(event, ensure_ascii=False)}\ndata: [DONE]\n" + + assert LiteLLMResponsesTransformationHandler._recover_output_items_from_raw_sse(raw_sse) == [item] + + def test_transform_response_prefers_completed_output_from_raw_sse(): from litellm.completion_extras.litellm_responses_transformation.transformation import ( LiteLLMResponsesTransformationHandler, 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 120c99b6e42..cbb9638b3d9 100644 --- a/tests/unit/google_genai/test_google_genai_streaming_iterator.py +++ b/tests/unit/google_genai/test_google_genai_streaming_iterator.py @@ -62,7 +62,9 @@ def _logging_obj( ], ) def test_streaming_logging_targets_the_provider_that_served_the_request( - iterator_cls: type, + iterator_cls: ( + type[AsyncGoogleGenAIGenerateContentStreamingIterator] | type[GoogleGenAIGenerateContentStreamingIterator] + ), custom_llm_provider: str, expected_endpoint_type: EndpointType, ): @@ -99,18 +101,15 @@ def _large_inline_data_event() -> str: return f"data: {json.dumps(payload)}" +def _sse_response(body: str) -> httpx.Response: + return httpx.Response(200, stream=httpx.ByteStream(body.encode())) + + @pytest.mark.asyncio async def test_async_streaming_iterator_yields_complete_sse_events(): """Large inlineData must not be split across byte-chunk boundaries.""" - mock_response = MagicMock() - - async def _aiter_lines(): - yield _large_inline_data_event() - - mock_response.aiter_lines = _aiter_lines - iterator = AsyncGoogleGenAIGenerateContentStreamingIterator( - response=mock_response, + response=_sse_response(_large_inline_data_event() + "\n\n"), model="gemini-3.1-flash-image-preview", logging_obj=_logging_obj(), generate_content_provider_config=MagicMock(), @@ -128,11 +127,8 @@ async def test_async_streaming_iterator_yields_complete_sse_events(): def test_sync_streaming_iterator_yields_complete_sse_events(): - mock_response = MagicMock() - mock_response.iter_lines.return_value = iter([_large_inline_data_event()]) - iterator = GoogleGenAIGenerateContentStreamingIterator( - response=mock_response, + response=_sse_response(_large_inline_data_event() + "\n\n"), model="gemini-3.1-flash-image-preview", logging_obj=_logging_obj(), generate_content_provider_config=MagicMock(), @@ -150,17 +146,8 @@ def test_sync_streaming_iterator_yields_complete_sse_events(): @pytest.mark.asyncio async def test_async_streaming_iterator_preserves_multi_field_sse_event(): - mock_response = MagicMock() - - async def _aiter_lines(): - yield "event: message" - yield 'data: {"text":"hi"}' - yield "" - - mock_response.aiter_lines = _aiter_lines - iterator = AsyncGoogleGenAIGenerateContentStreamingIterator( - response=mock_response, + response=_sse_response('event: message\ndata: {"text":"hi"}\n\n'), model="gemini-test", logging_obj=_logging_obj(), generate_content_provider_config=MagicMock(), @@ -174,16 +161,8 @@ async def test_async_streaming_iterator_preserves_multi_field_sse_event(): @pytest.mark.asyncio async def test_async_streaming_iterator_forwards_sse_comment_events(): - mock_response = MagicMock() - - async def _aiter_lines(): - yield ": keepalive" - yield "" - - mock_response.aiter_lines = _aiter_lines - iterator = AsyncGoogleGenAIGenerateContentStreamingIterator( - response=mock_response, + response=_sse_response(": keepalive\n\n"), model="gemini-test", logging_obj=_logging_obj(), generate_content_provider_config=MagicMock(), @@ -195,6 +174,46 @@ async def test_async_streaming_iterator_forwards_sse_comment_events(): assert chunk == b": keepalive\n\n" +_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after" + + +async def _first_chunk( + iterator: GoogleGenAIGenerateContentStreamingIterator | AsyncGoogleGenAIGenerateContentStreamingIterator, +) -> bytes: + if isinstance(iterator, AsyncGoogleGenAIGenerateContentStreamingIterator): + return await iterator.__anext__() + return next(iterator) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "iterator_cls", + [AsyncGoogleGenAIGenerateContentStreamingIterator, GoogleGenAIGenerateContentStreamingIterator], +) +async def test_streaming_iterator_keeps_unicode_line_separators_inside_the_event( + iterator_cls: ( + type[AsyncGoogleGenAIGenerateContentStreamingIterator] | type[GoogleGenAIGenerateContentStreamingIterator] + ), +) -> None: + payload: Final = {"candidates": [{"content": {"parts": [{"text": _LINE_SEPARATED_TEXT}]}}]} + event: Final = f"data: {json.dumps(payload, ensure_ascii=False)}\n\n" + iterator: Final = iterator_cls( + response=_sse_response(event), + model="gemini-3.8-flash", + logging_obj=_logging_obj(), + generate_content_provider_config=MagicMock(), + litellm_metadata={}, + custom_llm_provider="gemini", + ) + + chunk: Final = await _first_chunk(iterator) + + assert chunk == event.encode() + assert ( + json.loads(chunk[len(b"data: ") : -2])["candidates"][0]["content"]["parts"][0]["text"] == _LINE_SEPARATED_TEXT + ) + + def _gemini_event(text: str, finish_reason: str | None = None) -> str: candidate: Final = {"content": {"role": "model", "parts": [{"text": text}]}, "index": 0} payload: Final = { diff --git a/tests/unit/interactions/test_interactions_streaming_iterator.py b/tests/unit/interactions/test_interactions_streaming_iterator.py index afcd0d559d7..fb1fe92d637 100644 --- a/tests/unit/interactions/test_interactions_streaming_iterator.py +++ b/tests/unit/interactions/test_interactions_streaming_iterator.py @@ -16,7 +16,10 @@ import pytest import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.interactions import streaming_iterator as interactions_streaming_iterator_module -from litellm.interactions.streaming_iterator import InteractionsAPIStreamingIterator +from litellm.interactions.streaming_iterator import ( + InteractionsAPIStreamingIterator, + SyncInteractionsAPIStreamingIterator, +) 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 @@ -148,3 +151,43 @@ async def test_stream_logs_its_first_event_time_measured_from_the_call_start(): assert first_event_at is not None assert recorder.logged == [(call_start, first_event_at)] + + +_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after" + + +async def _first_event( + iterator: InteractionsAPIStreamingIterator | SyncInteractionsAPIStreamingIterator, +) -> InteractionsAPIStreamingResponse: + if isinstance(iterator, InteractionsAPIStreamingIterator): + return await iterator.__anext__() + return next(iterator) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("iterator_cls", [InteractionsAPIStreamingIterator, SyncInteractionsAPIStreamingIterator]) +async def test_stream_keeps_unicode_line_separators_inside_the_delta_text( + iterator_cls: type[InteractionsAPIStreamingIterator] | type[SyncInteractionsAPIStreamingIterator], +) -> None: + delta: Final = {"type": "text", "text": _LINE_SEPARATED_TEXT} + event: Final = {"event_type": "step.delta", "id": "interaction-1", "delta": delta} + iterator: Final = iterator_cls( + response=httpx.Response( + 200, stream=httpx.ByteStream(f"data: {json.dumps(event, ensure_ascii=False)}\n\n".encode()) + ), + model="gemini/gemini-3.1-pro-preview", + interactions_api_config=GoogleAIStudioInteractionsConfig(), + logging_obj=LitellmLogging( + model="gemini/gemini-3.1-pro-preview", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="ainteraction", + start_time=datetime(2026, 10, 9, 12, 0, 0), + litellm_call_id="interactions-separators", + function_id="interactions-separators", + ), + ) + + first: Final = await _first_event(iterator) + + assert first.delta == delta diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index a1b58521d7e..79978ecfe8e 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -506,6 +506,118 @@ def test_multi_chunk_reasoning_and_content( assert initialized_custom_stream_wrapper.sent_last_thinking_block is True +def _stream_chunk(delta: dict[str, object], finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream(choices=[{"index": 0, "delta": delta, "finish_reason": finish_reason}]) + + +def _content_through_first_call(deltas: Sequence[Delta]) -> str: + first_call: Final = next( + (index for index, delta in enumerate(deltas) if delta.tool_calls or delta.function_call), len(deltas) + ) + return "".join(delta.content or "" for delta in deltas[: first_call + 1]) + + +_GET_WEATHER_CALL: Final = { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.parametrize( + "chunks,expected_content,expected_tool_calls", + [ + ( + ( + _stream_chunk({"reasoning_content": "Check the weather"}), + _stream_chunk({"tool_calls": [_GET_WEATHER_CALL]}), + _stream_chunk({}, "tool_calls"), + ), + "Check the weather", + 1, + ), + ( + ( + _stream_chunk({"reasoning_content": "Pick a function"}), + _stream_chunk({"function_call": {"name": "get_weather", "arguments": "{}"}}), + _stream_chunk({}, "function_call"), + ), + "Pick a function", + 0, + ), + ( + ( + _stream_chunk({"reasoning_content": "Step one"}), + _stream_chunk({"reasoning_content": ", step two"}), + _stream_chunk({}, "length"), + ), + "Step one, step two", + 0, + ), + ( + ( + _stream_chunk({"reasoning_content": "a"}), + _stream_chunk({"content": "b"}), + _stream_chunk({"reasoning_content": "c"}), + _stream_chunk({"content": "d"}), + _stream_chunk({}, "stop"), + ), + "abcd", + 0, + ), + ( + (_stream_chunk({"reasoning_content": "Still thinking"}),), + "Still thinking", + 0, + ), + ( + (_stream_chunk({"content": "hi"}), _stream_chunk({}, "stop")), + "hi", + 0, + ), + ], + ids=[ + "reasoning_then_tool_call", + "reasoning_then_function_call", + "reasoning_cut_by_length", + "reasoning_after_answer", + "stream_ends_mid_reasoning", + "answer_without_reasoning", + ], +) +async def test_merged_reasoning_stream_keeps_think_tags_balanced( + logging_obj: Logging, + is_async: bool, + chunks: tuple[ModelResponseStream, ...], + expected_content: str, + expected_tool_calls: int, +) -> None: + response: Final = CustomStreamWrapper( + completion_stream=ModelResponseListIterator(model_responses=[chunk.model_copy(deep=True) for chunk in chunks]), + model="deepseek-reasoner", + logging_obj=logging_obj, + ) + response.merge_reasoning_content_in_choices = True + + deltas: Final = ( + tuple([chunk.choices[0].delta async for chunk in response]) + if is_async + else tuple(chunk.choices[0].delta for chunk in response) + ) + logged: Final = litellm.stream_chunk_builder(chunks=response.chunks) + + assert "".join(delta.content or "" for delta in deltas) == expected_content + assert _content_through_first_call(deltas) == expected_content + assert sum(len(delta.tool_calls or ()) for delta in deltas) == expected_tool_calls + assert logged.choices[0].message.content == (expected_content or None) + assert any(chunk.choices[0].finish_reason for chunk in response.chunks if chunk.choices) == any( + chunk.choices[0].finish_reason for chunk in chunks + ) + + def test_strip_sse_data_from_chunk(): """Test the static method that strips 'data: ' prefix from SSE chunks""" # Test with string inputs @@ -7849,7 +7961,7 @@ def gemini_mock_post_streaming(url, **kwargs): mock_response = MagicMock() mock_response.status_code = 200 mock_response.headers = {"Content-Type": "text/event-stream"} - mock_response.iter_lines = MagicMock(return_value=stream_response()) + mock_response.iter_text = MagicMock(return_value=stream_response()) return mock_response diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py index 1f4a5aaa256..355d407001c 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -15,7 +15,7 @@ import litellm from litellm._uuid import uuid from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt -from litellm.llms.anthropic.chat.handler import ModelResponseIterator, _envelope_verdict, make_call +from litellm.llms.anthropic.chat.handler import ModelResponseIterator, _envelope_verdict, make_call, make_sync_call from litellm.llms.base_llm.base_utils import response_format_schema, unwrap_response_format_envelope from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.types.llms.openai import ( @@ -27,24 +27,34 @@ from litellm.types.utils import ChatCompletionDeltaToolCall, ModelResponseStream from tests.unit.llms.anthropic.anthropic_docs_refusal import ANTHROPIC_DOCS_REFUSAL_MESSAGE +def _text_delta_stream(text: str) -> bytes: + events: Final = ( + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}), + ) + return "".join( + f"event: {name}\ndata: {json.dumps(payload, ensure_ascii=False)}\n\n" for name, payload in events + ).encode() + + +def _streamed_text(chunks: list[object]) -> str: + return "".join(chunk.choices[0].delta.content or "" for chunk in chunks if isinstance(chunk, ModelResponseStream)) + + @pytest.mark.asyncio async def test_make_call_passes_logging_obj_to_client_post(): """make_call must pass logging_obj to client.post so track_llm_api_timing can set llm_api_duration_ms for litellm_overhead_time_ms.""" mock_client = AsyncMock() - mock_response = MagicMock() - mock_response.aiter_lines = MagicMock( - return_value=iter([b'data: {"type":"message_start"}\n', b'data: {"type":"message_delta"}\n']) - ) - mock_client.post.return_value = mock_response + mock_client.post.return_value = httpx.Response(200, stream=httpx.ByteStream(_text_delta_stream("Hi there"))) logging_obj = MagicMock() - await make_call( + completion_stream, _ = await make_call( client=mock_client, api_base="https://api.anthropic.com/v1/messages", headers={}, data="{}", - model="claude-3-5-haiku", + model="claude-haiku-5-5", messages=[{"role": "user", "content": "Hi"}], logging_obj=logging_obj, timeout=60.0, @@ -54,6 +64,39 @@ async def test_make_call_passes_logging_obj_to_client_post(): mock_client.post.assert_called_once() call_kwargs = mock_client.post.call_args[1] assert call_kwargs.get("logging_obj") is logging_obj + assert _streamed_text([chunk async for chunk in completion_stream]) == "Hi there" + + +_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after" + + +async def _streamed_chunks(sync_stream: bool, text: str) -> list[object]: + transport: Final = httpx.MockTransport( + lambda request: httpx.Response(200, stream=httpx.ByteStream(_text_delta_stream(text))) + ) + call_args: Final = { + "api_base": "https://api.anthropic.com/v1/messages", + "headers": {}, + "data": "{}", + "model": "claude-haiku-5-5", + "messages": [{"role": "user", "content": "Hi"}], + "logging_obj": MagicMock(), + "timeout": 60.0, + "json_mode": False, + } + if sync_stream: + sync_iterator, _ = make_sync_call(client=HTTPHandler(client=httpx.Client(transport=transport)), **call_args) + return list(sync_iterator) + async_iterator, _ = await make_call(client=AsyncHTTPHandler(transport=transport), **call_args) + return [chunk async for chunk in async_iterator] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_stream", [True, False]) +async def test_streamed_text_keeps_unicode_line_separators(sync_stream: bool) -> None: + chunks: Final = await _streamed_chunks(sync_stream, _LINE_SEPARATED_TEXT) + + assert _streamed_text(chunks) == _LINE_SEPARATED_TEXT def test_anthropic_completion_does_not_send_deployment_default_limits(): diff --git a/tests/unit/llms/anthropic/messages/test_stream_assembly.py b/tests/unit/llms/anthropic/messages/test_stream_assembly.py index f7d6b8d4f8c..4993d2fa9bf 100644 --- a/tests/unit/llms/anthropic/messages/test_stream_assembly.py +++ b/tests/unit/llms/anthropic/messages/test_stream_assembly.py @@ -144,6 +144,30 @@ def test_anthropic_message_delta_fields_reads_the_fields_the_openai_shape_drops( assert dict(anthropic_message_delta_fields(GOOGLE_STREAM)) == {} +def test_anthropic_message_delta_fields_reads_a_message_delta_whose_json_holds_unicode_line_separators() -> None: + safeguard_results: Final = [ + { + "type": "dangerous_tool_use", + "status": { + "type": "available", + "tool_uses": { + "toolu_01": {"type": "evaluated", "outcome": "flagged", "explanation": "deletes\u2028the\u2029root"} + }, + }, + } + ] + message_delta: Final = { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "safeguard_results": safeguard_results}, + "usage": {"output_tokens": 2}, + } + raw_message_delta_frame: Final = f"event: message_delta\ndata: {json.dumps(message_delta, ensure_ascii=False)}\n\n" + + fields: Final = anthropic_message_delta_fields((*ANTHROPIC_STREAM[:5], raw_message_delta_frame.encode())) + + assert dict(fields) == {"safeguard_results": safeguard_results} + + @pytest.mark.parametrize( "chunks", [GOOGLE_STREAM, INVALID_UTF8, (ANTHROPIC_ERROR,), ()], diff --git a/tests/unit/llms/azure_ai/agents/test_handler.py b/tests/unit/llms/azure_ai/agents/test_handler.py index 71466bb172c..14b6094747f 100644 --- a/tests/unit/llms/azure_ai/agents/test_handler.py +++ b/tests/unit/llms/azure_ai/agents/test_handler.py @@ -173,7 +173,7 @@ async def test_handler_acompletion_with_thread_id_reuses_that_thread() -> None: def _sse(event: str, data: dict[str, object] | str) -> str: - return f"event: {event}\ndata: {data if isinstance(data, str) else json.dumps(data)}\n\n" + return f"event: {event}\ndata: {data if isinstance(data, str) else json.dumps(data, ensure_ascii=False)}\n\n" def _message_delta(text: str) -> dict[str, object]: @@ -222,3 +222,24 @@ async def test_acompletion_stream_yields_agent_message_deltas_then_stop(respx_mo "thread": {"messages": [{"role": "user", "content": "Hi Agent, what is 10 + 5?"}]}, } assert chunks == [("10 + 5", None), (" = 15", None), (None, "stop")] + + +@pytest.mark.asyncio +async def test_acompletion_stream_keeps_unicode_line_separators_inside_a_delta(respx_mock: respx.MockRouter) -> None: + separated_text: Final = "before\u2028between\u2029after" + _ = respx_mock.post(f"{API_BASE}/threads/runs").respond( + 200, + text=_sse("thread.message.delta", _message_delta(separated_text)) + _sse("done", "[DONE]"), + headers={"content-type": "text/event-stream"}, + ) + + response: Final = await litellm.acompletion( + model=f"azure_ai/agents/{AGENT_ID}", + messages=[{"role": "user", "content": "Say it with separators"}], + api_base=API_BASE, + api_key="azure-ad-token", + stream=True, + ) + chunks: Final = [(chunk.choices[0].delta.content, chunk.choices[0].finish_reason) async for chunk in response] + + assert chunks == [(separated_text, None), (None, "stop")] diff --git a/tests/unit/llms/azure_ai/agents/test_transformation.py b/tests/unit/llms/azure_ai/agents/test_transformation.py index 10676438858..413908c8313 100644 --- a/tests/unit/llms/azure_ai/agents/test_transformation.py +++ b/tests/unit/llms/azure_ai/agents/test_transformation.py @@ -1,7 +1,7 @@ import json from typing import Final -from unittest.mock import MagicMock +import httpx import pytest from litellm.llms.azure_ai.agents.transformation import AzureAIAgentsConfig @@ -428,15 +428,10 @@ async def test_azure_ai_agents_streaming_accumulates_annotations_from_multiple_t "data: [DONE]", ] - async def mock_aiter_lines(): - for line in sse_lines: - yield line - - mock_response = MagicMock() - mock_response.aiter_lines = MagicMock(return_value=mock_aiter_lines()) + sse_response = httpx.Response(200, stream=httpx.ByteStream("\n".join(sse_lines).encode())) chunks = [] - async for chunk in handler._process_sse_stream(mock_response, "azure_ai/agents/asst_123"): + async for chunk in handler._process_sse_stream(sse_response, "azure_ai/agents/asst_123"): chunks.append(chunk) final_chunk = chunks[-1] @@ -498,15 +493,10 @@ async def test_azure_ai_agents_streaming_annotations_from_completed_message(): "data: [DONE]", ] - async def mock_aiter_lines(): - for line in sse_lines: - yield line - - mock_response = MagicMock() - mock_response.aiter_lines = MagicMock(return_value=mock_aiter_lines()) + sse_response = httpx.Response(200, stream=httpx.ByteStream("\n".join(sse_lines).encode())) chunks = [] - async for chunk in handler._process_sse_stream(mock_response, "azure_ai/agents/asst_123"): + async for chunk in handler._process_sse_stream(sse_response, "azure_ai/agents/asst_123"): chunks.append(chunk) assert len(chunks) >= 1 diff --git a/tests/unit/llms/base_llm/guardrail_translation/test_utils.py b/tests/unit/llms/base_llm/guardrail_translation/test_utils.py index 65208a2bd3b..acd50616f90 100644 --- a/tests/unit/llms/base_llm/guardrail_translation/test_utils.py +++ b/tests/unit/llms/base_llm/guardrail_translation/test_utils.py @@ -1,9 +1,11 @@ +import json from typing import Final import pytest import litellm from litellm.llms.base_llm.guardrail_translation.utils import ( + blocked_response_usage, effective_skip_system_message_for_guardrail, effective_skip_tool_message_for_guardrail, ) @@ -32,3 +34,18 @@ def test_a_per_guardrail_skip_flag_wins_over_the_global_setting( assert effective_skip_system_message_for_guardrail(guardrail) is expected assert effective_skip_tool_message_for_guardrail(guardrail) is expected + + +def test_blocked_response_usage_reads_a_message_delta_whose_json_holds_unicode_line_separators() -> None: + message_start: Final = {"type": "message_start", "message": {"usage": {"input_tokens": 11, "output_tokens": 1}}} + message_delta: Final = { + "type": "message_delta", + "delta": {"stop_reason": "stop_sequence", "stop_sequence": "END\u2028OF\u2029TURN"}, + "usage": {"output_tokens": 7}, + } + stream: Final = "".join( + f"event: {event['type']}\ndata: {json.dumps(event, ensure_ascii=False)}\n\n" + for event in (message_start, message_delta) + ).encode() + + assert blocked_response_usage([stream]) == {"input_tokens": 11, "output_tokens": 7} diff --git a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py index 7e05ef500f5..9d9e39df944 100644 --- a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py @@ -4458,6 +4458,20 @@ class TestMalformedBatchJsonl: _VALID_BATCH_RECORD: Final = json.loads(_VALID_BATCH_LINE) +def test_file_content_conversion_keeps_a_line_whose_json_holds_unicode_line_separators() -> None: + text: Final = "Hello\u2028from\u2029batch" + body: Final = {**_VALID_BATCH_RECORD["body"], "messages": [{"role": "user", "content": text}]} + record: Final = {**_VALID_BATCH_RECORD, "body": body} + file_content: Final = f"{json.dumps(record, ensure_ascii=False)}\n".encode() + + bedrock_jsonl, _ = BedrockJsonlFilesTransformation().transform_openai_file_content_to_bedrock_file_content( + file_content + ) + + (bedrock_record,) = (json.loads(line) for line in bedrock_jsonl.split("\n")) + assert bedrock_record["modelInput"]["messages"] == [{"role": "user", "content": [{"text": text}]}] + + def _batch_line(body: object) -> str: return json.dumps({**_VALID_BATCH_RECORD, "body": body}) diff --git a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py index 5aaee1f0ee4..aec1672b335 100644 --- a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py +++ b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py @@ -18,7 +18,7 @@ import types from collections.abc import Callable, Mapping, Sequence from datetime import datetime from itertools import chain -from typing import ClassVar, Literal, TypeAlias +from typing import ClassVar, Final, Literal, TypeAlias from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -60,6 +60,7 @@ from litellm.llms.bedrock.passthrough.guardrail_translation.handler import ( BedrockPassthroughGuardrailHandler, _extract_converse_texts, _is_converse_endpoint, + _sse_events, _write_back_texts, ) @@ -3840,3 +3841,14 @@ class TestInvokeStreamGuardrailRelay: assert received_key is not None and received_key.texts == ("In Paris",) assert rewritten_key == received_key + + +def test_sse_events_keeps_a_data_line_whose_json_holds_unicode_line_separators() -> None: + event: Final = { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Card\u2028on\u2029file"}, + } + sse: Final = f"event: content_block_delta\ndata: {json.dumps(event, ensure_ascii=False)}\n\n".encode() + + assert _sse_events(sse) == (event,) diff --git a/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py b/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py index 664dbf87fe9..01cca1fc270 100644 --- a/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py +++ b/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py @@ -8,6 +8,7 @@ the sharded CI (coverage collection runs against this tree). import json import os +from typing import Final import httpx import pytest @@ -682,3 +683,17 @@ def test_transform_search_response_ignores_text_blocks_without_result_objects(bl body = {"result": {"content": [{"type": "text", "text": block_text}]}} assert _transform_text(json.dumps(body)).results == [] + + +def test_transform_search_response_keeps_a_result_whose_sse_json_holds_unicode_line_separators() -> None: + snippet: Final = "Snippet\u2028for\u2029result" + results: Final = [{"title": "Result", "url": "https://example.com/1", "text": snippet}] + body: Final = { + "jsonrpc": "2.0", + "id": 1, + "result": {"content": [{"type": "text", "text": json.dumps(results, ensure_ascii=False)}]}, + } + + response: Final = _transform_text(f"event: message\ndata: {json.dumps(body, ensure_ascii=False)}\n\n") + + assert _as_tuples(response) == [("Result", "https://example.com/1", snippet, None, None)] diff --git a/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py index a6f7d613f28..635363e77ea 100644 --- a/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py +++ b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py @@ -587,6 +587,37 @@ class TestChatGPTResponsesMultiAccount: assert headers_default["ChatGPT-Account-Id"] == "acct-default" +def test_chatgpt_non_stream_sse_keeps_an_output_item_whose_text_holds_unicode_line_separators() -> None: + text: Final = "Hello\u2028from\u2029stream" + output_item_done: Final = { + "type": "response.output_item.done", + "output_index": 0, + "item": {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": text}]}, + } + completed: Final = { + "type": "response.completed", + "response": { + "id": "resp_test", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-6.1-sol", + "output": [], + }, + } + sse_body: Final = "".join( + f"data: {json.dumps(event, ensure_ascii=False)}\n" for event in (output_item_done, completed) + ) + + parsed: Final = ChatGPTResponsesAPIConfig().transform_response_api_response( + model="chatgpt/gpt-6.1-sol", + raw_response=httpx.Response(200, headers={"content-type": "text/event-stream"}, text=sse_body), + logging_obj=MagicMock(), + ) + + assert parsed.output_text == text + + @pytest.mark.parametrize("call", ["responses", "completion"]) def test_prompt_cache_key_keeps_one_session_header_across_requests( call: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch diff --git a/tests/unit/llms/cloudflare/test_cloudflare_transformation.py b/tests/unit/llms/cloudflare/test_cloudflare_transformation.py index 3a45c83556c..ccbfc02f0e6 100644 --- a/tests/unit/llms/cloudflare/test_cloudflare_transformation.py +++ b/tests/unit/llms/cloudflare/test_cloudflare_transformation.py @@ -1,6 +1,6 @@ import asyncio import json -from collections.abc import Callable, Iterator +from collections.abc import AsyncIterator, Callable, Iterator from typing import Final from unittest.mock import MagicMock @@ -380,6 +380,35 @@ def test_completion_cloudflare_stream(sync_mode: bool, _cloudflare_httpx_transpo assert "language" in content.lower() +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"]) +async def test_cloudflare_stream_keeps_unicode_line_separators_inside_content( + sync_mode: bool, _cloudflare_httpx_transport: None +) -> None: + streamed_text: Final = ("alpha\u2028", "beta\u2029", "gamma") + base: Final = {"id": "chatcmpl-cf", "object": "chat.completion.chunk", "created": 1, "model": "@cf/glm"} + body: Final = "".join( + f"data: {json.dumps({**base, 'choices': [{'index': 0, 'delta': {'content': text}}]}, ensure_ascii=False)}\n\n" + for text in streamed_text + ) + with respx.mock(assert_all_called=True) as api: + api.post(f"{FAKE_API_BASE}/chat/completions").respond( + 200, headers={"content-type": "text/event-stream"}, content=(body + "data: [DONE]\n\n").encode() + ) + request: Final = { + "model": "cloudflare/@cf/glm", + "messages": [{"role": "user", "content": "repeat"}], + "stream": True, + "api_base": FAKE_API_BASE, + "api_key": FAKE_API_KEY, + } + chunks: Final = ( + tuple(completion(**request)) if sync_mode else tuple([c async for c in await acompletion(**request)]) + ) + + assert "".join(c.choices[0].delta.content or "" for c in chunks) == "".join(streamed_text) + + def _make_mock_response(json_data: Dict[str, Any]) -> MagicMock: mock = MagicMock(spec=httpx.Response) mock.status_code = 200 @@ -397,13 +426,13 @@ def test_acompletion_cloudflare_stream(sync_mode, monkeypatch): if sync_mode: - def _iter_lines(): + def _iter_text() -> Iterator[str]: for chunk in raw_chunks: - yield f"data: {chunk}" - yield "data: [DONE]" + yield f"data: {chunk}\n\n" + yield "data: [DONE]\n\n" mock_resp = MagicMock() - mock_resp.iter_lines.return_value = _iter_lines() + mock_resp.iter_text.return_value = _iter_text() mock_resp.status_code = 200 mock_resp.headers = {"content-type": "text/event-stream"} @@ -420,13 +449,13 @@ def test_acompletion_cloudflare_stream(sync_mode, monkeypatch): mock_post.assert_called_once() else: - async def _aiter_lines(): + async def _aiter_text() -> AsyncIterator[str]: for chunk in raw_chunks: - yield f"data: {chunk}" - yield "data: [DONE]" + yield f"data: {chunk}\n\n" + yield "data: [DONE]\n\n" mock_resp = MagicMock() - mock_resp.aiter_lines.return_value = _aiter_lines() + mock_resp.aiter_text.return_value = _aiter_text() mock_resp.status_code = 200 mock_resp.headers = {"content-type": "text/event-stream"} diff --git a/tests/unit/llms/codestral/__init__.py b/tests/unit/llms/codestral/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/codestral/completion/__init__.py b/tests/unit/llms/codestral/completion/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/codestral/completion/test_codestral_completion_handler.py b/tests/unit/llms/codestral/completion/test_codestral_completion_handler.py new file mode 100644 index 00000000000..9ce50d8ceca --- /dev/null +++ b/tests/unit/llms/codestral/completion/test_codestral_completion_handler.py @@ -0,0 +1,50 @@ +import json +from typing import Final + +import httpx +import pytest +import respx + +import litellm + +_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after" + + +def _fim_stream_body(text: str) -> bytes: + chunks: Final = ( + {"id": "c1", "object": "chat.completion.chunk", "created": 1, "model": "codestral-latest", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": text}, "finish_reason": None}]}, + {"id": "c1", "object": "chat.completion.chunk", "created": 1, "model": "codestral-latest", + "choices": [{"index": 0, "delta": {"content": ""}, "finish_reason": "stop"}]}, + ) + events: Final = "".join(f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n" for chunk in chunks) + return f"{events}data: [DONE]\n\n".encode() + + +async def _streamed_texts(sync_stream: bool) -> list[str | None]: + request: Final = { + "model": "text-completion-codestral/codestral-latest", + "prompt": "def f(", + "stream": True, + "api_key": "codestral-key", + "max_retries": 0, + } + if sync_stream: + return [chunk.choices[0].text for chunk in litellm.text_completion(**request)] + return [chunk.choices[0].text async for chunk in await litellm.atext_completion(**request)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_stream", [True, False]) +async def test_fim_stream_keeps_unicode_line_separators_inside_a_delta( + sync_stream: bool, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post("https://codestral.mistral.ai/v1/fim/completions").mock( + return_value=httpx.Response(200, content=_fim_stream_body(_LINE_SEPARATED_TEXT)) + ) + + texts: Final = await _streamed_texts(sync_stream) + + assert texts == [_LINE_SEPARATED_TEXT, None] + assert json.loads(route.calls.last.request.content)["stream"] is True diff --git a/tests/unit/llms/custom_httpx/test_aiohttp_handler.py b/tests/unit/llms/custom_httpx/test_aiohttp_handler.py index 83c38b65beb..bb66a70da84 100644 --- a/tests/unit/llms/custom_httpx/test_aiohttp_handler.py +++ b/tests/unit/llms/custom_httpx/test_aiohttp_handler.py @@ -1,11 +1,17 @@ +import json +from typing import Final from unittest.mock import AsyncMock, Mock, patch import aiohttp +import httpx import pytest import litellm +from litellm.llms.aiohttp_openai.chat.transformation import AiohttpOpenAIChatConfig from litellm.llms.custom_httpx.aiohttp_handler import BaseLLMAIOHTTPHandler from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.types.utils import ModelResponseStream class TestBaseLLMAIOHTTPHandler: @@ -445,3 +451,34 @@ class TestBaseLLMAIOHTTPHandler: # Should use transport, not connector mock_transport.get_valid_client_session.assert_called_once() assert result is mock_session_from_transport + + +def test_make_sync_call_stream_keeps_unicode_line_separators_inside_a_delta() -> None: + text: Final = "before\u2028between\u2029after" + chunk: Final = { + "id": "c1", + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-5.5", + "choices": [{"index": 0, "delta": {"content": text}, "finish_reason": None}], + } + body: Final = f"data: {json.dumps(chunk, ensure_ascii=False)}\n\ndata: [DONE]\n\n".encode() + transport: Final = httpx.MockTransport(lambda request: httpx.Response(200, stream=httpx.ByteStream(body))) + + completion_stream, _ = BaseLLMAIOHTTPHandler().make_sync_call( + provider_config=AiohttpOpenAIChatConfig(), + api_base="https://example.com/v1/chat/completions", + headers={}, + data={"model": "gpt-5.5", "messages": []}, + model="gpt-5.5", + messages=[], + logging_obj=Mock(), + litellm_params={}, + timeout=5.0, + client=HTTPHandler(client=httpx.Client(transport=transport)), + ) + + contents: Final = [ + chunk.choices[0].delta.content for chunk in completion_stream if isinstance(chunk, ModelResponseStream) + ] + assert contents == [text] diff --git a/tests/unit/llms/custom_httpx/test_sse_lines.py b/tests/unit/llms/custom_httpx/test_sse_lines.py new file mode 100644 index 00000000000..9bb03f6da98 --- /dev/null +++ b/tests/unit/llms/custom_httpx/test_sse_lines.py @@ -0,0 +1,64 @@ +from collections.abc import AsyncIterator, Sequence +from typing import Final + +import pytest + +from litellm.llms.custom_httpx.sse_lines import ( + aiter_sse_lines, + iter_sse_lines, + split_sse_lines, + split_sse_lines_keepends, +) + +CASES: Final = ( + pytest.param(("data: a\u2028b\r\n\r\n",), ("data: a\u2028b", ""), id="u2028-stays-inside-the-line"), + pytest.param( + ("data: \u0085\u2029\x0b\x0c\x1c\x1d\x1e\n",), + ("data: \u0085\u2029\x0b\x0c\x1c\x1d\x1e",), + id="no-splitlines-only-separator-breaks-a-line", + ), + pytest.param(("a\r\nb\nc\rd\n",), ("a", "b", "c", "d"), id="crlf-lf-and-cr-each-end-a-line"), + pytest.param(("a\r", "\nb\n"), ("a", "b"), id="crlf-split-across-chunks-is-one-break"), + pytest.param(("a\r", "\n\r\n"), ("a", ""), id="crlf-split-then-blank-line"), + pytest.param(("a\r", "", "\nb\n"), ("a", "b"), id="empty-chunk-between-cr-and-lf"), + pytest.param(("a\r", "\r"), ("a", ""), id="cr-then-cr-is-two-breaks"), + pytest.param(("a\r\r\nb\n",), ("a", "", "b"), id="cr-then-crlf-in-one-chunk-is-two-breaks"), + pytest.param(("a\r\nb\r\n\r\n",), ("a", "b", ""), id="crlf-only-chunk"), + pytest.param(("da", "ta: x", "y\n"), ("data: xy",), id="line-spanning-chunks-is-joined"), + pytest.param(("data: x",), ("data: x",), id="unterminated-last-line-is-kept"), + pytest.param(("",), (), id="empty-body"), +) + + +async def _async_chunks(chunks: Sequence[str]) -> AsyncIterator[str]: + for chunk in chunks: + yield chunk + + +@pytest.mark.parametrize(("chunks", "expected"), CASES) +def test_iter_sse_lines_splits_only_on_cr_lf_and_crlf(chunks: Sequence[str], expected: Sequence[str]) -> None: + assert tuple(iter_sse_lines(chunks)) == tuple(expected) + + +@pytest.mark.parametrize(("chunks", "expected"), CASES) +async def test_aiter_sse_lines_splits_only_on_cr_lf_and_crlf(chunks: Sequence[str], expected: Sequence[str]) -> None: + assert tuple([line async for line in aiter_sse_lines(_async_chunks(chunks))]) == tuple(expected) + + +@pytest.mark.parametrize(("chunks", "expected"), CASES) +def test_split_sse_lines_splits_a_whole_body_only_on_cr_lf_and_crlf( + chunks: Sequence[str], expected: Sequence[str] +) -> None: + assert split_sse_lines("".join(chunks)) == tuple(expected) + + +@pytest.mark.parametrize( + ("text", "expected"), + ( + pytest.param("a\r\nb\rc\nd", ("a\r\n", "b\r", "c\n", "d"), id="each-break-stays-on-its-line"), + pytest.param("data: a\u2028b\u2029c\n\n", ("data: a\u2028b\u2029c\n", "\n"), id="u2028-and-u2029-stay-inside"), + pytest.param("", (), id="empty-body"), + ), +) +def test_split_sse_lines_keepends_splits_only_on_cr_lf_and_crlf(text: str, expected: Sequence[str]) -> None: + assert split_sse_lines_keepends(text) == tuple(expected) diff --git a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py index b3a224209e9..eb2e1d8e347 100644 --- a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py @@ -1,6 +1,6 @@ import copy import json -from collections.abc import Iterator +from collections.abc import AsyncIterator, Iterator from typing import Final from unittest.mock import MagicMock, patch @@ -2160,13 +2160,13 @@ def mock_chat_response_nonanthropic_prompt_caching() -> Dict[str, Any]: def mock_http_handler_chat_streaming_response() -> MagicMock: mock_stream_chunks = mock_chat_streaming_response_chunks() - def mock_iter_lines(): + def mock_iter_text() -> Iterator[str]: for chunk in mock_stream_chunks: for line in chunk.splitlines(): - yield line + yield f"{line}\n" mock_response = MagicMock() - mock_response.iter_lines.side_effect = mock_iter_lines + mock_response.iter_text.side_effect = mock_iter_text mock_response.status_code = 200 return mock_response @@ -2175,13 +2175,13 @@ def mock_http_handler_chat_streaming_response() -> MagicMock: def mock_http_handler_chat_async_streaming_response() -> MagicMock: mock_stream_chunks = mock_chat_streaming_response_chunks() - async def mock_iter_lines(): + async def mock_aiter_text() -> AsyncIterator[str]: for chunk in mock_stream_chunks: for line in chunk.splitlines(): - yield line + yield f"{line}\n" mock_response = MagicMock() - mock_response.aiter_lines.return_value = mock_iter_lines() + mock_response.aiter_text.return_value = mock_aiter_text() mock_response.status_code = 200 return mock_response diff --git a/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 04b791cf289..61e5545b8bf 100644 --- a/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -1329,7 +1329,7 @@ def test_streaming_surfaces_fireworks_response_fields(): raw_response = MagicMock() raw_response.status_code = 200 raw_response.headers = {} - raw_response.iter_lines = lambda: iter(sse_lines) + raw_response.iter_text = lambda: iter(f"{line}\n" for line in sse_lines) client = HTTPHandler() with patch.object(client, "post", return_value=raw_response): @@ -1788,7 +1788,7 @@ def test_streaming_preserves_selected_model_for_private_accounting(): raw_response = MagicMock() raw_response.status_code = 200 raw_response.headers = {} - raw_response.iter_lines = lambda: iter(sse_lines) + raw_response.iter_text = lambda: iter(f"{line}\n" for line in sse_lines) client = HTTPHandler() with patch.object(client, "post", return_value=raw_response): @@ -2275,7 +2275,7 @@ def test_completion_forwards_litellm_user_id_as_user(call_kwargs: dict[str, obje def test_completion_forwards_litellm_user_id_when_streaming() -> None: client: Final = _fireworks_chat_client() - client.post.return_value.iter_lines = lambda: iter(()) + client.post.return_value.iter_text = lambda: iter(()) litellm.completion( model="fireworks_ai/accounts/fireworks/models/kimi-k3", messages=[{"role": "user", "content": "hi"}], diff --git a/tests/unit/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py b/tests/unit/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py index 4119ce99423..53ddb1135b4 100644 --- a/tests/unit/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py +++ b/tests/unit/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py @@ -4,6 +4,7 @@ Tests for the Google GenAI generateContent guardrail translation handler. import json from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -190,6 +191,21 @@ async def test_streaming_raw_sse_chunks_accumulate_text_across_split_frames(): assert guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] == ["harmful answer"] +@pytest.mark.asyncio +async def test_streaming_raw_sse_text_holding_unicode_line_separators_reaches_the_guardrail_whole() -> None: + guardrail: Final = _mock_guardrail(["clean"]) + text: Final = "harmful\u2028answer\u2029here" + payload: Final = {"candidates": [{"content": {"parts": [{"text": text}]}}]} + chunk: Final = f"data: {json.dumps(payload, ensure_ascii=False)}\r\n\r\n" + + await GoogleGenAIGenerateContentHandler().process_output_streaming_response( + responses_so_far=[chunk], + guardrail_to_apply=guardrail, + ) + + assert guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] == [text] + + @pytest.mark.asyncio async def test_streaming_blocking_guardrail_exception_propagates(): handler = GoogleGenAIGenerateContentHandler() diff --git a/tests/unit/llms/langgraph/chat/test_langgraph_sse_iterator.py b/tests/unit/llms/langgraph/chat/test_langgraph_sse_iterator.py new file mode 100644 index 00000000000..06485362162 --- /dev/null +++ b/tests/unit/llms/langgraph/chat/test_langgraph_sse_iterator.py @@ -0,0 +1,32 @@ +import json +from typing import Final + +import httpx +import pytest + +from litellm.llms.langgraph.chat.sse_iterator import LangGraphSSEStreamIterator + +_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after" + + +def _messages_event_stream(text: str) -> httpx.Response: + event: Final = ["messages", [[{"type": "AIMessageChunk", "content": text}, {}]]] + return httpx.Response(200, stream=httpx.ByteStream(f"data: {json.dumps(event, ensure_ascii=False)}\n\n".encode())) + + +async def _streamed_contents(iterator: LangGraphSSEStreamIterator, sync_stream: bool) -> list[str | None]: + if sync_stream: + return [chunk.choices[0].delta.content for chunk in iterator] + return [chunk.choices[0].delta.content async for chunk in iterator] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_stream", [True, False]) +async def test_stream_keeps_unicode_line_separators_inside_a_message_chunk(sync_stream: bool) -> None: + iterator: Final = LangGraphSSEStreamIterator( + response=_messages_event_stream(_LINE_SEPARATED_TEXT), model="langgraph/agent" + ) + + contents: Final = await _streamed_contents(iterator, sync_stream) + + assert contents == [_LINE_SEPARATED_TEXT, None] diff --git a/tests/unit/llms/oobabooga/chat/test_oobabooga.py b/tests/unit/llms/oobabooga/chat/test_oobabooga.py index 395a4fb5715..332b6913890 100644 --- a/tests/unit/llms/oobabooga/chat/test_oobabooga.py +++ b/tests/unit/llms/oobabooga/chat/test_oobabooga.py @@ -1,7 +1,12 @@ +import json +from typing import Final from unittest.mock import MagicMock, patch +import httpx +import respx import litellm +from litellm.llms.oobabooga.chat import oobabooga MOCK_COMPLETION_RESPONSE = { "choices": [{"message": {"role": "assistant", "content": "hi there"}}], @@ -50,3 +55,27 @@ def test_url_valued_model_still_targets_that_url(): mock_post.assert_called_once() called_url = mock_post.call_args[0][0] assert called_url == "https://sdk-user.example/v1/chat/completions" + + +def test_stream_keeps_unicode_line_separators_inside_a_data_line(respx_mock: respx.MockRouter) -> None: + text: Final = "before\u2028between\u2029after" + chunk: Final = {"choices": [{"index": 0, "delta": {"content": text}, "finish_reason": None}]} + data_line: Final = f"data: {json.dumps(chunk, ensure_ascii=False)}" + respx_mock.post("https://ooba.example/v1/chat/completions").mock( + return_value=httpx.Response(200, content=f"{data_line}\n\ndata: [DONE]\n\n".encode()) + ) + + lines: Final = oobabooga.completion( + model="my-model", + messages=[{"role": "user", "content": "hello"}], + api_base="https://ooba.example", + model_response=litellm.ModelResponse(), + print_verbose=print, + encoding=None, + api_key=None, + logging_obj=MagicMock(), + optional_params={"stream": True}, + litellm_params={}, + ) + + assert [line for line in lines if line] == [data_line, "data: [DONE]"] diff --git a/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py b/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py index 75b60b2dd6c..3df1efe311b 100644 --- a/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py +++ b/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py @@ -4,9 +4,10 @@ from typing import Final from unittest.mock import Mock import httpx +import pytest -from litellm.llms.custom_httpx.http_handler import HTTPHandler -from litellm.llms.openai_like.chat.handler import make_sync_call +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.openai_like.chat.handler import make_call, make_sync_call _FAKE_STREAMED_REPLY: Final = { "id": "chatcmpl-fake-stream", @@ -73,3 +74,35 @@ def test_fake_stream_keeps_reasoning_every_tool_call_and_token_detail(): assert [call.function.name for call in delta.tool_calls] == ["get_weather", "get_time"] assert chunks[0].usage.prompt_tokens_details.cached_tokens == 64 assert chunks[0].usage.completion_tokens_details.reasoning_tokens == 1920 + + +_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after" + + +def _streamed_reply(request: httpx.Request) -> httpx.Response: + chunk: Final = {"id": "chatcmpl-separators", "choices": [{"index": 0, "delta": {"content": _LINE_SEPARATED_TEXT}}]} + body: Final = f"data: {json.dumps(chunk, ensure_ascii=False)}\n\ndata: [DONE]\n\n" + return httpx.Response(200, stream=httpx.ByteStream(body.encode())) + + +async def _streamed_chunks(sync_stream: bool) -> list[dict[str, object]]: + call_args: Final = { + "api_base": "https://openai-like.example/v1/chat/completions", + "headers": {}, + "data": json.dumps({"model": "openai/gpt-oss-120b", "messages": [], "stream": True}), + "model": "openai/gpt-oss-120b", + "messages": [], + "logging_obj": Mock(), + } + transport: Final = httpx.MockTransport(_streamed_reply) + if sync_stream: + return list(make_sync_call(client=HTTPHandler(client=httpx.Client(transport=transport)), **call_args)) + return [chunk async for chunk in await make_call(client=AsyncHTTPHandler(transport=transport), **call_args)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_stream", [True, False]) +async def test_stream_keeps_unicode_line_separators_inside_the_delta(sync_stream: bool) -> None: + chunks: Final = await _streamed_chunks(sync_stream) + + assert "".join(str(chunk["text"]) for chunk in chunks) == _LINE_SEPARATED_TEXT diff --git a/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py b/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py index f9d12b7eeea..36c319d9524 100644 --- a/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py +++ b/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py @@ -902,20 +902,14 @@ class TestSnowflakeChatCompletion: @pytest.mark.parametrize("sync_mode", [True, False]) def test_chat_completion_snowflake_stream(self, sync_mode): - raw_chunks = _streaming_chunks() + sse_body = "".join(f"data: {chunk}\n\n" for chunk in _streaming_chunks()) + "data: [DONE]\n\n" + mock_resp = httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + stream=httpx.ByteStream(sse_body.encode()), + ) if sync_mode: - - def _iter_lines(): - for chunk in raw_chunks: - yield f"data: {chunk}" - yield "data: [DONE]" - - mock_resp = MagicMock() - mock_resp.iter_lines.return_value = _iter_lines() - mock_resp.status_code = 200 - mock_resp.headers = {"content-type": "text/event-stream"} - with patch.object(HTTPHandler, "post", return_value=mock_resp) as mock_post: response = completion( model="snowflake/mistral-7b", @@ -930,16 +924,6 @@ class TestSnowflakeChatCompletion: mock_post.assert_called_once() else: - async def _aiter_lines(): - for chunk in raw_chunks: - yield f"data: {chunk}" - yield "data: [DONE]" - - mock_resp = MagicMock() - mock_resp.aiter_lines.return_value = _aiter_lines() - mock_resp.status_code = 200 - mock_resp.headers = {"content-type": "text/event-stream"} - async def _run(): with patch.object( AsyncHTTPHandler, @@ -964,8 +948,8 @@ class TestSnowflakeChatCompletion: chunks_received = asyncio.run(_run()) - assert len(chunks_received) > 0 content = "".join(c.choices[0].delta.content for c in chunks_received if c.choices[0].delta.content) + assert content == "The sky is blue" async def test_snowflake_claude_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch): diff --git a/tests/unit/llms/test_predibase_transformation.py b/tests/unit/llms/test_predibase_transformation.py index a02840ed4ef..7afbd05b221 100644 --- a/tests/unit/llms/test_predibase_transformation.py +++ b/tests/unit/llms/test_predibase_transformation.py @@ -1,3 +1,4 @@ +import json from typing import Final from unittest.mock import AsyncMock, Mock @@ -639,3 +640,49 @@ def test_predibase_completion_passes_existing_config_to_async_completion(monkeyp assert result == "async-result" assert captured["async_kwargs"]["predibase_config"] is captured["config_instance"] + + +_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after" + + +def _generate_stream_body(text: str) -> bytes: + events: Final = ( + {"token": {"text": text}}, + {"token": {"text": ""}, "generated_text": text, "details": {"finish_reason": "eos_token"}}, + ) + return "".join(f"data:{json.dumps(event, ensure_ascii=False)}\n\n" for event in events).encode() + + +async def _streamed_contents(sync_stream: bool) -> list[str | None]: + request: Final = { + "model": "predibase/llama-3-8b-instruct", + "messages": [{"role": "user", "content": "hi"}], + "stream": True, + "predibase_tenant_id": "tenant-1", + "api_key": "pb-key", + "max_retries": 0, + } + if sync_stream: + return [chunk.choices[0].delta.content for chunk in litellm.completion(**request)] + return [chunk.choices[0].delta.content async for chunk in await litellm.acompletion(**request)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_stream", [True, False]) +async def test_predibase_stream_keeps_unicode_line_separators_inside_a_token( + sync_stream: bool, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_hf_tokenizer_download", True) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr( + litellm, + "custom_prompt_dict", + {"llama-3-8b-instruct": {"roles": {}, "initial_prompt_value": "", "final_prompt_value": ""}}, + ) + respx_mock.post( + "https://serving.app.predibase.com/tenant-1/deployments/v2/llms/llama-3-8b-instruct/generate_stream" + ).mock(return_value=httpx.Response(200, content=_generate_stream_body(_LINE_SEPARATED_TEXT))) + + contents: Final = await _streamed_contents(sync_stream) + + assert contents == [_LINE_SEPARATED_TEXT, None] diff --git a/tests/unit/llms/triton/test_triton.py b/tests/unit/llms/triton/test_triton.py index 7c0de603a3e..238c30816b2 100644 --- a/tests/unit/llms/triton/test_triton.py +++ b/tests/unit/llms/triton/test_triton.py @@ -1,5 +1,7 @@ import json import traceback +from collections.abc import Iterator +from typing import Final from unittest.mock import MagicMock, patch import httpx @@ -210,8 +212,8 @@ def test_completion_triton_generate_api(stream, monkeypatch, respx_mock): mock_response = MagicMock() if stream: - def mock_iter_lines(): - mock_output = "".join( + def mock_iter_text() -> Iterator[str]: + mock_output: Final = "".join( [ 'data: {"model_name":"ensemble","model_version":"1","sequence_end":false,"sequence_id":0,"sequence_start":false,"text_output":"' + t @@ -219,10 +221,9 @@ def test_completion_triton_generate_api(stream, monkeypatch, respx_mock): for t in ["I", " am", " an", " AI", " assistant"] ] ) - for out in mock_output.split("\n"): - yield out + yield mock_output - mock_response.iter_lines = mock_iter_lines + mock_response.iter_text = mock_iter_text else: def return_val(): diff --git a/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py b/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py index 06aaca44e8a..20333fd0465 100644 --- a/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py +++ b/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py @@ -14,7 +14,7 @@ import pytest from litellm.litellm_core_utils.litellm_logging import Logging -from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler, get_async_httpx_client from litellm.llms.vertex_ai.agent_engine.sse_iterator import ( VertexAgentEngineResponseIterator, ) @@ -170,3 +170,48 @@ async def test_async_stream_wrapper_without_a_client_posts_through_the_cached_ve assert type(wrapper) is CustomStreamWrapper assert [str(request.url) for request in sent_requests] == [api_base] assert json.loads(sent_requests[0].content) == {"class_method": "stream_query"} + + +_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after" + + +def _separated_text_reply(request: httpx.Request) -> httpx.Response: + chunk: Final = {"content": {"parts": [{"text": _LINE_SEPARATED_TEXT}], "role": "model"}, "finish_reason": "STOP"} + return httpx.Response(200, stream=httpx.ByteStream(f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n".encode())) + + +async def _streamed_text(sync_stream: bool) -> str: + messages: Final = [{"role": "user", "content": "hi"}] + call_args: Final = { + "model": "agent_engine/123", + "custom_llm_provider": "vertex_ai", + "logging_obj": Logging( + model="agent_engine/123", + messages=messages, + stream=True, + call_type="completion" if sync_stream else "acompletion", + start_time=datetime(2026, 1, 1), + litellm_call_id="call-separators", + function_id="fn-separators", + ), + "api_base": "https://us-central1-aiplatform.googleapis.com/v1/reasoningEngines/123:streamQuery", + "headers": {}, + "data": {"class_method": "stream_query"}, + "messages": messages, + "litellm_params": {}, + } + transport: Final = httpx.MockTransport(_separated_text_reply) + if sync_stream: + sync_wrapper: Final = VertexAgentEngineConfig().get_sync_custom_stream_wrapper( + client=HTTPHandler(client=httpx.Client(transport=transport)), **call_args + ) + return "".join(chunk.choices[0].delta.content or "" for chunk in sync_wrapper) + async_wrapper: Final = await VertexAgentEngineConfig().get_async_custom_stream_wrapper( + client=AsyncHTTPHandler(transport=transport), **call_args + ) + return "".join([chunk.choices[0].delta.content or "" async for chunk in async_wrapper]) + + +@pytest.mark.parametrize("sync_stream", [True, False]) +async def test_stream_keeps_unicode_line_separators_inside_the_text(sync_stream: bool) -> None: + assert await _streamed_text(sync_stream) == _LINE_SEPARATED_TEXT diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 3e1bca18f08..67c78b62151 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -3,8 +3,9 @@ import json import logging import re import threading -from collections.abc import Awaitable, Callable, Iterator, Mapping +from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping from copy import deepcopy +from itertools import pairwise from typing import Final, List, Literal, cast, get_args from unittest.mock import MagicMock, Mock, patch @@ -34,7 +35,7 @@ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( ) from litellm.types.llms.base import CachedTokensDetails from litellm.types.llms.vertex_ai import GeminiFinishReason, RequestBody, UsageMetadata -from litellm.types.utils import ChoiceLogprobs, ModelInfo, Usage +from litellm.types.utils import ChoiceLogprobs, ModelInfo, ModelResponseStream, Usage from litellm.utils import _invalidate_model_cost_lowercase_map, CustomStreamWrapper from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.integrations.custom_logger import CustomLogger @@ -10125,7 +10126,7 @@ def test_vertex_ai_streaming_response_id(): mock_response = MagicMock() mock_response.status_code = 200 - mock_response.iter_lines = MagicMock(return_value=stream_response()) + mock_response.iter_text = MagicMock(return_value=stream_response()) return mock_response logging_obj = MagicMock() @@ -10690,6 +10691,63 @@ async def test_gemini_sends_a_boolean_for_a_gateway_default_written_as_text_or_n assert sent["includeThoughts"] is expected +class _BodyInChunks(httpx.SyncByteStream, httpx.AsyncByteStream): + def __init__(self, chunks: tuple[bytes, ...]) -> None: + self._chunks: Final = chunks + + def __iter__(self) -> Iterator[bytes]: + yield from self._chunks + + async def __aiter__(self) -> AsyncIterator[bytes]: + for chunk in self._chunks: + yield chunk + + +def _split_at(body: bytes, cuts: tuple[int, ...]) -> tuple[bytes, ...]: + return tuple(body[start:end] for start, end in pairwise((0, *sorted(cuts), len(body)))) + + +async def _gemini_stream(is_async: bool, chunks: tuple[bytes, ...]) -> list[ModelResponseStream]: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, stream=_BodyInChunks(chunks), headers={"content-type": "text/event-stream"}) + + kwargs: Final = { + "model": "gemini/gemini-3.8-flash", + "messages": [{"role": "user", "content": "hi"}], + "stream": True, + "stream_options": {"include_usage": True}, + "api_key": "fake-gemini-key", + } + if is_async: + stream: Final = await litellm.acompletion( + **kwargs, client=AsyncHTTPHandler(transport=httpx.MockTransport(handler)) + ) + return [chunk async for chunk in stream] + return list( + litellm.completion(**kwargs, client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler)))) + ) + + +@pytest.mark.parametrize("is_async", [True, False]) +async def test_gemini_stream_keeps_text_and_usage_when_a_part_contains_a_unicode_line_separator(is_async: bool) -> None: + events: Final = ( + {"candidates": [{"content": {"role": "model", "parts": [{"text": "line one\u2028line two"}]}}]}, + { + "candidates": [{"content": {"role": "model", "parts": [{"text": " and three"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 8, "totalTokenCount": 18}, + }, + ) + body: Final = "".join(f"data: {json.dumps(event, ensure_ascii=False)}\r\n\r\n" for event in events).encode() + cut_inside_separator_and_crlf: Final = (body.index(b"\xe2\x80\xa8") + 1, body.index(b"\r") + 1) + + received: Final = await _gemini_stream(is_async, _split_at(body, cut_inside_separator_and_crlf)) + + usage: Final = next(chunk.usage for chunk in reversed(received) if getattr(chunk, "usage", None)) + text: Final = "".join(c.choices[0].delta.content or "" for c in received if c.choices) + assert text == "line one\u2028line two and three" + assert (usage.prompt_tokens, usage.completion_tokens) == (10, 8) + + @pytest.mark.parametrize("chunk_indices", [[0, 1], [2, 1, 0]]) def test_grounding_support_preserves_every_web_source(chunk_indices: list[int]) -> None: metadata: Final = { diff --git a/tests/unit/llms/vertex_ai/test_http_status_201.py b/tests/unit/llms/vertex_ai/test_http_status_201.py index 34502d5cebe..0ed469acfa7 100644 --- a/tests/unit/llms/vertex_ai/test_http_status_201.py +++ b/tests/unit/llms/vertex_ai/test_http_status_201.py @@ -39,8 +39,8 @@ class TestVertexAIHTTPStatus201(unittest.TestCase): # Create a mock response with status code 201 mock_response = MagicMock() mock_response.status_code = 201 - mock_response.aiter_lines = MagicMock() - mock_response.aiter_lines.return_value = ["test response"] + mock_response.aiter_text = MagicMock() + mock_response.aiter_text.return_value = ["test response"] # Setup mock client mock_client = MagicMock() @@ -76,8 +76,8 @@ class TestVertexAIHTTPStatus201(unittest.TestCase): # Create a mock response with status code 201 mock_response = MagicMock() mock_response.status_code = 201 - mock_response.iter_lines = MagicMock() - mock_response.iter_lines.return_value = ["test response"] + mock_response.iter_text = MagicMock() + mock_response.iter_text.return_value = ["test response"] mock_post.return_value = mock_response # Call the make_sync_call function diff --git a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py index a7d7ac5627c..7abe6ad7799 100644 --- a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py @@ -4,10 +4,11 @@ Mock tests for A2A endpoints. Tests that invoke_agent_a2a properly integrates with add_litellm_data_to_request. """ +import asyncio import json import socket import sys -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping from contextlib import AbstractContextManager, ExitStack from dataclasses import dataclass from typing import Final @@ -24,6 +25,21 @@ from litellm.types.agents import AgentCaller AddLiteLLMData = Callable[..., Awaitable[dict[str, object]]] +class _SSELineStream(httpx.AsyncByteStream): + def __init__(self, lines: tuple[str, ...], first_line_delay: float) -> None: + self._lines = lines + self._first_line_delay = first_line_delay + + async def __aiter__(self) -> AsyncIterator[bytes]: + await asyncio.sleep(self._first_line_delay) + for line in self._lines: + yield f"{line}\n".encode() + + +def _sse_upstream_response(*lines: str, first_line_delay: float = 0.0) -> httpx.Response: + return httpx.Response(200, stream=_SSELineStream(lines, first_line_delay)) + + @dataclass(frozen=True, slots=True) class CapturedAgentCall: request_id: object @@ -941,13 +957,7 @@ async def test_task_methods_extract_litellm_params_before_forwarding(method: str mock_http_response.json.return_value = upstream_response mock_http_response.is_success = True - async def fake_aiter_lines(): - yield 'data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1"}}' - - mock_resp = AsyncMock() - mock_resp.is_success = True - mock_resp.aiter_lines = fake_aiter_lines - mock_resp.aclose = AsyncMock() + mock_resp = _sse_upstream_response('data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1"}}') mock_async_client = MagicMock() mock_async_client.build_request = MagicMock(return_value=MagicMock()) @@ -982,8 +992,8 @@ async def test_task_methods_extract_litellm_params_before_forwarding(method: str user_api_key_dict=user_api_key_dict, ) if method == "tasks/resubscribe": - async for _ in response.body_iterator: - pass + streamed = "".join([chunk async for chunk in response.body_iterator]) + assert "task-1" in streamed if method == "tasks/resubscribe": forwarded_body = mock_async_client.build_request.call_args.kwargs["json"] @@ -1007,14 +1017,7 @@ async def test_subscribe_to_task_returns_sse_stream(): 'data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1","status":{"state":"completed"}}}', ] - async def fake_aiter_lines(): - for line in sse_lines: - yield line - - mock_resp = AsyncMock() - mock_resp.is_success = True - mock_resp.aiter_lines = fake_aiter_lines - mock_resp.aclose = AsyncMock() + mock_resp = _sse_upstream_response(*sse_lines) mock_async_client = MagicMock() mock_async_client.build_request = MagicMock(return_value=MagicMock()) @@ -1064,13 +1067,9 @@ async def test_subscribe_to_task_calls_pre_call_hook(): mock_request = _make_request_mock("tasks/resubscribe", {"id": "task-1"}) user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1") - async def fake_aiter_lines(): - yield 'data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1","status":{"state":"completed"}}}' - - mock_resp = AsyncMock() - mock_resp.is_success = True - mock_resp.aiter_lines = fake_aiter_lines - mock_resp.aclose = AsyncMock() + mock_resp = _sse_upstream_response( + 'data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1","status":{"state":"completed"}}}' + ) mock_async_client = MagicMock() mock_async_client.build_request = MagicMock(return_value=MagicMock()) @@ -1117,9 +1116,9 @@ async def test_subscribe_to_task_calls_pre_call_hook(): ) assert response.media_type == "text/event-stream" - async for _ in response.body_iterator: - pass + streamed = "".join([chunk async for chunk in response.body_iterator]) + assert "completed" in streamed mock_proxy_logging.pre_call_hook.assert_awaited_once() call_kwargs = mock_proxy_logging.pre_call_hook.await_args.kwargs assert call_kwargs.get("call_type") == "asend_message" @@ -1149,16 +1148,10 @@ async def test_subscribe_to_task_runs_post_call_streaming_guardrail(): mock_request = _make_request_mock("tasks/resubscribe", {"id": "task-1"}) user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1") - async def fake_aiter_lines(): - yield ( - 'data: {"jsonrpc":"2.0","id":"req-1","result":' - '{"kind":"message","parts":[{"kind":"text","text":"resubscribe-secret"}]}}' - ) - - mock_resp = AsyncMock() - mock_resp.is_success = True - mock_resp.aiter_lines = fake_aiter_lines - mock_resp.aclose = AsyncMock() + mock_resp = _sse_upstream_response( + 'data: {"jsonrpc":"2.0","id":"req-1","result":' + '{"kind":"message","parts":[{"kind":"text","text":"resubscribe-secret"}]}}' + ) mock_async_client = MagicMock() mock_async_client.build_request = MagicMock(return_value=MagicMock()) @@ -1600,14 +1593,7 @@ async def test_pascal_method_names_normalize_to_wire_format(pascal_method: str, mock_http_response.is_success = True mock_http_response.raise_for_status = MagicMock() - async def _empty_aiter_lines(): - return - yield # make it an async generator - - mock_sse_resp = AsyncMock() - mock_sse_resp.is_success = True - mock_sse_resp.aiter_lines = _empty_aiter_lines - mock_sse_resp.aclose = AsyncMock() + mock_sse_resp = _sse_upstream_response() mock_async_client = MagicMock() mock_async_client.build_request = MagicMock(return_value=MagicMock()) @@ -2349,13 +2335,7 @@ async def test_task_methods_forward_trace_header(method: str): mock_http_response.json.return_value = upstream_response mock_http_response.is_success = True - async def fake_aiter_lines(): - yield 'data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1"}}' - - mock_resp = AsyncMock() - mock_resp.is_success = True - mock_resp.aiter_lines = fake_aiter_lines - mock_resp.aclose = AsyncMock() + mock_resp = _sse_upstream_response('data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1"}}') mock_async_client = MagicMock() mock_async_client.build_request = MagicMock(return_value=MagicMock()) @@ -2390,8 +2370,8 @@ async def test_task_methods_forward_trace_header(method: str): user_api_key_dict=user_api_key_dict, ) if method == "tasks/resubscribe": - async for _ in response.body_iterator: - pass + streamed = "".join([chunk async for chunk in response.body_iterator]) + assert "task-1" in streamed if method == "tasks/resubscribe": forwarded_headers = mock_async_client.build_request.call_args.kwargs["headers"] @@ -2646,17 +2626,12 @@ def test_served_version_falls_back_to_header_when_unconfigured(): assert _served_version(_agent(None), _request_with_a2a_header(None)) == "0.3" -def _sse_agent_handler(lines): - mock_resp = AsyncMock() - mock_resp.is_success = True - mock_resp.aiter_lines = lines - mock_resp.aclose = AsyncMock() - - mock_async_client = MagicMock() +def _sse_agent_handler(upstream_response: httpx.Response) -> MagicMock: + mock_async_client: Final = MagicMock() mock_async_client.build_request = MagicMock(return_value=MagicMock()) - mock_async_client.send = AsyncMock(return_value=mock_resp) + mock_async_client.send = AsyncMock(return_value=upstream_response) - mock_handler = MagicMock() + mock_handler: Final = MagicMock() mock_handler.client = mock_async_client return mock_handler @@ -2685,13 +2660,13 @@ async def test_forward_jsonrpc_sse_pings_while_the_upstream_agent_is_still_silen monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", 0.05) - async def _slow_lines(): - await asyncio.sleep(0.3) - yield 'data: {"jsonrpc": "2.0", "id": "req-1", "result": {"kind": "task"}}' + slow_upstream = _sse_upstream_response( + 'data: {"jsonrpc": "2.0", "id": "req-1", "result": {"kind": "task"}}', first_line_delay=0.3 + ) with patch( "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", - return_value=_sse_agent_handler(_slow_lines), + return_value=_sse_agent_handler(slow_upstream), ): response = await _resubscribe_response() assert response.headers["x-accel-buffering"] == "no" @@ -2713,12 +2688,11 @@ async def test_forward_jsonrpc_sse_is_untouched_while_keepalives_are_unconfigure monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", None) - async def _lines(): - yield 'data: {"jsonrpc": "2.0", "id": "req-1", "result": {"kind": "task"}}' - with patch( "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", - return_value=_sse_agent_handler(_lines), + return_value=_sse_agent_handler( + _sse_upstream_response('data: {"jsonrpc": "2.0", "id": "req-1", "result": {"kind": "task"}}') + ), ): response = await _resubscribe_response() assert "x-accel-buffering" not in response.headers @@ -2728,6 +2702,28 @@ async def test_forward_jsonrpc_sse_is_untouched_while_keepalives_are_unconfigure assert json.loads(chunks[-1].removeprefix("data: "))["result"]["kind"] == "task" +@pytest.mark.asyncio +async def test_forward_jsonrpc_sse_keeps_unicode_line_separators_inside_an_event(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", None) + separated_text: Final = "before\u2028between\u2029after" + event: Final = { + "jsonrpc": "2.0", + "id": "req-1", + "result": {"kind": "message", "parts": [{"kind": "text", "text": separated_text}]}, + } + + with patch( + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", + return_value=_sse_agent_handler(_sse_upstream_response(f"data: {json.dumps(event, ensure_ascii=False)}")), + ): + response: Final = await _resubscribe_response() + chunks: Final = [chunk async for chunk in response.body_iterator] + + assert [json.loads(chunk.removeprefix("data: "))["result"]["parts"][0]["text"] for chunk in chunks] == [ + separated_text + ] + + async def _stream_message_response(): from litellm.proxy.agent_endpoints.a2a_endpoints import _handle_stream_message diff --git a/tests/unit/proxy/anthropic_endpoints/test_streaming_model_restamp.py b/tests/unit/proxy/anthropic_endpoints/test_streaming_model_restamp.py index 85bb8304b95..5911fcdc073 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_streaming_model_restamp.py +++ b/tests/unit/proxy/anthropic_endpoints/test_streaming_model_restamp.py @@ -4,6 +4,7 @@ Tests for restamping the public model on Anthropic Messages streaming chunks. import json from collections.abc import AsyncIterator, Sequence +from typing import Final import pytest @@ -46,6 +47,28 @@ def test_restamps_sse_bytes_frame(): assert b"event: message_start" in restamped +def test_restamps_sse_bytes_frame_holding_a_line_separator() -> None: + payload: Final = { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5-20251001", + "content": [{"type": "text", "text": "alpha\u2028beta"}], + }, + } + frame: Final = f"event: message_start\ndata: {json.dumps(payload, ensure_ascii=False)}\n\n".encode() + + restamped: Final = restamp_anthropic_stream_chunk_model(frame, "claude-auto-1") + + assert isinstance(restamped, bytes) + data_line: Final = next(line for line in restamped.decode().split("\n") if line.startswith("data:")) + message: Final = json.loads(data_line[len("data:") :])["message"] + assert message["model"] == "claude-auto-1" + assert message["content"] == [{"type": "text", "text": "alpha\u2028beta"}] + + def test_restamps_event_dict(): chunk = {"type": "message_start", "message": {"id": "msg_1", "model": "claude-sonnet-4-6"}} diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index e09a179eadf..47b6af9aeb3 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -2503,6 +2503,33 @@ def test_stream_was_interrupted_skips_data_lines_that_are_not_json_objects( assert AnthropicPassthroughLoggingHandler._stream_was_interrupted(all_chunks) is interrupted + +def test_stream_was_interrupted_reads_a_message_delta_whose_json_holds_unicode_line_separators() -> None: + message_delta: Final = { + "type": "message_delta", + "delta": {"stop_reason": "stop_sequence", "stop_sequence": "END\u2028OF\u2029TURN"}, + "usage": {"output_tokens": 3}, + } + all_chunks: Final = [ + 'data: {"type": "content_block_stop", "index": 0}', + f"data: {json.dumps(message_delta, ensure_ascii=False)}", + ] + + assert AnthropicPassthroughLoggingHandler._stream_was_interrupted(all_chunks) is False + + +def test_extract_model_reads_a_message_start_whose_json_holds_unicode_line_separators() -> None: + message_start: Final = { + "type": "message_start", + "message": { + "model": "claude-sonnet-5", + "content": [{"type": "text", "text": "Prefill\u2028line\u2029two"}], + }, + } + chunk: Final = f"event: message_start\ndata: {json.dumps(message_start, ensure_ascii=False)}\n\n" + + assert AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks([chunk]) == "claude-sonnet-5" + @pytest.mark.asyncio async def test_anthropic_passthrough_handler(mock_httpx_response, mock_response, mock_logging_obj): """ @@ -3004,3 +3031,24 @@ def test_anthropic_passthrough_stream_bills_the_bedrock_guardrail_trace_of_the_m (entry,) = _recorded_guardrail_entries(mock_logging_obj) assert entry["guardrail_usage"] == {"topicPolicyUnits": 3, "contentPolicyUnits": 3} assert entry["guardrail_cost"] == pytest.approx(6 * 0.00015) + + +def test_anthropic_passthrough_stream_bills_a_bedrock_guardrail_trace_whose_json_holds_unicode_line_separators( + monkeypatch: pytest.MonkeyPatch, mock_logging_obj: LiteLLMLoggingObj, all_chunks: list[str] +) -> None: + monkeypatch.setattr(litellm, "model_cost", _BEDROCK_TRACE_PRICING) + mock_logging_obj.model_call_details["litellm_params"] = {"metadata": {}} + trace: Final = {"guardrail": {**_BEDROCK_INVOKE_TRACE["guardrail"], "actionReason": "Blocked\u2028by\u2029policy"}} + stop_event: Final = json.dumps({"type": "message_stop", "amazon-bedrock-trace": trace}, ensure_ascii=False) + AnthropicPassthroughLoggingHandler.handle_logging_anthropic_collected_chunks( + litellm_logging_obj=mock_logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + request_body={"model": "claude", "messages": []}, + endpoint_type=MagicMock(), + start_time=datetime(2025, 1, 1), + all_chunks=[*all_chunks, f"event: message_stop\ndata: {stop_event}\n\n"], + end_time=datetime(2025, 1, 1), + ) + (entry,) = _recorded_guardrail_entries(mock_logging_obj) + assert entry["guardrail_usage"] == {"topicPolicyUnits": 3, "contentPolicyUnits": 3} diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py index 1b83c5140ca..b0f84f2ecbe 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py @@ -1,6 +1,7 @@ import asyncio import json from datetime import datetime +from typing import Final from unittest.mock import AsyncMock, MagicMock import httpx @@ -267,6 +268,12 @@ class TestRunIdFromSseFrames: frames = b'data: {"run_id": "run-7", "event": "INITIALIZED"}\n\ndata: {"run_id": "run-7", "event": "ACTION"}\n\n' assert run_id_from_sse_frames(frames) == "run-7" + def test_finds_run_id_in_a_frame_whose_json_holds_unicode_line_separators(self) -> None: + event: Final = {"run_id": "run-7", "event": "ACTION", "message": "Clicked\u2028Submit\u2029button"} + frames: Final = f"data: {json.dumps(event, ensure_ascii=False)}\n\n".encode() + + assert run_id_from_sse_frames(frames) == "run-7" + def test_skips_frames_without_run_id(self): frames = b': keepalive\n\ndata: not-json\n\ndata: {"event": "HEARTBEAT"}\n\n' assert run_id_from_sse_frames(frames) is None diff --git a/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py b/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py index 49bc2e9c29e..2496aa47d22 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py +++ b/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py @@ -30,7 +30,7 @@ from litellm.llms.base_llm.managed_resources.utils import( resolve_passthrough_managed_id_provider, ) from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome -from typing import Any +from typing import Any, Final def _user() -> UserAPIKeyAuth: @@ -254,6 +254,37 @@ async def test_streamed_response_with_cr_only_frame_delimiters_is_still_owned_an assert output == _response_stream_bytes(managed_id).replace(b"\n", b"\r") +def _instructed_response_stream_bytes(raw_id: str, instructions: str) -> bytes: + events: Final = ( + {"type": "response.created", "response": {"id": raw_id, "status": "in_progress", "instructions": instructions}}, + {"type": "response.completed", "response": {"id": raw_id, "status": "completed", "instructions": instructions}}, + ) + return b"".join( + f"event: {event['type']}\ndata: {json.dumps(event, ensure_ascii=False)}\n\n".encode() for event in events + ) + + +@pytest.mark.asyncio +async def test_streamed_response_whose_json_holds_unicode_line_separators_is_still_owned_and_rewritten() -> None: + instructions: Final = "Answer\u2028in\u2029verse" + pc: Final = _prisma_client() + + output: Final = await _collect( + rewrite_streamed_response_ids( + stream=_chunks(_instructed_response_stream_bytes(RAW_RESPONSE_ID, instructions), 4096), + provider="openai", + method="POST", + route="/openai_passthrough/v1/responses", + user_api_key_dict=_user(), + prisma_client=pc, + ) + ) + + pc.db.litellm_managedobjecttable.upsert.assert_awaited_once() + managed_id: Final = pc.db.litellm_managedobjecttable.upsert.await_args.kwargs["data"]["create"]["unified_object_id"] + assert output == _instructed_response_stream_bytes(managed_id, instructions) + + @pytest.mark.asyncio async def test_streamed_bytes_untouched_on_routes_without_a_response_id(): pc = _prisma_client() diff --git a/tests/unit/proxy/response_polling/test_response_polling_handler.py b/tests/unit/proxy/response_polling/test_response_polling_handler.py index 63e38ce533a..42acad28488 100644 --- a/tests/unit/proxy/response_polling/test_response_polling_handler.py +++ b/tests/unit/proxy/response_polling/test_response_polling_handler.py @@ -17,7 +17,7 @@ following the OpenAI Response API format. import asyncio import json from datetime import datetime, timezone -from typing import Any, Dict, Optional +from typing import Any, Dict, Final, Optional from unittest.mock import AsyncMock, Mock, patch import pytest @@ -1400,7 +1400,7 @@ def _make_sse_stream(events: list) -> Mock: async def _body_iterator(): for event in events: - yield f"data: {json.dumps(event)}" + yield f"data: {json.dumps(event, ensure_ascii=False)}" yield "data: [DONE]" mock_response = Mock() @@ -1639,6 +1639,33 @@ class TestBackgroundStreamingTerminalEvents: assert final_call.kwargs["status"] == "completed" assert final_call.kwargs["usage"] == {"input_tokens": 10, "output_tokens": 50} + @pytest.mark.asyncio + async def test_response_completed_whose_json_holds_unicode_line_separators_keeps_its_output(self) -> None: + from litellm.proxy.response_polling.background_streaming import ( + background_streaming_task, + ) + + item: Final = { + "id": "item_1", + "type": "message", + "content": [{"type": "output_text", "text": "Line\u2028one\u2029two"}], + } + events: Final = [ + {"type": "response.completed", "response": {"id": "resp_123", "status": "completed", "output": [item]}} + ] + handler: Final = AsyncMock(spec=ResponsePollingHandler) + + with patch( + "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" + ) as MockProcessor: + MockProcessor.return_value.base_process_llm_request = AsyncMock(return_value=_make_sse_stream(events)) + await background_streaming_task(**_make_background_streaming_kwargs("poll_u2028", handler)) + + calls: Final = handler.update_state.call_args_list + outputs: Final = [call.kwargs["output"] for call in calls if "output" in call.kwargs] + assert outputs[-1] == [item] + assert calls[-1].kwargs["status"] == "completed" + @pytest.mark.asyncio async def test_fallback_status_derived_from_event_type_when_status_field_missing( self, diff --git a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index ed90acbd2fa..0b971bc1696 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -1288,6 +1288,106 @@ async def test_reasoning_item_closes_before_a_tool_call_from_the_same_chunk_open ] +def _text_opening_a_tool_call_chunk(content: str, arguments: str = "") -> ModelResponseStream: + return ModelResponseStream( + id=CHAT_COMPLETION_ID, + created=1748575031, + model="claude-haiku-5-5", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta( + role="assistant", + content=content, + tool_calls=[ + { + "id": "call_weather", + "type": "function", + "function": {"name": "get_weather", "arguments": arguments}, + "index": 0, + } + ], + ), + finish_reason=None, + ) + ], + ) + + +def _tool_arguments_chunk(arguments: str, finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream( + id=CHAT_COMPLETION_ID, + created=1748575031, + model="claude-haiku-5-5", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta( + content=None, + tool_calls=[{"type": "function", "function": {"arguments": arguments}, "index": 0}], + ), + finish_reason=finish_reason, + ) + ], + ) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_tool_call_opened_in_a_text_chunk_streams_each_argument_delta(sync_mode: bool) -> None: + argument_pieces: Final = ('{"', "city", '"', ': "', "Paris", '"', "}") + iterator: Final = _build_iterator( + [ + _text_opening_a_tool_call_chunk(""), + *(_tool_arguments_chunk(piece) for piece in argument_pieces), + _chunk("", finish_reason="tool_calls"), + ] + ) + + events: Final = await _collect_events(iterator, sync_mode) + + assert _event_text(events, ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA) == "" + assert [ + event.delta + for event in events + if getattr(event, "type", None) == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA + ] == list(argument_pieces) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_tool_call_sharing_a_chunk_with_text_opens_after_the_whole_message(sync_mode: bool) -> None: + iterator: Final = _build_iterator( + [ + _text_opening_a_tool_call_chunk("Checking ", arguments='{"city": "Paris"}'), + _chunk("the "), + _chunk("weather now."), + _chunk("", finish_reason="tool_calls"), + ] + ) + + events: Final = await _collect_events(iterator, sync_mode) + + event_types: Final = [getattr(event, "type", None) for event in events] + function_call_added_at: Final = next( + position + for position, event in enumerate(events) + if event_types[position] == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED and event.item.type == "function_call" + ) + streamed: Final = [ + event + for event in events + if getattr(event, "type", None) + in (ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA) + ] + assert ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA not in event_types[function_call_added_at:] + assert [event.sequence_number for event in streamed] == sorted(event.sequence_number for event in streamed) + assert _event_text(events, ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA) == "Checking the weather now." + assert _event_text(events, ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA) == '{"city": "Paris"}' + + @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_tool_only_stream_puts_the_function_call_at_output_index_zero(sync_mode: bool): diff --git a/tests/unit/sandbox/test_e2b_sandbox.py b/tests/unit/sandbox/test_e2b_sandbox.py index e01b9120416..e58392ca266 100644 --- a/tests/unit/sandbox/test_e2b_sandbox.py +++ b/tests/unit/sandbox/test_e2b_sandbox.py @@ -7,6 +7,7 @@ integration tests live in tests/integration/sandbox/test_e2b_sandbox.py. """ import json +from typing import Final import httpx import pytest @@ -19,19 +20,18 @@ from litellm.llms.e2b.sandbox.transformation import ( ) +def _line_stream(lines: list[str]) -> httpx.Response: + return httpx.Response(200, stream=httpx.ByteStream("".join(f"{line}\n" for line in lines).encode())) + + class FakeResponse: - def __init__(self, *, json_data=None, lines=None, status_code=200): + def __init__(self, *, json_data: object = None, status_code: int = 200) -> None: self._json = json_data - self._lines = lines or [] self.status_code = status_code def json(self): return self._json - async def aiter_lines(self): - for line in self._lines: - yield line - class FakeHTTPClient: """Records outbound requests and returns canned responses keyed by URL.""" @@ -61,7 +61,7 @@ class FakeHTTPClient: if url.endswith("/execute"): if self.execute_raises is not None: raise self.execute_raises - return FakeResponse(lines=self.execute_lines) + return _line_stream(self.execute_lines) raise AssertionError(f"unexpected POST {url}") async def delete(self, url, headers=None, **kwargs): @@ -249,6 +249,20 @@ def test_parse_lines_skips_non_json_lines(): assert result.error is None +@pytest.mark.asyncio +async def test_run_code_keeps_unicode_line_separators_inside_stdout() -> None: + separated_text: Final = "before\u2028between\u2029after\n" + client: Final = FakeHTTPClient( + execute_lines=[json.dumps({"type": "stdout", "text": separated_text, "timestamp": 1}, ensure_ascii=False)] + ) + handle: Final = ContainerHandle(id="sbx_sep", provider="e2b", domain="e2b.app") + handle._hidden_params = {"envd_access_token": "tok"} + + result: Final = await E2BSandboxConfig().arun_code(container=handle, code="print(text)", client=client) + + assert result.stdout == separated_text + + @pytest.mark.asyncio async def test_run_code_aborts_on_output_over_cap(): big_line = "x" * (MAX_OUTPUT_BYTES + 1) diff --git a/tests/unit/sandbox/test_opensandbox_sandbox.py b/tests/unit/sandbox/test_opensandbox_sandbox.py index 2928dea100e..36ae848786c 100644 --- a/tests/unit/sandbox/test_opensandbox_sandbox.py +++ b/tests/unit/sandbox/test_opensandbox_sandbox.py @@ -27,10 +27,13 @@ def sse(data): return f"data: {json.dumps(data)}" +def _line_stream(lines: list[str]) -> httpx.Response: + return httpx.Response(200, stream=httpx.ByteStream("".join(f"{line}\n" for line in lines).encode())) + + class FakeResponse: - def __init__(self, *, json_data=None, lines=None, status_code=200): + def __init__(self, *, json_data: object = None, status_code: int = 200) -> None: self._json = json_data - self._lines = lines or [] self.status_code = status_code def json(self): @@ -40,10 +43,6 @@ class FakeResponse: if self.status_code >= 400: raise http_status_error(self.status_code) - async def aiter_lines(self): - for line in self._lines: - yield line - class FakeHTTPClient: def __init__( @@ -93,7 +92,7 @@ class FakeHTTPClient: if url.endswith("/code"): if self.execute_raises is not None: raise self.execute_raises - return FakeResponse(lines=self.execute_lines) + return _line_stream(self.execute_lines) raise AssertionError(f"unexpected POST {url}") async def get(self, url, headers=None, params=None, **kwargs): diff --git a/tests/unit/test_groq_streaming_encoding.py b/tests/unit/test_groq_streaming_encoding.py index ffe87fb9890..37f05272f67 100644 --- a/tests/unit/test_groq_streaming_encoding.py +++ b/tests/unit/test_groq_streaming_encoding.py @@ -6,148 +6,58 @@ UTF-8 encoded content in streaming responses, specifically fixing the ASCII encoding error described in issue #12660. """ -import asyncio -from unittest.mock import AsyncMock, Mock +import json +from typing import Final +from unittest.mock import Mock +import httpx import pytest +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.openai_like.chat.handler import make_call, make_sync_call - -class MockResponse: - """Mock httpx response for testing UTF-8 handling.""" - - def __init__(self, test_content: str): - self.test_content = test_content - self.status_code = 200 - - def iter_text(self, encoding="utf-8"): - """Mock iter_text that yields content with the specified encoding.""" - yield self.test_content - - async def aiter_text(self, encoding="utf-8"): - """Mock aiter_text that yields content with the specified encoding.""" - yield self.test_content - - def iter_lines(self): - """Mock iter_lines method for synchronous streaming.""" - yield self.test_content - - async def aiter_lines(self): - """Mock aiter_lines method for asynchronous streaming.""" - yield self.test_content - - def json(self): - return {"choices": [{"delta": {"content": "test"}}]} +_UNICODE_TEXTS: Final = ( + "The symbol µ represents micro", + "Testing © character", + "Testing ™ character", + "Testing € character", + "Testing 北京 character", + "Testing 🚀 character", + "Testing Ñoño character", +) -class MockSyncClient: - """Mock synchronous HTTP client for testing.""" - - def __init__(self, response_content: str): - self.response_content = response_content - - def post(self, *args, **kwargs): - return MockResponse(self.response_content) +def _utf8_stream_transport(text: str) -> httpx.MockTransport: + chunk: Final = {"choices": [{"index": 0, "delta": {"content": text}}]} + body: Final = f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n".encode() + return httpx.MockTransport(lambda request: httpx.Response(200, stream=httpx.ByteStream(body))) -class MockAsyncClient: - """Mock asynchronous HTTP client for testing.""" - - def __init__(self, response_content: str): - self.response_content = response_content - - async def post(self, *args, **kwargs): - return MockResponse(self.response_content) +def _call_args() -> dict[str, object]: + return { + "api_base": "https://test.com/v1/chat/completions", + "headers": {"Authorization": "Bearer test"}, + "data": '{"model": "test", "messages": []}', + "model": "test-model", + "messages": [], + "logging_obj": Mock(), + } -def test_utf8_streaming_sync(): - """Test that synchronous streaming handles UTF-8 characters correctly.""" - # Content with the µ character that was causing issues - test_content = ( - 'data: {"choices":[{"delta":{"content":"The symbol µ represents micro"}}]}\n\n' +@pytest.mark.parametrize("text", _UNICODE_TEXTS) +def test_utf8_streaming_sync(text: str) -> None: + completion_stream: Final = make_sync_call( + client=HTTPHandler(client=httpx.Client(transport=_utf8_stream_transport(text))), **_call_args() ) - mock_client = MockSyncClient(test_content) - mock_logging = Mock() - - # This should not raise an ASCII encoding error - completion_stream = make_sync_call( - client=mock_client, - api_base="https://test.com/v1/chat/completions", - headers={"Authorization": "Bearer test"}, - data='{"model": "test", "messages": []}', - model="test-model", - messages=[], - logging_obj=mock_logging, - ) - - # Verify we can iterate through the stream without encoding errors - assert completion_stream is not None + assert "".join(chunk["text"] for chunk in completion_stream) == text @pytest.mark.asyncio -async def test_utf8_streaming_async(): - """Test that asynchronous streaming handles UTF-8 characters correctly.""" - # Content with the µ character that was causing issues - test_content = ( - 'data: {"choices":[{"delta":{"content":"The symbol µ represents micro"}}]}\n\n' +@pytest.mark.parametrize("text", _UNICODE_TEXTS) +async def test_utf8_streaming_async(text: str) -> None: + completion_stream: Final = await make_call( + client=AsyncHTTPHandler(transport=_utf8_stream_transport(text)), **_call_args() ) - mock_client = MockAsyncClient(test_content) - mock_logging = Mock() - - # This should not raise an ASCII encoding error - completion_stream = await make_call( - client=mock_client, - api_base="https://test.com/v1/chat/completions", - headers={"Authorization": "Bearer test"}, - data='{"model": "test", "messages": []}', - model="test-model", - messages=[], - logging_obj=mock_logging, - ) - - # Verify we can iterate through the stream without encoding errors - assert completion_stream is not None - - -def test_various_unicode_characters(): - """Test streaming with various Unicode characters that could cause issues.""" - unicode_test_cases = [ - "µ", # Micro symbol (the original issue) - "©", # Copyright symbol - "™", # Trademark symbol - "€", # Euro symbol - "北京", # Chinese characters - "🚀", # Emoji - "Ñoño", # Spanish characters with tildes - ] - - for unicode_char in unicode_test_cases: - test_content = f'data: {{"choices":[{{"delta":{{"content":"Testing {unicode_char} character"}}}}]}}\n\n' - - mock_client = MockSyncClient(test_content) - mock_logging = Mock() - - # This should not raise an ASCII encoding error for any Unicode character - completion_stream = make_sync_call( - client=mock_client, - api_base="https://test.com/v1/chat/completions", - headers={"Authorization": "Bearer test"}, - data='{"model": "test", "messages": []}', - model="test-model", - messages=[], - logging_obj=mock_logging, - ) - - assert ( - completion_stream is not None - ), f"Failed to handle Unicode character: {unicode_char}" - - -if __name__ == "__main__": - test_utf8_streaming_sync() - asyncio.run(test_utf8_streaming_async()) - test_various_unicode_characters() - print("All UTF-8 streaming tests passed!") + assert "".join([chunk["text"] async for chunk in completion_stream]) == text From 7650f766b006077442abe758ee4b22b5c436ad18 Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Wed, 28 Oct 2026 17:01:52 -0700 Subject: [PATCH 118/128] refactor(router): fence the router surface and document the router-core boundary (#48839) Add rustdoc overviews for a new litellm-router-core crate and for litellm-model-catalog, a module docstring on litellm/router.py, and a shrink-only allowlist of the Router members and router_utils modules used outside the router, checked in CI by router_surface_check.py. Move the router-free helpers out of litellm/router_utils: batch_utils to litellm_core_utils, and the auto-router naming contract, tuning baseline and baseline request next to the auto-router plugin. No behaviour change. --- .github/workflows/test-linting.yml | 3 + litellm-rust/Cargo.lock | 4 + litellm-rust/crates/model-catalog/src/lib.rs | 33 +++ litellm-rust/crates/router-core/Cargo.toml | 9 + litellm-rust/crates/router-core/src/lib.rs | 63 +++++ .../batch_utils.py | 0 .../prompt_templates/common_utils.py | 2 +- .../llms/anthropic/prompt_cache_prediction.py | 2 +- litellm/proxy/common_request_processing.py | 2 +- .../guardrails/auto_router_compression.py | 2 +- litellm/proxy/health_check.py | 2 +- .../proxy/hooks/autorouter_baseline_cache.py | 2 +- .../proxy/hooks/model_max_budget_limiter.py | 2 +- litellm/proxy/litellm_pre_call_utils.py | 2 +- .../auto_router_endpoints.py | 12 +- .../model_management_endpoints.py | 31 +-- .../auto_router_availability.py | 4 +- .../auto_router_permissions.py | 5 +- litellm/proxy/proxy_server.py | 14 +- litellm/router.py | 61 +++-- .../auto_router}/auto_router_model_naming.py | 0 .../auto_router_tuning_baseline.py | 0 .../auto_router}/baseline_request.py | 0 litellm/router_strategy/budget_limiter.py | 2 +- litellm/router_strategy/least_busy.py | 2 +- litellm/router_strategy/lowest_cost.py | 2 +- litellm/router_strategy/lowest_latency.py | 2 +- litellm/router_strategy/lowest_tpm_rpm.py | 2 +- litellm/router_strategy/lowest_tpm_rpm_v2.py | 2 +- .../router_utils/fallback_event_handlers.py | 2 +- .../router_code_coverage.py | 4 +- .../router_surface_allowlist.json | 211 ++++++++++++++++ .../router_surface_check.py | 235 ++++++++++++++++++ .../unbounded_in_baseline.txt | 2 +- .../test_batch_utils.py | 13 +- .../test_auto_router_endpoints.py | 4 +- .../test_model_management_endpoints.py | 49 ++-- .../test_auto_router_availability.py | 7 +- .../unit/proxy/proxy_server/test_lifecycle.py | 12 +- .../test_auto_router_model_naming.py | 8 +- .../test_auto_router_tuning_baseline.py | 4 +- .../test_baseline_request.py | 2 +- .../router_strategy/test_complexity_router.py | 36 ++- tests/unit/router_strategy/test_llm_v2.py | 4 +- tests/unit/test_private_usage_aliases.py | 2 +- .../test_router/test_router_helper_utils.py | 15 +- .../test_router/test_router_surface_check.py | 58 +++++ 47 files changed, 797 insertions(+), 138 deletions(-) create mode 100644 litellm-rust/crates/router-core/Cargo.toml create mode 100644 litellm-rust/crates/router-core/src/lib.rs rename litellm/{router_utils => litellm_core_utils}/batch_utils.py (100%) rename litellm/{router_utils => router_strategy/auto_router}/auto_router_model_naming.py (100%) rename litellm/{router_utils => router_strategy/auto_router}/auto_router_tuning_baseline.py (100%) rename litellm/{router_utils => router_strategy/auto_router}/baseline_request.py (100%) create mode 100644 tests/code_coverage_tests/router_surface_allowlist.json create mode 100644 tests/code_coverage_tests/router_surface_check.py rename tests/unit/{router_utils => litellm_core_utils}/test_batch_utils.py (97%) rename tests/unit/{router_utils => router_strategy}/test_auto_router_model_naming.py (99%) rename tests/unit/{router_utils => router_strategy}/test_auto_router_tuning_baseline.py (99%) rename tests/unit/{router_utils => router_strategy}/test_baseline_request.py (95%) create mode 100644 tests/unit/test_router/test_router_surface_check.py diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index b9f8d7e7dfb..89ca60ad9d1 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -561,6 +561,9 @@ jobs: - name: router_code_coverage run: uv run --no-sync python ./tests/code_coverage_tests/router_code_coverage.py + - name: router_surface_check + run: uv run --no-sync python ./tests/code_coverage_tests/router_surface_check.py + - name: test_chat_completion_imports run: uv run --no-sync python ./tests/code_coverage_tests/test_chat_completion_imports.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 9c45c5c4b1c..10e590b5006 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4393,6 +4393,10 @@ dependencies = [ "rstest", ] +[[package]] +name = "litellm-router-core" +version = "0.1.0" + [[package]] name = "litellm-router-types" version = "0.1.0" diff --git a/litellm-rust/crates/model-catalog/src/lib.rs b/litellm-rust/crates/model-catalog/src/lib.rs index 066c2c83c6b..d42c05e6fe7 100644 --- a/litellm-rust/crates/model-catalog/src/lib.rs +++ b/litellm-rust/crates/model-catalog/src/lib.rs @@ -1,3 +1,36 @@ +//! The model catalog: per-model pricing, context limits and capabilities. +//! +//! This is the Rust form of `model_prices_and_context_window.json`, the map Python reads as +//! `litellm.model_cost` and through `litellm.get_model_info`. It answers "what does this model +//! cost and what can it do"; it knows nothing about configured deployments. The router's view of +//! deployments (groups, aliases, wildcard patterns, team mappings) is the router registry, defined +//! in `litellm-router-core`, which reads this catalog through its dynamic config. +//! +//! # Structure +//! - `Catalog`: the parsed map. `Catalog::parse` takes the JSON body and a `Provenance` (where +//! the body came from: the bundled copy in the repo or the remote copy, with revision and etag), +//! and keeps every entry in file order with its raw fields. +//! - `ModelEntry`: one model. `ModelEntry::fields` is the raw JSON object, so unknown fields are +//! never lost; `ModelEntry::info` is the typed mirror `ModelInfo` (pricing in `pricing`, +//! modalities and modes in `capabilities`). +//! - Aliases: an entry may list `aliases`; each alias resolves to its canonical key. Problems +//! found while indexing them (bad lists, collisions) are kept as `AliasIssue`s instead of +//! failing the parse. +//! - `FallbackGeneralizations` / `FallbackRule`: regex rules that map unknown model ids onto known +//! families. They are data in the file, not entries. +//! - Validation: `Catalog::validate` checks a catalog against `IntegrityLimits` (a remote copy that +//! lost most of its models is refused), and `validate_registry` checks every entry against the +//! checked-in contract. +//! +//! # Lookups +//! `Catalog::lookup` resolves a key exactly first (entry or alias), then case-insensitively, and +//! returns a `ModelMatch` with both the key that matched and the canonical key of the entry. Provider +//! prefixes are not stripped here; callers decide which keys to try. +//! +//! # Readers +//! Cost tracking, the router (through `RouterDynamicConfig`) and the gateway. Nothing in this crate +//! performs I/O; loading and refreshing the file is the caller's job. + mod capabilities; mod catalog; mod error; diff --git a/litellm-rust/crates/router-core/Cargo.toml b/litellm-rust/crates/router-core/Cargo.toml new file mode 100644 index 00000000000..af97bd52c57 --- /dev/null +++ b/litellm-rust/crates/router-core/Cargo.toml @@ -0,0 +1,9 @@ +[package] +name = "litellm-router-core" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true +description = "The router core behind litellm.Router: registry, settings, health, selection and the call loop" + +[dependencies] diff --git a/litellm-rust/crates/router-core/src/lib.rs b/litellm-rust/crates/router-core/src/lib.rs new file mode 100644 index 00000000000..05d95d02c84 --- /dev/null +++ b/litellm-rust/crates/router-core/src/lib.rs @@ -0,0 +1,63 @@ +//! Routing for LiteLLM: decide which deployment serves a request, and drive retries, fallbacks +//! and cooldowns around the attempts. +//! +//! This crate defines the boundary `litellm.Router` is moving behind. Today the Python router in +//! `litellm/router.py` implements all of it; this crate is where the Rust core will implement the +//! same interface. Until then it holds this overview and nothing else, so the interface is written +//! down in code before either side changes. +//! +//! # The public surface +//! Code outside the router uses `litellm.Router` through two kinds of members and nothing else: +//! - The endpoint methods (`acompletion`, `aembedding`, `aresponses`, `aanthropic_messages`, the +//! files, batches, assistants and vector store wrappers, 37 in all). Each is a thin wrapper over +//! `call(op, request)`. +//! - Eight entry points: `registry()`, `settings()`, `health()`, `select()`, `apply()`, +//! `observe()`, `register_plugin()` and `retire()`. +//! +//! Every other member that outside code uses today has a named replacement among those. The old +//! names stay available for about a month through a `DeprecatedRouterMembers` mixin, which +//! implements them over the entry points and warns on use. A CI check +//! (`tests/code_coverage_tests/router_surface_check.py`) keeps an allowlist of the members outside +//! code touches; the list may only shrink. +//! +//! # Structure +//! - `Registry`: the configured deployments, indexed by id, model group, alias, wildcard pattern +//! and team. Immutable and versioned. Seven queries: `deployment(id)`, `deployments(filter)`, +//! `names(filter)`, `resolve(name, scope)`, `group(name)`, `access_groups()` and +//! `credentials(id)`. "Registry" is the router's view of what is configured; per-model pricing +//! and capabilities are the model catalog (`litellm-model-catalog`), which the router reads +//! through its dynamic config. +//! - `RouterConfig` / `RouterDynamicConfig`: everything the router needs to know, passed in. The +//! router never reads process globals. Static config (model list, strategy, retry, timeout and +//! cooldown policy, Redis settings, plugins) is fixed for the router's life. Dynamic config +//! (pricing, model aliases, fallbacks, retry counts, callbacks) is a versioned snapshot read +//! once per call. +//! - `Settings`: one typed, frozen view of the configuration a caller may read. +//! - `Health`: cooldowns, blocked deployments and usage, kept in shared cache keys so Python and +//! Rust workers see the same state. `report(id, outcome)` records outcomes the router did not +//! see itself (a health check, an interceptor that marks a deployment failed). +//! - `select(model, request)`: resolve a name, filter eligible deployments, pick one with the +//! configured strategy, without calling it. +//! - `apply(changes)`: one `ChangeSet` (upserts, deletes, settings), applied atomically, producing +//! a new registry version. +//! - `call(op, request)`: select, run the attempt, and retry or fall back until it succeeds or the +//! policy is exhausted. +//! +//! # Execution +//! Each attempt resolves an `Executor` after a deployment is picked. A Rust-native operation runs +//! its `inference-*` route machine in-process; anything else goes through `Host::invoke`. A +//! started attempt is never replayed on a different executor. +//! +//! # Hosts +//! `Host` is what the embedding process provides: `invoke` for non-native operations, `select` +//! for strategies and sub-routers not yet ported, `allow_fallback` for access and budget checks, +//! `sleep`, `observe` for routing events, and `cache`, a DualCache-style store that holds all +//! routing state under the same keys, encodings and TTLs the Python router uses. `PythonHost` +//! serves them from the existing Python code when the router runs inside Python; `NativeHost` +//! serves them in the Rust server, where an operation without a Rust implementation makes the +//! config fail at build time, never partway through a request. +//! +//! # Plugins +//! Routing strategies, pre-call checks and sub-routers (auto, complexity, quality, adaptive) +//! register through `register_plugin` and see a `RoutingContext` (registry, settings, health and +//! the current request), never the router itself. diff --git a/litellm/router_utils/batch_utils.py b/litellm/litellm_core_utils/batch_utils.py similarity index 100% rename from litellm/router_utils/batch_utils.py rename to litellm/litellm_core_utils/batch_utils.py diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index cf6bdc69556..101f069c0f8 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -21,8 +21,8 @@ from pydantic import TypeAdapter import litellm from litellm import verbose_logger +from litellm.litellm_core_utils.batch_utils import InMemoryFile from litellm.litellm_core_utils.core_helpers import as_str_mapping -from litellm.router_utils.batch_utils import InMemoryFile from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionAssistantMessage, diff --git a/litellm/llms/anthropic/prompt_cache_prediction.py b/litellm/llms/anthropic/prompt_cache_prediction.py index 7a4c707860d..17c178d169e 100644 --- a/litellm/llms/anthropic/prompt_cache_prediction.py +++ b/litellm/llms/anthropic/prompt_cache_prediction.py @@ -29,7 +29,7 @@ from litellm.llms.anthropic.messages.native.transformation import ( AnthropicMessagesConfig, ) from litellm.llms.anthropic.messages.native.utils import AnthropicMessagesRequestUtils, prepare_native_messages -from litellm.router_utils.baseline_request import ( +from litellm.router_strategy.auto_router.baseline_request import ( BASELINE_PARAMETERS, capture_baseline_parameters, ) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 70e305f46fb..9bb451495c8 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -139,8 +139,8 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports ) from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect, spawn_logging_task from litellm.router import Router +from litellm.router_strategy.auto_router.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict -from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX from litellm.router_utils.common_utils import resolve_model_group_alias, resolve_served_model from litellm.types.extraction import ExtractionResponse from litellm.types.guardrails import GuardrailEventHooks diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index 1267fb15db8..3071bcbb1dc 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -14,7 +14,7 @@ from typing import TYPE_CHECKING, Final from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket -from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX +from litellm.router_strategy.auto_router.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 0d2f8dc9a0c..be32405154c 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -27,7 +27,7 @@ from litellm.constants import ( HEALTH_CHECK_TIMEOUT_SECONDS, ) from litellm.litellm_core_utils.health_check_helpers import native_health_check_mode -from litellm.router_utils.auto_router_model_naming import ( +from litellm.router_strategy.auto_router.auto_router_model_naming import ( StrategyRouterDependency, classify_strategy_router_model, strategy_router_dependencies, diff --git a/litellm/proxy/hooks/autorouter_baseline_cache.py b/litellm/proxy/hooks/autorouter_baseline_cache.py index c4232609124..1147ce7ec12 100644 --- a/litellm/proxy/hooks/autorouter_baseline_cache.py +++ b/litellm/proxy/hooks/autorouter_baseline_cache.py @@ -37,8 +37,8 @@ from litellm.proxy.spend_tracking.savings import ( _effective_model_info, # pyright: ignore[reportPrivateUsage] # existing deployment-price owner _proxy_llm_router, # pyright: ignore[reportPrivateUsage] # existing optional proxy-router owner ) +from litellm.router_strategy.auto_router.baseline_request import baseline_request from litellm.router_strategy.complexity_router.context_compaction import compaction_applied -from litellm.router_utils.baseline_request import baseline_request from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.router import BaselineRouteStamp from litellm.types.utils import CallTypes, ModelInfo, Usage diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 11a47b2d424..1c4a61cfad8 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -19,6 +19,7 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import Span +from litellm.litellm_core_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.llms.bedrock.common_utils import get_bedrock_base_model, get_bedrock_cross_region_inference_regions @@ -30,7 +31,6 @@ from litellm.proxy.hooks.model_budget_shared_windows import ( SharedWindows, ) from litellm.router_strategy.budget_limiter import RouterBudgetLimiting, budget_window_seconds -from litellm.router_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import BudgetConfig, LlmProviders, StandardLoggingPayload diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index d4528bc85e8..2c05f724374 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -3290,7 +3290,7 @@ def _per_user_credential_names_for_groups( GITHUB_COPILOT_AUTH_TYPE_KEY, GITHUB_COPILOT_PER_USER_AUTH_TYPE, ) - from litellm.router_utils.auto_router_model_naming import classify_strategy_router_model + from litellm.router_strategy.auto_router.auto_router_model_naming import classify_strategy_router_model names: Final[list[str]] = [] # mutable-ok: accumulates one name per per-user deployment for group in model_groups: diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 4f48f3c57b9..3dceab6b7f7 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -60,16 +60,16 @@ from litellm.repositories.table_repositories import PrismaTableRepository from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository from litellm.repositories.verification_token_repository import VerificationTokenRepository +from litellm.router_strategy.auto_router.auto_router_model_naming import ( + StrategyRouterDependencyRole, + classify_strategy_router_model, + strategy_router_dependencies, +) from litellm.router_strategy.complexity_router import ComplexityRouter from litellm.router_strategy.complexity_router.request_models import ( AutoRouterRoutingTestRequest, RequestComplexityRouterConfig, ) -from litellm.router_utils.auto_router_model_naming import ( - StrategyRouterDependencyRole, - classify_strategy_router_model, - strategy_router_dependencies, -) from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.management_endpoints.auto_router_endpoints import ( SHADOW_EVAL_TURN_VALVE, @@ -411,7 +411,7 @@ async def validate_complexity_router_config( auto_router_environment_reference_error(data.complexity_router_config, user_api_key_dict) ) - from litellm.router_utils.auto_router_model_naming import ( + from litellm.router_strategy.auto_router.auto_router_model_naming import ( validate_complexity_router_config_write, ) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 45006b14df6..8cd3559ba86 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -131,6 +131,23 @@ from litellm.repositories.project_repository import ProjectRepository from litellm.repositories.table_repositories import ModelTableRepository from litellm.repositories.team_repository import TeamRepository from litellm.router import Router +from litellm.router_strategy.auto_router.auto_router_model_naming import ( + GATED_AUTO_ROUTER_CAPABILITIES, + STRATEGY_ROUTER_PARAM_FIELDS, + GatedAutoRouterCapability, + capability_limit_violation, + carries_complexity_router_settings, + count_capability_routers, + gated_capabilities_of, + is_complexity_router_model, + validate_complexity_router_config_placement, + validate_complexity_router_config_write, + validate_strategy_router_model_write, +) +from litellm.router_strategy.auto_router.auto_router_tuning_baseline import ( + is_mutable_tuned_candidate, + tuning_quota_violation, +) from litellm.router_strategy.complexity_router import ( DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE, ClassificationRubric, @@ -144,20 +161,6 @@ from litellm.router_strategy.complexity_router import ( normalize_classification_prompt, ) from litellm.router_strategy.complexity_router.config import resolve_complexity_router_config_write -from litellm.router_utils.auto_router_model_naming import ( - GATED_AUTO_ROUTER_CAPABILITIES, - STRATEGY_ROUTER_PARAM_FIELDS, - GatedAutoRouterCapability, - capability_limit_violation, - carries_complexity_router_settings, - count_capability_routers, - gated_capabilities_of, - is_complexity_router_model, - validate_complexity_router_config_placement, - validate_complexity_router_config_write, - validate_strategy_router_model_write, -) -from litellm.router_utils.auto_router_tuning_baseline import is_mutable_tuned_candidate, tuning_quota_violation from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.bedrock import AwsSessionTag from litellm.types.proxy.management_endpoints.model_management_endpoints import ( diff --git a/litellm/proxy/management_helpers/auto_router_availability.py b/litellm/proxy/management_helpers/auto_router_availability.py index 9b511ce0817..ce7859aff2d 100644 --- a/litellm/proxy/management_helpers/auto_router_availability.py +++ b/litellm/proxy/management_helpers/auto_router_availability.py @@ -7,14 +7,14 @@ from typing import Final from pydantic import Json, TypeAdapter, ValidationError from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper -from litellm.router_utils.auto_router_model_naming import ( +from litellm.router_strategy.auto_router.auto_router_model_naming import ( GATED_AUTO_ROUTER_CAPABILITIES, capability_limit_violation, classify_strategy_router_model, count_capability_routers, gated_capabilities_of, ) -from litellm.router_utils.auto_router_tuning_baseline import ( +from litellm.router_strategy.auto_router.auto_router_tuning_baseline import ( is_mutable_tuned_candidate, mutable_tuned_identities, tuning_quota_violation, diff --git a/litellm/proxy/management_helpers/auto_router_permissions.py b/litellm/proxy/management_helpers/auto_router_permissions.py index 995446db7a6..8597f0d3667 100644 --- a/litellm/proxy/management_helpers/auto_router_permissions.py +++ b/litellm/proxy/management_helpers/auto_router_permissions.py @@ -34,13 +34,16 @@ from litellm.repositories.prisma_protocols import DatabaseClient from litellm.repositories.project_repository import ProjectRepository from litellm.repositories.table_repositories import TeamMembershipRepository from litellm.router import Router +from litellm.router_strategy.auto_router.auto_router_model_naming import ( + classify_strategy_router_model, + strategy_router_dependencies, +) from litellm.router_strategy.complexity_router.config import ( ComplexityRouterConfigWrite, normalize_classifier_config_aliases, resolve_complexity_router_config_write, ) from litellm.router_strategy.complexity_router.request_models import RequestComplexityRouterConfig -from litellm.router_utils.auto_router_model_naming import classify_strategy_router_model, strategy_router_dependencies from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.router import Deployment, updateDeployment diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 70e1c25a70a..b263412ce0f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -168,12 +168,7 @@ from litellm.proxy.common_utils.token_count_settings import ( ) from litellm.proxy.management_helpers.auto_router_availability import AutoRouterCatalogEntry, build_auto_router_catalog from litellm.proxy.management_helpers.model_group_custom_fields import model_group_custom_fields -from litellm.router_utils.access_windows import access_windows_config_error -from litellm.router_utils.add_retry_fallback_headers import ( - get_fallback_errors_from_headers, - get_hidden_params_dict, -) -from litellm.router_utils.auto_router_model_naming import ( +from litellm.router_strategy.auto_router.auto_router_model_naming import ( GATED_AUTO_ROUTER_CAPABILITIES, STRATEGY_ROUTER_PARAM_FIELDS, capability_limit_violation, @@ -181,12 +176,17 @@ from litellm.router_utils.auto_router_model_naming import ( count_capability_routers, validate_complexity_router_config_placement, ) -from litellm.router_utils.auto_router_tuning_baseline import ( +from litellm.router_strategy.auto_router.auto_router_tuning_baseline import ( TUNING_BASELINE_PARAM_NAME, mutable_tuned_identities, snapshot_tuning_baselines, tuning_limit_violation, ) +from litellm.router_utils.access_windows import access_windows_config_error +from litellm.router_utils.add_retry_fallback_headers import ( + get_fallback_errors_from_headers, + get_hidden_params_dict, +) from litellm.router_utils.common_utils import resolve_model_group_alias from litellm.router_utils.routing_groups import parse_routing_groups from litellm.types.caching import RedisPipelineIncrementOperation diff --git a/litellm/router.py b/litellm/router.py index 408cb20945f..fc888495e99 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -7,6 +7,29 @@ # # Thank you ! We ❤️ you! - Krrish & Ishaan +""" +litellm.Router: load balancing, retries, fallbacks and cooldowns across deployments. + +Public surface. The router is moving behind a narrow interface; code outside +the router should end up using a Router through two kinds of members only: + +- endpoint methods (acompletion, aembedding, aresponses, aanthropic_messages, + the files, batches, assistants and vector store wrappers, ...), each a thin + wrapper around one call; +- entry points (being added) that replace direct access to internals: registry() for + deployment lookups, settings() for configuration reads, health() for + cooldowns, blocks and usage, select() to pick a deployment without calling + it, apply() for runtime changes, observe() for routing events, + register_plugin() for strategies and sub-routers, and retire() to release + the router. + +Everything else on the class is internal. The members outside code uses today +are listed in tests/code_coverage_tests/router_surface_allowlist.json, and a CI +check (router_surface_check.py next to it) fails on any new outside use, so the +list only shrinks while callers move to the entry points. The same interface +is documented for the Rust core in litellm-rust/crates/router-core. +""" + import asyncio import contextlib import copy @@ -85,6 +108,15 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.otel.routing import routing_decision_attributes from litellm.integrations.otel.runtime import phase_attributes, phase_event, phase_span from litellm.litellm_core_utils.asyncify import run_async_function +from litellm.litellm_core_utils.batch_utils import ( + batch_create_routing, + deployment_supported_endpoints, + get_router_metadata_variable_name, + is_batch_retrieve_call_type, + replace_model_in_jsonl, + should_replace_model_in_jsonl, + without_deployments_that_cannot_batch, +) from litellm.litellm_core_utils.completion_timeout import CompletionTimeout from litellm.litellm_core_utils.core_helpers import ( coerce_token_limit, @@ -148,6 +180,15 @@ from litellm.llms.openai_like.model_info import ( get_openai_compatible_model_info, ) from litellm.responses.bridge_check import keep_registered_mode +from litellm.router_strategy.auto_router.auto_router_model_naming import ( + AUTO_ROUTER_MODEL_PREFIX, + STRATEGY_ROUTER_PARAM_FIELDS, + GatedAutoRouterCapability, + capability_limit_violation, + claimed_capabilities, + classify_strategy_router_model, + count_capability_routers, +) from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.router_strategy.complexity_router.context_compaction import ( @@ -187,24 +228,6 @@ from litellm.router_utils.add_retry_fallback_headers import ( replace_complexity_router_headers, response_total_token_count, ) -from litellm.router_utils.auto_router_model_naming import ( - AUTO_ROUTER_MODEL_PREFIX, - STRATEGY_ROUTER_PARAM_FIELDS, - GatedAutoRouterCapability, - capability_limit_violation, - claimed_capabilities, - classify_strategy_router_model, - count_capability_routers, -) -from litellm.router_utils.batch_utils import ( - batch_create_routing, - deployment_supported_endpoints, - get_router_metadata_variable_name, - is_batch_retrieve_call_type, - replace_model_in_jsonl, - should_replace_model_in_jsonl, - without_deployments_that_cannot_batch, -) from litellm.router_utils.client_initalization_utils import InitalizeCachedClient, MaxParallelRequestsLimit from litellm.router_utils.clientside_credential_handler import ( get_dynamic_litellm_params, @@ -16271,7 +16294,7 @@ class Router: to the deployment that actually served the request. Every attempt therefore writes or clears, never just writes. """ - from litellm.router_utils.baseline_request import capture_baseline_parameters + from litellm.router_strategy.auto_router.baseline_request import capture_baseline_parameters from litellm.types.router import BaselineRouteStamp phase_attributes(routing_decision_attributes(routing_decision)) diff --git a/litellm/router_utils/auto_router_model_naming.py b/litellm/router_strategy/auto_router/auto_router_model_naming.py similarity index 100% rename from litellm/router_utils/auto_router_model_naming.py rename to litellm/router_strategy/auto_router/auto_router_model_naming.py diff --git a/litellm/router_utils/auto_router_tuning_baseline.py b/litellm/router_strategy/auto_router/auto_router_tuning_baseline.py similarity index 100% rename from litellm/router_utils/auto_router_tuning_baseline.py rename to litellm/router_strategy/auto_router/auto_router_tuning_baseline.py diff --git a/litellm/router_utils/baseline_request.py b/litellm/router_strategy/auto_router/baseline_request.py similarity index 100% rename from litellm/router_utils/baseline_request.py rename to litellm/router_strategy/auto_router/baseline_request.py diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index eab232eae18..1d502dd8d95 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -36,12 +36,12 @@ from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.caching.redis_cache import RedisCache, RedisPipelineIncrementOperation, log_redis_failure from litellm.integrations.custom_logger import CustomLogger, Span +from litellm.litellm_core_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, ) from litellm.litellm_core_utils.duration_parser import duration_in_seconds_from from litellm.router_strategy.tag_based_routing import get_tags_from_request_kwargs -from litellm.router_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once from litellm.router_utils.cooldown_callbacks import ( get_prometheus_logger_from_callbacks, ) diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 6f3e0936641..8e0c3d6a7e7 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -10,7 +10,7 @@ from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.caching.redis_cache import log_redis_failure from litellm.integrations.custom_logger import CustomLogger -from litellm.router_utils.batch_utils import is_batch_retrieve_call_type +from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60 diff --git a/litellm/router_strategy/lowest_cost.py b/litellm/router_strategy/lowest_cost.py index 5894da9cb4f..17d03fd4ff8 100644 --- a/litellm/router_strategy/lowest_cost.py +++ b/litellm/router_strategy/lowest_cost.py @@ -10,7 +10,7 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger -from litellm.router_utils.batch_utils import is_batch_retrieve_call_type +from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type from litellm.router_utils.deployment_pricing import deployment_pricing_models, str_keyed diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index 4e994d3fdc6..c352934db6c 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -13,9 +13,9 @@ from litellm import ModelResponse, token_counter, verbose_logger from litellm._internal_context import with_service_target from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs, safe_divide_seconds from litellm.responses.utils import ResponseAPILoggingUtils -from litellm.router_utils.batch_utils import is_batch_retrieve_call_type from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import LiteLLMPydanticObjectBase diff --git a/litellm/router_strategy/lowest_tpm_rpm.py b/litellm/router_strategy/lowest_tpm_rpm.py index 2d373e0c266..9636d7f009d 100644 --- a/litellm/router_strategy/lowest_tpm_rpm.py +++ b/litellm/router_strategy/lowest_tpm_rpm.py @@ -9,7 +9,7 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger -from litellm.router_utils.batch_utils import is_batch_retrieve_call_type +from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type from litellm.types.utils import LiteLLMPydanticObjectBase from litellm.utils import print_verbose diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index 527ea830d55..0719b30d36d 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -15,8 +15,8 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_logger, verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs -from litellm.router_utils.batch_utils import is_batch_retrieve_call_type from litellm.types.router import RouterErrors from litellm.types.utils import LiteLLMPydanticObjectBase, StandardLoggingPayload from litellm.utils import get_utc_datetime, print_verbose diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 0a77ad5b9ea..6d1de1f9f6c 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -13,6 +13,7 @@ from pydantic import TypeAdapter import litellm from litellm._logging import verbose_router_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.batch_utils import get_router_metadata_variable_name from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs, safe_deep_copy from litellm.litellm_core_utils.credential_accessor import reject_api_key_secret_reference from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider @@ -22,7 +23,6 @@ from litellm.router_utils.add_retry_fallback_headers import ( add_fallback_headers_to_response, get_fallback_error_info, ) -from litellm.router_utils.batch_utils import get_router_metadata_variable_name from litellm.router_utils.cooldown_handlers import ( _first_present, # pyright: ignore[reportPrivateUsage] - shared internal helper, used across router_utils cast_exception_status_to_int, diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 4c9afd26dad..9fb3633d53c 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -187,13 +187,13 @@ ignored_function_names = [ def main(): router_file = [ "./litellm/router.py", - "./litellm/router_utils/batch_utils.py", + "./litellm/litellm_core_utils/batch_utils.py", "./litellm/router_utils/pattern_match_deployments.py", ] # router_file = [ # "../../litellm/router.py", # "../../litellm/router_utils/pattern_match_deployments.py", - # "../../litellm/router_utils/batch_utils.py", + # "../../litellm/litellm_core_utils/batch_utils.py", # ] ## LOCAL TESTING tests_dir = "./tests/" # Update this path if your tests directory is located elsewhere # tests_dir = "../../tests/" # LOCAL TESTING diff --git a/tests/code_coverage_tests/router_surface_allowlist.json b/tests/code_coverage_tests/router_surface_allowlist.json new file mode 100644 index 00000000000..428fd9b4011 --- /dev/null +++ b/tests/code_coverage_tests/router_surface_allowlist.json @@ -0,0 +1,211 @@ +{ + "_doc": "Router members and litellm.router_utils modules used outside the router. Shrink-only: see tests/code_coverage_tests/router_surface_check.py.", + "members_outside": [ + "_get_client", + "_zero_cost_cache", + "a_add_message", + "aadapter_completion", + "aanthropic_messages", + "abatch_completion", + "abatch_completion_fastest_response", + "acancel_batch", + "acancel_fine_tuning_job", + "acompletion", + "acreate_assistants", + "acreate_batch", + "acreate_file", + "acreate_fine_tuning_job", + "acreate_thread", + "adaptive_routers", + "adecisions", + "adelete_assistant", + "aembedding", + "afile_content", + "afile_delete", + "aget_assistants", + "aget_messages", + "aget_responses", + "aget_thread", + "alist_batches", + "alist_fine_tuning_jobs", + "allm_passthrough_route", + "amoderation", + "aocr", + "are_all_deployments_blocked", + "arefresh_model_info", + "aretrieve_batch", + "aretrieve_fine_tuning_job", + "arm_routing_read_prefetch", + "arun_thread", + "async_get_available_deployment", + "async_get_fully_unhealthy_model_names", + "async_get_healthy_deployments", + "async_pre_routing_hook", + "auto_routers", + "avector_store_create", + "avector_store_search", + "background_health_check_model_groups", + "bind_to_request_logging", + "cache", + "cache_responses", + "cached_deployment_model_info", + "cached_model_group_info", + "completion", + "complexity_routers", + "config_deployments", + "content_policy_fallbacks", + "context_window_fallbacks", + "cooldown_time", + "default_deployment", + "default_litellm_params", + "delete_deployment", + "deployment_latency_map", + "deployment_names", + "deployment_usable_by_team", + "deployments_for_request", + "deployments_named", + "discard", + "embedding", + "enable_health_check_routing", + "fallbacks", + "generate_model_id", + "get_all_deployments", + "get_available_deployment_for_pass_through", + "get_available_guardrail", + "get_candidate_model_ids_for_route", + "get_configured_display_name", + "get_configured_mode", + "get_configured_service_tiers", + "get_configured_token_limits", + "get_credential_deployment", + "get_deployment", + "get_deployment_by_model_group_name", + "get_deployment_credentials", + "get_deployment_credentials_with_provider", + "get_deployment_model_info", + "get_discovered_model_info", + "get_fully_blocked_model_names", + "get_model_access_groups", + "get_model_from_alias", + "get_model_group_info", + "get_model_group_usage", + "get_model_ids", + "get_model_info", + "get_model_list", + "get_model_list_from_model_alias", + "get_model_list_from_routing_groups", + "get_model_list_of_routed_group", + "get_model_listing_info", + "get_model_names", + "get_remaining_model_group_usage", + "get_routable_deployments", + "get_routable_upstream_model", + "get_router_deployment_budget_limiter", + "get_routing_group", + "get_settings", + "get_valid_args", + "guardrail_list", + "has_model_id", + "health_check_ignore_transient_errors", + "health_state_cache", + "is_deployment_blocked", + "is_model_access_group_for_wildcard_route", + "is_recognized_model", + "kubernetes_pod_discovery", + "map_team_model", + "model_group_alias", + "model_list", + "model_name_to_deployment_indices", + "model_names", + "pattern_router", + "priority_header_update", + "prompt_caching_pin_cache", + "provider_budget_config", + "quality_routers", + "resolve_model_name_from_model_id", + "resolve_served_model_name", + "resolved_litellm_models", + "routable_model_group", + "router_general_settings", + "routing_groups", + "routing_plugins", + "schedule_acompletion", + "search_tools", + "serve_through_request_logging", + "should_include_deployment", + "split_request_priority", + "team_model_to_deployment_indices", + "team_pattern_routers", + "team_ptu_ceiling", + "team_public_model_names", + "team_public_model_names_of", + "update_redis_cache", + "update_settings", + "upsert_deployment", + "wildcard_route_deployments", + "wildcard_routed_models" + ], + "members_router_family": [ + "_count_pre_call_check_tokens", + "_pre_call_checks", + "_update_kwargs_before_fallbacks", + "acompletion", + "adecisions", + "aembedding", + "allowed_fails", + "allowed_fails_policy", + "aresponses", + "async_function_with_fallbacks", + "async_get_healthy_deployments", + "cache", + "cached_deployment_model_info", + "cooldown_cache", + "cooldown_time", + "default_max_parallel_requests", + "deployments_named", + "disable_cooldowns", + "embedding", + "enable_pre_call_checks", + "enable_tag_filtering", + "explicit_allowed_fails", + "fallback_access_check", + "fallback_budget_check", + "fallbacks", + "get_all_deployments", + "get_allowed_fails_from_policy", + "get_deployment", + "get_model_group", + "get_model_ids", + "get_model_info", + "get_model_list", + "get_router_model_info", + "log_retry", + "model_list", + "provider_default_deployment_ids", + "routing_group_has_alternatives", + "routing_strategy", + "search_tools", + "set_custom_routing_strategy", + "slack_alerting_logger", + "tag_filtering_match_any", + "tag_routing_prefix", + "team_model_has_alternatives" + ], + "router_utils_imports_outside": [ + "litellm.router_utils.access_windows", + "litellm.router_utils.add_retry_fallback_headers", + "litellm.router_utils.clientside_credential_handler", + "litellm.router_utils.common_utils", + "litellm.router_utils.cooldown_handlers", + "litellm.router_utils.deployment_pricing", + "litellm.router_utils.fallback_event_handlers", + "litellm.router_utils.get_retry_from_policy", + "litellm.router_utils.health_state_cache", + "litellm.router_utils.pattern_match_deployments", + "litellm.router_utils.ptu_shares", + "litellm.router_utils.reasoning_effort_capability", + "litellm.router_utils.router_callbacks.track_deployment_metrics", + "litellm.router_utils.routing_groups", + "litellm.router_utils.search_api_router" + ] +} diff --git a/tests/code_coverage_tests/router_surface_check.py b/tests/code_coverage_tests/router_surface_check.py new file mode 100644 index 00000000000..0b35bac5240 --- /dev/null +++ b/tests/code_coverage_tests/router_surface_check.py @@ -0,0 +1,235 @@ +""" +Fence for the Router's public surface. + +Code outside the router may use only the Router members listed in +router_surface_allowlist.json, and may import only the litellm.router_utils +modules listed there. The lists hold what is used today and may only shrink: + +- a member or module used outside the router that is not listed fails the + check (use an existing entry point instead, or discuss widening the surface); +- a listed member or module that nothing outside the router uses any more + also fails, so the entry is removed in the same change. + +"Outside" means every non-test file under litellm/ and enterprise/ except +litellm/router.py and the router's own modules. Uses from the router family +(litellm/router_utils, litellm/router_strategy) are listed separately, because +those modules become private parts of the router core in a later step. + +Member uses are found the same way the boundary survey found them: attribute +access on a receiver whose name looks like a router (llm_router, +litellm_router_instance, router, ...) or on the Router class, plus +getattr/setattr/hasattr with a constant name on such a receiver. Members +include the container file endpoints Router.__init__ registers by name. + +Run: python tests/code_coverage_tests/router_surface_check.py [--write] +--write regenerates the allowlist from the current tree (for shrinking it). +""" + +import ast +import json +import os +import re +import sys +from typing import Dict, List, Set + +ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +ALLOWLIST = os.path.join(os.path.dirname(__file__), "router_surface_allowlist.json") + +ROUTER_FILES = { + "litellm/router.py", +} +ROUTER_CLASS_FILES = ["litellm/router.py"] +ROUTER_CLASS_NAMES = {"Router"} + +SCAN_ROOTS = ["litellm", "enterprise"] +SKIP_DIRS = { + ".git", + "node_modules", + ".venv", + "__pycache__", + "litellm-rust", + "ui", + "docs", + "cookbook", + "tests", + "test", +} + +RECEIVER = re.compile( + r"(^|[._])(llm_router|router|litellm_router|litellm_router_instance|router_instance" + r"|_router|proxy_router|llm_router_instance|router_obj|mock_router|Router)$" +) +# Router.__init__ registers these with setattr(self, name, ...), so they never appear as self.. +CONTAINER_FILE_ENDPOINTS = "litellm/containers/endpoints.json" +DYNAMIC_FUNCS = {"getattr", "setattr", "hasattr"} + + +def _area(rel: str) -> str: + if rel.startswith(("litellm/router_utils/", "litellm/router_strategy/")): + return "router_family" + return "outside" + + +def _is_test(rel: str) -> bool: + name = os.path.basename(rel) + return "/tests/" in rel or name.startswith("test_") or name.endswith("_test.py") + + +def router_members() -> Set[str]: + members: Set[str] = set() + for rel in ROUTER_CLASS_FILES: + with open(os.path.join(ROOT, rel), encoding="utf-8") as f: + tree = ast.parse(f.read()) + for cls in tree.body: + if not isinstance(cls, ast.ClassDef) or cls.name not in ROUTER_CLASS_NAMES: + continue + for node in cls.body: + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + members.add(node.name) + elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name): + members.add(node.target.id) + elif isinstance(node, ast.Assign): + members.update(t.id for t in node.targets if isinstance(t, ast.Name)) + for node in ast.walk(cls): + if ( + isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id == "self" + and isinstance(node.ctx, ast.Store) + ): + members.add(node.attr) + with open(os.path.join(ROOT, CONTAINER_FILE_ENDPOINTS), encoding="utf-8") as f: + for endpoint in json.load(f)["endpoints"]: + members.update((endpoint["name"], endpoint["async_name"])) + return {m for m in members if not (m.startswith("__") and m.endswith("__"))} + + +def _python_files(): + for top in SCAN_ROOTS: + base = os.path.join(ROOT, top) + for dirpath, dirnames, filenames in os.walk(base): + dirnames[:] = [d for d in dirnames if d not in SKIP_DIRS] + for fn in filenames: + if fn.endswith(".py"): + path = os.path.join(dirpath, fn) + yield os.path.relpath(path, ROOT).replace(os.sep, "/"), path + + +def scan() -> Dict[str, Dict[str, Set[str]]]: + """Return {"members": {area: {member: files}}, "router_utils_imports": {module: files}}.""" + members = router_members() + uses: Dict[str, Dict[str, Set[str]]] = {"outside": {}, "router_family": {}} + imports: Dict[str, Set[str]] = {} + for rel, path in _python_files(): + if rel in ROUTER_FILES or _is_test(rel): + continue + with open(path, encoding="utf-8") as f: + src = f.read() + has_router = "router" in src.lower() + if not has_router: + continue + try: + tree = ast.parse(src) + except SyntaxError: + continue + area = _area(rel) + for name in member_uses(tree, members): + uses[area].setdefault(name, set()).add(rel) + if area == "outside": + for node in ast.walk(tree): + for mod in _router_utils_modules(node): + imports.setdefault(mod, set()).add(rel) + return {"members": uses, "router_utils_imports": imports} + + +def member_uses(tree: ast.AST, members: Set[str]) -> Set[str]: + """The Router members a module uses, on an instance or on the Router class itself.""" + used: Set[str] = set() + for node in ast.walk(tree): + if isinstance(node, ast.Attribute) and node.attr in members: + if RECEIVER.search(ast.unparse(node.value)): + used.add(node.attr) + elif ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id in DYNAMIC_FUNCS + and len(node.args) >= 2 + and isinstance(node.args[1], ast.Constant) + and node.args[1].value in members + and RECEIVER.search(ast.unparse(node.args[0])) + ): + used.add(node.args[1].value) + return used + + +def _router_utils_modules(node: ast.AST) -> List[str]: + if isinstance(node, ast.ImportFrom) and node.module: + if node.module == "litellm.router_utils": + return ["litellm.router_utils." + a.name for a in node.names] + if node.module.startswith("litellm.router_utils."): + return [node.module] + if isinstance(node, ast.Import): + return [a.name for a in node.names if a.name.startswith("litellm.router_utils.")] + return [] + + +def build_allowlist(found) -> dict: + return { + "_doc": ( + "Router members and litellm.router_utils modules used outside the router. " + "Shrink-only: see tests/code_coverage_tests/router_surface_check.py." + ), + "members_outside": sorted(found["members"]["outside"]), + "members_router_family": sorted(found["members"]["router_family"]), + "router_utils_imports_outside": sorted(found["router_utils_imports"]), + } + + +def compare(allow: dict, found) -> List[str]: + errors: List[str] = [] + sections = [ + ("members_outside", found["members"]["outside"], "Router member"), + ("members_router_family", found["members"]["router_family"], "Router member (router family)"), + ("router_utils_imports_outside", found["router_utils_imports"], "litellm.router_utils import"), + ] + for key, used, what in sections: + listed = set(allow.get(key, [])) + for name in sorted(set(used) - listed): + files = ", ".join(sorted(used[name])) + errors.append( + f"new {what} outside the router: {name!r} in {files}. The router surface may only " + f"shrink: use an entry point instead." + ) + for name in sorted(listed - set(used)): + errors.append( + f"{what} {name!r} is listed in {os.path.basename(ALLOWLIST)} [{key}] but nothing uses it " + f"any more: remove it from the list." + ) + return errors + + +def main(argv: List[str]) -> int: + found = scan() + if "--write" in argv: + with open(ALLOWLIST, "w", encoding="utf-8") as f: + json.dump(build_allowlist(found), f, indent=2) + f.write("\n") + print(f"wrote {ALLOWLIST}") + return 0 + with open(ALLOWLIST, encoding="utf-8") as f: + allow = json.load(f) + errors = compare(allow, found) + for e in errors: + print(e) + if errors: + return 1 + print( + f"router surface ok: {len(allow['members_outside'])} members outside, " + f"{len(allow['members_router_family'])} from the router family, " + f"{len(allow['router_utils_imports_outside'])} router_utils modules imported outside" + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1:])) diff --git a/tests/code_coverage_tests/unbounded_in_baseline.txt b/tests/code_coverage_tests/unbounded_in_baseline.txt index 1da964681a9..d8c6e2241f4 100644 --- a/tests/code_coverage_tests/unbounded_in_baseline.txt +++ b/tests/code_coverage_tests/unbounded_in_baseline.txt @@ -129,4 +129,4 @@ litellm/proxy/utils.py prefetch_config_params prisma param_name.in `param_names` litellm/repositories/daily_activity_repository.py DailyActivityRepository.daily_rows prisma [scope.entity_id_field].in `list(scope.entity_ids)` 0 litellm/repositories/daily_activity_repository.py DailyActivityRepository.daily_rows prisma api_key.in `list(scope.api_keys)` 0 litellm/repositories/daily_activity_repository.py DailyActivityRepository.daily_rows prisma not.in `list(scope.exclude_entity_ids)` 0 -litellm/router_utils/auto_router_model_naming.py raw-sql classifier_type.IN `IN ({_LLM_CLASSIFIER_TYPES_SQL})` 0 +litellm/router_strategy/auto_router/auto_router_model_naming.py raw-sql classifier_type.IN `IN ({_LLM_CLASSIFIER_TYPES_SQL})` 0 diff --git a/tests/unit/router_utils/test_batch_utils.py b/tests/unit/litellm_core_utils/test_batch_utils.py similarity index 97% rename from tests/unit/router_utils/test_batch_utils.py rename to tests/unit/litellm_core_utils/test_batch_utils.py index 79fe86a37c0..cd37d635c43 100644 --- a/tests/unit/router_utils/test_batch_utils.py +++ b/tests/unit/litellm_core_utils/test_batch_utils.py @@ -7,7 +7,7 @@ from typing import Dict, Final, List import pytest import litellm -from litellm.router_utils.batch_utils import ( +from litellm.litellm_core_utils.batch_utils import ( InMemoryFile, get_router_metadata_variable_name, parse_jsonl_with_embedded_newlines, @@ -116,7 +116,7 @@ def test_router_metadata_variable_name(): def test_non_json_input(): """Test that replace_model_in_jsonl returns original content for non-JSON input""" - from litellm.router_utils.batch_utils import replace_model_in_jsonl + from litellm.litellm_core_utils.batch_utils import replace_model_in_jsonl # Test with non-JSON string non_json_str = "This is not a JSON string" @@ -138,7 +138,7 @@ def test_non_json_input(): def test_should_replace_model_in_jsonl(): """Test that should_replace_model_in_jsonl returns the correct value""" - from litellm.router_utils.batch_utils import should_replace_model_in_jsonl + from litellm.litellm_core_utils.batch_utils import should_replace_model_in_jsonl assert should_replace_model_in_jsonl(purpose="batch") is True assert should_replace_model_in_jsonl(purpose="batch", passthrough=True) is False @@ -340,7 +340,7 @@ def test_replace_model_in_jsonl_with_embedded_newlines(): def test_is_batch_retrieve_call_type_matches_only_batch_retrieves(): - from litellm.router_utils.batch_utils import is_batch_retrieve_call_type + from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type from litellm.types.utils import CallTypes assert is_batch_retrieve_call_type(CallTypes.aretrieve_batch.value) is True @@ -389,7 +389,8 @@ def setup_and_teardown(): ) def test_batch_charge_identity_requires_a_cost_bearing_retrieval(call_type: str, cost: float, eligible: bool) -> None: from openai.types import Batch - from litellm.router_utils.batch_utils import batch_id_to_charge_once + + from litellm.litellm_core_utils.batch_utils import batch_id_to_charge_once batch: Final = Batch( id="batch-budget", completion_window="24h", created_at=1, endpoint="/v1/chat/completions", @@ -401,7 +402,7 @@ def test_batch_charge_identity_requires_a_cost_bearing_retrieval(call_type: str, def test_batch_charge_markers_keep_counters_and_batches_independent() -> None: from litellm.caching.in_memory_cache import InMemoryCache - from litellm.router_utils.batch_utils import batch_charged_once_marker_key + from litellm.litellm_core_utils.batch_utils import batch_charged_once_marker_key cache: Final = InMemoryCache() cache.set_cache(batch_charged_once_marker_key("provider-spend", "first-batch"), 1) diff --git a/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py index 748b291d5ff..c13e3d5812b 100644 --- a/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py @@ -170,7 +170,7 @@ async def test_preview_and_serving_share_default_model_resolution( explicit_default: str | None, expected: str, ): - from litellm.router_utils.auto_router_model_naming import validate_complexity_router_config_write + from litellm.router_strategy.auto_router.auto_router_model_naming import validate_complexity_router_config_write from litellm.types.management_endpoints.auto_router_endpoints import ComplexityRouterConfigValidationRequest config: Final = { @@ -1240,7 +1240,7 @@ class TestAutoRouterBenchmarks: def test_the_listed_kinds_match_the_router_types_traffic_can_record(self): from typing import get_args, get_type_hints - from litellm.router_utils.auto_router_model_naming import StrategyRouterKind + from litellm.router_strategy.auto_router.auto_router_model_naming import StrategyRouterKind from litellm.types.utils import StandardLoggingRoutingDecision readonly_router_type: Final = get_type_hints(StandardLoggingRoutingDecision, include_extras=True)["router_type"] diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index fc6749aed0f..437d948f8f5 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -1,6 +1,6 @@ -import inspect import asyncio import contextlib +import inspect import json import socket from collections.abc import Iterator, Mapping @@ -1461,10 +1461,10 @@ class TestDeleteModelClearsRouterRegistry: add_deployment never restores, so an unguarded pop would make it permanently unroutable (the same cross-tenant DoS clear_cache was hardened against). """ + from litellm.proxy.management_endpoints.model_management_endpoints import ModelInfoDelete from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) - from litellm.proxy.management_endpoints.model_management_endpoints import ModelInfoDelete model_id = "regular-del-1" admin_user = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) @@ -2739,11 +2739,11 @@ class TestAddAndDeleteModelLifecycle: - Delete same model again → raises (model not found) """ from litellm.proxy.management_endpoints.model_management_endpoints import ( + ModelInfoDelete, add_new_model, - delete_model as delete_model_endpoint, ) from litellm.proxy.management_endpoints.model_management_endpoints import ( - ModelInfoDelete, + delete_model as delete_model_endpoint, ) model_id = "lifecycle-test-model-123" @@ -2833,6 +2833,8 @@ class TestDeleteTeamBYOKModelGhost: async def test_delete_strips_public_name_and_refreshes_cache(self): from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) @@ -2916,6 +2918,8 @@ class TestDeleteTeamBYOKModelGhost: alias cleanup (delete_team_model_alias), preserving legacy behavior.""" from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) @@ -2984,6 +2988,8 @@ class TestDeleteTeamBYOKModelGhost: team.models when one replica is deleted but a sibling still backs it.""" from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) @@ -3065,6 +3071,8 @@ class TestDeleteTeamBYOKModelGhost: deployment still serves it.""" from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) @@ -3145,6 +3153,8 @@ class TestDeleteTeamBYOKModelGhost: the alias would break routing that works.""" from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) @@ -3252,6 +3262,8 @@ class TestDeleteModelTeamAuth: async def test_proxy_admin_can_delete_model_when_team_deleted(self): from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) @@ -3287,6 +3299,8 @@ class TestDeleteModelTeamAuth: """A missing team must never let a non-admin delete the orphan (no fail-open).""" from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) from litellm.proxy.proxy_server import ProxyException @@ -3322,6 +3336,8 @@ class TestDeleteModelTeamAuth: """The auth check must not add a redundant team query on the live-team path.""" from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) from litellm.proxy.proxy_server import ProxyException @@ -4392,7 +4408,6 @@ class TestModelInfoCostMapEchoFilter: def test_echoed_cost_map_metadata_is_not_persisted(self): import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -4418,7 +4433,6 @@ class TestModelInfoCostMapEchoFilter: def test_an_edited_value_survives_the_echo_filter(self): import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -4483,7 +4497,6 @@ class TestModelInfoCostMapEchoFilter: def test_metadata_without_a_cost_map_key_is_persisted(self): import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo from litellm.types.utils import echoed_cost_map_fields @@ -4513,7 +4526,6 @@ class TestModelInfoCostMapEchoFilter: def test_a_stored_mode_survives_an_echoed_save(self): import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -4536,7 +4548,6 @@ class TestModelInfoCostMapEchoFilter: def test_resetting_an_override_to_the_cost_map_value_removes_it(self): import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -4560,7 +4571,6 @@ class TestModelInfoCostMapEchoFilter: def test_reset_is_recognised_after_the_router_registered_the_override(self, monkeypatch: pytest.MonkeyPatch): import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -4586,7 +4596,6 @@ class TestModelInfoCostMapEchoFilter: from types import MappingProxyType import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -4611,7 +4620,6 @@ class TestModelInfoCostMapEchoFilter: def test_echo_is_compared_against_the_deployments_lookup_not_the_key(self): import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -4653,7 +4661,6 @@ class TestModelInfoCostMapEchoFilter: def test_base_model_wins_over_litellm_params_model_for_the_lookup(self): import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -4681,7 +4688,6 @@ class TestModelInfoCostMapEchoFilter: def test_encrypted_stored_model_is_decrypted_for_the_lookup(self, monkeypatch): import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -6431,7 +6437,7 @@ class TestStrategyRouterWriteValidation: AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY, _auto_router_capability_slot, ) - from litellm.router_utils.auto_router_model_naming import gated_capabilities_of + from litellm.router_strategy.auto_router.auto_router_model_naming import gated_capabilities_of capabilities: Final = gated_capabilities_of(effective_params) @@ -6568,7 +6574,7 @@ class TestStrategyRouterWriteValidation: from fastapi import HTTPException from litellm.proxy.management_endpoints.model_management_endpoints import _auto_router_capability_slot - from litellm.router_utils.auto_router_tuning_baseline import snapshot_tuning_baselines + from litellm.router_strategy.auto_router.auto_router_tuning_baseline import snapshot_tuning_baselines configs = { "_TUNED_A": self._TUNED_A, @@ -6624,7 +6630,7 @@ class TestStrategyRouterWriteValidation: from litellm.proxy import proxy_server from litellm.proxy.management_endpoints.model_management_endpoints import _auto_router_capability_slot - from litellm.router_utils.auto_router_tuning_baseline import snapshot_tuning_baselines + from litellm.router_strategy.auto_router.auto_router_tuning_baseline import snapshot_tuning_baselines fake = self._FakeDb([]) fake.tx_obj.litellm_proxymodeltable.find_many = AsyncMock( @@ -6669,7 +6675,7 @@ class TestStrategyRouterWriteValidation: row id before the slot is entered, so a second tuned router is refused before its DB write.""" from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model - from litellm.router_utils.auto_router_tuning_baseline import snapshot_tuning_baselines + from litellm.router_strategy.auto_router.auto_router_tuning_baseline import snapshot_tuning_baselines admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) baselines = snapshot_tuning_baselines([self._db_router_row("a", self._TUNED_A)]) @@ -8980,7 +8986,10 @@ class TestTeamMemberAutoRouterWrites: from fastapi import HTTPException from litellm.proxy.management_endpoints.model_management_endpoints import _auto_router_capability_slot - from litellm.proxy.management_helpers.auto_router_permissions import MemberAutoRouterWrite, validate_member_auto_router_config + from litellm.proxy.management_helpers.auto_router_permissions import ( + MemberAutoRouterWrite, + validate_member_auto_router_config, + ) row: Final = self._row() database: Final = self._database(self._team(), row) @@ -10677,6 +10686,8 @@ class TestFederationGateScopesToWhatTheWriteTouches: and the team that owns the model must be able to take it off their page.""" from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) diff --git a/tests/unit/proxy/management_helpers/test_auto_router_availability.py b/tests/unit/proxy/management_helpers/test_auto_router_availability.py index a88dcd7f66d..0eaae60b4ea 100644 --- a/tests/unit/proxy/management_helpers/test_auto_router_availability.py +++ b/tests/unit/proxy/management_helpers/test_auto_router_availability.py @@ -1,16 +1,15 @@ from collections.abc import Mapping -from typing import Final from types import SimpleNamespace - -from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper +from typing import Final import pytest +from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper from litellm.proxy.management_helpers.auto_router_availability import ( auto_router_availability, build_auto_router_catalog, ) -from litellm.router_utils.auto_router_tuning_baseline import snapshot_tuning_baselines +from litellm.router_strategy.auto_router.auto_router_tuning_baseline import snapshot_tuning_baselines def deployment( diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index 11a44a0ec16..8b1d7c1222a 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/tests/unit/proxy/proxy_server/test_lifecycle.py @@ -33,9 +33,9 @@ from pydantic import BaseModel from typing_extensions import TypedDict import litellm.proxy.proxy_server as ps +from litellm.proxy.management_helpers.audit_logs import track_audit_task from litellm.proxy.proxy_server import ( ProxyStartupEvent, - initialize_shared_aiohttp_session, _resolve_pydantic_type, _resolve_typed_dict_type, cleanup_router_config_variables, @@ -43,6 +43,7 @@ from litellm.proxy.proxy_server import ( get_litellm_model_info, initialize, initialize_from_worker_config, + initialize_shared_aiohttp_session, load_from_azure_key_vault, proxy_shutdown_event, proxy_startup_event, @@ -50,8 +51,6 @@ from litellm.proxy.proxy_server import ( save_worker_config, ) -from litellm.proxy.management_helpers.audit_logs import track_audit_task - from .conftest import normalize # --------------------------------------------------------------------------- @@ -1165,7 +1164,7 @@ def test_proxy_startup_event_warns_for_global_budget_without_database(): @pytest.mark.asyncio async def test_tuning_baseline_v3_is_created_alongside_the_legacy_row(): - from litellm.router_utils.auto_router_tuning_baseline import DEFAULT_TUNING_FINGERPRINT + from litellm.router_strategy.auto_router.auto_router_tuning_baseline import DEFAULT_TUNING_FINGERPRINT prisma_client = MagicMock() prisma_client.db.litellm_config.find_unique = AsyncMock(return_value=None) @@ -1186,7 +1185,10 @@ async def test_tuning_baseline_v3_is_created_alongside_the_legacy_row(): @pytest.mark.asyncio async def test_scorer_baseline_upgrade_preserves_existing_routers_and_is_not_refreshed_on_restart(): - from litellm.router_utils.auto_router_tuning_baseline import mutable_tuned_identities, snapshot_tuning_baselines + from litellm.router_strategy.auto_router.auto_router_tuning_baseline import ( + mutable_tuned_identities, + snapshot_tuning_baselines, + ) deployments = [ { diff --git a/tests/unit/router_utils/test_auto_router_model_naming.py b/tests/unit/router_strategy/test_auto_router_model_naming.py similarity index 99% rename from tests/unit/router_utils/test_auto_router_model_naming.py rename to tests/unit/router_strategy/test_auto_router_model_naming.py index 1cfa3aa5c91..44fcfdaf5be 100644 --- a/tests/unit/router_utils/test_auto_router_model_naming.py +++ b/tests/unit/router_strategy/test_auto_router_model_naming.py @@ -3,9 +3,7 @@ from typing import Final import pytest -from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets -from litellm.router_strategy.complexity_router.jev_classifier import DEFAULT_JEV_INSTRUCTIONS -from litellm.router_utils.auto_router_model_naming import ( +from litellm.router_strategy.auto_router.auto_router_model_naming import ( GATED_AUTO_ROUTER_CAPABILITIES, capability_limit_violation, carries_complexity_router_settings, @@ -18,6 +16,8 @@ from litellm.router_utils.auto_router_model_naming import ( validate_complexity_router_config_write, validate_strategy_router_model_write, ) +from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets +from litellm.router_strategy.complexity_router.jev_classifier import DEFAULT_JEV_INSTRUCTIONS COMPLEXITY_FIELDS = frozenset({"complexity_router_config"}) SEMANTIC_FIELDS = frozenset({"auto_router_config", "auto_router_default_model", "auto_router_embedding_model"}) @@ -596,7 +596,7 @@ def test_custom_classifier_prompt_capability(config: Mapping[str, object], expec ], ) def test_is_complexity_router_model(model: str | None, expected: bool) -> None: - from litellm.router_utils.auto_router_model_naming import is_complexity_router_model + from litellm.router_strategy.auto_router.auto_router_model_naming import is_complexity_router_model assert is_complexity_router_model(model) is expected diff --git a/tests/unit/router_utils/test_auto_router_tuning_baseline.py b/tests/unit/router_strategy/test_auto_router_tuning_baseline.py similarity index 99% rename from tests/unit/router_utils/test_auto_router_tuning_baseline.py rename to tests/unit/router_strategy/test_auto_router_tuning_baseline.py index d0d98ce7042..4ee219a7784 100644 --- a/tests/unit/router_utils/test_auto_router_tuning_baseline.py +++ b/tests/unit/router_strategy/test_auto_router_tuning_baseline.py @@ -7,8 +7,7 @@ from typing import Final import pytest -from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig -from litellm.router_utils.auto_router_tuning_baseline import ( +from litellm.router_strategy.auto_router.auto_router_tuning_baseline import ( DEFAULT_TUNING_FINGERPRINT, HEURISTIC_V1_TUNING_FIELDS, heuristic_v1_router_fingerprint, @@ -19,6 +18,7 @@ from litellm.router_utils.auto_router_tuning_baseline import ( tuning_limit_violation, tuning_quota_violation, ) +from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig _TIERS = {"SIMPLE": "cheap", "MEDIUM": "mid", "COMPLEX": "strong", "REASONING": "top"} _ALT_TIERS = {**_TIERS, "COMPLEX": "other-strong"} diff --git a/tests/unit/router_utils/test_baseline_request.py b/tests/unit/router_strategy/test_baseline_request.py similarity index 95% rename from tests/unit/router_utils/test_baseline_request.py rename to tests/unit/router_strategy/test_baseline_request.py index 37d21597752..eb785a41d63 100644 --- a/tests/unit/router_utils/test_baseline_request.py +++ b/tests/unit/router_strategy/test_baseline_request.py @@ -1,6 +1,6 @@ from typing import Final -from litellm.router_utils.baseline_request import baseline_request, capture_baseline_parameters +from litellm.router_strategy.auto_router.baseline_request import baseline_request, capture_baseline_parameters def test_baseline_snapshot_owns_nested_caller_settings_and_overrides_routed_settings() -> None: diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index ca39c7f5c8f..018f93779a6 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -24,14 +24,6 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm import Router -from litellm.integrations.custom_logger import CustomLogger -from litellm.types.utils import StandardLoggingRoutingDecision -from litellm.router_utils.auto_router_model_naming import ( - CUSTOMIZATION_CAPABILITY, - GATED_AUTO_ROUTER_CAPABILITIES, - HEURISTIC_V2_CAPABILITY, - count_capability_routers, -) from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache @@ -42,7 +34,19 @@ from litellm.constants import ( RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY, ) +from litellm.integrations.custom_logger import CustomLogger +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.router import as_output_cap, stamped_client_ceiling +from litellm.router_strategy.auto_router.auto_router_model_naming import ( + CUSTOMIZATION_CAPABILITY, + GATED_AUTO_ROUTER_CAPABILITIES, + HEURISTIC_V2_CAPABILITY, + count_capability_routers, +) +from litellm.router_strategy.complexity_router.capability_classifier import ( + CAPABILITY_CLASSIFIER_SYSTEM_PROMPT, + CapabilityClassifierVerdict, +) from litellm.router_strategy.complexity_router.complexity_router import ( _CLASSIFICATION_CURRENT_MESSAGE_ONLY, _CLASSIFICATION_WITH_CONVERSATION, @@ -58,18 +62,14 @@ from litellm.router_strategy.complexity_router.complexity_router import ( classification_system_prompt, custom_tier_classification_prompt, ) -from litellm.router_strategy.complexity_router.capability_classifier import ( - CAPABILITY_CLASSIFIER_SYSTEM_PROMPT, - CapabilityClassifierVerdict, -) from litellm.router_strategy.complexity_router.config import ( - CapabilityCalibrationConfig, - CapabilityClassifierConfig, DEFAULT_CLASSIFICATION_RUBRIC, DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE, DEFAULT_COMPLEXITY_CONFIG, DEFAULT_TECHNICAL_KEYWORDS, TIER_SEVERITY_ORDER, + CapabilityCalibrationConfig, + CapabilityClassifierConfig, ClassificationRubric, ClassifierLLMConfig, ComplexityRouterConfig, @@ -83,10 +83,12 @@ from litellm.router_strategy.complexity_router.jev_classifier import ( JevUsage, ) from litellm.router_strategy.complexity_router.llm_v2 import LLM_V2_PROMPT_VERSION +from litellm.router_strategy.complexity_router.request_models import RequestComplexityRouterConfig from litellm.router_strategy.complexity_router.tier_predictor import ( TierGlobalStatistic, TrainedTierArtifact, ) +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.router import ( Deployment, LiteLLM_Params, @@ -94,10 +96,7 @@ from litellm.types.router import ( RouterErrors, TaggedPreRoutingStrategy, ) -from litellm.types.llms.openai import ResponsesAPIResponse -from litellm.router_strategy.complexity_router.request_models import RequestComplexityRouterConfig -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler - +from litellm.types.utils import StandardLoggingRoutingDecision requires_semantic_router = pytest.mark.skipif( sys.version_info >= (3, 14), reason="The semantic-router extra excludes Python 3.14" @@ -16299,7 +16298,6 @@ class TestTierHealthFailover: tests can prove the probe passes a kwargs copy and forwards the prompt arguments. """ import litellm as litellm_module - from litellm.types.router import RouterRateLimitError probed_kwargs = [] diff --git a/tests/unit/router_strategy/test_llm_v2.py b/tests/unit/router_strategy/test_llm_v2.py index 3fd2e8808e7..8c1e9e48661 100644 --- a/tests/unit/router_strategy/test_llm_v2.py +++ b/tests/unit/router_strategy/test_llm_v2.py @@ -4,11 +4,12 @@ from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest -import litellm from pydantic import ValidationError +import litellm from litellm import ModelResponse, Router from litellm.caching.dual_cache import DualCache +from litellm.router_strategy.auto_router.auto_router_model_naming import strategy_router_dependencies from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig, ComplexityTier from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets @@ -21,7 +22,6 @@ from litellm.router_strategy.complexity_router.llm_v2 import ( LLMV2Verdict, llm_v2_response_format, ) -from litellm.router_utils.auto_router_model_naming import strategy_router_dependencies from litellm.types.llms.openai import ResponsesAPIResponse diff --git a/tests/unit/test_private_usage_aliases.py b/tests/unit/test_private_usage_aliases.py index 1d9c2e87100..d0c5c8247d3 100644 --- a/tests/unit/test_private_usage_aliases.py +++ b/tests/unit/test_private_usage_aliases.py @@ -1969,7 +1969,7 @@ ALIAS_CASES: Final = ( ("litellm.router_utils.add_retry_fallback_headers", "HiddenParamsAsyncIteratorWrapper", "_inner", "inner", True), ("litellm.router_utils.add_retry_fallback_headers", "", "_HiddenParamsHost", "HiddenParamsHost", False), ( - "litellm.router_utils.batch_utils", + "litellm.litellm_core_utils.batch_utils", "", "_get_router_metadata_variable_name", "get_router_metadata_variable_name", diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index af8e7f53462..513136ae150 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -4,20 +4,21 @@ import gc import json import os import weakref -from datetime import datetime, timezone from collections.abc import Iterator, Mapping +from datetime import datetime, timezone from types import MethodType from typing import Final from unittest.mock import AsyncMock, MagicMock, patch -import litellm import pytest + +import litellm from litellm import Router from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY -from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook from litellm.types.litellm_params import RoutingStrategyName @@ -801,10 +802,11 @@ async def test_deployment_callback_on_failure(model_list): def test_deployment_callback_respects_cooldown_time(model_list): """Ensure per-model cooldown_time is honored even when exception headers are present.""" - import httpx import time from unittest.mock import patch + import httpx + router = Router(model_list=model_list) class FakeException(Exception): @@ -1811,8 +1813,8 @@ def test_get_allowed_fails_from_policy(model_list, exception_type, exception_nam def test_initialize_alerting(model_list): """Test if the 'initialize_alerting' function is working correctly""" - from litellm.types.router import AlertingConfig from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting + from litellm.types.router import AlertingConfig router = Router( model_list=model_list, alerting_config=AlertingConfig(webhook_url="test") @@ -2245,6 +2247,7 @@ def test_deployments_by_pattern(model_list): ) def test_pattern_match_deployment_set_model_name(user_request_model, model_name, litellm_model, expected_model): from re import Match + from litellm.router_utils.pattern_match_deployments import PatternMatchRouter pattern_router = PatternMatchRouter() @@ -2702,7 +2705,7 @@ def test_handle_clientside_credential_metadata_loading(model_list, function_name ) def test_handle_clientside_credential_metadata_variable_name(model_list, function_name, metadata_key): """Test that _handle_clientside_credential uses the correct metadata variable name based on function name""" - from litellm.router_utils.batch_utils import get_router_metadata_variable_name + from litellm.litellm_core_utils.batch_utils import get_router_metadata_variable_name router = Router(model_list=model_list) diff --git a/tests/unit/test_router/test_router_surface_check.py b/tests/unit/test_router/test_router_surface_check.py new file mode 100644 index 00000000000..a97fd029380 --- /dev/null +++ b/tests/unit/test_router/test_router_surface_check.py @@ -0,0 +1,58 @@ +""" +The router surface fence (tests/code_coverage_tests/router_surface_check.py) must see every way +outside code reaches a Router member, or new dependencies slip past the shrink-only allowlist. +""" + +import ast +import importlib.util +import pathlib + +import pytest + +_CHECK = pathlib.Path(__file__).parents[2] / "code_coverage_tests" / "router_surface_check.py" + + +@pytest.fixture(scope="module") +def surface_check(): + spec = importlib.util.spec_from_file_location("router_surface_check", _CHECK) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _uses(surface_check, src: str) -> set: + return surface_check.member_uses(ast.parse(src), surface_check.router_members()) + + +def test_instance_and_dynamic_access_are_uses(surface_check): + assert _uses(surface_check, "llm_router.get_model_list()\ngetattr(router, 'model_list')") == { + "get_model_list", + "model_list", + } + + +def test_class_level_calls_are_uses(surface_check): + assert _uses(surface_check, "litellm.Router.get_valid_args()") == {"get_valid_args"} + assert _uses(surface_check, "from litellm import Router\nRouter.get_valid_args()") == {"get_valid_args"} + + +def test_container_file_endpoints_registered_by_name_are_members(surface_check): + members = surface_check.router_members() + assert {"aupload_container_file", "upload_container_file", "alist_container_files"} <= members + assert _uses(surface_check, "await llm_router.aupload_container_file(file=f)") == {"aupload_container_file"} + + +def test_unlisted_class_level_use_fails_the_fence(surface_check): + found = { + "members": {"outside": {"get_valid_args": {"litellm/new_module.py"}}, "router_family": {}}, + "router_utils_imports": {}, + } + allow = {"members_outside": [], "members_router_family": [], "router_utils_imports_outside": []} + errors = surface_check.compare(allow, found) + assert len(errors) == 1 + assert "'get_valid_args' in litellm/new_module.py" in errors[0] + + +def test_unrelated_receivers_are_not_uses(surface_check): + assert _uses(surface_check, "client.get_model_list()\nRouterConfig.model_list") == set() From 49ccce7813725cc87b64d1b0f1b097dca5c681af Mon Sep 17 00:00:00 2001 From: paveln83 <349044532+paveln83@users.noreply.github.com> Date: Wed, 28 Oct 2026 17:12:07 -0700 Subject: [PATCH 119/128] test(extraction): audit #47473 (#48898) * test(extraction): audit #47473 * test(extraction): check over-budget extraction never reaches Cloudflare --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../management/test_model_health_check.py | 37 ++ ...test_cloudflare_clairic_extraction_wire.py | 418 ++++++++++++++++++ 2 files changed, 455 insertions(+) create mode 100644 tests/integration/providers/test_cloudflare_clairic_extraction_wire.py diff --git a/tests/integration/management/test_model_health_check.py b/tests/integration/management/test_model_health_check.py index 57a05daca94..992c1e8e875 100644 --- a/tests/integration/management/test_model_health_check.py +++ b/tests/integration/management/test_model_health_check.py @@ -26,6 +26,20 @@ _DECISIONS_PROBE_REPLY: Final = JsonResponse( "usage": {"input_tokens": 10, "output_tokens": 1}, }, ) +_EXTRACTION_PROBE_REPLY: Final = JsonResponse( + content_type="application/json", + body={ + "result": { + "model": "clairic", + "answers": {"reachable": {"type": "text", "value": "health", "confidence": 0.9}}, + "truncated": False, + "usage": {"input_tokens": 10, "output_tokens": 0}, + }, + "success": True, + "errors": [], + "messages": [], + }, +) _CONFIGURED_PROBE_REPLY: Final = JsonResponse( content_type="application/json", body={ @@ -287,6 +301,29 @@ def test_decisions_mode_health_check_resolves_the_mode_from_the_cost_map_and_sen ] +def test_extraction_mode_health_check_resolves_the_mode_from_the_cost_map_and_sends_the_extraction_probe( + gateway: Gateway, +) -> None: + with gateway.scenario() as scenario: + handle: Final = register_scenario(f"health-extraction-{uuid.uuid4().hex[:12]}", _EXTRACTION_PROBE_REPLY) + scenario.cleanups.callback(delete_scenario, handle) + model: Final = scenario.model( + model="cloudflare/@cf/cloudflare/clairic", api_base=handle.api_base(), api_key="synthetic-cloudflare-key" + ) + report: Final = _health_report(gateway, model) + assert (report["healthy_count"], report["unhealthy_count"]) == (1, 0), report + assert _probes_sent_to(gateway, handle) == [ + ( + f"/{handle.scenario_id}/ai/run/@cf/cloudflare/clairic", + { + "model": "clairic", + "state": os.environ.get("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm"), + "fields": {"reachable": {"type": "text", "instructions": "Any word from the text"}}, + }, + ) + ] + + @pytest.mark.parametrize("mode", ("decisions", "evaluation"), ids=("decisions", "legacy-evaluation")) def test_decisions_mode_health_check_sends_the_configured_state_and_questions(gateway: Gateway, mode: str) -> None: with gateway.scenario() as scenario: diff --git a/tests/integration/providers/test_cloudflare_clairic_extraction_wire.py b/tests/integration/providers/test_cloudflare_clairic_extraction_wire.py new file mode 100644 index 00000000000..7b2a48c9676 --- /dev/null +++ b/tests/integration/providers/test_cloudflare_clairic_extraction_wire.py @@ -0,0 +1,418 @@ +import json +import math +import re +import uuid +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Final, Literal + +import httpx +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value +from integration._support.database import read_rows +from integration._support.process import OwnedProxy, WorkerHealthcheck, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +Route = Literal["/v1/extraction", "/extraction"] + +_API_KEY: Final = "synthetic-cloudflare-key" +_CLAIRIC: Final = "cloudflare/@cf/cloudflare/clairic" +_CLAIRIC_PATH: Final = "/ai/run/@cf/cloudflare/clairic" +_INPUT_TOKENS: Final = 412 +_EMAIL: Final = "billing@acmetools.example" +_SCENARIO: Final = re.compile(r"scenario-([a-z-]+)") +_JSON_ROWS: Final = TypeAdapter(list[dict[str, JsonValue]]) +_FIELDS: Final[dict[str, JsonValue]] = { + "vendor": {"type": "text", "instructions": "Who sent the invoice?"}, + "total": {"type": "number", "unit": "USD"}, +} +_SPEND_QUERY: Final = ( + 'SELECT spend, status, call_type, model_group, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" ' + "WHERE request_id = %s" +) +_OWNED_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120) +_REQUEST_GUARDRAIL: Final = "extraction-mask-email-in" +_RESPONSE_GUARDRAIL: Final = "extraction-mask-email-out" +_OWNED_MODEL: Final = "invoice-reader" + + +def _state(marker: str, scenario: str = "plain") -> str: + return f"Invoice from Acme Tools, total $174. scenario-{scenario} marker-{marker}" + + +def _vendor_span(state: str) -> list[JsonValue]: + start: Final = state.index("Acme Tools") + return [start, start + len("Acme Tools")] + + +def _email_span(state: str) -> list[JsonValue]: + start: Final = state.index(_EMAIL) + return [start, start + len(_EMAIL)] + + +def _answers(state: str) -> dict[str, JsonValue]: + return { + "vendor": {"type": "text", "value": "Acme Tools", "confidence": 0.97, "span": _vendor_span(state)}, + "total": {"type": "number", "value": 174, "confidence": 0.95}, + } + + +def _envelope(result: Mapping[str, JsonValue]) -> Reply: + return Reply(body=json.dumps({"result": result, "success": True, "errors": [], "messages": []}).encode()) + + +def _clairic_reply(request: Request) -> Reply: + state: Final = JSON_OBJECT.validate_json(request.body)["state"] + assert isinstance(state, str), request.body + return _envelope( + { + "model": "clairic", + "answers": _answers(state), + "truncated": False, + "usage": {"input_tokens": _INPUT_TOKENS, "output_tokens": 0}, + } + ) + + +def _expected_response(model: str, state: str) -> dict[str, JsonValue]: + return { + "model": model, + "answers": _answers(state), + "truncated": False, + "usage": {"input_tokens": _INPUT_TOKENS, "output_tokens": 0}, + } + + +def _upstream_body(request: Request) -> dict[str, JsonValue]: + assert (request.method, request.target) == ("POST", _CLAIRIC_PATH), (request.method, request.target) + assert request.headers["authorization"] == f"Bearer {_API_KEY}", request.headers + return JSON_OBJECT.validate_json(request.body) + + +def _body(model: str, marker: str, **extra: JsonValue) -> dict[str, JsonValue]: + return {"model": model, "state": _state(marker), "fields": _FIELDS, **extra} + + +def _extract( + gateway: Gateway, + model: str, + marker: str, + *, + route: Route = "/v1/extraction", + key: str | None = None, + **extra: JsonValue, +) -> httpx.Response: + return gateway.request("POST", route, _body(model, marker, **extra), key=key) + + +def _number(value: JsonValue) -> float: + assert isinstance(value, (int, float)) and not isinstance(value, bool), value + return float(value) + + +def _expected_spend() -> float: + cost_map: Final = JSON_OBJECT.validate_json(Path("model_prices_and_context_window.json").read_bytes()) + priced: Final = cost_map.get(_CLAIRIC) + assert priced is not None, f"{_CLAIRIC} has no cost map entry" + return _INPUT_TOKENS * _number(object_value(priced)["input_cost_per_token"]) + + +def _spend_row(call_id: str) -> dict[str, JsonValue]: + rows: Final = eventually(lambda: read_rows(_SPEND_QUERY, (call_id,)), lambda found: len(found) == 1, seconds=70) + return rows[0] + + +def _assert_billed(gateway: Gateway, response: httpx.Response, model: str) -> None: + expected: Final = _expected_spend() + assert math.isclose(float(response.headers["x-litellm-response-cost"]), expected, rel_tol=1e-9), response.headers + call_id: Final = response.headers["x-litellm-call-id"] + row: Final = _spend_row(call_id) + assert (row["status"], row["call_type"], row["model_group"], row["prompt_tokens"], row["completion_tokens"]) == ( + "success", + "aextraction", + model, + _INPUT_TOKENS, + 0, + ), row + assert math.isclose(_number(row["spend"]), expected, rel_tol=1e-9), row + listed: Final = gateway.request("GET", "/spend/logs", params={"request_id": call_id}) + assert listed.status_code == 200, listed.text + (reported,) = _JSON_ROWS.validate_json(listed.content) + assert (reported["request_id"], reported["model_group"]) == (call_id, model), reported + assert math.isclose(_number(reported["spend"]), expected, rel_tol=1e-9), reported + + +@pytest.mark.parametrize("route", ("/v1/extraction", "/extraction")) +def test_extraction_reaches_clairic_on_both_routes_and_is_billed(gateway: Gateway, route: Route) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_clairic_reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY) + response: Final = _extract(gateway, model, marker, route=route, metadata={"tags": [f"audit-{marker}"]}) + assert response.status_code == 200, response.text + assert response.json() == _expected_response(model, _state(marker)) + (request,) = wire.drain() + assert _upstream_body(request) == {"model": "clairic", "state": _state(marker), "fields": _FIELDS} + _assert_billed(gateway, response, model) + + +_INVALID_BODIES: Final[tuple[tuple[str, dict[str, JsonValue]], ...]] = ( + ("unknown_type", {"fields": {"vendor": {"type": "paragraph"}}}), + ("list_without_items", {"fields": {"items": {"type": "list"}}}), + ("remote_image", {"images": ["https://images.example/receipt.png"]}), +) + + +@pytest.mark.parametrize(("label", "change"), _INVALID_BODIES, ids=[label for label, _ in _INVALID_BODIES]) +def test_invalid_extraction_bodies_are_refused_before_any_upstream_call( + gateway: Gateway, label: str, change: dict[str, JsonValue] +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_clairic_reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request("POST", "/v1/extraction", {**_body(model, marker), **change}) + assert response.status_code == 400, (label, response.text) + assert "Invalid Extraction request" in response.text, (label, response.text) + assert wire.drain() == (), label + + +@pytest.mark.parametrize(("param", "value"), (("temperature", 0.2), ("stream", True)), ids=("temperature", "stream")) +def test_unsupported_top_level_params_are_refused_unless_the_deployment_drops_params( + gateway: Gateway, param: str, value: JsonValue +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_clairic_reply) as wire, gateway.scenario() as scenario: + strict: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY) + dropping: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY, drop_params=True) + refused: Final = _extract(gateway, strict, marker, **{param: value}) + assert refused.status_code == 400, refused.text + assert param in refused.text, refused.text + assert wire.drain() == () + dropped: Final = _extract(gateway, dropping, marker, **{param: value}) + assert dropped.status_code == 200, dropped.text + assert dropped.headers["content-type"].startswith("application/json"), dropped.headers + assert dropped.json() == _expected_response(dropping, _state(marker)) + (request,) = wire.drain() + assert _upstream_body(request) == {"model": "clairic", "state": _state(marker), "fields": _FIELDS} + _assert_billed(gateway, dropped, dropping) + + +def test_key_checks_on_extraction_match_chat(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_clairic_reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY) + anonymous: Final = gateway.client.post("/v1/extraction", json=_body(model, marker)) + assert anonymous.status_code == 401, anonymous.text + restricted: Final = scenario.key(models=[f"other-{uuid.uuid4().hex}"]) + refused: Final = _extract(gateway, model, marker, key=restricted) + chat_refused: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": marker}]}, + key=restricted, + ) + assert 400 <= refused.status_code < 500, refused.text + assert refused.status_code == chat_refused.status_code, (refused.text, chat_refused.text) + assert wire.drain() == () + allowed: Final = scenario.key(models=[model], max_budget=1e-06) + first: Final = _extract(gateway, model, marker, key=allowed) + assert first.status_code == 200, first.text + assert len(wire.drain()) == 1 + blocked, upstream_calls = eventually( + lambda: (_extract(gateway, model, uuid.uuid4().hex, key=allowed), len(wire.drain())), + lambda answer: answer[0].status_code != 200, + seconds=70, + ) + chat_blocked: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": marker}]}, + key=allowed, + ) + assert 400 <= blocked.status_code < 500, blocked.text + assert blocked.status_code == chat_blocked.status_code, (blocked.text, chat_blocked.text) + assert upstream_calls == 0 + assert wire.drain() == () + + +def test_a_small_key_tpm_limit_serves_extraction_and_counts_its_input_tokens(gateway: Gateway) -> None: + with wire_server(_clairic_reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY) + key: Final = scenario.key(models=[model], tpm_limit=100) + first: Final = _extract(gateway, model, uuid.uuid4().hex, key=key) + assert first.status_code == 200, first.text + assert len(wire.drain()) == 1 + limited, upstream_calls = eventually( + lambda: (_extract(gateway, model, uuid.uuid4().hex, key=key), len(wire.drain())), + lambda answer: answer[0].status_code != 200, + seconds=70, + ) + assert limited.status_code == 429, limited.text + assert upstream_calls == 0 + + +def test_cloudflare_errors_keep_their_status_and_log_an_unbilled_failure(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + refusal: Final = json.dumps( + { + "result": None, + "success": False, + "errors": [{"code": 5006, "message": f"min above max {marker}"}], + "messages": [], + } + ).encode() + with wire_server(lambda _: Reply(status=422, body=refusal)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY) + response: Final = _extract(gateway, model, marker) + assert response.status_code == 422, response.text + assert f"min above max {marker}" in response.text, response.text + assert len(wire.drain()) >= 1 + row: Final = _spend_row(response.headers["x-litellm-call-id"]) + assert (row["status"], row["call_type"], row["model_group"], _number(row["spend"])) == ( + "failure", + "aextraction", + model, + 0.0, + ), row + + +def _owned_reply(request: Request) -> Reply: + body: Final = JSON_OBJECT.validate_json(request.body) + state: Final = body["state"] + assert isinstance(state, str), request.body + scenario: Final = _SCENARIO.search(state) + assert scenario is not None, state + if scenario.group(1) != "mask-out": + return _clairic_reply(request) + return _envelope( + { + "model": "clairic", + "answers": { + "vendor": {"type": "text", "value": "Acme Tools", "confidence": 0.97, "span": _vendor_span(state)}, + "contact_email": {"type": "text", "value": _EMAIL, "confidence": 0.9, "span": _email_span(state)}, + }, + "truncated": False, + "state_text": state, + "usage": {"input_tokens": _INPUT_TOKENS, "output_tokens": 0}, + } + ) + + +def _owned_config(upstream_url: str) -> dict[str, JsonValue]: + base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + email_filter: Final[list[JsonValue]] = [{"pattern_type": "prebuilt", "pattern_name": "email", "action": "MASK"}] + return { + **base, + "guardrails": [ + { + "guardrail_name": _REQUEST_GUARDRAIL, + "litellm_params": {"guardrail": "litellm_content_filter", "mode": "pre_call", "patterns": email_filter}, + }, + { + "guardrail_name": _RESPONSE_GUARDRAIL, + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "post_call", + "patterns": email_filter, + }, + }, + ], + "model_list": [ + { + "model_name": _OWNED_MODEL, + "litellm_params": {"model": _CLAIRIC, "api_base": upstream_url, "api_key": _API_KEY}, + } + ], + } + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: OwnedProxy + upstream: Wire + + +@pytest.fixture(scope="module") +def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with wire_server(_owned_reply) as upstream: + directory: Final = tmp_path_factory.mktemp("clairic-extraction-owned") + config_path: Final = directory / "clairic-extraction.yaml" + config_path.write_text(yaml.safe_dump(_owned_config(upstream.url))) + with owned_proxy_process( + session_gateway, + directory, + {}, + config=config_path, + workers=2, + extra_arguments=("--max_tokens", "256", "--temperature", "0.2"), + healthcheck=WorkerHealthcheck.LENIENT, + ) as proxy: + yield Rig(proxy, upstream) + + +def _owned_extract(rig: Rig, state: str, **extra: JsonValue) -> httpx.Response: + return rig.proxy.gateway.request( + "POST", "/v1/extraction", {"model": _OWNED_MODEL, "state": state, "fields": _FIELDS, **extra} + ) + + +def _applied_guardrails(response: httpx.Response) -> frozenset[str]: + return frozenset(name.strip() for name in response.headers.get("x-litellm-applied-guardrails", "").split(",")) + + +@pytest.mark.timeout(_OWNED_TIMEOUT) +def test_cli_generation_defaults_never_reach_an_extraction_call(rig: Rig) -> None: + state: Final = _state(uuid.uuid4().hex) + response: Final = _owned_extract(rig, state) + assert response.status_code == 200, response.text + assert response.json() == _expected_response(_OWNED_MODEL, state) + (request,) = rig.upstream.drain() + assert _upstream_body(request) == {"model": "clairic", "state": state, "fields": _FIELDS} + + +@pytest.mark.timeout(_OWNED_TIMEOUT) +def test_a_request_guardrail_masks_the_state_before_clairic_and_drops_spans_into_it(rig: Rig) -> None: + state: Final = f"{_state(uuid.uuid4().hex, 'mask-in')} Questions go to {_EMAIL}." + response: Final = _owned_extract(rig, state, guardrails=[_REQUEST_GUARDRAIL]) + assert response.status_code == 200, response.text + assert _REQUEST_GUARDRAIL in _applied_guardrails(response), response.headers + (request,) = rig.upstream.drain() + sent: Final = _upstream_body(request) + sent_state: Final = sent["state"] + assert isinstance(sent_state, str) and _EMAIL not in sent_state, sent + assert sent_state.startswith(state.removesuffix(f"{_EMAIL}.")) and "REDACTED" in sent_state, sent + assert sent == {"model": "clairic", "state": sent_state, "fields": _FIELDS} + answers: Final = object_value(response.json()["answers"]) + assert object_value(answers["vendor"]) == {"type": "text", "value": "Acme Tools", "confidence": 0.97}, answers + + +@pytest.mark.timeout(_OWNED_TIMEOUT) +def test_a_response_guardrail_masks_text_answers_and_state_text(rig: Rig) -> None: + state: Final = f"{_state(uuid.uuid4().hex, 'mask-out')} Questions go to {_EMAIL}." + fields: Final[dict[str, JsonValue]] = {**_FIELDS, "contact_email": {"type": "text"}} + response: Final = rig.proxy.gateway.request( + "POST", + "/v1/extraction", + { + "model": _OWNED_MODEL, + "state": state, + "fields": fields, + "guardrails": [_RESPONSE_GUARDRAIL], + "options": {"return_state_text": True}, + }, + ) + assert response.status_code == 200, response.text + assert _RESPONSE_GUARDRAIL in _applied_guardrails(response), response.headers + (request,) = rig.upstream.drain() + assert _upstream_body(request)["state"] == state + body: Final = response.json() + state_text: Final = body["state_text"] + assert isinstance(state_text, str) and _EMAIL not in state_text and "REDACTED" in state_text, body + answers: Final = object_value(body["answers"]) + contact: Final = object_value(answers["contact_email"]) + assert "span" not in contact, answers + masked: Final = contact["value"] + assert isinstance(masked, str) and _EMAIL not in masked and "REDACTED" in masked, answers + assert object_value(answers["vendor"]) == {"type": "text", "value": "Acme Tools", "confidence": 0.97}, answers From 93cfc5265eb4df6e8aa463913c684c08728b5561 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 28 Oct 2026 17:44:48 -0700 Subject: [PATCH 120/128] ci: run zizmor with --no-ignores and move the spend logs and cost map guards to pull_request (#48947) * ci(zizmor): run zizmor 1.31.1 without ignores and own the workflow zizmor-action has no input for --no-ignores, so the workflow now installs the zizmor 1.31.1 release binary directly, pinned by SHA-256 like grype in image-scan.yml, and runs it with --no-ignores and --no-config. Inline `zizmor: ignore` comments and config-file ignores no longer hide findings. Behavior otherwise matches the old action defaults (regular persona, online audits with the job token, medium and above, GitHub annotations) Adds a CODEOWNERS entry for .github/workflows/zizmor.yml * ci: run spend-logs-guard and cost-map-guard on pull_request Both guards ran on pull_request_target behind a zizmor ignore comment. They now run on pull_request and check out the base commit by sha, so the guard scripts still come from the base branch while the PR's head is only fetched and diffed. The zizmor ignore comments are gone --- .github/CODEOWNERS | 1 + .github/workflows/cost-map-guard.yml | 5 +++-- .github/workflows/spend-logs-guard.yml | 6 +++--- .github/workflows/zizmor.yml | 17 +++++++++++------ 4 files changed, 18 insertions(+), 11 deletions(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index f6c06f7d83f..158110cbbab 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -2,3 +2,4 @@ /litellm/model_prices_and_context_window_backup.json @mateo-berri @ryan-crabbe-berri @kerry-berri @joshua-berri /helm/ @yassin-berriai /terraform/ @yassin-berriai +/.github/workflows/zizmor.yml @yuneng-berri diff --git a/.github/workflows/cost-map-guard.yml b/.github/workflows/cost-map-guard.yml index f9172b9551a..05b22017cbb 100644 --- a/.github/workflows/cost-map-guard.yml +++ b/.github/workflows/cost-map-guard.yml @@ -1,7 +1,7 @@ name: Cost map guard -on: # zizmor: ignore[dangerous-triggers] runs the base branch's code only; the PR's cost map files are read as data and never executed - pull_request_target: +on: + pull_request: branches: - main - "litellm_**" @@ -19,6 +19,7 @@ jobs: steps: - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 with: + ref: ${{ github.event.pull_request.base.sha }} persist-credentials: false - name: Fetch the pull request head and its merge base id: revisions diff --git a/.github/workflows/spend-logs-guard.yml b/.github/workflows/spend-logs-guard.yml index 29e9f236202..cc3bda9b880 100644 --- a/.github/workflows/spend-logs-guard.yml +++ b/.github/workflows/spend-logs-guard.yml @@ -1,7 +1,7 @@ name: Spend logs guard -on: # zizmor: ignore[dangerous-triggers] runs the base branch's code only; the PR's diff is read as data and never executed - pull_request_target: +on: + pull_request: branches: - main @@ -19,7 +19,7 @@ jobs: steps: - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 with: - ref: ${{ github.event.pull_request.base.ref }} + ref: ${{ github.event.pull_request.base.sha }} persist-credentials: false - name: Fetch the pull request head, its merge base and its reviews id: revisions diff --git a/.github/workflows/zizmor.yml b/.github/workflows/zizmor.yml index 73f9efb8df9..eba23e795f3 100644 --- a/.github/workflows/zizmor.yml +++ b/.github/workflows/zizmor.yml @@ -27,10 +27,15 @@ jobs: with: persist-credentials: false + - name: Install zizmor + run: | + curl -fsSL --retry 3 -o "$RUNNER_TEMP/zizmor.tar.gz" \ + https://github.com/zizmorcore/zizmor/releases/download/v1.31.1/zizmor-x86_64-unknown-linux-gnu.tar.gz + echo "d185f726bdaf92df78273217dfdd24ebcd6b0ef48cf3ce776ff5762b6e138c99 $RUNNER_TEMP/zizmor.tar.gz" | sha256sum -c - + tar xzf "$RUNNER_TEMP/zizmor.tar.gz" -C "$RUNNER_TEMP" zizmor + - name: Run zizmor - uses: zizmorcore/zizmor-action@5f14fd08f7cf1cb1609c1e344975f152c7ee938d # v0.5.6 - with: - version: "1.24.1" - min-severity: medium - advanced-security: false - annotations: true + env: + GH_TOKEN: ${{ github.token }} + run: | + "$RUNNER_TEMP/zizmor" --no-ignores --no-config --min-severity=medium --format=github . From e1580b7b837b680d0a0770cfa5f6c88938bc6050 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 28 Oct 2026 17:48:09 -0700 Subject: [PATCH 121/128] refactor(images): remove dall-e-2 and dall-e-3 code paths (#48703) OpenAI shut down dall-e-2 and dall-e-3 on 2026-05-12. Drop their image generation and edit configs, the Azure dall-e-2 create-and-poll flow, and the silent dall-e-2 fallback when no model is passed. image_generation without a model or provider now fails locally with "model is required". image_edit and image_variation leave model out of the provider request instead of sending dall-e-2, so providers that ignore model (Topaz) keep working and OpenAI returns its own 400. --- litellm/__init__.py | 1 - litellm/constants.py | 1 - litellm/cost_calculator.py | 11 +- litellm/images/main.py | 32 ++- .../get_llm_provider_logic.py | 1 - litellm/llms/azure/azure.py | 196 +----------------- .../llms/azure/image_generation/__init__.py | 10 +- .../dall_e_2_transformation.py | 7 - .../dall_e_3_transformation.py | 7 - .../azure_ai/image_generation/__init__.py | 10 +- .../dall_e_2_transformation.py | 7 - .../dall_e_3_transformation.py | 7 - litellm/llms/custom_httpx/aiohttp_handler.py | 5 +- litellm/llms/openai/image_edit/__init__.py | 26 --- .../image_edit/dalle2_transformation.py | 98 --------- .../llms/openai/image_edit/transformation.py | 2 +- .../llms/openai/image_generation/__init__.py | 17 -- .../dall_e_2_transformation.py | 81 -------- .../dall_e_3_transformation.py | 81 -------- .../llms/openai/image_variations/handler.py | 2 +- litellm/utils.py | 11 +- tests/unit/images/test_image_edit.py | 18 ++ tests/unit/images/test_main.py | 23 +- .../test_azure_image_generation_init.py | 18 -- ...t_openai_image_generation_extra_headers.py | 2 +- .../llms/openai/test_openai_common_utils.py | 38 ++-- 26 files changed, 89 insertions(+), 623 deletions(-) delete mode 100644 litellm/llms/azure/image_generation/dall_e_2_transformation.py delete mode 100644 litellm/llms/azure/image_generation/dall_e_3_transformation.py delete mode 100644 litellm/llms/azure_ai/image_generation/dall_e_2_transformation.py delete mode 100644 litellm/llms/azure_ai/image_generation/dall_e_3_transformation.py delete mode 100644 litellm/llms/openai/image_edit/dalle2_transformation.py delete mode 100644 litellm/llms/openai/image_generation/dall_e_2_transformation.py delete mode 100644 litellm/llms/openai/image_generation/dall_e_3_transformation.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 8dfe5adca1c..c649a48211a 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1336,7 +1336,6 @@ all_embedding_models = ( ) ####### IMAGE GENERATION MODELS ################### -openai_image_generation_models = ["dall-e-2", "dall-e-3"] ####### VIDEO GENERATION MODELS ################### openai_video_generation_models = ["sora-2"] diff --git a/litellm/constants.py b/litellm/constants.py index 1d5c6404bfd..25770d960bb 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -764,7 +764,6 @@ ANTHROPIC_WEB_SEARCH_TOOL_MAX_USES: Final = { # Used for web search interception across providers LITELLM_WEB_SEARCH_TOOL_NAME: Final = "litellm_web_search" -DEFAULT_IMAGE_ENDPOINT_MODEL: Final = "dall-e-2" DEFAULT_VIDEO_ENDPOINT_MODEL: Final = "sora-2" DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS: Final = get_env_int("DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS", 8) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 86b22069a39..087d5b5b718 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1592,14 +1592,6 @@ def completion_cost( for part in ws_tier_parts ) - if ( - (call_type == "aimage_generation" or call_type == "image_generation") - and model is not None - and isinstance(model, str) - and len(model) == 0 - and custom_llm_provider == "azure" - ): - model = "dall-e-2" # for dall-e-2, azure expects an empty model name # Handle Inputs to completion_cost prompt_tokens = 0 prompt_characters: int | None = None @@ -1674,7 +1666,8 @@ def completion_cost( raw_candidates if custom_pricing is True else _priced_candidates_before_unpriced_rules(raw_candidates) ) - for idx, (model, candidate_pricing_provider) in enumerate(potential_model_names): + for idx, (candidate_model, candidate_pricing_provider) in enumerate(potential_model_names): + model = candidate_model try: if verbose_logger.isEnabledFor(logging.DEBUG): verbose_logger.debug("selected model name for cost calculation: %s", model) diff --git a/litellm/images/main.py b/litellm/images/main.py index d0618e5e804..6e77cbb7046 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -16,7 +16,6 @@ import litellm # client is imported from litellm as it's a decorator from litellm import client -from litellm.constants import DEFAULT_IMAGE_ENDPOINT_MODEL from litellm.constants import request_timeout as DEFAULT_REQUEST_TIMEOUT from litellm.exceptions import LiteLLMUnknownProvider from litellm.litellm_core_utils.get_llm_provider_logic import ( @@ -299,17 +298,14 @@ def image_generation( if custom_llm_provider == "chatgpt" or (model is not None and model.startswith("chatgpt/")) else None ) - dynamic_api_key: str | None = None - if model is not None or custom_llm_provider is not None: - model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider( - model=model, - custom_llm_provider=custom_llm_provider, - api_base=api_base, - litellm_params=chatgpt_params, - ) - else: - model = "dall-e-2" - custom_llm_provider = "openai" # default to dall-e-2 on openai + if model is None and custom_llm_provider is None: + raise litellm.BadRequestError(message="model is required", model="", llm_provider="") + model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider( + model=model, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + litellm_params=chatgpt_params, + ) model_response.hidden_params["model"] = model openai_params: Final = [ "user", @@ -712,7 +708,7 @@ async def aimage_variation(*args, **kwargs) -> ImageResponse: @client def image_variation( image: FileTypes, - model: str = "dall-e-2", # set to dall-e-2 by default - like OpenAI. + model: str | None = None, n: int = 1, response_format: Literal["url", "b64_json"] = "url", size: str | None = None, @@ -728,7 +724,7 @@ def image_variation( litellm_params: Final = get_litellm_params(**kwargs) # get the custom llm provider model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider( - model=model, + model=model or "", custom_llm_provider=litellm_params.get("custom_llm_provider", None), api_base=litellm_params.get("api_base", None), api_key=litellm_params.get("api_key", None), @@ -747,7 +743,7 @@ def image_variation( response: ImageResponse | None = None provider_config: Final = ProviderConfigManager.get_provider_model_info( - model=model or "", # openai defaults to dall-e-2 + model=model, provider=llm_provider, ) @@ -772,7 +768,7 @@ def image_variation( model_response=model_response, api_key=api_key, api_base=api_base, - model=model, + model=model or None, image=image, timeout=litellm_params.get("timeout", None), custom_llm_provider=custom_llm_provider, @@ -790,7 +786,7 @@ def image_variation( model_response=model_response, api_key=api_key, api_base=api_base, - model=model, + model=model or None, image=image, timeout=litellm_params.get("timeout", None) or DEFAULT_REQUEST_TIMEOUT, custom_llm_provider=custom_llm_provider, @@ -901,7 +897,7 @@ def image_edit( # get llm provider logic litellm_params: Final = GenericLiteLLMParams(**kwargs) model, custom_llm_provider, _, _ = get_llm_provider( - model=model or DEFAULT_IMAGE_ENDPOINT_MODEL, + model=model or "", custom_llm_provider=custom_llm_provider, ) diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 19c1d5c414c..d9933087334 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -437,7 +437,6 @@ def get_llm_provider( model in litellm.open_ai_chat_completion_models or "ft:gpt-3.5-turbo" in model or "ft:gpt-4" in model # catches ft:gpt-4-0613, ft:gpt-4o - or model in litellm.openai_image_generation_models or model.startswith("gpt-image") or model in litellm.openai_video_generation_models ): diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 02f1ba3472a..5a181faad11 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -13,13 +13,12 @@ from openai import ( BadRequestError, OpenAI, ) -from pydantic import BaseModel, TypeAdapter +from pydantic import BaseModel import litellm -from litellm.constants import AZURE_OPERATION_POLLING_TIMEOUT, DEFAULT_MAX_RETRIES +from litellm.constants import DEFAULT_MAX_RETRIES from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_utils import speech_request_body, track_llm_api_timing -from litellm.litellm_core_utils.url_utils import SSRFError, assert_same_origin from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -34,7 +33,6 @@ from litellm.types.utils import ( from litellm.utils import ( CustomStreamWrapper, convert_to_model_response_object, - modify_url, ) from ...types.llms.openai import HttpxBinaryResponseContent @@ -59,8 +57,6 @@ from .image_generation import ( ) from .image_generation.http_utils import azure_deployment_image_generation_json_body -_DECODED_JSON: Final = TypeAdapter(object) - class AzureOpenAIAssistantsAPIConfig: """ @@ -860,11 +856,6 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): headers: dict, deployment_name: str | None = None, ) -> httpx.Response: - """ - Implemented for azure dall-e-2 image gen calls - - Alternative to needing a custom transport implementation - """ if client is None: _params: Final = {} if timeout is not None: @@ -881,95 +872,6 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): else: async_handler = client - if ( - "images/generations" in api_base - and api_version - in [ # dall-e-3 starts from `2023-12-01-preview` so we should be able to avoid conflict - "2023-06-01-preview", - "2023-07-01-preview", - "2023-08-01-preview", - "2023-09-01-preview", - "2023-10-01-preview", - ] - ): # CREATE + POLL for azure dall-e-2 calls - api_base = modify_url(original_url=api_base, new_path="/openai/images/generations:submit") - - data.pop( - "model", None - ) # REMOVE 'model' from dall-e-2 arg https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#request-a-generated-image-dall-e-2-preview - response = await async_handler.post( - url=api_base, - data=json.dumps(data), - headers=headers, - ) - if "operation-location" in response.headers: - operation_location_url: Final = response.headers["operation-location"] - else: - raise AzureOpenAIError(status_code=500, message=response.text) - # Reject polling URLs that don't share an origin with ``api_base``. - # Without this an upstream-controlled or attacker-controlled - # value would receive the operator's Azure API key in the - # request headers below. VERIA-51. - try: - assert_same_origin(operation_location_url, api_base) - except SSRFError as ssrf_err: - raise AzureOpenAIError( - status_code=502, - message=f"Rejected polling URL: {ssrf_err}", - ) - response = await async_handler.get( - url=operation_location_url, - headers=headers, - ) - - await response.aread() - - timeout_secs: Final[int] = AZURE_OPERATION_POLLING_TIMEOUT - start_time: Final = time.time() - if "status" not in response.json(): - # Don't reflect the raw response body — when the polling - # URL points at an internal JSON API (cloud metadata - # service etc.) reflecting it here turns Blind SSRF into - # Full-Read SSRF. VERIA-51. - raise AzureOpenAIError( - status_code=502, - message="Polling response missing 'status' field", - ) - while response.json()["status"] not in ["succeeded", "failed"]: - if time.time() - start_time > timeout_secs: - raise AzureOpenAIError(status_code=408, message="Operation polling timed out.") - - await asyncio.sleep(int(response.headers.get("retry-after") or 10)) - response = await async_handler.get( - url=operation_location_url, - headers=headers, - ) - await response.aread() - - if response.json()["status"] == "failed": - error_data: Final = response.json() - # Preserve Azure error details (e.g. content_policy_violation, - # inner_error, content_filter_results) as structured body so - # exception_type() can route them correctly. - _error_body: Final = error_data.get("error", error_data) - _error_msg: Final = ( - _error_body.get("message", "Image generation failed") - if isinstance(_error_body, dict) - else json.dumps(error_data) - ) - raise AzureOpenAIError( - status_code=400, - message=_error_msg, - body=error_data, - ) - - result: Final = _DECODED_JSON.validate_python(response.json()["result"]) - return httpx.Response( - status_code=200, - headers=response.headers, - content=json.dumps(result).encode("utf-8"), - request=httpx.Request(method="POST", url="https://api.openai.com/v1"), - ) request_json: Final = azure_deployment_image_generation_json_body(api_base, data, deployment_name) return await async_handler.post( url=api_base, @@ -988,11 +890,6 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): headers: dict, deployment_name: str | None = None, ) -> httpx.Response: - """ - Implemented for azure dall-e-2 image gen calls - - Alternative to needing a custom transport implementation - """ if client is None: _params: Final = {} if timeout is not None: @@ -1006,87 +903,6 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): else: sync_handler = client - if ( - "images/generations" in api_base - and api_version - in [ # dall-e-3 starts from `2023-12-01-preview` so we should be able to avoid conflict - "2023-06-01-preview", - "2023-07-01-preview", - "2023-08-01-preview", - "2023-09-01-preview", - "2023-10-01-preview", - ] - ): # CREATE + POLL for azure dall-e-2 calls - api_base = modify_url(original_url=api_base, new_path="/openai/images/generations:submit") - - data.pop( - "model", None - ) # REMOVE 'model' from dall-e-2 arg https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#request-a-generated-image-dall-e-2-preview - response = sync_handler.post( - url=api_base, - data=json.dumps(data), - headers=headers, - ) - if "operation-location" in response.headers: - operation_location_url: Final = response.headers["operation-location"] - else: - raise AzureOpenAIError(status_code=500, message=response.text) - try: - assert_same_origin(operation_location_url, api_base) - except SSRFError as ssrf_err: - raise AzureOpenAIError( - status_code=502, - message=f"Rejected polling URL: {ssrf_err}", - ) - response = sync_handler.get( - url=operation_location_url, - headers=headers, - ) - - response.read() - - timeout_secs: Final[int] = AZURE_OPERATION_POLLING_TIMEOUT - start_time: Final = time.time() - if "status" not in response.json(): - raise AzureOpenAIError( - status_code=502, - message="Polling response missing 'status' field", - ) - while response.json()["status"] not in ["succeeded", "failed"]: - if time.time() - start_time > timeout_secs: - raise AzureOpenAIError(status_code=408, message="Operation polling timed out.") - - time.sleep(int(response.headers.get("retry-after") or 10)) - response = sync_handler.get( - url=operation_location_url, - headers=headers, - ) - response.read() - - if response.json()["status"] == "failed": - error_data: Final = response.json() - # Preserve Azure error details (e.g. content_policy_violation, - # inner_error, content_filter_results) as structured body so - # exception_type() can route them correctly. - _error_body: Final = error_data.get("error", error_data) - _error_msg: Final = ( - _error_body.get("message", "Image generation failed") - if isinstance(_error_body, dict) - else json.dumps(error_data) - ) - raise AzureOpenAIError( - status_code=400, - message=_error_msg, - body=error_data, - ) - - result: Final = _DECODED_JSON.validate_python(response.json()["result"]) - return httpx.Response( - status_code=200, - headers=response.headers, - content=json.dumps(result).encode("utf-8"), - request=httpx.Request(method="POST", url="https://api.openai.com/v1"), - ) request_json: Final = azure_deployment_image_generation_json_body(api_base, data, deployment_name) return sync_handler.post( url=api_base, @@ -1193,10 +1009,10 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): deployment_name=model, ) - provider_config: Final = get_azure_image_generation_config(data.get("model", "dall-e-2")) + provider_config: Final = get_azure_image_generation_config(data.get("model", "")) if provider_config is not None: return provider_config.transform_image_generation_response( - model=data.get("model", "dall-e-2"), + model=data.get("model", ""), raw_response=httpx_response, model_response=model_response or ImageResponse(), logging_obj=logging_obj, @@ -1329,10 +1145,10 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): headers=request_headers, deployment_name=model, ) - provider_config: Final = get_azure_image_generation_config(data.get("model", "dall-e-2")) + provider_config: Final = get_azure_image_generation_config(data.get("model", "")) if isinstance(provider_config, AzureFoundryMAIImageGenerationConfig): return provider_config.transform_image_generation_response( - model=data.get("model", "dall-e-2"), + model=data.get("model", ""), raw_response=httpx_response, model_response=model_response or ImageResponse(), logging_obj=logging_obj, diff --git a/litellm/llms/azure/image_generation/__init__.py b/litellm/llms/azure/image_generation/__init__.py index a2a905f2287..7e9b5f5d99c 100644 --- a/litellm/llms/azure/image_generation/__init__.py +++ b/litellm/llms/azure/image_generation/__init__.py @@ -4,14 +4,10 @@ from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) -from .dall_e_2_transformation import AzureDallE2ImageGenerationConfig -from .dall_e_3_transformation import AzureDallE3ImageGenerationConfig from .gpt_transformation import AzureGPTImageGenerationConfig from .http_utils import azure_deployment_image_generation_json_body __all__ = [ - "AzureDallE2ImageGenerationConfig", - "AzureDallE3ImageGenerationConfig", "AzureGPTImageGenerationConfig", "azure_deployment_image_generation_json_body", ] @@ -21,11 +17,7 @@ def get_azure_image_generation_config(model: str) -> BaseImageGenerationConfig: model = model.lower() model = model.replace("-", "") model = model.replace("_", "") - if model == "" or "dalle2" in model: # empty model is dall-e-2 - return AzureDallE2ImageGenerationConfig() - elif "dalle3" in model: - return AzureDallE3ImageGenerationConfig() - elif AzureFoundryMAIImageGenerationConfig.is_mai_model(model): + if AzureFoundryMAIImageGenerationConfig.is_mai_model(model): return AzureFoundryMAIImageGenerationConfig() else: verbose_logger.debug( diff --git a/litellm/llms/azure/image_generation/dall_e_2_transformation.py b/litellm/llms/azure/image_generation/dall_e_2_transformation.py deleted file mode 100644 index 788753e4ea9..00000000000 --- a/litellm/llms/azure/image_generation/dall_e_2_transformation.py +++ /dev/null @@ -1,7 +0,0 @@ -from litellm.llms.openai.image_generation import DallE2ImageGenerationConfig - - -class AzureDallE2ImageGenerationConfig(DallE2ImageGenerationConfig): - """ - Azure dall-e-2 image generation config - """ diff --git a/litellm/llms/azure/image_generation/dall_e_3_transformation.py b/litellm/llms/azure/image_generation/dall_e_3_transformation.py deleted file mode 100644 index e6974f9317e..00000000000 --- a/litellm/llms/azure/image_generation/dall_e_3_transformation.py +++ /dev/null @@ -1,7 +0,0 @@ -from litellm.llms.openai.image_generation import DallE3ImageGenerationConfig - - -class AzureDallE3ImageGenerationConfig(DallE3ImageGenerationConfig): - """ - Azure dall-e-3 image generation config - """ diff --git a/litellm/llms/azure_ai/image_generation/__init__.py b/litellm/llms/azure_ai/image_generation/__init__.py index 742922ad295..bd3c11d36f9 100644 --- a/litellm/llms/azure_ai/image_generation/__init__.py +++ b/litellm/llms/azure_ai/image_generation/__init__.py @@ -3,15 +3,11 @@ from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) -from .dall_e_2_transformation import AzureFoundryDallE2ImageGenerationConfig -from .dall_e_3_transformation import AzureFoundryDallE3ImageGenerationConfig from .flux_transformation import AzureFoundryFluxImageGenerationConfig from .gpt_transformation import AzureFoundryGPTImageGenerationConfig from .mai_transformation import AzureFoundryMAIImageGenerationConfig __all__ = [ - "AzureFoundryDallE2ImageGenerationConfig", - "AzureFoundryDallE3ImageGenerationConfig", "AzureFoundryFluxImageGenerationConfig", "AzureFoundryGPTImageGenerationConfig", "AzureFoundryMAIImageGenerationConfig", @@ -22,11 +18,7 @@ def get_azure_ai_image_generation_config(model: str) -> BaseImageGenerationConfi model = model.lower() model = model.replace("-", "") model = model.replace("_", "") - if model == "" or "dalle2" in model: # empty model is dall-e-2 - return AzureFoundryDallE2ImageGenerationConfig() - elif "dalle3" in model: - return AzureFoundryDallE3ImageGenerationConfig() - elif AzureFoundryMAIImageGenerationConfig.is_mai_model(model): + if AzureFoundryMAIImageGenerationConfig.is_mai_model(model): return AzureFoundryMAIImageGenerationConfig() elif "flux" in model: return AzureFoundryFluxImageGenerationConfig() diff --git a/litellm/llms/azure_ai/image_generation/dall_e_2_transformation.py b/litellm/llms/azure_ai/image_generation/dall_e_2_transformation.py deleted file mode 100644 index 2c8a5426628..00000000000 --- a/litellm/llms/azure_ai/image_generation/dall_e_2_transformation.py +++ /dev/null @@ -1,7 +0,0 @@ -from litellm.llms.openai.image_generation import DallE2ImageGenerationConfig - - -class AzureFoundryDallE2ImageGenerationConfig(DallE2ImageGenerationConfig): - """ - Azure dall-e-2 image generation config - """ diff --git a/litellm/llms/azure_ai/image_generation/dall_e_3_transformation.py b/litellm/llms/azure_ai/image_generation/dall_e_3_transformation.py deleted file mode 100644 index 3ceb9b405da..00000000000 --- a/litellm/llms/azure_ai/image_generation/dall_e_3_transformation.py +++ /dev/null @@ -1,7 +0,0 @@ -from litellm.llms.openai.image_generation import DallE3ImageGenerationConfig - - -class AzureFoundryDallE3ImageGenerationConfig(DallE3ImageGenerationConfig): - """ - Azure dall-e-3 image generation config - """ diff --git a/litellm/llms/custom_httpx/aiohttp_handler.py b/litellm/llms/custom_httpx/aiohttp_handler.py index 4060d84f7f4..e6544968d24 100644 --- a/litellm/llms/custom_httpx/aiohttp_handler.py +++ b/litellm/llms/custom_httpx/aiohttp_handler.py @@ -561,11 +561,8 @@ class BaseLLMAIOHTTPHandler: organization: str | None = None, headers: dict | None = None, ) -> ImageResponse: - if model is None: - raise ValueError("model is required for non-openai image variations") - provider_config: Final = ProviderConfigManager.get_provider_image_variation_config( - model=model, # openai defaults to dall-e-2 + model=model or "", provider=LlmProviders(custom_llm_provider), ) diff --git a/litellm/llms/openai/image_edit/__init__.py b/litellm/llms/openai/image_edit/__init__.py index 63a41b65e8a..261010071c9 100644 --- a/litellm/llms/openai/image_edit/__init__.py +++ b/litellm/llms/openai/image_edit/__init__.py @@ -1,31 +1,5 @@ -from typing import Final - -from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig - -from .dalle2_transformation import DallE2ImageEditConfig from .transformation import OpenAIImageEditConfig __all__ = [ - "DallE2ImageEditConfig", "OpenAIImageEditConfig", - "get_openai_image_edit_config", ] - - -def get_openai_image_edit_config(model: str) -> BaseImageEditConfig: - """ - Get the appropriate OpenAI image edit config based on the model. - - Args: - model: The model name (e.g., "dall-e-2", "gpt-image-1") - - Returns: - The appropriate config instance for the model - """ - model_normalized: Final = model.lower().replace("-", "").replace("_", "") - - if model_normalized == "dalle2": - return DallE2ImageEditConfig() - else: - # Default to standard OpenAI config for gpt-image-1 and other models - return OpenAIImageEditConfig() diff --git a/litellm/llms/openai/image_edit/dalle2_transformation.py b/litellm/llms/openai/image_edit/dalle2_transformation.py deleted file mode 100644 index f349959b14f..00000000000 --- a/litellm/llms/openai/image_edit/dalle2_transformation.py +++ /dev/null @@ -1,98 +0,0 @@ -from io import BufferedReader -from typing import TYPE_CHECKING, Any, Final, cast - -from httpx._types import RequestFiles - -import litellm -from litellm.images.utils import ImageEditRequestUtils -from litellm.types.images.main import ImageEditRequestParams -from litellm.types.llms.openai import FileTypes -from litellm.types.router import GenericLiteLLMParams - -from .transformation import OpenAIImageEditConfig - -if TYPE_CHECKING: - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj - - LiteLLMLoggingObj = _LiteLLMLoggingObj -else: - LiteLLMLoggingObj = Any - - -class DallE2ImageEditConfig(OpenAIImageEditConfig): - """ - DALL-E-2 specific configuration for image edit API. - - DALL-E-2 only supports editing a single image (not an array). - Uses "image" field name instead of "image[]". - """ - - def transform_image_edit_request( - self, - model: str, - prompt: str | None, - image: FileTypes | None, - image_edit_optional_request_params: dict, - litellm_params: GenericLiteLLMParams, - headers: dict, - ) -> tuple[dict, RequestFiles]: - """ - Transform image edit request for DALL-E-2. - - DALL-E-2 only accepts a single image with field name "image" (not "image[]"). - """ - request_params: Final = { - "model": model, - **image_edit_optional_request_params, - } - if image is not None: - request_params["image"] = image - if prompt is not None: - request_params["prompt"] = prompt - - request: Final = ImageEditRequestParams(**request_params) - request_dict: Final = cast(dict, request) - - ######################################################### - # Separate images and masks as `files` and send other parameters as `data` - ######################################################### - _image_list: Final = request_dict.get("image") - _mask = request_dict.get("mask") - data_without_files: Final = {k: v for k, v in request_dict.items() if k not in ["image", "mask"]} - files_list: Final[list[tuple[str, Any]]] = [] - - # Handle image parameter - DALL-E-2 only supports single image - if _image_list is not None: - image_list: Final = [_image_list] if not isinstance(_image_list, list) else _image_list - - # Validate only one image is provided - if len(image_list) > 1: - raise litellm.BadRequestError( - message="DALL-E-2 only supports editing a single image. Please provide one image.", - model=model, - llm_provider="openai", - ) - - # Use "image" field name (singular) for DALL-E-2 - for _image in image_list: - if _image is not None: - self._add_image_to_files( - files_list=files_list, - image=_image, - field_name="image", - ) - - # Handle mask parameter if provided - if _mask is not None: - # Handle case where mask can be a list (extract first mask) - if isinstance(_mask, list): - _mask = _mask[0] if _mask else None - - if _mask is not None: - mask_content_type: Final[str] = ImageEditRequestUtils.get_image_content_type(_mask) - if isinstance(_mask, BufferedReader): - files_list.append(("mask", (_mask.name, _mask, mask_content_type))) - else: - files_list.append(("mask", ("mask.png", _mask, mask_content_type))) - - return data_without_files, files_list diff --git a/litellm/llms/openai/image_edit/transformation.py b/litellm/llms/openai/image_edit/transformation.py index 9def833052b..0bba58096ef 100644 --- a/litellm/llms/openai/image_edit/transformation.py +++ b/litellm/llms/openai/image_edit/transformation.py @@ -96,7 +96,7 @@ class OpenAIImageEditConfig(BaseImageEditConfig): """ # Build request params, only including non-None values request_params: Final = { - "model": model, + **({"model": model} if model else {}), **image_edit_optional_request_params, } if image is not None: diff --git a/litellm/llms/openai/image_generation/__init__.py b/litellm/llms/openai/image_generation/__init__.py index e20c80f20bb..4cffbe81a0e 100644 --- a/litellm/llms/openai/image_generation/__init__.py +++ b/litellm/llms/openai/image_generation/__init__.py @@ -1,9 +1,3 @@ -from litellm.llms.base_llm.image_generation.transformation import ( - BaseImageGenerationConfig, -) - -from .dall_e_2_transformation import DallE2ImageGenerationConfig -from .dall_e_3_transformation import DallE3ImageGenerationConfig from .gpt_transformation import GPTImageGenerationConfig from .guardrail_translation import ( OpenAIImageGenerationHandler, @@ -11,18 +5,7 @@ from .guardrail_translation import ( ) __all__ = [ - "DallE2ImageGenerationConfig", - "DallE3ImageGenerationConfig", "GPTImageGenerationConfig", "OpenAIImageGenerationHandler", "guardrail_translation_mappings", ] - - -def get_openai_image_generation_config(model: str) -> BaseImageGenerationConfig: - if model.startswith("dall-e-2") or model == "": # empty model is dall-e-2 - return DallE2ImageGenerationConfig() - elif model.startswith("dall-e-3"): - return DallE3ImageGenerationConfig() - else: - return GPTImageGenerationConfig() diff --git a/litellm/llms/openai/image_generation/dall_e_2_transformation.py b/litellm/llms/openai/image_generation/dall_e_2_transformation.py deleted file mode 100644 index ffc6f1d5fe9..00000000000 --- a/litellm/llms/openai/image_generation/dall_e_2_transformation.py +++ /dev/null @@ -1,81 +0,0 @@ -from typing import TYPE_CHECKING, Final - -import httpx - -from litellm.llms.base_llm.image_generation.transformation import ( - BaseImageGenerationConfig, -) -from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams -from litellm.types.utils import ImageResponse -from litellm.utils import convert_to_model_response_object - -if TYPE_CHECKING: - from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj - - from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer - - -class DallE2ImageGenerationConfig(BaseImageGenerationConfig): - """ - OpenAI dall-e-2 image generation config - """ - - def get_supported_openai_params(self, model: str) -> list[OpenAIImageGenerationOptionalParams]: - return ["n", "response_format", "quality", "size", "user"] - - def map_openai_params( - self, - non_default_params: dict, - optional_params: dict, - model: str, - drop_params: bool, - ) -> dict: - supported_params: Final = self.get_supported_openai_params(model) - for k in non_default_params: - if k not in optional_params: - if k in supported_params: - optional_params[k] = non_default_params[k] - elif drop_params: - pass - else: - raise ValueError( - f"Parameter {k} is not supported for model {model}. Supported parameters are {supported_params}. Set drop_params=True to drop unsupported parameters." - ) - - return optional_params - - def transform_image_generation_response( - self, - model: str, - raw_response: httpx.Response, - model_response: ImageResponse, - logging_obj: "LiteLLMLoggingObj", - request_data: dict, - optional_params: dict, - litellm_params: dict, - encoding: "Tokenizer | None", - api_key: str | None = None, - json_mode: bool | None = None, - ) -> ImageResponse: - response: Final = raw_response.json() - - stringified_response: Final = response - ## LOGGING - logging_obj.post_call( - input=request_data.get("prompt", ""), - api_key=api_key, - additional_args={"complete_input_dict": request_data}, - original_response=stringified_response, - ) - image_response: Final[ImageResponse] = convert_to_model_response_object( - response_object=stringified_response, - model_response_object=model_response, - response_type="image_generation", - ) - - # set optional params - image_response.size = optional_params.get("size", "1024x1024") # default is always 1024x1024 - image_response.quality = optional_params.get("quality", "standard") # always standard for dall-e-2 - image_response.output_format = optional_params.get("output_format", "png") # always png for dall-e-2 - - return image_response diff --git a/litellm/llms/openai/image_generation/dall_e_3_transformation.py b/litellm/llms/openai/image_generation/dall_e_3_transformation.py deleted file mode 100644 index 90b7eaedf2f..00000000000 --- a/litellm/llms/openai/image_generation/dall_e_3_transformation.py +++ /dev/null @@ -1,81 +0,0 @@ -from typing import TYPE_CHECKING, Final - -import httpx - -from litellm.llms.base_llm.image_generation.transformation import ( - BaseImageGenerationConfig, -) -from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams -from litellm.types.utils import ImageResponse -from litellm.utils import convert_to_model_response_object - -if TYPE_CHECKING: - from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj - - from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer - - -class DallE3ImageGenerationConfig(BaseImageGenerationConfig): - """ - OpenAI dall-e-3 image generation config - """ - - def get_supported_openai_params(self, model: str) -> list[OpenAIImageGenerationOptionalParams]: - return ["n", "response_format", "quality", "size", "user", "style"] - - def map_openai_params( - self, - non_default_params: dict, - optional_params: dict, - model: str, - drop_params: bool, - ) -> dict: - supported_params: Final = self.get_supported_openai_params(model) - for k in non_default_params: - if k not in optional_params: - if k in supported_params: - optional_params[k] = non_default_params[k] - elif drop_params: - pass - else: - raise ValueError( - f"Parameter {k} is not supported for model {model}. Supported parameters are {supported_params}. Set drop_params=True to drop unsupported parameters." - ) - - return optional_params - - def transform_image_generation_response( - self, - model: str, - raw_response: httpx.Response, - model_response: ImageResponse, - logging_obj: "LiteLLMLoggingObj", - request_data: dict, - optional_params: dict, - litellm_params: dict, - encoding: "Tokenizer | None", - api_key: str | None = None, - json_mode: bool | None = None, - ) -> ImageResponse: - response: Final = raw_response.json() - - stringified_response: Final = response - ## LOGGING - logging_obj.post_call( - input=request_data.get("prompt", ""), - api_key=api_key, - additional_args={"complete_input_dict": request_data}, - original_response=stringified_response, - ) - image_response: Final[ImageResponse] = convert_to_model_response_object( - response_object=stringified_response, - model_response_object=model_response, - response_type="image_generation", - ) - - # set optional params - image_response.size = optional_params.get("size", "1024x1024") # default is always 1024x1024 - image_response.quality = optional_params.get("quality", "hd") # always hd for dall-e-3 - image_response.output_format = optional_params.get("output_format", "png") # always png for dall-e-3 - - return image_response diff --git a/litellm/llms/openai/image_variations/handler.py b/litellm/llms/openai/image_variations/handler.py index 89dcdc40e02..01f01d3de1d 100644 --- a/litellm/llms/openai/image_variations/handler.py +++ b/litellm/llms/openai/image_variations/handler.py @@ -151,7 +151,7 @@ class OpenAIImageVariationsHandler: ) -> ImageResponse: try: provider_config: Final = ProviderConfigManager.get_provider_image_variation_config( - model=model or "", # openai defaults to dall-e-2 + model=model or "", provider=LlmProviders.OPENAI, ) diff --git a/litellm/utils.py b/litellm/utils.py index 1ef39059e1b..be1e11dd39f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7664,7 +7664,6 @@ def validate_environment( model in litellm.open_ai_chat_completion_models or model in litellm.open_ai_text_completion_models or model in litellm.open_ai_embedding_models - or model in litellm.openai_image_generation_models or model.startswith("gpt-image") ): if "OPENAI_API_KEY" in os.environ: @@ -10311,11 +10310,9 @@ class ProviderConfigManager: provider: LlmProviders, ) -> BaseImageGenerationConfig | None: if LlmProviders.OPENAI == provider: - from litellm.llms.openai.image_generation import ( - get_openai_image_generation_config, - ) + from litellm.llms.openai.image_generation import GPTImageGenerationConfig - return get_openai_image_generation_config(model) + return GPTImageGenerationConfig() elif LlmProviders.AZURE == provider: from litellm.llms.azure.image_generation import ( get_azure_image_generation_config, @@ -10546,9 +10543,9 @@ class ProviderConfigManager: provider: LlmProviders, ) -> BaseImageEditConfig | None: if LlmProviders.OPENAI == provider: - from litellm.llms.openai.image_edit import get_openai_image_edit_config + from litellm.llms.openai.image_edit import OpenAIImageEditConfig - return get_openai_image_edit_config(model=model) + return OpenAIImageEditConfig() elif LlmProviders.HOSTED_VLLM == provider: from litellm.llms.hosted_vllm.image_edit import get_hosted_vllm_image_edit_config diff --git a/tests/unit/images/test_image_edit.py b/tests/unit/images/test_image_edit.py index fb7a2869b2a..4f858050e3e 100644 --- a/tests/unit/images/test_image_edit.py +++ b/tests/unit/images/test_image_edit.py @@ -105,6 +105,24 @@ def test_openai_image_edit_sync_accepts_bytesio_images(respx_mock: respx.MockRou assert _multipart_image_parts(route.calls[0].request) == (_FIRST_IMAGE,) +def test_openai_image_edit_without_a_model_leaves_model_out_of_the_request( + respx_mock: respx.MockRouter, httpx_transport: None +) -> None: + route: Final = respx_mock.post("https://api.openai.com/v1/images/edits").mock( + return_value=httpx.Response(200, json=_EDIT_RESPONSE) + ) + + litellm.image_edit( + prompt="combine the reference images", + image=io.BytesIO(_FIRST_IMAGE), + custom_llm_provider="openai", + api_key="fake-key", + ) + + assert route.call_count == 1 + assert b'name="model"' not in route.calls[0].request.read() + + @pytest.mark.asyncio async def test_openai_image_edit_accepts_bytesio_images(respx_mock: respx.MockRouter, httpx_transport: None) -> None: route: Final = respx_mock.post("https://api.openai.com/v1/images/edits").mock( diff --git a/tests/unit/images/test_main.py b/tests/unit/images/test_main.py index 29424e50079..61b2f28b727 100644 --- a/tests/unit/images/test_main.py +++ b/tests/unit/images/test_main.py @@ -300,7 +300,7 @@ async def test_aimage_generation_ignores_a_non_string_base_url_like_the_sync_pat ) response: Final = await litellm.aimage_generation( - model="dall-e-3", + model="gpt-image-2.5-sunburst", prompt="a red circle", api_base=api_base, base_url=httpx.URL(api_base), @@ -688,3 +688,24 @@ async def test_chatgpt_images_forward_client_login_and_workspace( assert "cookie" not in request.headers assert "x-litellm-api-key" not in request.headers assert json.loads(request.content) == {"model": "gpt-image-2", "prompt": "A blue circle"} + + +def test_image_generation_without_a_model_or_provider_fails_before_calling_any_provider( + respx_mock: respx.MockRouter, +) -> None: + with pytest.raises(litellm.BadRequestError, match="model is required"): + litellm.image_generation(prompt="a red circle", api_key="fake_openai_api_key") + + assert not respx_mock.calls + + +def test_topaz_image_variation_works_without_a_model(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post("https://api.topazlabs.com/image/v1/enhance").mock( + return_value=httpx.Response(200, content=b"enhanced-image", headers={"content-type": "image/png"}) + ) + + response: Final = litellm.image_variation(image=b"source-image", custom_llm_provider="topaz", api_key="fake-key") + + assert route.call_count == 1 + assert response.data is not None + assert response.data[0].b64_json == base64.b64encode(b"enhanced-image").decode() diff --git a/tests/unit/llms/azure/image_generation/test_azure_image_generation_init.py b/tests/unit/llms/azure/image_generation/test_azure_image_generation_init.py index eabd5c8427d..83e01fce642 100644 --- a/tests/unit/llms/azure/image_generation/test_azure_image_generation_init.py +++ b/tests/unit/llms/azure/image_generation/test_azure_image_generation_init.py @@ -20,27 +20,9 @@ from litellm.llms.azure.image_generation.http_utils import ( azure_deployment_image_generation_json_body, ) from litellm.llms.custom_httpx.http_handler import HTTPHandler -from litellm.llms.azure.image_generation import ( - AzureDallE3ImageGenerationConfig, - get_azure_image_generation_config, -) from litellm.utils import get_optional_params_image_gen -@pytest.mark.parametrize( - "received_model, expected_config", - [ - ("dall-e-3", AzureDallE3ImageGenerationConfig), - ("dalle-3", AzureDallE3ImageGenerationConfig), - ("openai_dall_e_3", AzureDallE3ImageGenerationConfig), - ], -) -def test_azure_image_generation_config(received_model, expected_config): - assert isinstance( - get_azure_image_generation_config(received_model), expected_config - ) - - def test_azure_deployment_image_generation_json_body(): """Deployment-scoped Azure image URL must not send ``model`` in JSON.""" api = ( diff --git a/tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py b/tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py index 11df07f5fea..e103412e0c3 100644 --- a/tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py +++ b/tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py @@ -265,7 +265,7 @@ class TestImageGenerationEntryPointHeaders: test_headers = {"cf-aig-authorization": "Bearer my-secret"} await litellm.aimage_generation( - model="dall-e-3", + model="gpt-image-2.5-sunburst", prompt="A white cat", extra_headers=test_headers, client=mock_openai_client, diff --git a/tests/unit/llms/openai/test_openai_common_utils.py b/tests/unit/llms/openai/test_openai_common_utils.py index 831df30c98a..8dc92e23fdf 100644 --- a/tests/unit/llms/openai/test_openai_common_utils.py +++ b/tests/unit/llms/openai/test_openai_common_utils.py @@ -46,12 +46,12 @@ API_FUNCTION_PARAMS = [ ( "image_generation", False, - {"model": "dall-e-3", "prompt": "A beautiful sunset over mountains"}, + {"model": "gpt-image-2.5-sunburst", "prompt": "A beautiful sunset over mountains"}, ), ( "image_generation", True, - {"model": "dall-e-3", "prompt": "A beautiful sunset over mountains"}, + {"model": "gpt-image-2.5-sunburst", "prompt": "A beautiful sunset over mountains"}, ), ( "speech", @@ -84,11 +84,7 @@ async def test_openai_client_reuse(function_name, is_async, args): """ # Determine which client class to mock based on whether the test is async - client_path = ( - "litellm.llms.openai.openai.AsyncOpenAI" - if is_async - else "litellm.llms.openai.openai.OpenAI" - ) + client_path = "litellm.llms.openai.openai.AsyncOpenAI" if is_async else "litellm.llms.openai.openai.OpenAI" # Create the appropriate patches with ( @@ -98,9 +94,7 @@ async def test_openai_client_reuse(function_name, is_async, args): ): # Setup the mock to return None first time (cache miss) then a client for subsequent calls mock_client = MagicMock() - mock_get_cache.side_effect = [None] + [ - mock_client - ] * 9 # First call returns None, rest return the mock client + mock_get_cache.side_effect = [None] + [mock_client] * 9 # First call returns None, rest return the mock client # Make 10 API calls for _ in range(10): @@ -118,9 +112,9 @@ async def test_openai_client_reuse(function_name, is_async, args): pass # Verify client was created only once - assert ( - mock_client_class.call_count == 1 - ), f"{'Async' if is_async else ''}OpenAI client should be created only once" + assert mock_client_class.call_count == 1, ( + f"{'Async' if is_async else ''}OpenAI client should be created only once" + ) # Verify the client was cached assert mock_set_cache.call_count == 1, "Client should be cached once" @@ -144,12 +138,8 @@ def test_precomputed_init_params_match_inspect_signature(): _OPENAI_INIT_PARAMS, ) - expected_openai = tuple( - p for p in inspect.signature(OpenAI.__init__).parameters if p != "self" - ) - expected_azure = tuple( - p for p in inspect.signature(AzureOpenAI.__init__).parameters if p != "self" - ) + expected_openai = tuple(p for p in inspect.signature(OpenAI.__init__).parameters if p != "self") + expected_azure = tuple(p for p in inspect.signature(AzureOpenAI.__init__).parameters if p != "self") assert _OPENAI_INIT_PARAMS == expected_openai assert _AZURE_OPENAI_INIT_PARAMS == expected_azure @@ -500,7 +490,10 @@ async def test_sdk_api_factories_keep_httpx_transport_and_request_timeouts(api: moved: Final = mock.get(host="sdk-default.example", path="/v1/moved-models").mock( return_value=httpx.Response( 200, - json={"object": "list", "data": [{"id": "sdk-compat", "object": "model", "created": 0, "owned_by": "t"}]}, + json={ + "object": "list", + "data": [{"id": "sdk-compat", "object": "model", "created": 0, "owned_by": "t"}], + }, ) ) sdk_client: Final = _sdk_api_client(api, is_async, timeout) @@ -542,7 +535,10 @@ async def test_azure_gateway_and_image_variations_use_the_callers_async_session( assert azure_client is not None assert azure_client._client is session response: Final = await aimage_variation( - image=BytesIO(b"image-bytes"), api_key="transport-only", api_base="https://sdk-default.example/v1" + image=BytesIO(b"image-bytes"), + custom_llm_provider="openai", + api_key="transport-only", + api_base="https://sdk-default.example/v1", ) assert response.data[0].url == "https://example.com/image.png" assert not session.is_closed From 77cad6e9032369fab248870d8ca9cfa94367484a Mon Sep 17 00:00:00 2001 From: zoveki <349028329+zoveki@users.noreply.github.com> Date: Wed, 28 Oct 2026 18:05:03 -0700 Subject: [PATCH 122/128] test(proxy): audit #46997 (#48918) * test(proxy): audit #46997 Add request-path cells for the global proxy spend counter: a key with no user spends toward litellm max_budget, and two proxies sharing Redis add every request to one counter that matches the proxy budget row * test(proxy): wait for the recorded budget spend and shared counter before checking them --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../spend/test_global_proxy_spend_counter.py | 159 +++++++++++++++++- 1 file changed, 158 insertions(+), 1 deletion(-) diff --git a/tests/integration/spend/test_global_proxy_spend_counter.py b/tests/integration/spend/test_global_proxy_spend_counter.py index 0185e48ffbd..e818487cf53 100644 --- a/tests/integration/spend/test_global_proxy_spend_counter.py +++ b/tests/integration/spend/test_global_proxy_spend_counter.py @@ -1,16 +1,113 @@ +import json import os import uuid +from itertools import product +from pathlib import Path from typing import Final +import httpx import pytest +import yaml +from integration._support.client import Gateway, eventually, object_value +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.redis_process import owned_redis +from pydantic import JsonValue, TypeAdapter +from redis import Redis import litellm from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache -from litellm.constants import GLOBAL_PROXY_SPEND_CACHE_KEY +from litellm.constants import GLOBAL_PROXY_SPEND_CACHE_KEY, LITELLM_PROXY_BUDGET_NAME from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.proxy_server import update_cache +BUDGET_MODEL: Final = "proxy-budget-model" +CONFIG: Final = TypeAdapter(dict[str, JsonValue]) +SECTION: Final = TypeAdapter(dict[str, JsonValue]) +OWNED_PROXY_SECONDS: Final = 2 * graceful_stop_seconds() + 120 + + +def _proxy_budget_config(gateway: Gateway, directory: Path, max_budget: float) -> Path: + loaded: Final = CONFIG.validate_python( + yaml.safe_load((Path(__file__).resolve().parents[1] / "proxy_config.yaml").read_text()) + ) + model: Final[JsonValue] = { + "model_name": BUDGET_MODEL, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "integration-provider-key", + "api_base": f"{gateway.upstream_url}/v1", + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + }, + } + path: Final = directory / "proxy_budget_config.yaml" + path.write_text( + yaml.safe_dump( + { + **loaded, + "model_list": [model], + "general_settings": { + **SECTION.validate_python(loaded["general_settings"]), + "user_api_key_cache_ttl": 600, + }, + "litellm_settings": { + **SECTION.validate_python(loaded["litellm_settings"]), + "max_budget": max_budget, + "budget_duration": "30d", + }, + } + ) + ) + return path + + +def _chat(proxy: Gateway, key: str, marker: str) -> httpx.Response: + return proxy.request( + "POST", + "/v1/chat/completions", + {"model": BUDGET_MODEL, "max_tokens": 20, "messages": [{"role": "user", "content": marker}]}, + key=key, + ) + + +def _served_id(proxy: Gateway, key: str) -> str: + response: Final = _chat(proxy, key, f"proxy budget {uuid.uuid4().hex}") + assert response.status_code == 200, response.text + identity: Final = response.json()["id"] + assert isinstance(identity, str), response.text + return identity + + +def _logged_spend(proxy: Gateway, request_ids: tuple[str, ...]) -> tuple[float, ...]: + rows: Final = eventually( + lambda: proxy.rows( + 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id = ANY(%s::text[])', + ("{" + ",".join(request_ids) + "}",), + ), + lambda found: len(found) == len(request_ids), + seconds=70, + ) + return tuple(float(str(row["spend"])) for row in rows) + + +def _proxy_budget_spend(proxy: Gateway) -> float | None: + rows: Final = proxy.rows('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id = %s', (LITELLM_PROXY_BUDGET_NAME,)) + return float(str(rows[0]["spend"])) if rows else None + + +def _counter_key(redis: Redis) -> str: + found: Final = tuple(str(key, "utf-8") for key in redis.scan_iter(match=f"*{GLOBAL_PROXY_SPEND_CACHE_KEY}")) + assert len(found) == 1, found + return found[0] + + +def _upstream_saw(proxy: Gateway, marker: str) -> bool: + with httpx.Client(base_url=proxy.upstream_url, timeout=5, trust_env=False) as upstream: + response: Final = upstream.get("/__observations") + response.raise_for_status() + return marker in json.dumps(object_value(response.json())["requests"]) + async def _add_request_cost(cache: UserApiKeyCache, response_cost: float) -> None: await update_cache( @@ -56,3 +153,63 @@ async def test_global_proxy_spend_adds_every_instance_cost_and_stays_cleared_aft assert await redis_cache.async_get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY) is None finally: await redis_cache.async_delete_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY) + + +@pytest.mark.timeout(OWNED_PROXY_SECONDS) +def test_a_key_with_no_user_spends_toward_the_proxy_max_budget(gateway: Gateway, tmp_path: Path) -> None: + with ( + owned_proxy(gateway, tmp_path, {}, config=_proxy_budget_config(gateway, tmp_path, max_budget=0.01)) as owned, + owned.scenario() as scenario, + ): + key: Final = scenario.key(models=[BUDGET_MODEL]) + (spent,) = _logged_spend(owned, (_served_id(owned, key),)) + assert spent > 0.01, spent + eventually( + lambda: _proxy_budget_spend(owned), + lambda recorded: recorded is not None and recorded > 0.01, + seconds=70, + ) + + marker: Final = f"over the proxy budget {uuid.uuid4().hex}" + refused: Final = _chat(owned, key, marker) + + assert refused.status_code != 200, refused.text + assert object_value(refused.json()["error"])["type"] == "budget_exceeded", refused.text + assert not _upstream_saw(owned, marker), refused.text + + +@pytest.mark.timeout(2 * OWNED_PROXY_SECONDS) +def test_two_proxies_add_every_request_to_one_proxy_spend_counter(gateway: Gateway, tmp_path: Path) -> None: + config: Final = _proxy_budget_config(gateway, tmp_path, max_budget=1000.0) + with owned_redis(tmp_path / "redis") as cache, Redis(host=cache.host, port=cache.port) as redis: + shared_redis: Final = {"REDIS_HOST": cache.host, "REDIS_PORT": str(cache.port)} + with owned_proxy(gateway, tmp_path / "first", shared_redis, config=config) as first: + assert first.database_url is not None + with ( + owned_proxy( + gateway, tmp_path / "second", {**shared_redis, "DATABASE_URL": first.database_url}, config=config + ) as second, + first.scenario() as scenario, + ): + user_key: Final = scenario.key(user_id=scenario.user(), models=[BUDGET_MODEL]) + userless_key: Final = scenario.key(models=[BUDGET_MODEL]) + counter: Final = _counter_key(redis) + ttl_before: Final = redis.ttl(counter) + + request_ids: Final = tuple( + _served_id(proxy, key) for key, proxy in product((user_key, userless_key), (first, second)) + ) + total: Final = sum(_logged_spend(first, request_ids)) + persisted: Final = eventually( + lambda: _proxy_budget_spend(first), + lambda spend: spend is not None and spend == pytest.approx(total), + seconds=70, + ) + + assert persisted == pytest.approx(total) + eventually( + lambda: redis.get(counter), + lambda shared: shared is not None and float(str(shared, "utf-8")) == pytest.approx(total), + seconds=70, + ) + assert 0 < redis.ttl(counter) <= ttl_before From 07d314beb598ae4cb5d263fda3e1ecc531633112 Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Wed, 28 Oct 2026 18:05:08 -0700 Subject: [PATCH 123/128] test(responses): audit #47445 (#48902) * test(responses): audit #47445 Integration cells for billing every round of a streamed MCP auto-execute on one spend row, on a native Responses upstream and the chat bridge, across tool rounds, a rejected or broken follow-up, a post-call guardrail block, a client hang-up, the response cache, and the next bridged turn * test(responses): check a broken follow-up bills its own nonzero usage --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../integration/mcp/test_mcp_llm_endpoints.py | 365 +++++++++++++++++- 1 file changed, 351 insertions(+), 14 deletions(-) diff --git a/tests/integration/mcp/test_mcp_llm_endpoints.py b/tests/integration/mcp/test_mcp_llm_endpoints.py index 394221708bf..f751d1f8f4c 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/tests/integration/mcp/test_mcp_llm_endpoints.py @@ -6,7 +6,7 @@ from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass from hashlib import sha256 -from itertools import chain +from itertools import chain, count from pathlib import Path from typing import Final, Literal, TypeVar @@ -50,6 +50,7 @@ def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gat directory: Final = tmp_path_factory.mktemp("mcp-llm-gateway") stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) config: Final = directory / "stored-responses.yaml" + (directory / f"{STREAM_BLOCK_MODULE}.py").write_text(STREAM_BLOCK_SOURCE) config.write_text( yaml.safe_dump( { @@ -57,8 +58,17 @@ def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gat "general_settings": {**object_value(stock["general_settings"]), "store_prompts_in_spend_logs": True}, "litellm_settings": { **object_value(stock["litellm_settings"]), - "cost_margin_config": {"hosted_vllm": {"fixed_amount": FIXED_MARGIN}}, + "cost_margin_config": { + "hosted_vllm": {"fixed_amount": FIXED_MARGIN}, + "deepseek": {"fixed_amount": FIXED_MARGIN}, + }, }, + "guardrails": [ + { + "guardrail_name": STREAM_BLOCK, + "litellm_params": {"guardrail": f"{STREAM_BLOCK_MODULE}.BlockGuard", "mode": "post_call"}, + } + ], } ) ) @@ -79,7 +89,7 @@ def gateway(_llm_gateway: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway yield private -Surface = Literal["chat", "responses", "messages", "messages_bridge"] +Surface = Literal["chat", "responses", "messages", "messages_bridge", "responses_bridge"] SURFACES: Final[tuple[Surface, ...]] = ("chat", "responses", "messages", "messages_bridge") ADD: Final = {"a": 2, "b": 3} ANSWER: Final = "the sum is 5" @@ -88,6 +98,19 @@ AUTO: Final = {**GATEWAY_REF, "require_approval": "never"} OUTAGE: Final = "bridge-outage" DIRECT: Final = "answer-without-tools" FIXED_MARGIN: Final = 0.25 +STREAM_BLOCK: Final = "mcp-stream-block" +STREAM_BLOCK_MODULE: Final = "mcp_stream_block" +STREAM_BLOCK_DETAIL: Final = "blocked by output policy" +STREAM_BLOCK_SOURCE: Final = ( + "from fastapi import HTTPException\n\n" + "from litellm.integrations.custom_guardrail import CustomGuardrail\n\n\n" + "class BlockGuard(CustomGuardrail):\n" + " async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data):\n" + " async for _ in response:\n" + " pass\n" + f' raise HTTPException(status_code=400, detail="{STREAM_BLOCK_DETAIL}")\n' + " yield\n" +) JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) @@ -128,6 +151,8 @@ class Billing: first: tuple[int, int] = (10, 5) follow_up: tuple[int, int] = (10, 5) follow_up_fails: bool = False + follow_up_breaks: bool = False + tool_rounds: int = 1 EVEN_ROUNDS: Final = Billing() @@ -138,6 +163,10 @@ def _fixed_turn(tool: str) -> Callable[[Mapping[str, JsonValue]], Turn]: return lambda _: Turn(tool, json.dumps(ADD), ANSWER) +def _slow_turn(tool: str, seconds: float) -> Callable[[Mapping[str, JsonValue]], Turn]: + return lambda _: Turn(tool, json.dumps({"seconds": seconds}), ANSWER) + + def _echoing_turn(body: Mapping[str, JsonValue]) -> Turn: names: Final = _tool_names(body) return Turn(names[0] if names else "", json.dumps({"query": _prompt(body)}), _tool_result_text(body) or "") @@ -183,6 +212,50 @@ def _item_events(index: int, item: Mapping[str, JsonValue]) -> tuple[dict[str, J ) +def _event_stream(events: Sequence[Mapping[str, JsonValue]]) -> Reply: + return Reply( + content_type="text/event-stream", + chunks=tuple( + f"event: {event['type']}\ndata: {json.dumps({**event, 'sequence_number': number})}\n\n".encode() + for number, event in enumerate(events) + ), + ) + + +def _broken_follow_up(target: str, identity: str, tokens: tuple[int, int]) -> Reply: + outage: Final = {"message": FOLLOW_UP_OUTAGE, "type": "server_error", "code": None} + if target.endswith("/chat/completions"): + delta: Final = {"role": "assistant", "content": "the sum"} + started: Final = { + "id": f"chatcmpl-{identity}", + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "delta": delta, "finish_reason": None}], + } + return Reply( + content_type="text/event-stream", + chunks=tuple(f"data: {json.dumps(frame)}\n\n".encode() for frame in (started, {"error": outage})), + ) + input_tokens, output_tokens = tokens + failed: Final[dict[str, JsonValue]] = { + "id": f"resp_{identity}", + "object": "response", + "created_at": 1, + "status": "failed", + "model": "gpt-4o-mini", + "output": [], + "error": {"code": "server_error", "message": FOLLOW_UP_OUTAGE}, + "usage": {"input_tokens": input_tokens, "output_tokens": output_tokens, "total_tokens": sum(tokens)}, + } + return _event_stream( + ( + {"type": "response.created", "response": {**failed, "status": "in_progress", "error": None, "usage": None}}, + {"type": "response.failed", "response": failed}, + ) + ) + + def _responses_stream(response: Mapping[str, JsonValue], items: Sequence[Mapping[str, JsonValue]]) -> Reply: item_events: Final = tuple(chain.from_iterable(_item_events(index, item) for index, item in enumerate(items))) events: Final = ( @@ -191,13 +264,7 @@ def _responses_stream(response: Mapping[str, JsonValue], items: Sequence[Mapping *item_events, {"type": "response.completed", "response": response}, ) - return Reply( - content_type="text/event-stream", - chunks=tuple( - f"event: {event['type']}\ndata: {json.dumps({**event, 'sequence_number': number})}\n\n".encode() - for number, event in enumerate(events) - ), - ) + return _event_stream(events) def _block_ids(message: JsonValue, kind: str, field: str) -> frozenset[str]: @@ -252,6 +319,8 @@ def _anthropic_tool_result_rule(body: Mapping[str, JsonValue]) -> Reply | None: def _model_double( plan: Callable[[Mapping[str, JsonValue]], Turn], billing: Billing = EVEN_ROUNDS, parallel: int = 1 ) -> Callable[[Request], Reply]: + served: Final = count() + def respond(request: Request) -> Reply: if request.method == "GET" and request.target.endswith("/models"): return _json({"object": "list", "data": []}) @@ -260,11 +329,12 @@ def _model_double( outage: Final = {"error": {"message": "scripted provider outage", "type": "server_error", "code": None}} return Reply(status=500, body=json.dumps(outage).encode()) turn: Final = plan(body) + call_number: Final = next(served) answered: Final = _has_tool_result(body) if answered and billing.follow_up_fails: failure: Final = {"error": {"message": FOLLOW_UP_OUTAGE, "type": "server_error", "code": None}} return Reply(status=500, body=json.dumps(failure).encode()) - done: Final = answered or DIRECT in _prompt(body) + done: Final = DIRECT in _prompt(body) or (answered and call_number >= billing.tool_rounds) calls: Final = range(1, parallel + 1) prompt_tokens, completion_tokens = billing.follow_up if answered else billing.first identity: Final = uuid.uuid4().hex[:12] @@ -273,9 +343,15 @@ def _model_double( "completion_tokens": completion_tokens, "total_tokens": prompt_tokens + completion_tokens, } + if answered and billing.follow_up_breaks: + return _broken_follow_up(request.target, identity, (prompt_tokens, completion_tokens)) if request.target.endswith("/chat/completions"): requested: Final = tuple( - {"id": f"call_{call}", "type": "function", "function": {"name": turn.tool, "arguments": turn.arguments}} + { + "id": f"call_{identity}_{call}", + "type": "function", + "function": {"name": turn.tool, "arguments": turn.arguments}, + } for call in calls ) message: Final = ( @@ -460,6 +536,7 @@ def _upstream_model(surface: Surface) -> str: "responses": "openai/responses/gpt-4o-mini", "messages": "anthropic/claude-sonnet-4-5", "messages_bridge": "hosted_vllm/gpt-4o-mini", + "responses_bridge": "deepseek/gpt-4o-mini", }[surface] @@ -475,11 +552,15 @@ def _rig( billing: Billing = EVEN_ROUNDS, prices: Mapping[str, JsonValue] | None = None, parallel: int = 1, + slow_tool_seconds: float | None = None, ) -> Iterator[Rig]: alias: Final = "llm" + uuid.uuid4().hex[:8] + plan: Final = ( + _fixed_turn(f"{alias}-add") if slow_tool_seconds is None else _slow_turn(f"{alias}-slow", slow_tool_seconds) + ) with ( - mcp_peer() as peer, - wire_server(_model_double(_fixed_turn(f"{alias}-add"), billing, parallel)) as wire, + mcp_peer(rich=slow_tool_seconds is not None) as peer, + wire_server(_model_double(plan, billing, parallel)) as wire, gateway.scenario() as scenario, ): server_id: Final = register_mcp(scenario, peer, alias) @@ -1048,6 +1129,262 @@ def test_bridged_next_turn_replays_the_mcp_turn_answer(gateway: Gateway) -> None assert ANSWER in _assistant_texts(opening), opening +StreamedUpstream = Literal["messages_bridge", "responses_bridge"] +STREAMED_UPSTREAMS: Final[tuple[StreamedUpstream, ...]] = ("messages_bridge", "responses_bridge") +STREAMED_IDS: Final = ("native", "bridged") +SLOW_TOOL_SECONDS: Final = 5.0 + + +def _events(response: httpx.Response) -> tuple[dict[str, JsonValue], ...]: + return tuple( + JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in response.text.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + + +def _completed(response: httpx.Response) -> dict[str, JsonValue]: + finished: Final = tuple( + object_value(event["response"]) for event in _events(response) if event["type"] == "response.completed" + ) + assert len(finished) == 1, response.text + return finished[0] + + +def _model_rows(call_id: str, settled: Callable[[list[dict[str, JsonValue]]], bool]) -> list[dict[str, JsonValue]]: + return eventually( + lambda: read_rows( + "SELECT request_id, status, spend, prompt_tokens, completion_tokens " + 'FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s AND call_type NOT IN (%s, %s)', + (call_id, "call_mcp_tool", "list_mcp_tools"), + ), + settled, + seconds=70, + return_last_on_timeout=True, + ) + + +def _billed(rows: Sequence[Mapping[str, JsonValue]]) -> float: + return sum(float(str(row["spend"])) for row in rows) + + +def _bills_at_least(expected: float) -> Callable[[list[dict[str, JsonValue]]], bool]: + return lambda rows: bool(rows) and _billed(rows) >= expected - 1e-9 + + +def _statuses(rows: Sequence[Mapping[str, JsonValue]]) -> list[str]: + return sorted(str(row["status"]) for row in rows) + + +def _logs_a_failure_billing(expected: float) -> Callable[[list[dict[str, JsonValue]]], bool]: + return lambda rows: "failure" in _statuses(rows) and _billed(rows) >= expected - 1e-9 + + +def _limited_key(rig: Rig) -> str: + return rig.scenario.key(object_permission={"mcp_servers": [rig.server_id]}, max_parallel_requests=1) + + +def _answers_directly(rig: Rig, key: str) -> int: + return rig.send(key, [AUTO], prompt=f"{DIRECT} {rig.alias}").status_code + + +STREAMED_ROUNDS: Final = ( + pytest.param("messages_bridge", 1, id="native"), + pytest.param("responses_bridge", 1, id="bridged"), + pytest.param("messages_bridge", 2, id="native-two-tool-rounds"), + pytest.param("responses_bridge", 2, id="bridged-two-tool-rounds"), + pytest.param("messages_bridge", 0, id="native-no-tool-call"), + pytest.param("responses_bridge", 0, id="bridged-no-tool-call"), +) + + +@pytest.mark.parametrize(("upstream", "tool_rounds"), STREAMED_ROUNDS) +def test_streaming_responses_auto_execute_bills_every_round_once_on_one_success_row( + gateway: Gateway, upstream: StreamedUpstream, tool_rounds: int +) -> None: + rounds: Final = (FIRST_ROUND, *(FOLLOW_UP_ROUND,) * tool_rounds) + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, tool_rounds=max(tool_rounds, 1)) + with _rig(gateway, "responses", upstream=upstream, billing=billing, prices=PRICES) as rig: + prompt: Final = None if tool_rounds else f"{DIRECT} {rig.alias}" + response: Final = rig.send(_granted_key(rig), [AUTO], prompt=prompt, stream=True) + assert response.status_code == 200, response.text + completed: Final = _completed(response) + assert rig.final_text(completed) == ANSWER, response.text + assert len(rig.upstream_tools()) == len(rounds) + assert len(_peer_add_calls(rig.peer)) == tool_rounds + expected: Final = _price(*rounds) + FIXED_MARGIN + rows: Final = _model_rows(response.headers["x-litellm-call-id"], _bills_at_least(expected)) + assert [row["status"] for row in rows] == ["success"], rows + row: Final = rows[0] + assert (row["prompt_tokens"], row["completion_tokens"]) == ( + sum(prompt_tokens for prompt_tokens, _ in rounds), + sum(completion_tokens for _, completion_tokens in rounds), + ), row + usage: Final = object_value(completed["usage"]) + assert (usage["input_tokens"], usage["output_tokens"]) == (row["prompt_tokens"], row["completion_tokens"]) + assert float(str(row["spend"])) == pytest.approx(expected), row + + +@pytest.mark.parametrize("upstream", STREAMED_UPSTREAMS, ids=STREAMED_IDS) +def test_streaming_responses_rejected_follow_up_bills_round_one_once( + gateway: Gateway, upstream: StreamedUpstream +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_fails=True) + with _rig(gateway, "responses", upstream=upstream, billing=billing, prices=PRICES) as rig: + response: Final = rig.send(_granted_key(rig), [AUTO], stream=True) + assert FOLLOW_UP_OUTAGE in response.text, response.text + assert ANSWER not in response.text, response.text + assert len(rig.upstream_tools()) == 2 + expected: Final = _price(FIRST_ROUND) + FIXED_MARGIN + rows: Final = _model_rows(response.headers["x-litellm-call-id"], _bills_at_least(expected)) + assert _billed(rows) == pytest.approx(expected), rows + + +BROKEN_FOLLOW_UPS: Final = ( + pytest.param("messages_bridge", FOLLOW_UP_ROUND, id="native"), + pytest.param("responses_bridge", None, id="bridged"), +) + + +@pytest.mark.parametrize(("upstream", "reported_usage"), BROKEN_FOLLOW_UPS) +def test_streaming_responses_follow_up_broken_mid_stream_bills_round_one_once_and_itself_on_a_failure_row( + gateway: Gateway, upstream: StreamedUpstream, reported_usage: tuple[int, int] | None +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_breaks=True) + with _rig(gateway, "responses", upstream=upstream, billing=billing, prices=PRICES) as rig: + response: Final = rig.send(_granted_key(rig), [AUTO], stream=True) + assert ANSWER not in response.text, response.text + assert len(rig.upstream_tools()) == 2 + round_one: Final = _price(FIRST_ROUND) + FIXED_MARGIN + rows: Final = _model_rows(response.headers["x-litellm-call-id"], _logs_a_failure_billing(round_one)) + assert _statuses(rows) == ["failure", "success"], rows + (success,) = (row for row in rows if row["status"] == "success") + (failure,) = (row for row in rows if row["status"] == "failure") + assert (success["prompt_tokens"], success["completion_tokens"]) == FIRST_ROUND, rows + assert float(str(success["spend"])) == pytest.approx(round_one), rows + failed_round: Final = (int(str(failure["prompt_tokens"])), int(str(failure["completion_tokens"]))) + assert min(failed_round) > 0, rows + assert reported_usage in (None, failed_round), rows + assert float(str(failure["spend"])) == pytest.approx(_price(failed_round)), rows + + +@pytest.mark.parametrize("upstream", STREAMED_UPSTREAMS, ids=STREAMED_IDS) +def test_streaming_responses_blocked_by_a_post_call_guardrail_bills_every_round_once( + gateway: Gateway, upstream: StreamedUpstream +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND) + with _rig(gateway, "responses", upstream=upstream, billing=billing, prices=PRICES) as rig: + key: Final = _granted_key(rig) + one_round: Final = rig.send(key, [], prompt=f"{DIRECT} {rig.alias}", stream=True, guardrails=[STREAM_BLOCK]) + assert STREAM_BLOCK_DETAIL in one_round.text, one_round.text + one_round_rows: Final = _model_rows( + one_round.headers["x-litellm-call-id"], _logs_a_failure_billing(_price(FIRST_ROUND) + FIXED_MARGIN) + ) + rig.wire.drain() + response: Final = rig.send(key, [AUTO], stream=True, guardrails=[STREAM_BLOCK]) + assert STREAM_BLOCK_DETAIL in response.text, response.text + assert len(rig.upstream_tools()) == 2 and len(_peer_add_calls(rig.peer)) == 1 + expected: Final = _price(FIRST_ROUND, FOLLOW_UP_ROUND) + FIXED_MARGIN + rows: Final = _model_rows(response.headers["x-litellm-call-id"], _logs_a_failure_billing(expected)) + assert _statuses(rows) == _statuses(one_round_rows), (rows, one_round_rows) + assert _billed(rows) == pytest.approx(expected), rows + + +@pytest.mark.parametrize( + ("surface", "stream"), + (("responses", True), ("responses", False), ("chat", False)), + ids=("responses-stream", "responses", "chat"), +) +def test_auto_executed_tool_call_runs_inside_the_requests_one_parallel_slot( + gateway: Gateway, surface: Surface, stream: bool +) -> None: + with _rig(gateway, surface, slow_tool_seconds=SLOW_TOOL_SECONDS) as rig: + key: Final = _limited_key(rig) + with ThreadPoolExecutor(max_workers=1) as pool: + running: Final = pool.submit(rig.send, key, [AUTO], None, stream=stream) + eventually(lambda: _peer_add_calls(rig.peer), lambda found: len(found) == 1, seconds=30) + during: Final = _answers_directly(rig, key) + served: Final = running.result(timeout=90) + assert during == 429 + assert served.status_code == 200 and ANSWER in served.text, served.text + eventually(lambda: _answers_directly(rig, key), lambda status: status == 200, seconds=30) + + +@pytest.mark.parametrize("upstream", STREAMED_UPSTREAMS, ids=STREAMED_IDS) +def test_streaming_responses_hung_up_during_the_tool_call_bills_each_served_round_once_and_frees_the_slot( + gateway: Gateway, upstream: StreamedUpstream +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND) + with _rig( + gateway, + "responses", + upstream=upstream, + billing=billing, + prices=PRICES, + slow_tool_seconds=SLOW_TOOL_SECONDS, + ) as rig: + key: Final = _limited_key(rig) + with rig.gateway.client.stream( + "POST", + "/v1/responses", + headers={"Authorization": f"Bearer {key}"}, + json={"model": rig.model, "input": f"add {rig.alias}", "tools": [AUTO], "stream": True}, + timeout=90, + ) as hung_up: + assert hung_up.status_code == 200 + call_id: Final = hung_up.headers["x-litellm-call-id"] + eventually(lambda: _peer_add_calls(rig.peer), lambda found: len(found) == 1, seconds=30) + eventually(lambda: _answers_directly(rig, key), lambda status: status == 200, seconds=60) + posted: Final = tuple( + JSON_OBJECT.validate_json(request.body) for request in rig.wire.drain() if request.method == "POST" + ) + served: Final = tuple( + FOLLOW_UP_ROUND if _has_tool_result(body) else FIRST_ROUND for body in posted if DIRECT not in _prompt(body) + ) + expected: Final = _price(*served) + FIXED_MARGIN + rows: Final = _model_rows(call_id, _bills_at_least(expected)) + assert served and rows, (served, rows) + assert _billed(rows) == pytest.approx(expected), (served, rows) + + +@pytest.mark.parametrize("upstream", STREAMED_UPSTREAMS, ids=STREAMED_IDS) +def test_streaming_responses_repeated_from_the_response_cache_bills_nothing( + gateway: Gateway, upstream: StreamedUpstream +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND) + with _rig(gateway, "responses", upstream=upstream, billing=billing, prices=PRICES) as rig: + key: Final = _granted_key(rig) + paid: Final = rig.send(key, [AUTO], stream=True) + assert paid.status_code == 200 and rig.final_text(_completed(paid)) == ANSWER, paid.text + paid_price: Final = _price(FIRST_ROUND, FOLLOW_UP_ROUND) + FIXED_MARGIN + _model_rows(paid.headers["x-litellm-call-id"], _bills_at_least(paid_price)) + assert len(rig.upstream_tools()) == 2 + cached: Final = rig.send(key, [AUTO], stream=True) + assert cached.status_code == 200 and rig.final_text(_completed(cached)) == ANSWER, cached.text + assert rig.upstream_tools() == () + rows: Final = _model_rows( + cached.headers["x-litellm-call-id"], lambda found: "success" in {row["status"] for row in found} + ) + assert "success" in {row["status"] for row in rows} and _billed(rows) == 0.0, rows + + +def test_bridged_next_turn_after_a_streamed_mcp_turn_replays_its_answer(gateway: Gateway) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND) + with _rig(gateway, "responses", upstream="responses_bridge", billing=billing, prices=PRICES) as rig: + key: Final = _granted_key(rig) + first: Final = rig.send(key, [AUTO], stream=True) + assert first.status_code == 200, first.text + first_id: Final = string_value(_completed(first)["id"]) + _model_rows(first.headers["x-litellm-call-id"], _bills_at_least(_price(FIRST_ROUND, FOLLOW_UP_ROUND))) + rig.wire.drain() + second: Final = rig.send(key, [AUTO], previous_response_id=first_id, cache={"no-cache": True}) + assert second.status_code == 200, second.text + opening: Final = next( + JSON_OBJECT.validate_json(request.body) for request in rig.wire.drain() if request.method == "POST" + ) + assert ANSWER in _assistant_texts(opening), opening + + Bridge = Literal["chat", "responses", "messages"] BRIDGES: Final[tuple[Bridge, ...]] = ("chat", "responses") Client = Literal["sync", "async"] From 3b9a59348f564d7746cb05a9e16341f42321a219 Mon Sep 17 00:00:00 2001 From: zoveki <349028329+zoveki@users.noreply.github.com> Date: Wed, 28 Oct 2026 18:06:58 -0700 Subject: [PATCH 124/128] fix(vertex_ai)!: send the caller's anthropic-version in the Vertex Claude /v1/messages body, not the header (#48814) * fix(vertex_ai): drop the caller's anthropic-version header on Claude /v1/messages Vertex reads the API version from the body's anthropic_version and answers 400 when a forwarded anthropic-version header names a different one, which anthropic SDK 2.x always does. Strip the header on the Python and LITELLM_RUST Vertex messages routes. * fix(vertex_ai): carry the caller's anthropic-version into the Claude /v1/messages body Vertex reads the API version from the body alone, so the caller's anthropic-version header now becomes the body's anthropic_version instead of a pinned vertex-2023-10-16. The 2023-01-01 and 2023-06-01 first-party versions keep vertex-2023-10-16, and the header itself is still never sent. Same change on the Python proxy and the Rust gateway. * fix(vertex_ai): follow drop_params for an anthropic-version Vertex does not accept Only versions Vertex accepts in the body (vertex-2023-10-16, 2026-10-21) are carried into anthropic_version. 2023-06-01, 2023-01-01, a blank header, or no header keep vertex-2023-10-16. Any other value drops to vertex-2023-10-16 under drop_params and answers a 400 naming drop_params without it, on the Python proxy and the Rust gateway. * refactor(vertex_ai): return the unsupported anthropic-version error as a value and keep Rust helpers and tests in place * fix(rust gateway): resolve os.environ drop_params and accept the same flag words as Python The gateway read drop_params: os.environ/NAME and words like yes as off, so the drop_params opt-out from the new anthropic-version 400 did not work there * docs(rust router): name Router::from_config in the crate README --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm-rust/Cargo.lock | 1 + .../crates/core-utils/src/serde_compat.rs | 41 ++++- litellm-rust/crates/gateway/src/lib.rs | 8 +- .../crates/inference-messages/src/handler.rs | 88 ++++++++- .../src/base_llm/messages/transformation.rs | 9 + .../llms/src/vertex_ai/messages/AGENTS.md | 3 +- .../src/vertex_ai/messages/transformation.rs | 173 ++++++++++++++++-- litellm-rust/crates/router/Cargo.toml | 1 + litellm-rust/crates/router/README.md | 2 +- litellm-rust/crates/router/src/lib.rs | 88 ++++++++- litellm-rust/crates/router/tests/router.rs | 8 +- .../anthropic/messages/transformation.py | 59 +++++- ...artner_models_anthropic_messages_config.py | 83 ++++++++- 13 files changed, 522 insertions(+), 42 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 10e590b5006..66a345fc438 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4387,6 +4387,7 @@ name = "litellm-router" version = "0.1.0" dependencies = [ "litellm-config", + "litellm-core-utils", "litellm-inference", "litellm-inference-messages", "litellm-router-types", diff --git a/litellm-rust/crates/core-utils/src/serde_compat.rs b/litellm-rust/crates/core-utils/src/serde_compat.rs index 3e4d82d3a3e..6f40a4af60a 100644 --- a/litellm-rust/crates/core-utils/src/serde_compat.rs +++ b/litellm-rust/crates/core-utils/src/serde_compat.rs @@ -1,13 +1,25 @@ -pub fn parse_str_bool(value: &str) -> Option { - let token = value.trim_matches(|character: char| { +fn python_strip(value: &str) -> &str { + value.trim_matches(|character: char| { character.is_whitespace() || matches!(character, '\u{1c}'..='\u{1f}') - }); + }) +} + +pub fn parse_str_bool(value: &str) -> Option { + let token = python_strip(value); if token.eq_ignore_ascii_case("true") { return Some(true); } token.eq_ignore_ascii_case("false").then_some(false) } +pub fn parse_pydantic_bool(value: &str) -> Option { + match python_strip(value).to_ascii_lowercase().as_str() { + "1" | "on" | "t" | "true" | "y" | "yes" => Some(true), + "0" | "off" | "f" | "false" | "n" | "no" => Some(false), + _ => None, + } +} + /// `redis-py` string Booleans: only `1`, `true`, and `yes` (case-insensitive) are true. pub fn parse_redis_bool(value: &str) -> bool { value == "1" || value.eq_ignore_ascii_case("true") || value.eq_ignore_ascii_case("yes") @@ -34,4 +46,27 @@ mod tests { ) { assert_eq!(parse_str_bool(input), expected, "{input:?}"); } + + #[rstest] + #[case::yes(" YES ", Some(true))] + #[case::on("On", Some(true))] + #[case::one("1", Some(true))] + #[case::t("t", Some(true))] + #[case::y("y", Some(true))] + #[case::control_whitespace_true("\u{1c}TRUE\u{1f}", Some(true))] + #[case::no("No", Some(false))] + #[case::off("OFF", Some(false))] + #[case::zero("0", Some(false))] + #[case::f("F", Some(false))] + #[case::n("n", Some(false))] + #[case::false_word("false", Some(false))] + #[case::two("2", None)] + #[case::empty("", None)] + #[case::unknown("maybe", None)] + fn boolean_words_follow_pydantic_lax_bool_parsing( + #[case] input: &str, + #[case] expected: Option, + ) { + assert_eq!(parse_pydantic_bool(input), expected, "{input:?}"); + } } diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index b8ff8e46033..12a9bf9110a 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -37,6 +37,7 @@ pub fn build_inference(config: &Config) -> Result, Error> { .map(|value| value.expose().to_owned()) .or_else(|| std::env::var(name).ok()) }; + let router = ModelRouter::from_config(config, &lookup); let settings = &config.litellm_settings; let http_settings = HttpSettings::from_layers([ HttpSettingsLayer::from_environment(&lookup), @@ -64,12 +65,7 @@ pub fn build_inference(config: &Config) -> Result, Error> { Arc::new(EnvironmentSecrets::python_compatible(client)), )); let resources = CoreResources::new(pool); - Ok(Arc::new(Gateway::new( - resources, - http, - secrets, - ModelRouter::from_model_list(&config.model_list), - )?)) + Ok(Arc::new(Gateway::new(resources, http, secrets, router)?)) } pub async fn build_redis(config: &Config) -> Result { diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index fa802d20587..9d0bf09197e 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -141,9 +141,13 @@ async fn provider_call( &env_lookup, )?; - let body = config.transform_anthropic_messages_request( - trimmed, - &MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params), + let body = config.request_body( + config.transform_anthropic_messages_request( + trimmed, + &MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params), + )?, + &validated.headers, + shaping.settings.drop_params, )?; let environment = ValidatedEnvironment { headers: config.request_headers( @@ -514,6 +518,84 @@ mod tests { ); } + fn vertex_call_requesting( + shaping: MessagesShaping, + requested: &str, + drop_params: bool, + ) -> MessagesCall { + MessagesCall { + body: body( + json!({"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}), + ), + api_key: None, + api_base: Some("https://vertex.test".into()), + custom_llm_provider: Some("vertex_ai".into()), + litellm_params: Default::default(), + extra_headers: Some(Map::from_iter([ + ("authorization".to_string(), json!("Bearer vertex-token")), + ("anthropic-version".to_string(), json!(requested)), + ])), + provider_specific_header: None, + timeout: None, + shaping: MessagesShaping { + settings: MessagesSettings { + drop_params, + ..shaping.settings + }, + ..shaping + }, + } + } + + fn vertex_project(name: &str) -> Option { + (name == "VERTEXAI_PROJECT").then(|| "version-project".to_string()) + } + + #[rstest] + #[case::sdk_2("2026-10-21", false, "2026-10-21")] + #[case::sdk_1("2023-06-01", false, "vertex-2023-10-16")] + #[case::unknown_under_drop_params("latest", true, "vertex-2023-10-16")] + #[tokio::test] + async fn vertex_gets_the_requested_version_in_the_body_alone( + shaping: MessagesShaping, + #[case] requested: &str, + #[case] drop_params: bool, + #[case] body_version: &str, + ) { + let call = provider_call_with_secrets( + vertex_call_requesting(shaping, requested, drop_params), + &vertex_project, + ) + .await + .unwrap(); + let version_headers: Vec<&str> = call + .wire + .headers + .iter() + .filter(|(name, _)| name.eq_ignore_ascii_case("anthropic-version")) + .map(|(_, value)| value.as_str()) + .collect(); + assert_eq!( + (version_headers, &call.wire.body["anthropic_version"]), + (Vec::<&str>::new(), &json!(body_version)) + ); + } + + #[rstest] + #[tokio::test] + async fn vertex_refuses_an_unsupported_version_without_drop_params(shaping: MessagesShaping) { + assert!(matches!( + provider_call_with_secrets( + vertex_call_requesting(shaping, "latest", false), + &vertex_project + ) + .await, + Err(Error::Llm(litellm_llms::Error::InvalidRequest( + litellm_llms::ErrorDetail::UnsupportedParameter { .. } + ))) + )); + } + #[rstest] #[tokio::test] async fn missing_credentials_are_reported_before_the_request_is_transformed( diff --git a/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs b/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs index 8a88af40d6b..744ce89afdb 100644 --- a/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs @@ -77,6 +77,15 @@ pub trait BaseMessagesConfig: Sync { headers } + fn request_body( + &self, + request: MessagesRequest, + _headers: &Headers, + _drop_params: bool, + ) -> Result { + Ok(request) + } + /// The JSON that goes on the wire, for a host whose body differs from the typed request: /// Python's configs `pop("model")` when the model is addressed by the URL. fn wire_body(&self, body: Value) -> Value { diff --git a/litellm-rust/crates/llms/src/vertex_ai/messages/AGENTS.md b/litellm-rust/crates/llms/src/vertex_ai/messages/AGENTS.md index b5d637d6bc7..d95da65ac9d 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/messages/AGENTS.md +++ b/litellm-rust/crates/llms/src/vertex_ai/messages/AGENTS.md @@ -1,6 +1,7 @@ # rules -- Claude on Vertex AI: reuses `anthropic/messages` shaping, moves the model into the publisher URL, puts the Vertex `anthropic_version` in the body, and authenticates with a Google access token +- Claude on Vertex AI: reuses `anthropic/messages` shaping, moves the model into the publisher URL, and authenticates with a Google access token +- Vertex reads the API version only from the body's `anthropic_version` and answers 400 when an `anthropic-version` header names a different one, so the caller's header never reaches Vertex. A version Vertex accepts in the body is carried there, `2023-06-01`, `2023-01-01`, or no header get `vertex-2023-10-16`, and any other version drops under `drop_params` and is refused without it - Applies Vertex's beta policy and strips `output_config.effort` for models that do not accept it # references diff --git a/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs index 2e77dbc1ca8..46e5a68dcb8 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs @@ -3,10 +3,10 @@ use litellm_auth_gcp::{ VertexConfig, get_vertex_ai_location, get_vertex_ai_project, get_vertex_ai_project_from_credentials, secret_names, }; -use litellm_http::request::has_header; +use litellm_http::request::{has_header, header_value, without_headers}; use litellm_llms_types::{ formats::messages::{MessagesOptionalParams, MessagesRequest, MessagesTool, OutputConfig}, - providers::anthropic::{AnthropicBeta, BetaProvider, BetaSet}, + providers::anthropic::{AnthropicBeta, BetaProvider, BetaSet, VERSION_HEADER}, recognized::Recognized, }; use litellm_router_types::LitellmParams; @@ -36,6 +36,8 @@ use litellm_llms_types::providers::vertex_ai::{ }; const WEB_SEARCH_TOOL_PREFIX: &str = "web_search"; +const VERSIONS_SERVED_AS_VERTEX_DEFAULT: [&str; 3] = ["", "2023-01-01", "2023-06-01"]; +const VERSIONS_VERTEX_ACCEPTS_IN_BODY: [&str; 2] = [ANTHROPIC_VERSION, "2026-10-21"]; /// Claude on Vertex AI, Python's `VertexAIPartnerModelsAnthropicMessagesConfig`: the Anthropic /// payload addressed to a project and location, authenticated with a Google access token. @@ -90,12 +92,6 @@ impl BaseMessagesConfig for VertexAiPartnerModelsAnthropicMessagesConfig { .params .output_config .and_then(|config| sanitize_output_config(config, accepts_effort)), - extra: request - .params - .extra - .into_iter() - .chain([("anthropic_version".to_string(), json!(ANTHROPIC_VERSION))]) - .collect(), ..request.params }, ..request @@ -133,7 +129,35 @@ impl BaseMessagesConfig for VertexAiPartnerModelsAnthropicMessagesConfig { } fn request_headers(&self, headers: Headers, request: &MessagesRequest) -> Headers { - merge_beta_headers(headers, vertex_feature_betas(request)) + merge_beta_headers( + without_headers(headers, &[VERSION_HEADER]), + vertex_feature_betas(request), + ) + } + + fn request_body( + &self, + request: MessagesRequest, + headers: &Headers, + drop_params: bool, + ) -> Result { + let version = vertex_version( + &request.model, + header_value(headers, VERSION_HEADER), + drop_params, + )?; + Ok(MessagesRequest { + params: MessagesOptionalParams { + extra: request + .params + .extra + .into_iter() + .chain([("anthropic_version".to_string(), json!(version))]) + .collect(), + ..request.params + }, + ..request + }) } /// Vertex addresses the model in the URL and rejects it in the body. @@ -225,6 +249,26 @@ fn vertex_feature_betas(request: &MessagesRequest) -> BetaSet { ) } +fn vertex_version<'a>( + model: &str, + requested: Option<&'a str>, + drop_params: bool, +) -> Result<&'a str, Error> { + match requested.map(str::trim).unwrap_or_default() { + version if VERSIONS_SERVED_AS_VERTEX_DEFAULT.contains(&version) => Ok(ANTHROPIC_VERSION), + version if VERSIONS_VERTEX_ACCEPTS_IN_BODY.contains(&version) => Ok(version), + _ if drop_params => Ok(ANTHROPIC_VERSION), + version => Err(Error::InvalidRequest( + crate::ErrorDetail::UnsupportedParameter { + model: model.into(), + param: VERSION_HEADER.into(), + value: version.into(), + hint: String::new(), + }, + )), + } +} + #[cfg(test)] mod tests { use rstest::rstest; @@ -471,6 +515,28 @@ mod tests { ); } + #[rstest] + #[case::sdk_spelling("anthropic-version")] + #[case::any_case("Anthropic-Version")] + fn a_forwarded_anthropic_version_never_reaches_vertex(#[case] name: &str) { + let request: MessagesRequest = serde_json::from_value(json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}], + })) + .unwrap(); + assert_eq!( + VERTEX_ANTHROPIC_MESSAGES_CONFIG.request_headers( + vec![ + (name.to_string(), "2026-10-21".to_string()), + ("x-request-id".to_string(), "req-1".to_string()), + ], + &request, + ), + vec![("x-request-id".to_string(), "req-1".to_string())] + ); + } + fn transformed(value: Value, context: &MessagesTransformContext) -> Value { let request: MessagesRequest = serde_json::from_value(value).unwrap(); serde_json::to_value( @@ -482,16 +548,87 @@ mod tests { } #[rstest] - fn the_body_carries_the_vertex_version_and_the_wire_drops_the_model() { - let body = transformed( - json!({ - "model": "claude-sonnet-4-5@20250929", - "max_tokens": 16, - "messages": [{"role": "user", "content": "hi"}] - }), - &MessagesTransformContext::default(), + #[case::sdk_2(&[("anthropic-version", "2026-10-21")], false, "2026-10-21")] + #[case::any_case(&[("Anthropic-Version", "2026-10-21")], false, "2026-10-21")] + #[case::vertex_name(&[("anthropic-version", "vertex-2023-10-16")], false, "vertex-2023-10-16")] + #[case::sdk_1(&[("anthropic-version", "2023-06-01")], false, "vertex-2023-10-16")] + #[case::first_version(&[("anthropic-version", "2023-01-01")], false, "vertex-2023-10-16")] + #[case::blank(&[("anthropic-version", " ")], false, "vertex-2023-10-16")] + #[case::no_header(&[], false, "vertex-2023-10-16")] + #[case::unknown_under_drop_params(&[("anthropic-version", "latest")], true, "vertex-2023-10-16")] + fn the_requested_version_goes_in_the_body( + #[case] forwarded: &[(&str, &str)], + #[case] drop_params: bool, + #[case] body_version: &str, + ) { + let request: MessagesRequest = serde_json::from_value(json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}], + })) + .unwrap(); + let headers: Headers = forwarded + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect(); + let body = serde_json::to_value( + VERTEX_ANTHROPIC_MESSAGES_CONFIG + .request_body(request, &headers, drop_params) + .unwrap(), + ) + .unwrap(); + assert_eq!(body["anthropic_version"], json!(body_version)); + } + + #[rstest] + #[case::unreleased("2025-01-01")] + #[case::alias("latest")] + fn an_unsupported_version_is_refused_without_drop_params(#[case] requested: &str) { + let request: MessagesRequest = serde_json::from_value(json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}], + })) + .unwrap(); + let headers: Headers = vec![("anthropic-version".to_string(), requested.to_string())]; + let Error::InvalidRequest(detail) = VERTEX_ANTHROPIC_MESSAGES_CONFIG + .request_body(request, &headers, false) + .unwrap_err() + else { + panic!("expected an invalid request"); + }; + assert_eq!( + detail.to_string(), + format!( + "claude-sonnet-4-5 does not support anthropic-version={requested}. \ + To drop unsupported params, set `litellm.drop_params = True`." + ) ); - assert_eq!(body["anthropic_version"], json!("vertex-2023-10-16")); + } + + #[rstest] + fn the_body_carries_the_vertex_version_and_the_wire_drops_the_model() { + let request: MessagesRequest = serde_json::from_value(json!({ + "model": "claude-sonnet-4-5@20250929", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}] + })) + .unwrap(); + let body = serde_json::to_value( + VERTEX_ANTHROPIC_MESSAGES_CONFIG + .request_body( + VERTEX_ANTHROPIC_MESSAGES_CONFIG + .transform_anthropic_messages_request( + request, + &MessagesTransformContext::default(), + ) + .unwrap(), + &Vec::new(), + false, + ) + .unwrap(), + ) + .unwrap(); assert_eq!(body["model"], json!("claude-sonnet-4-5@20250929")); let wire = VERTEX_ANTHROPIC_MESSAGES_CONFIG.wire_body(body); assert_eq!(wire.get("model"), None); diff --git a/litellm-rust/crates/router/Cargo.toml b/litellm-rust/crates/router/Cargo.toml index a356bedcf16..988d1b68805 100644 --- a/litellm-rust/crates/router/Cargo.toml +++ b/litellm-rust/crates/router/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true [dependencies] litellm-config.workspace = true +litellm-core-utils.workspace = true litellm-inference.workspace = true litellm-inference-messages.workspace = true litellm-router-types.workspace = true diff --git a/litellm-rust/crates/router/README.md b/litellm-rust/crates/router/README.md index 972d37d14be..a734e634104 100644 --- a/litellm-rust/crates/router/README.md +++ b/litellm-rust/crates/router/README.md @@ -1,4 +1,4 @@ -`litellm-router` scaffolds the model-list setup and deployment lookup portion of Python's `litellm.Router`. `Router::from_model_list(&config.model_list)` maps configured public names to provider deployments. Programmatic callers can collect `(String, Deployment)` entries into a `Router` +`litellm-router` scaffolds the model-list setup and deployment lookup portion of Python's `litellm.Router`. `Router::from_config(&config, &env)` maps configured public names to provider deployments and reads `drop_params` from `litellm_settings` and each deployment, resolving an `os.environ/NAME` value through `env`. Programmatic callers can collect `(String, Deployment)` entries into a `Router` Lookup is exact and returns `None` for an unknown name. This extraction preserves the gateway's existing behavior: the last entry wins when public names repeat. Multiple deployments per model group, routing strategies, retries, cooldowns, and fallbacks are not implemented yet diff --git a/litellm-rust/crates/router/src/lib.rs b/litellm-rust/crates/router/src/lib.rs index fb48213e8a7..b75c1a14446 100644 --- a/litellm-rust/crates/router/src/lib.rs +++ b/litellm-rust/crates/router/src/lib.rs @@ -2,16 +2,33 @@ mod deployment; use std::collections::HashMap; -use litellm_config::Model; +use litellm_config::{Config, Spelled}; +use litellm_core_utils::{serde_compat::parse_pydantic_bool, settings::Lookup}; +use litellm_inference_messages::{MessagesSettings, MessagesShaping}; pub use deployment::Deployment; #[derive(Clone, Debug, Default)] pub struct Router(HashMap); +fn is_on(flag: Option<&Spelled>, env: &impl Lookup) -> bool { + match flag { + Some(Spelled::Value(on)) => *on, + Some(Spelled::Text(text)) => match text.strip_prefix("os.environ/") { + Some(name) => env + .get(name) + .is_some_and(|value| parse_pydantic_bool(&value) == Some(true)), + None => parse_pydantic_bool(text) == Some(true), + }, + None => false, + } +} + impl Router { - pub fn from_model_list(model_list: &[Model]) -> Self { - model_list + pub fn from_config(config: &Config, env: &impl Lookup) -> Self { + let drops_params_globally = is_on(config.litellm_settings.drop_params.as_ref(), env); + config + .model_list .iter() .map(|model| { ( @@ -26,6 +43,14 @@ impl Router { api_base: model.litellm_params.api_base.clone(), custom_llm_provider: model.litellm_params.custom_llm_provider.clone(), litellm_params: model.litellm_params.clone(), + shaping: MessagesShaping { + settings: MessagesSettings { + drop_params: drops_params_globally + || is_on(model.litellm_params.drop_params.as_ref(), env), + ..MessagesSettings::default() + }, + ..MessagesShaping::default() + }, ..Deployment::default() }, ) @@ -43,3 +68,60 @@ impl FromIterator<(String, Deployment)> for Router { Self(entries.into_iter().collect()) } } + +#[cfg(test)] +mod tests { + use litellm_config::Config; + use rstest::rstest; + + use super::Router; + + #[rstest] + #[case::unset("", "", false)] + #[case::deployment_flag("", "\n drop_params: true", true)] + #[case::deployment_word("", "\n drop_params: \" True \"", true)] + #[case::deployment_off("", "\n drop_params: false", false)] + #[case::deployment_unknown_word("", "\n drop_params: maybe", false)] + #[case::deployment_yes("", "\n drop_params: yes", true)] + #[case::deployment_no("", "\n drop_params: \"no\"", false)] + #[case::deployment_env_reference("", "\n drop_params: os.environ/DROP_PARAMS_FLAG", true)] + #[case::deployment_unset_env_reference( + "", + "\n drop_params: os.environ/MISSING_FLAG", + false + )] + #[case::global_env_reference( + "litellm_settings:\n drop_params: os.environ/DROP_PARAMS_FLAG\n", + "", + true + )] + #[case::global_flag("litellm_settings:\n drop_params: true\n", "", true)] + #[case::global_word("litellm_settings:\n drop_params: \"true\"\n", "", true)] + #[case::global_off("litellm_settings:\n drop_params: false\n", "", false)] + #[case::global_over_deployment_off( + "litellm_settings:\n drop_params: true\n", + "\n drop_params: false", + true + )] + fn configured_drop_params_reach_the_deployment( + #[case] settings: &str, + #[case] deployment: &str, + #[case] expected: bool, + ) { + let config = Config::from_yaml(&format!( + "{settings}model_list:\n - model_name: public-model\n litellm_params:\n model: provider/model{deployment}" + )) + .unwrap(); + + assert_eq!( + Router::from_config(&config, &|name: &str| (name == "DROP_PARAMS_FLAG") + .then(|| " Yes ".to_owned())) + .get("public-model") + .unwrap() + .shaping + .settings + .drop_params, + expected + ); + } +} diff --git a/litellm-rust/crates/router/tests/router.rs b/litellm-rust/crates/router/tests/router.rs index bbaeb806b96..03e78ab9d11 100644 --- a/litellm-rust/crates/router/tests/router.rs +++ b/litellm-rust/crates/router/tests/router.rs @@ -16,7 +16,7 @@ fn configuration_preserves_deployment_parameters(#[case] parameters: &str) { "model_list:\n - model_name: public-model\n litellm_params:\n model: provider/model\n {parameters}" )) .unwrap(); - let router = Router::from_model_list(&config.model_list); + let router = Router::from_config(&config, &|_: &str| None); let deployment = router.get(&config.model_list[0].model_name).unwrap(); let params = &config.model_list[0].litellm_params; @@ -44,7 +44,7 @@ fn configuration_types_the_provider_params_of_a_deployment() { azure_ad_token: ignored-by-every-group", ) .unwrap(); - let router = Router::from_model_list(&config.model_list); + let router = Router::from_config(&config, &|_: &str| None); let params = &router.get("public-model").unwrap().litellm_params; assert_eq!( @@ -74,7 +74,7 @@ fn lookup_uses_public_names(#[case] name: &str, #[case] expected: Option<&str>) model: provider/b", ) .unwrap(); - let router = Router::from_model_list(&config.model_list); + let router = Router::from_config(&config, &|_: &str| None); assert_eq!(router.get(name).map(|entry| entry.model.as_str()), expected); } @@ -84,7 +84,7 @@ fn empty_configuration_has_no_deployment() { let config = Config::from_yaml("model_list: []").unwrap(); assert!( - Router::from_model_list(&config.model_list) + Router::from_config(&config, &|_: &str| None) .get("") .is_none() ); diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py index 6d8eb228786..1f0ee728cd1 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py @@ -1,6 +1,8 @@ from collections.abc import Mapping, Sequence from typing import Final +from litellm.exceptions import UnsupportedParamsError +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.anthropic.messages.native.transformation import ( AnthropicMessagesConfig, @@ -17,6 +19,41 @@ from litellm.types.router import GenericLiteLLMParams from ....vertex_llm_base import VertexBase from ..output_params_utils import sanitize_vertex_anthropic_output_params +_VERSION_HEADER_VERTEX_READS_FROM_BODY: Final = "anthropic-version" +_VERTEX_DEFAULT_VERSION: Final = "vertex-2023-10-16" +_VERSIONS_SERVED_AS_VERTEX_DEFAULT: Final = frozenset({"", "2023-01-01", "2023-06-01"}) +_VERSIONS_VERTEX_ACCEPTS_IN_BODY: Final = frozenset({_VERTEX_DEFAULT_VERSION, "2026-10-21"}) + + +def _requested_version(headers: Mapping[str, object]) -> str: + requested: Final = next( + (value for name, value in headers.items() if name.lower() == _VERSION_HEADER_VERTEX_READS_FROM_BODY), + "", + ) + return requested.strip() if isinstance(requested, str) else "" + + +def _vertex_version(headers: Mapping[str, object], model: str, drop_params: bool) -> str | UnsupportedParamsError: + requested: Final = _requested_version(headers) + if requested in _VERSIONS_SERVED_AS_VERTEX_DEFAULT: + return _VERTEX_DEFAULT_VERSION + if requested in _VERSIONS_VERTEX_ACCEPTS_IN_BODY: + return requested + if drop_params: + return _VERTEX_DEFAULT_VERSION + return UnsupportedParamsError( + message=( + f"{model} does not support {_VERSION_HEADER_VERTEX_READS_FROM_BODY}={requested}. " + "To drop unsupported params, set `litellm.drop_params = True`." + ), + llm_provider="vertex_ai", + model=model, + ) + + +def _without_version_header(headers: Mapping[str, str]) -> dict[str, str]: # mutable-ok: the caller sends these headers + return {name: value for name, value in headers.items() if name.lower() != _VERSION_HEADER_VERTEX_READS_FROM_BODY} + class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, VertexBase): @property @@ -198,7 +235,7 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert messages: list[dict], anthropic_messages_optional_request_params: dict, litellm_params: GenericLiteLLMParams, - headers: dict, + headers: dict[str, object], # mutable-ok: mirrors the provider transform contract ) -> dict: anthropic_messages_request: Final = super().transform_anthropic_messages_request( model=model, @@ -212,10 +249,28 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert self._remove_scope_from_cache_control(anthropic_messages_request) - anthropic_messages_request["anthropic_version"] = "vertex-2023-10-16" + version: Final = _vertex_version( + headers, model=model, drop_params=drops_unsupported_params(litellm_params.drop_params) + ) + if isinstance(version, UnsupportedParamsError): + raise version + anthropic_messages_request["anthropic_version"] = version anthropic_messages_request.pop("model", None) # do not pass model in request body to vertex ai sanitize_vertex_anthropic_output_params(anthropic_messages_request, model) return anthropic_messages_request + + def sign_request( + self, + headers: Mapping[str, str], + optional_params: Mapping[str, object], + request_data: Mapping[str, object], + api_base: str, + api_key: str | None = None, + model: str | None = None, + stream: bool | None = None, + fake_stream: bool | None = None, + ) -> tuple[dict[str, str], bytes | None]: # mutable-ok: the caller sends these headers + return _without_version_header(headers), None diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py index 6c10d524531..5916ce9f7a1 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py @@ -6,12 +6,14 @@ from datetime import datetime, timedelta from queue import SimpleQueue from threading import get_ident from typing import Final -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest import requests import responses +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.messages.transformation import ( VertexAIPartnerModelsAnthropicMessagesConfig, ) @@ -94,6 +96,85 @@ async def test_messages_auth_load_and_refresh_leave_the_event_loop(stream: bool, assert headers == {"Authorization": "Bearer incoming", "anthropic-beta": "retained-beta", "x-request-id": "req-1"} +def _vertex_client() -> MagicMock: + vertex_response: Final = MagicMock() + vertex_response.json.return_value = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "hi"}], + "model": "claude-sonnet-5", + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + } + client: Final = MagicMock(spec=AsyncHTTPHandler) + client.post = AsyncMock(return_value=vertex_response) + return client + + +async def _send_to_vertex(client: MagicMock, caller_headers: Mapping[str, str], drop_params: bool) -> None: + with responses.RequestsMock(assert_all_requests_are_fired=False) as oauth: + oauth.add(responses.POST, "https://oauth2.googleapis.com/token", json={"access_token": "vertex-token", "expires_in": 3600}) + await litellm.anthropic.messages.acreate( + model="vertex_ai/claude-sonnet-5", + messages=[{"role": "user", "content": "hi"}], + max_tokens=16, + vertex_ai_project="version-project", + vertex_ai_location="global", + vertex_credentials=json.dumps( + { + "type": "authorized_user", + "client_id": "version-header-client", + "client_secret": "synthetic-secret", + "refresh_token": "synthetic-refresh", + } + ), + client=client, + drop_params=drop_params, + provider_specific_header={ + "custom_llm_provider": "vertex_ai", + "extra_headers": {**caller_headers, "x-request-id": "req-1"}, + }, + ) + + +@pytest.mark.parametrize( + ("caller_headers", "drop_params", "body_version"), + ( + ({"anthropic-version": "2026-10-21"}, False, "2026-10-21"), + ({"Anthropic-Version": "2026-10-21"}, False, "2026-10-21"), + ({"anthropic-version": "vertex-2023-10-16"}, False, "vertex-2023-10-16"), + ({"anthropic-version": "2023-06-01"}, False, "vertex-2023-10-16"), + ({"anthropic-version": "2023-01-01"}, False, "vertex-2023-10-16"), + ({"anthropic-version": " "}, False, "vertex-2023-10-16"), + ({}, False, "vertex-2023-10-16"), + ({"anthropic-version": "latest"}, True, "vertex-2023-10-16"), + ), +) +@pytest.mark.asyncio +async def test_caller_anthropic_version_reaches_vertex_only_in_the_body( + caller_headers: dict[str, str], drop_params: bool, body_version: str +) -> None: + client: Final = _vertex_client() + await _send_to_vertex(client, caller_headers, drop_params) + + sent: Final = client.post.call_args.kwargs + sent_header_names: Final = {name.lower() for name in sent["headers"]} + assert "anthropic-version" not in sent_header_names + assert sent["headers"]["x-request-id"] == "req-1" + assert json.loads(sent["data"])["anthropic_version"] == body_version + + +@pytest.mark.parametrize("requested", ("2025-01-01", "latest")) +@pytest.mark.asyncio +async def test_unsupported_anthropic_version_is_refused_without_drop_params(requested: str) -> None: + client: Final = _vertex_client() + with pytest.raises(litellm.BadRequestError, match=f"anthropic-version={requested}. To drop unsupported params"): + await _send_to_vertex(client, {"anthropic-version": requested}, drop_params=False) + client.post.assert_not_awaited() + + def test_validate_environment_uses_vertex_ai_location(): config = VertexAIPartnerModelsAnthropicMessagesConfig() headers = {} From 57c1df8c33b01bbee21b6496649428124635301b Mon Sep 17 00:00:00 2001 From: yujonglee Date: Wed, 28 Oct 2026 18:24:40 -0700 Subject: [PATCH 125/128] refactor(model-catalog): parse, lower and validate the catalog into typed prices (#48853) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-rust.yml | 2 + .github/workflows/test-unit.yml | 2 +- litellm-rust/Cargo.lock | 34 +- litellm-rust/Cargo.toml | 4 + litellm-rust/crates/cost/Cargo.toml | 15 +- litellm-rust/crates/cost/README.md | 12 +- litellm-rust/crates/cost/benches/calculate.rs | 66 -- litellm-rust/crates/cost/examples/charge.rs | 40 - litellm-rust/crates/cost/src/call.rs | 42 + litellm-rust/crates/cost/src/error.rs | 47 +- litellm-rust/crates/cost/src/lib.rs | 388 +------- litellm-rust/crates/cost/src/money.rs | 257 ++++++ litellm-rust/crates/cost/src/price.rs | 505 +++++++++++ litellm-rust/crates/cost/src/price/meter.rs | 197 ++++ litellm-rust/crates/cost/src/priced.rs | 53 ++ litellm-rust/crates/cost/src/schedule.rs | 241 +++++ litellm-rust/crates/cost/src/usage.rs | 104 +++ litellm-rust/crates/cost/tests/calculation.rs | 535 ----------- .../cost/tests/generate_python_reference.py | 96 -- .../crates/cost/tests/python_reference.tsv | 9 - .../crates/model-catalog-refresh/Cargo.toml | 19 + .../crates/model-catalog-refresh/src/error.rs | 105 +++ .../crates/model-catalog-refresh/src/lib.rs | 9 + .../model-catalog-refresh/src/policy.rs | 20 + .../model-catalog-refresh/src/source.rs | 17 + .../crates/model-catalog-refresh/src/step.rs | 66 ++ litellm-rust/crates/model-catalog/AGENTS.md | 28 +- litellm-rust/crates/model-catalog/Cargo.toml | 10 +- .../crates/model-catalog/src/bundled.rs | 3 + .../crates/model-catalog/src/capabilities.rs | 67 -- .../crates/model-catalog/src/catalog.rs | 485 ++++++---- .../crates/model-catalog/src/error.rs | 391 +++++++- .../crates/model-catalog/src/fallback.rs | 9 +- .../crates/model-catalog/src/layers.rs | 32 + litellm-rust/crates/model-catalog/src/lib.rs | 46 +- .../crates/model-catalog/src/lower.rs | 855 ++++++++++++++++++ .../crates/model-catalog/src/model.rs | 22 + .../crates/model-catalog/src/model_info.rs | 734 ++++----------- .../crates/model-catalog/src/pricing.rs | 97 -- .../crates/model-catalog/src/register.rs | 26 + .../crates/model-catalog/src/resolve/image.rs | 18 + .../crates/model-catalog/src/resolve/mod.rs | 46 + .../crates/model-catalog/src/schema.rs | 197 ++-- .../crates/model-catalog/src/store/AGENTS.md | 35 + .../crates/model-catalog/src/store/mod.rs | 460 ++++++++++ .../model-catalog/src/store/name_index.rs | 35 + .../crates/model-catalog/src/validation.rs | 252 ------ .../crates/model-catalog/src/vocabulary.rs | 126 +++ .../model-catalog/src/wire/attributes.rs | 344 +++++++ .../crates/model-catalog/src/wire/blocks.rs | 182 ++++ .../crates/model-catalog/src/wire/entry.rs | 193 ++++ .../crates/model-catalog/src/wire/flat_key.rs | 359 ++++++++ .../crates/model-catalog/src/wire/lax.rs | 128 +++ .../crates/model-catalog/src/wire/mod.rs | 10 + .../crates/model-catalog/tests/catalog.rs | 632 +++++++------ .../tests/generate_price_lookup_reference.py | 91 ++ .../model-catalog/tests/price_lookup.rs | 89 ++ .../tests/price_lookup_reference.json | 89 ++ .../tests/registry_validation.rs | 93 -- .../crates/model-catalog/tests/schema.rs | 73 +- 60 files changed, 6345 insertions(+), 2797 deletions(-) delete mode 100644 litellm-rust/crates/cost/benches/calculate.rs delete mode 100644 litellm-rust/crates/cost/examples/charge.rs create mode 100644 litellm-rust/crates/cost/src/call.rs create mode 100644 litellm-rust/crates/cost/src/money.rs create mode 100644 litellm-rust/crates/cost/src/price.rs create mode 100644 litellm-rust/crates/cost/src/price/meter.rs create mode 100644 litellm-rust/crates/cost/src/priced.rs create mode 100644 litellm-rust/crates/cost/src/schedule.rs create mode 100644 litellm-rust/crates/cost/src/usage.rs delete mode 100644 litellm-rust/crates/cost/tests/calculation.rs delete mode 100644 litellm-rust/crates/cost/tests/generate_python_reference.py delete mode 100644 litellm-rust/crates/cost/tests/python_reference.tsv create mode 100644 litellm-rust/crates/model-catalog-refresh/Cargo.toml create mode 100644 litellm-rust/crates/model-catalog-refresh/src/error.rs create mode 100644 litellm-rust/crates/model-catalog-refresh/src/lib.rs create mode 100644 litellm-rust/crates/model-catalog-refresh/src/policy.rs create mode 100644 litellm-rust/crates/model-catalog-refresh/src/source.rs create mode 100644 litellm-rust/crates/model-catalog-refresh/src/step.rs create mode 100644 litellm-rust/crates/model-catalog/src/bundled.rs delete mode 100644 litellm-rust/crates/model-catalog/src/capabilities.rs create mode 100644 litellm-rust/crates/model-catalog/src/layers.rs create mode 100644 litellm-rust/crates/model-catalog/src/lower.rs create mode 100644 litellm-rust/crates/model-catalog/src/model.rs delete mode 100644 litellm-rust/crates/model-catalog/src/pricing.rs create mode 100644 litellm-rust/crates/model-catalog/src/register.rs create mode 100644 litellm-rust/crates/model-catalog/src/resolve/image.rs create mode 100644 litellm-rust/crates/model-catalog/src/resolve/mod.rs create mode 100644 litellm-rust/crates/model-catalog/src/store/AGENTS.md create mode 100644 litellm-rust/crates/model-catalog/src/store/mod.rs create mode 100644 litellm-rust/crates/model-catalog/src/store/name_index.rs delete mode 100644 litellm-rust/crates/model-catalog/src/validation.rs create mode 100644 litellm-rust/crates/model-catalog/src/vocabulary.rs create mode 100644 litellm-rust/crates/model-catalog/src/wire/attributes.rs create mode 100644 litellm-rust/crates/model-catalog/src/wire/blocks.rs create mode 100644 litellm-rust/crates/model-catalog/src/wire/entry.rs create mode 100644 litellm-rust/crates/model-catalog/src/wire/flat_key.rs create mode 100644 litellm-rust/crates/model-catalog/src/wire/lax.rs create mode 100644 litellm-rust/crates/model-catalog/src/wire/mod.rs create mode 100644 litellm-rust/crates/model-catalog/tests/generate_price_lookup_reference.py create mode 100644 litellm-rust/crates/model-catalog/tests/price_lookup.rs create mode 100644 litellm-rust/crates/model-catalog/tests/price_lookup_reference.json delete mode 100644 litellm-rust/crates/model-catalog/tests/registry_validation.rs diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 636a5031e5f..04af770b24f 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -4,6 +4,7 @@ on: push: paths: - "litellm-rust/**" + - "model_prices_and_context_window.json" - "litellm/rust_bridge/**" - "scripts/generate_trace_types.py" - "scripts/trace_codegen/**" @@ -38,6 +39,7 @@ on: - "litellm_**" paths: - "litellm-rust/**" + - "model_prices_and_context_window.json" - "litellm/rust_bridge/**" - "scripts/generate_trace_types.py" - "scripts/trace_codegen/**" diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 80d5018fe9a..4768eeb633e 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -42,7 +42,7 @@ jobs: id: rust-bridge-key if: steps.changes.outputs.decision != 'skip' env: - RUST_BRIDGE_CACHE_KEY: ${{ runner.os }}-rust-bridge-editable-${{ hashFiles('litellm-rust/Cargo.lock', 'litellm-rust/Cargo.toml', 'litellm-rust/crates/**', 'pyproject.toml') }} + RUST_BRIDGE_CACHE_KEY: ${{ runner.os }}-rust-bridge-editable-${{ hashFiles('litellm-rust/Cargo.lock', 'litellm-rust/Cargo.toml', 'litellm-rust/crates/**', 'model_prices_and_context_window.json', 'pyproject.toml') }} run: echo "key=$RUST_BRIDGE_CACHE_KEY" >> "$GITHUB_OUTPUT" - name: Set up Python diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 66a345fc438..f7d7e5f4dab 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3150,6 +3150,7 @@ dependencies = [ "defmt", "jiff-core", "jiff-static", + "jiff-tzdb", "log", "portable-atomic", "portable-atomic-util", @@ -3178,6 +3179,12 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "jiff-tzdb" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa8377070c6bae868759445e5a77f66d84f0b72f3a054bfb00e6d038b8282da7" + [[package]] name = "jni" version = "0.22.4" @@ -3729,9 +3736,12 @@ dependencies = [ name = "litellm-cost" version = "0.1.0" dependencies = [ - "criterion", - "proptest", + "jiff", "rstest", + "schemars 1.2.2", + "serde", + "serde_json", + "strum", "thiserror 2.0.19", ] @@ -4277,17 +4287,35 @@ name = "litellm-model-catalog" version = "0.1.0" dependencies = [ "indexmap 2.14.0", + "jiff", "jsonschema", + "litellm-cost", "litellm-error", - "litellm-llms-types", + "macro_rules_attribute", "rstest", "schemars 1.2.2", "serde", "serde_json", + "serde_with", + "strum", "thiserror 2.0.19", "time", ] +[[package]] +name = "litellm-model-catalog-refresh" +version = "0.1.0" +dependencies = [ + "litellm-error", + "litellm-http", + "litellm-model-catalog", + "reqwest", + "rstest", + "thiserror 2.0.19", + "tokio", + "url", +] + [[package]] name = "litellm-python-bridge" version = "0.1.0" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index b7b22cfd32c..49a33e08456 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -73,6 +73,9 @@ litellm-python-compat = { path = "crates/python-compat" } litellm-circuit-breaker = { path = "crates/circuit-breaker" } litellm-redis-client = { path = "crates/redis-client" } litellm-redis-testing = { path = "crates/redis-testing" } +litellm-cost = { path = "crates/cost" } +litellm-model-catalog = { path = "crates/model-catalog" } +litellm-model-catalog-refresh = { path = "crates/model-catalog-refresh" } askama = { version = "0.16.1", default-features = false, features = ["derive", "std"] } tracing = "0.1" @@ -109,6 +112,7 @@ serde = { version = "1.0", features = ["derive"] } serde_json = { version = "1.0", features = ["float_roundtrip"] } serde_with = { version = "=3.21.0", default-features = false, features = ["std", "macros"] } sha2 = "0.10" +sha1 = "0.10" syn = { version = "2", default-features = false } sqlx = { version = "0.9.0", default-features = false, features = ["json", "macros", "postgres", "runtime-tokio", "chrono", "tls-rustls-ring-native-roots"] } chrono = { version = "0.4.45", default-features = false, features = ["std"] } diff --git a/litellm-rust/crates/cost/Cargo.toml b/litellm-rust/crates/cost/Cargo.toml index 6921e6151b5..858287c7860 100644 --- a/litellm-rust/crates/cost/Cargo.toml +++ b/litellm-rust/crates/cost/Cargo.toml @@ -5,14 +5,17 @@ edition.workspace = true license.workspace = true repository.workspace = true +[features] +serde = ["dep:serde"] +schema = ["serde", "dep:schemars"] + [dependencies] +jiff = { version = "0.2.37", default-features = false, features = ["std"] } +schemars = { workspace = true, optional = true } +serde = { workspace = true, optional = true } thiserror.workspace = true +strum.workspace = true [dev-dependencies] -criterion.workspace = true -proptest.workspace = true rstest.workspace = true - -[[bench]] -name = "calculate" -harness = false +serde_json.workspace = true diff --git a/litellm-rust/crates/cost/README.md b/litellm-rust/crates/cost/README.md index ee4722b6438..f8fc9403917 100644 --- a/litellm-rust/crates/cost/README.md +++ b/litellm-rust/crates/cost/README.md @@ -1,13 +1,3 @@ # litellm-cost -This crate calculates text token charges from rates and usage supplied by its caller. It is standalone and has no Python bridge or proxy integration - -Call `compile(&pricing)` once for an immutable plan, then `plan.calculate(&request)` for each supported request. `calculate(&pricing, &request)` compiles on each call. A successful result exposes pre-multiplier component costs, selected rates, the multiplier, and derived `input()`, `output()`, and `total()` values - -The caller states whether `prompt_tokens` includes cache tokens. Threshold selection uses total input tokens for either convention and selects one rate for the whole request. Thresholds are sorted when compiled, and duplicate thresholds or tier overrides fail deterministically. `Fast` selects priority rates; unknown tiers use standard rates - -`Rate::Missing`, `Rate::Null`, and `Rate::Value(0.0)` remain distinct. Missing cache rates fall back to the selected input rate, and an absent one-hour write rate falls back to the selected write rate. Missing input or output rates return typed errors, including for zero usage. Python's sparse-entry behavior remains outside this native contract - -The supported off-peak shape is one non-wrapping UTC daily window. The caller supplies the applicable regional multiplier after provider-specific selection. Negative or non-finite rates, ambiguous rules, inconsistent cache counts, incomplete write splits, invalid windows and overflow return errors. Callers must decline unsupported inputs before native execution if their public contract accepts those shapes - -This crate does not select models, read catalogs, fetch provider prices, normalize multimodal usage, process provider-reported costs, or calculate non-token charges. It does not change proxy behavior. The reference fixture was generated by `tests/generate_python_reference.py` against the Python implementation at the commit recorded in `tests/python_reference.tsv`, using synthetic rates and fixed usage +This crate holds the price data model that `litellm-model-catalog` lowers each catalog entry into, plus the usage, call context and priced outcome types that `calculate` will compute over. There is no `calculate` until the token pricing PR lands. `UnitPrice` is USD per unit and `Usd` is an amount, and only the operations that make sense between them exist, so a price is never added to an amount. `Prices::band` and `Prices::rate` resolve a token rate the way Python's `_get_token_base_cost` does diff --git a/litellm-rust/crates/cost/benches/calculate.rs b/litellm-rust/crates/cost/benches/calculate.rs deleted file mode 100644 index 2561b606e6c..00000000000 --- a/litellm-rust/crates/cost/benches/calculate.rs +++ /dev/null @@ -1,66 +0,0 @@ -use criterion::{Criterion, criterion_group, criterion_main}; -use litellm_cost::{ - Pricing, PromptConvention, Rate, Rates, Request, ServiceTier, ThresholdPolicy, ThresholdRates, - Usage, calculate, compile, -}; -use std::hint::black_box; - -fn bench(c: &mut Criterion) { - let pricing = Pricing { - standard: Rates { - input: Rate::Value(0.000002), - output: Rate::Value(0.000008), - cache_read: Rate::Value(0.0000005), - cache_write: Rate::Missing, - cache_write_1h: Rate::Missing, - }, - tiers: &[], - thresholds: &[], - off_peak: None, - }; - let request = Request { - usage: Usage { - prompt_tokens: 1000, - completion_tokens: 200, - cache_read_tokens: 250, - cache_write_tokens: 0, - cache_write_5m_tokens: None, - cache_write_1h_tokens: None, - prompt_convention: PromptConvention::IncludesCache, - }, - service_tier: ServiceTier::Standard, - threshold_policy: ThresholdPolicy::Exclusive, - region_multiplier: None, - billed_at_utc_minute: None, - }; - let plan = compile(&pricing).unwrap(); - c.bench_function("native_compiled_calculation", |b| { - b.iter(|| black_box(plan.calculate(black_box(&request)).unwrap())) - }); - c.bench_function("native_full_wrapper", |b| { - b.iter(|| black_box(calculate(black_box(&pricing), black_box(&request)).unwrap())) - }); - c.bench_function("native_rate_compilation", |b| { - b.iter(|| black_box(compile(black_box(&pricing)).unwrap())) - }); - let threshold = ThresholdRates { - above_prompt_tokens: 1000, - standard: Rates { - input: Rate::Value(0.000004), - output: Rate::Value(0.000016), - ..Rates::EMPTY - }, - tiers: &[], - }; - let threshold_pricing = Pricing { - thresholds: &[threshold], - ..pricing - }; - let threshold_plan = compile(&threshold_pricing).unwrap(); - c.bench_function("native_threshold_boundary", |b| { - b.iter(|| black_box(threshold_plan.calculate(black_box(&request)).unwrap())) - }); -} - -criterion_group!(benches, bench); -criterion_main!(benches); diff --git a/litellm-rust/crates/cost/examples/charge.rs b/litellm-rust/crates/cost/examples/charge.rs deleted file mode 100644 index a3a83febd74..00000000000 --- a/litellm-rust/crates/cost/examples/charge.rs +++ /dev/null @@ -1,40 +0,0 @@ -use litellm_cost::{ - Pricing, PromptConvention, Rate, Rates, Request, ServiceTier, ThresholdPolicy, Usage, compile, -}; - -fn main() { - let pricing = Pricing { - standard: Rates { - input: Rate::Value(2.0), - output: Rate::Value(4.0), - cache_read: Rate::Value(0.5), - cache_write: Rate::Value(3.0), - cache_write_1h: Rate::Missing, - }, - tiers: &[], - thresholds: &[], - off_peak: None, - }; - let request = Request { - usage: Usage { - prompt_tokens: 100, - completion_tokens: 20, - cache_read_tokens: 25, - cache_write_tokens: 10, - cache_write_5m_tokens: None, - cache_write_1h_tokens: None, - prompt_convention: PromptConvention::IncludesCache, - }, - service_tier: ServiceTier::Standard, - threshold_policy: ThresholdPolicy::Exclusive, - region_multiplier: None, - billed_at_utc_minute: None, - }; - let cost = compile(&pricing).unwrap().calculate(&request).unwrap(); - println!( - "input={} output={} total={}", - cost.input(), - cost.output(), - cost.total() - ); -} diff --git a/litellm-rust/crates/cost/src/call.rs b/litellm-rust/crates/cost/src/call.rs new file mode 100644 index 00000000000..3f44ec03e9b --- /dev/null +++ b/litellm-rust/crates/cost/src/call.rs @@ -0,0 +1,42 @@ +use crate::price::{Channel, Prices, Residency, ServiceTier}; + +/// Facts about the call, never about the row or the provider +#[derive(Clone, Debug, PartialEq)] +pub struct CallContext { + /// Requested service tier, `"fast"` maps to `Priority` as on main + pub tier: ServiceTier, + /// Online, batch or decisions + pub channel: Channel, + /// OpenAI data residency, picks the residency uplift + pub residency: Option, + /// Vertex call sent to a regional endpoint, picks the regional uplift + pub non_global_endpoint: bool, + /// Anthropic `speed: fast`, picks the speed uplift + pub speed: Option, + /// Anthropic and Bedrock `inference_geo`, picks the geo uplift + pub inference_geo: Option, + /// When the call was billed, for off-peak windows + pub at: jiff::Timestamp, + /// Per-call custom pricing, lowered like a row and applied field by field over it + pub overlay: Option, +} + +/// Anthropic `speed` +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum Speed { + Fast, +} + +/// Anthropic and Bedrock `inference_geo`, lowercased as on main +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct InferenceGeo(Box); + +impl InferenceGeo { + pub fn new(geo: &str) -> Self { + Self(geo.to_lowercase().into()) + } + + pub fn as_str(&self) -> &str { + &self.0 + } +} diff --git a/litellm-rust/crates/cost/src/error.rs b/litellm-rust/crates/cost/src/error.rs index 1753080617f..7ef5eb87597 100644 --- a/litellm-rust/crates/cost/src/error.rs +++ b/litellm-rust/crates/cost/src/error.rs @@ -1,27 +1,24 @@ -#[derive(Clone, Copy, Debug, Eq, PartialEq, thiserror::Error)] +use thiserror::Error; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Error, strum::VariantArray)] pub enum Error { - #[error("model has no input rate")] - MissingInputRate, - #[error("model has no output rate")] - MissingOutputRate, - #[error("rate must be finite and nonnegative")] - InvalidRate, - #[error("region multiplier must be finite and positive")] - InvalidRegionMultiplier, - #[error("off-peak pricing needs a billing minute within the day")] - InvalidBillingTime, - #[error("off-peak window must start before it ends within the day")] - InvalidOffPeakWindow, - #[error("cached tokens exceed prompt tokens")] - CacheExceedsPrompt, - #[error("cache write details do not add up to the cache write tokens")] - InvalidCacheWriteDetails, - #[error("token count or cost overflows")] - TokenCountOverflow, - #[error("service tier is repeated or cannot carry its own rates")] - DuplicateTier, - #[error("prompt token threshold is repeated")] - DuplicateThreshold, - #[error("service tier is repeated or cannot carry its own rates within a threshold")] - DuplicateThresholdTier, + #[error("expected a UTC window HH:MM-HH:MM")] + InvalidUtcWindow, + #[error("expected an ISO weekday number 1-7 or an English day name")] + InvalidWeekday, +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::Error; + use crate::schedule::{UtcWindow, Weekday}; + + #[rstest] + #[case::window("not-a-window".parse::().unwrap_err(), Error::InvalidUtcWindow)] + #[case::weekday("not-a-weekday".parse::().unwrap_err(), Error::InvalidWeekday)] + fn schedule_parsers_share_the_crate_error(#[case] error: Error, #[case] expected: Error) { + assert_eq!(error, expected); + } } diff --git a/litellm-rust/crates/cost/src/lib.rs b/litellm-rust/crates/cost/src/lib.rs index 24c69afd9d4..192dfa361d6 100644 --- a/litellm-rust/crates/cost/src/lib.rs +++ b/litellm-rust/crates/cost/src/lib.rs @@ -1,383 +1,9 @@ +pub mod call; mod error; +pub mod money; +pub mod price; +pub mod priced; +pub mod schedule; +pub mod usage; -pub use error::Error; - -#[derive(Clone, Copy, Debug, PartialEq)] -pub enum Rate { - Missing, - Null, - Value(f64), -} - -impl Rate { - fn value(self) -> Option { - match self { - Self::Value(value) => Some(value), - Self::Missing | Self::Null => None, - } - } - - fn or(self, fallback: Self) -> Self { - if self.value().is_some() { - self - } else { - fallback - } - } -} - -#[derive(Clone, Copy, Debug, PartialEq)] -pub struct Rates { - pub input: Rate, - pub output: Rate, - pub cache_read: Rate, - pub cache_write: Rate, - pub cache_write_1h: Rate, -} - -impl Rates { - pub const EMPTY: Self = Self { - input: Rate::Missing, - output: Rate::Missing, - cache_read: Rate::Missing, - cache_write: Rate::Missing, - cache_write_1h: Rate::Missing, - }; - - fn overlay(self, base: Self) -> Self { - Self { - input: self.input.or(base.input), - output: self.output.or(base.output), - cache_read: self.cache_read.or(base.cache_read), - cache_write: self.cache_write.or(base.cache_write), - cache_write_1h: self.cache_write_1h.or(base.cache_write_1h), - } - } -} - -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub enum ServiceTier { - Standard, - Flex, - Priority, - Fast, - Ultrafast, - Unknown, -} - -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub enum ThresholdPolicy { - Exclusive, - Inclusive, -} - -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub enum PromptConvention { - IncludesCache, - ExcludesCache, -} - -#[derive(Clone, Copy, Debug, PartialEq)] -pub struct Usage { - pub prompt_tokens: u64, - pub completion_tokens: u64, - pub cache_read_tokens: u64, - pub cache_write_tokens: u64, - pub cache_write_5m_tokens: Option, - pub cache_write_1h_tokens: Option, - pub prompt_convention: PromptConvention, -} - -#[derive(Clone, Copy, Debug, PartialEq)] -pub struct TierRates { - pub tier: ServiceTier, - pub rates: Rates, -} - -#[derive(Clone, Copy, Debug, PartialEq)] -pub struct ThresholdRates<'a> { - pub above_prompt_tokens: u64, - pub standard: Rates, - pub tiers: &'a [TierRates], -} - -#[derive(Clone, Copy, Debug, PartialEq)] -pub struct OffPeakRates { - pub start_utc_minute: u16, - pub end_utc_minute: u16, - pub rates: Rates, -} - -#[derive(Clone, Copy, Debug, PartialEq)] -pub struct Pricing<'a> { - pub standard: Rates, - pub tiers: &'a [TierRates], - pub thresholds: &'a [ThresholdRates<'a>], - pub off_peak: Option, -} - -#[derive(Clone, Copy, Debug, PartialEq)] -pub struct Request { - pub usage: Usage, - pub service_tier: ServiceTier, - pub threshold_policy: ThresholdPolicy, - pub region_multiplier: Option, - pub billed_at_utc_minute: Option, -} - -#[derive(Clone, Copy, Debug, PartialEq)] -pub struct Cost { - pub uncached_input: f64, - pub cache_read: f64, - pub cache_write_5m: f64, - pub cache_write_1h: f64, - pub output: f64, - pub multiplier: f64, - pub rates: EffectiveRates, -} - -impl Cost { - pub fn input(self) -> f64 { - (self.uncached_input + self.cache_read + self.cache_write_5m + self.cache_write_1h) - * self.multiplier - } - - pub fn output(self) -> f64 { - self.output * self.multiplier - } - - pub fn total(self) -> f64 { - self.input() + self.output() - } -} - -#[derive(Clone, Copy, Debug, PartialEq)] -pub struct EffectiveRates { - pub input: f64, - pub output: f64, - pub cache_read: f64, - pub cache_write_5m: f64, - pub cache_write_1h: f64, -} - -fn selected_tier(tier: ServiceTier) -> ServiceTier { - if tier == ServiceTier::Fast { - ServiceTier::Priority - } else { - tier - } -} - -#[derive(Clone, Debug)] -struct CompiledThreshold { - above_prompt_tokens: u64, - standard: Rates, - tiers: Vec, -} - -#[derive(Clone, Debug)] -pub struct PricingPlan { - standard: Rates, - tiers: Vec, - thresholds: Vec, - off_peak: Option, -} - -fn valid_rates(rates: Rates) -> bool { - [ - rates.input, - rates.output, - rates.cache_read, - rates.cache_write, - rates.cache_write_1h, - ] - .into_iter() - .all(|rate| { - rate.value() - .is_none_or(|value| value.is_finite() && value >= 0.0) - }) -} - -fn validate_tiers(tiers: &[TierRates], duplicate: Error) -> Result<(), Error> { - if tiers.iter().any(|entry| !valid_rates(entry.rates)) { - return Err(Error::InvalidRate); - } - if tiers.iter().enumerate().any(|(index, entry)| { - matches!( - entry.tier, - ServiceTier::Standard | ServiceTier::Unknown | ServiceTier::Fast - ) || tiers[..index] - .iter() - .any(|previous| previous.tier == entry.tier) - }) { - return Err(duplicate); - } - Ok(()) -} - -pub fn compile(pricing: &Pricing<'_>) -> Result { - if !valid_rates(pricing.standard) { - return Err(Error::InvalidRate); - } - validate_tiers(pricing.tiers, Error::DuplicateTier)?; - if let Some(window) = pricing.off_peak { - if window.start_utc_minute >= 1440 - || window.end_utc_minute > 1440 - || window.start_utc_minute >= window.end_utc_minute - { - return Err(Error::InvalidOffPeakWindow); - } - if !valid_rates(window.rates) { - return Err(Error::InvalidRate); - } - } - let mut thresholds: Vec<_> = pricing - .thresholds - .iter() - .map(|entry| { - if !valid_rates(entry.standard) { - return Err(Error::InvalidRate); - } - validate_tiers(entry.tiers, Error::DuplicateThresholdTier)?; - Ok(CompiledThreshold { - above_prompt_tokens: entry.above_prompt_tokens, - standard: entry.standard, - tiers: entry.tiers.to_vec(), - }) - }) - .collect::>()?; - thresholds.sort_unstable_by_key(|entry| entry.above_prompt_tokens); - if thresholds - .windows(2) - .any(|pair| pair[0].above_prompt_tokens == pair[1].above_prompt_tokens) - { - return Err(Error::DuplicateThreshold); - } - Ok(PricingPlan { - standard: pricing.standard, - tiers: pricing.tiers.to_vec(), - thresholds, - off_peak: pricing.off_peak, - }) -} - -impl PricingPlan { - fn resolve_rates(&self, request: &Request, threshold_tokens: u64) -> Result { - let tier = selected_tier(request.service_tier); - let base = self - .tiers - .iter() - .find(|entry| tier != ServiceTier::Standard && entry.tier == tier) - .map_or(self.standard, |entry| entry.rates.overlay(self.standard)); - let threshold = self.thresholds.iter().rev().find(|entry| { - threshold_tokens > entry.above_prompt_tokens - || (request.threshold_policy == ThresholdPolicy::Inclusive - && threshold_tokens == entry.above_prompt_tokens) - }); - let selected = threshold.map_or(base, |entry| { - let standard = entry.standard.overlay(base); - entry - .tiers - .iter() - .find(|specific| tier != ServiceTier::Standard && specific.tier == tier) - .map_or(standard, |specific| specific.rates.overlay(standard)) - }); - match self.off_peak { - None => Ok(selected), - Some(window) => { - if window.start_utc_minute >= 1440 - || window.end_utc_minute > 1440 - || window.start_utc_minute >= window.end_utc_minute - { - return Err(Error::InvalidOffPeakWindow); - } - let minute = request - .billed_at_utc_minute - .ok_or(Error::InvalidBillingTime)?; - if minute >= 1440 { - return Err(Error::InvalidBillingTime); - } - if (window.start_utc_minute..window.end_utc_minute).contains(&minute) { - Ok(window.rates.overlay(selected)) - } else { - Ok(selected) - } - } - } - } - - fn checked_rate(rate: Rate, missing: Error) -> Result { - let value = rate.value().ok_or(missing)?; - if !value.is_finite() || value < 0.0 { - return Err(Error::InvalidRate); - } - Ok(value) - } - - pub fn calculate(&self, request: &Request) -> Result { - let usage = request.usage; - let cached = usage - .cache_read_tokens - .checked_add(usage.cache_write_tokens) - .ok_or(Error::TokenCountOverflow)?; - let (regular, threshold_tokens) = match usage.prompt_convention { - PromptConvention::IncludesCache => ( - usage - .prompt_tokens - .checked_sub(cached) - .ok_or(Error::CacheExceedsPrompt)?, - usage.prompt_tokens, - ), - PromptConvention::ExcludesCache => ( - usage.prompt_tokens, - usage - .prompt_tokens - .checked_add(cached) - .ok_or(Error::TokenCountOverflow)?, - ), - }; - let writes = match (usage.cache_write_5m_tokens, usage.cache_write_1h_tokens) { - (None, None) => (usage.cache_write_tokens, 0), - (Some(five), Some(one)) if five.checked_add(one) == Some(usage.cache_write_tokens) => { - (five, one) - } - _ => return Err(Error::InvalidCacheWriteDetails), - }; - let rates = self.resolve_rates(request, threshold_tokens)?; - let input = Self::checked_rate(rates.input, Error::MissingInputRate)?; - let output = Self::checked_rate(rates.output, Error::MissingOutputRate)?; - let read = Self::checked_rate(rates.cache_read.or(rates.input), Error::MissingInputRate)?; - let write = Self::checked_rate(rates.cache_write.or(rates.input), Error::MissingInputRate)?; - let write_1h = Self::checked_rate( - rates.cache_write_1h.or(rates.cache_write).or(rates.input), - Error::MissingInputRate, - )?; - let multiplier = request.region_multiplier.unwrap_or(1.0); - if !multiplier.is_finite() || multiplier <= 0.0 { - return Err(Error::InvalidRegionMultiplier); - } - let cost = Cost { - uncached_input: regular as f64 * input, - cache_read: usage.cache_read_tokens as f64 * read, - cache_write_5m: writes.0 as f64 * write, - cache_write_1h: writes.1 as f64 * write_1h, - output: usage.completion_tokens as f64 * output, - multiplier, - rates: EffectiveRates { - input, - output, - cache_read: read, - cache_write_5m: write, - cache_write_1h: write_1h, - }, - }; - if !cost.total().is_finite() { - return Err(Error::TokenCountOverflow); - } - Ok(cost) - } -} - -pub fn calculate(pricing: &Pricing<'_>, request: &Request) -> Result { - compile(pricing)?.calculate(request) -} +pub use error::*; diff --git a/litellm-rust/crates/cost/src/money.rs b/litellm-rust/crates/cost/src/money.rs new file mode 100644 index 00000000000..0140fff52d2 --- /dev/null +++ b/litellm-rust/crates/cost/src/money.rs @@ -0,0 +1,257 @@ +use std::iter::Sum; +use std::ops::{Add, Mul}; + +use crate::priced::Quantity; + +#[derive(Clone, Copy, Debug, PartialEq, PartialOrd)] +pub struct UnitPrice(f64); + +#[derive(Clone, Copy, Debug, Default, PartialEq, PartialOrd)] +pub struct Usd(f64); + +#[derive(Clone, Copy, Debug, PartialEq, PartialOrd)] +pub struct Multiplier(f64); + +impl UnitPrice { + pub fn new(usd_per_unit: f64) -> Option { + (usd_per_unit.is_finite() && usd_per_unit >= 0.0).then_some(Self(usd_per_unit)) + } + pub fn get(self) -> f64 { + self.0 + } +} + +impl Usd { + pub const ZERO: Self = Self(0.0); + + pub fn new(amount: f64) -> Option { + (amount.is_finite() && amount >= 0.0).then_some(Self(amount)) + } + pub fn get(self) -> f64 { + self.0 + } +} + +impl Multiplier { + pub fn new(factor: f64) -> Option { + (factor.is_finite() && factor >= 1.0).then_some(Self(factor)) + } + pub fn get(self) -> f64 { + self.0 + } +} + +impl Mul for UnitPrice { + type Output = Usd; + + fn mul(self, quantity: Quantity) -> Usd { + match quantity { + Quantity::Count(units) => Usd(self.0 * units as f64), + Quantity::Duration(duration) => Usd(self.0 * duration.as_secs_f64()), + } + } +} + +impl Mul for UnitPrice { + type Output = UnitPrice; + + fn mul(self, factor: Multiplier) -> UnitPrice { + UnitPrice(self.0 * factor.0) + } +} + +impl Mul for Usd { + type Output = Usd; + + fn mul(self, factor: Multiplier) -> Usd { + Usd(self.0 * factor.0) + } +} + +impl Add for Usd { + type Output = Usd; + + fn add(self, other: Usd) -> Usd { + Usd(self.0 + other.0) + } +} + +impl Sum for Usd { + fn sum>(amounts: I) -> Usd { + amounts.fold(Usd::ZERO, Add::add) + } +} + +#[cfg(feature = "serde")] +mod serde_impls { + use std::fmt; + + use serde::de::{self, Deserialize, Deserializer, Visitor}; + + use super::{Multiplier, UnitPrice}; + + #[derive(Clone, Copy)] + struct Number(&'static str); + + impl Visitor<'_> for Number { + type Value = f64; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(self.0) + } + + fn visit_i64(self, value: i64) -> Result { + Ok(value as f64) + } + + fn visit_u64(self, value: u64) -> Result { + Ok(value as f64) + } + + fn visit_f64(self, value: f64) -> Result { + Ok(value) + } + + fn visit_str(self, value: &str) -> Result { + value + .trim() + .parse() + .map_err(|_| E::invalid_value(de::Unexpected::Str(value), &self)) + } + } + + fn checked<'de, D: Deserializer<'de>, T>( + deserializer: D, + expected: Number, + new: fn(f64) -> Option, + ) -> Result { + let number = deserializer.deserialize_any(expected)?; + new(number) + .ok_or_else(|| de::Error::invalid_value(de::Unexpected::Float(number), &expected)) + } + + impl<'de> Deserialize<'de> for UnitPrice { + fn deserialize>(deserializer: D) -> Result { + checked(deserializer, Number("a non-negative USD amount"), Self::new) + } + } + + impl<'de> Deserialize<'de> for Multiplier { + fn deserialize>(deserializer: D) -> Result { + checked( + deserializer, + Number("a multiplier of at least 1"), + Self::new, + ) + } + } +} + +#[cfg(feature = "schema")] +mod schema { + use std::borrow::Cow; + + use schemars::{JsonSchema, Schema, SchemaGenerator, json_schema}; + + use super::{Multiplier, UnitPrice}; + + impl JsonSchema for UnitPrice { + fn inline_schema() -> bool { + true + } + fn schema_name() -> Cow<'static, str> { + "UnitPrice".into() + } + fn json_schema(_: &mut SchemaGenerator) -> Schema { + json_schema!({"type": "number", "minimum": 0}) + } + } + + impl JsonSchema for Multiplier { + fn inline_schema() -> bool { + true + } + fn schema_name() -> Cow<'static, str> { + "Multiplier".into() + } + fn json_schema(_: &mut SchemaGenerator) -> Schema { + json_schema!({"type": "number", "minimum": 1}) + } + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::price(0.5)] + #[case::free(0.0)] + fn unit_price_keeps_finite_non_negative_values(#[case] value: f64) { + assert_eq!(UnitPrice::new(value).map(UnitPrice::get), Some(value)); + } + + #[rstest] + #[case::negative(-1e-9)] + #[case::nan(f64::NAN)] + #[case::infinite(f64::INFINITY)] + fn unit_price_rejects_what_no_price_can_be(#[case] value: f64) { + assert_eq!(UnitPrice::new(value), None); + } + + #[rstest] + #[case::uplift(1.1, true)] + #[case::neutral(1.0, true)] + #[case::discount(0.5, false)] + #[case::nan(f64::NAN, false)] + fn multiplier_is_never_a_discount(#[case] value: f64, #[case] accepted: bool) { + assert_eq!(Multiplier::new(value).is_some(), accepted); + } + + #[rstest] + fn amounts_follow_units_times_price_then_uplift() { + let price = UnitPrice::new(0.25).unwrap(); + let uplift = Multiplier::new(2.0).unwrap(); + + let total: Usd = [ + price * Quantity::Count(4), + (price * uplift) * Quantity::Count(2), + ] + .into_iter() + .sum(); + + assert_eq!(total.get(), 2.0); + assert_eq!((total * uplift).get(), 4.0); + } + + #[rstest] + fn a_duration_is_billed_in_fractional_seconds() { + let per_second = UnitPrice::new(0.01).unwrap(); + + let amount = per_second * Quantity::Duration(std::time::Duration::from_millis(2_500)); + + assert!((amount.get() - 0.025).abs() < 1e-12); + } + + #[cfg(feature = "serde")] + #[rstest] + #[case::number(serde_json::json!(1.5e-6), Some(1.5e-6))] + #[case::integer(serde_json::json!(2), Some(2.0))] + #[case::numeric_string(serde_json::json!(" 0.000002 "), Some(2e-6))] + #[case::negative(serde_json::json!(-1), None)] + #[case::word(serde_json::json!("free"), None)] + #[case::boolean(serde_json::json!(true), None)] + fn unit_price_reads_numbers_and_numeric_strings( + #[case] written: serde_json::Value, + #[case] expected: Option, + ) { + use serde::Deserialize; + + assert_eq!( + UnitPrice::deserialize(&written).ok().map(UnitPrice::get), + expected + ); + } +} diff --git a/litellm-rust/crates/cost/src/price.rs b/litellm-rust/crates/cost/src/price.rs new file mode 100644 index 00000000000..a90bedacb0f --- /dev/null +++ b/litellm-rust/crates/cost/src/price.rs @@ -0,0 +1,505 @@ +mod meter; + +use std::collections::BTreeMap; +use std::num::NonZeroU64; + +pub use meter::*; + +use crate::call::{InferenceGeo, Speed}; +use crate::money::{Multiplier, UnitPrice}; +use crate::schedule::Schedule; + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum ServiceTier { + #[default] + Standard, + Priority, + Flex, + Ultrafast, + Balanced, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum Band { + #[default] + Base, + OverPromptTokens(NonZeroU64), +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct When { + pub tier: ServiceTier, + pub band: Band, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum Channel { + #[default] + Online, + Batch, + Decisions, +} + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct Price { + pub meter: Meter, + pub channel: Channel, + pub when: When, +} + +#[derive(Clone, Debug, Default, PartialEq)] +pub struct Prices { + rates: BTreeMap, + off_peak: Option, + billing: Billing, +} + +/// Row attributes pricing reads, lowered by model-catalog from the row and its provider +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct Billing { + /// Unit the row bills text in. `Character` only when the row prices characters and not tokens + pub text_unit: TextUnit, + /// Whether a band's floor is in the band. xAI rows bill the band from the floor itself + pub band_edge: BandEdge, + /// Rate a cache read borrows when the row has no cache-read price. Fireworks rows borrow half the input rate + pub cache_read_fallback: CacheReadFallback, + /// Whether a row with no token price is priced by wall clock, from its `mode` + pub wall_clock: WallClock, + /// From `web_search_billing_unit`, whether web search and Maps grounding bill per query or once per prompt + pub web_search: WebSearchUnit, + /// Whether `tiered_pricing` bills each meter in its own tier or the whole call at the tier the input picks. + /// Dashscope rows bill the whole call (inferred, 10 checks no other row has that shape) + pub tier_scope: TierScope, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum TextUnit { + #[default] + Token, + Character, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum BandEdge { + #[default] + Exclusive, + Inclusive, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum CacheReadFallback { + #[default] + InputRate, + HalfInputRate, +} + +/// `ResponseTime` for `chat`, `completion`, `embedding`, `responses` or no mode, `AudioFile` for +/// `audio_transcription`, which prices the input audio when it has one, `Never` for every other mode +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum WallClock { + Never, + #[default] + ResponseTime, + AudioFile, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum WebSearchUnit { + PerQuery, + #[default] + PerPrompt, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum TierScope { + #[default] + PerMeter, + WholeCall, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct OffPeak { + schedule: Schedule, + rates: OffPeakRates, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq)] +pub struct OffPeakRates { + pub input: Option, + pub output: Option, + pub cache_read: Option, + pub explicit_cache_read: Option, + pub cache_write: Option, + pub reasoning: Option, +} + +#[derive(Clone, Debug, Default, PartialEq)] +pub struct Adjustments(BTreeMap); + +/// Which call fact picks which multiplier in `Adjustments`, lowered by model-catalog from the row +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum Uplift { + /// OpenAI data residency + Residency(Residency), + /// Vertex non-global endpoint, from `regional_endpoint_uplift_multiplier` + NonGlobalEndpoint, + /// Anthropic `speed: fast`, from `provider_specific_entry.fast` + Speed(Speed), + /// Anthropic and Bedrock geo, from `provider_specific_entry.`, lowercased as on main + InferenceGeo(InferenceGeo), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum Residency { + Eu, + Us, +} + +impl Band { + pub fn over_prompt_tokens(tokens: u64) -> Self { + NonZeroU64::new(tokens).map_or(Self::Base, Self::OverPromptTokens) + } +} + +impl Price { + pub fn online(meter: Meter, when: When) -> Self { + Self { + meter, + channel: Channel::Online, + when, + } + } +} + +impl Prices { + pub fn new( + rates: BTreeMap, + off_peak: Option, + billing: Billing, + ) -> Self { + Self { + rates, + off_peak, + billing, + } + } + + pub fn get(&self, price: &Price) -> Option { + self.rates.get(price).copied() + } + + pub fn iter(&self) -> impl Iterator { + self.rates.iter().map(|(price, rate)| (price, *rate)) + } + + pub fn off_peak(&self) -> Option<&OffPeak> { + self.off_peak.as_ref() + } + + pub fn billing(&self) -> Billing { + self.billing + } + + /// The online prompt band a request falls in. Every token meter's `_above_` prices count, + /// and the highest band the prompt passes wins. `billing.band_edge` says whether the floor is in the band. + pub fn band(&self, tier: ServiceTier, prompt_tokens: u64) -> Band { + let inclusive = self.billing.band_edge == BandEdge::Inclusive; + self.rates + .keys() + .filter(|price| { + price.channel == Channel::Online + && price.meter.sets_prompt_band() + && (price.when.tier == ServiceTier::Standard || price.when.tier == tier) + }) + .filter_map(|price| match price.when.band { + Band::OverPromptTokens(floor) => Some(floor.get()), + Band::Base => None, + }) + .filter(|floor| prompt_tokens > *floor || (inclusive && prompt_tokens == *floor)) + .max() + .map_or(Band::Base, Band::over_prompt_tokens) + } + + /// The online rate for `meter`, tried at `(tier, band)`, `(standard, band)`, `(tier, base)` + /// and `(standard, base)` in that order. A missing cache price falls back to the input price + /// resolved the same way. + pub fn rate(&self, meter: &Meter, tier: ServiceTier, band: Band) -> Option { + [ + When { tier, band }, + When { + tier: ServiceTier::Standard, + band, + }, + When { + tier, + band: Band::Base, + }, + When::default(), + ] + .into_iter() + .find_map(|when| self.get(&Price::online(meter.clone(), when))) + .or_else(|| self.rate(&meter.fallback()?, tier, band)) + } +} + +impl OffPeak { + pub fn new(schedule: Schedule, rates: OffPeakRates) -> Option { + (rates != OffPeakRates::default()).then_some(Self { schedule, rates }) + } + + pub fn schedule(&self) -> &Schedule { + &self.schedule + } + + pub fn rates(&self) -> &OffPeakRates { + &self.rates + } +} + +impl Adjustments { + pub fn new(uplifts: impl IntoIterator) -> Self { + Self(uplifts.into_iter().collect()) + } + + pub fn get(&self, uplift: &Uplift) -> Option { + self.0.get(uplift).copied() + } +} + +#[cfg(test)] +mod tests { + use rstest::{fixture, rstest}; + + use super::*; + + const INPUT: Meter = Meter::TEXT_INPUT; + const OUTPUT: Meter = Meter::Output(OutputUnit::Token(TokenKind::Text)); + const CACHE_READ: Meter = Meter::CacheRead { + kind: TokenKind::Text, + mode: CacheReadMode::Implicit, + }; + const CACHE_WRITE_HOUR: Meter = Meter::CacheWrite { + kind: TokenKind::Text, + ttl: CacheTtl::OneHour, + }; + + fn over(tokens: u64) -> Band { + Band::over_prompt_tokens(tokens) + } + + fn prices(rates: &[(Meter, Channel, ServiceTier, Band, f64)]) -> Prices { + priced_with(rates, Billing::default()) + } + + fn priced_with(rates: &[(Meter, Channel, ServiceTier, Band, f64)], billing: Billing) -> Prices { + Prices::new( + rates + .iter() + .map(|(meter, channel, tier, band, rate)| { + ( + Price { + meter: meter.clone(), + channel: *channel, + when: When { + tier: *tier, + band: *band, + }, + }, + UnitPrice::new(*rate).unwrap(), + ) + }) + .collect(), + None, + billing, + ) + } + + #[fixture] + fn priority_with_a_standard_band() -> Prices { + use Channel::Online; + use ServiceTier::{Priority, Standard}; + prices(&[ + (INPUT, Online, Standard, Band::Base, 1.0), + (INPUT, Online, Priority, Band::Base, 2.0), + (INPUT, Online, Standard, over(200_000), 10.0), + (OUTPUT, Online, Standard, Band::Base, 3.0), + (OUTPUT, Online, Priority, over(200_000), 30.0), + ]) + } + + #[rstest] + #[case::below(1_000, false, Band::Base)] + #[case::at_floor_exclusive(200_000, false, Band::Base)] + #[case::at_floor_inclusive(200_000, true, over(200_000))] + #[case::above(200_001, false, over(200_000))] + fn the_band_starts_strictly_above_its_floor_unless_inclusive( + #[case] prompt_tokens: u64, + #[case] inclusive: bool, + #[case] expected: Band, + ) { + let band_edge = if inclusive { + BandEdge::Inclusive + } else { + BandEdge::Exclusive + }; + let prices = priced_with( + &[( + INPUT, + Channel::Online, + ServiceTier::Standard, + over(200_000), + 10.0, + )], + Billing { + band_edge, + ..Billing::default() + }, + ); + assert_eq!(prices.band(ServiceTier::Standard, prompt_tokens), expected); + } + + #[rstest] + fn a_band_only_another_tier_prices_does_not_apply() { + let prices = prices(&[( + OUTPUT, + Channel::Online, + ServiceTier::Flex, + over(100_000), + 1.0, + )]); + assert_eq!(prices.band(ServiceTier::Standard, 300_000), Band::Base); + assert_eq!(prices.band(ServiceTier::Flex, 300_000), over(100_000)); + } + + #[rstest] + fn batch_bands_never_pick_the_online_band() { + let prices = prices(&[( + INPUT, + Channel::Batch, + ServiceTier::Standard, + over(100_000), + 1.0, + )]); + assert_eq!(prices.band(ServiceTier::Standard, 300_000), Band::Base); + } + + #[rstest] + fn the_highest_band_passed_wins() { + let prices = prices(&[ + ( + INPUT, + Channel::Online, + ServiceTier::Standard, + over(100_000), + 1.0, + ), + ( + OUTPUT, + Channel::Online, + ServiceTier::Standard, + over(200_000), + 1.0, + ), + ]); + assert_eq!(prices.band(ServiceTier::Standard, 300_000), over(200_000)); + } + + #[rstest] + #[case::tier_base(INPUT, ServiceTier::Priority, Band::Base, Some(2.0))] + #[case::standard_band_beats_tier_base(INPUT, ServiceTier::Priority, over(200_000), Some(10.0))] + #[case::tier_band(OUTPUT, ServiceTier::Priority, over(200_000), Some(30.0))] + #[case::standard_base_when_nothing_else(OUTPUT, ServiceTier::Flex, over(200_000), Some(3.0))] + #[case::cache_read_falls_back_to_the_resolved_input( + CACHE_READ, + ServiceTier::Priority, + over(200_000), + Some(10.0) + )] + #[case::one_hour_write_falls_back_through_write_to_input( + CACHE_WRITE_HOUR, + ServiceTier::Priority, + Band::Base, + Some(2.0) + )] + #[case::twelve_hour_write_falls_back_through_write_to_input( + Meter::CacheWrite { kind: TokenKind::Text, ttl: CacheTtl::TwelveHours }, + ServiceTier::Standard, + Band::Base, + Some(1.0) + )] + #[case::no_fallback_for_other_meters( + Meter::Input(InputUnit::Token(TokenKind::Audio)), + ServiceTier::Standard, + Band::Base, + None + )] + fn rate_walks_tier_then_band( + priority_with_a_standard_band: Prices, + #[case] meter: Meter, + #[case] tier: ServiceTier, + #[case] band: Band, + #[case] expected: Option, + ) { + assert_eq!( + priority_with_a_standard_band + .rate(&meter, tier, band) + .map(UnitPrice::get), + expected + ); + } + + #[rstest] + fn an_explicit_cache_price_wins_over_the_input_fallback() { + let prices = prices(&[ + ( + INPUT, + Channel::Online, + ServiceTier::Standard, + Band::Base, + 1.0, + ), + ( + CACHE_READ, + Channel::Online, + ServiceTier::Standard, + Band::Base, + 0.0, + ), + ]); + assert_eq!( + prices + .rate(&CACHE_READ, ServiceTier::Standard, Band::Base) + .map(UnitPrice::get), + Some(0.0) + ); + } + + #[rstest] + fn off_peak_needs_a_price() { + let schedule = Schedule { + windows: Vec::new(), + weekday_timezone: None, + }; + assert_eq!( + OffPeak::new(schedule.clone(), OffPeakRates::default()), + None + ); + assert!( + OffPeak::new( + schedule, + OffPeakRates { + output: UnitPrice::new(1.0), + ..OffPeakRates::default() + } + ) + .is_some() + ); + } + + #[rstest] + #[case::zero(0, Band::Base)] + #[case::positive(272_000, Band::OverPromptTokens(NonZeroU64::new(272_000).unwrap()))] + fn a_zero_floor_is_the_base_band(#[case] tokens: u64, #[case] expected: Band) { + assert_eq!(Band::over_prompt_tokens(tokens), expected); + } +} diff --git a/litellm-rust/crates/cost/src/price/meter.rs b/litellm-rust/crates/cost/src/price/meter.rs new file mode 100644 index 00000000000..6e68981d011 --- /dev/null +++ b/litellm-rust/crates/cost/src/price/meter.rs @@ -0,0 +1,197 @@ +use std::num::NonZeroU64; + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum Meter { + Input(InputUnit), + Output(OutputUnit), + CacheRead { + kind: TokenKind, + mode: CacheReadMode, + }, + CacheWrite { + kind: TokenKind, + ttl: CacheTtl, + }, + Feature(Feature), + WallClockSecond, + WebSearch(SearchContextSize), + Guardrail(GuardrailUnit), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum TokenKind { + Text, + Audio, + Image, + Video, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum InputUnit { + Token(TokenKind), + Character, + Image, + Pixel, + Query { over_results: Option }, + Request, + Second, + TextMessage, + AudioSecond, + VideoSecond { over_seconds: Option }, +} + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum OutputUnit { + Token(TokenKind), + ReasoningToken, + Character, + Image(Option), + Pixel, + Second(Option), + VideoSecond, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum CacheReadMode { + Implicit, + Explicit, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum CacheTtl { + Default, + OneHour, + TwelveHours, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum Feature { + AnnotationPage, + CitationToken, + CodeInterpreterSession, + GoogleMapsGroundingQuery, + OcrCredit, + OcrPage, + /// `browser_open_cost_per_request`, Groq `browser.open` + BrowserOpen, + /// `cost_per_additional_result`, Parallel AI results past the ones the query price includes + AdditionalResult, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum SearchContextSize { + Low, + Medium, + High, +} + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct GuardrailUnit(Box); + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ImageSize(Box); + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct Resolution(Box); + +impl Meter { + pub(crate) const TEXT_INPUT: Self = Self::Input(InputUnit::Token(TokenKind::Text)); + + pub(crate) fn sets_prompt_band(&self) -> bool { + matches!( + self, + Self::Input(InputUnit::Token(TokenKind::Text)) + | Self::Output(OutputUnit::Token(TokenKind::Text)) + | Self::CacheRead { + kind: TokenKind::Text, + mode: CacheReadMode::Implicit, + } + | Self::CacheWrite { + kind: TokenKind::Text, + .. + } + ) + } + + pub(crate) fn fallback(&self) -> Option { + match self { + Self::CacheRead { + kind: TokenKind::Text, + mode: CacheReadMode::Implicit, + } + | Self::CacheWrite { + kind: TokenKind::Text, + ttl: CacheTtl::Default, + } => Some(Self::TEXT_INPUT), + Self::CacheWrite { + kind: TokenKind::Text, + ttl: CacheTtl::OneHour | CacheTtl::TwelveHours, + } => Some(Self::CacheWrite { + kind: TokenKind::Text, + ttl: CacheTtl::Default, + }), + _ => None, + } + } +} + +impl GuardrailUnit { + pub fn new(counter: &str) -> Self { + Self(counter.into()) + } + pub fn as_str(&self) -> &str { + &self.0 + } +} + +impl ImageSize { + pub fn new(label: &str) -> Option { + let number = label.strip_suffix('K').unwrap_or(label); + let (whole, fraction) = number.split_once('.').unwrap_or((number, "0")); + (is_digits(whole) && is_digits(fraction)).then(|| Self(label.into())) + } + pub fn as_str(&self) -> &str { + &self.0 + } +} + +impl Resolution { + pub fn new(label: &str) -> Option { + let digits = label.strip_suffix(['p', 'k'])?; + is_digits(digits).then(|| Self(label.into())) + } + pub fn as_str(&self) -> &str { + &self.0 + } +} + +fn is_digits(value: &str) -> bool { + !value.is_empty() && value.bytes().all(|byte| byte.is_ascii_digit()) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::kilo("1K", true)] + #[case::half_kilo("0.5K", true)] + #[case::pixels("1024", true)] + #[case::word("big", false)] + #[case::bare_suffix("K", false)] + #[case::trailing_dot("1.K", false)] + fn image_size_labels_are_numbers_with_an_optional_k(#[case] label: &str, #[case] valid: bool) { + assert_eq!(ImageSize::new(label).is_some(), valid); + } + + #[rstest] + #[case::progressive("720p", true)] + #[case::kilo("4k", true)] + #[case::bare("720", false)] + #[case::word("hd", false)] + fn resolution_labels_are_numbers_with_p_or_k(#[case] label: &str, #[case] valid: bool) { + assert_eq!(Resolution::new(label).is_some(), valid); + } +} diff --git a/litellm-rust/crates/cost/src/priced.rs b/litellm-rust/crates/cost/src/priced.rs new file mode 100644 index 00000000000..535560c1b9a --- /dev/null +++ b/litellm-rust/crates/cost/src/priced.rs @@ -0,0 +1,53 @@ +use std::time::Duration; + +use crate::money::{UnitPrice, Usd}; +use crate::price::Meter; + +#[derive(Clone, Debug, PartialEq)] +pub enum Outcome { + Priced(Priced), + /// No meter with usage had a price + Unpriced { + missing: Vec, + }, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct Priced { + /// One line per billed meter, fallbacks included + pub lines: Vec, + /// Meters with usage and no price and no fallback + pub missing: Vec, + /// Sum of `lines` + pub total: Usd, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct Line { + /// The meter this line bills + pub meter: Meter, + /// How much of the meter the call used + pub quantity: Quantity, + /// Rate after off-peak and before uplifts + pub unit_price: UnitPrice, + /// `unit_price * quantity` times every uplift that applies + pub amount: Usd, + /// The meter whose rate was borrowed, when this meter had none + pub fallback_from: Option, + /// Which half of main's `(prompt, completion)` pair the line belongs to + pub side: Side, +} + +/// A billed amount of one meter +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Quantity { + Count(u64), + Duration(Duration), +} + +/// Meters main bills as one total (video, Replicate, Amazon Transcribe) take a fixed side, chosen in 09 +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Side { + Prompt, + Completion, +} diff --git a/litellm-rust/crates/cost/src/schedule.rs b/litellm-rust/crates/cost/src/schedule.rs new file mode 100644 index 00000000000..adebc044cca --- /dev/null +++ b/litellm-rust/crates/cost/src/schedule.rs @@ -0,0 +1,241 @@ +use std::str::FromStr; + +use crate::Error; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Schedule { + pub windows: Vec, + pub weekday_timezone: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Window { + pub hours: Vec, + pub weekdays: Option>, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct UtcWindow { + pub start_minute: u16, + pub end_minute: u16, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum Weekday { + Monday, + Tuesday, + Wednesday, + Thursday, + Friday, + Saturday, + Sunday, +} + +impl FromStr for UtcWindow { + type Err = Error; + + fn from_str(window: &str) -> Result { + let (start, end) = window.split_once('-').ok_or(Error::InvalidUtcWindow)?; + Ok(Self { + start_minute: minute_of_day(start).ok_or(Error::InvalidUtcWindow)?, + end_minute: minute_of_day(end).ok_or(Error::InvalidUtcWindow)?, + }) + } +} + +fn minute_of_day(clock: &str) -> Option { + let (hour, minute) = clock.split_once(':')?; + if hour.len() != 2 || minute.len() != 2 { + return None; + } + let hour: u16 = hour.parse().ok().filter(|hour| *hour < 24)?; + let minute: u16 = minute.parse().ok().filter(|minute| *minute < 60)?; + Some(hour * 60 + minute) +} + +impl Weekday { + pub fn from_iso(number: u8) -> Option { + match number { + 1 => Some(Self::Monday), + 2 => Some(Self::Tuesday), + 3 => Some(Self::Wednesday), + 4 => Some(Self::Thursday), + 5 => Some(Self::Friday), + 6 => Some(Self::Saturday), + 7 => Some(Self::Sunday), + _ => None, + } + } +} + +impl FromStr for Weekday { + type Err = Error; + + fn from_str(name: &str) -> Result { + match name.trim().to_ascii_lowercase().as_str() { + "mon" | "monday" => Ok(Self::Monday), + "tue" | "tues" | "tuesday" => Ok(Self::Tuesday), + "wed" | "wednesday" => Ok(Self::Wednesday), + "thu" | "thur" | "thurs" | "thursday" => Ok(Self::Thursday), + "fri" | "friday" => Ok(Self::Friday), + "sat" | "saturday" => Ok(Self::Saturday), + "sun" | "sunday" => Ok(Self::Sunday), + _ => Err(Error::InvalidWeekday), + } + } +} + +#[cfg(feature = "serde")] +mod serde_impls { + use std::fmt; + + use serde::de::{self, Deserialize, Deserializer, Visitor}; + + use super::{UtcWindow, Weekday}; + + impl<'de> Deserialize<'de> for UtcWindow { + fn deserialize>(deserializer: D) -> Result { + let window = String::deserialize(deserializer)?; + window.parse().map_err(de::Error::custom) + } + } + + struct WrittenWeekday; + + impl Visitor<'_> for WrittenWeekday { + type Value = Weekday; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("an ISO weekday number 1-7 or an English day name") + } + + fn visit_u64(self, number: u64) -> Result { + u8::try_from(number) + .ok() + .and_then(Weekday::from_iso) + .ok_or_else(|| E::invalid_value(de::Unexpected::Unsigned(number), &self)) + } + + fn visit_i64(self, number: i64) -> Result { + u64::try_from(number) + .map_err(|_| E::invalid_value(de::Unexpected::Signed(number), &self)) + .and_then(|number| self.visit_u64(number)) + } + + fn visit_str(self, name: &str) -> Result { + name.parse() + .map_err(|_| E::invalid_value(de::Unexpected::Str(name), &self)) + } + } + + impl<'de> Deserialize<'de> for Weekday { + fn deserialize>(deserializer: D) -> Result { + deserializer.deserialize_any(WrittenWeekday) + } + } +} + +#[cfg(feature = "schema")] +mod schema { + use std::borrow::Cow; + + use schemars::{JsonSchema, Schema, SchemaGenerator, json_schema}; + + use super::{UtcWindow, Weekday}; + + impl JsonSchema for UtcWindow { + fn inline_schema() -> bool { + true + } + fn schema_name() -> Cow<'static, str> { + "UtcWindow".into() + } + fn json_schema(_: &mut SchemaGenerator) -> Schema { + json_schema!({ + "type": "string", + "pattern": r"^([01]\d|2[0-3]):[0-5]\d-([01]\d|2[0-3]):[0-5]\d$" + }) + } + } + + impl JsonSchema for Weekday { + fn inline_schema() -> bool { + true + } + fn schema_name() -> Cow<'static, str> { + "Weekday".into() + } + fn json_schema(_: &mut SchemaGenerator) -> Schema { + json_schema!({ + "anyOf": [ + {"type": "integer", "minimum": 1, "maximum": 7}, + { + "type": "string", + "pattern": r"(?i)^\s*(mon|monday|tue|tues|tuesday|wed|wednesday|thu|thur|thurs|thursday|fri|friday|sat|saturday|sun|sunday)\s*$" + } + ] + }) + } + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::same_day("10:00-18:30", Some((600, 1110)))] + #[case::wraps_midnight("22:00-02:00", Some((1320, 120)))] + #[case::whole_day("00:00-00:00", Some((0, 0)))] + #[case::hour_out_of_range("25:00-01:00", None)] + #[case::minute_out_of_range("10:60-11:00", None)] + #[case::one_digit_hour("9:00-10:00", None)] + #[case::no_separator("10:00", None)] + fn utc_window_reads_hh_mm_ranges(#[case] text: &str, #[case] expected: Option<(u16, u16)>) { + assert_eq!( + text.parse::() + .ok() + .map(|window| (window.start_minute, window.end_minute)), + expected + ); + } + + #[rstest] + #[case::short("thurs", Some(Weekday::Thursday))] + #[case::upper("SATURDAY", Some(Weekday::Saturday))] + #[case::padded(" mon ", Some(Weekday::Monday))] + #[case::unknown("someday", None)] + #[case::empty("", None)] + fn weekday_names_match_python_after_trimming( + #[case] name: &str, + #[case] expected: Option, + ) { + assert_eq!(name.parse().ok(), expected); + } + + #[rstest] + #[case::monday(1, Some(Weekday::Monday))] + #[case::sunday(7, Some(Weekday::Sunday))] + #[case::zero(0, None)] + #[case::eight(8, None)] + fn weekday_numbers_are_iso(#[case] number: u8, #[case] expected: Option) { + assert_eq!(Weekday::from_iso(number), expected); + } + + #[cfg(feature = "serde")] + #[rstest] + #[case::number(serde_json::json!(6), Some(Weekday::Saturday))] + #[case::name(serde_json::json!("Sun"), Some(Weekday::Sunday))] + #[case::negative(serde_json::json!(-1), None)] + #[case::boolean(serde_json::json!(true), None)] + fn weekday_reads_numbers_and_names( + #[case] written: serde_json::Value, + #[case] expected: Option, + ) { + use serde::Deserialize; + + assert_eq!(Weekday::deserialize(&written).ok(), expected); + } +} diff --git a/litellm-rust/crates/cost/src/usage.rs b/litellm-rust/crates/cost/src/usage.rs new file mode 100644 index 00000000000..8eed83ac87b --- /dev/null +++ b/litellm-rust/crates/cost/src/usage.rs @@ -0,0 +1,104 @@ +use std::collections::BTreeMap; +use std::time::Duration; + +use crate::price::{CacheReadMode, CacheTtl, ImageSize, Resolution, SearchContextSize, TokenKind}; + +/// What one call consumed. Python fills it from the provider response, `calculate` never reads the catalog for it +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct Usage { + /// Prompt size as the provider reported it. Only picks `_above_Nk_tokens` bands, never billed. + /// Filled even when the call bills characters, `0` when the provider reports none + pub prompt_tokens: u64, + /// Billed token quantities by meter + pub tokens: Tokens, + /// Billed character quantities. Used only when the row's `billing.text_unit` is `Character` + pub characters: Characters, + /// Billed time per second meter, never a bare number of seconds or milliseconds + pub durations: BTreeMap, + /// Images sent and returned + pub images: Images, + /// Search queries, web searches and grounding requests + pub search: Search, + /// Realtime text messages, xAI bills them per message + pub text_messages: u64, +} + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct Tokens { + /// Uncached prompt tokens by modality + pub input: BTreeMap, + /// Prompt tokens served from cache, by modality and implicit or explicit cache + pub cache_read: BTreeMap<(TokenKind, CacheReadMode), u64>, + /// Prompt tokens written to cache, by modality and cache lifetime + pub cache_write: BTreeMap<(TokenKind, CacheTtl), u64>, + /// Completion tokens by modality, reasoning excluded + pub output: BTreeMap, + /// Reasoning tokens, billed at the reasoning rate or the text output rate + pub reasoning: u64, + /// Perplexity citation tokens, prompt side + pub citation: u64, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct Characters { + /// Characters of the formatted prompt + pub prompt: u64, + /// Characters of the response text + pub completion: u64, +} + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum DurationMeter { + /// Response time of the call, for rows priced by wall clock + WallClock, + /// Length of the input audio, a transcription file or a realtime transcription + InputAudio, + /// Realtime voice session time + LiveVoice, + /// Length of generated video, already multiplied by the video count. `resolution` picks `output_cost_per_second_` + OutputVideo { resolution: Option }, +} + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct Images { + /// Images sent in the request + pub input: u64, + /// Pixels sent in the request, for rows priced per input pixel + pub input_pixels: u64, + /// One entry per returned image, since price can depend on each image's size + pub output: Vec, +} + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct ImageOut { + /// Size the request asked for, picks a size keyed price like fal's `output_cost_per_image_` + pub size: Option, + /// Pixel width the response reported, for per-pixel and nearest-size pricing + pub width: Option, + /// Pixel height the response reported, for per-pixel and nearest-size pricing + pub height: Option, +} + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct Search { + /// Search API queries, for `input_cost_per_query` rows + pub queries: u64, + /// Results the query asked for, picks the `max_results_range` tier + pub max_results: Option, + /// Web searches a model ran as a tool + pub web_search: Option, + /// Google Maps grounding requests + pub maps_grounding: u64, + /// Groq `browser.open` calls + pub browser_opens: u64, + /// Results past the ones the query price includes. Python counts them, since the included 10 is a request default + pub additional_results: u64, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct WebSearch { + /// Web search requests the model made + pub requests: u64, + /// Context size the request asked for, picks `search_context_size_` + pub context: SearchContextSize, +} diff --git a/litellm-rust/crates/cost/tests/calculation.rs b/litellm-rust/crates/cost/tests/calculation.rs deleted file mode 100644 index af1a73d8c61..00000000000 --- a/litellm-rust/crates/cost/tests/calculation.rs +++ /dev/null @@ -1,535 +0,0 @@ -use litellm_cost::{ - Error, OffPeakRates, Pricing, PromptConvention, Rate, Rates, Request, ServiceTier, - ThresholdPolicy, ThresholdRates, TierRates, Usage, calculate, compile, -}; -use rstest::rstest; - -fn rates(input: Rate, output: Rate) -> Rates { - Rates { - input, - output, - ..Rates::EMPTY - } -} - -fn request() -> Request { - Request { - usage: Usage { - prompt_tokens: 100, - completion_tokens: 20, - cache_read_tokens: 25, - cache_write_tokens: 10, - cache_write_5m_tokens: None, - cache_write_1h_tokens: None, - prompt_convention: PromptConvention::IncludesCache, - }, - service_tier: ServiceTier::Standard, - threshold_policy: ThresholdPolicy::Exclusive, - region_multiplier: None, - billed_at_utc_minute: None, - } -} - -fn pricing(standard: Rates) -> Pricing<'static> { - Pricing { - standard, - tiers: &[], - thresholds: &[], - off_peak: None, - } -} - -#[test] -fn breakdown_and_total_agree() { - let standard = Rates { - cache_read: Rate::Value(0.5), - cache_write: Rate::Value(3.0), - ..rates(Rate::Value(2.0), Rate::Value(4.0)) - }; - let result = calculate(&pricing(standard), &request()).unwrap(); - assert_eq!(result.uncached_input, 65.0 * 2.0); - assert_eq!(result.cache_read, 25.0 * 0.5); - assert_eq!(result.cache_write_5m, 10.0 * 3.0); - assert_eq!(result.output(), 20.0 * 4.0); - assert_eq!(result.total(), result.input() + result.output()); - assert_eq!(result.rates.cache_read, 0.5); -} - -#[rstest] -#[case::missing(Rate::Missing, Rate::Missing, 200.0)] -#[case::null(Rate::Null, Rate::Missing, 200.0)] -#[case::zero(Rate::Value(0.0), Rate::Value(0.0), 130.0)] -fn absent_null_and_zero_cache_rates_are_distinct( - #[case] cache_read: Rate, - #[case] cache_write: Rate, - #[case] expected_input: f64, -) { - let base = rates(Rate::Value(2.0), Rate::Value(4.0)); - let standard = Rates { - cache_read, - cache_write, - ..base - }; - assert_eq!( - calculate(&pricing(standard), &request()).unwrap().input(), - expected_input - ); -} - -#[test] -fn equivalent_prompt_conventions_select_the_same_threshold() { - let threshold = ThresholdRates { - above_prompt_tokens: 90, - standard: rates(Rate::Value(5.0), Rate::Value(8.0)), - tiers: &[], - }; - let specification = Pricing { - standard: rates(Rate::Value(2.0), Rate::Value(4.0)), - tiers: &[], - thresholds: &[threshold], - off_peak: None, - }; - let included = request(); - let excluded = Request { - usage: Usage { - prompt_tokens: 65, - prompt_convention: PromptConvention::ExcludesCache, - ..included.usage - }, - ..included - }; - let plan = compile(&specification).unwrap(); - assert_eq!(plan.calculate(&included), plan.calculate(&excluded)); - assert_eq!(plan.calculate(&included).unwrap().rates.input, 5.0); -} - -#[test] -fn split_writes_and_invalid_accounting() { - let standard = Rates { - cache_read: Rate::Value(0.5), - cache_write: Rate::Value(3.0), - cache_write_1h: Rate::Value(5.0), - ..rates(Rate::Value(2.0), Rate::Value(4.0)) - }; - let base = request(); - let split = Request { - usage: Usage { - cache_write_5m_tokens: Some(4), - cache_write_1h_tokens: Some(6), - ..base.usage - }, - ..base - }; - let result = calculate(&pricing(standard), &split).unwrap(); - assert_eq!(result.cache_write_5m, 12.0); - assert_eq!(result.cache_write_1h, 30.0); - let overlapping = Request { - usage: Usage { - prompt_tokens: 30, - ..split.usage - }, - ..split - }; - assert_eq!( - calculate(&pricing(standard), &overlapping), - Err(Error::CacheExceedsPrompt) - ); - let incomplete = Request { - usage: Usage { - cache_write_1h_tokens: None, - ..split.usage - }, - ..split - }; - assert_eq!( - calculate(&pricing(standard), &incomplete), - Err(Error::InvalidCacheWriteDetails) - ); -} - -#[test] -fn threshold_tiers_and_boundaries() { - let priority = TierRates { - tier: ServiceTier::Priority, - rates: rates(Rate::Value(3.0), Rate::Missing), - }; - let threshold = ThresholdRates { - above_prompt_tokens: 100, - standard: rates(Rate::Value(5.0), Rate::Value(8.0)), - tiers: &[ - TierRates { - tier: ServiceTier::Priority, - rates: rates(Rate::Value(7.0), Rate::Missing), - }, - TierRates { - tier: ServiceTier::Flex, - rates: rates(Rate::Value(6.0), Rate::Missing), - }, - ], - }; - let specification = Pricing { - standard: rates(Rate::Value(2.0), Rate::Value(4.0)), - tiers: &[priority], - thresholds: &[threshold], - off_peak: None, - }; - let base = request(); - let no_cache = Request { - usage: Usage { - cache_read_tokens: 0, - cache_write_tokens: 0, - ..base.usage - }, - ..base - }; - let fast = Request { - service_tier: ServiceTier::Fast, - ..no_cache - }; - let inclusive = Request { - threshold_policy: ThresholdPolicy::Inclusive, - ..fast - }; - let flex = Request { - service_tier: ServiceTier::Flex, - ..inclusive - }; - assert_eq!(calculate(&specification, &no_cache).unwrap().input(), 200.0); - assert_eq!(calculate(&specification, &fast).unwrap().input(), 300.0); - assert_eq!( - calculate(&specification, &inclusive).unwrap().input(), - 700.0 - ); - assert_eq!(calculate(&specification, &flex).unwrap().input(), 600.0); -} - -#[rstest] -#[case::ultrafast_above_threshold(ServiceTier::Ultrafast, 300_000, 9_301_000.0, 37_000.0)] -#[case::ultrafast_at_threshold(ServiceTier::Ultrafast, 272_000, 544_500.0, 5_000.0)] -#[case::standard_above_threshold(ServiceTier::Standard, 300_000, 3_300_600.0, 13_000.0)] -#[case::priority_above_threshold(ServiceTier::Priority, 300_000, 5_701_000.0, 23_000.0)] -fn tiered_long_context_rates_are_selected_by_service_tier( - #[case] service_tier: ServiceTier, - #[case] prompt_tokens: u64, - #[case] expected_input: f64, - #[case] expected_output: f64, -) { - let standard = Rates { - cache_read: Rate::Value(3.0), - ..rates(Rate::Value(1.0), Rate::Value(2.0)) - }; - let tiers = [ - TierRates { - tier: ServiceTier::Priority, - rates: Rates { - cache_read: Rate::Value(5.0), - ..rates(Rate::Value(3.0), Rate::Value(4.0)) - }, - }, - TierRates { - tier: ServiceTier::Ultrafast, - rates: Rates { - cache_read: Rate::Value(7.0), - ..rates(Rate::Value(2.0), Rate::Value(5.0)) - }, - }, - ]; - let threshold_tiers = [ - TierRates { - tier: ServiceTier::Priority, - rates: Rates { - cache_read: Rate::Value(29.0), - ..rates(Rate::Value(19.0), Rate::Value(23.0)) - }, - }, - TierRates { - tier: ServiceTier::Ultrafast, - rates: Rates { - cache_read: Rate::Value(41.0), - ..rates(Rate::Value(31.0), Rate::Value(37.0)) - }, - }, - ]; - let thresholds = [ThresholdRates { - above_prompt_tokens: 272_000, - standard: Rates { - cache_read: Rate::Value(17.0), - ..rates(Rate::Value(11.0), Rate::Value(13.0)) - }, - tiers: &threshold_tiers, - }]; - let pricing = Pricing { - standard, - tiers: &tiers, - thresholds: &thresholds, - off_peak: None, - }; - let base = request(); - let long_context_request = Request { - usage: Usage { - prompt_tokens, - completion_tokens: 1_000, - cache_read_tokens: 100, - cache_write_tokens: 0, - ..base.usage - }, - service_tier, - ..base - }; - let cost = calculate(&pricing, &long_context_request).unwrap(); - - assert_eq!(cost.input(), expected_input); - assert_eq!(cost.output(), expected_output); -} - -#[test] -fn compile_rejects_ambiguous_rates() { - let duplicate = ThresholdRates { - above_prompt_tokens: 100, - standard: Rates::EMPTY, - tiers: &[], - }; - let specification = Pricing { - standard: rates(Rate::Value(1.0), Rate::Value(1.0)), - tiers: &[], - thresholds: &[duplicate, duplicate], - off_peak: None, - }; - assert_eq!( - compile(&specification).err(), - Some(Error::DuplicateThreshold) - ); - let invalid = pricing(rates(Rate::Value(f64::NAN), Rate::Value(1.0))); - assert_eq!(compile(&invalid).err(), Some(Error::InvalidRate)); -} - -#[test] -fn off_peak_is_one_non_wrapping_utc_window() { - let specification = Pricing { - standard: rates(Rate::Value(2.0), Rate::Value(4.0)), - tiers: &[], - thresholds: &[], - off_peak: Some(OffPeakRates { - start_utc_minute: 60, - end_utc_minute: 120, - rates: rates(Rate::Value(1.0), Rate::Value(2.0)), - }), - }; - let base = request(); - let start = Request { - billed_at_utc_minute: Some(60), - ..base - }; - let end = Request { - billed_at_utc_minute: Some(120), - ..base - }; - assert_eq!( - calculate(&specification, &base), - Err(Error::InvalidBillingTime) - ); - assert_eq!(calculate(&specification, &start).unwrap().input(), 100.0); - assert_eq!(calculate(&specification, &end).unwrap().input(), 200.0); -} - -#[test] -fn missing_rates_and_free_rates_remain_distinct() { - let base = request(); - let empty = Request { - usage: Usage { - prompt_tokens: 0, - completion_tokens: 0, - cache_read_tokens: 0, - cache_write_tokens: 0, - ..base.usage - }, - ..base - }; - assert_eq!( - calculate(&pricing(Rates::EMPTY), &empty), - Err(Error::MissingInputRate) - ); - assert_eq!( - calculate(&pricing(rates(Rate::Value(0.0), Rate::Missing)), &empty), - Err(Error::MissingOutputRate) - ); - assert_eq!( - calculate(&pricing(rates(Rate::Value(0.0), Rate::Value(0.0))), &empty) - .unwrap() - .total(), - 0.0 - ); -} - -#[test] -fn matches_executed_python_reference_cases() { - for row in include_str!("python_reference.tsv") - .lines() - .filter(|line| !line.starts_with('#')) - { - let fields: Vec<_> = row.split('\t').collect(); - let count = |index: usize| fields[index].parse::().unwrap(); - let number = |index: usize| fields[index].parse::().unwrap(); - let optional_rate = |index: usize| { - if fields[index].is_empty() { - Rate::Missing - } else { - Rate::Value(number(index)) - } - }; - let threshold = ThresholdRates { - above_prompt_tokens: if fields[9].is_empty() { 0 } else { count(9) }, - standard: rates(optional_rate(10), optional_rate(11)), - tiers: &[], - }; - let thresholds = if fields[9].is_empty() { - &[][..] - } else { - std::slice::from_ref(&threshold) - }; - let specification = Pricing { - standard: Rates { - cache_read: optional_rate(7), - cache_write: optional_rate(8), - ..rates(Rate::Value(number(5)), Rate::Value(number(6))) - }, - tiers: &[], - thresholds, - off_peak: None, - }; - let base = request(); - let input = Request { - usage: Usage { - prompt_tokens: count(1), - completion_tokens: count(2), - cache_read_tokens: count(3), - cache_write_tokens: count(4), - ..base.usage - }, - ..base - }; - let actual = calculate(&specification, &input).unwrap(); - assert_eq!(actual.input(), number(12), "{}", fields[0]); - assert_eq!(actual.output(), number(13), "{}", fields[0]); - } -} - -proptest::proptest! { - #[test] - fn equivalent_usage_conventions_and_breakdown_agree( - regular in 0_u64..1000, - read in 0_u64..1000, - write in 0_u64..1000, - output in 0_u64..1000, - ) { - let threshold = ThresholdRates { - above_prompt_tokens: 1000, - standard: rates(Rate::Value(5.0), Rate::Value(8.0)), - tiers: &[], - }; - let specification = Pricing { - standard: rates(Rate::Value(2.0), Rate::Value(4.0)), - tiers: &[], - thresholds: &[threshold], - off_peak: None, - }; - let base = request(); - let included = Request { - usage: Usage { - prompt_tokens: regular + read + write, - completion_tokens: output, - cache_read_tokens: read, - cache_write_tokens: write, - ..base.usage - }, - ..base - }; - let excluded = Request { - usage: Usage { - prompt_tokens: regular, - prompt_convention: PromptConvention::ExcludesCache, - ..included.usage - }, - ..included - }; - let plan = compile(&specification).unwrap(); - let left = plan.calculate(&included).unwrap(); - let right = plan.calculate(&excluded).unwrap(); - proptest::prop_assert_eq!(left, right); - proptest::prop_assert_eq!(left.total(), left.input() + left.output()); - } -} - -#[test] -fn regional_multiplier_applies_after_input_components_are_summed() { - let standard = Rates { - cache_read: Rate::Value(0.5), - cache_write: Rate::Value(3.0), - ..rates(Rate::Value(2.0), Rate::Value(4.0)) - }; - let base = request(); - let regional = Request { - region_multiplier: Some(1.1), - ..base - }; - let result = calculate(&pricing(standard), ®ional).unwrap(); - assert_eq!(result.input(), (65.0 * 2.0 + 25.0 * 0.5 + 10.0 * 3.0) * 1.1); - assert_eq!(result.output(), 20.0 * 4.0 * 1.1); - let invalid = Request { - region_multiplier: Some(f64::NAN), - ..base - }; - assert_eq!( - calculate(&pricing(standard), &invalid), - Err(Error::InvalidRegionMultiplier) - ); -} - -#[test] -fn compilation_sorts_thresholds_and_rejects_duplicate_tiers() { - let high = ThresholdRates { - above_prompt_tokens: 200, - standard: rates(Rate::Value(7.0), Rate::Missing), - tiers: &[], - }; - let low = ThresholdRates { - above_prompt_tokens: 100, - standard: rates(Rate::Value(5.0), Rate::Missing), - tiers: &[], - }; - let specification = Pricing { - standard: rates(Rate::Value(2.0), Rate::Value(4.0)), - tiers: &[], - thresholds: &[high, low], - off_peak: None, - }; - let base = request(); - let above_both = Request { - usage: Usage { - prompt_tokens: 201, - cache_read_tokens: 0, - cache_write_tokens: 0, - ..base.usage - }, - ..base - }; - assert_eq!( - compile(&specification) - .unwrap() - .calculate(&above_both) - .unwrap() - .rates - .input, - 7.0 - ); - let duplicate = TierRates { - tier: ServiceTier::Flex, - rates: Rates::EMPTY, - }; - let invalid = Pricing { - tiers: &[duplicate, duplicate], - thresholds: &[], - ..specification - }; - assert_eq!(compile(&invalid).err(), Some(Error::DuplicateTier)); -} diff --git a/litellm-rust/crates/cost/tests/generate_python_reference.py b/litellm-rust/crates/cost/tests/generate_python_reference.py deleted file mode 100644 index 4a179555650..00000000000 --- a/litellm-rust/crates/cost/tests/generate_python_reference.py +++ /dev/null @@ -1,96 +0,0 @@ -import subprocess -from dataclasses import dataclass -from pathlib import Path - -from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token -from litellm.types.utils import Usage - - -@dataclass(frozen=True, slots=True) -class Case: - name: str - prompt: int - completion: int - cache_read: int - cache_write: int - input_rate: float - output_rate: float - cache_read_rate: float | None = None - cache_write_rate: float | None = None - threshold: int | None = None - threshold_input_rate: float | None = None - threshold_output_rate: float | None = None - - -CASES = ( - Case("ordinary", 100, 20, 0, 0, 2.0, 4.0), - Case("cache_fallback", 100, 20, 25, 10, 2.0, 4.0), - Case("cache_specific", 100, 20, 25, 10, 2.0, 4.0, 0.5, 3.0), - Case("free_cache", 100, 20, 25, 10, 2.0, 4.0, 0.0, 0.0), - Case("threshold_below", 99, 20, 0, 0, 2.0, 4.0, threshold=100, threshold_input_rate=5.0, threshold_output_rate=8.0), - Case("threshold_at", 100, 20, 0, 0, 2.0, 4.0, threshold=100, threshold_input_rate=5.0, threshold_output_rate=8.0), - Case( - "threshold_above", 101, 20, 0, 0, 2.0, 4.0, threshold=100, threshold_input_rate=5.0, threshold_output_rate=8.0 - ), - Case( - "cache_threshold_above", - 101, - 20, - 25, - 10, - 2.0, - 4.0, - threshold=100, - threshold_input_rate=5.0, - threshold_output_rate=8.0, - ), -) - - -def reference(case: Case) -> tuple[float, float]: - info = {"input_cost_per_token": case.input_rate, "output_cost_per_token": case.output_rate} - if case.cache_read_rate is not None: - info["cache_read_input_token_cost"] = case.cache_read_rate - if case.cache_write_rate is not None: - info["cache_creation_input_token_cost"] = case.cache_write_rate - if case.threshold is not None: - info[f"input_cost_per_token_above_{case.threshold}_tokens"] = case.threshold_input_rate - info[f"output_cost_per_token_above_{case.threshold}_tokens"] = case.threshold_output_rate - details = {"cached_tokens": case.cache_read, "cache_write_tokens": case.cache_write} - usage = Usage(prompt_tokens=case.prompt, completion_tokens=case.completion, prompt_tokens_details=details) - return generic_cost_per_token( - model="synthetic", - usage=usage, - custom_llm_provider="openai", - model_info=info, - ) - - -def main() -> None: - revision = subprocess.check_output(("git", "rev-parse", "HEAD"), text=True).strip() - rows = ("# Python reference commit: " + revision,) + tuple( - "\t".join( - str(value) if value is not None else "" - for value in ( - case.name, - case.prompt, - case.completion, - case.cache_read, - case.cache_write, - case.input_rate, - case.output_rate, - case.cache_read_rate, - case.cache_write_rate, - case.threshold, - case.threshold_input_rate, - case.threshold_output_rate, - *reference(case), - ) - ) - for case in CASES - ) - Path(__file__).with_name("python_reference.tsv").write_text("\n".join(rows) + "\n") - - -if __name__ == "__main__": - main() diff --git a/litellm-rust/crates/cost/tests/python_reference.tsv b/litellm-rust/crates/cost/tests/python_reference.tsv deleted file mode 100644 index 5abb1329a61..00000000000 --- a/litellm-rust/crates/cost/tests/python_reference.tsv +++ /dev/null @@ -1,9 +0,0 @@ -# Python reference commit: dc4be2fd987c993aefcf16444e34f960c12c8627 -ordinary 100 20 0 0 2.0 4.0 200.0 80.0 -cache_fallback 100 20 25 10 2.0 4.0 200.0 80.0 -cache_specific 100 20 25 10 2.0 4.0 0.5 3.0 172.5 80.0 -free_cache 100 20 25 10 2.0 4.0 0.0 0.0 130.0 80.0 -threshold_below 99 20 0 0 2.0 4.0 100 5.0 8.0 198.0 80.0 -threshold_at 100 20 0 0 2.0 4.0 100 5.0 8.0 200.0 80.0 -threshold_above 101 20 0 0 2.0 4.0 100 5.0 8.0 505.0 160.0 -cache_threshold_above 101 20 25 10 2.0 4.0 100 5.0 8.0 505.0 160.0 diff --git a/litellm-rust/crates/model-catalog-refresh/Cargo.toml b/litellm-rust/crates/model-catalog-refresh/Cargo.toml new file mode 100644 index 00000000000..27219042c6c --- /dev/null +++ b/litellm-rust/crates/model-catalog-refresh/Cargo.toml @@ -0,0 +1,19 @@ +[package] +name = "litellm-model-catalog-refresh" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-http.workspace = true +litellm-model-catalog.workspace = true +thiserror.workspace = true +url.workspace = true + +[dev-dependencies] +litellm-error.workspace = true +litellm-http = { workspace = true, features = ["test-support"] } +reqwest.workspace = true +rstest.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/model-catalog-refresh/src/error.rs b/litellm-rust/crates/model-catalog-refresh/src/error.rs new file mode 100644 index 00000000000..656f32370cb --- /dev/null +++ b/litellm-rust/crates/model-catalog-refresh/src/error.rs @@ -0,0 +1,105 @@ +use std::io; + +use thiserror::Error; + +#[derive(Debug, Error)] +#[error(transparent)] +pub struct Error(Box); + +#[derive(Debug, Error)] +enum Repr { + #[error(transparent)] + Http(#[from] litellm_http::Error), + #[error("could not read the catalog")] + Read(#[source] io::Error), + #[error("invalid catalog URL")] + InvalidUrl(#[source] url::ParseError), + #[error("unsupported catalog URL scheme {scheme}")] + UnsupportedScheme { scheme: String }, +} + +impl Error { + pub fn unsupported_scheme(scheme: String) -> Self { + Self(Box::new(Repr::UnsupportedScheme { scheme })) + } +} + +impl From for Error { + fn from(cause: litellm_http::Error) -> Self { + Self(Box::new(Repr::Http(cause))) + } +} + +impl From for Error { + fn from(cause: io::Error) -> Self { + Self(Box::new(Repr::Read(cause))) + } +} + +impl From for Error { + fn from(cause: url::ParseError) -> Self { + Self(Box::new(Repr::InvalidUrl(cause))) + } +} + +#[cfg(test)] +mod tests { + use litellm_error::{Report, find}; + use rstest::rstest; + + use super::*; + use crate::{DoneReason, Reply, Step}; + + #[rstest] + fn exhausted_transport_keeps_its_typed_cause() { + let reply = Reply::Transport { + error: io::Error::from(io::ErrorKind::ConnectionReset).into(), + }; + let Reply::Transport { error } = reply else { + panic!("transport failure") + }; + let step = Step::Done { + reason: DoneReason::Exhausted(error), + }; + let Step::Done { + reason: DoneReason::Exhausted(error), + } = step + else { + panic!("exhausted attempts") + }; + let cause = find::(&error).expect("typed I/O cause"); + assert_eq!(cause.kind(), io::ErrorKind::ConnectionReset); + assert_eq!( + Report::new(&error) + .to_string() + .matches(&cause.to_string()) + .count(), + 1 + ); + assert!(!error.to_string().contains(&cause.to_string())); + } + + #[rstest] + fn invalid_source_keeps_its_typed_cause() { + let error = Error::from(url::Url::parse("http://[invalid").unwrap_err()); + assert!(find::(&error).is_some()); + } + + #[rstest] + #[tokio::test] + async fn http_failures_keep_the_cause_without_the_url() { + let cause = litellm_http::Client::for_test() + .get("http://localhost:invalid/private?api_key=secret") + .send() + .await + .unwrap_err(); + let error = Error::from(litellm_http::Error::from(cause)); + assert!(find::(&error).is_some()); + let report = Report::new(&error).to_string(); + assert!(!report.contains("secret"), "{report}"); + assert!(!report.contains("private"), "{report}"); + let debug = format!("{error:?}"); + assert!(!debug.contains("secret"), "{debug}"); + assert!(!debug.contains("private"), "{debug}"); + } +} diff --git a/litellm-rust/crates/model-catalog-refresh/src/lib.rs b/litellm-rust/crates/model-catalog-refresh/src/lib.rs new file mode 100644 index 00000000000..006c4887fad --- /dev/null +++ b/litellm-rust/crates/model-catalog-refresh/src/lib.rs @@ -0,0 +1,9 @@ +mod error; +mod policy; +mod source; +mod step; + +pub use error::Error; +pub use policy::*; +pub use source::*; +pub use step::*; diff --git a/litellm-rust/crates/model-catalog-refresh/src/policy.rs b/litellm-rust/crates/model-catalog-refresh/src/policy.rs new file mode 100644 index 00000000000..bc32e40408e --- /dev/null +++ b/litellm-rust/crates/model-catalog-refresh/src/policy.rs @@ -0,0 +1,20 @@ +use std::time::Duration; + +use litellm_model_catalog::IntegrityLimits; + +/// Main's defaults, 3 attempts, 30 s max wait, 5 s timeout, backoff base 2.0, jitter 0.0 to 1.0 s +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct RetryPolicy { + pub max_attempts: u32, + pub max_wait: Duration, + pub timeout: Duration, + pub backoff_base_seconds: f64, + pub jitter_min_seconds: f64, + pub jitter_max_seconds: f64, + /// Main's 50 model floor and 0.5 shrink ratio + pub limits: IntegrityLimits, +} + +/// Main's default `LITELLM_MODEL_COST_MAP_URL` +pub const DEFAULT_REGISTRY_URL: &str = + "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"; diff --git a/litellm-rust/crates/model-catalog-refresh/src/source.rs b/litellm-rust/crates/model-catalog-refresh/src/source.rs new file mode 100644 index 00000000000..713e0386941 --- /dev/null +++ b/litellm-rust/crates/model-catalog-refresh/src/source.rs @@ -0,0 +1,17 @@ +use std::path::PathBuf; + +use url::Url; + +/// Where a fetched catalog comes from +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Source { + Http(Url), + File(PathBuf), + S3(S3Object), +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct S3Object { + pub bucket: String, + pub key: String, +} diff --git a/litellm-rust/crates/model-catalog-refresh/src/step.rs b/litellm-rust/crates/model-catalog-refresh/src/step.rs new file mode 100644 index 00000000000..822d1807916 --- /dev/null +++ b/litellm-rust/crates/model-catalog-refresh/src/step.rs @@ -0,0 +1,66 @@ +use std::time::Duration; + +use litellm_model_catalog::Catalog; + +use crate::Error; +use crate::source::Source; + +/// One attempt. Where to fetch from is the `Source`, so the request carries no URL +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct FetchRequest { + pub if_none_match: Option, + pub timeout: Duration, +} + +#[derive(Debug)] +pub enum Reply { + /// An HTTP, S3 or file answer. The body is read only on 200 + Response { + status: u16, + etag: Option, + retry_after: Option, + body: Option>, + }, + /// No answer reached us + Transport { error: Error }, +} + +/// What the fetch state machine asks its driver to do next +#[derive(Debug)] +pub enum Step { + Fetch(FetchRequest), + Sleep(Duration), + /// A parsed document to install, with its rejected row errors + Adopt { + catalog: Box, + installed: Installed, + errors: Vec, + }, + /// 304, a final status, or retries exhausted + Done { + reason: DoneReason, + }, +} + +#[derive(Debug)] +pub enum DoneReason { + NotModified, + Status(u16), + Exhausted(Error), +} + +/// What is installed now, the input to the next revalidation +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Installed { + pub origin: Origin, + pub etag: Option, + /// `git_blob_id` of the installed bytes + pub revision: String, + pub model_count: usize, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Origin { + Bundled, + Remote { source: Source }, +} diff --git a/litellm-rust/crates/model-catalog/AGENTS.md b/litellm-rust/crates/model-catalog/AGENTS.md index 0a380788c19..fc4b28d4ef0 100644 --- a/litellm-rust/crates/model-catalog/AGENTS.md +++ b/litellm-rust/crates/model-catalog/AGENTS.md @@ -1,5 +1,27 @@ -## Validation +## Layout -For `model_prices_and_context_window.json` validation, we should eventually: +`wire/` reads one entry as written. `EntryWire::read` makes one pass over the entry's keys and puts every non-null key in exactly one place, which is an attribute, a flat price parsed by `FlatKey::parse`, a price block, a `supports_*` capability or `unknown`. A known key with a bad value becomes an `Error` that names the key, and the rest of the entry is still read -- Stop skipping this crate's tests +`lower.rs` is the only code that knows both the wire types and `litellm-cost`. It lowers an `EntryWire` into `Prices`, `Adjustments` and `Model` and checks the rules that span keys, such as one price per meter and condition, contiguous tiers, no suffix inside a block, and off-peak needing hours and prices. Lowering only goes one way. Nothing turns `Prices` back into JSON, since the raw fields are kept + +`catalog.rs` reads the document, indexes aliases and runs the checks across entries. A bad entry keeps its raw fields and its errors are listed in `Loaded.errors`. Deciding what to do with them is the caller's job. `Catalog::written` gives back every top-level key and value as the document wrote them, in document order. `Catalog::served` gives back every model and alias as Python serves it, without `sample_spec` and `fallback_generalizations` + +`Billing` on `Prices` holds the row attributes pricing reads (text unit, band edge, cache-read fallback, wall clock, web search unit, tier scope). `lower.rs` derives it from the row and its `litellm_provider`, so `litellm-cost` holds no provider list + +`store/` is the catalog's in-memory store and the only place that indexes model names. Its rules are in [src/store/AGENTS.md](src/store/AGENTS.md) + +`layers.rs`, `resolve/`, `model_info.rs` and `register.rs` hold types only. The functions over them land with their first caller + +Catalog errors use one opaque `Error` with typed causes and JSON Pointer accessors (`path`, `key`, `field`). `Error::rejects` alone decides whether an error fails the document, rejects its row or neither, and `ErrorKind` is the stable name hosts report it by + +## Prices + +The price vocabulary lives in `litellm-cost`, and only this crate knows JSON key names. Every price is keyed by `Meter`, `Channel` and `When`. A variation only one meter has stays in the meter (cache TTL, image size, resolution, `_above_Ns_interval`, `max_results_range`). Service tier and prompt band make up `When`. Batch and decisions are a `Channel` because Python bills them with their own rules. Off-peak is an overlay on `Prices`, and regional uplifts are `Adjustments`, never prices + +A new `N`, image size or resolution in an existing key shape parses with no Rust change. A new unit, token kind, feature or tier needs a variant in `litellm-cost` plus its spelling in `wire/flat_key.rs` and `FLAT_KEY_PATTERN` in `schema.rs` + +`Prices::band` and `Prices::rate` are pinned to Python's `_get_token_base_cost` by `tests/price_lookup_reference.json`. After a billing change in Python, regenerate it from the repo root with `PYTHONPATH=. python litellm-rust/crates/model-catalog/tests/generate_price_lookup_reference.py` + +## Tests + +`BUNDLED_REGISTRY` embeds the root `model_prices_and_context_window.json`, so the tests over it fail a JSON change with a bad value or a key no type reads diff --git a/litellm-rust/crates/model-catalog/Cargo.toml b/litellm-rust/crates/model-catalog/Cargo.toml index 2c2866d7090..45d16515dcd 100644 --- a/litellm-rust/crates/model-catalog/Cargo.toml +++ b/litellm-rust/crates/model-catalog/Cargo.toml @@ -6,15 +6,19 @@ license.workspace = true repository.workspace = true [features] -schema = ["dep:schemars", "litellm-llms-types/schema"] +schema = ["dep:schemars", "litellm-cost/schema"] [dependencies] -litellm-llms-types.workspace = true +litellm-cost = { workspace = true, features = ["serde"] } indexmap = { version = "2.14.0", features = ["serde"] } +jiff = { version = "0.2.37", default-features = false, features = ["std", "tzdb-bundle-always"] } +macro_rules_attribute.workspace = true schemars = { workspace = true, optional = true } serde.workspace = true -serde_json.workspace = true +serde_json = { workspace = true, features = ["preserve_order"] } +serde_with.workspace = true +strum.workspace = true thiserror.workspace = true time.workspace = true diff --git a/litellm-rust/crates/model-catalog/src/bundled.rs b/litellm-rust/crates/model-catalog/src/bundled.rs new file mode 100644 index 00000000000..d71a2d57ebb --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/bundled.rs @@ -0,0 +1,3 @@ +/// The registry document this build serves when no other one is loaded. +pub const BUNDLED_REGISTRY: &[u8] = + include_bytes!("../../../../model_prices_and_context_window.json"); diff --git a/litellm-rust/crates/model-catalog/src/capabilities.rs b/litellm-rust/crates/model-catalog/src/capabilities.rs deleted file mode 100644 index 30724766aeb..00000000000 --- a/litellm-rust/crates/model-catalog/src/capabilities.rs +++ /dev/null @@ -1,67 +0,0 @@ -use serde::{Deserialize, Serialize}; - -/// Primary API surface / task type of the model. -#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(rename_all = "snake_case")] -pub enum Mode { - AudioSpeech, - AudioTranscription, - Chat, - Completion, - Decisions, - Embedding, - Evaluation, - Guardrail, - ImageEdit, - ImageGeneration, - Moderation, - Ocr, - Realtime, - Rerank, - Responses, - Search, - VectorStore, - VideoGeneration, -} - -/// Gemini audio generation API the model is served through. -#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(rename_all = "snake_case")] -pub enum VertexAiAudioApi { - LyriaPredict, - LyriaInteractions, -} - -/// Audio container format the model can return. -#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(rename_all = "snake_case")] -pub enum AudioFormat { - Mp3, - Wav, -} - -/// Input modality the model accepts. -#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(rename_all = "snake_case")] -pub enum InputModality { - Text, - Image, - Audio, - Video, -} - -/// Output modality the model can produce. -#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(rename_all = "snake_case")] -pub enum OutputModality { - Text, - Image, - Audio, - Video, - Code, -} diff --git a/litellm-rust/crates/model-catalog/src/catalog.rs b/litellm-rust/crates/model-catalog/src/catalog.rs index 37412ef1e14..ac5f4ffcf25 100644 --- a/litellm-rust/crates/model-catalog/src/catalog.rs +++ b/litellm-rust/crates/model-catalog/src/catalog.rs @@ -1,54 +1,29 @@ -use crate::error::Error; -use crate::fallback::{FallbackGeneralizations, FallbackRule}; -use crate::model_info::ModelInfo; +use std::borrow::Cow; +use std::collections::HashSet; + use indexmap::IndexMap; use serde::Deserialize; -use serde_json::{Map, Value}; -use std::collections::HashMap; +use serde_json::Value; -#[derive(Clone, Debug, Default, PartialEq, Eq)] -pub struct Provenance { - pub source: Option, - pub revision: Option, - pub etag: Option, -} +use crate::error::{Error, Loaded, Rejects, Repr}; +use crate::fallback::{FallbackGeneralizations, FallbackRule}; +use crate::model::Model; +use crate::store::{Rows, Store}; +use crate::wire::{EntryWire, Fields}; + +const SAMPLE_SPEC: &str = "sample_spec"; +const FALLBACK_GENERALIZATIONS: &str = "fallback_generalizations"; #[derive(Clone, Copy, Debug, PartialEq)] pub struct IntegrityLimits { - pub reference_model_count: usize, - pub min_model_count: usize, - pub min_reference_ratio: f64, -} - -#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] -pub enum AliasIssue { - #[error("model {model:?} has aliases that are not an array")] - InvalidList { model: String }, - #[error("model {model:?} has an alias that is not a string")] - InvalidName { model: String }, - #[error("alias {alias:?} of model {model:?} is already a model name")] - CanonicalCollision { model: String, alias: String }, - #[error("alias {alias:?} of model {model:?} is already an alias of another model")] - AliasCollision { model: String, alias: String }, + min_model_count: usize, + min_reference_ratio: f64, } #[derive(Clone, Debug)] pub struct ModelEntry { - fields: Map, - info: ModelInfo, -} - -impl ModelEntry { - pub fn field(&self, name: &str) -> Option<&Value> { - self.fields.get(name) - } - pub fn fields(&self) -> &Map { - &self.fields - } - /// The entry deserialized into the typed mirror of the catalog schema. - pub fn info(&self) -> &ModelInfo { - &self.info - } + fields: Fields, + model: Option, } #[derive(Clone, Copy, Debug)] @@ -58,142 +33,312 @@ pub struct ModelMatch<'a> { pub entry: &'a ModelEntry, } -#[derive(Debug)] +#[derive(Clone, Debug)] +struct Section { + index: usize, + key: String, + value: Value, +} + +/// One fetched key kept off the served map because its row failed validation +#[derive(Clone, Debug)] +pub struct Rejection<'a> { + pub key: String, + pub errors: Vec<&'a Error>, + /// The bundled row for the key is served instead. `false` means the key is left out + pub replaced_by_bundled: bool, +} + +#[derive(Clone, Copy, Debug, serde::Serialize)] +#[serde(untagged)] +pub enum Written<'a> { + Entry(&'a Fields), + Section(&'a Value), +} + +#[derive(Clone, Debug)] pub struct Catalog { + store: Store, + sections: Vec
, + fallback_generalizations: Option, +} + +/// Models and aliases, an alias served as its owner's row +#[derive(Clone, Debug)] +struct CatalogRows { entries: IndexMap, aliases: IndexMap, - lowercase_keys: HashMap, - sample_spec: Option, - fallback_generalizations: Option, - provenance: Provenance, - alias_issues: Vec, +} + +impl CatalogRows { + fn owner<'a>(&'a self, key: &'a str) -> &'a str { + self.aliases.get(key).map_or(key, String::as_str) + } + + fn fields(&self, key: &str) -> Option<&Fields> { + Some(self.entries.get(self.owner(key))?.fields()) + } +} + +/// Model names match ignoring case, and keys are grouped by their raw `litellm_provider` +impl Rows for CatalogRows { + type Row = Fields; + type Stamp = (); + + fn fold(name: &str) -> String { + name.to_lowercase() + } + + fn stamp(&self) -> Self::Stamp {} + + fn contains(&self, key: &str) -> bool { + self.entries.contains_key(key) || self.aliases.contains_key(key) + } + + fn row(&self, key: &str) -> Option> { + self.fields(key).map(Cow::Borrowed) + } + + fn scan(&self, visit: &mut dyn FnMut(&str, Option<&str>)) { + let keys = self.entries.keys().chain(self.aliases.keys()); + for key in keys { + let provider = self + .fields(key) + .and_then(|fields| fields.get("litellm_provider")?.as_str()); + visit(key, provider); + } + } +} + +impl IntegrityLimits { + pub fn new(min_model_count: usize, min_reference_ratio: f64) -> Option { + (0.0..=1.0).contains(&min_reference_ratio).then_some(Self { + min_model_count, + min_reference_ratio, + }) + } +} + +impl ModelEntry { + pub fn field(&self, name: &str) -> Option<&Value> { + self.fields.get(name) + } + + pub fn fields(&self) -> &Fields { + &self.fields + } + + pub fn model(&self) -> Option<&Model> { + self.model.as_ref() + } } impl Catalog { - pub fn parse(body: &[u8], provenance: Provenance) -> Result { - let root: IndexMap = serde_json::from_slice(body)?; - if root.is_empty() { - return Err(Error::Empty); - } + pub fn parse(body: &[u8]) -> Result, Error> { + let document = serde_json::from_slice::(body) + .map_err(|error| Error::new("", Repr::Json(error)))?; - let mut entries = IndexMap::with_capacity(root.len()); + let mut problems = Vec::new(); + let mut sections = Vec::new(); + let mut entries = IndexMap::with_capacity(document.len()); let mut alias_lists = Vec::new(); - let mut alias_issues = Vec::new(); - let mut sample_spec = None; - let mut fallback_generalizations = None; - for (name, value) in root { - match name.as_str() { - "sample_spec" => { - sample_spec = Some(value); + for (index, (name, value)) in document.into_iter().enumerate() { + let fields = match value { + Value::Object(fields) if !is_section(&name) => fields, + value => { + if !is_section(&name) { + problems.push(Error::from_key(&name, Repr::NotAnObject)); + } + sections.push(Section { + index, + key: name, + value, + }); continue; } - "fallback_generalizations" => { - fallback_generalizations = - Some(serde_json::from_value::(value)?); - continue; - } - _ => {} - } - let Value::Object(mut fields) = value else { - return Err(Error::EntryNotObject { model: name }); }; - let info = ModelInfo::deserialize(&fields)?; - if let Some(aliases) = fields.remove("aliases") { - match aliases { - Value::Array(names) => alias_lists.push((name.clone(), names)), - _ => alias_issues.push(AliasIssue::InvalidList { - model: name.clone(), - }), - } + let (model, aliases, row_problems) = read_row(&fields); + problems.extend( + row_problems + .into_iter() + .map(|problem| problem.prepend(&name)), + ); + if let Some(aliases) = aliases { + alias_lists.push((name.clone(), aliases)); } - entries.insert(name, ModelEntry { fields, info }); + entries.insert(name, ModelEntry { fields, model }); + } + if entries.is_empty() { + return Err(Error::new("", Repr::Empty)); } - let mut aliases = IndexMap::new(); + let mut aliases: IndexMap = IndexMap::new(); for (model, names) in alias_lists { - for name in names { - let Value::String(alias) = name else { - alias_issues.push(AliasIssue::InvalidName { - model: model.clone(), - }); - continue; - }; + for alias in names { if entries.contains_key(&alias) { - alias_issues.push(AliasIssue::CanonicalCollision { - model: model.clone(), - alias, - }); - } else if aliases.contains_key(&alias) { - alias_issues.push(AliasIssue::AliasCollision { - model: model.clone(), - alias, - }); + problems.push(Error::from_key(&model, Repr::AliasShadowsModel { alias })); + } else if let Some(taken_by) = aliases.get(&alias) { + problems.push(Error::from_key( + &model, + Repr::AliasTaken { + alias, + taken_by: taken_by.clone(), + }, + )); } else { aliases.insert(alias, model.clone()); } } } - let lowercase_keys = entries - .keys() - .chain(aliases.keys()) - .map(|key| (key.to_lowercase(), key.clone())) - .collect(); - Ok(Self { - entries, - aliases, - lowercase_keys, - sample_spec, - fallback_generalizations, - provenance, - alias_issues, + let fallback_generalizations = sections + .iter() + .find(|section| section.key == FALLBACK_GENERALIZATIONS) + .and_then(|section| { + FallbackGeneralizations::deserialize(§ion.value) + .map_err(|error| { + problems.push(Error::from_key(§ion.key, Repr::Invalid(error))) + }) + .ok() + }); + Ok(Loaded { + value: Self { + store: Store::new(CatalogRows { entries, aliases }), + sections, + fallback_generalizations, + }, + errors: problems, }) } - pub fn validate(&self, limits: IntegrityLimits) -> Result<(), Error> { - if !limits.min_reference_ratio.is_finite() - || !(0.0..=1.0).contains(&limits.min_reference_ratio) + pub fn row_errors(model: &str, fields: &Fields) -> Vec { + let (_, _, errors) = read_row(fields); + errors + .into_iter() + .map(|error| error.prepend(model)) + .collect() + } + + pub fn written(&self) -> impl Iterator)> { + let mut sections = self.sections.iter().peekable(); + let mut entries = self.rows().entries.iter(); + (0..self.rows().entries.len() + self.sections.len()).map_while(move |index| match sections + .next_if(|section| section.index == index) { - return Err(Error::InvalidRatio); - } - let actual = self.entries.len(); + Some(section) => Some((section.key.as_str(), Written::Section(§ion.value))), + None => entries + .next() + .map(|(name, entry)| (name.as_str(), Written::Entry(&entry.fields))), + }) + } + + pub fn served(&self) -> impl Iterator { + let rows = self.rows(); + rows.entries + .iter() + .map(|(name, entry)| (name.as_str(), entry.fields())) + .chain( + rows.aliases + .iter() + .map(|(alias, owner)| (alias.as_str(), rows.entries[owner].fields())), + ) + } + + pub fn check_size(&self, reference: usize, limits: IntegrityLimits) -> Result<(), Error> { + let actual = self.model_count(); if actual < limits.min_model_count { - return Err(Error::BelowMinimum { - actual, - minimum: limits.min_model_count, - }); + return Err(Error::new( + "", + Repr::BelowMinimum { + actual, + minimum: limits.min_model_count, + }, + )); } - if limits.reference_model_count > 0 - && (actual as f64) < (limits.reference_model_count as f64) * limits.min_reference_ratio - { - return Err(Error::Shrunk { - actual, - reference: limits.reference_model_count, - ratio: limits.min_reference_ratio, - }); + if (actual as f64) < (reference as f64) * limits.min_reference_ratio { + return Err(Error::new( + "", + Repr::Shrunk { + actual, + reference, + ratio: limits.min_reference_ratio, + }, + )); } Ok(()) } - pub fn lookup(&self, key: &str) -> Option> { - let matched_key = if self.entries.contains_key(key) || self.aliases.contains_key(key) { - key - } else { - self.lowercase_keys.get(&key.to_lowercase())?.as_str() - }; - let canonical_key = self - .aliases - .get(matched_key) - .map(String::as_str) - .unwrap_or(matched_key); - let (canonical_key, entry) = self.entries.get_key_value(canonical_key)?; - let matched_key = self + pub fn without_rejected(self, errors: &[Error]) -> (Self, Vec<&Error>) { + let row_names: HashSet = self + .rows() .entries - .get_key_value(matched_key) + .keys() + .cloned() + .chain( + self.sections + .iter() + .filter(|section| !is_section(§ion.key)) + .map(|section| section.key.clone()), + ) + .collect(); + let rejected: HashSet = errors + .iter() + .filter(|error| error.rejects() == Rejects::Row) + .filter_map(|error| crate::error::pointer_root(error.path())) + .filter(|root| row_names.contains(root)) + .collect(); + let removed_errors = errors + .iter() + .filter(|error| { + crate::error::pointer_root(error.path()) + .is_some_and(|root| rejected.contains(&root)) + }) + .collect(); + let removed_positions: HashSet = self + .written() + .enumerate() + .filter(|(_, (key, _))| rejected.contains(*key)) + .map(|(index, _)| index) + .collect(); + let mut catalog = self; + catalog.store.edit(|rows| { + rows.entries.retain(|name, _| !rejected.contains(name)); + rows.aliases.retain(|_, owner| !rejected.contains(owner)); + }); + catalog + .sections + .retain(|section| is_section(§ion.key) || !rejected.contains(§ion.key)); + for section in &mut catalog.sections { + let removed_before = removed_positions + .iter() + .filter(|index| **index < section.index) + .count(); + section.index -= removed_before; + } + (catalog, removed_errors) + } + + pub fn unrecognized(&self) -> impl Iterator { + self.rows().entries.iter().flat_map(|(name, entry)| { + entry + .model + .iter() + .flat_map(|model| &model.unrecognized) + .map(move |key| (name.as_str(), key.as_str())) + }) + } + + pub fn lookup(&self, key: &str) -> Option> { + let rows = self.rows(); + let found = self.store.find_key(key)?; + let (canonical_key, entry) = rows.entries.get_key_value(rows.owner(&found))?; + let matched_key = rows + .entries + .get_key_value(found.as_str()) .map(|(key, _)| key.as_str()) .or_else(|| { - self.aliases - .get_key_value(matched_key) + rows.aliases + .get_key_value(found.as_str()) .map(|(key, _)| key.as_str()) })?; Some(ModelMatch { @@ -204,30 +349,58 @@ impl Catalog { } pub fn model_count(&self) -> usize { - self.entries.len() + self.rows().entries.len() } + pub fn model_names(&self) -> impl Iterator { - self.entries.keys().map(String::as_str) - } - pub fn alias_count(&self) -> usize { - self.aliases.len() + self.rows().entries.keys().map(String::as_str) } + pub fn aliases(&self) -> &IndexMap { - &self.aliases + &self.rows().aliases } - pub fn alias_issues(&self) -> &[AliasIssue] { - &self.alias_issues + + fn rows(&self) -> &CatalogRows { + self.store.rows() } + pub fn sample_spec(&self) -> Option<&Value> { - self.sample_spec.as_ref() - } - pub fn fallback_generalizations(&self) -> Option<&FallbackGeneralizations> { - self.fallback_generalizations.as_ref() + self.sections + .iter() + .find(|section| section.key == SAMPLE_SPEC) + .map(|section| §ion.value) } + pub fn fallback_rules(&self) -> Option<&[FallbackRule]> { Some(self.fallback_generalizations.as_ref()?.rules.as_slice()) } - pub fn provenance(&self) -> &Provenance { - &self.provenance +} + +fn read_row(fields: &Fields) -> (Option, Option>, Vec) { + let wire = EntryWire::read(fields); + let aliases = wire.attributes.aliases.clone(); + let timezone_error = unknown_timezone(&wire); + match Model::try_from(wire) { + Ok(model) => (Some(model), aliases, timezone_error.into_iter().collect()), + Err(mut errors) => { + errors.extend(timezone_error); + (None, aliases, errors) + } } } + +fn is_section(key: &str) -> bool { + key == SAMPLE_SPEC || key == FALLBACK_GENERALIZATIONS +} + +fn unknown_timezone(wire: &EntryWire) -> Option { + let zone = wire.off_peak.as_ref()?.weekday_timezone.as_deref()?.trim(); + (!zone.is_empty() && jiff::tz::TimeZone::get(zone).is_err()).then(|| { + Error::from_path( + "off_peak_pricing.weekday_timezone", + Repr::UnknownTimezone { + zone: zone.to_owned(), + }, + ) + }) +} diff --git a/litellm-rust/crates/model-catalog/src/error.rs b/litellm-rust/crates/model-catalog/src/error.rs index 00902858be4..425ae47687f 100644 --- a/litellm-rust/crates/model-catalog/src/error.rs +++ b/litellm-rust/crates/model-catalog/src/error.rs @@ -1,21 +1,392 @@ +use std::num::ParseIntError; + +use strum::{EnumDiscriminants, IntoStaticStr, VariantArray}; use thiserror::Error; #[derive(Debug, Error)] -pub enum Error { - #[error("invalid catalog JSON")] - Json(#[from] serde_json::Error), - #[error("catalog is empty")] +#[error(transparent)] +pub struct Error(Box); + +#[derive(Debug, Error)] +enum Context { + #[error(transparent)] + Bare(Repr), + #[error("{}", display_path(path))] + At { + path: String, + #[source] + cause: Repr, + }, +} + +#[derive(Debug, Error)] +pub(crate) enum DateCause { + #[error("invalid date component")] + Number(#[from] ParseIntError), + #[error("calendar component out of range")] + Range(#[from] time::error::ComponentRange), +} + +#[derive(Debug, Error, EnumDiscriminants)] +#[strum_discriminants(name(ErrorKind), derive(IntoStaticStr, VariantArray), vis(pub))] +pub(crate) enum Repr { + #[error("invalid JSON")] + #[strum_discriminants(strum(serialize = "json"))] + Json(#[source] serde_json::Error), + #[error("has no models")] + #[strum_discriminants(strum(serialize = "empty"))] Empty, - #[error("model {model:?} must be an object")] - EntryNotObject { model: String }, - #[error("catalog has {actual} models, below minimum {minimum}")] + #[error("has {actual} models, below minimum {minimum}")] + #[strum_discriminants(strum(serialize = "below_minimum"))] BelowMinimum { actual: usize, minimum: usize }, - #[error("catalog has {actual} models, below {ratio} of reference count {reference}")] + #[error("has {actual} models, below {ratio} of reference count {reference}")] + #[strum_discriminants(strum(serialize = "shrunk"))] Shrunk { actual: usize, reference: usize, ratio: f64, }, - #[error("minimum reference ratio must be finite and between zero and one")] - InvalidRatio, + #[error("must be a JSON object")] + #[strum_discriminants(strum(serialize = "not_an_object"))] + NotAnObject, + #[error("invalid value")] + #[strum_discriminants(strum(serialize = "invalid"))] + Invalid(#[source] serde_json::Error), + #[error("not a field this block reads")] + #[strum_discriminants(strum(serialize = "unknown_block_key"))] + UnknownBlockKey, + #[error("required")] + #[strum_discriminants(strum(serialize = "missing_provider"))] + MissingProvider, + #[error("{zone:?} is not an IANA time zone")] + #[strum_discriminants(strum(serialize = "unknown_timezone"))] + UnknownTimezone { zone: String }, + #[error("alias {alias:?} is already a model name")] + #[strum_discriminants(strum(serialize = "alias_shadows_model"))] + AliasShadowsModel { alias: String }, + #[error("alias {alias:?} already belongs to {taken_by}")] + #[strum_discriminants(strum(serialize = "alias_taken"))] + AliasTaken { alias: String, taken_by: String }, + #[error("{first} and {second} price the same meter under the same conditions")] + #[strum_discriminants(strum(serialize = "duplicate_price"))] + DuplicatePrice { first: String, second: String }, + #[error("a price inside {block} takes no band or tier suffix")] + #[strum_discriminants(strum(serialize = "suffix_inside_block"))] + SuffixInsideBlock { block: &'static str }, + #[error("a tier needs exactly one of range and max_results_range")] + #[strum_discriminants(strum(serialize = "tier_range"))] + TierRange, + #[error("the range ends below its start")] + #[strum_discriminants(strum(serialize = "tier_reversed"))] + TierReversed, + #[error("starts at {start}, but the tiers before it end at {expected}")] + #[strum_discriminants(strum(serialize = "tier_gap"))] + TierGap { start: u64, expected: u64 }, + #[error("tiered_pricing mixes token ranges and result ranges")] + #[strum_discriminants(strum(serialize = "mixed_tier_basis"))] + MixedTierBasis, + #[error("a tier by result count only prices input_cost_per_query")] + #[strum_discriminants(strum(serialize = "not_priced_by_results"))] + NotPricedByResults, + #[error("tiered_pricing and prompt token bands are set in two places")] + #[strum_discriminants(strum(serialize = "bands_from_two_sources"))] + BandsFromTwoSources, + #[error("has no off-peak rate")] + #[strum_discriminants(strum(serialize = "not_an_off_peak_rate"))] + NotAnOffPeakRate, + #[error("off_peak_pricing has no prices")] + #[strum_discriminants(strum(serialize = "no_off_peak_prices"))] + NoOffPeakPrices, + #[error("off_peak_pricing needs hours_utc or windows")] + #[strum_discriminants(strum(serialize = "no_off_peak_hours"))] + NoOffPeakHours, + #[error("must list at least one entry")] + #[strum_discriminants(strum(serialize = "empty_list"))] + EmptyList, + #[error("expected a calendar date YYYY-MM-DD")] + #[strum_discriminants(strum(serialize = "invalid_date"))] + InvalidDate { + #[source] + cause: Option, + }, +} + +/// What a problem takes out of the served map +#[derive(Clone, Copy, Debug, PartialEq, Eq, VariantArray)] +pub enum Rejects { + Document, + Row, + Nothing, +} + +impl Error { + pub fn kind(&self) -> ErrorKind { + ErrorKind::from(self.repr()) + } + + pub fn rejects(&self) -> Rejects { + match self.repr() { + Repr::Json(_) | Repr::Empty | Repr::BelowMinimum { .. } | Repr::Shrunk { .. } => { + Rejects::Document + } + Repr::NotAnObject + | Repr::Invalid(_) + | Repr::UnknownBlockKey + | Repr::UnknownTimezone { .. } + | Repr::AliasShadowsModel { .. } + | Repr::AliasTaken { .. } + | Repr::DuplicatePrice { .. } + | Repr::SuffixInsideBlock { .. } + | Repr::TierRange + | Repr::TierReversed + | Repr::TierGap { .. } + | Repr::MixedTierBasis + | Repr::NotPricedByResults + | Repr::BandsFromTwoSources + | Repr::NotAnOffPeakRate + | Repr::NoOffPeakPrices + | Repr::NoOffPeakHours + | Repr::EmptyList + | Repr::InvalidDate { .. } => Rejects::Row, + Repr::MissingProvider => Rejects::Nothing, + } + } +} + +#[derive(Debug)] +pub struct Loaded { + pub value: T, + pub errors: Vec, +} + +impl From for Error { + fn from(cause: Repr) -> Self { + Self(Box::new(Context::Bare(cause))) + } +} + +impl From for Error { + fn from(cause: DateCause) -> Self { + Repr::InvalidDate { cause: Some(cause) }.into() + } +} + +impl Error { + fn repr(&self) -> &Repr { + match &*self.0 { + Context::Bare(cause) | Context::At { cause, .. } => cause, + } + } + + pub(crate) fn new(path: impl Into, cause: Repr) -> Self { + Self(Box::new(Context::At { + path: path.into(), + cause, + })) + } + + pub(crate) fn from_key(key: &str, cause: Repr) -> Self { + Self::new(pointer_for_key(key), cause) + } + + pub(crate) fn from_path(path: &str, cause: Repr) -> Self { + Self::new(path_to_pointer(path), cause) + } + + pub fn path(&self) -> &str { + match &*self.0 { + Context::Bare(_) => "", + Context::At { path, .. } => path, + } + } + + pub fn key(&self) -> Option { + pointer_root(self.path()) + } + + pub fn field(&self) -> Option { + let segments: Vec<_> = self + .path() + .split('/') + .skip(2) + .map(unescape_pointer_segment) + .collect(); + (!segments.is_empty()).then(|| segments.join("/")) + } + + pub(crate) fn prepend(self, key: &str) -> Self { + let (path, cause) = match *self.0 { + Context::Bare(cause) => (String::new(), cause), + Context::At { path, cause } => (path, cause), + }; + Self::new(format!("/{}{path}", escape_pointer_segment(key)), cause) + } +} + +pub(crate) fn pointer_root(path: &str) -> Option { + path.strip_prefix('/') + .and_then(|path| path.split('/').next()) + .map(unescape_pointer_segment) +} + +fn escape_pointer_segment(segment: &str) -> String { + segment.replace('~', "~0").replace('/', "~1") +} + +pub(crate) fn pointer_for_key(key: &str) -> String { + format!("/{}", escape_pointer_segment(key)) +} + +pub(crate) fn pointer_child(parent: &str, key: &str) -> String { + format!("{parent}/{}", escape_pointer_segment(key)) +} + +fn unescape_pointer_segment(segment: &str) -> String { + segment.replace("~1", "/").replace("~0", "~") +} + +fn path_to_pointer(key: &str) -> String { + let mut pointer = String::new(); + let mut segment = String::new(); + let mut chars = key.chars(); + let mut push_segment = |segment: &mut String| { + if !segment.is_empty() { + pointer.push('/'); + pointer.push_str(&escape_pointer_segment(segment)); + segment.clear(); + } + }; + while let Some(character) = chars.next() { + match character { + '.' => push_segment(&mut segment), + '[' => { + push_segment(&mut segment); + let mut index = String::new(); + for character in chars.by_ref() { + if character == ']' { + break; + } + index.push(character); + } + push_segment(&mut index); + } + character => segment.push(character), + } + } + push_segment(&mut segment); + pointer +} + +fn display_path(path: &str) -> String { + if path.is_empty() { + return "catalog".to_owned(); + } + path.split('/') + .skip(1) + .map(unescape_pointer_segment) + .enumerate() + .fold(String::new(), |mut display, (index, segment)| { + if index == 0 { + display.push_str(&segment); + } else if segment.bytes().all(|byte| byte.is_ascii_digit()) { + display.push('['); + display.push_str(&segment); + display.push(']'); + } else { + display.push('.'); + display.push_str(&segment); + } + display + }) +} + +#[cfg(test)] +mod tests { + use litellm_error::{Report, find}; + use rstest::rstest; + use serde_json::Value; + + use super::*; + use crate::{Catalog, wire::IsoDate}; + + #[rstest] + #[case::json(Repr::Json(serde_json::from_str::("{").unwrap_err()), Rejects::Document)] + #[case::empty(Repr::Empty, Rejects::Document)] + #[case::below_minimum(Repr::BelowMinimum { actual: 1, minimum: 2 }, Rejects::Document)] + #[case::shrunk(Repr::Shrunk { actual: 1, reference: 3, ratio: 0.5 }, Rejects::Document)] + #[case::not_an_object(Repr::NotAnObject, Rejects::Row)] + #[case::invalid(Repr::Invalid(serde_json::from_str::("x").unwrap_err()), Rejects::Row)] + #[case::unknown_block_key(Repr::UnknownBlockKey, Rejects::Row)] + #[case::missing_provider(Repr::MissingProvider, Rejects::Nothing)] + #[case::unknown_timezone(Repr::UnknownTimezone { zone: "invalid".into() }, Rejects::Row)] + #[case::alias_shadows_model(Repr::AliasShadowsModel { alias: "model".into() }, Rejects::Row)] + #[case::alias_taken(Repr::AliasTaken { alias: "alias".into(), taken_by: "owner".into() }, Rejects::Row)] + #[case::duplicate_price(Repr::DuplicatePrice { first: "a".into(), second: "b".into() }, Rejects::Row)] + #[case::suffix_inside_block(Repr::SuffixInsideBlock { block: "tiered_pricing" }, Rejects::Row)] + #[case::tier_range(Repr::TierRange, Rejects::Row)] + #[case::tier_reversed(Repr::TierReversed, Rejects::Row)] + #[case::tier_gap(Repr::TierGap { start: 1, expected: 0 }, Rejects::Row)] + #[case::mixed_tier_basis(Repr::MixedTierBasis, Rejects::Row)] + #[case::not_priced_by_results(Repr::NotPricedByResults, Rejects::Row)] + #[case::bands_from_two_sources(Repr::BandsFromTwoSources, Rejects::Row)] + #[case::not_an_off_peak_rate(Repr::NotAnOffPeakRate, Rejects::Row)] + #[case::no_off_peak_prices(Repr::NoOffPeakPrices, Rejects::Row)] + #[case::no_off_peak_hours(Repr::NoOffPeakHours, Rejects::Row)] + #[case::empty_list(Repr::EmptyList, Rejects::Row)] + #[case::invalid_date(Repr::InvalidDate { cause: None }, Rejects::Row)] + fn policy_classifies_each_error(#[case] error: Repr, #[case] expected: Rejects) { + assert_eq!(Error::from(error).rejects(), expected); + } + + #[rstest] + #[case::root("", "x", "/x")] + #[case::empty_model("/", "x", "//x")] + #[case::empty_parent_segment("/m/", "x", "/m//x")] + #[case::escaped("/m~1n", "a~/b", "/m~1n/a~0~1b")] + fn pointers_preserve_segments(#[case] parent: &str, #[case] key: &str, #[case] expected: &str) { + assert_eq!(pointer_child(parent, key), expected); + } + + #[rstest] + #[case::malformed(br#"{"#)] + #[case::bad_price(br#"{"m":{"litellm_provider":"test","input_cost_per_token":"free"}}"#)] + fn json_causes_are_typed_and_rendered_once(#[case] body: &[u8]) { + let error = match Catalog::parse(body) { + Ok(loaded) => loaded.errors.into_iter().next().expect("invalid price"), + Err(error) => error, + }; + let cause = find::(&error) + .expect("typed JSON cause") + .to_string(); + let report = Report::new(&error).to_string(); + assert_eq!(report.matches(&cause).count(), 1, "{report}"); + assert!(!error.to_string().contains(&cause)); + } + + #[rstest] + #[case::year("xxxx-01-01")] + #[case::month("2026-xx-01")] + #[case::day("2026-01-xx")] + fn date_number_errors_keep_the_cause(#[case] value: &str) { + let error = value.parse::().unwrap_err(); + let cause = find::(&error) + .expect("typed integer cause") + .to_string(); + assert_eq!(error.rejects(), Rejects::Row); + assert_eq!(Report::new(&error).to_string().matches(&cause).count(), 1); + assert!(!error.to_string().contains(&cause)); + } + + #[rstest] + #[case::month("2026-13-01")] + #[case::day("2026-02-30")] + fn date_range_errors_keep_the_cause(#[case] value: &str) { + let error = value.parse::().unwrap_err(); + let cause = find::(&error) + .expect("typed calendar cause") + .to_string(); + assert_eq!(error.rejects(), Rejects::Row); + assert_eq!(Report::new(&error).to_string().matches(&cause).count(), 1); + assert!(!error.to_string().contains(&cause)); + } } diff --git a/litellm-rust/crates/model-catalog/src/fallback.rs b/litellm-rust/crates/model-catalog/src/fallback.rs index 62291a84929..f5a0cadb6b9 100644 --- a/litellm-rust/crates/model-catalog/src/fallback.rs +++ b/litellm-rust/crates/model-catalog/src/fallback.rs @@ -1,10 +1,9 @@ -use serde::{Deserialize, Serialize}; use serde_json::Value; use std::collections::BTreeMap; /// One regex rule generalizing unknown model ids to known families. -#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[macro_rules_attribute::apply(crate::wire_type)] +#[derive(Eq, Hash)] pub struct FallbackRule { pub name: String, pub pattern: String, @@ -15,8 +14,8 @@ pub struct FallbackRule { } /// Regex rules that generalize unknown model ids to known families; not a model entry. -#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[macro_rules_attribute::apply(crate::wire_type)] +#[derive(Eq, Hash)] #[serde(deny_unknown_fields)] pub struct FallbackGeneralizations { pub rules: Vec, diff --git a/litellm-rust/crates/model-catalog/src/layers.rs b/litellm-rust/crates/model-catalog/src/layers.rs new file mode 100644 index 00000000000..881f0f6af9b --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/layers.rs @@ -0,0 +1,32 @@ +use std::sync::Arc; + +use indexmap::IndexMap; +use serde_json::Value; + +use crate::catalog::Catalog; +use crate::register::Registrations; +use crate::wire::Fields; + +/// One field of a fetched row over the layer below, RFC 7386. An absent field inherits, JSON `null` is `Remove` +#[derive(Clone, Debug, PartialEq)] +pub enum Patch { + Set(Value), + Remove, +} + +/// The four layers served rows are composed from, lowest first +#[derive(Clone, Debug)] +#[expect( + dead_code, + reason = "02 and 07 add the functions that read these layers" +)] +pub struct Layers { + /// 0, the wheel's embedded map + bundled: Arc, + /// 1, the last installed fetch, rejected rows left out + fetched: Option>, + /// 2, `register(persist=False)`, cleared by every install and refilled by the router replay + transient: IndexMap, + /// 3, `register(persist=True)`, applied over every install + persisted: Registrations, +} diff --git a/litellm-rust/crates/model-catalog/src/lib.rs b/litellm-rust/crates/model-catalog/src/lib.rs index d42c05e6fe7..fe419dde054 100644 --- a/litellm-rust/crates/model-catalog/src/lib.rs +++ b/litellm-rust/crates/model-catalog/src/lib.rs @@ -7,20 +7,18 @@ //! in `litellm-router-core`, which reads this catalog through its dynamic config. //! //! # Structure -//! - `Catalog`: the parsed map. `Catalog::parse` takes the JSON body and a `Provenance` (where -//! the body came from: the bundled copy in the repo or the remote copy, with revision and etag), -//! and keeps every entry in file order with its raw fields. +//! - `Catalog`: the parsed map. `Catalog::parse` takes the JSON body and returns `Loaded`, +//! keeping every entry in file order with its raw fields and collecting row and alias errors. //! - `ModelEntry`: one model. `ModelEntry::fields` is the raw JSON object, so unknown fields are -//! never lost; `ModelEntry::info` is the typed mirror `ModelInfo` (pricing in `pricing`, -//! modalities and modes in `capabilities`). +//! never lost; `ModelEntry::model` is the optional typed `Model` with attributes, prices, +//! adjustments and capabilities. Invalid rows retain their raw fields without a typed model. //! - Aliases: an entry may list `aliases`; each alias resolves to its canonical key. Problems -//! found while indexing them (bad lists, collisions) are kept as `AliasIssue`s instead of -//! failing the parse. +//! found while indexing them (bad lists, collisions) are kept in `Loaded::errors` instead of +//! failing the document. //! - `FallbackGeneralizations` / `FallbackRule`: regex rules that map unknown model ids onto known //! families. They are data in the file, not entries. -//! - Validation: `Catalog::validate` checks a catalog against `IntegrityLimits` (a remote copy that -//! lost most of its models is refused), and `validate_registry` checks every entry against the -//! checked-in contract. +//! - Validation: `Catalog::check_size` checks a catalog against `IntegrityLimits` (a remote copy +//! that lost most of its models is refused), and `Catalog::row_errors` validates one row. //! //! # Lookups //! `Catalog::lookup` resolves a key exactly first (entry or alias), then case-insensitively, and @@ -31,21 +29,37 @@ //! Cost tracking, the router (through `RouterDynamicConfig`) and the gateway. Nothing in this crate //! performs I/O; loading and refreshing the file is the caller's job. -mod capabilities; +macro_rules_attribute::attribute_alias! { + #[apply(wire_type)] = + #[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)] + #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]; +} + +mod bundled; mod catalog; mod error; mod fallback; +mod layers; +mod lower; +mod model; mod model_info; -mod pricing; -mod validation; +mod register; +mod resolve; +mod store; +mod vocabulary; +pub mod wire; -pub use capabilities::*; +pub use bundled::*; pub use catalog::*; pub use error::*; pub use fallback::*; +pub use layers::*; +pub use model::*; pub use model_info::*; -pub use pricing::*; -pub use validation::*; +pub use register::*; +pub use resolve::*; +pub use store::*; +pub use vocabulary::*; #[cfg(feature = "schema")] mod schema; diff --git a/litellm-rust/crates/model-catalog/src/lower.rs b/litellm-rust/crates/model-catalog/src/lower.rs new file mode 100644 index 00000000000..4c3fc38a01a --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/lower.rs @@ -0,0 +1,855 @@ +use std::collections::BTreeMap; +use std::collections::btree_map::Entry; +use std::num::NonZeroU64; + +use litellm_cost::money::UnitPrice; +use litellm_cost::price::{ + Adjustments, Band, BandEdge, Billing, CacheReadFallback, CacheReadMode, CacheTtl, + GuardrailUnit, InputUnit, Meter, OffPeak, OffPeakRates, OutputUnit, Price, Prices, Residency, + SearchContextSize, TextUnit, TierScope, TokenKind, Uplift, WallClock, WebSearchUnit, When, +}; +use litellm_cost::schedule::{Schedule, Window}; + +use crate::error::{Error, Repr, pointer_child, pointer_for_key}; +use crate::model::Model; +use crate::vocabulary::Mode; +use crate::wire::{EntryWire, FlatPrice, OffPeakBlock, TierRow, WebSearchBillingUnit}; + +impl TryFrom<&EntryWire> for Prices { + type Error = Vec; + + fn try_from(wire: &EntryWire) -> Result { + let mut problems = Vec::new(); + let sourced = flat_prices(&wire.prices) + .chain(tier_prices(wire, &mut problems)) + .chain(search_context_prices(wire)) + .chain(guardrail_prices(wire)); + let rates = dedupe(sourced, &mut problems); + let off_peak = wire + .off_peak + .as_ref() + .and_then(|block| off_peak(block, &mut problems)); + if problems.is_empty() { + let billing = billing(wire, &rates); + Ok(Prices::new(rates, off_peak, billing)) + } else { + Err(problems) + } + } +} + +fn billing(wire: &EntryWire, rates: &BTreeMap) -> Billing { + let provider = wire.attributes.litellm_provider.as_deref(); + let prices_meter = |priced: fn(&Meter) -> bool| rates.keys().any(|price| priced(&price.meter)); + let prices_tokens = prices_meter(|meter| { + matches!( + meter, + Meter::Input(InputUnit::Token(_)) | Meter::Output(OutputUnit::Token(_)) + ) + }); + let prices_characters = prices_meter(|meter| { + matches!( + meter, + Meter::Input(InputUnit::Character) | Meter::Output(OutputUnit::Character) + ) + }); + let prices_cache_read = prices_meter(|meter| { + matches!( + meter, + Meter::CacheRead { + kind: TokenKind::Text, + .. + } + ) + }); + Billing { + text_unit: if prices_characters && !prices_tokens { + TextUnit::Character + } else { + TextUnit::Token + }, + band_edge: if provider == Some("xai") { + BandEdge::Inclusive + } else { + BandEdge::Exclusive + }, + cache_read_fallback: if provider == Some("fireworks_ai") && !prices_cache_read { + CacheReadFallback::HalfInputRate + } else { + CacheReadFallback::InputRate + }, + wall_clock: match wire.attributes.mode { + None | Some(Mode::Chat | Mode::Completion | Mode::Embedding | Mode::Responses) => { + WallClock::ResponseTime + } + Some(Mode::AudioTranscription) => WallClock::AudioFile, + Some(_) => WallClock::Never, + }, + web_search: match wire.attributes.web_search_billing_unit { + Some(WebSearchBillingUnit::PerQuery) => WebSearchUnit::PerQuery, + Some(WebSearchBillingUnit::PerPrompt) | None => WebSearchUnit::PerPrompt, + }, + tier_scope: if provider == Some("dashscope") && wire.tier_rows.is_some() { + TierScope::WholeCall + } else { + TierScope::PerMeter + }, + } +} + +impl From<&EntryWire> for Adjustments { + fn from(wire: &EntryWire) -> Self { + Adjustments::new( + [ + (Uplift::NonGlobalEndpoint, wire.regional_endpoint_uplift), + (Uplift::Residency(Residency::Eu), wire.residency_uplift_eu), + (Uplift::Residency(Residency::Us), wire.residency_uplift_us), + ] + .into_iter() + .filter_map(|(uplift, factor)| Some((uplift, factor?))), + ) + } +} + +impl TryFrom for Model { + type Error = Vec; + + fn try_from(wire: EntryWire) -> Result { + let prices = Prices::try_from(&wire); + let adjustments = Adjustments::from(&wire); + let EntryWire { + attributes, + capabilities, + unknown, + problems, + .. + } = wire; + let missing_provider = attributes.litellm_provider.is_none().then_some(Error::new( + pointer_for_key("litellm_provider"), + Repr::MissingProvider, + )); + match prices { + Ok(prices) if problems.is_empty() && missing_provider.is_none() => Ok(Model { + attributes, + prices, + adjustments, + capabilities, + unrecognized: unknown.into_iter().map(|(key, _)| key).collect(), + }), + prices => Err(problems + .into_iter() + .chain(missing_provider) + .chain(prices.err().into_iter().flatten()) + .collect()), + } + } +} + +struct Sourced { + price: Price, + rate: UnitPrice, + source: String, + path: String, +} + +fn flat_prices(prices: &[FlatPrice]) -> impl Iterator + '_ { + prices.iter().map(|flat| { + let (channel, when) = flat.flat.channel_and_when(); + Sourced { + price: Price { + meter: flat.flat.meter.clone(), + channel, + when, + }, + rate: flat.price, + source: flat.key.clone(), + path: flat.path.clone(), + } + }) +} + +fn dedupe( + sourced: impl Iterator, + problems: &mut Vec, +) -> BTreeMap { + let mut seen: BTreeMap = BTreeMap::new(); + for Sourced { + price, + rate, + source, + path, + } in sourced + { + match seen.entry(price) { + Entry::Occupied(first) if first.get().0 == rate => {} + Entry::Occupied(first) => problems.push(Error::new( + path, + Repr::DuplicatePrice { + first: first.get().1.clone(), + second: source.clone(), + }, + )), + Entry::Vacant(slot) => { + slot.insert((rate, source, path)); + } + } + } + seen.into_iter() + .map(|(price, (rate, _, _))| (price, rate)) + .collect() +} + +#[derive(Clone, Copy, PartialEq, Eq)] +enum TierBasis { + PromptTokens, + Results, +} + +fn tier_basis(row: &TierRow) -> Result<(TierBasis, [u64; 2]), Error> { + match (row.range, row.max_results_range) { + (Some(range), None) => Ok((TierBasis::PromptTokens, range)), + (None, Some(range)) => Ok((TierBasis::Results, range)), + _ => Err(Error::new(row.path.clone(), Repr::TierRange)), + } +} + +fn tier_prices(wire: &EntryWire, problems: &mut Vec) -> Vec { + let Some(rows) = wire.tier_rows.as_deref().filter(|rows| !rows.is_empty()) else { + return Vec::new(); + }; + let based: Vec<_> = rows + .iter() + .filter_map(|row| match tier_basis(row) { + Ok((basis, range)) => Some((basis, range, row)), + Err(problem) => { + problems.push(problem); + None + } + }) + .collect(); + let Some(&(basis, ..)) = based.first() else { + return Vec::new(); + }; + if based.iter().any(|(other, ..)| *other != basis) { + problems.push(Error::from_key("tiered_pricing", Repr::MixedTierBasis)); + return Vec::new(); + } + if basis == TierBasis::PromptTokens { + problems.extend( + wire.prices + .iter() + .filter(|flat| flat.flat.band != Band::Base) + .map(|flat| Error::new(flat.path.clone(), Repr::BandsFromTwoSources)), + ); + } + let mut ordered = based; + ordered.sort_by_key(|(_, [start, _], _)| *start); + problems.extend(ladder_problems(basis, &ordered)); + ordered + .iter() + .flat_map(|(basis, [start, _], row)| { + row.prices.iter().map(move |flat| (*basis, *start, flat)) + }) + .filter_map( + |(basis, start, flat)| match tier_price(basis, start, flat) { + Ok(sourced) => Some(sourced), + Err(problem) => { + problems.push(problem); + None + } + }, + ) + .collect() +} + +fn ladder_problems(basis: TierBasis, ordered: &[(TierBasis, [u64; 2], &TierRow)]) -> Vec { + let step = match basis { + TierBasis::PromptTokens => 0, + TierBasis::Results => 1, + }; + let reversed = ordered + .iter() + .filter(|(_, [start, end], _)| end < start) + .map(|(.., row)| Error::new(row.path.clone(), Repr::TierReversed)); + let first = ordered + .first() + .filter(|(_, [start, _], _)| *start > step) + .map(|(_, [start, _], row)| { + Error::new( + row.path.clone(), + Repr::TierGap { + start: *start, + expected: 0, + }, + ) + }); + let gaps = ordered.windows(2).filter_map(|pair| { + let (_, [_, end], _) = pair[0]; + let (_, [start, _], row) = pair[1]; + (end.checked_add(step) != Some(start)).then(|| { + Error::new( + row.path.clone(), + Repr::TierGap { + start, + expected: end, + }, + ) + }) + }); + reversed.chain(first).chain(gaps).collect() +} + +fn tier_price(basis: TierBasis, start: u64, flat: &FlatPrice) -> Result { + if flat.flat.has_suffix() { + return Err(Error::new( + flat.path.clone(), + Repr::SuffixInsideBlock { + block: "tiered_pricing", + }, + )); + } + let (meter, band) = match (basis, &flat.flat.meter) { + (TierBasis::PromptTokens, meter) => (meter.clone(), Band::over_prompt_tokens(start)), + (TierBasis::Results, Meter::Input(InputUnit::Query { over_results: None })) => ( + Meter::Input(InputUnit::Query { + over_results: NonZeroU64::new(start.saturating_sub(1)), + }), + Band::Base, + ), + (TierBasis::Results, _) => { + return Err(Error::new(flat.path.clone(), Repr::NotPricedByResults)); + } + }; + Ok(Sourced { + price: Price::online( + meter, + When { + band, + ..When::default() + }, + ), + rate: flat.price, + source: flat.key.clone(), + path: flat.path.clone(), + }) +} + +fn search_context_prices(wire: &EntryWire) -> impl Iterator + '_ { + wire.search_context.iter().flat_map(|costs| { + [ + (SearchContextSize::Low, "low", costs.search_context_size_low), + ( + SearchContextSize::Medium, + "medium", + costs.search_context_size_medium, + ), + ( + SearchContextSize::High, + "high", + costs.search_context_size_high, + ), + ] + .into_iter() + .filter_map(|(size, name, rate)| { + Some(Sourced { + price: Price::online(Meter::WebSearch(size), When::default()), + rate: rate?, + source: format!("search_context_cost_per_query.search_context_size_{name}"), + path: pointer_child( + &pointer_for_key("search_context_cost_per_query"), + &format!("search_context_size_{name}"), + ), + }) + }) + }) +} + +fn guardrail_prices(wire: &EntryWire) -> impl Iterator + '_ { + wire.guardrail.iter().flatten().map(|(unit, rate)| Sourced { + price: Price::online(Meter::Guardrail(GuardrailUnit::new(unit)), When::default()), + rate: *rate, + source: format!("guardrail_cost_per_unit.{unit}"), + path: pointer_child(&pointer_for_key("guardrail_cost_per_unit"), unit), + }) +} + +fn off_peak(block: &OffPeakBlock, problems: &mut Vec) -> Option { + let rates = off_peak_rates(&block.prices, problems); + let schedule = schedule(block, problems)?; + let off_peak = OffPeak::new(schedule, rates); + if off_peak.is_none() { + problems.push(Error::from_key("off_peak_pricing", Repr::NoOffPeakPrices)); + } + off_peak +} + +fn off_peak_rates(prices: &[FlatPrice], problems: &mut Vec) -> OffPeakRates { + let mut rates = OffPeakRates::default(); + for flat in prices { + if flat.flat.has_suffix() { + problems.push(Error::new( + flat.path.clone(), + Repr::SuffixInsideBlock { + block: "off_peak_pricing", + }, + )); + continue; + } + let slot = match flat.flat.meter { + Meter::Input(InputUnit::Token(TokenKind::Text)) => &mut rates.input, + Meter::Output(OutputUnit::Token(TokenKind::Text)) => &mut rates.output, + Meter::Output(OutputUnit::ReasoningToken) => &mut rates.reasoning, + Meter::CacheRead { + kind: TokenKind::Text, + mode: CacheReadMode::Implicit, + } => &mut rates.cache_read, + Meter::CacheRead { + kind: TokenKind::Text, + mode: CacheReadMode::Explicit, + } => &mut rates.explicit_cache_read, + Meter::CacheWrite { + kind: TokenKind::Text, + ttl: CacheTtl::Default, + } => &mut rates.cache_write, + _ => { + problems.push(Error::new(flat.path.clone(), Repr::NotAnOffPeakRate)); + continue; + } + }; + *slot = Some(flat.price); + } + rates +} + +fn schedule(block: &OffPeakBlock, problems: &mut Vec) -> Option { + let empty = |key: String| Error::from_path(&key, Repr::EmptyList); + let every_day = block + .hours_utc + .iter() + .map(|hours| ("off_peak_pricing.hours_utc".to_owned(), hours, None)); + let windows = block + .windows + .iter() + .flatten() + .enumerate() + .map(|(index, window)| { + ( + format!("off_peak_pricing.windows[{index}].hours_utc"), + &window.hours_utc, + Some((index, window.weekdays.as_ref())), + ) + }); + let windows: Vec = every_day + .chain(windows) + .map(|(key, hours, weekdays)| { + if hours.windows().is_empty() { + problems.push(empty(key)); + } + if let Some((index, Some(days))) = weekdays + && days.is_empty() + { + problems.push(empty(format!("off_peak_pricing.windows[{index}].weekdays"))); + } + Window { + hours: hours.windows().to_vec(), + weekdays: weekdays.and_then(|(_, days)| days.cloned()), + } + }) + .collect(); + if block.windows.as_ref().is_some_and(Vec::is_empty) { + problems.push(empty("off_peak_pricing.windows".to_owned())); + } + if block.hours_utc.is_none() && block.windows.is_none() { + problems.push(Error::from_key("off_peak_pricing", Repr::NoOffPeakHours)); + return None; + } + Some(Schedule { + windows, + weekday_timezone: block.weekday_timezone.clone(), + }) +} + +#[cfg(test)] +mod tests { + use litellm_cost::money::Multiplier; + use litellm_cost::price::{Channel, ServiceTier}; + use litellm_cost::schedule::{UtcWindow, Weekday}; + use rstest::rstest; + use serde_json::{Value, json}; + + use super::*; + + const TEXT_INPUT: Meter = Meter::Input(InputUnit::Token(TokenKind::Text)); + + fn wire(entry: Value) -> EntryWire { + let Value::Object(fields) = entry else { + panic!("entry must be an object") + }; + let wire = EntryWire::read(&fields); + assert!(wire.problems.is_empty(), "{:?}", wire.problems); + wire + } + + fn prices(entry: Value) -> Result> { + Prices::try_from(&wire(entry)) + } + + fn online(meter: Meter, tier: ServiceTier, band: Band) -> Price { + Price::online(meter, When { tier, band }) + } + + #[rstest] + #[case::flat_key(json!({"input_cost_per_token": 1e-6}), online(TEXT_INPUT, ServiceTier::Standard, Band::Base))] + #[case::band_and_tier( + json!({"input_cost_per_token_above_272k_tokens_priority": 1e-6}), + online(TEXT_INPUT, ServiceTier::Priority, Band::over_prompt_tokens(272_000)) + )] + #[case::batch( + json!({"input_cost_per_token_batches": 1e-6}), + Price { meter: TEXT_INPUT, channel: Channel::Batch, when: When::default() } + )] + #[case::first_token_tier_is_the_base_band( + json!({"tiered_pricing": [{"range": [0.0, 256000.0], "input_cost_per_token": 1e-6}]}), + online(TEXT_INPUT, ServiceTier::Standard, Band::Base) + )] + #[case::result_tier_stays_in_its_meter( + json!({"tiered_pricing": [ + {"max_results_range": [0, 25], "input_cost_per_query": 2e-6}, + {"max_results_range": [26, 100], "input_cost_per_query": 1e-6} + ]}), + online( + Meter::Input(InputUnit::Query { over_results: NonZeroU64::new(25) }), + ServiceTier::Standard, + Band::Base + ) + )] + #[case::search_context( + json!({"search_context_cost_per_query": {"search_context_size_low": 1e-6}}), + online(Meter::WebSearch(SearchContextSize::Low), ServiceTier::Standard, Band::Base) + )] + #[case::character_band_is_a_prompt_token_band( + json!({"input_cost_per_character_above_128k_tokens": 1e-6}), + online(Meter::Input(InputUnit::Character), ServiceTier::Standard, Band::over_prompt_tokens(128_000)) + )] + #[case::guardrail( + json!({"guardrail_cost_per_unit": {"contentPolicyUnits": 1e-6}}), + online(Meter::Guardrail(GuardrailUnit::new("contentPolicyUnits")), ServiceTier::Standard, Band::Base) + )] + fn every_spelling_lowers_to_one_price(#[case] entry: Value, #[case] expected: Price) { + assert_eq!(prices(entry).unwrap().get(&expected), UnitPrice::new(1e-6)); + } + + #[rstest] + fn two_spellings_that_agree_are_one_price() { + let prices = prices(json!({ + "input_cost_per_token": 1e-6, + "tiered_pricing": [ + {"range": [0, 10], "input_cost_per_token": 1e-6}, + {"range": [10, 20], "input_cost_per_token": 2e-6} + ] + })) + .unwrap(); + assert_eq!( + prices.get(&Price::online(TEXT_INPUT, When::default())), + UnitPrice::new(1e-6) + ); + } + + #[rstest] + fn token_tiers_lower_to_prompt_bands_at_each_start() { + let prices = prices(json!({"tiered_pricing": [ + {"range": [0, 256000], "input_cost_per_token": 1e-6}, + {"range": [256000, 1000000], "input_cost_per_token": 2e-6} + ]})) + .unwrap(); + assert_eq!( + prices.band(ServiceTier::Standard, 256_001), + Band::over_prompt_tokens(256_000) + ); + assert_eq!( + prices + .rate( + &TEXT_INPUT, + ServiceTier::Standard, + Band::over_prompt_tokens(256_000) + ) + .map(UnitPrice::get), + Some(2e-6) + ); + } + + #[rstest] + #[case::two_spellings_of_one_price( + json!({ + "input_cost_per_token": 1e-6, + "tiered_pricing": [{"range": [0, 10], "input_cost_per_token": 2e-6}] + }), + "/tiered_pricing/0/input_cost_per_token", + "duplicate_price" + )] + #[case::zero_band_is_the_base_price( + json!({"input_cost_per_token": 1e-6, "input_cost_per_token_above_0k_tokens": 2e-6}), + "/input_cost_per_token_above_0k_tokens", + "duplicate_price" + )] + #[case::tier_without_range( + json!({"tiered_pricing": [{"input_cost_per_token": 1e-6}]}), + "/tiered_pricing/0", + "tier_range" + )] + #[case::tier_with_suffix( + json!({"tiered_pricing": [{"range": [0, 10], "input_cost_per_token_priority": 1e-6}]}), + "/tiered_pricing/0/input_cost_per_token_priority", + "suffix_inside_block" + )] + #[case::reversed_range( + json!({"tiered_pricing": [{"range": [0, 10], "input_cost_per_token": 1e-6}, {"range": [10, 5], "input_cost_per_token": 1e-6}]}), + "/tiered_pricing/1", + "tier_reversed" + )] + #[case::gap_between_tiers( + json!({"tiered_pricing": [ + {"range": [0, 10], "input_cost_per_token": 1e-6}, + {"range": [20, 30], "input_cost_per_token": 2e-6} + ]}), + "/tiered_pricing/1", + "tier_gap" + )] + #[case::ladder_not_from_zero( + json!({"tiered_pricing": [{"range": [5, 10], "input_cost_per_token": 1e-6}]}), + "/tiered_pricing/0", + "tier_gap" + )] + #[case::result_ladder_overlaps( + json!({"tiered_pricing": [ + {"max_results_range": [1, 10], "input_cost_per_query": 1e-6}, + {"max_results_range": [10, 20], "input_cost_per_query": 2e-6} + ]}), + "/tiered_pricing/1", + "tier_gap" + )] + #[case::mixed_basis( + json!({"tiered_pricing": [ + {"range": [0, 10], "input_cost_per_token": 1e-6}, + {"max_results_range": [11, 20], "input_cost_per_query": 2e-6} + ]}), + "/tiered_pricing", + "mixed_tier_basis" + )] + #[case::result_tier_prices_a_token( + json!({"tiered_pricing": [{"max_results_range": [0, 10], "input_cost_per_token": 1e-6}]}), + "/tiered_pricing/0/input_cost_per_token", + "not_priced_by_results" + )] + #[case::two_band_sources( + json!({ + "output_cost_per_token_above_200k_tokens": 1e-6, + "tiered_pricing": [{"range": [0, 10], "input_cost_per_token": 1e-6}] + }), + "/output_cost_per_token_above_200k_tokens", + "bands_from_two_sources" + )] + #[case::off_peak_without_hours( + json!({"off_peak_pricing": {"input_cost_per_token": 1e-6}}), + "/off_peak_pricing", + "no_off_peak_hours" + )] + #[case::off_peak_without_prices( + json!({"off_peak_pricing": {"hours_utc": "00:00-01:00"}}), + "/off_peak_pricing", + "no_off_peak_prices" + )] + #[case::off_peak_empty_windows( + json!({"off_peak_pricing": {"windows": [], "input_cost_per_token": 1e-6}}), + "/off_peak_pricing/windows", + "empty_list" + )] + #[case::off_peak_empty_hours( + json!({"off_peak_pricing": {"hours_utc": [], "input_cost_per_token": 1e-6}}), + "/off_peak_pricing/hours_utc", + "empty_list" + )] + #[case::off_peak_empty_weekdays( + json!({"off_peak_pricing": { + "windows": [{"hours_utc": "00:00-01:00", "weekdays": []}], + "input_cost_per_token": 1e-6 + }}), + "/off_peak_pricing/windows/0/weekdays", + "empty_list" + )] + #[case::off_peak_image_price( + json!({"off_peak_pricing": {"hours_utc": "00:00-01:00", "output_cost_per_image": 1}}), + "/off_peak_pricing/output_cost_per_image", + "not_an_off_peak_rate" + )] + #[case::off_peak_band( + json!({"off_peak_pricing": {"hours_utc": "00:00-01:00", "input_cost_per_token_above_200k_tokens": 1}}), + "/off_peak_pricing/input_cost_per_token_above_200k_tokens", + "suffix_inside_block" + )] + fn rules_across_keys_name_where_each_price_came_from( + #[case] entry: Value, + #[case] expected_path: &str, + #[case] expected_variant: &str, + ) { + let problems = prices(entry).unwrap_err(); + assert!( + problems + .iter() + .any(|problem| problem.path() == expected_path + && <&str>::from(problem.kind()) == expected_variant), + "{problems:?}" + ); + } + + #[rstest] + fn off_peak_hours_and_windows_form_one_schedule_with_its_rates() { + let prices = prices(json!({"off_peak_pricing": { + "hours_utc": ["00:00-01:00"], + "windows": [{"hours_utc": "10:00-12:00", "weekdays": [6, " sun "]}], + "weekday_timezone": "Asia/Shanghai", + "output_cost_per_token": 1e-6, + "explicit_cache_read_input_token_cost": 2e-6 + }})) + .unwrap(); + let off_peak = prices.off_peak().unwrap(); + let window = |text: &str| text.parse::().unwrap(); + + assert_eq!( + off_peak.schedule(), + &Schedule { + windows: vec![ + Window { + hours: vec![window("00:00-01:00")], + weekdays: None, + }, + Window { + hours: vec![window("10:00-12:00")], + weekdays: Some(vec![Weekday::Saturday, Weekday::Sunday]), + }, + ], + weekday_timezone: Some("Asia/Shanghai".into()), + } + ); + assert_eq!( + off_peak.rates(), + &OffPeakRates { + output: UnitPrice::new(1e-6), + explicit_cache_read: UnitPrice::new(2e-6), + ..OffPeakRates::default() + } + ); + assert_eq!(prices.iter().count(), 0); + } + + #[rstest] + fn uplifts_are_adjustments_not_prices() { + let wire = wire(json!({ + "regional_endpoint_uplift_multiplier": 1.1, + "regional_processing_uplift_multiplier_eu": "1.2" + })); + let adjustments = Adjustments::from(&wire); + + assert_eq!(Prices::try_from(&wire).unwrap().iter().count(), 0); + assert_eq!( + adjustments + .get(&Uplift::NonGlobalEndpoint) + .map(Multiplier::get), + Some(1.1) + ); + assert_eq!( + adjustments + .get(&Uplift::Residency(Residency::Eu)) + .map(Multiplier::get), + Some(1.2) + ); + assert_eq!(adjustments.get(&Uplift::Residency(Residency::Us)), None); + } + + #[rstest] + fn a_model_reports_field_provider_and_price_problems_together() { + let Value::Object(fields) = json!({ + "mode": "chatt", + "tiered_pricing": [{"input_cost_per_token": 1e-6}] + }) else { + unreachable!() + }; + let problems: Vec = Model::from_fields(&fields) + .unwrap_err() + .iter() + .map(|problem| litellm_error::Report::new(problem).to_string()) + .collect(); + + assert_eq!(problems.len(), 3, "{problems:?}"); + assert!(problems[0].starts_with("mode: invalid value: unknown variant `chatt`")); + assert_eq!(problems[1], "litellm_provider: required"); + assert_eq!( + problems[2], + "tiered_pricing[0]: a tier needs exactly one of range and max_results_range" + ); + } + + #[rstest] + fn a_null_value_is_absent() { + let Value::Object(fields) = json!({ + "litellm_provider": "openai", + "input_cost_per_token": null, + "mode": null + }) else { + unreachable!() + }; + let model = Model::from_fields(&fields).unwrap(); + + assert_eq!(model.prices.iter().count(), 0); + assert_eq!(model.attributes.mode, None); + assert!(model.unrecognized.is_empty()); + } + + #[rstest] + #[case::tokens_win_over_characters( + json!({"litellm_provider": "vertex_ai-language-models", "input_cost_per_character": 1e-6, "input_cost_per_token": 1e-6}), + Billing::default() + )] + #[case::characters_only( + json!({"litellm_provider": "vertex_ai-language-models", "input_cost_per_character": 1e-6}), + Billing { text_unit: TextUnit::Character, ..Billing::default() } + )] + #[case::xai_band_is_inclusive( + json!({"litellm_provider": "xai", "input_cost_per_token": 1e-6}), + Billing { band_edge: BandEdge::Inclusive, ..Billing::default() } + )] + #[case::fireworks_without_cache_read_borrows_half( + json!({"litellm_provider": "fireworks_ai", "input_cost_per_token": 1e-6}), + Billing { cache_read_fallback: CacheReadFallback::HalfInputRate, ..Billing::default() } + )] + #[case::fireworks_with_cache_read_keeps_it( + json!({"litellm_provider": "fireworks_ai", "input_cost_per_token": 1e-6, "cache_read_input_token_cost": 1e-7}), + Billing::default() + )] + #[case::transcription_prices_the_audio_file( + json!({"litellm_provider": "openai", "mode": "audio_transcription", "input_cost_per_second": 1e-4}), + Billing { wall_clock: WallClock::AudioFile, ..Billing::default() } + )] + #[case::provisioned_chat_prices_response_time( + json!({"litellm_provider": "bedrock", "mode": "chat", "cost_per_second": 1e-2}), + Billing::default() + )] + #[case::image_generation_is_never_wall_clock( + json!({"litellm_provider": "openai", "mode": "image_generation"}), + Billing { wall_clock: WallClock::Never, ..Billing::default() } + )] + #[case::gemini_per_query_search( + json!({"litellm_provider": "gemini", "web_search_billing_unit": "per_query"}), + Billing { web_search: WebSearchUnit::PerQuery, ..Billing::default() } + )] + #[case::dashscope_tiers_bill_the_whole_call( + json!({"litellm_provider": "dashscope", "tiered_pricing": [{"range": [0, 32000], "input_cost_per_token": 1e-6}]}), + Billing { tier_scope: TierScope::WholeCall, ..Billing::default() } + )] + fn billing_is_lowered_from_the_row_and_its_provider( + #[case] entry: Value, + #[case] expected: Billing, + ) { + assert_eq!(prices(entry).unwrap().billing(), expected); + } +} diff --git a/litellm-rust/crates/model-catalog/src/model.rs b/litellm-rust/crates/model-catalog/src/model.rs new file mode 100644 index 00000000000..c1acd1c14cb --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/model.rs @@ -0,0 +1,22 @@ +use std::collections::BTreeMap; + +use litellm_cost::price::{Adjustments, Prices}; + +use crate::error::Error; +use crate::wire::{Attributes, EntryWire, Fields}; + +/// One catalog entry read into types. +#[derive(Clone, Debug, PartialEq)] +pub struct Model { + pub attributes: Attributes, + pub prices: Prices, + pub adjustments: Adjustments, + pub capabilities: BTreeMap, + pub unrecognized: Vec, +} + +impl Model { + pub fn from_fields(fields: &Fields) -> Result> { + Self::try_from(EntryWire::read(fields)) + } +} diff --git a/litellm-rust/crates/model-catalog/src/model_info.rs b/litellm-rust/crates/model-catalog/src/model_info.rs index ebee2c64140..8f7549c1c6d 100644 --- a/litellm-rust/crates/model-catalog/src/model_info.rs +++ b/litellm-rust/crates/model-catalog/src/model_info.rs @@ -1,582 +1,194 @@ -use crate::capabilities::{AudioFormat, InputModality, Mode, OutputModality, VertexAiAudioApi}; -use crate::pricing::{OffPeakPricing, SearchContextCostPerQuery, TieredRate, WebSearchBillingUnit}; -use litellm_llms_types::formats::chat_completions::ReasoningEffort; -use serde::{Deserialize, Serialize}; +use indexmap::IndexMap; use serde_json::Value; -use std::collections::BTreeMap; -/// Typed mirror of one catalog model entry. -#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -pub struct ModelInfo { - #[serde(skip_serializing_if = "Option::is_none")] - pub annotation_cost_per_page: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub annotation_cost_per_page_batches: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub audio_transcription_config: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub bedrock_converse_supports_strict_tools: Option, - /// Highest reasoning effort the Bedrock output_config accepts for this model. - #[serde(skip_serializing_if = "Option::is_none")] - pub bedrock_output_config_effort_ceiling: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_creation_input_audio_token_cost: Option, - /// USD per token written to the provider's prompt cache. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_creation_input_token_cost: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_creation_input_token_cost_above_32k_tokens: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_creation_input_token_cost_above_100k_tokens: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_creation_input_token_cost_above_100k_tokens_batches: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_creation_input_token_cost_above_128k_tokens: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_creation_input_token_cost_above_1hr: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_creation_input_token_cost_above_1hr_above_100k_tokens: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_creation_input_token_cost_above_1hr_above_200k_tokens: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_creation_input_token_cost_above_200k_tokens: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_creation_input_token_cost_above_200k_tokens_batches: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_creation_input_token_cost_above_256k_tokens: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_creation_input_token_cost_above_272k_tokens: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_creation_input_token_cost_above_272k_tokens_batches: Option, - /// Flex service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_creation_input_token_cost_above_272k_tokens_flex: Option, - /// Priority service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_creation_input_token_cost_above_272k_tokens_priority: Option, - /// Ultrafast service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_creation_input_token_cost_above_272k_tokens_ultrafast: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_creation_input_token_cost_batches: Option, - /// Flex service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_creation_input_token_cost_flex: Option, - /// Priority service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_creation_input_token_cost_priority: Option, - /// Ultrafast service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_creation_input_token_cost_ultrafast: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_read_input_audio_token_cost: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_read_input_image_token_cost: Option, - /// USD per prompt token served from the provider's prompt cache. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_read_input_token_cost: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_read_input_token_cost_above_32k_tokens: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_read_input_token_cost_above_100k_tokens: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_read_input_token_cost_above_100k_tokens_batches: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_read_input_token_cost_above_128k_tokens: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_read_input_token_cost_above_200k_tokens: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_read_input_token_cost_above_200k_tokens_batches: Option, - /// Priority service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_read_input_token_cost_above_200k_tokens_priority: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_read_input_token_cost_above_256k_tokens: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_read_input_token_cost_above_272k_tokens: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_read_input_token_cost_above_272k_tokens_batches: Option, - /// Flex service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_read_input_token_cost_above_272k_tokens_flex: Option, - /// Priority service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_read_input_token_cost_above_272k_tokens_priority: Option, - /// Ultrafast service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_read_input_token_cost_above_272k_tokens_ultrafast: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_read_input_token_cost_above_512k_tokens: Option, - /// Balanced service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_read_input_token_cost_balanced: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_read_input_token_cost_batches: Option, - /// Flex service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_read_input_token_cost_flex: Option, - /// Priority service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_read_input_token_cost_priority: Option, - /// Ultrafast service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_read_input_token_cost_ultrafast: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub citation_cost_per_token: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub code_interpreter_cost_per_session: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub comment: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub computer_use_input_cost_per_1k_tokens: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub computer_use_output_cost_per_1k_tokens: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub cost_per_second: Option, - /// Reasoning effort the provider applies when the request omits reasoning_effort. Gates whether a non-default temperature or the top_p/logprobs sampling params are accepted, which hold only when the effort resolves to 'none'. - #[serde(skip_serializing_if = "Option::is_none")] - pub default_reasoning_effort: Option, - /// Date the provider deprecates the model, YYYY-MM-DD. - #[serde(skip_serializing_if = "Option::is_none")] - pub deprecation_date: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub file_search_cost_per_1k_calls: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub file_search_cost_per_gb_per_day: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub gemini_audio_only_live: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub gemini_native_audio: Option, - /// USD per Grounding with Google Maps request; billed per query or per prompt per web_search_billing_unit. - #[serde(skip_serializing_if = "Option::is_none")] - pub google_maps_grounding_cost_per_query: Option, - /// USD cost per billable guardrail unit, keyed by the provider's usage counter name (e.g. Bedrock's contentPolicyUnits). - #[serde(skip_serializing_if = "Option::is_none")] - pub guardrail_cost_per_unit: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_audio_per_second: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_audio_per_second_above_128k_tokens: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_audio_token: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_audio_token_batches: Option, - /// Priority service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_audio_token_priority: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_character: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_character_above_128k_tokens: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_image: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_image_above_128k_tokens: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_image_token: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_image_token_batches: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_pixel: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_query: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_request: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_second: Option, - /// USD per prompt token. - #[serde(skip_serializing_if = "Option::is_none")] +/// The fields main's `ModelInfoBase(...)` copies (`litellm/utils.py` at c7c895db45), one typed field each, +/// read from the raw row without lowering it, so building one never fails +#[derive(Clone, Debug, Default, PartialEq)] +pub struct ModelInfoFields { + pub key: String, + pub max_tokens: Option, + pub max_input_tokens: Option, + pub max_output_tokens: Option, pub input_cost_per_token: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_token_above_32k_tokens: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_token_above_100k_tokens: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_token_above_100k_tokens_batches: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_token_above_128k_tokens: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_token_above_200k_tokens: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_token_above_200k_tokens_batches: Option, - /// Priority service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_token_above_200k_tokens_priority: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_token_above_256k_tokens: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_token_above_272k_tokens: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_token_above_272k_tokens_batches: Option, - /// Flex service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_token_above_272k_tokens_flex: Option, - /// Priority service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_token_above_272k_tokens_priority: Option, - /// Ultrafast service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_token_above_272k_tokens_ultrafast: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_token_above_512k_tokens: Option, - /// Balanced service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_token_balanced: Option, - /// USD per prompt token via the provider's batch API. - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_token_batches: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_token_cache_hit: Option, - /// Flex service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_token_flex: Option, - /// Priority service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_token_priority: Option, - /// Ultrafast service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] + pub input_cost_per_token_balanced: Option, pub input_cost_per_token_ultrafast: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_video_per_second: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_video_per_second_above_128k_tokens: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_video_per_second_above_15s_interval: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_video_per_second_above_8s_interval: Option, - #[serde(skip_serializing_if = "Option::is_none")] + pub cache_creation_input_token_cost: Option, + pub cache_creation_input_token_cost_above_200k_tokens: Option, + pub cache_creation_input_token_cost_above_272k_tokens: Option, + pub cache_creation_input_token_cost_above_272k_tokens_priority: Option, + pub cache_creation_input_token_cost_above_272k_tokens_flex: Option, + pub cache_creation_input_token_cost_above_272k_tokens_ultrafast: Option, + pub cache_creation_input_token_cost_flex: Option, + pub cache_creation_input_token_cost_priority: Option, + pub cache_creation_input_token_cost_ultrafast: Option, + pub cache_read_input_token_cost: Option, + pub cache_read_input_audio_token_cost: Option, + pub cache_read_input_audio_token_cost_flex: Option, + pub cache_read_input_audio_token_cost_priority: Option, + pub cache_read_input_audio_token_cost_batches: Option, + pub explicit_cache_read_input_token_cost: Option, + pub prompt_cache_min_tokens: Option, + pub cache_read_input_token_cost_above_200k_tokens: Option, + pub cache_read_input_token_cost_above_200k_tokens_priority: Option, + pub cache_read_input_token_cost_above_272k_tokens: Option, + pub cache_read_input_token_cost_above_272k_tokens_priority: Option, + pub cache_read_input_token_cost_above_272k_tokens_flex: Option, + pub cache_read_input_token_cost_above_272k_tokens_ultrafast: Option, + pub cache_read_input_token_cost_above_512k_tokens: Option, + pub cache_read_input_token_cost_flex: Option, + pub cache_read_input_token_cost_priority: Option, + pub cache_read_input_token_cost_balanced: Option, + pub cache_read_input_token_cost_ultrafast: Option, + pub cache_read_input_token_cost_batches: Option, + pub cache_read_input_token_cost_above_200k_tokens_batches: Option, + pub cache_read_input_token_cost_above_272k_tokens_batches: Option, + pub cache_creation_input_token_cost_batches: Option, + pub cache_creation_input_token_cost_above_200k_tokens_batches: Option, + pub cache_creation_input_token_cost_above_272k_tokens_batches: Option, + pub cache_creation_input_token_cost_above_1hr: Option, + pub cache_creation_input_token_cost_above_12hr: Option, + pub cache_creation_input_token_cost_above_12hr_flex: Option, + pub cache_creation_input_token_cost_above_12hr_priority: Option, + pub cache_creation_input_token_cost_above_12hr_above_272k_tokens: Option, + pub cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex: Option, + pub cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority: Option, + pub cache_creation_input_token_cost_above_12hr_batches: Option, + pub cache_creation_input_token_cost_above_12hr_ultrafast: Option, + pub cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches: Option, + pub cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast: Option, + pub off_peak_pricing: Option, + pub input_cost_per_character: Option, + pub input_cost_per_token_above_128k_tokens: Option, + pub input_cost_per_token_above_200k_tokens: Option, + pub input_cost_per_token_above_200k_tokens_priority: Option, + pub input_cost_per_token_above_272k_tokens: Option, + pub input_cost_per_token_above_272k_tokens_priority: Option, + pub input_cost_per_token_above_272k_tokens_flex: Option, + pub input_cost_per_token_above_272k_tokens_ultrafast: Option, + pub input_cost_per_token_above_512k_tokens: Option, + pub input_cost_per_query: Option, + pub cost_per_second: Option, + pub input_cost_per_second: Option, + pub input_cost_per_text_message: Option, + pub input_cost_per_audio_token: Option, + pub input_cost_per_audio_token_flex: Option, + pub input_cost_per_audio_token_priority: Option, + pub input_cost_per_image_token: Option, pub input_cost_per_video_token: Option, - #[serde(skip_serializing_if = "Option::is_none")] + pub input_cost_per_audio_token_batches: Option, + pub input_cost_per_image_token_batches: Option, + pub input_cost_per_image: Option, + pub input_cost_per_audio_per_second: Option, + pub input_cost_per_video_per_second: Option, + pub input_cost_per_token_batches: Option, pub input_cost_per_video_token_batches: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub input_dbu_cost_per_token: Option, - /// LiteLLM provider slug; one of https://docs.litellm.ai/docs/providers. - #[serde(skip_serializing_if = "Option::is_none")] - pub litellm_provider: Option, - /// Maximum prompt/context tokens the model accepts. - #[serde(skip_serializing_if = "Option::is_none")] - pub max_input_tokens: Option, - /// Maximum tokens the model can generate in one response. - #[serde(skip_serializing_if = "Option::is_none")] - pub max_output_tokens: Option, - /// Legacy field: max output tokens if the provider specifies it, else max input tokens. - #[serde(skip_serializing_if = "Option::is_none")] - pub max_tokens: Option, - /// Free-form notes about the entry (e.g. pricing derivation). - #[serde(skip_serializing_if = "Option::is_none")] - pub metadata: Option>, - /// Primary API surface / task type of the model. - #[serde(skip_serializing_if = "Option::is_none")] - pub mode: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub ocr_cost_per_credit: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub ocr_cost_per_page: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub ocr_cost_per_page_batches: Option, - /// Rates that replace the same-named base fields while the request falls inside the stated UTC windows. - #[serde(skip_serializing_if = "Option::is_none")] - pub off_peak_pricing: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_audio_token: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_character: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_character_above_128k_tokens: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_image: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_image_1024: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_image_1536: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_image_512: Option, - #[serde( - rename = "output_cost_per_image_0.5K", - skip_serializing_if = "Option::is_none" - )] - pub output_cost_per_image_0_5k: Option, - #[serde( - rename = "output_cost_per_image_1K", - skip_serializing_if = "Option::is_none" - )] - pub output_cost_per_image_1k: Option, - #[serde( - rename = "output_cost_per_image_2K", - skip_serializing_if = "Option::is_none" - )] - pub output_cost_per_image_2k: Option, - #[serde( - rename = "output_cost_per_image_4K", - skip_serializing_if = "Option::is_none" - )] - pub output_cost_per_image_4k: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_image_token: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_pixel: Option, - /// USD per reasoning/thinking token, when billed separately. - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_reasoning_token: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_second: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_second_1080p: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_second_2k: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_second_480p: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_second_4k: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_second_720p: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_second_768p: Option, - /// USD per generated token. - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_token: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_token_above_32k_tokens: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_token_above_100k_tokens: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_token_above_100k_tokens_batches: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_token_above_128k_tokens: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_token_above_200k_tokens: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_token_above_200k_tokens_batches: Option, - /// Priority service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_token_above_200k_tokens_priority: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_token_above_256k_tokens: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_token_above_272k_tokens: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_token_above_272k_tokens_batches: Option, - /// Flex service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_token_above_272k_tokens_flex: Option, - /// Priority service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_token_above_272k_tokens_priority: Option, - /// Ultrafast service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_token_above_272k_tokens_ultrafast: Option, - /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_token_above_512k_tokens: Option, - /// Balanced service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_token_balanced: Option, - /// USD per generated token via the provider's batch API. - #[serde(skip_serializing_if = "Option::is_none")] + pub input_cost_per_token_above_200k_tokens_batches: Option, + pub input_cost_per_token_above_272k_tokens_batches: Option, pub output_cost_per_token_batches: Option, - /// Flex service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] + pub output_cost_per_token_above_200k_tokens_batches: Option, + pub output_cost_per_token_above_272k_tokens_batches: Option, + pub output_cost_per_token: Option, pub output_cost_per_token_flex: Option, - /// Priority service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_token_priority: Option, - /// Ultrafast service-tier rate for the same-named base field. - #[serde(skip_serializing_if = "Option::is_none")] + pub output_cost_per_token_balanced: Option, pub output_cost_per_token_ultrafast: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_video_per_second: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_video_token: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub output_dbu_cost_per_token: Option, - /// Embedding dimension for embedding models. - #[serde(skip_serializing_if = "Option::is_none")] - pub output_vector_size: Option, - /// Smallest prefix the provider will actually cache; absent means the provider default applies. - #[serde(skip_serializing_if = "Option::is_none")] - pub prompt_cache_min_tokens: Option, - /// Provider-internal routing hints (e.g. bedrock_invocation_schema). - #[serde(skip_serializing_if = "Option::is_none")] - pub provider_specific_entry: Option>, - /// Exact reasoning_effort levels this deployment accepts; wins over supports_* flags. - #[serde(skip_serializing_if = "Option::is_none")] - pub reasoning_effort_levels: Option>, - /// Multiplier applied to all token costs when served from a non-global Vertex AI endpoint (e.g. 1.10 = +10%). - #[serde(skip_serializing_if = "Option::is_none")] - pub regional_endpoint_uplift_multiplier: Option, - /// Multiplier applied to all token costs for EU data residency (e.g. 1.10 = +10%). - #[serde(skip_serializing_if = "Option::is_none")] pub regional_processing_uplift_multiplier_eu: Option, - /// Multiplier applied to all token costs for US data residency (e.g. 1.10 = +10%). - #[serde(skip_serializing_if = "Option::is_none")] pub regional_processing_uplift_multiplier_us: Option, - /// Provider default requests-per-minute limit. - #[serde(skip_serializing_if = "Option::is_none")] - pub rpm: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub rules: Option>, - /// USD cost per web search query, keyed by search context size. - #[serde(skip_serializing_if = "Option::is_none")] - pub search_context_cost_per_query: Option, - /// URL of the provider pricing/model page this entry was taken from. - #[serde(skip_serializing_if = "Option::is_none")] - pub source: Option, - /// Audio container formats the model can return. - #[serde(skip_serializing_if = "Option::is_none")] - pub supported_audio_formats: Option>, - /// OpenAI-style API routes this model can be called through, e.g. /v1/chat/completions. - #[serde(skip_serializing_if = "Option::is_none")] + pub regional_endpoint_uplift_multiplier: Option, + pub output_cost_per_audio_token: Option, + pub output_cost_per_character: Option, + pub output_cost_per_reasoning_token: Option, + pub output_cost_per_reasoning_token_flex: Option, + pub output_cost_per_reasoning_token_priority: Option, + pub output_cost_per_token_above_128k_tokens: Option, + pub output_cost_per_character_above_128k_tokens: Option, + pub output_cost_per_token_above_200k_tokens: Option, + pub output_cost_per_token_above_200k_tokens_priority: Option, + pub output_cost_per_token_above_272k_tokens: Option, + pub output_cost_per_token_above_272k_tokens_priority: Option, + pub output_cost_per_token_above_272k_tokens_flex: Option, + pub output_cost_per_token_above_272k_tokens_ultrafast: Option, + pub output_cost_per_token_above_512k_tokens: Option, + pub output_cost_per_second: Option, + pub output_cost_per_second_1080p: Option, + pub output_cost_per_second_480p: Option, + pub output_cost_per_second_720p: Option, + pub output_cost_per_second_768p: Option, + pub output_cost_per_second_2k: Option, + pub output_cost_per_second_4k: Option, + pub output_cost_per_video_per_second: Option, + pub output_cost_per_image: Option, + pub output_cost_per_pixel: Option, + pub output_cost_per_image_token: Option, + pub output_cost_per_image_token_batches: Option, + pub output_cost_per_video_token: Option, + pub output_vector_size: Option, + pub citation_cost_per_token: Option, + pub tiered_pricing: Option, + pub litellm_provider: Option, + pub mode: Option, pub supported_endpoints: Option>, - /// Input modalities the model accepts. - #[serde(skip_serializing_if = "Option::is_none")] - pub supported_modalities: Option>, - /// Output modalities the model can produce. - #[serde(skip_serializing_if = "Option::is_none")] - pub supported_output_modalities: Option>, - /// Cloud regions the model is available in ('global' or region ids). - #[serde(skip_serializing_if = "Option::is_none")] - pub supported_regions: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_adaptive_thinking: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_anthropic_compaction: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_anthropic_thinking_payload: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_assistant_prefill: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_audio_input: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_audio_output: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_bedrock_runtime_chat_completions_response_format: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_bedrock_runtime_chat_completions_tools_with_reasoning: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_computer_use: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_embedding_image_input: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_fast_mode: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_forced_tool_use: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_function_calling: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_image_input: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_image_size: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_legacy_thinking: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_low_reasoning_effort: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_max_reasoning_effort: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_mid_conversation_system: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_minimal_reasoning_effort: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_multimodal: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_native_streaming: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_native_structured_output: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_none_reasoning_effort: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_nova_canvas_image_edit: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_output_config: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_parallel_function_calling: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_parallel_tool_use_config: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_pdf_input: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_prompt_cache_breakpoint: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_prompt_caching: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_reasoning: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_response_schema: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_sampling_params: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_speed: Option, - #[serde(skip_serializing_if = "Option::is_none")] pub supports_system_messages: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_thinking_cache_preservation: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_tool_choice: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_tool_search: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_url_context: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_video_input: Option, - #[serde(skip_serializing_if = "Option::is_none")] + pub supports_response_schema: Option, pub supports_vision: Option, - #[serde(skip_serializing_if = "Option::is_none")] + pub supports_function_calling: Option, + pub supports_parallel_function_calling: Option, + pub supports_tool_choice: Option, + pub supports_assistant_prefill: Option, + pub supports_prompt_caching: Option, + pub supports_prompt_cache_breakpoint: Option, + pub supports_thinking_cache_preservation: Option, + pub supports_audio_input: Option, + pub supports_video_input: Option, + pub supports_box_questions: Option, + pub supports_audio_output: Option, + pub supports_pdf_input: Option, + pub supports_embedding_image_input: Option, + pub supports_native_streaming: Option, + pub supports_native_structured_output: Option, + pub supports_bedrock_native_structured_output: Option, pub supports_web_search: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub supports_xhigh_reasoning_effort: Option, - #[serde(skip_serializing_if = "Option::is_none")] + pub supports_url_context: Option, + pub supports_reasoning: Option, + pub supports_adaptive_thinking: Option, + pub supports_legacy_thinking: Option, pub thinking_always_on: Option, - /// Context-length or result-count tiered rates; each tier's costs apply within its range. - #[serde(skip_serializing_if = "Option::is_none")] - pub tiered_pricing: Option>, - /// Provider default tokens-per-minute limit. - #[serde(skip_serializing_if = "Option::is_none")] - pub tpm: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub use_openai_responses_path: Option, - #[serde(skip_serializing_if = "Option::is_none")] + pub supports_tool_search: Option, + pub supports_mid_conversation_system: Option, + pub supports_anthropic_thinking_payload: Option, + pub supports_anthropic_compaction: Option, + pub supports_none_reasoning_effort: Option, + pub supports_minimal_reasoning_effort: Option, + pub supports_low_reasoning_effort: Option, + pub supports_xhigh_reasoning_effort: Option, + pub supports_max_reasoning_effort: Option, + pub reasoning_effort_levels: Option, + pub default_reasoning_effort: Option, + pub bedrock_output_config_effort_ceiling: Option, + pub bedrock_converse_supports_strict_tools: Option, + pub supports_regex_lookaround: Option, + pub supports_computer_use: Option, + pub search_context_cost_per_query: Option, + pub web_search_billing_unit: Option, + pub search_free_requests: Option, + pub google_maps_grounding_cost_per_query: Option, + pub tpm: Option, + pub rpm: Option, + pub ocr_cost_per_page: Option, + pub ocr_cost_per_page_batches: Option, + pub ocr_cost_per_credit: Option, + pub annotation_cost_per_page: Option, + pub annotation_cost_per_page_batches: Option, + pub provider_specific_entry: Option, pub uses_embed_content: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub vector_store_cost_per_gb_per_day: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub vertex_ai_audio_api: Option, - /// Whether web search is billed per query or per prompt. - #[serde(skip_serializing_if = "Option::is_none")] - pub web_search_billing_unit: Option, + pub supports_image_size: Option, + pub supported_audio_formats: Option, + pub token_count_framing: Option, + pub token_count_model: Option, + pub token_count_tokenizer: Option, + pub vertex_ai_audio_api: Option, + /// Every key matching `ABOVE_THRESHOLD_COST_KEY_PATTERN`, copied as written + pub above_threshold_costs: IndexMap, } diff --git a/litellm-rust/crates/model-catalog/src/pricing.rs b/litellm-rust/crates/model-catalog/src/pricing.rs deleted file mode 100644 index b8ed652451c..00000000000 --- a/litellm-rust/crates/model-catalog/src/pricing.rs +++ /dev/null @@ -1,97 +0,0 @@ -use serde::{Deserialize, Serialize}; - -/// Whether web search is billed per query or per prompt. -#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(rename_all = "snake_case")] -pub enum WebSearchBillingUnit { - PerQuery, - PerPrompt, -} - -/// UTC "HH:MM-HH:MM" window, or a list of them; a window may wrap past midnight. -#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(untagged)] -pub enum UtcHours { - Single(String), - Multiple(Vec), -} - -/// ISO-8601 weekday number (1 = Monday .. 7 = Sunday) or English day name. -#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(untagged)] -pub enum Weekday { - Number(u8), - Name(String), -} - -/// One off-peak window entry inside `windows`. -#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(deny_unknown_fields)] -pub struct OffPeakWindow { - pub hours_utc: UtcHours, - #[serde(skip_serializing_if = "Option::is_none")] - pub weekdays: Option>, -} - -/// Rates that replace the same-named base fields inside the stated UTC windows. -#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(deny_unknown_fields)] -pub struct OffPeakPricing { - #[serde(skip_serializing_if = "Option::is_none")] - pub hours_utc: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub windows: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub weekday_timezone: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_token: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_token: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_reasoning_token: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_read_input_token_cost: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_creation_input_token_cost: Option, -} - -/// USD cost per web search query, keyed by search context size. -#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(deny_unknown_fields)] -pub struct SearchContextCostPerQuery { - #[serde(skip_serializing_if = "Option::is_none")] - pub search_context_size_low: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub search_context_size_medium: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub search_context_size_high: Option, -} - -/// One tier of a context-length or result-count tiered rate. -#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(deny_unknown_fields)] -pub struct TieredRate { - #[serde(skip_serializing_if = "Option::is_none")] - pub range: Option<[f64; 2]>, - #[serde(skip_serializing_if = "Option::is_none")] - pub max_results_range: Option<[f64; 2]>, - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_token: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_token: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub output_cost_per_reasoning_token: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_read_input_token_cost: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_creation_input_token_cost: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_query: Option, -} diff --git a/litellm-rust/crates/model-catalog/src/register.rs b/litellm-rust/crates/model-catalog/src/register.rs new file mode 100644 index 00000000000..ed34082ae57 --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/register.rs @@ -0,0 +1,26 @@ +use std::collections::{BTreeSet, HashMap}; + +use indexmap::IndexMap; + +use crate::wire::Fields; + +/// Where one `register` call writes and which row it inherits from +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RegistrationTarget { + /// Always the caller's key, never a case variant of it + pub write: String, + /// The first candidate that is an exact key whose row passes `provider_matches` + pub inherit_from: Option, +} + +/// Python's inheritance candidates in main's order, `get_model_info`'s then `_resolve_builtin_model_cost_entry`'s +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RegistrationCandidates(pub Vec); + +/// Layer 3, and the fields `register` last wrote, which litellm owns again +#[derive(Clone, Debug, Default)] +#[expect(dead_code, reason = "07 adds the functions that read registrations")] +pub struct Registrations { + persisted: IndexMap, + registered_fields: HashMap>, +} diff --git a/litellm-rust/crates/model-catalog/src/resolve/image.rs b/litellm-rust/crates/model-catalog/src/resolve/image.rs new file mode 100644 index 00000000000..d24e170b32c --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/resolve/image.rs @@ -0,0 +1,18 @@ +/// What an image call asks for. `resolve_image` builds main's eight key spellings from it +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ImageCandidates { + pub model: String, + pub provider: Option, + /// Width and height, `1024 x 1024` when the request names none, as on main + pub size: Option<(u32, u32)>, + pub quality: Option, +} + +/// A size and quality keyed row, `{quality}/{W}-x-{H}/{model}`, parsed once per index build for fal's nearest match +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ImageKey { + pub quality: Option, + pub width: u32, + pub height: u32, + pub model: String, +} diff --git a/litellm-rust/crates/model-catalog/src/resolve/mod.rs b/litellm-rust/crates/model-catalog/src/resolve/mod.rs new file mode 100644 index 00000000000..fbaead622aa --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/resolve/mod.rs @@ -0,0 +1,46 @@ +mod image; + +pub use image::*; + +/// The seven names main's `get_model_info` tries, built by Python, walked in this order by Rust +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct Candidates { + pub combined_model_name: Option, + pub model: String, + pub region_free_combined_model_name: Option, + pub split_model: Option, + pub combined_stripped_model_name: Option, + pub stripped_model_name: Option, + pub provider_prefixed_model_name: Option, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum CandidateSlot { + CombinedModelName, + Model, + RegionFreeCombinedModelName, + SplitModel, + CombinedStrippedModelName, + StrippedModelName, + ProviderPrefixedModelName, +} + +/// Which step matched, so a test can name the step that changed +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Via { + pub slot: CandidateSlot, + /// Matched through the case-insensitive index rather than an exact key + pub case_insensitive: bool, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Resolution { + Found { + key: String, + via: Via, + }, + /// Every candidate tried, in order + Missing { + tried: Vec, + }, +} diff --git a/litellm-rust/crates/model-catalog/src/schema.rs b/litellm-rust/crates/model-catalog/src/schema.rs index 7de988b3ff3..66f9933a3e3 100644 --- a/litellm-rust/crates/model-catalog/src/schema.rs +++ b/litellm-rust/crates/model-catalog/src/schema.rs @@ -1,9 +1,80 @@ -use schemars::Schema; +use std::collections::BTreeMap; + +use litellm_cost::money::{Multiplier, UnitPrice}; +use schemars::{JsonSchema, Schema}; use serde_json::{Map, Value, json}; +use crate::fallback::FallbackGeneralizations; +use crate::wire::{Attributes, OffPeakWindow, SearchContextCostPerQuery, UtcHours}; + +#[derive(JsonSchema)] +#[schemars(rename = "ModelEntry")] +struct ModelEntry { + #[serde(flatten)] + _attributes: Attributes, + #[serde(flatten)] + _blocks: PriceBlocks, +} + +#[derive(JsonSchema)] +#[allow(dead_code)] +struct PriceBlocks { + /// USD cost per billable guardrail unit, keyed by the provider's usage counter name (e.g. Bedrock's contentPolicyUnits). + guardrail_cost_per_unit: Option>, + /// Rates that replace the same-named base fields while the request falls inside the stated UTC windows. + off_peak_pricing: Option, + /// Multiplier applied to all token costs when served from a non-global Vertex AI endpoint (e.g. 1.10 = +10%). + regional_endpoint_uplift_multiplier: Option, + /// Multiplier applied to all token costs for EU data residency (e.g. 1.10 = +10%). + regional_processing_uplift_multiplier_eu: Option, + /// Multiplier applied to all token costs for US data residency (e.g. 1.10 = +10%). + regional_processing_uplift_multiplier_us: Option, + /// USD cost per web search query, keyed by search context size. + search_context_cost_per_query: Option, + /// Context-length or result-count tiered rates; each tier's costs apply within its range. + tiered_pricing: Option>, +} + +/// Rates that replace the same-named base fields inside the stated UTC windows. +#[derive(JsonSchema)] +#[allow(dead_code)] +struct OffPeakPricing { + hours_utc: Option, + windows: Option>, + weekday_timezone: Option, +} + +/// One tier of a context-length or result-count tiered rate. +#[derive(JsonSchema)] +#[allow(dead_code)] +struct TieredRate { + range: Option<[u64; 2]>, + max_results_range: Option<[u64; 2]>, +} + +#[derive(JsonSchema)] +#[allow(dead_code)] +struct RegistryDocument { + sample_spec: Option>, + fallback_generalizations: Option, +} + +pub(crate) const FLAT_KEY_PATTERN: &str = concat!( + r"^(", + r"input_cost_per_((audio_|image_|video_)?token|character|image|pixel|query|request|second|text_message|audio_per_second|video_per_second(_above_\d+s_interval)?)", + r"|output_cost_per_((audio_|image_|video_|reasoning_)?token|character|image(_\d+(\.\d+)?K?)?|pixel|second(_\d+[pk])?|video_per_second)", + r"|cache_read_input_(audio_|image_|video_)?token_cost", + r"|explicit_cache_read_input_token_cost", + r"|cache_creation_input_(audio_|image_|video_)?token_cost(_above_1hr|_above_12hr)?", + r"|annotation_cost_per_page|citation_cost_per_token|code_interpreter_cost_per_session", + r"|google_maps_grounding_cost_per_query|ocr_cost_per_credit|ocr_cost_per_page|cost_per_second", + r"|browser_open_cost_per_request|cost_per_additional_result", + r")(_above_\d+k?_tokens)?(_(priority|flex|ultrafast|balanced|batches|decisions))?$", +); + /// JSON Schema for one model entry, including registry validation constraints. pub fn model_entry_json_schema() -> Schema { - let mut schema = serde_json::to_value(schemars::schema_for!(crate::ModelInfo)) + let mut schema = serde_json::to_value(schemars::schema_for!(ModelEntry)) .expect("derived model schema serializes"); remove_nullable_optional_fields(&mut schema); decorate_model_entry(&mut schema); @@ -22,24 +93,21 @@ pub fn registry_json_schema() -> Schema { entry.as_object_mut().unwrap().remove("$schema"); definitions.insert("modelEntry".into(), entry); - let mut fallback = serde_json::to_value(schemars::schema_for!(crate::FallbackGeneralizations)) - .expect("derived fallback schema serializes"); - remove_nullable_optional_fields(&mut fallback); - definitions.extend(take_definitions(&mut fallback)); - fallback.as_object_mut().unwrap().remove("$schema"); - - let root = json!({ - "$schema": "https://json-schema.org/draft/2020-12/schema", - "title": "LiteLLM model prices and context window registry", - "type": "object", - "properties": { - "sample_spec": {"type": "object"}, - "fallback_generalizations": fallback - }, - "additionalProperties": {"$ref": "#/$defs/modelEntry"}, - "$defs": definitions - }); - Schema::from(root.as_object().unwrap().clone()) + let mut registry = serde_json::to_value(schemars::schema_for!(RegistryDocument)) + .expect("derived registry schema serializes"); + remove_nullable_optional_fields(&mut registry); + definitions.extend(take_definitions(&mut registry)); + let root = registry.as_object_mut().unwrap(); + root.insert( + "title".into(), + "LiteLLM model prices and context window registry".into(), + ); + root.insert( + "additionalProperties".into(), + json!({"$ref": "#/$defs/modelEntry"}), + ); + root.insert("$defs".into(), Value::Object(definitions)); + Schema::from(root.clone()) } fn take_definitions(schema: &mut Value) -> Map { @@ -80,42 +148,21 @@ fn remove_nullable_optional_fields(value: &mut Value) { } fn decorate_model_entry(schema: &mut Value) { + let rates = json!({FLAT_KEY_PATTERN: {"type": "number", "minimum": 0}}); let object = schema.as_object_mut().unwrap(); object.insert("required".into(), json!(["litellm_provider"])); object.insert("additionalProperties".into(), Value::Bool(true)); - let properties = object - .get_mut("properties") - .unwrap() - .as_object_mut() - .unwrap(); - properties.insert( - "aliases".into(), - json!({"type": "array", "items": {"type": "string"}}), + object.insert( + "patternProperties".into(), + json!({ + FLAT_KEY_PATTERN: {"type": "number", "minimum": 0}, + "^supports_": {"type": "boolean"}, + }), ); - properties.get_mut("deprecation_date").unwrap()["format"] = json!("date"); - properties.get_mut("deprecation_date").unwrap()["pattern"] = - json!(r"^\d{4}-(0[1-9]|1[0-2])-(0[1-9]|[12]\d|3[01])$"); - - properties.iter_mut().for_each(|(name, property)| { - if name.contains("cost") { - property["minimum"] = json!(0); - } else if name.contains("uplift_multiplier") { - property["minimum"] = json!(1); - } - }); - properties.get_mut("guardrail_cost_per_unit").unwrap()["additionalProperties"]["minimum"] = - json!(0); - let definitions = object.get_mut("$defs").unwrap().as_object_mut().unwrap(); - for definition in ["OffPeakPricing", "TieredRate", "SearchContextCostPerQuery"] { - let properties = definitions[definition]["properties"] - .as_object_mut() - .unwrap(); - properties.iter_mut().for_each(|(name, property)| { - if name.contains("cost") || definition == "SearchContextCostPerQuery" { - property["minimum"] = json!(0); - } - }); + for definition in ["OffPeakPricing", "TieredRate"] { + definitions[definition]["patternProperties"] = rates.clone(); + definitions[definition]["additionalProperties"] = Value::Bool(false); } definitions["OffPeakPricing"]["anyOf"] = json!([ {"required": ["hours_utc"]}, @@ -123,15 +170,43 @@ fn decorate_model_entry(schema: &mut Value) { ]); definitions["OffPeakPricing"]["properties"]["windows"]["minItems"] = json!(1); definitions["OffPeakWindow"]["properties"]["weekdays"]["minItems"] = json!(1); - definitions["TieredRate"]["properties"]["range"]["items"]["minimum"] = json!(0); - definitions["TieredRate"]["properties"]["max_results_range"]["items"]["minimum"] = json!(0); - definitions["Weekday"]["anyOf"][0]["minimum"] = json!(1); - definitions["Weekday"]["anyOf"][0]["maximum"] = json!(7); - definitions["Weekday"]["anyOf"][1]["pattern"] = json!( - r"(?i)^(mon|monday|tue|tues|tuesday|wed|wednesday|thu|thur|thurs|thursday|fri|friday|sat|saturday|sun|sunday)$" - ); - let window_pattern = json!(r"^([01]\d|2[0-3]):[0-5]\d-([01]\d|2[0-3]):[0-5]\d$"); - definitions["UtcHours"]["anyOf"][0]["pattern"] = window_pattern.clone(); - definitions["UtcHours"]["anyOf"][1]["items"]["pattern"] = window_pattern; definitions["UtcHours"]["anyOf"][1]["minItems"] = json!(1); } + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::FLAT_KEY_PATTERN; + use crate::wire::FlatKey; + + #[rstest] + #[case::input_token("input_cost_per_token")] + #[case::cache_write_one_hour_band( + "cache_creation_input_token_cost_above_1hr_above_200k_tokens" + )] + #[case::image_label("output_cost_per_image_0.5K")] + #[case::resolution_tier("output_cost_per_second_4k_flex")] + #[case::feature_batch("annotation_cost_per_page_batches")] + #[case::video_interval("input_cost_per_video_per_second_above_8s_interval")] + #[case::explicit_cache_read("explicit_cache_read_input_token_cost")] + #[case::wall_clock("cost_per_second")] + #[case::unused_meter("cost_per_token")] + #[case::other_currency("input_dbu_cost_per_token")] + #[case::typo("input_cost_per_tokne")] + #[case::read_has_no_ttl("cache_read_input_token_cost_above_1hr")] + #[case::interval_on_a_token_meter("input_cost_per_token_above_8s_interval")] + fn flat_key_pattern_agrees_with_the_parser(#[case] key: &str) { + let schema = serde_json::json!({ + "type": "object", + "patternProperties": {FLAT_KEY_PATTERN: {}}, + "additionalProperties": false, + }); + let validator = jsonschema::validator_for(&schema).unwrap(); + + assert_eq!( + validator.is_valid(&serde_json::json!({key: 1})), + FlatKey::parse(key).is_some() + ); + } +} diff --git a/litellm-rust/crates/model-catalog/src/store/AGENTS.md b/litellm-rust/crates/model-catalog/src/store/AGENTS.md new file mode 100644 index 00000000000..f1dacd2e217 --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/store/AGENTS.md @@ -0,0 +1,35 @@ +# Store + +`Store` is an in-memory store, a small database built for exactly three queries. It knows nothing about the catalog. A source picks its own `Rows::Row` type, how names fold for matching, and the group each key is listed under. The catalog's sources fold to lowercase and group by raw `litellm_provider` + +| Query | Method | Real caller | +|---|---|---| +| row by exact key | `Store::get` | reading `litellm.model_cost["gpt-4o"]` | +| actual key for a name that folds to the same form | `Store::find_key` | `get_model_info("GPT-4o")` finding `gpt-4o` | +| keys in a group | `Store::keys_in` | the provider sets, `litellm.open_ai_chat_completion_models` | + +## Rules + +- nothing in `store/` imports from the rest of the crate. Wire types, case rules and field names such as `litellm_provider` belong to the source that implements `Rows` +- the index in `name_index.rs` is private to the store. Never build a folded-name map or a provider grouping anywhere else in the workspace. A new lookup pattern is a new `Store` method backed by `NameIndex` +- the index holds keys, never rows. Rows are always read from `Rows`, so a price edit is seen on the next read and never rebuilds the index +- `find_key` and `keys_in` never return a key the rows lack. A hit on a removed key rebuilds the index once and asks again +- `keys_in` matches the group exactly as the source wrote it. Provider aliases such as `vertex_ai` over `vertex_ai-language-models` are Python's job +- when two keys fold to the same name the later key in row order wins, like main's `_model_cost_lowercase_map` + +## Stamps + +`Rows::stamp` is how the store learns that the index may be stale. The store keeps the stamp it built the index at and rebuilds when the source returns a different one. A source must change its stamp whenever a key or a key's group changes, since those are the two things the index holds + +| Source | Stamp | Blind spot | +|---|---|---| +| `Catalog`'s `CatalogRows` | `()`, every change goes through `Store::edit`, which drops the index | none | +| python-bridge's bound `litellm.model_cost` | the dict object and its `len` | a delete plus an add between two reads, or a `litellm_provider` rewrite on an existing key | + +A source with a blind spot calls `Store::invalidate` when it knows keys changed. Until then a key added in the blind spot is not found, and a key moved to another provider stays listed under the old one. Removed keys are still never returned + +## Tests + +Tests live in `mod.rs`'s `tests` module and drive the store through `ModelCost`, a fake `litellm.model_cost` with real bundled keys whose row is just its `litellm_provider`. Each Python write is a method (`set`, `del`, `rebind`) and its stamp is what python-bridge can see. Add a scenario as a named test that reads like the Python user code it stands for + +The two property tests compare every answer with a linear scan of the same dict. With a stamp that sees every key change, answers must match exactly. With the dict stamp, every answer must be sound (present, same folded name), and `invalidate` must make them match again. A store change that fails either is a bug in the store, not in the test diff --git a/litellm-rust/crates/model-catalog/src/store/mod.rs b/litellm-rust/crates/model-catalog/src/store/mod.rs new file mode 100644 index 00000000000..2220526f8b6 --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/store/mod.rs @@ -0,0 +1,460 @@ +mod name_index; + +use std::borrow::Cow; +use std::fmt; +use std::sync::{PoisonError, RwLock}; + +use name_index::NameIndex; + +/// Where a `Store`'s rows live. The store rebuilds its index whenever `stamp` changes, so a source must change its +/// stamp whenever a key or a key's group changes, or call `Store::invalidate` when it cannot +pub trait Rows { + type Row: ToOwned + ?Sized; + type Stamp: PartialEq; + + /// The name two spellings of one key share, which `Store::find_key` matches on + fn fold(name: &str) -> String; + fn stamp(&self) -> Self::Stamp; + fn contains(&self, key: &str) -> bool; + fn row(&self, key: &str) -> Option>; + /// Every key in row order with the group it is listed under + fn scan(&self, visit: &mut dyn FnMut(&str, Option<&str>)); +} + +/// An in-memory store for three queries (row by exact key, actual key for a name by its folded form, keys in a group) +/// that owns the only index behind them +pub struct Store { + rows: R, + index: RwLock>>, +} + +struct Built { + stamp: S, + names: NameIndex, +} + +impl Store { + pub fn new(rows: R) -> Self { + Self { + rows, + index: RwLock::new(None), + } + } + + pub fn rows(&self) -> &R { + &self.rows + } + + /// Changes the rows in place. The index is rebuilt on the next query + pub fn edit(&mut self, edit: impl FnOnce(&mut R) -> T) -> T { + let edited = edit(&mut self.rows); + *self.index.get_mut().unwrap_or_else(PoisonError::into_inner) = None; + edited + } + + pub fn get(&self, key: &str) -> Option> { + self.rows.row(key) + } + + /// The key itself when present, otherwise the actual key that folds to the same name. Never a key the rows lack + pub fn find_key(&self, name: &str) -> Option { + if self.rows.contains(name) { + return Some(name.to_owned()); + } + let folded = R::fold(name); + let hit = self.query(|names| names.key_for(&folded).map(str::to_owned))?; + if self.rows.contains(&hit) { + return Some(hit); + } + self.invalidate(); + self.query(|names| names.key_for(&folded).map(str::to_owned)) + .filter(|key| self.rows.contains(key)) + } + + /// Keys whose group is exactly `group`, in row order. Never a key the rows lack + pub fn keys_in(&self, group: &str) -> Vec { + self.query(|names| names.keys_in(group).to_vec()) + .into_iter() + .filter(|key| self.rows.contains(key)) + .collect() + } + + /// Drops the index, for a source whose stamp cannot see every key change + pub fn invalidate(&self) { + *self.index.write().unwrap_or_else(PoisonError::into_inner) = None; + } + + fn query(&self, read: impl Fn(&NameIndex) -> T) -> T { + let stamp = self.rows.stamp(); + if let Some(built) = self + .index + .read() + .unwrap_or_else(PoisonError::into_inner) + .as_ref() + .filter(|built| built.stamp == stamp) + { + return read(&built.names); + } + let names = NameIndex::scan(&self.rows); + let answer = read(&names); + *self.index.write().unwrap_or_else(PoisonError::into_inner) = Some(Built { stamp, names }); + answer + } +} + +impl Clone for Store { + fn clone(&self) -> Self { + Self::new(self.rows.clone()) + } +} + +impl fmt::Debug for Store { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Store") + .field("rows", &self.rows) + .finish_non_exhaustive() + } +} + +#[cfg(test)] +mod tests { + use std::cell::{Cell, RefCell}; + + use indexmap::IndexMap; + use rstest::{fixture, rstest}; + + use super::*; + + const BUNDLED: &[(&str, Option<&str>)] = &[ + ("gpt-4o", Some("openai")), + ("o3-pro", Some("openai")), + ("azure/gpt-4o", Some("azure")), + ("claude-sonnet-4-5", Some("anthropic")), + ("gemini-2.5-pro", Some("vertex_ai-language-models")), + ( + "fireworks_ai/accounts/fireworks/models/llama-v3p1-8b-instruct", + Some("fireworks_ai"), + ), + ( + "anyscale/meta-llama/Meta-Llama-3-8B-Instruct", + Some("anyscale"), + ), + ("sample_free_row", None), + ]; + + /// `litellm.model_cost` as Python sees it. The stamp is what python-bridge can observe for free, the dict's + /// identity and its length, so a write that keeps both hides a key change from the store + struct ModelCost { + dict: RefCell>>, + identity: Cell, + index_changes: Cell, + honest_stamp: bool, + scans: Cell, + } + + impl ModelCost { + fn new(rows: &[(&str, Option<&str>)], honest_stamp: bool) -> Self { + Self { + dict: RefCell::new( + rows.iter() + .map(|(key, provider)| ((*key).to_owned(), provider.map(str::to_owned))) + .collect(), + ), + identity: Cell::new(0), + index_changes: Cell::new(0), + honest_stamp, + scans: Cell::new(0), + } + } + + /// `model_cost[key] = {"litellm_provider": provider, ...}` + fn set(&self, key: &str, provider: Option<&str>) { + let previous = self + .dict + .borrow_mut() + .insert(key.to_owned(), provider.map(str::to_owned)); + if previous.as_ref().map(Option::as_deref) != Some(provider) { + self.index_changes.set(self.index_changes.get() + 1); + } + } + + /// `del model_cost[key]` + fn del(&self, key: &str) { + if self.dict.borrow_mut().shift_remove(key).is_some() { + self.index_changes.set(self.index_changes.get() + 1); + } + } + + /// `litellm.model_cost = {...}`, a new dict object + fn rebind(&self, rows: &[(&str, Option<&str>)]) { + *self.dict.borrow_mut() = rows + .iter() + .map(|(key, provider)| ((*key).to_owned(), provider.map(str::to_owned))) + .collect(); + self.identity.set(self.identity.get() + 1); + self.index_changes.set(self.index_changes.get() + 1); + } + } + + /// A row is its `litellm_provider`, which is also the group the store lists its key under + impl Rows for &ModelCost { + type Row = Option; + type Stamp = (u64, u64); + + fn fold(name: &str) -> String { + name.to_lowercase() + } + + fn stamp(&self) -> Self::Stamp { + if self.honest_stamp { + (self.identity.get(), self.index_changes.get()) + } else { + (self.identity.get(), self.dict.borrow().len() as u64) + } + } + + fn contains(&self, key: &str) -> bool { + self.dict.borrow().contains_key(key) + } + + fn row(&self, key: &str) -> Option>> { + Some(Cow::Owned(self.dict.borrow().get(key)?.clone())) + } + + fn scan(&self, visit: &mut dyn FnMut(&str, Option<&str>)) { + self.scans.set(self.scans.get() + 1); + for (key, provider) in self.dict.borrow().iter() { + visit(key, provider.as_deref()); + } + } + } + + #[fixture] + fn model_cost() -> ModelCost { + ModelCost::new(BUNDLED, false) + } + + /// Rows only `Store::edit` can change, so the stamp never moves + #[derive(Debug)] + struct Owned(IndexMap>); + + impl Rows for Owned { + type Row = Option; + type Stamp = (); + + fn fold(name: &str) -> String { + name.to_lowercase() + } + + fn stamp(&self) -> Self::Stamp {} + + fn contains(&self, key: &str) -> bool { + self.0.contains_key(key) + } + + fn row(&self, key: &str) -> Option>> { + self.0.get(key).map(Cow::Borrowed) + } + + fn scan(&self, visit: &mut dyn FnMut(&str, Option<&str>)) { + for (key, provider) in &self.0 { + visit(key, provider.as_deref()); + } + } + } + + fn owned(rows: &[(&str, Option<&str>)]) -> Store { + Store::new(Owned( + rows.iter() + .map(|(key, provider)| ((*key).to_owned(), provider.map(str::to_owned))) + .collect(), + )) + } + + #[rstest] + #[case::bundled_key_as_written("gpt-4o", Some("gpt-4o"))] + #[case::user_typed_upper_case("GPT-4o", Some("gpt-4o"))] + #[case::provider_prefixed("Azure/GPT-4o", Some("azure/gpt-4o"))] + #[case::fireworks_path_in_another_case( + "fireworks_ai/accounts/fireworks/models/Llama-v3p1-8B-Instruct", + Some("fireworks_ai/accounts/fireworks/models/llama-v3p1-8b-instruct") + )] + #[case::mixed_case_bundled_key( + "anyscale/meta-llama/meta-llama-3-8b-instruct", + Some("anyscale/meta-llama/Meta-Llama-3-8B-Instruct") + )] + #[case::unmapped_model("my-finetune", None)] + #[case::prefix_is_not_a_match("gpt-4", None)] + fn find_key_maps_a_requested_model_to_its_key( + model_cost: ModelCost, + #[case] model: &str, + #[case] expected: Option<&str>, + ) { + assert_eq!(Store::new(&model_cost).find_key(model).as_deref(), expected); + } + + #[rstest] + #[case::the_bundled_spelling("gpt-4o", "gpt-4o")] + #[case::the_registered_spelling("GPT-4o", "GPT-4o")] + #[case::another_spelling_takes_the_later_key("Gpt-4O", "GPT-4o")] + fn a_registered_case_variant_is_its_own_row( + model_cost: ModelCost, + #[case] model: &str, + #[case] expected: &str, + ) { + model_cost.set("GPT-4o", Some("openai")); + assert_eq!( + Store::new(&model_cost).find_key(model).as_deref(), + Some(expected) + ); + } + + #[rstest] + #[case::openai_in_map_order("openai", &["gpt-4o", "o3-pro"])] + #[case::azure("azure", &["azure/gpt-4o"])] + #[case::raw_vertex_family("vertex_ai-language-models", &["gemini-2.5-pro"])] + #[case::provider_alias_is_python_s_job("vertex_ai", &[])] + #[case::matched_as_written("OpenAI", &[])] + #[case::unknown_provider("acme", &[])] + fn keys_in_lists_the_keys_a_raw_provider_names( + model_cost: ModelCost, + #[case] provider: &str, + #[case] expected: &[&str], + ) { + assert_eq!(Store::new(&model_cost).keys_in(provider), expected); + } + + #[rstest] + fn a_model_written_straight_into_the_dict_is_found_in_any_case(model_cost: ModelCost) { + let store = Store::new(&model_cost); + assert_eq!(store.find_key("mycorp/custom-model"), None); + + model_cost.set("MyCorp/Custom-Model", Some("openai")); + + assert_eq!( + store.find_key("mycorp/custom-model").as_deref(), + Some("MyCorp/Custom-Model") + ); + assert_eq!( + store.keys_in("openai"), + ["gpt-4o", "o3-pro", "MyCorp/Custom-Model"] + ); + } + + #[rstest] + fn a_price_edit_reuses_the_index(model_cost: ModelCost) { + let store = Store::new(&model_cost); + assert_eq!(store.find_key("GPT-4o").as_deref(), Some("gpt-4o")); + + model_cost.set("gpt-4o", Some("openai")); + model_cost.set("o3-pro", Some("openai")); + model_cost.set("gpt-4o", Some("openai")); + + assert_eq!(store.find_key("O3-PRO").as_deref(), Some("o3-pro")); + assert_eq!(store.keys_in("openai"), ["gpt-4o", "o3-pro"]); + assert_eq!(model_cost.scans.get(), 1); + } + + #[rstest] + fn a_reload_that_rebinds_the_dict_is_seen_even_at_the_same_length(model_cost: ModelCost) { + let store = Store::new(&model_cost); + assert_eq!(store.find_key("GPT-4o").as_deref(), Some("gpt-4o")); + + let mut reloaded = BUNDLED.to_vec(); + reloaded[0] = ("gpt-4o-2024-11-20", Some("openai")); + model_cost.rebind(&reloaded); + + assert_eq!(store.find_key("GPT-4o"), None); + assert_eq!( + store.find_key("GPT-4o-2024-11-20").as_deref(), + Some("gpt-4o-2024-11-20") + ); + } + + #[rstest] + fn a_deleted_model_is_never_returned(model_cost: ModelCost) { + let store = Store::new(&model_cost); + assert_eq!(store.find_key("O3-PRO").as_deref(), Some("o3-pro")); + + model_cost.del("o3-pro"); + + assert_eq!(store.find_key("O3-PRO"), None); + assert_eq!(store.keys_in("openai"), ["gpt-4o"]); + } + + #[rstest] + fn a_respelled_key_at_the_same_length_is_found_after_one_rebuild(model_cost: ModelCost) { + let store = Store::new(&model_cost); + assert_eq!(store.find_key("o3-PRO").as_deref(), Some("o3-pro")); + + model_cost.del("o3-pro"); + model_cost.set("O3-Pro", Some("openai")); + + assert_eq!(store.find_key("o3-PRO").as_deref(), Some("O3-Pro")); + assert_eq!(model_cost.scans.get(), 2); + } + + #[rstest] + fn a_swap_at_the_same_length_hides_the_new_key_until_invalidate(model_cost: ModelCost) { + let store = Store::new(&model_cost); + assert_eq!(store.find_key("O3-PRO").as_deref(), Some("o3-pro")); + + model_cost.del("o3-pro"); + model_cost.set("MyCorp/Custom-Model", Some("openai")); + + assert_eq!(store.find_key("mycorp/custom-model"), None); + assert_eq!(store.keys_in("openai"), ["gpt-4o"]); + + store.invalidate(); + + assert_eq!( + store.find_key("mycorp/custom-model").as_deref(), + Some("MyCorp/Custom-Model") + ); + assert_eq!(store.keys_in("openai"), ["gpt-4o", "MyCorp/Custom-Model"]); + } + + #[rstest] + fn a_provider_rewrite_on_an_existing_key_needs_invalidate(model_cost: ModelCost) { + let store = Store::new(&model_cost); + assert_eq!(store.keys_in("openai"), ["gpt-4o", "o3-pro"]); + + model_cost.set("o3-pro", Some("azure")); + + assert_eq!(store.keys_in("openai"), ["gpt-4o", "o3-pro"]); + + store.invalidate(); + + assert_eq!(store.keys_in("openai"), ["gpt-4o"]); + assert_eq!(store.keys_in("azure"), ["o3-pro", "azure/gpt-4o"]); + } + + #[rstest] + fn an_edit_on_owned_rows_is_seen_by_the_next_query() { + let mut store = owned(BUNDLED); + assert_eq!(store.find_key("GPT-4o").as_deref(), Some("gpt-4o")); + + store.edit(|rows| { + rows.0.shift_remove("gpt-4o"); + rows.0 + .insert("MyCorp/Custom-Model".to_owned(), Some("openai".to_owned())); + }); + + assert_eq!(store.find_key("GPT-4o"), None); + assert_eq!( + store.find_key("mycorp/custom-model").as_deref(), + Some("MyCorp/Custom-Model") + ); + assert_eq!(store.keys_in("openai"), ["o3-pro", "MyCorp/Custom-Model"]); + } + + #[rstest] + fn get_reads_the_row_by_exact_key_only() { + let store = owned(BUNDLED); + assert_eq!( + store.get("azure/gpt-4o").as_deref(), + Some(&Some("azure".to_owned())) + ); + assert_eq!(store.get("Azure/GPT-4o"), None); + } +} diff --git a/litellm-rust/crates/model-catalog/src/store/name_index.rs b/litellm-rust/crates/model-catalog/src/store/name_index.rs new file mode 100644 index 00000000000..86f05cb31af --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/store/name_index.rs @@ -0,0 +1,35 @@ +use std::collections::HashMap; + +use super::Rows; + +/// Keys only, never rows, so a value write never makes it stale. A later key wins when two keys fold to one name +#[derive(Debug, Default)] +pub(super) struct NameIndex { + by_folded: HashMap, + by_group: HashMap>, +} + +impl NameIndex { + pub(super) fn scan(rows: &R) -> Self { + let mut index = Self::default(); + rows.scan(&mut |key, group| { + index.by_folded.insert(R::fold(key), key.to_owned()); + if let Some(group) = group { + index + .by_group + .entry(group.to_owned()) + .or_default() + .push(key.to_owned()); + } + }); + index + } + + pub(super) fn key_for(&self, folded: &str) -> Option<&str> { + self.by_folded.get(folded).map(String::as_str) + } + + pub(super) fn keys_in(&self, group: &str) -> &[String] { + self.by_group.get(group).map_or(&[], Vec::as_slice) + } +} diff --git a/litellm-rust/crates/model-catalog/src/validation.rs b/litellm-rust/crates/model-catalog/src/validation.rs deleted file mode 100644 index 6dd4174148c..00000000000 --- a/litellm-rust/crates/model-catalog/src/validation.rs +++ /dev/null @@ -1,252 +0,0 @@ -use std::collections::BTreeSet; - -use serde_json::{Map, Value}; -use thiserror::Error; - -use crate::{AliasIssue, Catalog, ModelInfo, UtcHours, Weekday}; - -#[derive(Debug, Error)] -pub enum RegistryValidationError { - #[error("invalid registry entry {model:?}")] - Entry { - model: String, - #[source] - reason: EntryViolation, - }, - #[error(transparent)] - Alias(AliasIssue), -} - -#[derive(Debug, Error)] -pub enum EntryViolation { - #[error("entry must be an object")] - NotObject, - #[error("aliases must be an array")] - AliasesNotArray, - #[error("aliases must contain strings")] - AliasNotString, - #[error("entry does not match the model schema")] - Schema(#[source] serde_json::Error), - #[error("litellm_provider is required")] - MissingProvider, - #[error("invalid deprecation_date")] - DeprecationDate(#[source] time::error::Parse), - #[error("off_peak_pricing needs hours or windows")] - OffPeakWithoutHours, - #[error("off_peak_pricing.hours_utc is invalid")] - InvalidHours, - #[error("off_peak_pricing.windows is empty")] - EmptyWindows, - #[error("off_peak_pricing.weekdays is invalid")] - InvalidWeekdays, - #[error("typed entry could not be serialized")] - Serialize(#[source] serde_json::Error), - #[error("entry has an unknown field, null, or changed value: {}", fields.join(", "))] - ChangedFields { fields: Vec }, - #[error("{field} must be nonnegative")] - Negative { field: String }, - #[error("{field} must be at least one")] - UpliftBelowOne { field: String }, -} - -pub fn validate_model_entry(model: &str, value: &Value) -> Result<(), RegistryValidationError> { - validate_entry_inner(value).map_err(|reason| RegistryValidationError::Entry { - model: model.to_owned(), - reason, - }) -} - -pub fn validate_registry(catalog: &Catalog) -> Result<(), RegistryValidationError> { - if let Some(issue) = catalog.alias_issues().first() { - return Err(RegistryValidationError::Alias(issue.clone())); - } - catalog.model_names().try_for_each(|name| { - let entry = catalog.lookup(name).expect("catalog name must resolve"); - validate_model_entry(name, &Value::Object(entry.entry.fields().clone())) - }) -} - -fn json_eq(left: &Value, right: &Value) -> bool { - match (left, right) { - (Value::Number(left), Value::Number(right)) => left.as_f64() == right.as_f64(), - (Value::Array(left), Value::Array(right)) => { - left.len() == right.len() && left.iter().zip(right).all(|(a, b)| json_eq(a, b)) - } - (Value::Object(left), Value::Object(right)) => { - left.len() == right.len() - && left - .iter() - .all(|(key, value)| right.get(key).is_some_and(|other| json_eq(value, other))) - } - _ => left == right, - } -} - -fn keys(value: &Map) -> BTreeSet { - value.keys().cloned().collect() -} - -fn symmetric_difference(left: &BTreeSet, right: &BTreeSet) -> Vec { - left.symmetric_difference(right).cloned().collect() -} - -fn validate_entry_inner(value: &Value) -> Result<(), EntryViolation> { - let object = value.as_object().ok_or(EntryViolation::NotObject)?; - if let Some(aliases) = object.get("aliases") { - let names = aliases.as_array().ok_or(EntryViolation::AliasesNotArray)?; - if names.iter().any(|name| !name.is_string()) { - return Err(EntryViolation::AliasNotString); - } - } - let info: ModelInfo = serde_json::from_value(value.clone()).map_err(EntryViolation::Schema)?; - if info.litellm_provider.is_none() { - return Err(EntryViolation::MissingProvider); - } - validate_dates_and_windows(&info)?; - let serialized = serde_json::to_value(info).map_err(EntryViolation::Serialize)?; - let mut expected = object.clone(); - expected.remove("aliases"); - if !json_eq(&Value::Object(expected.clone()), &serialized) { - let actual = serialized - .as_object() - .expect("ModelInfo serializes as an object"); - return Err(EntryViolation::ChangedFields { - fields: symmetric_difference(&keys(&expected), &keys(actual)), - }); - } - check_prices("", value) -} - -fn validate_dates_and_windows(info: &ModelInfo) -> Result<(), EntryViolation> { - if let Some(date) = &info.deprecation_date { - let format = time::format_description::parse_borrowed::<2>("[year]-[month]-[day]").unwrap(); - time::Date::parse(date, &format).map_err(EntryViolation::DeprecationDate)?; - } - let Some(pricing) = &info.off_peak_pricing else { - return Ok(()); - }; - if pricing.hours_utc.is_none() && pricing.windows.is_none() { - return Err(EntryViolation::OffPeakWithoutHours); - } - if let Some(hours) = &pricing.hours_utc { - validate_hours(hours)?; - } - if let Some(windows) = &pricing.windows { - if windows.is_empty() { - return Err(EntryViolation::EmptyWindows); - } - windows.iter().try_for_each(|window| { - validate_hours(&window.hours_utc)?; - if let Some(days) = &window.weekdays - && (days.is_empty() || days.iter().any(|day| !valid_weekday(day))) - { - return Err(EntryViolation::InvalidWeekdays); - } - Ok(()) - })?; - } - Ok(()) -} - -fn validate_hours(hours: &UtcHours) -> Result<(), EntryViolation> { - let values = match hours { - UtcHours::Single(value) => std::slice::from_ref(value), - UtcHours::Multiple(values) => values.as_slice(), - }; - if values.is_empty() || values.iter().any(|value| !valid_utc_window(value)) { - return Err(EntryViolation::InvalidHours); - } - Ok(()) -} - -fn valid_utc_window(value: &str) -> bool { - let Some((start, end)) = value.split_once('-') else { - return false; - }; - [start, end].into_iter().all(|clock| { - let Some((hour, minute)) = clock.split_once(':') else { - return false; - }; - hour.len() == 2 - && minute.len() == 2 - && hour.parse::().is_ok_and(|hour| hour < 24) - && minute.parse::().is_ok_and(|minute| minute < 60) - }) -} - -fn valid_weekday(day: &Weekday) -> bool { - match day { - Weekday::Number(number) => (1..=7).contains(number), - Weekday::Name(name) => matches!( - name.to_ascii_lowercase().as_str(), - "mon" - | "monday" - | "tue" - | "tues" - | "tuesday" - | "wed" - | "wednesday" - | "thu" - | "thur" - | "thurs" - | "thursday" - | "fri" - | "friday" - | "sat" - | "saturday" - | "sun" - | "sunday" - ), - } -} - -fn check_prices(path: &str, value: &Value) -> Result<(), EntryViolation> { - let Some(object) = value.as_object() else { - return Ok(()); - }; - object.iter().try_for_each(|(key, field)| { - let field_path = format!("{path}.{key}"); - if (key.contains("cost") - || path.ends_with(".guardrail_cost_per_unit") - || path.ends_with(".search_context_cost_per_query")) - && let Some(number) = field.as_f64() - && number < 0.0 - { - return Err(EntryViolation::Negative { - field: relative(&field_path), - }); - } - if key.contains("uplift_multiplier") - && let Some(number) = field.as_f64() - && number < 1.0 - { - return Err(EntryViolation::UpliftBelowOne { - field: relative(&field_path), - }); - } - if matches!(key.as_str(), "range" | "max_results_range") - && field.as_array().is_some_and(|values| { - values - .iter() - .any(|value| value.as_f64().is_some_and(|n| n < 0.0)) - }) - { - return Err(EntryViolation::Negative { - field: relative(&field_path), - }); - } - if matches!(key.as_str(), "metadata" | "provider_specific_entry") { - return Ok(()); - } - match field.as_array() { - Some(items) => items.iter().enumerate().try_for_each(|(index, item)| { - check_prices(&format!("{field_path}[{index}]"), item) - }), - None => check_prices(&field_path, field), - } - }) -} - -fn relative(field_path: &str) -> String { - field_path.trim_start_matches('.').to_owned() -} diff --git a/litellm-rust/crates/model-catalog/src/vocabulary.rs b/litellm-rust/crates/model-catalog/src/vocabulary.rs new file mode 100644 index 00000000000..2b7ac8bd1ab --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/vocabulary.rs @@ -0,0 +1,126 @@ +/// Primary API surface / task type of the model. +#[macro_rules_attribute::apply(crate::wire_type)] +#[derive(Copy, Eq, Hash)] +#[serde(rename_all = "snake_case")] +pub enum Mode { + AudioSpeech, + AudioTranscription, + Chat, + Completion, + Decisions, + Embedding, + Evaluation, + Extraction, + Guardrail, + ImageEdit, + ImageGeneration, + Moderation, + Ocr, + Realtime, + Rerank, + Responses, + Search, + VectorStore, + VideoGeneration, +} + +/// Gemini audio generation API the model is served through. +#[macro_rules_attribute::apply(crate::wire_type)] +#[derive(Copy, Eq, Hash)] +#[serde(rename_all = "snake_case")] +pub enum VertexAiAudioApi { + LyriaPredict, + LyriaInteractions, +} + +/// Audio container format the model can return. +#[macro_rules_attribute::apply(crate::wire_type)] +#[derive(Copy, Eq, Hash)] +#[serde(rename_all = "snake_case")] +pub enum AudioFormat { + Mp3, + Wav, +} + +/// Input modality the model accepts. +#[macro_rules_attribute::apply(crate::wire_type)] +#[derive(Copy, Eq, Hash)] +#[serde(rename_all = "snake_case")] +pub enum InputModality { + Text, + Image, + Audio, + Video, +} + +/// Output modality the model can produce. +#[macro_rules_attribute::apply(crate::wire_type)] +#[derive(Copy, Eq, Hash)] +#[serde(rename_all = "snake_case")] +pub enum OutputModality { + Text, + Image, + Audio, + Video, + Code, +} + +/// API route a model can be called through. +#[macro_rules_attribute::apply(crate::wire_type)] +#[derive(Copy, Eq, Hash)] +pub enum Endpoint { + #[serde(rename = "/v1/chat/completions")] + ChatCompletions, + #[serde(rename = "/v1/completions")] + Completions, + #[serde(rename = "/v1/responses")] + Responses, + #[serde(rename = "/v1/messages")] + Messages, + #[serde(rename = "/v1/batch")] + Batch, + #[serde(rename = "/v1/embeddings")] + Embeddings, + #[serde(rename = "/v1/images/generations")] + ImageGenerations, + #[serde(rename = "/v1/images/edits")] + ImageEdits, + #[serde(rename = "/v1/audio/speech")] + AudioSpeech, + #[serde(rename = "/v1/audio/transcriptions")] + AudioTranscriptions, + #[serde(rename = "/v1/realtime")] + Realtime, + #[serde(rename = "/v1/realtime/transcription_sessions")] + RealtimeTranscriptionSessions, + #[serde(rename = "/v1/videos")] + Videos, + #[serde(rename = "/v1/ocr")] + Ocr, + #[serde(rename = "/v1/listen")] + Listen, + #[serde(rename = "/v1/systemone")] + SystemOne, + #[serde(rename = "/v1/decisions")] + Decisions, + #[serde(rename = "/v1beta/interactions")] + Interactions, + #[serde(rename = "/v1beta/models")] + Models, + #[serde(rename = "/vertex_ai/live")] + VertexAiLive, +} + +/// Reasoning effort level a model accepts or applies. +#[macro_rules_attribute::apply(crate::wire_type)] +#[derive(Copy, Eq, Hash)] +#[serde(rename_all = "snake_case")] +pub enum EffortLevel { + None, + Minimal, + Low, + Medium, + High, + Xhigh, + Max, +} diff --git a/litellm-rust/crates/model-catalog/src/wire/attributes.rs b/litellm-rust/crates/model-catalog/src/wire/attributes.rs new file mode 100644 index 00000000000..93c10a81524 --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/wire/attributes.rs @@ -0,0 +1,344 @@ +use std::collections::BTreeMap; +use std::fmt; +use std::str::FromStr; + +use litellm_cost::money::UnitPrice; +use serde::de::{self, Deserializer}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use crate::error::{DateCause, Error, Repr}; +use crate::vocabulary::{ + AudioFormat, EffortLevel, Endpoint, InputModality, Mode, OutputModality, VertexAiAudioApi, +}; +use crate::wire::lax::Count; + +macro_rules! attributes { + ($($(#[$doc:meta])* $name:ident: $ty:ty $(as $via:ty)?),* $(,)?) => { + /// The named, non-price fields of a catalog entry. + #[derive(Clone, Debug, Default, PartialEq)] + #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] + pub struct Attributes { + $($(#[$doc])* pub $name: Option<$ty>,)* + } + + impl Attributes { + pub(crate) fn read( + &mut self, + key: &str, + value: &Value, + ) -> Option> { + match key { + $(stringify!($name) => Some( + attributes!(@read value, $ty $(, $via)?).map(|read| self.$name = Some(read)), + ),)* + _ => None, + } + } + } + }; + (@read $value:ident, $ty:ty) => { + <$ty as Deserialize>::deserialize($value) + }; + (@read $value:ident, $ty:ty, $via:ty) => { + serde_with::de::DeserializeAsWrap::<$ty, $via>::deserialize($value) + .map(serde_with::de::DeserializeAsWrap::into_inner) + }; +} + +attributes! { + /// Other names this entry is served under. + aliases: Vec, + audio_transcription_config: String, + bedrock_converse_supports_strict_tools: bool, + /// Highest reasoning effort the Bedrock output_config accepts for this model. + bedrock_output_config_effort_ceiling: EffortLevel, + comment: String, + /// Reasoning effort the provider applies when the request omits reasoning_effort. Gates whether a non-default temperature or the top_p/logprobs sampling params are accepted, which hold only when the effort resolves to 'none'. + default_reasoning_effort: EffortLevel, + /// Date the provider deprecates the model, YYYY-MM-DD. + deprecation_date: IsoDate, + gemini_audio_only_live: bool, + gemini_native_audio: bool, + input_cost_per_token_cache_hit: UnitPrice, + /// Databricks units, not USD. + input_dbu_cost_per_token: Dbu, + /// LiteLLM provider slug; one of https://docs.litellm.ai/docs/providers. + litellm_provider: String, + /// Maximum prompt/context tokens the model accepts. + max_input_tokens: u64 as Count, + /// Maximum tokens the model can generate in one response. + max_output_tokens: u64 as Count, + /// Legacy field: max output tokens if the provider specifies it, else max input tokens. + max_tokens: u64 as Count, + /// Free-form notes about the entry (e.g. pricing derivation). + metadata: BTreeMap, + /// Primary API surface / task type of the model. + mode: Mode, + /// Databricks units, not USD. + output_dbu_cost_per_token: Dbu, + /// Embedding dimension for embedding models. + output_vector_size: u64 as Count, + /// Smallest prefix the provider will actually cache; absent means the provider default applies. + prompt_cache_min_tokens: u64 as Count, + /// Provider-internal routing hints (e.g. bedrock_invocation_schema). + provider_specific_entry: BTreeMap, + reasoning_contexts: Vec, + /// Exact reasoning_effort levels this deployment accepts; wins over supports_* flags. + reasoning_effort_levels: Vec, + reasoning_modes: Vec, + /// Provider default requests-per-minute limit. + rpm: u64 as Count, + /// Grounded search requests the provider includes for free before search is billed. + search_free_requests: SearchFreeRequests, + /// URL of the provider pricing/model page this entry was taken from. + source: String, + /// Audio container formats the model can return. + supported_audio_formats: Vec, + /// OpenAI-style API routes this model can be called through, e.g. /v1/chat/completions. + supported_endpoints: Vec, + /// Input modalities the model accepts. + supported_modalities: Vec, + /// Output modalities the model can produce. + supported_output_modalities: Vec, + /// Cloud regions the model is available in ('global' or region ids). + supported_regions: Vec, + thinking_always_on: bool, + token_count_framing: u64 as Count, + token_count_model: TokenCountModel, + token_count_tokenizer: TokenCountTokenizer, + /// Provider default tokens-per-minute limit. + tpm: u64 as Count, + use_openai_responses_path: bool, + uses_embed_content: bool, + vertex_ai_audio_api: VertexAiAudioApi, + /// Whether web search is billed per query or per prompt. + web_search_billing_unit: WebSearchBillingUnit, +} + +/// The model whose provider token counter counts tokens for this entry. +#[macro_rules_attribute::apply(crate::wire_type)] +#[derive(Eq, Hash)] +#[serde(deny_unknown_fields)] +pub struct TokenCountModel { + pub model: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub litellm_provider: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub region: Option, +} + +/// Free grounded search requests per period, shared across the models of one group. +#[macro_rules_attribute::apply(crate::wire_type)] +#[derive(Eq, Hash)] +#[serde(deny_unknown_fields)] +pub struct SearchFreeRequests { + pub count: u64, + pub period: FreeRequestPeriod, + pub shared_group: String, +} + +#[macro_rules_attribute::apply(crate::wire_type)] +#[derive(Copy, Eq, Hash)] +#[serde(rename_all = "snake_case")] +pub enum FreeRequestPeriod { + Day, + Month, +} + +/// Local tokenizer that counts tokens for this entry. +#[macro_rules_attribute::apply(crate::wire_type)] +#[derive(Copy, Eq, Hash)] +#[serde(rename_all = "snake_case")] +pub enum TokenCountTokenizer { + KimiK3, +} + +/// Whether web search is billed per query or per prompt. +#[macro_rules_attribute::apply(crate::wire_type)] +#[derive(Copy, Eq, Hash)] +#[serde(rename_all = "snake_case")] +pub enum WebSearchBillingUnit { + PerQuery, + PerPrompt, +} + +/// A non-negative amount of Databricks units. +#[derive(Clone, Copy, Debug, PartialEq, PartialOrd, Serialize)] +#[serde(transparent)] +pub struct Dbu(f64); + +impl Dbu { + pub fn new(units: f64) -> Option { + (units.is_finite() && units >= 0.0).then_some(Self(units)) + } + pub fn get(self) -> f64 { + self.0 + } +} + +impl<'de> Deserialize<'de> for Dbu { + fn deserialize>(deserializer: D) -> Result { + let units = f64::deserialize(deserializer)?; + Self::new(units).ok_or_else(|| { + de::Error::invalid_value( + de::Unexpected::Float(units), + &"a non-negative number of Databricks units", + ) + }) + } +} + +/// A calendar date written YYYY-MM-DD. +#[derive( + Clone, + Copy, + Debug, + PartialEq, + Eq, + PartialOrd, + Ord, + Hash, + serde_with::DeserializeFromStr, + serde_with::SerializeDisplay, +)] +pub struct IsoDate(pub time::Date); + +impl FromStr for IsoDate { + type Err = Error; + + fn from_str(date: &str) -> Result { + let mut parts = date.split('-'); + let (Some(year), Some(month), Some(day), None) = + (parts.next(), parts.next(), parts.next(), parts.next()) + else { + return Err(Repr::InvalidDate { cause: None }.into()); + }; + if year.len() != 4 || month.len() != 2 || day.len() != 2 { + return Err(Repr::InvalidDate { cause: None }.into()); + } + let number = |part: &str| part.parse::().map_err(DateCause::from); + let month = time::Month::try_from(number(month)?).map_err(DateCause::from)?; + time::Date::from_calendar_date(year.parse().map_err(DateCause::from)?, month, number(day)?) + .map(Self) + .map_err(DateCause::from) + .map_err(Error::from) + } +} + +impl fmt::Display for IsoDate { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "{:04}-{:02}-{:02}", + self.0.year(), + u8::from(self.0.month()), + self.0.day() + ) + } +} + +#[cfg(feature = "schema")] +mod schema { + use std::borrow::Cow; + + use schemars::{JsonSchema, Schema, SchemaGenerator, json_schema}; + + use super::{Dbu, IsoDate}; + + impl JsonSchema for IsoDate { + fn inline_schema() -> bool { + true + } + fn schema_name() -> Cow<'static, str> { + "IsoDate".into() + } + fn json_schema(_: &mut SchemaGenerator) -> Schema { + json_schema!({ + "type": "string", + "format": "date", + "pattern": r"^\d{4}-(0[1-9]|1[0-2])-(0[1-9]|[12]\d|3[01])$" + }) + } + } + + impl JsonSchema for Dbu { + fn inline_schema() -> bool { + true + } + fn schema_name() -> Cow<'static, str> { + "Dbu".into() + } + fn json_schema(_: &mut SchemaGenerator) -> Schema { + json_schema!({"type": "number", "minimum": 0}) + } + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + #[rstest] + #[case::leap_day("2028-02-29", true)] + #[case::not_a_leap_year("2026-02-29", false)] + #[case::impossible_day("2026-02-31", false)] + #[case::one_digit_month("2026-2-01", false)] + #[case::extra_part("2026-02-01-01", false)] + fn iso_dates_are_real_calendar_dates(#[case] text: &str, #[case] valid: bool) { + assert_eq!(text.parse::().is_ok(), valid); + } + + #[rstest] + #[case::count_written_as_text("max_tokens", json!("8192"), true)] + #[case::negative_count("max_tokens", json!(-1), false)] + #[case::known_mode("mode", json!("chat"), true)] + #[case::unknown_mode("mode", json!("chatt"), false)] + #[case::negative_dbu("input_dbu_cost_per_token", json!(-1), false)] + #[case::reasoning_contexts("reasoning_contexts", json!(["reasoning", "non_reasoning"]), true)] + #[case::reasoning_modes("reasoning_modes", json!(["none", "adaptive"]), true)] + #[case::reasoning_context_not_a_list("reasoning_contexts", json!("reasoning"), false)] + #[case::reasoning_mode_not_a_string("reasoning_modes", json!([1]), false)] + fn each_attribute_reads_its_own_type( + #[case] key: &str, + #[case] value: Value, + #[case] valid: bool, + ) { + assert_eq!( + Attributes::default() + .read(key, &value) + .map(|read| read.is_ok()), + Some(valid) + ); + } + + #[rstest] + fn keys_that_are_not_attributes_are_left_to_the_caller() { + assert!( + Attributes::default() + .read("input_cost_per_token", &json!(1)) + .is_none() + ); + } + + #[rstest] + fn reasoning_metadata_is_retained() { + let mut attributes = Attributes::default(); + attributes + .read("reasoning_contexts", &json!(["reasoning", "non_reasoning"])) + .unwrap() + .unwrap(); + attributes + .read("reasoning_modes", &json!(["none", "adaptive"])) + .unwrap() + .unwrap(); + assert_eq!( + attributes.reasoning_contexts.unwrap(), + ["reasoning", "non_reasoning"] + ); + assert_eq!(attributes.reasoning_modes.unwrap(), ["none", "adaptive"]); + } +} diff --git a/litellm-rust/crates/model-catalog/src/wire/blocks.rs b/litellm-rust/crates/model-catalog/src/wire/blocks.rs new file mode 100644 index 00000000000..3a355686b7c --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/wire/blocks.rs @@ -0,0 +1,182 @@ +use litellm_cost::money::UnitPrice; +use litellm_cost::schedule::{UtcWindow, Weekday}; +use serde::Deserialize; +use serde::de::DeserializeOwned; +use serde_json::{Map, Value}; +use serde_with::de::DeserializeAsWrap; + +use crate::error::{Error, Repr, pointer_child, pointer_for_key}; +use crate::wire::flat_key::FlatKey; +use crate::wire::lax::Count; + +/// A flat price key with its spelling, kept so problems can name it. +#[derive(Clone, Debug, PartialEq)] +pub struct FlatPrice { + pub key: String, + pub path: String, + pub flat: FlatKey, + pub price: UnitPrice, +} + +#[derive(Clone, Debug, Default, PartialEq)] +pub struct TierRow { + pub path: String, + pub range: Option<[u64; 2]>, + pub max_results_range: Option<[u64; 2]>, + pub prices: Vec, +} + +#[derive(Clone, Debug, Default, PartialEq)] +pub struct OffPeakBlock { + pub hours_utc: Option, + pub windows: Option>, + pub weekday_timezone: Option, + pub prices: Vec, +} + +/// One UTC window or a list of them. +#[derive(Clone, Debug, PartialEq, Eq, Deserialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(untagged)] +pub enum UtcHours { + Single(UtcWindow), + Multiple(Vec), +} + +/// One off-peak window entry inside `windows`. +#[derive(Clone, Debug, PartialEq, Eq, Deserialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(deny_unknown_fields)] +pub struct OffPeakWindow { + pub hours_utc: UtcHours, + pub weekdays: Option>, +} + +/// USD cost per web search query, keyed by search context size. +#[derive(Clone, Debug, Default, PartialEq, Deserialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(deny_unknown_fields)] +pub struct SearchContextCostPerQuery { + pub search_context_size_low: Option, + pub search_context_size_medium: Option, + pub search_context_size_high: Option, +} + +impl UtcHours { + pub fn windows(&self) -> &[UtcWindow] { + match self { + Self::Single(window) => std::slice::from_ref(window), + Self::Multiple(windows) => windows, + } + } +} + +pub(crate) struct Reader<'a> { + pub(crate) problems: &'a mut Vec, +} + +impl Reader<'_> { + pub(crate) fn value(&mut self, path: String, value: &Value) -> Option { + T::deserialize(value) + .map_err(|error| self.problems.push(Error::new(path, Repr::Invalid(error)))) + .ok() + } + + pub(crate) fn price( + &mut self, + key: String, + path: String, + flat: FlatKey, + value: &Value, + ) -> Option { + let price = self.value(path.clone(), value)?; + Some(FlatPrice { + key, + path, + flat, + price, + }) + } + + fn object<'v>(&mut self, path: &str, value: &'v Value) -> Option<&'v Map> { + match value { + Value::Object(fields) => Some(fields), + other => self + .value::>(path.to_owned(), other) + .and(None), + } + } + + fn block_price( + &mut self, + name: &str, + key: String, + path: String, + value: &Value, + ) -> Option { + let Some(flat) = FlatKey::parse(name) else { + self.problems.push(Error::new(path, Repr::UnknownBlockKey)); + return None; + }; + self.price(key, path, flat, value) + } + + pub(crate) fn tier_rows(&mut self, value: &Value) -> Option> { + let path = pointer_for_key("tiered_pricing"); + let rows: Vec = self.value(path.clone(), value)?; + Some( + rows.iter() + .enumerate() + .filter_map(|(index, row)| self.tier_row(path.clone(), index, row)) + .collect(), + ) + } + + fn tier_row(&mut self, parent: String, index: usize, value: &Value) -> Option { + let path = format!("{parent}/{index}"); + let display_path = format!("tiered_pricing[{index}]"); + let fields = self.object(&path, value)?; + let mut row = TierRow { + path: path.clone(), + ..TierRow::default() + }; + for (name, value) in fields.iter().filter(|(_, value)| !value.is_null()) { + let key = format!("{display_path}.{name}"); + let child_path = pointer_child(&path, name); + match name.as_str() { + "range" => row.range = self.bounds(child_path, value), + "max_results_range" => row.max_results_range = self.bounds(child_path, value), + _ => row + .prices + .extend(self.block_price(name, key, child_path, value)), + } + } + Some(row) + } + + fn bounds(&mut self, path: String, value: &Value) -> Option<[u64; 2]> { + DeserializeAsWrap::<[u64; 2], [Count; 2]>::deserialize(value) + .map(DeserializeAsWrap::into_inner) + .map_err(|error| self.problems.push(Error::new(path, Repr::Invalid(error)))) + .ok() + } + + pub(crate) fn off_peak(&mut self, value: &Value) -> Option { + let path = pointer_for_key("off_peak_pricing"); + let fields = self.object(&path, value)?; + let mut block = OffPeakBlock::default(); + for (name, value) in fields.iter().filter(|(_, value)| !value.is_null()) { + let key = format!("off_peak_pricing.{name}"); + let child_path = pointer_child(&path, name); + match name.as_str() { + "hours_utc" => block.hours_utc = self.value(child_path, value), + "windows" => block.windows = self.value(child_path, value), + "weekday_timezone" => block.weekday_timezone = self.value(child_path, value), + _ => block + .prices + .extend(self.block_price(name, key, child_path, value)), + } + } + Some(block) + } +} diff --git a/litellm-rust/crates/model-catalog/src/wire/entry.rs b/litellm-rust/crates/model-catalog/src/wire/entry.rs new file mode 100644 index 00000000000..9c155c6496a --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/wire/entry.rs @@ -0,0 +1,193 @@ +use std::collections::BTreeMap; + +use litellm_cost::money::{Multiplier, UnitPrice}; +use serde_json::{Map, Value}; + +use crate::error::{Error, Repr, pointer_for_key}; +use crate::wire::attributes::Attributes; +use crate::wire::blocks::{FlatPrice, OffPeakBlock, Reader, SearchContextCostPerQuery, TierRow}; +use crate::wire::flat_key::FlatKey; + +pub type Fields = Map; + +/// One catalog entry as written. Every non-null key lands in exactly one field, and a known key +/// with a bad value is kept out and listed in `problems`. +#[derive(Debug, Default)] +pub struct EntryWire { + pub attributes: Attributes, + pub prices: Vec, + pub capabilities: BTreeMap, + pub tier_rows: Option>, + pub off_peak: Option, + pub search_context: Option, + pub guardrail: Option>, + pub regional_endpoint_uplift: Option, + pub residency_uplift_eu: Option, + pub residency_uplift_us: Option, + pub unknown: Fields, + pub problems: Vec, +} + +impl EntryWire { + pub fn read(fields: &Fields) -> Self { + let mut problems = Vec::new(); + let mut wire = Self::default(); + let mut reader = Reader { + problems: &mut problems, + }; + for (key, value) in fields.iter().filter(|(_, value)| !value.is_null()) { + wire.read_key(&mut reader, key, value); + } + Self { problems, ..wire } + } + + fn read_key(&mut self, reader: &mut Reader<'_>, key: &str, value: &Value) { + let path = || pointer_for_key(key); + match key { + "tiered_pricing" => self.tier_rows = reader.tier_rows(value), + "off_peak_pricing" => self.off_peak = reader.off_peak(value), + "search_context_cost_per_query" => self.search_context = reader.value(path(), value), + "guardrail_cost_per_unit" => self.guardrail = reader.value(path(), value), + "regional_endpoint_uplift_multiplier" => { + self.regional_endpoint_uplift = reader.value(path(), value) + } + "regional_processing_uplift_multiplier_eu" => { + self.residency_uplift_eu = reader.value(path(), value) + } + "regional_processing_uplift_multiplier_us" => { + self.residency_uplift_us = reader.value(path(), value) + } + _ => self.read_open_key(reader, key, value), + } + } + + fn read_open_key(&mut self, reader: &mut Reader<'_>, key: &str, value: &Value) { + if let Some(read) = self.attributes.read(key, value) { + if let Err(error) = read { + reader + .problems + .push(Error::from_key(key, Repr::Invalid(error))); + } + return; + } + if let Some(flat) = FlatKey::parse(key) { + self.prices + .extend(reader.price(key.to_owned(), pointer_for_key(key), flat, value)); + return; + } + if key.starts_with("supports_") { + if let Some(flag) = reader.value(pointer_for_key(key), value) { + self.capabilities.insert(key.to_owned(), flag); + } + return; + } + self.unknown.insert(key.to_owned(), value.clone()); + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + fn read(entry: Value) -> EntryWire { + let Value::Object(fields) = entry else { + panic!("entry must be an object") + }; + EntryWire::read(&fields) + } + + fn problem_keys(wire: &EntryWire) -> Vec { + wire.problems + .iter() + .map(|problem| problem.path().to_owned()) + .collect() + } + + #[rstest] + fn every_key_lands_in_exactly_one_place() { + let wire = read(json!({ + "litellm_provider": "openai", + "max_tokens": "8192", + "input_cost_per_token": "0.000001", + "tiered_pricing": [{"range": [0, 10], "output_cost_per_token": 2e-6}], + "supports_vision": true, + "id": "deployment-1", + "mode": null + })); + + assert_eq!(wire.attributes.litellm_provider.as_deref(), Some("openai")); + assert_eq!(wire.attributes.max_tokens, Some(8192)); + assert_eq!(wire.attributes.mode, None); + assert_eq!( + wire.prices + .iter() + .map(|price| price.key.as_str()) + .collect::>(), + ["input_cost_per_token"] + ); + assert_eq!(wire.tier_rows.unwrap()[0].prices.len(), 1); + assert_eq!( + wire.capabilities, + BTreeMap::from([("supports_vision".to_owned(), true)]) + ); + assert_eq!(wire.unknown.keys().collect::>(), ["id"]); + assert!(wire.problems.is_empty()); + } + + #[rstest] + fn every_bad_value_is_reported_by_its_key_at_once() { + let wire = read(json!({ + "mode": "chatt", + "max_tokens": -1, + "input_cost_per_token": "x", + "output_cost_per_image_0.5K": "x", + "supports_vision": "yes", + "regional_endpoint_uplift_multiplier": 0.5, + "tiered_pricing": [ + {"range": [0, 10], "input_cost_per_token": 1e-6}, + {"range": [10.5, 20], "typo_cost": 1, "output_cost_per_token": -1} + ], + "off_peak_pricing": {"hours_utc": "25:00-01:00", "input_cost_per_token": 1e-6} + })); + + assert_eq!( + problem_keys(&wire) + .into_iter() + .collect::>(), + [ + "/input_cost_per_token", + "/max_tokens", + "/mode", + "/off_peak_pricing/hours_utc", + "/output_cost_per_image_0.5K", + "/regional_endpoint_uplift_multiplier", + "/supports_vision", + "/tiered_pricing/1/output_cost_per_token", + "/tiered_pricing/1/range", + "/tiered_pricing/1/typo_cost", + ] + .map(str::to_owned) + .into_iter() + .collect() + ); + } + + #[rstest] + #[case::tiers_not_a_list(json!({"tiered_pricing": {"range": [0, 1]}}), "/tiered_pricing")] + #[case::row_not_an_object(json!({"tiered_pricing": [3]}), "/tiered_pricing/0")] + #[case::off_peak_not_an_object(json!({"off_peak_pricing": true}), "/off_peak_pricing")] + #[case::unknown_search_size( + json!({"search_context_cost_per_query": {"search_context_size_huge": 1}}), + "/search_context_cost_per_query" + )] + #[case::negative_guardrail( + json!({"guardrail_cost_per_unit": {"contentPolicyUnits": -1}}), + "/guardrail_cost_per_unit" + )] + fn a_bad_block_is_reported_at_the_block(#[case] entry: Value, #[case] expected: &str) { + assert_eq!(problem_keys(&read(entry)), [expected]); + } +} diff --git a/litellm-rust/crates/model-catalog/src/wire/flat_key.rs b/litellm-rust/crates/model-catalog/src/wire/flat_key.rs new file mode 100644 index 00000000000..906d917b487 --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/wire/flat_key.rs @@ -0,0 +1,359 @@ +use std::num::NonZeroU64; + +use litellm_cost::price::{ + Band, CacheReadMode, CacheTtl, Channel, Feature, ImageSize, InputUnit, Meter, OutputUnit, + Resolution, ServiceTier, TokenKind, When, +}; + +/// A flat price key as the grammar reads it: `{meter}[_above_{N}k_tokens][_{tier}]`. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct FlatKey { + pub meter: Meter, + pub band: Band, + pub tier: Option, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum TierSuffix { + Service(ServiceTier), + Batches, + Decisions, +} + +impl FlatKey { + pub fn parse(key: &str) -> Option { + let (rest, tier) = split_tier(key); + let (base, suffix) = split_above(rest)?; + let (meter, band) = match suffix { + Above::None => (meter(base)?, Band::Base), + Above::PromptTokens(tokens) => (meter(base)?, Band::over_prompt_tokens(tokens)), + Above::Seconds(seconds) => match meter(base)? { + Meter::Input(InputUnit::VideoSecond { over_seconds: None }) => ( + Meter::Input(InputUnit::VideoSecond { + over_seconds: NonZeroU64::new(seconds), + }), + Band::Base, + ), + _ => return None, + }, + }; + Some(Self { meter, band, tier }) + } + + pub fn has_suffix(&self) -> bool { + self.band != Band::Base || self.tier.is_some() + } + + pub fn channel_and_when(&self) -> (Channel, When) { + let (channel, tier) = match self.tier { + None => (Channel::Online, ServiceTier::Standard), + Some(TierSuffix::Service(tier)) => (Channel::Online, tier), + Some(TierSuffix::Batches) => (Channel::Batch, ServiceTier::Standard), + Some(TierSuffix::Decisions) => (Channel::Decisions, ServiceTier::Standard), + }; + ( + channel, + When { + tier, + band: self.band, + }, + ) + } +} + +enum Above { + None, + PromptTokens(u64), + Seconds(u64), +} + +fn split_tier(key: &str) -> (&str, Option) { + let Some((head, tail)) = key.rsplit_once('_') else { + return (key, None); + }; + let tier = match tail { + "priority" => TierSuffix::Service(ServiceTier::Priority), + "flex" => TierSuffix::Service(ServiceTier::Flex), + "ultrafast" => TierSuffix::Service(ServiceTier::Ultrafast), + "balanced" => TierSuffix::Service(ServiceTier::Balanced), + "batches" => TierSuffix::Batches, + "decisions" => TierSuffix::Decisions, + _ => return (key, None), + }; + (head, Some(tier)) +} + +fn split_above(key: &str) -> Option<(&str, Above)> { + let Some((head, tail)) = key.rsplit_once("_above_") else { + return Some((key, Above::None)); + }; + if tail == "1hr" || tail == "12hr" { + return Some((key, Above::None)); + } + if let Some(thousands) = tail.strip_suffix("k_tokens") { + return Some(( + head, + Above::PromptTokens(number(thousands)?.checked_mul(1000)?), + )); + } + if let Some(tokens) = tail.strip_suffix("_tokens") { + return Some((head, Above::PromptTokens(number(tokens)?))); + } + let seconds = tail.strip_suffix("s_interval")?; + Some((head, Above::Seconds(number(seconds)?))) +} + +fn number(digits: &str) -> Option { + (!digits.is_empty() && digits.bytes().all(|byte| byte.is_ascii_digit())) + .then(|| digits.parse().ok()) + .flatten() +} + +fn meter(base: &str) -> Option { + let fixed = match base { + "cost_per_second" => Some(Meter::WallClockSecond), + "explicit_cache_read_input_token_cost" => Some(Meter::CacheRead { + kind: TokenKind::Text, + mode: CacheReadMode::Explicit, + }), + "annotation_cost_per_page" => Some(Meter::Feature(Feature::AnnotationPage)), + "citation_cost_per_token" => Some(Meter::Feature(Feature::CitationToken)), + "code_interpreter_cost_per_session" => { + Some(Meter::Feature(Feature::CodeInterpreterSession)) + } + "google_maps_grounding_cost_per_query" => { + Some(Meter::Feature(Feature::GoogleMapsGroundingQuery)) + } + "ocr_cost_per_credit" => Some(Meter::Feature(Feature::OcrCredit)), + "ocr_cost_per_page" => Some(Meter::Feature(Feature::OcrPage)), + "browser_open_cost_per_request" => Some(Meter::Feature(Feature::BrowserOpen)), + "cost_per_additional_result" => Some(Meter::Feature(Feature::AdditionalResult)), + _ => None, + }; + if fixed.is_some() { + return fixed; + } + if let Some(unit) = base.strip_prefix("input_cost_per_") { + return input_unit(unit).map(Meter::Input); + } + if let Some(unit) = base.strip_prefix("output_cost_per_") { + return output_unit(unit).map(Meter::Output); + } + if let Some(kind) = base + .strip_prefix("cache_read_input_") + .and_then(|rest| rest.strip_suffix("_cost")) + { + return Some(Meter::CacheRead { + kind: token_kind(kind)?, + mode: CacheReadMode::Implicit, + }); + } + let rest = base.strip_prefix("cache_creation_input_")?; + let (kind, ttl) = if let Some(kind) = rest.strip_suffix("_cost_above_1hr") { + (kind, CacheTtl::OneHour) + } else if let Some(kind) = rest.strip_suffix("_cost_above_12hr") { + (kind, CacheTtl::TwelveHours) + } else { + (rest.strip_suffix("_cost")?, CacheTtl::Default) + }; + Some(Meter::CacheWrite { + kind: token_kind(kind)?, + ttl, + }) +} + +fn token_kind(unit: &str) -> Option { + match unit { + "token" => Some(TokenKind::Text), + "audio_token" => Some(TokenKind::Audio), + "image_token" => Some(TokenKind::Image), + "video_token" => Some(TokenKind::Video), + _ => None, + } +} + +fn input_unit(unit: &str) -> Option { + if let Some(kind) = token_kind(unit) { + return Some(InputUnit::Token(kind)); + } + match unit { + "character" => Some(InputUnit::Character), + "image" => Some(InputUnit::Image), + "pixel" => Some(InputUnit::Pixel), + "query" => Some(InputUnit::Query { over_results: None }), + "request" => Some(InputUnit::Request), + "second" => Some(InputUnit::Second), + "text_message" => Some(InputUnit::TextMessage), + "audio_per_second" => Some(InputUnit::AudioSecond), + "video_per_second" => Some(InputUnit::VideoSecond { over_seconds: None }), + _ => None, + } +} + +fn output_unit(unit: &str) -> Option { + if let Some(kind) = token_kind(unit) { + return Some(OutputUnit::Token(kind)); + } + match unit.split_once('_') { + Some(("image", label)) => { + return ImageSize::new(label).map(|size| OutputUnit::Image(Some(size))); + } + Some(("second", label)) => { + return Resolution::new(label).map(|resolution| OutputUnit::Second(Some(resolution))); + } + _ => {} + } + match unit { + "reasoning_token" => Some(OutputUnit::ReasoningToken), + "character" => Some(OutputUnit::Character), + "image" => Some(OutputUnit::Image(None)), + "pixel" => Some(OutputUnit::Pixel), + "second" => Some(OutputUnit::Second(None)), + "video_per_second" => Some(OutputUnit::VideoSecond), + _ => None, + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + const TEXT_INPUT: Meter = Meter::Input(InputUnit::Token(TokenKind::Text)); + const TEXT_OUTPUT: Meter = Meter::Output(OutputUnit::Token(TokenKind::Text)); + + fn key(meter: Meter, band: Band, tier: Option) -> FlatKey { + FlatKey { meter, band, tier } + } + + #[rstest] + #[case::input_token("input_cost_per_token", key(TEXT_INPUT, Band::Base, None))] + #[case::band_and_tier( + "output_cost_per_token_above_272k_tokens_priority", + key( + TEXT_OUTPUT, + Band::over_prompt_tokens(272_000), + Some(TierSuffix::Service(ServiceTier::Priority)) + ) + )] + #[case::band_without_k( + "input_cost_per_token_above_128_tokens", + key(TEXT_INPUT, Band::over_prompt_tokens(128), None) + )] + #[case::zero_band_is_base( + "input_cost_per_token_above_0k_tokens", + key(TEXT_INPUT, Band::Base, None) + )] + #[case::batch( + "input_cost_per_token_batches", + key(TEXT_INPUT, Band::Base, Some(TierSuffix::Batches)) + )] + #[case::reasoning( + "output_cost_per_reasoning_token", + key(Meter::Output(OutputUnit::ReasoningToken), Band::Base, None) + )] + #[case::cache_read_audio( + "cache_read_input_audio_token_cost", + key(Meter::CacheRead { kind: TokenKind::Audio, mode: CacheReadMode::Implicit }, Band::Base, None) + )] + #[case::explicit_cache_read( + "explicit_cache_read_input_token_cost", + key(Meter::CacheRead { kind: TokenKind::Text, mode: CacheReadMode::Explicit }, Band::Base, None) + )] + #[case::cache_write_one_hour_band( + "cache_creation_input_token_cost_above_1hr_above_200k_tokens", + key( + Meter::CacheWrite { kind: TokenKind::Text, ttl: CacheTtl::OneHour }, + Band::over_prompt_tokens(200_000), + None + ) + )] + #[case::cache_write_twelve_hours_priority( + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority", + key( + Meter::CacheWrite { kind: TokenKind::Text, ttl: CacheTtl::TwelveHours }, + Band::over_prompt_tokens(272_000), + Some(TierSuffix::Service(ServiceTier::Priority)) + ) + )] + #[case::image_label( + "output_cost_per_image_0.5K", + key( + Meter::Output(OutputUnit::Image(ImageSize::new("0.5K"))), + Band::Base, + None + ) + )] + #[case::image_token_is_a_token( + "output_cost_per_image_token", + key(Meter::Output(OutputUnit::Token(TokenKind::Image)), Band::Base, None) + )] + #[case::resolution( + "output_cost_per_second_720p", + key( + Meter::Output(OutputUnit::Second(Resolution::new("720p"))), + Band::Base, + None + ) + )] + #[case::video_interval( + "input_cost_per_video_per_second_above_8s_interval", + key(Meter::Input(InputUnit::VideoSecond { over_seconds: NonZeroU64::new(8) }), Band::Base, None) + )] + #[case::feature_batch( + "ocr_cost_per_page_batches", + key( + Meter::Feature(Feature::OcrPage), + Band::Base, + Some(TierSuffix::Batches) + ) + )] + #[case::wall_clock("cost_per_second", key(Meter::WallClockSecond, Band::Base, None))] + #[case::browser_open( + "browser_open_cost_per_request", + key(Meter::Feature(Feature::BrowserOpen), Band::Base, None) + )] + #[case::additional_result( + "cost_per_additional_result", + key(Meter::Feature(Feature::AdditionalResult), Band::Base, None) + )] + fn parses_every_key_shape(#[case] text: &str, #[case] expected: FlatKey) { + assert_eq!(FlatKey::parse(text), Some(expected)); + } + + #[rstest] + #[case::typo("input_cost_per_tokne")] + #[case::unused_meter("cost_per_token")] + #[case::other_currency("input_dbu_cost_per_token")] + #[case::cache_hit_alias("input_cost_per_token_cache_hit")] + #[case::non_numeric_band("input_cost_per_token_above_xk_tokens")] + #[case::interval_on_a_token_meter("input_cost_per_token_above_8s_interval")] + #[case::read_has_no_ttl("cache_read_input_token_cost_above_1hr")] + #[case::cache_unit_must_be_token("cache_read_input_character_cost")] + #[case::bad_image_size("output_cost_per_image_big")] + #[case::bad_resolution("output_cost_per_second_hd")] + #[case::unknown_feature("telepathy_cost_per_page")] + #[case::empty("")] + fn rejects_keys_outside_the_grammar(#[case] text: &str) { + assert_eq!(FlatKey::parse(text), None); + } + + #[rstest] + #[case::plain("input_cost_per_token", (Channel::Online, When::default()))] + #[case::service_tier( + "input_cost_per_token_flex", + (Channel::Online, When { tier: ServiceTier::Flex, band: Band::Base }) + )] + #[case::batch( + "input_cost_per_token_above_200k_tokens_batches", + (Channel::Batch, When { tier: ServiceTier::Standard, band: Band::over_prompt_tokens(200_000) }) + )] + #[case::decisions("input_cost_per_token_decisions", (Channel::Decisions, When::default()))] + fn batch_and_decisions_pick_a_channel_not_a_tier( + #[case] text: &str, + #[case] expected: (Channel, When), + ) { + assert_eq!(FlatKey::parse(text).unwrap().channel_and_when(), expected); + } +} diff --git a/litellm-rust/crates/model-catalog/src/wire/lax.rs b/litellm-rust/crates/model-catalog/src/wire/lax.rs new file mode 100644 index 00000000000..d514a6a82c1 --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/wire/lax.rs @@ -0,0 +1,128 @@ +use std::fmt; + +use serde::Deserializer; +use serde::de::{self, Visitor}; +use serde_with::DeserializeAs; + +pub(crate) struct LaxI64; + +pub(crate) struct Count; + +impl<'de> DeserializeAs<'de, i64> for LaxI64 { + fn deserialize_as>(deserializer: D) -> Result { + deserializer.deserialize_any(Self) + } +} + +impl Visitor<'_> for LaxI64 { + type Value = i64; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("an integer in the i64 range") + } + + fn visit_i64(self, value: i64) -> Result { + Ok(value) + } + + fn visit_u64(self, value: u64) -> Result { + i64::try_from(value).map_err(E::custom) + } + + fn visit_f64(self, value: f64) -> Result { + integral_float(value).ok_or_else(|| E::custom("expected an integer in the i64 range")) + } + + fn visit_str(self, value: &str) -> Result { + integer_string(value.trim()) + .ok_or_else(|| E::custom("expected an integer in the i64 range")) + } + + fn visit_bool(self, value: bool) -> Result { + Ok(i64::from(value)) + } +} + +impl<'de> DeserializeAs<'de, u64> for Count { + fn deserialize_as>(deserializer: D) -> Result { + let count = LaxI64::deserialize_as(deserializer)?; + u64::try_from(count).map_err(|_| { + de::Error::invalid_value(de::Unexpected::Signed(count), &"a count of zero or more") + }) + } +} + +fn integer_string(value: &str) -> Option { + let integer = match value.split_once('.') { + Some((integer, fraction)) => { + if fraction.is_empty() || !fraction.bytes().all(|byte| byte == b'0') { + return None; + } + integer + } + None => value, + }; + if integer.starts_with('_') || integer.ends_with('_') || integer.contains("__") { + return None; + } + let digits = integer.strip_prefix(['+', '-']).unwrap_or(integer); + if digits.is_empty() + || digits.starts_with('_') + || !digits + .bytes() + .all(|byte| byte.is_ascii_digit() || byte == b'_') + { + return None; + } + integer.replace('_', "").parse().ok() +} + +fn integral_float(value: f64) -> Option { + (value.is_finite() + && value.fract() == 0.0 + && value >= i64::MIN as f64 + && value < -(i64::MIN as f64)) + .then_some(value as i64) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::{Value, json}; + use serde_with::de::DeserializeAsWrap; + + use super::*; + + fn count(value: Value) -> Option { + DeserializeAsWrap::::deserialize(value) + .ok() + .map(DeserializeAsWrap::into_inner) + } + + use serde::Deserialize; + + #[rstest] + #[case::integer(json!(128000), Some(128_000))] + #[case::whole_float(json!(128000.0), Some(128_000))] + #[case::string(json!(" 8192 "), Some(8192))] + #[case::separated(json!("1_000"), Some(1000))] + #[case::trailing_zeros(json!("2.000"), Some(2))] + #[case::boolean(json!(true), Some(1))] + #[case::maximum(json!(i64::MAX.to_string()), Some(i64::MAX as u64))] + #[case::negative(json!(-5), None)] + #[case::fraction(json!(2.5), None)] + #[case::precise_fraction(json!("1.0000000000000001"), None)] + #[case::exponent(json!("1e3"), None)] + #[case::missing_fraction(json!("2."), None)] + #[case::repeated_separator(json!("2__0"), None)] + #[case::above_i64(json!(u64::MAX), None)] + #[case::float_above_i64(json!(9_223_372_036_854_775_808.0), None)] + #[case::word(json!("8k"), None)] + #[case::object(json!({}), None)] + fn counts_accept_whole_numbers_however_they_are_written( + #[case] written: Value, + #[case] expected: Option, + ) { + assert_eq!(count(written), expected); + } +} diff --git a/litellm-rust/crates/model-catalog/src/wire/mod.rs b/litellm-rust/crates/model-catalog/src/wire/mod.rs new file mode 100644 index 00000000000..e8f23b7bd0d --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/wire/mod.rs @@ -0,0 +1,10 @@ +mod attributes; +mod blocks; +mod entry; +mod flat_key; +mod lax; + +pub use attributes::*; +pub use blocks::*; +pub use entry::*; +pub use flat_key::*; diff --git a/litellm-rust/crates/model-catalog/tests/catalog.rs b/litellm-rust/crates/model-catalog/tests/catalog.rs index d384dfea9a8..30c070cfb52 100644 --- a/litellm-rust/crates/model-catalog/tests/catalog.rs +++ b/litellm-rust/crates/model-catalog/tests/catalog.rs @@ -1,303 +1,405 @@ -use std::path::{Path, PathBuf}; +use std::collections::BTreeSet; -use litellm_model_catalog::{AliasIssue, Catalog, Error, IntegrityLimits, Mode, Provenance}; +use litellm_model_catalog::{ + BUNDLED_REGISTRY, Catalog, Error, ErrorKind, IntegrityLimits, Mode, Rejects, Written, +}; use rstest::{fixture, rstest}; -use serde_json::json; +use serde_json::{Map, Value, json}; const ALPHA_FIXTURE: &[u8] = br#"{ - "sample_spec":{"explanation":"example"}, + "sample_spec":{"explanation":"example","input_cost_per_token":"not checked"}, "fallback_generalizations":{"rules":[{"name":"family","pattern":"^new-","model_info":{"mode":"chat"}}]}, "Alpha":{"litellm_provider":"test","aliases":["short"],"price":0,"enabled":false, "optional":null,"unknown":{"nested":[1,{"x":true}]}} }"#; -#[fixture] -fn repo_root() -> PathBuf { - Path::new(env!("CARGO_MANIFEST_DIR")).join("../../..") -} - -#[fixture] -fn current_catalog(repo_root: PathBuf) -> Catalog { - let body = std::fs::read(repo_root.join("model_prices_and_context_window.json")).unwrap(); - Catalog::parse(&body, Provenance::default()).unwrap() -} - -#[fixture] -fn backup_catalog(repo_root: PathBuf) -> Catalog { - let body = std::fs::read(repo_root.join("litellm/model_prices_and_context_window_backup.json")) - .unwrap(); - Catalog::parse(&body, Provenance::default()).unwrap() -} - #[fixture] fn fixture_catalog() -> Catalog { - Catalog::parse( - ALPHA_FIXTURE, - Provenance { - source: Some("fixture".into()), - revision: Some("rev".into()), - etag: None, - }, - ) - .unwrap() + Catalog::parse(ALPHA_FIXTURE).unwrap().value +} + +fn loaded(document: Value) -> litellm_model_catalog::Loaded { + Catalog::parse(document.to_string().as_bytes()).unwrap() +} + +fn catalog(document: Value) -> Catalog { + loaded(document).value +} + +fn written(catalog: &Catalog) -> String { + let document: Map = catalog + .written() + .map(|(key, value)| { + let value = match value { + Written::Entry(fields) => serde_json::to_value(fields).unwrap(), + Written::Section(value) => value.clone(), + }; + (key.to_owned(), value) + }) + .collect(); + serde_json::to_string(&document).unwrap() +} + +fn report(error: &Error) -> String { + litellm_error::Report::new(error).to_string() +} + +fn error_text(document: Value) -> Vec { + loaded(document).errors.iter().map(report).collect() } #[rstest] #[case::decisions("decisions", Mode::Decisions)] #[case::legacy_evaluation("evaluation", Mode::Evaluation)] fn parses_decision_modes(#[case] mode: &str, #[case] expected: Mode) { - let source = json!({"Test":{"litellm_provider":"test","mode":mode}}).to_string(); - let catalog = Catalog::parse(source.as_bytes(), Provenance::default()).unwrap(); - + let catalog = catalog(json!({"Test": {"litellm_provider": "test", "mode": mode}})); assert_eq!( - catalog.lookup("Test").unwrap().entry.info().mode, + catalog + .lookup("Test") + .unwrap() + .entry + .model() + .unwrap() + .attributes + .mode, Some(expected) ); - assert_eq!(serde_json::to_value(expected).unwrap(), json!(mode)); } #[rstest] -#[ignore] -fn preserves_fields_and_metadata(fixture_catalog: Catalog) { - let catalog = fixture_catalog; - let entry = catalog.lookup("SHORT").unwrap(); - assert_eq!(entry.canonical_key, "Alpha"); - assert_eq!(entry.matched_key, "short"); - assert_eq!(entry.entry.field("price"), Some(&json!(0))); - assert_eq!(entry.entry.field("enabled"), Some(&json!(false))); - assert_eq!(entry.entry.field("optional"), Some(&json!(null))); - assert_eq!(entry.entry.field("missing"), None); +fn unknown_keys_are_kept_raw_and_listed_as_unrecognized(fixture_catalog: Catalog) { + let entry = fixture_catalog.lookup("Alpha").unwrap().entry; assert_eq!( - entry.entry.field("unknown"), + entry.field("unknown"), Some(&json!({"nested":[1,{"x":true}]})) ); - assert_eq!(entry.entry.field("aliases"), None); - assert_eq!(entry.entry.info().litellm_provider.as_deref(), Some("test")); assert_eq!( - catalog.sample_spec(), - Some(&json!({"explanation":"example"})) - ); - assert_eq!(catalog.fallback_rules().unwrap().len(), 1); - assert_eq!(catalog.provenance().revision.as_deref(), Some("rev")); - assert_eq!(catalog.model_count(), 1); -} - -#[rstest] -#[ignore] -fn snapshot_does_not_borrow_source() { - let mut source = ALPHA_FIXTURE.to_vec(); - let catalog = Catalog::parse(&source, Provenance::default()).unwrap(); - source.fill(b' '); - - let entry = catalog.lookup("short").unwrap(); - assert_eq!(entry.canonical_key, "Alpha"); - assert_eq!(entry.entry.field("price"), Some(&json!(0))); -} - -#[rstest] -#[case("Shared", "First")] -#[case("Second", "Second")] -#[case("shared", "Second")] -#[case("FIRST", "First")] -#[case("sHaReD", "Second")] -#[ignore] -fn alias_collisions_and_case_fallback_follow_entry_order( - #[case] lookup: &str, - #[case] expected: &str, -) { - let catalog = Catalog::parse( - br#"{ - "First":{"aliases":["Shared","Second","first"],"value":1}, - "Second":{"aliases":["Shared","sHaReD"],"value":2}, - "SHARED":{"value":3} - }"#, - Provenance::default(), - ) - .unwrap(); - assert_eq!(catalog.lookup(lookup).unwrap().canonical_key, expected); - assert_eq!(catalog.alias_count(), 3); - assert!( - catalog - .alias_issues() - .contains(&AliasIssue::CanonicalCollision { - model: "First".into(), - alias: "Second".into(), - }) - ); - assert!( - catalog - .alias_issues() - .contains(&AliasIssue::AliasCollision { - model: "Second".into(), - alias: "Shared".into(), - }) + fixture_catalog.unrecognized().collect::>(), + BTreeSet::from([ + ("Alpha", "enabled"), + ("Alpha", "price"), + ("Alpha", "unknown") + ]) ); } -#[test] -#[ignore] -fn json_entry_order_controls_alias_ownership_and_case_fallback() { - let forward = Catalog::parse( - br#"{ - "Alpha":{"aliases":["shared"]}, - "Beta":{"aliases":["shared"]}, - "Foo":{}, - "fOO":{} - }"#, - Provenance::default(), - ) - .unwrap(); - let reversed = Catalog::parse( - br#"{ - "fOO":{}, - "Foo":{}, - "Beta":{"aliases":["shared"]}, - "Alpha":{"aliases":["shared"]} - }"#, - Provenance::default(), - ) - .unwrap(); - - assert_eq!(forward.lookup("shared").unwrap().canonical_key, "Alpha"); - assert_eq!(reversed.lookup("shared").unwrap().canonical_key, "Beta"); - assert_eq!(forward.lookup("foo").unwrap().canonical_key, "fOO"); - assert_eq!(reversed.lookup("foo").unwrap().canonical_key, "Foo"); -} - -#[derive(Debug)] -enum ValidationOutcome { - Ok, - Shrunk, - BelowMinimum, - InvalidRatio, -} - #[rstest] -#[case( - IntegrityLimits { - reference_model_count: 2, - min_model_count: 1, - min_reference_ratio: 0.5, - }, - ValidationOutcome::Ok -)] -#[case( - IntegrityLimits { - reference_model_count: 3, - min_model_count: 1, - min_reference_ratio: 0.5, - }, - ValidationOutcome::Shrunk -)] -#[case( - IntegrityLimits { - reference_model_count: 0, - min_model_count: 2, - min_reference_ratio: 0.5, - }, - ValidationOutcome::BelowMinimum -)] -#[case( - IntegrityLimits { - reference_model_count: 0, - min_model_count: 0, - min_reference_ratio: f64::NAN, - }, - ValidationOutcome::InvalidRatio -)] -#[ignore] -fn integrity_uses_canonical_count_and_strict_shrink_boundary( - #[case] limits: IntegrityLimits, - #[case] expected: ValidationOutcome, -) { - let catalog = Catalog::parse( - br#"{"sample_spec":{},"fallback_generalizations":{"rules":[]},"a":{"aliases":["b","c"]}}"#, - Provenance::default(), - ) - .unwrap(); - let actual = catalog.validate(limits); - match expected { - ValidationOutcome::Ok => assert!(actual.is_ok()), - ValidationOutcome::Shrunk => { - assert!(matches!(actual, Err(Error::Shrunk { actual: 1, .. }))) - } - ValidationOutcome::BelowMinimum => { - assert!(matches!(actual, Err(Error::BelowMinimum { actual: 1, .. }))) - } - ValidationOutcome::InvalidRatio => assert!(matches!(actual, Err(Error::InvalidRatio))), - } -} - -#[derive(Debug)] -enum MalformedOutcome { - Empty, - Json, - EntryNotObject, -} - -#[rstest] -#[case::empty(b"{}", MalformedOutcome::Empty)] -#[case::invalid_json(b"{", MalformedOutcome::Json)] -#[case::entry_not_object(br#"{"a":1}"#, MalformedOutcome::EntryNotObject)] -#[case::fallback_rules_missing( - br#"{"fallback_generalizations":{},"a":{}}"#, - MalformedOutcome::Json -)] -#[ignore] -fn malformed_input_and_aliases_have_typed_outcomes( - #[case] body: &[u8], - #[case] expected: MalformedOutcome, -) { - let actual = Catalog::parse(body, Provenance::default()); - match expected { - MalformedOutcome::Empty => assert!(matches!(actual, Err(Error::Empty))), - MalformedOutcome::Json => assert!(matches!(actual, Err(Error::Json(_)))), - MalformedOutcome::EntryNotObject => { - assert!(matches!(actual, Err(Error::EntryNotObject { .. }))) - } - } -} - -#[rstest] -#[ignore] -fn invalid_aliases_are_reported_not_fatal() { - let catalog = Catalog::parse( - br#"{"a":{"aliases":"bad"},"b":{"aliases":[9,"ok"]}}"#, - Provenance::default(), - ) - .unwrap(); +fn sections_aliases_and_raw_fields_are_read_in_their_document_views(fixture_catalog: Catalog) { + let found = fixture_catalog.lookup("SHORT").unwrap(); + assert_eq!(found.canonical_key, "Alpha"); + assert_eq!(found.matched_key, "short"); + assert_eq!(found.entry.field("optional"), Some(&json!(null))); + assert_eq!(found.entry.field("aliases"), Some(&json!(["short"]))); assert_eq!( - catalog.alias_issues(), - &[ - AliasIssue::InvalidList { model: "a".into() }, - AliasIssue::InvalidName { model: "b".into() }, - ] + fixture_catalog.sample_spec(), + Some(&json!({"explanation":"example","input_cost_per_token":"not checked"})) + ); + assert_eq!(fixture_catalog.fallback_rules().unwrap()[0].name, "family"); + assert_eq!(fixture_catalog.model_count(), 1); + assert_eq!( + fixture_catalog + .served() + .map(|(key, _)| key) + .collect::>(), + ["Alpha", "short"] ); - assert_eq!(catalog.lookup("ok").unwrap().canonical_key, "b"); - assert!(catalog.lookup("missing").is_none()); } #[rstest] -#[ignore] -fn parses_current_and_packaged_catalogs_against_independent_baseline( - current_catalog: Catalog, - backup_catalog: Catalog, -) { - assert!(current_catalog.model_count() > 0); - assert!(backup_catalog.model_count() > 0); - assert!(current_catalog.sample_spec().is_some()); - assert!(backup_catalog.sample_spec().is_some()); - // Snapshot from 2026-09-23; the backup file mirrors the current file and cannot detect shrinkage. - const REFERENCE_MODEL_COUNT: usize = 4303; - current_catalog - .validate(IntegrityLimits { - reference_model_count: REFERENCE_MODEL_COUNT, - min_model_count: 50, - min_reference_ratio: 0.9, - }) - .unwrap(); - assert!(current_catalog.model_names().all(|name| { - let entry = current_catalog.lookup(name).unwrap().entry; - entry.info().litellm_provider.is_some() == entry.field("litellm_provider").is_some() +fn a_bad_entry_is_reported_by_model_and_key_and_kept_raw() { + let parsed = loaded(json!({ + "Good": {"litellm_provider": "test"}, + "Bad": {"litellm_provider": "test", "mode": "chatt", "input_cost_per_token": "free"}, + "Worse": {"mode": "chat"} })); + let catalog = parsed.value; + let problems: Vec = parsed.errors.iter().map(report).collect(); + assert_eq!(problems.len(), 3, "{problems:?}"); + assert!(problems[0].starts_with("Bad.mode: invalid value: unknown variant `chatt`")); + assert!(problems[1].starts_with("Bad.input_cost_per_token: invalid value")); + assert_eq!(problems[2], "Worse.litellm_provider: required"); + assert!(catalog.lookup("Good").unwrap().entry.model().is_some()); + assert!(catalog.lookup("Bad").unwrap().entry.model().is_none()); + assert_eq!(catalog.model_count(), 3); + assert!( + parsed + .errors + .iter() + .find(|error| error.rejects() == Rejects::Document) + .is_none() + ); +} + +#[rstest] +#[case::alias_names_a_model( + json!({"First": {"litellm_provider": "t", "aliases": ["Second"]}, "Second": {"litellm_provider": "t"}}), + r#"First: alias "Second" is already a model name"# +)] +#[case::alias_claimed_twice( + json!({"First": {"litellm_provider": "t", "aliases": ["shared"]}, "Second": {"litellm_provider": "t", "aliases": ["shared"]}}), + r#"Second: alias "shared" already belongs to First"# +)] +#[case::aliases_not_a_list( + json!({"a": {"litellm_provider": "t", "aliases": "bad"}}), + "a.aliases: invalid value: invalid type: string \"bad\", expected a sequence" +)] +#[case::entry_not_object( + json!({"a": 1, "valid": {"litellm_provider": "test"}}), + "a: must be a JSON object" +)] +#[case::fallback_rules_missing( + json!({"fallback_generalizations": {}, "a": {"litellm_provider": "t"}}), + "fallback_generalizations: invalid value: missing field `rules`" +)] +#[case::unknown_timezone( + json!({"a": {"litellm_provider": "t", "off_peak_pricing": { + "windows": [{"hours_utc": "00:00-01:00"}], + "weekday_timezone": "Not/AZone", + "input_cost_per_token": 1e-6 + }}}), + r#"a.off_peak_pricing.weekday_timezone: "Not/AZone" is not an IANA time zone"# +)] +fn document_level_problems_are_reported(#[case] document: Value, #[case] expected: &str) { + assert_eq!(error_text(document), [expected]); +} + +#[rstest] +fn escaped_row_paths_and_row_validation_share_the_parse_rules() { + let document = json!({ + "open/router~v1": {"litellm_provider": "test", "input_cost_per_token": -1} + }); + let parsed = loaded(document); + let fields = parsed + .value + .lookup("open/router~v1") + .unwrap() + .entry + .fields(); + let row_errors = Catalog::row_errors("open/router~v1", fields); + assert_eq!( + row_errors + .iter() + .map(|error| (error.path(), report(error))) + .collect::>(), + parsed + .errors + .iter() + .map(|error| (error.path(), report(error))) + .collect::>() + ); + assert_eq!( + row_errors[0].path(), + "/open~1router~0v1/input_cost_per_token" + ); + assert!(row_errors[0].kind() == ErrorKind::Invalid); +} + +#[rstest] +fn malformed_fallback_is_collected_while_models_are_read() { + let parsed = loaded(json!({ + "fallback_generalizations": {"rules": "bad"}, + "a": {"litellm_provider": "test"} + })); + assert_eq!(parsed.value.model_count(), 1); + assert_eq!(parsed.errors.len(), 1); + assert_eq!(parsed.errors[0].path(), "/fallback_generalizations"); +} + +#[rstest] +fn rejected_rows_and_their_aliases_are_removed_without_disturbing_written_order() { + let body = r#"{"before":{"litellm_provider":"test","aliases":["before-alias"]},"bad":{"litellm_provider":"test","aliases":["bad-alias"],"mode":"chatt"},"not-object":4,"sample_spec":{"keep":true},"after":{"litellm_provider":"test"}}"#; + let parsed = Catalog::parse(body.as_bytes()).unwrap(); + let (catalog, removed) = parsed.value.without_rejected(&parsed.errors); + + assert_eq!(removed.len(), 2); + assert_eq!( + catalog.lookup("before-alias").unwrap().canonical_key, + "before" + ); + assert_eq!( + catalog.lookup("BEFORE-ALIAS").unwrap().matched_key, + "before-alias" + ); + assert!(catalog.lookup("bad").is_none()); + assert!(catalog.lookup("BAD").is_none()); + assert!(catalog.lookup("Bad-Alias").is_none()); + assert_eq!(catalog.model_count(), 2); + assert_eq!( + written(&catalog), + r#"{"before":{"litellm_provider":"test","aliases":["before-alias"]},"sample_spec":{"keep":true},"after":{"litellm_provider":"test"}}"# + ); +} + +#[rstest] +fn invalid_empty_key_rows_are_removed() { + let parsed = loaded(json!({ + "": {"litellm_provider": "test", "mode": "chatt"}, + "good": {"litellm_provider": "test"} + })); + assert_eq!(parsed.errors[0].path(), "//mode"); + + let (catalog, removed) = parsed.value.without_rejected(&parsed.errors); + assert_eq!(catalog.model_count(), 1); + assert!(catalog.lookup("").is_none()); + assert!(catalog.lookup("good").is_some()); + assert_eq!(removed.len(), 1); +} + +#[rstest] +fn timezone_errors_reject_providerless_rows_without_rejecting_other_providerless_rows() { + let parsed = loaded(json!({ + "good": {"mode": "chat"}, + "bad": {"off_peak_pricing": { + "windows": [{"hours_utc": "00:00-01:00"}], + "weekday_timezone": "Not/AZone", + "input_cost_per_token": 1e-6 + }} + })); + assert_eq!(parsed.errors.len(), 3); + assert!( + parsed + .errors + .iter() + .any(|error| error.kind() == ErrorKind::UnknownTimezone) + ); + + let (catalog, removed) = parsed.value.without_rejected(&parsed.errors); + assert_eq!(catalog.model_count(), 1); + assert!(catalog.lookup("good").is_some()); + assert!(catalog.lookup("bad").is_none()); + assert_eq!(removed.len(), 2); +} + +#[rstest] +fn bundled_per_row_errors_match_document_errors() { + let parsed = Catalog::parse(BUNDLED_REGISTRY).unwrap(); + let row_errors: Vec<_> = parsed + .value + .written() + .flat_map(|(key, written)| match written { + Written::Entry(fields) => Catalog::row_errors(key, fields), + Written::Section(_) => Vec::new(), + }) + .map(|error| (error.path().to_owned(), report(&error))) + .collect(); + let document_errors: Vec<_> = parsed + .errors + .iter() + .map(|error| (error.path().to_owned(), report(error))) + .collect(); + + assert_eq!(row_errors, document_errors); +} + +#[rstest] +fn three_hundred_rows_filter_one_row_with_two_invalid_prices() { + let valid_rows: Map = (0..299) + .map(|index| { + ( + format!("valid-{index:03}"), + json!({"litellm_provider": "test"}), + ) + }) + .collect(); + let mut document = valid_rows; + document.insert( + "bad".to_owned(), + json!({ + "litellm_provider": "test", + "input_cost_per_token": -1, + "output_cost_per_token": -1 + }), + ); + let parsed = Catalog::parse(serde_json::to_vec(&document).unwrap().as_slice()).unwrap(); + assert_eq!(parsed.value.model_count(), 300); + assert_eq!(parsed.errors.len(), 2); + + let (catalog, removed) = parsed.value.without_rejected(&parsed.errors); + assert_eq!(catalog.model_count(), 299); + assert_eq!(removed.len(), 2); +} + +#[rstest] +fn integrity_limits_count_canonical_models_not_aliases() { + let fetched = catalog(json!({"a": {"litellm_provider": "t", "aliases": ["b", "c"]}})); + assert!( + fetched + .check_size(2, IntegrityLimits::new(1, 0.5).unwrap()) + .is_ok() + ); + assert_eq!( + fetched + .check_size(3, IntegrityLimits::new(1, 0.5).unwrap()) + .unwrap_err() + .kind(), + ErrorKind::Shrunk + ); + assert_eq!( + fetched + .check_size(2, IntegrityLimits::new(2, 0.5).unwrap()) + .unwrap_err() + .kind(), + ErrorKind::BelowMinimum + ); +} + +#[rstest] +#[case::empty(b"{}", "empty")] +#[case::only_sections(br#"{"sample_spec":{}}"#, "empty")] +#[case::invalid_json(b"{", "json")] +#[case::array_root_is_json_error(b"[]", "json")] +fn malformed_documents_are_rejected_whole(#[case] body: &[u8], #[case] kind: &str) { + let error = Catalog::parse(body).unwrap_err(); + assert_eq!(error.rejects(), Rejects::Document); + assert_eq!(<&str>::from(error.kind()), kind); +} + +#[rstest] +fn written_is_the_document_in_order_including_sections_and_bad_entries() { + let body = concat!( + r#"{"b":{"litellm_provider":"t","z":1,"a":null},"sample_spec":{"x":"y"},"n":4,"#, + r#""fallback_generalizations":{"rules":"bad"},"a":{"mode":"chatt","price":1.0}}"# + ); + let catalog = Catalog::parse(body.as_bytes()).unwrap().value; + assert_eq!(written(&catalog), body); + assert_eq!(catalog.model_count(), 2); +} + +#[rstest] +fn the_bundled_registry_parses_with_no_problems_and_every_key_understood() { + let loaded = Catalog::parse(BUNDLED_REGISTRY).unwrap(); + let catalog = loaded.value; + assert!(loaded.errors.is_empty(), "{:#?}", loaded.errors); + let unrecognized: Vec<_> = catalog.unrecognized().collect(); + assert!(unrecognized.is_empty(), "{unrecognized:#?}"); + assert!( + catalog + .fallback_rules() + .is_some_and(|rules| !rules.is_empty()) + ); + const REFERENCE_MODEL_COUNT: usize = 4303; + assert!(catalog.model_count() as f64 >= REFERENCE_MODEL_COUNT as f64 * 0.9); +} + +#[rstest] +fn the_bundled_registry_is_written_back_as_it_was_read() { + let document: Value = serde_json::from_slice(BUNDLED_REGISTRY).unwrap(); + assert_eq!( + written(&Catalog::parse(BUNDLED_REGISTRY).unwrap().value), + serde_json::to_string(&document).unwrap() + ); +} + +#[rstest] +fn a_bad_value_keeps_its_json_error_as_the_source_and_names_its_key_and_field() { + let parsed = + loaded(json!({"open/router": {"litellm_provider": "t", "input_cost_per_token": "free"}})); + let error = &parsed.errors[0]; + let cause = litellm_error::find::(error) + .expect("the JSON error is in the chain") + .to_string(); + assert_eq!(report(error).matches(&cause).count(), 1); + assert!(!error.to_string().contains(&cause)); + assert_eq!(error.key().as_deref(), Some("open/router")); + assert_eq!(error.field().as_deref(), Some("input_cost_per_token")); } diff --git a/litellm-rust/crates/model-catalog/tests/generate_price_lookup_reference.py b/litellm-rust/crates/model-catalog/tests/generate_price_lookup_reference.py new file mode 100644 index 00000000000..e49b635ff4a --- /dev/null +++ b/litellm-rust/crates/model-catalog/tests/generate_price_lookup_reference.py @@ -0,0 +1,91 @@ +import itertools +import json +import re +import subprocess +from pathlib import Path +from typing import Final + +from litellm.litellm_core_utils.llm_cost_calc.utils import _get_token_base_cost +from litellm.types.utils import Usage + +REPO_ROOT: Final = Path(__file__).resolve().parents[4] +TOKEN_RATE_KEY: Final = re.compile( + r"^(input_cost_per_token|output_cost_per_token|cache_read_input_token_cost" + r"|cache_creation_input_token_cost(_above_1hr)?)" + r"(_above_(\d+)(k?)_tokens)?(_(priority|flex|ultrafast|balanced|batches|decisions))?$" +) +SERVICE_TIERS: Final = ("priority", "flex", "ultrafast", "balanced") +LARGEST_PROMPT: Final = 10_000_000 +INCLUSIVE_PROVIDERS: Final = frozenset({"xai"}) + + +def token_rates(entry: dict[str, object]) -> dict[str, object]: + return { + key: value + for key, value in entry.items() + if value is not None and (key == "tiered_pricing" or TOKEN_RATE_KEY.match(key)) + } + + +def band_floors(info: dict[str, object]) -> frozenset[int]: + matches: Final = (TOKEN_RATE_KEY.match(key) for key in info) + flat: Final = frozenset( + int(match.group(4)) * (1000 if match.group(5) else 1) + for match in matches + if match is not None and match.group(4) is not None + ) + rows: Final = info.get("tiered_pricing") + tiers: Final = ( + frozenset(int(row["range"][0]) for row in rows if isinstance(row, dict) and "range" in row) + if isinstance(rows, list) + else frozenset() + ) + return flat | tiers + + +def tiers_in(info: dict[str, object]) -> tuple[str | None, ...]: + return (None, *(tier for tier in SERVICE_TIERS if any(key.endswith(f"_{tier}") for key in info))) + + +def prompts(info: dict[str, object]) -> tuple[int, ...]: + floors: Final = band_floors(info) + return tuple(sorted(({1, LARGEST_PROMPT} | floors | {floor + 1 for floor in floors}) - {0})) + + +def check(info: dict[str, object], tier: str | None, prompt: int, inclusive: bool) -> list[object]: + usage: Final = Usage(prompt_tokens=prompt, completion_tokens=1, total_tokens=prompt + 1) + rates: Final = _get_token_base_cost(info, usage, service_tier=tier, threshold_is_inclusive=inclusive) + one_hour: Final = rates.cache_creation_rates_by_ttl.get("1hr", rates.cache_creation_rate) + return [tier, prompt, [rates.input_rate, rates.output_rate, rates.cache_creation_rate, one_hour, rates.cache_read_rate]] + + +def shape(info: dict[str, object], inclusive: bool) -> str: + rows: Final = info.get("tiered_pricing") + row_keys: Final = [sorted(row) for row in rows if isinstance(row, dict)] if isinstance(rows, list) else [] + return json.dumps([sorted(info), row_keys, inclusive]) + + +def checks(info: dict[str, object], inclusive: bool) -> list[list[object]]: + return [check(info, tier, prompt, inclusive) for tier, prompt in itertools.product(tiers_in(info), prompts(info))] + + +def main() -> None: + registry: Final = json.loads((REPO_ROOT / "model_prices_and_context_window.json").read_text()) + entries: Final = tuple( + (token_rates(entry), entry.get("litellm_provider") in INCLUSIVE_PROVIDERS) + for name, entry in registry.items() + if name not in ("sample_spec", "fallback_generalizations") and isinstance(entry, dict) + ) + by_shape: Final = {shape(info, inclusive): (info, inclusive) for info, inclusive in reversed(entries) if info} + models: Final = tuple( + json.dumps({"info": info, "inclusive": inclusive, "checks": checks(info, inclusive)}) + for info, inclusive in by_shape.values() + ) + revision: Final = subprocess.check_output(("git", "rev-parse", "HEAD"), text=True, cwd=REPO_ROOT).strip() + Path(__file__).with_name("price_lookup_reference.json").write_text( + f'{{"python_reference_commit": "{revision}", "models": [\n' + ",\n".join(models) + "\n]}\n" + ) + + +if __name__ == "__main__": + main() diff --git a/litellm-rust/crates/model-catalog/tests/price_lookup.rs b/litellm-rust/crates/model-catalog/tests/price_lookup.rs new file mode 100644 index 00000000000..cc16e9f23e9 --- /dev/null +++ b/litellm-rust/crates/model-catalog/tests/price_lookup.rs @@ -0,0 +1,89 @@ +use litellm_cost::money::UnitPrice; +use litellm_cost::price::{ + CacheReadMode, CacheTtl, InputUnit, Meter, OutputUnit, Prices, ServiceTier, TokenKind, +}; +use litellm_model_catalog::wire::{EntryWire, Fields}; +use rstest::rstest; +use serde::Deserialize; + +#[derive(Deserialize)] +struct Reference { + models: Vec, +} + +#[derive(Deserialize)] +struct ReferenceModel { + info: Fields, + inclusive: bool, + checks: Vec<(Option, u64, [f64; 5])>, +} + +const METERS: [Meter; 5] = [ + Meter::Input(InputUnit::Token(TokenKind::Text)), + Meter::Output(OutputUnit::Token(TokenKind::Text)), + Meter::CacheWrite { + kind: TokenKind::Text, + ttl: CacheTtl::Default, + }, + Meter::CacheWrite { + kind: TokenKind::Text, + ttl: CacheTtl::OneHour, + }, + Meter::CacheRead { + kind: TokenKind::Text, + mode: CacheReadMode::Implicit, + }, +]; + +fn service_tier(name: Option<&str>) -> ServiceTier { + match name { + None => ServiceTier::Standard, + Some("priority") => ServiceTier::Priority, + Some("flex") => ServiceTier::Flex, + Some("ultrafast") => ServiceTier::Ultrafast, + Some("balanced") => ServiceTier::Balanced, + Some(other) => panic!("the reference uses an unknown service tier {other}"), + } +} + +fn mismatch( + prices: &Prices, + (tier_name, prompt, python): &(Option, u64, [f64; 5]), +) -> Option { + let tier = service_tier(tier_name.as_deref()); + let band = prices.band(tier, *prompt); + let rust = METERS + .each_ref() + .map(|meter| prices.rate(meter, tier, band).map_or(0.0, UnitPrice::get)); + (rust != *python) + .then(|| format!("tier={tier_name:?} prompt={prompt}: rust {rust:?}, python {python:?}")) +} + +#[rstest] +fn band_and_rate_match_python_token_base_cost() { + let reference: Reference = + serde_json::from_str(include_str!("price_lookup_reference.json")).unwrap(); + + let mismatches: Vec = reference + .models + .iter() + .flat_map(|model| { + let provider = if model.inclusive { "xai" } else { "openai" }; + let fields: Fields = model + .info + .clone() + .into_iter() + .chain([("litellm_provider".to_owned(), provider.into())]) + .collect(); + let prices = Prices::try_from(&EntryWire::read(&fields)).unwrap(); + model + .checks + .iter() + .filter_map(move |check| mismatch(&prices, check)) + .map(move |found| format!("{:?} {found}", model.info)) + }) + .collect(); + + assert!(reference.models.len() > 50); + assert!(mismatches.is_empty(), "{mismatches:#?}"); +} diff --git a/litellm-rust/crates/model-catalog/tests/price_lookup_reference.json b/litellm-rust/crates/model-catalog/tests/price_lookup_reference.json new file mode 100644 index 00000000000..4b87b696723 --- /dev/null +++ b/litellm-rust/crates/model-catalog/tests/price_lookup_reference.json @@ -0,0 +1,89 @@ +{"python_reference_commit": "c7c895db456eda349a7be9661b4b21f578f9b589", "models": [ +{"info": {"cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, "cache_read_input_token_cost": 1e-06, "cache_read_input_token_cost_above_272k_tokens": 2e-06, "input_cost_per_token": 1e-05, "input_cost_per_token_above_272k_tokens": 2e-05, "output_cost_per_token": 5e-05, "output_cost_per_token_above_272k_tokens": 7.5e-05}, "inclusive": false, "checks": [[null, 1, [1e-05, 5e-05, 1.25e-05, 1.25e-05, 1e-06]], [null, 272000, [1e-05, 5e-05, 1.25e-05, 1.25e-05, 1e-06]], [null, 272001, [2e-05, 7.5e-05, 2.5e-05, 2.5e-05, 2e-06]], [null, 10000000, [2e-05, 7.5e-05, 2.5e-05, 2.5e-05, 2e-06]]]}, +{"info": {"cache_read_input_token_cost": 1e-07, "input_cost_per_token": 1e-06, "output_cost_per_token": 3e-06}, "inclusive": false, "checks": [[null, 1, [1e-06, 3e-06, 1e-06, 1e-06, 1e-07]], [null, 10000000, [1e-06, 3e-06, 1e-06, 1e-06, 1e-07]]]}, +{"info": {"input_cost_per_token": 1.25e-05, "output_cost_per_token": 1.25e-05}, "inclusive": false, "checks": [[null, 1, [1.25e-05, 1.25e-05, 1.25e-05, 1.25e-05, 1.25e-05]], [null, 10000000, [1.25e-05, 1.25e-05, 1.25e-05, 1.25e-05, 1.25e-05]]]}, +{"info": {"cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-05, "cache_creation_input_token_cost_flex": 1.25e-06, "cache_creation_input_token_cost_priority": 5e-06, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_272k_tokens": 4e-07, "cache_read_input_token_cost_above_272k_tokens_flex": 2e-07, "cache_read_input_token_cost_above_272k_tokens_priority": 8e-07, "cache_read_input_token_cost_batches": 1e-07, "cache_read_input_token_cost_above_272k_tokens_batches": 2e-07, "cache_creation_input_token_cost_batches": 1.25e-06, "cache_creation_input_token_cost_above_272k_tokens_batches": 2.5e-06, "cache_read_input_token_cost_flex": 1e-07, "cache_read_input_token_cost_priority": 4e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_272k_tokens": 4e-06, "input_cost_per_token_above_272k_tokens_flex": 2e-06, "input_cost_per_token_above_272k_tokens_priority": 8e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_above_272k_tokens_batches": 2e-06, "input_cost_per_token_flex": 1e-06, "input_cost_per_token_priority": 4e-06, "output_cost_per_token": 1e-05, "output_cost_per_token_above_272k_tokens": 1.5e-05, "output_cost_per_token_above_272k_tokens_flex": 7.5e-06, "output_cost_per_token_above_272k_tokens_priority": 3e-05, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_above_272k_tokens_batches": 7.5e-06, "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05}, "inclusive": false, "checks": [[null, 1, [2e-06, 1e-05, 2.5e-06, 2.5e-06, 2e-07]], [null, 272000, [2e-06, 1e-05, 2.5e-06, 2.5e-06, 2e-07]], [null, 272001, [4e-06, 1.5e-05, 5e-06, 5e-06, 4e-07]], [null, 10000000, [4e-06, 1.5e-05, 5e-06, 5e-06, 4e-07]], ["priority", 1, [4e-06, 2e-05, 5e-06, 5e-06, 4e-07]], ["priority", 272000, [4e-06, 2e-05, 5e-06, 5e-06, 4e-07]], ["priority", 272001, [8e-06, 3e-05, 1e-05, 1e-05, 8e-07]], ["priority", 10000000, [8e-06, 3e-05, 1e-05, 1e-05, 8e-07]], ["flex", 1, [1e-06, 5e-06, 1.25e-06, 1.25e-06, 1e-07]], ["flex", 272000, [1e-06, 5e-06, 1.25e-06, 1.25e-06, 1e-07]], ["flex", 272001, [2e-06, 7.5e-06, 2.5e-06, 2.5e-06, 2e-07]], ["flex", 10000000, [2e-06, 7.5e-06, 2.5e-06, 2.5e-06, 2e-07]]]}, +{"info": {"cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_creation_input_token_cost_batches": 6.25e-07, "cache_read_input_token_cost": 1e-07, "cache_read_input_token_cost_batches": 5e-08, "input_cost_per_token": 1e-06, "input_cost_per_token_batches": 5e-07, "output_cost_per_token": 5e-06, "output_cost_per_token_batches": 2.5e-06}, "inclusive": false, "checks": [[null, 1, [1e-06, 5e-06, 1.25e-06, 2e-06, 1e-07]], [null, 10000000, [1e-06, 5e-06, 1.25e-06, 2e-06, 1e-07]]]}, +{"info": {"input_cost_per_token_above_100k_tokens_batches": 2.5e-07, "output_cost_per_token_above_100k_tokens_batches": 1.25e-06, "cache_creation_input_token_cost_above_100k_tokens_batches": 3.125e-07, "cache_read_input_token_cost_above_100k_tokens_batches": 2.5e-08, "cache_creation_input_token_cost": 1.25e-07, "cache_creation_input_token_cost_above_1hr": 2e-07, "cache_creation_input_token_cost_batches": 6.25e-08, "cache_read_input_token_cost": 1e-08, "cache_read_input_token_cost_batches": 5e-09, "input_cost_per_token": 1e-07, "input_cost_per_token_batches": 5e-08, "output_cost_per_token": 5e-07, "output_cost_per_token_batches": 2.5e-07, "input_cost_per_token_above_100k_tokens": 5e-07, "output_cost_per_token_above_100k_tokens": 2.5e-06, "cache_creation_input_token_cost_above_100k_tokens": 6.25e-07, "cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1e-06, "cache_read_input_token_cost_above_100k_tokens": 5e-08}, "inclusive": false, "checks": [[null, 1, [1e-07, 5e-07, 1.25e-07, 2e-07, 1e-08]], [null, 100000, [1e-07, 5e-07, 1.25e-07, 2e-07, 1e-08]], [null, 100001, [5e-07, 2.5e-06, 6.25e-07, 1e-06, 5e-08]], [null, 10000000, [5e-07, 2.5e-06, 6.25e-07, 1e-06, 5e-08]]]}, +{"info": {"input_cost_per_token": 1.0000000000000001e-07, "output_cost_per_token": 5e-07, "cache_read_input_token_cost": 1e-08, "cache_creation_input_token_cost": 1.25e-07, "tiered_pricing": [{"range": [0, 272000], "input_cost_per_token": 1.0000000000000001e-07, "output_cost_per_token": 5e-07, "cache_read_input_token_cost": 1e-08, "cache_creation_input_token_cost": 1.25e-07}, {"range": [272000, 1050000], "input_cost_per_token": 2.0000000000000002e-07, "output_cost_per_token": 7.5e-07, "cache_read_input_token_cost": 2e-08, "cache_creation_input_token_cost": 2.5e-07}]}, "inclusive": false, "checks": [[null, 1, [1.0000000000000001e-07, 5e-07, 1.25e-07, 1.25e-07, 1e-08]], [null, 272000, [1.0000000000000001e-07, 5e-07, 1.25e-07, 1.25e-07, 1e-08]], [null, 272001, [2.0000000000000002e-07, 7.5e-07, 2.5e-07, 2.5e-07, 2e-08]], [null, 10000000, [2.0000000000000002e-07, 7.5e-07, 2.5e-07, 2.5e-07, 2e-08]]]}, +{"info": {"cache_creation_input_token_cost": 1e-06, "cache_read_input_token_cost": 8e-08, "input_cost_per_token": 8e-07, "output_cost_per_token": 4e-06}, "inclusive": false, "checks": [[null, 1, [8e-07, 4e-06, 1e-06, 1e-06, 8e-08]], [null, 10000000, [8e-07, 4e-06, 1e-06, 1e-06, 8e-08]]]}, +{"info": {"input_cost_per_token": 2e-06, "output_cost_per_token": 6e-06, "cache_read_input_token_cost": 5e-07, "tiered_pricing": [{"range": [0, 200000], "input_cost_per_token": 2e-06, "output_cost_per_token": 6e-06, "cache_read_input_token_cost": 5e-07}, {"range": [200000, 500000], "input_cost_per_token": 4e-06, "output_cost_per_token": 1.2e-05, "cache_read_input_token_cost": 1e-06}]}, "inclusive": false, "checks": [[null, 1, [2e-06, 6e-06, 2e-06, 2e-06, 5e-07]], [null, 200000, [2e-06, 6e-06, 2e-06, 2e-06, 5e-07]], [null, 200001, [4e-06, 1.2e-05, 4e-06, 4e-06, 1e-06]], [null, 10000000, [4e-06, 1.2e-05, 4e-06, 4e-06, 1e-06]]]}, +{"info": {"cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_priority": 1.5e-07, "input_cost_per_token": 7.5e-07, "input_cost_per_token_priority": 1.5e-06, "output_cost_per_token": 4.5e-06, "output_cost_per_token_priority": 9e-06}, "inclusive": false, "checks": [[null, 1, [7.5e-07, 4.5e-06, 7.5e-07, 7.5e-07, 7.5e-08]], [null, 10000000, [7.5e-07, 4.5e-06, 7.5e-07, 7.5e-07, 7.5e-08]], ["priority", 1, [1.5e-06, 9e-06, 1.5e-06, 1.5e-06, 1.5e-07]], ["priority", 10000000, [1.5e-06, 9e-06, 1.5e-06, 1.5e-06, 1.5e-07]]]}, +{"info": {"cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06}, "inclusive": false, "checks": [[null, 1, [1.25e-06, 2.5e-06, 1.25e-06, 1.25e-06, 2e-07]], [null, 200000, [1.25e-06, 2.5e-06, 1.25e-06, 1.25e-06, 2e-07]], [null, 200001, [2.5e-06, 5e-06, 2.5e-06, 2.5e-06, 4e-07]], [null, 10000000, [2.5e-06, 5e-06, 2.5e-06, 2.5e-06, 4e-07]]]}, +{"info": {"cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 1e-06, "output_cost_per_token": 5e-06}, "inclusive": false, "checks": [[null, 1, [1e-06, 5e-06, 1.25e-06, 2e-06, 1e-07]], [null, 10000000, [1e-06, 5e-06, 1.25e-06, 2e-06, 1e-07]]]}, +{"info": {"cache_read_input_token_cost": 8.25e-08, "cache_read_input_token_cost_flex": 4.125e-08, "cache_read_input_token_cost_priority": 1.44375e-07, "input_cost_per_token": 3.3e-07, "input_cost_per_token_batches": 1.65e-07, "input_cost_per_token_flex": 1.65e-07, "input_cost_per_token_priority": 5.775e-07, "output_cost_per_token": 2.75e-06, "output_cost_per_token_batches": 1.375e-06, "output_cost_per_token_flex": 1.375e-06, "output_cost_per_token_priority": 4.8125e-06}, "inclusive": false, "checks": [[null, 1, [3.3e-07, 2.75e-06, 3.3e-07, 3.3e-07, 8.25e-08]], [null, 10000000, [3.3e-07, 2.75e-06, 3.3e-07, 3.3e-07, 8.25e-08]], ["priority", 1, [5.775e-07, 4.8125e-06, 5.775e-07, 5.775e-07, 1.44375e-07]], ["priority", 10000000, [5.775e-07, 4.8125e-06, 5.775e-07, 5.775e-07, 1.44375e-07]], ["flex", 1, [1.65e-07, 1.375e-06, 1.65e-07, 1.65e-07, 4.125e-08]], ["flex", 10000000, [1.65e-07, 1.375e-06, 1.65e-07, 1.65e-07, 4.125e-08]]]}, +{"info": {"cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_batches": 3.75e-08, "cache_read_input_token_cost_flex": 3.75e-08, "cache_read_input_token_cost_priority": 1.5e-07, "input_cost_per_token": 7.5e-07, "output_cost_per_token": 4.5e-06, "input_cost_per_token_batches": 3.75e-07, "input_cost_per_token_flex": 3.75e-07, "input_cost_per_token_priority": 1.5e-06, "output_cost_per_token_batches": 2.25e-06, "output_cost_per_token_flex": 2.25e-06, "output_cost_per_token_priority": 9e-06}, "inclusive": false, "checks": [[null, 1, [7.5e-07, 4.5e-06, 7.5e-07, 7.5e-07, 7.5e-08]], [null, 10000000, [7.5e-07, 4.5e-06, 7.5e-07, 7.5e-07, 7.5e-08]], ["priority", 1, [1.5e-06, 9e-06, 1.5e-06, 1.5e-06, 1.5e-07]], ["priority", 10000000, [1.5e-06, 9e-06, 1.5e-06, 1.5e-06, 1.5e-07]], ["flex", 1, [3.75e-07, 2.25e-06, 3.75e-07, 3.75e-07, 3.75e-08]], ["flex", 10000000, [3.75e-07, 2.25e-06, 3.75e-07, 3.75e-07, 3.75e-08]]]}, +{"info": {"cache_read_input_token_cost": 5.5e-09, "input_cost_per_token": 5.5e-08, "input_cost_per_token_batches": 2.75e-08, "output_cost_per_token": 4.4e-07, "output_cost_per_token_batches": 2.2e-07, "cache_read_input_token_cost_batches": 2.75e-09}, "inclusive": false, "checks": [[null, 1, [5.5e-08, 4.4e-07, 5.5e-08, 5.5e-08, 5.5e-09]], [null, 10000000, [5.5e-08, 4.4e-07, 5.5e-08, 5.5e-08, 5.5e-09]]]}, +{"info": {"output_cost_per_token": 0.0}, "inclusive": false, "checks": [[null, 1, [0.0, 0.0, 0.0, 0.0, 0.0]], [null, 10000000, [0.0, 0.0, 0.0, 0.0, 0.0]]]}, +{"info": {"input_cost_per_token_above_100k_tokens": 5e-07, "output_cost_per_token_above_100k_tokens": 2.5e-06, "cache_creation_input_token_cost_above_100k_tokens": 6.25e-07, "cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1e-06, "cache_read_input_token_cost_above_100k_tokens": 5e-08, "cache_creation_input_token_cost": 1.25e-07, "cache_creation_input_token_cost_above_1hr": 2e-07, "cache_creation_input_token_cost_batches": 6.25e-08, "cache_read_input_token_cost": 1e-08, "cache_read_input_token_cost_batches": 5e-09, "input_cost_per_token": 1e-07, "input_cost_per_token_batches": 5e-08, "output_cost_per_token": 5e-07, "output_cost_per_token_batches": 2.5e-07}, "inclusive": false, "checks": [[null, 1, [1e-07, 5e-07, 1.25e-07, 2e-07, 1e-08]], [null, 100000, [1e-07, 5e-07, 1.25e-07, 2e-07, 1e-08]], [null, 100001, [5e-07, 2.5e-06, 6.25e-07, 1e-06, 5e-08]], [null, 10000000, [5e-07, 2.5e-06, 6.25e-07, 1e-06, 5e-08]]]}, +{"info": {"input_cost_per_token_above_100k_tokens": 5.5e-07, "output_cost_per_token_above_100k_tokens": 2.75e-06, "cache_creation_input_token_cost_above_100k_tokens": 6.875e-07, "cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1.1e-06, "cache_read_input_token_cost_above_100k_tokens": 5.5e-08, "cache_creation_input_token_cost": 1.375e-07, "cache_creation_input_token_cost_above_1hr": 2.2e-07, "cache_read_input_token_cost": 1.1e-08, "input_cost_per_token": 1.1e-07, "output_cost_per_token": 5.5e-07}, "inclusive": false, "checks": [[null, 1, [1.1e-07, 5.5e-07, 1.375e-07, 2.2e-07, 1.1e-08]], [null, 100000, [1.1e-07, 5.5e-07, 1.375e-07, 2.2e-07, 1.1e-08]], [null, 100001, [5.5e-07, 2.75e-06, 6.875e-07, 1.1e-06, 5.5e-08]], [null, 10000000, [5.5e-07, 2.75e-06, 6.875e-07, 1.1e-06, 5.5e-08]]]}, +{"info": {"cache_creation_input_token_cost": 4.125e-06, "cache_creation_input_token_cost_flex": 2.0625e-06, "cache_creation_input_token_cost_priority": 7.21875e-06, "cache_read_input_token_cost": 3.3e-07, "cache_read_input_token_cost_flex": 1.65e-07, "cache_read_input_token_cost_priority": 5.775e-07, "input_cost_per_token": 3.3e-06, "input_cost_per_token_flex": 1.65e-06, "input_cost_per_token_priority": 5.775e-06, "output_cost_per_token": 1.65e-05, "output_cost_per_token_flex": 8.25e-06, "output_cost_per_token_priority": 2.8875e-05}, "inclusive": false, "checks": [[null, 1, [3.3e-06, 1.65e-05, 4.125e-06, 4.125e-06, 3.3e-07]], [null, 10000000, [3.3e-06, 1.65e-05, 4.125e-06, 4.125e-06, 3.3e-07]], ["priority", 1, [5.775e-06, 2.8875e-05, 7.21875e-06, 7.21875e-06, 5.775e-07]], ["priority", 10000000, [5.775e-06, 2.8875e-05, 7.21875e-06, 7.21875e-06, 5.775e-07]], ["flex", 1, [1.65e-06, 8.25e-06, 2.0625e-06, 2.0625e-06, 1.65e-07]], ["flex", 10000000, [1.65e-06, 8.25e-06, 2.0625e-06, 2.0625e-06, 1.65e-07]]]}, +{"info": {"cache_read_input_token_cost": 1.375e-07, "cache_read_input_token_cost_priority": 2.75e-07, "input_cost_per_token": 1.375e-06, "input_cost_per_token_batches": 6.875e-07, "input_cost_per_token_priority": 2.75e-06, "output_cost_per_token": 1.1e-05, "output_cost_per_token_batches": 5.5e-06, "output_cost_per_token_priority": 2.2e-05, "cache_read_input_token_cost_batches": 6.875e-08}, "inclusive": false, "checks": [[null, 1, [1.375e-06, 1.1e-05, 1.375e-06, 1.375e-06, 1.375e-07]], [null, 10000000, [1.375e-06, 1.1e-05, 1.375e-06, 1.375e-06, 1.375e-07]], ["priority", 1, [2.75e-06, 2.2e-05, 2.75e-06, 2.75e-06, 2.75e-07]], ["priority", 10000000, [2.75e-06, 2.2e-05, 2.75e-06, 2.75e-06, 2.75e-07]]]}, +{"info": {"input_cost_per_token": 5e-06, "input_cost_per_token_batches": 2.5e-06, "output_cost_per_token": 1.5e-05, "output_cost_per_token_batches": 7.5e-06}, "inclusive": false, "checks": [[null, 1, [5e-06, 1.5e-05, 5e-06, 5e-06, 5e-06]], [null, 10000000, [5e-06, 1.5e-05, 5e-06, 5e-06, 5e-06]]]}, +{"info": {"cache_creation_input_token_cost": 2.1e-06, "cache_creation_input_token_cost_batches": 1.05e-06, "cache_creation_input_token_cost_flex": 1.05e-06, "cache_creation_input_token_cost_priority": 3.675e-06, "cache_read_input_token_cost": 3.12e-07, "cache_read_input_token_cost_batches": 1.56e-07, "cache_read_input_token_cost_flex": 1.56e-07, "cache_read_input_token_cost_priority": 5.46e-07, "input_cost_per_token": 1.68e-06, "input_cost_per_token_batches": 8.4e-07, "input_cost_per_token_flex": 8.4e-07, "input_cost_per_token_priority": 2.94e-06, "output_cost_per_token": 5.28e-06, "output_cost_per_token_batches": 2.64e-06, "output_cost_per_token_flex": 2.64e-06, "output_cost_per_token_priority": 9.24e-06}, "inclusive": false, "checks": [[null, 1, [1.68e-06, 5.28e-06, 2.1e-06, 2.1e-06, 3.12e-07]], [null, 10000000, [1.68e-06, 5.28e-06, 2.1e-06, 2.1e-06, 3.12e-07]], ["priority", 1, [2.94e-06, 9.24e-06, 3.675e-06, 3.675e-06, 5.46e-07]], ["priority", 10000000, [2.94e-06, 9.24e-06, 3.675e-06, 3.675e-06, 5.46e-07]], ["flex", 1, [8.4e-07, 2.64e-06, 1.05e-06, 1.05e-06, 1.56e-07]], ["flex", 10000000, [8.4e-07, 2.64e-06, 1.05e-06, 1.05e-06, 1.56e-07]]]}, +{"info": {"input_cost_per_token": 1.1e-05, "input_cost_per_token_above_272k_tokens": 2.2e-05, "cache_creation_input_token_cost": 1.375e-05, "cache_creation_input_token_cost_above_272k_tokens": 2.75e-05, "cache_creation_input_token_cost_above_272k_tokens_ultrafast": 0.000165, "cache_creation_input_token_cost_ultrafast": 8.25e-05, "cache_read_input_token_cost": 1.1e-06, "cache_read_input_token_cost_above_272k_tokens": 2.2e-06, "cache_read_input_token_cost_above_272k_tokens_ultrafast": 1.32e-05, "cache_read_input_token_cost_ultrafast": 6.6e-06, "input_cost_per_token_above_272k_tokens_ultrafast": 0.000132, "input_cost_per_token_ultrafast": 6.6e-05, "output_cost_per_token": 5.5e-05, "output_cost_per_token_above_272k_tokens": 8.25e-05, "output_cost_per_token_above_272k_tokens_ultrafast": 0.000495, "output_cost_per_token_ultrafast": 0.00033}, "inclusive": false, "checks": [[null, 1, [1.1e-05, 5.5e-05, 1.375e-05, 1.375e-05, 1.1e-06]], [null, 272000, [1.1e-05, 5.5e-05, 1.375e-05, 1.375e-05, 1.1e-06]], [null, 272001, [2.2e-05, 8.25e-05, 2.75e-05, 2.75e-05, 2.2e-06]], [null, 10000000, [2.2e-05, 8.25e-05, 2.75e-05, 2.75e-05, 2.2e-06]], ["ultrafast", 1, [6.6e-05, 0.00033, 8.25e-05, 8.25e-05, 6.6e-06]], ["ultrafast", 272000, [6.6e-05, 0.00033, 8.25e-05, 8.25e-05, 6.6e-06]], ["ultrafast", 272001, [0.000132, 0.000495, 0.000165, 0.000165, 1.32e-05]], ["ultrafast", 10000000, [0.000132, 0.000495, 0.000165, 0.000165, 1.32e-05]]]}, +{"info": {"cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05, "cache_creation_input_token_cost_flex": 6.25e-06, "cache_creation_input_token_cost_priority": 2.5e-05, "cache_read_input_token_cost": 1e-06, "cache_read_input_token_cost_above_272k_tokens": 2e-06, "cache_read_input_token_cost_above_272k_tokens_flex": 1e-06, "cache_read_input_token_cost_above_272k_tokens_priority": 4e-06, "cache_read_input_token_cost_batches": 5e-07, "cache_read_input_token_cost_above_272k_tokens_batches": 1e-06, "cache_creation_input_token_cost_batches": 6.25e-06, "cache_creation_input_token_cost_above_272k_tokens_batches": 1.25e-05, "cache_creation_input_token_cost_above_272k_tokens_ultrafast": 0.00015, "cache_creation_input_token_cost_ultrafast": 7.5e-05, "cache_read_input_token_cost_above_272k_tokens_ultrafast": 1.2e-05, "cache_read_input_token_cost_flex": 5e-07, "cache_read_input_token_cost_priority": 2e-06, "cache_read_input_token_cost_ultrafast": 6e-06, "input_cost_per_token": 1e-05, "input_cost_per_token_above_272k_tokens": 2e-05, "input_cost_per_token_above_272k_tokens_flex": 1e-05, "input_cost_per_token_above_272k_tokens_priority": 4e-05, "input_cost_per_token_batches": 5e-06, "input_cost_per_token_above_272k_tokens_batches": 1e-05, "input_cost_per_token_above_272k_tokens_ultrafast": 0.00012, "input_cost_per_token_flex": 5e-06, "input_cost_per_token_priority": 2e-05, "input_cost_per_token_ultrafast": 6e-05, "output_cost_per_token": 5e-05, "output_cost_per_token_above_272k_tokens": 7.5e-05, "output_cost_per_token_above_272k_tokens_flex": 3.75e-05, "output_cost_per_token_above_272k_tokens_priority": 0.00015, "output_cost_per_token_batches": 2.5e-05, "output_cost_per_token_above_272k_tokens_batches": 3.75e-05, "output_cost_per_token_above_272k_tokens_ultrafast": 0.00045, "output_cost_per_token_flex": 2.5e-05, "output_cost_per_token_priority": 0.0001, "output_cost_per_token_ultrafast": 0.0003}, "inclusive": false, "checks": [[null, 1, [1e-05, 5e-05, 1.25e-05, 1.25e-05, 1e-06]], [null, 272000, [1e-05, 5e-05, 1.25e-05, 1.25e-05, 1e-06]], [null, 272001, [2e-05, 7.5e-05, 2.5e-05, 2.5e-05, 2e-06]], [null, 10000000, [2e-05, 7.5e-05, 2.5e-05, 2.5e-05, 2e-06]], ["priority", 1, [2e-05, 0.0001, 2.5e-05, 2.5e-05, 2e-06]], ["priority", 272000, [2e-05, 0.0001, 2.5e-05, 2.5e-05, 2e-06]], ["priority", 272001, [4e-05, 0.00015, 5e-05, 5e-05, 4e-06]], ["priority", 10000000, [4e-05, 0.00015, 5e-05, 5e-05, 4e-06]], ["flex", 1, [5e-06, 2.5e-05, 6.25e-06, 6.25e-06, 5e-07]], ["flex", 272000, [5e-06, 2.5e-05, 6.25e-06, 6.25e-06, 5e-07]], ["flex", 272001, [1e-05, 3.75e-05, 1.25e-05, 1.25e-05, 1e-06]], ["flex", 10000000, [1e-05, 3.75e-05, 1.25e-05, 1.25e-05, 1e-06]], ["ultrafast", 1, [6e-05, 0.0003, 7.5e-05, 7.5e-05, 6e-06]], ["ultrafast", 272000, [6e-05, 0.0003, 7.5e-05, 7.5e-05, 6e-06]], ["ultrafast", 272001, [0.00012, 0.00045, 0.00015, 0.00015, 1.2e-05]], ["ultrafast", 10000000, [0.00012, 0.00045, 0.00015, 0.00015, 1.2e-05]]]}, +{"info": {"input_cost_per_token": 2.2e-06, "output_cost_per_token": 6.6e-06, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_flex": 2.75e-07, "cache_read_input_token_cost_priority": 9.625e-07, "input_cost_per_token_flex": 1.1e-06, "input_cost_per_token_priority": 3.85e-06, "output_cost_per_token_flex": 3.3e-06, "output_cost_per_token_priority": 1.155e-05}, "inclusive": false, "checks": [[null, 1, [2.2e-06, 6.6e-06, 2.2e-06, 2.2e-06, 5.5e-07]], [null, 10000000, [2.2e-06, 6.6e-06, 2.2e-06, 2.2e-06, 5.5e-07]], ["priority", 1, [3.85e-06, 1.155e-05, 3.85e-06, 3.85e-06, 9.625e-07]], ["priority", 10000000, [3.85e-06, 1.155e-05, 3.85e-06, 3.85e-06, 9.625e-07]], ["flex", 1, [1.1e-06, 3.3e-06, 1.1e-06, 1.1e-06, 2.75e-07]], ["flex", 10000000, [1.1e-06, 3.3e-06, 1.1e-06, 1.1e-06, 2.75e-07]]]}, +{"info": {"cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, "input_cost_per_token_above_272k_tokens": 6e-05, "output_cost_per_token": 0.00018, "output_cost_per_token_above_272k_tokens": 0.00027}, "inclusive": false, "checks": [[null, 1, [3e-05, 0.00018, 3e-05, 3e-05, 3e-06]], [null, 272000, [3e-05, 0.00018, 3e-05, 3e-05, 3e-06]], [null, 272001, [6e-05, 0.00027, 6e-05, 6e-05, 6e-06]], [null, 10000000, [6e-05, 0.00027, 6e-05, 6e-05, 6e-06]]]}, +{"info": {"input_cost_per_token": 1.5e-07, "input_cost_per_token_batches": 1.2e-07, "output_cost_per_token": 0}, "inclusive": false, "checks": [[null, 1, [1.5e-07, 0.0, 1.5e-07, 1.5e-07, 1.5e-07]], [null, 10000000, [1.5e-07, 0.0, 1.5e-07, 1.5e-07, 1.5e-07]]]}, +{"info": {"cache_read_input_token_cost_priority": 6e-07, "input_cost_per_token_priority": 1.2e-06, "output_cost_per_token_priority": 1.2e-06}, "inclusive": false, "checks": [[null, 1, [0.0, 0.0, 0.0, 0.0, 0.0]], [null, 10000000, [0.0, 0.0, 0.0, 0.0, 0.0]], ["priority", 1, [1.2e-06, 1.2e-06, 1.2e-06, 1.2e-06, 6e-07]], ["priority", 10000000, [1.2e-06, 1.2e-06, 1.2e-06, 1.2e-06, 6e-07]]]}, +{"info": {"cache_read_input_token_cost": 2e-07, "input_cost_per_token": 1.25e-06, "output_cost_per_token": 2.5e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07}, "inclusive": true, "checks": [[null, 1, [1.25e-06, 2.5e-06, 1.25e-06, 1.25e-06, 2e-07]], [null, 200000, [2.5e-06, 5e-06, 2.5e-06, 2.5e-06, 4e-07]], [null, 200001, [2.5e-06, 5e-06, 2.5e-06, 2.5e-06, 4e-07]], [null, 10000000, [2.5e-06, 5e-06, 2.5e-06, 2.5e-06, 4e-07]]]}, +{"info": {"cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, "input_cost_per_token_batches": 1e-06, "output_cost_per_token": 2.5e-06, "output_cost_per_token_batches": 2e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07}, "inclusive": true, "checks": [[null, 1, [1.25e-06, 2.5e-06, 1.25e-06, 1.25e-06, 2e-07]], [null, 200000, [2.5e-06, 5e-06, 2.5e-06, 2.5e-06, 4e-07]], [null, 200001, [2.5e-06, 5e-06, 2.5e-06, 2.5e-06, 4e-07]], [null, 10000000, [2.5e-06, 5e-06, 2.5e-06, 2.5e-06, 4e-07]]]}, +{"info": {"input_cost_per_token": 5e-06}, "inclusive": false, "checks": [[null, 1, [5e-06, 0.0, 5e-06, 5e-06, 5e-06]], [null, 10000000, [5e-06, 0.0, 5e-06, 5e-06, 5e-06]]]}, +{"info": {"cache_read_input_token_cost": 1.5e-08, "input_cost_per_token": 6e-08, "input_cost_per_token_batches": 3e-08, "output_cost_per_token": 2.4e-07, "output_cost_per_token_batches": 1.2e-07}, "inclusive": false, "checks": [[null, 1, [6e-08, 2.4e-07, 6e-08, 6e-08, 1.5e-08]], [null, 10000000, [6e-08, 2.4e-07, 6e-08, 6e-08, 1.5e-08]]]}, +{"info": {"input_cost_per_token": 3.3e-05, "input_cost_per_token_above_272k_tokens": 6.6e-05, "input_cost_per_token_above_272k_tokens_batches": 3.3e-05, "input_cost_per_token_batches": 1.65e-05, "output_cost_per_token": 0.000198, "output_cost_per_token_above_272k_tokens": 0.000297, "output_cost_per_token_above_272k_tokens_batches": 0.0001485, "output_cost_per_token_batches": 9.9e-05}, "inclusive": false, "checks": [[null, 1, [3.3e-05, 0.000198, 3.3e-05, 3.3e-05, 3.3e-05]], [null, 272000, [3.3e-05, 0.000198, 3.3e-05, 3.3e-05, 3.3e-05]], [null, 272001, [6.6e-05, 0.000297, 6.6e-05, 6.6e-05, 6.6e-05]], [null, 10000000, [6.6e-05, 0.000297, 6.6e-05, 6.6e-05, 6.6e-05]]]}, +{"info": {"cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_priority": 1.3125e-07, "input_cost_per_token": 3e-07, "input_cost_per_token_batches": 1.5e-07, "input_cost_per_token_priority": 5.25e-07, "output_cost_per_token": 2.5e-06, "output_cost_per_token_batches": 1.25e-06, "output_cost_per_token_priority": 4.375e-06}, "inclusive": false, "checks": [[null, 1, [3e-07, 2.5e-06, 3e-07, 3e-07, 7.5e-08]], [null, 10000000, [3e-07, 2.5e-06, 3e-07, 3e-07, 7.5e-08]], ["priority", 1, [5.25e-07, 4.375e-06, 5.25e-07, 5.25e-07, 1.3125e-07]], ["priority", 10000000, [5.25e-07, 4.375e-06, 5.25e-07, 5.25e-07, 1.3125e-07]]]}, +{"info": {"input_cost_per_token": 1.25e-06, "output_cost_per_token": 2.5e-06, "cache_read_input_token_cost": 2e-07, "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06}, "inclusive": false, "checks": [[null, 1, [1.25e-06, 2.5e-06, 1.25e-06, 1.25e-06, 2e-07]], [null, 200000, [1.25e-06, 2.5e-06, 1.25e-06, 1.25e-06, 2e-07]], [null, 200001, [2.5e-06, 5e-06, 2.5e-06, 2.5e-06, 2e-07]], [null, 10000000, [2.5e-06, 5e-06, 2.5e-06, 2.5e-06, 2e-07]]]}, +{"info": {"input_cost_per_token": 5e-06, "output_cost_per_token": 3e-05, "cache_read_input_token_cost": 5e-07, "input_cost_per_token_above_272k_tokens": 1e-05, "output_cost_per_token_above_272k_tokens": 4.5e-05}, "inclusive": false, "checks": [[null, 1, [5e-06, 3e-05, 5e-06, 5e-06, 5e-07]], [null, 272000, [5e-06, 3e-05, 5e-06, 5e-06, 5e-07]], [null, 272001, [1e-05, 4.5e-05, 1e-05, 1e-05, 5e-07]], [null, 10000000, [1e-05, 4.5e-05, 1e-05, 1e-05, 5e-07]]]}, +{"info": {"input_cost_per_token": 1.5e-07, "input_cost_per_token_batches": 7.5e-08, "input_cost_per_token_flex": 7.5e-08, "input_cost_per_token_priority": 2.625e-07, "output_cost_per_token": 6e-07, "output_cost_per_token_batches": 3e-07, "output_cost_per_token_flex": 3e-07, "output_cost_per_token_priority": 1.05e-06}, "inclusive": false, "checks": [[null, 1, [1.5e-07, 6e-07, 1.5e-07, 1.5e-07, 1.5e-07]], [null, 10000000, [1.5e-07, 6e-07, 1.5e-07, 1.5e-07, 1.5e-07]], ["priority", 1, [2.625e-07, 1.05e-06, 2.625e-07, 2.625e-07, 2.625e-07]], ["priority", 10000000, [2.625e-07, 1.05e-06, 2.625e-07, 2.625e-07, 2.625e-07]], ["flex", 1, [7.5e-08, 3e-07, 7.5e-08, 7.5e-08, 7.5e-08]], ["flex", 10000000, [7.5e-08, 3e-07, 7.5e-08, 7.5e-08, 7.5e-08]]]}, +{"info": {"tiered_pricing": [{"input_cost_per_token": 1.2e-06, "output_cost_per_token": 6e-06, "range": [0, 32000.0]}, {"input_cost_per_token": 2.4e-06, "output_cost_per_token": 1.2e-05, "range": [32000.0, 128000.0]}, {"input_cost_per_token": 3e-06, "output_cost_per_token": 1.5e-05, "range": [128000.0, 252000.0]}]}, "inclusive": false, "checks": [[null, 1, [1.2e-06, 6e-06, 1.2e-06, 1.2e-06, 1.2e-06]], [null, 32000, [1.2e-06, 6e-06, 1.2e-06, 1.2e-06, 1.2e-06]], [null, 32001, [2.4e-06, 1.2e-05, 2.4e-06, 2.4e-06, 2.4e-06]], [null, 128000, [2.4e-06, 1.2e-05, 2.4e-06, 2.4e-06, 2.4e-06]], [null, 128001, [3e-06, 1.5e-05, 3e-06, 3e-06, 3e-06]], [null, 10000000, [3e-06, 1.5e-05, 3e-06, 3e-06, 3e-06]]]}, +{"info": {"cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 1e-06, "output_cost_per_token": 5e-06, "input_cost_per_token_batches": 5e-07, "output_cost_per_token_batches": 2.5e-06}, "inclusive": false, "checks": [[null, 1, [1e-06, 5e-06, 1.25e-06, 2e-06, 1e-07]], [null, 10000000, [1e-06, 5e-06, 1.25e-06, 2e-06, 1e-07]]]}, +{"info": {"cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, "cache_read_input_token_cost_priority": 3.6e-07, "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_priority": 3.6e-06, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "output_cost_per_token_priority": 2.16e-05, "cache_read_input_token_cost_flex": 2e-07, "input_cost_per_token_flex": 1e-06, "output_cost_per_token_flex": 6e-06}, "inclusive": false, "checks": [[null, 1, [2e-06, 1.2e-05, 2e-06, 2e-06, 2e-07]], [null, 200000, [2e-06, 1.2e-05, 2e-06, 2e-06, 2e-07]], [null, 200001, [4e-06, 1.8e-05, 4e-06, 4e-06, 4e-07]], [null, 10000000, [4e-06, 1.8e-05, 4e-06, 4e-06, 4e-07]], ["priority", 1, [3.6e-06, 2.16e-05, 3.6e-06, 3.6e-06, 3.6e-07]], ["priority", 200000, [3.6e-06, 2.16e-05, 3.6e-06, 3.6e-06, 3.6e-07]], ["priority", 200001, [7.2e-06, 3.24e-05, 7.2e-06, 7.2e-06, 7.2e-07]], ["priority", 10000000, [7.2e-06, 3.24e-05, 7.2e-06, 7.2e-06, 7.2e-07]], ["flex", 1, [1e-06, 6e-06, 1e-06, 1e-06, 2e-07]], ["flex", 200000, [1e-06, 6e-06, 1e-06, 1e-06, 2e-07]], ["flex", 200001, [4e-06, 1.8e-05, 4e-06, 4e-06, 4e-07]], ["flex", 10000000, [4e-06, 1.8e-05, 4e-06, 4e-06, 4e-07]]]}, +{"info": {"input_cost_per_token": 5e-06, "input_cost_per_token_batches": 2.5e-06, "input_cost_per_token_priority": 8.75e-06, "output_cost_per_token": 1.5e-05, "output_cost_per_token_batches": 7.5e-06, "output_cost_per_token_priority": 2.625e-05}, "inclusive": false, "checks": [[null, 1, [5e-06, 1.5e-05, 5e-06, 5e-06, 5e-06]], [null, 10000000, [5e-06, 1.5e-05, 5e-06, 5e-06, 5e-06]], ["priority", 1, [8.75e-06, 2.625e-05, 8.75e-06, 8.75e-06, 8.75e-06]], ["priority", 10000000, [8.75e-06, 2.625e-05, 8.75e-06, 8.75e-06, 8.75e-06]]]}, +{"info": {"cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_priority": 3.6e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, "cache_read_input_token_cost_batches": 1e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_priority": 3.6e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, "output_cost_per_token": 1.2e-05, "output_cost_per_token_priority": 2.16e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06}, "inclusive": false, "checks": [[null, 1, [2e-06, 1.2e-05, 2e-06, 2e-06, 2e-07]], [null, 200000, [2e-06, 1.2e-05, 2e-06, 2e-06, 2e-07]], [null, 200001, [4e-06, 1.8e-05, 4e-06, 4e-06, 4e-07]], [null, 10000000, [4e-06, 1.8e-05, 4e-06, 4e-06, 4e-07]], ["priority", 1, [3.6e-06, 2.16e-05, 3.6e-06, 3.6e-06, 3.6e-07]], ["priority", 200000, [3.6e-06, 2.16e-05, 3.6e-06, 3.6e-06, 3.6e-07]], ["priority", 200001, [7.2e-06, 3.24e-05, 7.2e-06, 7.2e-06, 7.2e-07]], ["priority", 10000000, [7.2e-06, 3.24e-05, 7.2e-06, 7.2e-06, 7.2e-07]]]}, +{"info": {"cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, "cache_read_input_token_cost_batches": 1e-07, "cache_read_input_token_cost_flex": 1e-07, "cache_read_input_token_cost_priority": 3.6e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, "input_cost_per_token_priority": 3.6e-06, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "output_cost_per_token_flex": 6e-06, "output_cost_per_token_priority": 2.16e-05}, "inclusive": false, "checks": [[null, 1, [2e-06, 1.2e-05, 2e-06, 2e-06, 2e-07]], [null, 200000, [2e-06, 1.2e-05, 2e-06, 2e-06, 2e-07]], [null, 200001, [4e-06, 1.8e-05, 4e-06, 4e-06, 4e-07]], [null, 10000000, [4e-06, 1.8e-05, 4e-06, 4e-06, 4e-07]], ["priority", 1, [3.6e-06, 2.16e-05, 3.6e-06, 3.6e-06, 3.6e-07]], ["priority", 200000, [3.6e-06, 2.16e-05, 3.6e-06, 3.6e-06, 3.6e-07]], ["priority", 200001, [7.2e-06, 3.24e-05, 7.2e-06, 7.2e-06, 7.2e-07]], ["priority", 10000000, [7.2e-06, 3.24e-05, 7.2e-06, 7.2e-06, 7.2e-07]], ["flex", 1, [1e-06, 6e-06, 1e-06, 1e-06, 1e-07]], ["flex", 200000, [1e-06, 6e-06, 1e-06, 1e-06, 1e-07]], ["flex", 200001, [4e-06, 1.8e-05, 4e-06, 4e-06, 4e-07]], ["flex", 10000000, [4e-06, 1.8e-05, 4e-06, 4e-06, 4e-07]]]}, +{"info": {"input_cost_per_token": 3e-07, "input_cost_per_token_batches": 1.5e-07, "input_cost_per_token_flex": 1.5e-07, "output_cost_per_token": 2.5e-06, "output_cost_per_token_batches": 1.25e-06, "output_cost_per_token_flex": 1.25e-06}, "inclusive": false, "checks": [[null, 1, [3e-07, 2.5e-06, 3e-07, 3e-07, 3e-07]], [null, 10000000, [3e-07, 2.5e-06, 3e-07, 3e-07, 3e-07]], ["flex", 1, [1.5e-07, 1.25e-06, 1.5e-07, 1.5e-07, 1.5e-07]], ["flex", 10000000, [1.5e-07, 1.25e-06, 1.5e-07, 1.5e-07, 1.5e-07]]]}, +{"info": {"cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, "input_cost_per_token_above_200k_tokens": 6e-06, "output_cost_per_token_above_200k_tokens": 2.25e-05, "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "cache_read_input_token_cost_batches": 1.5e-07, "input_cost_per_token_batches": 1.5e-06, "output_cost_per_token": 1.5e-05, "output_cost_per_token_batches": 7.5e-06, "cache_creation_input_token_cost_batches": 1.88e-06}, "inclusive": false, "checks": [[null, 1, [3e-06, 1.5e-05, 3.75e-06, 6e-06, 3e-07]], [null, 200000, [3e-06, 1.5e-05, 3.75e-06, 6e-06, 3e-07]], [null, 200001, [6e-06, 2.25e-05, 7.5e-06, 1.2e-05, 6e-07]], [null, 10000000, [6e-06, 2.25e-05, 7.5e-06, 1.2e-05, 6e-07]]]}, +{"info": {"input_cost_per_token": 7.4e-07, "input_cost_per_token_flex": 3.7e-07, "output_cost_per_token": 2.22e-06, "output_cost_per_token_flex": 1.11e-06}, "inclusive": false, "checks": [[null, 1, [7.4e-07, 2.22e-06, 7.4e-07, 7.4e-07, 7.4e-07]], [null, 10000000, [7.4e-07, 2.22e-06, 7.4e-07, 7.4e-07, 7.4e-07]], ["flex", 1, [3.7e-07, 1.11e-06, 3.7e-07, 3.7e-07, 3.7e-07]], ["flex", 10000000, [3.7e-07, 1.11e-06, 3.7e-07, 3.7e-07, 3.7e-07]]]}, +{"info": {"cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, "input_cost_per_token_above_200k_tokens": 6e-06, "output_cost_per_token_above_200k_tokens": 2.25e-05, "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "output_cost_per_token": 1.5e-05}, "inclusive": false, "checks": [[null, 1, [3e-06, 1.5e-05, 3.75e-06, 3.75e-06, 3e-07]], [null, 200000, [3e-06, 1.5e-05, 3.75e-06, 3.75e-06, 3e-07]], [null, 200001, [6e-06, 2.25e-05, 7.5e-06, 7.5e-06, 6e-07]], [null, 10000000, [6e-06, 2.25e-05, 7.5e-06, 7.5e-06, 6e-07]]]}, +{"info": {"cache_creation_input_token_cost": 4.5e-06, "cache_creation_input_token_cost_above_1hr": 7.2e-06, "cache_read_input_token_cost": 3.6e-07, "input_cost_per_token": 3.6e-06, "input_cost_per_token_above_200k_tokens": 7.2e-06, "output_cost_per_token_above_200k_tokens": 2.7e-05, "cache_creation_input_token_cost_above_200k_tokens": 9e-06, "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.44e-05, "cache_read_input_token_cost_above_200k_tokens": 7.2e-07, "output_cost_per_token": 1.8e-05}, "inclusive": false, "checks": [[null, 1, [3.6e-06, 1.8e-05, 4.5e-06, 7.2e-06, 3.6e-07]], [null, 200000, [3.6e-06, 1.8e-05, 4.5e-06, 7.2e-06, 3.6e-07]], [null, 200001, [7.2e-06, 2.7e-05, 9e-06, 1.44e-05, 7.2e-07]], [null, 10000000, [7.2e-06, 2.7e-05, 9e-06, 1.44e-05, 7.2e-07]]]}, +{"info": {"cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, "input_cost_per_token_above_200k_tokens": 6e-06, "output_cost_per_token_above_200k_tokens": 2.25e-05, "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "output_cost_per_token": 1.5e-05, "input_cost_per_token_batches": 1.5e-06, "output_cost_per_token_batches": 7.5e-06}, "inclusive": false, "checks": [[null, 1, [3e-06, 1.5e-05, 3.75e-06, 6e-06, 3e-07]], [null, 200000, [3e-06, 1.5e-05, 3.75e-06, 6e-06, 3e-07]], [null, 200001, [6e-06, 2.25e-05, 7.5e-06, 1.2e-05, 6e-07]], [null, 10000000, [6e-06, 2.25e-05, 7.5e-06, 1.2e-05, 6e-07]]]}, +{"info": {"input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token": 1e-05, "output_cost_per_token_above_200k_tokens": 1.5e-05}, "inclusive": false, "checks": [[null, 1, [1.25e-06, 1e-05, 1.25e-06, 1.25e-06, 1.25e-06]], [null, 200000, [1.25e-06, 1e-05, 1.25e-06, 1.25e-06, 1.25e-06]], [null, 200001, [2.5e-06, 1.5e-05, 2.5e-06, 2.5e-06, 2.5e-06]], [null, 10000000, [2.5e-06, 1.5e-05, 2.5e-06, 2.5e-06, 2.5e-06]]]}, +{"info": {"cache_read_input_token_cost": 1.25e-06, "input_cost_per_token": 5e-06}, "inclusive": false, "checks": [[null, 1, [5e-06, 0.0, 5e-06, 5e-06, 1.25e-06]], [null, 10000000, [5e-06, 0.0, 5e-06, 5e-06, 1.25e-06]]]}, +{"info": {"input_cost_per_token": 3e-07, "input_cost_per_token_above_512k_tokens": 6e-07, "output_cost_per_token": 1.2e-06, "output_cost_per_token_above_512k_tokens": 2.4e-06, "cache_read_input_token_cost": 6e-08, "cache_read_input_token_cost_above_512k_tokens": 1.2e-07}, "inclusive": false, "checks": [[null, 1, [3e-07, 1.2e-06, 3e-07, 3e-07, 6e-08]], [null, 512000, [3e-07, 1.2e-06, 3e-07, 3e-07, 6e-08]], [null, 512001, [6e-07, 2.4e-06, 6e-07, 6e-07, 1.2e-07]], [null, 10000000, [6e-07, 2.4e-06, 6e-07, 6e-07, 1.2e-07]]]}, +{"info": {"cache_read_input_token_cost": 1.25e-06, "cache_read_input_token_cost_batches": 6.25e-07, "input_cost_per_token": 5e-06, "input_cost_per_token_batches": 2.5e-06}, "inclusive": false, "checks": [[null, 1, [5e-06, 0.0, 5e-06, 5e-06, 1.25e-06]], [null, 10000000, [5e-06, 0.0, 5e-06, 5e-06, 1.25e-06]]]}, +{"info": {"cache_read_input_token_cost": 5e-09, "cache_read_input_token_cost_batches": 2.5e-09, "cache_read_input_token_cost_flex": 2.5e-09, "input_cost_per_token": 5e-08, "input_cost_per_token_batches": 2.5e-08, "input_cost_per_token_flex": 2.5e-08, "input_cost_per_token_priority": 2.5e-06, "output_cost_per_token": 4e-07, "output_cost_per_token_batches": 2e-07, "output_cost_per_token_flex": 2e-07}, "inclusive": false, "checks": [[null, 1, [5e-08, 4e-07, 5e-08, 5e-08, 5e-09]], [null, 10000000, [5e-08, 4e-07, 5e-08, 5e-08, 5e-09]], ["priority", 1, [2.5e-06, 4e-07, 2.5e-06, 2.5e-06, 5e-09]], ["priority", 10000000, [2.5e-06, 4e-07, 2.5e-06, 2.5e-06, 5e-09]], ["flex", 1, [2.5e-08, 2e-07, 2.5e-08, 2.5e-08, 2.5e-09]], ["flex", 10000000, [2.5e-08, 2e-07, 2.5e-08, 2.5e-08, 2.5e-09]]]}, +{"info": {"cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_batches": 1e-08, "cache_read_input_token_cost_flex": 1e-08, "input_cost_per_token": 2e-07, "output_cost_per_token": 1.25e-06, "input_cost_per_token_batches": 1e-07, "input_cost_per_token_flex": 1e-07, "output_cost_per_token_batches": 6.25e-07, "output_cost_per_token_flex": 6.25e-07}, "inclusive": false, "checks": [[null, 1, [2e-07, 1.25e-06, 2e-07, 2e-07, 2e-08]], [null, 10000000, [2e-07, 1.25e-06, 2e-07, 2e-07, 2e-08]], ["flex", 1, [1e-07, 6.25e-07, 1e-07, 1e-07, 1e-08]], ["flex", 10000000, [1e-07, 6.25e-07, 1e-07, 1e-07, 1e-08]]]}, +{"info": {"input_cost_per_token": 3e-05, "input_cost_per_token_above_272k_tokens": 6e-05, "input_cost_per_token_flex": 1.5e-05, "input_cost_per_token_batches": 1.5e-05, "input_cost_per_token_above_272k_tokens_batches": 3e-05, "output_cost_per_token": 0.00018, "output_cost_per_token_above_272k_tokens": 0.00027, "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, "output_cost_per_token_above_272k_tokens_batches": 0.000135, "input_cost_per_token_above_272k_tokens_flex": 3e-05, "output_cost_per_token_above_272k_tokens_flex": 0.000135}, "inclusive": false, "checks": [[null, 1, [3e-05, 0.00018, 3e-05, 3e-05, 3e-05]], [null, 272000, [3e-05, 0.00018, 3e-05, 3e-05, 3e-05]], [null, 272001, [6e-05, 0.00027, 6e-05, 6e-05, 6e-05]], [null, 10000000, [6e-05, 0.00027, 6e-05, 6e-05, 6e-05]], ["flex", 1, [1.5e-05, 9e-05, 1.5e-05, 1.5e-05, 1.5e-05]], ["flex", 272000, [1.5e-05, 9e-05, 1.5e-05, 1.5e-05, 1.5e-05]], ["flex", 272001, [3e-05, 0.000135, 3e-05, 3e-05, 3e-05]], ["flex", 10000000, [3e-05, 0.000135, 3e-05, 3e-05, 3e-05]]]}, +{"info": {"cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_batches": 2.5e-07, "cache_read_input_token_cost_above_272k_tokens_batches": 5e-07, "cache_read_input_token_cost_flex": 2.5e-07, "cache_read_input_token_cost_priority": 1.25e-06, "input_cost_per_token": 5e-06, "input_cost_per_token_above_272k_tokens": 1e-05, "input_cost_per_token_flex": 2.5e-06, "input_cost_per_token_batches": 2.5e-06, "input_cost_per_token_above_272k_tokens_batches": 5e-06, "input_cost_per_token_priority": 1.25e-05, "output_cost_per_token": 3e-05, "output_cost_per_token_above_272k_tokens": 4.5e-05, "output_cost_per_token_flex": 1.5e-05, "output_cost_per_token_batches": 1.5e-05, "output_cost_per_token_above_272k_tokens_batches": 2.25e-05, "output_cost_per_token_priority": 7.5e-05, "input_cost_per_token_above_272k_tokens_flex": 5e-06, "output_cost_per_token_above_272k_tokens_flex": 2.25e-05, "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07}, "inclusive": false, "checks": [[null, 1, [5e-06, 3e-05, 5e-06, 5e-06, 5e-07]], [null, 272000, [5e-06, 3e-05, 5e-06, 5e-06, 5e-07]], [null, 272001, [1e-05, 4.5e-05, 1e-05, 1e-05, 1e-06]], [null, 10000000, [1e-05, 4.5e-05, 1e-05, 1e-05, 1e-06]], ["priority", 1, [1.25e-05, 7.5e-05, 1.25e-05, 1.25e-05, 1.25e-06]], ["priority", 272000, [1.25e-05, 7.5e-05, 1.25e-05, 1.25e-05, 1.25e-06]], ["priority", 272001, [1e-05, 4.5e-05, 1e-05, 1e-05, 1e-06]], ["priority", 10000000, [1e-05, 4.5e-05, 1e-05, 1e-05, 1e-06]], ["flex", 1, [2.5e-06, 1.5e-05, 2.5e-06, 2.5e-06, 2.5e-07]], ["flex", 272000, [2.5e-06, 1.5e-05, 2.5e-06, 2.5e-06, 2.5e-07]], ["flex", 272001, [5e-06, 2.25e-05, 5e-06, 5e-06, 5e-07]], ["flex", 10000000, [5e-06, 2.25e-05, 5e-06, 5e-06, 5e-07]]]}, +{"info": {"input_cost_per_token": 3e-05, "input_cost_per_token_above_272k_tokens": 6e-05, "input_cost_per_token_flex": 1.5e-05, "input_cost_per_token_batches": 1.5e-05, "input_cost_per_token_above_272k_tokens_batches": 3e-05, "output_cost_per_token": 0.00018, "output_cost_per_token_above_272k_tokens": 0.00027, "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, "output_cost_per_token_above_272k_tokens_batches": 0.000135}, "inclusive": false, "checks": [[null, 1, [3e-05, 0.00018, 3e-05, 3e-05, 3e-05]], [null, 272000, [3e-05, 0.00018, 3e-05, 3e-05, 3e-05]], [null, 272001, [6e-05, 0.00027, 6e-05, 6e-05, 6e-05]], [null, 10000000, [6e-05, 0.00027, 6e-05, 6e-05, 6e-05]], ["flex", 1, [1.5e-05, 9e-05, 1.5e-05, 1.5e-05, 1.5e-05]], ["flex", 272000, [1.5e-05, 9e-05, 1.5e-05, 1.5e-05, 1.5e-05]], ["flex", 272001, [6e-05, 0.00027, 6e-05, 6e-05, 6e-05]], ["flex", 10000000, [6e-05, 0.00027, 6e-05, 6e-05, 6e-05]]]}, +{"info": {"cache_creation_input_token_cost": 1.25e-07, "cache_creation_input_token_cost_decisions": 1.25e-07, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-07, "cache_creation_input_token_cost_above_272k_tokens_decisions": 2.5e-07, "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-07, "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-07, "cache_creation_input_token_cost_flex": 6.25e-08, "cache_creation_input_token_cost_priority": 2.5e-07, "cache_read_input_token_cost": 1e-08, "cache_read_input_token_cost_decisions": 1e-08, "cache_read_input_token_cost_above_272k_tokens": 2e-08, "cache_read_input_token_cost_above_272k_tokens_decisions": 2e-08, "cache_read_input_token_cost_above_272k_tokens_flex": 1e-08, "cache_read_input_token_cost_above_272k_tokens_priority": 4e-08, "cache_read_input_token_cost_batches": 5e-09, "cache_read_input_token_cost_above_272k_tokens_batches": 1e-08, "cache_creation_input_token_cost_batches": 6.25e-08, "cache_creation_input_token_cost_above_272k_tokens_batches": 1.25e-07, "cache_read_input_token_cost_flex": 5e-09, "cache_read_input_token_cost_priority": 2e-08, "input_cost_per_token": 1e-07, "input_cost_per_token_decisions": 1e-07, "input_cost_per_token_above_272k_tokens": 2e-07, "input_cost_per_token_above_272k_tokens_decisions": 2e-07, "input_cost_per_token_above_272k_tokens_flex": 1e-07, "input_cost_per_token_above_272k_tokens_priority": 4e-07, "input_cost_per_token_batches": 5e-08, "input_cost_per_token_above_272k_tokens_batches": 1e-07, "input_cost_per_token_flex": 5e-08, "input_cost_per_token_priority": 2e-07, "output_cost_per_token": 5e-07, "output_cost_per_token_decisions": 0.0, "output_cost_per_token_above_272k_tokens": 7.5e-07, "output_cost_per_token_above_272k_tokens_decisions": 0.0, "output_cost_per_token_above_272k_tokens_flex": 3.75e-07, "output_cost_per_token_above_272k_tokens_priority": 1.5e-06, "output_cost_per_token_batches": 2.5e-07, "output_cost_per_token_above_272k_tokens_batches": 3.75e-07, "output_cost_per_token_flex": 2.5e-07, "output_cost_per_token_priority": 1e-06}, "inclusive": false, "checks": [[null, 1, [1e-07, 5e-07, 1.25e-07, 1.25e-07, 1e-08]], [null, 272000, [1e-07, 5e-07, 1.25e-07, 1.25e-07, 1e-08]], [null, 272001, [2e-07, 7.5e-07, 2.5e-07, 2.5e-07, 2e-08]], [null, 10000000, [2e-07, 7.5e-07, 2.5e-07, 2.5e-07, 2e-08]], ["priority", 1, [2e-07, 1e-06, 2.5e-07, 2.5e-07, 2e-08]], ["priority", 272000, [2e-07, 1e-06, 2.5e-07, 2.5e-07, 2e-08]], ["priority", 272001, [4e-07, 1.5e-06, 5e-07, 5e-07, 4e-08]], ["priority", 10000000, [4e-07, 1.5e-06, 5e-07, 5e-07, 4e-08]], ["flex", 1, [5e-08, 2.5e-07, 6.25e-08, 6.25e-08, 5e-09]], ["flex", 272000, [5e-08, 2.5e-07, 6.25e-08, 6.25e-08, 5e-09]], ["flex", 272001, [1e-07, 3.75e-07, 1.25e-07, 1.25e-07, 1e-08]], ["flex", 10000000, [1e-07, 3.75e-07, 1.25e-07, 1.25e-07, 1e-08]]]}, +{"info": {"cache_creation_input_token_cost": 1.25e-07, "cache_creation_input_token_cost_above_100k_tokens": 6.25e-07, "cache_read_input_token_cost": 1e-08, "cache_read_input_token_cost_above_100k_tokens": 5e-08, "input_cost_per_token": 1e-07, "input_cost_per_token_above_100k_tokens": 5e-07, "output_cost_per_token": 5e-07, "output_cost_per_token_above_100k_tokens": 2.5e-06}, "inclusive": false, "checks": [[null, 1, [1e-07, 5e-07, 1.25e-07, 1.25e-07, 1e-08]], [null, 100000, [1e-07, 5e-07, 1.25e-07, 1.25e-07, 1e-08]], [null, 100001, [5e-07, 2.5e-06, 6.25e-07, 6.25e-07, 5e-08]], [null, 10000000, [5e-07, 2.5e-06, 6.25e-07, 6.25e-07, 5e-08]]]}, +{"info": {"input_cost_per_token": 0, "input_cost_per_token_above_128k_tokens": 0, "output_cost_per_token": 0, "output_cost_per_token_above_128k_tokens": 0}, "inclusive": false, "checks": [[null, 1, [0.0, 0.0, 0.0, 0.0, 0.0]], [null, 128000, [0.0, 0.0, 0.0, 0.0, 0.0]], [null, 128001, [0.0, 0.0, 0.0, 0.0, 0.0]], [null, 10000000, [0.0, 0.0, 0.0, 0.0, 0.0]]]}, +{"info": {"cache_read_input_token_cost": 5e-08, "cache_read_input_token_cost_priority": 9e-08, "input_cost_per_token": 5e-07, "input_cost_per_token_batches": 2.5e-07, "input_cost_per_token_flex": 2.5e-07, "input_cost_per_token_priority": 9e-07, "output_cost_per_token": 3e-06, "output_cost_per_token_batches": 1.5e-06, "output_cost_per_token_flex": 1.5e-06, "output_cost_per_token_priority": 5.4e-06}, "inclusive": false, "checks": [[null, 1, [5e-07, 3e-06, 5e-07, 5e-07, 5e-08]], [null, 10000000, [5e-07, 3e-06, 5e-07, 5e-07, 5e-08]], ["priority", 1, [9e-07, 5.4e-06, 9e-07, 9e-07, 9e-08]], ["priority", 10000000, [9e-07, 5.4e-06, 9e-07, 9e-07, 9e-08]], ["flex", 1, [2.5e-07, 1.5e-06, 2.5e-07, 2.5e-07, 5e-08]], ["flex", 10000000, [2.5e-07, 1.5e-06, 2.5e-07, 2.5e-07, 5e-08]]]}, +{"info": {"cache_read_input_token_cost": 3e-08, "input_cost_per_token": 3e-07, "input_cost_per_token_batches": 1.5e-07, "input_cost_per_token_flex": 1.5e-07, "input_cost_per_token_priority": 5.4e-07, "output_cost_per_token": 2.5e-06}, "inclusive": false, "checks": [[null, 1, [3e-07, 2.5e-06, 3e-07, 3e-07, 3e-08]], [null, 10000000, [3e-07, 2.5e-06, 3e-07, 3e-07, 3e-08]], ["priority", 1, [5.4e-07, 2.5e-06, 5.4e-07, 5.4e-07, 3e-08]], ["priority", 10000000, [5.4e-07, 2.5e-06, 5.4e-07, 5.4e-07, 3e-08]], ["flex", 1, [1.5e-07, 2.5e-06, 1.5e-07, 1.5e-07, 3e-08]], ["flex", 10000000, [1.5e-07, 2.5e-06, 1.5e-07, 1.5e-07, 3e-08]]]}, +{"info": {"cache_read_input_token_cost": 8.3e-08, "input_cost_per_token": 1.65e-07, "input_cost_per_token_batches": 8.3e-08, "output_cost_per_token": 6.6e-07}, "inclusive": false, "checks": [[null, 1, [1.65e-07, 6.6e-07, 1.65e-07, 1.65e-07, 8.3e-08]], [null, 10000000, [1.65e-07, 6.6e-07, 1.65e-07, 1.65e-07, 8.3e-08]]]}, +{"info": {"cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "input_cost_per_token_batches": 1e-06, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_batches": 6e-06, "input_cost_per_token_priority": 3.6e-06, "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "output_cost_per_token_priority": 2.16e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "cache_read_input_token_cost_priority": 3.6e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07}, "inclusive": false, "checks": [[null, 1, [2e-06, 1.2e-05, 2e-06, 2e-06, 2e-07]], [null, 200000, [2e-06, 1.2e-05, 2e-06, 2e-06, 2e-07]], [null, 200001, [4e-06, 1.8e-05, 2.5e-07, 2.5e-07, 4e-07]], [null, 10000000, [4e-06, 1.8e-05, 2.5e-07, 2.5e-07, 4e-07]], ["priority", 1, [3.6e-06, 2.16e-05, 3.6e-06, 3.6e-06, 3.6e-07]], ["priority", 200000, [3.6e-06, 2.16e-05, 3.6e-06, 3.6e-06, 3.6e-07]], ["priority", 200001, [7.2e-06, 3.24e-05, 2.5e-07, 2.5e-07, 7.2e-07]], ["priority", 10000000, [7.2e-06, 3.24e-05, 2.5e-07, 2.5e-07, 7.2e-07]]]}, +{"info": {"cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, "cache_read_input_token_cost_priority": 3.6e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, "input_cost_per_token_priority": 3.6e-06, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "output_cost_per_token_flex": 6e-06, "output_cost_per_token_priority": 2.16e-05}, "inclusive": false, "checks": [[null, 1, [2e-06, 1.2e-05, 2e-06, 2e-06, 2e-07]], [null, 200000, [2e-06, 1.2e-05, 2e-06, 2e-06, 2e-07]], [null, 200001, [4e-06, 1.8e-05, 2.5e-07, 2.5e-07, 4e-07]], [null, 10000000, [4e-06, 1.8e-05, 2.5e-07, 2.5e-07, 4e-07]], ["priority", 1, [3.6e-06, 2.16e-05, 3.6e-06, 3.6e-06, 3.6e-07]], ["priority", 200000, [3.6e-06, 2.16e-05, 3.6e-06, 3.6e-06, 3.6e-07]], ["priority", 200001, [7.2e-06, 3.24e-05, 2.5e-07, 2.5e-07, 7.2e-07]], ["priority", 10000000, [7.2e-06, 3.24e-05, 2.5e-07, 2.5e-07, 7.2e-07]], ["flex", 1, [1e-06, 6e-06, 1e-06, 1e-06, 2e-07]], ["flex", 200000, [1e-06, 6e-06, 1e-06, 1e-06, 2e-07]], ["flex", 200001, [4e-06, 1.8e-05, 2.5e-07, 2.5e-07, 4e-07]], ["flex", 10000000, [4e-06, 1.8e-05, 2.5e-07, 2.5e-07, 4e-07]]]}, +{"info": {"cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "input_cost_per_token_batches": 1e-06, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_batches": 6e-06}, "inclusive": false, "checks": [[null, 1, [2e-06, 1.2e-05, 2e-06, 2e-06, 2e-07]], [null, 200000, [2e-06, 1.2e-05, 2e-06, 2e-06, 2e-07]], [null, 200001, [4e-06, 1.8e-05, 2.5e-07, 2.5e-07, 4e-07]], [null, 10000000, [4e-06, 1.8e-05, 2.5e-07, 2.5e-07, 4e-07]]]}, +{"info": {"cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 4.5e-07, "cache_read_input_token_cost_flex": 1.25e-07, "cache_read_input_token_cost_priority": 2.25e-07, "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token": 1e-05, "output_cost_per_token_above_200k_tokens": 1.5e-05, "input_cost_per_token_above_200k_tokens_priority": 4.5e-06, "input_cost_per_token_batches": 6.25e-07, "input_cost_per_token_flex": 6.25e-07, "input_cost_per_token_priority": 2.25e-06, "output_cost_per_token_above_200k_tokens_priority": 2.7e-05, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 1.8e-05}, "inclusive": false, "checks": [[null, 1, [1.25e-06, 1e-05, 1.25e-06, 1.25e-06, 1.25e-07]], [null, 200000, [1.25e-06, 1e-05, 1.25e-06, 1.25e-06, 1.25e-07]], [null, 200001, [2.5e-06, 1.5e-05, 2.5e-07, 2.5e-07, 2.5e-07]], [null, 10000000, [2.5e-06, 1.5e-05, 2.5e-07, 2.5e-07, 2.5e-07]], ["priority", 1, [2.25e-06, 1.8e-05, 2.25e-06, 2.25e-06, 2.25e-07]], ["priority", 200000, [2.25e-06, 1.8e-05, 2.25e-06, 2.25e-06, 2.25e-07]], ["priority", 200001, [4.5e-06, 2.7e-05, 2.5e-07, 2.5e-07, 4.5e-07]], ["priority", 10000000, [4.5e-06, 2.7e-05, 2.5e-07, 2.5e-07, 4.5e-07]], ["flex", 1, [6.25e-07, 5e-06, 6.25e-07, 6.25e-07, 1.25e-07]], ["flex", 200000, [6.25e-07, 5e-06, 6.25e-07, 6.25e-07, 1.25e-07]], ["flex", 200001, [2.5e-06, 1.5e-05, 2.5e-07, 2.5e-07, 2.5e-07]], ["flex", 10000000, [2.5e-06, 1.5e-05, 2.5e-07, 2.5e-07, 2.5e-07]]]}, +{"info": {"cache_creation_input_token_cost": 1.875e-06, "input_cost_per_token": 3.75e-06, "output_cost_per_token": 1.5e-05}, "inclusive": false, "checks": [[null, 1, [3.75e-06, 1.5e-05, 1.875e-06, 1.875e-06, 3.75e-06]], [null, 10000000, [3.75e-06, 1.5e-05, 1.875e-06, 1.875e-06, 3.75e-06]]]}, +{"info": {"cache_read_input_token_cost": 2.1e-07, "cache_read_input_token_cost_flex": 1.05e-07, "input_cost_per_token": 8.4e-07, "input_cost_per_token_flex": 4.2e-07, "output_cost_per_token": 3.36e-06, "output_cost_per_token_flex": 1.68e-06}, "inclusive": false, "checks": [[null, 1, [8.4e-07, 3.36e-06, 8.4e-07, 8.4e-07, 2.1e-07]], [null, 10000000, [8.4e-07, 3.36e-06, 8.4e-07, 8.4e-07, 2.1e-07]], ["flex", 1, [4.2e-07, 1.68e-06, 4.2e-07, 4.2e-07, 1.05e-07]], ["flex", 10000000, [4.2e-07, 1.68e-06, 4.2e-07, 4.2e-07, 1.05e-07]]]}, +{"info": {"input_cost_per_token": 2.5e-06, "output_cost_per_token": 1.25e-05, "cache_read_input_token_cost": 2.5e-07, "input_cost_per_token_balanced": 2e-06, "output_cost_per_token_balanced": 1e-05, "cache_read_input_token_cost_balanced": 2e-07, "input_cost_per_token_flex": 1.25e-06, "output_cost_per_token_flex": 6.25e-06, "cache_read_input_token_cost_flex": 1.5e-07}, "inclusive": false, "checks": [[null, 1, [2.5e-06, 1.25e-05, 2.5e-06, 2.5e-06, 2.5e-07]], [null, 10000000, [2.5e-06, 1.25e-05, 2.5e-06, 2.5e-06, 2.5e-07]], ["flex", 1, [1.25e-06, 6.25e-06, 1.25e-06, 1.25e-06, 1.5e-07]], ["flex", 10000000, [1.25e-06, 6.25e-06, 1.25e-06, 1.25e-06, 1.5e-07]], ["balanced", 1, [2e-06, 1e-05, 2e-06, 2e-06, 2e-07]], ["balanced", 10000000, [2e-06, 1e-05, 2e-06, 2e-06, 2e-07]]]}, +{"info": {"tiered_pricing": [{"input_cost_per_query": 0.00166, "max_results_range": [1, 10]}, {"input_cost_per_query": 0.00332, "max_results_range": [11, 20]}, {"input_cost_per_query": 0.00498, "max_results_range": [21, 30]}, {"input_cost_per_query": 0.00664, "max_results_range": [31, 40]}, {"input_cost_per_query": 0.0083, "max_results_range": [41, 50]}, {"input_cost_per_query": 0.00996, "max_results_range": [51, 60]}, {"input_cost_per_query": 0.01162, "max_results_range": [61, 70]}, {"input_cost_per_query": 0.01328, "max_results_range": [71, 80]}, {"input_cost_per_query": 0.01494, "max_results_range": [81, 90]}, {"input_cost_per_query": 0.0166, "max_results_range": [91, 100]}]}, "inclusive": false, "checks": [[null, 1, [0.0, 0.0, 0.0, 0.0, 0.0]], [null, 10000000, [0.0, 0.0, 0.0, 0.0, 0.0]]]}, +{"info": {"tiered_pricing": [{"input_cost_per_query": 0.005, "max_results_range": [0, 25]}, {"input_cost_per_query": 0.025, "max_results_range": [26, 100]}]}, "inclusive": false, "checks": [[null, 1, [0.0, 0.0, 0.0, 0.0, 0.0]], [null, 10000000, [0.0, 0.0, 0.0, 0.0, 0.0]]]}, +{"info": {"tiered_pricing": [{"cache_read_input_token_cost": 8e-08, "input_cost_per_token": 4e-07, "output_cost_per_token": 1.6e-06, "range": [0, 256000.0]}, {"cache_read_input_token_cost": 2.4e-07, "input_cost_per_token": 1.2e-06, "output_cost_per_token": 4.8e-06, "range": [256000.0, 1000000.0]}]}, "inclusive": false, "checks": [[null, 1, [4e-07, 1.6e-06, 4e-07, 4e-07, 8e-08]], [null, 256000, [4e-07, 1.6e-06, 4e-07, 4e-07, 8e-08]], [null, 256001, [1.2e-06, 4.8e-06, 1.2e-06, 1.2e-06, 2.4e-07]], [null, 10000000, [1.2e-06, 4.8e-06, 1.2e-06, 1.2e-06, 2.4e-07]]]}, +{"info": {"tiered_pricing": [{"input_cost_per_token": 5e-08, "output_cost_per_token": 4e-07, "range": [0, 256000.0]}, {"input_cost_per_token": 2.5e-07, "output_cost_per_token": 2e-06, "range": [256000.0, 1000000.0]}]}, "inclusive": false, "checks": [[null, 1, [5e-08, 4e-07, 5e-08, 5e-08, 5e-08]], [null, 256000, [5e-08, 4e-07, 5e-08, 5e-08, 5e-08]], [null, 256001, [2.5e-07, 2e-06, 2.5e-07, 2.5e-07, 2.5e-07]], [null, 10000000, [2.5e-07, 2e-06, 2.5e-07, 2.5e-07, 2.5e-07]]]}, +{"info": {"tiered_pricing": [{"input_cost_per_token": 3e-07, "output_cost_per_token": 1.5e-06, "range": [0, 32000.0]}, {"input_cost_per_token": 5e-07, "output_cost_per_token": 2.5e-06, "range": [32000.0, 128000.0]}, {"input_cost_per_token": 8e-07, "output_cost_per_token": 4e-06, "range": [128000.0, 256000.0]}, {"input_cost_per_token": 1.6e-06, "output_cost_per_token": 9.6e-06, "range": [256000.0, 1000000.0]}]}, "inclusive": false, "checks": [[null, 1, [3e-07, 1.5e-06, 3e-07, 3e-07, 3e-07]], [null, 32000, [3e-07, 1.5e-06, 3e-07, 3e-07, 3e-07]], [null, 32001, [5e-07, 2.5e-06, 5e-07, 5e-07, 5e-07]], [null, 128000, [5e-07, 2.5e-06, 5e-07, 5e-07, 5e-07]], [null, 128001, [8e-07, 4e-06, 8e-07, 8e-07, 8e-07]], [null, 256000, [8e-07, 4e-06, 8e-07, 8e-07, 8e-07]], [null, 256001, [1.6e-06, 9.6e-06, 1.6e-06, 1.6e-06, 1.6e-06]], [null, 10000000, [1.6e-06, 9.6e-06, 1.6e-06, 1.6e-06, 1.6e-06]]]}, +{"info": {"tiered_pricing": [{"cache_read_input_token_cost": 8e-08, "input_cost_per_token": 3e-07, "output_cost_per_token": 1.5e-06, "range": [0, 32000.0]}, {"cache_read_input_token_cost": 1.2e-07, "input_cost_per_token": 5e-07, "output_cost_per_token": 2.5e-06, "range": [32000.0, 128000.0]}, {"cache_read_input_token_cost": 2e-07, "input_cost_per_token": 8e-07, "output_cost_per_token": 4e-06, "range": [128000.0, 256000.0]}, {"cache_read_input_token_cost": 4e-07, "input_cost_per_token": 1.6e-06, "output_cost_per_token": 9.6e-06, "range": [256000.0, 1000000.0]}]}, "inclusive": false, "checks": [[null, 1, [3e-07, 1.5e-06, 3e-07, 3e-07, 8e-08]], [null, 32000, [3e-07, 1.5e-06, 3e-07, 3e-07, 8e-08]], [null, 32001, [5e-07, 2.5e-06, 5e-07, 5e-07, 1.2e-07]], [null, 128000, [5e-07, 2.5e-06, 5e-07, 5e-07, 1.2e-07]], [null, 128001, [8e-07, 4e-06, 8e-07, 8e-07, 2e-07]], [null, 256000, [8e-07, 4e-06, 8e-07, 8e-07, 2e-07]], [null, 256001, [1.6e-06, 9.6e-06, 1.6e-06, 1.6e-06, 4e-07]], [null, 10000000, [1.6e-06, 9.6e-06, 1.6e-06, 1.6e-06, 4e-07]]]}, +{"info": {"tiered_pricing": [{"input_cost_per_token": 4e-07, "output_cost_per_reasoning_token": 4e-06, "output_cost_per_token": 1.2e-06, "range": [0, 256000.0]}, {"input_cost_per_token": 1.2e-06, "output_cost_per_reasoning_token": 1.2e-05, "output_cost_per_token": 3.6e-06, "range": [256000.0, 1000000.0]}]}, "inclusive": false, "checks": [[null, 1, [4e-07, 1.2e-06, 4e-07, 4e-07, 4e-07]], [null, 256000, [4e-07, 1.2e-06, 4e-07, 4e-07, 4e-07]], [null, 256001, [1.2e-06, 3.6e-06, 1.2e-06, 1.2e-06, 1.2e-06]], [null, 10000000, [1.2e-06, 3.6e-06, 1.2e-06, 1.2e-06, 1.2e-06]]]}, +{"info": {"cache_creation_input_token_cost": 1.375e-06, "cache_read_input_token_cost": 1.1e-07, "input_cost_per_token": 1.1e-06, "output_cost_per_token": 5.5e-06, "input_cost_per_token_batches": 5.5e-07, "output_cost_per_token_batches": 2.75e-06}, "inclusive": false, "checks": [[null, 1, [1.1e-06, 5.5e-06, 1.375e-06, 1.375e-06, 1.1e-07]], [null, 10000000, [1.1e-06, 5.5e-06, 1.375e-06, 1.375e-06, 1.1e-07]]]}, +{"info": {"cache_read_input_token_cost": 2.75e-07, "cache_read_input_token_cost_above_272k_tokens": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens_batches": 2.75e-07, "cache_read_input_token_cost_batches": 1.43e-07, "cache_read_input_token_cost_priority": 5.5e-07, "input_cost_per_token": 2.75e-06, "input_cost_per_token_priority": 5.5e-06, "output_cost_per_token": 1.65e-05, "output_cost_per_token_priority": 3.3e-05, "input_cost_per_token_above_272k_tokens": 5.5e-06, "input_cost_per_token_above_272k_tokens_batches": 2.75e-06, "input_cost_per_token_batches": 1.375e-06, "output_cost_per_token_above_272k_tokens": 2.475e-05, "output_cost_per_token_above_272k_tokens_batches": 1.2375e-05, "output_cost_per_token_batches": 8.25e-06}, "inclusive": false, "checks": [[null, 1, [2.75e-06, 1.65e-05, 2.75e-06, 2.75e-06, 2.75e-07]], [null, 272000, [2.75e-06, 1.65e-05, 2.75e-06, 2.75e-06, 2.75e-07]], [null, 272001, [5.5e-06, 2.475e-05, 5.5e-06, 5.5e-06, 5.5e-07]], [null, 10000000, [5.5e-06, 2.475e-05, 5.5e-06, 5.5e-06, 5.5e-07]], ["priority", 1, [5.5e-06, 3.3e-05, 5.5e-06, 5.5e-06, 5.5e-07]], ["priority", 272000, [5.5e-06, 3.3e-05, 5.5e-06, 5.5e-06, 5.5e-07]], ["priority", 272001, [5.5e-06, 2.475e-05, 5.5e-06, 5.5e-06, 5.5e-07]], ["priority", 10000000, [5.5e-06, 2.475e-05, 5.5e-06, 5.5e-06, 5.5e-07]]]}, +{"info": {"cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_above_272k_tokens_batches": 5e-07, "cache_read_input_token_cost_priority": 1.25e-06, "cache_read_input_token_cost_above_272k_tokens_priority": 2e-06, "cache_read_input_token_cost_batches": 2.5e-07, "cache_read_input_token_cost_flex": 2.5e-07, "input_cost_per_token": 5e-06, "input_cost_per_token_above_272k_tokens": 1e-05, "input_cost_per_token_above_272k_tokens_batches": 5e-06, "input_cost_per_token_above_272k_tokens_flex": 5e-06, "input_cost_per_token_priority": 1.25e-05, "input_cost_per_token_above_272k_tokens_priority": 2e-05, "input_cost_per_token_batches": 2.5e-06, "input_cost_per_token_flex": 2.5e-06, "output_cost_per_token": 3e-05, "output_cost_per_token_above_272k_tokens": 4.5e-05, "output_cost_per_token_above_272k_tokens_batches": 2.25e-05, "output_cost_per_token_priority": 7.5e-05, "output_cost_per_token_above_272k_tokens_priority": 9e-05, "output_cost_per_token_batches": 1.5e-05}, "inclusive": false, "checks": [[null, 1, [5e-06, 3e-05, 5e-06, 5e-06, 5e-07]], [null, 272000, [5e-06, 3e-05, 5e-06, 5e-06, 5e-07]], [null, 272001, [1e-05, 4.5e-05, 1e-05, 1e-05, 1e-06]], [null, 10000000, [1e-05, 4.5e-05, 1e-05, 1e-05, 1e-06]], ["priority", 1, [1.25e-05, 7.5e-05, 1.25e-05, 1.25e-05, 1.25e-06]], ["priority", 272000, [1.25e-05, 7.5e-05, 1.25e-05, 1.25e-05, 1.25e-06]], ["priority", 272001, [2e-05, 9e-05, 2e-05, 2e-05, 2e-06]], ["priority", 10000000, [2e-05, 9e-05, 2e-05, 2e-05, 2e-06]], ["flex", 1, [2.5e-06, 3e-05, 2.5e-06, 2.5e-06, 2.5e-07]], ["flex", 272000, [2.5e-06, 3e-05, 2.5e-06, 2.5e-06, 2.5e-07]], ["flex", 272001, [5e-06, 4.5e-05, 5e-06, 5e-06, 1e-06]], ["flex", 10000000, [5e-06, 4.5e-05, 5e-06, 5e-06, 1e-06]]]}, +{"info": {"cache_creation_input_token_cost": 5e-06, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, "cache_creation_input_token_cost_priority": 1e-05, "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, "cache_read_input_token_cost": 4e-07, "cache_read_input_token_cost_above_272k_tokens": 8e-07, "cache_read_input_token_cost_priority": 8e-07, "cache_read_input_token_cost_above_272k_tokens_priority": 1.6e-06, "input_cost_per_token": 4e-06, "input_cost_per_token_above_272k_tokens": 8e-06, "input_cost_per_token_priority": 8e-06, "input_cost_per_token_above_272k_tokens_priority": 1.6e-05, "output_cost_per_token": 2e-05, "output_cost_per_token_above_272k_tokens": 3e-05, "output_cost_per_token_priority": 4e-05, "output_cost_per_token_above_272k_tokens_priority": 6e-05}, "inclusive": false, "checks": [[null, 1, [4e-06, 2e-05, 5e-06, 5e-06, 4e-07]], [null, 272000, [4e-06, 2e-05, 5e-06, 5e-06, 4e-07]], [null, 272001, [8e-06, 3e-05, 1e-05, 1e-05, 8e-07]], [null, 10000000, [8e-06, 3e-05, 1e-05, 1e-05, 8e-07]], ["priority", 1, [8e-06, 4e-05, 1e-05, 1e-05, 8e-07]], ["priority", 272000, [8e-06, 4e-05, 1e-05, 1e-05, 8e-07]], ["priority", 272001, [1.6e-05, 6e-05, 2e-05, 2e-05, 1.6e-06]], ["priority", 10000000, [1.6e-05, 6e-05, 2e-05, 2e-05, 1.6e-06]]]}, +{"info": {"cache_creation_input_token_cost": 5e-06, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 5e-06, "cache_creation_input_token_cost_priority": 1e-05, "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, "cache_creation_input_token_cost_flex": 2.5e-06, "cache_read_input_token_cost": 4e-07, "cache_read_input_token_cost_above_272k_tokens": 8e-07, "cache_read_input_token_cost_above_272k_tokens_flex": 4e-07, "cache_read_input_token_cost_priority": 8e-07, "cache_read_input_token_cost_above_272k_tokens_priority": 1.6e-06, "cache_read_input_token_cost_flex": 2e-07, "input_cost_per_token": 4e-06, "input_cost_per_token_above_272k_tokens": 8e-06, "input_cost_per_token_above_272k_tokens_flex": 4e-06, "input_cost_per_token_priority": 8e-06, "input_cost_per_token_above_272k_tokens_priority": 1.6e-05, "input_cost_per_token_flex": 2e-06, "output_cost_per_token": 2e-05, "output_cost_per_token_above_272k_tokens": 3e-05, "output_cost_per_token_above_272k_tokens_flex": 1.5e-05, "output_cost_per_token_priority": 4e-05, "output_cost_per_token_above_272k_tokens_priority": 6e-05, "output_cost_per_token_flex": 1e-05}, "inclusive": false, "checks": [[null, 1, [4e-06, 2e-05, 5e-06, 5e-06, 4e-07]], [null, 272000, [4e-06, 2e-05, 5e-06, 5e-06, 4e-07]], [null, 272001, [8e-06, 3e-05, 1e-05, 1e-05, 8e-07]], [null, 10000000, [8e-06, 3e-05, 1e-05, 1e-05, 8e-07]], ["priority", 1, [8e-06, 4e-05, 1e-05, 1e-05, 8e-07]], ["priority", 272000, [8e-06, 4e-05, 1e-05, 1e-05, 8e-07]], ["priority", 272001, [1.6e-05, 6e-05, 2e-05, 2e-05, 1.6e-06]], ["priority", 10000000, [1.6e-05, 6e-05, 2e-05, 2e-05, 1.6e-06]], ["flex", 1, [2e-06, 1e-05, 2.5e-06, 2.5e-06, 2e-07]], ["flex", 272000, [2e-06, 1e-05, 2.5e-06, 2.5e-06, 2e-07]], ["flex", 272001, [4e-06, 1.5e-05, 5e-06, 5e-06, 4e-07]], ["flex", 10000000, [4e-06, 1.5e-05, 5e-06, 5e-06, 4e-07]]]}, +{"info": {"cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, "input_cost_per_token_above_272k_tokens": 6e-05, "input_cost_per_token_above_272k_tokens_batches": 3e-05, "input_cost_per_token_above_272k_tokens_flex": 3e-05, "input_cost_per_token_batches": 1.5e-05, "input_cost_per_token_flex": 1.5e-05, "output_cost_per_token": 0.00018, "output_cost_per_token_above_272k_tokens": 0.00027, "output_cost_per_token_above_272k_tokens_batches": 0.000135, "output_cost_per_token_above_272k_tokens_flex": 0.000135, "output_cost_per_token_batches": 9e-05, "output_cost_per_token_flex": 9e-05}, "inclusive": false, "checks": [[null, 1, [3e-05, 0.00018, 3e-05, 3e-05, 3e-06]], [null, 272000, [3e-05, 0.00018, 3e-05, 3e-05, 3e-06]], [null, 272001, [6e-05, 0.00027, 6e-05, 6e-05, 6e-06]], [null, 10000000, [6e-05, 0.00027, 6e-05, 6e-05, 6e-06]], ["flex", 1, [1.5e-05, 9e-05, 1.5e-05, 1.5e-05, 3e-06]], ["flex", 272000, [1.5e-05, 9e-05, 1.5e-05, 1.5e-05, 3e-06]], ["flex", 272001, [3e-05, 0.000135, 3e-05, 3e-05, 6e-06]], ["flex", 10000000, [3e-05, 0.000135, 3e-05, 3e-05, 6e-06]]]}, +{"info": {"cache_read_input_token_cost": 2.5e-07, "cache_read_input_token_cost_above_272k_tokens": 5e-07, "cache_read_input_token_cost_above_272k_tokens_batches": 2.5e-07, "cache_read_input_token_cost_above_272k_tokens_flex": 2.5e-07, "cache_read_input_token_cost_priority": 5e-07, "cache_read_input_token_cost_above_272k_tokens_priority": 1e-06, "cache_read_input_token_cost_batches": 1.3e-07, "input_cost_per_token": 2.5e-06, "input_cost_per_token_above_272k_tokens": 5e-06, "input_cost_per_token_priority": 5e-06, "input_cost_per_token_above_272k_tokens_priority": 1e-05, "output_cost_per_token": 1.5e-05, "output_cost_per_token_above_272k_tokens": 2.25e-05, "output_cost_per_token_priority": 3e-05, "output_cost_per_token_above_272k_tokens_priority": 4.5e-05, "input_cost_per_token_above_272k_tokens_batches": 2.5e-06, "input_cost_per_token_above_272k_tokens_flex": 2.5e-06, "input_cost_per_token_batches": 1.25e-06, "input_cost_per_token_flex": 1.25e-06, "output_cost_per_token_above_272k_tokens_batches": 1.125e-05, "output_cost_per_token_above_272k_tokens_flex": 1.125e-05, "output_cost_per_token_batches": 7.5e-06, "output_cost_per_token_flex": 7.5e-06}, "inclusive": false, "checks": [[null, 1, [2.5e-06, 1.5e-05, 2.5e-06, 2.5e-06, 2.5e-07]], [null, 272000, [2.5e-06, 1.5e-05, 2.5e-06, 2.5e-06, 2.5e-07]], [null, 272001, [5e-06, 2.25e-05, 5e-06, 5e-06, 5e-07]], [null, 10000000, [5e-06, 2.25e-05, 5e-06, 5e-06, 5e-07]], ["priority", 1, [5e-06, 3e-05, 5e-06, 5e-06, 5e-07]], ["priority", 272000, [5e-06, 3e-05, 5e-06, 5e-06, 5e-07]], ["priority", 272001, [1e-05, 4.5e-05, 1e-05, 1e-05, 1e-06]], ["priority", 10000000, [1e-05, 4.5e-05, 1e-05, 1e-05, 1e-06]], ["flex", 1, [1.25e-06, 7.5e-06, 1.25e-06, 1.25e-06, 2.5e-07]], ["flex", 272000, [1.25e-06, 7.5e-06, 1.25e-06, 1.25e-06, 2.5e-07]], ["flex", 272001, [2.5e-06, 1.125e-05, 2.5e-06, 2.5e-06, 2.5e-07]], ["flex", 10000000, [2.5e-06, 1.125e-05, 2.5e-06, 2.5e-06, 2.5e-07]]]}, +{"info": {"cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_priority": 1e-06, "cache_read_input_token_cost_above_272k_tokens_priority": 2e-06, "input_cost_per_token": 5e-06, "input_cost_per_token_above_272k_tokens": 1e-05, "input_cost_per_token_priority": 1e-05, "input_cost_per_token_above_272k_tokens_priority": 2e-05, "output_cost_per_token": 3e-05, "output_cost_per_token_above_272k_tokens": 4.5e-05, "output_cost_per_token_priority": 6e-05, "output_cost_per_token_above_272k_tokens_priority": 9e-05}, "inclusive": false, "checks": [[null, 1, [5e-06, 3e-05, 5e-06, 5e-06, 5e-07]], [null, 272000, [5e-06, 3e-05, 5e-06, 5e-06, 5e-07]], [null, 272001, [1e-05, 4.5e-05, 1e-05, 1e-05, 1e-06]], [null, 10000000, [1e-05, 4.5e-05, 1e-05, 1e-05, 1e-06]], ["priority", 1, [1e-05, 6e-05, 1e-05, 1e-05, 1e-06]], ["priority", 272000, [1e-05, 6e-05, 1e-05, 1e-05, 1e-06]], ["priority", 272001, [2e-05, 9e-05, 2e-05, 2e-05, 2e-06]], ["priority", 10000000, [2e-05, 9e-05, 2e-05, 2e-05, 2e-06]]]}, +{"info": {"cache_read_input_token_cost": 3.125e-07, "input_cost_per_token": 1.25e-06, "input_cost_per_token_batches": 6.25e-07, "input_cost_per_token_priority": 2.1875e-06, "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_priority": 1.75e-05}, "inclusive": false, "checks": [[null, 1, [1.25e-06, 1e-05, 1.25e-06, 1.25e-06, 3.125e-07]], [null, 10000000, [1.25e-06, 1e-05, 1.25e-06, 1.25e-06, 3.125e-07]], ["priority", 1, [2.1875e-06, 1.75e-05, 2.1875e-06, 2.1875e-06, 3.125e-07]], ["priority", 10000000, [2.1875e-06, 1.75e-05, 2.1875e-06, 2.1875e-06, 3.125e-07]]]} +]} diff --git a/litellm-rust/crates/model-catalog/tests/registry_validation.rs b/litellm-rust/crates/model-catalog/tests/registry_validation.rs deleted file mode 100644 index 3b65ef13d78..00000000000 --- a/litellm-rust/crates/model-catalog/tests/registry_validation.rs +++ /dev/null @@ -1,93 +0,0 @@ -use std::path::{Path, PathBuf}; - -use litellm_model_catalog::{ - Catalog, FallbackGeneralizations, Provenance, validate_model_entry, validate_registry, -}; -use rstest::{fixture, rstest}; -use serde_json::{Map, Value}; - -#[fixture] -fn repo_root() -> PathBuf { - Path::new(env!("CARGO_MANIFEST_DIR")).join("../../..") -} - -#[rstest] -#[case("model_prices_and_context_window.json")] -#[case("litellm/model_prices_and_context_window_backup.json")] -#[ignore] -fn checked_in_registry_passes_strict_validation(repo_root: PathBuf, #[case] filename: &str) { - let body = std::fs::read(repo_root.join(filename)).unwrap(); - let catalog = Catalog::parse(&body, Provenance::default()).unwrap(); - validate_registry(&catalog).unwrap(); -} - -#[rstest] -#[ignore] -fn fallback_generalizations_are_typed(repo_root: PathBuf) { - let body = std::fs::read(repo_root.join("model_prices_and_context_window.json")).unwrap(); - let document: Map = serde_json::from_slice(&body).unwrap(); - let Some(raw_rules) = document.get("fallback_generalizations") else { - return; - }; - let _: FallbackGeneralizations = serde_json::from_value(raw_rules.clone()).unwrap(); - let catalog = Catalog::parse(&body, Provenance::default()).unwrap(); - assert!( - catalog - .fallback_rules() - .is_some_and(|rules| !rules.is_empty()) - ); -} - -#[rstest] -#[case::missing_provider(serde_json::json!({"mode": "chat"}), "litellm_provider")] -#[case::unknown_field(serde_json::json!({"litellm_provider": "test", "typo": true}), "unknown field")] -#[case::negative_price(serde_json::json!({"litellm_provider": "test", "input_cost_per_token": -1}), "nonnegative")] -#[case::negative_nested_price(serde_json::json!({"litellm_provider": "test", "guardrail_cost_per_unit": {"unit": -1}}), "nonnegative")] -#[case::invalid_mode(serde_json::json!({"litellm_provider": "test", "mode": "invalid"}), "unknown variant")] -#[case::invalid_date(serde_json::json!({"litellm_provider": "test", "deprecation_date": "2026-02-31"}), "deprecation_date")] -#[case::invalid_hours(serde_json::json!({"litellm_provider": "test", "off_peak_pricing": {"hours_utc": "25:00-01:00"}}), "hours_utc")] -#[case::empty_windows(serde_json::json!({"litellm_provider": "test", "off_peak_pricing": {"windows": []}}), "windows is empty")] -#[case::invalid_weekday(serde_json::json!({"litellm_provider": "test", "off_peak_pricing": {"windows": [{"hours_utc": "00:00-01:00", "weekdays": [0]}]}}), "weekdays is invalid")] -#[case::invalid_aliases(serde_json::json!({"litellm_provider": "test", "aliases": ["good", 7]}), "aliases must contain strings")] -#[case::null_aliases(serde_json::json!({"litellm_provider": "test", "aliases": null}), "aliases must be an array")] -#[ignore] -fn registry_validation_rejects_malformed_entries(#[case] entry: Value, #[case] expected: &str) { - let error = validate_model_entry("test", &entry).unwrap_err(); - assert!( - litellm_error::Report::new(&error) - .to_string() - .contains(expected) - ); -} - -#[test] -#[ignore] -fn checked_in_catalog_and_backup_match() { - let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../../.."); - let current = std::fs::read(root.join("model_prices_and_context_window.json")).unwrap(); - let backup = - std::fs::read(root.join("litellm/model_prices_and_context_window_backup.json")).unwrap(); - assert_eq!(current, backup); - let catalog = Catalog::parse(¤t, Provenance::default()).unwrap(); - assert!( - catalog.alias_issues().is_empty(), - "invalid registry aliases" - ); -} - -#[rstest] -#[case::input("input_cost_per_token_above_100k_tokens")] -#[case::input_batches("input_cost_per_token_above_100k_tokens_batches")] -#[case::output("output_cost_per_token_above_100k_tokens")] -#[case::output_batches("output_cost_per_token_above_100k_tokens_batches")] -#[case::cache_creation("cache_creation_input_token_cost_above_100k_tokens")] -#[case::cache_creation_batches("cache_creation_input_token_cost_above_100k_tokens_batches")] -#[case::cache_creation_1hr("cache_creation_input_token_cost_above_1hr_above_100k_tokens")] -#[case::cache_read("cache_read_input_token_cost_above_100k_tokens")] -#[case::cache_read_batches("cache_read_input_token_cost_above_100k_tokens_batches")] -fn registry_validation_keeps_above_100k_tier_rates(#[case] field: &str) { - let mut entry = Map::new(); - entry.insert("litellm_provider".into(), "anthropic".into()); - entry.insert(field.into(), 5e-7.into()); - validate_model_entry("test", &Value::Object(entry)).unwrap(); -} diff --git a/litellm-rust/crates/model-catalog/tests/schema.rs b/litellm-rust/crates/model-catalog/tests/schema.rs index f9a028a78ab..464cf069cac 100644 --- a/litellm-rust/crates/model-catalog/tests/schema.rs +++ b/litellm-rust/crates/model-catalog/tests/schema.rs @@ -1,16 +1,17 @@ #![cfg(feature = "schema")] -use std::collections::BTreeSet; -use std::path::Path; - -use litellm_model_catalog::{model_entry_json_schema, registry_json_schema}; +use litellm_model_catalog::{BUNDLED_REGISTRY, model_entry_json_schema, registry_json_schema}; use rstest::rstest; -use serde_json::{Value, json}; +use serde_json::{Map, Value, json}; fn schema() -> Value { serde_json::to_value(model_entry_json_schema()).expect("generated schema serializes") } +fn checked_in_catalog() -> Value { + serde_json::from_slice(BUNDLED_REGISTRY).unwrap() +} + fn registry_validator() -> jsonschema::Validator { let schema = serde_json::to_value(registry_json_schema()).unwrap(); jsonschema::options() @@ -19,26 +20,21 @@ fn registry_validator() -> jsonschema::Validator { .expect("generated registry schema is valid") } -#[rstest] -#[case("model_prices_and_context_window.json")] -#[case("litellm/model_prices_and_context_window_backup.json")] -#[ignore] -fn generated_registry_schema_validates_checked_in_catalog(#[case] path: &str) { - let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../../.."); - let catalog: Value = serde_json::from_slice(&std::fs::read(root.join(path)).unwrap()).unwrap(); +#[test] +fn generated_registry_schema_validates_checked_in_catalog() { + let catalog = checked_in_catalog(); let validator = registry_validator(); let errors: Vec<_> = validator .iter_errors(&catalog) .map(|error| error.to_string()) .collect(); - assert!(errors.is_empty(), "{path}: {errors:?}"); + assert!(errors.is_empty(), "{errors:?}"); } #[rstest] #[case(json!({"example": {"litellm_provider": "test"}}))] #[case(json!({"example": {"litellm_provider": "test", "future_field": true}}))] #[case(json!({"sample_spec": {"litellm_provider": "placeholder"}}))] -#[ignore] fn generated_registry_schema_keeps_reader_compatibility(#[case] document: Value) { assert!(registry_validator().is_valid(&document)); } @@ -59,47 +55,55 @@ fn generated_registry_schema_keeps_reader_compatibility(#[case] document: Value) #[case::invalid_weekday(json!({"litellm_provider": "test", "off_peak_pricing": {"windows": [{"hours_utc": "00:00-01:00", "weekdays": [0]}]}}))] #[case::invalid_aliases(json!({"litellm_provider": "test", "aliases": "wrong"}))] #[case::non_object_model(json!(4))] -#[ignore] fn generated_registry_schema_rejects_invalid_entries(#[case] entry: Value) { assert!(!registry_validator().is_valid(&json!({"example": entry}))); } -#[rstest] -#[case("model_prices_and_context_window.json")] -#[case("litellm/model_prices_and_context_window_backup.json")] -#[ignore] -fn generated_schema_covers_catalog_fields(#[case] path: &str) { - let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../../.."); - let catalog: Value = serde_json::from_slice(&std::fs::read(root.join(path)).unwrap()).unwrap(); +#[test] +fn generated_schema_covers_catalog_fields() { + let catalog = checked_in_catalog(); let schema = schema(); - let properties = schema["properties"] + let any_value = |keys: &Value| -> Map { + keys.as_object() + .expect("schema keys are an object") + .keys() + .map(|key| (key.clone(), json!({}))) + .collect() + }; + let keys_only = json!({ + "type": "object", + "properties": any_value(&schema["properties"]), + "patternProperties": any_value(&schema["patternProperties"]), + "additionalProperties": false, + }); + let registry = serde_json::to_value(registry_json_schema()).unwrap(); + let sections = registry["properties"] .as_object() - .expect("ModelInfo schema has properties"); - let fields: BTreeSet<&str> = catalog + .expect("registry has sections"); + let fields: Map = catalog .as_object() .expect("catalog is an object") .iter() - .filter(|(name, _)| *name != "sample_spec" && *name != "fallback_generalizations") + .filter(|(name, _)| !sections.contains_key(*name)) .flat_map(|(_, entry)| entry.as_object().expect("model entry is an object").keys()) - .map(String::as_str) - .filter(|name| *name != "aliases") + .map(|name| (name.clone(), Value::Null)) .collect(); - let missing: Vec<_> = fields - .into_iter() - .filter(|name| !properties.contains_key(*name)) + let missing: Vec<_> = jsonschema::validator_for(&keys_only) + .unwrap() + .iter_errors(&Value::Object(fields)) + .map(|error| error.to_string()) .collect(); assert!( missing.is_empty(), - "{path}: fields missing from schema: {missing:?}" + "fields missing from schema: {missing:?}" ); } #[rstest] #[case("Mode", "chat")] -#[case("ReasoningEffort", "high")] +#[case("EffortLevel", "high")] #[case("InputModality", "image")] -#[ignore] fn generated_schema_includes_enum_values(#[case] definition: &str, #[case] value: &str) { let schema = schema(); let variants = schema["$defs"][definition]["enum"] @@ -110,7 +114,6 @@ fn generated_schema_includes_enum_values(#[case] definition: &str, #[case] value } #[test] -#[ignore] fn generated_schema_includes_nested_pricing_types() { let schema = schema(); let definitions = schema["$defs"].as_object().expect("schema has definitions"); From e4932a453953340e9df7b839ea20a9eff28429b5 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 28 Oct 2026 18:35:54 -0700 Subject: [PATCH 126/128] test(router): set string fallback targets after Router init in the sdk cooldown cases (#48954) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../sdk/test_router_cooldown_fallback_sdk.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/integration/sdk/test_router_cooldown_fallback_sdk.py b/tests/integration/sdk/test_router_cooldown_fallback_sdk.py index e7508a2670e..0a5bfd339fc 100644 --- a/tests/integration/sdk/test_router_cooldown_fallback_sdk.py +++ b/tests/integration/sdk/test_router_cooldown_fallback_sdk.py @@ -122,12 +122,12 @@ async def test_cooldown_skips_only_the_failed_duplicate_model_deployment() -> No @pytest.mark.parametrize( - "router_kwargs", - [{"fallbacks": [{"primary": "backup"}]}, {"default_fallbacks": "backup"}], - ids=["string_fallback_target", "string_default_fallbacks"], + "fallbacks", + [[{"primary": "backup"}], [{"*": "backup"}]], + ids=["string_fallback_target", "string_catch_all_target"], ) @pytest.mark.asyncio -async def test_string_valued_fallbacks_still_cool_down_the_failing_primary(router_kwargs: dict[str, object]) -> None: +async def test_string_valued_fallbacks_still_cool_down_the_failing_primary(fallbacks: list[dict[str, object]]) -> None: with wire_server(_peer) as wire: router: Final = Router( model_list=[ @@ -136,8 +136,8 @@ async def test_string_valued_fallbacks_still_cool_down_the_failing_primary(route ], allowed_fails=0, num_retries=0, - **router_kwargs, ) + router.fallbacks = fallbacks with pytest.raises(litellm.AuthenticationError): await router.acompletion( model="primary", From 7aa7948a472f0d5dbb455ec5ced740cb3338819a Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 18:37:29 -0700 Subject: [PATCH 127/128] fix(health): ask OpenAI and Azure audio chat models for audio output in the health check probe (#47730) * fix(health): ask OpenAI and Azure audio chat models for audio output in the health check probe * fix(health): check the probed model when health_check_model is set, and mark Azure audio models as audio output A deployment with an audio base_model and a text health_check_model got the audio fields sent to its text probe model. The Azure gpt-audio and gpt-4o audio preview cost-map entries now set supports_audio_output, so a dated Azure name like azure/gpt-audio-1.5-2026-02-23 gets the audio probe too * fix(health): read litellm_params.base_model, and drop the stored base_model when Test Connect swaps the model * fix(health): judge the audio probe on a swapped-in health_check_params model, and keep Test Connect's stored base_model * fix(health): look up a provider-prefixed base_model under its own provider and keep text-only operator modalities --- ...odel_prices_and_context_window_backup.json | 6 + litellm/proxy/health_check.py | 68 +++++++++- .../health_endpoints/_health_endpoints.py | 2 + model_prices_and_context_window.json | 6 + .../health_endpoints/test_health_endpoints.py | 47 ++++++- .../proxy/test_health_check_max_tokens.py | 119 ++++++++++++++++++ 6 files changed, 243 insertions(+), 5 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 8719565a96f..60385d940d8 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -6101,6 +6101,7 @@ "text", "audio" ], + "supports_audio_output": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_parallel_function_calling": true, @@ -6133,6 +6134,7 @@ "text", "audio" ], + "supports_audio_output": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_parallel_function_calling": true, @@ -6164,6 +6166,7 @@ "text", "audio" ], + "supports_audio_output": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_parallel_function_calling": true, @@ -6230,6 +6233,7 @@ "text", "audio" ], + "supports_audio_output": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_parallel_function_calling": true, @@ -6293,6 +6297,7 @@ "text", "audio" ], + "supports_audio_output": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_parallel_function_calling": true, @@ -6369,6 +6374,7 @@ "text", "audio" ], + "supports_audio_output": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_parallel_function_calling": true, diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index be32405154c..60d2abe0ab4 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -131,6 +131,52 @@ def _should_inject_health_check_max_tokens(model_info: Mapping[str, object], mod return (mode or "chat") in _MAX_TOKEN_SUPPORT_MODES +_AUDIO_OUTPUT_PROBE_PROVIDERS: Final[frozenset[str]] = frozenset({"openai", "azure"}) + + +def _health_check_requires_audio_output( + model_info: Mapping[str, object], probe_params: Mapping[str, object], mode: str | None, deployment_model: str | None +) -> bool: + if (mode or "chat") != "chat": + return False + operator_modalities: Final = probe_params.get("modalities") + if isinstance(operator_modalities, list) and "audio" not in operator_modalities: + return False + model: Final = probe_params.get("model") + if not isinstance(model, str): + return False + provider_param: Final = probe_params.get("custom_llm_provider") + try: + _, provider, _, _ = litellm.get_llm_provider( + model=model, custom_llm_provider=provider_param if isinstance(provider_param, str) else None + ) + except litellm.exceptions.BadRequestError: + return False + if provider not in _AUDIO_OUTPUT_PROBE_PROVIDERS: + return False + base_model: Final = model_info.get("base_model") or probe_params.get("base_model") + checked_model: Final = base_model if isinstance(base_model, str) and model == deployment_model else model + checked_model_prefix: Final = checked_model.partition("/")[0] + return litellm.supports_audio_output( + model=checked_model, + custom_llm_provider=checked_model_prefix if checked_model_prefix in _AUDIO_OUTPUT_PROBE_PROVIDERS else provider, + ) + + +def _health_check_audio_output_params( + model_info: Mapping[str, object], probe_params: Mapping[str, object], mode: str | None, deployment_model: str | None +) -> Mapping[str, object]: + if not _health_check_requires_audio_output(model_info, probe_params, mode, deployment_model): + return MappingProxyType({}) + voice: Final = model_info.get("health_check_voice") + return MappingProxyType( + { + "modalities": ["text", "audio"], + "audio": {"voice": voice if isinstance(voice, str) else "alloy", "format": "wav"}, + } + ) + + # Health-check modes that forward `reasoning_effort` to the provider (chat-style calls). _HEALTH_CHECK_MODES_SUPPORTING_REASONING_EFFORT: Final = frozenset((None, "chat", "completion")) @@ -782,9 +828,12 @@ def _resolve_health_check_max_tokens(model_info: dict, litellm_params: dict) -> _PROBE_PARAMS_ADAPTER: Final = TypeAdapter(dict[str, object]) +_MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object]) -def update_litellm_params_for_health_check(model_info: dict, litellm_params: dict) -> dict: +def update_litellm_params_for_health_check( + model_info: dict, litellm_params: dict, deployment_model: str | None = None +) -> dict: """ Update the litellm params for health check. @@ -800,7 +849,15 @@ def update_litellm_params_for_health_check(model_info: dict, litellm_params: dic - updates the `model` param with the `health_check_model` if it exists Doc: https://docs.litellm.ai/docs/proxy/health#wildcard-routes - updates the `voice` param with the `health_check_voice` for `audio_speech` mode if it exists Doc: https://docs.litellm.ai/docs/proxy/health#text-to-speech-models - for Bedrock models with region routing (bedrock/region/model), strips the litellm routing prefix but preserves the model ID, and pins `custom_llm_provider` to `bedrock` (only when the deployment hasn't already set one, so an explicit `bedrock_converse` survives) so the bare model id still resolves to the provider (e.g. cross-region ids like `us.cohere.embed-v4:0`) + - asks OpenAI and Azure chat models that support audio output for `modalities` text + audio and a wav + `audio` in the `health_check_voice` (default "alloy"), since those models reject text-only requests. + `base_model` (from `model_info`, then `litellm_params`) decides it only while the probe still targets + `deployment_model` (default: the incoming `litellm_params` model), so a `health_check_model` or a `model` + in `health_check_params` is judged on its own. Keys already set through `health_check_params` are kept, + and `modalities` set there without "audio" turns the audio fields off """ + incoming_model: Final = _PROBE_PARAMS_ADAPTER.validate_python(litellm_params).get("model") + base_model_owner: Final = deployment_model or (incoming_model if isinstance(incoming_model, str) else None) mode: Final = resolve_health_check_mode( model_info, litellm_params, # any-ok: untyped router config dict @@ -875,7 +932,14 @@ def update_litellm_params_for_health_check(model_info: dict, litellm_params: dic ) probe_params: Final = _PROBE_PARAMS_ADAPTER.validate_python(litellm_params) - return {key: value for key, value in probe_params.items() if not is_model_capability_field(key)} + audio_output_params: Final = _health_check_audio_output_params( + _MODEL_INFO_ADAPTER.validate_python(model_info), probe_params, mode, base_model_owner + ) + return { + key: value + for key, value in {**audio_output_params, **probe_params}.items() + if not is_model_capability_field(key) + } _update_litellm_params_for_health_check: Final = update_litellm_params_for_health_check diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 53c54fe15bf..869e4953df0 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -2330,9 +2330,11 @@ async def test_model_connection( stored_params=_OBJECT_MAPPING.validate_python(config_litellm_params), request_params=_OBJECT_MAPPING.validate_python(request_litellm_params), ) + stored_model: Final = _OBJECT_MAPPING.validate_python(config_litellm_params).get("model") litellm_params = update_litellm_params_for_health_check( model_info=dict(probe_model_info), litellm_params=litellm_params, + deployment_model=stored_model if isinstance(stored_model, str) else None, ) ## Auth check, on the final probe params so health_check_params cannot retarget it afterwards diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8719565a96f..60385d940d8 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -6101,6 +6101,7 @@ "text", "audio" ], + "supports_audio_output": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_parallel_function_calling": true, @@ -6133,6 +6134,7 @@ "text", "audio" ], + "supports_audio_output": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_parallel_function_calling": true, @@ -6164,6 +6166,7 @@ "text", "audio" ], + "supports_audio_output": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_parallel_function_calling": true, @@ -6230,6 +6233,7 @@ "text", "audio" ], + "supports_audio_output": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_parallel_function_calling": true, @@ -6293,6 +6297,7 @@ "text", "audio" ], + "supports_audio_output": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_parallel_function_calling": true, @@ -6369,6 +6374,7 @@ "text", "audio" ], + "supports_audio_output": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_parallel_function_calling": true, diff --git a/tests/unit/proxy/health_endpoints/test_health_endpoints.py b/tests/unit/proxy/health_endpoints/test_health_endpoints.py index 67b280eb379..49183a143c0 100644 --- a/tests/unit/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/unit/proxy/health_endpoints/test_health_endpoints.py @@ -394,7 +394,7 @@ async def test_test_model_connection_loads_config_from_router(): mock_run_with_timeout = AsyncMock(return_value=mock_health_check_result) # Mock _update_litellm_params_for_health_check - def mock_update_params(model_info, litellm_params): + def mock_update_params(model_info, litellm_params, deployment_model: str | None = None): # Just return params with messages added params = litellm_params.copy() params["messages"] = [{"role": "user", "content": "test"}] @@ -545,7 +545,7 @@ async def test_test_model_connection_uses_model_info_id_to_disambiguate_duplicat mock_ahealth_check = AsyncMock(return_value=mock_health_check_result) mock_run_with_timeout = AsyncMock(return_value=mock_health_check_result) - def mock_update_params(model_info, litellm_params): + def mock_update_params(model_info, litellm_params, deployment_model: str | None = None): params = litellm_params.copy() params["messages"] = [{"role": "user", "content": "test"}] return params @@ -656,7 +656,7 @@ async def test_test_model_connection_falls_back_to_deployments_zero_without_id() mock_ahealth_check = AsyncMock(return_value=mock_health_check_result) mock_run_with_timeout = AsyncMock(return_value=mock_health_check_result) - def mock_update_params(model_info, litellm_params): + def mock_update_params(model_info, litellm_params, deployment_model: str | None = None): params = litellm_params.copy() params["messages"] = [{"role": "user", "content": "test"}] return params @@ -902,6 +902,47 @@ async def test_test_model_connection_overridden_model_probe_params_follow_the_pr assert "max_tokens" not in ahealth_check.call_args.kwargs["model_params"] +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("deployment_params", "deployment_model_info"), + [ + ({"model": "azure/my-audio-deployment"}, {"id": "audio-id", "base_model": "azure/gpt-audio-1.5"}), + ({"model": "azure/my-audio-deployment", "base_model": "azure/gpt-audio-1.5"}, {"id": "audio-id"}), + ], + ids=["model_info.base_model", "litellm_params.base_model"], +) +@pytest.mark.parametrize( + ("request_params", "expects_audio_probe"), + [({}, True), ({"model": "azure/my-text-deployment"}, False)], + ids=["own model", "swapped model"], +) +async def test_test_model_connection_stored_audio_base_model_follows_the_stored_model( + deployment_params: Mapping[str, str], + deployment_model_info: Mapping[str, str], + request_params: Mapping[str, str], + expects_audio_probe: bool, +): + deployment: Final = MappingProxyType( + { + "model_name": "custom-name-audio", + "litellm_params": {**deployment_params, "api_key": "fake-azure-key", "api_base": "https://x.openai.azure.com"}, + "model_info": deployment_model_info, + } + ) + with _test_connection_probe(deployment) as ahealth_check: + await health_test_model_connection( + request=MagicMock(), + mode=None, + litellm_params=dict(request_params), + model_info={"id": "audio-id"}, + user_api_key_dict=UserAPIKeyAuth(user_id="test-user", token="test-token"), + ) + + probe_params: Final = ahealth_check.call_args.kwargs["model_params"] + assert ("modalities" in probe_params) is expects_audio_probe + assert probe_params.get("base_model") == deployment_params.get("base_model") + + @pytest.mark.asyncio @pytest.mark.parametrize("params_mode", [123, ["chat"], {"mode": "chat"}, False], ids=["int", "list", "dict", "bool"]) async def test_test_model_connection_non_string_params_mode_is_a_bad_request(params_mode: object): diff --git a/tests/unit/proxy/test_health_check_max_tokens.py b/tests/unit/proxy/test_health_check_max_tokens.py index 24117f45106..e48bf4edc09 100644 --- a/tests/unit/proxy/test_health_check_max_tokens.py +++ b/tests/unit/proxy/test_health_check_max_tokens.py @@ -840,6 +840,125 @@ def test_health_check_params_apply_to_non_chat_modes(): assert "max_tokens" not in updated +@pytest.mark.parametrize( + ("model_info", "litellm_params"), + [ + ({}, {"model": "azure/gpt-audio-mini"}), + ({}, {"model": "openai/gpt-audio-mini"}), + ({"base_model": "azure/gpt-audio-mini"}, {"model": "azure/my-audio-deployment"}), + ({}, {"model": "azure/my-audio-deployment", "base_model": "azure/gpt-audio-1.5"}), + ( + {"base_model": "azure/gpt-4o-mini", "health_check_model": "azure/gpt-audio-mini"}, + {"model": "azure/my-text-deployment"}, + ), + ( + {"base_model": "gpt-4o-mini", "health_check_params": {"model": "openai/gpt-audio-mini"}}, + {"model": "openai/my-text-model"}, + ), + ({"base_model": "openai/gpt-audio-1.5"}, {"model": "azure/my-audio-deployment"}), + ], +) +def test_openai_and_azure_audio_chat_probe_asks_for_audio_output( + model_info: dict[str, str], litellm_params: dict[str, str] +): + updated: Final = update_litellm_params_for_health_check(model_info, litellm_params) + + assert updated["modalities"] == ["text", "audio"] + assert updated["audio"] == {"voice": "alloy", "format": "wav"} + assert updated["max_tokens"] == 16 + + +@pytest.mark.parametrize( + ("model_info", "litellm_params"), + [ + ({}, {"model": "azure/gpt-4o-mini"}), + ({}, {"model": "gemini/lyria-3-clip-preview"}), + ({}, {"model": "azure/my-audio-deployment"}), + ({}, {"model": "unmapped-model-name"}), + ({"mode": "responses"}, {"model": "openai/gpt-audio-mini"}), + ({"mode": "completion"}, {"model": "openai/gpt-audio-mini"}), + ({"base_model": "gemini/lyria-3-clip-preview"}, {"model": "azure/my-audio-deployment"}), + ( + {"base_model": "azure/gpt-audio-mini", "health_check_model": "azure/gpt-4o-mini"}, + {"model": "azure/my-audio-deployment"}, + ), + ( + {"health_check_model": "azure/gpt-4o-mini"}, + {"model": "azure/my-audio-deployment", "base_model": "azure/gpt-audio-mini"}, + ), + ( + {"base_model": "gpt-audio-mini", "health_check_params": {"model": "openai/gpt-5.4-mini"}}, + {"model": "openai/gpt-audio-mini"}, + ), + ( + {"health_check_params": {"model": "openai/gpt-5.4-mini"}}, + {"model": "openai/gpt-audio-mini", "base_model": "gpt-audio-mini"}, + ), + ], +) +def test_probe_without_openai_audio_chat_model_asks_for_no_audio_output( + model_info: dict[str, str], litellm_params: dict[str, str] +): + updated: Final = update_litellm_params_for_health_check(model_info, litellm_params) + + assert "modalities" not in updated + assert "audio" not in updated + + +@pytest.mark.parametrize( + "operator_params", + [ + {"modalities": ["audio", "text"], "audio": {"voice": "verse", "format": "mp3"}}, + {"audio": {"voice": "verse", "format": "mp3"}}, + ], +) +def test_audio_chat_probe_keeps_audio_params_set_in_health_check_params(operator_params: dict[str, object]): + updated: Final = update_litellm_params_for_health_check( + {"health_check_params": operator_params}, {"model": "azure/gpt-audio-mini"} + ) + + assert {key: updated[key] for key in operator_params} == operator_params + assert updated["modalities"] == operator_params.get("modalities", ["text", "audio"]) + + +_AUDIO_OUTPUT_CHAT_MODELS: Final = tuple( + sorted( + name + for name, info in litellm.model_cost.items() + if info.get("litellm_provider") in ("openai", "azure") + and info.get("mode") == "chat" + and "audio" in (info.get("supported_output_modalities") or []) + ) +) + + +@pytest.mark.parametrize("cost_map_model", _AUDIO_OUTPUT_CHAT_MODELS) +def test_every_audio_output_chat_model_in_the_cost_map_gets_the_audio_probe(cost_map_model: str): + provider: Final = litellm.model_cost[cost_map_model]["litellm_provider"] + updated: Final = update_litellm_params_for_health_check( + {"base_model": cost_map_model}, {"model": f"{provider}/my-audio-deployment"} + ) + + assert updated["modalities"] == ["text", "audio"] + + +def test_text_only_modalities_in_health_check_params_get_no_default_audio(): + updated: Final = update_litellm_params_for_health_check( + {"health_check_params": {"modalities": ["text"]}}, {"model": "openai/gpt-audio-mini"} + ) + + assert updated["modalities"] == ["text"] + assert "audio" not in updated + + +def test_audio_chat_probe_uses_health_check_voice(): + updated: Final = update_litellm_params_for_health_check( + {"health_check_voice": "verse"}, {"model": "openai/gpt-audio-mini"} + ) + + assert updated["audio"] == {"voice": "verse", "format": "wav"} + + async def _pegasus_health_check_request_body( model_info: dict[str, object], monkeypatch: pytest.MonkeyPatch ) -> dict[str, object]: From d9c9d842982afc3129f3295709e0bb8c7ad7942c Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 28 Oct 2026 18:45:24 -0700 Subject: [PATCH 128/128] fix(proxy): stop log lines from leaving terminal colors on (#48966) * fix(proxy): stop log lines from leaving terminal colors on * test(proxy): drop source scan for unclosed terminal colors 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> --- litellm/proxy/common_request_processing.py | 4 ++-- litellm/proxy/example_config_yaml/custom_callbacks.py | 3 +-- litellm/proxy/proxy_server.py | 6 +++--- 3 files changed, 6 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 9bb451495c8..89886c43525 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -4234,7 +4234,7 @@ class ProxyBaseLLMRequestProcessing: e = transformed_exception litellm_debug_info: Final = getattr(e, "litellm_debug_info", "") verbose_proxy_logger.debug( - "\033[1;31mAn error occurred: %s %s\n\n Debug this by setting `--debug`, e.g. `litellm --model gpt-3.5-turbo --debug`", + "An error occurred: %s %s", e, litellm_debug_info, ) @@ -4597,7 +4597,7 @@ class ProxyBaseLLMRequestProcessing: if transformed_exception is not None: e = transformed_exception verbose_proxy_logger.debug( - "\x1b[1;31mAn error occurred: %s\n\n Debug this by setting `--debug`, e.g. `litellm --model gpt-3.5-turbo --debug`", + "An error occurred: %s", e, ) diff --git a/litellm/proxy/example_config_yaml/custom_callbacks.py b/litellm/proxy/example_config_yaml/custom_callbacks.py index 5ae55060930..d11e27e0718 100644 --- a/litellm/proxy/example_config_yaml/custom_callbacks.py +++ b/litellm/proxy/example_config_yaml/custom_callbacks.py @@ -28,7 +28,7 @@ class MyCustomHandler(CustomLogger): def __init__(self): blue_color_code: Final = "\033[94m" reset_color_code: Final = "\033[0m" - print_verbose(f"{blue_color_code}Initialized LiteLLM custom logger") + print_verbose(f"{blue_color_code}Initialized LiteLLM custom logger{reset_color_code}") try: print_verbose("Logger Initialized with following methods:") methods: Final = [ @@ -40,7 +40,6 @@ class MyCustomHandler(CustomLogger): # Pretty print_verbose the methods for method in methods: print_verbose(f" - {method}") - print_verbose(f"{reset_color_code}") except Exception: pass diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b263412ce0f..42816e7db61 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6606,7 +6606,7 @@ class ProxyConfig: reset_color_code: Final = "\033[0m" for key, value in litellm_settings.items(): if key == "cache" and value is True: - print_startup_message(f"{blue_color_code}\nSetting Cache on Proxy") + print_startup_message(f"{blue_color_code}\nSetting Cache on Proxy{reset_color_code}") from litellm.caching.caching import Cache cache_params = {} @@ -9825,7 +9825,7 @@ async def async_assistants_data_generator(response, user_api_key_dict: UserAPIKe request_data=request_data, ) verbose_proxy_logger.debug( - "\x1b[1;31mAn error occurred: %s\n\n Debug this by setting `--debug`, e.g. `litellm --model gpt-3.5-turbo --debug`", + "An error occurred: %s", e, ) if isinstance(e, HTTPException): @@ -10655,7 +10655,7 @@ async def async_data_generator( ) client_error: Final[Exception] = e if replacement is None else replacement verbose_proxy_logger.debug( - "\x1b[1;31mAn error occurred: %s\n\n Debug this by setting `--debug`, e.g. `litellm --model gpt-3.5-turbo --debug`", + "An error occurred: %s", e, )