From d10de278c1bb65a271c032e4e17359bb3f0c64f9 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 9 Nov 2026 06:58:46 -0800 Subject: [PATCH 01/46] test(integration): recognize cached statements in database relay barriers (#50651) Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/_support/database_relay.py | 54 ++++++++++++++++--- .../database/test_harness_lifecycle.py | 43 +++++++++++++++ .../test_database_relay.py | 54 ++++++++++++++++++- 3 files changed, 142 insertions(+), 9 deletions(-) diff --git a/tests/integration/_support/database_relay.py b/tests/integration/_support/database_relay.py index 0ac1e1d2521..0ac060dc867 100644 --- a/tests/integration/_support/database_relay.py +++ b/tests/integration/_support/database_relay.py @@ -91,7 +91,6 @@ class HeldStatementRelay: self._upstream_port: Final = upstream_port self._trigger: Final = trigger self._after: Final = after - self._marker_length: Final = max(len(trigger), len(after or b"")) self._server: Final = LoopbackServer(self.port, self._serve) self._loop: Final = self._server.loop self._released: Final = asyncio.Event() @@ -115,27 +114,25 @@ class HeldStatementRelay: self.release() self._server.stop() - def _holds(self, window: bytes) -> bool: + def _holds(self, markers: frozenset[bytes]) -> bool: if not self._armed.is_set() or self.held.is_set(): return False if self._after is not None and not self._after_seen.is_set(): - if self._after in window: + if self._after in markers: self._after_seen.set() return False - return self._trigger in window + return self._trigger in markers async def _serve(self, client_reader: asyncio.StreamReader, client_writer: asyncio.StreamWriter) -> None: server_reader, server_writer = await asyncio.open_connection(self._upstream_host, self._upstream_port) async def forward(reader: asyncio.StreamReader, writer: asyncio.StreamWriter, inspect: bool) -> None: - tail = b"" # rebind-ok: carries the previous read's end so a trigger split across reads still matches + scanner: Final = StatementMarkerScanner((self._trigger,) + ((self._after,) if self._after else ())) try: while chunk := await reader.read(65536): - window: Final = tail + chunk - if inspect and self._holds(window): + if inspect and any(self._holds(markers) for markers in scanner.feed(chunk)): self.held.set() await self._released.wait() - tail = window[-(self._marker_length - 1) :] writer.write(chunk) await writer.drain() except (ConnectionError, asyncio.IncompleteReadError): @@ -149,6 +146,47 @@ class HeldStatementRelay: ) +class StatementMarkerScanner: + def __init__(self, markers: tuple[bytes, ...]) -> None: + self._markers: Final = markers + self._statements: Final[dict[bytes, frozenset[bytes]]] = {} + self._buffer = b"" + self._startup = True + + def feed(self, chunk: bytes) -> Generator[frozenset[bytes]]: + self._buffer += chunk + while len(self._buffer) >= (4 if self._startup else 5): + offset: Final = 0 if self._startup else 1 + length: Final = int.from_bytes(self._buffer[offset : offset + 4], "big") + if length < 4: + raise ValueError("Invalid PostgreSQL message length") + if len(self._buffer) < length + offset: + return + message: Final = self._buffer[: length + offset] + self._buffer = self._buffer[length + offset :] + if self._startup: + self._startup = message[4:8] == b"\x04\xd2\x16/" + continue + payload: Final = message[5:] + if message[:1] == b"P": + name, _, query = payload.partition(b"\x00") + matches: Final = self._matches(query.split(b"\x00", 1)[0]) + if name: + self._statements[name] = matches + yield matches + elif message[:1] == b"B": + fields: Final = payload.split(b"\x00", 2) + prepared: Final = self._statements.get(fields[1], frozenset()) if len(fields) > 1 else frozenset() + yield prepared | self._matches(payload) + elif message[:1] == b"Q": + yield self._matches(payload) + elif message[:1] == b"C" and payload[:1] == b"S": + self._statements.pop(payload[1:].split(b"\x00", 1)[0], None) + + def _matches(self, query: bytes) -> frozenset[bytes]: + return frozenset(marker for marker in self._markers if marker in query) + + class TriggerScanner: def __init__(self, trigger: bytes) -> None: self._trigger: Final = trigger diff --git a/tests/integration/database/test_harness_lifecycle.py b/tests/integration/database/test_harness_lifecycle.py index d600a92c5a2..4cb01dfeef6 100644 --- a/tests/integration/database/test_harness_lifecycle.py +++ b/tests/integration/database/test_harness_lifecycle.py @@ -219,6 +219,49 @@ async def _unused_peer(reader: asyncio.StreamReader, writer: asyncio.StreamWrite await reader.read() +async def _echo_protocol(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + while chunk := await reader.read(65536): + writer.write(chunk) + await writer.drain() + + +def _protocol_message(kind: bytes, payload: bytes) -> bytes: + return kind + (len(payload) + 4).to_bytes(4, "big") + payload + + +@pytest.mark.parametrize("prepared_trigger", (False, True)) +def test_relay_holds_cached_statements_parsed_before_arming(prepared_trigger: bool) -> None: + upstream_port: Final = allocate_port() + upstream: Final = LoopbackServer(upstream_port, _echo_protocol) + relay: Final = HeldStatementRelay("127.0.0.1", upstream_port, b"trigger", after=b"snapshot", armed=False) + startup: Final = b"\x00\x00\x00\x08\x00\x03\x00\x00" + snapshot: Final = _protocol_message(b"P", b"s1\x00SELECT snapshot\x00\x00\x00") + trigger: Final = _protocol_message(b"P", b"s2\x00SELECT trigger\x00\x00\x00") + bind_snapshot: Final = _protocol_message(b"B", b"\x00s1\x00\x00\x00\x00\x00\x00\x00") + bind_trigger: Final = _protocol_message(b"B", b"\x00s2\x00\x00\x00\x00\x00\x00\x00") + upstream.start() + try: + relay.start() + with ( + socket.create_connection(("127.0.0.1", relay.port), timeout=5) as client, + client.makefile("rb") as replies, + ): + before: Final = startup + snapshot + bind_snapshot + (trigger if prepared_trigger else b"") + client.sendall(before) + assert replies.read(len(before)) == before + relay.arm() + client.sendall(bind_snapshot) + assert replies.read(len(bind_snapshot)) == bind_snapshot + held: Final = bind_trigger if prepared_trigger else _protocol_message(b"Q", b"SELECT trigger\x00") + client.sendall(held) + assert relay.held.wait(timeout=5), "cached snapshot did not open the relay barrier" + relay.release() + assert replies.read(len(held)) == held + finally: + relay.stop() + upstream.stop() + + def test_relay_bind_failure_closes_its_loop_and_allows_repeated_cleanup() -> None: port: Final = allocate_port() with socket.socket() as occupied: diff --git a/tests/unit/integration_support/test_database_relay.py b/tests/unit/integration_support/test_database_relay.py index 98501f84058..de2bf1f5d2b 100644 --- a/tests/unit/integration_support/test_database_relay.py +++ b/tests/unit/integration_support/test_database_relay.py @@ -4,7 +4,7 @@ from typing import Final import pytest -from tests.integration._support.database_relay import TriggerScanner +from tests.integration._support.database_relay import StatementMarkerScanner, TriggerScanner TRIGGER: Final = b'SELECT "startTime" FROM "LiteLLM_SpendLogs"' @@ -26,3 +26,55 @@ def test_trigger_scanner_reports_a_match_once() -> None: scanner: Final = TriggerScanner(TRIGGER) assert scanner.feed(b"x" + TRIGGER + b"y") assert not scanner.feed(b"z") + + +STARTUP: Final = b"\x00\x00\x00\x08\x00\x03\x00\x00" + + +def _message(kind: bytes, payload: bytes) -> bytes: + return kind + (len(payload) + 4).to_bytes(4, "big") + payload + + +PARSE: Final = _message(b"P", b"s1\x00" + TRIGGER + b"\x00\x00\x00") +BIND: Final = _message(b"B", b"\x00s1\x00\x00\x00\x00\x00\x00\x00") + + +@pytest.mark.parametrize("split_at", range(1, len(STARTUP + PARSE + BIND))) +def test_statement_markers_survive_every_protocol_frame_split(split_at: int) -> None: + scanner: Final = StatementMarkerScanner((TRIGGER,)) + stream: Final = STARTUP + PARSE + BIND + assert (*scanner.feed(stream[:split_at]), *scanner.feed(stream[split_at:])) == ( + frozenset({TRIGGER}), + frozenset({TRIGGER}), + ) + + +def test_statement_markers_are_local_to_a_connection_and_forget_closed_statements() -> None: + scanner: Final = StatementMarkerScanner((TRIGGER,)) + other: Final = StatementMarkerScanner((TRIGGER,)) + assert tuple(scanner.feed(STARTUP + PARSE)) == (frozenset({TRIGGER}),) + assert tuple(other.feed(STARTUP + BIND)) == (frozenset(),) + assert tuple(scanner.feed(_message(b"C", b"Ss1\x00") + BIND)) == (frozenset(),) + + +def test_replaced_and_unnamed_statements_do_not_reuse_old_markers() -> None: + scanner: Final = StatementMarkerScanner((TRIGGER,)) + assert tuple(scanner.feed(STARTUP + PARSE)) == (frozenset({TRIGGER}),) + replacement: Final = _message(b"P", b"s1\x00SELECT 1\x00\x00\x00") + unnamed: Final = _message(b"P", b"\x00" + TRIGGER + b"\x00\x00\x00") + assert tuple(scanner.feed(replacement + BIND + unnamed)) == (frozenset(), frozenset(), frozenset({TRIGGER})) + + +def test_statement_scanner_rejects_invalid_frame_lengths() -> None: + scanner: Final = StatementMarkerScanner((TRIGGER,)) + assert tuple(scanner.feed(STARTUP)) == () + with pytest.raises(ValueError, match="Invalid PostgreSQL message length"): + tuple(scanner.feed(b"P\x00\x00\x00\x03")) + + +def test_statement_scanner_matches_bound_parameter_markers() -> None: + scanner: Final = StatementMarkerScanner((TRIGGER,)) + bind: Final = _message( + b"B", b"\x00s2\x00\x00\x00\x00\x01" + len(TRIGGER).to_bytes(4, "big") + TRIGGER + b"\x00\x00" + ) + assert tuple(scanner.feed(STARTUP + bind)) == (frozenset({TRIGGER}),) From 001126cd1c4710f75d7407e406f3d502b5904442 Mon Sep 17 00:00:00 2001 From: quietjules <349048734+quietjules@users.noreply.github.com> Date: Mon, 9 Nov 2026 10:54:20 -0800 Subject: [PATCH 02/46] feat(proxy): add OpenSearch index pass-through with per-key read/write grants (#50573) * feat(proxy): add OpenSearch index pass-through with per-key read/write grants Adds /opensearch// for LiteLLM vector store indexes on OpenSearch stores. A key or team needs read or write on the index in allowed_vector_store_indexes, and the proxy signs each call with the vector store's AWS identity, so callers never hold AWS credentials The route maps the LiteLLM index name to the store's provider index, refuses any path, query parameter, or body field that could reach another index, and sends the exact signed bytes upstream with only the signed headers * fix(proxy): let OpenSearch _bulk and _msearch NDJSON bodies through auth opensearch-py sends _bulk and _msearch NDJSON with content-type application/json, and the auth-time body read parsed it as one JSON document and answered 400 before the route ran. Auth now skips parsing OpenSearch NDJSON bodies the same way it skips Azure Speech audio. * feat(proxy): accept a _bulk _index naming the path index and allow named search pipelines A _bulk action line may name the index in the URL path, and LiteLLM rewrites it to the real index, which is what opensearch-py's helpers.bulk sends. The search_pipeline query param now goes through so hybrid search works, while an inline pipeline stays refused. Also types the type-checking Router stub with _router_init's parameters again, which make lint's tests/e2e check needs since the Rust core change. * fix(opensearch): answer 401 when the AWS credential source cannot read its token file * refactor(proxy): check OpenSearch request bodies without recursive helpers * fix(opensearch): answer 401 when AWS credentials fail to refresh at signing time * fix(opensearch): keep auth off OpenSearch bodies and widen the read list to index checks Auth no longer parses any OpenSearch request body, so a document field named model, api_base or base_url is written instead of answering 401 or 403. GET and HEAD on the index and HEAD on _doc and _source are reads, each operation checks the exact grant the list names, an index whose store name two stores share answers 400, and an _msearch header that names the path index is sent to the mapped index * fix(opensearch): need read for writes that return documents and keep document fields out of LiteLLM params * fix(opensearch): keep documents out of LiteLLM hooks, keep main's credential statuses, make _msearch admin only, and bill only searches * fix(proxy): bill only an exact /opensearch//_search path A document read such as /opensearch/docs/_doc/_search ended in /_search and was counted as a billable search. The check now matches the index segment and _search exactly * fix(proxy): only proxy admins set vector store index grants, and bill no OpenSearch document call An index grant now opens document writes through /opensearch, so key generate, update, bulk update, regenerate, team new and update, /user/new, and bulk user creation answer 403 when a non-admin changes allowed_vector_store_indexes, top-level or under metadata. Resending the stored grants is not a change An /opensearch path is matched before the LLM route suffixes, so a document id such as search or embeddings never makes a document call count as a billable request * fix(proxy): treat an empty index grant list from a non-admin as a change An explicit empty top-level allowed_vector_store_indexes passed the admin-only check as unchanged and then replaced the stored grants. It now counts as a change, so a non-admin gets 403. The new grant tests also carry parameter types --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/constants.py | 1 + litellm/llms/bedrock/base_aws_llm.py | 26 +- .../bedrock/web_identity_session_policy.py | 15 + .../vector_stores/index_passthrough.py | 379 ++++++++ .../vector_stores/transformation.py | 68 +- litellm/proxy/_lazy_features.py | 1 + litellm/proxy/_lazy_openapi_snapshot.json | 380 ++++++++ litellm/proxy/_types.py | 1 + litellm/proxy/auth/auth_utils.py | 44 +- litellm/proxy/auth/user_api_key_auth.py | 4 +- .../proxy/common_utils/http_parsing_utils.py | 7 +- litellm/proxy/litellm_pre_call_utils.py | 2 +- .../internal_user_endpoints.py | 6 +- .../key_management_endpoints.py | 35 +- .../management_endpoints/team_endpoints.py | 19 +- .../management_helpers/bulk_user_creation.py | 6 +- .../billable_request_metrics_middleware.py | 10 +- .../llm_passthrough_endpoints.py | 133 ++- .../pass_through_endpoints.py | 24 +- litellm/proxy/vector_store_endpoints/utils.py | 56 +- .../test_opensearch_transformation.py | 129 ++- .../test_key_management_endpoints.py | 144 +++ ...est_billable_request_metrics_middleware.py | 13 + .../test_llm_pass_through_endpoints.py | 919 +++++++++++++++++- .../test_pass_through_endpoints.py | 1 + .../unit/proxy/test_litellm_pre_call_utils.py | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 308 +++++- 27 files changed, 2658 insertions(+), 74 deletions(-) create mode 100644 litellm/llms/opensearch/vector_stores/index_passthrough.py diff --git a/litellm/constants.py b/litellm/constants.py index b43f6d9f358..fe605294c50 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1787,6 +1787,7 @@ AZURE_SPEECH_PASS_THROUGH_ROUTE_PREFIX: Final = "/azure_speech" AZURE_SPEECH_SHORT_AUDIO_PATH_PREFIX: Final = "/speech/" AZURE_SPEECH_BATCH_PATH_PREFIX: Final = "/speechtotext/" AZURE_SPEECH_FAST_TRANSCRIPTION_PATH: Final = "/speechtotext/transcriptions:transcribe" +OPENSEARCH_PASS_THROUGH_ROUTE_PREFIX: Final = "/opensearch" AZURE_SPEECH_STT_DOMAIN: Final = "stt.speech.microsoft.com" AZURE_SPEECH_COGNITIVE_SERVICES_DOMAIN: Final = "api.cognitive.microsoft.com" AZURE_SPEECH_SUBSCRIPTION_KEY_HEADER: Final = "Ocp-Apim-Subscription-Key" diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 8ab8a43ba3a..b0ba7470433 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -1712,6 +1712,24 @@ class BaseAWSLLM(SignsRequestsWithAWS): return request_headers_dict, request.body +def sign_aws_request( + method: str, + get_credentials: Callable[[], Credentials], + service_name: str, + aws_region_name: str | None, + url: str, + body: str | bytes, + headers: Mapping[str, str], +) -> AWSPreparedRequest: + ensure_optional_import("botocore") + from botocore.auth import SigV4Auth + from botocore.awsrequest import AWSRequest + + aws_request: Final = AWSRequest(method=method, url=url, data=body, headers=headers) + SigV4Auth(get_credentials(), service_name, aws_region_name).add_auth(aws_request) + return aws_request.prepare() + + def sign_aws_json_post( get_credentials: Callable[[], Credentials], service_name: str, @@ -1720,13 +1738,7 @@ def sign_aws_json_post( body: str, headers: Mapping[str, str], ) -> AWSPreparedRequest: - ensure_optional_import("botocore") - from botocore.auth import SigV4Auth - from botocore.awsrequest import AWSRequest - - aws_request: Final = AWSRequest(method="POST", url=url, data=body, headers=headers) - SigV4Auth(get_credentials(), service_name, aws_region_name).add_auth(aws_request) - return aws_request.prepare() + return sign_aws_request("POST", get_credentials, service_name, aws_region_name, url, body, headers) _SignParams = ParamSpec("_SignParams") diff --git a/litellm/llms/bedrock/web_identity_session_policy.py b/litellm/llms/bedrock/web_identity_session_policy.py index d567038cd1b..31187953dbc 100644 --- a/litellm/llms/bedrock/web_identity_session_policy.py +++ b/litellm/llms/bedrock/web_identity_session_policy.py @@ -266,6 +266,21 @@ POLLY_SESSION_POLICY: Final = _policy(_allow("PollyLiteLLM", ("polly:SynthesizeS OPENSEARCH_SESSION_POLICY: Final = _policy(_allow("OpenSearchLiteLLM", ("es:ESHttpPost", "aoss:APIAccessAll"))) +OPENSEARCH_INDEX_SESSION_POLICY: Final = _policy( + _allow( + "OpenSearchIndexLiteLLM", + ( + "es:ESHttpGet", + "es:ESHttpHead", + "es:ESHttpPost", + "es:ESHttpPut", + "es:ESHttpDelete", + "es:ESHttpPatch", + "aoss:APIAccessAll", + ), + ), +) + S3_VECTORS_QUERY_SESSION_POLICY: Final = _policy( _allow("S3VectorsLiteLLM", ("s3vectors:QueryVectors", "s3vectors:GetVectors")), ) diff --git a/litellm/llms/opensearch/vector_stores/index_passthrough.py b/litellm/llms/opensearch/vector_stores/index_passthrough.py new file mode 100644 index 00000000000..9432af4c1b3 --- /dev/null +++ b/litellm/llms/opensearch/vector_stores/index_passthrough.py @@ -0,0 +1,379 @@ +import json +from collections.abc import Iterable, Iterator, Mapping, Sequence +from dataclasses import dataclass +from itertools import accumulate, chain, repeat, takewhile +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal, TypeAlias +from urllib.parse import quote + +from pydantic import TypeAdapter + +from litellm.exceptions import AuthenticationError, BadRequestError, InternalServerError +from litellm.llms.bedrock.base_aws_llm import run_aws_signing + +if TYPE_CHECKING: + from litellm.llms.opensearch.vector_stores.transformation import OpenSearchVectorStoreConfig + +IndexHttpMethod: TypeAlias = Literal["GET", "HEAD", "POST", "PUT", "DELETE", "PATCH"] +IndexPermission: TypeAlias = Literal["read", "write"] +BodyKind: TypeAlias = Literal["query", "mget", "bulk", "update", "unscanned"] + +DEFAULT_CONTENT_TYPE: Final = "application/json" +NDJSON_CONTENT_TYPE: Final = "application/x-ndjson" +ACCEPTED_MEDIA_TYPES: Final = frozenset({DEFAULT_CONTENT_TYPE, NDJSON_CONTENT_TYPE}) +_ACCEPTED_MEDIA_TYPE_PARAMETERS: Final = frozenset({"", "charset=utf-8"}) +_INDEX_KEYS: Final = frozenset({"index", "_index", "indices"}) +_OUT_OF_INDEX_KEYS: Final = _INDEX_KEYS | {"indexed_shape", "wrapper", "search_pipeline"} +_BULK_METADATA_KEYS: Final = _INDEX_KEYS | {"pipeline"} +_OUT_OF_INDEX_QUERY_PARAMETERS: Final = frozenset({"pipeline", "source", "source_content_type"}) +_BULK_ACTIONS: Final = frozenset({"index", "create", "update", "delete"}) +_DOCUMENT_READING_KEYS: Final = frozenset({"_source", "script"}) +_JSON_BLANK: Final = b" \t\r\n" +_PATH_HELP: Final = "Use /opensearch//, for example /opensearch//_search." + + +@dataclass(frozen=True, slots=True) +class IndexRefusal: + status_code: int + message: str + + +@dataclass(frozen=True, slots=True) +class IndexPath: + index_name: str + api: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class IndexOperation: + permission: IndexPermission + methods: tuple[IndexHttpMethod, ...] + api: str + takes_id: bool + body: BodyKind + + def matches(self, method: str, api: tuple[str, ...]) -> bool: + named: Final = (self.api,) if self.api else () + return method in self.methods and api[: len(named)] == named and len(api) == len(named) + self.takes_id + + +@dataclass(frozen=True, slots=True) +class SignedIndexRequest: + url: str + media_type: str + headers: Mapping[str, str] + body: bytes + + +INDEX_OPERATIONS: Final = ( + IndexOperation("read", ("GET", "HEAD"), "", False, "unscanned"), + IndexOperation("write", ("POST",), "_doc", False, "unscanned"), + IndexOperation("write", ("PUT", "POST", "DELETE"), "_doc", True, "unscanned"), + IndexOperation("write", ("PUT", "POST"), "_create", True, "unscanned"), + IndexOperation("write", ("POST",), "_update", True, "update"), + IndexOperation("write", ("POST", "PUT"), "_bulk", False, "bulk"), + IndexOperation("write", ("POST",), "_delete_by_query", False, "query"), + IndexOperation("write", ("POST",), "_update_by_query", False, "query"), + IndexOperation("read", ("GET", "POST"), "_search", False, "query"), + IndexOperation("read", ("GET", "POST"), "_count", False, "query"), + IndexOperation("read", ("GET", "HEAD"), "_doc", True, "unscanned"), + IndexOperation("read", ("GET", "HEAD"), "_source", True, "unscanned"), + IndexOperation("read", ("GET", "POST"), "_mget", False, "mget"), + IndexOperation("read", ("GET", "POST"), "_explain", True, "query"), + IndexOperation("read", ("GET",), "_mapping", False, "unscanned"), +) + + +def index_operation(method: str, api: tuple[str, ...]) -> IndexOperation | None: + return next((operation for operation in INDEX_OPERATIONS if operation.matches(method, api)), None) + + +def parse_index_path(endpoint: str) -> IndexPath | IndexRefusal: + segments: Final = tuple(endpoint.split("/")) + if any(segment in ("", ".", "..") for segment in segments): + return IndexRefusal(400, f"The OpenSearch path '{endpoint}' has an empty, '.' or '..' segment. {_PATH_HELP}") + return IndexPath(index_name=segments[0], api=segments[1:]) + + +@dataclass(frozen=True, slots=True) +class _JsonObject: + pairs: tuple[tuple[str, object], ...] + + +@dataclass(frozen=True, slots=True) +class _Json: + value: object + + +@dataclass(frozen=True, slots=True) +class _BulkAction: + name: str + metadata: _JsonObject + has_source: bool + + +@dataclass(frozen=True, slots=True) +class _IndexScope: + name: str + provider_name: str + can_read: bool + + +def _json_object(pairs: Sequence[tuple[str, object]]) -> _JsonObject: + return _JsonObject(tuple(pairs)) + + +_DECODER: Final = json.JSONDecoder(object_pairs_hook=_json_object) +_JSON_ARRAY: Final = TypeAdapter(list[object]) +_JSON_VALUE: Final = TypeAdapter(object) + + +def _parse(raw: bytes) -> _Json | None: + try: + return _Json(_JSON_VALUE.validate_python(_DECODER.decode(raw.decode("utf-8")))) + except ValueError: + return None + + +def _own_keys(value: object) -> tuple[str, ...]: + return tuple(key for key, _ in value.pairs) if isinstance(value, _JsonObject) else () + + +def _children(value: object) -> Sequence[object]: + if isinstance(value, _JsonObject): + return tuple(child for _, child in value.pairs) + if isinstance(value, list): + return _JSON_ARRAY.validate_python(value) + return () + + +def _next_level(level: Sequence[object], _: None) -> Sequence[object]: + return tuple(chain.from_iterable(map(_children, level))) + + +def _forbidden_key(value: object, forbidden: frozenset[str]) -> str | None: + levels: Final = takewhile(bool, accumulate(repeat(None), _next_level, initial=(value,))) + keys: Final[Iterable[str]] = chain.from_iterable(map(_own_keys, chain.from_iterable(levels))) + return next((key for key in keys if key in forbidden), None) + + +def _plain(value: _JsonObject) -> Mapping[str, object]: + return dict(value.pairs) + + +def _object_json(pairs: Sequence[tuple[str, object]]) -> str: + fields: Final = ",".join( + f"{json.dumps(key)}:{json.dumps(child, default=_plain, separators=(',', ':'))}" for key, child in pairs + ) + return f"{{{fields}}}" + + +def _unparseable(api: str) -> IndexRefusal: + return IndexRefusal(400, f"LiteLLM could not parse the {api} request body as JSON, so it did not send it.") + + +def _index_reference(value: object, forbidden: frozenset[str]) -> IndexRefusal | None: + key: Final = _forbidden_key(value, forbidden) + if key is None: + return None + return IndexRefusal( + 400, + f"The request body uses '{key}', which can reach beyond the OpenSearch index in the URL path. " + f"Remove '{key}' from the body.", + ) + + +def _read_needed(cause: str, index_name: str) -> IndexRefusal: + return IndexRefusal( + 403, + f"{cause}, so the key needs read on index '{index_name}' as well as write. " + "Ask your administrator to add read to the key or its team.", + ) + + +def _document_reading_payload(api: str, payload: bytes, index: _IndexScope) -> IndexRefusal | None: + if index.can_read: + return None + parsed: Final = _parse(payload) + if parsed is None: + return _unparseable(api) + key: Final = next((key for key in _own_keys(parsed.value) if key in _DOCUMENT_READING_KEYS), None) + return None if key is None else _read_needed(f"'{key}' in an {api} body reads stored documents", index.name) + + +def _line_refusal(api: str, line: bytes, forbidden: frozenset[str]) -> IndexRefusal | None: + parsed: Final = _parse(line) + if parsed is None: + return _unparseable(api) + return _index_reference(parsed.value, forbidden) + + +def _whole_body_refusal(api: str, body: bytes) -> IndexRefusal | None: + if not body.strip(_JSON_BLANK): + return None + return _line_refusal(api, body, _OUT_OF_INDEX_KEYS) + + +def _pinned_index(metadata: _JsonObject, index: _IndexScope) -> _JsonObject | IndexRefusal | None: + unnamed: Final = _JsonObject(tuple(pair for pair in metadata.pairs if pair[0] != "_index")) + refusal: Final = _index_reference(unnamed, _BULK_METADATA_KEYS) + if refusal is not None: + return refusal + named: Final = tuple(value for name, value in metadata.pairs if name == "_index") + if not named: + return None + if any(value != index.name for value in named): + return IndexRefusal( + 400, + f"A _bulk action line sets '_index' to another index than '{index.name}' in the URL path. " + f"Remove '_index' or set it to '{index.name}'.", + ) + return _JsonObject( + tuple((name, index.provider_name if name == "_index" else value) for name, value in metadata.pairs) + ) + + +def _bulk_action(line: bytes) -> _BulkAction | None: + parsed: Final = _parse(line) + if parsed is None: + return None + match parsed.value: + case _JsonObject(pairs=((name, _JsonObject() as metadata),)) if name in _BULK_ACTIONS: + return _BulkAction(name=name, metadata=metadata, has_source=name != "delete") + case _: + return None + + +def _bulk_action_line(action: _BulkAction | None, line: bytes, index: _IndexScope) -> bytes | IndexRefusal: + if action is None: + return IndexRefusal( + 400, + "A _bulk action line must be one JSON object naming index, create, update, or delete. " + "LiteLLM did not send the request.", + ) + if not index.can_read and "_source" in _own_keys(action.metadata): + return _read_needed("'_source' in a _bulk action line returns stored documents", index.name) + pinned: Final = _pinned_index(action.metadata, index) + if pinned is None: + return line + if isinstance(pinned, IndexRefusal): + return pinned + return f"{{{json.dumps(action.name)}:{_object_json(pinned.pairs)}}}".encode() + + +def _bulk_entry(line: bytes, lines: Iterator[bytes], index: _IndexScope) -> tuple[bytes | IndexRefusal, ...]: + if not line.strip(_JSON_BLANK): + return (line,) + action: Final = _bulk_action(line) + rendered: Final = _bulk_action_line(action, line, index) + if isinstance(rendered, IndexRefusal) or action is None or not action.has_source: + return (rendered,) + source: Final = next(lines, None) + if source is None: + return (rendered,) + if action.name != "update": + return (rendered, source) + return (rendered, _document_reading_payload("_bulk update", source, index) or source) + + +def _bulk_lines(lines: Iterator[bytes], index: _IndexScope) -> Iterator[bytes | IndexRefusal]: + for line in lines: + yield from (entry := _bulk_entry(line, lines, index)) + if any(isinstance(rendered, IndexRefusal) for rendered in entry): + return + + +def _rendered_body(lines: Iterator[bytes | IndexRefusal]) -> bytes | IndexRefusal: + rendered: Final = tuple(lines) + refusal: Final = next((line for line in rendered if isinstance(line, IndexRefusal)), None) + if refusal is not None: + return refusal + return b"\n".join(line for line in rendered if isinstance(line, bytes)) + + +def _checked_body(operation: IndexOperation, body: bytes, index: _IndexScope) -> bytes | IndexRefusal: + match operation.body: + case "query" if operation.permission == "write" and not index.can_read: + return _read_needed(f"{operation.api} searches stored documents", index.name) + case "query" | "mget": + return _whole_body_refusal(operation.api, body) or body + case "update": + return _document_reading_payload(operation.api, body, index) or body + case "bulk": + return _rendered_body(_bulk_lines(iter(body.split(b"\n")), index)) + case "unscanned": + return body + + +def accepted_media_type(content_type: str | None) -> str | None: + media, _, parameters = (content_type or DEFAULT_CONTENT_TYPE).partition(";") + normalized: Final = media.strip().lower() + if normalized in ACCEPTED_MEDIA_TYPES and parameters.strip().lower() in _ACCEPTED_MEDIA_TYPE_PARAMETERS: + return normalized + return None + + +def _checked_request_body( + operation: IndexOperation | None, + query: Sequence[tuple[str, str]], + body: bytes, + index: _IndexScope, +) -> bytes | IndexRefusal: + if operation is None: + return body + parameter: Final = next((key for key, _ in query if key in _OUT_OF_INDEX_QUERY_PARAMETERS), None) + if parameter is not None: + return IndexRefusal( + 400, f"The {parameter} query parameter can reach beyond the OpenSearch index in the URL path. Remove it." + ) + source_parameter: Final = next((key for key, _ in query if key.startswith("_source")), None) + if operation.permission == "write" and not index.can_read and source_parameter is not None: + return _read_needed(f"The {source_parameter} query parameter returns stored documents", index.name) + try: + return _checked_body(operation, body, index) + except RecursionError: + return IndexRefusal(400, f"The {operation.api} request body is nested too deeply for LiteLLM to check.") + + +def upstream_url(api_base: str, provider_index: str, api: tuple[str, ...], query: Sequence[tuple[str, str]]) -> str: + path: Final = "/".join(quote(segment, safe="") for segment in (provider_index, *api)) + query_string: Final = "&".join(f"{quote(key, safe='')}={quote(value, safe='')}" for key, value in query) + return f"{api_base}/{path}?{query_string}" if query_string else f"{api_base}/{path}" + + +async def prepare_index_request( + config: "OpenSearchVectorStoreConfig", + method: str, + path: IndexPath, + operation: IndexOperation | None, + provider_index: str, + content_type: str | None, + query: Sequence[tuple[str, str]], + body: bytes, + store_params: Mapping[str, object], + can_read: bool, +) -> SignedIndexRequest | IndexRefusal: + media: Final = accepted_media_type(content_type) + if media is None: + return IndexRefusal( + 400, "The OpenSearch pass-through takes application/json or application/x-ndjson bodies in UTF-8." + ) + checked: Final = _checked_request_body( + operation, query, body, _IndexScope(name=path.index_name, provider_name=provider_index, can_read=can_read) + ) + if isinstance(checked, IndexRefusal): + return checked + api_base: Final = store_params.get("api_base") + try: + url: Final = upstream_url( + config.get_complete_url(api_base if isinstance(api_base, str) else None, store_params), + provider_index, + path.api, + query, + ) + signed: Final = await run_aws_signing(config.sign_http_request, method, url, checked, media, store_params) + except (AuthenticationError, BadRequestError, InternalServerError) as error: + return IndexRefusal(error.status_code, error.message) + return SignedIndexRequest( + url=url, media_type=media, headers=MappingProxyType(dict(signed.headers.items())), body=checked + ) diff --git a/litellm/llms/opensearch/vector_stores/transformation.py b/litellm/llms/opensearch/vector_stores/transformation.py index e153cdf43d8..f7ba28e7f3f 100644 --- a/litellm/llms/opensearch/vector_stores/transformation.py +++ b/litellm/llms/opensearch/vector_stores/transformation.py @@ -10,7 +10,7 @@ import httpx from pydantic import ConfigDict, Field, TypeAdapter, ValidationError from litellm._logging import verbose_logger -from litellm.exceptions import AuthenticationError, BadRequestError +from litellm.exceptions import AuthenticationError, BadRequestError, InternalServerError from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.optional_imports import ensure_optional_import from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -18,8 +18,12 @@ from litellm.llms.base_llm.vector_store.transformation import ( BaseQueryEmbeddingVectorStoreConfig, VectorStoreEmbeddingExecutor, ) -from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, SignsRequestsWithAWS, sign_aws_json_post -from litellm.llms.bedrock.web_identity_session_policy import OPENSEARCH_SESSION_POLICY, WebIdentitySessionPolicy +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, SignsRequestsWithAWS, sign_aws_request +from litellm.llms.bedrock.web_identity_session_policy import ( + OPENSEARCH_INDEX_SESSION_POLICY, + OPENSEARCH_SESSION_POLICY, + WebIdentitySessionPolicy, +) from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.bedrock import AwsAuthParams from litellm.types.router import GenericLiteLLMParams @@ -35,6 +39,7 @@ from litellm.types.vector_stores import ( ) if TYPE_CHECKING: + from botocore.awsrequest import AWSPreparedRequest from botocore.credentials import Credentials from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -81,6 +86,16 @@ def _unauthenticated(reason: str) -> AuthenticationError: ) +def _unavailable_credentials(error: Exception) -> InternalServerError: + verbose_logger.warning("OpenSearch vector store AWS credentials could not be loaded or refreshed: %s", error) + return InternalServerError( + message="LiteLLM could not load or refresh the AWS credentials that sign the OpenSearch request, " + "see the LiteLLM logs.", + llm_provider="opensearch", + model=None, + ) + + def _usable_credentials(credentials: object) -> "Credentials": from botocore.credentials import Credentials @@ -213,6 +228,10 @@ class _OpenSearchAWS(BaseAWSLLM): _web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = OPENSEARCH_SESSION_POLICY +class OpenSearchIndexAWS(BaseAWSLLM): + _web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = OPENSEARCH_INDEX_SESSION_POLICY + + class OpenSearchVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, SignsRequestsWithAWS): def __init__(self, aws: BaseAWSLLM | None = None) -> None: self._aws: Final = aws if aws is not None else _OpenSearchAWS() @@ -243,6 +262,7 @@ class OpenSearchVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, SignsRequ if bool(auth.aws_access_key_id) != bool(auth.aws_secret_access_key): raise _unauthenticated("aws_access_key_id and aws_secret_access_key must be set together") ensure_optional_import("botocore") + from botocore.credentials import Credentials from botocore.exceptions import ( CredentialRetrievalError, NoCredentialsError, @@ -251,12 +271,21 @@ class OpenSearchVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, SignsRequ ) try: - return _usable_credentials(self._aws.resolve_credentials(auth, region)) + credentials: Final = _usable_credentials(self._aws.resolve_credentials(auth, region)) except CredentialRetrievalError as error: verbose_logger.warning("OpenSearch vector store AWS credential source failed: %s", error) raise _unauthenticated("the AWS credential source failed, see the LiteLLM logs") from None + except OSError as error: + raise _unavailable_credentials(error) from None except (NoCredentialsError, PartialCredentialsError, ProfileNotFound) as error: raise _unauthenticated(str(error)) from None + try: + frozen: Final = credentials.get_frozen_credentials() + except (CredentialRetrievalError, OSError, RuntimeError) as error: + raise _unavailable_credentials(error) from None + if frozen.access_key is None or frozen.secret_key is None: + raise _unauthenticated("none were found") + return Credentials(frozen.access_key, frozen.secret_key, frozen.token, account_id=frozen.account_id) def get_complete_url(self, api_base: str | None, litellm_params: Mapping[str, object]) -> str: parsed: Final = urlsplit(api_base or "") @@ -368,25 +397,28 @@ class OpenSearchVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, SignsRequ api_base: str, api_key: str | None = None, ) -> tuple[dict[str, str], bytes]: # mutable-ok: the shared HTTP handler requires writable headers - region: Final = self._signing_region(optional_params, api_base) - body: Final = json.dumps(request_data) - signed: Final = sign_aws_json_post( - get_credentials=partial(self._signing_credentials, optional_params, region), - service_name=aws_signing_target(api_base)[0], - aws_region_name=region, - url=api_base, - body=body, - headers={ - "Content-Type": "application/json", - "x-amz-content-sha256": hashlib.sha256(body.encode()).hexdigest(), - }, - ) + body: Final = json.dumps(request_data).encode() + signed: Final = self.sign_http_request("POST", api_base, body, "application/json", optional_params) signed_headers: Final = dict(signed.headers.items()) signed_names: Final = frozenset(name.lower() for name in signed_headers) return { **{name: value for name, value in headers.items() if name.lower() not in signed_names}, **signed_headers, - }, body.encode() + }, body + + def sign_http_request( + self, method: str, url: str, body: bytes, content_type: str, litellm_params: Mapping[str, object] + ) -> "AWSPreparedRequest": + region: Final = self._signing_region(litellm_params, url) + return sign_aws_request( + method=method, + get_credentials=partial(self._signing_credentials, litellm_params, region), + service_name=aws_signing_target(url)[0], + aws_region_name=region, + url=url, + body=body, + headers={"Content-Type": content_type, "x-amz-content-sha256": hashlib.sha256(body).hexdigest()}, + ) def transform_search_vector_store_response( self, response: httpx.Response, litellm_logging_obj: "LiteLLMLoggingObj" diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index 680b0e6ede4..85f71defe3b 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -243,6 +243,7 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = ( "/mistral/", "/nvidia_nim/", "/openai/", + "/opensearch/", "/openai_passthrough/", "/tinyfish/", "/transcribe", diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index fde5be049eb..746d6a1e8d7 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -35987,6 +35987,386 @@ ] } }, + "/opensearch/{endpoint}": { + "delete": { + "description": "Call an OpenSearch index through the proxy, e.g. `POST /opensearch//_search`.\n\n`` is a LiteLLM vector store index on an OpenSearch vector store. The key or its team needs `read` or\n`write` on it in `allowed_vector_store_indexes`, and LiteLLM signs the call with the vector store's AWS identity.", + "operationId": "opensearch_proxy_route_opensearch__endpoint__delete", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + }, + "application/x-ndjson": { + "schema": { + "type": "string" + } + } + }, + "required": false + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + } + }, + "description": "OpenSearch's response, unmodified" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Opensearch Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "get": { + "description": "Call an OpenSearch index through the proxy, e.g. `POST /opensearch//_search`.\n\n`` is a LiteLLM vector store index on an OpenSearch vector store. The key or its team needs `read` or\n`write` on it in `allowed_vector_store_indexes`, and LiteLLM signs the call with the vector store's AWS identity.", + "operationId": "opensearch_proxy_route_opensearch__endpoint__get", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + }, + "application/x-ndjson": { + "schema": { + "type": "string" + } + } + }, + "required": false + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + } + }, + "description": "OpenSearch's response, unmodified" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Opensearch Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "head": { + "description": "Call an OpenSearch index through the proxy, e.g. `POST /opensearch//_search`.\n\n`` is a LiteLLM vector store index on an OpenSearch vector store. The key or its team needs `read` or\n`write` on it in `allowed_vector_store_indexes`, and LiteLLM signs the call with the vector store's AWS identity.", + "operationId": "opensearch_proxy_route_opensearch__endpoint__head", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + }, + "application/x-ndjson": { + "schema": { + "type": "string" + } + } + }, + "required": false + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + } + }, + "description": "OpenSearch's response, unmodified" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Opensearch Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "patch": { + "description": "Call an OpenSearch index through the proxy, e.g. `POST /opensearch//_search`.\n\n`` is a LiteLLM vector store index on an OpenSearch vector store. The key or its team needs `read` or\n`write` on it in `allowed_vector_store_indexes`, and LiteLLM signs the call with the vector store's AWS identity.", + "operationId": "opensearch_proxy_route_opensearch__endpoint__patch", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + }, + "application/x-ndjson": { + "schema": { + "type": "string" + } + } + }, + "required": false + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + } + }, + "description": "OpenSearch's response, unmodified" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Opensearch Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "post": { + "description": "Call an OpenSearch index through the proxy, e.g. `POST /opensearch//_search`.\n\n`` is a LiteLLM vector store index on an OpenSearch vector store. The key or its team needs `read` or\n`write` on it in `allowed_vector_store_indexes`, and LiteLLM signs the call with the vector store's AWS identity.", + "operationId": "opensearch_proxy_route_opensearch__endpoint__post", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + }, + "application/x-ndjson": { + "schema": { + "type": "string" + } + } + }, + "required": false + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + } + }, + "description": "OpenSearch's response, unmodified" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Opensearch Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "put": { + "description": "Call an OpenSearch index through the proxy, e.g. `POST /opensearch//_search`.\n\n`` is a LiteLLM vector store index on an OpenSearch vector store. The key or its team needs `read` or\n`write` on it in `allowed_vector_store_indexes`, and LiteLLM signs the call with the vector store's AWS identity.", + "operationId": "opensearch_proxy_route_opensearch__endpoint__put", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + }, + "application/x-ndjson": { + "schema": { + "type": "string" + } + } + }, + "required": false + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + } + }, + "description": "OpenSearch's response, unmodified" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Opensearch Proxy Route", + "tags": [ + "llm_passthrough" + ] + } + }, "/tinyfish/{endpoint}": { "get": { "description": "Pass-through for the TinyFish Agent API (goal-based web automation).\n\nForwarded endpoints:\n- POST /v1/automation/run \u2014 run to completion (blocking)\n- POST /v1/automation/run-async \u2014 submit a run, poll GET /v1/runs/{id} for the result\n- POST /v1/automation/run-sse \u2014 run with SSE progress events\n- GET /v1/runs/{id} \u2014 run status / result\n- POST /v1/runs/{id}/cancel \u2014 cancel a run\n\nEvery other Agent API endpoint (vault, wallet, browser profiles, and the GET /v1/runs\nlisting, which would let any caller discover other callers' run ids) returns 403: all\nproxy callers share one upstream key.\n\nCredential lookup order:\n1. passthrough_endpoint_router (config.yaml deployments with use_in_pass_through)\n2. TINYFISH_API_KEY environment variable\n\n[Docs](https://docs.litellm.ai/docs/pass_through/tinyfish)", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e53b5d4ad52..b3357579829 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -427,6 +427,7 @@ class LiteLLMRoutes(enum.Enum): "/bespoke", "/openrouter", "/milvus", + "/opensearch", "/gigachat", "/watsonx", "/nvidia_nim", diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 230dc7f6c58..d3978a58c43 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -2,7 +2,7 @@ import importlib.util import os import re import sys -from collections.abc import Collection, Iterator, Mapping +from collections.abc import Collection, Iterator, Mapping, Sequence from functools import lru_cache from itertools import accumulate, chain, repeat from logging import Logger @@ -1556,6 +1556,48 @@ def enforce_batch_limits_are_admin_only( ) +VECTOR_STORE_INDEX_GRANTS_METADATA_KEY: Final = "allowed_vector_store_indexes" + + +def _stored_vector_store_index_grants(existing_metadata: Mapping[str, object] | None) -> object: + if existing_metadata is None: + return None + return existing_metadata.get(VECTOR_STORE_INDEX_GRANTS_METADATA_KEY) or None + + +def _requested_vector_store_index_grants( + data: BatchLimitRequest, + top_level_grants: Sequence[AllowedVectorStoreIndexItem] | None, + stored_grants: object, +) -> object: + if top_level_grants is not None: + return [grant.model_dump() for grant in top_level_grants] or None + if "metadata" in data.model_fields_set: + return (data.metadata or EMPTY_MAPPING).get(VECTOR_STORE_INDEX_GRANTS_METADATA_KEY) or None + return stored_grants + + +def enforce_vector_store_index_grants_are_admin_only( + data: BatchLimitRequest, + top_level_grants: Sequence[AllowedVectorStoreIndexItem] | None, + existing_metadata: Mapping[str, object] | None, + user_api_key_dict: UserAPIKeyAuth, + entity: Literal["key", "team"], +) -> None: + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: + return + stored_grants: Final = _stored_vector_store_index_grants(existing_metadata) + if _requested_vector_store_index_grants(data, top_level_grants, stored_grants) == stored_grants: + return + raise HTTPException( + status_code=403, + detail={ + "error": f"Only proxy admins can set {VECTOR_STORE_INDEX_GRANTS_METADATA_KEY} on a {entity}. " + "It grants read and write access to vector store indexes, so the holder cannot widen it." + }, + ) + + def get_model_rate_limit_from_metadata( user_api_key_dict: UserAPIKeyAuth, metadata_accessor_key: Literal["team_metadata", "organization_metadata", "project_metadata"], diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index a52ad14e293..61ee4fd2439 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -144,7 +144,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg _safe_get_request_headers, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export _safe_get_request_query_params, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export _safe_set_request_parsed_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export - is_opaque_audio_pass_through_request, + is_opaque_pass_through_request, populate_request_with_path_params, read_raw_json_body, read_request_body, @@ -2128,7 +2128,7 @@ async def _read_request_body_deferring_parse_failure( must run (resolving identity onto the request's trace) before the 400 goes out; the caller re-raises the returned exception once identity is seeded. """ - if is_opaque_audio_pass_through_request( + if is_opaque_pass_through_request( route=get_request_route(request=request), content_type=safe_get_request_headers(request=request).get("content-type", ""), ): diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 08461aa3a44..80ba2f49745 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -19,6 +19,7 @@ from litellm.constants import ( LITELLM_CLIENT_DISCONNECT_DETAIL, LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB, + OPENSEARCH_PASS_THROUGH_ROUTE_PREFIX, ) from litellm.integrations.otel.runtime import phase_event from litellm.litellm_core_utils.core_helpers import as_object_list, has_attributes, require_attributes @@ -404,8 +405,10 @@ async def read_request_body(request: Request | None) -> dict: _read_request_body: Final = read_request_body -def is_opaque_audio_pass_through_request(route: str, content_type: str) -> bool: - """Azure Speech bodies (raw audio, multipart uploads) are forwarded byte for byte, so auth must not consume them.""" +def is_opaque_pass_through_request(route: str, content_type: str) -> bool: + """Azure Speech audio and every OpenSearch body are forwarded byte for byte, so auth must not parse them.""" + if route.startswith(f"{OPENSEARCH_PASS_THROUGH_ROUTE_PREFIX}/"): + return True media_type: Final = _normalize_media_type(content_type) return route.startswith(f"{AZURE_SPEECH_PASS_THROUGH_ROUTE_PREFIX}/") and ( media_type.startswith("audio/") or media_type == "multipart/form-data" diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 8a23262acc1..15747da63bc 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -105,7 +105,7 @@ _REDACTED_HEADER_VALUE: Final = "***REDACTED***" _DECODED_JSON: Final = TypeAdapter(object) _OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) _CREDENTIAL_HEADER_NAMES: Final = SpecialHeaders.litellm_credential_header_names() | frozenset( - {"cookie", "proxy-authorization"} + {"cookie", "proxy-authorization", "x-amz-security-token"} ) _TRANSPORT_ONLY_CREDENTIAL_KEYS: Final = frozenset({"provider_specific_header", "headers", "api_key"}) _REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object]) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 6c6118b2505..cea62d6c2e4 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -39,7 +39,10 @@ from litellm.proxy.auth.auth_checks import ( get_team_object, get_user_object, ) -from litellm.proxy.auth.auth_utils import enforce_batch_limits_are_admin_only +from litellm.proxy.auth.auth_utils import ( + enforce_batch_limits_are_admin_only, + enforce_vector_store_index_grants_are_admin_only, +) from litellm.proxy.auth.password_policy import ( validate_password_not_breached, validate_password_policy, @@ -594,6 +597,7 @@ async def new_user( if data.auto_create_key and isinstance(user_api_key_dict, UserAPIKeyAuth): enforce_batch_limits_are_admin_only(data, None, user_api_key_dict, "key") + enforce_vector_store_index_grants_are_admin_only(data, None, None, user_api_key_dict, "key") check_permissions_caller_permission( data=data, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index d54eda45160..69ae2c8f64a 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -62,6 +62,7 @@ from litellm.proxy.auth.auth_utils import ( abbreviate_api_key, enforce_batch_limits_are_admin_only, enforce_output_token_estimates_are_admin_only, + enforce_vector_store_index_grants_are_admin_only, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( @@ -1509,6 +1510,13 @@ async def _common_key_generation_helper( user_api_key_dict=user_api_key_dict, entity="key", ) + enforce_vector_store_index_grants_are_admin_only( + data=data, + top_level_grants=data.allowed_vector_store_indexes, + existing_metadata=None, + user_api_key_dict=user_api_key_dict, + entity="key", + ) if data.metadata is not None and data.metadata.get("service_account_id") is not None and data.team_id is None: await validate_team_id_used_in_service_account_request( @@ -2418,7 +2426,7 @@ async def generate_key_fn( - prompts: Optional[List[str]] - List of allowed prompts for the key. If specified, the key will only be able to use these specific prompts. - auto_rotate: Optional[bool] - Whether this key should be automatically rotated (regenerated) - rotation_interval: Optional[str] - How often to auto-rotate this key (e.g., '30s', '30m', '30h', '30d'). Required if auto_rotate=True. - - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint. + - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Only a proxy admin can set or change it. Create index, using `/v1/indexes` endpoint. - router_settings: Optional[UpdateRouterConfig] - key-specific router settings. Example - {"model_group_retry_policy": {"gpt-4": {"RateLimitErrorRetries": 5}}}. IF null or {} then no router settings. - access_group_ids: Optional[List[str]] - List of access group IDs to associate with the key. Access groups define which models a key can access. Example - ["access_group_1", "access_group_2"]. - budget_limits: Optional[list] - List of concurrent budget windows for the key. Each window specifies a budget_limit, time_period, and optional budget_duration. Example - [{"budget_limit": 10.0, "time_period": "1d"}, {"budget_limit": 50.0, "time_period": "7d"}]. @@ -2620,7 +2628,7 @@ async def generate_service_account_key_fn( - allowed_routes: Optional[list] - List of allowed routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/chat/completions", "/embeddings", "/keys/*"] - object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission. Examples: - - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint. + - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Only a proxy admin can set or change it. Create index, using `/v1/indexes` endpoint. 1. Allow users to turn on/off pii masking @@ -3339,6 +3347,13 @@ async def _process_single_key_update( user_api_key_dict=user_api_key_dict, entity="key", ) + enforce_vector_store_index_grants_are_admin_only( + data=update_key_request, + top_level_grants=update_key_request.allowed_vector_store_indexes, + existing_metadata=existing_key_row.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_VerificationToken.metadata is a bare dict + user_api_key_dict=user_api_key_dict, + entity="key", + ) await _check_single_key_update_team_permissions( user_api_key_dict=user_api_key_dict, @@ -3847,6 +3862,13 @@ async def _validate_update_key_data( user_api_key_dict=user_api_key_dict, entity="key", ) + enforce_vector_store_index_grants_are_admin_only( + data=data, + top_level_grants=data.allowed_vector_store_indexes, + existing_metadata=_existing_metadata if isinstance(_existing_metadata, dict) else None, # pyright: ignore[reportUnknownArgumentType] # stored metadata is a bare dict + user_api_key_dict=user_api_key_dict, + entity="key", + ) # Personal-key bypass: the caller both created the key AND still owns it # (user_id == caller). Checking only created_by would let a demoted admin @@ -4090,7 +4112,7 @@ async def update_key_fn( - object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission. - auto_rotate: Optional[bool] - Whether this key should be automatically rotated - rotation_interval: Optional[str] - How often to rotate this key (e.g., '30d', '90d'). Required if auto_rotate=True - - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint. + - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Only a proxy admin can set or change it. Create index, using `/v1/indexes` endpoint. - router_settings: Optional[UpdateRouterConfig] - key-specific router settings. Example - {"model_group_retry_policy": {"gpt-4": {"RateLimitErrorRetries": 5}}}. IF null or {} then no router settings. - access_group_ids: Optional[List[str]] - List of access group IDs to associate with the key. Access groups define which models a key can access. Example - ["access_group_1", "access_group_2"]. - budget_limits: Optional[list] - List of concurrent budget windows for the key. Each window specifies a budget_limit, time_period, and optional budget_duration. Example - [{"budget_limit": 10.0, "time_period": "1d"}, {"budget_limit": 50.0, "time_period": "7d"}]. @@ -6312,6 +6334,13 @@ async def _execute_virtual_key_regeneration( user_api_key_dict=user_api_key_dict, entity="key", ) + enforce_vector_store_index_grants_are_admin_only( + data=data, + top_level_grants=data.allowed_vector_store_indexes, + existing_metadata=_existing_key_metadata if isinstance(_existing_key_metadata, dict) else None, # pyright: ignore[reportUnknownArgumentType] # stored metadata is a bare dict + user_api_key_dict=user_api_key_dict, + entity="key", + ) await _validate_end_user_budget_id_change( requested_budget_id=_requested_end_user_budget_id(data), existing_budget_id=get_key_end_user_budget_id( diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 59fa8a18ac7..71258049adf 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -122,6 +122,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo from litellm.proxy.auth.auth_utils import ( enforce_batch_limits_are_admin_only, enforce_output_token_estimates_are_admin_only, + enforce_vector_store_index_grants_are_admin_only, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( @@ -1674,7 +1675,7 @@ async def new_team( - team_member_key_duration: Optional[str] - The duration for a team member's key. e.g. "1d", "1w", "1mo" - allowed_passthrough_routes: Optional[List[str]] - List of allowed pass through routes for the team. - denied_passthrough_routes: Optional[List[str]] - List of pass through routes the team's keys may not call, even if allowed. Applies together with each key's `denied_passthrough_routes`. - - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint. + - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Only a proxy admin can set or change it. Create index, using `/v1/indexes` endpoint. - secret_manager_settings: Optional[dict] - Secret manager settings for the team. [Docs](https://docs.litellm.ai/docs/secret_managers/overview) - router_settings: Optional[UpdateRouterConfig] - team-specific router settings. Example - {"model_group_retry_policy": {"gpt-4": {"RateLimitErrorRetries": 5}}}. IF null or {} then no router settings. - access_group_ids: Optional[List[str]] - List of access group IDs to associate with the team. Access groups define which models the team can access. Example - ["access_group_1", "access_group_2"]. @@ -1785,6 +1786,13 @@ async def new_team( user_api_key_dict=user_api_key_dict, entity="team", ) + enforce_vector_store_index_grants_are_admin_only( + data=data, + top_level_grants=data.allowed_vector_store_indexes, + existing_metadata=None, + user_api_key_dict=user_api_key_dict, + entity="team", + ) # Check if license is over limit total_teams: Final = await _team_db(prisma_client).count() @@ -2459,7 +2467,7 @@ async def update_team( - default_estimated_output_tokens_per_model: Optional[Dict[str, int]] - Per-model override of the above. Example: {"gpt-4": 4096, "gpt-3.5-turbo": 1024} - mcp_rpm_limit: Optional[Dict[str, int]] - Per-MCP-server RPM limit for this team, keyed by MCP server name (alias if set, else the configured name). Example: {"github": 100, "slack": 200}. Applied across all keys for this team. Example - update team TPM Limit - - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint. + - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Only a proxy admin can set or change it. Create index, using `/v1/indexes` endpoint. - secret_manager_settings: Optional[dict] - Secret manager settings for the team. [Docs](https://docs.litellm.ai/docs/secret_managers/overview) - router_settings: Optional[UpdateRouterConfig] - team-specific router settings. Example - {"model_group_retry_policy": {"gpt-4": {"RateLimitErrorRetries": 5}}}. IF null or {} then no router settings. - access_group_ids: Optional[List[str]] - List of access group IDs to associate with the team. Access groups define which models the team can access. Example - ["access_group_1", "access_group_2"]. @@ -2597,6 +2605,13 @@ async def update_team( user_api_key_dict=user_api_key_dict, entity="team", ) + enforce_vector_store_index_grants_are_admin_only( + data=data, + top_level_grants=data.allowed_vector_store_indexes, + existing_metadata=_existing_team_metadata if isinstance(_existing_team_metadata, dict) else None, # pyright: ignore[reportUnknownArgumentType] # stored metadata is a bare dict + user_api_key_dict=user_api_key_dict, + entity="team", + ) check_passthrough_routes_caller_permission( data, diff --git a/litellm/proxy/management_helpers/bulk_user_creation.py b/litellm/proxy/management_helpers/bulk_user_creation.py index cca37ebd7ed..e900941f67b 100644 --- a/litellm/proxy/management_helpers/bulk_user_creation.py +++ b/litellm/proxy/management_helpers/bulk_user_creation.py @@ -29,7 +29,10 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.auth_checks import invalidate_team_member_spend_state -from litellm.proxy.auth.auth_utils import enforce_batch_limits_are_admin_only +from litellm.proxy.auth.auth_utils import ( + enforce_batch_limits_are_admin_only, + enforce_vector_store_index_grants_are_admin_only, +) from litellm.proxy.auth.litellm_license import LicenseCheck from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler @@ -219,6 +222,7 @@ def _row_error(item: BulkNewUserItem, user_api_key_dict: UserAPIKeyAuth) -> str check_permissions_caller_permission(data=item, user_api_key_dict=user_api_key_dict) if item.auto_create_key: enforce_batch_limits_are_admin_only(item, None, user_api_key_dict, "key") + enforce_vector_store_index_grants_are_admin_only(item, None, None, user_api_key_dict, "key") except Exception as exc: # noqa: BLE001 # any validation failure is reported on this row only return _error_message(exc) return None diff --git a/litellm/proxy/middleware/billable_request_metrics_middleware.py b/litellm/proxy/middleware/billable_request_metrics_middleware.py index fd78b6b58da..95e5c9d1293 100644 --- a/litellm/proxy/middleware/billable_request_metrics_middleware.py +++ b/litellm/proxy/middleware/billable_request_metrics_middleware.py @@ -24,6 +24,7 @@ from typing import Final, Protocol, runtime_checkable from starlette.types import ASGIApp, Message, Receive, Scope, Send from litellm._logging import verbose_proxy_logger +from litellm.constants import OPENSEARCH_PASS_THROUGH_ROUTE_PREFIX from litellm.proxy._types import LiteLLMRoutes from litellm.proxy.auth.route_checks import RouteChecks @@ -101,16 +102,21 @@ _LLM_ROUTE_EXACT: Final[tuple[str, ...]] = ( # inference calls that write SpendLogs rows, so they bill. Anchored to the # routes enum so new providers are picked up without touching this module. # /langfuse forwards observability traffic, not inference: it writes no -# SpendLogs row and must not bill. -_NON_BILLABLE_PASSTHROUGH_PREFIXES: Final = frozenset({"/langfuse"}) +# SpendLogs row and must not bill. /opensearch bills only index searches, like +# /v1/vector_stores/{id}/search, and never its document writes or lookups. +_NON_BILLABLE_PASSTHROUGH_PREFIXES: Final = frozenset({"/langfuse", OPENSEARCH_PASS_THROUGH_ROUTE_PREFIX}) _PASSTHROUGH_PREFIXES: Final[tuple[str, ...]] = tuple( prefix for prefix in LiteLLMRoutes.mapped_pass_through_routes.value if prefix not in _NON_BILLABLE_PASSTHROUGH_PREFIXES ) +_OPENSEARCH_INDEX_SEARCH: Final = re.compile(rf"{re.escape(OPENSEARCH_PASS_THROUGH_ROUTE_PREFIX)}/[^/]+/_search") + def _classify_llm_route(path: str) -> str | None: + if path.startswith(f"{OPENSEARCH_PASS_THROUGH_ROUTE_PREFIX}/"): + return OPENSEARCH_PASS_THROUGH_ROUTE_PREFIX if _OPENSEARCH_INDEX_SEARCH.fullmatch(path) else None exact_match: Final = next((route for route in _LLM_ROUTE_EXACT if path == route), None) if exact_match is not None: return exact_match diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index ba8153ebb6d..5bdb2219194 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -159,11 +159,15 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( ) from litellm.proxy.utils import ProxyLogging as ProxyLoggingType from litellm.proxy.utils import is_known_model, normalize_route_for_root_path +from litellm.proxy.vector_store_endpoints.endpoints import build_request_data_from_managed_vector_store from litellm.proxy.vector_store_endpoints.utils import ( assert_proxy_admin_for_vector_store_index_management, assert_user_can_access_vector_store, + assert_vector_store_index_access, get_litellm_managed_vector_store, + has_vector_store_index_access, is_allowed_to_call_vector_store_endpoint, + normalize_litellm_params, ) from litellm.secret_managers.main import get_secret_str, normalize_nonempty_secret_str, str_to_bool from litellm.types.passthrough_endpoints.eu_jev import EuJevDecisionRequest, EuJevDecisionResponse @@ -182,8 +186,8 @@ from litellm.types.passthrough_endpoints.tinyfish import ( from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials from litellm.types.router import LiteLLMParamsTypedDict from litellm.types.utils import CallTypes, LlmProviders -from litellm.types.vector_stores import LiteLLM_ManagedVectorStore -from litellm.utils import ProviderConfigManager +from litellm.types.vector_stores import LiteLLM_ManagedVectorStore, LiteLLM_ManagedVectorStoreIndex +from litellm.utils import ProviderConfigManager, with_named_credential from .passthrough_endpoint_router import PassthroughEndpointRouter @@ -1056,6 +1060,131 @@ async def milvus_proxy_route( return received_value +def _registered_vector_store_index(index_name: str) -> LiteLLM_ManagedVectorStoreIndex: + registry: Final = litellm.vector_store_index_registry + index_object: Final = registry.get_vector_store_index_by_name(index_name) if registry is not None else None + if index_object is None: + raise HTTPException( + status_code=400, + detail=f"Index {index_name} is not a litellm managed vector store index. " + "Only litellm managed vector store indexes are supported.", + ) + return index_object + + +async def _opensearch_store_params( + index_object: LiteLLM_ManagedVectorStoreIndex, user_api_key_dict: UserAPIKeyAuth +) -> Mapping[str, object]: + registry: Final = litellm.vector_store_registry + store_name: Final = index_object.litellm_params.vector_store_name + named_stores: Final = tuple( + store + for store in (registry.vector_stores if registry is not None else ()) + if store.get("vector_store_name") == store_name + ) + if len(named_stores) > 1: + raise HTTPException( + status_code=400, + detail=f"{len(named_stores)} vector stores are named {store_name}, so LiteLLM cannot tell which one " + f"index {index_object.index_name} is on. Give each vector store a unique name.", + ) + if not named_stores: + raise HTTPException( + status_code=400, + detail=f"Vector store {store_name} of index {index_object.index_name} is not registered.", + ) + vector_store: Final = named_stores[0] + params: Final = with_named_credential( + build_request_data_from_managed_vector_store(normalize_litellm_params(vector_store)) + ) + provider: Final = params.get("custom_llm_provider") + if not isinstance(provider, str) or provider.split("/", 1)[0] != LlmProviders.OPENSEARCH.value: + raise HTTPException( + status_code=400, + detail=f"Index {index_object.index_name} is on vector store {store_name}, " + "which is not an OpenSearch vector store.", + ) + await assert_user_can_access_vector_store(vector_store=vector_store, user_api_key_dict=user_api_key_dict) + return params + + +@router.api_route( + "/opensearch/{endpoint:path}", + methods=["GET", "HEAD", "POST", "PUT", "DELETE", "PATCH"], + tags=["OpenSearch Pass-through", "pass-through"], + openapi_extra={ + "requestBody": { + "required": False, + "content": { + "application/json": {"schema": {"type": "object", "additionalProperties": True}}, + "application/x-ndjson": {"schema": {"type": "string"}}, + }, + } + }, + responses={ + 200: { + "description": "OpenSearch's response, unmodified", + "content": {"application/json": {"schema": {"type": "object", "additionalProperties": True}}}, + } + }, +) +async def opensearch_proxy_route( # noqa: ANN201 # returns create_pass_through_route's untyped handler result + endpoint: str, + request: Request, + fastapi_response: Response, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +): + """ + Call an OpenSearch index through the proxy, e.g. `POST /opensearch//_search`. + + `` is a LiteLLM vector store index on an OpenSearch vector store. The key or its team needs `read` or + `write` on it in `allowed_vector_store_indexes`, and LiteLLM signs the call with the vector store's AWS identity. + """ + from litellm.llms.opensearch.vector_stores.index_passthrough import ( + IndexRefusal, + index_operation, + parse_index_path, + prepare_index_request, + ) + from litellm.llms.opensearch.vector_stores.transformation import OpenSearchIndexAWS, OpenSearchVectorStoreConfig + + path: Final = parse_index_path(endpoint) + if isinstance(path, IndexRefusal): + raise HTTPException(status_code=path.status_code, detail=path.message) + operation: Final = index_operation(request.method, path.api) + assert_vector_store_index_access( + user_api_key_dict, path.index_name, operation.permission if operation is not None else None + ) + index_object: Final = _registered_vector_store_index(path.index_name) + store_params: Final = await _opensearch_store_params(index_object, user_api_key_dict) + prepared: Final = await prepare_index_request( + config=OpenSearchVectorStoreConfig(aws=OpenSearchIndexAWS()), + method=request.method, + path=path, + operation=operation, + provider_index=index_object.litellm_params.vector_store_index, + content_type=request.headers.get("content-type"), + query=tuple(request.query_params.multi_items()), + body=await request.body(), + store_params=store_params, + can_read=has_vector_store_index_access(user_api_key_dict, path.index_name, "read"), + ) + if isinstance(prepared, IndexRefusal): + raise HTTPException(status_code=prepared.status_code, detail=prepared.message) + endpoint_func: Final = create_pass_through_route( + endpoint=endpoint, + target=prepared.url, + custom_headers=prepared.headers, + custom_llm_provider=LlmProviders.OPENSEARCH.value, + keep_target_query=True, + upstream_headers_filter=lambda _forwarded: prepared.headers, + ) + setattr(request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, prepared.body) + setattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, {}) + safe_set_request_parsed_body(request, {}) + return await endpoint_func(request, fastapi_response, user_api_key_dict) + + async def is_streaming_request_fn(request: Request) -> bool: if request.method == "POST": content_type: Final = request.headers.get("content-type", None) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 53317010af4..cdcb42ab761 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -119,6 +119,7 @@ from litellm.proxy.common_utils.error_body_call_id import JSON_OBJECT, error_bod 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 _safe_get_request_headers, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export + is_opaque_pass_through_request, read_request_body, safe_get_request_headers, ) @@ -423,6 +424,10 @@ async def chat_completion_pass_through_endpoint( ) +def _is_provider_payload(request: Request) -> bool: + return is_opaque_pass_through_request(get_request_route(request), request.headers.get("content-type", "")) + + class HttpPassThroughEndpointHelpers(BasePassthroughUtils): @staticmethod def get_response_headers( @@ -1209,6 +1214,14 @@ def _unfiltered_headers(headers: Mapping[str, str]) -> Mapping[str, str]: return headers +def _caller_query_source( + keep_target_query: bool, query_params: Mapping[str, object] | None, request: Request +) -> Mapping[str, object]: + if keep_target_query: + return MappingProxyType({}) + return query_params or request.query_params + + def _json_object(value: object) -> Mapping[str, object]: return cast("Mapping[str, object]", value) if isinstance(value, Mapping) else MappingProxyType({}) # cast-ok: json @@ -1266,6 +1279,7 @@ async def pass_through_request( timeout: float | None = None, upstream_headers_filter: Callable[[Mapping[str, str]], Mapping[str, str]] = _unfiltered_headers, operation: GenAIOperation | None = None, + keep_target_query: bool = False, ): """ Pass through endpoint handler, makes the httpx request for pass-through endpoints and ensures logging hooks are called @@ -1288,6 +1302,7 @@ async def pass_through_request( 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 + keep_target_query: Send the target URL's query exactly as given, adding none of the caller's query params """ from litellm.exceptions import ModifyResponseException from litellm.litellm_core_utils.litellm_logging import Logging @@ -1326,7 +1341,10 @@ async def pass_through_request( parent_span=user_api_key_dict.parent_otel_span, ) - requested_query_params: dict | None = query_params or dict(request.query_params) or None + requested_query_params: dict | None = ( + dict(_caller_query_source(keep_target_query, query_params, request)) # pyright: ignore[reportUnknownArgumentType] # query_params keeps its legacy untyped dict annotation + or None + ) endpoint_type: Final[EndpointType] = HttpPassThroughEndpointHelpers.get_endpoint_type( str(url), custom_llm_provider @@ -2263,6 +2281,8 @@ async def _parse_request_data_by_content_type( Returns: Tuple of (query_params_data, custom_body_data, file_data, stream) """ + if _is_provider_payload(request): + return None, None, None, None content_type: Final = request.headers.get("content-type", "") query_params_data = None @@ -2333,6 +2353,7 @@ def create_pass_through_route( timeout: float | None = None, upstream_headers_filter: Callable[[Mapping[str, str]], Mapping[str, str]] = _unfiltered_headers, operation: GenAIOperation | None = None, + keep_target_query: bool = False, ): # check if target is an adapter.py or a url from litellm._uuid import uuid @@ -2482,6 +2503,7 @@ def create_pass_through_route( timeout=cast(float | None, param_timeout), upstream_headers_filter=upstream_headers_filter, operation=operation, + keep_target_query=keep_target_query, ) finally: if hasattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY): diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py index 8ce48fc3516..4dec4c7b38b 100644 --- a/litellm/proxy/vector_store_endpoints/utils.py +++ b/litellm/proxy/vector_store_endpoints/utils.py @@ -25,7 +25,7 @@ from litellm.utils import ProviderConfigManager _JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) -def _normalize_litellm_params( +def normalize_litellm_params( vector_store: LiteLLM_ManagedVectorStore, ) -> LiteLLM_ManagedVectorStore: litellm_params: Final = vector_store.get("litellm_params") @@ -46,6 +46,10 @@ def _is_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> bool: ) +def _key_metadata(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, object]: + return user_api_key_dict.metadata + + def assert_proxy_admin_for_vector_store_index_management( user_api_key_dict: UserAPIKeyAuth, *, @@ -278,7 +282,7 @@ async def get_litellm_managed_vector_store( vector_store_id=vector_store_id ) if vector_store is not None: - return _normalize_litellm_params(vector_store) + return normalize_litellm_params(vector_store) except Exception as e: verbose_proxy_logger.warning( "Failed to resolve vector store id=%s from registry: %s", @@ -310,7 +314,7 @@ async def get_litellm_managed_vector_store( ) if not rows: return None - return _normalize_litellm_params(LiteLLM_ManagedVectorStore(**rows[0].model_dump())) + return normalize_litellm_params(LiteLLM_ManagedVectorStore(**rows[0].model_dump())) except Exception as e: verbose_proxy_logger.warning( "Failed to resolve vector store id=%s from shared cache: %s", @@ -413,6 +417,34 @@ def check_vector_store_permission( return False +def vector_store_endpoint_forbidden(index_name: str) -> HTTPException: + return HTTPException( + status_code=403, + detail=( + f"User does not have permission to call vector store endpoint {index_name}. " + "Ask your administrator to add the necessary permissions to your API key/Team." + ), + ) + + +def has_vector_store_index_access( + user_api_key_dict: UserAPIKeyAuth, index_name: str, permission: Literal["read", "write"] | None +) -> bool: + return _is_proxy_admin(user_api_key_dict) or ( + permission is not None + and check_vector_store_permission( + index_name, permission, _key_metadata(user_api_key_dict), user_api_key_dict.team_metadata + ) + ) + + +def assert_vector_store_index_access( + user_api_key_dict: UserAPIKeyAuth, index_name: str, permission: Literal["read", "write"] | None +) -> None: + if not has_vector_store_index_access(user_api_key_dict, index_name, permission): + raise vector_store_endpoint_forbidden(index_name) + + def is_allowed_to_call_vector_store_endpoint( provider: LlmProviders, index_name: str, @@ -432,7 +464,7 @@ def is_allowed_to_call_vector_store_endpoint( ): return True # check what allowed permissions are for the key - key_metadata: Final = user_api_key_dict.metadata + key_metadata: Final = _key_metadata(user_api_key_dict) team_metadata: Final = user_api_key_dict.team_metadata provider_config: Final = ProviderConfigManager.get_provider_vector_stores_config(provider=provider) @@ -478,14 +510,7 @@ def is_allowed_to_call_vector_store_endpoint( break if permission_type is None: - raise HTTPException( - status_code=403, - detail=( - f"User does not have permission to call vector store endpoint " - f"{index_name}. Ask your administrator to add the necessary " - "permissions to your API key/Team." - ), - ) + raise vector_store_endpoint_forbidden(index_name) # Check if key has specific permission for allowed_vector_store_indexes has_permission: Final = check_vector_store_permission( @@ -496,10 +521,7 @@ def is_allowed_to_call_vector_store_endpoint( ) if not has_permission: - raise HTTPException( - status_code=403, - detail=f"User does not have permission to call vector store endpoint {index_name}. Ask your administrator to add the necessary permissions to your API key/Team.", - ) + raise vector_store_endpoint_forbidden(index_name) return has_permission @@ -516,7 +538,7 @@ def is_allowed_to_call_vector_store_files_endpoint( ): return True - key_metadata: Final = user_api_key_dict.metadata + key_metadata: Final = _key_metadata(user_api_key_dict) team_metadata: Final = user_api_key_dict.team_metadata provider_config: Final = ProviderConfigManager.get_provider_vector_store_files_config(provider=provider) diff --git a/tests/unit/llms/opensearch/vector_stores/test_opensearch_transformation.py b/tests/unit/llms/opensearch/vector_stores/test_opensearch_transformation.py index 07836d57af0..c55393f2e06 100644 --- a/tests/unit/llms/opensearch/vector_stores/test_opensearch_transformation.py +++ b/tests/unit/llms/opensearch/vector_stores/test_opensearch_transformation.py @@ -4,6 +4,7 @@ import re from collections.abc import Callable, Mapping from datetime import datetime, timezone from pathlib import Path +from types import MappingProxyType from typing import Final from unittest.mock import MagicMock @@ -12,7 +13,7 @@ import openai import pytest from botocore.auth import SigV4Auth from botocore.awsrequest import AWSRequest -from botocore.credentials import Credentials +from botocore.credentials import Credentials, RefreshableCredentials from botocore.exceptions import CredentialRetrievalError import litellm @@ -20,6 +21,12 @@ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.opensearch.vector_stores.index_passthrough import ( + IndexPath, + IndexRefusal, + index_operation, + prepare_index_request, +) from litellm.llms.opensearch.vector_stores.transformation import OpenSearchVectorStoreConfig from litellm.types.llms.bedrock import AwsAuthParams from litellm.types.router import GenericLiteLLMParams @@ -286,6 +293,17 @@ _AMBIENT_AWS_ENV: Final = ( ) + +def _isolate_from_ambient_aws(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + for name in _AMBIENT_AWS_ENV: + monkeypatch.delenv(name, raising=False) + for name in ("AWS_CONFIG_FILE", "AWS_SHARED_CREDENTIALS_FILE", "BOTO_CONFIG"): + config_file: Path = tmp_path / name + config_file.write_text("") + monkeypatch.setenv(name, str(config_file)) + monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true") + + @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.parametrize( "credentials,env", @@ -313,13 +331,7 @@ async def test_missing_aws_credentials_are_a_401_before_embedding_or_any_request monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: - for name in _AMBIENT_AWS_ENV: - monkeypatch.delenv(name, raising=False) - for name in ("AWS_CONFIG_FILE", "AWS_SHARED_CREDENTIALS_FILE", "BOTO_CONFIG"): - config_file: Path = tmp_path / name - config_file.write_text("") - monkeypatch.setenv(name, str(config_file)) - monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true") + _isolate_from_ambient_aws(monkeypatch, tmp_path) for name, value in env.items(): monkeypatch.setenv(name, value) executor: Final = RecordingEmbeddingExecutor() @@ -353,6 +365,107 @@ def test_failing_credential_source_is_a_401_that_keeps_its_output_out_of_the_rep assert CREDENTIAL_SOURCE_OUTPUT not in str(error.value) +@pytest.mark.parametrize( + "sign", + [ + pytest.param( + lambda config: config.validate_environment({}, GenericLiteLLMParams(api_base=SERVERLESS)), + id="vector-store-search", + ), + pytest.param( + lambda config: config.sign_http_request( + "GET", f"{SERVERLESS}/docs/_mapping", b"", "application/json", {"api_base": SERVERLESS} + ), + id="index-pass-through", + ), + ], +) +def test_a_missing_container_credentials_token_file_is_a_500( + sign: Callable[[OpenSearchVectorStoreConfig], object], monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + _isolate_from_ambient_aws(monkeypatch, tmp_path) + monkeypatch.setenv("AWS_CONTAINER_CREDENTIALS_FULL_URI", "http://127.0.0.1:9/v1/credentials") + monkeypatch.setenv("AWS_CONTAINER_AUTHORIZATION_TOKEN_FILE", str(tmp_path / "missing-token")) + + with pytest.raises(litellm.InternalServerError, match="could not load or refresh the AWS credentials") as error: + sign(OpenSearchVectorStoreConfig()) + + assert error.value.status_code == 500 + + +EXPIRED_METADATA: Final = MappingProxyType( + {"access_key": "AKIDEXAMPLE", "secret_key": "secret", "token": "token", "expiry_time": "2020-01-01T00:00:00Z"} +) + + +def _refresh_returning_expired_credentials() -> dict[str, str]: + return dict(EXPIRED_METADATA) + + +def _refresh_after_the_token_file_is_gone() -> dict[str, str]: + raise CredentialRetrievalError(provider="container-role", error_msg=CREDENTIAL_SOURCE_OUTPUT) + + +class _CredentialSourceThatCannotRefresh(BaseAWSLLM): + def __init__(self, refresh: Callable[[], dict[str, str]]) -> None: + super().__init__() + self._refresh = refresh + + def resolve_credentials(self, auth_params: AwsAuthParams, aws_region_name: str | None) -> Credentials: + return RefreshableCredentials.create_from_metadata( + metadata=dict(EXPIRED_METADATA), refresh_using=self._refresh, method="container-role" + ) + + +@pytest.mark.parametrize( + "refresh", [_refresh_returning_expired_credentials, _refresh_after_the_token_file_is_gone] +) +@pytest.mark.parametrize( + "sign", + [ + pytest.param( + lambda config: config.validate_environment({}, GenericLiteLLMParams(api_base=SERVERLESS)), + id="vector-store-search", + ), + pytest.param( + lambda config: config.sign_http_request( + "GET", f"{SERVERLESS}/docs/_mapping", b"", "application/json", {"api_base": SERVERLESS} + ), + id="index-pass-through", + ), + ], +) +def test_credentials_that_fail_to_refresh_at_signing_are_a_500( + sign: Callable[[OpenSearchVectorStoreConfig], object], refresh: Callable[[], dict[str, str]] +) -> None: + config: Final = OpenSearchVectorStoreConfig(aws=_CredentialSourceThatCannotRefresh(refresh)) + + with pytest.raises(litellm.InternalServerError, match="could not load or refresh the AWS credentials") as error: + sign(config) + + assert error.value.status_code == 500 + assert CREDENTIAL_SOURCE_OUTPUT not in str(error.value) + + +@pytest.mark.asyncio +async def test_the_index_pass_through_answers_500_when_credentials_fail_to_refresh() -> None: + refused: Final = await prepare_index_request( + config=OpenSearchVectorStoreConfig(aws=_CredentialSourceThatCannotRefresh(_refresh_after_the_token_file_is_gone)), + method="GET", + path=IndexPath(index_name="docs", api=("_mapping",)), + operation=index_operation("GET", ("_mapping",)), + provider_index="prod-docs", + content_type=None, + query=(), + body=b"", + store_params={"api_base": SERVERLESS}, + can_read=True, + ) + + assert isinstance(refused, IndexRefusal) + assert (refused.status_code, CREDENTIAL_SOURCE_OUTPUT in refused.message) == (500, False) + + @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.asyncio async def test_opensearch_index_name_picks_the_index_so_the_store_id_can_differ(asynchronous: bool) -> None: diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index f83d53ac2e5..cde93915e0c 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -19793,6 +19793,150 @@ async def test_bulk_key_update_batch_enqueued_token_limit_rejected_for_non_admin assert "Only proxy admins can set" in str(exc.value.detail) + +_INDEX_GRANT = [{"index_name": "docs", "index_permissions": ["read", "write"]}] + + +@pytest.mark.parametrize( + "label, top_level, sent, stored, allowed", + [ + ("granted top-level on a key with none stored", _INDEX_GRANT, _UNTOUCHED, None, False), + ("granted in metadata on a key with none stored", None, _INDEX_GRANT, None, False), + ("widened from read to read and write", None, _INDEX_GRANT, [{"index_name": "docs", "index_permissions": ["read"]}], False), + ("cleared by replacing the blob", None, None, _INDEX_GRANT, False), + ("resent unchanged in metadata", None, _INDEX_GRANT, _INDEX_GRANT, True), + ("resent unchanged top-level", _INDEX_GRANT, _UNTOUCHED, _INDEX_GRANT, True), + ("left untouched", None, _UNTOUCHED, _INDEX_GRANT, True), + ("cleared top-level with an empty list", [], _UNTOUCHED, _INDEX_GRANT, False), + ("empty top-level list on a key with none stored", [], _UNTOUCHED, None, True), + ], +) +def test_vector_store_index_grants_admin_gate_matrix( + label: str, + top_level: list[dict[str, object]] | None, + sent: object, + stored: list[dict[str, object]] | None, + allowed: bool, +) -> None: + from litellm.proxy.auth.auth_utils import ( + enforce_vector_store_index_grants_are_admin_only, + ) + + request_body = {} if sent is _UNTOUCHED else {"metadata": {} if sent is None else {"allowed_vector_store_indexes": sent}} + if top_level is not None: + request_body["allowed_vector_store_indexes"] = top_level + data = UpdateKeyRequest(key="sk-1", **request_body) + existing_metadata = None if stored is None else {"allowed_vector_store_indexes": stored} + + def _call(caller: UserAPIKeyAuth) -> None: + enforce_vector_store_index_grants_are_admin_only( + data=data, + top_level_grants=data.allowed_vector_store_indexes, + existing_metadata=existing_metadata, + user_api_key_dict=caller, + entity="key", + ) + + non_admin = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-non-admin", user_id="alice") + if allowed: + _call(non_admin) + else: + with pytest.raises(HTTPException) as exc: + _call(non_admin) + assert exc.value.status_code == 403 + assert "Only proxy admins can set allowed_vector_store_indexes" in str(exc.value.detail) + + _call(UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_fields", + [{"allowed_vector_store_indexes": _INDEX_GRANT}, {"metadata": {"allowed_vector_store_indexes": _INDEX_GRANT}}], + ids=["top-level", "metadata"], +) +async def test_generate_key_vector_store_index_grants_rejected_for_non_admin(request_fields: dict[str, object]) -> None: + """An index grant on a key opens document reads and writes on that index, so a + non-admin minting their own key cannot pick its grants.""" + with patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()): + with pytest.raises(HTTPException) as exc: + await _common_key_generation_helper( + data=GenerateKeyRequest(**request_fields), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-alice", + user_id="alice", + ), + litellm_changed_by=None, + team_table=None, + ) + assert int(getattr(exc.value, "status_code", 0)) == 403 + assert "Only proxy admins can set allowed_vector_store_indexes" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_update_key_vector_store_index_grants_unchanged_allows_non_admin_edit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A key owner renaming a key that carries admin-set index grants resends them + with the rest of the form, which must not count as setting them.""" + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + token = "b3b2c3d4e5f6789012345678901234567890123456789012345678901234abcd" + _wire_update_key_fn(monkeypatch, _estimate_key_row(token, {"allowed_vector_store_indexes": _INDEX_GRANT})) + + mock_request = MagicMock() + mock_request.query_params = {} + + result = await update_key_fn( + request=mock_request, + data=UpdateKeyRequest(key=token, key_alias="my-alias", metadata={"allowed_vector_store_indexes": _INDEX_GRANT}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-internal", + user_id="internal_user", + ), + litellm_changed_by=None, + ) + + assert result is not None + + +@pytest.mark.asyncio +async def test_update_key_vector_store_index_grants_widened_rejected_for_non_admin( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + token = "c3b2c3d4e5f6789012345678901234567890123456789012345678901234abcd" + _wire_update_key_fn( + monkeypatch, + _estimate_key_row(token, {"allowed_vector_store_indexes": [{"index_name": "docs", "index_permissions": ["read"]}]}), + ) + + mock_request = MagicMock() + mock_request.query_params = {} + + with pytest.raises(ProxyException) as exc: + await update_key_fn( + request=mock_request, + data=UpdateKeyRequest(key=token, allowed_vector_store_indexes=_INDEX_GRANT), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-internal", + user_id="internal_user", + ), + litellm_changed_by=None, + ) + + assert str(exc.value.code) == "403" + assert "Only proxy admins can set allowed_vector_store_indexes" in str(exc.value.message) + + @pytest.mark.asyncio async def test_execute_virtual_key_regeneration_stamps_settings_updated_at(): """Regenerate rewrites the key's config, so it must move settings_updated_at.""" diff --git a/tests/unit/proxy/middleware/test_billable_request_metrics_middleware.py b/tests/unit/proxy/middleware/test_billable_request_metrics_middleware.py index d2e961ad158..2488af7952c 100644 --- a/tests/unit/proxy/middleware/test_billable_request_metrics_middleware.py +++ b/tests/unit/proxy/middleware/test_billable_request_metrics_middleware.py @@ -105,6 +105,7 @@ def test_is_pure_asgi_not_base_http_middleware(): # SpendLogs-producing routes surfaced by the route-inventory audit ("/v1/search", (BillableCategory.LLM, "/search")), ("/v1/vector_stores/vs_1/search", (BillableCategory.LLM, "/search")), + ("/opensearch/docs/_search", (BillableCategory.LLM, "/opensearch")), ("/v1/rag/query", (BillableCategory.LLM, "/rag/query")), ("/rag/ingest", (BillableCategory.LLM, "/rag/ingest")), # Provider passthrough carries real inference and writes SpendLogs @@ -164,6 +165,18 @@ def test_classify_billable(path: str, expected: Tuple[BillableCategory, str]): "/v1beta/interactions/int_123/cancel", # observability passthrough writes no SpendLogs row "/langfuse/api/public/ingestion", + "/opensearch/docs/_doc", + "/opensearch/docs/_doc/_search", + "/opensearch/docs/_doc/1", + "/opensearch/docs/_bulk", + "/opensearch/docs/_update/1", + "/opensearch/docs/_count", + "/opensearch/docs/_mget", + "/opensearch/docs/_doc/search", + "/opensearch/docs/_update/search", + "/opensearch/docs/_create/embeddings", + "/opensearch/docs/_doc/completions", + "/opensearch/docs/_explain/rerank", # a bare provider prefix is not an inference call "/bedrock", "/eu_jev/v1/systemone/extra", 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 73eb1088d20..96d0a2ea8e0 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 @@ -1,13 +1,16 @@ import asyncio import base64 import contextlib +import hashlib import itertools import json import logging import os +import re import traceback from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping from datetime import datetime, timezone +from pathlib import Path from types import MappingProxyType, SimpleNamespace from typing import Final, Literal from unittest import mock @@ -17,6 +20,9 @@ from urllib.parse import parse_qs import httpx import pytest import respx +from botocore.auth import SigV4Auth +from botocore.awsrequest import AWSRequest +from botocore.credentials import Credentials from fastapi import FastAPI, HTTPException, Request, Response from fastapi.routing import APIRoute from fastapi.responses import StreamingResponse @@ -35,7 +41,7 @@ from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessin from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException, ModifyResponseException from litellm.proxy._types import AwsErrorTypeProxyException, ProxyException from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id -from litellm.proxy.utils import ProxyLogging +from litellm.proxy.utils import ProxyLogging, hash_token from litellm.passthrough.main import raise_for_provider_status from litellm.llms.bedrock.common_utils import BedrockError from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe @@ -10169,3 +10175,914 @@ class TestTinyFishRouteTimeout: monkeypatch.setattr(proxy_server, "general_settings", {"pass_through_request_timeout": 30}, raising=False) assert _tinyfish_route_timeout() is None + + +OPENSEARCH_HOST: Final = "abc123.us-east-1.aoss.amazonaws.com" +OPENSEARCH_ACCESS_KEY: Final = "AKIDOPENSEARCHPASSTHROUGH" +OPENSEARCH_SECRET_KEY: Final = "opensearch-passthrough-secret" +DOCS_UPSTREAM: Final = f"https://{OPENSEARCH_HOST}/prod-docs-v2" +JSON_TYPE: Final = "application/json" +NDJSON_TYPE: Final = "application/x-ndjson" +_SIGV4_AUTHORIZATION: Final = re.compile( + r"\AAWS4-HMAC-SHA256 Credential=(?P[^/]+)/\d{8}/(?P[^/]+)/(?P[^/]+)/aws4_request, " + r"SignedHeaders=(?P[^,]+), Signature=(?P[0-9a-f]{64})\Z" +) + + +def _verified_sigv4_scope(sent: httpx.Request) -> tuple[str, str]: + authorization: Final = _SIGV4_AUTHORIZATION.match(sent.headers["authorization"]) + assert authorization is not None, sent.headers["authorization"] + signed_names: Final = authorization["signed"].split(";") + assert {"content-type", "host", "x-amz-content-sha256", "x-amz-date"} <= set(signed_names), signed_names + assert sent.headers["x-amz-content-sha256"] == hashlib.sha256(sent.content).hexdigest() + received: Final = AWSRequest( + method=sent.method, + url=str(sent.url), + data=sent.content, + headers={name: sent.headers[name] for name in signed_names}, + ) + received.context["timestamp"] = sent.headers["x-amz-date"] + signer: Final = SigV4Auth( + Credentials(OPENSEARCH_ACCESS_KEY, OPENSEARCH_SECRET_KEY), authorization["service"], authorization["region"] + ) + expected: Final = signer.signature(signer.string_to_sign(received, signer.canonical_request(received)), received) + assert (authorization["key"], authorization["signature"]) == (OPENSEARCH_ACCESS_KEY, expected) + return authorization["service"], authorization["region"] + + +def _opensearch_caller(*permissions: str, index_name: str = "docs") -> UserAPIKeyAuth: + grant: Final = { + "allowed_vector_store_indexes": [{"index_name": index_name, "index_permissions": list(permissions)}] + } + return UserAPIKeyAuth(api_key="sk-virtual", metadata=grant) + + +def _act_as(monkeypatch: pytest.MonkeyPatch, caller: UserAPIKeyAuth) -> None: + from litellm.proxy.proxy_server import app + + monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, lambda: caller) + + +@pytest.fixture +def opensearch_client(monkeypatch: pytest.MonkeyPatch) -> Iterator[TestClient]: + from litellm.proxy.proxy_server import app + from litellm.types.vector_stores import IndexCreateLiteLLMParams, LiteLLM_ManagedVectorStoreIndex + from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry, VectorStoreRegistry + + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + indexes: Final = ( + ("docs", "prod-docs-v2", "opensearch-store"), + ("hr", "hr-private", "opensearch-store"), + ("no-credentials", "prod-docs-v3", "opensearch-store-without-credentials"), + ("azure-docs", "azure-docs", "azure-store"), + ) + monkeypatch.setattr( + litellm, + "vector_store_index_registry", + VectorStoreIndexRegistry( + [ + LiteLLM_ManagedVectorStoreIndex( + id=f"index-{name}", + index_name=name, + litellm_params=IndexCreateLiteLLMParams(vector_store_index=provider_index, vector_store_name=store), + ) + for name, provider_index, store in indexes + ] + ), + ) + monkeypatch.setattr( + litellm, + "vector_store_registry", + VectorStoreRegistry( + [ + { + "vector_store_id": "opensearch-store", + "vector_store_name": "opensearch-store", + "custom_llm_provider": "opensearch", + "litellm_params": { + "api_base": f"https://{OPENSEARCH_HOST}", + "aws_access_key_id": OPENSEARCH_ACCESS_KEY, + "aws_secret_access_key": OPENSEARCH_SECRET_KEY, + }, + }, + { + "vector_store_id": "opensearch-store-without-credentials", + "vector_store_name": "opensearch-store-without-credentials", + "custom_llm_provider": "opensearch", + "litellm_params": { + "api_base": f"https://{OPENSEARCH_HOST}", + "aws_profile_name": "opensearch-missing-profile", + }, + }, + { + "vector_store_id": "azure-store", + "vector_store_name": "azure-store", + "custom_llm_provider": "azure_ai", + "litellm_params": {"api_base": "https://search.example.search.windows.net"}, + }, + ] + ), + ) + yield TestClient(app) + + +class TestOpenSearchProxyRoute: + SEARCH_BODY: Final = b'{"query":{"match":{"text":"remote work"}}}' + HITS: Final = MappingProxyType( + { + "took": 3, + "timed_out": False, + "hits": { + "total": {"value": 1, "relation": "eq"}, + "hits": [{"_index": "prod-docs-v2", "_id": "1", "_score": 1.2, "_source": {"text": "Remote work"}}], + }, + } + ) + + @pytest.mark.parametrize("content_type", [JSON_TYPE, "Application/JSON; charset=UTF-8"]) + def test_a_read_key_search_is_signed_for_the_mapped_index_and_returns_the_hits( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, content_type: str + ) -> None: + _act_as(monkeypatch, _opensearch_caller("read")) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post(f"{DOCS_UPSTREAM}/_search").mock( + return_value=httpx.Response(200, json=dict(self.HITS)) + ) + response = opensearch_client.post( + "/opensearch/docs/_search", + content=self.SEARCH_BODY, + headers={"Content-Type": content_type, "Authorization": "Bearer sk-virtual"}, + ) + + assert (response.status_code, response.json()) == (200, dict(self.HITS)) + sent: Final = route.calls.last.request + assert (sent.content, sent.headers["content-type"]) == (self.SEARCH_BODY, JSON_TYPE) + assert _verified_sigv4_scope(sent) == ("aoss", "us-east-1") + + def test_a_team_grant_lets_the_key_search( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: + team_grant: Final = {"allowed_vector_store_indexes": [{"index_name": "docs", "index_permissions": ["read"]}]} + _act_as(monkeypatch, UserAPIKeyAuth(api_key="sk-virtual", team_id="team-a", team_metadata=team_grant)) + with respx.mock(assert_all_called=True) as upstream: + upstream.post(f"{DOCS_UPSTREAM}/_search").mock(return_value=httpx.Response(200, json=dict(self.HITS))) + response = opensearch_client.post( + "/opensearch/docs/_search", content=self.SEARCH_BODY, headers={"Content-Type": JSON_TYPE} + ) + + assert (response.status_code, response.json()) == (200, dict(self.HITS)) + + @pytest.mark.parametrize( + "method, api", + [ + ("POST", "_doc"), + ("PUT", "_doc/1"), + ("POST", "_doc/1"), + ("DELETE", "_doc/1"), + ("PUT", "_create/1"), + ("POST", "_update/1"), + ("POST", "_bulk"), + ("POST", "_delete_by_query"), + ("POST", "_update_by_query"), + ], + ) + def test_a_read_only_key_cannot_write( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, method: str, api: str + ) -> None: + _act_as(monkeypatch, _opensearch_caller("read")) + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route(host=OPENSEARCH_HOST).mock(return_value=httpx.Response(201, json={})) + response = opensearch_client.request( + method, + f"/opensearch/docs/{api}", + content=b'{"title":"Remote work"}', + headers={"Content-Type": JSON_TYPE}, + ) + + assert (response.status_code, route.called) == (403, False) + + @pytest.mark.parametrize( + "document", + [b'{"title":"Remote work","index":"hr-private"}', b'{"title":"Remote work","stream":true}'], + ids=["field-named-index", "field-named-stream"], + ) + def test_a_read_write_key_writes_the_document_as_sent( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, document: bytes + ) -> None: + _act_as(monkeypatch, _opensearch_caller("read", "write")) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post(f"{DOCS_UPSTREAM}/_doc").mock( + return_value=httpx.Response(201, json={"_id": "abc", "result": "created"}) + ) + response = opensearch_client.post( + "/opensearch/docs/_doc", content=document, headers={"Content-Type": JSON_TYPE} + ) + + assert (response.status_code, response.json()["result"]) == (201, "created") + assert route.calls.last.request.content == document + assert _verified_sigv4_scope(route.calls.last.request) == ("aoss", "us-east-1") + + @pytest.mark.parametrize( + "path, body, status_code", + [ + pytest.param("_update/1", b'{"doc":{},"_source":true}', 403, id="update-body-source"), + pytest.param("_update/1?_source=text", b'{"doc":{"title":"b"}}', 403, id="update-source-param"), + pytest.param( + "_update/1", + b'{"script":{"source":"throw new IllegalArgumentException(ctx._source.text)"}}', + 403, + id="update-script", + ), + pytest.param("_update/1", b'{"doc":{} /* c */ ,"_source":true}', 400, id="update-body-json-with-comment"), + pytest.param("_bulk", b'{"update":{"_id":"1","_source":true}}\n{"doc":{}}\n', 403, id="bulk-action-source"), + pytest.param( + "_bulk", + b'{"index":{"_id":"2"}}\n{"title":"a"}\n{"update":{"_id":"1"}}\n{"doc":{},"_source":true}\n', + 403, + id="bulk-update-payload-source", + ), + pytest.param("_bulk?_source_includes=text", b'{"update":{"_id":"1"}}\n{"doc":{}}\n', 403, id="bulk-param"), + pytest.param("_delete_by_query", b'{"query":{"prefix":{"text":"a"}}}', 403, id="delete-by-query"), + pytest.param("_update_by_query", b'{"query":{"match_all":{}}}', 403, id="update-by-query"), + ], + ) + def test_a_write_only_key_cannot_read_documents_back_through_a_write( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, path: str, body: bytes, status_code: int + ) -> None: + _act_as(monkeypatch, _opensearch_caller("write")) + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route(host=OPENSEARCH_HOST).mock( + return_value=httpx.Response(200, json={"get": {"_source": {"text": "Remote work"}}}) + ) + response = opensearch_client.post( + f"/opensearch/docs/{path}", content=body, headers={"Content-Type": NDJSON_TYPE} + ) + + assert (response.status_code, route.called) == (status_code, False) + assert "Remote work" not in response.text + + @pytest.mark.parametrize( + "caller, path, body", + [ + pytest.param( + _opensearch_caller("write"), "_update/1", b'{"doc":{"script":"a field","_source":"b"}}', id="write-only" + ), + pytest.param( + _opensearch_caller("read", "write"), "_update/1?_source=text", b'{"doc":{},"_source":true}', id="update" + ), + pytest.param( + _opensearch_caller("read", "write"), + "_bulk", + b'{"update":{"_id":"1","_source":true}}\n{"doc":{},"script":{"source":"ctx._source.n += 1"}}\n', + id="bulk", + ), + pytest.param( + _opensearch_caller("read", "write"), "_delete_by_query", b'{"query":{"term":{"a":1}}}', id="by-query" + ), + ], + ) + def test_a_write_that_reads_no_document_or_has_the_read_grant_is_sent_as_is( + self, + opensearch_client: TestClient, + monkeypatch: pytest.MonkeyPatch, + caller: UserAPIKeyAuth, + path: str, + body: bytes, + ) -> None: + _act_as(monkeypatch, caller) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post(url__startswith=f"{DOCS_UPSTREAM}/{path.partition('?')[0]}").mock( + return_value=httpx.Response(200, json={"result": "updated"}) + ) + response = opensearch_client.post( + f"/opensearch/docs/{path}", content=body, headers={"Content-Type": NDJSON_TYPE} + ) + + assert (response.status_code, route.calls.last.request.content) == (200, body) + + @pytest.mark.parametrize( + "caller, index_name", + [ + (_opensearch_caller("read", "write", index_name="hr"), "docs"), + (UserAPIKeyAuth(api_key="sk-virtual"), "docs"), + (UserAPIKeyAuth(api_key="sk-virtual"), "unregistered"), + ], + ids=["grant-on-another-index", "no-grant", "no-grant-unregistered-index"], + ) + def test_a_key_without_a_grant_on_the_index_is_refused( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, caller: UserAPIKeyAuth, index_name: str + ) -> None: + _act_as(monkeypatch, caller) + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route(host=OPENSEARCH_HOST).mock(return_value=httpx.Response(200, json={})) + response = opensearch_client.post( + f"/opensearch/{index_name}/_search", content=self.SEARCH_BODY, headers={"Content-Type": JSON_TYPE} + ) + + assert (response.status_code, route.called) == (403, False) + assert index_name in response.json()["detail"] + + @pytest.mark.parametrize( + "method, path", + [ + ("POST", "/opensearch/docs/_search/template"), + ("POST", "/opensearch/docs/_msearch"), + ("GET", "/opensearch/docs/_msearch"), + ("PUT", "/opensearch/docs/_mapping"), + ("GET", "/opensearch/docs/_settings"), + ("POST", "/opensearch/docs/_close"), + ("POST", "/opensearch/docs/_refresh"), + ("GET", "/opensearch/docs/_doc/1/x"), + ("PATCH", "/opensearch/docs/_doc/1"), + ("PUT", "/opensearch/docs"), + ("DELETE", "/opensearch/docs"), + ], + ) + def test_an_unlisted_api_is_for_proxy_admins_only( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, method: str, path: str + ) -> None: + _act_as(monkeypatch, _opensearch_caller("read", "write")) + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route(host=OPENSEARCH_HOST).mock(return_value=httpx.Response(200, json={})) + response = opensearch_client.request( + method, path, content=b'{"source":"{}"}', headers={"Content-Type": JSON_TYPE} + ) + + assert (response.status_code, route.called) == (403, False) + + @pytest.mark.parametrize( + "method, path, upstream_path, status", + [ + pytest.param("GET", "/opensearch/docs", "/prod-docs-v2", 200, id="get-index"), + pytest.param("HEAD", "/opensearch/docs", "/prod-docs-v2", 200, id="index-exists"), + pytest.param("HEAD", "/opensearch/docs/_doc/1", "/prod-docs-v2/_doc/1", 404, id="doc-exists"), + pytest.param("HEAD", "/opensearch/docs/_source/1", "/prod-docs-v2/_source/1", 200, id="source-exists"), + pytest.param( + "POST", "/opensearch/docs/_explain/_bulk1", "/prod-docs-v2/_explain/_bulk1", 200, id="id-named-bulk" + ), + ], + ) + def test_a_read_key_reaches_the_index_and_its_existence_checks( + self, + opensearch_client: TestClient, + monkeypatch: pytest.MonkeyPatch, + method: str, + path: str, + upstream_path: str, + status: int, + ) -> None: + _act_as(monkeypatch, _opensearch_caller("read")) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.route(method=method, host=OPENSEARCH_HOST, path=upstream_path).mock( + return_value=httpx.Response(status, content=b"" if method == "HEAD" else b"{}") + ) + response = opensearch_client.request( + method, path, content=b"" if method == "HEAD" else self.SEARCH_BODY, headers={"Content-Type": JSON_TYPE} + ) + + assert response.status_code == status + assert _verified_sigv4_scope(route.calls.last.request) == ("aoss", "us-east-1") + + def test_a_proxy_admin_can_call_an_unlisted_api( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: + settings: Final = {"prod-docs-v2": {"settings": {"index": {"number_of_shards": "2"}}}} + _act_as(monkeypatch, UserAPIKeyAuth(api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN)) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.get(f"{DOCS_UPSTREAM}/_settings").mock(return_value=httpx.Response(200, json=settings)) + response = opensearch_client.get("/opensearch/docs/_settings") + + assert (response.status_code, response.json()) == (200, settings) + assert _verified_sigv4_scope(route.calls.last.request) == ("aoss", "us-east-1") + + @pytest.mark.parametrize( + "index_name, detail", + [ + ("unknown", "Index unknown is not a litellm managed vector store index"), + ("azure-docs", "not an OpenSearch vector store"), + ], + ) + def test_an_index_that_is_not_a_managed_opensearch_index_is_a_400( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, index_name: str, detail: str + ) -> None: + _act_as(monkeypatch, _opensearch_caller("read", index_name=index_name)) + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route().mock(return_value=httpx.Response(200, json={})) + response = opensearch_client.post( + f"/opensearch/{index_name}/_search", content=self.SEARCH_BODY, headers={"Content-Type": JSON_TYPE} + ) + + assert (response.status_code, route.called) == (400, False) + assert detail in response.json()["detail"] + + def test_an_index_whose_store_name_two_stores_share_is_a_400( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.vector_stores.vector_store_registry import VectorStoreRegistry + + same_name: Final = ( + ("opensearch-store", f"https://{OPENSEARCH_HOST}"), + ("opensearch-store-copy", "https://other123.us-east-1.aoss.amazonaws.com"), + ) + monkeypatch.setattr( + litellm, + "vector_store_registry", + VectorStoreRegistry( + [ + { + "vector_store_id": store_id, + "vector_store_name": "opensearch-store", + "custom_llm_provider": "opensearch", + "litellm_params": {"api_base": api_base}, + } + for store_id, api_base in same_name + ] + ), + ) + _act_as(monkeypatch, _opensearch_caller("read")) + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route().mock(return_value=httpx.Response(200, json={})) + response = opensearch_client.post( + "/opensearch/docs/_search", content=self.SEARCH_BODY, headers={"Content-Type": JSON_TYPE} + ) + + assert (response.status_code, route.called) == (400, False) + assert "2 vector stores are named opensearch-store" in response.json()["detail"] + + @pytest.mark.parametrize( + "path", + [ + "/opensearch/docs/", + "/opensearch/docs//_search", + "/opensearch/docs/%2E/_search", + "/opensearch/docs/%2E%2E/hr/_search", + "/opensearch/docs/_search/%2E%2E/%2E%2E/hr/_search", + ], + ) + def test_a_path_with_an_empty_or_dot_segment_is_a_400( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, path: str + ) -> None: + _act_as(monkeypatch, _opensearch_caller("read", "write")) + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route(host=OPENSEARCH_HOST).mock(return_value=httpx.Response(200, json={})) + response = opensearch_client.post(path, content=self.SEARCH_BODY, headers={"Content-Type": JSON_TYPE}) + + assert (response.status_code, route.called) == (400, False) + + @pytest.mark.parametrize( + "api, content_type, body, detail", + [ + pytest.param( + "_search", + JSON_TYPE, + b'{"query":{"terms":{"user":{"index":"hr-private","id":"1","path":"followers"}}}}', + "'index'", + id="search-terms-lookup", + ), + pytest.param( + "_search", + JSON_TYPE, + b'{"query":{"more_like_this":{"like":[{"_index":"hr-private","_id":"1"}]}}}', + "'_index'", + id="search-more-like-this", + ), + pytest.param( + "_count", + JSON_TYPE, + b'{"query":{"terms":{"user":{"\\u0069ndex":"hr-private","id":"1","path":"followers"}}}}', + "'index'", + id="count-escaped-key", + ), + pytest.param( + "_search", + JSON_TYPE, + b'{"query":{"terms":{"user":{"index":"hr-private","id":"1","path":"f"}}},"query":{"match_all":{}}}', + "'index'", + id="search-duplicate-key", + ), + pytest.param( + "_delete_by_query", + JSON_TYPE, + b'{"query":{"terms":{"user":{"index":"hr-private","id":"1","path":"followers"}}}}', + "'index'", + id="delete-by-query-terms-lookup", + ), + pytest.param( + "_mget", + JSON_TYPE, + b'{"docs":[{"_id":"1"},{"_index":"hr-private","_id":"2"}]}', + "'_index'", + id="mget-docs", + ), + pytest.param( + "_bulk", + NDJSON_TYPE, + b'{"index":{"_id":"1"}}\n{"title":"a"}\n{"delete":{"_index":"hr-private","_id":"2"}}\n', + "'_index'", + id="bulk-action-index", + ), + pytest.param( + "_bulk", + NDJSON_TYPE, + b'{"index":{"_index":"prod-docs-v2","_id":"1"}}\n{"title":"a"}\n', + "'_index'", + id="bulk-action-names-the-provider-index", + ), + pytest.param( + "_bulk", + NDJSON_TYPE, + b'{"index":{"_index":"docs","_index":"hr-private","_id":"1"}}\n{"title":"a"}\n', + "'_index'", + id="bulk-action-duplicate-index", + ), + pytest.param( + "_bulk", + NDJSON_TYPE, + b'{"index":{"_id":"1"}}\n\n{"delete":{"_index":"hr-private","_id":"2"}}\n', + "'_index'", + id="bulk-blank-source-line", + ), + pytest.param( + "_bulk", + NDJSON_TYPE, + b'{"index":{"_id":"1","pipeline":"reroute"}}\n{"title":"a"}\n', + "'pipeline'", + id="bulk-action-pipeline", + ), + pytest.param("_bulk", NDJSON_TYPE, b"not json\n", "action line", id="bulk-unparseable-action"), + pytest.param( + "_search", + JSON_TYPE, + b'{"profile":true,"query":{"wrapper":{"query":"eyJ0ZXJtcyI6e319"}}}', + "'wrapper'", + id="search-wrapper-query", + ), + pytest.param( + "_search", + JSON_TYPE, + b'{"query":{"geo_shape":{"location":{"indexed_shape":{"id":"deu","path":"location"}}}}}', + "'indexed_shape'", + id="search-indexed-shape", + ), + pytest.param( + "_search", + JSON_TYPE, + b'{"query":{"match_all":{}},"search_pipeline":{"request_processors":[]}}', + "'search_pipeline'", + id="search-temporary-pipeline", + ), + pytest.param( + "_mget?source=%7B%22docs%22%3A%5B%7B%22_index%22%3A%22hr-private%22%2C%22_id%22%3A%221%22%7D%5D%7D" + "&source_content_type=application%2Fjson", + JSON_TYPE, + b"", + "source query parameter", + id="mget-source-query", + ), + pytest.param( + "_search", + "application/x-ndjson; multipart/form-data", + b'{"stream":true,"custom_body":{}}', + "application/json", + id="multipart-parameter", + ), + pytest.param( + "_doc", "application/json; boundary=x", b'{"title":"a"}', "application/json", id="content-type-parameter" + ), + pytest.param("_doc?pipeline=reroute", JSON_TYPE, b'{"title":"a"}', "pipeline", id="pipeline-query"), + pytest.param("_search", JSON_TYPE, b'{/*c*/"query":{"match_all":{}}}', "could not parse", id="comment"), + pytest.param("_search", "application/smile", b"\x3a\x29\x0a\x01", "application/json", id="smile"), + pytest.param("_search", "application/yaml", b"query:\n match_all: {}\n", "application/json", id="yaml"), + pytest.param( + "_search", JSON_TYPE, b'{"query":' + b"[" * 100_000 + b"]" * 100_000 + b"}", "nested", id="deep-nesting" + ), + ], + ) + def test_a_body_or_query_that_could_reach_another_index_is_a_400( + self, + opensearch_client: TestClient, + monkeypatch: pytest.MonkeyPatch, + api: str, + content_type: str, + body: bytes, + detail: str, + ) -> None: + _act_as(monkeypatch, _opensearch_caller("read", "write")) + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route(host=OPENSEARCH_HOST).mock(return_value=httpx.Response(200, json={})) + response = opensearch_client.post( + f"/opensearch/docs/{api}", content=body, headers={"Content-Type": content_type} + ) + + assert (response.status_code, route.called) == (400, False) + assert detail in response.json()["detail"] + + @pytest.mark.parametrize( + "api, body", + [ + pytest.param( + "_bulk", + b'{"index":{"_id":"1"}}\n{"title":"a","index":"hr-private"}\n{"delete":{"_id":"2"}}\n\n' + b'{"update":{"_id":"3"}}\r\n{"doc":{"_index":"hr-private"}}\n', + id="bulk", + ), + ], + ) + @pytest.mark.parametrize("content_type", [NDJSON_TYPE, "application/json"]) + def test_an_ndjson_body_without_index_references_is_sent_byte_for_byte( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, api: str, body: bytes, content_type: str + ) -> None: + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", SHARED_MASTER_KEY) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post(f"{DOCS_UPSTREAM}/{api}").mock( + return_value=httpx.Response(200, json={"errors": False, "items": []}) + ) + response = opensearch_client.post( + f"/opensearch/docs/{api}", + content=body, + headers={"Content-Type": content_type, "Authorization": f"Bearer {SHARED_MASTER_KEY}"}, + ) + + assert response.status_code == 200 + sent: Final = route.calls.last.request + assert (sent.content, sent.headers["content-type"]) == (body, content_type) + assert _verified_sigv4_scope(sent) == ("aoss", "us-east-1") + + @pytest.mark.parametrize( + "document", + [ + b'{"title":"Remote work","model":"another-model"}', + b'{"title":"Remote work","api_base":"https://elsewhere.example"}', + b'{"title":"Remote work","base_url":"https://elsewhere.example"}', + ], + ids=["model", "api_base", "base_url"], + ) + def test_a_document_field_named_like_an_llm_param_is_written_for_a_model_limited_key( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, document: bytes + ) -> None: + from litellm.proxy import proxy_server + + key: Final = "sk-opensearch-model-limited-writer" + grant: Final = {"allowed_vector_store_indexes": [{"index_name": "docs", "index_permissions": ["write"]}]} + cache: Final = DualCache() + cache.set_cache( + key=hash_token(key), value=UserAPIKeyAuth(token=hash_token(key), models=["gpt-x"], metadata=grant) + ) + monkeypatch.setattr(proxy_server, "master_key", SHARED_MASTER_KEY) + monkeypatch.setattr(proxy_server, "prisma_client", "test-database") + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post(f"{DOCS_UPSTREAM}/_doc").mock( + return_value=httpx.Response(201, json={"_id": "abc", "result": "created"}) + ) + response = opensearch_client.post( + "/opensearch/docs/_doc", + content=document, + headers={"Content-Type": JSON_TYPE, "Authorization": f"Bearer {key}"}, + ) + + assert (response.status_code, route.calls.last.request.content) == (201, document) + + def test_a_bulk_action_naming_the_path_index_is_sent_to_the_mapped_index( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: + _act_as(monkeypatch, _opensearch_caller("write")) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post(f"{DOCS_UPSTREAM}/_bulk").mock( + return_value=httpx.Response(200, json={"errors": False, "items": []}) + ) + response = opensearch_client.post( + "/opensearch/docs/_bulk", + content=b'{"index":{"_index":"docs","_id":"1"}}\n{"title":"a","_index":"docs"}\n' + b'{"delete":{"_id":"2"}}\n{"update":{"_id":"3","_index":"docs"}}\n{"doc":{"title":"b"}}\n', + headers={"Content-Type": NDJSON_TYPE}, + ) + + assert response.status_code == 200 + sent: Final = route.calls.last.request + assert sent.content == ( + b'{"index":{"_index":"prod-docs-v2","_id":"1"}}\n{"title":"a","_index":"docs"}\n' + b'{"delete":{"_id":"2"}}\n{"update":{"_id":"3","_index":"prod-docs-v2"}}\n{"doc":{"title":"b"}}\n' + ) + assert _verified_sigv4_scope(sent) == ("aoss", "us-east-1") + + @pytest.mark.parametrize( + "method, path, body, sent_query", + [ + pytest.param( + "GET", + "/opensearch/docs/_search?q=title:remote+work&routing=a&routing=b", + b"", + b"q=title%3Aremote%20work&routing=a&routing=b", + id="space-and-repeated-key", + ), + pytest.param( + "POST", + "/opensearch/docs/_search", + b'{"size":1,"query_params":{"routing":"z"}}', + b"", + id="body-query-params-ignored", + ), + pytest.param( + "POST", + "/opensearch/docs/_search?search_pipeline=hybrid-min-max", + b'{"query":{"hybrid":{"queries":[{"match":{"title":"remote"}},' + b'{"knn":{"embedding":{"vector":[0.1,0.2],"k":2,"filter":{"term":{"lang":"en"}}}}}]}}}', + b"search_pipeline=hybrid-min-max", + id="hybrid-query-with-named-search-pipeline", + ), + ], + ) + def test_the_query_string_is_signed_as_sent( + self, + opensearch_client: TestClient, + monkeypatch: pytest.MonkeyPatch, + method: str, + path: str, + body: bytes, + sent_query: bytes, + ) -> None: + _act_as(monkeypatch, _opensearch_caller("read")) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.route(method=method, host=OPENSEARCH_HOST, path="/prod-docs-v2/_search").mock( + return_value=httpx.Response(200, json=dict(self.HITS)) + ) + response = opensearch_client.request(method, path, content=body, headers={"Content-Type": JSON_TYPE}) + + assert response.status_code == 200 + sent: Final = route.calls.last.request + assert sent.url.query == sent_query + assert _verified_sigv4_scope(sent) == ("aoss", "us-east-1") + + def test_document_fields_named_like_request_params_leave_the_rate_limiter_alone( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 + from litellm.proxy.utils import InternalUsageCache + + limiter: Final = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(DualCache())) + + class RateLimit(CustomLogger): + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: + await limiter.async_pre_call_hook(user_api_key_dict, cache, data, call_type) + return data + + monkeypatch.setattr(litellm, "callbacks", [RateLimit()]) + grant: Final = {"index_name": "docs", "index_permissions": ["write"]} + _act_as( + monkeypatch, + UserAPIKeyAuth( + api_key="sk-opensearch-rate-limited-writer", + metadata={"allowed_vector_store_indexes": [grant], "model_rpm_limit": {"text-embedding-3-small": 1}}, + ), + ) + documents: Final = ( + b'{"text":"a","metadata":"handbook"}', + b'{"text":"b","metadata":["handbook"]}', + b'{"text":"c","model":"text-embedding-3-small"}', + b'{"text":"d","model":"text-embedding-3-small"}', + ) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post(f"{DOCS_UPSTREAM}/_doc").mock( + return_value=httpx.Response(201, json={"result": "created"}) + ) + statuses: Final = tuple( + opensearch_client.post( + "/opensearch/docs/_doc", content=document, headers={"Content-Type": JSON_TYPE} + ).status_code + for document in documents + ) + + assert statuses == (201, 201, 201, 201) + assert tuple(call.request.content for call in route.calls) == documents + + @pytest.mark.parametrize( + "document", [b'[{"text":"array body"}]', b'{"text":"\xff\xfe"}'], ids=["json-array", "invalid-utf-8"] + ) + def test_a_malformed_document_gets_the_opensearch_answer( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, document: bytes + ) -> None: + rejected: Final = {"error": {"type": "mapper_parsing_exception"}, "status": 400} + _act_as(monkeypatch, _opensearch_caller("write")) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post(f"{DOCS_UPSTREAM}/_doc").mock(return_value=httpx.Response(400, json=rejected)) + response = opensearch_client.post( + "/opensearch/docs/_doc", content=document, headers={"Content-Type": JSON_TYPE} + ) + + assert (response.status_code, response.json()) == (400, rejected) + assert route.calls.last.request.content == document + + def test_an_opensearch_error_keeps_its_status_and_body( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: + not_found: Final = {"error": {"type": "index_not_found_exception", "index": "prod-docs-v2"}, "status": 404} + _act_as(monkeypatch, _opensearch_caller("read")) + with respx.mock(assert_all_called=True) as upstream: + upstream.post(f"{DOCS_UPSTREAM}/_search").mock(return_value=httpx.Response(404, json=not_found)) + response = opensearch_client.post( + "/opensearch/docs/_search", content=self.SEARCH_BODY, headers={"Content-Type": JSON_TYPE} + ) + + assert (response.status_code, response.json()) == (404, not_found) + + def test_missing_aws_credentials_are_a_401_before_anything_is_sent( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + for name in ("AWS_CONFIG_FILE", "AWS_SHARED_CREDENTIALS_FILE"): + empty_file: Path = tmp_path / name + empty_file.write_text("") + monkeypatch.setenv(name, str(empty_file)) + _act_as(monkeypatch, _opensearch_caller("read", index_name="no-credentials")) + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route(host=OPENSEARCH_HOST).mock(return_value=httpx.Response(200, json={})) + response = opensearch_client.post( + "/opensearch/no-credentials/_search", content=self.SEARCH_BODY, headers={"Content-Type": JSON_TYPE} + ) + + assert (response.status_code, route.called) == (401, False) + + @pytest.mark.parametrize( + "method, api, action", + [ + ("GET", "_doc/1", "es:ESHttpGet"), + ("HEAD", "_doc/1", "es:ESHttpHead"), + ("PUT", "_doc/1", "es:ESHttpPut"), + ("DELETE", "_doc/1", "es:ESHttpDelete"), + ("POST", "_search", "es:ESHttpPost"), + ], + ) + def test_a_web_identity_session_allows_the_verb_sent_to_a_managed_domain( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, method: str, api: str, action: str + ) -> None: + from litellm.types.vector_stores import IndexCreateLiteLLMParams, LiteLLM_ManagedVectorStoreIndex + from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry, VectorStoreRegistry + + domain: Final = "search-docs-abc123.us-east-1.es.amazonaws.com" + monkeypatch.setattr( + litellm, + "vector_store_index_registry", + VectorStoreIndexRegistry( + [ + LiteLLM_ManagedVectorStoreIndex( + id="index-domain-docs", + index_name="docs", + litellm_params=IndexCreateLiteLLMParams( + vector_store_index="prod-docs-v2", vector_store_name="opensearch-domain" + ), + ) + ] + ), + ) + monkeypatch.setattr( + litellm, + "vector_store_registry", + VectorStoreRegistry( + [ + { + "vector_store_id": "opensearch-domain", + "vector_store_name": "opensearch-domain", + "custom_llm_provider": "opensearch", + "litellm_params": { + "api_base": f"https://{domain}", + "aws_web_identity_token": "oidc/google/litellm", + "aws_role_name": "arn:aws:iam::123456789012:role/litellm-role", + "aws_session_name": f"opensearch-index-{method}", + "aws_region_name": "us-east-1", + }, + } + ] + ), + ) + sts: Final = MagicMock() + sts.assume_role_with_web_identity.return_value = { + "Credentials": { + "AccessKeyId": "k", + "SecretAccessKey": "s", + "SessionToken": "t", + "Expiration": datetime(2099, 1, 1, tzinfo=timezone.utc), + }, + "PackedPolicySize": 0, + } + _act_as(monkeypatch, _opensearch_caller("read", "write")) + with ( + patch("boto3.client", return_value=sts), + patch("litellm.llms.bedrock.base_aws_llm.get_secret", return_value="oidc-jwt-token"), + respx.mock(assert_all_called=True) as upstream, + ): + upstream.route(method=method, host=domain).mock(return_value=httpx.Response(200, json={})) + response = opensearch_client.request( + method, f"/opensearch/docs/{api}", content=b"{}", headers={"Content-Type": JSON_TYPE} + ) + + policy: Final = json.loads(sts.assume_role_with_web_identity.call_args.kwargs["Policy"]) + granted: Final = frozenset(action for statement in policy["Statement"] for action in statement["Action"]) + assert (response.status_code, action in granted) == (200, True) diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 45abb880f92..0e3f1f33e19 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -108,6 +108,7 @@ PROTOCOL_CONSTRAINED_PASS_THROUGH_ROUTES: Final = MappingProxyType( "/laya/v1/systemone": frozenset({"POST"}), "/bespoke/v1/systemone": frozenset({"POST"}), "/eu_jev/v1/systemone": frozenset({"POST"}), + "/opensearch/{endpoint:path}": frozenset({"GET", "HEAD", "POST", "PUT", "DELETE", "PATCH"}), } ) diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index af463fa19d4..557ac72572f 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -8083,6 +8083,7 @@ async def test_add_litellm_data_to_request_refuses_a_model_config_credential_hol ("API-Key", True), ("Cookie", True), ("Proxy-Authorization", True), + ("X-Amz-Security-Token", True), ("anthropic-version", False), ("user-agent", False), ], diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index e2f932b44a2..00c1b855033 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -8683,7 +8683,7 @@ export interface paths { * - prompts: Optional[List[str]] - List of allowed prompts for the key. If specified, the key will only be able to use these specific prompts. * - auto_rotate: Optional[bool] - Whether this key should be automatically rotated (regenerated) * - rotation_interval: Optional[str] - How often to auto-rotate this key (e.g., '30s', '30m', '30h', '30d'). Required if auto_rotate=True. - * - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint. + * - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Only a proxy admin can set or change it. Create index, using `/v1/indexes` endpoint. * - router_settings: Optional[UpdateRouterConfig] - key-specific router settings. Example - {"model_group_retry_policy": {"gpt-4": {"RateLimitErrorRetries": 5}}}. IF null or {} then no router settings. * - access_group_ids: Optional[List[str]] - List of access group IDs to associate with the key. Access groups define which models a key can access. Example - ["access_group_1", "access_group_2"]. * - budget_limits: Optional[list] - List of concurrent budget windows for the key. Each window specifies a budget_limit, time_period, and optional budget_duration. Example - [{"budget_limit": 10.0, "time_period": "1d"}, {"budget_limit": 50.0, "time_period": "7d"}]. @@ -8998,7 +8998,7 @@ export interface paths { * - allowed_routes: Optional[list] - List of allowed routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/chat/completions", "/embeddings", "/keys/*"] * - object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission. * Examples: - * - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint. + * - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Only a proxy admin can set or change it. Create index, using `/v1/indexes` endpoint. * * * 1. Allow users to turn on/off pii masking @@ -9150,7 +9150,7 @@ export interface paths { * - object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission. * - auto_rotate: Optional[bool] - Whether this key should be automatically rotated * - rotation_interval: Optional[str] - How often to rotate this key (e.g., '30d', '90d'). Required if auto_rotate=True - * - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint. + * - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Only a proxy admin can set or change it. Create index, using `/v1/indexes` endpoint. * - router_settings: Optional[UpdateRouterConfig] - key-specific router settings. Example - {"model_group_retry_policy": {"gpt-4": {"RateLimitErrorRetries": 5}}}. IF null or {} then no router settings. * - access_group_ids: Optional[List[str]] - List of access group IDs to associate with the key. Access groups define which models a key can access. Example - ["access_group_1", "access_group_2"]. * - budget_limits: Optional[list] - List of concurrent budget windows for the key. Each window specifies a budget_limit, time_period, and optional budget_duration. Example - [{"budget_limit": 10.0, "time_period": "1d"}, {"budget_limit": 50.0, "time_period": "7d"}]. @@ -11715,6 +11715,64 @@ export interface paths { patch: operations["openrouter_proxy_route_openrouter__endpoint__patch"]; trace?: never; }; + "/opensearch/{endpoint}": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Opensearch Proxy Route + * @description Call an OpenSearch index through the proxy, e.g. `POST /opensearch//_search`. + * + * `` is a LiteLLM vector store index on an OpenSearch vector store. The key or its team needs `read` or + * `write` on it in `allowed_vector_store_indexes`, and LiteLLM signs the call with the vector store's AWS identity. + */ + get: operations["opensearch_proxy_route_opensearch__endpoint__get"]; + /** + * Opensearch Proxy Route + * @description Call an OpenSearch index through the proxy, e.g. `POST /opensearch//_search`. + * + * `` is a LiteLLM vector store index on an OpenSearch vector store. The key or its team needs `read` or + * `write` on it in `allowed_vector_store_indexes`, and LiteLLM signs the call with the vector store's AWS identity. + */ + put: operations["opensearch_proxy_route_opensearch__endpoint__put"]; + /** + * Opensearch Proxy Route + * @description Call an OpenSearch index through the proxy, e.g. `POST /opensearch//_search`. + * + * `` is a LiteLLM vector store index on an OpenSearch vector store. The key or its team needs `read` or + * `write` on it in `allowed_vector_store_indexes`, and LiteLLM signs the call with the vector store's AWS identity. + */ + post: operations["opensearch_proxy_route_opensearch__endpoint__post"]; + /** + * Opensearch Proxy Route + * @description Call an OpenSearch index through the proxy, e.g. `POST /opensearch//_search`. + * + * `` is a LiteLLM vector store index on an OpenSearch vector store. The key or its team needs `read` or + * `write` on it in `allowed_vector_store_indexes`, and LiteLLM signs the call with the vector store's AWS identity. + */ + delete: operations["opensearch_proxy_route_opensearch__endpoint__delete"]; + options?: never; + /** + * Opensearch Proxy Route + * @description Call an OpenSearch index through the proxy, e.g. `POST /opensearch//_search`. + * + * `` is a LiteLLM vector store index on an OpenSearch vector store. The key or its team needs `read` or + * `write` on it in `allowed_vector_store_indexes`, and LiteLLM signs the call with the vector store's AWS identity. + */ + head: operations["opensearch_proxy_route_opensearch__endpoint__head"]; + /** + * Opensearch Proxy Route + * @description Call an OpenSearch index through the proxy, e.g. `POST /opensearch//_search`. + * + * `` is a LiteLLM vector store index on an OpenSearch vector store. The key or its team needs `read` or + * `write` on it in `allowed_vector_store_indexes`, and LiteLLM signs the call with the vector store's AWS identity. + */ + patch: operations["opensearch_proxy_route_opensearch__endpoint__patch"]; + trace?: never; + }; "/organization/daily/activity": { parameters: { query?: never; @@ -17853,7 +17911,7 @@ export interface paths { * - team_member_key_duration: Optional[str] - The duration for a team member's key. e.g. "1d", "1w", "1mo" * - allowed_passthrough_routes: Optional[List[str]] - List of allowed pass through routes for the team. * - denied_passthrough_routes: Optional[List[str]] - List of pass through routes the team's keys may not call, even if allowed. Applies together with each key's `denied_passthrough_routes`. - * - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint. + * - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Only a proxy admin can set or change it. Create index, using `/v1/indexes` endpoint. * - secret_manager_settings: Optional[dict] - Secret manager settings for the team. [Docs](https://docs.litellm.ai/docs/secret_managers/overview) * - router_settings: Optional[UpdateRouterConfig] - team-specific router settings. Example - {"model_group_retry_policy": {"gpt-4": {"RateLimitErrorRetries": 5}}}. IF null or {} then no router settings. * - access_group_ids: Optional[List[str]] - List of access group IDs to associate with the team. Access groups define which models the team can access. Example - ["access_group_1", "access_group_2"]. @@ -18089,7 +18147,7 @@ export interface paths { * - default_estimated_output_tokens_per_model: Optional[Dict[str, int]] - Per-model override of the above. Example: {"gpt-4": 4096, "gpt-3.5-turbo": 1024} * - mcp_rpm_limit: Optional[Dict[str, int]] - Per-MCP-server RPM limit for this team, keyed by MCP server name (alias if set, else the configured name). Example: {"github": 100, "slack": 200}. Applied across all keys for this team. * Example - update team TPM Limit - * - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint. + * - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Only a proxy admin can set or change it. Create index, using `/v1/indexes` endpoint. * - secret_manager_settings: Optional[dict] - Secret manager settings for the team. [Docs](https://docs.litellm.ai/docs/secret_managers/overview) * - router_settings: Optional[UpdateRouterConfig] - team-specific router settings. Example - {"model_group_retry_policy": {"gpt-4": {"RateLimitErrorRetries": 5}}}. IF null or {} then no router settings. * - access_group_ids: Optional[List[str]] - List of access group IDs to associate with the team. Access groups define which models the team can access. Example - ["access_group_1", "access_group_2"]. @@ -70085,6 +70143,246 @@ export interface operations { }; }; }; + opensearch_proxy_route_opensearch__endpoint__get: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: { + content: { + "application/json": { + [key: string]: unknown; + }; + "application/x-ndjson": string; + }; + }; + responses: { + /** @description OpenSearch's response, unmodified */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": { + [key: string]: unknown; + }; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + opensearch_proxy_route_opensearch__endpoint__put: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: { + content: { + "application/json": { + [key: string]: unknown; + }; + "application/x-ndjson": string; + }; + }; + responses: { + /** @description OpenSearch's response, unmodified */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": { + [key: string]: unknown; + }; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + opensearch_proxy_route_opensearch__endpoint__post: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: { + content: { + "application/json": { + [key: string]: unknown; + }; + "application/x-ndjson": string; + }; + }; + responses: { + /** @description OpenSearch's response, unmodified */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": { + [key: string]: unknown; + }; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + opensearch_proxy_route_opensearch__endpoint__delete: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: { + content: { + "application/json": { + [key: string]: unknown; + }; + "application/x-ndjson": string; + }; + }; + responses: { + /** @description OpenSearch's response, unmodified */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": { + [key: string]: unknown; + }; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + opensearch_proxy_route_opensearch__endpoint__head: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: { + content: { + "application/json": { + [key: string]: unknown; + }; + "application/x-ndjson": string; + }; + }; + responses: { + /** @description OpenSearch's response, unmodified */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": { + [key: string]: unknown; + }; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + opensearch_proxy_route_opensearch__endpoint__patch: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: { + content: { + "application/json": { + [key: string]: unknown; + }; + "application/x-ndjson": string; + }; + }; + responses: { + /** @description OpenSearch's response, unmodified */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": { + [key: string]: unknown; + }; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; get_organization_daily_activity_organization_daily_activity_get: { parameters: { query?: { From 1d06a1ca69335cb4c925b4f1092e1deb6d345bfb Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 9 Nov 2026 11:03:39 -0800 Subject: [PATCH 03/46] test(e2e): repin the Together reasoning replay backend to MiniMax-M3.1-Flash (#50601) MiniMax-M3 reached its 2026-11-06 deprecation date in the cost map, which names MiniMax-M3.1-Flash as its successor --- tests/e2e/llm_translation/test_together_ai_e2e.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/e2e/llm_translation/test_together_ai_e2e.py b/tests/e2e/llm_translation/test_together_ai_e2e.py index 2d5d5f2d93a..7775c5c6ed2 100644 --- a/tests/e2e/llm_translation/test_together_ai_e2e.py +++ b/tests/e2e/llm_translation/test_together_ai_e2e.py @@ -7,9 +7,9 @@ backends are pinned because the registry has no flag for what they prove: ``enab and the ``{"reasoning": {"enabled": false}}`` toggle that ``reasoning_effort="none"`` maps to are Qwen hybrid-model contracts, and the structured-output case runs on that hybrid model with reasoning off, since a reasoning-only model can spend the whole token budget -thinking and return no content. MiniMax-M3 is the serverless model whose +thinking and return no content. MiniMax-M3.1-Flash is the serverless model whose template renders a replayed ``reasoning_content`` back into the prompt (Qwen and -DeepSeek silently drop it). MiniMax-M3 honors that replayed field on nearly every call, not +DeepSeek silently drop it). MiniMax-M3.1-Flash honors that replayed field on nearly every call, not every call (one miss in dozens of otherwise identical calls), so the replay case asks up to ``REPLAY_ATTEMPTS`` times and fails only when no answer carries the secret, which a proxy that strips the field guarantees. Requires TOGETHER_API_KEY on the proxy; no @@ -56,7 +56,7 @@ from pydantic import BaseModel pytestmark = pytest.mark.e2e HYBRID_REASONING_BACKEND = "together_ai/Qwen/Qwen3.5-9B" -REASONING_REPLAY_BACKEND = "together_ai/MiniMaxAI/MiniMax-M3" +REASONING_REPLAY_BACKEND = "together_ai/MiniMaxAI/MiniMax-M3.1-Flash" SECRET_PROMPT = "Remember this for later and reply with just OK." SECRET_REASONING = "The user told me their favorite color is chartreuse. I must remember it." From 457365f6273b66df4e0dcaa7582a4640cecc966b Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 9 Nov 2026 11:07:57 -0800 Subject: [PATCH 04/46] fix(proxy): build plugin session claims from FIPS provider services in fips_mode (#50544) * fix(proxy): build plugin session claims from FIPS provider services in fips_mode With general_settings.fips_mode on, the per-plugin key comes from the salt key KBKDF and the Fernet token is built from ssl.RAND_bytes, AES-128-CBC and hmac.digest. Standard mode is unchanged. * test(proxy): freeze the clock in the fips_mode plugin claim test --- litellm/proxy/plugin_routes.py | 30 ++++++++++++ tests/unit/proxy/test_plugin_routes.py | 63 +++++++++++++++++++++++++- 2 files changed, 92 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/plugin_routes.py b/litellm/proxy/plugin_routes.py index cfbcfefe9e8..966136f6fa7 100644 --- a/litellm/proxy/plugin_routes.py +++ b/litellm/proxy/plugin_routes.py @@ -18,6 +18,10 @@ Plugin iframe auth: litellm bearer token, so a compromised plugin learns only the caller's identity, never their credential. LITELLM_SALT_KEY itself is never shared with plugins — each plugin holds only its own derived key. + + With general_settings.fips_mode on, the per-plugin key is instead + derive_key_from_salt_key(b"litellm-plugin:session-claim:", plugin_name), and + the Fernet token is built from the FIPS provider's DRBG, AES-CBC and HMAC. """ import base64 @@ -25,11 +29,14 @@ import hashlib import hmac as _hmac import json import os +import ssl import time from collections.abc import Mapping from typing import Final from cryptography.fernet import Fernet, InvalidToken +from cryptography.hazmat.primitives import padding +from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes from fastapi import APIRouter, Depends, HTTPException, Request, Response from pydantic import ConfigDict, TypeAdapter, with_config from typing_extensions import NotRequired, ReadOnly, TypedDict @@ -37,6 +44,8 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy._types import PluginConfig, SpecialHeaders, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.encrypt_decrypt_utils import derive_key_from_salt_key +from litellm.proxy.common_utils.fips import fips_mode_active from litellm.types.llms.custom_http import httpxSpecialProvider router: Final = APIRouter() @@ -145,11 +154,30 @@ def _plugin_fernet(plugin_name: str) -> Fernet: A plugin possessing its own key cannot derive the master salt or forge claims intended for a different plugin. """ + if fips_mode_active(): + return Fernet(base64.urlsafe_b64encode(_fips_plugin_key(plugin_name))) salt: Final = os.getenv("LITELLM_SALT_KEY", "").encode() derived: Final = _hmac.new(salt, plugin_name.encode(), hashlib.sha256).digest() return Fernet(base64.urlsafe_b64encode(derived)) +_FIPS_PLUGIN_KEY_LABEL: Final = b"litellm-plugin:session-claim:" + + +def _fips_plugin_key(plugin_name: str) -> bytes: + return bytes.fromhex(derive_key_from_salt_key(_FIPS_PLUGIN_KEY_LABEL, plugin_name.encode())) + + +def _fips_fernet_token(key: bytes, data: bytes) -> str: + """Encrypt data into a standard Fernet token using only the FIPS provider's DRBG, AES-128-CBC and HMAC-SHA256.""" + iv: Final = ssl.RAND_bytes(16) + padder: Final = padding.PKCS7(128).padder() + encryptor: Final = Cipher(algorithms.AES(key[16:]), modes.CBC(iv)).encryptor() + ciphertext: Final = encryptor.update(padder.update(data) + padder.finalize()) + encryptor.finalize() + body: Final = b"\x80" + int(time.time()).to_bytes(8, "big") + iv + ciphertext + return base64.urlsafe_b64encode(body + _hmac.digest(key[:16], body, "sha256")).decode() + + _CLAIM_TTL_SECONDS: Final = 30 # identity claims expire after 30 s @@ -166,6 +194,8 @@ def issue_plugin_session_claim(plugin_name: str, user_id: str | None, user_role: "user_role": user_role or "", "exp": int(time.time()) + _CLAIM_TTL_SECONDS, } + if fips_mode_active(): + return _fips_fernet_token(_fips_plugin_key(plugin_name), json.dumps(claim).encode()) return _plugin_fernet(plugin_name).encrypt(json.dumps(claim).encode()).decode() diff --git a/tests/unit/proxy/test_plugin_routes.py b/tests/unit/proxy/test_plugin_routes.py index 7d3fbda14fd..23145f13156 100644 --- a/tests/unit/proxy/test_plugin_routes.py +++ b/tests/unit/proxy/test_plugin_routes.py @@ -12,10 +12,19 @@ Covers three bugs: """ import asyncio +import base64 +import hashlib +import hmac +import json +import os +import time +from typing import Final, NoReturn from unittest.mock import MagicMock +import _hashlib import pytest import respx +from cryptography.fernet import Fernet from starlette.requests import Request import litellm @@ -25,7 +34,12 @@ from litellm.proxy._types import ( PluginConfig, UserAPIKeyAuth, ) -from litellm.proxy.plugin_routes import list_plugins, register_plugins_from_config +from litellm.proxy.plugin_routes import ( + issue_plugin_session_claim, + list_plugins, + register_plugins_from_config, + verify_plugin_session_claim, +) def _admin() -> UserAPIKeyAuth: @@ -320,3 +334,50 @@ def test_plugin_proxy_refuses_non_admin_callers_without_contacting_the_plugin( assert response.status_code == 403 assert response.body == b"Plugin proxy access requires proxy_admin role." assert not upstream.called + + +def test_plugin_session_claim_opens_under_the_documented_hmac_plugin_key(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-salt") + plugin_key: Final = base64.urlsafe_b64encode(hmac.new(b"sk-salt", b"chat-ui", hashlib.sha256).digest()) + + claim: Final[dict[str, object]] = json.loads( + Fernet(plugin_key).decrypt(issue_plugin_session_claim("chat-ui", "user-1", "proxy_admin")) + ) + + assert (claim["plugin"], claim["user_id"], claim["user_role"]) == ("chat-ui", "user-1", "proxy_admin") + + +@pytest.mark.skipif( + not hasattr(_hashlib, "__file__"), reason="this Python links OpenSSL statically, out of ctypes reach" +) +@pytest.mark.usefixtures("fips_mode") +def test_fips_mode_plugin_session_claim_opens_under_the_kbkdf_key_without_os_urandom_or_hmac_new( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def outside_the_fips_provider(*_: object) -> NoReturn: + raise AssertionError("fips_mode must not draw on os.urandom or hmac.new") + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-" + bytes(range(32)).hex()) + monkeypatch.setattr(os, "urandom", outside_the_fips_provider) + monkeypatch.setattr(hmac, "new", outside_the_fips_provider) + monkeypatch.setattr(time, "time", lambda: 1_800_000_000.0) + fixed_input: Final = ( + (1).to_bytes(4, "big") + b"litellm-plugin:session-claim:" + b"\x00" + b"chat-ui" + (256).to_bytes(4, "big") + ) + plugin_key: Final = base64.urlsafe_b64encode(hmac.digest(bytes(range(32)), fixed_input, "sha256")) + token: Final = issue_plugin_session_claim("chat-ui", "user-1", "proxy_admin") + + claim: Final[dict[str, object]] = json.loads(Fernet(plugin_key).decrypt(token)) + + assert (claim["plugin"], claim["user_id"], claim["user_role"]) == ("chat-ui", "user-1", "proxy_admin") + assert verify_plugin_session_claim("chat-ui", token) == claim + + +@pytest.mark.usefixtures("fips_mode") +def test_fips_mode_refuses_a_plugin_session_claim_under_a_salt_key_that_is_not_hex( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-salt") + + with pytest.raises(ValueError, match="LITELLM_SALT_KEY must be at least 32 bytes of hex"): + issue_plugin_session_claim("chat-ui", "user-1", "proxy_admin") From 8b96ab28d06c39659e8e70a8ca0db7a20f3b6de9 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 9 Nov 2026 19:35:55 +0000 Subject: [PATCH 05/46] feat(anthropic): add claude-fable-5-2 and its Claude Platform on AWS mirror (#50715) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 80 +++++++++++++++++++ model_prices_and_context_window.json | 80 +++++++++++++++++++ 2 files changed, 160 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 78f5f8c29b7..b0997982e16 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -87352,5 +87352,85 @@ "supports_parallel_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true + }, + "claude-fable-5-2": { + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_creation_input_token_cost_batches": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 8e-06, + "input_cost_per_token_batches": 4e-06, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "output_cost_per_token_batches": 2e-05, + "prompt_cache_min_tokens": 512.0, + "provider_specific_entry": {"us":1.1,"fast":2}, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://platform.claude.com/docs/en/about-claude/pricing", + "supports_adaptive_thinking": true, + "supports_anthropic_compaction": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_fast_mode": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_speed": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "bedrock/claude_platform/claude-fable-5-2": { + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_creation_input_token_cost_batches": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 8e-06, + "input_cost_per_token_batches": 4e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "output_cost_per_token_batches": 2e-05, + "prompt_cache_min_tokens": 512.0, + "provider_specific_entry": {"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}, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing", + "supports_adaptive_thinking": true, + "supports_anthropic_compaction": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 78f5f8c29b7..b0997982e16 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -87352,5 +87352,85 @@ "supports_parallel_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true + }, + "claude-fable-5-2": { + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_creation_input_token_cost_batches": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 8e-06, + "input_cost_per_token_batches": 4e-06, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "output_cost_per_token_batches": 2e-05, + "prompt_cache_min_tokens": 512.0, + "provider_specific_entry": {"us":1.1,"fast":2}, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://platform.claude.com/docs/en/about-claude/pricing", + "supports_adaptive_thinking": true, + "supports_anthropic_compaction": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_fast_mode": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_speed": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "bedrock/claude_platform/claude-fable-5-2": { + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_creation_input_token_cost_batches": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 8e-06, + "input_cost_per_token_batches": 4e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "output_cost_per_token_batches": 2e-05, + "prompt_cache_min_tokens": 512.0, + "provider_specific_entry": {"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}, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing", + "supports_adaptive_thinking": true, + "supports_anthropic_compaction": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true } } From 60d679798310892c67113d5158051e18d688e004 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 9 Nov 2026 11:55:52 -0800 Subject: [PATCH 06/46] feat(cost-map): rolling registry update (Anthropic and Bedrock Fable 5.2, Azure realtime cache, Vertex, Friendli) (#50707) * feat(bedrock): add Claude Fable 5.2 and Mythos 5.2 offer index prices * fix(cost-map): absorb verified prices from #50559 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> * fix(cost-map): remove unsupported Gemini cache-write rates from #50504 Co-authored-by: Chester <42577861+superdiaodiao@users.noreply.github.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(cost-map): remove retired Friendli models from #50692 Co-authored-by: Lee-Si-Yoon <76679207+Lee-Si-Yoon@users.noreply.github.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(cost-map): give the eu and us-gov Mythos 5.2 profiles the base model's native structured output flag Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(cost-map): give the eu and us-gov Mythos 5.2 profiles the base model's native structured output flag * feat(anthropic): add claude-fable-5-2 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(anthropic): add tool and caching flags to claude-fable-5-2 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(anthropic): carry supports_mid_conversation_system on claude-fable-5-2 * fix(anthropic): declare Fable 5.2 mid-conversation system support Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(cost-map): deduplicate concurrently added Fable capability Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(cost-map): resolve merge conflict with main, keep main's claude-fable-5-2 row --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Chester <42577861+superdiaodiao@users.noreply.github.com> Co-authored-by: Lee-Si-Yoon <76679207+Lee-Si-Yoon@users.noreply.github.com> Co-authored-by: kerry --- ...odel_prices_and_context_window_backup.json | 417 +++++++++++++++--- model_prices_and_context_window.json | 417 +++++++++++++++--- .../llm_cost_calc/test_utils.py | 40 ++ 3 files changed, 752 insertions(+), 122 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b0997982e16..92b6ae6c469 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -5343,7 +5343,7 @@ "supports_tool_choice": true }, "azure/eu/gpt-4o-realtime-preview-2024-12-17": { - "cache_read_input_audio_token_cost": 2.5e-06, + "cache_read_input_audio_token_cost": 2.75e-06, "cache_read_input_token_cost": 2.75e-06, "input_cost_per_audio_token": 4.4e-05, "input_cost_per_token": 5.5e-06, @@ -5353,7 +5353,7 @@ "mode": "realtime", "output_cost_per_audio_token": 8.8e-05, "output_cost_per_token": 2.2e-05, - "source": "https://azure.microsoft.com/en-us/pricing/details/azure-openai/", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_modalities": [ "text", "audio" @@ -6691,6 +6691,7 @@ "supports_tool_choice": true }, "azure/gpt-realtime": { + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -6726,9 +6727,11 @@ "supports_parallel_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure" + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/" }, "azure/gpt-realtime-2025-08-28": { + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -6766,6 +6769,7 @@ "supports_tool_choice": true }, "azure/gpt-realtime-1.5": { + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -6802,9 +6806,11 @@ "supports_parallel_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure" + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/" }, "azure/gpt-realtime-1.5-2026-02-23": { + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -6925,6 +6931,8 @@ "supports_tool_choice": true }, "azure/gpt-realtime-mini": { + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "cache_read_input_image_token_cost": 8e-08, "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 6e-08, @@ -7033,6 +7041,7 @@ ] }, "azure/gpt-4o-realtime-preview-2024-12-17": { + "cache_read_input_audio_token_cost": 2.5e-06, "cache_read_input_token_cost": 2.5e-06, "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 5e-06, @@ -7042,7 +7051,7 @@ "mode": "realtime", "output_cost_per_audio_token": 8e-05, "output_cost_per_token": 2e-05, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_modalities": [ "text", "audio" @@ -12430,7 +12439,7 @@ "supports_tool_choice": true }, "azure/us/gpt-4o-realtime-preview-2024-12-17": { - "cache_read_input_audio_token_cost": 2.5e-06, + "cache_read_input_audio_token_cost": 2.75e-06, "cache_read_input_token_cost": 2.75e-06, "input_cost_per_audio_token": 4.4e-05, "input_cost_per_token": 5.5e-06, @@ -12440,7 +12449,7 @@ "mode": "realtime", "output_cost_per_audio_token": 8.8e-05, "output_cost_per_token": 2.2e-05, - "source": "https://azure.microsoft.com/en-us/pricing/details/azure-openai/", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_modalities": [ "text", "audio" @@ -29363,52 +29372,6 @@ "comment": "Open flagship GLM for long-horizon coding agents and million-token context work", "source": "https://api.friendli.ai/serverless/v1/models" }, - "friendliai/deepseek-ai/DeepSeek-V3.2": { - "litellm_provider": "friendliai", - "max_input_tokens": 163840, - "max_output_tokens": 163840, - "input_cost_per_token": 5e-07, - "output_cost_per_token": 1.5e-06, - "cache_read_input_token_cost": 2.5e-07, - "supports_prompt_caching": true, - "supports_reasoning": true, - "reasoning_effort_levels": [], - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_native_structured_output": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": false, - "supports_image_input": false, - "supports_video_input": false, - "mode": "chat", - "comment": "DeepSeek chat model for instruction following, coding, and analysis", - "source": "https://api.friendli.ai/serverless/v1/models" - }, - "friendliai/MiniMaxAI/MiniMax-M2.5": { - "litellm_provider": "friendliai", - "max_input_tokens": 196608, - "max_output_tokens": 196608, - "input_cost_per_token": 3e-07, - "output_cost_per_token": 1.2e-06, - "cache_read_input_token_cost": 6e-08, - "supports_prompt_caching": true, - "supports_reasoning": true, - "reasoning_effort_levels": [], - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_native_structured_output": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": false, - "supports_image_input": false, - "supports_video_input": false, - "mode": "chat", - "comment": "Prior MiniMax coding model for agent workflows, office edits, and automation", - "source": "https://api.friendli.ai/serverless/v1/models" - }, "friendliai/zai-org/GLM-5.1": { "litellm_provider": "friendliai", "max_input_tokens": 202752, @@ -30575,7 +30538,6 @@ "successor_model": "gemini-3.8-flash", "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, @@ -30634,7 +30596,6 @@ "output_cost_per_token_priority": 1.8e-05 }, "gemini-3.1-pro-preview": { - "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, @@ -30698,7 +30659,6 @@ "prompt_cache_min_tokens": 4096, "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, @@ -30750,7 +30710,6 @@ "vertex_ai/gemini-3-pro-preview": { "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, @@ -31194,7 +31153,6 @@ "supports_web_search": true }, "vertex_ai/gemini-3.1-pro-preview": { - "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, @@ -31258,7 +31216,6 @@ "prompt_cache_min_tokens": 4096, "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, @@ -67761,7 +67718,7 @@ "deepinfra/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B": { "max_output_tokens": 262144, "max_input_tokens": 262144, - "input_cost_per_token": 8.5e-08, + "input_cost_per_token": 7.5e-08, "output_cost_per_token": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", @@ -72924,6 +72881,7 @@ "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-realtime-2": { + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-07, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -72942,7 +72900,7 @@ "provider_specific_entry": { "realtime_ga_only": 1 }, - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/realtime" ], @@ -87353,6 +87311,343 @@ "supports_system_messages": true, "supports_tool_choice": true }, + "anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "eu.anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.1e-05, + "cache_creation_input_token_cost_above_1hr": 1.76e-05, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 8.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output_by_default": true, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "eu.anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.1e-05, + "cache_creation_input_token_cost_above_1hr": 1.76e-05, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 8.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "global.anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output_by_default": true, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "global.anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "us-gov.anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "us-gov.anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "us.anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.1e-05, + "cache_creation_input_token_cost_above_1hr": 1.76e-05, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 8.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output_by_default": true, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "us.anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.1e-05, + "cache_creation_input_token_cost_above_1hr": 1.76e-05, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 8.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, "claude-fable-5-2": { "cache_creation_input_token_cost": 1e-05, "cache_creation_input_token_cost_above_1hr": 1.6e-05, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b0997982e16..92b6ae6c469 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -5343,7 +5343,7 @@ "supports_tool_choice": true }, "azure/eu/gpt-4o-realtime-preview-2024-12-17": { - "cache_read_input_audio_token_cost": 2.5e-06, + "cache_read_input_audio_token_cost": 2.75e-06, "cache_read_input_token_cost": 2.75e-06, "input_cost_per_audio_token": 4.4e-05, "input_cost_per_token": 5.5e-06, @@ -5353,7 +5353,7 @@ "mode": "realtime", "output_cost_per_audio_token": 8.8e-05, "output_cost_per_token": 2.2e-05, - "source": "https://azure.microsoft.com/en-us/pricing/details/azure-openai/", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_modalities": [ "text", "audio" @@ -6691,6 +6691,7 @@ "supports_tool_choice": true }, "azure/gpt-realtime": { + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -6726,9 +6727,11 @@ "supports_parallel_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure" + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/" }, "azure/gpt-realtime-2025-08-28": { + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -6766,6 +6769,7 @@ "supports_tool_choice": true }, "azure/gpt-realtime-1.5": { + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -6802,9 +6806,11 @@ "supports_parallel_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure" + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/" }, "azure/gpt-realtime-1.5-2026-02-23": { + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -6925,6 +6931,8 @@ "supports_tool_choice": true }, "azure/gpt-realtime-mini": { + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "cache_read_input_image_token_cost": 8e-08, "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 6e-08, @@ -7033,6 +7041,7 @@ ] }, "azure/gpt-4o-realtime-preview-2024-12-17": { + "cache_read_input_audio_token_cost": 2.5e-06, "cache_read_input_token_cost": 2.5e-06, "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 5e-06, @@ -7042,7 +7051,7 @@ "mode": "realtime", "output_cost_per_audio_token": 8e-05, "output_cost_per_token": 2e-05, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_modalities": [ "text", "audio" @@ -12430,7 +12439,7 @@ "supports_tool_choice": true }, "azure/us/gpt-4o-realtime-preview-2024-12-17": { - "cache_read_input_audio_token_cost": 2.5e-06, + "cache_read_input_audio_token_cost": 2.75e-06, "cache_read_input_token_cost": 2.75e-06, "input_cost_per_audio_token": 4.4e-05, "input_cost_per_token": 5.5e-06, @@ -12440,7 +12449,7 @@ "mode": "realtime", "output_cost_per_audio_token": 8.8e-05, "output_cost_per_token": 2.2e-05, - "source": "https://azure.microsoft.com/en-us/pricing/details/azure-openai/", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_modalities": [ "text", "audio" @@ -29363,52 +29372,6 @@ "comment": "Open flagship GLM for long-horizon coding agents and million-token context work", "source": "https://api.friendli.ai/serverless/v1/models" }, - "friendliai/deepseek-ai/DeepSeek-V3.2": { - "litellm_provider": "friendliai", - "max_input_tokens": 163840, - "max_output_tokens": 163840, - "input_cost_per_token": 5e-07, - "output_cost_per_token": 1.5e-06, - "cache_read_input_token_cost": 2.5e-07, - "supports_prompt_caching": true, - "supports_reasoning": true, - "reasoning_effort_levels": [], - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_native_structured_output": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": false, - "supports_image_input": false, - "supports_video_input": false, - "mode": "chat", - "comment": "DeepSeek chat model for instruction following, coding, and analysis", - "source": "https://api.friendli.ai/serverless/v1/models" - }, - "friendliai/MiniMaxAI/MiniMax-M2.5": { - "litellm_provider": "friendliai", - "max_input_tokens": 196608, - "max_output_tokens": 196608, - "input_cost_per_token": 3e-07, - "output_cost_per_token": 1.2e-06, - "cache_read_input_token_cost": 6e-08, - "supports_prompt_caching": true, - "supports_reasoning": true, - "reasoning_effort_levels": [], - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_native_structured_output": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": false, - "supports_image_input": false, - "supports_video_input": false, - "mode": "chat", - "comment": "Prior MiniMax coding model for agent workflows, office edits, and automation", - "source": "https://api.friendli.ai/serverless/v1/models" - }, "friendliai/zai-org/GLM-5.1": { "litellm_provider": "friendliai", "max_input_tokens": 202752, @@ -30575,7 +30538,6 @@ "successor_model": "gemini-3.8-flash", "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, @@ -30634,7 +30596,6 @@ "output_cost_per_token_priority": 1.8e-05 }, "gemini-3.1-pro-preview": { - "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, @@ -30698,7 +30659,6 @@ "prompt_cache_min_tokens": 4096, "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, @@ -30750,7 +30710,6 @@ "vertex_ai/gemini-3-pro-preview": { "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, @@ -31194,7 +31153,6 @@ "supports_web_search": true }, "vertex_ai/gemini-3.1-pro-preview": { - "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, @@ -31258,7 +31216,6 @@ "prompt_cache_min_tokens": 4096, "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, @@ -67761,7 +67718,7 @@ "deepinfra/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B": { "max_output_tokens": 262144, "max_input_tokens": 262144, - "input_cost_per_token": 8.5e-08, + "input_cost_per_token": 7.5e-08, "output_cost_per_token": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", @@ -72924,6 +72881,7 @@ "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-realtime-2": { + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-07, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -72942,7 +72900,7 @@ "provider_specific_entry": { "realtime_ga_only": 1 }, - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/realtime" ], @@ -87353,6 +87311,343 @@ "supports_system_messages": true, "supports_tool_choice": true }, + "anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "eu.anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.1e-05, + "cache_creation_input_token_cost_above_1hr": 1.76e-05, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 8.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output_by_default": true, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "eu.anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.1e-05, + "cache_creation_input_token_cost_above_1hr": 1.76e-05, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 8.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "global.anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output_by_default": true, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "global.anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "us-gov.anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "us-gov.anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "us.anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.1e-05, + "cache_creation_input_token_cost_above_1hr": 1.76e-05, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 8.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output_by_default": true, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "us.anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.1e-05, + "cache_creation_input_token_cost_above_1hr": 1.76e-05, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 8.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, "claude-fable-5-2": { "cache_creation_input_token_cost": 1e-05, "cache_creation_input_token_cost_above_1hr": 1.6e-05, diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py index 1d2a5e4cc9b..f007d3740aa 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py @@ -1,4 +1,5 @@ import asyncio +import math import uuid from collections.abc import Mapping from types import MappingProxyType @@ -483,3 +484,42 @@ async def test_router_prices_balanced_columns_by_the_same_rules_as_flex( assert balanced_cost == pytest.approx(flex_cost) if not overridden_columns: assert balanced_cost == pytest.approx(_cost_at(PARITY_ROW, "_balanced")) + + +def _vertex_cost_with_cache_writes(model: str, prompt_tokens: int, cache_creation_tokens: int) -> float: + response: Final = ModelResponse( + model=model, + usage=Usage( + prompt_tokens=prompt_tokens, + completion_tokens=100, + total_tokens=prompt_tokens + 100, + cache_creation_input_tokens=cache_creation_tokens, + ), + ) + return litellm.completion_cost(completion_response=response, model=model, custom_llm_provider="vertex_ai") + + +@pytest.mark.parametrize( + "model", + ( + "gemini-2.5-pro", + "gemini-3.1-pro-preview", + "gemini-3.1-pro-preview-customtools", + "vertex_ai/gemini-3-pro-preview", + "vertex_ai/gemini-3.1-pro-preview", + "vertex_ai/gemini-3.1-pro-preview-customtools", + ), +) +def test_vertex_gemini_pro_bills_cache_writes_above_200k_at_the_above_200k_input_rate( + local_model_cost_map: None, model: str +) -> None: + rates: Final = litellm.get_model_info(model=model, custom_llm_provider="vertex_ai") + input_rate_above_200k: Final = rates.get("input_cost_per_token_above_200k_tokens") + output_rate_above_200k: Final = rates.get("output_cost_per_token_above_200k_tokens") + cache_creation_tokens: Final = 150_000 + below: Final = _vertex_cost_with_cache_writes(model, 190_000, cache_creation_tokens) + above: Final = _vertex_cost_with_cache_writes(model, 210_000, cache_creation_tokens) + + assert input_rate_above_200k is not None and output_rate_above_200k is not None + assert above > below + assert math.isclose(above, 210_000 * input_rate_above_200k + 100 * output_rate_above_200k) From abe3b11f7acd525ea926e79cd6b13f2d562b943a Mon Sep 17 00:00:00 2001 From: zoveki <349028329+zoveki@users.noreply.github.com> Date: Mon, 9 Nov 2026 12:05:37 -0800 Subject: [PATCH 07/46] fix(proxy): show committed spend in the remaining budget gauges (#49199) * fix(proxy): report live key and team spend to budget gauges within the key cache TTL * fix(prometheus): carry live spend for the user, org and customer budget gauges too The remaining budget gauges for users, orgs and customers read the same cached rows as the key and team gauges, so they read high for about a minute after a charge. Auth now carries the live counter for each budgeted entity in one snapshot that only the gauges read, floored at the spend auth loaded. The x-litellm-key-spend header and logging spend fields keep their current values. * fix(prometheus): show committed spend in the remaining budget gauges Reads twin counters that take only charged cost, so a worst-case budget hold never shows as spent and a request that finishes last cannot set a gauge back. The twins exist only while the Prometheus logger is registered, and every path that resets a spend counter drops them. Failed requests now set the org gauge from the carried org row and committed spend, not the cached org row. * fix(prometheus): count a charge the cached customer row already holds once * fix(prometheus): keep committed spend only for budgeted entities and drop it on spend resets Auth now reads committed spend through the admission MGET and names the counters a charge may move, so the master key and unbudgeted entities get no twin. Organization spend edits and the global spend reset drop their twins, and a long request no longer reads as charged. * fix(prometheus): read the committed org spend on failed requests without the auth object * test(prometheus): cover a failed request with no carried org * fix(prometheus): move the committed counters for charges with no auth on the row A billed background response is charged from its creator's stored attribution, after the request's auth is gone. Its charge moved the live counters but not the committed ones, so every gauge read one charge behind. It now moves the committed counters that already exist, and its own gauge adds its cost to them. * fix(prometheus): keep the gauges right for pass-through rows, overlapping charges and slow DB spend writes Pass-through rows carry the auth object, so a cold committed counter gets seeded there too. The gauges read a charge still waiting in the post-call pipeline as not counted, whatever other charges reached the counter. Committed counters outlive two DB spend writes. Cold seeds read the DB once per worker, and the global spend reset reads keys and teams in bounded pages. * fix(prometheus): count a charged failure's own cost in the org gauge A stream cut mid-response still charges for what it delivered, and that charge is usually still queued when the failure logs. The failure path passed a cost of 0, so the org gauge left that charge out until the worker's next request * fix(prometheus): skip twins for window-named ids, drop twins on delete, keep auth out of PostHog * fix(prometheus): keep the auth object off pass-through rows and drop user twins on every user delete Pass-through rows now carry only the committed spend in litellm_params, so loggers that dump the row's metadata never get the auth object. Every user delete path, SCIM and bulk delete included, drops the deleted users' committed spend twins * fix(prometheus): keep committed spend values out of logger dumps of pass-through rows CommittedBudgetSpend fields are repr=False, so Athina's payload, Sentry breadcrumbs and debug logs that str() a pass-through row's litellm_params show no spend values * fix(prometheus): keep the gauge fallback on a slow committed read and refresh committed twin TTLs on cluster The success-path committed read now gets half the per-request budget, so a slow Redis falls back to stored spend before the outer timeout cancels the gauges. Post-call charges refresh a committed twin's TTL on Redis Cluster too, so a twin never expires while its charges are still on their way to the DB. Also applies cargo fmt to the PR's Rust files * fix(proxy): key the global reset's committed-spend page queries so the Prisma span check reads them --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm-rust/crates/python-bridge/src/lib.rs | 5 +- .../src/redis/store/spend_counters.rs | 27 +- litellm-rust/crates/redis-store/AGENTS.md | 2 +- .../redis-store/src/spend_counters/keys.rs | 16 + .../redis-store/src/spend_counters/mod.rs | 72 ++- .../redis-store/src/spend_counters/types.rs | 16 + .../redis-store/tests/spend_counters.rs | 165 +++++- litellm/caching/redis_batch.py | 17 + litellm/constants.py | 1 + litellm/integrations/prometheus.py | 238 ++++++-- litellm/proxy/auth/user_api_key_auth.py | 26 +- .../proxy/common_utils/reset_budget_job.py | 5 +- litellm/proxy/db/spend_counter_reseed.py | 4 + .../proxy/hooks/proxy_track_cost_callback.py | 54 ++ litellm/proxy/litellm_pre_call_utils.py | 10 +- .../customer_endpoints.py | 7 +- .../internal_user_endpoints.py | 5 +- .../key_management_endpoints.py | 6 +- .../organization_endpoints.py | 22 +- .../management_endpoints/team_endpoints.py | 7 +- .../management_helpers/bulk_user_deletion.py | 5 + .../pass_through_endpoints.py | 3 + litellm/proxy/proxy_server.py | 57 +- .../spend_tracking/carried_budget_state.py | 40 +- .../proxy/spend_tracking/committed_spend.py | 181 ++++++ .../spend_management_endpoints.py | 50 +- litellm/proxy/utils.py | 3 +- .../redis/stores/spend_counters/__init__.py | 20 +- litellm/redis/stores/spend_counters/handle.py | 7 +- litellm/redis/stores/spend_counters/types.py | 6 + litellm/rust_bridge/_native.pyi | 9 +- litellm/types/proxy/auth/user_api_key_auth.py | 17 + litellm/types/proxy/carried_budget_state.py | 58 +- .../test_request_redis_batch_post_call.py | 23 +- .../generic_api/test_generic_api_callback.py | 26 + .../test_prometheus_carried_budget_state.py | 517 ++++++++++++++++++ .../test_prometheus_user_team_metrics.py | 73 +++ .../test_user_api_key_auth_request_flow.py | 64 +++ .../common_utils/test_reset_budget_job.py | 46 +- .../hooks/test_proxy_track_cost_callback.py | 133 ++++- .../test_customer_endpoints.py | 23 + .../test_internal_user_endpoints.py | 38 +- .../test_key_management_endpoints.py | 6 +- .../test_organization_endpoints.py | 64 +++ .../test_team_endpoints.py | 31 ++ .../test_bulk_user_deletion.py | 32 +- .../test_pass_through_endpoints.py | 42 +- .../proxy/proxy_server/test_spend_counters.py | 98 +++- .../test_carried_budget_state.py | 101 ++++ .../spend_tracking/test_committed_spend.py | 274 ++++++++++ .../test_spend_management_endpoints.py | 110 +++- ...test_proxy_server_endpoints_and_startup.py | 43 ++ .../redis/stores/spend_counters_double.py | 36 +- .../unit/redis/stores/test_spend_counters.py | 26 +- 54 files changed, 2846 insertions(+), 121 deletions(-) create mode 100644 litellm/proxy/spend_tracking/committed_spend.py create mode 100644 tests/unit/proxy/spend_tracking/test_committed_spend.py diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 09b787b2c94..0883dfa7d31 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -55,7 +55,9 @@ mod _native { #[pymodule_export] use crate::redis::store::sensitive_route_pins::NativeSensitiveRoutePins; #[pymodule_export] - use crate::redis::store::spend_counters::{NativeSpendCounters, PySpendRead}; + use crate::redis::store::spend_counters::{ + NativeSpendCounters, PyPostCallSettlement, PySpendRead, + }; #[pymodule_export] use crate::redis::store::{ NativePending, NativeRedisStoreError, router_cooldowns::NativeRouterCooldowns, @@ -154,6 +156,7 @@ mod tests { "NativeSensitiveRoutePins", "NativeRedisStoreError", "NativeSpendCounters", + "PostCallSettlement", "SpendRead", "LocalSizes", "native_redis_breaker_states", diff --git a/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs b/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs index ba016821e48..9da6a0e00e5 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs @@ -2,12 +2,12 @@ use std::{collections::HashMap, sync::Arc}; use litellm_redis_store::{ Number, - spend_counters::{LostSeed, Reading, SpendCounter, SpendCounters}, + spend_counters::{LostSeed, PostCallSettlement, Reading, SpendCounter, SpendCounters}, }; use pyo3::prelude::*; use super::{ - NativePending, PyData, PyNumber, kind_error, pending, redis_async, remote, + NativePending, PyData, PyNumber, duration, kind_error, pending, redis_async, remote, request::{NativeRequestRedis, NativeSend}, ttl, unwrap_all, }; @@ -29,6 +29,18 @@ impl PySpendRead { } } +/// Whether every increment one `increment_many_post_call` declared has its reply. +#[pyclass(frozen, name = "PostCallSettlement")] +pub(crate) struct PyPostCallSettlement(PostCallSettlement); + +#[pymethods] +impl PyPostCallSettlement { + #[getter] + fn settled(&self) -> bool { + self.0.is_settled() + } +} + fn increments(increments: Vec<(PyData, PyNumber)>) -> Vec<(SpendCounter, Number)> { increments .into_iter() @@ -73,6 +85,12 @@ impl NativeSpendCounters { self.0.redis_ttl().map(|ttl| ttl.as_secs_f64()) } + /// The least TTL a Redis write gives a committed spend twin. + fn set_committed_ttl(&self, ttl_seconds: f64) -> PyResult<()> { + self.0.set_committed_ttl(duration(ttl_seconds)?); + Ok(()) + } + /// `("redis" | "memory", value)`: where the answer came from, a Redis miss included. fn read<'py>( &self, @@ -125,14 +143,15 @@ impl NativeSpendCounters { }) } - /// Declares the increments on the request's post-call batch; `False` when there is no Redis. + /// Declares the increments on the request's post-call batch; `None` when there is no Redis. fn increment_many_post_call( &self, request: &NativeRequestRedis, increments: Vec<(PyData, PyNumber)>, - ) -> bool { + ) -> Option { self.0 .increment_many_post_call(&request.0, &self::increments(increments)) + .map(PyPostCallSettlement) } fn set_max<'py>( diff --git a/litellm-rust/crates/redis-store/AGENTS.md b/litellm-rust/crates/redis-store/AGENTS.md index cdd2bd7ca64..963f7ff168f 100644 --- a/litellm-rust/crates/redis-store/AGENTS.md +++ b/litellm-rust/crates/redis-store/AGENTS.md @@ -1,5 +1,5 @@ - One folder per Redis family (`router_cooldowns/`, `router_usage/`, `spend_counters/`, `auth_cache/`, `rate_limits/`), each with `mod.rs` (the store and its ops), `keys.rs` (key types and spellings) and `types.rs` (the family's values, with the `FromValue` / `IntoValue` codecs for what it stores in Redis). It is the only place that family's keys, value codec and scripts are written, and it takes domain ids, never keys. Shared pieces live once at the crate root: `Error` and `ErrorKind` (`error.rs`), `Memory` and `Clock` (`memory.rs`), `Pending` (`request.rs`), `Remote` and its attach-once `RemoteSlot`, which knows nothing of memory (`remote.rs`), `DualCache` (`dual_cache/`) and the Python codecs (`codec.rs`, `value.rs`) -- An op declared on a `Send` returns `Pending`, a lazy future built when the op is declared. An op that only means something on Redis (an `MGET` the caller falls back from, a script) returns `None` without one, and an op that replaced a `DualCache` call answers from memory. A post-call method returns `false` when there was no Redis to declare on +- An op declared on a `Send` returns `Pending`, a lazy future built when the op is declared. An op that only means something on Redis (an `MGET` the caller falls back from, a script) returns `None` without one, and an op that replaced a `DualCache` call answers from memory. A post-call method returns `false` when there was no Redis to declare on, or `None` when it returns a settlement the caller waits on, like `SpendCounterStore::increment_many_post_call` - A store that sends a pipeline of its own uses `RedisBatch::pipelined`, never `Pipeline` and raw `Value`s - Python and Rust pods share these keys during a rolling deploy, so a module writes the bytes `main`'s Python writes and reads everything `main` reads. Python formats go through `litellm-python-compat` (`repr`, `literal_eval`, `json`), never a hand-rolled codec - A family holds one `DualCache` per `DualCache` its Python held on `main`, and each `DualCache` method is the Python method it names, so a family op that was a `DualCache` call is one call here. Where `main`'s Python reached past `DualCache` to `in_memory_cache` or `redis_cache`, the family does the same through `memory()` and `redis()`, which keeps that reach visible. Caches a family's Python attached to one `RedisCache` share one `Arc`. Only a family that batch-reads builds its cache with `DualCache::batch_reading` from a `TierPolicy` the host passes in, and the Python host fills it from what `main` reads from the environment or `litellm` globals (`litellm/rust_bridge/host/settings.py`, `router_tiers`). Callers pass a `Remote` (or none), never a cache. Redis-only families are plain functions over `&RedisClient` diff --git a/litellm-rust/crates/redis-store/src/spend_counters/keys.rs b/litellm-rust/crates/redis-store/src/spend_counters/keys.rs index 061e19ed1de..5221f8ab66c 100644 --- a/litellm-rust/crates/redis-store/src/spend_counters/keys.rs +++ b/litellm-rust/crates/redis-store/src/spend_counters/keys.rs @@ -1,6 +1,9 @@ /// One budget's spend counter. `parse` and `key` round-trip `main`'s `spend:*` keys exactly. #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub enum SpendCounter { + /// The charges-only twin of a key, team, user, org or end user lifetime counter, which budget + /// holds never touch: `spend:committed:` and the twin's key past its `spend:`. + Committed(Box), Key { token: String, window: Option, @@ -33,6 +36,9 @@ impl SpendCounter { None => format!("spend:{scope}:{id}"), }; match self { + Self::Committed(counter) => { + format!("spend:committed:{}", &counter.key()["spend:".len()..]) + } Self::Key { token, window } => windowed("key", token, window), Self::Team { team_id, window } => windowed("team", team_id, window), Self::TeamMember { user_id, team_id } => { @@ -51,6 +57,16 @@ impl SpendCounter { pub fn parse(key: &str) -> Result { let invalid = || NotASpendCounter(key.to_owned()); let rest = key.strip_prefix("spend:").ok_or_else(invalid)?; + if let Some(twin) = rest.strip_prefix("committed:") { + return match Self::parse(&format!("spend:{twin}")).map_err(|_| invalid())? { + counter @ (Self::Key { window: None, .. } + | Self::Team { window: None, .. } + | Self::User(_) + | Self::Org(_) + | Self::EndUser(_)) => Ok(Self::Committed(Box::new(counter))), + _ => Err(invalid()), + }; + } let (scope, id) = rest.split_once(':').ok_or_else(invalid)?; let windowed = || match id.rsplit_once(":window:") { Some((id, window)) if !window.is_empty() => (id.to_owned(), Some(window.to_owned())), diff --git a/litellm-rust/crates/redis-store/src/spend_counters/mod.rs b/litellm-rust/crates/redis-store/src/spend_counters/mod.rs index 98f619cc501..425f83ef85c 100644 --- a/litellm-rust/crates/redis-store/src/spend_counters/mod.rs +++ b/litellm-rust/crates/redis-store/src/spend_counters/mod.rs @@ -1,13 +1,19 @@ mod keys; mod types; -use std::{sync::Arc, time::Duration}; +use std::{ + sync::{ + Arc, + atomic::{AtomicBool, AtomicU64, Ordering}, + }, + time::Duration, +}; use litellm_python_compat::repr::to_str; use litellm_redis_client::{Declared, Expiry, RedisBatch, commands}; pub use self::keys::{NotASpendCounter, SpendCounter}; -pub use self::types::{LostSeed, Reading}; +pub use self::types::{LostSeed, PostCallSettlement, Reading}; use crate::{ Error, Number, Pending, RequestRedis, Send, dual_cache::{DualCache, MemorySettings}, @@ -31,6 +37,7 @@ const SET_MAX: &str = concat!( /// `DualCache` to its Redis and its memory. pub struct SpendCounters { cache: DualCache, + committed_ttl_millis: AtomicU64, } impl SpendCounters { @@ -45,6 +52,33 @@ impl SpendCounters { MemorySettings::default(), now, ), + committed_ttl_millis: AtomicU64::new(0), + } + } + + /// The least TTL a Redis write gives a [`SpendCounter::Committed`] twin, so it outlives the DB + /// write that holds its charges and its next seed never reads a total still missing them. + pub fn set_committed_ttl(&self, ttl: Duration) { + let millis = u64::try_from(ttl.as_millis()).unwrap_or(u64::MAX); + self.committed_ttl_millis.store(millis, Ordering::Relaxed); + } + + fn ttl(&self, counter: &SpendCounter, default_ttl: Duration) -> Duration { + match counter { + SpendCounter::Committed(_) => default_ttl.max(Duration::from_millis( + self.committed_ttl_millis.load(Ordering::Relaxed), + )), + _ => default_ttl, + } + } + + /// A committed twin's expiry moves out on every charge, on a cluster too, so it never expires + /// while a charge it holds is still on its way to the DB. + fn post_call_expiry(&self, counter: &SpendCounter, default_ttl: Duration) -> Expiry { + let ttl = self.ttl(counter, default_ttl); + match counter { + SpendCounter::Committed(_) => Expiry::Refresh(ttl), + _ => Expiry::KeepExisting(ttl), } } @@ -127,7 +161,7 @@ impl SpendCounters { let increment = commands::incr_by_float( &key, amount.into(), - Some(Expiry::Refresh(remote.default_ttl)), + Some(Expiry::Refresh(self.ttl(counter, remote.default_ttl))), ); match remote.client.send(increment).await { Ok(total) => { @@ -159,7 +193,7 @@ impl SpendCounters { pipeline.add(commands::incr_by_float( &counter.key(), (*amount).into(), - Some(Expiry::Refresh(remote.default_ttl)), + Some(Expiry::Refresh(self.ttl(counter, remote.default_ttl))), )) }); let totals = Declared::all(totals.collect::>()).await; @@ -175,15 +209,14 @@ impl SpendCounters { } /// `increment_many` on the request's post-call batch. Each total reaches memory once the batch - /// goes out, and a counter whose increment failed is dropped everywhere. `false` without Redis. + /// goes out, and a counter whose increment failed is dropped everywhere. `None` without Redis, + /// else a settlement that reports when every reply is back, whoever flushed the batch. pub fn increment_many_post_call( self: &Arc, request: &RequestRedis, increments: &[(SpendCounter, Number)], - ) -> bool { - let Some(remote) = self.cache.redis().get() else { - return false; - }; + ) -> Option { + let remote = self.cache.redis().get()?; let batch = Send::PostCall(request).batch(&remote.client); let declared = increments .iter() @@ -191,12 +224,14 @@ impl SpendCounters { let total = batch.add(commands::incr_by_float( &counter.key(), (*amount).into(), - Some(Expiry::KeepExisting(remote.default_ttl)), + Some(self.post_call_expiry(counter, remote.default_ttl)), )); (counter.clone(), total) }) .collect::>(); let store = Arc::clone(self); + let settled = Arc::new(AtomicBool::new(false)); + let settlement = PostCallSettlement(Arc::clone(&settled)); request.spawn(async move { for (counter, total) in declared { match total.settled().await { @@ -215,8 +250,9 @@ impl SpendCounters { } } } + settled.store(true, Ordering::Release); }); - true + Some(settlement) } /// `_repair_stale_spend_counter`: raise memory and Redis to `spend`, never lowering either. @@ -243,7 +279,7 @@ impl SpendCounters { }; let args = ( to_str(&spend.to_value()), - remote.default_ttl.as_secs().to_string(), + self.ttl(counter, remote.default_ttl).as_secs().to_string(), ); let stored: Option = remote .client @@ -276,8 +312,9 @@ impl SpendCounters { }) })?); }; + let ttl = self.ttl(counter, *default_ttl); let current = if client - .send(commands::set_nx(&key, spend.to_json(), *default_ttl)) + .send(commands::set_nx(&key, spend.to_json(), ttl)) .await? { spend @@ -290,7 +327,7 @@ impl SpendCounters { let increment = commands::incr_by_float( &key, spend.into(), - Some(Expiry::KeepExisting(*default_ttl)), + Some(Expiry::KeepExisting(ttl)), ); Number::Float(client.send(increment).await?) } @@ -313,7 +350,7 @@ impl SpendCounters { let Some(remote) = self.cache.redis().get() else { return Ok(()); }; - let ttl = ttl.unwrap_or(remote.default_ttl); + let ttl = ttl.unwrap_or_else(|| self.ttl(counter, remote.default_ttl)); remote .client .send(commands::set(&key, spend.to_json(), ttl)) @@ -334,7 +371,10 @@ impl SpendCounters { }; Ok(remote .client - .send(commands::expire(&counter.key(), remote.default_ttl)) + .send(commands::expire( + &counter.key(), + self.ttl(counter, remote.default_ttl), + )) .await?) } diff --git a/litellm-rust/crates/redis-store/src/spend_counters/types.rs b/litellm-rust/crates/redis-store/src/spend_counters/types.rs index 61c211ec7e4..0509046d8c8 100644 --- a/litellm-rust/crates/redis-store/src/spend_counters/types.rs +++ b/litellm-rust/crates/redis-store/src/spend_counters/types.rs @@ -1,3 +1,8 @@ +use std::sync::{ + Arc, + atomic::{AtomicBool, Ordering}, +}; + use litellm_python_compat::Value; use super::keys::SpendCounter; @@ -32,6 +37,17 @@ impl Reading { } } +/// One `increment_many_post_call`: settled once every increment it declared has its reply and +/// memory mirrors it, whether `flush_post` or the post-call deadline sent the batch. +#[derive(Clone, Debug)] +pub struct PostCallSettlement(pub(super) Arc); + +impl PostCallSettlement { + pub fn is_settled(&self) -> bool { + self.0.load(Ordering::Acquire) + } +} + /// Who wins when a seed's `SET NX` lost and the winner's value is gone by the time it is read. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum LostSeed { diff --git a/litellm-rust/crates/redis-store/tests/spend_counters.rs b/litellm-rust/crates/redis-store/tests/spend_counters.rs index 75b7af2c321..7c89fb5b41e 100644 --- a/litellm-rust/crates/redis-store/tests/spend_counters.rs +++ b/litellm-rust/crates/redis-store/tests/spend_counters.rs @@ -16,6 +16,7 @@ use support::{COOLDOWN_POLICY, connect, connect_observed, remote, within}; use tokio::runtime::Handle; const TTL: Duration = Duration::from_secs(60); +const COMMITTED_TTL: Duration = Duration::from_secs(300); const LONG_DEADLINE: Duration = Duration::from_secs(60); const READS: [(&str, Number); 4] = [ @@ -45,6 +46,10 @@ fn team(id: &str) -> SpendCounter { } } +fn committed(counter: SpendCounter) -> SpendCounter { + SpendCounter::Committed(Box::new(counter)) +} + fn store(client: Option) -> Arc { Arc::new(SpendCounters::new(client.map(remote))) } @@ -114,6 +119,26 @@ fn a_reading_names_where_it_was_answered( "spend:model_access_group:gpt-group" )] #[case(SpendCounter::ShadowEval("job-7".to_owned()), "spend:shadow_eval:job-7")] +#[case::committed_key( + committed(SpendCounter::Key { + token: "88dc28d0".to_owned(), + window: None, + }), + "spend:committed:key:88dc28d0" +)] +#[case::committed_team(committed(team("team-a")), "spend:committed:team:team-a")] +#[case::committed_user( + committed(SpendCounter::User("user@example.com".to_owned())), + "spend:committed:user:user@example.com" +)] +#[case::committed_org( + committed(SpendCounter::Org("org-1".to_owned())), + "spend:committed:org:org-1" +)] +#[case::committed_end_user( + committed(SpendCounter::EndUser("customer é 1".to_owned())), + "spend:committed:end_user:customer é 1" +)] fn every_counter_formats_and_parses_back_to_mains_key( #[case] counter: SpendCounter, #[case] key: &str, @@ -128,6 +153,11 @@ fn every_counter_formats_and_parses_back_to_mains_key( #[case("spend:unknown:abc")] #[case("spend:team_member:no-team")] #[case("spend:")] +#[case::committed_window("spend:committed:key:abc:window:1d")] +#[case::committed_tag("spend:committed:tag:prod")] +#[case::committed_team_member("spend:committed:team_member:user:team-a")] +#[case::committed_twice("spend:committed:committed:key:abc")] +#[case::committed_nothing("spend:committed:")] fn keys_outside_the_spend_counter_keyspace_are_refused(#[case] key: &str) { assert!(SpendCounter::parse(key).is_err(), "{key}"); } @@ -307,7 +337,7 @@ async fn a_post_call_increment_from_a_thread_without_a_runtime_lands_in_redis_an .expect("a Python thread has no Tokio context, and declaring there must not panic"); request.flush_post().await; - assert!(declared); + assert!(declared.is_some()); let mirrored = litellm_redis_testing::eventually(async || store.get_local(&counter).unwrap()).await; assert_eq!(mirrored, Number::Float(3.5)); @@ -331,9 +361,12 @@ async fn post_call_increments_wait_for_the_post_batch_then_land_in_memory( let request = RequestRedis::new(Handle::current(), LONG_DEADLINE); let counter = team("t"); - assert!(store.increment_many_post_call(&request, &[(counter.clone(), Number::Float(3.5))])); + let settlement = store + .increment_many_post_call(&request, &[(counter.clone(), Number::Float(3.5))]) + .unwrap(); request.flush_pre().await; + assert!(!settlement.is_settled()); assert!(trips.seen().is_empty()); assert_eq!( raw.send(commands::get::(&counter.key())) @@ -345,6 +378,7 @@ async fn post_call_increments_wait_for_the_post_batch_then_land_in_memory( request.flush_post().await; tokio::time::sleep(Duration::from_millis(100)).await; + assert!(settlement.is_settled()); let increment = match topology { Topology::Standalone => Operation::Pipeline { ops: 1 }, Topology::Cluster => Operation::Command("INCRBYFLOAT"), @@ -557,13 +591,132 @@ async fn a_failed_post_call_increment_drops_the_counter_from_memory() { let request = RequestRedis::new(Handle::current(), LONG_DEADLINE); let counter = team("t"); store.increment(&counter, Number::Float(1.0)).await.unwrap(); - assert!(store.increment_many_post_call(&request, &[(counter.clone(), Number::Float(1.0))])); + let settlement = store + .increment_many_post_call(&request, &[(counter.clone(), Number::Float(1.0))]) + .unwrap(); deployment.stop().await.unwrap(); request.flush_post().await; tokio::time::sleep(Duration::from_millis(200)).await; assert_eq!(store.get_local(&counter).unwrap(), None); + assert!(settlement.is_settled()); +} + +#[rstest] +#[tokio::test] +async fn a_post_call_settlement_settles_when_the_deadline_sends_the_batch( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let store = store(Some(connect(&deployment))); + let request = RequestRedis::new(Handle::current(), Duration::from_millis(50)); + let counter = committed(team("t")); + + let settlement = store + .increment_many_post_call(&request, &[(counter.clone(), Number::Float(2.5))]) + .unwrap(); + + assert!(!settlement.is_settled()); + litellm_redis_testing::eventually(async || settlement.is_settled().then_some(())).await; + assert_eq!( + raw.send(commands::get::(&counter.key())) + .await + .unwrap(), + Some("2.5".to_owned()) + ); +} + +#[rstest] +#[case::committed_twin(committed, 300)] +#[case::enforcement_counter(std::convert::identity, 60)] +#[tokio::test] +async fn every_write_gives_a_committed_twin_the_committed_ttl_and_other_counters_the_redis_ttl( + #[case] counter_of: fn(SpendCounter) -> SpendCounter, + #[case] ttl_seconds: u64, + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let store = store(Some(connect(&deployment))); + store.set_committed_ttl(COMMITTED_TTL); + let request = RequestRedis::new(Handle::current(), LONG_DEADLINE); + let [ + incremented, + pipelined, + post_call, + seeded, + repaired, + refreshed, + ] = ["a", "b", "c", "d", "e", "f"].map(|id| counter_of(team(id))); + raw.send(commands::set( + &refreshed.key(), + "1.0", + Duration::from_secs(5), + )) + .await + .unwrap(); + + store + .increment(&incremented, Number::Float(1.0)) + .await + .unwrap(); + store + .increment_many(&[(pipelined.clone(), Number::Float(1.0))]) + .await + .unwrap(); + let settlement = store + .increment_many_post_call(&request, &[(post_call.clone(), Number::Float(1.0))]) + .unwrap(); + request.flush_post().await; + store + .seed(&seeded, Number::Float(1.0), LostSeed::Answer) + .await + .unwrap(); + store.set_max(&repaired, Number::Float(1.0)).await.unwrap(); + assert!(store.refresh_ttl(&refreshed).await.unwrap()); + + assert!(settlement.is_settled()); + let remaining = + async |counter: &SpendCounter| match raw.send(commands::ttl(&counter.key())).await.unwrap() + { + Ttl::Expires(remaining) => remaining, + other => panic!("{} has no TTL: {other:?}", counter.key()), + }; + assert!(within(remaining(&incremented).await, ttl_seconds)); + assert!(within(remaining(&pipelined).await, ttl_seconds)); + assert!(within(remaining(&post_call).await, ttl_seconds)); + assert!(within(remaining(&seeded).await, ttl_seconds)); + assert!(within(remaining(&repaired).await, ttl_seconds)); + assert!(within(remaining(&refreshed).await, ttl_seconds)); +} + +#[rstest] +#[tokio::test] +async fn a_post_call_charge_moves_a_committed_twins_expiry_out_to_the_committed_ttl( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let store = store(Some(connect(&deployment))); + store.set_committed_ttl(COMMITTED_TTL); + let request = RequestRedis::new(Handle::current(), LONG_DEADLINE); + let counter = committed(team("t")); + raw.send(commands::set(&counter.key(), "1.0", Duration::from_secs(5))) + .await + .unwrap(); + + let settlement = store + .increment_many_post_call(&request, &[(counter.clone(), Number::Float(1.0))]) + .unwrap(); + request.flush_post().await; + + assert!(settlement.is_settled()); + match raw.send(commands::ttl(&counter.key())).await.unwrap() { + Ttl::Expires(remaining) => assert!(within(remaining, COMMITTED_TTL.as_secs())), + other => panic!("{} has no TTL: {other:?}", counter.key()), + } } #[tokio::test] @@ -592,7 +745,11 @@ async fn without_redis_counters_live_in_memory() { store.read(&counter).await.unwrap(), Reading::Memory(Some(Number::Float(7.0))) ); - assert!(!store.increment_many_post_call(&request, &[(counter.clone(), Number::Int(1))])); + assert!( + store + .increment_many_post_call(&request, &[(counter.clone(), Number::Int(1))]) + .is_none() + ); assert!( store .read_many(RedisSend::PreCall(&request), &[counter]) diff --git a/litellm/caching/redis_batch.py b/litellm/caching/redis_batch.py index d58ddb9c629..b3e6d2eb291 100644 --- a/litellm/caching/redis_batch.py +++ b/litellm/caching/redis_batch.py @@ -452,6 +452,13 @@ class NativePending(Protocol[_T_co]): def resolve(self) -> Awaitable[_T_co]: ... +class PostCallSettlement(Protocol): + """Post-call writes declared together: settled once every reply is back, whoever sent their batch.""" + + @property + def settled(self) -> bool: ... + + class NativeBatchResult(Generic[_T]): """Awaitable handle for a read declared on the request's native pre-call batch; awaiting it runs the native flush hooks, flushes that batch, then answers from what it read.""" @@ -526,6 +533,7 @@ class RequestRedisBatches: "_native_misses", "_native_send_factory", "_post_call", + "_post_call_settlements", "post_call_deadline", "prefetched", ) @@ -538,6 +546,7 @@ class RequestRedisBatches: ) -> None: self._batches: Final[dict[object, RedisBatch]] = {} # mutable-ok: lazily filled per backend self._post_call: Final[dict[object, RedisBatch]] = {} # mutable-ok: lazily filled per backend + self._post_call_settlements: tuple[PostCallSettlement, ...] = () self.post_call_deadline: Final = post_call_deadline self._native_factory: Final = native_factory self._native_send_factory: Final = native_send_factory @@ -641,6 +650,14 @@ class RequestRedisBatches: if not any(batch.pending for batch in self._post_call.values()): _open_post_call.discard(self) + def track_post_call_settlement(self, settlement: PostCallSettlement) -> None: + """Post-call writes a later reader in this request must know have landed, such as its committed spend.""" + self._post_call_settlements = (*self._post_call_settlements, settlement) + + @property + def post_call_settlements(self) -> tuple[PostCallSettlement, ...]: + return self._post_call_settlements + @property def batches(self) -> tuple[RedisBatch, ...]: return tuple(self._batches.values()) diff --git a/litellm/constants.py b/litellm/constants.py index fe605294c50..3a49e8a7fc4 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1941,6 +1941,7 @@ SPEND_ROLLUP_LOCK_TIMEOUT_MS: Final = max(1, get_env_int("SPEND_ROLLUP_LOCK_TIME SPEND_LOG_QUEUE_SIZE_THRESHOLD: Final = get_env_int("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100) SPEND_LOG_QUEUE_MAX_BYTES: Final = max(1, get_env_int("SPEND_LOG_QUEUE_MAX_BYTES", 64000000)) SPEND_LOG_QUEUE_POLL_INTERVAL: Final = get_env_float("SPEND_LOG_QUEUE_POLL_INTERVAL", 2.0) +SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS: Final = 30.0 RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS: Final = max(1, get_env_int("RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS", 3)) RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL: Final = get_env_float("RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL", 0.2) SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = get_env_int("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index fc1a601f9da..7984c7d04f3 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -21,6 +21,7 @@ from typing_extensions import ReadOnly, TypedDict, assert_never import litellm from litellm._internal_context import with_service_target from litellm._logging import print_verbose, verbose_logger +from litellm.caching.redis_batch import active_request_redis_batches from litellm.constants import ( AUTH_OBJECTS_TARGET, LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL, @@ -75,8 +76,10 @@ from litellm.types.integrations.prometheus import ( sanitize_prometheus_label_value, validate_prometheus_deployment_and_latency_caller_identity, ) -from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth, committed_budget_spend_of_row from litellm.types.proxy.carried_budget_state import ( + CommittedBudgetSpend, + EntitySpendCounterKeys, KeyBudgetSnapshot, OrgBudgetSnapshot, TeamBudgetSnapshot, @@ -281,6 +284,11 @@ def _logged_fields(value: object) -> Mapping[str, object]: return _NO_LOGGED_FIELDS +def _response_cost_of(logged: Mapping[str, object]) -> float: + cost: Final = logged.get("response_cost") + return float(cost) if isinstance(cost, (int, float)) else 0.0 + + _GOOGLE_PASS_THROUGH_PROVIDERS: Final = frozenset( (litellm.LlmProviders.GEMINI.value, litellm.LlmProviders.VERTEX_AI.value) ) @@ -365,6 +373,104 @@ def _get_budget_metrics_per_request_timeout() -> float: return parsed +def _stored_spend(value: object) -> float: + return float(value) if isinstance(value, (int, float)) else 0.0 + + +_OWN_CHARGE_MATCH_RATIO: Final = 1 - 1e-9 + + +def _committed_floor(stored: float, committed: float | None) -> float: + return stored if committed is None else max(stored, committed) + + +def _customer_spend_after_request(cached_spend: float, committed_spend: float | None, response_cost: float) -> float: + if committed_spend is None: + return cached_spend + response_cost + return max(committed_spend + response_cost, cached_spend) + + +def _budgeted_org_counter(org_id: str | None, org: OrgBudgetSnapshot | None) -> EntitySpendCounterKeys: + if org is None or org.max_budget is None: + return EntitySpendCounterKeys() + return EntitySpendCounterKeys.of(token=None, team_id=None, user_id=None, org_id=org_id, end_user_id=None) + + +async def _committed_spend_now(counter_keys: EntitySpendCounterKeys, timeout: float) -> CommittedBudgetSpend: + if not counter_keys.tracked(): + return CommittedBudgetSpend() + try: + from litellm.proxy.proxy_server import spend_counter_store + from litellm.proxy.spend_tracking.committed_spend import read_committed_budget_spend + + return await asyncio.wait_for( + read_committed_budget_spend(spend_counter_store, counter_keys), + timeout=timeout, + ) + except Exception as e: # noqa: BLE001 # a gauge read never fails the request; the gauges keep the stored spend + verbose_logger.debug("[Non-Blocking] Prometheus: committed spend read failed: %s", e) + return CommittedBudgetSpend() + + +OwnCharge: TypeAlias = Literal["unknown", "pending", "counted"] + + +def _own_charge() -> OwnCharge: + """Whether this request's committed spend charge is in Redis yet, known only inside a request with Redis.""" + try: + from litellm.proxy.proxy_server import spend_counter_store + except ImportError: + return "unknown" + batches: Final = active_request_redis_batches() + if batches is None or not spend_counter_store.has_redis: + return "unknown" + settlements: Final = batches.post_call_settlements + return "counted" if settlements and all(settlement.settled for settlement in settlements) else "pending" + + +def _committed_spend_before_request( + at_auth: CommittedBudgetSpend, now: CommittedBudgetSpend, response_cost: float, own_charge: OwnCharge +) -> CommittedBudgetSpend: + def before(at_auth_spend: float | None, now_spend: float | None) -> float | None: + if now_spend is None: + return at_auth_spend + match own_charge: + case "pending": + return now_spend + case "counted": + without_own_charge: Final = now_spend - response_cost + return without_own_charge if at_auth_spend is None else max(at_auth_spend, without_own_charge) + case "unknown": + if at_auth_spend is None: + return now_spend - response_cost + own_charge_counted: Final = now_spend - at_auth_spend >= response_cost * _OWN_CHARGE_MATCH_RATIO + return max(at_auth_spend, now_spend - response_cost) if own_charge_counted else now_spend + case _: + assert_never(own_charge) + + return CommittedBudgetSpend( + key=before(at_auth.key, now.key), + team=before(at_auth.team, now.team), + user=before(at_auth.user, now.user), + org=before(at_auth.org, now.org), + end_user=before(at_auth.end_user, now.end_user), + ) + + +async def _committed_spend_without_own_charge( + counter_keys: EntitySpendCounterKeys, + at_auth: CommittedBudgetSpend | None, + response_cost: float, + read_timeout: float, +) -> CommittedBudgetSpend: + own_charge_before_read: Final = _own_charge() + now: Final = await _committed_spend_now(counter_keys, read_timeout) + own_charge: Final = own_charge_before_read if _own_charge() == own_charge_before_read else "unknown" + return _committed_spend_before_request( + at_auth=now if at_auth is None else at_auth, now=now, response_cost=response_cost, own_charge=own_charge + ) + + def _get_proxy_llm_router() -> Router | None: try: from litellm.proxy.proxy_server import llm_router @@ -2229,6 +2335,10 @@ class PrometheusLogger(CustomLogger): _user_max_budget: Final = _metadata.get("user_api_key_user_max_budget", None) _user_email: Final = _metadata.get("user_api_key_user_email", None) _org_alias: Final = _metadata.get("user_api_key_org_alias", None) + _carried_metadata: Final = _REQUEST_METADATA_ADAPTER.validate_python(_metadata) + _committed_at_auth: Final = committed_budget_spend_of_row( + _carried_metadata, _REQUEST_METADATA_ADAPTER.validate_python(litellm_params) + ) # Bound the per-request budget-metric emission so that slow Redis/DB # lookups under load cannot consume the whole LoggingWorker watchdog @@ -2238,45 +2348,65 @@ class PrometheusLogger(CustomLogger): # so dropping one slow per-request emission only loses sub-cron real-time # detail, not correctness. budget_metrics_timeout: Final = _get_budget_metrics_per_request_timeout() - gather_coro: Final = asyncio.gather( - self._set_api_key_budget_metrics_after_api_request( - user_api_key=user_api_key, - user_api_key_alias=user_api_key_alias, + committed_read_timeout: Final = budget_metrics_timeout / 2 + + async def emit_budget_metrics() -> tuple[object, ...]: + committed: Final = await _committed_spend_without_own_charge( + counter_keys=_committed_at_auth.counter_keys + if _committed_at_auth is not None + else EntitySpendCounterKeys.of( + token=user_api_key, + team_id=user_api_team, + user_id=user_id, + org_id=user_api_key_org_id, + end_user_id=end_user_id, + ), + at_auth=_committed_at_auth, response_cost=response_cost, - key_max_budget=_api_key_max_budget, - key_spend=_api_key_spend, - carried=KeyBudgetSnapshot.from_metadata(_metadata), - ), - self._set_team_budget_metrics_after_api_request( - user_api_team=user_api_team, - user_api_team_alias=user_api_team_alias, - team_spend=_team_spend, - team_max_budget=_team_max_budget, - response_cost=response_cost, - carried=TeamBudgetSnapshot.from_metadata(_metadata), - ), - self._set_user_budget_metrics_after_api_request( - user_id=user_id, - user_spend=_user_spend, - user_max_budget=_user_max_budget, - response_cost=response_cost, - carried=UserBudgetSnapshot.from_metadata(_metadata), - user_email=_user_email if isinstance(_user_email, str) else None, - ), - self._set_org_budget_metrics_after_api_request( - org_id=user_api_key_org_id, - response_cost=response_cost, - carried=OrgBudgetSnapshot.from_metadata(_metadata), - org_alias=_org_alias if isinstance(_org_alias, str) else None, - ), - self._set_customer_budget_metrics_after_api_request( - end_user_id=end_user_id, - response_cost=response_cost, - ), - return_exceptions=True, - ) + read_timeout=committed_read_timeout, + ) + return await asyncio.gather( + self._set_api_key_budget_metrics_after_api_request( + user_api_key=user_api_key, + user_api_key_alias=user_api_key_alias, + response_cost=response_cost, + key_max_budget=_api_key_max_budget, + key_spend=_committed_floor(_stored_spend(_api_key_spend), committed.key), + carried=KeyBudgetSnapshot.from_metadata(_carried_metadata), + ), + self._set_team_budget_metrics_after_api_request( + user_api_team=user_api_team, + user_api_team_alias=user_api_team_alias, + team_spend=_committed_floor(_stored_spend(_team_spend), committed.team), + team_max_budget=_team_max_budget, + response_cost=response_cost, + carried=TeamBudgetSnapshot.from_metadata(_carried_metadata), + ), + self._set_user_budget_metrics_after_api_request( + user_id=user_id, + user_spend=_committed_floor(_stored_spend(_user_spend), committed.user), + user_max_budget=_user_max_budget, + response_cost=response_cost, + carried=UserBudgetSnapshot.from_metadata(_carried_metadata), + user_email=_user_email if isinstance(_user_email, str) else None, + ), + self._set_org_budget_metrics_after_api_request( + org_id=user_api_key_org_id, + response_cost=response_cost, + carried=OrgBudgetSnapshot.from_metadata(_carried_metadata), + org_alias=_org_alias if isinstance(_org_alias, str) else None, + committed_spend=committed.org, + ), + self._set_customer_budget_metrics_after_api_request( + end_user_id=end_user_id, + response_cost=response_cost, + committed_spend=committed.end_user, + ), + return_exceptions=True, + ) + try: - results: Final = await asyncio.wait_for(gather_coro, timeout=budget_metrics_timeout) + results: Final = await asyncio.wait_for(emit_budget_metrics(), timeout=budget_metrics_timeout) except asyncio.TimeoutError: verbose_logger.debug( "[Non-Blocking] Prometheus: per-request budget metric emission " @@ -2745,7 +2875,7 @@ class PrometheusLogger(CustomLogger): _labels, ) - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: verbose_logger.debug( "prometheus Logging - Enters failure logging function (kwargs keys: %s)", list(kwargs.keys()) if isinstance(kwargs, dict) else type(kwargs).__name__, @@ -2799,9 +2929,27 @@ class PrometheusLogger(CustomLogger): label_context=PrometheusLabelFactoryContext(enum_values), ) self.set_llm_deployment_failure_metrics(kwargs) + failure_metadata: Final = _REQUEST_METADATA_ADAPTER.validate_python( + get_litellm_metadata_from_kwargs(kwargs) + ) + failure_org_alias: Final = failure_metadata.get("user_api_key_org_alias") + failure_org: Final = OrgBudgetSnapshot.from_metadata(failure_metadata) + failure_cost: Final = _response_cost_of(_LOGGED_FIELDS.validate_python(standard_logging_payload)) + failure_committed: Final = await _committed_spend_without_own_charge( + counter_keys=_budgeted_org_counter(user_api_key_org_id, failure_org), + at_auth=committed_budget_spend_of_row( + failure_metadata, _REQUEST_METADATA_ADAPTER.validate_python(litellm_params) + ) + or CommittedBudgetSpend(), + response_cost=failure_cost, + read_timeout=_get_budget_metrics_per_request_timeout(), + ) await self._set_org_budget_metrics_after_api_request( org_id=user_api_key_org_id, - response_cost=0, + response_cost=failure_cost, + carried=failure_org, + org_alias=failure_org_alias if isinstance(failure_org_alias, str) else None, + committed_spend=failure_committed.org, ) except Exception as e: verbose_logger.exception("prometheus Layer Error(): Exception occured - %s", e) @@ -4516,6 +4664,7 @@ class PrometheusLogger(CustomLogger): response_cost: float, carried: OrgBudgetSnapshot | None = None, org_alias: str | None = None, + committed_spend: float | None = None, ): """ Set org budget metrics after an LLM API request @@ -4533,7 +4682,7 @@ class PrometheusLogger(CustomLogger): self._set_org_budget_metrics( org_id=org_id, org_alias=org_alias or "", - spend=carried.spend + response_cost, + spend=_committed_floor(carried.spend, committed_spend) + response_cost, max_budget=carried.max_budget, budget_reset_at=None, ) @@ -4559,7 +4708,7 @@ class PrometheusLogger(CustomLogger): if org_info is None: return - _total_org_spend: Final = (org_info.spend or 0.0) + response_cost + _total_org_spend: Final = _committed_floor(org_info.spend or 0.0, committed_spend) + response_cost budget_table: Final = org_info.litellm_budget_table max_budget: Final = budget_table.max_budget if budget_table else None budget_reset_at: Final = getattr(budget_table, "budget_reset_at", None) if budget_table else None @@ -4626,6 +4775,7 @@ class PrometheusLogger(CustomLogger): self, end_user_id: str | None, response_cost: float, + committed_spend: float | None = None, ): if self._customer_budget_gauges_are_noop() or not _customer_budget_metrics_enabled(): return @@ -4651,7 +4801,9 @@ class PrometheusLogger(CustomLogger): budget_table: Final = cached_customer.litellm_budget_table self._set_customer_budget_metrics( end_user_id=end_user_id, - spend=cached_customer.spend + response_cost, + spend=_customer_spend_after_request( + cached_spend=cached_customer.spend, committed_spend=committed_spend, response_cost=response_cost + ), max_budget=budget_table.max_budget if budget_table is not None else None, budget_reset_at=None, ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 61ee4fd2439..0148cafa9e9 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -14,6 +14,7 @@ import secrets from collections.abc import Mapping, Sequence from datetime import datetime, timezone from enum import Enum +from functools import partial from types import MappingProxyType from typing import Final, NamedTuple, Protocol, Union, cast @@ -170,7 +171,15 @@ from litellm.proxy.db.db_lookup_gate import bounded_db_lookup from litellm.proxy.db.db_span import db_span from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup -from litellm.proxy.spend_tracking.carried_budget_state import carry_team_and_user_budget_state +from litellm.proxy.spend_tracking.carried_budget_state import ( + carry_committed_budget_spend, + carry_team_and_user_budget_state, +) +from litellm.proxy.spend_tracking.committed_spend import ( + bind_committed_admission_counters, + committed_spend_tracked, + read_admission_committed_spend, +) from litellm.proxy.spend_tracking.spend_counter_batch import ( bind_admission_counter_keys, release_spend_counter_batch, @@ -3650,6 +3659,7 @@ async def run_centralized_common_checks( model_max_budget_limiter, prisma_client, proxy_logging_obj, + spend_counter_store, user_api_key_cache, user_custom_auth, ) @@ -3960,6 +3970,7 @@ async def run_centralized_common_checks( ) bind_admission_counter_keys(user_api_key_auth_obj, end_user_id=end_user_id) + bind_committed_admission_counters(user_api_key_auth_obj, end_user_id=end_user_id) try: _ = await common_checks( request=request, @@ -4035,6 +4046,18 @@ async def run_centralized_common_checks( user_obj=user_object, ) + if committed_spend_tracked(): + await carry_committed_budget_spend( + valid_token=user_api_key_auth_obj, + team_object=typed_team_object, + user_object=user_object, + end_user_object=end_user_object, + read_committed=( + None + if route_skips_budget_checks(route=route) + else partial(read_admission_committed_spend, spend_counter_store) + ), + ) await _reserve_budget_after_common_checks( user_api_key_auth_obj=user_api_key_auth_obj, request=request, @@ -4348,6 +4371,7 @@ async def _prefetch_referenced_auth_objects( """Warm every object and spend counter the checks below will read, in one MGET each (one DB query when cold). Runs after the key's model access check so a denied request costs no more than it did before.""" bind_admission_counter_keys(valid_token, end_user_id=end_user_id or None) + bind_committed_admission_counters(valid_token, end_user_id=end_user_id or None) await prefetch_auth_objects( refs=AuthObjectRefs.from_token(valid_token), user_api_key_cache=user_api_key_cache, diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 131f57568e8..9c06e3afc18 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -48,6 +48,7 @@ from litellm.proxy.db.budget_window_spend_writer import roll_window_spend_row from litellm.proxy.db.db_span import db_span, db_spanned from litellm.proxy.db.db_transaction_queue.pod_lock_manager import POD_LOCK_TARGET, PodLockManager from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry +from litellm.proxy.spend_tracking.committed_spend import committed_spend_keys from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.redis.stores.spend_counters import ( SPEND_COUNTERS_TARGET, @@ -621,7 +622,7 @@ class ResetBudgetJob: from litellm.proxy.proxy_server import spend_counter_store try: - await spend_counter_store.delete((counter_key,)) + await spend_counter_store.delete((counter_key, *committed_spend_keys((counter_key,)))) except Exception as e: verbose_proxy_logger.warning( "Failed to reset spend counter %s in Redis: %s. Budget may be over-enforced until counter expires.", @@ -680,7 +681,7 @@ class ResetBudgetJob: match cache: case "spend counter": - await spend_counter_store.delete(keys) + await spend_counter_store.delete((*keys, *committed_spend_keys(keys))) case "user_api_key_cache": await user_api_key_cache.async_delete_cache_keys(keys) case _: diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 4ed3ac30be7..ae45b43f15f 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -94,6 +94,10 @@ class SpendCounterReseed: _locks: ClassVar["OrderedDict[str, asyncio.Lock]"] = OrderedDict() _registry_lock: ClassVar[asyncio.Lock | None] = None + @staticmethod + async def counter_lock(counter_key: str) -> asyncio.Lock: + return await SpendCounterReseed._get_lock(counter_key) + @staticmethod async def _get_lock(counter_key: str) -> asyncio.Lock: if SpendCounterReseed._registry_lock is None: diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 68a6177cabc..133731fc335 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -47,6 +47,7 @@ from litellm.proxy.db.db_spend_update_writer import ( from litellm.proxy.hooks.search_free_tier import apply_search_free_tier from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, metadata_variable_name_for_request from litellm.proxy.route_llm_request import ProxyModelNotFoundError +from litellm.proxy.spend_tracking.committed_spend import committed_spend_tracked, warm_committed_counter_keys from litellm.proxy.spend_tracking.spend_counter_batch import post_call_counter_keys, spend_counter_batch_scope from litellm.proxy.spend_tracking.spend_event import ( ObjectMapping, @@ -75,6 +76,8 @@ from litellm.proxy.utils import ( BilledFailedAttempt, ProxyUpdateSpend, ) +from litellm.types.proxy.auth.user_api_key_auth import committed_budget_spend_of_row +from litellm.types.proxy.carried_budget_state import EntitySpendCounterKeys from litellm.types.router import Deployment from litellm.types.utils import ( CallTypes, @@ -268,6 +271,11 @@ def _proxy_spend_writer() -> DBSpendUpdateWriter: return proxy_logging_obj.db_spend_update_writer +def _committed_spend_scope(user_api_key_dict: UserAPIKeyAuth) -> EntitySpendCounterKeys: + committed: Final = user_api_key_dict.committed_budget_spend + return committed.counter_keys if committed is not None else EntitySpendCounterKeys() + + def _proxy_spend_counters() -> "_IncrementSpendCounters": from litellm.proxy.proxy_server import increment_spend_counters @@ -639,6 +647,7 @@ class _ProxyDBLogger(CustomLogger): team_id=user_api_key_dict.team_id, org_id=user_api_key_dict.org_id, project_id=user_api_key_dict.project_id, + committed_spend_scope=_committed_spend_scope(user_api_key_dict), kwargs=dict(row_kwargs), completion_response=original_exception, start_time=start_time, @@ -757,6 +766,7 @@ class _ProxyDBLogger(CustomLogger): router=get_llm_router(), ), project_id=user_api_key_dict.project_id, + committed_spend_scope=_committed_spend_scope(user_api_key_dict), **_ledger_counter_options(row_kwargs.get(ACCOUNTING_LEDGER_KWARG)), ) except Exception: # noqa: BLE001 # one row's counter failure must not keep the other rows' counters from landing @@ -1269,10 +1279,38 @@ class _IncrementSpendCounters(Protocol): request_started_at: datetime | None = None, model_access_groups: Sequence[str] | None = None, project_id: str | None = None, + committed_spend_scope: EntitySpendCounterKeys | None = None, windows_from_ledger: bool = False, ) -> None: ... +async def _committed_spend_scope_of_row( + kwargs: Mapping[str, object], + response_cost: float, + token: str | None, + team_id: str | None, + user_id: str | None, + org_id: str | None, + end_user_id: str | None, +) -> EntitySpendCounterKeys: + metadata: Final = _OBJECT_MAPPING.validate_python(get_litellm_metadata_from_kwargs(kwargs=dict(kwargs))) + carried: Final = committed_budget_spend_of_row( + metadata, _OBJECT_MAPPING.validate_python(kwargs.get("litellm_params") or {}) + ) + if carried is not None: + return carried.counter_keys + if response_cost <= 0 or not committed_spend_tracked(): + return EntitySpendCounterKeys() + from litellm.proxy.proxy_server import spend_counter_store + + return await warm_committed_counter_keys( + spend_counter_store, + EntitySpendCounterKeys.of( + token=token, team_id=team_id, user_id=user_id, org_id=org_id, end_user_id=end_user_id + ), + ) + + async def _update_database_and_spend_counters( proxy_logging_obj: "ProxyLogging", increment_spend_counters: _IncrementSpendCounters, @@ -1291,6 +1329,7 @@ async def _update_database_and_spend_counters( model_access_groups: Sequence[str] | None = None, project_id: str | None = None, spend_writer: DBSpendUpdateWriter | None = None, + committed_spend_scope: EntitySpendCounterKeys | None = None, ) -> bool: """The reservation is reconciled before the spend is persisted, from its own read. One spend counter batch then spans the database write and the counter update, so the post-call counters are read with a single MGET after the @@ -1333,6 +1372,7 @@ async def _update_database_and_spend_counters( request_tags=request_tags, model_access_groups=model_access_groups, project_id=project_id, + committed_spend_scope=committed_spend_scope, ) @@ -1354,6 +1394,7 @@ async def _update_database_and_spend_counters_in_batch( model_access_groups: Sequence[str] | None, project_id: str | None, spend_writer: DBSpendUpdateWriter | None = None, + committed_spend_scope: EntitySpendCounterKeys | None = None, ) -> bool: try: writer: Final = spend_writer if spend_writer is not None else proxy_logging_obj.db_spend_update_writer @@ -1407,6 +1448,19 @@ async def _update_database_and_spend_counters_in_batch( request_started_at=start_time, model_access_groups=model_access_groups, project_id=project_id, + committed_spend_scope=( + committed_spend_scope + if committed_spend_scope is not None + else await _committed_spend_scope_of_row( + kwargs=kwargs, + response_cost=response_cost, + token=user_api_key, + team_id=team_id, + user_id=user_id, + org_id=org_id, + end_user_id=end_user_id, + ) + ), **_ledger_counter_options(_BATCH_CALLBACK_FIELDS.validate_python(kwargs).get(ACCOUNTING_LEDGER_KWARG)), ) except Exception: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 15747da63bc..e2367ab8179 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -2156,7 +2156,14 @@ class LiteLLMProxyRequestSetup: user_api_key_dict.matched_model_access_groups ) # UserAPIKeyAuth object for MCP server access control - data[_metadata_variable_name]["user_api_key_auth"] = user_api_key_dict.model_copy( + data[_metadata_variable_name]["user_api_key_auth"] = LiteLLMProxyRequestSetup.auth_for_request_metadata( + user_api_key_dict + ) + return data + + @staticmethod + def auth_for_request_metadata(user_api_key_dict: UserAPIKeyAuth) -> UserAPIKeyAuth: + return user_api_key_dict.model_copy( update={ "metadata": strip_callback_config(user_api_key_dict.metadata), "team_metadata": strip_callback_config(user_api_key_dict.team_metadata), @@ -2164,7 +2171,6 @@ class LiteLLMProxyRequestSetup: "organization_metadata": strip_callback_config(user_api_key_dict.organization_metadata), } ) - return data @staticmethod def add_management_endpoint_metadata_to_request_metadata( diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index 0640bf0fa4b..672bbdb164a 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -42,7 +42,9 @@ from litellm.proxy.management_helpers.object_permission_utils import ( # noqa: handle_update_object_permission_common, set_object_permission, ) +from litellm.proxy.spend_tracking.committed_spend import drop_committed_spend from litellm.proxy.utils import handle_exception_on_proxy +from litellm.redis.stores.spend_counters import end_user_counter from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.chunked_in import find_many_in from litellm.repositories.table_repositories import EndUserRepository @@ -812,7 +814,7 @@ async def delete_end_user( See below for all params ``` """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import prisma_client, spend_counter_store try: if prisma_client is None: @@ -842,6 +844,9 @@ async def delete_end_user( verbose_proxy_logger.debug("received response from updating prisma client. response=%s", response) await _evict_end_user_cache_keys(_end_user_cache_keys(data.user_ids)) + await drop_committed_spend( + spend_counter_store, tuple(end_user_counter(user_id) for user_id in data.user_ids) + ) return DeleteCustomersResponse( deleted_customers=response, diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index cea62d6c2e4..e35a65d8ba8 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -1483,8 +1483,11 @@ async def _invalidate_user_spend_counter_if_changed( """ if non_default_values.get("spend") is not None: from litellm.proxy.proxy_server import invalidate_spend_counter + from litellm.proxy.spend_tracking.committed_spend import committed_spend_key - await invalidate_spend_counter(counter_key=user_counter(str(non_default_values["user_id"]))) + counter_key: Final = user_counter(str(non_default_values["user_id"])) + await invalidate_spend_counter(counter_key=counter_key) + await invalidate_spend_counter(counter_key=committed_spend_key(counter_key)) def _clears_object_permission(user_request: UpdateUserRequest) -> bool: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 69ae2c8f64a..50517eeefcf 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -158,6 +158,7 @@ from litellm.proxy.mcp_server.outbound_credentials.sso_assertion_store import ( ) from litellm.proxy.search_endpoints.search_tool_registry import rotate_search_tools_master_key from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start +from litellm.proxy.spend_tracking.committed_spend import committed_spend_key from litellm.proxy.spend_tracking.spend_tracking_utils import ( # noqa: F401 # legacy module exports _is_master_key, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export is_master_key, @@ -4288,7 +4289,7 @@ async def update_key_fn( ) if data.spend is not None: - from litellm.proxy.proxy_server import spend_counter_store + from litellm.proxy.proxy_server import invalidate_spend_counter, spend_counter_store counter_key: Final = key_counter(hash_token_if_needed(key)) try: @@ -4300,6 +4301,7 @@ async def update_key_fn( counter_key, redis_err, ) + await invalidate_spend_counter(counter_key=committed_spend_key(counter_key)) asyncio.create_task( KeyManagementEventHooks.async_key_updated_hook( @@ -6967,6 +6969,7 @@ async def reset_key_spend_fn( try: from litellm.proxy.proxy_server import ( hash_token, + invalidate_spend_counter, prisma_client, proxy_logging_obj, set_spend_counter_with_floor_and_broadcast, @@ -7022,6 +7025,7 @@ async def reset_key_spend_fn( # pre-reset value. _counter_key: Final = key_counter(hashed_api_key) await set_spend_counter_with_floor_and_broadcast(counter_key=_counter_key, value=reset_to) + await invalidate_spend_counter(counter_key=committed_spend_key(_counter_key)) await _reset_key_budget_windows( prisma_client=prisma_client, hashed_api_key=hashed_api_key, diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 15d05633679..bc1963af0bb 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -66,7 +66,9 @@ from litellm.proxy.management_helpers.utils import ( get_new_internal_user_defaults, management_endpoint_wrapper, ) +from litellm.proxy.spend_tracking.committed_spend import drop_committed_spend from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.redis.stores.spend_counters import org_counter, team_counter from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.chunked_in import find_many_in from litellm.repositories.object_permission_repository import ObjectPermissionRepository @@ -93,6 +95,7 @@ if TYPE_CHECKING: LiteLLM_OrganizationMembership as PrismaOrganizationMembership, ) from prisma.models import LiteLLM_OrganizationTable as PrismaOrganizationTable + from prisma.models import LiteLLM_TeamTable as PrismaTeamTable from prisma.models import LiteLLM_UserTable as PrismaUserTable from prisma.models import LiteLLM_VerificationToken as PrismaVerificationToken @@ -180,6 +183,8 @@ class _OrganizationMembershipTableClient(Protocol): class _TeamTableClient(Protocol): + async def find_many(self, *, where: Mapping[str, object]) -> "Sequence[PrismaTeamTable]": ... + async def delete_many(self, where: Mapping[str, object]) -> int: ... @@ -819,6 +824,11 @@ async def update_organization( include={"members": True, "teams": True, "litellm_budget_table": True}, ) + if data.spend is not None: + from litellm.proxy.proxy_server import spend_counter_store + + await drop_committed_spend(spend_counter_store, (org_counter(data.organization_id),)) + return response @@ -1045,7 +1055,8 @@ async def delete_organization( teams_in_deleted_orgs: Final = await find_many_in( _table(TeamRepository(prisma_client)), "organization_id", data.organization_ids ) - raise_if_teams_hold_ptu_shares(tuple(team.team_id for team in teams_in_deleted_orgs), llm_router) + team_ids_in_deleted_orgs: Final = tuple(team.team_id for team in teams_in_deleted_orgs) + raise_if_teams_hold_ptu_shares(team_ids_in_deleted_orgs, llm_router) deleted_orgs: Final = [] for organization_id in data.organization_ids: @@ -1073,6 +1084,15 @@ async def delete_organization( ) deleted_orgs.append(deleted_org) + from litellm.proxy.proxy_server import spend_counter_store + + await drop_committed_spend( + spend_counter_store, + ( + *(org_counter(organization_id) for organization_id in data.organization_ids), + *(team_counter(team_id) for team_id in team_ids_in_deleted_orgs), + ), + ) return deleted_orgs diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 71258049adf..210aa1aaa6f 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -226,12 +226,14 @@ from litellm.proxy.management_helpers.utils import ( add_new_member, management_endpoint_wrapper, ) +from litellm.proxy.spend_tracking.committed_spend import drop_committed_spend from litellm.proxy.utils import ( PrismaClient, ProxyLogging, handle_exception_on_proxy, query_raw_within_statement_timeout, ) +from litellm.redis.stores.spend_counters import team_counter from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.chunked_in import delete_many_in, find_many_in from litellm.repositories.daily_activity_buckets import api_key_from_bucket @@ -4863,6 +4865,7 @@ async def delete_validated_team_rows( """ from litellm.proxy.proxy_server import ( proxy_logging_obj, + spend_counter_store, user_api_key_cache, ) @@ -4934,7 +4937,9 @@ async def delete_validated_team_rows( _emit_team_members_metric(deleted_team.model_copy(update={"members_with_roles": ()})) await sync_team_access_group_membership(prisma_client=prisma_client, team_id=deleted_team.team_id) - return DeletedTeams(team_ids=tuple(team.team_id for team in locked)) + deleted_team_ids: Final = tuple(team.team_id for team in locked) + await drop_committed_spend(spend_counter_store, tuple(team_counter(team_id) for team_id in deleted_team_ids)) + return DeletedTeams(team_ids=deleted_team_ids) def _locked_roster_user_ids(locked: Sequence[LiteLLM_TeamTable]) -> tuple[str, ...]: diff --git a/litellm/proxy/management_helpers/bulk_user_deletion.py b/litellm/proxy/management_helpers/bulk_user_deletion.py index 04dd0b3f57f..cc60f58efb0 100644 --- a/litellm/proxy/management_helpers/bulk_user_deletion.py +++ b/litellm/proxy/management_helpers/bulk_user_deletion.py @@ -46,7 +46,9 @@ from litellm.proxy.management_helpers.access_group_team_sync import ( TEAM_ADVISORY_LOCK_SQL, USER_ROWS_LOCK_SQL, ) +from litellm.proxy.spend_tracking.committed_spend import drop_committed_spend from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.redis.stores.spend_counters import user_counter from litellm.repositories.table_repositories import ( OrganizationMembershipRepository, TeamMembershipRepository, @@ -585,6 +587,9 @@ async def _after_users_deleted( ) await evict_and_broadcast(cache_keys=deletion.jwt_mapping_cache_keys, user_api_key_cache=user_api_key_cache) await evict_and_broadcast(cache_keys=sorted(user_ids), user_api_key_cache=user_api_key_cache) + from litellm.proxy.proxy_server import spend_counter_store + + await drop_committed_spend(spend_counter_store, tuple(user_counter(user_id) for user_id in sorted(user_ids))) for removal in deletion.removals.values(): _emit_team_members_metric(removal.team) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index cdcb42ab761..3e25f7ce9af 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -165,6 +165,8 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( WebSocketSessionUsage, WebsocketUpstreamFrameKind, ) +from litellm.types.proxy.auth.user_api_key_auth import COMMITTED_BUDGET_SPEND_PARAM +from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD, CallTypes, Usage from .llm_provider_handlers.tinyfish_passthrough_logging_handler import ( @@ -779,6 +781,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): "litellm_params": { **litellm_params_in_body, "metadata": _metadata, + COMMITTED_BUDGET_SPEND_PARAM: user_api_key_dict.committed_budget_spend or CommittedBudgetSpend(), "proxy_server_request": { "url": str(request.url), "method": request.method, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d21b6169133..61f81907ff4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -93,6 +93,7 @@ from litellm.constants import ( LITELLM_UI_SESSION_DURATION, OPENAI_LIVE_SESSION_START_TIMEOUT_SECONDS, RUNTIME_UPDATABLE_ROUTER_SETTINGS, + SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS, VLLM_CHOICE_FIELDS, ) from litellm.litellm_core_utils.access_windows import access_windows_config_error @@ -338,7 +339,7 @@ from litellm._logging import ( verbose_router_logger, ) from litellm.caching.caching import DualCache, RedisCache -from litellm.caching.redis_batch import drain_post_call_redis_batches +from litellm.caching.redis_batch import active_request_redis_batches, drain_post_call_redis_batches from litellm.caching.redis_cache import RedisCircuitBreakerOpenError, is_redis_timeout_failure from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.constants import ( @@ -995,12 +996,20 @@ from litellm.proxy.spend_tracking.budget_reservation import ( release_unbound_budget_reservation, stamp_budget_reservation_actual_cost, ) +from litellm.proxy.spend_tracking.committed_spend import ( + committed_spend_key, + committed_spend_keys, + has_committed_twin, + is_committed_spend_key, + seed_committed_spend, +) from litellm.proxy.spend_tracking.spend_capture_rate import ( run_scheduled_spend_capture_rate_check, ) from litellm.proxy.spend_tracking.spend_counter_batch import ( PendingSpendIncrement, active_spend_counter_batch, + bind_spend_counter_keys, forget_spend_counter, post_call_counter_keys, read_batched_spend_counter, @@ -1105,6 +1114,7 @@ from litellm.types.llms.openai import ( ResponseInProgressEvent, ResponsesAPIResponse, ) +from litellm.types.proxy.carried_budget_state import EntitySpendCounterKeys from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry from litellm.types.proxy.management_endpoints.model_management_endpoints import ( ModelGroupInfoProxy, @@ -3240,6 +3250,7 @@ proxy_budget_rescheduler_min_time = PROXY_BUDGET_RESCHEDULER_MIN_TIME proxy_budget_rescheduler_max_time = PROXY_BUDGET_RESCHEDULER_MAX_TIME proxy_batch_polling_interval = PROXY_BATCH_POLLING_INTERVAL proxy_batch_write_at = PROXY_BATCH_WRITE_AT +_BATCH_WRITE_JITTER_MAX_SECONDS: Final = 5 proxy_config_reload_interval_seconds = PROXY_CONFIG_RELOAD_INTERVAL_SECONDS litellm_master_key_hash = None disable_spend_logs = False @@ -3638,6 +3649,7 @@ async def increment_spend_counters( request_started_at: datetime | None = None, model_access_groups: Sequence[str] | None = None, project_id: str | None = None, + committed_spend_scope: EntitySpendCounterKeys | None = None, windows_from_ledger: bool = False, ): """ @@ -3675,6 +3687,7 @@ async def increment_spend_counters( request_started_at=request_started_at, model_access_groups=model_access_groups, project_id=project_id, + committed_spend_scope=committed_spend_scope, windows_from_ledger=windows_from_ledger, ) @@ -3790,10 +3803,17 @@ async def _increment_spend_counters_batched( request_started_at: datetime | None, model_access_groups: Sequence[str] | None, project_id: str | None = None, + committed_spend_scope: EntitySpendCounterKeys | None = None, windows_from_ledger: bool = False, ): """Runs inside one spend counter batch: the reservation reconcile and the warm checks share a single MGET, and the reconcile adjustments go out in the same INCRBYFLOAT pipeline as the counter increments.""" + committed_counter_keys: Final = ( + tuple(counter_key for counter_key in committed_spend_scope.tracked() if has_committed_twin(counter_key)) + if committed_spend_scope is not None and response_cost is not None and response_cost > 0 + else () + ) + bind_spend_counter_keys(frozenset(committed_spend_keys(committed_counter_keys))) reservation_update: Final = await _reconcile_budget_reservation_for_counter_update( budget_reservation=budget_reservation, response_cost=response_cost, @@ -3980,6 +4000,9 @@ async def _increment_spend_counters_batched( ) if project_id is not None else None, + _prepare_committed_spend_increments(counter_keys=committed_counter_keys, increment=cost) + if committed_counter_keys + else None, ) if coro is not None ) @@ -4082,6 +4105,23 @@ async def _reconcile_budget_reservation_for_counter_update( return _ReservationCounterUpdate(reserved_counter_keys=frozenset(reserved_counter_keys), pending=pending) +async def _committed_spend_seeded(counter_key: str) -> bool: + return await _is_spend_counter_cache_warm(committed_spend_key(counter_key)) or await seed_committed_spend( + spend_counter_store=spend_counter_store, prisma_client=prisma_client, counter_key=counter_key + ) + + +async def _prepare_committed_spend_increments( + counter_keys: tuple[str, ...], increment: float +) -> tuple[PendingSpendIncrement, ...]: + seeded: Final = await asyncio.gather(*(_committed_spend_seeded(counter_key) for counter_key in counter_keys)) + return tuple( + PendingSpendIncrement(counter_key=committed_spend_key(counter_key), increment=increment) + for counter_key, is_seeded in zip(counter_keys, seeded) + if is_seeded + ) + + async def _prepare_end_user_and_tag_spend_increments( end_user_id: str | None, tags: list[str] | None, @@ -4487,7 +4527,15 @@ def _defer_spend_counter_increments(pending: Sequence[PendingSpendIncrement]) -> """Post-call increments ride the request's native post-call pipeline with the other counters. Each counter's new value lands in memory when the pipeline settles; a failed one is dropped everywhere so no reader trusts a counter whose increment may not have applied, as ``increment_spend_counters_pipeline`` does.""" - return spend_counter_store.increment_many_post_call(tuple((item.counter_key, item.increment) for item in pending)) + settlement: Final = spend_counter_store.increment_many_post_call( + tuple((item.counter_key, item.increment) for item in pending) + ) + if settlement is None: + return False + batches: Final = active_request_redis_batches() + if batches is not None and any(is_committed_spend_key(item.counter_key) for item in pending): + batches.track_post_call_settlement(settlement) + return True async def increment_spend_counters_pipeline(pending: Sequence[PendingSpendIncrement]) -> tuple[float | None, ...]: @@ -11651,7 +11699,10 @@ class ProxyStartupEvent: ) # Ensure minimum interval of 30 seconds for batch writing to prevent memory issues - batch_writing_interval: Final = proxy_batch_write_at + random.randint(0, 5) + batch_writing_interval: Final = proxy_batch_write_at + random.randint(0, _BATCH_WRITE_JITTER_MAX_SECONDS) + spend_counter_store.set_committed_ttl( + 2 * (SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS + proxy_batch_write_at + _BATCH_WRITE_JITTER_MAX_SECONDS) + ) ### PROXY WORKER HEARTBEAT ### worker_heartbeat: Final = ProxyWorkerHeartbeat(prisma_client=prisma_client) diff --git a/litellm/proxy/spend_tracking/carried_budget_state.py b/litellm/proxy/spend_tracking/carried_budget_state.py index 0dfb38272b0..150ddd1efe9 100644 --- a/litellm/proxy/spend_tracking/carried_budget_state.py +++ b/litellm/proxy/spend_tracking/carried_budget_state.py @@ -1,14 +1,17 @@ """Pins the budget state auth resolved onto ``UserAPIKeyAuth`` and emits it as request metadata.""" -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Mapping from types import MappingProxyType from typing import Final +from litellm.models.end_user import LiteLLM_EndUserTable from litellm.models.organization import LiteLLM_OrganizationTable from litellm.models.team import LiteLLM_TeamTable from litellm.models.user import LiteLLM_UserTable from litellm.proxy._types import UserAPIKeyAuth from litellm.types.proxy.carried_budget_state import ( + CommittedBudgetSpend, + EntitySpendCounterKeys, OrgBudgetSnapshot, TeamBudgetSnapshot, UserBudgetSnapshot, @@ -34,6 +37,41 @@ def carry_team_and_user_budget_state( ) +def _budgeted_org_id(valid_token: UserAPIKeyAuth, team_object: LiteLLM_TeamTable | None) -> str | None: + org_snapshot: Final = valid_token.org_budget_snapshot + if org_snapshot is None or org_snapshot.max_budget is None: + return None + return valid_token.org_id or (team_object.organization_id if team_object is not None else None) + + +def _budgeted_end_user_id(end_user_object: LiteLLM_EndUserTable | None) -> str | None: + budget_table: Final = end_user_object.litellm_budget_table if end_user_object is not None else None + if end_user_object is None or budget_table is None or budget_table.max_budget is None: + return None + return end_user_object.user_id + + +async def carry_committed_budget_spend( + valid_token: UserAPIKeyAuth, + team_object: LiteLLM_TeamTable | None, + user_object: LiteLLM_UserTable | None, + end_user_object: LiteLLM_EndUserTable | None, + read_committed: Callable[[EntitySpendCounterKeys], Awaitable[CommittedBudgetSpend]] | None, +) -> None: + counter_keys: Final = EntitySpendCounterKeys.of( + token=valid_token.token if valid_token.max_budget is not None else None, + team_id=team_object.team_id if team_object is not None and team_object.max_budget is not None else None, + user_id=user_object.user_id if user_object is not None and user_object.max_budget is not None else None, + org_id=_budgeted_org_id(valid_token=valid_token, team_object=team_object), + end_user_id=_budgeted_end_user_id(end_user_object=end_user_object), + ) + valid_token.committed_budget_spend = ( # rebind-ok: request-scoped credential + CommittedBudgetSpend(counter_keys=counter_keys) + if read_committed is None + else await read_committed(counter_keys) + ) + + def carry_organization_budget_state(valid_token: UserAPIKeyAuth, org_table: LiteLLM_OrganizationTable) -> None: budget_table: Final = org_table.litellm_budget_table valid_token.organization_alias = org_table.organization_alias diff --git a/litellm/proxy/spend_tracking/committed_spend.py b/litellm/proxy/spend_tracking/committed_spend.py new file mode 100644 index 00000000000..76066339cf1 --- /dev/null +++ b/litellm/proxy/spend_tracking/committed_spend.py @@ -0,0 +1,181 @@ +"""Committed spend counters: twins of the entity spend counters that only ever take charged cost, never budget holds. + +The remaining budget gauges read them, so a request's worst-case reservation never shows as spent. They are kept +only for entities with a budget, while a ``PrometheusLogger`` is registered, in the same ``SpendCounterStore`` as the +counters they twin. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Iterable, Mapping, Sequence +from typing import TYPE_CHECKING, Final + +import litellm +from litellm._internal_context import with_service_target +from litellm._logging import verbose_proxy_logger +from litellm.caching.redis_batch import Batch +from litellm.integrations.prometheus import PrometheusLogger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed +from litellm.proxy.spend_tracking.spend_counter_batch import ( + admission_counter_keys, + bind_spend_counter_keys, + read_batched_spend_counter, +) +from litellm.redis.stores.spend_counters import ( + COMMITTED_COUNTER_PREFIX, + END_USER_COUNTER_PREFIX, + SPEND_COUNTERS_TARGET, + SpendCounterStore, + committed_counter, +) +from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + +_BASE_COUNTER_PREFIXES: Final = ("spend:key:", "spend:team:", "spend:user:", "spend:org:", END_USER_COUNTER_PREFIX) + + +def committed_spend_key(counter_key: str) -> str: + return committed_counter(counter_key) + + +def has_committed_twin(counter_key: str) -> bool: + return counter_key.startswith(_BASE_COUNTER_PREFIXES) and ":window:" not in counter_key + + +def committed_spend_keys(counter_keys: Iterable[str]) -> tuple[str, ...]: + return tuple(committed_spend_key(counter_key) for counter_key in counter_keys if has_committed_twin(counter_key)) + + +def is_committed_spend_key(counter_key: str) -> bool: + return counter_key.startswith(COMMITTED_COUNTER_PREFIX) + + +def committed_spend_tracked() -> bool: + return litellm.logging_callback_manager.callback_is_active(callback_type=PrometheusLogger) + + +def bind_committed_admission_counters(token: UserAPIKeyAuth, end_user_id: str | None) -> None: + if committed_spend_tracked(): + bind_spend_counter_keys(frozenset(committed_spend_keys(admission_counter_keys(token, end_user_id)))) + + +async def _read_committed_values( + spend_counter_store: SpendCounterStore, twin_keys: tuple[str, ...] +) -> Mapping[str, float]: + batched: Final = await spend_counter_store.read_many(twin_keys, batch=Batch.NONE) + if batched is not None: + return {twin_key: value for twin_key, value in batched.items() if value is not None} + reads: Final = await asyncio.gather(*(spend_counter_store.read(twin_key) for twin_key in twin_keys)) + return {twin_key: value for twin_key, (value, _) in zip(twin_keys, reads, strict=True) if value is not None} + + +def _committed_budget_spend( + counter_keys: EntitySpendCounterKeys, committed: Mapping[str, float] +) -> CommittedBudgetSpend: + def spend_of(counter_key: str | None) -> float | None: + return committed.get(committed_spend_key(counter_key)) if counter_key is not None else None + + return CommittedBudgetSpend( + counter_keys=counter_keys, + key=spend_of(counter_keys.key), + team=spend_of(counter_keys.team), + user=spend_of(counter_keys.user), + org=spend_of(counter_keys.org), + end_user=spend_of(counter_keys.end_user), + ) + + +@with_service_target(SPEND_COUNTERS_TARGET) +async def read_committed_budget_spend( + spend_counter_store: SpendCounterStore, counter_keys: EntitySpendCounterKeys +) -> CommittedBudgetSpend: + twin_keys: Final = committed_spend_keys(counter_keys.tracked()) + if not twin_keys: + return CommittedBudgetSpend(counter_keys=counter_keys) + try: + committed: Final = await _read_committed_values(spend_counter_store, twin_keys) + except Exception as e: # noqa: BLE001 # a gauge read never fails the request; the gauges keep the stored spend + verbose_proxy_logger.debug("committed spend read failed for %s: %s", twin_keys, e) + return CommittedBudgetSpend(counter_keys=counter_keys) + return _committed_budget_spend(counter_keys=counter_keys, committed=committed) + + +async def warm_committed_counter_keys( + spend_counter_store: SpendCounterStore, counter_keys: EntitySpendCounterKeys +) -> EntitySpendCounterKeys: + """For a charge with no auth to name the budgeted entities: only a budgeted entity's twin exists.""" + committed: Final = await read_committed_budget_spend(spend_counter_store, counter_keys) + return EntitySpendCounterKeys( + key=counter_keys.key if committed.key is not None else None, + team=counter_keys.team if committed.team is not None else None, + user=counter_keys.user if committed.user is not None else None, + org=counter_keys.org if committed.org is not None else None, + end_user=counter_keys.end_user if committed.end_user is not None else None, + ) + + +async def read_admission_committed_spend( + spend_counter_store: SpendCounterStore, counter_keys: EntitySpendCounterKeys +) -> CommittedBudgetSpend: + """Rides the admission MGET: a counter the batch cannot answer stays unknown here and the emit-time read has it.""" + if not spend_counter_store.has_redis: + return await read_committed_budget_spend(spend_counter_store, counter_keys) + twin_keys: Final = committed_spend_keys(counter_keys.tracked()) + batched: Final = await asyncio.gather(*(read_batched_spend_counter(twin_key) for twin_key in twin_keys)) + return _committed_budget_spend( + counter_keys=counter_keys, + committed={ + twin_key: read[0] for twin_key, read in zip(twin_keys, batched) if read is not None and read[0] is not None + }, + ) + + +async def drop_committed_spend(spend_counter_store: SpendCounterStore, counter_keys: Sequence[str]) -> None: + """After a direct write to the stored spend the twin holds the old total, so the next charge seeds it again.""" + twin_keys: Final = committed_spend_keys(counter_keys) + if not twin_keys: + return + try: + await spend_counter_store.delete(twin_keys) + except Exception as e: # noqa: BLE001 # the write already committed; a kept twin lasts until charges stop for its TTL + verbose_proxy_logger.warning("committed spend drop failed: %s", e) + + +@with_service_target(SPEND_COUNTERS_TARGET) +async def seed_committed_spend( + spend_counter_store: SpendCounterStore, prisma_client: PrismaClient | None, counter_key: str +) -> bool: + """One DB read per worker for a cold twin: charges waiting on the lock find the twin the first one wrote.""" + twin_key: Final = committed_spend_key(counter_key) + async with await SpendCounterReseed.counter_lock(twin_key): + try: + stored, _ = await spend_counter_store.read(twin_key) + except Exception as e: # noqa: BLE001 # an unseeded twin only skips this charge; the next one seeds it again + verbose_proxy_logger.debug("committed spend seed read failed for %s: %s", twin_key, e) + return False + return stored is not None or await _seed_committed_spend_from_db( + spend_counter_store, prisma_client, counter_key + ) + + +async def _seed_committed_spend_from_db( + spend_counter_store: SpendCounterStore, prisma_client: PrismaClient | None, counter_key: str +) -> bool: + db_spend: Final = await ( + SpendCounterReseed.end_user_from_db(prisma_client=prisma_client, counter_key=counter_key) + if counter_key.startswith(END_USER_COUNTER_PREFIX) + else SpendCounterReseed.from_db(prisma_client=prisma_client, counter_key=counter_key) + ) + if db_spend is None: + return False + twin_key: Final = committed_spend_key(counter_key) + try: + await spend_counter_store.seed(twin_key, db_spend, increment_if_lost=False) + except Exception as e: # noqa: BLE001 # an unseeded twin only skips this charge; the next one seeds it again + verbose_proxy_logger.debug("committed spend seed failed for %s: %s", twin_key, e) + return False + return True diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 0348d8e1f82..f3dc672a072 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -2,7 +2,7 @@ import collections import json import os -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass from datetime import date, datetime, timedelta, timezone from functools import partial @@ -53,6 +53,7 @@ from litellm.proxy.auth.authorization_dependencies import ( from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.openai_error_payload import error_status_code, openai_error_param, openai_error_type from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup +from litellm.proxy.spend_tracking.committed_spend import committed_spend_tracked, drop_committed_spend from litellm.proxy.spend_tracking.session_page import ( AGENT_CALL_TYPE_SQL, MCP_CALL_TYPES_SQL, @@ -74,6 +75,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( get_spend_by_team_and_customer, ) from litellm.proxy.utils import handle_exception_on_proxy, query_raw_within_statement_timeout +from litellm.redis.stores.spend_counters import key_counter, team_counter from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import SpendLogsRepository from litellm.repositories.team_repository import TeamRepository @@ -3595,6 +3597,50 @@ async def view_spend_logs( ) +class _EntityIdRow(TypedDict): + id: ReadOnly[str] + + +_COMMITTED_DROP_PAGE_SIZE: Final = 1000 +_CommittedDropEntity: TypeAlias = Literal["key", "team"] +_COMMITTED_DROP_PAGE_QUERIES: Final[Mapping[_CommittedDropEntity, str]] = MappingProxyType( + { + "key": 'SELECT token AS id FROM "LiteLLM_VerificationToken" WHERE token > $1 ORDER BY token LIMIT $2', + "team": 'SELECT team_id AS id FROM "LiteLLM_TeamTable" WHERE team_id > $1 ORDER BY team_id LIMIT $2', + } +) +_COMMITTED_DROP_COUNTERS: Final[Mapping[_CommittedDropEntity, Callable[[str], str]]] = MappingProxyType( + {"key": key_counter, "team": team_counter} +) + + +async def _drop_committed_spend_page( + prisma_client: PrismaClient, entity: _CommittedDropEntity, after_id: str +) -> str | None: + from litellm.proxy.proxy_server import spend_counter_store + + rows: Final[Sequence[_EntityIdRow]] = await _query_raw( + prisma_client, _COMMITTED_DROP_PAGE_QUERIES[entity], after_id, _COMMITTED_DROP_PAGE_SIZE + ) + counter_of: Final = _COMMITTED_DROP_COUNTERS[entity] + await drop_committed_spend(spend_counter_store, tuple(counter_of(row["id"]) for row in rows)) + return rows[-1]["id"] if len(rows) == _COMMITTED_DROP_PAGE_SIZE else None + + +async def _drop_committed_spend_pages(prisma_client: PrismaClient, entity: _CommittedDropEntity) -> None: + after_id: str | None = "" # rebind-ok: the keyset cursor moves one page at a time + while after_id is not None: + after_id = await _drop_committed_spend_page(prisma_client, entity, after_id) + + +async def _drop_all_key_and_team_committed_spend(prisma_client: PrismaClient) -> None: + try: + await _drop_committed_spend_pages(prisma_client, "key") + await _drop_committed_spend_pages(prisma_client, "team") + except Exception as e: # noqa: BLE001 # the reset already committed; a kept twin lasts until charges stop for its TTL + verbose_proxy_logger.warning("committed spend drop after a global reset failed: %s", e) + + @router.post( "/global/spend/reset", tags=["Budget & Spend Tracking"], @@ -3626,6 +3672,8 @@ async def global_spend_reset(): _ = await transaction.litellm_teamtable.update_many(data={"spend": 0.0}, where={}) _ = await transaction.query_raw(_LOCK_ALL_KEYS_SQL) _ = await transaction.litellm_verificationtoken.update_many(data={"spend": 0.0}, where={}) + if committed_spend_tracked(): + await _drop_all_key_and_team_committed_spend(prisma_client) return { "message": "Spend for all API Keys and Teams reset successfully", diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8b9189f143f..97c327cbe4e 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -9263,13 +9263,14 @@ async def monitor_spend_logs_queue( proxy_logging_obj: Proxy logging object """ from litellm.constants import ( + SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS, SPEND_LOG_QUEUE_POLL_INTERVAL, SPEND_LOG_QUEUE_SIZE_THRESHOLD, ) threshold: Final = SPEND_LOG_QUEUE_SIZE_THRESHOLD base_interval: Final = SPEND_LOG_QUEUE_POLL_INTERVAL - max_backoff: Final = 30.0 # Maximum backoff interval in seconds + max_backoff: Final = SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS backoff_multiplier: Final = 1.5 # Exponential backoff multiplier current_interval = base_interval flush_requested: Final = asyncio.Event() diff --git a/litellm/redis/stores/spend_counters/__init__.py b/litellm/redis/stores/spend_counters/__init__.py index 3ad06589f53..6a6dc85586e 100644 --- a/litellm/redis/stores/spend_counters/__init__.py +++ b/litellm/redis/stores/spend_counters/__init__.py @@ -6,15 +6,17 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Final from litellm._internal_context import with_service_target -from litellm.caching.redis_batch import Batch +from litellm.caching.redis_batch import Batch, PostCallSettlement from litellm.redis.stores._route import post_call_request, ready, route from litellm.redis.stores.spend_counters.handle import SpendCountersHandle from litellm.redis.stores.spend_counters.types import ( + COMMITTED_COUNTER_PREFIX, END_USER_COUNTER_PREFIX, SPEND_COUNTER_PREFIX, SPEND_COUNTERS_TARGET, SpendIncrement, SpendReading, + committed_counter, end_user_counter, is_unavailable, key_counter, @@ -34,12 +36,14 @@ if TYPE_CHECKING: from litellm.redis.stores.connection import RedisStores __all__ = [ + "COMMITTED_COUNTER_PREFIX", "END_USER_COUNTER_PREFIX", "SPEND_COUNTERS_TARGET", "SPEND_COUNTER_PREFIX", "SpendCounterStore", "SpendIncrement", "SpendReading", + "committed_counter", "end_user_counter", "is_unavailable", "key_counter", @@ -81,6 +85,11 @@ class SpendCounterStore: def has_redis(self) -> bool: return self._handle.has_redis + def set_committed_ttl(self, ttl: float) -> None: + """The least TTL a Redis write gives a committed twin, which keeps it past the DB write that holds its + charges.""" + self._handle.set_committed_ttl(ttl) + @with_service_target(SPEND_COUNTERS_TARGET) async def read(self, counter: str) -> SpendReading: reading: Final = await self._handle.read(counter) @@ -104,12 +113,13 @@ class SpendCounterStore: return tuple(await self._handle.increment_many(increments)) @with_service_target(SPEND_COUNTERS_TARGET) - def increment_many_post_call(self, increments: Sequence[SpendIncrement]) -> bool: - """Declares the increments on the request's post-call batch, False when there is no Redis or request. - Each total lands in memory once the batch goes out; a counter whose increment failed is dropped everywhere.""" + def increment_many_post_call(self, increments: Sequence[SpendIncrement]) -> PostCallSettlement | None: + """Declares the increments on the request's post-call batch, None when there is no Redis or request. + Each total lands in memory once the batch goes out; a counter whose increment failed is dropped everywhere. + The settlement reports when every reply is back, whoever sent the batch.""" request: Final = post_call_request() if self.has_redis and increments else None if request is None: - return False + return None return self._handle.increment_many_post_call(request, increments) @with_service_target(SPEND_COUNTERS_TARGET) diff --git a/litellm/redis/stores/spend_counters/handle.py b/litellm/redis/stores/spend_counters/handle.py index e4bd789c341..5ee29b7ad8e 100644 --- a/litellm/redis/stores/spend_counters/handle.py +++ b/litellm/redis/stores/spend_counters/handle.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import Awaitable, Mapping, Sequence from typing import Protocol -from litellm.caching.redis_batch import NativePending +from litellm.caching.redis_batch import NativePending, PostCallSettlement from litellm.redis.stores.spend_counters.types import SpendIncrement from litellm.rust_bridge._native import NativeRedisClient, NativeRequestRedis, NativeSend, SpendRead @@ -13,6 +13,7 @@ class SpendCountersHandle(Protocol): def has_redis(self) -> bool: ... @property def redis_ttl_seconds(self) -> float | None: ... + def set_committed_ttl(self, ttl_seconds: float, /) -> None: ... def attach(self, client: NativeRedisClient, /) -> bool: ... def read(self, counter: str) -> Awaitable[SpendRead]: ... def read_many( @@ -20,7 +21,9 @@ class SpendCountersHandle(Protocol): ) -> NativePending[Sequence[float | None] | None] | None: ... def increment(self, counter: str, amount: float) -> Awaitable[float]: ... def increment_many(self, increments: Sequence[SpendIncrement]) -> Awaitable[Sequence[float]]: ... - def increment_many_post_call(self, request: NativeRequestRedis, increments: Sequence[SpendIncrement]) -> bool: ... + def increment_many_post_call( + self, request: NativeRequestRedis, increments: Sequence[SpendIncrement] + ) -> PostCallSettlement | None: ... def set_max(self, counter: str, spend: float) -> Awaitable[float | None]: ... def seed(self, counter: str, spend: float, *, increment_if_lost: bool) -> Awaitable[float]: ... def set(self, counter: str, spend: float, ttl_seconds: float | None = None) -> Awaitable[None]: ... diff --git a/litellm/redis/stores/spend_counters/types.py b/litellm/redis/stores/spend_counters/types.py index 7dc846dabaa..9be029b6953 100644 --- a/litellm/redis/stores/spend_counters/types.py +++ b/litellm/redis/stores/spend_counters/types.py @@ -7,6 +7,7 @@ from litellm.rust_bridge._native import NativeRedisStoreError SPEND_COUNTERS_TARGET: Final = "spend_counters" SPEND_COUNTER_PREFIX: Final = "spend:" END_USER_COUNTER_PREFIX: Final = "spend:end_user:" +COMMITTED_COUNTER_PREFIX: Final = "spend:committed:" SpendIncrement: TypeAlias = tuple[str, float] SpendReading: TypeAlias = tuple[float | None, bool] @@ -61,6 +62,11 @@ def shadow_eval_counter(job_id: str) -> str: return f"spend:shadow_eval:{job_id}" +def committed_counter(counter_key: str) -> str: + """The charges-only twin of a key, team, user, org or end user lifetime counter.""" + return f"{COMMITTED_COUNTER_PREFIX}{counter_key.removeprefix(SPEND_COUNTER_PREFIX)}" + + def is_unavailable(error: BaseException) -> bool: """Redis timed out or its breaker is open, the failures a post-call increment drops.""" return isinstance(error, NativeRedisStoreError) and error.args[:1] in (("timeout",), ("breaker_open",)) diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 77af87f5eb3..114d2df026e 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -358,6 +358,11 @@ class SpendRead: @property def from_redis(self) -> bool: ... +@final +class PostCallSettlement: + @property + def settled(self) -> bool: ... + @final class NativeSpendCounters: def __new__(cls, client: NativeRedisClient | None) -> NativeSpendCounters: ... @@ -366,6 +371,7 @@ class NativeSpendCounters: def has_redis(self) -> bool: ... @property def redis_ttl_seconds(self) -> float | None: ... + def set_committed_ttl(self, ttl_seconds: float) -> None: ... def read(self, counter: str) -> Future[SpendRead]: ... def read_many( self, counters: Sequence[str], send: NativeSend @@ -374,7 +380,7 @@ class NativeSpendCounters: def increment_many(self, increments: Sequence[tuple[str, float]]) -> Future[list[float]]: ... def increment_many_post_call( self, request: NativeRequestRedis, increments: Sequence[tuple[str, float]] - ) -> bool: ... + ) -> PostCallSettlement | None: ... def set_max(self, counter: str, spend: float) -> Future[float | None]: ... def seed(self, counter: str, spend: float, *, increment_if_lost: bool) -> Future[float]: ... def set(self, counter: str, spend: float, ttl_seconds: float | None = None) -> Future[None]: ... @@ -596,6 +602,7 @@ __all__ = [ "NativeSend", "NativeSensitiveRoutePins", "NativeSpendCounters", + "PostCallSettlement", "ProcessReservedForForking", "ResponsesWebSocketConnection", "RouterCall", diff --git a/litellm/types/proxy/auth/user_api_key_auth.py b/litellm/types/proxy/auth/user_api_key_auth.py index 1183c36eda7..e08459ea14d 100644 --- a/litellm/types/proxy/auth/user_api_key_auth.py +++ b/litellm/types/proxy/auth/user_api_key_auth.py @@ -16,6 +16,7 @@ from litellm.types.agents import AgentCaller, AgentResponse from litellm.types.proxy.agent_identity import ManagedAgentContext from litellm.types.proxy.auth.user_roles import LitellmUserRoles from litellm.types.proxy.carried_budget_state import ( + CommittedBudgetSpend, OrgBudgetSnapshot, TeamBudgetSnapshot, UserBudgetSnapshot, @@ -206,6 +207,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob team_budget_snapshot: TeamBudgetSnapshot | None = Field(default=None, exclude=True) user_budget_snapshot: UserBudgetSnapshot | None = Field(default=None, exclude=True) org_budget_snapshot: OrgBudgetSnapshot | None = Field(default=None, exclude=True) + committed_budget_spend: CommittedBudgetSpend | None = Field(default=None, exclude=True, repr=False) matched_model_access_groups: list[str] | None = Field(default=None, exclude=True) budget_throttle_pct: float | None = Field(default=None, exclude=True) user: object | None = None # Expanded user object when expand=user is used @@ -329,6 +331,21 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob ) +COMMITTED_BUDGET_SPEND_PARAM: Final = "committed_budget_spend" + + +def committed_budget_spend_of_row( + metadata: Mapping[str, object], litellm_params: Mapping[str, object] +) -> CommittedBudgetSpend | None: + auth: Final = metadata.get("user_api_key_auth") + if isinstance(auth, UserAPIKeyAuth) and auth.committed_budget_spend is not None: + return auth.committed_budget_spend + if "user_api_key_auth" in metadata: + return CommittedBudgetSpend() + carried: Final = litellm_params.get(COMMITTED_BUDGET_SPEND_PARAM) + return carried if isinstance(carried, CommittedBudgetSpend) else None + + def user_api_key_has_admin_view(user_api_key_dict: UserAPIKeyAuth) -> bool: """Return True if the caller's role grants unscoped read access to all tenant resources (managed files, batches, vector stores, spend rows, etc). diff --git a/litellm/types/proxy/carried_budget_state.py b/litellm/types/proxy/carried_budget_state.py index 1a0c55a38ff..83ef6a21188 100644 --- a/litellm/types/proxy/carried_budget_state.py +++ b/litellm/types/proxy/carried_budget_state.py @@ -9,10 +9,18 @@ reads them back with ``from_metadata``. ``None`` means this request never carrie from collections.abc import Mapping from datetime import datetime from types import MappingProxyType +from typing import Final from pydantic import ConfigDict, Field, ValidationError from typing_extensions import Self +from litellm.redis.stores.spend_counters.types import ( + end_user_counter, + key_counter, + org_counter, + team_counter, + user_counter, +) from litellm.types.llms.base import LiteLLMBaseModel @@ -25,7 +33,7 @@ class _BudgetSnapshot(LiteLLMBaseModel): @classmethod def from_metadata(cls, metadata: Mapping[str, object]) -> Self | None: try: - return cls.model_validate(metadata) + return cls.model_validate(metadata, by_name=False) except ValidationError: return None @@ -64,3 +72,51 @@ class OrgBudgetSnapshot(_BudgetSnapshot): max_budget: float | None = Field( validation_alias="user_api_key_org_max_budget", serialization_alias="user_api_key_org_max_budget" ) + + +class EntitySpendCounterKeys(LiteLLMBaseModel): + """Spend counter of each entity whose remaining budget a gauge reports. ``None`` for an entity with no budget.""" + + model_config = ConfigDict(frozen=True) + + key: str | None = None + team: str | None = None + user: str | None = None + org: str | None = None + end_user: str | None = None + + @classmethod + def of( + cls, + token: str | None, + team_id: str | None, + user_id: str | None, + org_id: str | None, + end_user_id: str | None, + ) -> Self: + return cls( + key=key_counter(token) if token else None, + team=team_counter(team_id) if team_id else None, + user=user_counter(user_id) if user_id else None, + org=org_counter(org_id) if org_id else None, + end_user=end_user_counter(end_user_id) if end_user_id else None, + ) + + def tracked(self) -> tuple[str, ...]: + counter_keys: Final = (self.key, self.team, self.user, self.org, self.end_user) + return tuple(counter_key for counter_key in counter_keys if counter_key is not None) + + +class CommittedBudgetSpend(LiteLLMBaseModel): + """Spend each entity was charged, budget holds excluded, read from the committed spend counters for the + remaining budget gauges only. ``None`` when the entity has no committed counter. ``counter_keys`` names the + budgeted entities whose committed counters this request charges and reads.""" + + model_config = ConfigDict(frozen=True) + + counter_keys: EntitySpendCounterKeys = Field(default=EntitySpendCounterKeys(), repr=False) + key: float | None = Field(default=None, repr=False) + team: float | None = Field(default=None, repr=False) + user: float | None = Field(default=None, repr=False) + org: float | None = Field(default=None, repr=False) + end_user: float | None = Field(default=None, repr=False) diff --git a/tests/unit/caching/test_request_redis_batch_post_call.py b/tests/unit/caching/test_request_redis_batch_post_call.py index 41d508a0eea..01feee0be54 100644 --- a/tests/unit/caching/test_request_redis_batch_post_call.py +++ b/tests/unit/caching/test_request_redis_batch_post_call.py @@ -424,7 +424,7 @@ async def test_spend_counter_increments_ride_the_native_post_call_batch(monkeypa with request_redis_batch_scope(native_factory=natives, native_send_factory=fake_send): # pyright: ignore[reportArgumentType] # test double await proxy_server._apply_spend_counter_increments(pending) - declared = list(spend.post_call) + declared = [(request, increments) for request, increments, _ in spend.post_call] assert natives.built[0].events == [] await flush_post_call_redis_batches() @@ -433,6 +433,27 @@ async def test_spend_counter_increments_ride_the_native_post_call_batch(monkeypa assert spend.commands == [] +@pytest.mark.asyncio +async def test_the_request_sees_its_committed_spend_charge_settle_whoever_sends_the_batch(monkeypatch): + from litellm.proxy import proxy_server + + spend = FakeSpendCounters(redis={}) + monkeypatch.setattr(proxy_server, "spend_counter_store", spend.store()) + + with request_redis_batch_scope(native_factory=FakeNatives(), native_send_factory=fake_send) as request: # pyright: ignore[reportArgumentType] # test double + await proxy_server._apply_spend_counter_increments([PendingSpendIncrement("spend:key:k1", 0.5)]) + without_committed_spend = request.post_call_settlements + await proxy_server._apply_spend_counter_increments( + [PendingSpendIncrement("spend:key:k1", 0.5), PendingSpendIncrement("spend:committed:key:k1", 0.5)] + ) + declared = [settlement.settled for settlement in request.post_call_settlements] + spend.settle_post_call() + sent = [settlement.settled for settlement in request.post_call_settlements] + + assert without_committed_spend == () + assert (declared, sent) == ([False], [True]) + + @pytest.mark.asyncio async def test_native_post_call_ops_go_out_with_the_post_call_flush_not_the_request_boundary(): natives = FakeNatives() diff --git a/tests/unit/integrations/generic_api/test_generic_api_callback.py b/tests/unit/integrations/generic_api/test_generic_api_callback.py index afaae02321d..8f1e79562db 100644 --- a/tests/unit/integrations/generic_api/test_generic_api_callback.py +++ b/tests/unit/integrations/generic_api/test_generic_api_callback.py @@ -20,6 +20,8 @@ from litellm.integrations.gcs_pubsub.pub_sub import * from litellm.integrations.generic_api.generic_api_callback import LOG_FORMAT_TYPES, GenericAPILogger from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys from litellm.types.utils import StandardLoggingPayload from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -795,3 +797,27 @@ async def test_generic_api_concurrent_flushes_deliver_each_event_once(log_format task.cancel() await asyncio.gather(*owned_tasks, return_exceptions=True) await handler.close() + + +@pytest.mark.asyncio +async def test_v1_payload_of_a_chat_row_never_shows_the_committed_spend_auth_read() -> None: + counter_keys: Final = EntitySpendCounterKeys.of( + token="hashed-token", team_id="t1", user_id=None, org_id=None, end_user_id=None + ) + auth: Final = UserAPIKeyAuth( + api_key="hashed-token", + team_id="t1", + committed_budget_spend=CommittedBudgetSpend(counter_keys=counter_keys, key=0.123456, team=0.654321), + ) + generic_logger: Final = GenericAPILogger(endpoint="https://example.com/api/logs", flush_interval=3600) + + payload: Final = generic_logger._get_v1_logging_payload( + kwargs={"litellm_params": {"metadata": {"user_api_key_auth": auth}}}, + response_obj={"usage": {}}, + start_time=None, + end_time=None, + ) + + assert "team_id='t1'" in payload["metadata"] + assert "0.123456" not in payload["metadata"] + assert str(counter_keys.team) not in payload["metadata"] diff --git a/tests/unit/integrations/test_prometheus_carried_budget_state.py b/tests/unit/integrations/test_prometheus_carried_budget_state.py index e49e6baa8b0..9e006c20079 100644 --- a/tests/unit/integrations/test_prometheus_carried_budget_state.py +++ b/tests/unit/integrations/test_prometheus_carried_budget_state.py @@ -4,13 +4,22 @@ from request metadata. get_*_object only runs when that state is missing (custom auth, SDK callers, failure paths) """ +import asyncio +from collections.abc import Sequence from datetime import datetime, timezone +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest from prometheus_client import REGISTRY +from litellm.caching.redis_batch import ( + active_post_call_redis_batch, + flush_post_call_redis_batches, + request_redis_batch_scope, +) from litellm.integrations.prometheus import PrometheusLogger +from litellm.proxy import proxy_server from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_OrganizationTable, @@ -24,11 +33,16 @@ from litellm.proxy.spend_tracking.carried_budget_state import ( carry_organization_budget_state, carry_team_and_user_budget_state, ) +from litellm.proxy.spend_tracking.spend_counter_batch import PendingSpendIncrement from litellm.types.proxy.carried_budget_state import ( + CommittedBudgetSpend, + EntitySpendCounterKeys, KeyBudgetSnapshot, TeamBudgetSnapshot, UserBudgetSnapshot, ) +from tests.unit.caching.test_redis_batch import FakeNatives, fake_send, make +from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters TEAM_RESET_AT = datetime(2026, 10, 1, tzinfo=timezone.utc) USER_RESET_AT = datetime(2026, 11, 1, tzinfo=timezone.utc) @@ -55,6 +69,15 @@ def prometheus_logger(): return PrometheusLogger() +@pytest.fixture(autouse=True) +def spend_counters(): + spend: Final = FakeSpendCounters(redis=None) + with patch( # test-quality-ok: prometheus reads the proxy_server spend counters at emit time, no injection seam + "litellm.proxy.proxy_server.spend_counter_store", spend.store() + ): + yield spend + + @pytest.fixture def getters(): """Every response-path object getter, patched where prometheus imports them from.""" @@ -137,6 +160,7 @@ def _request_metadata(token: UserAPIKeyAuth) -> dict: return { **LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(token), **carried_budget_metadata(token), + "user_api_key_auth": token, } @@ -198,6 +222,499 @@ async def test_authed_request_sets_every_gauge_without_any_object_getter(prometh prometheus_logger.litellm_user_budget_remaining_hours_metric.labels().set.assert_called_once() +_COUNTER_KEYS: Final = ("spend:key:hashed", "spend:team:t1", "spend:user:u1", "spend:org:o1") +_TWIN_KEYS: Final = ( + "spend:committed:key:hashed", + "spend:committed:team:t1", + "spend:committed:user:u1", + "spend:committed:org:o1", +) + + +def _store(spend: FakeSpendCounters, counter_keys: Sequence[str], values: Sequence[float]) -> None: + tier: Final = spend.memory if spend.redis is None else spend.redis + tier.update(zip(counter_keys, values, strict=True)) + + +def _store_counters(spend: FakeSpendCounters, key: float, team: float, user: float, org: float) -> None: + _store(spend, _COUNTER_KEYS, (key, team, user, org)) + + +def _store_committed(spend: FakeSpendCounters, key: float, team: float, user: float, org: float) -> None: + _store(spend, _TWIN_KEYS, (key, team, user, org)) + + +def _remaining(prometheus_logger: PrometheusLogger) -> dict[str, float]: + return { + "key": prometheus_logger.litellm_remaining_api_key_budget_metric.labels().set.call_args[0][0], + "team": prometheus_logger.litellm_remaining_team_budget_metric.labels().set.call_args[0][0], + "user": prometheus_logger.litellm_remaining_user_budget_metric.labels().set.call_args[0][0], + "org": prometheus_logger.litellm_remaining_org_budget_metric.labels().set.call_args[0][0], + } + + +def _token_committed_at_auth() -> UserAPIKeyAuth: + token = _authed_token() + token.committed_budget_spend = CommittedBudgetSpend( + counter_keys=EntitySpendCounterKeys.of( + token="hashed", team_id="t1", user_id="u1", org_id="o1", end_user_id=None + ), + key=3.0, + team=30.0, + user=8.0, + org=60.0, + ) + return token + + +@pytest.mark.asyncio +async def test_gauges_use_the_committed_spend_at_emit_when_other_requests_finished_after_auth( + prometheus_logger, getters, spend_counters +): + """This request's 2.0 plus 1.0 from another request that finished first, so the gauge is never set back.""" + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=6.0, team=33.0, user=11.0, org=63.0) + + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(4.0), + "team": pytest.approx(267.0), + "user": pytest.approx(39.0), + "org": pytest.approx(437.0), + } + + +@pytest.mark.asyncio +async def test_gauges_add_the_cost_to_another_requests_charge_while_the_own_charge_is_not_in_yet( + prometheus_logger, getters, spend_counters +): + """1.0 from another request reached the counters during this long request, this request's 2.0 has not.""" + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=4.0, team=31.0, user=9.0, org=61.0) + + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(4.0), + "team": pytest.approx(267.0), + "user": pytest.approx(39.0), + "org": pytest.approx(437.0), + } + + +@pytest.mark.asyncio +async def test_gauges_read_no_committed_counter_auth_did_not_name_as_budgeted( + prometheus_logger, getters, spend_counters +): + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=9.0, team=99.0, user=19.0, org=99.0) + + await _emit(prometheus_logger, _request_metadata(_authed_token())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(7.0), + "team": pytest.approx(278.0), + "user": pytest.approx(43.0), + "org": 458.0, + } + + +def _billed_row_metadata() -> dict: + return {key: value for key, value in _request_metadata(_authed_token()).items() if key != "user_api_key_auth"} + + +@pytest.mark.asyncio +async def test_row_with_no_auth_adds_its_cost_to_the_committed_spend_of_the_entities_it_is_charged_to( + prometheus_logger, getters, spend_counters +): + """Counter increments flush after the success callbacks, so the row's own charge is not in the counters yet.""" + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=6.0, team=33.0, user=11.0, org=63.0) + + await _emit(prometheus_logger, _billed_row_metadata()) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(2.0), + "team": pytest.approx(265.0), + "user": pytest.approx(37.0), + "org": 435.0, + } + + +@pytest.mark.asyncio +async def test_row_with_no_auth_and_no_committed_counters_keeps_the_stored_spend( + prometheus_logger, getters, spend_counters +): + _stub_gauges(prometheus_logger) + + await _emit(prometheus_logger, _billed_row_metadata()) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(7.0), + "team": pytest.approx(278.0), + "user": pytest.approx(43.0), + "org": 458.0, + } + + +@pytest.mark.asyncio +async def test_slow_committed_read_on_a_failed_request_is_bounded_and_keeps_the_auth_time_spend( + prometheus_logger, getters, monkeypatch +): + async def never_answers(counter: str) -> None: + await asyncio.sleep(30) + + monkeypatch.setattr( + "litellm.proxy.proxy_server.spend_counter_store", + FakeSpendCounters(redis={}, batch_reads=False, read_hook=never_answers).store(), + ) + monkeypatch.setenv("PROMETHEUS_BUDGET_METRICS_PER_REQUEST_TIMEOUT", "0.01") + _stub_gauges(prometheus_logger) + + await asyncio.wait_for( + prometheus_logger.async_log_failure_event( + _failed_request(_request_metadata(_token_committed_at_auth())), None, KEY_RESET_AT, KEY_RESET_AT + ), + timeout=5, + ) + + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(440.0)) + + +@pytest.mark.asyncio +async def test_slow_committed_read_on_a_successful_request_still_sets_the_gauges_from_the_auth_time_spend( + prometheus_logger, getters, monkeypatch +): + async def never_answers(counter: str) -> None: + await asyncio.sleep(30) + + monkeypatch.setattr( + "litellm.proxy.proxy_server.spend_counter_store", + FakeSpendCounters(redis={}, batch_reads=False, read_hook=never_answers).store(), + ) + monkeypatch.setenv("PROMETHEUS_BUDGET_METRICS_PER_REQUEST_TIMEOUT", "0.5") + _stub_gauges(prometheus_logger) + + await asyncio.wait_for(_emit(prometheus_logger, _request_metadata(_token_committed_at_auth())), timeout=5) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(5.0), + "team": pytest.approx(268.0), + "user": pytest.approx(40.0), + "org": pytest.approx(438.0), + } + + +@pytest.fixture +def redis_spend_counters(): + spend: Final = FakeSpendCounters(redis={}) + with patch( # test-quality-ok: prometheus reads the proxy_server spend counters at emit time, no injection seam + "litellm.proxy.proxy_server.spend_counter_store", spend.store() + ): + yield spend + + +def _request_scope() -> request_redis_batch_scope: + return request_redis_batch_scope( + native_factory=FakeNatives(), # pyright: ignore[reportArgumentType] # test double + native_send_factory=fake_send, # pyright: ignore[reportArgumentType] # test double + ) + + +async def _declare_the_own_two_dollar_charge() -> None: + await proxy_server._apply_spend_counter_increments( + tuple(PendingSpendIncrement(counter_key, 2.0) for counter_key in (*_COUNTER_KEYS, *_TWIN_KEYS)) + ) + + +_BOTH_TWO_DOLLAR_CHARGES_COUNTED_ONCE: Final = { + "key": pytest.approx(3.0), + "team": pytest.approx(266.0), + "user": pytest.approx(38.0), + "org": pytest.approx(436.0), +} + + +@pytest.mark.asyncio +async def test_gauges_add_the_own_charge_to_a_same_size_charge_of_another_request_while_the_own_waits_in_the_pipeline( + prometheus_logger, getters, redis_spend_counters +): + """Another request's 2.0 reached the counters and this request's 2.0 is queued, so both count.""" + _stub_gauges(prometheus_logger) + _store_committed(redis_spend_counters, key=5.0, team=32.0, user=10.0, org=62.0) + + with _request_scope(): + await _declare_the_own_two_dollar_charge() + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + await flush_post_call_redis_batches() + + assert _remaining(prometheus_logger) == _BOTH_TWO_DOLLAR_CHARGES_COUNTED_ONCE + + +@pytest.mark.asyncio +async def test_gauges_do_not_add_the_own_charge_twice_once_the_pipeline_went_out( + prometheus_logger, getters, redis_spend_counters +): + """The 2.0 over the auth time spend is this request's own charge, settled before the gauges read it.""" + _stub_gauges(prometheus_logger) + _store_committed(redis_spend_counters, key=3.0, team=30.0, user=8.0, org=60.0) + + with _request_scope(): + await _declare_the_own_two_dollar_charge() + redis_spend_counters.settle_post_call() + await flush_post_call_redis_batches() + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(5.0), + "team": pytest.approx(268.0), + "user": pytest.approx(40.0), + "org": pytest.approx(438.0), + } + + +@pytest.mark.asyncio +async def test_an_unrelated_post_call_flush_does_not_make_the_gauge_drop_its_own_unsent_charge( + prometheus_logger, getters, redis_spend_counters +): + """An unrelated write went out first, then another request's 2.0 reached the counters. This request's 2.0 counts + once whether it is undeclared, declared but unsent, or settled.""" + _stub_gauges(prometheus_logger) + _store_committed(redis_spend_counters, key=5.0, team=32.0, user=10.0, org=62.0) + unrelated_cache, _ = make() + metadata: Final = _request_metadata(_token_committed_at_auth()) + + with _request_scope(): + active_post_call_redis_batch(unrelated_cache).increment("unrelated-write", 1.0) + await flush_post_call_redis_batches() + await _emit(prometheus_logger, metadata) + undeclared: Final = _remaining(prometheus_logger) + await _declare_the_own_two_dollar_charge() + await _emit(prometheus_logger, metadata) + declared_unsent: Final = _remaining(prometheus_logger) + redis_spend_counters.settle_post_call() + await _emit(prometheus_logger, metadata) + settled: Final = _remaining(prometheus_logger) + await flush_post_call_redis_batches() + + assert {"undeclared": undeclared, "declared_unsent": declared_unsent, "settled": settled} == { + "undeclared": _BOTH_TWO_DOLLAR_CHARGES_COUNTED_ONCE, + "declared_unsent": _BOTH_TWO_DOLLAR_CHARGES_COUNTED_ONCE, + "settled": _BOTH_TWO_DOLLAR_CHARGES_COUNTED_ONCE, + } + + +@pytest.mark.asyncio +async def test_gauges_count_the_own_charge_once_when_it_settles_while_the_gauges_read_the_counters( + prometheus_logger, getters, redis_spend_counters +): + """Another request's 2.0 reached the counters, and this request's 2.0 goes out during the gauges' read.""" + _stub_gauges(prometheus_logger) + _store_committed(redis_spend_counters, key=5.0, team=32.0, user=10.0, org=62.0) + + async def settle_during_read(_counter: str) -> None: + redis_spend_counters.settle_post_call() + + redis_spend_counters.batch_reads = False + redis_spend_counters.read_hook = settle_during_read + + with _request_scope(): + await _declare_the_own_two_dollar_charge() + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + await flush_post_call_redis_batches() + + assert _remaining(prometheus_logger) == _BOTH_TWO_DOLLAR_CHARGES_COUNTED_ONCE + + +@pytest.mark.asyncio +async def test_gauges_without_redis_treat_a_charge_equal_to_the_own_as_the_own( + prometheus_logger, getters, spend_counters +): + """Without Redis the charge increments in memory at once, so a request scope says nothing about it.""" + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=5.0, team=32.0, user=10.0, org=62.0) + + with request_redis_batch_scope(): + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(5.0), + "team": pytest.approx(268.0), + "user": pytest.approx(40.0), + "org": pytest.approx(438.0), + } + + +@pytest.mark.asyncio +async def test_gauges_add_the_cost_to_the_auth_time_spend_before_the_charge_reaches_the_counters( + prometheus_logger, getters, spend_counters +): + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=3.0, team=30.0, user=8.0, org=60.0) + + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(5.0), + "team": pytest.approx(268.0), + "user": pytest.approx(40.0), + "org": pytest.approx(438.0), + } + + +@pytest.mark.asyncio +async def test_gauges_keep_the_auth_time_spend_when_the_proxy_spend_counters_cannot_be_read( + prometheus_logger, getters, monkeypatch +): + _stub_gauges(prometheus_logger) + monkeypatch.delattr("litellm.proxy.proxy_server.spend_counter_store") + + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(5.0), + "team": pytest.approx(268.0), + "user": pytest.approx(40.0), + "org": pytest.approx(438.0), + } + + +@pytest.mark.asyncio +async def test_a_budget_hold_on_the_enforcement_counters_never_shows_in_the_gauges( + prometheus_logger, getters, spend_counters +): + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=5.0, team=32.0, user=10.0, org=62.0) + _store_counters(spend_counters, key=9.0, team=36.0, user=14.0, org=66.0) + + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(5.0), + "team": pytest.approx(268.0), + "user": pytest.approx(40.0), + "org": pytest.approx(438.0), + } + + +def _failed_request(metadata: dict, response_cost: float = 0.0) -> dict: + return { + "model": "openai/gpt-5.4-nano", + "litellm_params": {"metadata": metadata}, + "standard_logging_object": { + "model": "openai/gpt-5.4-nano", + "model_id": "dep-1", + "model_group": "chat", + "api_base": "https://api.openai.com", + "custom_llm_provider": "openai", + "status": "failure", + "response_cost": response_cost, + "request_tags": [], + "hidden_params": {"litellm_overhead_time_ms": None, "additional_headers": None}, + "metadata": { + "user_api_key_hash": "hashed", + "user_api_key_alias": "key-alias", + "user_api_key_team_id": "t1", + "user_api_key_team_alias": "team-alias", + "user_api_key_user_id": "u1", + "user_api_key_org_id": "o1", + "requester_metadata": None, + "user_api_key_end_user_id": None, + }, + }, + "exception": Exception("upstream failed"), + } + + +@pytest.mark.asyncio +async def test_failed_request_sets_the_org_gauge_from_the_carried_org_and_committed_spend( + prometheus_logger, getters, spend_counters +): + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=3.0, team=30.0, user=8.0, org=65.0) + _store_counters(spend_counters, key=9.0, team=36.0, user=14.0, org=90.0) + + await prometheus_logger.async_log_failure_event( + _failed_request(_request_metadata(_token_committed_at_auth())), None, KEY_RESET_AT, KEY_RESET_AT + ) + + assert getters["get_org_object"].await_count == 0 + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(435.0)) + + +@pytest.mark.asyncio +async def test_failed_request_without_the_auth_object_still_reads_the_committed_org_spend( + prometheus_logger, getters, spend_counters +): + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=3.0, team=30.0, user=8.0, org=65.0) + metadata = _request_metadata(_token_committed_at_auth()) + del metadata["user_api_key_auth"] + + await prometheus_logger.async_log_failure_event(_failed_request(metadata), None, KEY_RESET_AT, KEY_RESET_AT) + + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(435.0)) + + +@pytest.mark.asyncio +async def test_failed_request_without_a_carried_org_reads_no_committed_org_spend( + prometheus_logger, getters, spend_counters +): + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=3.0, team=30.0, user=8.0, org=65.0) + token = _token_committed_at_auth().model_copy(update={"org_budget_snapshot": None}) + metadata = {key: value for key, value in _request_metadata(token).items() if key != "user_api_key_auth"} + + await prometheus_logger.async_log_failure_event(_failed_request(metadata), None, KEY_RESET_AT, KEY_RESET_AT) + + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(460.0)) + + +@pytest.mark.asyncio +async def test_failed_request_that_was_charged_counts_its_own_charge_while_it_waits_in_the_pipeline( + prometheus_logger, getters, redis_spend_counters +): + """A stream cut after 40 events still costs 2.0, and that charge is queued when the failure logs.""" + _stub_gauges(prometheus_logger) + _store_committed(redis_spend_counters, key=5.0, team=32.0, user=10.0, org=62.0) + + with _request_scope(): + await _declare_the_own_two_dollar_charge() + await prometheus_logger.async_log_failure_event( + _failed_request(_request_metadata(_token_committed_at_auth()), response_cost=2.0), + None, + KEY_RESET_AT, + KEY_RESET_AT, + ) + await flush_post_call_redis_batches() + + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(436.0)) + + +@pytest.mark.asyncio +async def test_ordinary_metadata_keys_named_like_snapshot_fields_change_no_gauge(prometheus_logger, getters): + _stub_gauges(prometheus_logger) + metadata = { + "user_api_key_team_spend": 20.0, + "user_api_key_team_max_budget": 300.0, + "spend": 1000.0, + "max_budget": 1.0, + "budget_reset_at": None, + "user_alias": "not-the-user", + } + + await _emit(prometheus_logger, metadata) + + assert {name: getter.await_count for name, getter in getters.items()} == { + "get_key_object": 1, + "get_team_object": 1, + "get_user_object": 1, + "get_org_object": 1, + } + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(458.0) + + @pytest.mark.asyncio async def test_metadata_without_carried_state_still_fetches_each_object_once(prometheus_logger, getters): _stub_gauges(prometheus_logger) diff --git a/tests/unit/integrations/test_prometheus_user_team_metrics.py b/tests/unit/integrations/test_prometheus_user_team_metrics.py index ab0fb67b52f..acec1c882ef 100644 --- a/tests/unit/integrations/test_prometheus_user_team_metrics.py +++ b/tests/unit/integrations/test_prometheus_user_team_metrics.py @@ -1065,6 +1065,79 @@ async def test_set_customer_budget_metrics_after_api_request_uses_cached_default assert _customer_sample("litellm_customer_max_budget_metric", "cust-default") == pytest.approx(3.0) +@pytest.mark.asyncio +async def test_customer_gauge_subtracts_the_committed_spend_instead_of_the_cached_spend( + prometheus_logger, customer_metrics_enabled +): + """Auth saw 1.5 committed, and by emit the counter holds 2.5 with this request's 0.5 and another request's.""" + import sys + + from litellm.models.budget import LiteLLM_BudgetTable + from litellm.models.end_user import LiteLLM_EndUserTable + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys + from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + mock_proxy_server = MagicMock() + mock_proxy_server.user_api_key_cache.async_get_cache = AsyncMock( + return_value=LiteLLM_EndUserTable( + user_id="cust-live", + blocked=False, + spend=0.5, + litellm_budget_table=LiteLLM_BudgetTable(budget_id="b", max_budget=3.0), + ) + ) + mock_proxy_server.spend_counter_store = FakeSpendCounters( + redis=None, memory={"spend:committed:end_user:cust-live": 2.5} + ).store() + token = UserAPIKeyAuth(token="hashed") + token.committed_budget_spend = CommittedBudgetSpend( + counter_keys=EntitySpendCounterKeys(end_user="spend:end_user:cust-live"), end_user=1.5 + ) + + with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}): + await prometheus_logger._increment_remaining_budget_metrics( + user_api_team=None, + user_api_team_alias=None, + user_api_key=None, + user_api_key_alias=None, + litellm_params={"metadata": {"user_api_key_auth": token}}, + response_cost=0.5, + end_user_id="cust-live", + ) + + assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-live") == pytest.approx(0.5) + + +@pytest.mark.asyncio +async def test_customer_gauge_counts_a_charge_the_cached_customer_row_already_holds_once( + prometheus_logger, customer_metrics_enabled +): + import sys + + from litellm.models.budget import LiteLLM_BudgetTable + from litellm.models.end_user import LiteLLM_EndUserTable + + mock_proxy_server = MagicMock() + mock_proxy_server.user_api_key_cache.async_get_cache = AsyncMock( + return_value=LiteLLM_EndUserTable( + user_id="cust-charged", + blocked=False, + spend=0.5, + litellm_budget_table=LiteLLM_BudgetTable(budget_id="b", max_budget=3.0), + ) + ) + + with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}): + await prometheus_logger._set_customer_budget_metrics_after_api_request( + end_user_id="cust-charged", + response_cost=0.5, + committed_spend=0.0, + ) + + assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-charged") == pytest.approx(2.5) + + @pytest.mark.asyncio async def test_set_customer_budget_metrics_after_api_request_without_budget_only_emits_remaining( prometheus_logger, customer_metrics_enabled 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 8440d3d2ab2..70208bec0e5 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 @@ -5094,6 +5094,70 @@ async def test_centralized_common_checks_carries_team_and_user_budget_state_on_t } +@pytest.mark.asyncio +@pytest.mark.parametrize("prometheus_registered", [True, False], ids=["prometheus_logger", "no_prometheus_logger"]) +async def test_centralized_common_checks_carry_the_committed_spend_beside_the_stale_cached_spend( + monkeypatch: pytest.MonkeyPatch, prometheus_registered: bool +): + """The enforcement counters also hold budget holds, the committed twins only charged spend.""" + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + from litellm.integrations.prometheus import PrometheusLogger + from litellm.proxy.spend_tracking.committed_spend import committed_spend_key + from litellm.proxy.utils import ProxyLogging + from litellm.redis.stores.connection import connect + from litellm.redis.stores.spend_counters import SpendCounterStore + from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys + + for callback_list in ("success_callback", "failure_callback", "_async_success_callback", "_async_failure_callback"): + monkeypatch.setattr(litellm, callback_list, []) + monkeypatch.setattr(litellm, "callbacks", [MagicMock(spec=PrometheusLogger)] if prometheus_registered else []) + token: Final = UserAPIKeyAuth(api_key="sk-live-spend", max_budget=10.0, spend=0.0, team_id="t-live", team_spend=0.0) + user_api_key_cache: Final = UserApiKeyCache() + user_api_key_cache.set_cache( + key="team_id:t-live", value=LiteLLM_TeamTableCachedObj(team_id="t-live", max_budget=20.0, spend=0.0) + ) + spend_counter_store: Final = SpendCounterStore(connect(None).spend_counters()) + for counter_key, held, committed in ((f"spend:key:{token.token}", 0.67, 0.37), ("spend:team:t-live", 1.55, 1.25)): + spend_counter_store.set_local(counter_key, held) + spend_counter_store.set_local(committed_spend_key(counter_key), committed) + attrs: Final = { + **_proxy_attrs_for_centralized_checks(), + "llm_router": None, + "prisma_client": MagicMock(), + "user_api_key_cache": user_api_key_cache, + "spend_counter_store": spend_counter_store, + "proxy_logging_obj": ProxyLogging(user_api_key_cache=user_api_key_cache), + } + originals: Final = {attribute: getattr(_proxy_server_mod, attribute, None) for attribute in attrs} + request: Final = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + try: + for attribute, value in attrs.items(): + setattr(_proxy_server_mod, attribute, value) + await run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 1000}, + route="/chat/completions", + ) + finally: + for attribute, value in originals.items(): + setattr(_proxy_server_mod, attribute, value) + + assert token.committed_budget_spend == ( + CommittedBudgetSpend( + counter_keys=EntitySpendCounterKeys(key=f"spend:key:{token.token}", team=f"spend:team:{token.team_id}"), + key=0.37, + team=1.25, + ) + if prometheus_registered + else None + ) + assert (token.spend, token.team_spend) == (0.0, 0.0) + + def _end_user_budget_row(budget_id: str, max_budget: float) -> MagicMock: row = MagicMock() row.dict = lambda: {"budget_id": budget_id, "max_budget": max_budget} diff --git a/tests/unit/proxy/common_utils/test_reset_budget_job.py b/tests/unit/proxy/common_utils/test_reset_budget_job.py index 4a4490afc33..aaa686f2a98 100644 --- a/tests/unit/proxy/common_utils/test_reset_budget_job.py +++ b/tests/unit/proxy/common_utils/test_reset_budget_job.py @@ -1438,7 +1438,47 @@ def test_reset_budget_for_keys_invalidates_redis_counter(reset_budget_job, mock_ asyncio.run(reset_budget_job.reset_budget_for_litellm_keys()) - assert "spend:key:sk-abc" in counter_cache.deleted() + assert {"spend:key:sk-abc", "spend:committed:key:sk-abc"} <= counter_cache.deleted() + + +def _expired_key_row() -> object: + return type( + "Key", + (), + { + "spend": 100.0, + "budget_duration": "30d", + "budget_reset_at": datetime.now(timezone.utc), + "id": "key-1", + "token": "sk-abc", + }, + ) + + +def test_reset_budget_for_keys_without_redis_clears_the_counter_and_its_committed_twin( + reset_budget_job, mock_prisma_client, monkeypatch +): + counter_cache: Final = _make_counter_invalidation_job(monkeypatch) + counter_cache.redis = None + counter_cache.memory.update({"spend:key:sk-abc": 100.0, "spend:committed:key:sk-abc": 100.0}) + mock_prisma_client.data["key"] = [_expired_key_row()] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_keys()) + + assert counter_cache.memory == {} + + +def test_a_redis_failure_on_the_key_reset_still_drops_its_committed_twin_locally( + reset_budget_job, mock_prisma_client, monkeypatch +): + counter_cache: Final = _make_counter_invalidation_job(monkeypatch) + counter_cache.failing = frozenset({"delete"}) + counter_cache.memory.update({"spend:key:sk-abc": 100.0, "spend:committed:key:sk-abc": 100.0}) + mock_prisma_client.data["key"] = [_expired_key_row()] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_keys()) + + assert counter_cache.memory == {} def test_reset_budget_for_users_invalidates_redis_counter(reset_budget_job, mock_prisma_client, monkeypatch): @@ -1790,7 +1830,9 @@ def test_enduser_invalidation_is_paged_and_batched(reset_budget_job, mock_prisma assert counter_cache.user_api_key_cache.async_delete_cache_keys.await_count == 6 invalidated: Final = counter_cache.deleted() - assert invalidated == {f"spend:end_user:cust-{i:06d}" for i in range(population)} + assert invalidated == {f"spend:end_user:cust-{i:06d}" for i in range(population)} | { + f"spend:committed:end_user:cust-{i:06d}" for i in range(population) + } evicted: Final = { key for call in counter_cache.user_api_key_cache.async_delete_cache_keys.await_args_list for key in call.args[0] } diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index d06f95111aa..0283c90f3d6 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -44,6 +44,7 @@ from litellm.proxy.utils import ( BilledFailedAttempt, ProxyUpdateSpend, ) +from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys from litellm.types.router import Deployment, LiteLLM_Params from litellm.types.utils import CallTypes, LiteLLMBatch, ModelResponse, PromptTokensDetailsWrapper, Usage from litellm.types.videos.main import VideoObject @@ -1129,6 +1130,95 @@ def test_get_budget_reservation_from_metadata_handles_dict_auth_object(): ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("metadata", "expected_scope"), + [ + ( + { + "user_api_key_auth": UserAPIKeyAuth( + committed_budget_spend=CommittedBudgetSpend( + counter_keys=EntitySpendCounterKeys(key="spend:key:test_api_key") + ) + ) + }, + EntitySpendCounterKeys(key="spend:key:test_api_key"), + ), + ({}, EntitySpendCounterKeys()), + ], +) +async def test_charge_moves_only_the_committed_counters_auth_named_as_budgeted(metadata, expected_scope): + proxy_logging_obj = MagicMock() + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock() + increment_spend_counters = AsyncMock() + + await _update_database_and_spend_counters( + proxy_logging_obj=proxy_logging_obj, + increment_spend_counters=increment_spend_counters, + user_api_key="test_api_key", + user_id="test_user_id", + end_user_id=None, + team_id="test_team_id", + org_id="test_org_id", + kwargs={"litellm_params": {"metadata": metadata}}, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.2, + budget_reservation=None, + ) + + assert increment_spend_counters.await_args.kwargs["committed_spend_scope"] == expected_scope + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("tracked", "expected_key", "expected_org"), + [(True, pytest.approx(1.2), pytest.approx(2.2)), (False, 1.0, 2.0)], +) +async def test_charge_with_no_auth_on_its_row_moves_the_committed_counters_that_exist( + tracked, expected_key, expected_org +): + """A background response is billed from its creator's stored attribution, after the request's auth is gone.""" + from litellm.proxy.proxy_server import increment_spend_counters + + proxy_logging_obj = MagicMock() + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock() + spend_counter_store: Final = SpendCounterStore(connect(None).spend_counters()) + spend_counter_store.set_local("spend:committed:key:test_api_key", 1.0) + spend_counter_store.set_local("spend:committed:org:test_org_id", 2.0) + + with ( + patch("litellm.proxy.hooks.proxy_track_cost_callback.committed_spend_tracked", return_value=tracked), + patch("litellm.proxy.proxy_server.spend_counter_store", spend_counter_store), + patch( + "litellm.proxy.spend_tracking.committed_spend.SpendCounterReseed.from_db", + AsyncMock(return_value=50.0), + ), + ): + await _update_database_and_spend_counters( + proxy_logging_obj=proxy_logging_obj, + increment_spend_counters=increment_spend_counters, + user_api_key="test_api_key", + user_id="test_user_id", + end_user_id=None, + team_id="test_team_id", + org_id="test_org_id", + kwargs={"litellm_params": {"metadata": {"user_api_key_hash": "test_api_key"}}}, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.2, + budget_reservation=None, + ) + + cached: Final = spend_counter_store.get_local + assert cached("spend:committed:key:test_api_key") == expected_key + assert cached("spend:committed:org:test_org_id") == expected_org + assert cached("spend:committed:team:test_team_id") is None + assert cached("spend:committed:user:test_user_id") is None + + @pytest.mark.asyncio async def test_update_database_and_spend_counters_releases_reservation_when_db_update_fails(): proxy_logging_obj = MagicMock() @@ -1257,6 +1347,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda request_started_at=start_time, model_access_groups=("premium",), project_id=None, + committed_spend_scope=EntitySpendCounterKeys(), ) @@ -2776,7 +2867,9 @@ _LIVE_COUNTERS: Final = { } -async def _fail_request(logger: ProxyDBLogger, failure: Mapping[str, object]) -> None: +async def _fail_request( + logger: ProxyDBLogger, failure: Mapping[str, object], committed: CommittedBudgetSpend | None = None +) -> None: await logger.async_post_call_failure_hook( request_data={ "model": "gpt-4.1", @@ -2793,6 +2886,7 @@ async def _fail_request(logger: ProxyDBLogger, failure: Mapping[str, object]) -> end_user_id="end-user-1", org_id="org-1", project_id="project-1", + committed_budget_spend=committed, ), ) @@ -2831,6 +2925,43 @@ async def test_failure_spend_rows_raise_the_live_spend_counters_by_their_cost( ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("failure", "charged"), + [ + ({"combined_usage_object": Usage(prompt_tokens=30), "response_cost": 0.125}, 0.125), + (_failure_with_billed_attempts(), 1.25), + ({"standard_logging_object": {"cost_finalized": True, "response_cost": 0.125}}, 0.125), + ], + ids=["recovered_stream_cost", "billed_retries", "completed_cost"], +) +async def test_a_failed_charge_moves_only_the_committed_counters_auth_named_as_budgeted( + live_spend_counters: SpendCounterStore, failure: Mapping[str, object], charged: float +): + """The remaining budget gauges read the committed counters, so a charge missing from them shows as budget the + request never left behind, until the next reseed from the database.""" + budgeted: Final = EntitySpendCounterKeys.of( + token="hashed-key", team_id="team-1", user_id=None, org_id=None, end_user_id=None + ) + starting: Final = { + "spend:committed:key:hashed-key": 0.2, + "spend:committed:team:team-1": 0.3, + "spend:committed:user:user-1": 0.4, + } + for key, spend in starting.items(): + live_spend_counters.set_local(key, spend) + + await _fail_request(ProxyDBLogger(), failure, CommittedBudgetSpend(counter_keys=budgeted)) + + assert {key: live_spend_counters.get_local(key) for key in starting} == pytest.approx( + { + "spend:committed:key:hashed-key": 0.2 + charged, + "spend:committed:team:team-1": 0.3 + charged, + "spend:committed:user:user-1": 0.4, + } + ) + + @pytest.mark.asyncio async def test_a_completed_cost_failure_raises_the_live_spend_counters_once(live_spend_counters: SpendCounterStore): """A failure whose cost was already finalized settles through the success path's write, which raises the diff --git a/tests/unit/proxy/management_endpoints/test_customer_endpoints.py b/tests/unit/proxy/management_endpoints/test_customer_endpoints.py index e9f10e8c0db..c42690050da 100644 --- a/tests/unit/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_customer_endpoints.py @@ -249,6 +249,29 @@ def test_delete_customer_not_found(mock_prisma_client, mock_user_api_key_auth): assert response_json["error"]["code"] == "404" +def test_delete_customer_drops_its_committed_spend_so_a_recreated_customer_starts_from_its_own_row( + monkeypatch, mock_prisma_client, mock_user_api_key_auth +): + from types import SimpleNamespace + + from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + spend: Final = FakeSpendCounters( + redis=None, + memory={"spend:committed:end_user:cust-gone": 7.0, "spend:committed:end_user:cust-stays": 3.0}, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store()) + mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock(return_value=[SimpleNamespace(user_id="cust-gone")]) + mock_prisma_client.db.litellm_endusertable.delete_many = AsyncMock(return_value=1) + + response: Final = client.post( + "/customer/delete", json={"user_ids": ["cust-gone"]}, headers={"Authorization": "Bearer test-key"} + ) + + assert response.status_code == 200 + assert spend.memory == {"spend:committed:end_user:cust-stays": 3.0} + + def test_error_schema_consistency(mock_prisma_client, mock_user_api_key_auth): """ Test that all customer endpoints return the same error schema format. diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py index 120cf8a2b84..b7be80da31a 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -4222,7 +4222,10 @@ async def test_admin_user_update_spend_invalidates_counter(mocker): admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) await _update_single_user_helper(user_request=user_request, user_api_key_dict=admin_caller) - mock_invalidate.assert_awaited_once_with(counter_key="spend:user:target-user") + assert [call.kwargs["counter_key"] for call in mock_invalidate.await_args_list] == [ + "spend:user:target-user", + "spend:committed:user:target-user", + ] @pytest.mark.asyncio @@ -5223,6 +5226,39 @@ async def test_delete_user_evicts_cached_user_rows(mocker: MockerFixture) -> Non broadcast.assert_awaited_once_with(cache_key=deleted.user_id) +@pytest.mark.asyncio +async def test_delete_user_drops_the_committed_spend_of_the_users_it_deleted(mocker: MockerFixture) -> None: + from litellm.proxy._types import DeleteUserRequest, LiteLLM_UserTable + from litellm.proxy.management_endpoints.internal_user_endpoints import delete_user + from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + spend: Final = FakeSpendCounters( + redis=None, memory={"spend:committed:user:user-gone": 7.0, "spend:committed:user:user-stays": 3.0} + ) + mocker.patch("litellm.proxy.proxy_server.spend_counter_store", spend.store()) # test-quality-ok: substitute the spend counter dependency + deleted: Final = LiteLLM_UserTable(user_id="user-gone", user_email="gone@example.test", teams=[]) + prisma_client: Final = mocker.MagicMock() + prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(return_value=deleted) + prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock(return_value=[]) + prisma_client.db.litellm_jwtkeymapping.find_many = mocker.AsyncMock(return_value=[]) + prisma_client.db.litellm_verificationtoken.find_many = mocker.AsyncMock(return_value=[]) + prisma_client.db.litellm_verificationtoken.delete_many = mocker.AsyncMock(return_value=0) + prisma_client.db.litellm_invitationlink.delete_many = mocker.AsyncMock(return_value=0) + prisma_client.db.litellm_organizationmembership.delete_many = mocker.AsyncMock(return_value=0) + prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock(return_value=0) + prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1) + _serve_locked_user_delete(prisma_client, deleted) + mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency + mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", None) # test-quality-ok: delete_user reads it off proxy_server at call time + + await delete_user( + data=DeleteUserRequest(user_ids=[deleted.user_id]), + user_api_key_dict=UserAPIKeyAuth(user_id="proxy-admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert spend.memory == {"spend:committed:user:user-stays": 3.0} + + _DB_OUTAGE_503_BODY: Final = { "error": { "message": "Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly.", diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index cde93915e0c..650bc89c917 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -8693,6 +8693,7 @@ async def test_reset_key_spend_success(monkeypatch): mock_spend_counter_cache.in_memory_cache.set_cache.assert_any_call( key=f"spend_db_floor:spend:key:{hashed_key}", value=50.0, ttl=5 ) + assert spend.deleted() == {f"spend:committed:key:{hashed_key}"} @pytest.mark.asyncio @@ -9013,7 +9014,10 @@ async def test_update_key_spend_updates_counter(monkeypatch): mock_delete_cache.assert_awaited_once() assert spend.memory == {f"spend:key:{hashed_key}": 0.0} - assert spend.commands == [f"SET spend:key:{hashed_key} 0.0 60"] + assert spend.commands == [ + f"SET spend:key:{hashed_key} 0.0 60", + f"DEL spend:committed:key:{hashed_key}", + ] @pytest.mark.asyncio diff --git a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py index 8d2c6c0054c..70cfaf4140d 100644 --- a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py @@ -1188,6 +1188,27 @@ async def _run_legacy_update_organization( return mock_prisma_client +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("body", "twin_after"), + [ + ({"organization_id": "org-1", "spend": 0.0}, None), + ({"organization_id": "org-1", "organization_alias": "renamed"}, 7.0), + ], +) +async def test_legacy_update_drops_the_committed_org_spend_only_when_it_writes_the_stored_spend( + monkeypatch, body, twin_after +): + from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + spend: Final = FakeSpendCounters(redis=None, memory={"spend:committed:org:org-1": 7.0}) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store()) + + await _run_legacy_update_organization(monkeypatch, body=body, existing_budget_id="budget-1") + + assert spend.memory.get("spend:committed:org:org-1") == twin_after + + @pytest.mark.asyncio @pytest.mark.parametrize( "body", @@ -1582,6 +1603,49 @@ async def test_delete_organization_evicts_the_cache_of_the_keys_it_deletes(monke assert jwt_table.rows == (kept_row,) +@pytest.mark.asyncio +async def test_delete_organization_drops_the_committed_spend_of_the_org_and_its_teams(monkeypatch): + from litellm.proxy._types import DeleteOrganizationRequest, LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints.organization_endpoints import delete_organization + from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + spend: Final = FakeSpendCounters( + redis=None, + memory={ + "spend:committed:org:org-doomed": 7.0, + "spend:committed:team:team-in-org": 5.0, + "spend:committed:org:org-kept": 3.0, + "spend:committed:team:team-elsewhere": 2.0, + }, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store()) + prisma_client: Final = AsyncMock() + prisma_client.db.litellm_teamtable.find_many = AsyncMock( + return_value=[SimpleNamespace(team_id="team-in-org", organization_id="org-doomed")] + ) + prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + prisma_client.db.litellm_jwtkeymapping.find_many = AsyncMock(return_value=[]) + prisma_client.db.litellm_organizationtable.delete = AsyncMock(return_value=MagicMock()) + + @asynccontextmanager + async def transaction(): + yield prisma_client.db + + prisma_client.tx = transaction + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True, raising=False) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", None) + + await delete_organization( + data=DeleteOrganizationRequest(organization_ids=["org-doomed"]), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert spend.memory == {"spend:committed:org:org-kept": 3.0, "spend:committed:team:team-elsewhere": 2.0} + + @pytest.mark.asyncio async def test_delete_organization_is_refused_while_one_of_its_teams_holds_a_ptu_share(monkeypatch): """Deleting an organization deletes its teams, so a team still holding a PTU share stops it the diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 94a8022fe96..0881c902a9d 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -9966,6 +9966,37 @@ async def test_delete_validated_team_rows_deletes_only_the_keys_its_rostered_mem assert deleted_keys.kwargs["where"] == {"token": {"in": ["hashed-member"]}} +@pytest.mark.asyncio +async def test_delete_validated_team_rows_drops_the_committed_spend_of_the_teams_it_deleted( + monkeypatch, + disable_audit_logging_for_mocked_team, +): + from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + spend: Final = FakeSpendCounters( + redis=None, memory={"spend:committed:team:team-doomed": 7.0, "spend:committed:team:team-kept": 3.0} + ) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store()) + team = LiteLLM_TeamTable(team_id="team-doomed", members_with_roles=[], metadata={}, model_max_budget={}, model_spend={}) + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_teammembership.delete_many = AsyncMock() + mock_access_group_tx_cm = MagicMock() + mock_access_group_tx_cm.__aenter__ = AsyncMock(return_value=AsyncMock()) + mock_access_group_tx_cm.__aexit__ = AsyncMock(return_value=False) + mock_prisma_client.db.tx = MagicMock(return_value=mock_access_group_tx_cm) + _wire_team_delete_tx(mock_prisma_client, locked_teams=(team,), team_keys=()) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + + await delete_validated_team_rows( + teams=(team,), + prisma_client=mock_prisma_client, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="hashed-admin"), + litellm_changed_by=None, + ) + + assert spend.memory == {"spend:committed:team:team-kept": 3.0} + + @pytest.mark.asyncio async def test_delete_validated_teams_writes_the_deleted_audit_row_only_for_a_team_still_there_under_the_lock( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/proxy/management_helpers/test_bulk_user_deletion.py b/tests/unit/proxy/management_helpers/test_bulk_user_deletion.py index e212a433393..880e5746150 100644 --- a/tests/unit/proxy/management_helpers/test_bulk_user_deletion.py +++ b/tests/unit/proxy/management_helpers/test_bulk_user_deletion.py @@ -2,11 +2,12 @@ import asyncio import copy import itertools import json -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from contextlib import asynccontextmanager from typing import Final, cast import pytest +from pytest_mock import MockerFixture from pydantic import BaseModel, ConfigDict, ValidationError import litellm @@ -22,6 +23,7 @@ from litellm.proxy.management_helpers.bulk_user_deletion import ( ) from litellm.types.proxy.management_endpoints.internal_user_endpoints import BulkDeleteUserRequest from litellm.types.proxy.management_endpoints.team_endpoints import BulkTeamMemberDeleteRequest, TeamMemberRef +from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters ADMIN: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin") INTERNAL: Final = UserAPIKeyAuth(user_id="someone", user_role=LitellmUserRoles.INTERNAL_USER) @@ -697,6 +699,34 @@ async def test_bulk_delete_evicts_jwt_key_mappings_of_the_deleted_users_keys(): assert cache.get_cache(key=kept) is not None +async def _delete_through_scim_or_user_delete(prisma: _FakePrisma, user_ids: Sequence[str]) -> object: + return await delete_users_from_current_teams( + prisma_client=prisma, # pyright: ignore[reportArgumentType] # fake stands in for PrismaClient + user_ids=user_ids, + user_api_key_dict=ADMIN, + user_api_key_cache=UserApiKeyCache(), + proxy_logging_obj=None, + litellm_changed_by=None, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("delete", [_delete, _delete_through_scim_or_user_delete]) +async def test_user_deletes_drop_the_committed_spend_of_the_users_they_deleted( + delete: Callable[[_FakePrisma, Sequence[str]], Awaitable[object]], mocker: MockerFixture +) -> None: + spend: Final = FakeSpendCounters( + redis=None, memory={"spend:committed:user:u1": 7.0, "spend:committed:user:keep": 3.0} + ) + mocker.patch("litellm.proxy.proxy_server.spend_counter_store", spend.store()) # test-quality-ok: substitute the spend counter dependency + prisma: Final = _FakePrisma(users=[_user("u1"), _user("keep")]) + + _ = await delete(prisma, ["u1"]) + + assert sorted(prisma.db.litellm_usertable.rows) == ["keep"] + assert spend.memory == {"spend:committed:user:keep": 3.0} + + @pytest.mark.asyncio async def test_bulk_delete_rejects_non_admin_callers_before_touching_the_db(): prisma = _FakePrisma(users=[_user("u1")]) diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 0e3f1f33e19..464c39d0ba4 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -44,6 +44,8 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitErro from litellm.proxy.hooks.parallel_request_limiter import PROXY_MaxParallelRequestsHandler from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup +from litellm.types.proxy.auth.user_api_key_auth import COMMITTED_BUDGET_SPEND_PARAM +from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import router as llm_passthrough_router from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import ( VertexPassthroughLoggingHandler, @@ -9693,6 +9695,31 @@ async def _track_cost_for_passthrough_kwargs(kwargs: dict) -> AsyncMock: return increment_spend_counters +@pytest.mark.asyncio +async def test_passthrough_charge_moves_the_committed_spend_of_the_entities_auth_named_budgeted(): + """A forged body value stays out, and loggers that dump the row never get the auth object or the spend it read.""" + counter_keys: Final = EntitySpendCounterKeys.of( + token="hashed-token", team_id="t1", user_id="u1", org_id=None, end_user_id=None + ) + user_api_key_dict: Final = UserAPIKeyAuth( + api_key="hashed-token", + user_id="u1", + committed_budget_spend=CommittedBudgetSpend(counter_keys=counter_keys, key=0.123456), + metadata={"logging": [{"callback_name": "langfuse"}], "kept": "yes"}, + ) + + kwargs: Final = _passthrough_kwargs_for_reservation( + user_api_key_dict, + parsed_body={"litellm_metadata": {"user_api_key_auth": {"committed_budget_spend": None}}}, + ) + increment_spend_counters: Final = await _track_cost_for_passthrough_kwargs(kwargs) + + assert increment_spend_counters.await_args.kwargs["committed_spend_scope"] == counter_keys + assert "user_api_key_auth" not in kwargs["litellm_params"]["metadata"] + assert "committed_budget_spend" not in repr(kwargs["litellm_params"]["metadata"]) + assert "0.123456" not in json.dumps(kwargs, default=str) + + @pytest.mark.asyncio async def test_passthrough_success_reconciles_budget_reservation(): """ @@ -11135,6 +11162,9 @@ def _split_pass_through_body( GEMINI_BODY: Final = '{"contents": [{"parts": [{"text": "hi"}]}], "generationConfig": {"temperature": 0}}' +PROXY_SET_PARAMS: Final = ("metadata", "proxy_server_request", COMMITTED_BUDGET_SPEND_PARAM) + + def _metadata_of(split: _PassThroughSplit) -> Mapping[str, object]: return MappingProxyType(_LITELLM_PARAMS.validate_python(split.litellm_params["metadata"])) @@ -11146,7 +11176,7 @@ def test_passthrough_moves_every_litellm_owned_key_from_the_forwarded_body_into_ ) assert frozenset(split.litellm_params) == frozenset( - ("ttl", "num_retries", "litellm_trace_id", "metadata", "proxy_server_request") + ("ttl", "num_retries", "litellm_trace_id", *PROXY_SET_PARAMS) ) assert tuple(split.litellm_params[k] for k in ("ttl", "num_retries", "litellm_trace_id")) == (30, 2, "trace-a") assert split.forwarded_body == json.loads(GEMINI_BODY) @@ -11180,7 +11210,7 @@ PROXY_STAMPED_NAMES: Final = frozenset( def test_passthrough_keeps_each_registered_litellm_owned_name_out_of_the_forwarded_body(name: str) -> None: split: Final = _split_pass_through_body(json.dumps({name: "owned", **json.loads(GEMINI_BODY)})) - assert frozenset(split.litellm_params) == frozenset((name, "metadata", "proxy_server_request")) + assert frozenset(split.litellm_params) == frozenset((name, *PROXY_SET_PARAMS)) assert split.litellm_params[name] == "owned" assert split.forwarded_body == json.loads(GEMINI_BODY) @@ -11224,7 +11254,7 @@ def test_passthrough_keeps_a_client_model_info_object_for_keys_that_allow_client def test_passthrough_drops_a_client_supplied_proxy_stamped_name(name: str) -> None: split: Final = _split_pass_through_body(json.dumps({name: {"forged": "by-client"}, **json.loads(GEMINI_BODY)})) - assert frozenset(split.litellm_params) == frozenset(("metadata", "proxy_server_request")) + assert frozenset(split.litellm_params) == frozenset(PROXY_SET_PARAMS) assert split.litellm_params["proxy_server_request"] != {"forged": "by-client"}, split.litellm_params assert split.forwarded_body == json.loads(GEMINI_BODY) @@ -11235,7 +11265,7 @@ def test_passthrough_merges_both_metadata_carriers_from_the_body_into_one_metada ' "litellm_metadata": {"lm": "b"}, "generationConfig": {"temperature": 0}}' ) - assert frozenset(split.litellm_params) == frozenset(("ttl", "metadata", "proxy_server_request")) + assert frozenset(split.litellm_params) == frozenset(("ttl", *PROXY_SET_PARAMS)) assert _metadata_of(split) == {**_metadata_of(_split_pass_through_body(GEMINI_BODY)), "client_tag": "a", "lm": "b"} assert split.forwarded_body == json.loads(GEMINI_BODY) @@ -11274,7 +11304,7 @@ def test_passthrough_sees_a_name_appended_to_the_public_list_after_import() -> N finally: litellm.all_litellm_params.remove("registered_later") - assert frozenset(split.litellm_params) == frozenset(("registered_later", "metadata", "proxy_server_request")) + assert frozenset(split.litellm_params) == frozenset(("registered_later", *PROXY_SET_PARAMS)) assert split.forwarded_body == {"contents": [{"parts": [{"text": "hi"}]}]} @@ -11283,7 +11313,7 @@ def test_passthrough_sees_the_public_list_rebound_after_import(monkeypatch: pyte split: Final = _split_pass_through_body(LATE_REGISTERED_BODY) - assert frozenset(split.litellm_params) == frozenset(("registered_later", "metadata", "proxy_server_request")) + assert frozenset(split.litellm_params) == frozenset(("registered_later", *PROXY_SET_PARAMS)) assert split.forwarded_body == {"contents": [{"parts": [{"text": "hi"}]}]} diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index 19db80b8930..d04021758ef 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -38,10 +38,13 @@ from litellm.models.end_user import LiteLLM_EndUserTable from litellm.models.tag import LiteLLM_TagTable from litellm.models.team import LiteLLM_TeamTable, LiteLLM_TeamTableCachedObj from litellm.models.user import LiteLLM_UserTable +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_local from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, end_user_cache_key, tag_cache_key from litellm.proxy.management_endpoints.team_endpoints import refresh_cached_team +from litellm.proxy.spend_tracking.committed_spend import committed_spend_key +from litellm.types.proxy.carried_budget_state import EntitySpendCounterKeys from .conftest import normalize from litellm.rust_bridge._native import NativeRedisStoreError @@ -679,7 +682,7 @@ async def test_increment_spend_counters_pipelines_all_scopes_in_one_redis_call( "spend:org:org-1", } spend = _spend_counters(monkeypatch, redis=dict.fromkeys(expected_keys, 1.0)) - monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache()) + monkeypatch.setattr(ps, "user_api_key_cache", UserApiKeyCache()) monkeypatch.setattr(ps, "prisma_client", None) await ps.increment_spend_counters( @@ -697,6 +700,97 @@ async def test_increment_spend_counters_pipelines_all_scopes_in_one_redis_call( assert spend.memory == dict.fromkeys(expected_keys, 1.5) +class _VerificationTokenTable: + def __init__(self, spend: float) -> None: + self._spend: Final = spend + self.lookups: list[dict[str, str]] = [] + + async def find_unique(self, where: dict[str, str]) -> MagicMock: + self.lookups.append(where) + return MagicMock(token=where["token"], spend=self._spend) + + +def _prisma_with_key_spend(spend: float) -> MagicMock: + prisma: Final = MagicMock() + prisma.db.litellm_verificationtoken = _VerificationTokenTable(spend) + return prisma + + +@pytest.mark.asyncio +async def test_committed_twin_takes_the_charge_but_never_a_budget_hold(monkeypatch): + """The key counter holds the DB spend 10.0, this request's 2.0 reservation and another request's 1.0 hold.""" + spend: Final = _spend_counters(monkeypatch, redis=None, memory={"spend:key:hashed": 13.0}) + monkeypatch.setattr(ps, "user_api_key_cache", UserApiKeyCache()) + monkeypatch.setattr(ps, "prisma_client", _prisma_with_key_spend(10.0)) + reservation: Final = {"reserved_cost": 2.0, "entries": [{"counter_key": "spend:key:hashed"}], "finalized": False} + scope: Final = EntitySpendCounterKeys(key="spend:key:hashed") + + await ps.increment_spend_counters( + token="hashed", + team_id=None, + user_id=None, + response_cost=0.5, + budget_reservation=reservation, + committed_spend_scope=scope, + ) + await ps.increment_spend_counters( + token="hashed", team_id=None, user_id=None, response_cost=0.25, committed_spend_scope=scope + ) + + assert spend.memory == pytest.approx({"spend:key:hashed": 11.75, committed_spend_key("spend:key:hashed"): 10.75}) + + +@pytest.mark.asyncio +async def test_a_charge_outside_the_committed_scope_keeps_no_twin_and_reads_no_stored_spend(monkeypatch): + spend: Final = _spend_counters(monkeypatch, redis=None, memory={"spend:key:hashed": 10.0}) + prisma: Final = _prisma_with_key_spend(10.0) + key_cache: Final = UserApiKeyCache() + key_cache.set_cache(key="hashed", value=UserAPIKeyAuth(token="hashed")) + monkeypatch.setattr(ps, "user_api_key_cache", key_cache) + monkeypatch.setattr(ps, "prisma_client", prisma) + + await ps.increment_spend_counters(token="hashed", team_id=None, user_id=None, response_cost=0.5) + + assert spend.memory == pytest.approx({"spend:key:hashed": 10.5}) + assert prisma.db.litellm_verificationtoken.lookups == [] + + +@pytest.mark.asyncio +async def test_only_the_entities_in_the_committed_scope_take_the_charge_on_their_twin(monkeypatch): + spend: Final = _spend_counters(monkeypatch, redis=None, memory={"spend:key:hashed": 10.0, "spend:team:t1": 10.0}) + monkeypatch.setattr(ps, "user_api_key_cache", UserApiKeyCache()) + monkeypatch.setattr(ps, "prisma_client", _prisma_with_key_spend(10.0)) + + await ps.increment_spend_counters( + token="hashed", + team_id="t1", + user_id=None, + response_cost=0.5, + committed_spend_scope=EntitySpendCounterKeys(key="spend:key:hashed"), + ) + + assert spend.memory.get(committed_spend_key("spend:key:hashed")) == pytest.approx(10.5) + assert committed_spend_key("spend:team:t1") not in spend.memory + + +@pytest.mark.asyncio +async def test_a_budgeted_team_whose_id_names_a_window_keeps_its_charge_and_gets_no_twin(monkeypatch): + team_counter_key: Final = "spend:team:t1:window:1d" + spend: Final = _spend_counters(monkeypatch, redis=None, memory={team_counter_key: 1.0}) + monkeypatch.setattr(ps, "user_api_key_cache", UserApiKeyCache()) + monkeypatch.setattr(ps, "prisma_client", MagicMock()) + + await ps.increment_spend_counters( + token=None, + team_id="t1:window:1d", + user_id=None, + response_cost=0.5, + committed_spend_scope=EntitySpendCounterKeys(team=team_counter_key), + ) + + assert spend.memory == pytest.approx({team_counter_key: 1.5}) + + _PENDING_SPEND_COUNTER_KEYS: tuple[str, ...] = ( "spend:key:hashed", "spend:team:team-1", @@ -717,7 +811,7 @@ def _warm_counters_behind_a_failing_pipeline(monkeypatch, failure: str) -> FakeS fail_pipelines=True, failure=failure, ) - monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache()) + monkeypatch.setattr(ps, "user_api_key_cache", UserApiKeyCache()) monkeypatch.setattr(ps, "prisma_client", None) return spend diff --git a/tests/unit/proxy/spend_tracking/test_carried_budget_state.py b/tests/unit/proxy/spend_tracking/test_carried_budget_state.py index 0bdf43b396c..75b9e422330 100644 --- a/tests/unit/proxy/spend_tracking/test_carried_budget_state.py +++ b/tests/unit/proxy/spend_tracking/test_carried_budget_state.py @@ -2,6 +2,9 @@ from datetime import datetime, timezone +import pytest + +from litellm.models.end_user import LiteLLM_EndUserTable from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_OrganizationTable, @@ -11,10 +14,13 @@ from litellm.proxy._types import ( ) from litellm.proxy.spend_tracking.carried_budget_state import ( carried_budget_metadata, + carry_committed_budget_spend, carry_organization_budget_state, carry_team_and_user_budget_state, ) from litellm.types.proxy.carried_budget_state import ( + CommittedBudgetSpend, + EntitySpendCounterKeys, KeyBudgetSnapshot, OrgBudgetSnapshot, TeamBudgetSnapshot, @@ -123,10 +129,105 @@ def test_snapshots_never_reach_the_serialized_token(): user_object=LiteLLM_UserTable(user_id="u1", user_alias="Alice"), ) token.org_budget_snapshot = OrgBudgetSnapshot(spend=1.0, max_budget=2.0) + metadata_without_committed_spend = dict(carried_budget_metadata(token)) + token.committed_budget_spend = CommittedBudgetSpend(key=1.0, team=2.0, user=3.0, org=4.0, end_user=5.0) dumped = token.model_dump() assert "team_budget_snapshot" not in dumped assert "user_budget_snapshot" not in dumped assert "org_budget_snapshot" not in dumped + assert "committed_budget_spend" not in dumped + assert dict(carried_budget_metadata(token)) == metadata_without_committed_spend assert UserAPIKeyAuth(**dumped).team_budget_snapshot is None + + +def _committed_reader(committed: CommittedBudgetSpend): + reads: list[EntitySpendCounterKeys] = [] + + async def read_committed(counter_keys: EntitySpendCounterKeys) -> CommittedBudgetSpend: + reads.append(counter_keys) + return committed + + return read_committed, reads + + +@pytest.mark.asyncio +async def test_committed_spend_is_read_once_for_every_budgeted_entity(): + token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1", org_id="o1", spend=1.0, max_budget=10.0) + token.org_budget_snapshot = OrgBudgetSnapshot(spend=40.0, max_budget=500.0) + committed = CommittedBudgetSpend(key=1.3, team=20.3, user=5.3, org=40.3, end_user=0.8) + read_committed, reads = _committed_reader(committed) + + await carry_committed_budget_spend( + valid_token=token, + team_object=LiteLLM_TeamTable(team_id="t1", max_budget=300.0, spend=20.0), + user_object=LiteLLM_UserTable(user_id="u1", max_budget=50.0, spend=5.0), + end_user_object=LiteLLM_EndUserTable( + user_id="cust-1", blocked=False, spend=0.5, litellm_budget_table=LiteLLM_BudgetTable(max_budget=3.0) + ), + read_committed=read_committed, + ) + + assert reads == [ + EntitySpendCounterKeys( + key="spend:key:hashed", + team="spend:team:t1", + user="spend:user:u1", + org="spend:org:o1", + end_user="spend:end_user:cust-1", + ) + ] + assert token.committed_budget_spend == committed + assert (token.spend, token.team_spend) == (1.0, None) + + +@pytest.mark.asyncio +async def test_route_without_a_budget_check_carries_the_counters_it_charges_and_reads_nothing(): + token = UserAPIKeyAuth(token="hashed", max_budget=10.0) + + await carry_committed_budget_spend( + valid_token=token, team_object=None, user_object=None, end_user_object=None, read_committed=None + ) + + assert token.committed_budget_spend == CommittedBudgetSpend( + counter_keys=EntitySpendCounterKeys(key="spend:key:hashed") + ) + + +@pytest.mark.asyncio +async def test_committed_spend_skips_entities_without_a_max_budget(): + token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1", org_id="o1", spend=1.0) + token.org_budget_snapshot = OrgBudgetSnapshot(spend=40.0, max_budget=None) + read_committed, reads = _committed_reader(CommittedBudgetSpend()) + + await carry_committed_budget_spend( + valid_token=token, + team_object=LiteLLM_TeamTable(team_id="t1", spend=20.0), + user_object=LiteLLM_UserTable(user_id="u1", spend=5.0), + end_user_object=LiteLLM_EndUserTable( + user_id="cust-1", blocked=False, spend=0.5, litellm_budget_table=LiteLLM_BudgetTable(max_budget=None) + ), + read_committed=read_committed, + ) + + assert reads == [EntitySpendCounterKeys()] + assert token.committed_budget_spend == CommittedBudgetSpend() + + +@pytest.mark.asyncio +async def test_committed_spend_reads_the_team_org_when_the_key_has_none(): + token = UserAPIKeyAuth(token="hashed", team_id="t1") + token.org_budget_snapshot = OrgBudgetSnapshot(spend=40.0, max_budget=500.0) + read_committed, reads = _committed_reader(CommittedBudgetSpend(org=41.0)) + + await carry_committed_budget_spend( + valid_token=token, + team_object=LiteLLM_TeamTable(team_id="t1", organization_id="team-org"), + user_object=None, + end_user_object=None, + read_committed=read_committed, + ) + + assert reads == [EntitySpendCounterKeys(org="spend:org:team-org")] + assert token.committed_budget_spend == CommittedBudgetSpend(org=41.0) diff --git a/tests/unit/proxy/spend_tracking/test_committed_spend.py b/tests/unit/proxy/spend_tracking/test_committed_spend.py new file mode 100644 index 00000000000..1ad5a2e7ed4 --- /dev/null +++ b/tests/unit/proxy/spend_tracking/test_committed_spend.py @@ -0,0 +1,274 @@ +"""Committed spend twins hold only charged cost, so the remaining budget gauges never count a budget hold.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Callable, Mapping +from types import MappingProxyType, SimpleNamespace +from typing import Final +from unittest.mock import MagicMock +from uuid import uuid4 + +import pytest + +import litellm +from litellm.integrations.prometheus import PrometheusLogger +from litellm.proxy.spend_tracking.committed_spend import ( + bind_committed_admission_counters, + drop_committed_spend, + read_admission_committed_spend, + read_committed_budget_spend, + seed_committed_spend, +) +from litellm.proxy.spend_tracking.spend_counter_batch import ( + active_spend_counter_batch, + read_batched_spend_counter, + spend_counter_batch_scope, +) +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys +from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + +def _no_side_effect() -> None: + return None + + +class _FindUniqueTable: + def __init__(self, row: SimpleNamespace | None, on_read: Callable[[], None]) -> None: + self._row: Final = row + self._on_read: Final = on_read + self.reads = 0 + + async def find_unique(self, where: dict[str, str]) -> SimpleNamespace | None: + self.reads += 1 + await asyncio.sleep(0) + self._on_read() + return self._row + + +def _prisma_with_spend(spend: float | None, on_read: Callable[[], None] = _no_side_effect) -> SimpleNamespace: + row: Final = None if spend is None else SimpleNamespace(spend=spend) + return SimpleNamespace( + db=SimpleNamespace( + litellm_verificationtoken=_FindUniqueTable(row, on_read), + litellm_teamtable=_FindUniqueTable(row, on_read), + litellm_endusertable=_FindUniqueTable(row, on_read), + ) + ) + + +@pytest.fixture(params=[False, True], ids=["in_memory", "redis"]) +def has_redis(request: pytest.FixtureRequest) -> bool: + return request.param + + +def _spend_counters(has_redis: bool, counters: Mapping[str, float] = MappingProxyType({})) -> FakeSpendCounters: + if has_redis: + return FakeSpendCounters(redis=dict(counters)) + return FakeSpendCounters(redis=None, memory=dict(counters)) + + +@pytest.mark.asyncio +async def test_read_returns_each_committed_twin_and_none_for_one_never_charged(has_redis: bool): + spend: Final = _spend_counters( + has_redis, + {"spend:committed:key:hashed": 1.5, "spend:committed:end_user:cust-1": 0.25, "spend:team:t1": 99.0}, + ) + counter_keys: Final = EntitySpendCounterKeys.of( + token="hashed", team_id="t1", user_id=None, org_id=None, end_user_id="cust-1" + ) + + committed: Final = await read_committed_budget_spend(spend.store(), counter_keys) + + assert committed == CommittedBudgetSpend(counter_keys=counter_keys, key=1.5, end_user=0.25) + + +@pytest.mark.asyncio +async def test_a_twin_holding_garbage_reads_as_unknown_without_hiding_the_other_twins(): + spend: Final = FakeSpendCounters( + redis={"spend:committed:key:hashed": "not-a-number", "spend:committed:team:t1": 4.0} + ) + counter_keys: Final = EntitySpendCounterKeys.of( + token="hashed", team_id="t1", user_id=None, org_id=None, end_user_id=None + ) + + committed: Final = await read_committed_budget_spend(spend.store(), counter_keys) + + assert committed == CommittedBudgetSpend(counter_keys=counter_keys, team=4.0) + + +@pytest.mark.parametrize( + ("counter_keys", "expected_spend"), + [ + (EntitySpendCounterKeys(key="spend:key:hashed"), {"key": 12.0}), + (EntitySpendCounterKeys(end_user="spend:end_user:cust-1"), {"end_user": 12.0}), + ], +) +@pytest.mark.asyncio +async def test_seed_writes_the_db_spend_when_the_twin_is_absent( + has_redis: bool, counter_keys: EntitySpendCounterKeys, expected_spend: dict[str, float] +): + store: Final = _spend_counters(has_redis).store() + (counter_key,) = counter_keys.tracked() + + seeded: Final = await seed_committed_spend(store, _prisma_with_spend(12.0), counter_key) + + assert seeded is True + assert await read_committed_budget_spend(store, counter_keys) == CommittedBudgetSpend( + counter_keys=counter_keys, **expected_spend + ) + + +@pytest.mark.asyncio +async def test_concurrent_charges_on_a_cold_twin_read_the_db_once(has_redis: bool): + store: Final = _spend_counters(has_redis).store() + prisma_client: Final = _prisma_with_spend(12.0) + counter_key: Final = f"spend:key:{uuid4()}" + + seeded: Final = await asyncio.gather(*(seed_committed_spend(store, prisma_client, counter_key) for _ in range(5))) + + assert seeded == [True] * 5 + assert prisma_client.db.litellm_verificationtoken.reads == 1 + + +@pytest.mark.asyncio +async def test_seed_reports_no_twin_when_the_twin_cannot_be_read(): + async def unreachable(counter: str) -> None: + raise ConnectionError("redis unreachable") + + spend: Final = FakeSpendCounters(read_hook=unreachable) + + assert await seed_committed_spend(spend.store(), _prisma_with_spend(12.0), "spend:team:t1") is False + + +@pytest.mark.asyncio +async def test_seed_never_lowers_a_twin_another_replica_already_advanced(has_redis: bool): + store: Final = _spend_counters(has_redis, {"spend:committed:team:t1": 20.0}).store() + + seeded: Final = await seed_committed_spend(store, _prisma_with_spend(12.0), "spend:team:t1") + + committed: Final = await read_committed_budget_spend(store, EntitySpendCounterKeys(team="spend:team:t1")) + assert seeded is True + assert committed.team == 20.0 + + +@pytest.mark.asyncio +async def test_seed_keeps_the_twin_another_worker_wrote_while_this_one_read_the_db(): + spend: Final = FakeSpendCounters(redis={}) + + def another_worker_seeds_and_charges() -> None: + assert spend.redis is not None + spend.redis["spend:committed:team:t1"] = 20.0 + + seeded: Final = await seed_committed_spend( + spend.store(), _prisma_with_spend(12.0, on_read=another_worker_seeds_and_charges), "spend:team:t1" + ) + + committed: Final = await read_committed_budget_spend(spend.store(), EntitySpendCounterKeys(team="spend:team:t1")) + assert seeded is True + assert committed.team == 20.0 + + +@pytest.mark.asyncio +async def test_read_without_any_counter_key_sends_redis_nothing(): + spend: Final = FakeSpendCounters(redis={}) + + committed: Final = await read_committed_budget_spend(spend.store(), EntitySpendCounterKeys()) + + assert committed == CommittedBudgetSpend() + assert spend.commands == [] + + +@pytest.mark.asyncio +async def test_seed_reports_no_twin_when_redis_rejects_the_write(): + spend: Final = FakeSpendCounters(redis={}, fail=True) + + seeded: Final = await seed_committed_spend(spend.store(), _prisma_with_spend(12.0), "spend:team:t1") + + committed: Final = await read_committed_budget_spend(spend.store(), EntitySpendCounterKeys(team="spend:team:t1")) + assert seeded is False + assert committed.team is None + + +@pytest.mark.asyncio +async def test_seed_creates_no_twin_when_the_db_has_no_spend_for_the_entity(has_redis: bool): + store: Final = _spend_counters(has_redis).store() + + seeded: Final = await seed_committed_spend(store, _prisma_with_spend(None), "spend:key:hashed") + + committed: Final = await read_committed_budget_spend(store, EntitySpendCounterKeys(key="spend:key:hashed")) + assert seeded is False + assert committed.key is None + + +@pytest.mark.asyncio +async def test_admission_read_rides_the_admission_mget_with_the_enforcement_counters( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "callbacks", [MagicMock(spec=PrometheusLogger)]) + spend: Final = FakeSpendCounters(redis={"spend:committed:key:hashed": 1.5}) + store: Final = spend.store() + token: Final = UserAPIKeyAuth(token="hashed", team_id="t1") + counter_keys: Final = EntitySpendCounterKeys(key="spend:key:hashed", team="spend:team:t1") + + with spend_counter_batch_scope(store, counter_keys=frozenset(("spend:key:hashed",))): + bind_committed_admission_counters(token, end_user_id=None) + await read_batched_spend_counter("spend:key:hashed") + committed: Final = await read_admission_committed_spend(store, counter_keys) + + assert committed == CommittedBudgetSpend(counter_keys=counter_keys, key=1.5) + assert len(spend.commands) == 1 + assert {"spend:committed:key:hashed", "spend:committed:team:t1"} <= set(spend.commands[0].split()[1:]) + + +@pytest.mark.asyncio +async def test_admission_read_outside_a_batch_leaves_redis_alone_and_the_spend_unknown(): + spend: Final = FakeSpendCounters(redis={"spend:committed:key:hashed": 1.5}) + counter_keys: Final = EntitySpendCounterKeys(key="spend:key:hashed") + + committed: Final = await read_admission_committed_spend(spend.store(), counter_keys) + + assert committed == CommittedBudgetSpend(counter_keys=counter_keys) + assert spend.commands == [] + + +@pytest.mark.asyncio +async def test_admission_read_without_redis_reads_the_in_memory_twin(): + spend: Final = FakeSpendCounters(redis=None, memory={"spend:committed:key:hashed": 1.5}) + + committed: Final = await read_admission_committed_spend( + spend.store(), EntitySpendCounterKeys(key="spend:key:hashed") + ) + + assert committed.key == 1.5 + + +def test_no_committed_counter_is_bound_without_a_prometheus_logger(): + with spend_counter_batch_scope(FakeSpendCounters(redis={}).store()): + bind_committed_admission_counters(UserAPIKeyAuth(token="hashed"), end_user_id=None) + batch: Final = active_spend_counter_batch() + assert batch is not None + bound: Final = batch.counter_keys + + assert bound == frozenset() + + +@pytest.mark.asyncio +async def test_drop_removes_the_twin_and_leaves_the_enforcement_counter(has_redis: bool): + store: Final = _spend_counters(has_redis, {"spend:committed:org:o1": 5.0, "spend:org:o1": 5.0}).store() + counter_keys: Final = EntitySpendCounterKeys(org="spend:org:o1") + + await drop_committed_spend(store, ("spend:org:o1",)) + + assert (await read_committed_budget_spend(store, counter_keys)).org is None + assert (await store.read("spend:org:o1"))[0] == 5.0 + + +@pytest.mark.asyncio +async def test_a_failed_drop_never_fails_the_write_that_already_committed(): + spend: Final = FakeSpendCounters(redis={"spend:committed:org:o1": 5.0}, failing=frozenset({"delete"})) + + await drop_committed_spend(spend.store(), ("spend:org:o1",)) + + assert spend.deleted() == frozenset({"spend:committed:org:o1"}) diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index c3e916d581d..05667629305 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -1,12 +1,14 @@ import asyncio import collections +import contextlib import datetime import hashlib import json import re import sqlite3 -from collections.abc import Mapping, Sequence +from collections.abc import AsyncIterator, Mapping, Sequence from datetime import timezone +from types import SimpleNamespace from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -17,14 +19,17 @@ from fastapi.testclient import TestClient import litellm import litellm.proxy.proxy_server as ps +from litellm.integrations.prometheus import PrometheusLogger from litellm.proxy.auth.authorization import OwnedRows from litellm.proxy.auth.authorization_dependencies import get_log_team_lookup, load_permitted_log_team_ids from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.spend_tracking.spend_management_endpoints import global_spend_reset from tests.unit.proxy.db.fake_prisma_engine import ( ReadTargetOf, statement_timeout_error, statement_timeout_reads, ) +from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters def _default_date_range(): @@ -8336,6 +8341,109 @@ async def test_global_spend_refresh_runs_the_refresh_on_the_writer_when_a_reader ] +class _ResetTableRows: + def __init__(self) -> None: + self.update_many: Final = AsyncMock() + + +class _ResetDb: + def __init__(self) -> None: + self.litellm_verificationtoken: Final = _ResetTableRows() + self.litellm_teamtable: Final = _ResetTableRows() + + async def query_raw(self, sql_query: str, after_id: str, limit: int) -> list[dict[str, str]]: + rows: Final = ["hashed"] if "VerificationToken" in sql_query else ["t1"] + return [{"id": row} for row in rows if row > after_id][:limit] + + +class _ResetTransaction: + def __init__(self, db: _ResetDb) -> None: + self.litellm_verificationtoken: Final = db.litellm_verificationtoken + self.litellm_teamtable: Final = db.litellm_teamtable + + async def query_raw(self, sql_query: str) -> list[dict[str, str]]: + return [] + + +def _reset_prisma(db: _ResetDb) -> SimpleNamespace: + @contextlib.asynccontextmanager + async def tx() -> AsyncIterator[_ResetTransaction]: + yield _ResetTransaction(db) + + return SimpleNamespace(db=db, tx=tx) + + +def _committed_spend(monkeypatch: pytest.MonkeyPatch, twins: Sequence[str]) -> FakeSpendCounters: + spend: Final = FakeSpendCounters(redis=None, memory=dict.fromkeys(twins, 7.0)) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store()) + return spend + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("registered_callbacks", "twin_after"), + [ + ([MagicMock(spec=PrometheusLogger)], None), + ([], 7.0), + ], +) +async def test_global_spend_reset_drops_the_committed_key_and_team_spend_the_gauges_read( + monkeypatch, registered_callbacks, twin_after +): + twins: Final = ("spend:committed:key:hashed", "spend:committed:team:t1") + spend: Final = _committed_spend(monkeypatch, twins) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _reset_prisma(_ResetDb())) + monkeypatch.setattr(litellm, "callbacks", registered_callbacks) + + response: Final = await global_spend_reset() + + assert response["status"] == "success" + assert [spend.memory.get(twin) for twin in twins] == [twin_after, twin_after] + + +class _PagedResetDb(_ResetDb): + def __init__(self, token_count: int) -> None: + super().__init__() + self._tokens: Final = sorted(f"token-{index}" for index in range(token_count)) + self.largest_page = 0 + + async def query_raw(self, sql_query: str, after_id: str, limit: int) -> list[dict[str, str]]: + rows: Final = self._tokens if "VerificationToken" in sql_query else [] + page: Final = [{"id": row} for row in rows if row > after_id][:limit] + self.largest_page = max(self.largest_page, len(page)) + return page + + +@pytest.mark.asyncio +@pytest.mark.parametrize("token_count", [999, 1000, 1001, 2500]) +async def test_global_spend_reset_drops_every_key_twin_reading_the_keys_in_bounded_pages(monkeypatch, token_count): + db: Final = _PagedResetDb(token_count) + twins: Final = tuple(f"spend:committed:key:token-{index}" for index in range(token_count)) + spend: Final = _committed_spend(monkeypatch, twins) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _reset_prisma(db)) + monkeypatch.setattr(litellm, "callbacks", [MagicMock(spec=PrometheusLogger)]) + + response: Final = await global_spend_reset() + + assert response["status"] == "success" + assert spend.memory == {} + assert db.largest_page <= 1000 + + +@pytest.mark.asyncio +async def test_global_spend_reset_still_succeeds_when_listing_the_committed_counters_fails(monkeypatch): + db: Final = _ResetDb() + db.query_raw = AsyncMock(side_effect=ConnectionError("db unreachable")) + _committed_spend(monkeypatch, ("spend:committed:team:t1",)) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _reset_prisma(db)) + monkeypatch.setattr(litellm, "callbacks", [MagicMock(spec=PrometheusLogger)]) + + response: Final = await global_spend_reset() + + assert response["status"] == "success" + db.litellm_teamtable.update_many.assert_awaited_once() + + def test_spend_calculate_rejection_body_is_openai_shaped(client): app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user", api_key="hashed-k" diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 1cfb14e6bf9..c62e73232d0 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -918,6 +918,49 @@ async def test_initialize_scheduled_jobs_loads_credentials_only_through_add_depl scheduler.shutdown(wait=False) +@pytest.mark.asyncio +@pytest.mark.parametrize("proxy_batch_write_at", [10, 60]) +async def test_committed_spend_twins_outlive_the_slowest_spend_write_to_the_db( + monkeypatch: pytest.MonkeyPatch, proxy_batch_write_at: int +) -> None: + """A twin that expires before its charges reach the DB seeds again from a total still missing them.""" + monkeypatch.delenv("DISABLE_PRISMA_SCHEMA_UPDATE", raising=False) + monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False) + from apscheduler.schedulers.asyncio import AsyncIOScheduler + + from litellm.constants import SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS + from litellm.proxy.proxy_server import _BATCH_WRITE_JITTER_MAX_SECONDS, ProxyStartupEvent + from litellm.proxy.utils import ProxyLogging + from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + spend: Final = FakeSpendCounters() + mock_proxy_logging: Final = MagicMock(spec=ProxyLogging) + mock_proxy_logging.slack_alerting_instance = MagicMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() + scheduler: Final = AsyncIOScheduler() + try: + with ( + patch("litellm.proxy.proxy_server.proxy_config", _mock_scheduled_proxy_config()), + patch("litellm.proxy.proxy_server.store_model_in_db", False), + patch("litellm.proxy.proxy_server.AsyncIOScheduler", return_value=scheduler), + patch("litellm.proxy.proxy_server.spend_counter_store", spend.store()), + ): + await ProxyStartupEvent.initialize_scheduled_background_jobs( + general_settings={}, + prisma_client=MagicMock(), + proxy_budget_rescheduler_min_time=1, + proxy_budget_rescheduler_max_time=2, + proxy_batch_write_at=proxy_batch_write_at, + proxy_logging_obj=mock_proxy_logging, + ) + finally: + scheduler.shutdown(wait=False) + + slowest_db_write: Final = SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS + proxy_batch_write_at + _BATCH_WRITE_JITTER_MAX_SECONDS + assert spend.committed_ttl is not None + assert spend.committed_ttl >= slowest_db_write + + @pytest.mark.asyncio async def test_periodic_reload_job_scheduled_without_store_model_in_db(monkeypatch): """ diff --git a/tests/unit/redis/stores/spend_counters_double.py b/tests/unit/redis/stores/spend_counters_double.py index 1f3c4b94d45..c50133f7076 100644 --- a/tests/unit/redis/stores/spend_counters_double.py +++ b/tests/unit/redis/stores/spend_counters_double.py @@ -35,9 +35,22 @@ def _is_number(value: object) -> bool: return value is None or (isinstance(value, (int, float)) and not isinstance(value, bool)) +_TWIN_SCOPES: Final = frozenset({"key", "team", "user", "org", "end_user"}) +_WINDOWED_SCOPES: Final = frozenset({"key", "team"}) + + +def _native_store_takes_twin(twin: str) -> bool: + scope, _, counter_id = twin.partition(":") + _, separator, window = counter_id.rpartition(":window:") + return scope in _TWIN_SCOPES and not (scope in _WINDOWED_SCOPES and separator and window) + + def _check(counter: str) -> None: if not counter.startswith("spend:"): raise ValueError(f"{counter!r} is not a spend counter") + twin: Final = counter.removeprefix("spend:committed:") + if twin != counter and not _native_store_takes_twin(twin): + raise ValueError(f"{counter!r} is not a spend counter") @dataclass @@ -53,6 +66,11 @@ class DeclaredRead: return self.owner.mget(self.counters) +@dataclass +class FakeSettlement: + settled: bool = False + + @dataclass class FakeSpendCounters: """``redis`` is None for a worker without Redis, and ``redis_default`` answers reads of a key it lacks. ``fail`` @@ -68,11 +86,12 @@ class FakeSpendCounters: failure: str = "transport" redis_default: float | None = None redis_ttl: float = 60 + committed_ttl: float | None = None batch_reads: bool = True read_hook: Callable[[str], Awaitable[None]] | None = None on_delete: Callable[[tuple[str, ...]], None] | None = None deletes: list[tuple[str, ...]] = field(default_factory=list) - post_call: list[tuple[object, tuple[SpendIncrement, ...]]] = field(default_factory=list) + post_call: list[tuple[object, tuple[SpendIncrement, ...], FakeSettlement]] = field(default_factory=list) declared: list[DeclaredRead] = field(default_factory=list) def store(self) -> SpendCounterStore: @@ -86,6 +105,9 @@ class FakeSpendCounters: def redis_ttl_seconds(self) -> float | None: return None if self.redis is None else self.redis_ttl + def set_committed_ttl(self, ttl_seconds: float, /) -> None: + self.committed_ttl = ttl_seconds + def attach(self, client: NativeRedisClient, /) -> bool: if self.redis is not None: return False @@ -139,17 +161,19 @@ class FakeSpendCounters: self.commands.append(f"PIPELINE {' '.join(counter for counter, _ in increments)}") return _done([self._incr(counter, amount) for counter, amount in increments]) - def increment_many_post_call(self, request: object, increments: Sequence[SpendIncrement]) -> bool: + def increment_many_post_call(self, request: object, increments: Sequence[SpendIncrement]) -> FakeSettlement | None: if self.redis is None: - return False - self.post_call.append((request, tuple(increments))) - return True + return None + settlement: Final = FakeSettlement() + self.post_call.append((request, tuple(increments), settlement)) + return settlement def settle_post_call(self) -> None: """What the native task does once the post-call batch goes out.""" pending: Final = tuple(self.post_call) self.post_call.clear() - for _, increments in pending: + for _, increments, settlement in pending: + settlement.settled = True if self.fail: for counter, _ in increments: self.memory.pop(counter, None) diff --git a/tests/unit/redis/stores/test_spend_counters.py b/tests/unit/redis/stores/test_spend_counters.py index c54dc6cc81e..1d39cffc933 100644 --- a/tests/unit/redis/stores/test_spend_counters.py +++ b/tests/unit/redis/stores/test_spend_counters.py @@ -10,6 +10,7 @@ from litellm.redis.stores.connection import RedisStores, connect from litellm.redis.stores.spend_counters import ( SPEND_COUNTERS_TARGET, SpendCounterStore, + committed_counter, end_user_counter, is_unavailable, key_counter, @@ -28,7 +29,7 @@ from litellm.redis.store_settings import RedisStoreSettings from litellm.rust_bridge._native import NativeRedisClient, NativeRedisStoreError, NativeSpendCounters, SpendRead from tests.unit.caching.test_redis_batch import FakeNatives, fake_send -from .spend_counters_double import FakeSpendCounters +from .spend_counters_double import FakeSettlement, FakeSpendCounters _KEYS: Final = ( ( @@ -46,6 +47,11 @@ _KEYS: Final = ( (project_counter("proj-1"), "spend:project:proj-1"), (model_access_group_counter("gpt-group"), "spend:model_access_group:gpt-group"), (shadow_eval_counter("job-7"), "spend:shadow_eval:job-7"), + (committed_counter(key_counter("88dc28d0")), "spend:committed:key:88dc28d0"), + (committed_counter(team_counter("team-a")), "spend:committed:team:team-a"), + (committed_counter(user_counter("user@example.com")), "spend:committed:user:user@example.com"), + (committed_counter(org_counter("org-1")), "spend:committed:org:org-1"), + (committed_counter(end_user_counter("customer é 1")), "spend:committed:end_user:customer é 1"), ) @@ -81,7 +87,17 @@ def test_a_read_decodes_what_redis_cache_decodes(wire: str, spend: float) -> Non assert type(decoded) is type(spend) -@pytest.mark.parametrize("key", ["spend_db_floor:spend:key:k", "virtual_key_spend:k:gpt:1d", "spend:unknown:k"]) +@pytest.mark.parametrize( + "key", + [ + "spend_db_floor:spend:key:k", + "virtual_key_spend:k:gpt:1d", + "spend:unknown:k", + "spend:committed:key:k:window:1d", + "spend:committed:tag:prod", + "spend:committed:committed:key:k", + ], +) def test_the_native_store_refuses_keys_outside_the_spend_counter_keyspace(key: str) -> None: with pytest.raises(ValueError, match="a spend counter is a spend:\* key"): NativeSpendCounters(None).get_local(key) @@ -153,7 +169,7 @@ async def test_post_call_increments_go_on_the_requests_native_post_call_batch_un targets: Final[list[str | None]] = [] class Recording(FakeSpendCounters): - def increment_many_post_call(self, request: object, increments: object) -> bool: + def increment_many_post_call(self, request: object, increments: object) -> FakeSettlement | None: targets.append(current_service_target()) return super().increment_many_post_call(request, increments) # pyright: ignore[reportArgumentType] # recorded as given @@ -165,8 +181,8 @@ async def test_post_call_increments_go_on_the_requests_native_post_call_batch_un inside: Final = store.increment_many_post_call((("spend:key:k", 1.0),)) empty: Final = store.increment_many_post_call(()) - assert (outside, inside, empty) == (False, True, False) - assert spend.post_call == [(natives.built[0], (("spend:key:k", 1.0),))] + assert (outside, empty) == (None, None) + assert spend.post_call == [(natives.built[0], (("spend:key:k", 1.0),), inside)] assert targets == [SPEND_COUNTERS_TARGET] From 31a9a99eb9b227f41350d2ce0682d44df3dc8590 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 9 Nov 2026 20:34:38 +0000 Subject: [PATCH 08/46] feat(azure): add azure_ai/claude-fable-5-2 and claude-fable-5-2-fast from the Foundry Claude page (#50723) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 62 +++++++++++++++++++ model_prices_and_context_window.json | 62 +++++++++++++++++++ 2 files changed, 124 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 92b6ae6c469..79e4d92aa9d 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -87727,5 +87727,67 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true, "thinking_always_on": true + }, + "azure_ai/claude-fable-5-2": { + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://platform.claude.com/docs/en/about-claude/pricing", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": false, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "azure_ai/claude-fable-5-2-fast": { + "cache_creation_input_token_cost": 2e-05, + "cache_creation_input_token_cost_above_1hr": 3.2e-05, + "cache_read_input_token_cost": 4e-07, + "input_cost_per_token": 1.6e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 8e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://platform.claude.com/docs/en/about-claude/pricing", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": false, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 92b6ae6c469..79e4d92aa9d 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -87727,5 +87727,67 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true, "thinking_always_on": true + }, + "azure_ai/claude-fable-5-2": { + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://platform.claude.com/docs/en/about-claude/pricing", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": false, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "azure_ai/claude-fable-5-2-fast": { + "cache_creation_input_token_cost": 2e-05, + "cache_creation_input_token_cost_above_1hr": 3.2e-05, + "cache_read_input_token_cost": 4e-07, + "input_cost_per_token": 1.6e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 8e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://platform.claude.com/docs/en/about-claude/pricing", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": false, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true } } From 6162ebab857bd686f9d424450a0c0c900272d0da Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 9 Nov 2026 12:40:40 -0800 Subject: [PATCH 09/46] build(ui): fail lint on hardcoded English in dashboard JSX (#50593) * build(ui): fail lint on hardcoded English in dashboard JSX Adds eslint-plugin-i18next no-literal-string in jsx-only mode with type info for dashboard source files, skipping tests and the two untranslated files, plus a no-restricted-syntax selector that rejects an English literal as the first argument to toast.success, error, warning or info. * test(ui): lint the fetch rule fixture at a real source path Type-aware linting needs the file to be in the TypeScript project, so a virtual path under src now fails to parse and no rule runs. The fixture now uses src/app/layout.tsx with the same code and assertions. --- ui/litellm-dashboard/eslint.config.mjs | 178 +++++++++++++++++- ui/litellm-dashboard/package-lock.json | 24 +++ ui/litellm-dashboard/package.json | 1 + .../tests/fetch-location-rule.test.ts | 2 +- 4 files changed, 196 insertions(+), 9 deletions(-) diff --git a/ui/litellm-dashboard/eslint.config.mjs b/ui/litellm-dashboard/eslint.config.mjs index acf6f624f96..828c0830cd0 100644 --- a/ui/litellm-dashboard/eslint.config.mjs +++ b/ui/litellm-dashboard/eslint.config.mjs @@ -5,8 +5,157 @@ import prettier from "eslint-config-prettier/flat"; import unusedImports from "eslint-plugin-unused-imports"; import testingLibrary from "eslint-plugin-testing-library"; import jestDom from "eslint-plugin-jest-dom"; +import i18next from "eslint-plugin-i18next"; +import htmlEntities from "eslint-plugin-i18next/lib/options/htmlEntities.js"; import local from "./scripts/eslint-rules/index.mjs"; +const rawFetchRestriction = { + selector: "CallExpression[callee.name='fetch']", + message: + "Raw fetch() is only allowed in src/lib/http/. Use the shared client (createApiClient / apiClient) from @/lib/http/client instead.", +}; + +const toastLiteralRestriction = { + selector: + "CallExpression[callee.object.name='toast'][callee.property.name=/^(success|error|warning|info)$/] > Literal.arguments:first-child", + message: "Pass toast a t() string, not an English literal, so the message can be translated.", +}; + +const noLiteralStringOptions = { + mode: "jsx-only", + "should-validate-template": true, + "jsx-components": { exclude: ["Trans", "code", "pre", "kbd", "Kbd", "CodeBlock", "SyntaxHighlighter"] }, + "jsx-attributes": { + exclude: [ + "className", + ".*ClassName", + "style", + ".*Style", + "key", + "id", + ".*Id", + ".*ID", + "htmlFor", + "name", + "value", + "defaultValue", + "data-.*", + "aria-(describedby|labelledby|controls|owns|haspopup|live|current|hidden|expanded|selected)", + "type", + "role", + "href", + "to", + "src", + "target", + "rel", + "accept", + "autoComplete", + "inputMode", + "step", + "min", + "max", + "width", + "height", + "size", + "color", + ".*Color", + "variant", + "language", + "index", + ".*Key", + "keys", + "feature", + "fill", + "stroke", + "scope", + "delay", + "stackId", + "category", + "categories", + "gridCols", + "defaultTheme", + "defaultRole", + "schemaComponent", + "storageKey", + "code", + ], + }, + callees: { + exclude: [ + "i18n(ext)?", + "t", + "i18n\\.t", + "require", + "addEventListener", + "removeEventListener", + "postMessage", + "getElementById", + "dispatch", + "commit", + "includes", + "indexOf", + "endsWith", + "startsWith", + "cn", + "clsx", + "cva", + ".*\\.(get|set|has|delete|getAll|append|getColumn|format|split|replace|closest|querySelector|getItem|setItem|removeItem|toLocaleDateString|toLocaleString|hasOwnProperty)", + "Object\\.hasOwn", + "uiHref", + "hasVisited", + "toggleKey", + "emitLocalStorageChange", + "setLocalStorageItem", + "removeLocalStorageItem", + "getLocalStorageItem", + "window\\.open", + "Error", + "new Error", + ], + }, + words: { + exclude: [ + "[0-9!-/:-@[-`{-~]+", + "[A-Z_-]+", + "v", + htmlEntities, + /^\p{Emoji}+$/u, + /^[\p{P}\p{S}\p{N}\p{M}\s]+$/u, + /^[a-z0-9]+(?:[_\-.:/][a-z0-9*[\]]+)+:?$/, + /^(?:https?:\/\/|\/|~\/|\.\/)\S*$/, + ], + }, + "object-properties": { + exclude: [ + "[A-Z_-]+", + "value", + "key", + "id", + "name", + "dataIndex", + "className", + "width", + "fontSize", + "maxHeight", + "borderRadius", + "bg", + "color", + "background", + "padding", + "fill", + "role", + "class", + "spellcheck", + "toast", + "scope", + "origin", + "classifierType", + "defaultValue", + "1", + ], + }, +}; + const eslintConfig = [ { ignores: [".next/**", "out/**", "build/**", "coverage/**", "next-env.d.ts", "src/lib/http/schema.d.ts"], @@ -41,14 +190,7 @@ const eslintConfig = [ "max-depth": ["warn", 4], "max-params": ["error", 4], "max-nested-callbacks": ["error", 4], - "no-restricted-syntax": [ - "error", - { - selector: "CallExpression[callee.name='fetch']", - message: - "Raw fetch() is only allowed in src/lib/http/. Use the shared client (createApiClient / apiClient) from @/lib/http/client instead.", - }, - ], + "no-restricted-syntax": ["error", rawFetchRestriction], "no-restricted-imports": [ "error", { @@ -81,6 +223,26 @@ const eslintConfig = [ "max-lines": ["error", { max: 800, skipBlankLines: true, skipComments: true }], }, }, + { + files: ["src/**/*.{ts,tsx}"], + ignores: [ + "src/**/*.test.{ts,tsx}", + "src/**/*.spec.{ts,tsx}", + "src/**/*.test-d.ts", + "src/**/__tests__/**", + "src/**/*TestHarness.tsx", + ], + languageOptions: { parserOptions: { projectService: true, tsconfigRootDir: import.meta.dirname } }, + plugins: { i18next }, + rules: { + "i18next/no-literal-string": ["error", noLiteralStringOptions], + "no-restricted-syntax": ["error", rawFetchRestriction, toastLiteralRestriction], + }, + }, + { + files: ["src/i18n/untranslated/**"], + rules: { "i18next/no-literal-string": "off" }, + }, { files: ["src/lib/http/**"], rules: { diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index d506d114d8c..5c530c4ac80 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -77,6 +77,7 @@ "eslint": "9.39.2", "eslint-config-next": "16.3.3", "eslint-config-prettier": "10.1.8", + "eslint-plugin-i18next": "6.1.5", "eslint-plugin-jest-dom": "5.10.1", "eslint-plugin-testing-library": "7.16.2", "eslint-plugin-unused-imports": "4.3.0", @@ -6647,6 +6648,29 @@ "ms": "^2.1.1" } }, + "node_modules/eslint-plugin-i18next": { + "version": "6.1.5", + "resolved": "https://registry.npmjs.org/eslint-plugin-i18next/-/eslint-plugin-i18next-6.1.5.tgz", + "integrity": "sha512-xCTfstbK9ZpQ6UFT5S1s6zbZMhKn2o2jRSQYiU2UkV7wt8y1m3WjkUYGkGlipgRdFInYI9+LAqE+JQU/L73HmA==", + "dev": true, + "license": "ISC", + "dependencies": { + "requireindex": "~1.1.0" + }, + "engines": { + "node": ">=18.10.0" + } + }, + "node_modules/eslint-plugin-i18next/node_modules/requireindex": { + "version": "1.1.0", + "resolved": "https://registry.npmjs.org/requireindex/-/requireindex-1.1.0.tgz", + "integrity": "sha512-LBnkqsDE7BZKvqylbmn7lTIVdpx4K/QCduRATpO5R+wtPmky/a8pN1bO2D6wXppn1497AJF9mNjqAXr6bdl9jg==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=0.10.5" + } + }, "node_modules/eslint-plugin-import": { "version": "2.32.0", "resolved": "https://registry.npmjs.org/eslint-plugin-import/-/eslint-plugin-import-2.32.0.tgz", diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index 9fc0242f040..2174d7cc2ad 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -94,6 +94,7 @@ "eslint": "9.39.2", "eslint-config-next": "16.3.3", "eslint-config-prettier": "10.1.8", + "eslint-plugin-i18next": "6.1.5", "eslint-plugin-jest-dom": "5.10.1", "eslint-plugin-testing-library": "7.16.2", "eslint-plugin-unused-imports": "4.3.0", diff --git a/ui/litellm-dashboard/tests/fetch-location-rule.test.ts b/ui/litellm-dashboard/tests/fetch-location-rule.test.ts index bd495ad3c40..329012724e8 100644 --- a/ui/litellm-dashboard/tests/fetch-location-rule.test.ts +++ b/ui/litellm-dashboard/tests/fetch-location-rule.test.ts @@ -18,7 +18,7 @@ describe("location-based fetch() rule", () => { }); it("flags a raw fetch() in a normal source file", async () => { - const messages = await fetchMessages("src/components/some_feature.tsx"); + const messages = await fetchMessages("src/app/layout.tsx"); expect(messages).toHaveLength(1); expect(messages[0].message).toMatch(/@\/lib\/http\/client/); }); From 1c807aef19f224d15e25e4beae365efdcaa97aab Mon Sep 17 00:00:00 2001 From: zoveki <349028329+zoveki@users.noreply.github.com> Date: Mon, 9 Nov 2026 12:52:52 -0800 Subject: [PATCH 10/46] fix(prompt_templates): keep moved private helpers importable from common_utils (#50721) * fix(prompt_templates): keep moved private helpers importable from common_utils The split moved 83 private helpers out of common_utils, so code importing them by their old path got ImportError. A module __getattr__ now serves each one from its new module. Also renames the new test files to the directory's naming convention * fix(prompt_templates): hide the old-path getattr from type checkers and serve only released names Type checkers stop seeing the module __getattr__, so a typo'd import from common_utils fails lint again. The old path now serves only the 69 moved helpers that common_utils had in a tagged release, since the other 14 were never importable from it * fix(prompt_templates): let patches through the old common_utils path reach moved helpers and warn on old-path use * fix(prompt_templates): sync old-path patches of re-exported helpers and restore moved helpers when warnings are errors * fix(prompt_templates): reset a moved helper on old-path delete instead of removing it, and warn on old-path writes --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../prompt_templates/common_utils.py | 175 +++++++++++++++++ ...ore_utils_prompt_templates_common_utils.py | 183 ++++++++++++++++++ ...e_utils_prompt_templates_file_id_utils.py} | 0 ...ils_prompt_templates_json_schema_utils.py} | 0 ...emplates_prompt_cache_breakpoint_utils.py} | 0 ...utils_prompt_templates_reasoning_utils.py} | 0 ...prompt_templates_tool_call_json_repair.py} | 0 ...utils_prompt_templates_tpm_media_utils.py} | 0 tests/unit/llms/test_file_content_block.py | 4 +- 9 files changed, 361 insertions(+), 1 deletion(-) rename tests/unit/litellm_core_utils/prompt_templates/{test_file_id_utils.py => test_litellm_core_utils_prompt_templates_file_id_utils.py} (100%) rename tests/unit/litellm_core_utils/prompt_templates/{test_json_schema_utils.py => test_litellm_core_utils_prompt_templates_json_schema_utils.py} (100%) rename tests/unit/litellm_core_utils/prompt_templates/{test_prompt_cache_breakpoint_utils.py => test_litellm_core_utils_prompt_templates_prompt_cache_breakpoint_utils.py} (100%) rename tests/unit/litellm_core_utils/prompt_templates/{test_reasoning_utils.py => test_litellm_core_utils_prompt_templates_reasoning_utils.py} (100%) rename tests/unit/litellm_core_utils/prompt_templates/{test_tool_call_json_repair.py => test_litellm_core_utils_prompt_templates_tool_call_json_repair.py} (100%) rename tests/unit/litellm_core_utils/prompt_templates/{test_tpm_media_utils.py => test_litellm_core_utils_prompt_templates_tpm_media_utils.py} (100%) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 0572ca66fc9..2e6f865b575 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -2,10 +2,14 @@ Common utility functions used for translating messages across providers """ +import importlib import json +import sys +import warnings from collections.abc import Hashable, Iterable, Mapping, Sequence from dataclasses import dataclass from itertools import chain, groupby +from types import MappingProxyType, ModuleType from typing import TYPE_CHECKING, Final, Literal, TypeAlias, cast from pydantic import TypeAdapter @@ -222,6 +226,108 @@ _extract_reasoning_content: Final = extract_reasoning_content _get_image_mime_type_from_url: Final = get_image_mime_type_from_url _parse_content_for_reasoning: Final = parse_content_for_reasoning +_MOVED_PRIVATE_NAME_MODULES: Final[Mapping[str, str]] = MappingProxyType( + { + **dict.fromkeys( + ("_decode_vector_store_ids_in_tools",), + "file_id_utils", + ), + **dict.fromkeys( + ( + "_EMPTY_SCHEMA", + "_LOCAL_SCHEMA_REF_PREFIXES", + "_MAX_SCHEMA_FLATTEN_DEPTH", + "_MAX_SCHEMA_NESTING", + "_OPENAI_REJECTED_TOP_LEVEL_SCHEMA_KEYS", + "_REGEX_LOOKAROUND_RE", + "_SCHEMA_LIST", + "_SUBSCHEMA_KEYWORDS", + "_SUBSCHEMA_LIST_KEYWORDS", + "_SUBSCHEMA_MAP_KEYWORDS", + "_TOP_LEVEL_SCHEMA_COMBINATORS", + "_any_of", + "_combinator_required_names", + "_dropped_pattern_properties", + "_flatten_schema_against_root", + "_grammar_branches", + "_is_not_python_regex", + "_is_object_schema", + "_is_open_object", + "_is_union", + "_keyword_value_rebuilt", + "_list_at", + "_may_be_string", + "_mergeable_branch", + "_named_schemas", + "_node_without_rejected_regex", + "_optional_property_names", + "_property_schemas", + "_refs_in", + "_resolve_local_schema_ref", + "_root_definitions", + "_schema_at", + "_schema_branches", + "_schema_levels", + "_schema_nodes", + "_schema_properties", + "_schema_required_names", + "_schema_without_rejected_regex", + "_subschemas", + "_uses_regex_lookaround", + ), + "json_schema_utils", + ), + **dict.fromkeys( + ( + "_BREAKPOINT_BLOCKS", + "_BREAKPOINT_BLOCK_LIST_KEYS", + "_BREAKPOINT_FIELDS", + "_MarkedT", + "_PROMPT_CACHE_BREAKPOINT", + "_breakpoints_stripped_from", + "_without_prompt_cache_breakpoint", + ), + "prompt_cache_breakpoint_utils", + ), + **dict.fromkeys( + ( + "_carries_encrypted_reasoning", + "_converse_wrapped_payload", + "_encrypted_reasoning_field", + "_provider_text_or_none", + "_readable_thinking_text", + "_reasoning_item_from_block_group", + "_reasoning_replay_group_key", + "_strip_encrypted_reasoning_from_blocks", + ), + "reasoning_utils", + ), + **dict.fromkeys( + ( + "_DECODED_JSON", + "_attempt_json_repair", + ), + "tool_call_json_repair", + ), + **dict.fromkeys( + ( + "_TPMMediaProjection", + "_tpm_blank_fields", + "_tpm_code_output_media", + "_tpm_document_media", + "_tpm_file_media", + "_tpm_image_part", + "_tpm_input_media", + "_tpm_media_leaf", + "_tpm_media_list", + "_tpm_replace_field", + "_tpm_tool_result_media", + ), + "tpm_media_utils", + ), + } +) + if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LoggingClass @@ -1578,3 +1684,72 @@ def invalid_type_message(field_path: str, value: object) -> str: def is_unhashable_value_behind(error: TypeError, value: object) -> bool: return not isinstance(value, Hashable) and f"unhashable type: '{type(value).__name__}'" in str(error) + + +_UNBOUND: Final = object() + + +def _module(module_name: str) -> ModuleType: + return importlib.import_module(f"litellm.litellm_core_utils.prompt_templates.{module_name}") + + +def _namespace(module: ModuleType) -> Mapping[str, object]: + return vars(module) + + +def _moved_name_home(name: str) -> ModuleType: + module_name: Final = _MOVED_PRIVATE_NAME_MODULES.get(name) + if module_name is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + return _module(module_name) + + +def _split_modules_sharing(name: str, value: object) -> tuple[ModuleType, ...]: + if value is _UNBOUND: + return () + split_modules: Final = (_module(module_name) for module_name in frozenset(_MOVED_PRIVATE_NAME_MODULES.values())) + return tuple(module for module in split_modules if _namespace(module).get(name, _UNBOUND) is value) + + +def _warn_old_path(name: str, home: ModuleType) -> None: + warnings.warn( + f"{__name__}.{name} moved to {home.__name__}, import or patch it there. " + "The old path will be removed in a future release", + DeprecationWarning, + stacklevel=3, + ) + + +_RELEASED_MOVED_VALUES: Final = MappingProxyType( + {name: _namespace(_moved_name_home(name))[name] for name in _MOVED_PRIVATE_NAME_MODULES} +) + + +class _CommonUtilsModule(ModuleType): + def __getattr__(self, name: str) -> object: + home: Final = _moved_name_home(name) + value: Final = _namespace(home).get(name, _UNBOUND) + if value is _UNBOUND: + raise AttributeError(f"module {home.__name__!r} has no attribute {name!r}") + _warn_old_path(name, home) + return value + + def __setattr__(self, name: str, value: object) -> None: + if name in _MOVED_PRIVATE_NAME_MODULES: + home: Final = _moved_name_home(name) + _warn_old_path(name, home) + setattr(home, name, value) + return + for module in _split_modules_sharing(name, _namespace(self).get(name, _UNBOUND)): + setattr(module, name, value) + super().__setattr__(name, value) + + def __delattr__(self, name: str) -> None: + released: Final = _RELEASED_MOVED_VALUES.get(name, _UNBOUND) + if released is _UNBOUND: + super().__delattr__(name) + return + setattr(_moved_name_home(name), name, released) + + +sys.modules[__name__].__class__ = _CommonUtilsModule 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 9d3c19b7460..777bef7afcc 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 @@ -1,10 +1,15 @@ +import contextlib import copy +import importlib import json +import warnings from collections.abc import Mapping from typing import Final, cast +from unittest import mock import pytest +from litellm.litellm_core_utils.prompt_templates import common_utils from litellm.litellm_core_utils.prompt_templates.common_utils import ( TOOL_RESULT_IMAGE_BOUNDARY, TOOL_RESULT_IMAGE_PLACEHOLDER, @@ -21,6 +26,10 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( strip_litellm_response_only_fields, system_messages_first, ) +from litellm.litellm_core_utils.prompt_templates.json_schema_utils import ( + flatten_combinators_and_drop_non_python_regex_patterns, +) +from litellm.litellm_core_utils.prompt_templates.tool_call_json_repair import parse_tool_call_arguments from litellm.types.llms.openai import ChatCompletionAudioObject, ChatCompletionFileObject, ChatCompletionImageObject from litellm.types.utils import ChatCompletionMessageToolCall, Choices, Function, Message, ModelResponse @@ -987,6 +996,180 @@ def test_content_list_flattening_returns_new_messages_and_leaves_the_callers_unt assert messages == messages_before +_MOVED_PRIVATE_NAMES_BY_MODULE: Final = ( + ( + "file_id_utils", + ("_decode_vector_store_ids_in_tools",), + ), + ( + "json_schema_utils", + ( + "_EMPTY_SCHEMA", + "_LOCAL_SCHEMA_REF_PREFIXES", + "_MAX_SCHEMA_FLATTEN_DEPTH", + "_MAX_SCHEMA_NESTING", + "_OPENAI_REJECTED_TOP_LEVEL_SCHEMA_KEYS", + "_REGEX_LOOKAROUND_RE", + "_SCHEMA_LIST", + "_SUBSCHEMA_KEYWORDS", + "_SUBSCHEMA_LIST_KEYWORDS", + "_SUBSCHEMA_MAP_KEYWORDS", + "_TOP_LEVEL_SCHEMA_COMBINATORS", + "_any_of", + "_combinator_required_names", + "_dropped_pattern_properties", + "_flatten_schema_against_root", + "_grammar_branches", + "_is_not_python_regex", + "_is_object_schema", + "_is_open_object", + "_is_union", + "_keyword_value_rebuilt", + "_list_at", + "_may_be_string", + "_mergeable_branch", + "_named_schemas", + "_node_without_rejected_regex", + "_optional_property_names", + "_property_schemas", + "_refs_in", + "_resolve_local_schema_ref", + "_root_definitions", + "_schema_at", + "_schema_branches", + "_schema_levels", + "_schema_nodes", + "_schema_properties", + "_schema_required_names", + "_schema_without_rejected_regex", + "_subschemas", + "_uses_regex_lookaround", + ), + ), + ( + "prompt_cache_breakpoint_utils", + ( + "_BREAKPOINT_BLOCKS", + "_BREAKPOINT_BLOCK_LIST_KEYS", + "_BREAKPOINT_FIELDS", + "_MarkedT", + "_PROMPT_CACHE_BREAKPOINT", + "_breakpoints_stripped_from", + "_without_prompt_cache_breakpoint", + ), + ), + ( + "reasoning_utils", + ( + "_carries_encrypted_reasoning", + "_converse_wrapped_payload", + "_encrypted_reasoning_field", + "_provider_text_or_none", + "_readable_thinking_text", + "_reasoning_item_from_block_group", + "_reasoning_replay_group_key", + "_strip_encrypted_reasoning_from_blocks", + ), + ), + ( + "tool_call_json_repair", + ( + "_DECODED_JSON", + "_attempt_json_repair", + ), + ), + ( + "tpm_media_utils", + ( + "_TPMMediaProjection", + "_tpm_blank_fields", + "_tpm_code_output_media", + "_tpm_document_media", + "_tpm_file_media", + "_tpm_image_part", + "_tpm_input_media", + "_tpm_media_leaf", + "_tpm_media_list", + "_tpm_replace_field", + "_tpm_tool_result_media", + ), + ), +) + + +@pytest.mark.parametrize(("module_name", "names"), _MOVED_PRIVATE_NAMES_BY_MODULE) +def test_moved_private_helpers_still_import_from_common_utils(module_name: str, names: tuple[str, ...]) -> None: + new_home: Final = importlib.import_module(f"litellm.litellm_core_utils.prompt_templates.{module_name}") + unreachable: Final = [name for name in names if getattr(common_utils, name, None) is not getattr(new_home, name)] + assert unreachable == [] + + +def test_unknown_private_name_still_raises_attribute_error_on_common_utils() -> None: + missing_name: Final = "_not_a_moved_helper" + with pytest.raises(AttributeError, match=missing_name): + getattr(common_utils, missing_name) + + +def test_old_path_access_warns_from_the_callers_line_with_the_new_home() -> None: + moved_name: Final = "_any_of" + with pytest.warns(DeprecationWarning, match=r"common_utils\._any_of moved to .*\.json_schema_utils") as caught: + getattr(common_utils, moved_name) + assert [warning.filename for warning in caught] == [__file__] + + +def test_patching_a_moved_helper_at_its_old_path_reaches_the_callers_in_its_new_home() -> None: + patched_result: Final = {"patched": True} + with mock.patch( + "litellm.litellm_core_utils.prompt_templates.common_utils._attempt_json_repair", return_value=patched_result + ): + assert parse_tool_call_arguments('{"a": 1') == patched_result + assert parse_tool_call_arguments('{"a": 1') == {"a": 1} + + +def test_patching_a_re_exported_helper_at_its_old_path_reaches_the_callers_in_its_new_home() -> None: + schema: Final = {"anyOf": [{"type": "object", "properties": {"a": {"type": "string"}}}]} + unpatched: Final = flatten_combinators_and_drop_non_python_regex_patterns(schema) + patched_result: Final = {"patched": True} + with mock.patch( + "litellm.litellm_core_utils.prompt_templates.common_utils.flatten_top_level_schema_combinators", + return_value=patched_result, + ): + assert flatten_combinators_and_drop_non_python_regex_patterns(schema) == patched_result + assert flatten_combinators_and_drop_non_python_regex_patterns(schema) == unpatched + + +def test_old_path_patch_restores_the_moved_helper_when_warnings_become_errors_inside_it() -> None: + with warnings.catch_warnings(), contextlib.suppress(DeprecationWarning): + with mock.patch( + "litellm.litellm_core_utils.prompt_templates.common_utils._attempt_json_repair", return_value=None + ): + warnings.simplefilter("error", DeprecationWarning) + assert parse_tool_call_arguments('{"a": 1') == {"a": 1} + + +def test_old_path_patch_with_create_leaves_the_moved_helper_in_place() -> None: + with mock.patch( + "litellm.litellm_core_utils.prompt_templates.common_utils._attempt_json_repair", return_value=None, create=True + ): + assert parse_tool_call_arguments("{}") == {} + assert parse_tool_call_arguments('{"a": 1') == {"a": 1} + + +def test_overriding_a_moved_helper_through_the_old_path_warns_and_deleting_it_restores_the_released_one() -> None: + moved_name: Final = "_attempt_json_repair" + with pytest.warns(DeprecationWarning, match=r"common_utils\._attempt_json_repair moved to .*\.tool_call_json_repair"): + setattr(common_utils, moved_name, mock.Mock(return_value={"overridden": True})) + assert parse_tool_call_arguments('{"a": 1') == {"overridden": True} + delattr(common_utils, moved_name) + assert parse_tool_call_arguments('{"a": 1') == {"a": 1} + + +def test_patching_a_name_common_utils_still_defines_stays_on_common_utils() -> None: + with mock.patch.object(common_utils, "invalid_type_message", return_value="patched"): + assert common_utils.invalid_type_message("field", 1) == "patched" + assert common_utils.invalid_type_message("field", 1) != "patched" + + def test_get_content_from_model_response_includes_server_tool_calls_before_client_calls() -> None: server_call: Final = ChatCompletionMessageToolCall( id="srvtoolu_01", type="function", function=Function(name="web_search", arguments='{"query": "persimmon"}') diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_file_id_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_file_id_utils.py similarity index 100% rename from tests/unit/litellm_core_utils/prompt_templates/test_file_id_utils.py rename to tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_file_id_utils.py diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_json_schema_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_json_schema_utils.py similarity index 100% rename from tests/unit/litellm_core_utils/prompt_templates/test_json_schema_utils.py rename to tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_json_schema_utils.py diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_prompt_cache_breakpoint_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_prompt_cache_breakpoint_utils.py similarity index 100% rename from tests/unit/litellm_core_utils/prompt_templates/test_prompt_cache_breakpoint_utils.py rename to tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_prompt_cache_breakpoint_utils.py diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_reasoning_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_reasoning_utils.py similarity index 100% rename from tests/unit/litellm_core_utils/prompt_templates/test_reasoning_utils.py rename to tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_reasoning_utils.py diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_tool_call_json_repair.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_tool_call_json_repair.py similarity index 100% rename from tests/unit/litellm_core_utils/prompt_templates/test_tool_call_json_repair.py rename to tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_tool_call_json_repair.py diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_tpm_media_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_tpm_media_utils.py similarity index 100% rename from tests/unit/litellm_core_utils/prompt_templates/test_tpm_media_utils.py rename to tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_tpm_media_utils.py diff --git a/tests/unit/llms/test_file_content_block.py b/tests/unit/llms/test_file_content_block.py index 297b3b6f7d6..6e605eee033 100644 --- a/tests/unit/llms/test_file_content_block.py +++ b/tests/unit/llms/test_file_content_block.py @@ -6,7 +6,9 @@ Regression tests for: - litellm/llms/vertex_ai/gemini/transformation.py - litellm/llms/gemini/chat/transformation.py - litellm/litellm_core_utils/prompt_templates/common_utils.py - (migrate_file_to_image_url raises on missing `file`; file-id helpers skip non-OpenAI shapes) + (migrate_file_to_image_url raises on missing `file`) +- litellm/litellm_core_utils/prompt_templates/file_id_utils.py + (file-id helpers skip non-OpenAI shapes) - litellm/litellm_core_utils/prompt_templates/factory.py (Bedrock + Anthropic) - litellm/llms/openai/chat/gpt_transformation.py """ From 32a27427ba9cfdb104fc5c34140b5f70e84dd16f Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Mon, 9 Nov 2026 13:08:56 -0800 Subject: [PATCH 11/46] perf(router): look up host step answers by id instead of scanning every candidate (#50704) answered_by_id and the select answer check compared every returned id against every candidate, about 1M string compares per pick at 1,000 deployments --- .../crates/python-bridge/src/router/host.rs | 46 ++++++++++++++++--- 1 file changed, 39 insertions(+), 7 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/router/host.rs b/litellm-rust/crates/python-bridge/src/router/host.rs index 0ded98f27c5..27f9fe347e0 100644 --- a/litellm-rust/crates/python-bridge/src/router/host.rs +++ b/litellm-rust/crates/python-bridge/src/router/host.rs @@ -357,7 +357,11 @@ impl HostCallResponse { Self::Select(ids) | Self::Filtered(Ok((ids, _))), HostCall::Select { candidates, .. }, ) = (self, call) - && let Some(id) = ids.iter().find(|id| !candidates.contains(id)) + && let Some(id) = { + let known: std::collections::HashSet<&str> = + candidates.iter().map(String::as_str).collect(); + ids.iter().find(|id| !known.contains(id.as_str())) + } { return Err(DriveError::UnknownCandidate { id: id.clone() }); } @@ -688,14 +692,22 @@ fn shares_ids(candidates: &[Arc]) -> bool { /// The candidates a host answer names, in its order. An id names every candidate that carries /// it, so deployments sharing an id stay together through a filter, as Python's list keeps them. fn answered_by_id(candidates: &[Arc], ids: &[String]) -> Vec> { - let mut taken = vec![false; candidates.len()]; + let mut positions: std::collections::HashMap<&str, Vec> = + std::collections::HashMap::new(); + for (position, deployment) in candidates.iter().enumerate() { + positions + .entry(deployment.id.as_str()) + .or_default() + .push(position); + } let mut kept = Vec::new(); for id in ids { - for (position, deployment) in candidates.iter().enumerate() { - if !taken[position] && deployment.id == *id { - taken[position] = true; - kept.push(deployment.clone()); - } + if let Some(found) = positions.remove(id.as_str()) { + kept.extend( + found + .into_iter() + .map(|position| candidates[position].clone()), + ); } } kept @@ -1540,6 +1552,26 @@ mod tests { ); } + #[rstest] + fn an_answer_by_id_keeps_its_order_and_every_candidate_sharing_an_id() { + let candidates: Vec> = ["a", "b", "a", "c"] + .into_iter() + .map(|id| { + let spec: litellm_router_core::DeploymentSpec = + serde_json::from_value(deployment("g", id)).unwrap(); + Arc::new(Deployment::from_spec(spec).unwrap()) + }) + .collect(); + let ids = ["c", "a", "x", "a"].map(String::from); + + let kept = answered_by_id(&candidates, &ids); + + let positions: Vec = kept + .iter() + .map(|d| candidates.iter().position(|c| Arc::ptr_eq(c, d)).unwrap()) + .collect(); + assert_eq!(positions, [3, 0, 2]); + } #[rstest] fn a_pick_among_deployments_sharing_an_id_is_answered_by_position() { let mut second = deployment("primary", "a1"); From 07f5baba73f4f1310ce21a0d03217b1fb91b89cf Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Mon, 9 Nov 2026 13:20:47 -0800 Subject: [PATCH 12/46] feat(cost-map): add Claude Fable 5.2 on Vertex AI and repoint claude-fable-latest (#50713) * feat(cost-map): add Claude Fable 5.2 on the Anthropic API, Vertex AI and Claude Platform on AWS Adds claude-fable-5-2, vertex_ai/claude-fable-5-2 (plus @default) and bedrock/claude_platform/claude-fable-5-2 at $8/$40 per MTok, $10 and $16 cache writes, $0.20 cache reads, 50% batch discount and a 2x fast mode multiplier, and points claude-fable-latest and vertex_ai/claude-fable-latest at the new rows. Bedrock converse rows ship separately in #50707. * fix(cost-map): drop fast mode from the Vertex Fable 5.2 rows, the Vertex cost path does not bill it * Merge main, drop the Anthropic API and Claude Platform rows that #50715 already landed 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 | 86 ++++++++++++++++++- model_prices_and_context_window.json | 86 ++++++++++++++++++- tests/unit/test_claude_fable_5_config.py | 66 +++++++++++++- 3 files changed, 233 insertions(+), 5 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 79e4d92aa9d..600c195c2c7 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -17091,7 +17091,7 @@ "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-fable-latest": { - "alias_of": "claude-fable-5-1" + "alias_of": "claude-fable-5-2" }, "claude-opus-5-5": { "supports_web_search": true, @@ -53119,8 +53119,49 @@ "output_cost_per_token_batches": 2.5e-05, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" }, + "vertex_ai/claude-fable-5-2": { + "regional_endpoint_uplift_multiplier": 1.1, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_creation_input_token_cost_batches": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "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 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 512, + "deprecation_date": "2027-11-09", + "input_cost_per_token_batches": 4e-06, + "output_cost_per_token_batches": 2e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, "vertex_ai/claude-fable-latest": { - "alias_of": "vertex_ai/claude-fable-5-1" + "alias_of": "vertex_ai/claude-fable-5-2" }, "vertex_ai/claude-fable-5@default": { "deprecation_date": "2027-06-08", @@ -53202,6 +53243,47 @@ "output_cost_per_token_batches": 2.5e-05, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" }, + "vertex_ai/claude-fable-5-2@default": { + "regional_endpoint_uplift_multiplier": 1.1, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_creation_input_token_cost_batches": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "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 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 512, + "deprecation_date": "2027-11-09", + "input_cost_per_token_batches": 4e-06, + "output_cost_per_token_batches": 2e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, "vertex_ai/claude-opus-5": { "deprecation_date": "2027-01-24", "regional_endpoint_uplift_multiplier": 1.1, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 79e4d92aa9d..600c195c2c7 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -17091,7 +17091,7 @@ "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-fable-latest": { - "alias_of": "claude-fable-5-1" + "alias_of": "claude-fable-5-2" }, "claude-opus-5-5": { "supports_web_search": true, @@ -53119,8 +53119,49 @@ "output_cost_per_token_batches": 2.5e-05, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" }, + "vertex_ai/claude-fable-5-2": { + "regional_endpoint_uplift_multiplier": 1.1, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_creation_input_token_cost_batches": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "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 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 512, + "deprecation_date": "2027-11-09", + "input_cost_per_token_batches": 4e-06, + "output_cost_per_token_batches": 2e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, "vertex_ai/claude-fable-latest": { - "alias_of": "vertex_ai/claude-fable-5-1" + "alias_of": "vertex_ai/claude-fable-5-2" }, "vertex_ai/claude-fable-5@default": { "deprecation_date": "2027-06-08", @@ -53202,6 +53243,47 @@ "output_cost_per_token_batches": 2.5e-05, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" }, + "vertex_ai/claude-fable-5-2@default": { + "regional_endpoint_uplift_multiplier": 1.1, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_creation_input_token_cost_batches": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "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 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 512, + "deprecation_date": "2027-11-09", + "input_cost_per_token_batches": 4e-06, + "output_cost_per_token_batches": 2e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, "vertex_ai/claude-opus-5": { "deprecation_date": "2027-01-24", "regional_endpoint_uplift_multiplier": 1.1, diff --git a/tests/unit/test_claude_fable_5_config.py b/tests/unit/test_claude_fable_5_config.py index 7ec76742bb8..c9e46dd6887 100644 --- a/tests/unit/test_claude_fable_5_config.py +++ b/tests/unit/test_claude_fable_5_config.py @@ -1,5 +1,5 @@ """ -Validate Claude Fable 5 and Claude Fable 5.1 model configuration entries. +Validate Claude Fable 5, Claude Fable 5.1 and Claude Fable 5.2 model configuration entries. Fable 5 is a new tier above Opus ($10/$50 per MTok) with the same adaptive-only API surface as Opus 4.7/4.8. The cost-map entries below are what make the model @@ -11,9 +11,11 @@ shape, which Fable 5 rejects with a 400. import json import os +from typing import Final import pytest +import litellm from litellm.constants import BEDROCK_CONVERSE_MODELS from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap @@ -92,3 +94,65 @@ def test_adaptive_thinking_detected_for_fable_5_1(local_model_cost_map, model): from litellm.llms.anthropic.common_utils import AnthropicModelInfo assert AnthropicModelInfo.is_adaptive_thinking_model(model, "anthropic") is True + + +FABLE_5_2_VARIANTS = ( + "claude-fable-5-2", + "vertex_ai/claude-fable-5-2", + "vertex_ai/claude-fable-5-2@default", + "bedrock/claude_platform/claude-fable-5-2", +) + + +def test_fable_5_2_present_in_bundled_backup() -> None: + backup = GetModelCostMap.load_local_model_cost_map() + root = _load_root_cost_map() + for model_name in FABLE_5_2_VARIANTS: + assert model_name in backup, f"Missing from backup cost map: {model_name}" + assert backup[model_name] == root[model_name], model_name + + +@pytest.mark.parametrize( + "model", + [ + "claude-fable-5-2", + "anthropic/claude-fable-5-2", + "claude-fable-latest", + "vertex_ai/claude-fable-5-2", + "vertex_ai/claude-fable-5-2@default", + "bedrock/claude_platform/claude-fable-5-2", + ], +) +def test_adaptive_thinking_detected_for_fable_5_2(local_model_cost_map, model: str) -> None: + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + assert AnthropicModelInfo.is_adaptive_thinking_model(model, "anthropic") is True + + +def test_fable_5_2_fast_mode_doubles_every_token_price_on_the_anthropic_api(local_model_cost_map) -> None: + """Anthropic bills fast mode at 2x the standard rate on every token type + (platform.claude.com/docs/en/about-claude/pricing, read 2026-11-09), and the + Anthropic cost path is the one that reads the row's `fast` multiplier.""" + from litellm.litellm_core_utils.llm_cost_calc.utils import get_provider_specific_speed_multiplier + from litellm.types.utils import Usage + + info: Final = litellm.get_model_info("claude-fable-5-2") + assert get_provider_specific_speed_multiplier(info, Usage(speed="fast")) == 2.0 + assert get_provider_specific_speed_multiplier(info, Usage()) == 1.0 + + +def test_fable_5_2_vertex_rows_do_not_advertise_fast_mode(local_model_cost_map) -> None: + """Vertex cost calculation ignores the `fast` multiplier, so a Vertex row that + advertised fast mode would bill fast calls at the standard rate.""" + for model in ("vertex_ai/claude-fable-5-2", "vertex_ai/claude-fable-5-2@default"): + info: Final = litellm.get_model_info(model) + assert info.get("supports_fast_mode") is not True, model + assert "fast" not in (info.get("provider_specific_entry") or {}), model + + +def test_fable_5_2_accepts_the_speed_param_on_the_anthropic_api(local_model_cost_map) -> None: + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + assert AnthropicModelInfo.supports_fast_mode("claude-fable-5-2", "anthropic") is True + assert AnthropicConfig._model_supports_speed_param("claude-fable-5-2", "anthropic") is True From d7dc5d7fa3e26af00dcea3329265f6f17f5cea36 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Mon, 9 Nov 2026 13:24:37 -0800 Subject: [PATCH 13/46] test(integration): cover claude-fable-5-2 capability flags and prompt-cache pricing on the scripted Anthropic upstream (#50722) * test(e2e): cover claude-fable-5-2 capability flags and prompt-cache pricing over Anthropic direct Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): replace the fable-5-2 e2e suite with scripted-upstream wire and cost cases Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): add claude-fable-5-2 wire and cost cases 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> --- .../cost_calculation/cases/anthropic.py | 594 +++++++++++++++++- .../integration/cost_calculation/cost_map.py | 13 + ...t_anthropic_fable_capability_flags_wire.py | 208 ++++++ 3 files changed, 808 insertions(+), 7 deletions(-) create mode 100644 tests/integration/providers/test_anthropic_fable_capability_flags_wire.py diff --git a/tests/integration/cost_calculation/cases/anthropic.py b/tests/integration/cost_calculation/cases/anthropic.py index 776c5b51dd4..565d521aa95 100644 --- a/tests/integration/cost_calculation/cases/anthropic.py +++ b/tests/integration/cost_calculation/cases/anthropic.py @@ -4,8 +4,8 @@ PARITY pairs the streamed case with its plain twin where both bill the same row. from typing import Final +from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase, bridge_pairs from integration.cost_calculation.cost_tracking_case import ( - with_usage, CostTrackingTestCase, ExactExpected, FailureDetails, @@ -14,11 +14,10 @@ from integration.cost_calculation.cost_tracking_case import ( RecountExpected, RecountRates, SseResponse, + with_usage, ) -from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase, bridge_pairs from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames - # claude-haiku-4-5 CLAUDE_HAIKU_4_5_INPUT_TEXT: Final = CostTrackingTestCase( name="claude-haiku-4-5-input_text", @@ -5344,6 +5343,66 @@ CLAUDE_FABLE_5_1_INPUT_TEXT: Final = CostTrackingTestCase( ), ) +CLAUDE_FABLE_5_2_INPUT_TEXT: Final = CostTrackingTestCase( + name="claude-fable-5-2-input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "7d8c0a2d1ecc summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5-2", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer 7d8c0a2d1ecc"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + "speed": "standard", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.03432, input_cost=0.016192, output_cost=0.018128, prompt_tokens=1840, completion_tokens=412 + ), +) + CLAUDE_FABLE_5_1_STREAM: Final = CostTrackingTestCase( name="claude-fable-5-1-stream", covers="quota_management.spend_tracking.scripted_wire.logs_cost", @@ -5429,6 +5488,91 @@ CLAUDE_FABLE_5_1_STREAM: Final = CostTrackingTestCase( ), ) +CLAUDE_FABLE_5_2_STREAM: Final = CostTrackingTestCase( + name="claude-fable-5-2-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "459736de6b4a summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "stream_options": {"include_usage": True}, + "allowed_openai_params": [], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "type": "message_start", + "message": { + "model": "claude-fable-5-2", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [], + "container": None, + "stop_reason": None, + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 1, + "service_tier": "standard", + "inference_geo": "global", + "speed": "standard", + }, + "diagnostics": None, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "ping"}, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "scripted answer 459736de6b4a"}, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None, "stop_details": None, "container": None}, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + }, + }, + {"type": "message_stop"}, + ), + ), + expected=ExactExpected( + spend=0.03432, input_cost=0.016192, output_cost=0.018128, prompt_tokens=1840, completion_tokens=412 + ), +) + CLAUDE_FABLE_5_1_CACHE_READ: Final = CostTrackingTestCase( name="claude-fable-5-1-cache_read", covers="quota_management.spend_tracking.cost_matrix.logs_cost", @@ -5490,6 +5634,67 @@ CLAUDE_FABLE_5_1_CACHE_READ: Final = CostTrackingTestCase( ), ) +CLAUDE_FABLE_5_2_CACHE_READ: Final = CostTrackingTestCase( + name="claude-fable-5-2-cache_read", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "c5a9841fc31b summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5-2", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer c5a9841fc31b"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 640, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 12288, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 380, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + "speed": "standard", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.02505536, input_cost=0.00833536, output_cost=0.01672, prompt_tokens=12928, completion_tokens=380 + ), +) + CLAUDE_FABLE_5_1_CACHE_WRITE_5M: Final = CostTrackingTestCase( name="claude-fable-5-1-cache_write_5m", covers="quota_management.spend_tracking.cost_matrix.logs_cost", @@ -5550,6 +5755,66 @@ CLAUDE_FABLE_5_1_CACHE_WRITE_5M: Final = CostTrackingTestCase( ), ) +CLAUDE_FABLE_5_2_CACHE_WRITE_5M: Final = CostTrackingTestCase( + name="claude-fable-5-2-cache_write_5m", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "95e59856831e summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5-2", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer 95e59856831e"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 512, + "cache_creation_input_tokens": 9216, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 9216, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 350, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.1212816, input_cost=0.1058816, output_cost=0.0154, prompt_tokens=9728, completion_tokens=350 + ), +) + CLAUDE_FABLE_5_1_CACHE_WRITE_1H: Final = CostTrackingTestCase( name="claude-fable-5-1-cache_write_1h", covers="quota_management.spend_tracking.cost_matrix.logs_cost", @@ -5610,6 +5875,66 @@ CLAUDE_FABLE_5_1_CACHE_WRITE_1H: Final = CostTrackingTestCase( ), ) +CLAUDE_FABLE_5_2_CACHE_WRITE_1H: Final = CostTrackingTestCase( + name="claude-fable-5-2-cache_write_1h", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "0057eda38d9c summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5-2", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer 0057eda38d9c"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 512, + "cache_creation_input_tokens": 9216, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 9216}, + "output_tokens": 350, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.1821072, input_cost=0.1667072, output_cost=0.0154, prompt_tokens=9728, completion_tokens=350 + ), +) + CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT: Final = CostTrackingTestCase( name="claude-fable-5-1-messages_input_text", covers="quota_management.spend_tracking.cost_matrix.logs_cost", @@ -5651,6 +5976,47 @@ CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT: Final = CostTrackingTestCase( ), ) +CLAUDE_FABLE_5_2_MESSAGES_INPUT_TEXT: Final = CostTrackingTestCase( + name="claude-fable-5-2-messages_input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5-2", + endpoint="/v1/messages", + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "summarize the attached material in one line"}], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5-2", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted response"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + "speed": "standard", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.03432, input_cost=0.016192, output_cost=0.018128, prompt_tokens=1840, completion_tokens=412 + ), +) + CLAUDE_FABLE_5_1_MESSAGES_STREAM: Final = CostTrackingTestCase( name="claude-fable-5-1-messages_stream", @@ -5680,6 +6046,35 @@ CLAUDE_FABLE_5_1_MESSAGES_STREAM: Final = CostTrackingTestCase( spend=0.0624, input_cost=0.02944, output_cost=0.03296, prompt_tokens=1840, completion_tokens=412 ), ) + +CLAUDE_FABLE_5_2_MESSAGES_STREAM: Final = CostTrackingTestCase( + name="claude-fable-5-2-messages_stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="claude-fable-5-2", + endpoint="/v1/messages", + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "summarize the attached material in one line"}], + "stream": True, + }, + response=SseResponse( + content_type="text/event-stream", + frames=( + 'event: message_start\ndata: {"type":"message_start","message":{"model":"claude-fable-5-2","id":"msg_$REQUEST_ID","type":"message","role":"assistant","content":[],"container":null,"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":1840,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"cache_creation":{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":0},"output_tokens":1,"service_tier":"standard","inference_geo":"global","speed":"standard"},"diagnostics":null}}', + 'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}', + 'event: ping\ndata: {"type":"ping"}', + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"scripted "}}', + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"response"}}', + 'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}', + 'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null,"stop_details":null,"container":null},"usage":{"input_tokens":1840,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"output_tokens":412,"output_tokens_details":{"thinking_tokens":0}}}', + 'event: message_stop\ndata: {"type":"message_stop"}', + ), + ), + expected=ExactExpected( + spend=0.03432, input_cost=0.016192, output_cost=0.018128, prompt_tokens=1840, completion_tokens=412 + ), +) CLAUDE_OPUS_5_5_STREAM: Final = CostTrackingTestCase( name="claude-opus-5-5-stream", covers="quota_management.spend_tracking.scripted_wire.logs_cost", @@ -6560,6 +6955,92 @@ CLAUDE_FABLE_5_1_TOOL_CALL: Final = CostTrackingTestCase( ), ) +CLAUDE_FABLE_5_2_TOOL_CALL: Final = CostTrackingTestCase( + name="claude-fable-5-2-tool_call", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "6c08522a6809 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather and a short forecast for a city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"}, + "days": {"type": "integer", "description": "Forecast horizon in days"}, + "units": {"type": "string", "enum": ["metric", "imperial"]}, + }, + "required": ["city"], + }, + }, + } + ], + "tool_choice": "auto", + "allowed_openai_params": ["tool_choice"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5-2", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_$REQUEST_ID", + "name": "get_weather", + "input": {"city": "Berlin", "days": 7, "units": "metric"}, + "caller": {"type": "direct"}, + } + ], + "container": None, + "stop_reason": "tool_use", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.03432, input_cost=0.016192, output_cost=0.018128, prompt_tokens=1840, completion_tokens=412 + ), +) + CLAUDE_FABLE_5_1_THINKING: Final = CostTrackingTestCase( name="claude-fable-5-1-thinking", covers="quota_management.spend_tracking.cost_matrix.logs_cost", @@ -6621,6 +7102,67 @@ CLAUDE_FABLE_5_1_THINKING: Final = CostTrackingTestCase( ), ) +CLAUDE_FABLE_5_2_THINKING: Final = CostTrackingTestCase( + name="claude-fable-5-2-thinking", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "697b4f57a5df summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5-2", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer 697b4f57a5df"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1240, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + "speed": "standard", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.02904, input_cost=0.010912, output_cost=0.018128, prompt_tokens=1240, completion_tokens=412 + ), +) + CLAUDE_OPUS_5_5_TOOL_CALL: Final = CostTrackingTestCase( name="claude-opus-5-5-tool_call", covers="quota_management.spend_tracking.cost_matrix.logs_cost", @@ -6895,6 +7437,14 @@ CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON: Final = CLAUDE_FABLE_5_1_MESSAGES_INPUT } ) +CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_JSON: Final = CLAUDE_FABLE_5_2_MESSAGES_INPUT_TEXT.model_copy( + update={ + "name": "claude-fable-5-2-responses_bridged_json", + "endpoint": "/v1/responses", + "request": {"model": "$MODEL", "input": "summarize the attached material in one line"}, + } +) + CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON: Final = CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT.model_copy( update={ "name": "claude-opus-5-5-responses_bridged_json", @@ -6903,6 +7453,7 @@ CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON: Final = CLAUDE_OPUS_5_5_MESSAGES_INPUT_T } ) + def _responses_bridged_stream_case( sibling: CostTrackingTestCase, stream_sibling: CostTrackingTestCase, name: str ) -> CostTrackingTestCase: @@ -6916,19 +7467,29 @@ def _responses_bridged_stream_case( CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM: Final = _responses_bridged_stream_case( - CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON, CLAUDE_FABLE_5_1_MESSAGES_STREAM, + CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON, + CLAUDE_FABLE_5_1_MESSAGES_STREAM, "claude-fable-5-1-responses_bridged_stream", ) + +CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_STREAM: Final = _responses_bridged_stream_case( + CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_JSON, + CLAUDE_FABLE_5_2_MESSAGES_STREAM, + "claude-fable-5-2-responses_bridged_stream", +) CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM: Final = _responses_bridged_stream_case( - CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, CLAUDE_HAIKU_5_5_MESSAGES_STREAM, + CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, + CLAUDE_HAIKU_5_5_MESSAGES_STREAM, "claude-haiku-5-5-responses_bridged_stream", ) CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_STREAM: Final = _responses_bridged_stream_case( - CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON, CLAUDE_OPUS_5_5_MESSAGES_STREAM, + CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON, + CLAUDE_OPUS_5_5_MESSAGES_STREAM, "claude-opus-5-5-responses_bridged_stream", ) CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM: Final = _responses_bridged_stream_case( - CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_JSON, CLAUDE_SONNET_5_5_MESSAGES_STREAM, + CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_JSON, + CLAUDE_SONNET_5_5_MESSAGES_STREAM, "claude-sonnet-5-5-responses_bridged_stream", ) @@ -7028,11 +7589,13 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_JSON, CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM, + CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_STREAM, CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM, CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_STREAM, CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM, CLAUDE_FABLE_5_RESPONSES_BRIDGED_JSON, CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON, + CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_JSON, CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON, CLAUDE_HAIKU_4_5_INPUT_TEXT, CLAUDE_HAIKU_4_5_CACHE_READ, @@ -7129,12 +7692,19 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( CLAUDE_FABLE_5_CACHE_WRITE_1H, CLAUDE_FABLE_5_MESSAGES_INPUT_TEXT, CLAUDE_FABLE_5_1_INPUT_TEXT, + CLAUDE_FABLE_5_2_INPUT_TEXT, CLAUDE_FABLE_5_1_STREAM, + CLAUDE_FABLE_5_2_STREAM, CLAUDE_FABLE_5_1_CACHE_READ, + CLAUDE_FABLE_5_2_CACHE_READ, CLAUDE_FABLE_5_1_CACHE_WRITE_5M, + CLAUDE_FABLE_5_2_CACHE_WRITE_5M, CLAUDE_FABLE_5_1_CACHE_WRITE_1H, + CLAUDE_FABLE_5_2_CACHE_WRITE_1H, CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, + CLAUDE_FABLE_5_2_MESSAGES_INPUT_TEXT, CLAUDE_FABLE_5_1_MESSAGES_STREAM, + CLAUDE_FABLE_5_2_MESSAGES_STREAM, CLAUDE_OPUS_5_5_STREAM, CLAUDE_OPUS_5_5_CACHE_READ, CLAUDE_OPUS_5_5_CACHE_WRITE_5M, @@ -7148,7 +7718,9 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( CLAUDE_FABLE_5_TOOL_CALL, CLAUDE_FABLE_5_THINKING, CLAUDE_FABLE_5_1_TOOL_CALL, + CLAUDE_FABLE_5_2_TOOL_CALL, CLAUDE_FABLE_5_1_THINKING, + CLAUDE_FABLE_5_2_THINKING, CLAUDE_OPUS_5_5_TOOL_CALL, CLAUDE_OPUS_5_5_THINKING, CLAUDE_OPUS_5_5_UPSTREAM_404_ZERO_SPEND, @@ -7159,6 +7731,9 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase( plain=CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON, streamed=CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM ), + StreamParityTestCase( + plain=CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_JSON, streamed=CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_STREAM + ), StreamParityTestCase( plain=CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, streamed=CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM ), @@ -7190,11 +7765,13 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase(plain=CLAUDE_HAIKU_5_5_INPUT_TEXT, streamed=CLAUDE_HAIKU_5_5_STREAM), StreamParityTestCase(plain=CLAUDE_FABLE_5_INPUT_TEXT, streamed=CLAUDE_FABLE_5_STREAM), StreamParityTestCase(plain=CLAUDE_FABLE_5_1_INPUT_TEXT, streamed=CLAUDE_FABLE_5_1_STREAM), + StreamParityTestCase(plain=CLAUDE_FABLE_5_2_INPUT_TEXT, streamed=CLAUDE_FABLE_5_2_STREAM), StreamParityTestCase(plain=CLAUDE_OPUS_5_5_INPUT_TEXT, streamed=CLAUDE_OPUS_5_5_STREAM), StreamParityTestCase(plain=CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, streamed=CLAUDE_OPUS_5_5_MESSAGES_STREAM), StreamParityTestCase(plain=CLAUDE_SONNET_5_5_MESSAGES_INPUT_TEXT, streamed=CLAUDE_SONNET_5_5_MESSAGES_STREAM), StreamParityTestCase(plain=CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, streamed=CLAUDE_HAIKU_5_5_MESSAGES_STREAM), StreamParityTestCase(plain=CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, streamed=CLAUDE_FABLE_5_1_MESSAGES_STREAM), + StreamParityTestCase(plain=CLAUDE_FABLE_5_2_MESSAGES_INPUT_TEXT, streamed=CLAUDE_FABLE_5_2_MESSAGES_STREAM), ) BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( @@ -7220,6 +7797,9 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( BridgeCostParityTestCase( native=CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, bridged=CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON ), + BridgeCostParityTestCase( + native=CLAUDE_FABLE_5_2_MESSAGES_INPUT_TEXT, bridged=CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_JSON + ), BridgeCostParityTestCase( native=CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, bridged=CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON ), diff --git a/tests/integration/cost_calculation/cost_map.py b/tests/integration/cost_calculation/cost_map.py index ffdb6752bf2..37dc86e2601 100644 --- a/tests/integration/cost_calculation/cost_map.py +++ b/tests/integration/cost_calculation/cost_map.py @@ -1398,6 +1398,19 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "output_cost_per_token": 8e-05, "supports_function_calling": True, }, + "claude-fable-5-2": { + "cache_creation_input_token_cost": 1.1e-05, + "cache_creation_input_token_cost_above_1hr": 1.76e-05, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 8.8e-06, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-05, + "supports_function_calling": True, + }, "us.anthropic.claude-sonnet-5": { "cache_creation_input_token_cost": 4.5e-06, "cache_creation_input_token_cost_above_1hr": 7.2e-06, diff --git a/tests/integration/providers/test_anthropic_fable_capability_flags_wire.py b/tests/integration/providers/test_anthropic_fable_capability_flags_wire.py new file mode 100644 index 00000000000..0b140e4626f --- /dev/null +++ b/tests/integration/providers/test_anthropic_fable_capability_flags_wire.py @@ -0,0 +1,208 @@ +"""Claude Fable 5.2 over Anthropic direct. Its cost-map row says `supports_sampling_params: false`, +`supports_forced_tool_use: false` and `thinking_always_on: true`, and Anthropic answers 400 to a +request that ignores any of them, so the proxy has to rewrite or refuse those requests before they +leave. The scripted upstream here is what proves the outbound body, a leaked field is invisible in +the response.""" + +import json +import uuid +from typing import Final + +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_MODEL: Final = "claude-fable-5-2" +_API_KEY: Final = "synthetic-anthropic-key" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_WEATHER_TOOL: Final[JsonValue] = { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, +} +_NAMED_TOOL_CHOICE: Final[JsonValue] = {"type": "function", "function": {"name": "get_weather"}} + + +def _anthropic_reply(identity: str, usage: dict[str, JsonValue] | None = None) -> bytes: + return json.dumps( + { + "id": identity, + "type": "message", + "role": "assistant", + "model": _MODEL, + "content": [ + {"type": "thinking", "thinking": "scripted thought", "signature": "scripted-signature"}, + {"type": "text", "text": "done"}, + ], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": usage or {"input_tokens": 12, "output_tokens": 3}, + } + ).encode() + + +def _chat_body(model: str, **extra: JsonValue) -> dict[str, JsonValue]: + return { + "model": model, + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}], + "cache": {"no-cache": True}, + **extra, + } + + +def _outbound(request: Request) -> dict[str, JsonValue]: + assert request.method == "POST" and request.target == "/v1/messages", (request.method, request.target) + assert request.headers["x-api-key"] == _API_KEY + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["model"] == _MODEL, body + return body + + +def test_drop_params_removes_temperature_and_top_p_before_anthropic(gateway: Gateway) -> None: + identity: Final = f"fable-sampling-{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + body: Final = _outbound(request) + assert "temperature" not in body and "top_p" not in body, body + return Reply(body=_anthropic_reply(identity)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, drop_params=True + ) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, temperature=0.2, top_p=0.5)) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 + + +def test_sampling_params_are_refused_before_reaching_anthropic_without_drop_params(gateway: Gateway) -> None: + def respond(request: Request) -> Reply: + raise AssertionError(f"a request with sampling params reached Anthropic: {request.body!r}") + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, temperature=0.2)) + assert response.status_code == 400, response.text + assert "temperature" in response.text and "drop_params" in response.text, response.text + assert wire.drain() == () + + +def test_drop_params_downgrades_forced_and_named_tool_choice_to_auto(gateway: Gateway) -> None: + identity: Final = f"fable-tool-choice-{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + body: Final = _outbound(request) + assert body["tool_choice"] == {"type": "auto"}, body + tools: Final = body["tools"] + assert isinstance(tools, list) and len(tools) == 1, body + first: Final = tools[0] + assert isinstance(first, dict) and first["name"] == "get_weather", body + return Reply(body=_anthropic_reply(identity)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, drop_params=True + ) + responses: Final = tuple( + gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, tools=[_WEATHER_TOOL], tool_choice=tool_choice) + ) + for tool_choice in ("required", _NAMED_TOOL_CHOICE) + ) + assert [response.status_code for response in responses] == [200, 200], [r.text for r in responses] + assert len(wire.drain()) == 2 + + +def test_forced_tool_choice_is_refused_before_reaching_anthropic_without_drop_params(gateway: Gateway) -> None: + def respond(request: Request) -> Reply: + raise AssertionError(f"a forced tool_choice reached Anthropic: {request.body!r}") + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, tools=[_WEATHER_TOOL], tool_choice="required") + ) + assert response.status_code == 400, response.text + assert "tool" in response.text and "drop_params" in response.text, response.text + assert wire.drain() == () + + +def test_thinking_disabled_is_dropped_for_an_always_on_model(gateway: Gateway) -> None: + identity: Final = f"fable-thinking-off-{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + body: Final = _outbound(request) + assert "thinking" not in body, body + return Reply(body=_anthropic_reply(identity)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, thinking={"type": "disabled"}) + ) + assert response.status_code == 200, response.text + payload: Final = _JSON_OBJECT.validate_json(response.content) + choices: Final = payload["choices"] + assert isinstance(choices, list) and isinstance(choices[0], dict), response.text + message: Final = choices[0]["message"] + assert isinstance(message, dict) and message["reasoning_content"] == "scripted thought", response.text + assert len(wire.drain()) == 1 + + +def test_reasoning_effort_becomes_adaptive_thinking_with_output_config_effort(gateway: Gateway) -> None: + identity: Final = f"fable-effort-{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + body: Final = _outbound(request) + assert body["thinking"] == {"type": "adaptive", "display": "summarized"}, body + assert body["output_config"] == {"effort": "max"}, body + assert "budget_tokens" not in json.dumps(body), body + return Reply(body=_anthropic_reply(identity)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, reasoning_effort="max")) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 + + +def test_messages_one_hour_cache_control_reaches_anthropic_and_its_usage_comes_back(gateway: Gateway) -> None: + identity: Final = f"fable-cache-1h-{uuid.uuid4().hex}" + policy: Final = f"policy {identity}" + upstream_usage: Final[dict[str, JsonValue]] = { + "input_tokens": 12, + "output_tokens": 3, + "cache_creation_input_tokens": 9216, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 9216}, + } + + def respond(request: Request) -> Reply: + body: Final = _outbound(request) + assert body["system"] == [ + {"type": "text", "text": policy, "cache_control": {"type": "ephemeral", "ttl": "1h"}} + ], body + return Reply(body=_anthropic_reply(identity, upstream_usage)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 16, + "system": [{"type": "text", "text": policy, "cache_control": {"type": "ephemeral", "ttl": "1h"}}], + "messages": [{"role": "user", "content": "hi"}], + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + usage: Final = _JSON_OBJECT.validate_json(response.content)["usage"] + assert isinstance(usage, dict), response.text + assert usage["cache_creation_input_tokens"] == 9216, usage + assert usage["cache_creation"] == {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 9216}, usage + assert len(wire.drain()) == 1 From dfb1aac8b1c2a8b6e80d0ff7449049693563f69f Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 9 Nov 2026 13:37:55 -0800 Subject: [PATCH 14/46] fix(cost)!: bill transcriptions by the seconds the provider reports (#50727) * fix(cost): bill transcriptions by the seconds the provider reports * Update litellm/cost_calculator.py sure Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --------- Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/cost_calculator.py | 40 ++------ .../transcriptions/test_gpt_transcribe.py | 4 +- .../test_transcription_duration_hidden.py | 4 +- tests/unit/test_cost_calculator.py | 97 ++++++++----------- 4 files changed, 50 insertions(+), 95 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 5f7d646e150..58947f9d0c2 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -103,9 +103,6 @@ from litellm.llms.openai.cost_calculation import ( from litellm.llms.openai.cost_calculation import ( video_output_cost_per_second, ) -from litellm.llms.openai.transcriptions.gpt_transformation import ( - OpenAIGPTAudioTranscriptionConfig, -) from litellm.llms.perplexity.cost_calculator import ( cost_per_token as perplexity_cost_per_token, ) @@ -381,38 +378,15 @@ def _transcription_uses_token_pricing( return (prompt_tokens_val > 0) or (completion_tokens_val > 0) -def _bills_transcription_usage_duration(pricing_model: str | None, custom_llm_provider: str | None) -> bool: - if pricing_model is None or custom_llm_provider is None: - return False - transcription_info: Final = _lookup_model_info_or_none(pricing_model, custom_llm_provider) - transcription_metadata: Final = ( - as_str_mapping(transcription_info.get("provider_specific_entry")) if transcription_info is not None else None - ) - return transcription_metadata is not None and transcription_metadata.get("transcription_usage_duration") == 1 - - -def _transcription_duration_for_cost( - completion_response: object, pricing_model: str | None, model: str | None, custom_llm_provider: str | None -) -> float: +def _transcription_duration_for_cost(completion_response: object) -> float: + billed_duration: Final = _get_transcription_usage_duration(completion_response) + if billed_duration is not None and billed_duration > 0: + return billed_duration hidden_params: Final = as_str_mapping(attribute_or_none(completion_response, "_hidden_params")) or {} measured_duration: Final = normalized_audio_duration_seconds( hidden_params.get("audio_transcription_duration", attribute_or_none(completion_response, "duration")) ) - if measured_duration is not None and measured_duration > 0: - return measured_duration - provider_duration: Final = _get_transcription_usage_duration(completion_response) - if ( - provider_duration is not None - and provider_duration > 0 - and any(_bills_transcription_usage_duration(name, custom_llm_provider) for name in (pricing_model, model)) - ): - return provider_duration - if model is None or custom_llm_provider != LlmProviders.OPENAI.value: - return 0.0 - transcription_config: Final = ProviderConfigManager.get_provider_audio_transcription_config( - model=model.removeprefix(f"{LlmProviders.OPENAI.value}/"), provider=LlmProviders.OPENAI - ) - return (provider_duration or 0.0) if isinstance(transcription_config, OpenAIGPTAudioTranscriptionConfig) else 0.0 + return measured_duration if measured_duration is not None and measured_duration > 0 else 0.0 OCRPricingField = Literal["ocr_cost_per_page", "ocr_cost_per_credit", "annotation_cost_per_page"] @@ -2132,9 +2106,7 @@ def completion_cost( prompt, ) elif call_type in _TRANSCRIPTION_CALL_TYPES: - audio_transcription_file_duration = _transcription_duration_for_cost( - completion_response, model, requested_model, custom_llm_provider - ) + audio_transcription_file_duration = _transcription_duration_for_cost(completion_response) elif call_type in _RERANK_CALL_TYPES: if completion_response is not None and isinstance(completion_response, RerankResponse): meta_obj = completion_response.meta diff --git a/tests/unit/llms/openai/transcriptions/test_gpt_transcribe.py b/tests/unit/llms/openai/transcriptions/test_gpt_transcribe.py index 9bb5fb6d300..4eff0651941 100644 --- a/tests/unit/llms/openai/transcriptions/test_gpt_transcribe.py +++ b/tests/unit/llms/openai/transcriptions/test_gpt_transcribe.py @@ -885,7 +885,7 @@ async def _spend_for_one_second_wav( @pytest.mark.asyncio -async def test_custom_named_gpt_transcribe_deployment_bills_the_measured_upload_over_reported_seconds( +async def test_custom_named_gpt_transcribe_deployment_bills_the_reported_seconds_over_the_measured_upload( local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch ) -> None: router: Final = litellm.Router( @@ -902,7 +902,7 @@ async def test_custom_named_gpt_transcribe_deployment_bills_the_measured_upload_ monkeypatch, {"text": "hello", "usage": {"type": "duration", "seconds": 2}}, router ) - assert spend.cost == pytest.approx(_gpt_transcribe_cost_per_second()) + assert spend.cost == pytest.approx(2 * _gpt_transcribe_cost_per_second()) @pytest.mark.asyncio diff --git a/tests/unit/llms/openai/transcriptions/test_transcription_duration_hidden.py b/tests/unit/llms/openai/transcriptions/test_transcription_duration_hidden.py index 61798b49be0..9c0b9bcec13 100644 --- a/tests/unit/llms/openai/transcriptions/test_transcription_duration_hidden.py +++ b/tests/unit/llms/openai/transcriptions/test_transcription_duration_hidden.py @@ -129,7 +129,7 @@ class TestCostCalculatorReadsDurationFromHiddenParams: @patch( # test-quality-ok: isolates duration selection from the model pricing table "litellm.cost_calculator.openai_cost_per_second" ) - def test_completion_cost_prefers_measured_duration_over_provider_usage_duration(self, mock_cost_fn): + def test_completion_cost_prefers_provider_usage_duration_over_measured_duration(self, mock_cost_fn): mock_cost_fn.return_value = (0.001, 0.0) response = TranscriptionResponse( text="test", @@ -148,7 +148,7 @@ class TestCostCalculatorReadsDurationFromHiddenParams: mock_cost_fn.assert_called_once() _, kwargs = mock_cost_fn.call_args - assert kwargs["duration"] == 17.5 + assert kwargs["duration"] == 60.0 @pytest.mark.parametrize( ("model", "provider"), diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 49d51e99790..1d9f79aafc6 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -7532,82 +7532,65 @@ def test_whisper_azure(): assert round(cost, 5) == round(expected_cost, 5) -@pytest.mark.parametrize( - ("provider", "requested_model", "response_extras", "expected_seconds"), - [ - ("openai", "openai/gpt-fake-transcribe", {}, 2.0), - ("openai", "openai/gpt-fake-transcribe", {"_audio_transcription_duration": 1.35}, 1.35), - ("openai", "openai/gpt-fake-transcribe", {"duration": 1.34}, 1.34), - ("openai", "openai/whisper-fake", {}, 0.0), - ("azure", "gpt-fake-transcribe", {}, 0.0), - ], - ids=[ - "undecodable_upload_bills_provider_seconds", - "measured_duration_still_wins", - "provider_duration_still_wins", - "whisper_models_never_bill_provider_seconds", - "azure_never_bills_provider_seconds", - ], -) -def test_transcription_cost_falls_back_to_provider_billed_seconds( - monkeypatch: pytest.MonkeyPatch, - provider: str, - requested_model: str, - response_extras: Mapping[str, float], - expected_seconds: float, -) -> None: - model: Final = requested_model.removeprefix(f"{provider}/") +def _per_second_transcription( + monkeypatch: pytest.MonkeyPatch, provider: str, model: str, response_object: Mapping[str, object] +) -> TranscriptionResponse: monkeypatch.setitem( litellm.model_cost, model, {"input_cost_per_second": 0.01, "litellm_provider": provider, "mode": "audio_transcription"}, ) - transcription: Final = convert_to_model_response_object( - response_object={ - "text": "What is the weather in Paris?", - "usage": {"type": "duration", "seconds": 2}, - **response_extras, - }, + return convert_to_model_response_object( + response_object={"text": "What is the weather in Paris?", **response_object}, model_response_object=TranscriptionResponse(), hidden_params={"model": model, "custom_llm_provider": provider}, response_type="audio_transcription", ) - cost: Final = completion_cost( - completion_response=transcription, model=requested_model, call_type="atranscription" + +@pytest.mark.parametrize( + ("provider", "model", "measured"), + [ + ("openai", "whisper-fake", {}), + ("openai", "whisper-fake", {"_audio_transcription_duration": 2.4}), + ("openai", "whisper-fake", {"duration": 2.4}), + ("openai", "gpt-fake-transcribe", {"_audio_transcription_duration": 2.4}), + ("azure", "whisper-fake", {"_audio_transcription_duration": 2.4}), + ], + ids=[ + "unmeasured_upload", + "measured_upload", + "verbose_json_duration", + "gpt_transcribe_model", + "azure", + ], +) +def test_transcription_cost_bills_the_seconds_the_provider_billed( + monkeypatch: pytest.MonkeyPatch, provider: str, model: str, measured: Mapping[str, float] +) -> None: + transcription: Final = _per_second_transcription( + monkeypatch, provider, model, {"usage": {"type": "duration", "seconds": 3}, **measured} ) - assert cost == pytest.approx(0.01 * expected_seconds) + cost: Final = completion_cost( + completion_response=transcription, model=f"{provider}/{model}", call_type="atranscription" + ) + + assert cost == pytest.approx(0.01 * 3) @pytest.mark.parametrize( - ("provider", "measured_seconds", "expected_seconds"), - [("openai", 1.0, 1.0), ("azure", 1.0, 1.0), ("azure", None, 7.0)], - ids=["openai_measured_upload_wins", "azure_measured_upload_wins", "unmeasured_upload_bills_provider_seconds"], + ("measured", "expected_seconds"), + [({"_audio_transcription_duration": 2.4}, 2.4), ({"duration": 2.4}, 2.4), ({}, 0.0)], + ids=["measured_upload", "reply_duration", "nothing_known"], ) -def test_transcription_priced_by_provider_duration_still_bills_the_measured_upload( - monkeypatch: pytest.MonkeyPatch, provider: str, measured_seconds: float | None, expected_seconds: float +def test_transcription_cost_without_provider_seconds_bills_the_measured_length( + monkeypatch: pytest.MonkeyPatch, measured: Mapping[str, float], expected_seconds: float ) -> None: - monkeypatch.setitem( - litellm.model_cost, - "gpt-fake-transcribe", - { - "input_cost_per_second": 0.01, - "litellm_provider": provider, - "mode": "audio_transcription", - "provider_specific_entry": {"transcription_usage_duration": 1}, - }, - ) - measured: Final = {} if measured_seconds is None else {"audio_transcription_duration": measured_seconds} - transcription: Final = convert_to_model_response_object( - response_object={"text": "What is the weather in Paris?", "usage": {"type": "duration", "seconds": 7}}, - model_response_object=TranscriptionResponse(), - hidden_params={"model": "gpt-fake-transcribe", "custom_llm_provider": provider, **measured}, - response_type="audio_transcription", - ) + transcription: Final = _per_second_transcription(monkeypatch, "openai", "whisper-fake", measured) cost: Final = completion_cost( - completion_response=transcription, model=f"{provider}/gpt-fake-transcribe", call_type="atranscription" + completion_response=transcription, model="openai/whisper-fake", call_type="atranscription" ) assert cost == pytest.approx(0.01 * expected_seconds) From 0b3558d9247f91192a5461f580afc03c04757bdb Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Mon, 9 Nov 2026 13:48:22 -0800 Subject: [PATCH 15/46] fix(logging): bill a refused response the router falls back from (#50735) When the proxy parks the success log of a non-streaming response and the router then discards that response (a safeguard refusal it falls back from), the parked log was sent in place of the serving attempt's, so the spend log and cost header reported only the refused attempt and the serving attempt was never billed. Carry the parked response as a billed attempt owed by whichever attempt ends the request, the same way a superseded stream is already billed, so one spend row and the cost header cover every billed attempt. Fixes #49103 Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 21 +++++++ .../test_litellm_logging.py | 63 +++++++++++++++++++ 2 files changed, 84 insertions(+) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index fcc521cf6cf..e9342a765bc 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2761,6 +2761,8 @@ class Logging(LiteLLMLoggingBaseClass): def carry_failed_attempt_usage(self) -> None: """Move what a failed earlier attempt stashed off the retry sharing this logging object, so the retry logs its own usage while the request's failure log still bills every attempt.""" + if self._carry_parked_success(): + return usage: Final = self._model_call_details_view.get("combined_usage_object") response_cost: Final = self._model_call_details_view.get("response_cost") billed_by_continuation: Final = self._stash_billed_by_continuation @@ -2783,6 +2785,25 @@ class Logging(LiteLLMLoggingBaseClass): return self._failed_attempts_owed_on_success = (*self._failed_attempts_owed_on_success, attempt) + def _carry_parked_success(self) -> bool: + """A billed response the router discards after its success log was parked (a safeguard refusal it falls + back from) is owed by whichever attempt ends the request, instead of being logged in its place.""" + parked: Final = self.deferred_success_result + if parked is None or self.enqueue_deferred_logging is None: + return False + self.deferred_success_result = None + self.enqueue_deferred_logging = None + self._async_success_scheduled = False + served: Final = self.served_round(parked[0]) + if served is None: + return True + self.model_call_details["failed_attempts_usage"] = ( + *self.model_call_details.get("failed_attempts_usage", ()), + served, + ) + self._failed_attempts_owed_on_success = (*self._failed_attempts_owed_on_success, served) + return True + def mark_stash_billed_by_continuation(self) -> None: self._stash_billed_by_continuation = True diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 58fd03e0cbc..2f5edbd6350 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -14690,3 +14690,66 @@ def test_standard_logging_payload_reports_the_per_call_prices_the_request_was_bi model_map_value: Final = payload["model_map_information"]["model_map_value"] assert model_map_value is not None assert (model_map_value["input_cost_per_token"], model_map_value["output_cost_per_token"]) == (3e-06, 7e-06) + + +@pytest.mark.asyncio +async def test_a_response_discarded_after_its_success_log_was_parked_is_billed_once_on_the_serving_attempt( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.responses.streaming_iterator import drain_logging_tasks + from litellm.utils import _schedule_async_success_logging + + class Recorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.events: tuple[tuple[dict[str, object], object], ...] = () + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self.events = (*self.events, (dict(kwargs), response_obj)) + + recorder: Final = Recorder() + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "success_callback", []) + logging_obj: Final = LitellmLogging( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="refused-then-served", + function_id="refused-then-served", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model="gpt-4o-mini", user="", optional_params={}, litellm_params={"metadata": {}}, custom_llm_provider="openai" + ) + logging_obj.defer_async_logging = True + refused: Final = ModelResponse( + id="refused", model="gpt-4o-mini", usage=Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + ) + served: Final = ModelResponse( + id="served", model="gpt-4o-mini", usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + ) + now: Final = datetime.datetime.now() + + _schedule_async_success_logging(logging_obj, refused, now, now, is_completion_with_fallbacks=False) + logging_obj.carry_failed_attempt_usage() + _schedule_async_success_logging(logging_obj, served, now, now, is_completion_with_fallbacks=False) + prepared: Final = logging_obj.prepare_response_cost(served, now) + enqueue: Final = logging_obj.enqueue_deferred_logging + assert enqueue is not None + enqueue() + await drain_logging_tasks() + await GLOBAL_LOGGING_WORKER.flush() + + both_attempts: Final = litellm.completion_cost(refused, model="gpt-4o-mini") + litellm.completion_cost( + served, model="gpt-4o-mini" + ) + ((kwargs, response_obj),) = recorder.events + assert getattr(response_obj, "id", None) == "served" + assert kwargs["response_cost"] == pytest.approx(both_attempts) + assert prepared is not None + assert prepared.amount == pytest.approx(both_attempts) From 74cad3a2c358b1d1b731696eed5f654172876535 Mon Sep 17 00:00:00 2001 From: Raunak Kumar Pandey <65064180+raunakpandey1@users.noreply.github.com> Date: Mon, 9 Nov 2026 16:54:57 -0500 Subject: [PATCH 16/46] fix(audio): measure duration for PathLike multipart tuple contents (#49806) * fix(audio): measure duration for PathLike multipart tuple contents * fix(audio): read a PathLike tuple's file directly instead of recursing * fix(audio): read a tuple's path in the existing path branch --------- Co-authored-by: ryan-crabbe-berri --- .../litellm_core_utils/audio_utils/utils.py | 9 ++++--- .../litellm_core_utils/test_audio_utils.py | 27 ++++++++++++------- 2 files changed, 24 insertions(+), 12 deletions(-) diff --git a/litellm/litellm_core_utils/audio_utils/utils.py b/litellm/litellm_core_utils/audio_utils/utils.py index 808a424815e..a89ab5b0c6c 100644 --- a/litellm/litellm_core_utils/audio_utils/utils.py +++ b/litellm/litellm_core_utils/audio_utils/utils.py @@ -301,9 +301,12 @@ def calculate_request_duration(file: FileTypes) -> float | None: "Pass bytes, an open file handle, a (filename, content) " "tuple, or a pathlib.Path." ) - elif isinstance(file, os.PathLike): - # File path (PathLike): SDK convenience. - with open(str(file), "rb") as f: + elif isinstance(file, os.PathLike) or ( + isinstance(file, tuple) and len(file) >= 2 and isinstance(file[1], os.PathLike) + ): + # File path (PathLike), bare or as the content of a (filename, path) tuple: SDK convenience. + path: Final = file if isinstance(file, os.PathLike) else file[1] + with open(str(path), "rb") as f: file_content = f.read() elif isinstance(file, tuple): # Tuple format: (filename, content, optional content_type) diff --git a/tests/unit/litellm_core_utils/test_audio_utils.py b/tests/unit/litellm_core_utils/test_audio_utils.py index 122445e8e30..520eeab12e9 100644 --- a/tests/unit/litellm_core_utils/test_audio_utils.py +++ b/tests/unit/litellm_core_utils/test_audio_utils.py @@ -6,6 +6,7 @@ import base64 import io import os import tempfile +from pathlib import Path from typing import Final from unittest.mock import mock_open, patch @@ -329,16 +330,12 @@ class TestCalculateRequestDuration: duration = calculate_request_duration(file_obj) # Verify it succeeded (returns a duration, not None) - assert ( - duration is not None - ), "Duration should be calculated even when BytesIO is at end" + assert duration is not None, "Duration should be calculated even when BytesIO is at end" assert isinstance(duration, float), "Duration should be a float" assert duration > 0, "Duration should be positive" # Verify the file position was restored - assert file_obj.tell() == len( - wav_header - ), "File position should be restored to original position" + assert file_obj.tell() == len(wav_header), "File position should be restored to original position" class TestGetAudioFileContentHash: @@ -364,9 +361,7 @@ class TestGetAudioFileContentHash: hash1 = get_audio_file_content_hash((filename1, content)) hash2 = get_audio_file_content_hash((filename2, content)) - assert ( - hash1 == hash2 - ), "Same content should produce same hash regardless of filename" + assert hash1 == hash2, "Same content should produce same hash regardless of filename" def test_bytes_input(self): """Test that raw bytes input works""" @@ -473,3 +468,17 @@ class TestSpeechMediaTypeFromAudioBytes: from litellm.litellm_core_utils.audio_utils.utils import speech_media_type_from_audio_bytes assert speech_media_type_from_audio_bytes(audio) == expected + + +@pytest.mark.parametrize("include_content_type", [False, True]) +def test_duration_reads_pathlike_content_in_multipart_tuple(tmp_path: Path, include_content_type: bool) -> None: + audio_path: Final = tmp_path / "sample.wav" + audio_path.write_bytes(base64.b64decode(_encoded_audio(2, 8000, 1, "WAV"))) + upload: Final = ("sample.wav", audio_path, "audio/wav") if include_content_type else ("sample.wav", audio_path) + assert calculate_request_duration(upload) == 2.0 + + +def test_duration_does_not_open_bare_string_tuple_content(tmp_path: Path) -> None: + audio_path: Final = tmp_path / "sample.wav" + audio_path.write_bytes(base64.b64decode(_encoded_audio(2, 8000, 1, "WAV"))) + assert calculate_request_duration(("sample.wav", str(audio_path))) is None From 70f399b8c5502b938aba09f0c81c1d677c3d5f3a Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 9 Nov 2026 13:59:40 -0800 Subject: [PATCH 17/46] feat(auth): grant access groups through principal attribute rules (#50455) * feat(auth): grant access groups through principal attribute rules Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(auth): satisfy cache tracing and recursion CI checks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(auth): close access group rule authorization bypasses Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(auth): reject self-editable identity rule attributes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * Centralize JWT model authorization Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * Cover centralized key checks and name denied JWT models Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * Test centralized authorization for custom auth and JWT mappings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * Preserve model normalization before budget validation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * Use a controlled clock for WebSocket session deadline tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * Unify JWT credential admission and simplify team selection Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * Simplify request-time access group authorization Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * Use real requests in custom auth budget fixtures Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Krrish Dholakia Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 8 + litellm/proxy/auth/access_group_rules.py | 201 +++++++ litellm/proxy/auth/auth_checks.py | 2 +- litellm/proxy/auth/handle_jwt.py | 48 ++ litellm/proxy/auth/user_api_key_auth.py | 73 ++- litellm/proxy/proxy_server.py | 1 + litellm/repositories/table_repositories.py | 5 + litellm/types/proxy/auth/user_api_key_auth.py | 11 + pyproject.toml | 1 + .../authorization/test_access_group_rules.py | 221 ++++++++ .../proxy/auth/test_access_group_rules.py | 245 ++++++++ .../auth/test_custom_auth_end_user_budget.py | 10 +- tests/unit/proxy/auth/test_handle_jwt.py | 522 ++++++++++++------ .../unit/proxy/auth/test_user_api_key_auth.py | 170 +++++- .../test_user_api_key_auth_request_flow.py | 271 ++++----- .../mcp_server/test_idp_token_exchange.py | 86 ++- .../test_codex_passthrough_endpoints.py | 16 +- .../response_api_endpoints/test_endpoints.py | 11 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 13 + uv.lock | 23 + 20 files changed, 1599 insertions(+), 339 deletions(-) create mode 100644 litellm/proxy/auth/access_group_rules.py create mode 100644 tests/integration/authorization/test_access_group_rules.py create mode 100644 tests/unit/proxy/auth/test_access_group_rules.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b3357579829..27186800e2b 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -37,6 +37,7 @@ from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( validate_langfuse_span_scope_value, validate_no_callback_env_reference, ) +from litellm.proxy.auth.access_group_rules import AccessGroupRule, _validate_rules from litellm.proxy.mcp_server.stdio_gate import MCP_STDIO_DISABLED_MESSAGE, is_mcp_stdio_enabled from litellm.types.agents import ( AgentCaller as AgentCaller, # noqa: PLC0414 # public re-export @@ -2507,6 +2508,13 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): Documents all the fields supported by `general_settings` in config.yaml """ + access_group_rules: tuple[AccessGroupRule, ...] | None = None + + @field_validator("access_group_rules", mode="before") + @classmethod + def validate_access_group_rules(cls, value: object) -> tuple[AccessGroupRule, ...] | None: + return _validate_rules(value) if value is not None else None + use_spend_logs_ledger: bool = Field( True, description="Apply request accounting from stored SpendLogs deliveries. Set to false to roll back to the previous in-memory and Redis spend queues.", diff --git a/litellm/proxy/auth/access_group_rules.py b/litellm/proxy/auth/access_group_rules.py new file mode 100644 index 00000000000..b8e3d23f2cc --- /dev/null +++ b/litellm/proxy/auth/access_group_rules.py @@ -0,0 +1,201 @@ +from __future__ import annotations + +import json +from collections.abc import Callable, Mapping, Sequence +from functools import lru_cache, reduce +from importlib import import_module +from typing import TYPE_CHECKING, Final, Literal, Protocol, runtime_checkable + +from pydantic import BaseModel, ConfigDict, TypeAdapter, model_validator + +from litellm._internal_context import with_service_target +from litellm._logging import verbose_proxy_logger +from litellm.constants import AUTH_OBJECTS_TARGET + +_TRUSTED_ATTRIBUTES: Final = frozenset(("jwt_claims", "user_id", "user_role")) + +if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.repositories.prisma_protocols import DatabaseClient + + +class AccessGroupRule(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid", strict=True) + + attribute: str + contains: str | int | float | bool | None = None + equals: str | int | float | bool | None = None + access_group: str + + @model_validator(mode="after") + def validate_operator(self) -> AccessGroupRule: + if ("contains" in self.model_fields_set) == ("equals" in self.model_fields_set): + raise ValueError("An access group rule requires exactly one of contains or equals") + from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth + + if not self.attribute or any(not part for part in self.attribute.split(".")) or not self.access_group: + raise ValueError("An access group rule requires a nonempty attribute path and access_group") + if self.attribute.split(".")[0] not in _TRUSTED_ATTRIBUTES: + raise ValueError( + "An access group rule attribute must start with an administrator- or IdP-controlled field: " + + ", ".join(sorted(_TRUSTED_ATTRIBUTES)) + ) + if self.attribute.split(".")[0] not in UserAPIKeyAuth.model_fields: + raise ValueError("An access group rule attribute must start with a UserAPIKeyAuth field") + return self + + +class _Enforcer(Protocol): + def enforce(self, subject: str, resource: str) -> bool: ... + + def add_policy(self, subject: str, resource: str) -> bool: ... + + +class _CasbinModel(Protocol): + def load_model_from_text(self, text: str) -> None: ... + + +@runtime_checkable +class _Casbin(Protocol): + Model: Callable[[], _CasbinModel] + Enforcer: Callable[[_CasbinModel], _Enforcer] + + +_MODEL: Final = """ +[request_definition] +r = sub, obj +[policy_definition] +p = sub, obj +[policy_effect] +e = some(where (p.eft == allow)) +[matchers] +m = r.sub == p.sub && r.obj == p.obj +""" +_RULES: Final = TypeAdapter(tuple[AccessGroupRule, ...]) +_ATTRIBUTES: Final = TypeAdapter(dict[str, object]) +_LIST: Final = TypeAdapter(list[object]) +_CASBIN: Final = TypeAdapter(_Casbin, config=ConfigDict(arbitrary_types_allowed=True)) +_MISSING: Final = object() + + +def _subject(operator: Literal["contains", "equals"], claim: str, value: object) -> str: + return json.dumps((operator, claim, value), ensure_ascii=True, separators=(",", ":")) + + +def _attribute_segment(value: object, segment: str) -> object: + if not isinstance(value, Mapping): + return _MISSING + return _ATTRIBUTES.validate_python(value).get(segment, _MISSING) + + +def _claim_value(attributes: Mapping[str, object], path: Sequence[str]) -> object: + return reduce(_attribute_segment, path, attributes) + + +def _subjects(rule: AccessGroupRule, identity: UserAPIKeyAuth) -> tuple[str, ...]: + first, *nested = rule.attribute.split(".") + root: Final[object] = getattr(identity, first, _MISSING) if first in type(identity).model_fields else _MISSING + value: Final = ( + _claim_value(_ATTRIBUTES.validate_python(root), nested) + if nested and isinstance(root, Mapping) + else _MISSING + if nested + else root + ) + if "contains" in rule.model_fields_set: + if not isinstance(value, list): + return () + values: Final = _LIST.validate_python(value) + return tuple(_subject("contains", rule.attribute, item) for item in values if _is_scalar(item)) + if not _is_scalar(value): + return () + return (_subject("equals", rule.attribute, value),) + + +def _is_scalar(value: object) -> bool: + return isinstance(value, (str, int, float, bool)) or value is None + + +@lru_cache(maxsize=1) +def _enforcer(policies: tuple[tuple[str, str], ...]) -> _Enforcer: + casbin: Final = _CASBIN.validate_python(import_module("casbin")) + model: Final = casbin.Model() + model.load_model_from_text(_MODEL) + enforcer: Final = casbin.Enforcer(model) + for subject, resource in policies: + enforcer.add_policy(subject, resource) + return enforcer + + +def _policy(rule: AccessGroupRule) -> tuple[str, str]: + operator: Final = "contains" if "contains" in rule.model_fields_set else "equals" + value: Final = rule.contains if operator == "contains" else rule.equals + return _subject(operator, rule.attribute, value), rule.access_group + + +def _validate_rules(value: object) -> tuple[AccessGroupRule, ...]: + rules: Final = _RULES.validate_python(value) + if rules: + _enforcer(tuple(_policy(rule) for rule in rules)) + return rules + + +@with_service_target(AUTH_OBJECTS_TARGET) +async def _resolve_name(name: str, prisma_client: DatabaseClient | None, cache: UserApiKeyCache) -> str | None: + from litellm.constants import DEFAULT_ACCESS_GROUP_CACHE_TTL + from litellm.proxy._types import LiteLLM_AccessGroupTable + from litellm.proxy.auth.auth_checks import get_access_group_tables + from litellm.repositories.table_repositories import AccessGroupRepository + + if prisma_client is None: + verbose_proxy_logger.debug("Access group rule cannot resolve %s without a database", name) + return None + cache_key: Final = f"access_group_rule_name:{name}" + cached: Final = await cache.async_get_cache(key=cache_key, model_type=LiteLLM_AccessGroupTable) + if cached is not None: + groups: Final = await get_access_group_tables( + (cached.access_group_id,), prisma_client=prisma_client, user_api_key_cache=cache + ) + if groups and groups[0].access_group_name == name: + return cached.access_group_id + group: Final = await AccessGroupRepository(prisma_client).find_by_name(name) + if group is None: + verbose_proxy_logger.debug("Access group rule references unknown access group %s", name) + return None + await cache.async_set_cache( + key=cache_key, value=group, model_type=LiteLLM_AccessGroupTable, ttl=DEFAULT_ACCESS_GROUP_CACHE_TTL + ) + return group.access_group_id + + +async def access_group_ids_for_identity( + identity: UserAPIKeyAuth, + *, + rules: tuple[AccessGroupRule, ...] | None = None, + prisma_client: DatabaseClient | None = None, + cache: UserApiKeyCache | None = None, +) -> tuple[str, ...]: + if rules is None: + from litellm.proxy import proxy_server + + settings: Final = _ATTRIBUTES.validate_python(proxy_server.general_settings) # pyright: ignore[reportUnknownMemberType] # legacy settings are validated at this boundary + configured: Final = settings.get("access_group_rules") + if not configured: + return () + return await access_group_ids_for_identity( + identity, + rules=_validate_rules(configured), + prisma_client=proxy_server.prisma_client, + cache=proxy_server.user_api_key_cache, + ) + if not rules or cache is None: + return () + enforcer: Final = _enforcer(tuple(_policy(rule) for rule in rules)) + names: Final = frozenset( + rule.access_group + for rule in rules + if any(enforcer.enforce(subject, rule.access_group) for subject in _subjects(rule, identity)) + ) + ids: Final = tuple([await _resolve_name(name, prisma_client, cache) for name in sorted(names)]) + return tuple(group_id for group_id in ids if group_id is not None) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 4514882a45f..e546c093473 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -5839,7 +5839,7 @@ async def get_authorized_resources_from_key_access_groups( continue team_authorized = bool(key_team_id and key_team_id in (ag.assigned_team_ids or [])) key_authorized = bool(key_token and key_token in (ag.assigned_key_ids or [])) - if team_authorized or key_authorized: + if team_authorized or key_authorized or ag_id in valid_token.access_group_rule_grants: authorized_resources.extend(getattr(ag, resource_field, []) or []) return list(set(authorized_resources)) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 9e11179c686..4c4e506ce38 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -1768,6 +1768,7 @@ class JWTAuthManager: denied_pass_through_team: LiteLLM_TeamTable | None = None route_denied_team_id: str | None = None # rebind-ok: retain route-denied team while checking all claims + fallback_team: tuple[str, LiteLLM_TeamTable] | None = None # rebind-ok: retain first route-allowed candidate if not team_ids: if ( @@ -1796,6 +1797,21 @@ class JWTAuthManager: if team_object and team_object.models is not None: team_models = team_object.models + if ( + fallback_team is None + and allowed_routes_check( + user_role=LitellmUserRoles.TEAM, + user_route=route, + litellm_proxy_roles=jwt_handler.litellm_jwtauth, + ) + and JWTAuthManager._team_has_passthrough_route_access( + team_object=team_object, + route=route, + request_method=request_method, + team_allowed_routes=jwt_handler.litellm_jwtauth.team_allowed_routes, + ) + ): + fallback_team = (team_id, team_object) if isinstance(team_models, list) and ( not requested_model or await can_team_access_model( @@ -1828,6 +1844,9 @@ class JWTAuthManager: except Exception: continue + if fallback_team is not None: + return fallback_team + if denied_pass_through_team is not None: JWTAuthManager._raise_team_passthrough_route_denial(route=route, team_object=denied_pass_through_team) @@ -2380,6 +2399,7 @@ class JWTAuthManager: route=route, request_method=request_method, jwt_handler=jwt_handler ) any_team_resolved = False + fallback_team: tuple[str, LiteLLM_TeamTable] | None = None # rebind-ok: retain first route-allowed membership for candidate_team_id in user_team_ids: try: team_object = await get_team_object( @@ -2393,6 +2413,17 @@ class JWTAuthManager: except HTTPException: continue any_team_resolved = True + if ( + team_route_allowed + and fallback_team is None + and JWTAuthManager._team_has_passthrough_route_access( + team_object=team_object, + route=route, + request_method=request_method, + team_allowed_routes=jwt_handler.litellm_jwtauth.team_allowed_routes, + ) + ): + fallback_team = (candidate_team_id, team_object) if requested_model: try: await can_team_access_model( @@ -2424,6 +2455,23 @@ class JWTAuthManager: ) return candidate_team_id, team_object, None + if fallback_team is not None: + fallback_id, fallback_object = fallback_team + return ( + fallback_id, + fallback_object, + await get_team_membership( + user_id=user_id, + team_id=fallback_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + if user_id + else None, + ) + if enforce_team_based_model_access: if requested_model and any_team_resolved: raise HTTPException( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 0148cafa9e9..9a3b18cf95c 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -45,6 +45,7 @@ from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_ from litellm.models.end_user import LiteLLM_EndUserTable from litellm.proxy._types import * from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_from_headers +from litellm.proxy.auth.access_group_rules import access_group_ids_for_identity from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module exports ExperimentalUIJWTToken, TeamNotFoundError, @@ -406,6 +407,17 @@ def _get_model_names_for_budget_checks( return model +async def _normalize_model_for_auth( + request_data: dict[str, object], # mutable-ok: normalization updates dispatch data + valid_token: UserAPIKeyAuth, + request: Request | None, + route: str, +) -> None: + await _normalize_claude_model(request_data, valid_token, request, route) + await _resolve_router_settings_model_group_alias(request_data, valid_token, request, route) + _resolve_codex_native_chatgpt_model(request_data, valid_token, request, route) + + def _budget_fallback_target_model(target: object) -> str | None: if isinstance(target, str): return target @@ -3053,14 +3065,7 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e param=None, code=status.HTTP_401_UNAUTHORIZED, ) - await enforce_key_and_fallback_model_access( - valid_token=valid_token, - request_data=request_data, - route=route, - request=request, - llm_model_list=llm_model_list, - llm_router=llm_router, - ) + await _normalize_model_for_auth(request_data, valid_token, request, route) await _prefetch_referenced_auth_objects( valid_token, end_user_id=end_user_id, user_api_key_cache=user_api_key_cache, prisma_client=prisma_client ) @@ -3414,6 +3419,7 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e start_time=start_time, ) virtual_key_auth_obj.via_virtual_key = True + request.state.requires_key_model_check = True # rebind-ok: defer until central grants return virtual_key_auth_obj @@ -4251,7 +4257,33 @@ async def _authorize_authenticated_request( prepare_agent_invocation, ) from litellm.proxy.agent_endpoints.identity_store import AgentIdentityStore - from litellm.proxy.proxy_server import general_settings, prisma_client, user_model + from litellm.proxy.proxy_server import ( + general_settings, + llm_model_list, # pyright: ignore[reportUnknownVariableType] # proxy global has untyped deployment entries + llm_router, + prisma_client, + user_model, + ) + + grants: Final = await access_group_ids_for_identity(user_api_key_auth_obj) + user_api_key_auth_obj.apply_access_group_rule_grants(grants) + key_check_required: Final = ( + getattr(request.state, "requires_key_model_check", False) or force_virtual_key_checks + ) + if key_check_required or grants: + model_identity: Final = ( + user_api_key_auth_obj + if key_check_required + else user_api_key_auth_obj.model_copy(update={"models": list(user_api_key_auth_obj.team_models or ())}) + ) + await enforce_key_and_fallback_model_access( + valid_token=model_identity, + request_data=request_data, + route=route, + request=request, + llm_model_list=llm_model_list, + llm_router=llm_router, + ) store: Final = AgentIdentityStore.from_client(prisma_client) if prisma_client is not None else None if user_api_key_auth_obj.agent_id is not None: @@ -4368,8 +4400,7 @@ async def _prefetch_referenced_auth_objects( user_api_key_cache: UserApiKeyCache, prisma_client: PrismaClient | None, ) -> None: - """Warm every object and spend counter the checks below will read, in one MGET each (one DB query when cold). - Runs after the key's model access check so a denied request costs no more than it did before.""" + """Warm every object and spend counter the checks below will read, in one MGET each (one DB query when cold).""" bind_admission_counter_keys(valid_token, end_user_id=end_user_id or None) bind_committed_admission_counters(valid_token, end_user_id=end_user_id or None) await prefetch_auth_objects( @@ -4462,7 +4493,7 @@ async def user_api_key_auth( spend_counter_batch_scope(_spend_counter_store()), ): try: - user_api_key_auth_obj: Final = await user_api_key_auth_builder( + authenticated: Final = await user_api_key_auth_builder( request=request, api_key=api_key, azure_api_key_header=azure_api_key_header, @@ -4479,6 +4510,14 @@ async def user_api_key_auth( if body_parse_exception is not None: raise body_parse_exception raise + from litellm.proxy.proxy_server import ( + general_settings, # pyright: ignore[reportUnknownVariableType] # proxy global is validated below + ) + + settings: Final = _GENERAL_SETTINGS_ADAPTER.validate_python(general_settings) + user_api_key_auth_obj: Final = ( + authenticated.model_copy() if settings.get("access_group_rules") else authenticated + ) user_api_key_auth_obj.budget_reservation = None user_api_key_auth_obj.agent_caller = agent_caller_from_headers( safe_get_request_headers(request), user_api_key_auth_obj @@ -4817,14 +4856,8 @@ async def _run_post_custom_auth_checks( ) if general_settings.get("custom_auth_run_common_checks", False): - await enforce_key_and_fallback_model_access( - valid_token=valid_token, - request_data=request_data, - route=route, - request=request, - llm_model_list=llm_model_list, - llm_router=llm_router, - ) + await _normalize_model_for_auth(request_data, valid_token, request, route) + request.state.requires_key_model_check = True # rebind-ok: defer until central grants current_model = _get_model_from_request_context( request_data=request_data, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 61f81907ff4..046510d026d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7272,6 +7272,7 @@ class ProxyConfig: "claude_code_marketplace_auth_required", "vertex_ai_pass_through_drop_anthropic_betas", "request_tag_mappings", + "access_group_rules", ) if name in typed_general_settings } diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index aa33da54572..fbf9e5d07f0 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -16,6 +16,7 @@ from pydantic import TypeAdapter from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.models.access_group import LiteLLM_AccessGroupTable from litellm.repositories.accounting_delivery_repository import LedgerDatabase, LedgerTransaction from litellm.repositories.prisma_protocols import RowT_co, TableActions @@ -300,6 +301,10 @@ class ModelTableRepository(PrismaTableRepository["prisma_models.LiteLLM_ModelTab class AccessGroupRepository(PrismaTableRepository["prisma_models.LiteLLM_AccessGroupTable"]): table_name = "litellm_accessgrouptable" + async def find_by_name(self, name: str) -> LiteLLM_AccessGroupTable | None: + row: Final = await self.table.find_unique(where={"access_group_name": name}) + return LiteLLM_AccessGroupTable.model_validate(row, from_attributes=True) if row is not None else None + class SSOConfigRepository(PrismaTableRepository["prisma_models.LiteLLM_SSOConfig"]): table_name = "litellm_ssoconfig" diff --git a/litellm/types/proxy/auth/user_api_key_auth.py b/litellm/types/proxy/auth/user_api_key_auth.py index e08459ea14d..8975a4057eb 100644 --- a/litellm/types/proxy/auth/user_api_key_auth.py +++ b/litellm/types/proxy/auth/user_api_key_auth.py @@ -133,6 +133,17 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob """ api_key: str | None = None + _access_group_rule_grants: frozenset[str] = PrivateAttr(default=frozenset()) + + @property + def access_group_rule_grants(self) -> frozenset[str]: + return self._access_group_rule_grants + + def apply_access_group_rule_grants(self, grants: tuple[str, ...]) -> None: + if grants: + self._access_group_rule_grants = frozenset(grants) + self.access_group_ids = list(dict.fromkeys((*(self.access_group_ids or ()), *grants))) + user_role: LitellmUserRoles | None = None allowed_model_region: AllowedModelRegion | None = None parent_otel_span: Span | None = None diff --git a/pyproject.toml b/pyproject.toml index 02834f604d2..de1ea342017 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,6 +56,7 @@ Documentation = "https://docs.litellm.ai" [project.optional-dependencies] switchyard = ["nemo-switchyard==0.4.0a1"] proxy = [ + "casbin==1.43.0", "gunicorn>=23.0.0,<24.0", "uvicorn>=0.33.0,<1.0", "granian>=2.7.4,<3.0; python_version < '3.15'", diff --git a/tests/integration/authorization/test_access_group_rules.py b/tests/integration/authorization/test_access_group_rules.py new file mode 100644 index 00000000000..68decdbf7aa --- /dev/null +++ b/tests/integration/authorization/test_access_group_rules.py @@ -0,0 +1,221 @@ +import json +import time +import uuid +from pathlib import Path +from typing import Final + +import jwt +import pytest +import yaml +from cryptography.hazmat.primitives.asymmetric import rsa + +from litellm.proxy.mcp_server.idp_token_exchange import REJECTED_SUBJECT_TOKEN +from tests.integration._support.client import Gateway, string_value +from tests.integration._support.database import read_rows +from tests.integration._support.process import owned_proxy +from tests.integration._support.wire import Reply, Request, wire_server + +pytestmark: Final = pytest.mark.shares_worker_state + + +@pytest.mark.parametrize( + "enabled,mapping,claimed", ((True, False, True), (False, False, True), (True, True, False), (True, True, True)) +) +def test_rules_add_models_to_jwt_and_scoped_virtual_key_without_opening_other_models( + gateway: Gateway, tmp_path: Path, enabled: bool, mapping: bool, claimed: bool +) -> None: + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + public_key: Final = json.loads(jwt.algorithms.RSAAlgorithm.to_jwk(private_key.public_key())) + key_id: Final = "rule-signing-key" + + def respond(request: Request) -> Reply: + assert request.method == "GET", request + return Reply(body=json.dumps({"keys": [{**public_key, "kid": key_id}]}).encode()) + + with wire_server(respond) as jwks, gateway.scenario() as scenario: + explicit: Final = scenario.model() + group_model: Final = scenario.model() + outside: Final = scenario.model() + group_name: Final = f"integration-rule-{uuid.uuid4().hex}" + created: Final = gateway.request( + "POST", "/v1/access_group", {"access_group_name": group_name, "access_model_names": [group_model]} + ) + assert created.status_code == 201, created.text + group_id: Final = string_value(created.json()["access_group_id"]) + scenario.cleanups.callback(gateway.request, "DELETE", f"/v1/access_group/{group_id}") + virtual_user: Final = scenario.user() + virtual_key: Final = scenario.key(models=[explicit], user_id=virtual_user) + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["general_settings"] = { + **config["general_settings"], + "enable_jwt_auth": True, + "litellm_jwtauth": { + "user_id_jwt_field": "sub", + "team_ids_jwt_field": "groups", + "user_id_upsert": True, + "enforce_team_based_model_access": True, + **( + {"virtual_key_claim_field": "sub", "unregistered_jwt_client_behavior": "auto_register"} + if mapping + else {} + ), + }, + **( + { + "access_group_rules": [ + {"attribute": "jwt_claims.groups", "contains": "science", "access_group": group_name}, + {"attribute": "user_id", "equals": virtual_user, "access_group": group_name}, + ] + } + if enabled + else {} + ), + } + path: Final = tmp_path / "access_group_rules.yaml" + path.write_text(yaml.safe_dump(config)) + subject: Final = f"integration-rule-user-{uuid.uuid4().hex}" + gateway.post("/user/new", {"user_id": subject, "user_role": "internal_user"}) + scenario.cleanups.callback(scenario.delete_user, subject) + team_id: Final = ( + scenario.team(models=[explicit], members_with_roles=[{"role": "user", "user_id": subject}]) + if claimed + else None + ) + team_claims: Final = [team_id] if team_id is not None else [] + token: Final = jwt.encode( + {"sub": subject, "groups": ["science", *team_claims], "exp": int(time.time()) + 300}, + private_key, + algorithm="RS256", + headers={"kid": key_id}, + ) + unmatched: Final = jwt.encode( + {"sub": subject, "groups": ["other", *team_claims], "exp": int(time.time()) + 300}, + private_key, + algorithm="RS256", + headers={"kid": key_id}, + ) + with owned_proxy(gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks.url}, config=path, workers=2) as candidate: + fallback_response: Final = candidate.request( + "POST", + "/v1/chat/completions", + { + "model": group_model, + "fallbacks": [outside], + "messages": [{"role": "user", "content": "ungranted fallback control"}], + }, + key=token, + ) + assert fallback_response.status_code == 403, fallback_response.text + for key, model, allowed in ( + (token, group_model, enabled and claimed), + *(((token, group_model, enabled and claimed),) if mapping else ()), + (token, outside, False), + (unmatched, group_model, False), + (virtual_key, explicit, True), + (virtual_key, group_model, enabled), + (virtual_key, outside, False), + ): + response: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "rule authorization control"}]}, + key=key, + ) + assert response.status_code == (200 if allowed else 403), response.text + if allowed: + assert ( + response.json()["choices"][0]["message"]["content"] + == "Hello! This is a mock response from the fake OpenAI endpoint." + ), response.text + + if mapping and claimed: + rows: Final = read_rows( + 'SELECT models, access_group_ids FROM "LiteLLM_VerificationToken" ' + 'WHERE user_id=%s AND team_id=%s', + (subject, team_id), + ) + assert len(rows) == 1, rows + assert group_model not in rows[0]["models"], rows + assert not rows[0]["access_group_ids"], rows + + +@pytest.mark.parametrize("consumer", ("auto_register", "token_exchange")) +@pytest.mark.parametrize("matching", (True, False)) +def test_teamless_jwt_cannot_provision_credentials_without_native_team_access( + gateway: Gateway, tmp_path: Path, consumer: str, matching: bool +) -> None: + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + public_key: Final = json.loads(jwt.algorithms.RSAAlgorithm.to_jwk(private_key.public_key())) + key_id: Final = "credential-rule-signing-key" + + def respond(request: Request) -> Reply: + assert request.method == "GET", request + return Reply(body=json.dumps({"keys": [{**public_key, "kid": key_id}]}).encode()) + + with wire_server(respond) as jwks, gateway.scenario() as scenario: + model: Final = scenario.model() + subject: Final = f"integration-credential-user-{uuid.uuid4().hex}" + gateway.post("/user/new", {"user_id": subject, "user_role": "internal_user"}) + scenario.cleanups.callback(scenario.delete_user, subject) + group_name: Final = f"credential-rule-{uuid.uuid4().hex}" + created_group: Final = gateway.request( + "POST", "/v1/access_group", {"access_group_name": group_name, "access_model_names": [model]} + ) + assert created_group.status_code == 201, created_group.text + group_id: Final = string_value(created_group.json()["access_group_id"]) + scenario.cleanups.callback(gateway.request, "DELETE", f"/v1/access_group/{group_id}") + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["general_settings"] = { + **config["general_settings"], + "enable_jwt_auth": True, + "litellm_jwtauth": { + "user_id_jwt_field": "sub", + "enforce_team_based_model_access": True, + **( + {"virtual_key_claim_field": "sub", "unregistered_jwt_client_behavior": "auto_register"} + if consumer == "auto_register" + else {} + ), + }, + "access_group_rules": [ + {"attribute": "jwt_claims.groups", "contains": "science", "access_group": group_name} + ], + } + path: Final = tmp_path / "credential_access_group_rules.yaml" + path.write_text(yaml.safe_dump(config)) + token: Final = jwt.encode( + {"sub": subject, "groups": ["science" if matching else "other"], "exp": int(time.time()) + 300}, + private_key, + algorithm="RS256", + headers={"kid": key_id}, + ) + with owned_proxy(gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks.url}, config=path, workers=2) as candidate: + if consumer == "auto_register": + for _ in range(2): + response: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "credential rejection control"}]}, + key=token, + ) + assert response.status_code == 403, response.text + assert "No teams found in token" in response.text, response.text + else: + registered: Final = candidate.client.post( + "/register", json={"redirect_uris": ["http://localhost/callback"], "client_name": "rule-control"} + ) + assert registered.status_code in (200, 201), registered.text + exchanged: Final = candidate.client.post( + "/token", + data={ + "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", + "client_id": registered.json()["client_id"], + "subject_token": token, + "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", + "requested_token_type": "urn:ietf:params:oauth:token-type:access_token", + }, + ) + assert exchanged.status_code == 400, exchanged.text + assert exchanged.json() == {"error": "invalid_request", "error_description": REJECTED_SUBJECT_TOKEN}, ( + exchanged.text + ) diff --git a/tests/unit/proxy/auth/test_access_group_rules.py b/tests/unit/proxy/auth/test_access_group_rules.py new file mode 100644 index 00000000000..e73acda7173 --- /dev/null +++ b/tests/unit/proxy/auth/test_access_group_rules.py @@ -0,0 +1,245 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest +from pydantic import ValidationError + +from litellm.proxy._types import ConfigGeneralSettings, LiteLLM_AccessGroupTable, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.access_group_rules import AccessGroupRule, access_group_ids_for_identity +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + +@dataclass(frozen=True) +class _Database: + db: object + + +@pytest.mark.parametrize( + "claims,expected", + ( + ({"groups": ["science"], "tenant": {"id": "production"}}, ("science-id", "standard-id")), + ({"groups": ["other"]}, ()), + ({"groups": "science", "tenant": {"id": ["production"]}}, ()), + ({}, ()), + ({"groups": ["science", "science"]}, ("science-id",)), + ({"groups": [{"name": "science"}]}, ()), + ), +) +@pytest.mark.asyncio +async def test_verified_attributes_grant_union_with_strict_claim_shapes( + claims: Mapping[str, object], expected: tuple[str, ...] +) -> None: + rules: Final = ( + AccessGroupRule(attribute="jwt_claims.groups", contains="science", access_group="science"), + AccessGroupRule(attribute="jwt_claims.tenant.id", equals="production", access_group="standard"), + ) + cache: Final = UserApiKeyCache() + for name in ("science", "standard"): + group: Final = LiteLLM_AccessGroupTable(access_group_id=f"{name}-id", access_group_name=name) + cache.set_cache(f"access_group_rule_name:{name}", group, model_type=LiteLLM_AccessGroupTable) + cache.set_cache(f"access_group_id:{name}-id", group, model_type=LiteLLM_AccessGroupTable) + assert ( + await access_group_ids_for_identity( + UserAPIKeyAuth(jwt_claims=dict(claims)), rules=rules, cache=cache, prisma_client=_Database(object()) + ) + == expected + ) + + +@pytest.mark.parametrize("rule", ({"contains": "science", "equals": "science"}, {})) +def test_invalid_rule_fails_general_settings_validation(rule: Mapping[str, str]) -> None: + with pytest.raises(ValidationError, match="exactly one"): + ConfigGeneralSettings.model_validate( + {"access_group_rules": [{"attribute": "jwt_claims.groups", "access_group": "science", **rule}]} + ) + + +@pytest.mark.parametrize("attribute", ("metadata.department", "user_email", "user_email.domain", "team_id", "org_id")) +def test_caller_editable_attributes_are_rejected_as_grant_sources(attribute: str) -> None: + with pytest.raises(ValidationError, match="administrator- or IdP-controlled"): + ConfigGeneralSettings.model_validate( + {"access_group_rules": [{"attribute": attribute, "equals": "science", "access_group": "science"}]} + ) + + +@pytest.mark.asyncio +async def test_missing_group_and_no_database_grant_nothing() -> None: + rules: Final = (AccessGroupRule(attribute="user_id", equals="subject", access_group="absent"),) + table: Final = MagicMock() + table.find_unique = AsyncMock(return_value=None) + cache: Final = UserApiKeyCache() + assert ( + await access_group_ids_for_identity( + UserAPIKeyAuth(user_id="subject"), + rules=rules, + cache=cache, + prisma_client=_Database(MagicMock(litellm_accessgrouptable=table)), + ) + == () + ) + assert await access_group_ids_for_identity(UserAPIKeyAuth(user_id="subject"), rules=rules, cache=cache) == () + assert await access_group_ids_for_identity( + UserAPIKeyAuth(user_id="subject"), rules=rules, prisma_client=_Database(object()) + ) == () + + +@pytest.mark.asyncio +async def test_database_failure_does_not_turn_into_an_unrestricted_grant() -> None: + table: Final = MagicMock() + table.find_unique = AsyncMock(side_effect=RuntimeError("database unavailable")) + with pytest.raises(RuntimeError, match="database unavailable"): + await access_group_ids_for_identity( + UserAPIKeyAuth(user_id="subject"), + rules=(AccessGroupRule(attribute="user_id", equals="subject", access_group="science"),), + cache=UserApiKeyCache(), + prisma_client=_Database(MagicMock(litellm_accessgrouptable=table)), + ) + + +@pytest.mark.asyncio +async def test_virtual_key_user_grants_are_additive_and_do_not_mutate_cached_identity() -> None: + identity: Final = UserAPIKeyAuth(models=["explicit"], user_id="science", access_group_ids=["existing"]) + group: Final = LiteLLM_AccessGroupTable( + access_group_id="science-id", access_group_name="science", access_model_names=["group-model"] + ) + cache: Final = UserApiKeyCache() + cache.set_cache("access_group_rule_name:science", group, model_type=LiteLLM_AccessGroupTable) + cache.set_cache("access_group_id:science-id", group, model_type=LiteLLM_AccessGroupTable) + ids: Final = await access_group_ids_for_identity( + identity, + rules=(AccessGroupRule(attribute="user_id", equals="science", access_group="science"),), + cache=cache, + prisma_client=_Database(object()), + ) + granted: Final = identity.model_copy() + granted.apply_access_group_rule_grants(ids) + assert granted.access_group_ids == ["existing", "science-id"] + assert granted.access_group_rule_grants == frozenset({"science-id"}) + assert identity.access_group_ids == ["existing"] + identity.apply_access_group_rule_grants(()) + assert identity.access_group_ids == ["existing"] + + +@pytest.mark.asyncio +async def test_null_equals_does_not_match_an_absent_claim() -> None: + rules: Final = (AccessGroupRule(attribute="jwt_claims.absent", equals=None, access_group="science"),) + assert ( + await access_group_ids_for_identity(UserAPIKeyAuth(jwt_claims={}), rules=rules, cache=UserApiKeyCache()) == () + ) + + +@pytest.mark.parametrize( + "claims,local_email,expected", + ( + ({"email": "science@example.com"}, "other@example.com", ("science-id",)), + ({"email": "other@example.com"}, "science@example.com", ()), + ({}, "science@example.com", ()), + ({"email": ["science@example.com"]}, "science@example.com", ()), + ), +) +@pytest.mark.asyncio +async def test_email_grants_read_verified_claims_not_the_locally_editable_email( + claims: Mapping[str, object], local_email: str, expected: tuple[str, ...] +) -> None: + group: Final = LiteLLM_AccessGroupTable(access_group_id="science-id", access_group_name="science") + cache: Final = UserApiKeyCache() + cache.set_cache("access_group_rule_name:science", group, model_type=LiteLLM_AccessGroupTable) + cache.set_cache("access_group_id:science-id", group, model_type=LiteLLM_AccessGroupTable) + assert ( + await access_group_ids_for_identity( + UserAPIKeyAuth(jwt_claims=dict(claims), user_email=local_email), + rules=( + AccessGroupRule(attribute="jwt_claims.email", equals="science@example.com", access_group="science"), + ), + cache=cache, + prisma_client=_Database(object()), + ) + == expected + ) + + +@pytest.mark.parametrize("operator", ("contains", "equals")) +@pytest.mark.asyncio +async def test_null_operator_remains_distinct_when_rules_are_reloaded(operator: str) -> None: + table: Final = MagicMock() + table.find_unique = AsyncMock(return_value={"access_group_id": "science-id", "access_group_name": "science"}) + cache: Final = UserApiKeyCache() + database: Final = _Database(MagicMock(litellm_accessgrouptable=table)) + rule: Final = AccessGroupRule.model_validate( + {"attribute": "jwt_claims.value", operator: None, "access_group": "science"} + ) + principal: Final = UserAPIKeyAuth(jwt_claims={"value": [None] if operator == "contains" else None}) + assert await access_group_ids_for_identity(principal, rules=(rule,), cache=cache, prisma_client=database) == ( + "science-id", + ) + + +@pytest.mark.asyncio +async def test_authenticated_identity_fields_and_claims_are_available_without_metadata_flattening() -> None: + table: Final = MagicMock() + table.find_unique = AsyncMock(return_value={"access_group_id": "science-id", "access_group_name": "science"}) + identity: Final = UserAPIKeyAuth( + user_id="verified-user", user_role=LitellmUserRoles.INTERNAL_USER, jwt_claims={"groups": ["science"]} + ) + rules: Final = tuple( + AccessGroupRule.model_validate({"attribute": attribute, operator: value, "access_group": "science"}) + for attribute, operator, value in ( + ("user_id", "equals", "verified-user"), + ("jwt_claims.groups", "contains", "science"), + ("user_role", "equals", "internal_user"), + ) + ) + database: Final = _Database(MagicMock(litellm_accessgrouptable=table)) + for index, expected in ((0, ("science-id",)), (1, ("science-id",)), (2, ("science-id",))): + assert ( + await access_group_ids_for_identity( + identity, rules=(rules[index],), cache=UserApiKeyCache(), prisma_client=database + ) + == expected + ) + + +@pytest.mark.asyncio +async def test_name_resolution_is_cached_and_renamed_group_does_not_keep_granting_old_name() -> None: + table: Final = MagicMock() + table.find_unique = AsyncMock(return_value={"access_group_id": "science-id", "access_group_name": "science"}) + database: Final = _Database(MagicMock(litellm_accessgrouptable=table)) + cache: Final = UserApiKeyCache() + rules: Final = (AccessGroupRule(attribute="user_id", equals="subject", access_group="science"),) + assert await access_group_ids_for_identity( + UserAPIKeyAuth(user_id="subject"), rules=rules, cache=cache, prisma_client=database + ) == ("science-id",) + cache.set_cache( + "access_group_id:science-id", + LiteLLM_AccessGroupTable(access_group_id="science-id", access_group_name="science"), + model_type=LiteLLM_AccessGroupTable, + ) + assert await access_group_ids_for_identity( + UserAPIKeyAuth(user_id="subject"), rules=rules, cache=cache, prisma_client=database + ) == ("science-id",) + assert table.find_unique.await_count == 1 + cache.set_cache( + "access_group_id:science-id", + LiteLLM_AccessGroupTable(access_group_id="science-id", access_group_name="renamed"), + model_type=LiteLLM_AccessGroupTable, + ) + table.find_unique.return_value = None + assert ( + await access_group_ids_for_identity( + UserAPIKeyAuth(user_id="subject"), rules=rules, cache=cache, prisma_client=database + ) + == () + ) + + +@pytest.mark.parametrize( + "attribute", + ("groups", "department", "metadata.department", "models", "access_group_rule_grants", "jwt_claims..groups"), +) +def test_attribute_requires_a_real_principal_field(attribute: str) -> None: + with pytest.raises(ValidationError, match="attribute"): + ConfigGeneralSettings.model_validate( + {"access_group_rules": [{"attribute": attribute, "contains": "science", "access_group": "science"}]} + ) diff --git a/tests/unit/proxy/auth/test_custom_auth_end_user_budget.py b/tests/unit/proxy/auth/test_custom_auth_end_user_budget.py index 0c3487255f6..e371139dced 100644 --- a/tests/unit/proxy/auth/test_custom_auth_end_user_budget.py +++ b/tests/unit/proxy/auth/test_custom_auth_end_user_budget.py @@ -1,5 +1,7 @@ import pytest +from typing import Final from unittest.mock import AsyncMock, patch +from starlette.requests import Request import litellm from litellm.proxy.auth.user_api_key_auth import ( _run_post_custom_auth_checks, @@ -19,6 +21,7 @@ async def test_custom_auth_run_post_custom_auth_checks_without_end_user_id(): # custom_auth_run_common_checks for custom-auth deployments. The # helper itself no longer calls common_checks. valid_token = UserAPIKeyAuth(token="test_token") + http_request: Final = Request({"type": "http", "method": "POST", "path": "/v1/chat/completions"}) # Default: common_checks should NOT be called inside the helper with patch( @@ -27,7 +30,7 @@ async def test_custom_auth_run_post_custom_auth_checks_without_end_user_id(): mock_common.return_value = True result = await _run_post_custom_auth_checks( valid_token=valid_token, - request=None, + request=http_request, request_data={}, route="/v1/chat/completions", parent_otel_span=None, @@ -50,7 +53,7 @@ async def test_custom_auth_run_post_custom_auth_checks_without_end_user_id(): mock_common.return_value = True result = await _run_post_custom_auth_checks( valid_token=valid_token, - request=None, + request=http_request, request_data={}, route="/v1/chat/completions", parent_otel_span=None, @@ -137,6 +140,7 @@ async def test_custom_auth_defers_end_user_budget_to_common_checks_when_enabled( # With custom_auth_run_common_checks set, the wrapper's common_checks() # enforces the end-user budget, so the helper must not double-enforce it. valid_token = UserAPIKeyAuth(token="test_token", end_user_id="customer-1") + http_request: Final = Request({"type": "http", "method": "POST", "path": "/v1/chat/completions"}) end_user_obj = LiteLLM_EndUserTable( user_id="customer-1", blocked=False, @@ -165,7 +169,7 @@ async def test_custom_auth_defers_end_user_budget_to_common_checks_when_enabled( ): await _run_post_custom_auth_checks( valid_token=valid_token, - request=None, + request=http_request, request_data={"model": "gpt-4"}, route="/v1/chat/completions", parent_otel_span=None, diff --git a/tests/unit/proxy/auth/test_handle_jwt.py b/tests/unit/proxy/auth/test_handle_jwt.py index ae6bf95c13f..36ba8a734b8 100644 --- a/tests/unit/proxy/auth/test_handle_jwt.py +++ b/tests/unit/proxy/auth/test_handle_jwt.py @@ -1,4 +1,5 @@ import asyncio +import json import re import sys import time @@ -8,26 +9,32 @@ from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx +import jwt import pytest +from cryptography.hazmat.primitives.asymmetric import rsa from fastapi import HTTPException +from starlette.requests import Request +import litellm.proxy.proxy_server as proxy_server from litellm.caching.dual_cache import DualCache from litellm.proxy._types import ( DEFAULT_JWKS_STALE_TTL, - JWTLiteLLMRoleMap, JWTIssuerConfig, - LiteLLMRoutes, + JWTLiteLLMRoleMap, + LiteLLM_AccessGroupTable, LiteLLM_JWTAuth, LiteLLM_ModelTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, LiteLLM_UserTable, + LiteLLMRoutes, LitellmUserRoles, Member, ProxyErrorTypes, ProxyException, RoleBasedPermissions, ScopeMapping, + UserAPIKeyAuth, ) from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry from litellm.proxy.auth.auth_checks import TeamNotFoundError @@ -44,10 +51,315 @@ from litellm.proxy.auth.handle_jwt import ( NoMatchingJWTPublicKeyError, ) from litellm.proxy.auth.model_access_denied import ModelAccessDeniedHTTPException +from litellm.proxy.auth.user_api_key_auth import ( + user_api_key_auth, +) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.types.agents import AgentResponse +async def _jwt_team_selection_identity( + monkeypatch: pytest.MonkeyPatch, + teams: tuple[LiteLLM_TeamTable, ...], + claimed_teams: tuple[str, ...], + membership_teams: tuple[str, ...], + requested_model: str | None, + *, + authorize: bool = False, + fallback_to_db_teams: bool | None = None, + matching_rule: bool | None = None, + scope_models: tuple[str, ...] | None = None, + role_models: tuple[str, ...] | None = None, +) -> UserAPIKeyAuth: + cache: Final = UserApiKeyCache() + database: Final = MagicMock() + handler: Final = JWTHandler() + handler.update_environment( + prisma_client=database, + user_api_key_cache=cache, + litellm_jwtauth=LiteLLM_JWTAuth( + user_id_jwt_field="sub", + team_ids_jwt_field="groups", + enforce_team_based_model_access=True, + enforce_rbac=role_models is not None, + fallback_to_db_teams=not claimed_teams if fallback_to_db_teams is None else fallback_to_db_teams, + enforce_scope_based_access=scope_models is not None, + scope_mappings=[ScopeMapping(scope="inference", models=list(scope_models))] + if scope_models is not None + else None, + ), + ) + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + jwks_url: Final = "https://identity.invalid/selection/jwks" + cache.set_cache( + f"litellm_jwt_auth_keys_{jwks_url}", + [json.loads(jwt.algorithms.RSAAlgorithm.to_jwk(private_key.public_key()))], + ) + cache.set_cache( + "selection-user", + LiteLLM_UserTable(user_id="selection-user", user_role="internal_user", teams=list(membership_teams)), + model_type=LiteLLM_UserTable, + ) + for team in teams: + cache.set_cache(f"team_id:{team.team_id}", team, model_type=LiteLLM_TeamTable) + cache.set_cache( + f"team_membership:selection-user:{team.team_id}", + LiteLLM_TeamMembership(user_id="selection-user", team_id=team.team_id), + model_type=LiteLLM_TeamMembership, + ) + monkeypatch.setenv("JWT_PUBLIC_KEY_URL", jwks_url) + monkeypatch.delenv("JWT_ISSUER", raising=False) + monkeypatch.delenv("JWT_AUDIENCE", raising=False) + monkeypatch.setattr(proxy_server, "llm_router", None) + monkeypatch.setattr(proxy_server, "jwt_handler", handler) + monkeypatch.setattr(proxy_server, "prisma_client", database) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "premium_user", True) + monkeypatch.setattr(proxy_server, "master_key", "selection-master") + monkeypatch.setattr(proxy_server, "user_custom_auth", None) + monkeypatch.setattr( + proxy_server, + "general_settings", + { + "enable_jwt_auth": True, + **({"role_permissions": [RoleBasedPermissions( + role=LitellmUserRoles.INTERNAL_USER, models=list(role_models), routes=["/chat/completions", "/models"] + )]} if role_models is not None else {}), + **( + { + "access_group_rules": [ + { + "attribute": "jwt_claims.sub", + "equals": "selection-user" if matching_rule else "another-user", + "access_group": "selection-group", + } + ] + } + if matching_rule is not None + else {} + ), + }, + ) + group: Final = LiteLLM_AccessGroupTable( + access_group_id="selection-group-id", + access_group_name="selection-group", + access_model_names=["selection-model"], + ) + cache.set_cache("access_group_rule_name:selection-group", group, model_type=LiteLLM_AccessGroupTable) + cache.set_cache("access_group_id:selection-group-id", group, model_type=LiteLLM_AccessGroupTable) + token: Final = jwt.encode( + {"sub": "selection-user", "groups": list(claimed_teams), "scope": "inference"}, private_key, algorithm="RS256" + ) + request_data: Final = {"model": requested_model} if requested_model is not None else {} + route: Final = "/chat/completions" if requested_model is not None else "/models" + try: + if authorize: + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": route, + "headers": [(b"content-type", b"application/json")], + "query_string": b"", + } + ) + request._body = json.dumps(request_data).encode() + return await user_api_key_auth(request=request, api_key=f"Bearer {token}") + result: Final = await JWTAuthManager.auth_builder( + api_key=token, + jwt_handler=handler, + request_data=request_data, + general_settings=proxy_server.general_settings, + route=route, + prisma_client=database, + user_api_key_cache=cache, + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + identity: Final = JWTAuthManager.user_api_key_auth_from_result(result) + return identity + finally: + await handler.close() + + +@pytest.mark.asyncio +async def test_unresolved_claim_keeps_model_denial_after_single_membership_is_selected( + monkeypatch: pytest.MonkeyPatch, +) -> None: + membership_team: Final = LiteLLM_TeamTable( + team_id="membership-team", + models=["selection-model"], + members_with_roles=[Member(user_id="selection-user", role="user")], + ) + with pytest.raises(ProxyException) as denied: + await _jwt_team_selection_identity( + monkeypatch, + (membership_team,), + ("unresolved-claim",), + ("membership-team",), + "selection-model", + authorize=True, + ) + assert int(denied.value.code) == 403 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("claimed_teams", (("team-c", "team-a", "team-b"), ("team-b", "team-c", "team-a"))) +@pytest.mark.parametrize("requested_model,expected_team", (("selection-model", "team-b"), (None, "team-a"))) +async def test_auth_builder_claimed_teams_attribute_to_first_sorted_eligible_team( + monkeypatch: pytest.MonkeyPatch, + claimed_teams: tuple[str, ...], + requested_model: str | None, + expected_team: str, +) -> None: + teams: Final = tuple( + LiteLLM_TeamTable( + team_id=team_id, + models=["other-model"] if team_id == "team-a" else ["selection-model"], + members_with_roles=[Member(user_id="selection-user", role="user")], + ) + for team_id in ("team-a", "team-b", "team-c") + ) + identity: Final = await _jwt_team_selection_identity( + monkeypatch, teams, claimed_teams, ("team-c", "team-a", "team-b"), requested_model + ) + assert identity.team_id == expected_team + assert identity.user_id == "selection-user" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("membership_teams", (("team-a", "team-c", "team-b"), ("team-a", "team-b", "team-c"))) +async def test_auth_builder_db_teams_attribute_to_first_model_allowed_membership( + monkeypatch: pytest.MonkeyPatch, membership_teams: tuple[str, ...] +) -> None: + teams: Final = tuple( + LiteLLM_TeamTable( + team_id=team_id, + models=["other-model"] if team_id == "team-a" else ["selection-model"], + members_with_roles=[Member(user_id="selection-user", role="user")], + ) + for team_id in membership_teams + ) + identity: Final = await _jwt_team_selection_identity(monkeypatch, teams, (), membership_teams, "selection-model") + assert identity.team_id == membership_teams[1] + assert identity.user_id == "selection-user" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("matching", (True, False)) +async def test_jwt_rule_grants_a_claimed_teams_missing_model_without_losing_attribution( + monkeypatch: pytest.MonkeyPatch, + matching: bool, +) -> None: + team: Final = LiteLLM_TeamTable( + team_id="claimed-team", + models=["other-model"], + members_with_roles=[Member(user_id="selection-user", role="user")], + ) + if matching: + identity: Final = await _jwt_team_selection_identity( + monkeypatch, + (team,), + ("claimed-team",), + ("claimed-team",), + "selection-model", + authorize=True, + matching_rule=matching, + ) + assert identity.team_id == "claimed-team" + assert identity.user_id == "selection-user" + else: + with pytest.raises(ProxyException) as denied: + await _jwt_team_selection_identity( + monkeypatch, + (team,), + ("claimed-team",), + ("claimed-team",), + "selection-model", + authorize=True, + matching_rule=matching, + ) + assert denied.value.code == "403" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("requested_model", (None, "selection-model")) +async def test_jwt_native_team_access_with_rules_configured( + monkeypatch: pytest.MonkeyPatch, + requested_model: str | None, +) -> None: + team: Final = LiteLLM_TeamTable( + team_id="native-team", + models=["selection-model"], + members_with_roles=[Member(user_id="selection-user", role="user")], + ) + identity: Final = await _jwt_team_selection_identity( + monkeypatch, + (team,), + ("native-team",), + ("native-team",), + requested_model, + authorize=True, + matching_rule=True, + ) + assert identity.team_id == "native-team" + assert identity.user_id == "selection-user" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("granted", (True, False)) +@pytest.mark.parametrize("policy_kind", ("scope", "role")) +async def test_jwt_role_and_scope_model_restrictions_remain_in_builder( + monkeypatch: pytest.MonkeyPatch, + granted: bool, + policy_kind: str, +) -> None: + team: Final = LiteLLM_TeamTable( + team_id="scope-team", + models=["selection-model"], + members_with_roles=[Member(user_id="selection-user", role="user")], + ) + scope_models: Final = ("selection-model",) if granted else () + if granted: + identity: Final = await _jwt_team_selection_identity( + monkeypatch, + (team,), + ("scope-team",), + ("scope-team",), + "selection-model", + authorize=True, + scope_models=scope_models if policy_kind == "scope" else None, + role_models=scope_models if policy_kind == "role" else None, + ) + assert identity.team_id == "scope-team" + else: + with pytest.raises(ProxyException) as denied: + await _jwt_team_selection_identity( + monkeypatch, + (team,), + ("scope-team",), + ("scope-team",), + "selection-model", + authorize=True, + scope_models=scope_models if policy_kind == "scope" else None, + role_models=scope_models if policy_kind == "role" else None, + ) + assert denied.value.code == "403" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("matching", (True, False)) +async def test_teamless_jwt_is_rejected_before_rule_grants( + monkeypatch: pytest.MonkeyPatch, matching: bool, +) -> None: + with pytest.raises(ProxyException, match="No teams found in token") as denied: + await _jwt_team_selection_identity( + monkeypatch, (), (), (), "selection-model", authorize=True, + fallback_to_db_teams=False, matching_rule=matching, + ) + assert denied.value.code == "403" + + @pytest.mark.asyncio async def test_map_user_to_teams_user_already_in_team(): """Test that no action is taken when user is already in team""" @@ -213,7 +525,7 @@ async def test_find_team_with_model_access_reports_passthrough_allowlist_denial( assert exc_info.value.status_code == 403 assert "allowed_passthrough_routes" in exc_info.value.detail assert "requested model" not in exc_info.value.detail - mock_is_auth_enforced_pass_through_route.assert_called_once_with( + mock_is_auth_enforced_pass_through_route.assert_called_with( route="/my-pass-through", method="POST" ) @@ -3846,57 +4158,22 @@ async def test_find_team_with_model_access_enforce_team_based_access_still_raise @pytest.mark.asyncio async def test_find_team_with_model_access_resolved_team_without_model_still_raises_403( - monkeypatch, -): - """Regression guard: when the JWT group claim DOES resolve to a real - LiteLLM team but that team does not grant the requested model, keep the - original 403. Only the unresolved-claim case is softened.""" - import sys - import types - - from fastapi import HTTPException - - from litellm.router import Router - - router = Router( - model_list=[ - {"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}}, - { - "model_name": "gpt-3.5-turbo", - "litellm_params": {"model": "gpt-3.5-turbo"}, - }, - ] + monkeypatch: pytest.MonkeyPatch, +) -> None: + team: Final = LiteLLM_TeamTable( + team_id="real-team", + models=["other-model"], + members_with_roles=[Member(user_id="selection-user", role="user")], ) - proxy_server_module = types.ModuleType("proxy_server") - proxy_server_module.llm_router = router - monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) - - team = LiteLLM_TeamTable(team_id="real-team", models=["gpt-3.5-turbo"]) - - async def mock_get_team_object(*_args, **_kwargs): - return team - - monkeypatch.setattr( - "litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object + identity: Final = await _jwt_team_selection_identity( + monkeypatch, (team,), ("real-team",), ("real-team",), "selection-model" ) - - jwt_handler = JWTHandler() - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() - - with pytest.raises(HTTPException) as exc_info: - await JWTAuthManager.find_team_with_model_access( - team_ids={"real-team"}, - requested_model="gpt-4o-mini", - route="/chat/completions", - jwt_handler=jwt_handler, - prisma_client=None, - user_api_key_cache=None, - parent_otel_span=None, - proxy_logging_obj=None, + assert identity.team_id == "real-team" + with pytest.raises(ProxyException) as denied: + await _jwt_team_selection_identity( + monkeypatch, (team,), ("real-team",), ("real-team",), "selection-model", authorize=True ) - - assert exc_info.value.status_code == 403 - assert "No team has access to the requested model" in str(exc_info.value.detail) + assert denied.value.code == "403" @pytest.mark.asyncio @@ -5441,44 +5718,22 @@ async def test_resolve_team_from_header_accepts_db_teams_provisionally_under_fal @pytest.mark.asyncio -async def test_find_team_with_model_access_defers_no_team_403_under_db_fallback(): - """find_team_with_model_access raises the early "no teams in token" 403 when - enforcement is on, but defers (returns no team) so auth_builder's DB fallback - can run when fallback_to_db_teams is enabled.""" - jwt_handler = JWTHandler() - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( - enforce_team_based_model_access=True, - fallback_to_db_teams=False, - ) - with pytest.raises(HTTPException) as exc_info: - await JWTAuthManager.find_team_with_model_access( - team_ids=set(), - requested_model="gpt-4", - route="/chat/completions", - jwt_handler=jwt_handler, - prisma_client=None, - user_api_key_cache=MagicMock(), - parent_otel_span=None, - proxy_logging_obj=MagicMock(), +@pytest.mark.parametrize("fallback", (False, True)) +async def test_auth_builder_preserves_teamless_enforcement_denial( + monkeypatch: pytest.MonkeyPatch, + fallback: bool, +) -> None: + with pytest.raises(HTTPException) as builder_denied: + await _jwt_team_selection_identity( + monkeypatch, (), (), (), "selection-model", fallback_to_db_teams=fallback ) - assert exc_info.value.status_code == 403 - - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( - enforce_team_based_model_access=True, - fallback_to_db_teams=True, - ) - team_id, team_object = await JWTAuthManager.find_team_with_model_access( - team_ids=set(), - requested_model="gpt-4", - route="/chat/completions", - jwt_handler=jwt_handler, - prisma_client=None, - user_api_key_cache=MagicMock(), - parent_otel_span=None, - proxy_logging_obj=MagicMock(), - ) - assert team_id is None - assert team_object is None + assert builder_denied.value.status_code == 403 + with pytest.raises(ProxyException) as denied: + await _jwt_team_selection_identity( + monkeypatch, (), (), (), "selection-model", authorize=True, fallback_to_db_teams=fallback + ) + assert denied.value.code == "403" + assert ("not a member of any team" if fallback else "No teams found in token") in denied.value.message def _db_fallback_handler(litellm_jwtauth: LiteLLM_JWTAuth | None = None) -> JWTHandler: @@ -6265,82 +6520,23 @@ async def test_auth_builder_db_fallback_does_not_validate_rbac_team_against_db_m @pytest.mark.asyncio -async def test_resolve_db_team_fallback_distinguishes_no_membership_vs_model_denied(): - """When enforce_team_based_model_access is on, a user with no DB memberships - and a user with memberships that all fail the model-access check must surface - different 403s; collapsing both into the no-membership message hides the real - cause and diverges from find_team_with_model_access's claim-based message.""" - membership_user = LiteLLM_UserTable( - user_id="u_no_model", - user_role=LitellmUserRoles.INTERNAL_USER, - teams=["only_team"], +async def test_resolve_db_team_fallback_distinguishes_no_membership_vs_model_denied( + monkeypatch: pytest.MonkeyPatch, +) -> None: + team: Final = LiteLLM_TeamTable( + team_id="only-team", + models=["other-model"], + members_with_roles=[Member(user_id="selection-user", role="user")], ) - no_membership_user = LiteLLM_UserTable( - user_id="u_empty", - user_role=LitellmUserRoles.INTERNAL_USER, - teams=[], - ) - - async def fake_get_team(team_id, **kwargs): - return LiteLLM_TeamTable(team_id=team_id, models=["other"]) - - async def fake_can_access(model, team_object, llm_router, team_model_aliases=None): - raise ProxyException( - message="team not allowed to access model", - type=ProxyErrorTypes.team_model_access_denied, - param="model", - code=403, - ) - - with ( - patch( - "litellm.proxy.auth.handle_jwt.get_team_object", - new_callable=AsyncMock, - side_effect=fake_get_team, - ), - patch( - "litellm.proxy.auth.handle_jwt.can_team_access_model", - new_callable=AsyncMock, - side_effect=fake_can_access, - ), - ): - with pytest.raises(HTTPException) as model_denied: - await JWTAuthManager._resolve_db_team_fallback( - user_object=membership_user, - user_id=None, - requested_model="gpt-4", - route="/chat/completions", - jwt_handler=_db_fallback_handler(), - enforce_team_based_model_access=True, - team_id_upsert=False, - prisma_client=None, - user_api_key_cache=MagicMock(), - parent_otel_span=None, - proxy_logging_obj=MagicMock(), - ) - - with pytest.raises(HTTPException) as no_member: - await JWTAuthManager._resolve_db_team_fallback( - user_object=no_membership_user, - user_id=None, - requested_model="gpt-4", - route="/chat/completions", - jwt_handler=_db_fallback_handler(), - enforce_team_based_model_access=True, - team_id_upsert=False, - prisma_client=None, - user_api_key_cache=MagicMock(), - parent_otel_span=None, - proxy_logging_obj=MagicMock(), - ) - - assert model_denied.value.status_code == 403 - assert "requested model" in model_denied.value.detail - assert "gpt-4" in model_denied.value.detail - assert "only_team" not in model_denied.value.detail - - assert no_member.value.status_code == 403 - assert "not a member of any team" in no_member.value.detail + identity: Final = await _jwt_team_selection_identity(monkeypatch, (team,), (), ("only-team",), "selection-model") + assert identity.team_id == "only-team" + with pytest.raises(ProxyException) as model_denied: + await _jwt_team_selection_identity(monkeypatch, (team,), (), ("only-team",), "selection-model", authorize=True) + assert model_denied.value.code == "403" + assert "not a member of any team" not in model_denied.value.message + with pytest.raises(ProxyException, match="not a member of any team") as no_member: + await _jwt_team_selection_identity(monkeypatch, (), (), (), "selection-model", authorize=True) + assert no_member.value.code == "403" @pytest.mark.asyncio diff --git a/tests/unit/proxy/auth/test_user_api_key_auth.py b/tests/unit/proxy/auth/test_user_api_key_auth.py index 19f9d0dcc73..ce4a844550b 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -2,27 +2,29 @@ ## Unit tests for user_api_key_auth helper functions -import litellm.proxy -import litellm.proxy.proxy_server - +import json +import logging from typing import Dict, Final, List, Mapping, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fastapi import HTTPException, WebSocket, status from starlette.datastructures import URL +from starlette.requests import Request as StarletteRequest from starlette.types import Message -from litellm._logging import verbose_proxy_logger -import json -import logging -import litellm -from litellm.proxy.auth.user_api_key_auth import ( - user_api_key_auth, - UserAPIKeyAuth, - get_api_key_from_custom_header, -) -from fastapi import WebSocket, HTTPException, status -from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles +import litellm +import litellm.proxy +import litellm.proxy.proxy_server +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_AccessGroupTable, LiteLLM_UserTable, LitellmUserRoles, ProxyException +from litellm.proxy.auth.user_api_key_auth import ( + UserAPIKeyAuth, + _authorize_authenticated_request, + get_api_key_from_custom_header, + user_api_key_auth, +) +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from tests._master_key import MASTER_KEY @@ -33,6 +35,146 @@ class Request: self.headers: Dict[str, str] = {} +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model,fallbacks,allowed", (("model-a", [], True), ("model-b", [], False), ("model-a", ["model-b"], False)) +) +async def test_virtual_key_model_policy_on_first_lookup_and_cache_hit( + monkeypatch: pytest.MonkeyPatch, model: str, fallbacks: list[str], allowed: bool +) -> None: + server: Final = litellm.proxy.proxy_server + cache: Final = UserApiKeyCache() + key: Final = "sk-central-model-check" + stored: Final = UserAPIKeyAuth(token=server.hash_token(key), models=["model-a"]) + database: Final = MagicMock(get_data=AsyncMock(side_effect=(stored, AssertionError("unexpected cache miss")))) + monkeypatch.setattr(server, "master_key", MASTER_KEY) + monkeypatch.setattr(server, "general_settings", {}) + monkeypatch.setattr(server, "user_api_key_cache", cache) + monkeypatch.setattr(server, "prisma_client", database) + monkeypatch.setattr(server, "llm_router", None) + monkeypatch.setattr(server, "user_custom_auth", None) + for _ in range(2): + request: Final = StarletteRequest( + {"type": "http", "method": "POST", "path": "/chat/completions", "headers": []}, + receive=AsyncMock( + return_value={ + "type": "http.request", + "body": json.dumps({"model": model, "fallbacks": fallbacks}).encode(), + "more_body": False, + } + ), + ) + if allowed: + identity: Final = await user_api_key_auth(request=request, api_key=f"Bearer {key}") + assert identity.models == ["model-a"] + else: + with pytest.raises(ProxyException) as denied: + await user_api_key_auth(request=request, api_key=f"Bearer {key}") + assert denied.value.code == "403" + assert "model-b" in denied.value.message + + +@pytest.mark.asyncio +async def test_scim_deactivation_precedes_central_key_model_denial(monkeypatch: pytest.MonkeyPatch) -> None: + server: Final = litellm.proxy.proxy_server + key: Final = "sk-deactivated-owner" + cache: Final = UserApiKeyCache() + stored: Final = UserAPIKeyAuth(token=server.hash_token(key), models=["model-a"], user_id="inactive-owner") + cache.set_cache( + "inactive-owner", + LiteLLM_UserTable(user_id="inactive-owner", metadata={"scim_active": False}), + model_type=LiteLLM_UserTable, + ) + monkeypatch.setattr(server, "master_key", MASTER_KEY) + monkeypatch.setattr(server, "general_settings", {}) + monkeypatch.setattr(server, "user_api_key_cache", cache) + monkeypatch.setattr(server, "prisma_client", MagicMock(get_data=AsyncMock(return_value=stored))) + monkeypatch.setattr(server, "llm_router", None) + monkeypatch.setattr(server, "user_custom_auth", None) + request: Final = StarletteRequest( + {"type": "http", "method": "POST", "path": "/chat/completions", "headers": []}, + receive=AsyncMock(return_value={"type": "http.request", "body": b'{"model":"model-b"}'}), + ) + with pytest.raises(ProxyException) as denied: + await user_api_key_auth(request=request, api_key=f"Bearer {key}") + assert denied.value.code == "401" + assert "deactivated via SCIM" in denied.value.message + + +@pytest.mark.asyncio +async def test_restricted_key_model_list_denial_names_the_denied_model( + monkeypatch: pytest.MonkeyPatch, +) -> None: + server: Final = litellm.proxy.proxy_server + cache: Final = UserApiKeyCache() + group: Final = LiteLLM_AccessGroupTable( + access_group_id="science-id", access_group_name="science", access_model_names=["model-a"] + ) + cache.set_cache("access_group_rule_name:science", group, model_type=LiteLLM_AccessGroupTable) + cache.set_cache("access_group_id:science-id", group, model_type=LiteLLM_AccessGroupTable) + monkeypatch.setattr(server, "user_api_key_cache", cache) + monkeypatch.setattr(server, "prisma_client", MagicMock()) + monkeypatch.setattr( + server, + "general_settings", + {"access_group_rules": [{"attribute": "user_id", "equals": "science", "access_group": "science"}]}, + ) + monkeypatch.setattr(server, "llm_router", None) + identity: Final = UserAPIKeyAuth(user_id="science", models=["model-a"]) + request: Final = StarletteRequest({"type": "http", "method": "POST", "path": "/chat/completions", "headers": []}) + with pytest.raises(ProxyException) as denied: + await _authorize_authenticated_request( + identity, request, {"model": ["model-a", "model-b"]}, "/chat/completions", "synthetic", + force_virtual_key_checks=True, + ) + assert denied.value.code == "403" + assert "model-b" in denied.value.message + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model,enabled,allowed", + (("group-model", True, True), ("outside", True, False), ("group-model", False, False), ("explicit", False, True)), +) +async def test_user_access_group_rules_are_enforced_at_central_authorization( + monkeypatch: pytest.MonkeyPatch, model: str, enabled: bool, allowed: bool +) -> None: + server: Final = litellm.proxy.proxy_server + cache: Final = UserApiKeyCache() + group: Final = LiteLLM_AccessGroupTable( + access_group_id="science-id", access_group_name="science", access_model_names=["group-model"] + ) + cache.set_cache("access_group_rule_name:science", group, model_type=LiteLLM_AccessGroupTable) + cache.set_cache("access_group_id:science-id", group, model_type=LiteLLM_AccessGroupTable) + cache.set_cache("science", LiteLLM_UserTable(user_id="science"), model_type=LiteLLM_UserTable) + monkeypatch.setattr( + server, + "general_settings", + {"access_group_rules": [{"attribute": "user_id", "equals": "science", "access_group": "science"}]} + if enabled + else {}, + ) + monkeypatch.setattr(server, "user_api_key_cache", cache) + monkeypatch.setattr(server, "prisma_client", MagicMock()) + monkeypatch.setattr(server, "llm_router", None) + monkeypatch.setattr(server, "user_custom_auth", None) + principal: Final = UserAPIKeyAuth(models=["explicit"], user_id="science") + request: Final = StarletteRequest({"type": "http", "method": "POST", "path": "/chat/completions", "headers": []}) + if allowed: + assert ( + await _authorize_authenticated_request( + principal, request, {"model": model}, "/chat/completions", "synthetic", force_virtual_key_checks=True + ) + is None + ) + assert principal.access_group_ids == (["science-id"] if enabled else None) + else: + with pytest.raises(ProxyException, match="model"): + await _authorize_authenticated_request( + principal, request, {"model": model}, "/chat/completions", "synthetic", force_virtual_key_checks=True + ) + + @pytest.mark.parametrize( "allowed_ips, client_ip, expected_result", [ 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 70208bec0e5..1489afc02af 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 @@ -87,6 +87,7 @@ from litellm.proxy.auth.user_api_key_auth import ( _reserve_budget_after_common_checks, _route_requires_auth_despite_public, _routing_selector_matches_claim, + _authorize_authenticated_request, run_centralized_common_checks, _run_post_custom_auth_checks, user_api_key_auth_builder, @@ -446,29 +447,47 @@ async def test_should_not_reuse_cached_key_object_for_request_state(): assert second_request_key.request_route is None -@pytest.mark.asyncio -async def test_custom_auth_does_not_enforce_key_model_access_by_default(): - valid_token = UserAPIKeyAuth(token="test_token", models=["gpt-4o-mini"]) - request_data = {"model": "gpt-4o"} - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - new_callable=AsyncMock, - ) as mock_can_key, - patch( - "litellm.proxy.proxy_server.general_settings", - {}, - ), - ): - await _run_post_custom_auth_checks( - valid_token=valid_token, - request=None, - request_data=request_data, - route="/v1/chat/completions", - parent_otel_span=None, +async def _custom_auth_model_request( + monkeypatch: pytest.MonkeyPatch, requested_model: str, *, opt_in: bool, from_header: bool = False +) -> UserAPIKeyAuth: + async def custom_auth(request: Request, api_key: str) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + models=["allowed-model"], + user_role=LitellmUserRoles.INTERNAL_USER, + request_route=request.url.path, + api_key=api_key, ) - mock_can_key.assert_not_awaited() + + server: Final = litellm.proxy.proxy_server + monkeypatch.setattr(litellm, "enable_post_custom_auth_checks", True, raising=False) + monkeypatch.setattr(server, "user_custom_auth", custom_auth) + monkeypatch.setattr(server, "general_settings", {"custom_auth_run_common_checks": opt_in}) + monkeypatch.setattr(server, "prisma_client", None) + monkeypatch.setattr(server, "llm_router", None) + monkeypatch.setattr(server, "llm_model_list", None) + monkeypatch.setattr(server, "user_api_key_cache", UserApiKeyCache()) + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/files" if from_header else "/v1/chat/completions", + "headers": [(b"x-litellm-model", requested_model.encode())] if from_header else [], + }, + receive=AsyncMock( + return_value={ + "type": "http.request", + "more_body": False, + "body": json.dumps({} if from_header else {"model": requested_model}).encode(), + } + ), + ) + return await user_api_key_auth(request=request, api_key="Bearer external-token") + + +@pytest.mark.asyncio +async def test_custom_auth_does_not_enforce_key_model_access_by_default(monkeypatch: pytest.MonkeyPatch): + identity: Final = await _custom_auth_model_request(monkeypatch, "restricted-model", opt_in=False) + assert identity.models == ["allowed-model"] @pytest.mark.asyncio @@ -492,99 +511,30 @@ async def test_post_custom_auth_expired_key_returns_unauthorized(): @pytest.mark.asyncio -async def test_custom_auth_honors_key_level_model_access_restriction_allowed_with_opt_in(): - valid_token = UserAPIKeyAuth(token="test_token", models=["gpt-4o-mini"]) - request_data = {"model": "gpt-4o-mini"} - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - new_callable=AsyncMock, - ) as mock_can_key, - patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock), - patch( - "litellm.proxy.proxy_server.general_settings", - {"custom_auth_run_common_checks": True}, - ), - ): - await _run_post_custom_auth_checks( - valid_token=valid_token, - request=None, - request_data=request_data, - route="/v1/chat/completions", - parent_otel_span=None, - ) - mock_can_key.assert_awaited_once_with( - model="gpt-4o-mini", - llm_model_list=ANY, - valid_token=valid_token, - llm_router=ANY, - ) +async def test_custom_auth_honors_key_level_model_access_restriction_allowed_with_opt_in( + monkeypatch: pytest.MonkeyPatch, +): + identity: Final = await _custom_auth_model_request(monkeypatch, "allowed-model", opt_in=True) + assert identity.models == ["allowed-model"] @pytest.mark.asyncio -async def test_custom_auth_enforces_key_model_access_from_file_route_header_with_opt_in(): - valid_token = UserAPIKeyAuth(token="test_token", models=["allowed-model"]) - request = _RoutingRequest(headers={"x-litellm-model": "restricted-model"}) - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - new_callable=AsyncMock, - ) as mock_can_key, - patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock), - patch( - "litellm.proxy.proxy_server.general_settings", - {"custom_auth_run_common_checks": True}, - ), - ): - await _run_post_custom_auth_checks( - valid_token=valid_token, - request=request, - request_data={}, - route="/v1/files", - parent_otel_span=None, - ) - mock_can_key.assert_awaited_once_with( - model="restricted-model", - llm_model_list=ANY, - valid_token=valid_token, - llm_router=ANY, - ) +async def test_custom_auth_enforces_key_model_access_from_file_route_header_with_opt_in( + monkeypatch: pytest.MonkeyPatch, +): + with pytest.raises(ProxyException) as denied: + await _custom_auth_model_request(monkeypatch, "restricted-model", opt_in=True, from_header=True) + assert denied.value.code == "403" @pytest.mark.asyncio -async def test_custom_auth_honors_key_level_model_access_restriction_denied_with_opt_in(): - valid_token = UserAPIKeyAuth(token="test_token", models=["gpt-4o-mini"]) - request_data = {"model": "gpt-4o"} - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - new_callable=AsyncMock, - ) as mock_can_key, - patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock), - patch( - "litellm.proxy.proxy_server.general_settings", - {"custom_auth_run_common_checks": True}, - ), - ): - mock_can_key.side_effect = ProxyException( - message="Key not allowed to access model", - type=ProxyErrorTypes.key_model_access_denied, - param="model", - code=401, - ) - with pytest.raises(ProxyException) as exc: - await _run_post_custom_auth_checks( - valid_token=valid_token, - request=None, - request_data=request_data, - route="/v1/chat/completions", - parent_otel_span=None, - ) - - assert exc.value.type == ProxyErrorTypes.key_model_access_denied +async def test_custom_auth_honors_key_level_model_access_restriction_denied_with_opt_in( + monkeypatch: pytest.MonkeyPatch, +): + with pytest.raises(ProxyException) as denied: + await _custom_auth_model_request(monkeypatch, "restricted-model", opt_in=True) + assert denied.value.type == ProxyErrorTypes.key_model_access_denied + assert denied.value.code == "403" def _proxy_server_attrs_for_custom_auth(*, user_custom_auth): @@ -1830,8 +1780,7 @@ async def test_db_virtual_key_auth_sets_via_virtual_key_marker(): @pytest.mark.asyncio @pytest.mark.parametrize("model_allowed", [True, False]) -async def test_auth_prefetches_referenced_objects_only_after_the_key_may_call_the_model(model_allowed): - """A request denied by the key's model list must not pay for the team/user/org MGET or DB join.""" +async def test_auth_prefetches_referenced_objects_before_central_model_authorization(model_allowed): from fastapi import Request from starlette.datastructures import URL @@ -1841,9 +1790,9 @@ async def test_auth_prefetches_referenced_objects_only_after_the_key_may_call_th api_key = "sk-prefetch-order-test" valid_token = UserAPIKeyAuth(api_key=api_key, token=hash_token(api_key), user_id="u1", team_id="t1") - mock_cache = AsyncMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.delete_cache = MagicMock() + mock_cache = UserApiKeyCache() + mock_cache.set_cache("team_id:t1", LiteLLM_TeamTable(team_id="t1")) + mock_cache.set_cache("team_membership:u1:t1", LiteLLM_TeamMembership(user_id="u1", team_id="t1")) mock_proxy_logging_obj = MagicMock() mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) @@ -1873,7 +1822,7 @@ async def test_auth_prefetches_referenced_objects_only_after_the_key_may_call_th try: for attr, val in _attrs_to_set.items(): setattr(_proxy_server_mod, attr, val) - request = Request(scope={"type": "http"}) + request = Request(scope={"type": "http", "method": "POST", "path": "/chat/completions"}) request._url = URL(url="/chat/completions") with ( @@ -1894,7 +1843,7 @@ async def test_auth_prefetches_referenced_objects_only_after_the_key_may_call_th "litellm.proxy.auth.user_api_key_auth.get_user_object", new_callable=AsyncMock, return_value=None ), ): - call = user_api_key_auth_builder( + resolved: Final = await user_api_key_auth_builder( request=request, api_key=f"Bearer {api_key}", azure_api_key_header="", @@ -1903,15 +1852,18 @@ async def test_auth_prefetches_referenced_objects_only_after_the_key_may_call_th azure_apim_header=None, request_data={"model": "gpt-4o"}, ) + assert isinstance(resolved, UserAPIKeyAuth) + mock_prefetch.assert_awaited_once() + assert mock_prefetch.await_args.kwargs["refs"].team_id == "t1" + call: Final = _authorize_authenticated_request( + resolved, request, {"model": "gpt-4o"}, "/chat/completions", api_key + ) if model_allowed: - assert isinstance(await call, UserAPIKeyAuth) - mock_prefetch.assert_awaited_once() - assert mock_prefetch.await_args.kwargs["refs"].team_id == "t1" + await call else: with pytest.raises(ProxyException) as exc: await call assert exc.value.type == ProxyErrorTypes.key_model_access_denied - mock_prefetch.assert_not_awaited() finally: for attr, val in _original_values.items(): setattr(_proxy_server_mod, attr, val) @@ -2601,6 +2553,48 @@ async def test_auto_register_map_existing_key_uses_the_issuers_own_user_field_ov assert generate_key.await_count == (0 if expect_reuse else 1) +def _validated_jwt_admission_cache(*, org_id: str | None = None, end_user_id: str | None = None) -> UserApiKeyCache: + cache: Final = UserApiKeyCache() + cache.set_cache( + "validated-user", + LiteLLM_UserTable( + user_id="validated-user", + user_email="validated@example.com", + user_role="internal_user", + teams=["validated-team"], + ), + model_type=LiteLLM_UserTable, + ) + cache.set_cache( + "team_id:validated-team", + LiteLLM_TeamTableCachedObj( + team_id="validated-team", + models=["gpt-4o-mini", "gpt-5.6"], + members_with_roles=[{"user_id": "validated-user", "role": "user"}], + organization_id=org_id, + ), + model_type=LiteLLM_TeamTableCachedObj, + ) + cache.set_cache( + team_membership_reservation_cache_key(user_id="validated-user", team_id="validated-team"), + LiteLLM_TeamMembership(user_id="validated-user", team_id="validated-team"), + model_type=LiteLLM_TeamMembership, + ) + if org_id is not None: + org: Final = LiteLLM_OrganizationTable( + organization_id=org_id, budget_id="org-budget", created_by="admin", updated_by="admin" + ) + cache.set_cache(f"org_id:{org_id}", org, model_type=LiteLLM_OrganizationTable) + cache.set_cache(f"org_id:{org_id}:with_budget", org, model_type=LiteLLM_OrganizationTable) + if end_user_id is not None: + cache.set_cache( + f"end_user_id:{end_user_id}", + LiteLLM_EndUserTable(user_id=end_user_id, blocked=False), + model_type=LiteLLM_EndUserTable, + ) + return cache + + @pytest.mark.asyncio @pytest.mark.parametrize( ("map_existing_key", "master_key", "reused_key_models", "expect_denied"), @@ -2615,7 +2609,7 @@ async def test_auto_register_map_existing_key_first_request_runs_key_checks( map_existing_key: bool, master_key: str | None, reused_key_models: list[str], expect_denied: bool ) -> None: jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" - user_api_key_cache = DualCache() + user_api_key_cache: Final = _validated_jwt_admission_cache() prisma_client = MagicMock() jwt_handler = MagicMock() jwt_handler.is_jwt.return_value = True @@ -2648,12 +2642,21 @@ async def test_auto_register_map_existing_key_first_request_runs_key_checks( "jwt_claims": {"sub": "user1"}, } - mock_request = MagicMock() - mock_request.url.path = "/v1/chat/completions" - mock_request.method = "POST" - mock_request.headers = {"authorization": f"Bearer {jwt_token}"} - mock_request.query_params = {} - mock_request.state = SimpleNamespace() + mock_request: Final = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/chat/completions", + "headers": [(b"authorization", f"Bearer {jwt_token}".encode())], + }, + receive=AsyncMock( + return_value={ + "type": "http.request", + "more_body": False, + "body": json.dumps({"model": "gpt-4o-mini"}).encode(), + } + ), + ) with ( patch("litellm.proxy.proxy_server.general_settings", {"enable_jwt_auth": True}), @@ -2686,14 +2689,13 @@ async def test_auto_register_map_existing_key_first_request_runs_key_checks( return_value=reused_key, ), ): - call = user_api_key_auth_builder( + call = user_api_key_auth( request=mock_request, api_key=jwt_token, azure_api_key_header="", anthropic_api_key_header=None, google_ai_studio_api_key_header=None, azure_apim_header=None, - request_data={"model": "gpt-4o-mini"}, ) if expect_denied: with pytest.raises(ProxyException, match="not available for this API key"): @@ -2719,7 +2721,7 @@ async def test_auto_register_first_request_propagates_user_email(active: bool) - """ jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" general_settings = {"enable_jwt_auth": True} - user_api_key_cache = DualCache() + user_api_key_cache: Final = _validated_jwt_admission_cache() prisma_client = MagicMock() jwt_handler = MagicMock() jwt_handler.is_jwt.return_value = True @@ -2952,7 +2954,10 @@ async def test_jwt_auto_register_forwards_bound_agent_id(): id auth_builder resolved must reach the key creation, not be dropped when valid_token is swapped for the freshly registered key.""" jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" - user_api_key_cache = DualCache() + user_api_key_cache: Final = _validated_jwt_admission_cache() + prisma_client: Final = MagicMock() + prisma_client.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=None) + prisma_client.writer_db.litellm_retiredagent.find_unique = AsyncMock(return_value=None) jwt_handler = MagicMock() jwt_handler.is_jwt.return_value = True jwt_handler.auth_jwt = AsyncMock(return_value={"sub": "user1", "appid": "2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21"}) @@ -3001,7 +3006,7 @@ async def test_jwt_auto_register_forwards_bound_agent_id(): general_settings={"enable_jwt_auth": True}, premium_user=True, master_key="sk-master", - prisma_client=MagicMock(), + prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=MagicMock(), jwt_handler=jwt_handler, @@ -3274,7 +3279,7 @@ class TestJWTOAuth2Coexistence: async def test_auto_register_passes_validated_org_context_to_generated_key(self): jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" general_settings = {"enable_jwt_auth": True} - user_api_key_cache = DualCache() + user_api_key_cache: Final = _validated_jwt_admission_cache(org_id="validated-org", end_user_id="validated-end-user") prisma_client = MagicMock() jwt_handler = MagicMock() jwt_handler.is_jwt.return_value = True diff --git a/tests/unit/proxy/mcp_server/test_idp_token_exchange.py b/tests/unit/proxy/mcp_server/test_idp_token_exchange.py index 036ee4a332b..75b0a5cd308 100644 --- a/tests/unit/proxy/mcp_server/test_idp_token_exchange.py +++ b/tests/unit/proxy/mcp_server/test_idp_token_exchange.py @@ -1,20 +1,34 @@ +import json import logging from typing import Final +import jwt import pytest -from fastapi import HTTPException +from cryptography.hazmat.primitives.asymmetric import rsa +from fastapi import HTTPException, Request from prisma.engine.errors import BinaryNotFoundError from prisma.errors import DataError from litellm.caching.caching import DualCache -from litellm.proxy._types import JWTIssuerConfig, LiteLLM_JWTAuth, ProxyException +from litellm.proxy import proxy_server +from litellm.proxy._types import ( + JWTIssuerConfig, + LiteLLM_JWTAuth, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LiteLLM_UserTable, + Member, + ProxyException, +) from litellm.proxy.auth.handle_jwt import JWKSUnreachableError, JWTHandler, jwks_unavailable_exception +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.mcp_server.gateway_dcr_flow import SubjectIdentity, SubjectTokenRefusal from litellm.proxy.mcp_server.idp_token_exchange import ( REJECTED_SUBJECT_TOKEN, SUBJECT_TOKEN_CHECK_FAULTED, SUBJECT_TOKEN_CHECK_UNAVAILABLE, TokenExchangePrerequisites, + exchange_idp_subject_token, identity_from_subject_token, token_exchange_available, ) @@ -246,3 +260,71 @@ async def test_agent_token_cannot_be_exchanged_for_a_user_identity(user_id: str assert isinstance(result, SubjectTokenRefusal) assert result.error == "invalid_request" assert "direct JWT authentication" in result.description + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("enforce_team", "team_id", "expected"), + ( + (True, "exchange-team", SubjectIdentity(user_id="exchange-user", team_id="exchange-team")), + (False, None, SubjectIdentity(user_id="exchange-user", team_id=None)), + (True, None, SubjectTokenRefusal(error="invalid_request", description=REJECTED_SUBJECT_TOKEN)), + ), + ids=("native-team", "teamless-without-enforcement", "teamless-under-enforcement"), +) +async def test_subject_exchange_applies_team_admission_before_returning_an_identity( + monkeypatch: pytest.MonkeyPatch, + enforce_team: bool, + team_id: str | None, + expected: SubjectIdentity | SubjectTokenRefusal, +) -> None: + cache: Final = UserApiKeyCache() + database: Final = object() + handler: Final = JWTHandler() + handler.update_environment( + prisma_client=database, + user_api_key_cache=cache, + litellm_jwtauth=LiteLLM_JWTAuth( + user_id_jwt_field="sub", + team_ids_jwt_field="groups", + enforce_team_based_model_access=enforce_team, + team_allowed_routes=["/token"], + ), + ) + signing_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + jwks_url: Final = "https://identity.invalid/exchange/jwks" + cache.set_cache( + f"litellm_jwt_auth_keys_{jwks_url}", + [json.loads(jwt.algorithms.RSAAlgorithm.to_jwk(signing_key.public_key()))], + ) + cache.set_cache( + "exchange-user", + LiteLLM_UserTable(user_id="exchange-user", user_role="internal_user", teams=[team_id] if team_id else []), + model_type=LiteLLM_UserTable, + ) + if team_id is not None: + cache.set_cache( + f"team_id:{team_id}", + LiteLLM_TeamTable(team_id=team_id, members_with_roles=[Member(user_id="exchange-user", role="user")]), + model_type=LiteLLM_TeamTable, + ) + cache.set_cache( + f"team_membership:exchange-user:{team_id}", + LiteLLM_TeamMembership(user_id="exchange-user", team_id=team_id), + model_type=LiteLLM_TeamMembership, + ) + monkeypatch.setenv("JWT_PUBLIC_KEY_URL", jwks_url) + monkeypatch.delenv("JWT_ISSUER", raising=False) + monkeypatch.delenv("JWT_AUDIENCE", raising=False) + monkeypatch.setattr(proxy_server, "jwt_handler", handler) + monkeypatch.setattr(proxy_server, "prisma_client", database) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "general_settings", {"enable_jwt_auth": True}) + monkeypatch.setattr(proxy_server, "premium_user", True) + monkeypatch.setattr(proxy_server, "llm_router", None) + monkeypatch.setattr(proxy_server, "llm_model_list", None) + token: Final = jwt.encode( + {"sub": "exchange-user", "groups": [team_id] if team_id else []}, signing_key, algorithm="RS256" + ) + request: Final = Request({"type": "http", "method": "POST", "path": "/token", "headers": []}) + assert await exchange_idp_subject_token(token, request) == expected diff --git a/tests/unit/proxy/pass_through_endpoints/test_codex_passthrough_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_codex_passthrough_endpoints.py index 1a4751b4598..322752782fd 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_codex_passthrough_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_codex_passthrough_endpoints.py @@ -11,6 +11,7 @@ import jwt import pytest import respx from cryptography.hazmat.primitives.asymmetric import rsa +from fastapi import Request from fastapi.testclient import TestClient from jwt.algorithms import RSAAlgorithm @@ -704,10 +705,23 @@ def _key_with_alias_to_budgeted_target(monkeypatch: pytest.MonkeyPatch, target_b return key +@pytest.mark.parametrize("custom_auth", (False, True), ids=("virtual-key", "custom-auth")) def test_native_route_refuses_a_saved_alias_when_its_target_is_over_budget( - client: TestClient, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch + client: TestClient, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, custom_auth: bool ) -> None: key: Final = _key_with_alias_to_budgeted_target(monkeypatch, target_budget=0.0) + if custom_auth: + monkeypatch.setattr(litellm, "enable_post_custom_auth_checks", True, raising=False) + identity: Final = proxy_server.user_api_key_cache.get_cache(key=hash_token(key), model_type=UserAPIKeyAuth) + assert identity is not None + + async def authenticate(request: Request, api_key: str) -> UserAPIKeyAuth: + return identity + + monkeypatch.setattr(proxy_server, "user_custom_auth", authenticate) + monkeypatch.setattr( + proxy_server, "general_settings", {**proxy_server.general_settings, "custom_auth_run_common_checks": True} + ) _account_catalog(respx_mock) provider: Final = _open_weights_provider(respx_mock) response: Final = client.post( diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 24ad068e6df..2c23305d5c9 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -3691,12 +3691,19 @@ class TestResponsesWSFirstFrameModelAuth: "standard_logging_guardrail_information": [{"guardrail_cost": 0.2}, {"guardrail_cost": 0.4}], } - @pytest.mark.asyncio @pytest.mark.parametrize( "relay_outcome", ["still_relaying", "already_returned", "returns_despite_the_cancel", "crashes_despite_the_cancel"], ) - async def test_session_limit_frees_the_slot_at_once_and_never_cuts_the_failure_row(self, relay_outcome: str): + def test_session_limit_frees_the_slot_at_once_and_never_cuts_the_failure_row(self, relay_outcome: str) -> None: + loop: Final = AutojumpClockLoop() + try: + loop.run_until_complete(self._assert_session_limit_releases_slot_before_failure_row(relay_outcome)) + finally: + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.close() + + async def _assert_session_limit_releases_slot_before_failure_row(self, relay_outcome: str) -> None: from litellm.exceptions import RateLimitError from litellm.proxy.response_api_endpoints.endpoints import responses_websocket_endpoint diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 00c1b855033..6f8c49cc0f9 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -26239,6 +26239,17 @@ export interface components { /** Updated By */ updated_by?: string | null; }; + /** AccessGroupRule */ + AccessGroupRule: { + /** Access Group */ + access_group: string; + /** Attribute */ + attribute: string; + /** Contains */ + contains?: string | number | boolean | null; + /** Equals */ + equals?: string | number | boolean | null; + }; /** AccessGroupUpdateRequest */ AccessGroupUpdateRequest: { /** Access Agent Ids */ @@ -30808,6 +30819,8 @@ export interface components { * @description Documents all the fields supported by `general_settings` in config.yaml */ ConfigGeneralSettings: { + /** Access Group Rules */ + access_group_rules?: components["schemas"]["AccessGroupRule"][] | null; /** * Admission Queue Timeout Seconds * @description maximum time a request waits for a worker slot diff --git a/uv.lock b/uv.lock index 8712b78bbd6..62187d596ec 100644 --- a/uv.lock +++ b/uv.lock @@ -1117,6 +1117,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/bf/0f/f897abe4ea0a8c408ae65c8c83bffab4936ad65d6032d4fb4cd35bbdc3ee/cachetools-7.1.1-py3-none-any.whl", hash = "sha256:0335cd7a0952d2b22327441fb0628139e234c565559eeb91a8a4ac7551c5353d", size = 16775, upload-time = "2026-05-03T20:00:27.857Z" }, ] +[[package]] +name = "casbin" +version = "1.43.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "simpleeval" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ad/df/ff2aa55cf0d7c14622ce4f9252cdc34c828c81d4213965d73207ac5434ae/casbin-1.43.0.tar.gz", hash = "sha256:d2e90ce8e72f912877851e94d37999f32c558c6ba7aba0437d483275262e86e0", size = 425727, upload-time = "2025-05-10T06:57:18.902Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/08/07/facef6abd81378e6153b757dc1848621675d971fbc88ebb5d182ebc1c37f/casbin-1.43.0-py3-none-any.whl", hash = "sha256:63a3d1228870250e859ccd94133fe478821093f71dd37f05e0baa0c6fea26623", size = 475059, upload-time = "2025-05-10T06:57:16.89Z" }, +] + [[package]] name = "certifi" version = "2026.10.19" @@ -5096,6 +5108,7 @@ proxy = [ { name = "backoff" }, { name = "backports-zstd", marker = "python_full_version < '3.14'" }, { name = "boto3" }, + { name = "casbin" }, { name = "cryptography" }, { name = "expression" }, { name = "fastapi" }, @@ -5312,6 +5325,7 @@ requires-dist = [ { name = "backports-zstd", marker = "python_full_version < '3.14' and extra == 'proxy'", specifier = ">=1.0.0,<2.0" }, { name = "boto3", specifier = ">=1.43.117,<2.0" }, { name = "boto3", marker = "extra == 'proxy'", specifier = ">=1.43.117,<2.0" }, + { name = "casbin", marker = "extra == 'proxy'", specifier = "==1.43.0" }, { name = "certifi", specifier = ">=2026.10.19" }, { name = "click", specifier = ">=8.0.0,<9.0" }, { name = "cryptography", marker = "extra == 'proxy'", specifier = ">=49.0.0,<51.0" }, @@ -10377,6 +10391,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/52/59/0782e51887ac6b07ffd1570e0364cf901ebc36345fea669969d2084baebb/simple_websocket-1.1.0-py3-none-any.whl", hash = "sha256:4af6069630a38ed6c561010f0e11a5bc0d4ca569b36306eb257cd9a192497c8c", size = 13842, upload-time = "2024-10-10T22:39:29.645Z" }, ] +[[package]] +name = "simpleeval" +version = "1.0.8" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/89/3e/25326f373871a34e78cac231c0d3c18a0a5c3a80aa4813b50b9f26a5d460/simpleeval-1.0.8.tar.gz", hash = "sha256:b417f869c11f2de034eb919820a825f35d12ce99f63a43971947ce73679eb9ba", size = 31846, upload-time = "2026-09-12T10:16:08.103Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/12/33/d1ecd9f4f5f272b36ad7b268e307ba0a59c109d8ad134a231cda780baac9/simpleeval-1.0.8-py3-none-any.whl", hash = "sha256:c9c4ff42879ebf0c05f6562cb81031869948916be343bcc950b81fcc081fe9ce", size = 18925, upload-time = "2026-09-12T10:16:06.862Z" }, +] + [[package]] name = "six" version = "1.17.0" From 9de20b6cf880931286a64362aaf7b5d1a0386216 Mon Sep 17 00:00:00 2001 From: Jeevan Babu Gotru <114808955+jeevan1098@users.noreply.github.com> Date: Mon, 9 Nov 2026 16:03:50 -0600 Subject: [PATCH 18/46] fix(minimax): forward reasoning_effort instead of dropping it (#49583) * fix(minimax): forward reasoning_effort instead of dropping it * fix(minimax): forward reasoning_effort instead of dropping it * fix(minimax): gate reasoning_effort on the cost map and map MiniMax-M3.1-Flash Forward reasoning_effort for every MiniMax model unless its cost map entry says it does not reason, so a MiniMax model newer than the map keeps the caller's effort instead of losing it. Add a minimax/MiniMax-M3.1-Flash entry with MiniMax's published pay-as-you-go pricing and its accepted effort levels (low, medium, high, xhigh, max), and a minimax-reasoning-baseline fallback rule that marks unmapped minimax/MiniMax-M* ids as reasoning-capable, so thinking is kept for them too. * refactor(minimax): gate reasoning_effort on supports_reasoning like thinking Forward reasoning_effort when the cost map, or the MiniMax M-series fallback rule, says the model supports reasoning, the same check thinking already uses. An unlisted MiniMax model outside the M series now drops it, with the existing drop warning. * refactor(minimax): build supported params in one expression and type the new tests Build the MiniMax supported params list without mutating it, give the new test helpers precise types, and assert invariants (both cost map files agree, prices are positive and ordered, every declared effort level is forwarded) instead of pinning MiniMax's prices and effort levels as literals. * fix(minimax): treat every unlisted minimax/ model as reasoning-capable Widen the MiniMax fallback rule from M-series ids to the whole minimax/ namespace, matching the W&B rule, so an unlisted MiniMax model keeps reasoning_effort and thinking. Exact cost map entries still win, so mapped speech, video and image models are unaffected. --------- Co-authored-by: ryan-crabbe-berri --- litellm/llms/minimax/chat/transformation.py | 23 +++-- ...odel_prices_and_context_window_backup.json | 31 ++++++ model_prices_and_context_window.json | 31 ++++++ .../test_fallback_generalizations.py | 29 ++++++ .../llms/minimax/chat/test_transformation.py | 99 ++++++++++++++++++- .../llms/test_cache_control_and_reasoning.py | 13 +++ 6 files changed, 211 insertions(+), 15 deletions(-) diff --git a/litellm/llms/minimax/chat/transformation.py b/litellm/llms/minimax/chat/transformation.py index f397bf7fbb9..84ed8a56430 100644 --- a/litellm/llms/minimax/chat/transformation.py +++ b/litellm/llms/minimax/chat/transformation.py @@ -11,6 +11,13 @@ from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam +def _supports_reasoning(model: str) -> bool: + try: + return litellm.supports_reasoning(model=model, custom_llm_provider="minimax") + except Exception: + return False + + class MinimaxChatConfig(OpenAIGPTConfig): """ MiniMax OpenAI configuration that extends OpenAIGPTConfig. @@ -81,16 +88,8 @@ class MinimaxChatConfig(OpenAIGPTConfig): def get_supported_openai_params(self, model: str) -> list: """ Get supported OpenAI parameters for MiniMax. - Adds reasoning_split and thinking to the list of supported params. + Adds reasoning_split, plus thinking and reasoning_effort when the cost map, or its + minimax/ fallback rule for unlisted models, says the model supports reasoning. """ - base_params: Final = super().get_supported_openai_params(model=model) - additional_params: Final = ["reasoning_split"] - - # Add thinking parameter if model supports reasoning - try: - if litellm.supports_reasoning(model=model, custom_llm_provider="minimax"): - additional_params.append("thinking") - except Exception: - pass - - return base_params + additional_params + reasoning_params: Final = ("thinking", "reasoning_effort") if _supports_reasoning(model) else () + return [*super().get_supported_openai_params(model=model), "reasoning_split", *reasoning_params] diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 600c195c2c7..e1cbb953c5b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -43107,6 +43107,29 @@ "max_input_tokens": 1000000, "max_output_tokens": 128000 }, + "minimax/MiniMax-M3.1-Flash": { + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "cache_read_input_token_cost": 1.5e-08, + "litellm_provider": "minimax", + "mode": "chat", + "source": "https://platform.minimax.io/docs/guides/pricing-paygo", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_none_reasoning_effort": false, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "reasoning_effort_levels": [ + "low", + "medium", + "high", + "xhigh", + "max" + ], + "supports_vision": true, + "max_input_tokens": 1000000 + }, "minimax/minimax-latest": { "alias_of": "minimax/MiniMax-M3" }, @@ -65211,6 +65234,14 @@ "supports_reasoning": true } }, + { + "name": "minimax-reasoning-baseline", + "pattern": "^minimax/", + "description": "Any MiniMax model id under the minimax/ namespace, anchored so only that provider's ids match. MiniMax ships reasoning models faster than this registry names them, so an id the map has not described yet is treated as reasoning-capable and keeps the caller's reasoning_effort and thinking instead of dropping them. Rules lose to exact entries, so mapped speech, video and image models are unaffected. Carries no mode and no pricing, so cost stays on the standard unpriced behavior.", + "model_info": { + "supports_reasoning": true + } + }, { "name": "openai-reasoning-family-baseline", "pattern": "^(?!.*search-api)(?:[a-z0-9_.-]+/)*(?:ft:)?(?:o[1-9]\\d*(?![a-z0-9])|gpt-[5-9](?:\\.\\d+)?(?![0-9.])|(?:gpt-\\d+(?:\\.\\d+)?(?:-[a-z0-9]+)*-)?(?:codex|deep-research|chat-latest)(?![a-z0-9]))", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 600c195c2c7..e1cbb953c5b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -43107,6 +43107,29 @@ "max_input_tokens": 1000000, "max_output_tokens": 128000 }, + "minimax/MiniMax-M3.1-Flash": { + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "cache_read_input_token_cost": 1.5e-08, + "litellm_provider": "minimax", + "mode": "chat", + "source": "https://platform.minimax.io/docs/guides/pricing-paygo", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_none_reasoning_effort": false, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "reasoning_effort_levels": [ + "low", + "medium", + "high", + "xhigh", + "max" + ], + "supports_vision": true, + "max_input_tokens": 1000000 + }, "minimax/minimax-latest": { "alias_of": "minimax/MiniMax-M3" }, @@ -65211,6 +65234,14 @@ "supports_reasoning": true } }, + { + "name": "minimax-reasoning-baseline", + "pattern": "^minimax/", + "description": "Any MiniMax model id under the minimax/ namespace, anchored so only that provider's ids match. MiniMax ships reasoning models faster than this registry names them, so an id the map has not described yet is treated as reasoning-capable and keeps the caller's reasoning_effort and thinking instead of dropping them. Rules lose to exact entries, so mapped speech, video and image models are unaffected. Carries no mode and no pricing, so cost stays on the standard unpriced behavior.", + "model_info": { + "supports_reasoning": true + } + }, { "name": "openai-reasoning-family-baseline", "pattern": "^(?!.*search-api)(?:[a-z0-9_.-]+/)*(?:ft:)?(?:o[1-9]\\d*(?![a-z0-9])|gpt-[5-9](?:\\.\\d+)?(?![0-9.])|(?:gpt-\\d+(?:\\.\\d+)?(?:-[a-z0-9]+)*-)?(?:codex|deep-research|chat-latest)(?![a-z0-9]))", diff --git a/tests/unit/litellm_core_utils/test_fallback_generalizations.py b/tests/unit/litellm_core_utils/test_fallback_generalizations.py index d8e56c0ae51..b22edaa171a 100644 --- a/tests/unit/litellm_core_utils/test_fallback_generalizations.py +++ b/tests/unit/litellm_core_utils/test_fallback_generalizations.py @@ -810,6 +810,35 @@ def test_shipped_wandb_rule_is_anchored_to_the_wandb_namespace(shipped_cost_map) assert matched is None or not matched.get("supports_reasoning"), foreign +@pytest.mark.parametrize("model", ["minimax/MiniMax-M3.2", "minimax/MiniMax-Text-02"]) +def test_shipped_rules_flag_unmapped_minimax_ids_as_reasoning(shipped_cost_map, model): + """MiniMax ships reasoning models faster than the registry names them. The rule + carries no mode and no pricing, so cost stays on the standard unpriced behavior.""" + assert model not in litellm.model_cost + + info = litellm.get_model_info(model, custom_llm_provider="minimax") + assert info["litellm_provider"] == "minimax" + assert info["supports_reasoning"] is True + assert info.get("mode") is None + assert not info.get("input_cost_per_token") + assert not info.get("output_cost_per_token") + + +def test_shipped_minimax_rule_is_anchored_to_the_minimax_namespace(shipped_cost_map): + """``^minimax/`` is anchored, so it cannot leak onto another provider's ids.""" + assert match_capability_generalizations("minimax/some-new-model") == {"supports_reasoning": True} + for foreign in ("openrouter/minimax/minimax-m3", "notminimax/MiniMax-M3", "together_ai/minimax/MiniMax-M3"): + matched = match_capability_generalizations(foreign) + assert matched is None or not matched.get("supports_reasoning"), foreign + + +@pytest.mark.parametrize("model", ["speech-02-hd", "MiniMax-H3"]) +def test_shipped_minimax_rule_leaves_mapped_entries_alone(shipped_cost_map, model): + assert f"minimax/{model}" in litellm.model_cost + assert match_fill_missing_generalizations(f"minimax/{model}", "minimax") is None + assert litellm.supports_reasoning(model=model, custom_llm_provider="minimax") is False + + def test_shipped_wandb_rule_keeps_reasoning_effort_on_an_unmapped_model(shipped_cost_map): """End to end through the provider config: the gate WandbConfig applies reads the rule, so reasoning_effort is advertised and survives get_optional_params rather than diff --git a/tests/unit/llms/minimax/chat/test_transformation.py b/tests/unit/llms/minimax/chat/test_transformation.py index 2645b2832aa..d289bdba9bf 100644 --- a/tests/unit/llms/minimax/chat/test_transformation.py +++ b/tests/unit/llms/minimax/chat/test_transformation.py @@ -2,10 +2,18 @@ Test MiniMax OpenAI-compatible API support """ +import json +from collections.abc import Mapping +from pathlib import Path +from typing import Final, get_args from unittest.mock import MagicMock, patch +import pytest +import respx + import litellm from litellm.llms.minimax.chat.transformation import MinimaxChatConfig +from litellm.types.llms.openai import REASONING_EFFORT def test_minimax_chat_config(): @@ -94,9 +102,7 @@ def test_minimax_provider_config_manager(): from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager - config = ProviderConfigManager.get_provider_chat_config( - model="MiniMax-M2.1", provider=LlmProviders.MINIMAX - ) + config = ProviderConfigManager.get_provider_chat_config(model="MiniMax-M2.1", provider=LlmProviders.MINIMAX) assert config is not None assert isinstance(config, MinimaxChatConfig) @@ -121,3 +127,90 @@ if __name__ == "__main__": print("✓ Provider config manager test passed") print("\n✅ All basic tests passed!") + + +@pytest.fixture +def minimax_request_mock( + local_model_cost_map: None, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> respx.Route: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + return respx_mock.post("https://api.minimax.io/v1/chat/completions").respond( + json={ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": "MiniMax-M3.1-Flash", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "Done"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 9, "completion_tokens": 12, "total_tokens": 21}, + } + ) + + +def _sent_body(model: str, request_mock: respx.Route, effort: str, drop_params: bool = True) -> Mapping[str, object]: + litellm.completion( + model=f"minimax/{model}", + messages=[{"role": "user", "content": "What is the weather in Oslo?"}], + api_key="fake-minimax-key", + reasoning_effort=effort, + drop_params=drop_params, + ) + body: Final[Mapping[str, object]] = json.loads(request_mock.calls.last.request.content) + return body + + +def _flash_effort_levels() -> tuple[str, ...]: + levels: Final = litellm.get_model_info("MiniMax-M3.1-Flash", custom_llm_provider="minimax")[ + "reasoning_effort_levels" + ] + assert levels + return tuple(levels) + + +def test_minimax_flash_forwards_every_declared_reasoning_effort(minimax_request_mock: respx.Route) -> None: + """Issue #49574: reasoning_effort sent to MiniMax-M3.1-Flash was stripped before the upstream call.""" + levels: Final = _flash_effort_levels() + sent: Final = tuple( + _sent_body("MiniMax-M3.1-Flash", minimax_request_mock, effort)["reasoning_effort"] for effort in levels + ) + assert sent == levels + assert minimax_request_mock.call_count == len(levels) + + +def test_minimax_flash_cost_map_entry_is_a_priced_reasoning_model(local_model_cost_map: None) -> None: + repo_root: Final = Path(litellm.__file__).parent.parent + entries: Final = tuple( + json.loads(path.read_text())["minimax/MiniMax-M3.1-Flash"] + for path in ( + repo_root / "model_prices_and_context_window.json", + repo_root / "litellm/model_prices_and_context_window_backup.json", + ) + ) + assert entries[0] == entries[1] + + info: Final = litellm.get_model_info("MiniMax-M3.1-Flash", custom_llm_provider="minimax") + assert info["supports_reasoning"] is True + assert set(_flash_effort_levels()) <= set(get_args(REASONING_EFFORT)) + assert 0 < info["cache_read_input_token_cost"] < info["input_cost_per_token"] < info["output_cost_per_token"] + + +@pytest.mark.parametrize("model", ["MiniMax-M3.2", "MiniMax-Text-02"]) +def test_minimax_forwards_reasoning_effort_for_an_unmapped_model(minimax_request_mock: respx.Route, model: str) -> None: + """A MiniMax model newer than the cost map keeps reasoning_effort and thinking through the fallback rule.""" + assert f"minimax/{model}" not in litellm.model_cost + assert "thinking" in MinimaxChatConfig().get_supported_openai_params(model=model) + assert _sent_body(model, minimax_request_mock, "low", drop_params=False)["reasoning_effort"] == "low" + + +@pytest.mark.respx(assert_all_called=False) +def test_minimax_drops_reasoning_effort_for_a_known_non_reasoning_model( + minimax_request_mock: respx.Route, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setitem(litellm.model_cost["minimax/MiniMax-M3"], "supports_reasoning", False) + litellm.get_model_info.cache_clear() + + assert "reasoning_effort" not in MinimaxChatConfig().get_supported_openai_params(model="MiniMax-M3") + assert "reasoning_effort" not in _sent_body("MiniMax-M3", minimax_request_mock, "low") + + with pytest.raises(litellm.UnsupportedParamsError, match="reasoning_effort"): + _sent_body("MiniMax-M3", minimax_request_mock, "low", drop_params=False) + assert minimax_request_mock.call_count == 1 diff --git a/tests/unit/llms/test_cache_control_and_reasoning.py b/tests/unit/llms/test_cache_control_and_reasoning.py index be1ba1e7dbd..82406d114a9 100644 --- a/tests/unit/llms/test_cache_control_and_reasoning.py +++ b/tests/unit/llms/test_cache_control_and_reasoning.py @@ -76,6 +76,19 @@ def test_minimax_supports_thinking_param(): assert "reasoning_split" in supported_params +def test_minimax_forwards_reasoning_effort(): + config = MinimaxChatConfig() + + optional_params = config.map_openai_params( + non_default_params={"reasoning_effort": "low"}, + optional_params={}, + model="minimax/MiniMax-M3.1-Flash", + drop_params=True, + ) + + assert optional_params == {"reasoning_effort": "low"} + + def test_zai_preserves_cache_control_in_messages(): """ZAI should NOT strip cache_control from messages.""" config = ZAIChatConfig() From 6a232c90e0865aa4c01f252986a4bfb1a4075795 Mon Sep 17 00:00:00 2001 From: Ted Swinyar <8061019+tedswinyar@users.noreply.github.com> Date: Mon, 9 Nov 2026 14:13:40 -0800 Subject: [PATCH 19/46] fix(ollama_chat): give each streamed tool call its own index (#49183) * fix(ollama_chat): give each streamed tool call its own index Ollama streams parallel tool calls one per chunk with no top-level index, and the delta numbers the calls of each chunk from zero, so every call left the proxy as index 0 and OpenAI-style clients appended the arguments of each later call onto the first one. Number the calls with a counter that lives for the whole stream instead. Fixes #33678 * fix(ollama_chat): answer 400 when replayed tool call arguments are not JSON A chat history that already holds a merged tool call from an older stream used to fail with a 500 APIConnectionError. It now answers 400 naming the tool call. The streaming tests use Ollama's real tool call shape and drop a case the merge base already passed. * fix(ollama_chat): parse replayed tool call arguments with the shared repair parser Empty, null, and truncated arguments in a replayed assistant tool call now go through parse_tool_call_arguments, so they are forwarded as an object instead of answering 500. Arguments that still are not a JSON object answer 400 naming the call id. * fix(ollama_chat): send {} for null replayed arguments and keep argument text out of the repair warning Replayed tool call arguments of the JSON string "null" used to reach Ollama and return 200. The shared parser turns them into None, so they now go out as {} like empty or missing arguments instead of answering 400. The parser's repair warning logged up to 200 characters of the arguments. It now logs only the tool name, context, and length. * fix(ollama): keep replayed tool arguments out of a rejected replay's error chain * fix(ollama): keep the repair warning's argument preview unless message logging is off The shared repair warning lost its 200-character argument preview for every provider. It now logs the same preview as before by default, and logs redacted-by-litellm instead when turn_off_message_logging is on. * test(ollama): cover short and boundary-length repair warning previews --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../prompt_templates/tool_call_json_repair.py | 16 +- litellm/llms/ollama/chat/transformation.py | 34 ++++- ..._prompt_templates_tool_call_json_repair.py | 54 +++++++ .../ollama/test_ollama_chat_transformation.py | 137 +++++++++++++++++- 4 files changed, 231 insertions(+), 10 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/tool_call_json_repair.py b/litellm/litellm_core_utils/prompt_templates/tool_call_json_repair.py index 317ffa7b3b1..ace6d25c9f6 100644 --- a/litellm/litellm_core_utils/prompt_templates/tool_call_json_repair.py +++ b/litellm/litellm_core_utils/prompt_templates/tool_call_json_repair.py @@ -3,10 +3,21 @@ from typing import Any, Final # noqa: TID251 # Existing annotations retained u from pydantic import TypeAdapter +import litellm from litellm import verbose_logger +from litellm.constants import REDACTED_BY_LITELLM from litellm.litellm_core_utils.core_helpers import as_str_mapping _DECODED_JSON: Final = TypeAdapter(object) +_ARGUMENT_PREVIEW_CHARS: Final = 200 + + +def _logged_arguments(arguments: str) -> str: + if litellm.turn_off_message_logging: + return REDACTED_BY_LITELLM + if len(arguments) <= _ARGUMENT_PREVIEW_CHARS: + return arguments + return f"{arguments[:_ARGUMENT_PREVIEW_CHARS]}..." def _attempt_json_repair(s: str) -> object | None: @@ -104,12 +115,11 @@ def parse_tool_call_arguments( repaired: Final = _attempt_json_repair(arguments) if repaired is not None: verbose_logger.warning( - "Repaired truncated tool call arguments for tool '%s' (%s). Original (%d chars): %.200s%s", + "Repaired truncated tool call arguments for tool '%s' (%s). Original (%d chars): %s", tool_name or "", context or "unknown context", len(arguments), - arguments, - "..." if len(arguments) > 200 else "", + _logged_arguments(arguments), ) return repaired diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index c6f3a7a8280..1e3e6186dd1 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -16,6 +16,9 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.litellm_core_utils.prompt_templates.reasoning_utils import ( extract_reasoning_content, ) +from litellm.litellm_core_utils.prompt_templates.tool_call_json_repair import ( + parse_tool_call_arguments, +) from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE from litellm.litellm_core_utils.tool_fields import tool_params_without_hosted_mcp from litellm.litellm_core_utils.url_utils import ensure_api_route @@ -44,6 +47,27 @@ else: LiteLLMLoggingObj = Any _REQUEST_PARAMS: Final = TypeAdapter(dict[str, object]) +_TOOL_CALL_ARGUMENTS: Final = TypeAdapter(dict[str, object]) + + +def _to_ollama_tool_call_function(model: str, tool_call: ChatCompletionAssistantToolCall) -> OllamaToolCallFunction: + function: Final = tool_call["function"] + name: Final = function.get("name") or "" + raw_arguments: Final[object] = function.get("arguments") + try: + parsed: Final[object] = ( + parse_tool_call_arguments(raw_arguments, tool_name=name, context="Ollama chat request") + if isinstance(raw_arguments, str) or raw_arguments is None + else raw_arguments + ) + arguments: Final = _TOOL_CALL_ARGUMENTS.validate_python({} if parsed is None else parsed) + except ValueError: + raise litellm.BadRequestError( + message=f"OllamaException - arguments of assistant tool call {tool_call.get('id')!r} are not a JSON object", + model=model, + llm_provider="ollama_chat", + ) from None + return OllamaToolCallFunction(name=name, arguments=arguments) class OllamaChatConfig(BaseConfig): @@ -247,12 +271,7 @@ class OllamaChatConfig(BaseConfig): tools: Final = optional_params.pop("tools", None) def transform_tool_call(typed_tool: ChatCompletionAssistantToolCall) -> OllamaToolCall: - function: Final = OllamaToolCallFunction( - name=typed_tool["function"].get("name") or "", - arguments=json.loads(typed_tool["function"]["arguments"]) - if "arguments" in typed_tool["function"] - else {}, - ) + function: Final = _to_ollama_tool_call_function(model, typed_tool) tool_id: Final = typed_tool.get("id") ollama_tool_call: Final = ( OllamaToolCall(function=function) if tool_id is None else OllamaToolCall(id=tool_id, function=function) @@ -442,6 +461,7 @@ class OllamaChatCompletionResponseIterator(BaseModelResponseIterator): started_reasoning_content: bool = False finished_reasoning_content: bool = False seen_tool_calls: bool = False + _next_tool_call_index: int = 0 def _is_function_call_complete(self, function_args: str | dict) -> bool: if isinstance(function_args, dict): @@ -489,6 +509,8 @@ class OllamaChatCompletionResponseIterator(BaseModelResponseIterator): if tool_calls is not None: self.seen_tool_calls = True for tool_call in tool_calls: + tool_call["index"] = self._next_tool_call_index + self._next_tool_call_index += 1 function_args = tool_call.get("function").get("arguments") if function_args is not None and len(function_args) > 0: is_function_call_complete = self._is_function_call_complete(function_args) diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_tool_call_json_repair.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_tool_call_json_repair.py index 75f4e686137..4bd86d80c36 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_tool_call_json_repair.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_tool_call_json_repair.py @@ -1,7 +1,11 @@ import json +import logging +from typing import Final import pytest +import litellm +from litellm.constants import REDACTED_BY_LITELLM from litellm.litellm_core_utils.prompt_templates.tool_call_json_repair import ( parse_tool_call_arguments, salvage_concatenated_tool_arguments, @@ -115,3 +119,53 @@ def _distinct_json_objects(count: int) -> str: ) def test_salvage_concatenated_tool_arguments(raw: str, expected: tuple[dict[str, object], ...]) -> None: assert salvage_concatenated_tool_arguments(raw) == expected + + +def test_parse_tool_call_arguments_repair_warning_previews_the_first_200_characters( + caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "turn_off_message_logging", False) + arguments: Final = '{"patient": "Jane Doe", "note": "' + "x" * 300 + '"' + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + result: Final = parse_tool_call_arguments(arguments, tool_name="lookup", context="test") + + assert result == {"patient": "Jane Doe", "note": "x" * 300} + assert ( + f"Repaired truncated tool call arguments for tool 'lookup' (test). Original ({len(arguments)} chars): " + f"{arguments[:200]}..." in caplog.text + ) + assert arguments not in caplog.text + + +@pytest.mark.parametrize( + ("arguments", "expected"), + ( + ('{"patient": "Jane Doe"', {"patient": "Jane Doe"}), + ('{"note": "' + "x" * 189 + '"', {"note": "x" * 189}), + ), +) +def test_parse_tool_call_arguments_repair_warning_logs_arguments_up_to_200_characters_whole( + caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, + arguments: str, + expected: dict[str, str], +) -> None: + monkeypatch.setattr(litellm, "turn_off_message_logging", False) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + result: Final = parse_tool_call_arguments(arguments, tool_name="lookup", context="test") + + assert result == expected + assert f"Original ({len(arguments)} chars): {arguments}\n" in caplog.text + + +def test_parse_tool_call_arguments_repair_warning_redacts_arguments_when_message_logging_is_off( + caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "turn_off_message_logging", True) + arguments: Final = '{"patient": "Jane Doe"' + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + result: Final = parse_tool_call_arguments(arguments, tool_name="lookup", context="test") + + assert result == {"patient": "Jane Doe"} + assert f"Original ({len(arguments)} chars): {REDACTED_BY_LITELLM}" in caplog.text + assert "Jane Doe" not in caplog.text diff --git a/tests/unit/llms/ollama/test_ollama_chat_transformation.py b/tests/unit/llms/ollama/test_ollama_chat_transformation.py index 98ca5aee407..36df47deaab 100644 --- a/tests/unit/llms/ollama/test_ollama_chat_transformation.py +++ b/tests/unit/llms/ollama/test_ollama_chat_transformation.py @@ -1,8 +1,10 @@ import asyncio import inspect +import itertools import logging import os import sys +import traceback from typing import Final, cast import pytest @@ -26,7 +28,14 @@ from unittest.mock import ANY, MagicMock, patch import litellm from litellm._logging import verbose_logger -from litellm.types.utils import Choices, EmbeddingResponse, Message, ModelResponse, ModelResponseStream +from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Choices, + EmbeddingResponse, + Message, + ModelResponse, + ModelResponseStream, +) from unittest import mock @@ -1043,6 +1052,132 @@ class TestOllamaToolCallTransformation: assert "tool_call_id" in tool_msg, "tool_call_id must be forwarded to Ollama" assert tool_msg["tool_call_id"] == "call_abc123" + @staticmethod + def _replayed_history(arguments: object) -> list[AllMessageValues]: + return cast( + list[AllMessageValues], + [ + {"role": "user", "content": "Read a.rs and b.rs"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_abc123", + "type": "function", + "function": {"name": "read_file", "arguments": arguments}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_abc123", "content": "fn main() {}"}, + ], + ) + + @pytest.mark.parametrize( + ("arguments", "expected"), + [ + ("", {}), + (None, {}), + ("null", {}), + ('{"path": "a.rs"', {"path": "a.rs"}), + ('{"path": "a.rs", "limit": 2}', {"path": "a.rs", "limit": 2}), + ({"path": "a.rs"}, {"path": "a.rs"}), + ], + ids=["empty", "null", "null_string", "truncated", "non_string_value", "object"], + ) + def test_replayed_tool_call_arguments_are_forwarded_as_an_object(self, arguments: object, expected: object): + result: Final = OllamaChatConfig().transform_request( + model="qwen3:8b", + messages=self._replayed_history(arguments), + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert result["messages"][1]["tool_calls"] == [ + {"id": "call_abc123", "function": {"name": "read_file", "arguments": expected}} + ] + + @pytest.mark.parametrize( + "arguments", + ['{"path": "a.rs"}{"path": "b.rs"}', '["a.rs", "b.rs"]', '[["path", "a.rs"]]', "[]"], + ids=["glued", "array", "pairs", "empty_array"], + ) + def test_replayed_tool_call_arguments_that_are_not_an_object_answer_400_without_echoing_them( + self, arguments: str + ): + with pytest.raises(litellm.BadRequestError) as exc_info: + litellm.completion( + model="ollama_chat/qwen3:8b", + messages=self._replayed_history(arguments), + api_base="http://127.0.0.1:9", + ) + + assert exc_info.value.status_code == 400 + assert "call_abc123" in str(exc_info.value) + assert "a.rs" not in "".join(traceback.format_exception(exc_info.value)) + + +class TestOllamaStreamingParallelToolCallIndex: + """Regression for https://github.com/BerriAI/litellm/issues/33678. + + Ollama streams each parallel tool call in its own chunk without an OpenAI-style top-level + index. Every delta used to carry index 0, so index-accumulating clients appended the + arguments of every later call onto the first one. + """ + + @staticmethod + def _tool_call(index: int, command: str) -> dict[str, object]: + return { + "id": f"call_{index}", + "function": {"index": index, "name": "bash", "arguments": {"command": command}}, + } + + @staticmethod + def _chunk(*tool_calls: dict[str, object]) -> dict[str, object]: + return { + "model": "qwen3:8b", + "message": {"role": "assistant", "content": "", "tool_calls": list(tool_calls)}, + "done": False, + } + + @staticmethod + def _parsed_tool_calls( + iterator: OllamaChatCompletionResponseIterator, *chunks: dict[str, object] + ) -> tuple[ChatCompletionDeltaToolCall, ...]: + return tuple( + itertools.chain.from_iterable( + iterator.chunk_parser(chunk).choices[0].delta.tool_calls or () for chunk in chunks + ) + ) + + def test_tool_calls_in_separate_chunks_get_distinct_sequential_indices(self): + iterator: Final = OllamaChatCompletionResponseIterator(streaming_response=iter([]), sync_stream=True) + chunks: Final = [ + self._chunk(self._tool_call(index, command)) + for index, command in enumerate(("date", "uname -a", "whoami")) + ] + + tool_calls: Final = self._parsed_tool_calls(iterator, *chunks) + + assert [tool_call.index for tool_call in tool_calls] == [0, 1, 2] + assert [json.loads(tool_call.function.arguments) for tool_call in tool_calls] == [ + {"command": "date"}, + {"command": "uname -a"}, + {"command": "whoami"}, + ] + + def test_index_counter_continues_across_chunks(self): + iterator: Final = OllamaChatCompletionResponseIterator(streaming_response=iter([]), sync_stream=True) + chunks: Final = [ + self._chunk(self._tool_call(0, "date"), self._tool_call(1, "uname -a")), + self._chunk(self._tool_call(2, "whoami")), + ] + + tool_calls: Final = self._parsed_tool_calls(iterator, *chunks) + + assert [tool_call.index for tool_call in tool_calls] == [0, 1, 2] + class TestOllamaStreamingUsage: @staticmethod From b09061b6402f967b05e1a275a585ddb2fcb83faf Mon Sep 17 00:00:00 2001 From: yujonglee Date: Tue, 10 Nov 2026 07:31:12 +0900 Subject: [PATCH 20/46] fix(redis): accept a bare service account email for GCP Memorystore IAM (#50676) The Rust credential source passed gcp_service_account straight into the IAM credentials URL, so a bare email got a 404 and the stores fell back to per-worker memory. Address it as projects/-/serviceAccounts/ when it is not already a resource name, like the Python gRPC path did. Co-authored-by: Claude Opus 5.5 --- .../crates/redis-connect/src/sources/gcp.rs | 36 ++++++++++++++++--- 1 file changed, 31 insertions(+), 5 deletions(-) diff --git a/litellm-rust/crates/redis-connect/src/sources/gcp.rs b/litellm-rust/crates/redis-connect/src/sources/gcp.rs index 9da60fe9c08..9fb238709a1 100644 --- a/litellm-rust/crates/redis-connect/src/sources/gcp.rs +++ b/litellm-rust/crates/redis-connect/src/sources/gcp.rs @@ -41,6 +41,15 @@ impl GcpIam { } } +fn token_url(service_account: &str) -> String { + match service_account.starts_with("projects/") { + true => format!("{IAM_CREDENTIALS}/{service_account}:generateAccessToken"), + false => format!( + "{IAM_CREDENTIALS}/projects/-/serviceAccounts/{service_account}:generateAccessToken" + ), + } +} + impl CredentialSource for GcpIam { fn fetch(&self) -> SourceFuture<'_> { Box::pin(async move { @@ -52,10 +61,7 @@ impl CredentialSource for GcpIam { .map_err(Error::credential(CredentialStage::Identity))?; let response = self .http - .post(format!( - "{IAM_CREDENTIALS}/{}:generateAccessToken", - self.service_account - )) + .post(token_url(&self.service_account)) .bearer_auth(caller.expose()) .json(&serde_json::json!({ "scope": [CLOUD_PLATFORM_SCOPE] })) .send() @@ -92,9 +98,29 @@ impl CredentialSource for GcpIam { mod tests { use rstest::rstest; - use super::GcpIam; + use super::{GcpIam, token_url}; use crate::CredentialSource; + #[rstest] + #[case::bare_email( + "redis@project.iam.gserviceaccount.com", + "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/redis@project.iam.gserviceaccount.com:generateAccessToken" + )] + #[case::resource_name( + "projects/-/serviceAccounts/redis@project.iam.gserviceaccount.com", + "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/redis@project.iam.gserviceaccount.com:generateAccessToken" + )] + #[case::project_scoped_resource_name( + "projects/project/serviceAccounts/redis@project.iam.gserviceaccount.com", + "https://iamcredentials.googleapis.com/v1/projects/project/serviceAccounts/redis@project.iam.gserviceaccount.com:generateAccessToken" + )] + fn a_bare_service_account_email_is_addressed_as_a_resource_name( + #[case] service_account: &str, + #[case] expected: &str, + ) { + assert_eq!(token_url(service_account), expected); + } + #[rstest] #[tokio::test] async fn the_caller_identity_comes_from_application_default_credentials_not_vertex_settings() { From 3876802160f4298f59f52d91df1da0aef2bf3ddf Mon Sep 17 00:00:00 2001 From: yujonglee Date: Tue, 10 Nov 2026 07:31:36 +0900 Subject: [PATCH 21/46] fix(rust): install the rustls crypto provider at the Python host boundary (#50502) python-bridge now installs ring in its pymodule_init, matching the gateway's startup install. redis-connect no longer sets the process default as a side effect of connecting, so a dependency that calls ClientConfig::builder() without a provider cannot panic on the Python host depending on whether Redis has connected yet Co-authored-by: Claude Opus 5.5 --- litellm-rust/Cargo.lock | 1 + litellm-rust/crates/python-bridge/Cargo.toml | 1 + litellm-rust/crates/python-bridge/src/lib.rs | 15 +++++++++++++++ litellm-rust/crates/redis-connect/Cargo.toml | 1 - litellm-rust/crates/redis-connect/src/connect.rs | 3 --- 5 files changed, 17 insertions(+), 4 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 48b6730b0c6..c20454c7058 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4423,6 +4423,7 @@ dependencies = [ "pyo3-async-runtimes", "reqwest", "rstest", + "rustls 0.23.42", "serde", "serde_json", "serde_with", diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 6ecfebf0096..79773ca5559 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -19,6 +19,7 @@ huggingface = ["litellm-token-counter/huggingface"] tiktoken = ["litellm-token-counter/tiktoken"] [dependencies] +rustls.workspace = true litellm-secret-value.workspace = true fancy-regex.workspace = true litellm-tracing.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 0883dfa7d31..7513ff6b87b 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -92,6 +92,7 @@ mod _native { #[pymodule_init] fn init(module: &Bound<'_, PyModule>) -> PyResult<()> { + let _ = rustls::crypto::ring::default_provider().install_default(); let py = module.py(); let dict = module.dict(); dict.set_item( @@ -189,4 +190,18 @@ mod tests { assert_eq!(public_names, expected); }); } + #[rstest::rstest] + fn module_import_installs_ring_as_the_process_crypto_provider() { + Python::initialize(); + Python::attach(|py| { + native_module(py); + }); + + let installed = rustls::crypto::CryptoProvider::get_default() + .expect("importing _native should install a process crypto provider"); + assert_eq!( + format!("{installed:?}"), + format!("{:?}", rustls::crypto::ring::default_provider()) + ); + } } diff --git a/litellm-rust/crates/redis-connect/Cargo.toml b/litellm-rust/crates/redis-connect/Cargo.toml index 72a453d36de..29b3293b856 100644 --- a/litellm-rust/crates/redis-connect/Cargo.toml +++ b/litellm-rust/crates/redis-connect/Cargo.toml @@ -17,7 +17,6 @@ litellm-circuit-breaker.workspace = true litellm-redis-types.workspace = true litellm-core-utils.workspace = true litellm-http.workspace = true -rustls.workspace = true litellm-redis-client.workspace = true futures-util.workspace = true serde.workspace = true diff --git a/litellm-rust/crates/redis-connect/src/connect.rs b/litellm-rust/crates/redis-connect/src/connect.rs index cf7fe2334d5..d11f6ba47ff 100644 --- a/litellm-rust/crates/redis-connect/src/connect.rs +++ b/litellm-rust/crates/redis-connect/src/connect.rs @@ -42,9 +42,6 @@ pub fn connect( host: HostServices, breaker: Option, ) -> Result { - if rustls::crypto::CryptoProvider::get_default().is_none() { - let _ = rustls::crypto::ring::default_provider().install_default(); - } tracing::info!(target = ?resolved.target, namespace = ?resolved.namespace, "redis client configured"); let credentials = credentials(resolved.auth, &host)?; let events = host.events; From 1df1c4074b23d1ec5555e17f1c9dc53659bbcf4e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 9 Nov 2026 22:32:34 +0000 Subject: [PATCH 22/46] feat(cost-map): add GovCloud regional rows for Bedrock Claude Fable 5.2 and Mythos 5.2 (#50746) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 136 ++++++++++++++++++ model_prices_and_context_window.json | 136 ++++++++++++++++++ 2 files changed, 272 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index e1cbb953c5b..cc920c5fe25 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -87902,5 +87902,141 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "thinking_always_on": true + }, + "bedrock/us-gov-east-1/anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "bedrock/us-gov-west-1/anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "bedrock/us-gov-east-1/anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "bedrock/us-gov-west-1/anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index e1cbb953c5b..cc920c5fe25 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -87902,5 +87902,141 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "thinking_always_on": true + }, + "bedrock/us-gov-east-1/anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "bedrock/us-gov-west-1/anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "bedrock/us-gov-east-1/anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "bedrock/us-gov-west-1/anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true } } From 0021d5493cf203f56cac04308a6b46c351bcfdbe Mon Sep 17 00:00:00 2001 From: paveln83 <349044532+paveln83@users.noreply.github.com> Date: Mon, 9 Nov 2026 15:04:36 -0800 Subject: [PATCH 23/46] fix(bedrock): answer 400, not 500, for malformed count-tokens requests (#50708) * fix(bedrock): answer 400, not 500, for malformed native count-tokens requests The Bedrock count-tokens handler now checks the body before any AWS call and raises a 400 BedrockError for a missing or empty messages list, a message without role or content, or a missing model. The token counter marks that rejection as local, so it skips the Mantle fallback and the framing estimate, and /v1/messages/count_tokens answers invalid_request_error get_bedrock_base_model returns a bare cross-region prefix such as us or global unchanged instead of raising IndexError, so Bedrock judges the model id and its own 400 reaches the caller * fix(bedrock): answer 400 for a malformed message on Mantle token counts * fix(bedrock): limit the region-only model id fix to token counting * fix(bedrock): count a routing-prefixed region-only id and leave other providers' count errors as they were --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../count_tokens/bedrock_token_counter.py | 12 ++- litellm/llms/bedrock/count_tokens/handler.py | 11 ++- .../bedrock/count_tokens/mantle_handler.py | 5 + .../bedrock/count_tokens/transformation.py | 74 +++++++-------- .../test_bedrock_count_tokens_handler.py | 37 ++++++++ ...est_bedrock_count_tokens_transformation.py | 34 +++++++ .../test_bedrock_token_counter.py | 53 +++++++++++ .../llms/bedrock/test_bedrock_common_utils.py | 15 +++ .../anthropic_endpoints/test_endpoints.py | 44 +++++++++ .../test_llm_pass_through_endpoints.py | 91 ++++++++++++++++++- tests/unit/proxy/test_proxy_token_counter.py | 49 +++++----- 11 files changed, 351 insertions(+), 74 deletions(-) diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 6f0a18b0088..3866240030b 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -21,7 +21,7 @@ from litellm.llms.bedrock.common_utils import ( BedrockError, get_bedrock_base_model, ) -from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler +from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler, BedrockCountTokensRequestError from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler from litellm.llms.bedrock_mantle.common_utils import is_mantle_claude_model from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -54,6 +54,7 @@ class CountTokensFailure: status_code: int message: str tokenizer_type: str + rejected_before_provider_call: bool = False CountTokensOutcome = CountedTokens | CountTokensFailure @@ -74,6 +75,7 @@ def _claude_count_translator(litellm_params: Mapping[str, object]) -> AnthropicC def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model: str) -> bool: return ( isinstance(outcome, CountTokensFailure) + and not outcome.rejected_before_provider_call and outcome.status_code in _RUNTIME_FALLBACK_STATUSES and is_mantle_claude_model(resolved_model) ) @@ -162,6 +164,13 @@ class BedrockTokenCounter(BaseTokenCounter): client=self._client, ) reply: Final = _CountTokensReply.model_validate(result) + except BedrockCountTokensRequestError as e: + return CountTokensFailure( + status_code=e.status_code, + message=e.message, + tokenizer_type=tokenizer_type, + rejected_before_provider_call=True, + ) except BedrockError as e: verbose_logger.debug( "%s CountTokens API error: status=%s, message=%s", tokenizer_type, e.status_code, e.message @@ -268,6 +277,7 @@ class BedrockTokenCounter(BaseTokenCounter): error=True, error_message=outcome.message, status_code=outcome.status_code, + rejected_before_provider_call=outcome.rejected_before_provider_call, ) case _: assert_never(outcome) diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py index 84067923788..eb29459f8b1 100644 --- a/litellm/llms/bedrock/count_tokens/handler.py +++ b/litellm/llms/bedrock/count_tokens/handler.py @@ -20,6 +20,11 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_h _DECODED_JSON: Final = TypeAdapter(object) +class BedrockCountTokensRequestError(BedrockError): + def __init__(self, problem: str) -> None: + super().__init__(status_code=400, message=problem) + + class BedrockCountTokensHandler(BedrockCountTokensConfig): """ Simplified handler for AWS Bedrock CountTokens API requests. @@ -45,10 +50,10 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): Returns: Dictionary containing token count response """ + problem: Final = self.count_tokens_request_problem(request_data) + if problem is not None: + raise BedrockCountTokensRequestError(problem) try: - # Validate the request - self.validate_count_tokens_request(request_data) - verbose_logger.debug("Processing CountTokens request for resolved model: %s", resolved_model) # Get AWS region using existing LiteLLM function diff --git a/litellm/llms/bedrock/count_tokens/mantle_handler.py b/litellm/llms/bedrock/count_tokens/mantle_handler.py index 5e5c5516b94..0b589eda896 100644 --- a/litellm/llms/bedrock/count_tokens/mantle_handler.py +++ b/litellm/llms/bedrock/count_tokens/mantle_handler.py @@ -12,6 +12,8 @@ from litellm.llms.anthropic.common_utils import context_management_betas, tools_ from litellm.llms.anthropic.count_tokens.transformation import COUNT_TOKEN_OPTION_NAMES, AnthropicCountTokensConfig from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing from litellm.llms.bedrock.common_utils import BedrockError, build_mantle_messages_url +from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensRequestError +from litellm.llms.bedrock.count_tokens.transformation import count_tokens_messages_problem from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client MANTLE_COUNT_TOKENS_SUFFIX: Final = "/count_tokens" @@ -58,6 +60,9 @@ class BedrockMantleCountTokensHandler(BaseAWSLLM): ) -> dict[str, JsonValue]: try: request: Final = _COUNT_REQUEST.validate_python(request_data) + problem: Final = count_tokens_messages_problem(request["messages"]) + if problem is not None: + raise BedrockCountTokensRequestError(problem) aws_region_name: Final = await run_aws_signing( self.get_aws_region_name, optional_params=litellm_params, model=resolved_model, model_id=None ) diff --git a/litellm/llms/bedrock/count_tokens/transformation.py b/litellm/llms/bedrock/count_tokens/transformation.py index cd59e714ba0..aaeb57303f9 100644 --- a/litellm/llms/bedrock/count_tokens/transformation.py +++ b/litellm/llms/bedrock/count_tokens/transformation.py @@ -6,7 +6,7 @@ to AWS Bedrock's CountTokens API format and vice versa. """ import re -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import Final, Literal from pydantic import JsonValue @@ -14,7 +14,11 @@ from pydantic import JsonValue from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.common_utils import tools_without_custom_field from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM -from litellm.llms.bedrock.common_utils import get_bedrock_base_model +from litellm.llms.bedrock.common_utils import ( + get_bedrock_base_model, + get_bedrock_cross_region_inference_regions, + strip_bedrock_routing_prefix, +) # Placeholder satisfying the Anthropic InvokeModel schema's required # max_tokens field; CountTokens only counts input, so it has no effect @@ -46,6 +50,23 @@ def _to_converse_message(message: JsonValue) -> dict[str, JsonValue]: } +def _converse_message_problem(index: int, message: JsonValue) -> str | None: + if not isinstance(message, dict): + return f"Message {index} must be a dictionary" + if "role" not in message: + return f"Message {index} must have a 'role' field" + if "content" not in message and "tool_calls" in message: + return f"Message {index} has tool_calls and no content, which token counting cannot convert for this model" + if "content" not in message: + return f"Message {index} must have a 'content' field" + return None + + +def count_tokens_messages_problem(messages: Sequence[JsonValue]) -> str | None: + problems: Final = (_converse_message_problem(index, message) for index, message in enumerate(messages)) + return next((problem for problem in problems if problem is not None), None) + + def _sanitized_bedrock_tool_name(raw_name: JsonValue) -> str: name: Final = re.sub(r"[^a-zA-Z0-9_]", "_", raw_name if isinstance(raw_name, str) else "") prefixed: Final = name if not name or name[0].isalpha() else f"t_{name}" @@ -217,6 +238,9 @@ class BedrockCountTokensConfig(BaseAWSLLM): @staticmethod def count_tokens_model_id(model: str) -> str: + model_id: Final = strip_bedrock_routing_prefix(model) + if model_id in get_bedrock_cross_region_inference_regions(): + return model_id return get_bedrock_base_model(model).removeprefix("bedrock/") def get_bedrock_count_tokens_endpoint( @@ -265,42 +289,12 @@ class BedrockCountTokensConfig(BaseAWSLLM): return {"input_tokens": input_tokens} - def validate_count_tokens_request(self, request_data: Mapping[str, JsonValue]) -> None: - """ - Validate the incoming count tokens request. - Supports both Converse and InvokeModel input formats. - - Args: - request_data: The request payload - - Raises: - ValueError: If the request is invalid - """ + def count_tokens_request_problem(self, request_data: Mapping[str, JsonValue]) -> str | None: if not request_data.get("model"): - raise ValueError("model parameter is required") - - input_type: Final = self._detect_input_type(request_data) - - if input_type == "converse": - # Validate Converse format (messages-based) - messages: Final = request_data.get("messages", []) - if not messages: - raise ValueError("messages parameter is required for Converse input") - - if not isinstance(messages, list): - raise ValueError("messages must be a list") - - for i, message in enumerate(messages): - if not isinstance(message, dict): - raise ValueError(f"Message {i} must be a dictionary") - - if "role" not in message: - raise ValueError(f"Message {i} must have a 'role' field") - - if "content" not in message: - raise ValueError(f"Message {i} must have a 'content' field") - else: - # For InvokeModel format, we need at least some content to count tokens - # The content structure varies by model, so we do minimal validation - if len(request_data) <= 1: # Only has 'model' field - raise ValueError("Request must contain content to count tokens") + return "model parameter is required" + if self._detect_input_type(request_data) == "invokeModel": + return "Request must contain content to count tokens" if len(request_data) <= 1 else None + messages: Final = _json_list(request_data.get("messages")) + if not messages: + return "messages parameter is required for Converse input" + return count_tokens_messages_problem(messages) diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py index 591205b4602..19c9a042afb 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py @@ -6,6 +6,7 @@ import httpx import pytest from botocore.credentials import RefreshableCredentials +from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.llms.bedrock.count_tokens.transformation import BedrockCountTokensConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -85,6 +86,42 @@ async def test_handle_count_tokens_request_looks_up_its_region_off_the_event_loo assert result == {"input_tokens": 7} +_COUNTED_MODEL: Final = "us.anthropic.claude-opus-4-8" + + +@pytest.mark.parametrize( + ("request_data", "problem"), + [ + ({"model": _COUNTED_MODEL}, "Request must contain content to count tokens"), + ({"model": _COUNTED_MODEL, "messages": []}, "messages parameter is required for Converse input"), + ({"model": _COUNTED_MODEL, "messages": [{"role": "user"}]}, "Message 0 must have a 'content' field"), + ({"model": _COUNTED_MODEL, "messages": [{"content": "hi"}]}, "Message 0 must have a 'role' field"), + ({"model": _COUNTED_MODEL, "messages": ["hi"]}, "Message 0 must be a dictionary"), + ( + {"model": _COUNTED_MODEL, "messages": [{"role": "user", "content": "hi"}, {"role": "assistant"}]}, + "Message 1 must have a 'content' field", + ), + ({"messages": [{"role": "user", "content": "hi"}]}, "model parameter is required"), + ], +) +@pytest.mark.asyncio +async def test_handle_count_tokens_request_answers_400_for_a_malformed_body_without_calling_bedrock( + request_data: dict[str, object], problem: str +) -> None: + client: Final = AsyncMock(spec=AsyncHTTPHandler) + + with pytest.raises(BedrockError) as exc_info: + await BedrockCountTokensHandler().handle_count_tokens_request( + request_data=request_data, + litellm_params={"aws_region_name": "us-east-1"}, + resolved_model=_COUNTED_MODEL, + client=client, + ) + + assert (exc_info.value.status_code, exc_info.value.message) == (400, problem) + client.post.assert_not_awaited() + + class TestBedrockCountTokensEndpoint: def _make_handler(self) -> BedrockCountTokensConfig: return BedrockCountTokensConfig() diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py index d5b57e42146..018b0fde8a0 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py @@ -311,3 +311,37 @@ def test_count_tokens_endpoint_encodes_model_id(monkeypatch): endpoint == "https://bedrock-runtime.us-east-1.amazonaws.com/model/..%2F..%2Fmodel%2Fother%3Fx%3D1%23frag/count-tokens" ) + + +@pytest.mark.parametrize( + "request_data", + [ + {"model": "us.anthropic.claude-opus-4-8", "messages": [{"role": "user", "content": "hi"}]}, + {"model": "us.anthropic.claude-opus-4-8", "messages": "hi"}, + { + "model": "us.anthropic.claude-opus-4-8", + "messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}, {"content": "hi"}], + }, + ], +) +def test_count_tokens_request_problem_leaves_bodies_bedrock_can_judge_to_bedrock(request_data): + assert BedrockCountTokensConfig().count_tokens_request_problem(request_data) is None + + +def test_count_tokens_request_problem_names_tool_calls_instead_of_asking_for_content(): + problem = BedrockCountTokensConfig().count_tokens_request_problem( + { + "model": "us.amazon.nova-pro-v1:0", + "messages": [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "tool_calls": [{"id": "1", "type": "function", "function": {"name": "t", "arguments": "{}"}}], + }, + ], + } + ) + + assert problem is not None + assert "tool_calls" in problem + assert "must have a 'content' field" not in problem 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 30063936b2b..de0c14558b3 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 @@ -74,6 +74,59 @@ async def test_claude_model_bedrock_runtime_cannot_count_is_counted_on_mantle() assert mantle_body["system"] == "You are a terse assistant." +@pytest.mark.asyncio +async def test_claude_model_message_without_content_is_rejected_without_calling_runtime_or_mantle() -> None: + bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 1353}), mantle=(200, {"input_tokens": 1336})) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="global.anthropic.claude-opus-4-8", + messages=[{"role": "user"}], + contents=None, + deployment=DEPLOYMENT, + request_model="claude-opus-4-8", + ) + + assert result is not None + assert (result.error, result.status_code, result.error_message, result.rejected_before_provider_call) == ( + True, + 400, + "Message 0 must have a 'content' field", + True, + ) + assert bedrock.posted_hosts() == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("message", "problem"), + [ + ({"role": "user"}, "Message 0 must have a 'content' field"), + ({"content": "hi"}, "Message 0 must have a 'role' field"), + ], +) +async def test_mantle_counted_model_rejects_a_malformed_message_without_calling_mantle( + message: dict[str, str], problem: str +) -> None: + bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 1353}), mantle=(200, {"input_tokens": 1336})) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="bedrock_mantle/anthropic.claude-opus-5-5", + messages=[message], + contents=None, + deployment=DEPLOYMENT, + request_model="claude-opus-5-5", + ) + + assert result is not None + assert (result.error, result.status_code, result.error_message, result.rejected_before_provider_call) == ( + True, + 400, + problem, + True, + ) + assert bedrock.posted_hosts() == () + + @pytest.mark.asyncio async def test_bedrock_runtime_count_is_kept_when_it_answers() -> None: bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 1353}), mantle=(200, {"input_tokens": 1336})) diff --git a/tests/unit/llms/bedrock/test_bedrock_common_utils.py b/tests/unit/llms/bedrock/test_bedrock_common_utils.py index 31999ee0379..0799ec15225 100644 --- a/tests/unit/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/unit/llms/bedrock/test_bedrock_common_utils.py @@ -1663,6 +1663,21 @@ class TestGetBedrockBaseModel: assert get_bedrock_base_model(input_model) == expected +def test_region_only_deployment_takes_no_traffic_from_a_valid_model_group() -> None: + valid_model: Final = "bedrock/us.anthropic.claude-opus-4-6-v1" + router: Final = litellm.Router( + model_list=[ + {"model_name": "group", "litellm_params": {"model": "bedrock/us", "aws_region_name": "us-east-1"}}, + {"model_name": "group", "litellm_params": {"model": valid_model, "aws_region_name": "us-east-1"}}, + ], + ignore_invalid_deployments=True, + ) + + assert [deployment["litellm_params"]["model"] for deployment in router.get_model_list(model_name="group")] == [ + valid_model + ] + + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestBedrockModelInfoWrappers: """Tests that BedrockModelInfo methods correctly wrap standalone functions.""" diff --git a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py index 80e7e793f6e..117d148c287 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py @@ -13,6 +13,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +import respx from fastapi import Request, Response from fastapi.responses import StreamingResponse from fastapi.testclient import TestClient @@ -1222,6 +1223,49 @@ def test_count_tokens_rejected_before_counting_logs_the_400_it_answers( ) +def test_count_tokens_on_bedrock_answers_400_for_a_message_without_content_without_calling_bedrock( + monkeypatch: pytest.MonkeyPatch, httpx_transport: None, respx_mock: respx.MockRouter +) -> None: + 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.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.setattr( + proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "count-bedrock-claude", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-zeta-9", + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKIATESTACCESSKEY", + "aws_secret_access_key": "test-secret", + }, + "model_info": {"id": "count-bedrock-claude-deployment", "token_count_framing": 3}, + } + ] + ), + ) + app: Final = FastAPI() + app.include_router(router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth() + + response: Final = TestClient(app).post( + "/v1/messages/count_tokens", json={"model": "count-bedrock-claude", "messages": [{"role": "user"}]} + ) + + assert response.status_code == 400, response.text + assert response.json()["error"]["type"] == "invalid_request_error" + assert response.json()["error"]["message"] == "Message 0 must have a 'content' field" + assert len(respx_mock.calls) == 0 + + class TestEventLoggingBatchEndpoint: """Test the stubbed event logging batch endpoint""" 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 96d0a2ea8e0..70f5e855e16 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 @@ -7659,8 +7659,16 @@ _BEDROCK_COUNT_ALIAS: Final = "count-log-bedrock" _BEDROCK_COUNT_MODEL: Final = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" +_BEDROCK_COUNT_BODY: Final = MappingProxyType( + {"model": _BEDROCK_COUNT_ALIAS, "messages": [{"role": "user", "content": "hi"}]} +) + + async def _bedrock_native_count( - monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, aws_reply: httpx.Response + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + aws_reply: httpx.Response, + request_body: Mapping[str, object] = _BEDROCK_COUNT_BODY, ) -> dict[str, object]: from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( handle_bedrock_count_tokens, @@ -7700,7 +7708,7 @@ async def _bedrock_native_count( request=request, fastapi_response=Response(), user_api_key_dict=MagicMock(), - request_body={"model": _BEDROCK_COUNT_ALIAS, "messages": [{"role": "user", "content": "hi"}]}, + request_body=dict(request_body), ) @@ -7754,6 +7762,81 @@ async def test_bedrock_native_count_tokens_logs_the_provider_status_when_the_cou ) +@pytest.mark.parametrize( + ("request_body", "problem"), + [ + ({"model": _BEDROCK_COUNT_ALIAS}, "Request must contain content to count tokens"), + ({"model": _BEDROCK_COUNT_ALIAS, "messages": []}, "messages parameter is required for Converse input"), + ({"model": _BEDROCK_COUNT_ALIAS, "messages": [{"role": "user"}]}, "Message 0 must have a 'content' field"), + ({"model": _BEDROCK_COUNT_ALIAS, "messages": [{"content": "hi"}]}, "Message 0 must have a 'role' field"), + ], +) +@pytest.mark.asyncio +async def test_bedrock_native_count_tokens_answers_400_for_a_malformed_body_without_calling_bedrock( + monkeypatch: pytest.MonkeyPatch, + httpx_transport: None, + respx_mock: respx.MockRouter, + request_body: Mapping[str, object], + problem: str, +) -> None: + with pytest.raises(HTTPException) as exc_info: + await _bedrock_native_count( + monkeypatch, respx_mock, httpx.Response(200, json={"inputTokens": 12}), request_body + ) + + assert (exc_info.value.status_code, exc_info.value.detail) == (400, {"error": problem}) + assert len(respx_mock.calls) == 0 + + +@pytest.mark.parametrize( + ("model", "model_id"), + [ + ("us", "us"), + ("global", "global"), + ("apac", "apac"), + ("bedrock/us", "us"), + ("converse/us", "us"), + ("invoke/us", "us"), + ("bedrock/converse/global", "global"), + ], +) +@pytest.mark.asyncio +async def test_bedrock_native_count_tokens_for_a_bare_region_prefix_answers_bedrocks_own_400( + monkeypatch: pytest.MonkeyPatch, + httpx_transport: None, + respx_mock: respx.MockRouter, + model: str, + model_id: str, +) -> None: + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + handle_bedrock_count_tokens, + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "AKIATESTACCESSKEY") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "test-secret") + monkeypatch.setenv("AWS_REGION_NAME", "us-east-1") + invalid_model: Final = {"message": "The provided model identifier is invalid."} + respx_mock.post(host="bedrock-runtime.us-east-1.amazonaws.com", path=f"/model/{model_id}/count-tokens").mock( + return_value=httpx.Response(400, json=invalid_model) + ) + request: Final = Request({"type": "http", "method": "POST", "path": f"/bedrock/model/{model}/count-tokens", "headers": []}) + + with pytest.raises(HTTPException) as exc_info: + await handle_bedrock_count_tokens( + endpoint=f"model/{model}/count-tokens", + request=request, + fastapi_response=Response(), + user_api_key_dict=MagicMock(), + request_body={"model": model, "messages": [{"role": "user", "content": "hi"}]}, + ) + + assert exc_info.value.status_code == 400 + assert json.loads(exc_info.value.detail["error"]) == invalid_model + assert len(respx_mock.calls) == 1 + + @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], ...]], @@ -7820,9 +7903,7 @@ async def test_bedrock_native_count_tokens_with_a_region_only_model_answers_the_ request_body={"model": "us"}, ) - assert exc_info.value.detail == { - "error": "CountTokens processing error: Request must contain content to count tokens" - } + assert exc_info.value.detail == {"error": "Request must contain content to count tokens"} assert token_count_log_lines() == ( { "litellm_call_id": "bedrock-native-region-only", diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 2f85d1c7516..694b257b1f9 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -1080,33 +1080,32 @@ async def test_bedrock_handler_httpx_error_status_code_propagation(): response=mock_response, ) - with patch.object(handler, "validate_count_tokens_request"): - with patch.object(handler, "_get_aws_region_name", return_value="us-west-2"): - with patch.object(handler, "transform_anthropic_to_bedrock_count_tokens", return_value={}): - with patch.object( - handler, - "get_bedrock_count_tokens_endpoint", - return_value="https://example.com", - ): - with patch.object(handler, "_sign_request", return_value=({}, "{}")): - with patch("litellm.llms.bedrock.count_tokens.handler.get_async_httpx_client") as mock_client: - mock_async_client = AsyncMock() - mock_async_client.post = AsyncMock(side_effect=http_error) - mock_client.return_value = mock_async_client + with patch.object(handler, "_get_aws_region_name", return_value="us-west-2"): + with patch.object(handler, "transform_anthropic_to_bedrock_count_tokens", return_value={}): + with patch.object( + handler, + "get_bedrock_count_tokens_endpoint", + return_value="https://example.com", + ): + with patch.object(handler, "_sign_request", return_value=({}, "{}")): + with patch("litellm.llms.bedrock.count_tokens.handler.get_async_httpx_client") as mock_client: + mock_async_client = AsyncMock() + mock_async_client.post = AsyncMock(side_effect=http_error) + mock_client.return_value = mock_async_client - with pytest.raises(BedrockError) as exc_info: - await handler.handle_count_tokens_request( - request_data={ - "model": "test", - "messages": [{"role": "user", "content": "hello"}], - }, - litellm_params={}, - resolved_model="anthropic.claude-3-sonnet", - ) + with pytest.raises(BedrockError) as exc_info: + await handler.handle_count_tokens_request( + request_data={ + "model": "test", + "messages": [{"role": "user", "content": "hello"}], + }, + litellm_params={}, + resolved_model="anthropic.claude-3-sonnet", + ) - assert exc_info.value.status_code == 403 - # Message should be the raw response text - assert exc_info.value.message == "Forbidden - Invalid credentials" + assert exc_info.value.status_code == 403 + # Message should be the raw response text + assert exc_info.value.message == "Forbidden - Invalid credentials" @pytest.mark.asyncio From 5d6f32a6f66b4d2c332f42d2607959be7f93681e Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Mon, 9 Nov 2026 15:19:09 -0800 Subject: [PATCH 24/46] feat(router): let deployments declare a shared encrypted content boundary (#50750) * feat(router): let deployments declare a shared encrypted content boundary Deployments that set the same model_info.encrypted_content_boundary are treated as encryption boundary peers by encrypted_content_affinity, so regional endpoints of one provider organization with different api_base values keep encrypted reasoning when a follow-up moves between them. Deployments without the key keep the api_base and api_key match. * refactor(router): type the encryption boundary peer lookup's deployments * test(router): cover order fallback hops onto a declared encryption boundary peer * chore(ui): regenerate API types for encrypted_content_boundary --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> --- .../encrypted_content_affinity_check.py | 61 +++++- litellm/types/router.py | 4 + .../test_encrypted_content_affinity_check.py | 191 ++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 + 4 files changed, 247 insertions(+), 11 deletions(-) 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 9badee854fb..1564c0da634 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 @@ -38,7 +38,7 @@ Safe to enable globally: import hashlib from collections.abc import Callable, Iterator, Mapping, Sequence -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from functools import cache, lru_cache, partial from time import monotonic from typing import TYPE_CHECKING, Final, Optional, cast @@ -62,7 +62,7 @@ from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.router_deprecated import RenamedRouterAttribute, renamed_router_param from litellm.secret_managers.main import secret_manager_would_be_consulted from litellm.types.llms.openai import AllMessageValues -from litellm.types.router import Deployment, GenericLiteLLMParams +from litellm.types.router import Deployment, GenericLiteLLMParams, ModelInfo from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager, get_order_filtered_deployments, with_named_credential @@ -91,8 +91,11 @@ class _EncryptionBoundary: model: Callable[[], str | None] = field(compare=False) account_scope: frozenset[tuple[str, object]] explicit_pair: bool + declared: str | None = None def is_shared_with(self, other: "_EncryptionBoundary") -> bool: + if self.declared is not None and other.declared is not None: + return self.declared == other.declared if (self.api_base, self.api_key) != (other.api_base, other.api_key): return False if self.explicit_pair and other.explicit_pair: @@ -290,6 +293,37 @@ class EncryptedContentAffinityCheck(CustomLogger): ) ) + @staticmethod + def _deployment_encryption_boundary( + litellm_params: object, + model_info: object, + clock: Callable[[], float] = monotonic, + ) -> _EncryptionBoundary | None: + """ + ``_encryption_boundary_key``, overridden by an operator-declared + ``model_info.encrypted_content_boundary``. Two deployments that both declare one are peers exactly + when the names match, which lets endpoints of one provider organization with different ``api_base`` + values share encrypted reasoning. + """ + declared: Final = EncryptedContentAffinityCheck._declared_encryption_boundary(model_info) + resolved: Final = EncryptedContentAffinityCheck._encryption_boundary_key(litellm_params, clock) + if declared is None: + return resolved + if resolved is None: + return _EncryptionBoundary(None, None, lambda: None, frozenset(), explicit_pair=False, declared=declared) + return replace(resolved, declared=declared) + + @staticmethod + def _declared_encryption_boundary(model_info: object) -> str | None: + if isinstance(model_info, ModelInfo): + return model_info.encrypted_content_boundary + if not isinstance(model_info, Mapping): + return None + declared: Final = cast(Mapping[str, object], model_info).get( # cast-ok: narrowed by isinstance + "encrypted_content_boundary" + ) + return declared if isinstance(declared, str) and declared else None + @staticmethod @lru_cache(maxsize=1024) def _encryption_boundary_in_refresh_window(key: _ParamsInRefreshWindow) -> _EncryptionBoundary | None: @@ -408,9 +442,9 @@ class EncryptedContentAffinityCheck(CustomLogger): def _find_deployments_on_same_encryption_boundary( self, - healthy_deployments: list[dict], + healthy_deployments: list[dict[str, object]], model_id: str, - ) -> tuple[list[dict], Deployment | None]: + ) -> tuple[list[dict[str, object]], Deployment | None]: """ Deployments in ``healthy_deployments`` sharing the originating deployment's encryption boundary, alongside the originating @@ -421,15 +455,20 @@ class EncryptedContentAffinityCheck(CustomLogger): originating: Final = self.ctx.registry().deployment(model_id) if originating is None: return [], None - boundary: Final = self._encryption_boundary_key( - originating.litellm_params.model_dump(exclude_none=True), self._clock + boundary: Final = self._deployment_encryption_boundary( + originating.litellm_params.model_dump(exclude_none=True), originating.model_info, self._clock ) if boundary is None: return [], originating matches: Final = [ d for d in healthy_deployments - if (candidate := self._encryption_boundary_key(d.get("litellm_params", {}), self._clock)) is not None + if ( + candidate := self._deployment_encryption_boundary( + d.get("litellm_params", {}), d.get("model_info"), self._clock + ) + ) + is not None and boundary.is_shared_with(candidate) ] return matches, originating @@ -464,8 +503,8 @@ class EncryptedContentAffinityCheck(CustomLogger): boundary for target in target_deployments if ( - boundary := EncryptedContentAffinityCheck._encryption_boundary_key( - target.get("litellm_params"), clock + boundary := EncryptedContentAffinityCheck._deployment_encryption_boundary( + target.get("litellm_params"), target.get("model_info"), clock ) ) is not None @@ -482,8 +521,8 @@ class EncryptedContentAffinityCheck(CustomLogger): return False origin: Final = ctx.registry().deployment(origin_model_id) origin_boundary: Final = ( - EncryptedContentAffinityCheck._encryption_boundary_key( - origin.litellm_params.model_dump(exclude_none=True), clock + EncryptedContentAffinityCheck._deployment_encryption_boundary( + origin.litellm_params.model_dump(exclude_none=True), origin.model_info, clock ) if origin is not None else None diff --git a/litellm/types/router.py b/litellm/types/router.py index c0b182f0f20..6f5652db771 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -353,6 +353,10 @@ class ModelInfo(MirroredPricingParams): max_embedding_batch_size: Annotated[int, Field(gt=0)] | None = None + # deployments that set the same value share one provider encryption boundary, so encrypted reasoning + # minted on one can be replayed to another (e.g. regional endpoints of one OpenAI organization) + encrypted_content_boundary: Annotated[str, Field(min_length=1)] | None = None + def __init__(self, id: str | int | None = None, **params) -> None: if id is None: id = str(uuid.uuid4()) # Generate a UUID if id is None or not provided diff --git a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index 76a574bae24..adf08e43c95 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 @@ -3293,3 +3293,194 @@ async def test_a_changed_named_credential_moves_the_pin_within_the_secret_manage assert pin_before_the_change == [peer] assert pin_after_the_change == [other_org] assert reasoning_after_the_change == _reasoning_input_minted_by("origin")[0] + + +# --------------------------------------------------------------------------- +# Operator-declared encryption boundary (model_info.encrypted_content_boundary) +# --------------------------------------------------------------------------- + + +def _regional_deployment(model_name: str, region: str, boundary: str | None) -> dict: + return { + "model_name": model_name, + "litellm_params": { + "model": "openai/gpt-5.6", + "api_base": f"https://{region}.example.invalid/v1", + "api_key": "sk-one-org", + }, + "model_info": {"id": region} | ({} if boundary is None else {"encrypted_content_boundary": boundary}), + } + + +async def _forwarded_reasoning_after_model_switch(origin_boundary: str | None, peer_boundary: str | None) -> tuple: + router: Final = litellm.Router( + model_list=[ + _regional_deployment("reasoning-a", "region-1", origin_boundary), + _regional_deployment("reasoning-b", "region-2", peer_boundary), + ], + optional_pre_call_checks=["encrypted_content_affinity"], + num_retries=0, + ) + handler: Final = AsyncMock(return_value=_build_mock_response(output_items=[], response_id="resp_second")) + with patch( + "litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler.async_response_api_handler", + new=handler, + ): + response: Final = await router.aresponses( + model="reasoning-b", + input=[ + { + "type": "reasoning", + "summary": [], + "encrypted_content": ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id( + "blob-from-region-1", "region-1" + ), + }, + {"role": "user", "content": "follow-up"}, + ], + litellm_metadata={}, + ) + forwarded_reasoning: Final = handler.await_args.kwargs["input"][0] + return response._hidden_params["model_id"], forwarded_reasoning.get("encrypted_content") + + +@pytest.mark.asyncio +async def test_declared_boundary_keeps_reasoning_across_endpoints_with_different_api_bases(): + assert await _forwarded_reasoning_after_model_switch("org-1", "org-1") == ("region-2", "blob-from-region-1") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("origin_boundary", "peer_boundary"), + [(None, None), ("org-1", "org-2"), ("org-1", None), (None, "org-1")], + ids=["undeclared", "different-boundaries", "only-origin-declared", "only-peer-declared"], +) +async def test_reasoning_is_stripped_unless_both_deployments_declare_the_same_boundary( + origin_boundary, peer_boundary +): + assert await _forwarded_reasoning_after_model_switch(origin_boundary, peer_boundary) == ("region-2", None) + + +def test_declared_boundary_overrides_a_shared_api_base_and_key(): + origin: Final = _make_originating_mock("https://region-1.example.invalid/v1", "sk-one-org") + origin.model_info = litellm.types.router.ModelInfo(id="region-1", encrypted_content_boundary="org-1") + check: Final = EncryptedContentAffinityCheck(ctx=_make_router_mock_with_cooldown(origin)) + other_org: Final = { + "model_info": {"id": "region-1-org-2", "encrypted_content_boundary": "org-2"}, + "litellm_params": {"api_base": "https://region-1.example.invalid/v1", "api_key": "sk-one-org"}, + } + + matches, _ = check._find_deployments_on_same_encryption_boundary([other_org], "region-1") + + assert matches == [] + + +@pytest.mark.asyncio +async def test_boundary_pin_only_selects_deployments_that_declare_the_origins_boundary(): + origin: Final = _make_originating_mock("https://region-1.example.invalid/v1", "sk-one-org") + origin.model_info = litellm.types.router.ModelInfo(id="region-1", encrypted_content_boundary="org-1") + check: Final = EncryptedContentAffinityCheck(ctx=_make_router_mock_with_cooldown(origin)) + peer: Final = _regional_deployment("reasoning-b", "region-2", "org-1") + stranger: Final = _regional_deployment("reasoning-b", "region-3", "org-2") + request_kwargs: Final = { + "input": [ + { + "type": "reasoning", + "encrypted_content": ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id( + "blob-from-region-1", "region-1" + ), + } + ] + } + + result: Final = await check.async_filter_deployments( + model="reasoning-b", + healthy_deployments=[stranger, peer], + messages=None, + request_kwargs=request_kwargs, + ) + + assert result == [peer] + assert request_kwargs["input"][0]["encrypted_content"] == ( + ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id("blob-from-region-1", "region-1") + ) + + +def test_a_boundary_declared_on_one_side_keeps_the_shared_api_base_and_key_match(): + origin: Final = _make_originating_mock("https://region-1.example.invalid/v1", "sk-one-org") + origin.model_info = litellm.types.router.ModelInfo(id="region-1", encrypted_content_boundary="org-1") + check: Final = EncryptedContentAffinityCheck(ctx=_make_router_mock_with_cooldown(origin)) + undeclared_twin: Final = { + "model_info": {"id": "region-1-twin"}, + "litellm_params": {"api_base": "https://region-1.example.invalid/v1", "api_key": "sk-one-org"}, + } + + matches, _ = check._find_deployments_on_same_encryption_boundary([undeclared_twin], "region-1") + + assert matches == [undeclared_twin] + + +def test_declared_boundary_matches_deployments_whose_credentials_cannot_be_resolved(): + origin: Final = _make_originating_mock("https://region-1.example.invalid/v1", "") + origin.litellm_params.model_dump.return_value = {"api_base": "https://region-1.example.invalid/v1"} + origin.model_info = litellm.types.router.ModelInfo(id="region-1", encrypted_content_boundary="org-1") + check: Final = EncryptedContentAffinityCheck(ctx=_make_router_mock_with_cooldown(origin)) + peer: Final = { + "model_info": {"id": "region-2", "encrypted_content_boundary": "org-1"}, + "litellm_params": {"api_base": "https://region-2.example.invalid/v1"}, + } + + matches, _ = check._find_deployments_on_same_encryption_boundary([peer], "region-1") + + assert matches == [peer] + + +async def _forwarded_reasoning_after_order_fallback(boundary: str | None) -> tuple: + primary: Final = _regional_deployment("reasoning", "region-1", boundary) + backup: Final = _regional_deployment("reasoning", "region-3", boundary) + router: Final = litellm.Router( + model_list=[ + primary | {"litellm_params": primary["litellm_params"] | {"order": 1}}, + backup | {"litellm_params": backup["litellm_params"] | {"order": 2}}, + ], + optional_pre_call_checks=["encrypted_content_affinity"], + disable_cooldowns=True, + num_retries=0, + ) + + async def upstream(*, litellm_params, **_kwargs): + if "region-1" in litellm_params.api_base: + raise litellm.InternalServerError("upstream 500", llm_provider="openai", model="gpt-5.6") + return _build_mock_response(output_items=[], response_id="resp_second") + + handler: Final = AsyncMock(side_effect=upstream) + with patch( + "litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler.async_response_api_handler", + new=handler, + ): + response: Final = await router.aresponses( + model="reasoning", + input=[ + { + "type": "reasoning", + "summary": [], + "encrypted_content": ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id( + "blob-from-region-1", "region-1" + ), + }, + {"role": "user", "content": "follow-up"}, + ], + litellm_metadata={}, + ) + forwarded_reasoning: Final = handler.await_args.kwargs["input"][0] + return response._hidden_params["model_id"], forwarded_reasoning.get("encrypted_content") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("boundary", "expected_reasoning"), + [("org-1", "blob-from-region-1"), (None, None)], + ids=["declared", "undeclared"], +) +async def test_order_fallback_keeps_reasoning_only_on_a_declared_boundary_peer(boundary, expected_reasoning): + assert await _forwarded_reasoning_after_order_fallback(boundary) == ("region-3", expected_reasoning) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 6f8c49cc0f9..b171a72dcf9 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -52372,6 +52372,8 @@ export interface components { discoverable?: boolean | null; /** Enable Tag Filtering */ enable_tag_filtering?: boolean | null; + /** Encrypted Content Boundary */ + encrypted_content_boundary?: string | null; /** Enforce Supported Endpoints */ enforce_supported_endpoints?: boolean | null; /** Explicit Cache Read Input Token Cost */ From df20f90d12a7d3ab40f02aa66f6442c2ec7ec247 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Mon, 9 Nov 2026 15:26:40 -0800 Subject: [PATCH 25/46] test(integration): cover claude-fable-5-2 capability flags and prompt-cache pricing on Vertex AI and Bedrock Converse (#50742) * test(integration): cover claude-fable-5-2 on Vertex AI and Bedrock Converse Scripted-upstream wire tests for the Fable capability flags on both routes, the vertex_ai and us.anthropic cost cases every other Fable row has, Fable 5.2 in BEDROCK_CONVERSE_MODELS, and bedrock/claude-fable-latest repointed to 5.2 like the Anthropic and Vertex aliases already are. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(cost-map): add GovCloud Bedrock rows for claude-fable-5-2 bedrock/us-gov-west-1 and bedrock/us-gov-east-1 rows copy the Fable 5.1 field set with the GovCloud rates from the Bedrock price list. 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/constants.py | 1 + ...odel_prices_and_context_window_backup.json | 141 +++--- model_prices_and_context_window.json | 141 +++--- .../cases/bedrock_converse.py | 432 ++++++++++++++++++ .../cost_calculation/cases/vertex_ai.py | 149 +++++- .../integration/cost_calculation/cost_map.py | 35 ++ ...ck_converse_fable_capability_flags_wire.py | 234 ++++++++++ ...x_anthropic_fable_capability_flags_wire.py | 242 ++++++++++ tests/unit/test_claude_fable_5_config.py | 16 + 9 files changed, 1250 insertions(+), 141 deletions(-) create mode 100644 tests/integration/providers/test_bedrock_converse_fable_capability_flags_wire.py create mode 100644 tests/integration/providers/test_vertex_anthropic_fable_capability_flags_wire.py diff --git a/litellm/constants.py b/litellm/constants.py index 3a49e8a7fc4..6cca43f8c8b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1533,6 +1533,7 @@ BEDROCK_CONVERSE_MODELS: Final = [ "openai.gpt-oss-120b-1:0", "anthropic.claude-haiku-4-5-20251001-v1:0", "anthropic.claude-sonnet-4-5-20250929-v1:0", + "anthropic.claude-fable-5-2", "anthropic.claude-fable-5-1", "anthropic.claude-fable-5", "anthropic.claude-sonnet-5", diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index cc920c5fe25..0d35c643561 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1970,7 +1970,7 @@ "source": "https://aws.amazon.com/bedrock/pricing/" }, "bedrock/claude-fable-latest": { - "alias_of": "global.anthropic.claude-fable-5-1" + "alias_of": "global.anthropic.claude-fable-5-2" }, "us.anthropic.claude-fable-5": { "cache_creation_input_token_cost": 1.375e-05, @@ -69958,6 +69958,44 @@ "region": "us-gov-west-1" } }, + "bedrock/us-gov-west-1/anthropic.claude-fable-5-2": { + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh", + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512, + "token_count_model": { + "model": "anthropic.claude-fable-5-2", + "litellm_provider": "bedrock_mantle", + "region": "us-gov-west-1" + } + }, "bedrock/us-gov-east-1/nvidia.nemotron-nano-3-30b": { "input_cost_per_token": 7.2e-08, "input_cost_per_token_batches": 3.6e-08, @@ -70266,6 +70304,39 @@ "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512 }, + "bedrock/us-gov-east-1/anthropic.claude-fable-5-2": { + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh", + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, "bedrock_mantle/us-gov-west-1/openai.gpt-5.6-terra": { "litellm_provider": "bedrock_mantle", "max_input_tokens": 922000, @@ -87903,74 +87974,6 @@ "supports_xhigh_reasoning_effort": true, "thinking_always_on": true }, - "bedrock/us-gov-east-1/anthropic.claude-fable-5-2": { - "bedrock_output_config_effort_ceiling": "xhigh", - "cache_creation_input_token_cost": 1.2e-05, - "cache_creation_input_token_cost_above_1hr": 1.92e-05, - "cache_read_input_token_cost": 2.4e-07, - "input_cost_per_token": 9.6e-06, - "litellm_provider": "bedrock", - "max_input_tokens": 1000000, - "max_output_tokens": 128000, - "mode": "chat", - "output_cost_per_token": 4.8e-05, - "prompt_cache_min_tokens": 512.0, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", - "supports_adaptive_thinking": true, - "supports_assistant_prefill": false, - "supports_bedrock_native_structured_output": false, - "supports_computer_use": true, - "supports_forced_tool_use": false, - "supports_function_calling": true, - "supports_max_reasoning_effort": true, - "supports_mid_conversation_system": true, - "supports_native_structured_output": false, - "supports_output_config": true, - "supports_parallel_tool_use_config": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_sampling_params": false, - "supports_tool_choice": true, - "supports_vision": true, - "supports_xhigh_reasoning_effort": true, - "thinking_always_on": true - }, - "bedrock/us-gov-west-1/anthropic.claude-fable-5-2": { - "bedrock_output_config_effort_ceiling": "xhigh", - "cache_creation_input_token_cost": 1.2e-05, - "cache_creation_input_token_cost_above_1hr": 1.92e-05, - "cache_read_input_token_cost": 2.4e-07, - "input_cost_per_token": 9.6e-06, - "litellm_provider": "bedrock", - "max_input_tokens": 1000000, - "max_output_tokens": 128000, - "mode": "chat", - "output_cost_per_token": 4.8e-05, - "prompt_cache_min_tokens": 512.0, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", - "supports_adaptive_thinking": true, - "supports_assistant_prefill": false, - "supports_bedrock_native_structured_output": false, - "supports_computer_use": true, - "supports_forced_tool_use": false, - "supports_function_calling": true, - "supports_max_reasoning_effort": true, - "supports_mid_conversation_system": true, - "supports_native_structured_output": false, - "supports_output_config": true, - "supports_parallel_tool_use_config": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_sampling_params": false, - "supports_tool_choice": true, - "supports_vision": true, - "supports_xhigh_reasoning_effort": true, - "thinking_always_on": true - }, "bedrock/us-gov-east-1/anthropic.claude-mythos-5-2": { "bedrock_output_config_effort_ceiling": "xhigh", "cache_creation_input_token_cost": 1.2e-05, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index cc920c5fe25..0d35c643561 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1970,7 +1970,7 @@ "source": "https://aws.amazon.com/bedrock/pricing/" }, "bedrock/claude-fable-latest": { - "alias_of": "global.anthropic.claude-fable-5-1" + "alias_of": "global.anthropic.claude-fable-5-2" }, "us.anthropic.claude-fable-5": { "cache_creation_input_token_cost": 1.375e-05, @@ -69958,6 +69958,44 @@ "region": "us-gov-west-1" } }, + "bedrock/us-gov-west-1/anthropic.claude-fable-5-2": { + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh", + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512, + "token_count_model": { + "model": "anthropic.claude-fable-5-2", + "litellm_provider": "bedrock_mantle", + "region": "us-gov-west-1" + } + }, "bedrock/us-gov-east-1/nvidia.nemotron-nano-3-30b": { "input_cost_per_token": 7.2e-08, "input_cost_per_token_batches": 3.6e-08, @@ -70266,6 +70304,39 @@ "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512 }, + "bedrock/us-gov-east-1/anthropic.claude-fable-5-2": { + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh", + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, "bedrock_mantle/us-gov-west-1/openai.gpt-5.6-terra": { "litellm_provider": "bedrock_mantle", "max_input_tokens": 922000, @@ -87903,74 +87974,6 @@ "supports_xhigh_reasoning_effort": true, "thinking_always_on": true }, - "bedrock/us-gov-east-1/anthropic.claude-fable-5-2": { - "bedrock_output_config_effort_ceiling": "xhigh", - "cache_creation_input_token_cost": 1.2e-05, - "cache_creation_input_token_cost_above_1hr": 1.92e-05, - "cache_read_input_token_cost": 2.4e-07, - "input_cost_per_token": 9.6e-06, - "litellm_provider": "bedrock", - "max_input_tokens": 1000000, - "max_output_tokens": 128000, - "mode": "chat", - "output_cost_per_token": 4.8e-05, - "prompt_cache_min_tokens": 512.0, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", - "supports_adaptive_thinking": true, - "supports_assistant_prefill": false, - "supports_bedrock_native_structured_output": false, - "supports_computer_use": true, - "supports_forced_tool_use": false, - "supports_function_calling": true, - "supports_max_reasoning_effort": true, - "supports_mid_conversation_system": true, - "supports_native_structured_output": false, - "supports_output_config": true, - "supports_parallel_tool_use_config": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_sampling_params": false, - "supports_tool_choice": true, - "supports_vision": true, - "supports_xhigh_reasoning_effort": true, - "thinking_always_on": true - }, - "bedrock/us-gov-west-1/anthropic.claude-fable-5-2": { - "bedrock_output_config_effort_ceiling": "xhigh", - "cache_creation_input_token_cost": 1.2e-05, - "cache_creation_input_token_cost_above_1hr": 1.92e-05, - "cache_read_input_token_cost": 2.4e-07, - "input_cost_per_token": 9.6e-06, - "litellm_provider": "bedrock", - "max_input_tokens": 1000000, - "max_output_tokens": 128000, - "mode": "chat", - "output_cost_per_token": 4.8e-05, - "prompt_cache_min_tokens": 512.0, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", - "supports_adaptive_thinking": true, - "supports_assistant_prefill": false, - "supports_bedrock_native_structured_output": false, - "supports_computer_use": true, - "supports_forced_tool_use": false, - "supports_function_calling": true, - "supports_max_reasoning_effort": true, - "supports_mid_conversation_system": true, - "supports_native_structured_output": false, - "supports_output_config": true, - "supports_parallel_tool_use_config": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_sampling_params": false, - "supports_tool_choice": true, - "supports_vision": true, - "supports_xhigh_reasoning_effort": true, - "thinking_always_on": true - }, "bedrock/us-gov-east-1/anthropic.claude-mythos-5-2": { "bedrock_output_config_effort_ceiling": "xhigh", "cache_creation_input_token_cost": 1.2e-05, diff --git a/tests/integration/cost_calculation/cases/bedrock_converse.py b/tests/integration/cost_calculation/cases/bedrock_converse.py index 858db06cafc..96033a8b49a 100644 --- a/tests/integration/cost_calculation/cases/bedrock_converse.py +++ b/tests/integration/cost_calculation/cases/bedrock_converse.py @@ -4672,6 +4672,57 @@ US_ANTHROPIC_CLAUDE_FABLE_5_1_INPUT_TEXT: Final = CostTrackingTestCase( spend=0.0702, input_cost=0.03312, output_cost=0.03708, prompt_tokens=1840, completion_tokens=412 ), ) +US_ANTHROPIC_CLAUDE_FABLE_5_2_INPUT_TEXT: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-2-input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "a5b0ca93d2e1 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer a5b0ca93d2e1"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + "inputTokens": 1840, + "outputTokens": 412, + "serverToolUsage": {}, + "totalTokens": 2252, + }, + }, + ), + expected=ExactExpected( + spend=0.05616, input_cost=0.026496, output_cost=0.029664, prompt_tokens=1840, completion_tokens=412 + ), +) US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM: Final = CostTrackingTestCase( name="us.anthropic.claude-fable-5-1-stream", @@ -4726,6 +4777,59 @@ US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM: Final = CostTrackingTestCase( spend=0.0702, input_cost=0.03312, output_cost=0.03708, prompt_tokens=1840, completion_tokens=412 ), ) +US_ANTHROPIC_CLAUDE_FABLE_5_2_STREAM: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-2-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="us.anthropic.claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "d62b383256ff summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "stream_options": {"include_usage": True}, + "allowed_openai_params": [], + }, + response=EventStreamResponse( + content_type="application/vnd.amazon.eventstream", + events=( + EventStreamEvent(event_type="messageStart", payload={"role": "assistant"}), + EventStreamEvent( + event_type="contentBlockDelta", + payload={"contentBlockIndex": 0, "delta": {"text": "scripted answer d62b383256ff"}}, + ), + EventStreamEvent(event_type="contentBlockStop", payload={"contentBlockIndex": 0}), + EventStreamEvent(event_type="messageStop", payload={"stopReason": "end_turn"}), + EventStreamEvent( + event_type="metadata", + payload={ + "metrics": {"latencyMs": 883}, + "usage": {"inputTokens": 1840, "outputTokens": 412, "serverToolUsage": {}, "totalTokens": 2252}, + }, + ), + ), + ), + expected=ExactExpected( + spend=0.05616, input_cost=0.026496, output_cost=0.029664, prompt_tokens=1840, completion_tokens=412 + ), +) US_ANTHROPIC_CLAUDE_FABLE_5_1_CACHE_READ: Final = CostTrackingTestCase( name="us.anthropic.claude-fable-5-1-cache_read", @@ -4779,6 +4883,58 @@ US_ANTHROPIC_CLAUDE_FABLE_5_1_CACHE_READ: Final = CostTrackingTestCase( spend=0.0678384, input_cost=0.0336384, output_cost=0.0342, prompt_tokens=12928, completion_tokens=380 ), ) +US_ANTHROPIC_CLAUDE_FABLE_5_2_CACHE_READ: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-2-cache_read", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "1a5a397f1134 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer 1a5a397f1134"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "cacheReadInputTokenCount": 12288, + "cacheReadInputTokens": 12288, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + "inputTokens": 640, + "outputTokens": 380, + "serverToolUsage": {}, + "totalTokens": 13308, + }, + }, + ), + expected=ExactExpected( + spend=0.05427072, input_cost=0.02691072, output_cost=0.02736, prompt_tokens=12928, completion_tokens=380 + ), +) US_ANTHROPIC_CLAUDE_FABLE_5_1_CACHE_WRITE_5M: Final = CostTrackingTestCase( name="us.anthropic.claude-fable-5-1-cache_write_5m", @@ -4830,6 +4986,56 @@ US_ANTHROPIC_CLAUDE_FABLE_5_1_CACHE_WRITE_5M: Final = CostTrackingTestCase( spend=0.248076, input_cost=0.216576, output_cost=0.0315, prompt_tokens=9728, completion_tokens=350 ), ) +US_ANTHROPIC_CLAUDE_FABLE_5_2_CACHE_WRITE_5M: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-2-cache_write_5m", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "3481bca67e1a summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer 3481bca67e1a"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 512, + "outputTokens": 350, + "totalTokens": 10078, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 9216, + "cacheDetails": [{"ttl": "5m", "inputTokens": 9216}], + }, + }, + ), + expected=ExactExpected( + spend=0.1984608, input_cost=0.1732608, output_cost=0.0252, prompt_tokens=9728, completion_tokens=350 + ), +) US_ANTHROPIC_CLAUDE_FABLE_5_1_CACHE_WRITE_1H: Final = CostTrackingTestCase( name="us.anthropic.claude-fable-5-1-cache_write_1h", @@ -4881,6 +5087,56 @@ US_ANTHROPIC_CLAUDE_FABLE_5_1_CACHE_WRITE_1H: Final = CostTrackingTestCase( spend=0.372492, input_cost=0.340992, output_cost=0.0315, prompt_tokens=9728, completion_tokens=350 ), ) +US_ANTHROPIC_CLAUDE_FABLE_5_2_CACHE_WRITE_1H: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-2-cache_write_1h", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "3f037d5317b0 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer 3f037d5317b0"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 512, + "outputTokens": 350, + "totalTokens": 10078, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 9216, + "cacheDetails": [{"ttl": "1h", "inputTokens": 9216}], + }, + }, + ), + expected=ExactExpected( + spend=0.2979936, input_cost=0.2727936, output_cost=0.0252, prompt_tokens=9728, completion_tokens=350 + ), +) US_ANTHROPIC_CLAUDE_SONNET_5_TOOL_CALL: Final = CostTrackingTestCase( name="us.anthropic.claude-sonnet-5-tool_call", @@ -5832,6 +6088,87 @@ US_ANTHROPIC_CLAUDE_FABLE_5_1_TOOL_CALL: Final = CostTrackingTestCase( spend=0.0702, input_cost=0.03312, output_cost=0.03708, prompt_tokens=1840, completion_tokens=412 ), ) +US_ANTHROPIC_CLAUDE_FABLE_5_2_TOOL_CALL: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-2-tool_call", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "7f97ab74e808 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather and a short forecast for a city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"}, + "days": {"type": "integer", "description": "Forecast horizon in days"}, + "units": {"type": "string", "enum": ["metric", "imperial"]}, + }, + "required": ["city"], + }, + }, + } + ], + "tool_choice": "auto", + "allowed_openai_params": ["tool_choice"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "output": { + "message": { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "tooluse_$REQUEST_ID", + "name": "get_weather", + "input": {"city": "Berlin", "days": 7, "units": "metric"}, + "type": "tool_use", + } + } + ], + } + }, + "stopReason": "tool_use", + "usage": { + "inputTokens": 1840, + "outputTokens": 412, + "totalTokens": 2252, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 0, + }, + "metrics": {"latencyMs": 2071}, + }, + ), + expected=ExactExpected( + spend=0.05616, input_cost=0.026496, output_cost=0.029664, prompt_tokens=1840, completion_tokens=412 + ), +) US_ANTHROPIC_CLAUDE_FABLE_5_1_THINKING: Final = CostTrackingTestCase( name="us.anthropic.claude-fable-5-1-thinking", @@ -5885,6 +6222,58 @@ US_ANTHROPIC_CLAUDE_FABLE_5_1_THINKING: Final = CostTrackingTestCase( spend=0.0594, input_cost=0.02232, output_cost=0.03708, prompt_tokens=1240, completion_tokens=412 ), ) +US_ANTHROPIC_CLAUDE_FABLE_5_2_THINKING: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-2-thinking", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "bca700fd52a2 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "output": {"message": {"role": "assistant", "content": [{"text": "scripted answer bca700fd52a2"}]}}, + "stopReason": "end_turn", + "usage": { + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + "inputTokens": 1240, + "outputTokens": 412, + "serverToolUsage": {}, + "totalTokens": 1652, + }, + "metrics": {"latencyMs": 2071}, + }, + ), + expected=ExactExpected( + spend=0.04752, input_cost=0.017856, output_cost=0.029664, prompt_tokens=1240, completion_tokens=412 + ), +) US_ANTHROPIC_CLAUDE_OPUS_5_5_UPSTREAM_400_ZERO_SPEND: Final = CostTrackingTestCase( name="us.anthropic.claude-opus-5-5-upstream_400_zero_spend", @@ -6014,6 +6403,13 @@ US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON: Final = US_ANTHROPIC_CLAUD "request": {"model": "$MODEL", "input": "summarize the attached material in one line"}, } ) +US_ANTHROPIC_CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_JSON: Final = US_ANTHROPIC_CLAUDE_FABLE_5_2_INPUT_TEXT.model_copy( + update={ + "name": "us.anthropic.claude-fable-5-2-responses_bridged_json", + "endpoint": "/v1/responses", + "request": {"model": "$MODEL", "input": "summarize the attached material in one line"}, + } +) US_ANTHROPIC_CLAUDE_HAIKU_5_5_TIERED_INPUT_ABOVE_100K: Final = with_usage( US_ANTHROPIC_CLAUDE_HAIKU_5_5_INPUT_TEXT, @@ -6088,6 +6484,7 @@ US_ANTHROPIC_CLAUDE_HAIKU_5_5_TIERED_CACHE_WRITE_1H_ABOVE_100K: Final = with_usa ), ) + def _us_anthropic_messages_bridge_pairs( input_case: CostTrackingTestCase, stream_case: CostTrackingTestCase ) -> tuple[BridgeCostParityTestCase, ...]: @@ -6125,6 +6522,11 @@ US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM: Final = _us_anthropic_re US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM, "us.anthropic.claude-fable-5-1-responses_bridged_stream", ) +US_ANTHROPIC_CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_STREAM: Final = _us_anthropic_responses_stream_case( + US_ANTHROPIC_CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_JSON, + US_ANTHROPIC_CLAUDE_FABLE_5_2_STREAM, + "us.anthropic.claude-fable-5-2-responses_bridged_stream", +) US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM: Final = _us_anthropic_responses_stream_case( US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, US_ANTHROPIC_CLAUDE_HAIKU_5_5_STREAM, @@ -6143,6 +6545,9 @@ US_ANTHROPIC_CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM: Final = _us_anthropic_r _US_ANTHROPIC_FABLE_5_1_MESSAGES_PAIRS: Final = _us_anthropic_messages_bridge_pairs( US_ANTHROPIC_CLAUDE_FABLE_5_1_INPUT_TEXT, US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM ) +_US_ANTHROPIC_FABLE_5_2_MESSAGES_PAIRS: Final = _us_anthropic_messages_bridge_pairs( + US_ANTHROPIC_CLAUDE_FABLE_5_2_INPUT_TEXT, US_ANTHROPIC_CLAUDE_FABLE_5_2_STREAM +) _US_ANTHROPIC_HAIKU_5_5_MESSAGES_PAIRS: Final = _us_anthropic_messages_bridge_pairs( US_ANTHROPIC_CLAUDE_HAIKU_5_5_INPUT_TEXT, US_ANTHROPIC_CLAUDE_HAIKU_5_5_STREAM ) @@ -6161,6 +6566,7 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, US_ANTHROPIC_CLAUDE_FABLE_5_RESPONSES_BRIDGED_JSON, US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON, + US_ANTHROPIC_CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_JSON, US_ANTHROPIC_CLAUDE_HAIKU_5_5_TIERED_INPUT_ABOVE_100K, US_ANTHROPIC_CLAUDE_HAIKU_5_5_TIERED_AT_100K_BOUNDARY, US_ANTHROPIC_CLAUDE_HAIKU_5_5_TIERED_CACHE_READ_ABOVE_100K, @@ -6255,10 +6661,15 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( US_ANTHROPIC_CLAUDE_FABLE_5_CACHE_WRITE_5M, US_ANTHROPIC_CLAUDE_FABLE_5_CACHE_WRITE_1H, US_ANTHROPIC_CLAUDE_FABLE_5_1_INPUT_TEXT, + US_ANTHROPIC_CLAUDE_FABLE_5_2_INPUT_TEXT, US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM, + US_ANTHROPIC_CLAUDE_FABLE_5_2_STREAM, US_ANTHROPIC_CLAUDE_FABLE_5_1_CACHE_READ, + US_ANTHROPIC_CLAUDE_FABLE_5_2_CACHE_READ, US_ANTHROPIC_CLAUDE_FABLE_5_1_CACHE_WRITE_5M, + US_ANTHROPIC_CLAUDE_FABLE_5_2_CACHE_WRITE_5M, US_ANTHROPIC_CLAUDE_FABLE_5_1_CACHE_WRITE_1H, + US_ANTHROPIC_CLAUDE_FABLE_5_2_CACHE_WRITE_1H, US_ANTHROPIC_CLAUDE_SONNET_5_TOOL_CALL, US_ANTHROPIC_CLAUDE_SONNET_5_THINKING, US_ANTHROPIC_CLAUDE_SONNET_5_5_TOOL_CALL, @@ -6272,14 +6683,18 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( US_ANTHROPIC_CLAUDE_FABLE_5_TOOL_CALL, US_ANTHROPIC_CLAUDE_FABLE_5_THINKING, US_ANTHROPIC_CLAUDE_FABLE_5_1_TOOL_CALL, + US_ANTHROPIC_CLAUDE_FABLE_5_2_TOOL_CALL, US_ANTHROPIC_CLAUDE_FABLE_5_1_THINKING, + US_ANTHROPIC_CLAUDE_FABLE_5_2_THINKING, US_ANTHROPIC_CLAUDE_OPUS_5_5_UPSTREAM_400_ZERO_SPEND, US_ANTHROPIC_CLAUDE_OPUS_5_5_INPUT_TEXT_ROLLUPS, US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM, + US_ANTHROPIC_CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_STREAM, US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM, US_ANTHROPIC_CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_STREAM, US_ANTHROPIC_CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM, *tuple(pair.bridged for pair in _US_ANTHROPIC_FABLE_5_1_MESSAGES_PAIRS), + *tuple(pair.bridged for pair in _US_ANTHROPIC_FABLE_5_2_MESSAGES_PAIRS), *tuple(pair.bridged for pair in _US_ANTHROPIC_HAIKU_5_5_MESSAGES_PAIRS), *tuple(pair.bridged for pair in _US_ANTHROPIC_OPUS_5_5_MESSAGES_PAIRS), *tuple(pair.bridged for pair in _US_ANTHROPIC_SONNET_5_5_MESSAGES_PAIRS), @@ -6308,10 +6723,15 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase(plain=US_ANTHROPIC_CLAUDE_HAIKU_5_5_INPUT_TEXT, streamed=US_ANTHROPIC_CLAUDE_HAIKU_5_5_STREAM), StreamParityTestCase(plain=US_ANTHROPIC_CLAUDE_FABLE_5_INPUT_TEXT, streamed=US_ANTHROPIC_CLAUDE_FABLE_5_STREAM), StreamParityTestCase(plain=US_ANTHROPIC_CLAUDE_FABLE_5_1_INPUT_TEXT, streamed=US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM), + StreamParityTestCase(plain=US_ANTHROPIC_CLAUDE_FABLE_5_2_INPUT_TEXT, streamed=US_ANTHROPIC_CLAUDE_FABLE_5_2_STREAM), StreamParityTestCase( plain=US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON, streamed=US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM, ), + StreamParityTestCase( + plain=US_ANTHROPIC_CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_JSON, + streamed=US_ANTHROPIC_CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_STREAM, + ), StreamParityTestCase( plain=US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, streamed=US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM, @@ -6328,6 +6748,10 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( plain=_US_ANTHROPIC_FABLE_5_1_MESSAGES_PAIRS[0].bridged, streamed=_US_ANTHROPIC_FABLE_5_1_MESSAGES_PAIRS[1].bridged, ), + StreamParityTestCase( + plain=_US_ANTHROPIC_FABLE_5_2_MESSAGES_PAIRS[0].bridged, + streamed=_US_ANTHROPIC_FABLE_5_2_MESSAGES_PAIRS[1].bridged, + ), StreamParityTestCase( plain=_US_ANTHROPIC_HAIKU_5_5_MESSAGES_PAIRS[0].bridged, streamed=_US_ANTHROPIC_HAIKU_5_5_MESSAGES_PAIRS[1].bridged, @@ -6372,10 +6796,17 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( BridgeCostParityTestCase( native=US_ANTHROPIC_CLAUDE_FABLE_5_1_INPUT_TEXT, bridged=US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON ), + BridgeCostParityTestCase( + native=US_ANTHROPIC_CLAUDE_FABLE_5_2_INPUT_TEXT, bridged=US_ANTHROPIC_CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_JSON + ), BridgeCostParityTestCase( native=US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM, bridged=US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM, ), + BridgeCostParityTestCase( + native=US_ANTHROPIC_CLAUDE_FABLE_5_2_STREAM, + bridged=US_ANTHROPIC_CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_STREAM, + ), BridgeCostParityTestCase( native=US_ANTHROPIC_CLAUDE_HAIKU_5_5_STREAM, bridged=US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM, @@ -6389,6 +6820,7 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( bridged=US_ANTHROPIC_CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM, ), *_US_ANTHROPIC_FABLE_5_1_MESSAGES_PAIRS, + *_US_ANTHROPIC_FABLE_5_2_MESSAGES_PAIRS, *_US_ANTHROPIC_HAIKU_5_5_MESSAGES_PAIRS, *_US_ANTHROPIC_OPUS_5_5_MESSAGES_PAIRS, *_US_ANTHROPIC_SONNET_5_5_MESSAGES_PAIRS, diff --git a/tests/integration/cost_calculation/cases/vertex_ai.py b/tests/integration/cost_calculation/cases/vertex_ai.py index 8590a19646b..47a484e972b 100644 --- a/tests/integration/cost_calculation/cases/vertex_ai.py +++ b/tests/integration/cost_calculation/cases/vertex_ai.py @@ -5,8 +5,7 @@ PARITY pairs the streamed case with its plain twin where both bill the same row. import json from typing import Final, Literal, cast -from pydantic import JsonValue - +from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase, bridge_pairs from integration.cost_calculation.cases.anthropic import ( CLAUDE_FABLE_5_1_CACHE_READ, CLAUDE_FABLE_5_1_CACHE_WRITE_1H, @@ -17,6 +16,15 @@ from integration.cost_calculation.cases.anthropic import ( CLAUDE_FABLE_5_1_STREAM, CLAUDE_FABLE_5_1_THINKING, CLAUDE_FABLE_5_1_TOOL_CALL, + CLAUDE_FABLE_5_2_CACHE_READ, + CLAUDE_FABLE_5_2_CACHE_WRITE_1H, + CLAUDE_FABLE_5_2_CACHE_WRITE_5M, + CLAUDE_FABLE_5_2_INPUT_TEXT, + CLAUDE_FABLE_5_2_MESSAGES_INPUT_TEXT, + CLAUDE_FABLE_5_2_MESSAGES_STREAM, + CLAUDE_FABLE_5_2_STREAM, + CLAUDE_FABLE_5_2_THINKING, + CLAUDE_FABLE_5_2_TOOL_CALL, CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, CLAUDE_HAIKU_5_5_MESSAGES_STREAM, CLAUDE_OPUS_5_5_CACHE_READ, @@ -38,7 +46,6 @@ from integration.cost_calculation.cases.anthropic import ( CLAUDE_SONNET_5_5_THINKING, CLAUDE_SONNET_5_5_TOOL_CALL, ) -from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase, bridge_pairs from integration.cost_calculation.cost_tracking_case import ( CostTrackingTestCase, Deployment, @@ -49,6 +56,7 @@ from integration.cost_calculation.cost_tracking_case import ( SseResponse, ) from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames +from pydantic import JsonValue def _gemini_case_response( @@ -3009,6 +3017,37 @@ VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT: Final = CLAUDE_FABLE_5_1_INPUT_TEXT.model ), } ) +VERTEX_AI_CLAUDE_FABLE_5_2_INPUT_TEXT: Final = CLAUDE_FABLE_5_2_INPUT_TEXT.model_copy( + update={ + "name": "vertex_ai/claude-fable-5-2-input_text", + "model": "vertex_ai/claude-fable-5-2", + "response": JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5-2", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "thinking", "thinking": "scripted thought", "signature": "scripted signature"}], + "container": None, + "stop_reason": "max_tokens", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + }, + }, + ), + "expected": ExactExpected( + spend=0.05928, input_cost=0.027968, output_cost=0.031312, prompt_tokens=1840, completion_tokens=412 + ), + } +) VERTEX_AI_CLAUDE_FABLE_5_1_STREAM: Final = CLAUDE_FABLE_5_1_STREAM.model_copy( update={ "name": "vertex_ai/claude-fable-5-1-stream", @@ -3017,6 +3056,14 @@ VERTEX_AI_CLAUDE_FABLE_5_1_STREAM: Final = CLAUDE_FABLE_5_1_STREAM.model_copy( "expected": VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT.expected, } ) +VERTEX_AI_CLAUDE_FABLE_5_2_STREAM: Final = CLAUDE_FABLE_5_2_STREAM.model_copy( + update={ + "name": "vertex_ai/claude-fable-5-2-stream", + "model": "vertex_ai/claude-fable-5-2", + "response": _vertex_claude_stream_response("claude-fable-5-2", "max_tokens", ()), + "expected": VERTEX_AI_CLAUDE_FABLE_5_2_INPUT_TEXT.expected, + } +) VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_READ: Final = _vertex_claude_shape_case( CLAUDE_SONNET_5_5_CACHE_READ, @@ -3135,6 +3182,15 @@ VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_READ: Final = _vertex_claude_shape_case( spend=0.013567425, input_cost=0.013092425, output_cost=0.000475, prompt_tokens=27017, completion_tokens=5 ), ) +VERTEX_AI_CLAUDE_FABLE_5_2_CACHE_READ: Final = _vertex_claude_shape_case( + CLAUDE_FABLE_5_2_CACHE_READ, + "vertex_ai/claude-fable-5-2-cache_read", + "vertex_ai/claude-fable-5-2", + _vertex_claude_usage(input_tokens=14, output_tokens=5, cache_read_input_tokens=27003), + ExactExpected( + spend=0.01085394, input_cost=0.01047394, output_cost=0.00038, prompt_tokens=27017, completion_tokens=5 + ), +) VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_5M: Final = _vertex_claude_shape_case( CLAUDE_FABLE_5_1_CACHE_WRITE_5M, "vertex_ai/claude-fable-5-1-cache_write_5m", @@ -3146,6 +3202,17 @@ VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_5M: Final = _vertex_claude_shape_case( spend=0.64225225, input_cost=0.64177725, output_cost=0.000475, prompt_tokens=27025, completion_tokens=5 ), ) +VERTEX_AI_CLAUDE_FABLE_5_2_CACHE_WRITE_5M: Final = _vertex_claude_shape_case( + CLAUDE_FABLE_5_2_CACHE_WRITE_5M, + "vertex_ai/claude-fable-5-2-cache_write_5m", + "vertex_ai/claude-fable-5-2", + _vertex_claude_usage( + input_tokens=14, output_tokens=5, cache_creation_input_tokens=27011, ephemeral_5m_input_tokens=27011 + ), + ExactExpected( + spend=0.5138018, input_cost=0.5134218, output_cost=0.00038, prompt_tokens=27025, completion_tokens=5 + ), +) VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_1H: Final = _vertex_claude_shape_case( CLAUDE_FABLE_5_1_CACHE_WRITE_1H, "vertex_ai/claude-fable-5-1-cache_write_1h", @@ -3157,6 +3224,17 @@ VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_1H: Final = _vertex_claude_shape_case( spend=1.027159, input_cost=1.026684, output_cost=0.000475, prompt_tokens=27025, completion_tokens=5 ), ) +VERTEX_AI_CLAUDE_FABLE_5_2_CACHE_WRITE_1H: Final = _vertex_claude_shape_case( + CLAUDE_FABLE_5_2_CACHE_WRITE_1H, + "vertex_ai/claude-fable-5-2-cache_write_1h", + "vertex_ai/claude-fable-5-2", + _vertex_claude_usage( + input_tokens=14, output_tokens=5, cache_creation_input_tokens=27011, ephemeral_1h_input_tokens=27011 + ), + ExactExpected( + spend=0.8217272, input_cost=0.8213472, output_cost=0.00038, prompt_tokens=27025, completion_tokens=5 + ), +) VERTEX_AI_CLAUDE_FABLE_5_1_THINKING: Final = _vertex_claude_shape_case( CLAUDE_FABLE_5_1_THINKING, "vertex_ai/claude-fable-5-1-thinking", @@ -3171,6 +3249,20 @@ VERTEX_AI_CLAUDE_FABLE_5_1_THINKING: Final = _vertex_claude_shape_case( ], adaptive_thinking=True, ) +VERTEX_AI_CLAUDE_FABLE_5_2_THINKING: Final = _vertex_claude_shape_case( + CLAUDE_FABLE_5_2_THINKING, + "vertex_ai/claude-fable-5-2-thinking", + "vertex_ai/claude-fable-5-2", + _vertex_claude_usage(input_tokens=18, output_tokens=98, thinking_tokens=33), + ExactExpected( + spend=0.0077216, input_cost=0.0002736, output_cost=0.007448, prompt_tokens=18, completion_tokens=98 + ), + [ + {"type": "thinking", "thinking": "scripted thought", "signature": "scripted signature"}, + {"type": "text", "text": "scripted answer for cost fixtures"}, + ], + adaptive_thinking=True, +) VERTEX_AI_CLAUDE_FABLE_5_1_TOOL_CALL: Final = _vertex_claude_shape_case( CLAUDE_FABLE_5_1_TOOL_CALL, "vertex_ai/claude-fable-5-1-tool_call", @@ -3180,6 +3272,15 @@ VERTEX_AI_CLAUDE_FABLE_5_1_TOOL_CALL: Final = _vertex_claude_shape_case( spend=0.012065, input_cost=0.007315, output_cost=0.00475, prompt_tokens=385, completion_tokens=50 ), ) +VERTEX_AI_CLAUDE_FABLE_5_2_TOOL_CALL: Final = _vertex_claude_shape_case( + CLAUDE_FABLE_5_2_TOOL_CALL, + "vertex_ai/claude-fable-5-2-tool_call", + "vertex_ai/claude-fable-5-2", + _vertex_claude_usage(input_tokens=385, output_tokens=50), + ExactExpected( + spend=0.009652, input_cost=0.005852, output_cost=0.0038, prompt_tokens=385, completion_tokens=50 + ), +) VERTEX_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT: Final = _vertex_claude_shape_case( VERTEX_AI_CLAUDE_SONNET_5_5_INPUT_TEXT, @@ -3325,6 +3426,15 @@ VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT: Final = CLAUDE_FABLE_5_1_MESSAGE "expected": VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT.expected, } ) +VERTEX_AI_CLAUDE_FABLE_5_2_MESSAGES_INPUT_TEXT: Final = CLAUDE_FABLE_5_2_MESSAGES_INPUT_TEXT.model_copy( + update={ + "name": "vertex_ai/claude-fable-5-2-messages_input_text", + "model": "vertex_ai/claude-fable-5-2", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_FABLE_5_2_INPUT_TEXT.response, + "expected": VERTEX_AI_CLAUDE_FABLE_5_2_INPUT_TEXT.expected, + } +) VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM: Final = CLAUDE_FABLE_5_1_MESSAGES_STREAM.model_copy( update={ "name": "vertex_ai/claude-fable-5-1-messages_stream", @@ -3334,10 +3444,23 @@ VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM: Final = CLAUDE_FABLE_5_1_MESSAGES_ST "expected": VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT.expected, } ) +VERTEX_AI_CLAUDE_FABLE_5_2_MESSAGES_STREAM: Final = CLAUDE_FABLE_5_2_MESSAGES_STREAM.model_copy( + update={ + "name": "vertex_ai/claude-fable-5-2-messages_stream", + "model": "vertex_ai/claude-fable-5-2", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_FABLE_5_2_STREAM.response, + "expected": VERTEX_AI_CLAUDE_FABLE_5_2_INPUT_TEXT.expected, + } +) VERTEX_AI_CLAUDE_FABLE_5_1_BRIDGE_PAIRS: Final = bridge_pairs( (VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM), ("/v1/responses",), ) +VERTEX_AI_CLAUDE_FABLE_5_2_BRIDGE_PAIRS: Final = bridge_pairs( + (VERTEX_AI_CLAUDE_FABLE_5_2_MESSAGES_INPUT_TEXT, VERTEX_AI_CLAUDE_FABLE_5_2_MESSAGES_STREAM), + ("/v1/responses",), +) VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT: Final = CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT.model_copy( update={ "name": "vertex_ai/claude-opus-5-5-messages_input_text", @@ -3915,7 +4038,9 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( VERTEX_AI_CLAUDE_OPUS_5_5_INPUT_TEXT, VERTEX_AI_CLAUDE_OPUS_5_5_STREAM, VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT, + VERTEX_AI_CLAUDE_FABLE_5_2_INPUT_TEXT, VERTEX_AI_CLAUDE_FABLE_5_1_STREAM, + VERTEX_AI_CLAUDE_FABLE_5_2_STREAM, VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_READ, VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_WRITE_5M, VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_WRITE_1H, @@ -3927,10 +4052,15 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( VERTEX_AI_CLAUDE_OPUS_5_5_THINKING, VERTEX_AI_CLAUDE_OPUS_5_5_TOOL_CALL, VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_READ, + VERTEX_AI_CLAUDE_FABLE_5_2_CACHE_READ, VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_5M, + VERTEX_AI_CLAUDE_FABLE_5_2_CACHE_WRITE_5M, VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_1H, + VERTEX_AI_CLAUDE_FABLE_5_2_CACHE_WRITE_1H, VERTEX_AI_CLAUDE_FABLE_5_1_THINKING, + VERTEX_AI_CLAUDE_FABLE_5_2_THINKING, VERTEX_AI_CLAUDE_FABLE_5_1_TOOL_CALL, + VERTEX_AI_CLAUDE_FABLE_5_2_TOOL_CALL, VERTEX_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT, VERTEX_AI_CLAUDE_HAIKU_5_5_STREAM, VERTEX_AI_CLAUDE_HAIKU_5_5_CACHE_READ, @@ -3944,8 +4074,11 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( VERTEX_AI_CLAUDE_HAIKU_5_5_MESSAGES_STREAM, *(pair.bridged for pair in VERTEX_AI_CLAUDE_HAIKU_5_5_BRIDGE_PAIRS), VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, + VERTEX_AI_CLAUDE_FABLE_5_2_MESSAGES_INPUT_TEXT, VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM, + VERTEX_AI_CLAUDE_FABLE_5_2_MESSAGES_STREAM, *(pair.bridged for pair in VERTEX_AI_CLAUDE_FABLE_5_1_BRIDGE_PAIRS), + *(pair.bridged for pair in VERTEX_AI_CLAUDE_FABLE_5_2_BRIDGE_PAIRS), VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_STREAM, *(pair.bridged for pair in VERTEX_AI_CLAUDE_OPUS_5_5_BRIDGE_PAIRS), @@ -3995,6 +4128,7 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( ), StreamParityTestCase(plain=VERTEX_AI_CLAUDE_OPUS_5_5_INPUT_TEXT, streamed=VERTEX_AI_CLAUDE_OPUS_5_5_STREAM), StreamParityTestCase(plain=VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT, streamed=VERTEX_AI_CLAUDE_FABLE_5_1_STREAM), + StreamParityTestCase(plain=VERTEX_AI_CLAUDE_FABLE_5_2_INPUT_TEXT, streamed=VERTEX_AI_CLAUDE_FABLE_5_2_STREAM), StreamParityTestCase( plain=VERTEX_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT, streamed=VERTEX_AI_CLAUDE_HAIKU_5_5_STREAM ), @@ -4010,10 +4144,18 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( plain=VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, streamed=VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM, ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_FABLE_5_2_MESSAGES_INPUT_TEXT, + streamed=VERTEX_AI_CLAUDE_FABLE_5_2_MESSAGES_STREAM, + ), StreamParityTestCase( plain=VERTEX_AI_CLAUDE_FABLE_5_1_BRIDGE_PAIRS[0].bridged, streamed=VERTEX_AI_CLAUDE_FABLE_5_1_BRIDGE_PAIRS[1].bridged, ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_FABLE_5_2_BRIDGE_PAIRS[0].bridged, + streamed=VERTEX_AI_CLAUDE_FABLE_5_2_BRIDGE_PAIRS[1].bridged, + ), StreamParityTestCase( plain=VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, streamed=VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_STREAM, @@ -4068,6 +4210,7 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( *GEMINI_4_ARGON_ROUTE_BRIDGES, *VERTEX_AI_CLAUDE_HAIKU_5_5_BRIDGE_PAIRS, *VERTEX_AI_CLAUDE_FABLE_5_1_BRIDGE_PAIRS, + *VERTEX_AI_CLAUDE_FABLE_5_2_BRIDGE_PAIRS, *VERTEX_AI_CLAUDE_OPUS_5_5_BRIDGE_PAIRS, *VERTEX_AI_CLAUDE_SONNET_5_5_BRIDGE_PAIRS, ) diff --git a/tests/integration/cost_calculation/cost_map.py b/tests/integration/cost_calculation/cost_map.py index 37dc86e2601..4f87caf7854 100644 --- a/tests/integration/cost_calculation/cost_map.py +++ b/tests/integration/cost_calculation/cost_map.py @@ -906,6 +906,28 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "max_output_tokens": 128000, "max_tokens": 128000, }, + "vertex_ai/claude-fable-5-2": { + "litellm_provider": "vertex_ai-anthropic_models", + "mode": "chat", + "input_cost_per_token": 1.52e-05, + "input_cost_per_token_batches": 7.6e-06, + "output_cost_per_token": 7.6e-05, + "output_cost_per_token_batches": 3.8e-05, + "cache_creation_input_token_cost": 1.9e-05, + "cache_creation_input_token_cost_above_1hr": 3.04e-05, + "cache_read_input_token_cost": 3.8e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.019, + "search_context_size_low": 0.019, + "search_context_size_medium": 0.019, + }, + "supports_function_calling": True, + "supports_prompt_caching": True, + "supports_reasoning": True, + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + }, "vertex_ai/claude-haiku-5-5": { "litellm_provider": "vertex_ai-anthropic_models", "mode": "chat", @@ -1494,6 +1516,19 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "output_cost_per_token": 9e-05, "supports_function_calling": True, }, + "us.anthropic.claude-fable-5-2": { + "cache_creation_input_token_cost": 1.8e-05, + "cache_creation_input_token_cost_above_1hr": 2.88e-05, + "cache_read_input_token_cost": 1.44e-06, + "input_cost_per_token": 1.44e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.2e-05, + "supports_function_calling": True, + }, "gpt-6-sol": { "cache_read_input_token_cost": 2.2e-07, "input_cost_per_token": 2.2e-06, diff --git a/tests/integration/providers/test_bedrock_converse_fable_capability_flags_wire.py b/tests/integration/providers/test_bedrock_converse_fable_capability_flags_wire.py new file mode 100644 index 00000000000..815c5ba750f --- /dev/null +++ b/tests/integration/providers/test_bedrock_converse_fable_capability_flags_wire.py @@ -0,0 +1,234 @@ +"""us.anthropic.claude-fable-5-2 on Bedrock Converse: the row's capability flags shape the Converse body. + +Fable rejects sampling params and forced tool choice and always thinks, so the proxy must strip +or refuse those before the request reaches Bedrock, the same way it does on the Anthropic API. +""" + +import json +from collections.abc import Callable +from typing import Final + +from integration._support.client import Gateway, Scenario +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_MODEL_ID: Final = "us.anthropic.claude-fable-5-2" +_TARGET: Final = f"/model/{_MODEL_ID}/converse" +_TOKEN: Final = "synthetic-bedrock-bearer" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_FORCED_TOOL_CHOICES: Final[tuple[JsonValue, ...]] = ("required", {"type": "function", "function": {"name": "lookup"}}) +_TOOLS: Final[list[JsonValue]] = [ + { + "type": "function", + "function": { + "name": "lookup", + "description": "Look up a synthetic record.", + "parameters": { + "type": "object", + "properties": {"record_id": {"type": "string"}}, + "required": ["record_id"], + }, + }, + } +] + + +def _reply(usage: dict[str, JsonValue] | None = None) -> bytes: + return json.dumps( + { + "output": { + "message": { + "role": "assistant", + "content": [ + { + "reasoningContent": { + "reasoningText": {"text": "scripted thought", "signature": "scripted-signature"} + } + }, + {"text": "done"}, + ], + } + }, + "stopReason": "end_turn", + "usage": usage or {"inputTokens": 12, "outputTokens": 3, "totalTokens": 15}, + "metrics": {"latencyMs": 1}, + } + ).encode() + + +def _deployment(scenario: Scenario, api_base: str, *, drop_params: bool = False) -> str: + if drop_params: + return scenario.model( + model=f"bedrock/converse/{_MODEL_ID}", + api_key=_TOKEN, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=api_base, + num_retries=0, + drop_params=True, + ) + return scenario.model( + model=f"bedrock/converse/{_MODEL_ID}", + api_key=_TOKEN, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=api_base, + num_retries=0, + ) + + +def _peer(on_body: Callable[[dict[str, JsonValue]], bytes]) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == _TARGET, (request.method, request.target) + assert request.headers["authorization"] == f"Bearer {_TOKEN}", request.headers + return Reply(body=on_body(_JSON_OBJECT.validate_json(request.body))) + + return respond + + +def _chat_body(model: str, **extra: JsonValue) -> dict[str, JsonValue]: + return { + "model": model, + "messages": [{"role": "user", "content": "synthetic fable request"}], + "max_tokens": 64, + **extra, + } + + +def test_drop_params_removes_temperature_and_top_p_before_bedrock(gateway: Gateway) -> None: + def on_body(body: dict[str, JsonValue]) -> bytes: + assert body["inferenceConfig"] == {"maxTokens": 64}, body + return _reply() + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url, drop_params=True) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, temperature=0.2, top_p=0.9)) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == "done", response.text + assert len(wire.drain()) == 1 + + +def test_sampling_params_are_refused_before_reaching_bedrock_without_drop_params(gateway: Gateway) -> None: + def on_body(body: dict[str, JsonValue]) -> bytes: + raise AssertionError(f"temperature must not reach Bedrock: {body}") + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, temperature=0.2)) + assert response.status_code == 400, response.text + assert "temperature" in response.text and "drop_params" in response.text, response.text + assert len(wire.drain()) == 0 + + +def test_drop_params_downgrades_forced_and_named_tool_choice_to_auto(gateway: Gateway) -> None: + def on_body(body: dict[str, JsonValue]) -> bytes: + tool_config: Final = body["toolConfig"] + assert isinstance(tool_config, dict), body + assert tool_config.get("toolChoice") in (None, {"auto": {}}), body + tools: Final = tool_config["tools"] + assert isinstance(tools, list) and len(tools) == 1, body + return _reply() + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url, drop_params=True) + responses: Final = tuple( + gateway.request("POST", "/v1/chat/completions", _chat_body(model, tools=_TOOLS, tool_choice=tool_choice)) + for tool_choice in _FORCED_TOOL_CHOICES + ) + assert [response.status_code for response in responses] == [200, 200], [response.text for response in responses] + assert len(wire.drain()) == 2 + + +def test_forced_tool_choice_is_refused_before_reaching_bedrock_without_drop_params(gateway: Gateway) -> None: + def on_body(body: dict[str, JsonValue]) -> bytes: + raise AssertionError(f"forced tool_choice must not reach Bedrock: {body}") + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, tools=_TOOLS, tool_choice="required") + ) + assert response.status_code == 400, response.text + assert "tool_choice" in response.text, response.text + assert len(wire.drain()) == 0 + + +def test_thinking_disabled_is_dropped_for_an_always_on_model(gateway: Gateway) -> None: + def on_body(body: dict[str, JsonValue]) -> bytes: + assert "thinking" not in json.dumps(body.get("additionalModelRequestFields", {})), body + return _reply() + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, thinking={"type": "disabled"}) + ) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["reasoning_content"] == "scripted thought", response.text + assert len(wire.drain()) == 1 + + +def test_reasoning_effort_becomes_adaptive_thinking_with_output_config_effort(gateway: Gateway) -> None: + def on_body(body: dict[str, JsonValue]) -> bytes: + extra: Final = body["additionalModelRequestFields"] + assert isinstance(extra, dict), body + assert extra["thinking"] == {"type": "adaptive", "display": "summarized"}, body + assert extra["output_config"] == {"effort": "max"}, body + assert "budget_tokens" not in json.dumps(body), body + return _reply() + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, reasoning_effort="max")) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 + + +def test_one_hour_cache_point_reaches_bedrock_and_its_usage_comes_back(gateway: Gateway) -> None: + usage: Final[dict[str, JsonValue]] = { + "inputTokens": 9, + "outputTokens": 3, + "totalTokens": 7237, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 7225, + "cacheDetails": [{"ttl": "1h", "inputTokens": 7225}], + } + + def on_body(body: dict[str, JsonValue]) -> bytes: + assert body["system"] == [ + {"text": "synthetic cached system prompt"}, + {"cachePoint": {"type": "default", "ttl": "1h"}}, + ], body + return _reply(usage) + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + _chat_body( + model, + messages=[ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "synthetic cached system prompt", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + }, + {"role": "user", "content": "synthetic fable request"}, + ], + ), + ) + assert response.status_code == 200, response.text + usage_out: Final = _JSON_OBJECT.validate_json(response.text)["usage"] + assert isinstance(usage_out, dict), response.text + prompt_details: Final = usage_out["prompt_tokens_details"] + assert isinstance(prompt_details, dict), response.text + assert prompt_details["cache_creation_token_details"] == { + "ephemeral_5m_input_tokens": 0, + "ephemeral_1h_input_tokens": 7225, + }, response.text + assert usage_out["cache_creation_input_tokens"] == 7225, response.text + assert len(wire.drain()) == 1 diff --git a/tests/integration/providers/test_vertex_anthropic_fable_capability_flags_wire.py b/tests/integration/providers/test_vertex_anthropic_fable_capability_flags_wire.py new file mode 100644 index 00000000000..8376cd92de7 --- /dev/null +++ b/tests/integration/providers/test_vertex_anthropic_fable_capability_flags_wire.py @@ -0,0 +1,242 @@ +"""claude-fable-5-2 on Vertex AI: the row's capability flags shape the rawPredict body the proxy sends. + +Fable rejects sampling params and forced tool choice and always thinks, so the proxy must strip +or refuse those before the request reaches Vertex, the same way it does on the Anthropic API. +""" + +import json +import uuid +from collections.abc import Callable +from typing import Final +from urllib.parse import parse_qs + +from integration._support.client import Gateway, Scenario +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_MODEL: Final = "claude-fable-5-2" +_PROJECT: Final = "scripted-project" +_LOCATION: Final = "global" +_TARGET: Final = f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/anthropic/models/{_MODEL}:rawPredict" +_TOKEN: Final = "scripted-vertex-access-token" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_FORCED_TOOL_CHOICES: Final[tuple[JsonValue, ...]] = ("required", {"type": "function", "function": {"name": "lookup"}}) +_TOOLS: Final[list[JsonValue]] = [ + { + "type": "function", + "function": { + "name": "lookup", + "description": "Look up a synthetic record.", + "parameters": { + "type": "object", + "properties": {"record_id": {"type": "string"}}, + "required": ["record_id"], + }, + }, + } +] + + +def _reply(identity: str, usage: dict[str, JsonValue] | None = None) -> bytes: + return json.dumps( + { + "id": identity, + "type": "message", + "role": "assistant", + "model": _MODEL, + "content": [ + {"type": "thinking", "thinking": "scripted thought", "signature": "scripted-signature"}, + {"type": "text", "text": "done"}, + ], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": usage or {"input_tokens": 12, "output_tokens": 3}, + } + ).encode() + + +def _deployment(scenario: Scenario, api_base: str, *, drop_params: bool = False) -> str: + credentials: Final = service_account_json(_PROJECT, api_base) + if drop_params: + return scenario.model( + model=f"vertex_ai/{_MODEL}", + api_base=api_base, + api_key=None, + vertex_project=_PROJECT, + vertex_location=_LOCATION, + vertex_credentials=credentials, + num_retries=0, + drop_params=True, + ) + return scenario.model( + model=f"vertex_ai/{_MODEL}", + api_base=api_base, + api_key=None, + vertex_project=_PROJECT, + vertex_location=_LOCATION, + vertex_credentials=credentials, + num_retries=0, + ) + + +def _peer(on_body: Callable[[dict[str, JsonValue]], bytes]) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.target == "/_oauth/token": + form: Final = parse_qs(request.body.decode()) + assert form["grant_type"] == ["urn:ietf:params:oauth:grant-type:jwt-bearer"], form + return Reply(body=json.dumps({"access_token": _TOKEN, "token_type": "Bearer", "expires_in": 3600}).encode()) + assert request.method == "POST" and request.target == _TARGET, (request.method, request.target) + assert request.headers["authorization"] == f"Bearer {_TOKEN}", request.headers + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["anthropic_version"] == "vertex-2023-10-16", body + return Reply(body=on_body(body)) + + return respond + + +def _chat_body(model: str, **extra: JsonValue) -> dict[str, JsonValue]: + return { + "model": model, + "messages": [{"role": "user", "content": "synthetic fable request"}], + "max_tokens": 64, + **extra, + } + + +def _model_calls(requests: tuple[Request, ...]) -> int: + return sum(request.target == _TARGET for request in requests) + + +def test_drop_params_removes_temperature_and_top_p_before_vertex(gateway: Gateway) -> None: + identity: Final = f"msg_{uuid.uuid4().hex}" + + def on_body(body: dict[str, JsonValue]) -> bytes: + assert "temperature" not in body and "top_p" not in body, body + assert body["max_tokens"] == 64, body + return _reply(identity) + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url, drop_params=True) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, temperature=0.2, top_p=0.9)) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == "done", response.text + assert _model_calls(wire.drain()) == 1 + + +def test_sampling_params_are_refused_before_reaching_vertex_without_drop_params(gateway: Gateway) -> None: + def on_body(body: dict[str, JsonValue]) -> bytes: + raise AssertionError(f"temperature must not reach Vertex: {body}") + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, temperature=0.2)) + assert response.status_code == 400, response.text + assert "temperature" in response.text and "drop_params" in response.text, response.text + assert _model_calls(wire.drain()) == 0 + + +def test_drop_params_downgrades_forced_and_named_tool_choice_to_auto(gateway: Gateway) -> None: + identity: Final = f"msg_{uuid.uuid4().hex}" + + def on_body(body: dict[str, JsonValue]) -> bytes: + assert body["tool_choice"] == {"type": "auto"}, body + tools: Final = body["tools"] + assert isinstance(tools, list) and [tool["name"] for tool in tools if isinstance(tool, dict)] == ["lookup"], ( + body + ) + return _reply(identity) + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url, drop_params=True) + responses: Final = tuple( + gateway.request("POST", "/v1/chat/completions", _chat_body(model, tools=_TOOLS, tool_choice=tool_choice)) + for tool_choice in _FORCED_TOOL_CHOICES + ) + assert [response.status_code for response in responses] == [200, 200], [response.text for response in responses] + assert _model_calls(wire.drain()) == 2 + + +def test_forced_tool_choice_is_refused_before_reaching_vertex_without_drop_params(gateway: Gateway) -> None: + def on_body(body: dict[str, JsonValue]) -> bytes: + raise AssertionError(f"forced tool_choice must not reach Vertex: {body}") + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, tools=_TOOLS, tool_choice="required") + ) + assert response.status_code == 400, response.text + assert "tool_choice" in response.text, response.text + assert _model_calls(wire.drain()) == 0 + + +def test_thinking_disabled_is_dropped_for_an_always_on_model(gateway: Gateway) -> None: + identity: Final = f"msg_{uuid.uuid4().hex}" + + def on_body(body: dict[str, JsonValue]) -> bytes: + assert "thinking" not in body, body + return _reply(identity) + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, thinking={"type": "disabled"}) + ) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["reasoning_content"] == "scripted thought", response.text + assert _model_calls(wire.drain()) == 1 + + +def test_reasoning_effort_becomes_adaptive_thinking_with_output_config_effort(gateway: Gateway) -> None: + identity: Final = f"msg_{uuid.uuid4().hex}" + + def on_body(body: dict[str, JsonValue]) -> bytes: + assert body["thinking"] == {"type": "adaptive", "display": "summarized"}, body + assert body["output_config"] == {"effort": "max"}, body + assert "budget_tokens" not in json.dumps(body), body + return _reply(identity) + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, reasoning_effort="max")) + assert response.status_code == 200, response.text + assert _model_calls(wire.drain()) == 1 + + +def test_messages_one_hour_cache_control_reaches_vertex_and_its_usage_comes_back(gateway: Gateway) -> None: + identity: Final = f"msg_{uuid.uuid4().hex}" + system: Final[list[JsonValue]] = [ + {"type": "text", "text": "synthetic cached system prompt", "cache_control": {"type": "ephemeral", "ttl": "1h"}} + ] + usage: Final[dict[str, JsonValue]] = { + "input_tokens": 9, + "output_tokens": 3, + "cache_creation_input_tokens": 7225, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 7225}, + } + + def on_body(body: dict[str, JsonValue]) -> bytes: + assert body["system"] == system, body + return _reply(identity, usage) + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 64, + "system": system, + "messages": [{"role": "user", "content": "synthetic fable request"}], + }, + ) + assert response.status_code == 200, response.text + assert response.json()["usage"]["cache_creation"] == { + "ephemeral_5m_input_tokens": 0, + "ephemeral_1h_input_tokens": 7225, + }, response.text + assert response.json()["usage"]["cache_creation_input_tokens"] == 7225, response.text + assert _model_calls(wire.drain()) == 1 diff --git a/tests/unit/test_claude_fable_5_config.py b/tests/unit/test_claude_fable_5_config.py index c9e46dd6887..2157364efb3 100644 --- a/tests/unit/test_claude_fable_5_config.py +++ b/tests/unit/test_claude_fable_5_config.py @@ -98,6 +98,13 @@ def test_adaptive_thinking_detected_for_fable_5_1(local_model_cost_map, model): FABLE_5_2_VARIANTS = ( "claude-fable-5-2", + "anthropic.claude-fable-5-2", + "global.anthropic.claude-fable-5-2", + "us.anthropic.claude-fable-5-2", + "eu.anthropic.claude-fable-5-2", + "us-gov.anthropic.claude-fable-5-2", + "bedrock/us-gov-west-1/anthropic.claude-fable-5-2", + "bedrock/us-gov-east-1/anthropic.claude-fable-5-2", "vertex_ai/claude-fable-5-2", "vertex_ai/claude-fable-5-2@default", "bedrock/claude_platform/claude-fable-5-2", @@ -118,6 +125,11 @@ def test_fable_5_2_present_in_bundled_backup() -> None: "claude-fable-5-2", "anthropic/claude-fable-5-2", "claude-fable-latest", + "anthropic.claude-fable-5-2", + "bedrock/us.anthropic.claude-fable-5-2", + "bedrock/invoke/eu.anthropic.claude-fable-5-2", + "bedrock/global.anthropic.claude-fable-5-2", + "bedrock/claude-fable-latest", "vertex_ai/claude-fable-5-2", "vertex_ai/claude-fable-5-2@default", "bedrock/claude_platform/claude-fable-5-2", @@ -129,6 +141,10 @@ def test_adaptive_thinking_detected_for_fable_5_2(local_model_cost_map, model: s assert AnthropicModelInfo.is_adaptive_thinking_model(model, "anthropic") is True +def test_fable_5_2_registered_for_bedrock_converse() -> None: + assert "anthropic.claude-fable-5-2" in BEDROCK_CONVERSE_MODELS + + def test_fable_5_2_fast_mode_doubles_every_token_price_on_the_anthropic_api(local_model_cost_map) -> None: """Anthropic bills fast mode at 2x the standard rate on every token type (platform.claude.com/docs/en/about-claude/pricing, read 2026-11-09), and the From 92c2e69ac534fc3588b2bff7f544003c0a7f25ae Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Mon, 9 Nov 2026 15:38:59 -0800 Subject: [PATCH 26/46] fix(proxy): fail-closed budget enforcement returns 403 naming /team/member_add for a team user with no membership row (#50603) * fix(proxy): read an absent team membership row as zero spend under fail-closed budget enforcement A team user under a default member budget has no LiteLLM_TeamMembership row until spend is persisted. SpendCounterReseed.from_db returned None for that absent row, the same value as a failed DB read, so fail_closed_budget_enforcement rejected every request from that user with 503. The DB answering that the row is absent is a verified zero, so from_db now returns 0.0 for a team-member counter without a row and still returns None when the read fails. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): only read an absent membership row as zero when the team member key parses unambiguously A team_id or user_id containing ":" makes the counter key split onto the wrong membership row, so that miss keeps returning None instead of a verified zero. Type the new unit tests. * fix(proxy): refuse a fail-closed team member with no membership row with a 403 that names /team/member_add Replaces the earlier reading of an absent LiteLLM_TeamMembership row as verified zero spend. A key minted with team_id and user_id but no /team/member_add has no row and the spend writer never creates one, so admitting it at zero would leave its member spend tracked only in Redis. With fail_closed_budget_enforcement on, _check_team_member_budget now rejects that key with a 403 that names the user, the team and the member_add call, before any spend read. With the flag off nothing changes. The membership load raises on a database failure and returns None only for a genuinely absent row, so the 403 never fires during an outage. spend_counter_reseed.py and the counter prefix constant go back to main. * test(proxy): drive the no-membership 403 tests through the real membership and spend readers Fake Postgres rows and the FakeSpendCounters double stand in for the external stores instead of patching get_team_membership and get_current_spend. Strip CR/LF from the ids in the fail-closed warning log. * fix(proxy): read fail_closed_budget_enforcement from general_settings in the membership check Importing proxy_server's private helper grew reportPrivateUsage past the type gate ceiling. Type the integration helper's return as JsonValue. * fix(proxy): expose fail_closed_budget_enforcement_enabled for the membership check Reading general_settings directly from auth_checks grew the unknown-type counts past the type gate ceiling; a small public wrapper in proxy_server keeps the gate flat. * test(proxy): drop docstrings that restate the no-membership tests --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 28 +++++++- litellm/proxy/proxy_server.py | 4 ++ .../test_fail_closed_team_member_budget.py | 66 ++++++++++++++++++ ...st_auth_checks_object_access_and_lookup.py | 69 ++++++++++++++++++- 4 files changed, 164 insertions(+), 3 deletions(-) create mode 100644 tests/integration/spend/test_fail_closed_team_member_budget.py diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e546c093473..f0c9632837d 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -7076,6 +7076,27 @@ def team_member_max_budget_alert_check( _team_member_max_budget_alert_check: Final = team_member_max_budget_alert_check +def _raise_team_member_without_membership(user_id: str, team_id: str) -> None: + from litellm.proxy.litellm_pre_call_utils import sanitize_for_log + + verbose_proxy_logger.warning( + "fail_closed_budget_enforcement: rejecting request, user_id=%s has no membership row in team_id=%s " + "so spend against the team member budget cannot be verified", + sanitize_for_log(user_id), + sanitize_for_log(team_id), + ) + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": ( + f"User={user_id} has no membership in Team={team_id}, so their spend against the team member " + "budget cannot be verified and fail_closed_budget_enforcement is enabled. Add them to the team " + "with POST /team/member_add." + ) + }, + ) + + async def _check_team_member_budget( team_object: LiteLLM_TeamTable | None, user_object: LiteLLM_UserTable | None, @@ -7125,11 +7146,14 @@ async def _check_team_member_budget( ) if team_member_budget is not None: + from litellm.proxy.proxy_server import fail_closed_budget_enforcement_enabled, get_current_spend + + if loaded_membership is None and fail_closed_budget_enforcement_enabled(): + _raise_team_member_without_membership(user_id=valid_token.user_id, team_id=team_object.team_id) + team_member_spend = (loaded_membership.spend if loaded_membership is not None else 0.0) or 0.0 # Read from cross-pod counter (Redis-first) if available - from litellm.proxy.proxy_server import get_current_spend - team_member_spend = await get_current_spend( counter_key=team_member_counter(valid_token.user_id, team_object.team_id), fallback_spend=team_member_spend, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 046510d026d..37dc135d4d7 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3389,6 +3389,10 @@ def _fail_closed_budget_enforcement() -> bool: return general_settings.get("fail_closed_budget_enforcement") is True +def fail_closed_budget_enforcement_enabled() -> bool: + return _fail_closed_budget_enforcement() + + def _raise_budget_unverifiable(counter_key: str) -> None: verbose_proxy_logger.warning( "fail_closed_budget_enforcement: rejecting request — spend for %s could " diff --git a/tests/integration/spend/test_fail_closed_team_member_budget.py b/tests/integration/spend/test_fail_closed_team_member_budget.py new file mode 100644 index 00000000000..52b2e8dd0f3 --- /dev/null +++ b/tests/integration/spend/test_fail_closed_team_member_budget.py @@ -0,0 +1,66 @@ +import uuid +from collections.abc import Iterator +from contextlib import ExitStack +from pathlib import Path +from typing import Final + +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, gateway_from_environment, object_value, string_value +from integration._support.process import owned_proxy_process +from pydantic import JsonValue + + +def _fail_closed_config(directory: Path) -> Path: + base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final = { + **base, + "general_settings": {**object_value(base["general_settings"]), "fail_closed_budget_enforcement": True}, + } + path: Final = directory / f"fail-closed-{uuid.uuid4().hex[:8]}.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.fixture(scope="module") +def fail_closed_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("fail-closed-team-member") + with ExitStack() as stack: + rig_gateway: Final = stack.enter_context(gateway_from_environment()) + owned: Final = stack.enter_context( + owned_proxy_process(rig_gateway, directory, {}, config=_fail_closed_config(directory)) + ) + yield owned.gateway + + +def _chat(gateway: Gateway, model: str, key: str) -> tuple[int, dict[str, JsonValue]]: + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "integration control"}]}, + key=key, + ) + return response.status_code, JSON_OBJECT.validate_json(response.content) + + +def test_team_member_without_membership_row_is_told_to_be_added_then_admitted(fail_closed_gateway: Gateway) -> None: + with fail_closed_gateway.scenario() as scenario: + model: Final = scenario.model() + team_id: Final = scenario.team(models=[model], team_member_budget=100.0) + user_id: Final = scenario.user(max_budget=100.0) + key: Final = scenario.key(team_id=team_id, user_id=user_id, models=[model]) + + status, body = _chat(fail_closed_gateway, model, key) + + assert status == 403, body + message: Final = string_value(object_value(body["error"])["message"]) + assert user_id in message and team_id in message and "/team/member_add" in message + + fail_closed_gateway.post( + "/team/member_add", {"team_id": team_id, "member": {"role": "user", "user_id": user_id}} + ) + + admitted_status, admitted = _chat(fail_closed_gateway, model, key) + + assert admitted_status == 200, admitted + assert object_value(admitted["usage"])["total_tokens"] diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index 78ce9728f3b..e7bcc696c0d 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -6080,6 +6080,71 @@ class TestGuardrailModificationCheck: self._call({"metadata": '"just a string"'}) +def _member_without_row_prisma(max_budget: float) -> MagicMock: + budget_row: Final = MagicMock() + budget_row.max_budget = max_budget + budget_row.dict = MagicMock(return_value={"budget_id": "budget-default", "max_budget": max_budget}) + prisma_client: Final = MagicMock() + prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=budget_row) + prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=None) + return prisma_client + + +async def _check_member_without_membership_row(prisma_client: MagicMock) -> None: + from litellm.caching.dual_cache import DualCache + from litellm.proxy.utils import ProxyLogging + + await _check_team_member_budget( + team_object=LiteLLM_TeamTable(team_id="team-1", metadata={"team_member_budget_id": "budget-default"}), + user_object=LiteLLM_UserTable(user_id="user-1"), + valid_token=UserAPIKeyAuth(token="tok-1", user_id="user-1", team_id="team-1"), + prisma_client=prisma_client, + user_api_key_cache=DualCache(), + proxy_logging_obj=ProxyLogging(user_api_key_cache=None), + ) + + +@pytest.mark.asyncio +async def test_fail_closed_rejects_member_without_membership_row_with_a_403_naming_member_add( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from fastapi import HTTPException + + import litellm.proxy.proxy_server as ps + from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + spend: Final = FakeSpendCounters(redis={}) + prisma_client: Final = _member_without_row_prisma(100.0) + monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True}) + monkeypatch.setattr(ps, "spend_counter_store", spend.store()) + monkeypatch.setattr(ps, "prisma_client", prisma_client) + + with pytest.raises(HTTPException) as exc_info: + await _check_member_without_membership_row(prisma_client) + + assert exc_info.value.status_code == 403 + message: Final = str(exc_info.value.detail) + assert "user-1" in message and "team-1" in message and "/team/member_add" in message + prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once() + assert spend.commands == [] + + +@pytest.mark.asyncio +async def test_fail_closed_off_still_admits_member_without_membership_row(monkeypatch: pytest.MonkeyPatch) -> None: + import litellm.proxy.proxy_server as ps + from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + spend: Final = FakeSpendCounters(redis={}) + prisma_client: Final = _member_without_row_prisma(100.0) + monkeypatch.setattr(ps, "general_settings", {}) + monkeypatch.setattr(ps, "spend_counter_store", spend.store()) + monkeypatch.setattr(ps, "prisma_client", prisma_client) + + await _check_member_without_membership_row(prisma_client) + + assert spend.commands, "the member budget check must still read the spend counter with the flag off" + + @pytest.mark.asyncio async def test_team_member_budget_check_falls_back_to_team_default_budget_id(): """When a member's TeamMembership has no linked budget row, the check @@ -11850,7 +11915,9 @@ async def test_baggage_mapping_rejects_exactly_the_client_tags_rejected_without_ unmapped_request = _client_tag_request(client_tags, baggage=False) unmapped = _client_tag_body(client_tags, client_slot) LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route(unmapped, route) - LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth(unmapped_request, unmapped, UserAPIKeyAuth(api_key="qa-key")) + LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth( + unmapped_request, unmapped, UserAPIKeyAuth(api_key="qa-key") + ) mapped_request = _client_tag_request(client_tags, baggage=baggage) mapped = _client_tag_body(client_tags, client_slot) From 35386e3f4a196b29f007eca4c45d461f9f6d63dd Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Mon, 9 Nov 2026 15:52:38 -0800 Subject: [PATCH 27/46] fix(complexity_router): keep Claude Code subagent and side-request traffic off other agents' tier pins (#50741) * fix(complexity_router): keep Claude Code subagent and side-request traffic off other agents' tier pins Claude Code sends one x-claude-code-session-id for the main agent, every subagent and their side requests. In user_turn mode a subagent's new ask rewrote the parent's pin (#50099), and the agent progress summaries Claude Code fires during a subagent run rewrote the subagent's pin to the cheapest tier (#50444). Subagent requests now get their own tier pin keyed on session id plus x-claude-code-agent-id. Requests marked x-claude-code-request-class: auxiliary route as housekeeping, and only main and subagent turns may write a pin. Without the hint header, the summary and prompt-suggestion prompts are matched as housekeeping sentinels. * chore(ui): regenerate schema.d.ts for the housekeeping_patterns description * fix(complexity_router): let the request-class hint only mark housekeeping, not suppress pin writes A client-sent request class must not decide whether a decision is saved. Tagging a real ask as workflow or auxiliary dropped its routing decision, so later continuations replayed an older pin with the new history. The auxiliary hint now only routes the request through the housekeeping path, which is already never pinned and defers to custom classifiers and keyword rules. * test(complexity_router): cover the auxiliary hint on a prompt no sentinel matches --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> --- litellm/proxy/litellm_pre_call_utils.py | 11 +- .../complexity_router/complexity_router.py | 32 +++- .../complexity_router/config.py | 10 +- .../unit/proxy/test_litellm_pre_call_utils.py | 35 ++++ .../router_strategy/test_complexity_router.py | 157 ++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 6 files changed, 240 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index e2367ab8179..9cc18294cd1 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -223,12 +223,21 @@ def _caller_session_beside_trace_header( return get_chain_id_from_headers(session_headers) or _caller_body_session_id(data, metadata_variable_name) +def _claude_code_subagent_tier_session(normalized_headers: Mapping[str, str]) -> str | None: + session: Final = normalized_headers.get("x-claude-code-session-id", "") + agent: Final = normalized_headers.get("x-claude-code-agent-id", "") + if not _SESSION_ID_VALUE_RE.fullmatch(session) or not _SESSION_ID_VALUE_RE.fullmatch(agent): + return None + return f"{session}:{agent}" + + 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 {} + claude_code_subagent: Final = _claude_code_subagent_tier_session(normalized_headers) + return {COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY: claude_code_subagent} if claude_code_subagent else {} own_headers: Final = { key: value for key, value in normalized_headers.items() diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 600b9d7c455..4fcf555035d 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -801,6 +801,27 @@ def _newest_turn_is_human_ask( return _last_human_ask_index(messages, marker_pairs) == newest_activity +_CLAUDE_CODE_REQUEST_CLASS_HEADER: Final = "x-claude-code-request-class" +_CLAUDE_CODE_AUXILIARY_SENTINEL: Final = f"{_CLAUDE_CODE_REQUEST_CLASS_HEADER}: auxiliary" + + +def _claude_code_request_class(request_kwargs: Mapping[str, object]) -> str | None: + proxy_request: Final = request_kwargs.get("proxy_server_request") + if not isinstance(proxy_request, Mapping): + return None + headers: Final = cast(Mapping[str, object], proxy_request).get("headers") # cast-ok: proxy request is JSON-shaped + if not isinstance(headers, Mapping): + return None + return next( + ( + value.strip().lower() + for key, value in cast(Mapping[object, object], headers).items() # cast-ok: header map from the proxy + if isinstance(key, str) and key.lower() == _CLAUDE_CODE_REQUEST_CLASS_HEADER and isinstance(value, str) + ), + None, + ) + + _CODEX_TURN_METADATA_KEY: Final = "x-codex-turn-metadata" _CODEX_COMPACTION_SUMMARY_PREFIX: Final = ( "Another language model started to solve this problem and produced a summary of its thinking process." @@ -3475,8 +3496,11 @@ class ComplexityRouter(CustomLogger): self._reminder_markers_for_request(request_kwargs), ) - def _matched_housekeeping_sentinel(self, newest_ask: str | None) -> str | None: - """The client housekeeping sentinel on this request's newest ask, or None. + def _matched_housekeeping_sentinel( + self, newest_ask: str | None, request_kwargs: Mapping[str, object] + ) -> str | None: + """The client housekeeping sentinel on this request's newest ask, or None. Claude Code's + `auxiliary` request-class hint counts as one, so it keeps working when the prompts change. Read from the newest ask alone, never the whole history, for the reason `_newest_turn_ask` exists: a title request quoted into a later turn's context would otherwise keep matching and @@ -3490,6 +3514,8 @@ class ComplexityRouter(CustomLogger): """ if self.config.classifier_type == "custom" or not self.config.route_housekeeping_to_cheapest_tier: return None + if _claude_code_request_class(request_kwargs) == "auxiliary": + return _CLAUDE_CODE_AUXILIARY_SENTINEL if not newest_ask: return None return next( @@ -4970,7 +4996,7 @@ class ComplexityRouter(CustomLogger): ), ) - housekeeping_sentinel: Final = self._matched_housekeeping_sentinel(newest_ask) + housekeeping_sentinel: Final = self._matched_housekeeping_sentinel(newest_ask, request_kwargs) housekeeping_tier: Final = self._cheapest_configured_tier() if housekeeping_sentinel is not None else None outcome: Final = ( ClassificationOutcome(tier=housekeeping_tier, score=None, signals=("housekeeping",), cause="housekeeping") diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index d5870bb6510..e42fa6e426b 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -414,9 +414,14 @@ PLAN_MODE_SYSTEM_SENTINELS: Final[tuple[str, ...]] = ('You are currently running # never a paraphrase: a sentinel that matches nothing costs a substring scan per request and reads # as coverage the router does not have. These are client-owned strings that drift with client # releases, so operators extend coverage via housekeeping_patterns rather than editing these. +# The agent progress summary and the next-prompt suggestion were captured from Claude Code 2.1.316 +# on 2026-11-09 (#50444). Both ride the conversation they describe, so only the sentinel tells them +# apart when the client does not send the x-claude-code-request-class hint. HOUSEKEEPING_ASK_SENTINELS: Final[tuple[str, ...]] = ( "Write the title in the predominant language of the session", "You are coming up with a succinct title for a coding session", + "Describe your most recent action in 3-5 words using present tense", + "[SUGGESTION MODE: Suggest what the user might naturally type next into Claude Code.]", ) PLAN_MODE_TOOL_NAME: Final[str] = "exit_plan_mode" @@ -1496,8 +1501,9 @@ class ComplexityRouterConfig(LiteLLMBaseModel): default=None, description=( "Additional case-sensitive literal sentinels that mark a request as client " - "housekeeping, on top of the built-in conversation-title ones. For clients whose " - "wording the built-ins don't cover, or after a client release changes its strings." + "housekeeping, on top of the built-in conversation-title, agent-summary and " + "prompt-suggestion ones. For clients whose wording the built-ins don't cover, or " + "after a client release changes its strings." ), ) diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 557ac72572f..fbd4de8ef5c 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -3895,6 +3895,41 @@ async def test_non_child_requests_cannot_forge_a_child_tier_identity(surface: st assert updated[get_metadata_variable_name(request)]["session_id"] == "child-01" +@pytest.mark.asyncio +@pytest.mark.parametrize( + "agent_id,session_id,child_id", + [ + ( + "a7256c04dfd99c9b9", + "11111111-2222-3333-4444-555555555555", + "11111111-2222-3333-4444-555555555555:a7256c04dfd99c9b9", + ), + (None, "11111111-2222-3333-4444-555555555555", None), + ("short", "11111111-2222-3333-4444-555555555555", None), + ("agent id with spaces", "11111111-2222-3333-4444-555555555555", None), + ("a7256c04dfd99c9b9", "bad", None), + ], +) +async def test_claude_code_subagents_get_their_own_tier_identity_under_the_parent_session( + agent_id: str | None, session_id: str, child_id: str | None +) -> None: + headers: Final = { + "x-claude-code-session-id": session_id, + **({"x-claude-code-agent-id": agent_id} if agent_id is not None else {}), + } + request: Final = _make_request_mock("/v1/messages", headers) + updated: Final = await add_litellm_data_to_request( + data={"model": "test-router", "messages": []}, + 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.get(COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY) == child_id + assert COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY not in metadata.get("requester_metadata", {}) + + 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 444cd2ba6dd..f06cb0d282d 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -7772,6 +7772,163 @@ class TestChildSessionTierPins: router.discard() +class TestClaudeCodeTierPins: + SESSION: Final = "11111111-2222-3333-4444-555555555555" + AGENT_A: Final = "a7256c04dfd99c9b9" + AGENT_B: Final = "a71aeeb9c67f32cee" + TOOL_USE: Final[dict[str, object]] = { + "role": "assistant", + "content": [{"type": "tool_use", "id": "t1", "name": "Read", "input": {}}], + } + TOOL_RESULT: Final[dict[str, object]] = { + "type": "tool_result", + "tool_use_id": "t1", + "content": "def add(a, b): return a + b", + } + # Captured from Claude Code 2.1.316 on 2026-11-09, trimmed to the sentinel and a little context. + SUMMARY_ASK: Final = ( + "Describe your most recent action in 3-5 words using present tense (-ing). " + "Name the file or function, not the branch. Do not use tools." + ) + SUGGESTION_ASK: Final = ( + "[SUGGESTION MODE: Suggest what the user might naturally type next into Claude Code.]\n\n" + "Reply with ONLY the suggestion, no quotes or explanation." + ) + + @staticmethod + def _router(route_housekeeping_to_cheapest_tier: bool = True) -> Router: + return Router( + model_list=[ + { + "model_name": "claude-code-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": {"SIMPLE": "simple", "REASONING": "reasoning"}, + "classification_mode": "user_turn", + "keyword_tier_rules": [{"keywords": ["reason carefully"], "tier": "REASONING"}], + "route_housekeeping_to_cheapest_tier": route_housekeeping_to_cheapest_tier, + }, + }, + }, + *[ + { + "model_name": group, + "litellm_params": {"model": "openai/test-model", "api_key": "test-only"}, + "model_info": {"id": group}, + } + for group in ("simple", "reasoning") + ], + ], + num_retries=0, + ) + + @staticmethod + def _ask(prompt: str) -> list[dict[str, object]]: + return [{"role": "user", "content": prompt}] + + @classmethod + def _continuation(cls) -> list[dict[str, object]]: + return [*cls._ask("Hello!"), cls.TOOL_USE, {"role": "user", "content": [cls.TOOL_RESULT]}] + + @classmethod + def _side_request(cls, side_ask: str) -> list[dict[str, object]]: + return [ + *cls._ask("Hello!"), + cls.TOOL_USE, + {"role": "user", "content": [cls.TOOL_RESULT, {"type": "text", "text": side_ask}]}, + ] + + async def _route( + self, + router: Router, + messages: list[dict[str, object]], + agent_id: str | None = None, + request_class: str | None = None, + ) -> tuple[str, str]: + from litellm.proxy.litellm_pre_call_utils import ( + LiteLLMProxyRequestSetup, + metadata_variable_name_for_route, + ) + + headers: Final = { + "x-claude-code-session-id": self.SESSION, + **({"x-claude-code-agent-id": agent_id} if agent_id is not None else {}), + **({"x-claude-code-request-class": request_class} if request_class is not None else {}), + } + bucket: Final = metadata_variable_name_for_route("/v1/messages") + data: Final[dict[str, object]] = { + bucket: {"user_api_key_hash": "caller-a", "user_api_key_request_route": "/v1/messages"}, + "proxy_server_request": {"headers": headers}, + } + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name=bucket + ) + await router.async_get_available_deployment(model="claude-code-router", request_kwargs=data, messages=messages) + metadata: Final = data[bucket] + assert isinstance(metadata, dict) + assert metadata["session_id"] == self.SESSION + decision: Final = TypeAdapter(StandardLoggingRoutingDecision).validate_python(metadata["routing_decision"]) + return str(decision.get("tier")), decision["cause"] + + @pytest.mark.asyncio + async def test_parent_and_each_subagent_keep_their_own_pin(self) -> None: + router: Final = self._router() + try: + for messages, agent_id, expected in ( + (self._ask("reason carefully"), None, ("REASONING", "literal_keyword_match")), + (self._ask("Hello!"), self.AGENT_A, ("SIMPLE", "heuristic_scorer")), + (self._continuation(), self.AGENT_A, ("SIMPLE", "user_turn_continuation")), + (self._continuation(), None, ("REASONING", "user_turn_continuation")), + (self._ask("reason carefully"), self.AGENT_B, ("REASONING", "literal_keyword_match")), + (self._continuation(), self.AGENT_A, ("SIMPLE", "user_turn_continuation")), + (self._continuation(), self.AGENT_B, ("REASONING", "user_turn_continuation")), + ): + assert await self._route(router, messages, agent_id) == expected + finally: + router.discard() + + @pytest.mark.asyncio + @pytest.mark.parametrize("hint_header", [True, False], ids=["request-class-hint", "sentinel-fallback"]) + @pytest.mark.parametrize( + "agent_id,side_ask", + [(AGENT_A, SUMMARY_ASK), (None, SUGGESTION_ASK)], + ids=["subagent-summary", "main-prompt-suggestion"], + ) + async def test_side_requests_route_cheapest_and_never_move_the_pin( + self, hint_header: bool, agent_id: str | None, side_ask: str + ) -> None: + own_class: Final = ("subagent" if agent_id is not None else "main") if hint_header else None + side_class: Final = "auxiliary" if hint_header else None + router: Final = self._router() + try: + for messages, request_class, expected in ( + (self._ask("reason carefully"), own_class, ("REASONING", "literal_keyword_match")), + (self._side_request(side_ask), side_class, ("SIMPLE", "housekeeping")), + (self._continuation(), own_class, ("REASONING", "user_turn_continuation")), + ): + assert await self._route(router, messages, agent_id, request_class) == expected + finally: + router.discard() + + @pytest.mark.asyncio + async def test_the_auxiliary_hint_routes_housekeeping_when_no_sentinel_matches(self) -> None: + router: Final = self._router() + try: + for messages, request_class, expected in ( + (self._ask("reason carefully"), "subagent", ("REASONING", "literal_keyword_match")), + ( + self._side_request("Summarize what you just did in a few words."), + "auxiliary", + ("SIMPLE", "housekeeping"), + ), + (self._continuation(), "subagent", ("REASONING", "user_turn_continuation")), + ): + assert await self._route(router, messages, self.AGENT_A, request_class) == expected + 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.""" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b171a72dcf9..634ef66f36f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -44032,7 +44032,7 @@ export interface components { heuristic_v2_success_threshold?: number | null; /** * Housekeeping Patterns - * @description Additional case-sensitive literal sentinels that mark a request as client housekeeping, on top of the built-in conversation-title ones. For clients whose wording the built-ins don't cover, or after a client release changes its strings. + * @description Additional case-sensitive literal sentinels that mark a request as client housekeeping, on top of the built-in conversation-title, agent-summary and prompt-suggestion ones. For clients whose wording the built-ins don't cover, or after a client release changes its strings. */ housekeeping_patterns?: string[] | null; /** From 408a67efd7d0fe1978a7360bdf105ffaca454912 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 9 Nov 2026 23:55:58 +0000 Subject: [PATCH 28/46] test(audio): bill reported transcription seconds for measurable uploads (#50767) Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../compatibility/test_audio_transcriptions_wire.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/integration/compatibility/test_audio_transcriptions_wire.py b/tests/integration/compatibility/test_audio_transcriptions_wire.py index faebfda715b..f943400bb15 100644 --- a/tests/integration/compatibility/test_audio_transcriptions_wire.py +++ b/tests/integration/compatibility/test_audio_transcriptions_wire.py @@ -832,12 +832,13 @@ def test_gpt_transcribe_unmeasurable_upload_bills_the_seconds_openai_reports( assert _uploaded_files(requests[0]) == (("file", "a.m4a", upload),) -def test_gpt_transcribe_measured_upload_bills_the_measured_length_over_reported_seconds(gateway: Gateway) -> None: +def test_gpt_transcribe_measured_upload_bills_reported_seconds_over_the_measured_length(gateway: Gateway) -> None: measured_seconds: Final = 1.0 + reported_seconds: Final = 7.0 upload: Final = _silent_wav(measured_seconds) with ( - wire_server(lambda request: _openai_gpt_transcribe_reply(request, 7)) as wire, + wire_server(lambda request: _openai_gpt_transcribe_reply(request, reported_seconds)) as wire, gateway.scenario() as scenario, ): model: Final = scenario.model( @@ -851,9 +852,13 @@ def test_gpt_transcribe_measured_upload_bills_the_measured_length_over_reported_ {"model": model, "response_format": "json"}, {"file": ("a.wav", upload, _WAV_CONTENT_TYPE)}, ) - expected_cost: Final = measured_seconds * _DEPLOYMENT_COST_PER_SECOND + expected_cost: Final = reported_seconds * _DEPLOYMENT_COST_PER_SECOND assert response.status_code == 200, response.text + assert JSON_OBJECT.validate_json(response.content) == { + "text": "hello world", + "usage": {"type": "duration", "seconds": reported_seconds}, + }, response.text assert math.isclose(_response_cost(response), expected_cost), (_response_cost(response), expected_cost) assert math.isclose(_logged_spend(response.headers["x-litellm-call-id"]), expected_cost) requests: Final = wire.drain() From e5ea9502bada03084b40585186c30cc39af31578 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 10 Nov 2026 00:01:34 +0000 Subject: [PATCH 29/46] chore(cost-map): add openai gpt-6-cyber long context prices from the pricing page (#50770) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 5 ++++- model_prices_and_context_window.json | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 0d35c643561..66d375b70ff 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -83449,13 +83449,16 @@ "cache_creation_input_token_cost": 2.5e-05, "cache_creation_input_token_cost_above_12hr": 5e-05, "cache_read_input_token_cost": 2e-06, + "cache_read_input_token_cost_above_272k_tokens": 4e-06, "input_cost_per_token": 2e-05, + "input_cost_per_token_above_272k_tokens": 4e-05, "litellm_provider": "openai", "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "output_cost_per_token": 0.0001, - "source": "https://developers.openai.com/api/docs/models/gpt-6-cyber", + "output_cost_per_token_above_272k_tokens": 0.00015, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 0d35c643561..66d375b70ff 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -83449,13 +83449,16 @@ "cache_creation_input_token_cost": 2.5e-05, "cache_creation_input_token_cost_above_12hr": 5e-05, "cache_read_input_token_cost": 2e-06, + "cache_read_input_token_cost_above_272k_tokens": 4e-06, "input_cost_per_token": 2e-05, + "input_cost_per_token_above_272k_tokens": 4e-05, "litellm_provider": "openai", "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "output_cost_per_token": 0.0001, - "source": "https://developers.openai.com/api/docs/models/gpt-6-cyber", + "output_cost_per_token_above_272k_tokens": 0.00015, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" From c76a95fc7e83f5e0a326ec9a60d138ca821711ff Mon Sep 17 00:00:00 2001 From: jesus-berri Date: Mon, 9 Nov 2026 16:02:46 -0800 Subject: [PATCH 30/46] feat(proxy): let a key's logging entry reference a proxy credential (#50572) * feat(proxy): let a key's logging entry reference a proxy credential A virtual key's metadata.logging entry can name a proxy credential with litellm_credential_name. At request time the credential's values, which config credential_list resolves through the secret manager at boot, are merged under the entry's callback_vars, with the entry's own values winning. An entry whose credential is not defined is dropped, so it never falls back to the proxy's own callback environment. Only proxy admins may change the logging entries of a key when the stored or requested list holds a reference (403), and a reference to an undefined credential is a 400. The check runs on key generate, update, team bulk update, regenerate and /user/new, and compares decrypted entries so re-saving metadata read back from /key/info is not a change. Resolves LIT-10061 * fix(proxy): drop a key logging entry whose referenced credential holds no logging values * fix(proxy): only require newly added credential references to exist on key writes * fix(proxy): log a dropped credential reference as an error so /key/health reports it * fix(proxy): gate credential references on bulk-created user keys * fix(proxy): keep failing requests on keys with non-object logging metadata A key whose stored metadata.logging is not a list of objects failed every request with a 500 before this branch. The credential reference rewrite turned those values into no entries, so the request quietly succeeded without the key's logging. Raise a TypeError for those shapes again so the request path behaves as it did on main, while empty values and objects with unknown fields still go through the existing skip. * test(proxy): audit cells for credential references across surfaces, v1, dropped entries and malformed logging The malformed logging cell expects the base behavior that 7ac4ab1356 restored, a 500 for a key whose stored metadata.logging is not a list of objects and a 200 for an empty list * fix(proxy): keep a 500 for teams with non-object logging metadata On main, a team whose stored metadata.logging held non-object items failed chat requests with a 500, because AddTeamCallback(**item) raised a TypeError. This branch switched to AddTeamCallback.model_validate, whose before-validator raises AttributeError on those items, and the proxy maps AttributeError to a 400 invalid_request_error. Raise a TypeError for non-mapping items before validation so the status matches main again, as the key path already does. --- litellm/proxy/_types.py | 14 +- .../callback_config_validation.py | 81 +++- litellm/proxy/common_utils/callback_utils.py | 45 +- litellm/proxy/litellm_pre_call_utils.py | 15 +- .../internal_user_endpoints.py | 2 + .../key_management_endpoints.py | 36 +- .../management_helpers/bulk_user_creation.py | 10 +- .../observability/_openinference_support.py | 4 + .../test_arize_key_credential_reference.py | 428 ++++++++++++++++++ .../test_callback_config_validation.py | 76 ++++ .../proxy/common_utils/test_callback_utils.py | 114 +++++ .../test_key_management_endpoints.py | 29 ++ .../test_bulk_user_creation.py | 35 ++ .../unit/proxy/test_litellm_pre_call_utils.py | 64 ++- 14 files changed, 929 insertions(+), 24 deletions(-) create mode 100644 tests/integration/observability/test_arize_key_credential_reference.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 27186800e2b..6b7aa4d15d2 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2101,6 +2101,11 @@ class BlockModelRequest(LiteLLMPydanticObjectBase): model_id: str # required +CALLBACK_VAR_NAMES: Final = frozenset( + {*StandardCallbackDynamicParams.__annotations__, CAPTURE_MESSAGE_CONTENT_VAR, EXCLUDED_SERVICES_VAR} +) + + class AddTeamCallback(LiteLLMPydanticObjectBase): callback_name: str callback_type: Literal["success", "failure", "success_and_failure"] | None = "success_and_failure" @@ -2110,14 +2115,9 @@ class AddTeamCallback(LiteLLMPydanticObjectBase): @classmethod def validate_callback_vars(cls, values): callback_vars: Final = values.get("callback_vars", {}) - valid_keys: Final = { - *StandardCallbackDynamicParams.__annotations__, - CAPTURE_MESSAGE_CONTENT_VAR, - EXCLUDED_SERVICES_VAR, - } for key, value in callback_vars.items(): - if key not in valid_keys: - raise ValueError(f"Invalid callback variable: {key}. Must be one of {valid_keys}") + if key not in CALLBACK_VAR_NAMES: + raise ValueError(f"Invalid callback variable: {key}. Must be one of {set(CALLBACK_VAR_NAMES)}") if key in (CAPTURE_MESSAGE_CONTENT_VAR, EXCLUDED_SERVICES_VAR) and not isinstance(value, str): raise ValueError(f"{key} must be a string, got {type(value).__name__}") callback_vars[key] = str(value) diff --git a/litellm/proxy/common_utils/callback_config_validation.py b/litellm/proxy/common_utils/callback_config_validation.py index 302978d17a4..a486d6665fc 100644 --- a/litellm/proxy/common_utils/callback_config_validation.py +++ b/litellm/proxy/common_utils/callback_config_validation.py @@ -7,9 +7,13 @@ per-integration checks here. import math from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass from types import MappingProxyType -from typing import Final, cast # noqa: TID251 # preserve arbitrary keys from validated callback mappings +from typing import Final, Literal, cast # noqa: TID251 # preserve arbitrary keys from validated callback mappings +from pydantic import TypeAdapter, ValidationError + +from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.types.utils import CAPTURE_MESSAGE_CONTENT_VAR, EXCLUDED_SERVICES_VAR _NEWRELIC_CALLBACK: Final = "newrelic" @@ -276,7 +280,7 @@ def logging_metadata_config_error(metadata: Mapping[str, object] | None) -> str """Validate every ``logging`` entry of a team/key metadata payload.""" if not metadata: return None - entries: Final = _object_sequence(metadata.get("logging")) + entries: Final = object_sequence(metadata.get("logging")) if entries is None: return None entry_error: Final = next( @@ -302,11 +306,11 @@ def logging_metadata_config_error(metadata: Mapping[str, object] | None) -> str def stored_callback_entries(logging_entries: object) -> tuple[tuple[str | None, Mapping[str, str]], ...]: """The callback name and vars of every stored ``logging`` entry.""" - entries: Final = _object_sequence(logging_entries) or () + entries: Final = object_sequence(logging_entries) or () return tuple((_entry_callback_name(entry), _entry_callback_vars(entry)) for entry in entries) -def _object_sequence(value: object) -> Sequence[object] | None: +def object_sequence(value: object) -> Sequence[object] | None: return value if isinstance(value, Sequence) and not isinstance(value, (str, bytes)) else None @@ -445,3 +449,72 @@ def _newrelic_config_error(callback_vars: Mapping[str, str]) -> str | None: if region is not None and (not api_key or api_key == "None"): return "newrelic_region requires newrelic_api_key; the region rides the team's own key." return None + + +@dataclass(frozen=True, slots=True) +class CredentialReferenceRejection: + status_code: Literal[400, 403] + message: str + + +def credential_is_defined(credential_name: str) -> bool: + return CredentialAccessor.find_credential(credential_name) is not None + + +_STRING_KEYED: Final = TypeAdapter(Mapping[str, object]) + + +def string_keyed(value: object) -> Mapping[str, object]: + try: + return _STRING_KEYED.validate_python(value) + except ValidationError: + return MappingProxyType({}) + + +def _logging_entries(metadata: object) -> object: + return string_keyed(metadata).get("logging") + + +def _entry_credential_name(entry: object) -> str | None: + match entry: + case {"litellm_credential_name": str() as credential_name}: + return credential_name + case _: + return None + + +def logging_credential_names(logging_entries: object) -> tuple[str, ...]: + return tuple( + name for entry in object_sequence(logging_entries) or () if (name := _entry_credential_name(entry)) is not None + ) + + +def credential_reference_rejection( + requested_metadata: object, + stored_logging: object, + *, + is_proxy_admin: bool, + is_defined: Callable[[str], bool] = credential_is_defined, +) -> CredentialReferenceRejection | None: + """Gate a key metadata write whose logging entries reference a proxy credential. + + ``requested_metadata`` is None when the write leaves the key's metadata untouched. The admin check runs before + the existence check so a non-admin cannot probe which credential names are defined. + """ + if requested_metadata is None: + return None + requested_logging: Final = _logging_entries(requested_metadata) + requested_names: Final = logging_credential_names(requested_logging) + stored_names: Final = logging_credential_names(stored_logging) + references_involved: Final = bool(requested_names or stored_names) + if references_involved and not is_proxy_admin and requested_logging != stored_logging: + return CredentialReferenceRejection( + 403, + "Only proxy admins can change the logging entries of a key that uses litellm_credential_name", + ) + undefined: Final = sorted({name for name in set(requested_names) - set(stored_names) if not is_defined(name)}) + if undefined: + return CredentialReferenceRejection( + 400, f"metadata.logging references undefined credential(s): {', '.join(undefined)}" + ) + return None diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index c195779aa0a..0aa8d2291ef 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -26,8 +26,10 @@ from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, get_or_create_metadata_bucket, ) +from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker -from litellm.proxy._types import CommonProxyErrors, LiteLLMPromptInjectionParams +from litellm.proxy._types import CALLBACK_VAR_NAMES, CommonProxyErrors, LiteLLMPromptInjectionParams +from litellm.proxy.common_utils.callback_config_validation import object_sequence, string_keyed from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, @@ -736,6 +738,47 @@ def decrypt_callback_vars(metadata: object) -> Any: return _transform_callback_vars(metadata, _decrypt_or_passthrough) +def decrypted_callback_metadata(metadata: object) -> object: + return _transform_callback_vars(metadata, _decrypt_or_passthrough) + + +def decrypted_logging_entries(metadata: object) -> object: + return string_keyed(decrypted_callback_metadata(metadata)).get("logging") + + +def key_logging_entries(metadata: object) -> tuple[Mapping[str, object], ...] | None: + entries: Final = decrypted_logging_entries(metadata) + if entries is None: + return None + return tuple(resolved for entry in _logging_objects(entries) if (resolved := _with_credential(entry)) is not None) + + +def _logging_objects(entries: object) -> tuple[Mapping[str, object], ...]: + items: Final = () if entries in ("", {}) else object_sequence(entries) + if items is None or not all(isinstance(item, Mapping) for item in items): + raise TypeError(f"key metadata.logging must be a list of objects, got {type(entries).__name__}") + return tuple(string_keyed(item) for item in items) + + +def _with_credential(entry: Mapping[str, object]) -> Mapping[str, object] | None: + credential_name: Final = entry.get("litellm_credential_name") + if not isinstance(credential_name, str): + return entry + credential_vars: Final = { + name: str(value) + for name, value in CredentialAccessor.get_credential_values(credential_name).items() + if name in CALLBACK_VAR_NAMES and value is not None + } + if not credential_vars: + verbose_proxy_logger.error( + "Ignoring key logging entry for %s: credential %r is not defined or holds no logging values", + entry.get("callback_name"), + credential_name, + ) + return None + return {**entry, "callback_vars": {**credential_vars, **string_keyed(entry.get("callback_vars"))}} + + def _transform_callback_vars(metadata: object, transform: Callable[[str, object], object]) -> object: if not isinstance(metadata, dict): return metadata diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 9cc18294cd1..0bb4a803609 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -88,6 +88,7 @@ from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.callback_utils import ( decrypt_callback_vars, get_metadata_variable_name_from_kwargs, + key_logging_entries, strip_callback_config, ) from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports @@ -1246,9 +1247,11 @@ def convert_key_logging_metadata_to_callback( return team_callback_settings_obj -def get_validated_callback_metadata(item: dict, *, source: str) -> AddTeamCallback | None: +def get_validated_callback_metadata(item: object, *, source: str) -> AddTeamCallback | None: + if not isinstance(item, Mapping): + raise TypeError(f"{source} metadata.logging must be a list of objects, got an item of {type(item).__name__}") try: - return AddTeamCallback(**item) + return AddTeamCallback.model_validate(item) except (PydanticValidationError, ValueError) as e: verbose_proxy_logger.warning( "Ignoring invalid %s callback metadata: %s", @@ -1267,9 +1270,9 @@ class KeyAndTeamLoggingSettings: """ @staticmethod - def get_key_dynamic_logging_settings(user_api_key_dict: UserAPIKeyAuth): + def get_key_dynamic_logging_settings(user_api_key_dict: UserAPIKeyAuth) -> tuple[Mapping[str, object], ...] | None: if user_api_key_dict.metadata is not None and "logging" in user_api_key_dict.metadata: - return decrypt_callback_vars(user_api_key_dict.metadata).get("logging") + return key_logging_entries(user_api_key_dict.metadata) return None @staticmethod @@ -1283,9 +1286,7 @@ def get_dynamic_logging_metadata( user_api_key_dict: UserAPIKeyAuth, proxy_config: ProxyConfig ) -> TeamCallbackMetadata | None: callback_settings_obj: TeamCallbackMetadata | None = None - key_dynamic_logging_settings: Final[dict | None] = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings( - user_api_key_dict - ) + key_dynamic_logging_settings: Final = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict) team_dynamic_logging_settings: Final[dict | None] = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings( user_api_key_dict ) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index e35a65d8ba8..f88ac445a0b 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -79,6 +79,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( # noq check_permissions_caller_permission, generate_key_helper_fn, prepare_metadata_fields, + raise_on_logging_credential_change, ) from litellm.proxy.management_helpers.bulk_user_deletion import delete_users_from_current_teams from litellm.proxy.management_helpers.object_permission_utils import ( # noqa: F401 # legacy module exports @@ -598,6 +599,7 @@ async def new_user( if data.auto_create_key and isinstance(user_api_key_dict, UserAPIKeyAuth): enforce_batch_limits_are_admin_only(data, None, user_api_key_dict, "key") enforce_vector_store_index_grants_are_admin_only(data, None, None, user_api_key_dict, "key") + raise_on_logging_credential_change(data.metadata, None, user_api_key_dict) check_permissions_caller_permission( data=data, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 50517eeefcf..038c1bdf5a6 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -26,7 +26,7 @@ from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, TypeVar, c import fastapi import yaml from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, status -from pydantic import TypeAdapter +from pydantic import ConfigDict, TypeAdapter from typing_extensions import ReadOnly, TypedDict import litellm @@ -68,9 +68,14 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( evict_and_broadcast, ) -from litellm.proxy.common_utils.callback_config_validation import logging_metadata_config_error +from litellm.proxy.common_utils.callback_config_validation import ( + credential_reference_rejection, + logging_metadata_config_error, +) from litellm.proxy.common_utils.callback_utils import ( decrypt_callback_vars, + decrypted_callback_metadata, + decrypted_logging_entries, encrypt_callback_vars, ) from litellm.proxy.common_utils.config_sync_pubsub import ( @@ -266,6 +271,12 @@ class _StoredKeyRouterSettings(LiteLLMBaseModel): router_settings: Mapping[str, object] | None = None +class _StoredKeyMetadata(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + metadata: object = None + + class _KeyRowWhere(TypedDict): token: ReadOnly[str] @@ -934,6 +945,23 @@ def raise_on_invalid_key_logging_config(metadata: Mapping[str, object] | None) - raise HTTPException(status_code=400, detail={"error": error}) +def raise_on_logging_credential_change( + requested_metadata: object, existing_key_row: LiteLLM_VerificationToken | None, user_api_key_dict: UserAPIKeyAuth +) -> None: + stored_metadata: Final = ( + None + if existing_key_row is None + else _StoredKeyMetadata.model_validate(existing_key_row, from_attributes=True).metadata + ) + rejection: Final = credential_reference_rejection( + decrypted_callback_metadata(requested_metadata), + decrypted_logging_entries(stored_metadata), + is_proxy_admin=user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN, + ) + if rejection is not None: + raise HTTPException(status_code=rejection.status_code, detail={"error": rejection.message}) + + def common_key_access_checks( user_api_key_dict: UserAPIKeyAuth, data: GenerateKeyRequest | UpdateKeyRequest, @@ -1485,6 +1513,7 @@ async def _common_key_generation_helper( validate_budget_duration(data.budget_duration) raise_on_invalid_key_logging_config(data.metadata) + raise_on_logging_credential_change(data.metadata, None, user_api_key_dict) if data.throttle_on_budget_exceeded is True and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: raise HTTPException( @@ -3430,6 +3459,7 @@ async def _process_single_key_update( ) # Prepare update data + raise_on_logging_credential_change(key_request.metadata, existing_key_row, user_api_key_dict) non_default_values = await prepare_key_update_data( data=key_request, existing_key_row=existing_key_row, prisma_client=prisma_client, llm_router=llm_router ) @@ -4176,6 +4206,7 @@ async def update_key_fn( # Enforce upperbound key params on update (don't fill defaults) _enforce_upperbound_key_params(data, fill_defaults=False) + raise_on_logging_credential_change(data.metadata, existing_key_row, user_api_key_dict) non_default_values: Final = await prepare_key_update_data( data=data, existing_key_row=existing_key_row, prisma_client=prisma_client, llm_router=llm_router ) @@ -6363,6 +6394,7 @@ async def _execute_virtual_key_regeneration( await _enforce_custom_key_update_policy(hook=_custom_key_update_hook(proxy_server), data=update_request) # Enforce upperbound key params on regenerate (don't fill defaults) _enforce_upperbound_key_params(data, fill_defaults=False) + raise_on_logging_credential_change(data.metadata, key_in_db, user_api_key_dict) non_default_values = await prepare_key_update_data( data=data, existing_key_row=key_in_db, prisma_client=prisma_client, llm_router=llm_router ) diff --git a/litellm/proxy/management_helpers/bulk_user_creation.py b/litellm/proxy/management_helpers/bulk_user_creation.py index e900941f67b..77e37cb11cb 100644 --- a/litellm/proxy/management_helpers/bulk_user_creation.py +++ b/litellm/proxy/management_helpers/bulk_user_creation.py @@ -34,6 +34,7 @@ from litellm.proxy.auth.auth_utils import ( enforce_vector_store_index_grants_are_admin_only, ) from litellm.proxy.auth.litellm_license import LicenseCheck +from litellm.proxy.common_utils.callback_config_validation import credential_reference_rejection from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks @@ -225,7 +226,14 @@ def _row_error(item: BulkNewUserItem, user_api_key_dict: UserAPIKeyAuth) -> str enforce_vector_store_index_grants_are_admin_only(item, None, None, user_api_key_dict, "key") except Exception as exc: # noqa: BLE001 # any validation failure is reported on this row only return _error_message(exc) - return None + rejection: Final = ( + credential_reference_rejection( + item.metadata, None, is_proxy_admin=user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + ) + if item.auto_create_key + else None + ) + return None if rejection is None else rejection.message def _normalized_email(email: str | None) -> str | None: diff --git a/tests/integration/observability/_openinference_support.py b/tests/integration/observability/_openinference_support.py index 58785c89033..3fc4b2e7628 100644 --- a/tests/integration/observability/_openinference_support.py +++ b/tests/integration/observability/_openinference_support.py @@ -672,10 +672,12 @@ def _write_config( callback_settings: Mapping[str, JsonValue] | None = None, litellm_settings: Mapping[str, JsonValue] | None = None, general_settings: Mapping[str, JsonValue] | None = None, + extra_config: Mapping[str, JsonValue] | None = None, ) -> Path: loaded: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) config: Final = { **loaded, + **(extra_config or {}), "litellm_settings": { **loaded["litellm_settings"], "callbacks": list(callbacks), @@ -752,6 +754,7 @@ def _rig( callback_settings: Mapping[str, JsonValue] | None = None, litellm_settings: Mapping[str, JsonValue] | None = None, general_settings: Mapping[str, JsonValue] | None = None, + extra_config: Mapping[str, JsonValue] | None = None, environment: Mapping[str, str] | None = None, remove_environment: tuple[str, ...] = (), disabled_environment: tuple[str, ...] = (), @@ -782,6 +785,7 @@ def _rig( callback_settings=resolved_settings, litellm_settings=litellm_settings, general_settings=general_settings, + extra_config=extra_config, ) preset_environment: Final = { **({"PHOENIX_COLLECTOR_ENDPOINT": destination.url + "/v1/traces"} if "arize_phoenix" in callbacks else {}), diff --git a/tests/integration/observability/test_arize_key_credential_reference.py b/tests/integration/observability/test_arize_key_credential_reference.py new file mode 100644 index 00000000000..b1087c4f4c5 --- /dev/null +++ b/tests/integration/observability/test_arize_key_credential_reference.py @@ -0,0 +1,428 @@ +from __future__ import annotations + +import operator +import time +import uuid +from collections.abc import Callable, Iterator +from itertools import accumulate +from pathlib import Path +from typing import Final + +import httpx +import pytest +from _openinference_support import ( + _anthropic_response, + _chat_request_marker, + _chat_response, + _chat_stream_response, + _chat_tool_call, + _matching_marker_spans, + _responses_response, + _rig, +) +from integration._support.client import Gateway, eventually +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +CREDENTIAL_NAME: Final = "arize-space-a" +SECRET_API_KEY: Final = "space-a-arize-key-" + uuid.uuid4().hex +SECRET_SPACE_ID: Final = "space-a-id-" + uuid.uuid4().hex +AMBIENT_ARIZE_ENVIRONMENT: Final = ("ARIZE_API_KEY", "ARIZE_SPACE_ID") +CREDENTIAL_CONFIG: Final[dict[str, JsonValue]] = { + "credential_list": [ + { + "credential_name": CREDENTIAL_NAME, + "credential_values": { + "arize_api_key": "os.environ/SPACE_A_ARIZE_API_KEY", + "arize_space_id": "os.environ/SPACE_A_ARIZE_SPACE_ID", + }, + "credential_info": {}, + }, + { + "credential_name": "model-only", + "credential_values": {"api_key": "sk-not-a-logging-value"}, + "credential_info": {}, + }, + ] +} +SECRET_ENVIRONMENT: Final = {"SPACE_A_ARIZE_API_KEY": SECRET_API_KEY, "SPACE_A_ARIZE_SPACE_ID": SECRET_SPACE_ID} + + +def _reference_entry(credential_name: str = CREDENTIAL_NAME) -> dict[str, JsonValue]: + return { + "callback_name": "arize", + "callback_type": "success", + "litellm_credential_name": credential_name, + "callback_vars": {}, + } + + +def _echo_marker(request: Request) -> Reply: + return _chat_response(_chat_request_marker(request)) + + +def _send_chat(proxy: Gateway, model: str, key: str, marker: str) -> httpx.Response: + return proxy.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": marker}], + "metadata": {"trace_marker": marker}, + "cache": {"no-cache": True}, + }, + key=key, + ) + + +def _chat(proxy: Gateway, model: str, key: str, marker: str) -> None: + response: Final = _send_chat(proxy, model, key, marker) + assert response.status_code == 200, response.text + + +def _drains(destination: Wire) -> Iterator[tuple[Request, ...]]: + while True: + yield destination.drain() + time.sleep(0.1) + + +def _carries(request: Request, marker: str) -> bool: + return any(True for _ in _matching_marker_spans((request,), marker)) + + +def _received_until(destination: Wire, marker: str, seconds: float = 30) -> tuple[Request, ...]: + deadline: Final = time.monotonic() + seconds + received: Final = next( + batch + for batch in accumulate(_drains(destination), operator.add) + if any(_carries(request, marker) for request in batch) or time.monotonic() >= deadline + ) + assert any(_carries(request, marker) for request in received), f"no span for {marker} reached the destination" + return received + + +def test_key_referencing_a_credential_exports_with_the_secret_manager_values(gateway: Gateway, tmp_path: Path) -> None: + unreferenced_marker: Final = "plain-" + uuid.uuid4().hex + referenced_marker: Final = "ref-" + uuid.uuid4().hex + with ( + _rig( + gateway, + tmp_path, + _echo_marker, + extra_config=CREDENTIAL_CONFIG, + environment=SECRET_ENVIRONMENT, + remove_environment=AMBIENT_ARIZE_ENVIRONMENT, + disabled_environment=AMBIENT_ARIZE_ENVIRONMENT, + ) as rig, + rig.proxy.scenario() as scenario, + ): + team: Final = scenario.team() + referencing_key: Final = scenario.key(team_id=team, metadata={"logging": [_reference_entry()]}) + unreferenced_key: Final = scenario.key(team_id=team) + + _chat(rig.proxy, rig.model, unreferenced_key, unreferenced_marker) + _chat(rig.proxy, rig.model, referencing_key, referenced_marker) + received: Final = _received_until(rig.destination, referenced_marker) + + exports: Final = tuple(request for request in received if _carries(request, referenced_marker)) + for export in exports: + headers = {name.lower(): value for name, value in export.headers.items()} + assert headers.get("api_key") == SECRET_API_KEY, headers + assert headers.get("arize-space-id") == SECRET_SPACE_ID, headers + assert not any(_carries(request, unreferenced_marker) for request in received) + + info: Final = rig.proxy.request("GET", "/key/info", params={"key": referencing_key}) + assert info.status_code == 200, info.text + assert CREDENTIAL_NAME in info.text + assert SECRET_API_KEY not in info.text + assert SECRET_SPACE_ID not in info.text + + +def test_only_proxy_admins_change_logging_on_a_key_that_references_a_credential( + gateway: Gateway, tmp_path: Path +) -> None: + with ( + _rig(gateway, tmp_path, _echo_marker, extra_config=CREDENTIAL_CONFIG) as rig, + rig.proxy.scenario() as scenario, + ): + team: Final = scenario.team() + team_admin: Final = scenario.member(team, role="admin") + team_admin_key: Final = scenario.key(user_id=team_admin, team_id=team) + langfuse_entry: Final[dict[str, JsonValue]] = { + "callback_name": "langfuse", + "callback_type": "success", + "callback_vars": {"langfuse_public_key": "pk-lf-integration", "langfuse_secret_key": "sk-lf-integration"}, + } + referencing_key: Final = scenario.key(team_id=team, metadata={"logging": [_reference_entry(), langfuse_entry]}) + + def as_team_admin(path: str, body: dict[str, JsonValue]) -> httpx.Response: + return rig.proxy.request("POST", path, body, key=team_admin_key) + + def assert_gated(response: httpx.Response) -> None: + assert response.status_code == 403, response.text + assert "litellm_credential_name" in response.text, response.text + + assert_gated(as_team_admin("/key/generate", {"team_id": team, "metadata": {"logging": [_reference_entry()]}})) + + sibling: Final = as_team_admin( + "/key/update", + { + "key": referencing_key, + "metadata": { + "logging": [ + _reference_entry(), + { + "callback_name": "arize", + "callback_type": "success", + "callback_vars": {"arize_endpoint": "http://127.0.0.1:9/v1/traces"}, + }, + ] + }, + }, + ) + assert_gated(sibling) + assert_gated(as_team_admin("/key/update", {"key": referencing_key, "metadata": {"team": "unrelated"}})) + assert_gated(as_team_admin("/key/regenerate", {"key": referencing_key, "metadata": {"team": "unrelated"}})) + + bulk: Final = as_team_admin( + "/team/key/bulk_update", + {"team_id": team, "key_ids": [referencing_key], "update_fields": {"metadata": {"team": "unrelated"}}}, + ) + assert bulk.status_code == 200, bulk.text + assert bulk.json()["successful_updates"] == [], bulk.text + assert "litellm_credential_name" in bulk.json()["failed_updates"][0]["failed_reason"], bulk.text + + renamed: Final = as_team_admin( + "/key/update", {"key": referencing_key, "key_alias": "alias-" + uuid.uuid4().hex} + ) + assert renamed.status_code == 200, renamed.text + + read_back: Final = rig.proxy.request("GET", "/key/info", params={"key": referencing_key}, key=team_admin_key) + assert read_back.status_code == 200, read_back.text + resaved: Final = as_team_admin( + "/key/update", + {"key": referencing_key, "metadata": {**read_back.json()["info"]["metadata"], "note": "unrelated edit"}}, + ) + assert resaved.status_code == 200, resaved.text + + stored: Final = rig.proxy.get("/key/info", params={"key": referencing_key}) + logging_entries: Final = stored["info"]["metadata"]["logging"] + assert [entry["callback_name"] for entry in logging_entries] == ["arize", "langfuse"], stored + assert logging_entries[0] == _reference_entry(), stored + + undefined: Final = rig.proxy.request( + "POST", "/key/generate", {"metadata": {"logging": [_reference_entry("no-such-credential")]}} + ) + assert undefined.status_code == 400, undefined.text + assert "no-such-credential" in undefined.text, undefined.text + + +def test_only_proxy_admins_create_a_user_whose_key_references_a_credential(gateway: Gateway, tmp_path: Path) -> None: + with ( + _rig(gateway, tmp_path, _echo_marker, extra_config=CREDENTIAL_CONFIG) as rig, + rig.proxy.scenario() as scenario, + ): + organization: Final = scenario.organization() + org_admin_key: Final = scenario.key(user_id=scenario.org_member(organization, "org_admin")) + referencing_metadata: Final = {"logging": [_reference_entry()]} + + single: Final = rig.proxy.request( + "POST", + "/user/new", + {"user_id": "user-" + uuid.uuid4().hex, "organizations": [organization], "metadata": referencing_metadata}, + key=org_admin_key, + ) + assert single.status_code == 403, single.text + assert "litellm_credential_name" in single.text, single.text + + delegate_key: Final = scenario.key( + user_id=scenario.user(user_role="internal_user"), allowed_routes=["/management/v1/users/bulk"] + ) + bulk: Final = rig.proxy.request( + "POST", + "/management/v1/users/bulk", + { + "users": [ + {"user_id": "user-" + uuid.uuid4().hex, "auto_create_key": True, "metadata": referencing_metadata} + ] + }, + key=delegate_key, + ) + assert bulk.status_code == 200, bulk.text + assert [row["success"] for row in bulk.json()["data"]] == [False], bulk.text + assert "litellm_credential_name" in bulk.json()["data"][0]["error"], bulk.text + + +def _carries_credential(request: Request) -> bool: + headers: Final = {name.lower(): value for name, value in request.headers.items()} + return headers.get("api_key") == SECRET_API_KEY and headers.get("arize-space-id") == SECRET_SPACE_ID + + +def _credential_exports(destination: Wire, seconds: float = 30) -> tuple[Request, ...]: + return eventually( + lambda: tuple(request for request in destination.drain() if _carries_credential(request)), + bool, + seconds=seconds, + ) + + +def _surface_reply(marker: str) -> Callable[[Request], Reply]: + def reply(request: Request) -> Reply: + if request.target.endswith("/responses"): + return _responses_response(marker) + if request.target.endswith("/messages"): + return _anthropic_response(marker) + return _chat_stream_response(marker, (_chat_tool_call(marker),)) + + return reply + + +_SURFACES: Final[dict[str, tuple[str, str, dict[str, JsonValue]]]] = { + "responses": ("/v1/responses", "openai/gpt-4o-mini", {"input": "weather in Paris?"}), + "messages": ( + "/v1/messages", + "anthropic/claude-opus-5-5", + {"max_tokens": 64, "messages": [{"role": "user", "content": "weather in Paris?"}]}, + ), + "chat-stream": ( + "/v1/chat/completions", + "openai/gpt-4o-mini", + {"stream": True, "messages": [{"role": "user", "content": "weather in Paris?"}]}, + ), +} + + +@pytest.mark.parametrize("surface", sorted(_SURFACES)) +def test_referencing_key_exports_on_each_surface(surface: str, gateway: Gateway, tmp_path: Path) -> None: + marker: Final = f"{surface}-" + uuid.uuid4().hex + path, model_name, body = _SURFACES[surface] + with ( + _rig( + gateway, + tmp_path, + _surface_reply(marker), + extra_config=CREDENTIAL_CONFIG, + environment=SECRET_ENVIRONMENT, + remove_environment=AMBIENT_ARIZE_ENVIRONMENT, + disabled_environment=AMBIENT_ARIZE_ENVIRONMENT, + model_name=model_name, + api_base_suffix="" if surface == "messages" else "/v1", + ) as rig, + rig.proxy.scenario() as scenario, + ): + key: Final = scenario.key(metadata={"logging": [_reference_entry()]}) + response: Final = rig.proxy.request( + "POST", + path, + {"model": rig.model, **body, "metadata": {"trace_marker": marker}, "cache": {"no-cache": True}}, + key=key, + ) + assert response.status_code == 200, response.text + received: Final = _received_until(rig.destination, marker) + exports: Final = tuple(request for request in received if _carries(request, marker)) + assert exports and all(_carries_credential(request) for request in exports), [ + dict(request.headers) for request in exports + ] + + +def test_referencing_key_exports_through_anthropic_pass_through(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "passthrough-" + uuid.uuid4().hex + with wire_server(lambda _request: _anthropic_response(marker)) as anthropic: + with ( + _rig( + gateway, + tmp_path, + _echo_marker, + extra_config=CREDENTIAL_CONFIG, + environment={ + **SECRET_ENVIRONMENT, + "ANTHROPIC_API_BASE": anthropic.url, + "ANTHROPIC_API_KEY": "sk-ant-synthetic", + }, + remove_environment=AMBIENT_ARIZE_ENVIRONMENT, + disabled_environment=AMBIENT_ARIZE_ENVIRONMENT, + ) as rig, + rig.proxy.scenario() as scenario, + ): + key: Final = scenario.key(metadata={"logging": [_reference_entry()]}) + response: Final = rig.proxy.request( + "POST", + "/anthropic/v1/messages", + {"model": "claude-opus-5-5", "max_tokens": 16, "messages": [{"role": "user", "content": marker}]}, + key=key, + headers={"anthropic-version": "2023-06-01"}, + ) + assert response.status_code == 200, response.text + assert response.json()["id"] == marker, response.text + assert anthropic.drain(), "the pass-through never reached the scripted Anthropic upstream" + assert _credential_exports(rig.destination) + + +def test_referencing_key_exports_on_otel_v1(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "otel-v1-" + uuid.uuid4().hex + with ( + _rig( + gateway, + tmp_path, + _echo_marker, + extra_config=CREDENTIAL_CONFIG, + environment={**SECRET_ENVIRONMENT, "LITELLM_OTEL_V2": "false"}, + remove_environment=AMBIENT_ARIZE_ENVIRONMENT, + disabled_environment=AMBIENT_ARIZE_ENVIRONMENT, + ) as rig, + rig.proxy.scenario() as scenario, + ): + key: Final = scenario.key(metadata={"logging": [_reference_entry()]}) + _chat(rig.proxy, rig.model, key, marker) + assert _credential_exports(rig.destination) + + +def test_reference_to_a_credential_without_logging_values_exports_nothing_and_is_unhealthy( + gateway: Gateway, tmp_path: Path +) -> None: + dropped_marker: Final = "model-only-" + uuid.uuid4().hex + sentinel_marker: Final = "sentinel-" + uuid.uuid4().hex + with ( + _rig( + gateway, + tmp_path, + _echo_marker, + extra_config=CREDENTIAL_CONFIG, + environment=SECRET_ENVIRONMENT, + remove_environment=AMBIENT_ARIZE_ENVIRONMENT, + disabled_environment=AMBIENT_ARIZE_ENVIRONMENT, + ) as rig, + rig.proxy.scenario() as scenario, + ): + dropped_key: Final = scenario.key(metadata={"logging": [_reference_entry("model-only")]}) + sentinel_key: Final = scenario.key(metadata={"logging": [_reference_entry()]}) + + _chat(rig.proxy, rig.model, dropped_key, dropped_marker) + _chat(rig.proxy, rig.model, sentinel_key, sentinel_marker) + received: Final = _received_until(rig.destination, sentinel_marker) + assert not any(_carries(request, dropped_marker) for request in received) + + health: Final = rig.proxy.request("POST", "/key/health", {}, key=dropped_key) + assert health.status_code == 200, health.text + assert health.json()["key"] == "unhealthy", health.text + + +def test_malformed_key_logging_fails_requests_like_main(gateway: Gateway, tmp_path: Path) -> None: + stored_logging: Final[dict[str, JsonValue]] = { + "string": "abc", + "list-with-string": ["x"], + "object": {"callback_name": "arize"}, + "empty-list": [], + } + with ( + _rig(gateway, tmp_path, _echo_marker, extra_config=CREDENTIAL_CONFIG) as rig, + rig.proxy.scenario() as scenario, + ): + statuses: Final = { + label: _send_chat( + rig.proxy, rig.model, scenario.key(metadata={"logging": value}), "malformed-" + uuid.uuid4().hex + ).status_code + for label, value in stored_logging.items() + } + + assert statuses == {"string": 500, "list-with-string": 500, "object": 500, "empty-list": 200}, statuses diff --git a/tests/unit/proxy/common_utils/test_callback_config_validation.py b/tests/unit/proxy/common_utils/test_callback_config_validation.py index f42387d3e76..ac1fe2d785e 100644 --- a/tests/unit/proxy/common_utils/test_callback_config_validation.py +++ b/tests/unit/proxy/common_utils/test_callback_config_validation.py @@ -7,6 +7,7 @@ from litellm.proxy.common_utils.callback_config_validation import ( callback_config_error, conflicting_callback_var_error, conflicting_span_scope_error, + credential_reference_rejection, cross_entry_family_error, logging_metadata_config_error, ) @@ -499,3 +500,78 @@ def test_key_logging_entries_compare_excluded_services_by_parsed_value( def test_a_failure_only_entry_rejects_excluded_services() -> None: error: Final = callback_config_error("langfuse_otel", {"excluded_services": "redis"}, "failure") assert error is not None and "excluded_services" in error and "success_and_failure" in error + + +_REFERENCE: Final = {"callback_name": "arize", "callback_type": "success", "litellm_credential_name": "space-a"} +_INLINE: Final = {"callback_name": "langfuse", "callback_vars": {"langfuse_public_key": "pk"}} +_SIBLING: Final = {"callback_name": "arize", "callback_vars": {"arize_endpoint": "http://attacker.example/v1/traces"}} + + +@pytest.mark.parametrize( + ("requested", "stored_logging", "is_proxy_admin", "expected_status"), + [ + pytest.param({"logging": [_REFERENCE]}, None, False, 403, id="non-admin-adds-reference"), + pytest.param({"logging": [_INLINE, _REFERENCE]}, None, False, 403, id="non-admin-adds-reference-second"), + pytest.param({"team": "x"}, [_REFERENCE], False, 403, id="non-admin-removes-reference"), + pytest.param({"logging": [_REFERENCE, _SIBLING]}, [_REFERENCE], False, 403, id="non-admin-adds-sibling"), + pytest.param({"logging": [_REFERENCE]}, [_REFERENCE], False, None, id="non-admin-resends-same"), + pytest.param(None, [_REFERENCE], False, None, id="non-admin-leaves-metadata-alone"), + pytest.param({"logging": [_INLINE]}, [], False, None, id="non-admin-without-references"), + pytest.param({"logging": [_REFERENCE]}, [], True, None, id="admin-adds-defined-reference"), + pytest.param({"team": "x"}, [_REFERENCE], True, None, id="admin-removes-reference"), + ], +) +def test_credential_reference_rejection_lets_only_proxy_admins_change_referencing_logging( + requested: dict[str, object] | None, + stored_logging: list[object] | None, + is_proxy_admin: bool, + expected_status: int | None, +) -> None: + rejection: Final = credential_reference_rejection( + requested, stored_logging, is_proxy_admin=is_proxy_admin, is_defined=lambda name: name == "space-a" + ) + assert (None if rejection is None else rejection.status_code) == expected_status + + +def test_credential_reference_rejection_names_undefined_credentials_only_to_admins() -> None: + requested: Final = {"logging": [_INLINE, {**_REFERENCE, "litellm_credential_name": "missing"}]} + + admin: Final = credential_reference_rejection(requested, None, is_proxy_admin=True, is_defined=lambda _: False) + non_admin: Final = credential_reference_rejection(requested, None, is_proxy_admin=False, is_defined=lambda _: False) + + assert admin is not None and admin.status_code == 400 and "missing" in admin.message + assert non_admin is not None and non_admin.status_code == 403 and "missing" not in non_admin.message + + +@pytest.mark.parametrize("is_proxy_admin", [True, False], ids=["admin", "non-admin"]) +def test_credential_reference_rejection_lets_a_stored_reference_to_a_deleted_credential_stay( + is_proxy_admin: bool, +) -> None: + stale: Final = {**_REFERENCE, "litellm_credential_name": "deleted"} + + kept: Final = credential_reference_rejection( + {"logging": [stale], "note": "unrelated edit"}, + [stale], + is_proxy_admin=is_proxy_admin, + is_defined=lambda _: False, + ) + added: Final = credential_reference_rejection( + {"logging": [stale, {**_REFERENCE, "litellm_credential_name": "also-missing"}]}, + [stale], + is_proxy_admin=True, + is_defined=lambda _: False, + ) + + assert kept is None + assert added is not None and added.status_code == 400 + assert "also-missing" in added.message and "deleted" not in added.message + + +@pytest.mark.parametrize("name", [5, ["space-a"], None], ids=["int", "list", "null"]) +def test_a_non_string_credential_name_is_not_a_reference(name: object) -> None: + entry: Final = {**_REFERENCE, "litellm_credential_name": name} + + assert ( + credential_reference_rejection({"logging": [entry]}, None, is_proxy_admin=False, is_defined=lambda _: False) + is None + ) diff --git a/tests/unit/proxy/common_utils/test_callback_utils.py b/tests/unit/proxy/common_utils/test_callback_utils.py index c8d909ac2eb..c7739b9cc2f 100644 --- a/tests/unit/proxy/common_utils/test_callback_utils.py +++ b/tests/unit/proxy/common_utils/test_callback_utils.py @@ -1,5 +1,6 @@ import copy import json +import logging import sys from types import ModuleType, SimpleNamespace from typing import Final @@ -21,6 +22,7 @@ from litellm.proxy.common_utils.callback_utils import ( get_logging_caching_headers, get_remaining_tokens_and_requests_from_request_data, initialize_callbacks_on_proxy, + key_logging_entries, normalize_callback_names, process_callback, sanitize_openai_provider_metadata, @@ -28,6 +30,7 @@ from litellm.proxy.common_utils.callback_utils import ( ) from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import CredentialItem def test_get_remaining_tokens_and_requests_from_request_data(): @@ -775,3 +778,114 @@ def test_initialize_callbacks_on_proxy_accepts_instance_non_list_value(probe_con assert len(litellm.callbacks) == 1 assert isinstance(litellm.callbacks[0], CustomLogger) + + +def _credential(name: str, values: dict[str, object]) -> CredentialItem: + return CredentialItem(credential_name=name, credential_values=values, credential_info={}) + + +def test_key_logging_entries_merges_each_referenced_credential_under_its_entry(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + litellm, + "credential_list", + [ + _credential("space-a", {"arize_api_key": "key-a", "arize_space_id": "space-a", "api_base": "x"}), + _credential("space-b", {"arize_api_key": "key-b", "arize_space_id": "space-b"}), + ], + ) + entries: Final = [ + {"callback_name": "langfuse", "callback_vars": {"langfuse_public_key": "pk"}}, + {"callback_name": "arize", "litellm_credential_name": "space-a", "callback_vars": {"arize_space_id": "own"}}, + {"callback_name": "arize", "litellm_credential_name": "space-b", "callback_vars": {}}, + ] + + resolved: Final = key_logging_entries({"logging": entries}) + + assert resolved is not None + assert [entry["callback_vars"] for entry in resolved] == [ + {"langfuse_public_key": "pk"}, + {"arize_api_key": "key-a", "arize_space_id": "own"}, + {"arize_api_key": "key-b", "arize_space_id": "space-b"}, + ] + + +@pytest.mark.parametrize( + "credential_values", + [{"api_key": "sk-model-key"}, {"arize_api_key": None}], + ids=["model-credential", "null-values"], +) +def test_key_logging_entries_drops_an_entry_whose_credential_supplies_no_logging_values( + monkeypatch: pytest.MonkeyPatch, credential_values: dict[str, object] +) -> None: + monkeypatch.setattr(litellm, "credential_list", [_credential("model-only", credential_values)]) + kept: Final = {"callback_name": "langfuse", "callback_vars": {"langfuse_public_key": "pk"}} + entries: Final = [ + kept, + {"callback_name": "arize", "litellm_credential_name": "model-only", "callback_vars": {}}, + ] + + assert key_logging_entries({"logging": entries}) == (kept,) + + +def test_a_dropped_credential_reference_is_logged_as_an_error_so_key_health_reports_it( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + monkeypatch.setattr(litellm, "credential_list", []) + entry: Final = {"callback_name": "arize", "litellm_credential_name": "missing", "callback_vars": {}} + + with caplog.at_level(logging.WARNING): + key_logging_entries({"logging": [entry]}) + + assert [record.levelno for record in caplog.records if "missing" in record.getMessage()] == [logging.ERROR] + + +def test_key_logging_entries_drops_an_entry_whose_credential_is_not_defined(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "credential_list", [_credential("space-a", {"arize_api_key": "key-a"})]) + kept: Final = {"callback_name": "langfuse", "callback_vars": {"langfuse_public_key": "pk"}} + entries: Final = [ + kept, + {"callback_name": "arize", "litellm_credential_name": "missing", "callback_vars": {}}, + ] + + assert key_logging_entries({"logging": entries}) == (kept,) + + +@pytest.mark.parametrize( + "logging_value", + [ + pytest.param("abc", id="string"), + pytest.param({"callback_name": "arize"}, id="object"), + pytest.param(5, id="number"), + pytest.param(["x"], id="list-with-string"), + pytest.param([{"callback_name": "langfuse", "callback_vars": {}}, None], id="list-with-null"), + ], +) +def test_key_logging_entries_fails_on_non_object_logging_like_the_request_path_always_has( + logging_value: object, +) -> None: + with pytest.raises(TypeError): + key_logging_entries({"logging": logging_value}) + + +@pytest.mark.parametrize( + ("logging_value", "expected"), + [ + pytest.param("", (), id="empty-string"), + pytest.param({}, (), id="empty-object"), + pytest.param([], (), id="empty-list"), + pytest.param([{"bad": 1}], ({"bad": 1},), id="object-with-unknown-fields"), + ], +) +def test_key_logging_entries_passes_empty_and_object_entries_through( + logging_value: object, expected: tuple[dict[str, object], ...] +) -> None: + assert key_logging_entries({"logging": logging_value}) == expected + + +def test_key_logging_entries_leaves_an_entry_with_a_non_string_credential_name_as_written( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "credential_list", [_credential("5", {"arize_api_key": "k"})]) + entry: Final = {"callback_name": "arize", "litellm_credential_name": 5, "callback_vars": {}} + + assert key_logging_entries({"logging": [entry]}) == (entry,) diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 650bc89c917..2e18a0c1099 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -51,6 +51,7 @@ from litellm.proxy.auth.auth_checks import ( ) from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import AuthCacheInvalidationSubscriber +from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, budget_windows_absent_cache_key, @@ -101,6 +102,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( persist_deleted_verification_tokens, prepare_key_update_data, prepare_metadata_fields, + raise_on_logging_credential_change, reset_key_spend_fn, update_key_fn, validate_key_list_check, @@ -112,6 +114,7 @@ from litellm.types.proxy.management_endpoints.key_management_endpoints import ( BulkUpdateKeyResponse, CustomKeyPolicyRequest, ) +from litellm.types.utils import CredentialItem from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters client = TestClient(app) @@ -24407,3 +24410,29 @@ async def test_key_aliases_answers_503_when_the_database_is_unreachable(monkeypa ) assert (exc_info.value.code, exc_info.value.type) == ("503", "no_db_connection") + + +def test_non_admin_resaving_referencing_key_metadata_read_back_encrypted_is_not_a_logging_change(monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", "test-salt-32-bytes-aaaaaaaaaaaaaa") + monkeypatch.setattr( + litellm, + "credential_list", + [CredentialItem(credential_name="space-a", credential_values={"arize_api_key": "k"}, credential_info={})], + ) + stored = encrypt_callback_vars( + { + "logging": [ + {"callback_name": "arize", "litellm_credential_name": "space-a", "callback_vars": {}}, + {"callback_name": "langfuse", "callback_vars": {"langfuse_secret_key": "sk-lf-secret"}}, + ] + } + ) + assert stored["logging"][1]["callback_vars"]["langfuse_secret_key"] != "sk-lf-secret" + key_row = LiteLLM_VerificationToken(token="hashed", metadata=stored) + internal_user = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER) + + raise_on_logging_credential_change({**stored, "note": "unrelated edit"}, key_row, internal_user) + with pytest.raises(HTTPException) as removed: + raise_on_logging_credential_change({"note": "unrelated edit"}, key_row, internal_user) + + assert removed.value.status_code == 403 diff --git a/tests/unit/proxy/management_helpers/test_bulk_user_creation.py b/tests/unit/proxy/management_helpers/test_bulk_user_creation.py index 2b966a8ebeb..96c84578111 100644 --- a/tests/unit/proxy/management_helpers/test_bulk_user_creation.py +++ b/tests/unit/proxy/management_helpers/test_bulk_user_creation.py @@ -7,6 +7,7 @@ import pytest from prisma.errors import UniqueViolationError from pydantic import BaseModel, ConfigDict, ValidationError +import litellm from litellm.caching.caching import DualCache from litellm.proxy._types import LiteLLM_TeamTable, LitellmUserRoles, Member, UserAPIKeyAuth from litellm.proxy.list_api.common import ManagementProblem @@ -15,6 +16,7 @@ from litellm.types.proxy.management_endpoints.internal_user_endpoints import ( BulkNewUserItem, BulkNewUserRequest, ) +from litellm.types.utils import CredentialItem ADMIN: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) INTERNAL: Final = UserAPIKeyAuth(user_id="someone", user_role=LitellmUserRoles.INTERNAL_USER) @@ -434,6 +436,39 @@ async def test_only_proxy_admin_sets_batch_limits_on_a_created_key(): assert calls[-1]["metadata"] == limits +@pytest.mark.asyncio +async def test_only_proxy_admin_creates_a_key_whose_logging_references_a_credential( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + litellm, + "credential_list", + [CredentialItem(credential_name="space-a", credential_values={"arize_api_key": "k"}, credential_info={})], + ) + reference: Final = {"logging": [{"callback_name": "arize", "litellm_credential_name": "space-a", "callback_vars": {}}]} + calls: list[dict[str, object]] = [] + + async def generate_key(**kwargs: object) -> dict[str, object]: + calls.append(kwargs) + return {"token": f"sk-{kwargs['user_id']}"} + + rows: Final = [ + {"user_id": "u1", "auto_create_key": True, "metadata": reference}, + {"user_id": "u2", "auto_create_key": True}, + ] + + response = await _run(_FakePrisma(), rows, caller=INTERNAL, generate_key=generate_key) + + assert [r.success for r in response.data] == [False, True] + assert "litellm_credential_name" in (response.data[0].error or "") + assert [call["user_id"] for call in calls] == ["u2"] + + admin_response = await _run(_FakePrisma(), rows[:1], caller=ADMIN, generate_key=generate_key) + + assert [r.key for r in admin_response.data] == ["sk-u1"] + assert calls[-1]["metadata"] == reference + + @pytest.mark.asyncio async def test_license_is_checked_once_against_the_whole_batch(): prisma = _FakePrisma() diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index fbd4de8ef5c..e4a560adca8 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -26,6 +26,7 @@ from litellm.constants import ( SESSION_ID_GENERATED_METADATA_KEY, SESSION_ID_OMITTED_METADATA_KEY, ) +from litellm.integrations.otel.model.config import is_otel_v2_enabled from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs from litellm.litellm_core_utils.get_provider_specific_headers import ( ProviderSpecificHeaderUtils, @@ -59,6 +60,7 @@ from litellm.proxy.litellm_pre_call_utils import ( get_dynamic_logging_metadata, get_metadata_variable_name, move_guardrails_to_metadata, + resolve_tenant_otel_destinations, ) from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload from litellm.types.litellm_params import AGENTIC_LOOP_KWARG_NAMES @@ -2779,7 +2781,7 @@ def test_key_dynamic_logging_settings(): team_metadata={}, ) result = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(key_with_arize) - assert result == [{"callback_name": "arize", "callback_type": "success"}] + assert result == ({"callback_name": "arize", "callback_type": "success"},) # Test with langfuse logging key_with_langfuse = UserAPIKeyAuth( @@ -2788,7 +2790,7 @@ def test_key_dynamic_logging_settings(): team_metadata={}, ) result = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(key_with_langfuse) - assert result == [{"callback_name": "langfuse", "callback_type": "success"}] + assert result == ({"callback_name": "langfuse", "callback_type": "success"},) # Test with no logging metadata key_without_logging = UserAPIKeyAuth(api_key="test-key", metadata={}, team_metadata={}) @@ -2912,6 +2914,22 @@ def test_get_dynamic_logging_metadata_with_arize_team_logging(): assert result.callback_vars["arize_space_id"] == "test_arize_space_id" +@pytest.mark.parametrize( + "logging_value", + [ + pytest.param({"callback_name": "langsmith"}, id="object"), + pytest.param("langsmith", id="string"), + pytest.param(["langsmith"], id="list-with-string"), + pytest.param([None], id="list-with-null"), + ], +) +def test_non_object_team_logging_metadata_fails_as_a_server_error_not_a_bad_request(logging_value: object) -> None: + user_api_key_dict: Final = UserAPIKeyAuth(api_key="test-key", metadata={}, team_metadata={"logging": logging_value}) + + with pytest.raises(TypeError): + get_dynamic_logging_metadata(user_api_key_dict=user_api_key_dict, proxy_config=MagicMock()) + + def test_add_team_callback_rejects_env_reference(): with pytest.raises(PydanticValidationError) as exc_info: AddTeamCallback( @@ -10267,6 +10285,48 @@ def test_arize_otlp_protocol_on_a_key_logging_entry_reaches_the_destination(monk is_otel_v2_enabled.cache_clear() +def test_a_key_referencing_a_credential_routes_to_that_credential_with_its_own_excluded_services(monkeypatch): + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + monkeypatch.setenv("ARIZE_HTTP_ENDPOINT", "https://arize.internal.example/v1/traces") + monkeypatch.delenv("ARIZE_API_KEY", raising=False) + monkeypatch.delenv("ARIZE_SPACE_ID", raising=False) + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="space-a", + credential_values={"arize_api_key": "k", "arize_space_id": "s"}, + credential_info={}, + ) + ], + ) + is_otel_v2_enabled.cache_clear() + try: + auth = UserAPIKeyAuth( + api_key="hashed-key", + metadata={ + "logging": [ + { + "callback_name": "arize", + "callback_type": "success", + "litellm_credential_name": "space-a", + "callback_vars": {"excluded_services": "redis,postgres"}, + } + ] + }, + team_metadata={}, + ) + + destinations = resolve_tenant_otel_destinations(auth) + + assert [(d.headers, d.excluded_db_systems) for d in destinations] == [ + ({"arize-space-id": "s", "api_key": "k"}, frozenset({"redis", "postgresql"})) + ] + finally: + is_otel_v2_enabled.cache_clear() + + class TestResolveUserProviderCredentials: """``_resolve_user_provider_credentials_for_request`` injects the calling user's stored provider tokens into secret_fields for deployments carrying a per_user_oauth credential.""" From c6794ac08b0e14d20760b8f6dd00eae79ed3ebb6 Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Mon, 9 Nov 2026 16:06:52 -0800 Subject: [PATCH 31/46] fix(messages): let proxy /v1/messages requests run on the rust path (#50335) The proxy attaches its shared aiohttp session to every request and a default stream_timeout to every stream, and the messages detectors treated both as Python-only features, so no proxy request reached Rust. A shared session is only a connection-reuse hint and no native client reads it, so the common caller_client detector now fires only for an explicit client. The native messages call now takes the one timeout Python's messages handler resolves from timeout, request_timeout, stream_timeout and an httpx.Timeout, so no timeout shape keeps a messages call on Python --- litellm/messages/dispatch.py | 20 +++++-- litellm/rust_bridge/catalog.py | 4 +- litellm/rust_bridge/detectors.py | 4 +- .../messages/test_dispatch.py | 58 +++++++++++++++++-- .../responses/test_dispatch.py | 1 - tests/unit/rust_bridge/test_detectors.py | 2 +- 6 files changed, 75 insertions(+), 14 deletions(-) diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index 4c94f26c6c1..40d21c5bad9 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -6,6 +6,7 @@ from typing import Final, TypeAlias, cast # noqa: TID251 # native binding sele from litellm.exceptions import BadRequestError from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.llms.anthropic.messages.native import handler as main +from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.rust_bridge import detectors from litellm.rust_bridge.catalog import Feature, Route, RouteContext, Target from litellm.rust_bridge.dispatch import PublicDispatch @@ -25,6 +26,7 @@ from litellm.rust_bridge.public_call import ( signature, ) from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse +from litellm.types.router import GenericLiteLLMParams __all__ = ("anthropic_messages", "anthropic_messages_handler") @@ -55,10 +57,20 @@ _AMESSAGES: Final = signature(_PYTHON_AMESSAGES) def _native_kwargs(fields: Mapping[str, object], kwargs: Mapping[str, object]) -> Mapping[str, object]: - stream_timeout: Final = kwargs.get("stream_timeout") - if fields.get("stream") is not True or stream_timeout is None: + """The native call takes one timeout, resolved from timeout, request_timeout and stream_timeout as Python's.""" + _, provider = _resolved_model(str(fields["model"]), optional_str(fields.get("custom_llm_provider"))) + timeout: Final = BaseLLMHTTPHandler._resolve_anthropic_messages_timeout( # pyright: ignore[reportPrivateUsage] # the handler's own timeout resolution + litellm_params=GenericLiteLLMParams( + timeout=kwargs.get("timeout"), + request_timeout=kwargs.get("request_timeout"), + stream_timeout=kwargs.get("stream_timeout"), + ), + stream=fields.get("stream") is True, + custom_llm_provider=provider or "", + ) + if timeout is None: return kwargs - return {**kwargs, "timeout": stream_timeout} + return {**kwargs, "timeout": timeout} def _public_request( @@ -122,7 +134,7 @@ def messages_target(model: str, custom_llm_provider: str | None, model_info: obj _FEATURES: Final[detectors.Detectors] = MappingProxyType( { - **detectors.COMMON, + **{feature: detect for feature, detect in detectors.COMMON.items() if feature is not Feature.TIMEOUT_VARIANTS}, Feature.PROVIDER_AFFINITY: route_detectors.provider_affinity, Feature.SAFEGUARDS_CLASSIFIER: route_detectors.safeguards_classifier, Feature.RESPONSE_CACHE: detectors.response_cache, diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 78040ba7225..c493617003d 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -39,7 +39,7 @@ class Feature(str, Enum): PROXY_MCP_TOOLS = "proxy_mcp_tools" """litellm_proxy MCP tools that Python expands and may auto-execute.""" CALLER_CLIENT = "caller_client" - """A caller-supplied HTTP client or shared session that Python sends through.""" + """A caller-supplied HTTP client that Python sends through.""" TIMEOUT_VARIANTS = "timeout_variants" """request_timeout, stream_timeout or httpx.Timeout phases beyond a plain timeout.""" DROP_PARAMS = "drop_params" @@ -172,7 +172,7 @@ CHAT_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | { Feature.MODEL_ALIAS_MAP, } -MESSAGES_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | { +MESSAGES_PYTHON_FEATURES: Final = (COMMON_PYTHON_FEATURES - {Feature.TIMEOUT_VARIANTS}) | { Feature.PROVIDER_AFFINITY, Feature.FOUNDRY_BETA_ROW, Feature.SAFEGUARDS_CLASSIFIER, diff --git a/litellm/rust_bridge/detectors.py b/litellm/rust_bridge/detectors.py index f2591bdbcd1..c38c3ea7135 100644 --- a/litellm/rust_bridge/detectors.py +++ b/litellm/rust_bridge/detectors.py @@ -44,8 +44,8 @@ def latest_alias(request: NativeCall) -> bool: def caller_client(request: NativeCall) -> bool: - """Python sends through the caller's HTTP client or shared session.""" - return present(request, "client") or present(request, "shared_session") + """Python sends through the caller's HTTP client; a shared session only reuses connections, which native pools.""" + return present(request, "client") def extra_body(request: NativeCall) -> bool: diff --git a/tests/test_litellm_rust/messages/test_dispatch.py b/tests/test_litellm_rust/messages/test_dispatch.py index 6d65f289524..b580eca2e42 100644 --- a/tests/test_litellm_rust/messages/test_dispatch.py +++ b/tests/test_litellm_rust/messages/test_dispatch.py @@ -2,10 +2,12 @@ import inspect from collections.abc import Awaitable, Callable, Mapping from typing import Final, cast # noqa: TID251 # narrows legacy callable signatures for inspect +import httpx import pytest import litellm from litellm.caching.caching import Cache +from litellm.constants import COMPLETION_HTTP_FALLBACK_SECONDS from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger from litellm.llms.anthropic.messages.native import handler as python_messages @@ -172,8 +174,25 @@ def test_native_receives_normalized_request_and_original_call_shape() -> None: assert call_kwargs["litellm_metadata"] is metadata -@pytest.mark.parametrize("stream, expected_timeout", ((True, 0.5), (False, 30.0))) -def test_native_streams_are_bounded_by_stream_timeout(stream: bool, expected_timeout: float) -> None: +@pytest.mark.parametrize( + ("kwargs", "expected_timeout"), + ( + pytest.param({"stream": True, "timeout": 30.0, "stream_timeout": 0.5}, 0.5, id="stream-timeout-bounds-stream"), + pytest.param( + {"stream": False, "timeout": 30.0, "stream_timeout": 0.5}, 30.0, id="stream-timeout-ignored-unary" + ), + pytest.param({"request_timeout": 5}, 5.0, id="request-timeout"), + pytest.param({"timeout": 30.0, "request_timeout": 5}, 30.0, id="timeout-before-request-timeout"), + pytest.param( + {"stream": True, "request_timeout": 5, "stream_timeout": 0.5}, 0.5, id="stream-timeout-over-request" + ), + pytest.param({"timeout": httpx.Timeout(60.0, connect=2.0)}, 60.0, id="httpx-timeout-read"), + pytest.param( + {"timeout": httpx.Timeout(None, connect=2.0)}, COMPLETION_HTTP_FALLBACK_SECONDS, id="httpx-timeout-no-read" + ), + ), +) +def test_native_timeout_resolves_like_python(kwargs: Mapping[str, object], expected_timeout: float) -> None: captured: Final[list[NativeCall]] = [] def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: rejected fallback @@ -185,7 +204,7 @@ def test_native_streams_are_bounded_by_stream_timeout(stream: bool, expected_tim _ = _DISPATCH.run( (16, MESSAGES, "anthropic/claude-sonnet-4-5"), - {"stream": stream, "timeout": 30.0, "stream_timeout": 0.5}, + kwargs, python=python, binding=messages_binding(native), native=lambda hook, request, call_args, call_kwargs: hook(request), @@ -456,7 +475,6 @@ STRICT_TOOL: Final = {"name": "a", "input_schema": {"type": "object"}, "strict": {}, ), (Feature.CALLER_CLIENT, "anthropic/claude-sonnet-4-5", {"client": object()}, {}), - (Feature.TIMEOUT_VARIANTS, "anthropic/claude-sonnet-4-5", {"request_timeout": 5}, {}), (Feature.AGENTIC_LOOP_HOOK, "anthropic/claude-sonnet-4-5", {}, {"callbacks": [_Agentic()]}), (Feature.RESPONSE_CACHE, "anthropic/claude-sonnet-4-5", {}, {"cache": Cache()}), (Feature.TLS_POLICY, "anthropic/claude-sonnet-4-5", {}, {"ssl_security_level": "DEFAULT@SECLEVEL=1"}), @@ -528,6 +546,38 @@ def test_each_detected_feature_reaches_its_rule( assert served(kwargs) is python_result +@pytest.mark.parametrize( + "kwargs", + ( + pytest.param({"shared_session": object()}, id="proxy-shared-session"), + pytest.param({"stream": True, "timeout": 600.0, "stream_timeout": 600.0}, id="proxy-stream-timeout"), + pytest.param({"request_timeout": 5}, id="request-timeout"), + pytest.param({"timeout": httpx.Timeout(60.0, connect=2.0)}, id="httpx-timeout"), + ), +) +def test_proxy_request_shape_stays_native(kwargs: Mapping[str, object]) -> None: + rules: Final[Rules] = ( + RouteRule( + Route.MESSAGES, Rollout.PYTHON_ONLY, features=frozenset({Feature.CALLER_CLIENT, Feature.TIMEOUT_VARIANTS}) + ), + RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED), + ) + native_result: Final = response() + + def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: rejected fallback + pytest.fail("A proxy-shaped request must not fall back to Python") + + served: Final = _DISPATCH.run( + (16, MESSAGES, "anthropic/claude-sonnet-4-5"), + kwargs, + python=python, + binding=messages_binding(lambda request: native_result), + native=lambda hook, request, call_args, native_kwargs: hook(request), + rules=rules, + ) + assert served is native_result + + @pytest.mark.asyncio @pytest.mark.parametrize( "kwargs", diff --git a/tests/test_litellm_rust/responses/test_dispatch.py b/tests/test_litellm_rust/responses/test_dispatch.py index 4498dbdf39d..19255379dcf 100644 --- a/tests/test_litellm_rust/responses/test_dispatch.py +++ b/tests/test_litellm_rust/responses/test_dispatch.py @@ -260,7 +260,6 @@ ENCODED_ID: Final = ResponsesAPIRequestUtils._build_encrypted_item_id("deploymen (Feature.MOCK_RESPONSE, {"mock_response": "hi"}, {}), (Feature.PROXY_MCP_TOOLS, {"tools": [{"type": "mcp", "server_label": "g", "server_url": "litellm_proxy"}]}, {}), (Feature.CALLER_CLIENT, {"client": object()}, {}), - pytest.param(Feature.CALLER_CLIENT, {"shared_session": object()}, {}, id="shared-session"), pytest.param(Feature.EXTRA_BODY, {"extra_body": {"stream": True}}, {}, id="extra-body-stream"), pytest.param( Feature.PROVIDER_AFFINITY, diff --git a/tests/unit/rust_bridge/test_detectors.py b/tests/unit/rust_bridge/test_detectors.py index a6e6b2c0dca..fccc01dd8e6 100644 --- a/tests/unit/rust_bridge/test_detectors.py +++ b/tests/unit/rust_bridge/test_detectors.py @@ -111,7 +111,7 @@ def _wrap_nested_level(value: object, depth: int) -> object: (detectors.prompt_management, {"model": "litellm_agent/gpt-5"}, True), (detectors.prompt_management, {"model": "gpt-5"}, False), (detectors.caller_client, {"client": object()}, True), - (detectors.caller_client, {"shared_session": object()}, True), + (detectors.caller_client, {"shared_session": object()}, False), (detectors.caller_client, {"client": None, "shared_session": None}, False), (detectors.caller_client, {}, False), (detectors.extra_body, {"extra_body": {"stream": True}}, True), From 1ddc91300259360e635e76d30262fadc26e48af2 Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Mon, 9 Nov 2026 16:06:53 -0800 Subject: [PATCH 32/46] fix(messages): log native messages success once, after the proxy releases it (#50420) * fix(messages): log native messages success once, after the proxy releases it The proxy defers non-stream success logging past post-call guardrails, but the native host parked its log in its own slot and the release only ran for OCR, so a unary /v1/messages call served by Rust logged no success event, no SpendLogs row and no key, team, user or end-user spend. A deferred native success now parks in enqueue_deferred_logging with deferred_success_result, the same slots Python uses, so the proxy releases it like any other call and its guardrail step pruning and cost preparation see it too Python runs sync success callbacks inside the deferred logging closure, so a post-call hook that blocks the response drops them. A parked native success now carries its sync callbacks and runs them on release A refused first attempt and its fallback share one logger. Python logs the async success once, for the first attempt, so the native adapter now claims the async success log once across fallback attempts. Like Python, an attempt that loses the claim runs its sync callbacks only when logging is not deferred, and an internal call logs nothing * test(messages): cover the deferred native success park and release --- .../python_contract.json | 11 +- .../callbacks-legacy-python/src/adapter.rs | 131 +++++++++++++++--- .../callbacks-legacy-python/src/callbacks.rs | 29 +++- .../callbacks-legacy-python/src/deferred.rs | 59 ++++---- .../callbacks-legacy-python/src/python.rs | 4 + .../src/test_support.rs | 11 +- litellm/proxy/common_request_processing.py | 11 -- .../host/callbacks_legacy_python.py | 28 +++- tests/test_litellm_rust/ocr/test_lifecycle.py | 39 +++++- .../test_deferred_guardrail_logging.py | 32 ----- .../host/test_callbacks_legacy_python.py | 16 +++ 11 files changed, 269 insertions(+), 102 deletions(-) diff --git a/litellm-rust/crates/callbacks-legacy-python/python_contract.json b/litellm-rust/crates/callbacks-legacy-python/python_contract.json index ab887bbc014..1d4a2f7f7c2 100644 --- a/litellm-rust/crates/callbacks-legacy-python/python_contract.json +++ b/litellm-rust/crates/callbacks-legacy-python/python_contract.json @@ -47,9 +47,18 @@ "defers_async_logging": [ "logger" ], + "claim_async_success_log": [ + "logger" + ], + "mark_async_success_enqueued": [ + "logger" + ], "defer_success": [ "logger", - "pending" + "pending", + "response", + "start", + "end" ], "sync_success_for_async_call": [ "logger", diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index 1bb84184e4c..8d1946861aa 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -251,26 +251,40 @@ impl LegacyLogging { if !self.asynchronous { return pending().sync(py); } - if !self.internal - && self + if self.internal { + return Ok(()); + } + let defers = logger.defers_async_logging(py); + if logger.claim_async_success_log(py)? { + let enqueues = self .call .kwargs() .bind(py) .get_item("fallbacks")? - .is_none_or(|value| value.is_none()) - { - if logger.defers_async_logging(py) { + .is_none_or(|value| value.is_none()); + if defers { let pending = Py::new( py, PendingLogging { pending: Some(pending()), + enqueues, }, )?; - logger.defer_success(py, pending.bind(py).as_any())?; - } else { + return logger.defer_success( + py, + pending.bind(py).as_any(), + &self.response, + &self.start, + &self.end, + ); + } + if enqueues { pending().asynchronous(py)?; } } + if defers { + return Ok(()); + } logger.sync_success_for_async_call(py, &self.response, &self.start, &self.end) } @@ -1873,10 +1887,15 @@ mod terminal_tests { #[case::async_listened( true, c"", - &["async_success_handler", "enqueued", "sync_success_for_async_call"] + &["async_success_handler", "mark_async_success_enqueued", "enqueued", "sync_success_for_async_call"] )] - #[case::async_deferred(true, c"logger._defer_async_logging = True", &["sync_success_for_async_call"])] + #[case::async_deferred(true, c"logger._defer_async_logging = True", &[])] #[case::async_with_fallbacks(true, c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])] + #[case::async_deferred_with_fallbacks( + true, + c"logger._defer_async_logging = True\nkwargs = {'fallbacks': ['other']}", + &[] + )] fn success_reaches_the_logging_handlers( #[case] asynchronous: bool, #[case] script: &CStr, @@ -1902,12 +1921,88 @@ mod terminal_tests { &locals, c" assert all(value is response for name, value in logger.calls if name.endswith('_handler')) -assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_async_logging', False) +assert hasattr(logger, 'enqueue_deferred_logging') == getattr(logger, '_defer_async_logging', False) +assert getattr(logger, 'deferred_success_result', (response,))[0] is response ", ); }); } + #[rstest] + #[case::accepted(c"", &["sync_success_for_async_call", "async_success_handler", "mark_async_success_enqueued", "enqueued"])] + #[case::accepted_with_fallbacks(c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])] + fn a_deferred_success_runs_sync_callbacks_only_once_released( + #[case] script: &CStr, + #[case] expected: &[&str], + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c"response = object()\nlogger._defer_async_logging = True", + ); + run(py, &locals, script); + let mut logging = logged(py, &locals, true); + succeed(py, &locals, &mut logging); + run( + py, + &locals, + c" +assert logger.names() == [], logger.calls +logger.enqueue_deferred_logging() +logger.enqueue_deferred_logging() +", + ); + let names: Vec = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert_eq!(names, expected); + }); + } + + #[rstest] + #[case::deferred( + c"logger._defer_async_logging = True", + c"logger.enqueue_deferred_logging()", + &["sync_success_for_async_call", "async_success_handler", "mark_async_success_enqueued", "enqueued"] + )] + #[case::deferred_and_blocked(c"logger._defer_async_logging = True", c"", &[])] + #[case::immediate( + c"", + c"", + &["async_success_handler", "mark_async_success_enqueued", "enqueued", "sync_success_for_async_call", "sync_success_for_async_call"] + )] + fn a_fallback_attempt_on_the_same_logger_leaves_the_first_success_claimed( + #[case] script: &CStr, + #[case] release: &CStr, + #[case] expected: &[&str], + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, c"refused = object()\nresponse = refused"); + run(py, &locals, script); + let mut first = logged(py, &locals, true); + succeed(py, &locals, &mut first); + run(py, &locals, c"response = object()"); + let mut fallback = logged(py, &locals, true); + succeed(py, &locals, &mut fallback); + run(py, &locals, release); + let names: Vec = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert_eq!(names, expected); + run( + py, + &locals, + c"assert all(value is refused for name, value in logger.calls if name == 'async_success_handler')", + ); + }); + } + #[rstest] fn dropped_observations_preserve_deferred_success_and_response_identity() { Python::initialize(); @@ -1933,10 +2028,10 @@ assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_asy )); assert_eq!(sender.dropped_events(), 1); run(py, &locals, c" -assert logger.names() == ['sync_success_for_async_call'], logger.calls -logger._native_pending_logging.release(True) -logger._native_pending_logging.release(True) -assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'enqueued'], logger.calls +assert logger.names() == [], logger.calls +logger.enqueue_deferred_logging() +logger.enqueue_deferred_logging() +assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'mark_async_success_enqueued', 'enqueued'], logger.calls assert logger.calls[0][1] is response assert logger.calls[1][1] is response "); @@ -2009,7 +2104,7 @@ assert chunks[1] is last } #[rstest::rstest] - fn internal_async_calls_skip_the_async_success_fan_out() { + fn internal_async_calls_skip_success_logging() { Python::initialize(); Python::attach(|py| { let locals = namespace(py, c"response = object()"); @@ -2018,11 +2113,7 @@ assert chunks[1] is last ..logged(py, &locals, true) }; succeed(py, &locals, &mut logging); - run( - py, - &locals, - c"assert logger.names() == ['sync_success_for_async_call'], logger.calls", - ); + run(py, &locals, c"assert logger.calls == [], logger.calls"); }); } diff --git a/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs b/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs index 1951c9ebfb4..a428afa9ece 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs @@ -59,7 +59,16 @@ pub trait LegacyCallbacks { fn defers_async_logging(&self, py: Python<'_>) -> bool; - fn defer_success(&self, py: Python<'_>, pending: &Bound<'_, PyAny>) -> PyResult<()>; + fn claim_async_success_log(&self, py: Python<'_>) -> PyResult; + + fn defer_success( + &self, + py: Python<'_>, + pending: &Bound<'_, PyAny>, + response: &Option>, + start: &Py, + end: &Option>, + ) -> PyResult<()>; fn sync_success_for_async_call( &self, @@ -203,8 +212,21 @@ impl LegacyCallbacks for PythonLogger { .unwrap_or(false) } - fn defer_success(&self, py: Python<'_>, pending: &Bound<'_, PyAny>) -> PyResult<()> { - Logging::DeferSuccess.call(py, (self.object(py), pending))?; + fn claim_async_success_log(&self, py: Python<'_>) -> PyResult { + Logging::ClaimAsyncSuccess + .call(py, (self.object(py),))? + .extract() + } + + fn defer_success( + &self, + py: Python<'_>, + pending: &Bound<'_, PyAny>, + response: &Option>, + start: &Py, + end: &Option>, + ) -> PyResult<()> { + Logging::DeferSuccess.call(py, (self.object(py), pending, response, start, end))?; Ok(()) } @@ -252,6 +274,7 @@ impl LegacyCallbacks for PythonLogger { ) -> PyResult<()> { let coroutine = Logging::AsyncSuccessHandler.call(py, (self.object(py), response, start, end))?; + Logging::MarkAsyncSuccessEnqueued.call(py, (self.object(py),))?; let enqueue = Logging::Enqueue.call(py, (&coroutine,)); if enqueue.is_err() && let Err(error) = coroutine.call_method0("close") diff --git a/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs b/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs index 292648ba69c..ea72191ddb2 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs @@ -22,21 +22,35 @@ impl PendingSuccess { self.logger .enqueue_success(py, &self.response, &self.start, &self.end) } + + pub(crate) fn sync_for_async_call(&self, py: Python<'_>) -> PyResult<()> { + self.logger + .sync_success_for_async_call(py, &self.response, &self.start, &self.end) + } } #[pyclass] pub(crate) struct PendingLogging { pub(crate) pending: Option, + pub(crate) enqueues: bool, } #[pymethods] impl PendingLogging { - fn release(slf: &Bound<'_, Self>, py: Python<'_>, success: bool) -> PyResult<()> { - let pending = slf.borrow_mut().pending.take(); - if let Some(pending) = pending - && success - { - match pending.asynchronous(py) { + fn __call__(slf: &Bound<'_, Self>, py: Python<'_>) -> PyResult<()> { + let (pending, enqueues) = { + let mut slot = slf.borrow_mut(); + (slot.pending.take(), slot.enqueues) + }; + if let Some(pending) = pending { + let released = pending.sync_for_async_call(py).and_then(|()| { + if enqueues { + pending.asynchronous(py) + } else { + Ok(()) + } + }); + match released { Err(error) if error.is_instance_of::(py) => { error.write_unraisable(py, Some(pending.logger.object(py))); } @@ -87,6 +101,7 @@ mod tests { start: py.None(), end: Some(py.None()), }), + enqueues: true, }, ) .unwrap(); @@ -108,7 +123,7 @@ observed = [] def on_enqueue(coroutine): observed.append(marker.get()) - pending.release(True) + pending() logger.on_enqueue = on_enqueue ", @@ -118,28 +133,12 @@ logger.on_enqueue = on_enqueue &locals, c" marker.set('release') -pending.release(True) -pending.release(True) +pending() +pending() assert observed == ['release'], observed -assert logger.names() == ['async_success_handler', 'enqueued'], logger.calls +assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'mark_async_success_enqueued', 'enqueued'], logger.calls assert logger.calls[0][1] is response -", - ); - }); - } - - #[test] - fn a_blocked_release_drops_the_success_for_good() { - Python::initialize(); - Python::attach(|py| { - let locals = defer(py, c""); - run( - py, - &locals, - c" -pending.release(False) -pending.release(True) -assert logger.calls == [], logger.calls +assert logger.calls[1][1] is response ", ); }); @@ -168,7 +167,7 @@ logger.on_enqueue = on_enqueue locals .set_item("failure", py.eval(failure, None, Some(&locals)).unwrap()) .unwrap(); - let released = local(&locals, "pending").call_method1("release", (true,)); + let released = local(&locals, "pending").call0(); match released { Ok(_) => assert!(!propagates), Err(error) => { @@ -181,8 +180,8 @@ logger.on_enqueue = on_enqueue py, &locals, c" -pending.release(True) -assert logger.names() == ['async_success_handler', 'enqueued', 'closed'], logger.calls +pending() +assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'mark_async_success_enqueued', 'enqueued', 'closed'], logger.calls assert unraisable_from(logger) == ([] if propagates else [failure]) ", ); diff --git a/litellm-rust/crates/callbacks-legacy-python/src/python.rs b/litellm-rust/crates/callbacks-legacy-python/src/python.rs index 6d910d35f53..32b11988703 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/python.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/python.rs @@ -50,6 +50,8 @@ pub(crate) enum Logging { RecordApiWindow, #[strum(serialize = "defers_async_logging")] DefersAsync, + #[strum(serialize = "claim_async_success_log")] + ClaimAsyncSuccess, #[strum(serialize = "defer_success")] DeferSuccess, #[strum(serialize = "sync_success_for_async_call")] @@ -58,6 +60,8 @@ pub(crate) enum Logging { SubmitSuccess, #[strum(serialize = "async_success_handler")] AsyncSuccessHandler, + #[strum(serialize = "mark_async_success_enqueued")] + MarkAsyncSuccessEnqueued, #[strum(serialize = "enqueue_logging")] Enqueue, #[strum(serialize = "failure_handler")] diff --git a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs index 8acfa7a898b..9e69d0c30e2 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs @@ -72,7 +72,11 @@ FAKES = { llm_api_duration_ms=(end - start) * 1000 ), 'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)), - 'defer_success': lambda logger, pending: setattr(logger, '_native_pending_logging', pending), + 'claim_async_success_log': lambda logger: logger.claim_async_success_log(), + 'defer_success': lambda logger, pending, response, start, end: ( + setattr(logger, 'deferred_success_result', (response, start, end)), + setattr(logger, 'enqueue_deferred_logging', pending), + ), 'sync_success_for_async_call': lambda logger, response, start, end: logger.handle_sync_success_callbacks_for_async_calls( response, start, end ), @@ -81,6 +85,7 @@ FAKES = { )(error, ''.join(traceback.format_exception(error)), start, end), 'submit_success': lambda logger, response, start, end: logger.record('submit', (response, start, end)), 'async_success_handler': lambda logger, response, start, end: logger.async_success_handler(response, start, end), + 'mark_async_success_enqueued': lambda logger: logger.record('mark_async_success_enqueued', None), 'enqueue_logging': lambda coroutine: coroutine.enqueue(), 'restore_context': lambda logger: logger.record('restore', None), 'custom_pricing_fields': lambda: ('ocr_cost_per_page',), @@ -156,6 +161,10 @@ class StubLogger: self.record('async_success_handler', response) return StubCoroutine(self) + def claim_async_success_log(self): + claimed, self.claimed = getattr(self, 'claimed', False), True + return not claimed + def handle_sync_success_callbacks_for_async_calls(self, response, start, end): self.record('sync_success_for_async_call', response) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 8f73d47d96b..eb356491014 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -64,7 +64,6 @@ from litellm.litellm_core_utils.bug_report import ( strip_bug_report_notice, ) from litellm.litellm_core_utils.core_helpers import ( - attribute_or_none, finite_cost, get_or_create_metadata_bucket, has_attributes, @@ -72,7 +71,6 @@ from litellm.litellm_core_utils.core_helpers import ( is_expected_client_error, proxy_request_headers, redact_nested_match_and_regex_keys, - require_attributes, stream_cost_injection_requested, ) from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer @@ -724,10 +722,6 @@ class _ReadsHeaderValues(Protocol): def get(self, key: str, default: str = "") -> str: ... -class _NativeLoggingRelease(Protocol): - def release(self, success: bool, /) -> object: ... - - def _as_header_reader(headers: _ReadsHeaderValues) -> _ReadsHeaderValues: return headers @@ -4384,11 +4378,6 @@ class ProxyBaseLLMRequestProcessing: Extracted as a static method so tests can exercise the production gating logic directly rather than reimplementing the finally block. """ - if getattr(logging_obj, "call_type", None) in ("ocr", "aocr"): - pending: Final = attribute_or_none(logging_obj, "_native_pending_logging") - if pending is not None: - logging_obj._native_pending_logging = None # rebind-ok: consume the native OCR release signal once - require_attributes(pending, _NativeLoggingRelease, "release").release(not exception_raised) _enqueue_fn: Final = getattr(logging_obj, "enqueue_deferred_logging", None) if _enqueue_fn is None: return diff --git a/litellm/rust_bridge/host/callbacks_legacy_python.py b/litellm/rust_bridge/host/callbacks_legacy_python.py index 92590531cdc..92e5436fbb8 100644 --- a/litellm/rust_bridge/host/callbacks_legacy_python.py +++ b/litellm/rust_bridge/host/callbacks_legacy_python.py @@ -12,7 +12,7 @@ import datetime import traceback import uuid import weakref -from collections.abc import Awaitable, Coroutine, Mapping, Sequence +from collections.abc import Awaitable, Callable, Coroutine, Mapping, Sequence from dataclasses import dataclass from typing import ( TYPE_CHECKING, @@ -111,6 +111,9 @@ def finalize( class LoggingSurface(Protocol): + deferred_success_result: tuple[object, datetime.datetime, datetime.datetime] | None + enqueue_deferred_logging: Callable[[], None] | None + @property def model_call_details(self) -> Mapping[str, object]: ... @@ -139,6 +142,10 @@ class LoggingSurface(Protocol): additional_args: dict[str, object] = ..., ) -> object: ... + def claim_async_success_log(self) -> bool: ... + + def mark_async_success_log_enqueued(self) -> None: ... + def handle_sync_success_callbacks_for_async_calls( self, result: object, start_time: datetime.datetime, end_time: datetime.datetime, cache_hit: bool | None = None ) -> None: ... @@ -269,8 +276,23 @@ def defers_async_logging(logger: LoggingSurface) -> bool: return bool(getattr(logger, "defer_async_logging", False)) -def defer_success(logger: LoggingSurface, pending: object) -> None: - setattr(logger, "_native_pending_logging", pending) +def claim_async_success_log(logger: LoggingSurface) -> bool: + return logger.claim_async_success_log() + + +def mark_async_success_enqueued(logger: LoggingSurface) -> None: + logger.mark_async_success_log_enqueued() + + +def defer_success( + logger: LoggingSurface, + pending: Callable[[], None], + response: object, + start: datetime.datetime, + end: datetime.datetime, +) -> None: + logger.deferred_success_result = (response, start, end) # rebind-ok: capture the native attempt before dispatch + logger.enqueue_deferred_logging = pending # rebind-ok: the proxy releases logging here def _cache_hit(logger: LoggingSurface) -> Literal[True] | None: diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py index 611316fb52c..c3846d894bb 100644 --- a/tests/test_litellm_rust/ocr/test_lifecycle.py +++ b/tests/test_litellm_rust/ocr/test_lifecycle.py @@ -211,8 +211,12 @@ async def test_deferred_logging_requires_release_and_runs_at_most_once(server: R response: Final = await call_aocr(server, litellm_logging_obj=logger) await drain_logging() assert "async_log_success_event" not in recorder.names + assert logger.deferred_success_result is not None + assert logger.deferred_success_result[0] is response + assert not logger.async_success_log_enqueued ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, blocked) ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, blocked) + assert logger.async_success_log_enqueued is not blocked await drain_logging() events: Final = tuple(event for event in recorder.events if event.name == "async_log_success_event") assert len(events) == int(not blocked) @@ -220,6 +224,39 @@ async def test_deferred_logging_requires_release_and_runs_at_most_once(server: R assert events[0].response is response +@pytest.mark.asyncio +@pytest.mark.parametrize("blocked", [False, True]) +async def test_deferred_sync_success_callbacks_wait_for_release(server: RecordingServer, blocked: bool) -> None: + finished: Final = threading.Event() + observations: Final = [] + + class Observe(CustomLogger): + def log_success_event(self, kwargs, response_obj, start_time, end_time): + observations.append(response_obj) + finished.set() + + logger: Final = Logging( + model="mistral-ocr-latest", + messages=[], + stream=False, + call_type="aocr", + start_time=datetime.datetime.now(), + litellm_call_id="deferred-sync", + function_id="deferred-sync", + dynamic_success_callbacks=[Observe()], + ) + logger._defer_async_logging = True + response: Final = await call_aocr(server, litellm_logging_obj=logger) + await drain_logging() + assert not await asyncio.to_thread(finished.wait, 0.2) + assert observations == [] + ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, blocked) + ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, blocked) + await drain_logging() + assert await asyncio.to_thread(finished.wait, 0.5 if blocked else 5) is not blocked + assert observations == ([] if blocked else [response]) + + def test_sync_success_uses_executor_and_copied_caller_context(server: RecordingServer) -> None: context: Final = ContextVar("sync-lifecycle", default="missing") context.set("caller") @@ -454,7 +491,7 @@ async def test_empty_callbacks_run_deployment_hooks_and_defer_like_the_python_cl assert logger.model_call_details["first_api_call_start_time"] <= logger.model_call_details["end_time"] assert logger.model_call_details["response_cost"] == 0 else: - assert getattr(logger, "_native_pending_logging", None) is not None + assert logger.enqueue_deferred_logging is not None assert "end_time" not in logger.model_call_details diff --git a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py index e9e0a585949..d2161fddeb0 100644 --- a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py @@ -416,38 +416,6 @@ async def test_deferred_anthropic_messages_bridged_to_the_responses_api_logs_the # --------------------------------------------------------------------------- -@pytest.mark.parametrize("call_type", ["ocr", "aocr", "completion", "acompletion", "embedding", "responses"]) -@pytest.mark.parametrize("exception_raised", [False, True]) -def test_native_pending_logging_is_released_only_for_ocr(call_type: str, exception_raised: bool) -> None: - pending: Final = MagicMock() - enqueue: Final = MagicMock() - logger: Final = MagicMock( - call_type=call_type, - _native_pending_logging=pending, - enqueue_deferred_logging=enqueue, - ) - - ProxyBaseLLMRequestProcessing._flush_deferred_async_logging( - logging_obj=logger, - exception_raised=exception_raised, - ) - ProxyBaseLLMRequestProcessing._flush_deferred_async_logging( - logging_obj=logger, - exception_raised=exception_raised, - ) - - if call_type in ("ocr", "aocr"): - pending.release.assert_called_once_with(not exception_raised) - assert logger._native_pending_logging is None - else: - pending.release.assert_not_called() - assert logger._native_pending_logging is pending - if exception_raised: - enqueue.assert_not_called() - else: - enqueue.assert_called_once_with() - - def test_flush_deferred_async_logging_fires_on_success(): """ Happy path: with no exception, the production flush helper invokes the diff --git a/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py index 1ba0188e76d..29a6b6f56a8 100644 --- a/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py +++ b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py @@ -101,6 +101,22 @@ def _kwargs_with_a_budget_reservation(reservation: dict) -> dict[str, object]: return {**_OCR_KWARGS, "metadata": {"user_api_key_budget_reservation": reservation}} +def test_a_deferred_success_waits_for_the_proxy_release() -> None: + logger: Final = _supplied_logger() + released: Final[list[str]] = [] + start: Final = datetime.datetime.now() + end: Final = start + datetime.timedelta(seconds=1) + + legacy.defer_success(logger, lambda: released.append("released"), "response", start, end) + + assert logger.deferred_success_result == ("response", start, end) + assert released == [] + release: Final = logger.enqueue_deferred_logging + assert release is not None + release() + assert released == ["released"] + + def test_setup_claims_the_budget_reservation_for_an_async_call() -> None: reservation: Final = _budget_reservation() From d631f3cf05a58df6017bf60840838309dce530fc Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Mon, 9 Nov 2026 16:24:19 -0800 Subject: [PATCH 33/46] fix(messages): log native messages request fields like python (#50421) Python logs the requested optional params, taken before the provider transform, as model_parameters. The native handler logged the transformed body, so a reasoning_effort request was logged as thinking and output_config. The request context now carries the requested params Python updates model_call_details with the transformed request body before pre_call, so the logged messages are what was sent. The native adapter kept the caller's messages and compared params as JSON values, so an int 0 logged as 0.0 Python only records custom_llm_provider in model_call_details. The native adapter also wrote it into litellm_params, and the proxy reuses one logger across router hops, so a fallback hop's pre_call log showed the provider where Python shows null Python only stamps cache_hit when a result came from the cache. The native adapter wrote False on every provider result Python calls pre_call with the request body as a json user message, and a stream's spend-log cache_key is computed over that input. The native adapter logged the string Messages, so a streamed call's cache_key differed --- .../python_contract.json | 4 + .../callbacks-legacy-python/src/adapter.rs | 146 +++++++++++++++--- .../callbacks-legacy-python/src/callbacks.rs | 4 +- .../callbacks-legacy-python/src/python.rs | 2 + .../src/test_support.rs | 2 + .../crates/inference-messages/src/handler.rs | 3 +- .../inference-messages/tests/messages/host.rs | 57 ++++++- .../host/callbacks_legacy_python.py | 27 +++- .../messages/test_callbacks.py | 73 ++++++++- .../messages/test_streaming.py | 40 +++++ .../host/test_callbacks_legacy_python.py | 24 +++ 11 files changed, 356 insertions(+), 26 deletions(-) diff --git a/litellm-rust/crates/callbacks-legacy-python/python_contract.json b/litellm-rust/crates/callbacks-legacy-python/python_contract.json index 1d4a2f7f7c2..c0fe1c7b52f 100644 --- a/litellm-rust/crates/callbacks-legacy-python/python_contract.json +++ b/litellm-rust/crates/callbacks-legacy-python/python_contract.json @@ -27,6 +27,10 @@ "api_key", "additional_args" ], + "messages_input": [ + "body", + "kwargs" + ], "post_call": [ "logger", "original_response", diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index 8d1946861aa..28bd5f2cefb 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -14,7 +14,7 @@ use pyo3::{ exceptions::{PyBaseException, PyException}, gc::{PyTraverseError, PyVisit}, prelude::*, - types::{PyDateTime, PyDict, PyList}, + types::{PyDateTime, PyDict, PyList, PyString}, }; use serde_json::Value; @@ -22,7 +22,7 @@ use crate::{ DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger, deferred::{PendingLogging, PendingSuccess}, finalize, is_internal_call, - python::Streaming, + python::{Logging, Streaming}, setup, }; @@ -137,6 +137,40 @@ impl LegacyLogging { Some(params) } + // Compared as Python compares them, so a caller's `0` is kept over a typed float's `0.0`. + fn caller_value<'py>( + &self, + py: Python<'py>, + name: &str, + sent: &Value, + ) -> PyResult>> { + match self.call.lookup(py, name)? { + Some(value) if from_py::(&value).is_ok() && value.eq(to_py(py, sent)?)? => { + Ok(Some(value)) + } + _ => Ok(None), + } + } + + // Python's Messages handler logs the caller's own optional params, not their typed copies. + fn as_called(&self, py: Python<'_>, context: &RequestContext) -> PyResult { + let mut context = context.clone(); + if self.operation != LoggingOperation::Messages { + return Ok(context); + } + for (name, sent) in context + .optional_params + .as_object_mut() + .into_iter() + .flatten() + { + if let Some(value) = self.caller_value(py, name, sent)? { + *sent = from_py(&value)?; + } + } + Ok(context) + } + /// Python times the provider exchange inside the async `AsyncHTTPHandler.post`; its sync /// client is untimed. fn close_api_window(&mut self, py: Python<'_>, end: f64) -> PyResult<()> { @@ -158,13 +192,20 @@ impl LegacyLogging { } } - fn input_description(&self) -> &'static str { - match self.operation { + fn logged_input<'py>( + &self, + py: Python<'py>, + body: &Bound<'py, PyDict>, + ) -> PyResult> { + let description = match self.operation { LoggingOperation::Completion => "Chat completions", LoggingOperation::Responses => "Responses", - LoggingOperation::Messages => "Messages", + LoggingOperation::Messages => { + return Logging::MessagesInput.call(py, (body, self.call.kwargs().bind(py))); + } LoggingOperation::Ocr => "OCR document processing", - } + }; + Ok(PyString::new(py, description).into_any()) } /// The key `pre_call` logs. Python's Messages handler logs an empty string, never the key. @@ -456,6 +497,7 @@ impl LegacyLogging { wire: Box, context: &RequestContext, ) -> PyResult>> { + let context = &self.as_called(py, context)?; let logger = self.logger()?; logger.update_from_kwargs( py, @@ -468,12 +510,17 @@ impl LegacyLogging { .into_bound(py) .cast_into::()?; for (name, sent) in wire.body.as_object().into_iter().flatten() { - if let Some(value) = self.call.lookup(py, name)? - && from_py::(&value).is_ok_and(|caller| caller == *sent) - { + if let Some(value) = self.caller_value(py, name, sent)? { body.set_item(name, value)?; } } + if self.operation == LoggingOperation::Messages { + // Python's Messages handler merges its transformed request body into the call details. + logger + .object(py) + .getattr("model_call_details")? + .call_method1("update", (&body,))?; + } let headers = PyDict::new(py); for (name, value) in &wire.headers { headers.set_item(name, value)?; @@ -483,9 +530,10 @@ impl LegacyLogging { headers: headers.clone().unbind(), context: context.clone(), }); + let input = self.logged_input(py, &body)?; self.logger()?.pre_call( py, - self.input_description(), + &input, self.logged_api_key(context), &body, &headers, @@ -539,7 +587,11 @@ impl LegacyLogging { .getattr("litellm_params")? .cast_into::()? .copy()?; - params.set_item("custom_llm_provider", &facts.provider.provider)?; + if self.operation != LoggingOperation::Messages { + // Python's Messages handler hands the provider to the logger beside its params, never + // in them, and the proxy's logger carries over to the next fallback hop. + params.set_item("custom_llm_provider", &facts.provider.provider)?; + } crate::python::Logging::Update.call( py, ( @@ -562,8 +614,17 @@ impl LegacyLogging { ResultSource::Provider => None, ResultSource::Cache { key } => Some(key.clone()), }; - details.set_item("cache_hit", self.cache_key.is_some())?; + if self.cache_key.is_some() { + details.set_item("cache_hit", true)?; + } else { + details.call_method1("pop", ("cache_hit", py.None()))?; + } details.set_item("cache_key", self.cache_key.as_deref())?; + if let (LoggingOperation::Messages, ResultSource::Provider, Some(request)) = + (self.operation, &facts.source, &self.request) + { + details.call_method1("update", (request.body.bind(py),))?; + } Ok(HookStep::Ready(())) } @@ -1128,10 +1189,10 @@ logger._provider_reported_cost = None }; let previous = litellm_host::interceptors::ExecutionFacts { provider: facts.provider.clone(), - cost: litellm_host::interceptors::Cost::Reported { - amount: serde_json::Number::from_f64(9.0).unwrap(), + cost: litellm_host::interceptors::Cost::CacheHit, + source: litellm_host::interceptors::ResultSource::Cache { + key: "earlier-hit".into(), }, - source: litellm_host::interceptors::ResultSource::Provider, }; logging.result_ready(py, &previous).unwrap(); local(&locals, "logger") @@ -1149,13 +1210,59 @@ logger._provider_reported_cost = None c" assert logger.model_call_details.get('response_cost') is None, logger.model_call_details assert logger._provider_reported_cost == reported, logger._provider_reported_cost -assert logger.model_call_details['cache_hit'] is cache_hit, logger.model_call_details +assert logger.model_call_details.get('cache_hit') is (True if cache_hit else None), logger.model_call_details assert logger.update['custom_llm_provider'] == 'edenai', logger.update ", ); }); } + #[rstest] + #[case::messages(crate::LoggingOperation::Messages, false)] + #[case::ocr(crate::LoggingOperation::Ocr, true)] + fn result_ready_puts_the_provider_in_the_logged_params_except_for_messages( + #[case] operation: crate::LoggingOperation, + #[case] in_params: bool, + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, PAYLOAD_LOGGER); + run( + py, + &locals, + c" +logger.litellm_params = {} +logger.optional_params = {} +logger.model_call_details = {} +logger._provider_reported_cost = None +", + ); + let mut logging = LegacyLogging { + logger: Some(PythonLogger::new(local(&locals, "logger").unbind())), + operation, + ..legacy_call(py, &locals, false) + }; + let facts = litellm_host::interceptors::ExecutionFacts { + provider: litellm_host::interceptors::ProviderIdentity { + model: "native-test".into(), + provider: "anthropic".into(), + }, + cost: litellm_host::interceptors::Cost::Deferred, + source: litellm_host::interceptors::ResultSource::Provider, + }; + logging.result_ready(py, &facts).unwrap(); + locals.set_item("in_params", in_params).unwrap(); + run( + py, + &locals, + c" +assert logger.update['custom_llm_provider'] == 'anthropic', logger.update +assert ('custom_llm_provider' in logger.update['litellm_params']) is in_params, logger.update +", + ); + }); + } + const DOCUMENT: &str = "data:application/pdf;base64,YWJj"; const EDITED: &str = "data:application/pdf;base64,ZWRpdGVk"; @@ -1290,7 +1397,10 @@ assert logger.update['custom_llm_provider'] == 'edenai', logger.update #[rstest] #[case::completion(crate::LoggingOperation::Completion, "Chat completions")] #[case::responses(crate::LoggingOperation::Responses, "Responses")] - #[case::messages(crate::LoggingOperation::Messages, "Messages")] + #[case::messages( + crate::LoggingOperation::Messages, + r#"[{'role': 'user', 'content': '{"pages": [1]}'}]"# + )] #[case::ocr(crate::LoggingOperation::Ocr, "OCR document processing")] fn prepared_arguments_replace_the_legacy_view_without_losing_callback_aliases( #[case] operation: crate::LoggingOperation, @@ -1331,7 +1441,7 @@ prepared = {'pages': replacement} &locals, c" assert logger.pre['complete_input_dict']['pages'] is replacement -assert logger.pre_input == description +assert str(logger.pre_input) == description assert original == [0] ", ); diff --git a/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs b/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs index a428afa9ece..093e86c43dc 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs @@ -26,7 +26,7 @@ pub trait LegacyCallbacks { fn pre_call( &self, py: Python<'_>, - input: &str, + input: &Bound<'_, PyAny>, api_key: Option<&str>, body: &Bound<'_, PyDict>, headers: &Bound<'_, PyDict>, @@ -157,7 +157,7 @@ impl LegacyCallbacks for PythonLogger { fn pre_call( &self, py: Python<'_>, - input: &str, + input: &Bound<'_, PyAny>, api_key: Option<&str>, body: &Bound<'_, PyDict>, headers: &Bound<'_, PyDict>, diff --git a/litellm-rust/crates/callbacks-legacy-python/src/python.rs b/litellm-rust/crates/callbacks-legacy-python/src/python.rs index 32b11988703..5b6faa3dbc2 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/python.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/python.rs @@ -42,6 +42,8 @@ pub(crate) enum Logging { Update, #[strum(serialize = "pre_call")] PreCall, + #[strum(serialize = "messages_input")] + MessagesInput, #[strum(serialize = "post_call")] PostCall, #[strum(serialize = "record_response")] diff --git a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs index 9e69d0c30e2..c88a6496423 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs @@ -62,6 +62,7 @@ FAKES = { custom_llm_provider=provider, ), 'pre_call': lambda logger, input, api_key, additional_args: logger.pre_call(input, api_key, additional_args), + 'messages_input': lambda body, kwargs: [{'role': 'user', 'content': json.dumps(body)}], 'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call( original_response, api_key, additional_args ), @@ -134,6 +135,7 @@ class StubCoroutine: class StubLogger: def __init__(self): self.calls = [] + self.model_call_details = {} self.hooks = {} self.on_enqueue = lambda coroutine: None diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index e0a7aa5adb8..5ab171076ff 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -131,6 +131,7 @@ async fn provider_call( let transform_context = MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params); let filtered = requested_optional_params(sanitized, provider, &transform_context)?; + let requested_params = serde_json::to_value(&filtered.params).map_err(serialize_failure)?; let trimmed = without_additional_drop_params(filtered, &shaping.settings.additional_drop_params)?; @@ -173,7 +174,7 @@ async fn provider_call( let request_context = RequestContext { model: body.model.clone(), custom_llm_provider: provider.as_str().to_string(), - optional_params: serde_json::to_value(&body.params).map_err(serialize_failure)?, + optional_params: requested_params, secret_fields: Vec::new(), api_key: api_key.map(SecretValue::new), }; diff --git a/litellm-rust/crates/inference-messages/tests/messages/host.rs b/litellm-rust/crates/inference-messages/tests/messages/host.rs index c3d06027c32..e8df20e7dcb 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/host.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/host.rs @@ -8,7 +8,9 @@ use litellm_host::{ }; use litellm_inference_messages::{MessagesCallResponse, route::Messages}; use litellm_inference_testing::{RecordingSecrets, no_secrets}; -use litellm_llms::base_llm::messages::context::MessagesModelCapabilities as AnthropicModelCapabilities; +use litellm_llms::base_llm::messages::context::{ + MessagesModelCapabilities as AnthropicModelCapabilities, SupportedEffortTiers, +}; use rstest::rstest; use super::*; @@ -419,7 +421,7 @@ async fn no_raw_response_is_emitted_for_a_stream_or_a_failure( assert!(host.raw_responses().is_empty()); } -/// Python logs `optional_params` as what it is about to send, so a dropped param must +/// Python logs `optional_params` after its requested-param gate, so a dropped param must /// not resurface in callbacks. #[rstest] #[tokio::test] @@ -450,3 +452,54 @@ async fn the_request_context_carries_the_shaped_params_without_model_or_messages .unwrap_or_else(|seen| panic!("before_provider_request runs once, saw {}", seen.len())); assert_eq!(optional_params, json!({"max_tokens": 16})); } + +#[rstest] +#[case::reasoning_effort_on_an_adaptive_model( + AnthropicModelCapabilities { + supports_reasoning: true, + supports_adaptive_thinking: true, + supports_output_config: true, + effort_tiers: SupportedEffortTiers { low: true, ..SupportedEffortTiers::default() }, + ..AnthropicModelCapabilities::default() + }, + json!({"reasoning_effort": "low"}) +)] +#[case::adaptive_payload_on_a_legacy_model( + AnthropicModelCapabilities { + supports_reasoning: true, + ..AnthropicModelCapabilities::default() + }, + json!({"max_tokens": 3000, "thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}}) +)] +#[tokio::test] +async fn the_request_context_carries_the_requested_reasoning_not_its_translation( + call: MessagesCall, + #[case] capabilities: AnthropicModelCapabilities, + #[case] fields: Value, +) { + let upstream = upstream([message_response()]).await; + let host = RecordingHost::passthrough(authenticated( + MessagesCall { + shaping: MessagesShaping { + capabilities, + ..MessagesShaping::default() + }, + ..with_fields(call, fields.clone()) + }, + upstream.uri(), + )); + + run_through(&host).await.expect("messages call succeeds"); + + let [optional_params] = <[Value; 1]>::try_from(host.optional_params.into_inner().unwrap()) + .unwrap_or_else(|seen| panic!("before_provider_request runs once, saw {}", seen.len())); + assert_eq!( + optional_params, + Value::Object( + object(json!({"max_tokens": 16})) + .into_iter() + .chain(object(fields)) + .collect() + ) + ); +} diff --git a/litellm/rust_bridge/host/callbacks_legacy_python.py b/litellm/rust_bridge/host/callbacks_legacy_python.py index 92e5436fbb8..8c281c73842 100644 --- a/litellm/rust_bridge/host/callbacks_legacy_python.py +++ b/litellm/rust_bridge/host/callbacks_legacy_python.py @@ -9,6 +9,8 @@ from __future__ import annotations import asyncio import contextvars import datetime +import inspect +import json import traceback import uuid import weakref @@ -241,10 +243,33 @@ def update_logging( ) -def pre_call(logger: LoggingSurface, input: str, api_key: str | None, additional_args: dict[str, object]) -> None: +def pre_call( + logger: LoggingSurface, + input: str | Sequence[Mapping[str, str]], + api_key: str | None, + additional_args: dict[str, object], +) -> None: logger.pre_call(input=input, api_key=api_key, additional_args=additional_args) +def messages_input(body: Mapping[str, object], kwargs: Mapping[str, object]) -> Sequence[Mapping[str, str]]: + from litellm.messages import anthropic_messages_handler + + # Python's Messages handler builds its body as messages, max_tokens, model, then its own parameters in + # signature order and the caller's other keywords in their order, and logs that body as serialized + order: Final = ( + "messages", + "max_tokens", + "model", + *inspect.signature(anthropic_messages_handler).parameters, + *kwargs, + *body, + ) + return [ + {"role": "user", "content": json.dumps({name: body[name] for name in dict.fromkeys(order) if name in body})} + ] + + def post_call( logger: LoggingSurface, original_response: str, api_key: str | None, additional_args: dict[str, object] ) -> None: diff --git a/tests/test_litellm_rust/messages/test_callbacks.py b/tests/test_litellm_rust/messages/test_callbacks.py index a3c03898a12..3073c1cf7dd 100644 --- a/tests/test_litellm_rust/messages/test_callbacks.py +++ b/tests/test_litellm_rust/messages/test_callbacks.py @@ -10,14 +10,15 @@ from pydantic import TypeAdapter import litellm from litellm.integrations.argilla import ArgillaLogger from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.types.integrations.argilla import ArgillaItem from litellm.types.utils import ModelResponse, StandardLoggingPayload -from tests.test_litellm_rust.messages.data import BAD_REQUEST, MESSAGES_RESPONSE, RATE_LIMITED, call +from tests.test_litellm_rust.messages.data import BAD_REQUEST, MESSAGES_RESPONSE, RATE_LIMITED, REQUEST, call from tests.test_litellm_rust.support.callback_recorder import HookRecorder, RecordingLogger, request_body from tests.test_litellm_rust.support.isolation import rebound from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed from tests.test_litellm_rust.support.results import response_id, with_id -from tests.test_litellm_rust.support.routing import assert_served_natively +from tests.test_litellm_rust.support.routing import CALL_ID, assert_served_natively pytestmark = pytest.mark.requires_rust_extension @@ -110,6 +111,33 @@ async def test_pre_api_call_sees_the_wire_body(server: RecordingServer) -> None: assert request_body(pre_call) == server.requests[0].body +@pytest.mark.asyncio +async def test_pre_api_call_of_a_retried_call_keeps_the_provider_out_of_the_litellm_params( + server: RecordingServer, +) -> None: + recorder: Final = RecordingLogger() + logging_obj: Final = Logging( + model=REQUEST["model"], + messages=[], + stream=False, + start_time=datetime.now(), + call_type="anthropic_messages", + litellm_call_id=CALL_ID, + function_id=CALL_ID, + dynamic_input_callbacks=[recorder], + ) + + await call(server, litellm_logging_obj=logging_obj) + await call(server, litellm_logging_obj=logging_obj) + + pre_calls: Final = recorder.wait_for("log_pre_api_call", count=2) + + assert_served_natively(server, count=2) + for pre_call in pre_calls: + assert isinstance(pre_call.kwargs, dict) + assert "custom_llm_provider" not in _OBJECT.validate_python(pre_call.kwargs["litellm_params"]) + + @pytest.mark.asyncio async def test_pre_api_call_edits_reach_the_provider(server: RecordingServer) -> None: class Edit(CustomLogger): @@ -140,6 +168,47 @@ async def test_pre_api_call_sees_the_shaped_optional_params(server: RecordingSer assert optional_params["max_tokens"] == sent["max_tokens"] +@pytest.mark.asyncio +async def test_success_event_logs_the_requested_reasoning_effort_not_its_translation(server: RecordingServer) -> None: + recorder: Final = RecordingLogger() + + with rebound(litellm, "callbacks", [recorder]): + response: Final = await call(server, reasoning_effort="low") + success: Final = await recorder.wait_for_async("async_log_success_event") + + assert_served_natively(server, response) + sent: Final = server.requests[0].body + assert isinstance(sent, dict) + assert sent["output_config"] == {"effort": "low"} + assert isinstance(success[0].kwargs, dict) + logged: Final = _OBJECT.validate_python(success[0].kwargs["standard_logging_object"]) + assert logged["model_parameters"] == {"max_tokens": 64, "stream": False, "reasoning_effort": "low"} + + +@pytest.mark.asyncio +async def test_success_event_logs_the_sent_messages_and_the_callers_params(server: RecordingServer) -> None: + recorder: Final = RecordingLogger() + tool_use: Final = {"type": "tool_use", "id": "functions.Bash:0", "name": "bash", "input": {"command": "ls"}} + messages: Final = [ + {"role": "user", "content": "List /tmp"}, + {"role": "assistant", "content": [{"type": "text", "text": " \n"}, tool_use]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "functions.Bash:0", "content": "a"}]}, + ] + + with rebound(litellm, "callbacks", [recorder]): + response: Final = await call(server, messages=messages, temperature=1) + success: Final = await recorder.wait_for_async("async_log_success_event") + + assert_served_natively(server, response) + sent: Final = server.requests[0].body + assert isinstance(sent, dict) + assert sent["messages"][1]["content"] == [{**tool_use, "id": "functions_Bash_0"}] + assert isinstance(success[0].kwargs, dict) + logged: Final = _OBJECT.validate_python(success[0].kwargs["standard_logging_object"]) + assert logged["messages"] == sent["messages"] + assert repr(_OBJECT.validate_python(logged["model_parameters"])["temperature"]) == "1" + + @pytest.mark.asyncio async def test_logging_hook_and_redaction_edits_reach_the_async_success_event(server: RecordingServer) -> None: class Rewrite(HookRecorder): diff --git a/tests/test_litellm_rust/messages/test_streaming.py b/tests/test_litellm_rust/messages/test_streaming.py index e57631f48b9..491d257b482 100644 --- a/tests/test_litellm_rust/messages/test_streaming.py +++ b/tests/test_litellm_rust/messages/test_streaming.py @@ -1,10 +1,13 @@ +import json from collections.abc import AsyncIterator, Iterator from typing import Final import pytest +import litellm from tests.test_litellm_rust.messages.data import MESSAGES_EVENTS, MESSAGES_RESPONSE, call from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging +from tests.test_litellm_rust.support.isolation import rebound from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed from tests.test_litellm_rust.support.routing import assert_served_natively @@ -68,3 +71,40 @@ async def test_stream_closed_early_logs_success_once(server: RecordingServer) -> assert len(await recorder.wait_for_async("async_log_success_event")) == 1 with pytest.raises(StopAsyncIteration): await anext(stream) + + +async def _streamed_pre_call_input(server: RecordingServer) -> object: + server.enqueue(STREAM) + recorder: Final = RecordingLogger() + with rebound(litellm, "callbacks", [recorder]): + stream: Final = await call( + server, stream=True, system="Be brief", metadata={"user_id": "user-1"}, stop_sequences=["END"] + ) + assert isinstance(stream, AsyncIterator) + _ = [chunk async for chunk in stream] + pre_call: Final = recorder.wait_for("log_pre_api_call")[0].kwargs + assert isinstance(pre_call, dict) + return pre_call["input"] + + +@pytest.mark.asyncio +async def test_stream_logs_the_request_body_as_input_like_python(server: RecordingServer) -> None: + logged: Final = await _streamed_pre_call_input(server) + + assert_served_natively(server) + assert logged == [ + { + "role": "user", + "content": json.dumps( + { + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 64, + "model": "claude-sonnet-5", + "metadata": {"user_id": "user-1"}, + "stop_sequences": ["END"], + "stream": True, + "system": "Be brief", + } + ), + } + ] diff --git a/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py index 29a6b6f56a8..87b9e70612c 100644 --- a/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py +++ b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py @@ -263,6 +263,30 @@ def test_the_rust_contract_matches_the_shim_signatures() -> None: assert contract == {name: list(inspect.signature(getattr(legacy, name)).parameters) for name in contract} +def test_messages_input_is_the_body_as_the_python_handler_builds_and_serializes_it() -> None: + body: Final = { + "model": "claude-sonnet-4-5", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 16, + "system": "brief", + "stream": True, + "context_management": {"edits": []}, + "output_config": {"effort": "low"}, + } + + logged: Final = legacy.messages_input( + body, {"output_config": {"effort": "low"}, "context_management": {"edits": []}} + ) + + assert logged == [ + { + "role": "user", + "content": '{"messages": [{"role": "user", "content": "hi"}], "max_tokens": 16, "model": "claude-sonnet-4-5", ' + '"stream": true, "system": "brief", "output_config": {"effort": "low"}, "context_management": {"edits": []}}', + } + ] + + def _messages_logger() -> Logging: return Logging( model="claude-sonnet-4-5", From ee032f10d66623da162eae53b8fb535486d04a33 Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Mon, 9 Nov 2026 16:24:20 -0800 Subject: [PATCH 34/46] fix(messages): pass non-message bodies through and refuse non-string headers as bad requests (#50422) Python's messages transform takes any JSON object from a 200 response and only adds stop_sequence, so a provider answering 200 with an error object reaches the client as 200 and the router counts a success. The native decoder turned it into a 500 APIConnectionError and a router failure. A JSON object that is not a message now passes through unchanged A non-string extra_headers or deployment header value surfaced from the native handler as an APIConnectionError. It is now refused as an invalid request when the headers are collected, before the provider request is built or logged, so the caller gets a 400 BadRequestError naming the header and its type --- .../inference-messages/src/common_utils.rs | 83 +++++++++++++++---- .../crates/inference-messages/src/handler.rs | 34 +++++--- .../crates/inference-messages/src/lib.rs | 4 +- .../crates/inference-messages/src/route.rs | 7 +- .../crates/inference-messages/src/types.rs | 11 ++- .../tests/bedrock_native/main.rs | 9 +- .../inference-messages/tests/caching.rs | 72 +++++++++------- .../inference-messages/tests/live/support.rs | 9 +- .../inference-messages/tests/messages/host.rs | 7 +- .../inference-messages/tests/messages/main.rs | 7 +- .../tests/messages/response.rs | 48 ++++++++++- .../python-bridge/src/routes/messages/host.rs | 4 +- .../messages/test_callbacks.py | 17 ++++ .../messages/test_requests.py | 17 +++- 14 files changed, 250 insertions(+), 79 deletions(-) diff --git a/litellm-rust/crates/inference-messages/src/common_utils.rs b/litellm-rust/crates/inference-messages/src/common_utils.rs index e5ece2837d8..d0f109c1e99 100644 --- a/litellm-rust/crates/inference-messages/src/common_utils.rs +++ b/litellm-rust/crates/inference-messages/src/common_utils.rs @@ -1,5 +1,5 @@ use litellm_core_utils::get_llm_provider_logic::LlmProviders; -use litellm_http::request::string_headers as shared_string_headers; +use litellm_llms::ErrorDetail; use litellm_llms::tencent::messages::transformation::TENCENT_MESSAGES_CONFIG; use litellm_llms::{ anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG, @@ -19,8 +19,6 @@ use serde_json::{Map, Value}; use super::Error; -const HEADER_CONTEXT: &str = "messages"; - #[derive(Clone, Copy)] pub(crate) struct MessagesProvider { provider: LlmProviders, @@ -75,7 +73,52 @@ pub(crate) fn messages_provider(provider: LlmProviders, model: &str) -> Option>, ) -> Result, Error> { - shared_string_headers(HEADER_CONTEXT, extra_headers).map_err(Error::from) + extra_headers + .unwrap_or_default() + .into_iter() + .map(|(name, value)| match value { + Value::String(value) => Ok((name, value)), + value => Err(Error::invalid_request(ErrorDetail::invalid( + "extra_headers", + NonStringHeader { + actual: python_type_name(&value), + name, + }, + ))), + }) + .collect() +} + +/// A header value the caller set to something other than a string, named by its Python type +/// since the caller wrote it in Python. +#[derive(Debug)] +struct NonStringHeader { + name: String, + actual: &'static str, +} + +impl std::fmt::Display for NonStringHeader { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + formatter, + "header {} must be a string, got {}", + self.name, self.actual + ) + } +} + +impl std::error::Error for NonStringHeader {} + +fn python_type_name(value: &Value) -> &'static str { + match value { + Value::Null => "NoneType", + Value::Bool(_) => "bool", + Value::Number(number) if number.is_f64() => "float", + Value::Number(_) => "int", + Value::String(_) => "str", + Value::Array(_) => "list", + Value::Object(_) => "dict", + } } #[cfg(test)] @@ -85,7 +128,6 @@ mod tests { use rstest::rstest; use super::{messages_provider, string_headers}; - use crate::Error; use litellm_core_utils::get_llm_provider_logic::LlmProviders; #[rstest] @@ -113,17 +155,28 @@ mod tests { assert!(messages_provider(provider, model).is_none()); } - #[test] - fn string_headers_rejects_non_string_values() { - let headers = json!({"x-count": 3}).as_object().unwrap().clone(); - let err = string_headers(Some(headers)).expect_err("non-string header rejected"); + #[rstest] + #[case::int(json!(12345), "int")] + #[case::float(json!(1.5), "float")] + #[case::bool(json!(true), "bool")] + #[case::null(json!(null), "NoneType")] + #[case::list(json!([1]), "list")] + #[case::dict(json!({"a": 1}), "dict")] + fn non_string_header_is_a_bad_request_naming_the_header_and_its_type( + #[case] value: serde_json::Value, + #[case] python_type: &str, + ) { + let headers = json!({"x-kept": "v", "x-count": value}) + .as_object() + .unwrap() + .clone(); + let error = string_headers(Some(headers)).expect_err("non-string header refused"); + assert!(error.is_request(), "{error:?}"); + assert_eq!(error.upstream_response(), None); + let rejection = error.rejected_request().expect("a rejected request"); assert_eq!( - err, - Error::Headers(litellm_http::request::HeaderError { - context: "messages", - name: "x-count".to_string(), - actual: "number", - }) + litellm_error::Report::new(rejection).to_string(), + format!("invalid extra_headers: header x-count must be a string, got {python_type}") ); } } diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index 5ab171076ff..5eb017ddfc5 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -26,10 +26,10 @@ use litellm_llms_types::formats::messages::{ }; use litellm_secrets::source::Secrets; use litellm_tracing::ByteChunk; -use serde_json::Value; +use serde_json::{Map, Value}; use super::{ - Error, MessagesCall, MessagesCallResponse, MessagesRoute, + Error, MessagesBody, MessagesCall, MessagesCallResponse, MessagesRoute, common_utils::{MessagesProvider, merge_messages_header_sources, string_headers}, constants::MESSAGES_TIMEOUT_SECS, route::Messages, @@ -85,7 +85,10 @@ pub(super) async fn execute( ), }; let reported = match &output { - MessagesCallResponse::Complete(response) => provider.config().reported_cost(response), + MessagesCallResponse::Complete(response) => match response { + MessagesBody::Message(message) => provider.config().reported_cost(message), + MessagesBody::Unchecked(_) => None, + }, MessagesCallResponse::Stream { .. } => None, }; context @@ -347,8 +350,7 @@ async fn call_provider( // body it cannot parse, and the public mapping classifies it by that status. return Err(litellm_http::Error::upstream(status, text).into()); } - decode_response(config, &identity.model, &text) - .map(|message| MessagesCallResponse::Complete(Box::new(message))) + decode_response(config, &identity.model, &text).map(MessagesCallResponse::Complete) } fn serialize_failure(err: serde_json::Error) -> Error { @@ -412,15 +414,23 @@ fn decode_response( config: &dyn BaseMessagesConfig, model: &str, text: &str, -) -> Result { - let response = serde_json::from_str(text).map_err(|err| { - Error::invalid_response(litellm_llms::ErrorDetail::invalid( - "messages response JSON", - err, - )) - })?; +) -> Result { + let response = match serde_json::from_str::(text) { + Ok(response) => response, + Err(err) => { + let Ok(mut body) = serde_json::from_str::>(text) else { + return Err(Error::invalid_response(litellm_llms::ErrorDetail::invalid( + "messages response JSON", + err, + ))); + }; + body.entry("stop_sequence").or_insert(Value::Null); + return Ok(MessagesBody::Unchecked(body)); + } + }; config .transform_anthropic_messages_response(model, response) + .map(|message| MessagesBody::Message(Box::new(message))) .map_err(Error::from) .map_err(Error::after_send) } diff --git a/litellm-rust/crates/inference-messages/src/lib.rs b/litellm-rust/crates/inference-messages/src/lib.rs index 745442d0314..25ebeaacb49 100644 --- a/litellm-rust/crates/inference-messages/src/lib.rs +++ b/litellm-rust/crates/inference-messages/src/lib.rs @@ -17,8 +17,8 @@ use crate::common_utils::messages_provider; pub use common_utils::merge_messages_header_sources; pub use litellm_inference::Error; pub use types::{ - MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping, litellm_params, - messages_body, + MessagesBody, MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping, + litellm_params, messages_body, }; #[derive(Clone)] diff --git a/litellm-rust/crates/inference-messages/src/route.rs b/litellm-rust/crates/inference-messages/src/route.rs index 7ff886689da..5f06d16d49f 100644 --- a/litellm-rust/crates/inference-messages/src/route.rs +++ b/litellm-rust/crates/inference-messages/src/route.rs @@ -5,11 +5,10 @@ use litellm_host::{ call::{HostedCompletion, HostedMachine, hosted_call}, protocol::Protocol, }; -use litellm_llms_types::formats::messages::MessagesResponse; -use super::{Error, MessagesCall}; +use super::{Error, MessagesBody, MessagesCall}; -pub type MessagesOutput = HostedCompletion>; +pub type MessagesOutput = HostedCompletion; /// The upstream response as the caller sees it at stream hand-off, before any chunk. pub struct MessagesStreamHead { @@ -19,7 +18,7 @@ pub struct MessagesStreamHead { pub struct Messages; impl Protocol for Messages { - type Response = Box; + type Response = MessagesBody; type Error = Error; type Request = MessagesCall; type HostCall = Infallible; diff --git a/litellm-rust/crates/inference-messages/src/types.rs b/litellm-rust/crates/inference-messages/src/types.rs index ece17b3380c..772ea6f7292 100644 --- a/litellm-rust/crates/inference-messages/src/types.rs +++ b/litellm-rust/crates/inference-messages/src/types.rs @@ -42,7 +42,16 @@ pub(super) fn invalid_request(err: serde_json::Error) -> Error { } pub type MessagesCallResponse = - CallOutput, super::route::MessagesStreamHead, Bytes, Error>; + CallOutput; + +// Python's `transform_anthropic_messages_response` returns any JSON object the provider sends +// with a 200, message or not, so one that is not a message passes through as sent. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum MessagesBody { + Message(Box), + Unchecked(Map), +} #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] pub struct MessagesShaping { 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 f91d5c79f7c..5207f4e9d58 100644 --- a/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs +++ b/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs @@ -1,7 +1,9 @@ use std::{sync::Arc, time::Duration}; use futures_util::TryStreamExt; -use litellm_inference_messages::{Error, MessagesCall, MessagesCallResponse, MessagesShaping}; +use litellm_inference_messages::{ + Error, MessagesBody, MessagesCall, MessagesCallResponse, MessagesShaping, +}; use litellm_inference_testing::{RecordingSecrets, no_secrets}; use litellm_llms_types::formats::messages::{MessagesRequest, MessagesResponse}; use rstest::{fixture, rstest}; @@ -74,7 +76,10 @@ async fn run(call: MessagesCall) -> Result { async fn run_message(call: MessagesCall) -> MessagesResponse { match run(call).await.unwrap() { - MessagesCallResponse::Complete(message) => *message, + MessagesCallResponse::Complete(message) => match message { + MessagesBody::Message(message) => *message, + MessagesBody::Unchecked(_) => panic!("expected a message"), + }, MessagesCallResponse::Stream { .. } => panic!("expected a completed message"), } } diff --git a/litellm-rust/crates/inference-messages/tests/caching.rs b/litellm-rust/crates/inference-messages/tests/caching.rs index 663c87453a9..787501a8fed 100644 --- a/litellm-rust/crates/inference-messages/tests/caching.rs +++ b/litellm-rust/crates/inference-messages/tests/caching.rs @@ -50,42 +50,47 @@ async fn messages_cache_identity_includes_provider_native_parameters( #[case] original: Value, #[case] changed: Value, ) { - use litellm_inference_messages::route::Messages; + use litellm_inference_messages::{MessagesBody, route::Messages}; use litellm_llms_types::formats::messages::MessagesResponse; let calls = AtomicUsize::new(0); for (value, expected_call) in [(original.clone(), 0), (changed, 1), (original, 0)] { - let response = - execute_unary::( - CacheRequest::from_wire( - ProviderIdentity { - model: "test".into(), - provider: "anthropic".into(), - }, - Some(&WireRequest { - url: "https://example.test/v1/messages".into(), - headers: vec![], - body: json!({ - "model":"test", "messages":[{"role":"user","content":"hello"}], - "max_tokens":32, (field):value - }), - }), - ), - Some(cache.clone()), - Some(CacheOptions::new(CacheScope::Shared)), - &(), - None, - || async { - let call = calls.fetch_add(1, Ordering::SeqCst); - Ok(Box::new(serde_json::from_value::(json!({ - "id":call.to_string(), "type":"message", "role":"assistant", "model":"test", - "content":[{"type":"text","text":format!("answer {call}")}], - "stop_reason":"end_turn", "stop_sequence":null - })).unwrap())) + let response = execute_unary::( + CacheRequest::from_wire( + ProviderIdentity { + model: "test".into(), + provider: "anthropic".into(), }, - ) - .await - .unwrap(); + Some(&WireRequest { + url: "https://example.test/v1/messages".into(), + headers: vec![], + body: json!({ + "model":"test", "messages":[{"role":"user","content":"hello"}], + "max_tokens":32, (field):value + }), + }), + ), + Some(cache.clone()), + Some(CacheOptions::new(CacheScope::Shared)), + &(), + None, + || async { + let call = calls.fetch_add(1, Ordering::SeqCst); + Ok(MessagesBody::Message(Box::new( + serde_json::from_value::(json!({ + "id":call.to_string(), "type":"message", "role":"assistant", "model":"test", + "content":[{"type":"text","text":format!("answer {call}")}], + "stop_reason":"end_turn", "stop_sequence":null + })) + .unwrap(), + ))) + }, + ) + .await + .unwrap(); + let MessagesBody::Message(response) = response else { + panic!("expected a message"); + }; assert_eq!(response.id, expected_call.to_string()); assert_eq!( response.content[0]["text"], @@ -253,7 +258,7 @@ async fn provider_reported_cost_is_delivered_once_and_never_charged_for_a_cache_ #[case] cost: Value, #[case] expected: Option, ) { - use litellm_inference_messages::{MessagesCall, MessagesCallResponse}; + use litellm_inference_messages::{MessagesBody, MessagesCall, MessagesCallResponse}; let response = json!({"id": "msg_cost", "type": "message", "role": "assistant", "model": "native-test", "content": [{"type": "text", "text": "answer"}], "stop_reason": "end_turn", "stop_sequence": null, "usage": {"input_tokens": 3, "output_tokens": 2}, "cost": cost}); @@ -277,6 +282,9 @@ async fn provider_reported_cost_is_delivered_once_and_never_charged_for_a_cache_ panic!("expected Messages responses"); }; assert_eq!(first, second); + let MessagesBody::Message(first) = first else { + panic!("expected a message"); + }; assert_eq!(first.extra.get("cost"), Some(&cost)); assert_eq!(support::received(&upstream).await.len(), 1); let facts = hooks.facts.lock().unwrap(); diff --git a/litellm-rust/crates/inference-messages/tests/live/support.rs b/litellm-rust/crates/inference-messages/tests/live/support.rs index 81e295e34ad..52684928f1f 100644 --- a/litellm-rust/crates/inference-messages/tests/live/support.rs +++ b/litellm-rust/crates/inference-messages/tests/live/support.rs @@ -4,7 +4,9 @@ use futures_util::{FutureExt, TryStreamExt}; use litellm_host::interceptors::{ ExecutionFacts, Interceptors, RawResponse, RequestContext, ResultSource, WireRequest, }; -use litellm_inference_messages::{Error, MessagesCall, MessagesCallResponse, MessagesRoute}; +use litellm_inference_messages::{ + Error, MessagesBody, MessagesCall, MessagesCallResponse, MessagesRoute, +}; use litellm_inference_testing::live::{LiveResources, model, within_deadline}; use litellm_llms_types::formats::messages::MessagesResponse; use serde_json::{Value, json}; @@ -111,7 +113,10 @@ pub(super) fn call(provider: &str, payload: Value) -> MessagesCall { pub(super) fn complete(response: MessagesCallResponse) -> Box { match response { - MessagesCallResponse::Complete(message) => message, + MessagesCallResponse::Complete(message) => match message { + MessagesBody::Message(message) => message, + MessagesBody::Unchecked(_) => panic!("a complete request returned no message"), + }, MessagesCallResponse::Stream { .. } => panic!("a complete request returned a stream"), } } diff --git a/litellm-rust/crates/inference-messages/tests/messages/host.rs b/litellm-rust/crates/inference-messages/tests/messages/host.rs index e8df20e7dcb..20f6a97aa37 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/host.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/host.rs @@ -6,7 +6,7 @@ use litellm_host::{ interceptors::{ExecutionFacts, RequestContext, ResultSource, WireRequest}, lifecycle::CallEvent, }; -use litellm_inference_messages::{MessagesCallResponse, route::Messages}; +use litellm_inference_messages::{MessagesBody, MessagesCallResponse, route::Messages}; use litellm_inference_testing::{RecordingSecrets, no_secrets}; use litellm_llms::base_llm::messages::context::{ MessagesModelCapabilities as AnthropicModelCapabilities, SupportedEffortTiers, @@ -340,7 +340,10 @@ async fn response_mode_follows_the_intercepted_request( } MessagesCallResponse::Complete(message) => { assert_eq!(expected_stream, Some(false)); - assert_eq!(*message, serde_json::from_value(message_body()).unwrap()); + assert_eq!( + message, + MessagesBody::Message(serde_json::from_value(message_body()).unwrap()) + ); } } Ok::<_, Error>(()) diff --git a/litellm-rust/crates/inference-messages/tests/messages/main.rs b/litellm-rust/crates/inference-messages/tests/messages/main.rs index a01e571ec75..fe203287766 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/main.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/main.rs @@ -5,7 +5,7 @@ use std::{ use litellm_http::{HttpSettings, Resolution}; use litellm_inference_messages::{ - Error, MessagesCall, MessagesSettings, MessagesShaping, + Error, MessagesBody, MessagesCall, MessagesSettings, MessagesShaping, route::{Messages, MessagesMachine, MessagesOutput}, }; use litellm_inference_testing::RecordingSecrets; @@ -119,7 +119,10 @@ async fn run(call: MessagesCall) -> Result { async fn run_message(call: MessagesCall) -> MessagesResponse { match run(call).await.expect("messages call succeeds") { - MessagesOutput::Complete(message) => *message, + MessagesOutput::Complete(message) => match message { + MessagesBody::Message(message) => *message, + MessagesBody::Unchecked(_) => panic!("a call returned no message"), + }, MessagesOutput::StreamEnded | MessagesOutput::Detached => { panic!("a non-streaming call returned a stream") } diff --git a/litellm-rust/crates/inference-messages/tests/messages/response.rs b/litellm-rust/crates/inference-messages/tests/messages/response.rs index 76813fe0a34..f978da5e995 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/response.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/response.rs @@ -2,7 +2,7 @@ use litellm_host::{ interceptors::{ExecutionFacts, ResultSource}, lifecycle::ExecutionEvent, }; -use litellm_inference_messages::{MessagesCallResponse, messages_body}; +use litellm_inference_messages::{MessagesBody, MessagesCallResponse, messages_body}; use litellm_inference_testing::{ RecordingSecrets, http_config, no_secrets, provider_http, resources, }; @@ -48,6 +48,9 @@ async fn calls_defer_execution_until_polled( let MessagesCallResponse::Complete(response) = future.await.unwrap() else { panic!("expected a completed message"); }; + let MessagesBody::Message(response) = response else { + panic!("expected a message"); + }; assert_eq!( response.content, message_body()["content"].as_array().unwrap().as_slice() @@ -240,10 +243,42 @@ async fn an_unparseable_success_body_is_reported_as_the_upstream_answer(call: Me assert_eq!(error.upstream_response(), Some((200, "not json"))); } +#[rstest] +#[case::error_object( + json!({"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}), + json!({"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}, + "stop_sequence": null}), +)] +#[case::message_missing_id( + json!({"type": "message", "content": [], "stop_sequence": "###"}), + json!({"type": "message", "content": [], "stop_sequence": "###"}), +)] +#[tokio::test] +async fn a_success_body_that_is_not_a_message_passes_through( + call: MessagesCall, + #[case] body: Value, + #[case] expected: Value, +) { + let upstream = upstream([json_response(body)]).await; + + let response = run(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }) + .await + .expect("any JSON object is returned"); + + let MessagesOutput::Complete(response) = response else { + panic!("expected a completed response"); + }; + assert_eq!(serde_json::to_value(response).unwrap(), expected); +} + #[rstest] #[tokio::test] async fn an_unreadable_success_body_is_an_invalid_response(call: MessagesCall) { - let upstream = upstream([json_response(json!({"unexpected": true}))]).await; + let upstream = upstream([json_response(json!(["unexpected"]))]).await; let error = run(MessagesCall { api_key: Some("sk".into()), @@ -320,6 +355,9 @@ async fn the_facade_sends_through_the_injected_http_pool_configuration(call: Mes let MessagesCallResponse::Complete(message) = response else { panic!("a non-streaming request returns a message"); }; + let MessagesBody::Message(message) = message else { + panic!("expected a message"); + }; assert_eq!(message.id, "msg_1"); let sent = only_request(&upstream).await; assert_eq!(sent.header("x-api-key"), Some("sk-ant")); @@ -412,6 +450,9 @@ async fn route_uses_injected_dependencies_and_optional_cache( else { panic!("expected a completed message"); }; + let MessagesBody::Message(response) = response else { + panic!("expected a message"); + }; assert_eq!( response.content, message_body()["content"].as_array().unwrap().as_slice() @@ -484,6 +525,9 @@ async fn cache_overrides_preserve_the_routes_isolated_scope(call: MessagesCall) else { panic!("expected a completed message"); }; + let MessagesBody::Message(response) = response else { + panic!("expected a message"); + }; assert_eq!(response.id, expected["id"].as_str().unwrap()); } assert_eq!(received(&upstream).await.len(), 2); 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 af614f17360..e92daa67be4 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -239,9 +239,9 @@ impl PythonBinding for MessagesPythonHost { fn encode_response( &mut self, py: Python<'_>, - response: Box, + response: litellm_inference_messages::MessagesBody, ) -> PyResult> { - public_response(py, ROUTE_HOST_MODULE, response.as_ref()) + public_response(py, ROUTE_HOST_MODULE, &response) } fn encode_stream_head( diff --git a/tests/test_litellm_rust/messages/test_callbacks.py b/tests/test_litellm_rust/messages/test_callbacks.py index 3073c1cf7dd..0550b530913 100644 --- a/tests/test_litellm_rust/messages/test_callbacks.py +++ b/tests/test_litellm_rust/messages/test_callbacks.py @@ -284,6 +284,23 @@ async def test_pre_call_deployment_rejection_logs_failure_without_provider_io(se assert not {"log_pre_api_call", "log_success_event", "async_log_success_event"} & set(recorder.names) +@pytest.mark.asyncio +async def test_non_string_header_is_a_bad_request_refused_before_the_provider_call(server: RecordingServer) -> None: + recorder: Final = RecordingLogger() + + with rebound(litellm, "callbacks", [recorder]): + with pytest.raises( + litellm.BadRequestError, match=r"invalid extra_headers: header x-org must be a string, got int" + ) as raised: + await call(server, extra_headers={"x-org": 12345}) + failures: Final = await recorder.wait_for_async("async_log_failure_event") + + assert raised.value.status_code == 400 + assert not server.requests + assert "log_pre_api_call" not in recorder.names + assert _OBJECT.validate_python(failures[0].kwargs)["exception"] is raised.value + + @pytest.mark.asyncio async def test_pre_call_deployment_hook_can_delete_keywords(server: RecordingServer) -> None: class DropKeywords(CustomLogger): diff --git a/tests/test_litellm_rust/messages/test_requests.py b/tests/test_litellm_rust/messages/test_requests.py index 92031e72aee..3e486bbe9e6 100644 --- a/tests/test_litellm_rust/messages/test_requests.py +++ b/tests/test_litellm_rust/messages/test_requests.py @@ -4,7 +4,7 @@ import pytest import litellm from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_RESPONSE, call -from tests.test_litellm_rust.support.recording_server import RecordingServer +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec from tests.test_litellm_rust.support.routing import assert_served_natively pytestmark = pytest.mark.requires_rust_extension @@ -57,6 +57,21 @@ async def test_replayed_tool_ids_are_sanitized_before_they_reach_the_provider(se ] +@pytest.mark.asyncio +async def test_a_success_body_that_is_not_a_message_passes_through(server: RecordingServer) -> None: + overloaded: Final = {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}} + server.default_response = ResponseSpec(body=overloaded) + + response: Final = await call(server) + + assert_served_natively(server, response) + assert isinstance(response, dict) + assert {name: value for name, value in response.items() if name != "_hidden_params"} == { + **overloaded, + "stop_sequence": None, + } + + @pytest.mark.asyncio async def test_provider_specific_headers_scoped_to_anthropic_reach_the_wire(server: RecordingServer) -> None: await call( From 1467b13fd2a777b0824fe3eb1140f117d733e5cf Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Mon, 9 Nov 2026 16:24:20 -0800 Subject: [PATCH 35/46] fix(rust): run native async calls through python's async wrapper steps (#50423) Python's async client wrapper stamps the effective timeout and num_retries on every failure, which the proxy reads for x-litellm-timeout and the router reads for retries. The native async paths skipped that wrapper, so error responses from Rust-served async calls had no x-litellm-timeout header Python's async chat completions, responses, embeddings and messages entries hand the call to run_in_executor before they build the provider body, so a during_call guardrail started next to the call masks the request first. The native paths built their request without yielding, so masked text never reached the provider The shared native_acall_hook now yields once and stamps the failure, and every async dispatch passes it as its native hook. OCR passes native_acall_hook_on_loop, which only stamps the failure, since its Python entry went through the client wrapper but built its request on the loop --- litellm/chat_completions/dispatch.py | 3 +- litellm/embeddings/dispatch.py | 3 +- litellm/messages/dispatch.py | 3 +- litellm/ocr/dispatch.py | 3 +- litellm/responses/dispatch.py | 3 +- litellm/rust_bridge/public_call.py | 36 +++++++++- .../test_litellm_rust/test_public_entries.py | 36 +++++++++- tests/unit/rust_bridge/test_public_call.py | 70 ++++++++++++++++++- 8 files changed, 147 insertions(+), 10 deletions(-) diff --git a/litellm/chat_completions/dispatch.py b/litellm/chat_completions/dispatch.py index 06ba8e9ad76..05572ee8837 100644 --- a/litellm/chat_completions/dispatch.py +++ b/litellm/chat_completions/dispatch.py @@ -16,6 +16,7 @@ from litellm.rust_bridge.public_call import ( AsyncFlag, NativeCall, bind, + native_acall_hook, native_call, native_call_hook, optional_sequence, @@ -131,7 +132,7 @@ async def acompletion(*args: object, **kwargs: object) -> ChatResult: # kwargs- _FLAG.dropped(kwargs), python=python, binding=NATIVE_ACOMPLETION, - native=native_call_hook, + native=native_acall_hook, ) diff --git a/litellm/embeddings/dispatch.py b/litellm/embeddings/dispatch.py index 772ac68737d..9547757ff8f 100644 --- a/litellm/embeddings/dispatch.py +++ b/litellm/embeddings/dispatch.py @@ -13,6 +13,7 @@ from litellm.rust_bridge.public_call import ( AsyncFlag, NativeCall, bind, + native_acall_hook, native_call, native_call_hook, optional_str, @@ -89,7 +90,7 @@ async def aembedding(*args: object, **kwargs: object) -> EmbeddingResponse: # k _FLAG.dropped(kwargs), python=_PYTHON_AEMBEDDING, binding=NATIVE_AEMBEDDING, - native=native_call_hook, + native=native_acall_hook, ) diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index 40d21c5bad9..faa604a25ce 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -19,6 +19,7 @@ from litellm.rust_bridge.public_call import ( AsyncFlag, NativeCall, bind, + native_acall_hook, native_call, native_call_hook, optional_sequence, @@ -207,7 +208,7 @@ async def anthropic_messages(*args: object, **kwargs: object) -> MessagesResult: _FLAG.dropped(kwargs), python=python, binding=NATIVE_AMESSAGES, - native=native_call_hook, + native=native_acall_hook, ) diff --git a/litellm/ocr/dispatch.py b/litellm/ocr/dispatch.py index 529c3df2bb9..f90a42bb1c9 100644 --- a/litellm/ocr/dispatch.py +++ b/litellm/ocr/dispatch.py @@ -12,6 +12,7 @@ from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR from litellm.rust_bridge.public_call import ( AsyncFlag, NativeCall, + native_acall_hook_on_loop, native_call, native_call_hook, optional_str, @@ -97,5 +98,5 @@ async def aocr(*args: object, **kwargs: object) -> OCRResponse: # kwargs-ok: pr _FLAG.dropped(kwargs), python=runtime.NO_PYTHON, binding=NATIVE_AOCR, - native=native_call_hook, + native=native_acall_hook_on_loop, ) diff --git a/litellm/responses/dispatch.py b/litellm/responses/dispatch.py index 0dda3f93660..acf6610d06e 100644 --- a/litellm/responses/dispatch.py +++ b/litellm/responses/dispatch.py @@ -14,6 +14,7 @@ from litellm.rust_bridge.public_call import ( AsyncFlag, NativeCall, bind, + native_acall_hook, native_call, native_call_hook, optional_str, @@ -150,7 +151,7 @@ async def aresponses(*args: object, **kwargs: object) -> ResponsesResult: # kwa _FLAG.dropped(kwargs), python=python, binding=NATIVE_ARESPONSES, - native=native_call_hook, + native=native_acall_hook, ) diff --git a/litellm/rust_bridge/public_call.py b/litellm/rust_bridge/public_call.py index 8625ca71bb2..1ab3d5e10c9 100644 --- a/litellm/rust_bridge/public_call.py +++ b/litellm/rust_bridge/public_call.py @@ -2,12 +2,17 @@ from __future__ import annotations +import asyncio import inspect -from collections.abc import Callable, Collection, Mapping, Sequence +from collections.abc import Awaitable, Callable, Collection, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import Final, TypeVar, cast # noqa: TID251 # narrows caller-owned containers without copying them +from litellm.utils import ( + _get_wrapper_timeout, # pyright: ignore[reportPrivateUsage] # the @client wrapper's own failure stamp +) + def signature(legacy: Callable[..., object]) -> inspect.Signature: return inspect.signature(legacy) @@ -115,3 +120,32 @@ def native_call_hook( _kwargs: Mapping[str, object], ) -> NativeResultT: return hook(call) + + +async def native_acall_hook( + hook: Callable[[NativeCall], Awaitable[NativeResultT]], + call: NativeCall, + args: tuple[object, ...], + kwargs: Mapping[str, object], +) -> NativeResultT: + # Python's async entries leave the event loop (run_in_executor) before they build the provider body, so + # tasks scheduled next to the call, such as the proxy's during_call guardrails, edit the request first + await asyncio.sleep(0) + return await native_acall_hook_on_loop(hook, call, args, kwargs) + + +async def native_acall_hook_on_loop( + hook: Callable[[NativeCall], Awaitable[NativeResultT]], + call: NativeCall, + _args: tuple[object, ...], + kwargs: Mapping[str, object], +) -> NativeResultT: + """``native_acall_hook`` for an async entry whose Python path built its request on the loop, as OCR's did.""" + try: + return await hook(call) + except Exception as error: + # utils.client stamps every async failure, which the proxy reads for x-litellm-timeout and the router for retries + if kwargs.get("num_retries") is not None: + setattr(error, "num_retries", kwargs.get("num_retries")) # noqa: B010 # Exception declares no such attribute, as in utils.client + setattr(error, "timeout", _get_wrapper_timeout(kwargs=dict(kwargs))) # noqa: B010 # Exception declares no such attribute, as in utils.client + raise diff --git a/tests/test_litellm_rust/test_public_entries.py b/tests/test_litellm_rust/test_public_entries.py index 2f5c33e3ced..24a3419c5a0 100644 --- a/tests/test_litellm_rust/test_public_entries.py +++ b/tests/test_litellm_rust/test_public_entries.py @@ -1,3 +1,4 @@ +import asyncio import inspect from collections.abc import Awaitable, Callable, Iterator, Mapping from contextlib import contextmanager @@ -29,7 +30,10 @@ Served: TypeAlias = Callable[[object], bool] @dataclass(frozen=True, slots=True) class Entry: - """``python`` is ``None`` for a route with no Python implementation, which only RUST_REQUIRED rules serve.""" + """``python`` is ``None`` for a route with no Python implementation, which only RUST_REQUIRED rules serve. + + ``executor_hop`` is whether the route's Python async entry left the loop through run_in_executor before it built the request. + """ route: Route sync: Callable[..., object] @@ -40,6 +44,7 @@ class Entry: arguments: Mapping[str, object] positional: int python: Served | None + executor_hop: bool def args(self) -> tuple[object, ...]: return tuple(self.arguments.values())[: self.positional] @@ -79,6 +84,7 @@ ENTRIES: Final = ( arguments={"model": "openai/gpt-4o", "messages": MESSAGES, "mock_response": "python"}, positional=2, python=_chat, + executor_hop=True, ), Entry( route=Route.EMBEDDINGS, @@ -90,6 +96,7 @@ ENTRIES: Final = ( arguments={"model": "openai/text-embedding-3-small", "input": ["hi"], "mock_response": [0.5]}, positional=2, python=_embedding, + executor_hop=True, ), Entry( route=Route.RESPONSES, @@ -101,6 +108,7 @@ ENTRIES: Final = ( arguments={"input": "hi", "model": "openai/gpt-4o", "mock_response": "python"}, positional=2, python=_responses, + executor_hop=True, ), Entry( route=Route.MESSAGES, @@ -117,6 +125,7 @@ ENTRIES: Final = ( }, positional=0, python=_messages, + executor_hop=True, ), Entry( route=Route.OCR, @@ -128,6 +137,7 @@ ENTRIES: Final = ( arguments={"model": "mistral/mistral-ocr-latest", "document": {"type": "file", "file": b"pdf"}}, positional=2, python=None, + executor_hop=False, ), ) @@ -285,3 +295,27 @@ async def test_sync_entry_with_flag_true_raises_setup_errors_at_await(entry: Ent assert inspect.iscoroutine(pending) with pytest.raises(RuntimeError, match="is unavailable"): await pending + + +@pytest.mark.parametrize("entry", ENTRIES, ids=entry_id) +async def test_async_entry_runs_native_through_the_client_wrapper_steps(entry: Entry) -> None: + scheduled: Final[list[str]] = [] + + async def failing(native_call: NativeCall) -> ModelResponse: + scheduled.append("native") + raise litellm.Timeout(message="timed out", model="model", llm_provider="openai") + + async def started_before() -> None: + scheduled.append("task") + + entry.async_binding.override(failing) + task: Final = asyncio.create_task(started_before()) + try: + with scenario_rules(entry, SCENARIOS["required"]), pytest.raises(litellm.Timeout) as raised: + await entry.asynchronous(*entry.args(), **entry.kwargs(timeout=30.0, num_retries=0)) + finally: + entry.async_binding.reset() + await task + + assert (getattr(raised.value, "timeout", None), getattr(raised.value, "num_retries", None)) == (30.0, 0) + assert scheduled == (["task", "native"] if entry.executor_hop else ["native", "task"]) diff --git a/tests/unit/rust_bridge/test_public_call.py b/tests/unit/rust_bridge/test_public_call.py index 5f50db6ea5a..27c3c119c1b 100644 --- a/tests/unit/rust_bridge/test_public_call.py +++ b/tests/unit/rust_bridge/test_public_call.py @@ -1,9 +1,22 @@ -from collections.abc import Mapping, Sequence -from typing import Final +import asyncio +from collections.abc import Awaitable, Callable, Mapping, Sequence +from typing import Final, TypeAlias import pytest -from litellm.rust_bridge.public_call import AsyncFlag, native_call, signature +import litellm +from litellm.rust_bridge.public_call import ( + AsyncFlag, + NativeCall, + native_acall_hook, + native_acall_hook_on_loop, + native_call, + signature, +) + +NativeAcallHook: TypeAlias = Callable[ + [Callable[[NativeCall], Awaitable[None]], NativeCall, tuple[object, ...], Mapping[str, object]], Awaitable[None] +] def _messages( @@ -76,3 +89,54 @@ def test_flag_keywords_name_positionals_and_drop_only_the_flag() -> None: def test_flag_keywords_reject_a_parameter_given_twice() -> None: with pytest.raises(TypeError): _FLAG.keywords((128,), {"max_tokens": 64}) + + +_CALL: Final = native_call(_SIGNATURE, (128, [], "model"), {}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("kwargs", "expected_timeout", "expected_num_retries"), + ( + ({"timeout": 30.0, "stream_timeout": 0.5, "num_retries": 0}, 30.0, 0), + ({"stream": True, "timeout": 30.0, "stream_timeout": 0.5}, 0.5, None), + ({"stream": True, "timeout": 30.0}, 30.0, None), + ({}, None, None), + ), +) +@pytest.mark.parametrize("hook", (native_acall_hook, native_acall_hook_on_loop)) +async def test_async_hook_failures_carry_the_client_wrapper_stamp( + hook: NativeAcallHook, + kwargs: Mapping[str, object], + expected_timeout: float | None, + expected_num_retries: int | None, +) -> None: + async def native(call: NativeCall) -> None: + raise litellm.Timeout(message="timed out", model="model", llm_provider="anthropic") + + with pytest.raises(litellm.Timeout) as raised: + await hook(native, _CALL, (), kwargs) + + assert getattr(raised.value, "timeout", "unset") == expected_timeout + assert getattr(raised.value, "num_retries", None) == expected_num_retries + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("hook", "yields"), ((native_acall_hook, True), (native_acall_hook_on_loop, False))) +async def test_only_the_executor_hook_lets_a_task_started_before_the_call_run_first( + hook: NativeAcallHook, yields: bool +) -> None: + messages: Final = [{"role": "user", "content": "during about secret"}] + sent: Final[list[object]] = [] + + async def mask() -> None: + messages[0]["content"] = "during about [REDACTED]" + + async def native(call: NativeCall) -> None: + sent.append(messages[0]["content"]) + + guardrail: Final = asyncio.create_task(mask()) + await hook(native, _CALL, (), {}) + await guardrail + + assert sent == ["during about [REDACTED]" if yields else "during about secret"] From ee34a3c2d4ce425ac532cc344799e5c4bdabaaff Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Mon, 9 Nov 2026 16:24:20 -0800 Subject: [PATCH 36/46] fix(messages): raise native messages failures as python's errors (#50424) A stream whose upstream body ends mid-frame reached proxy callers as a mapped APIConnectionError with status 500. Python's transport raises an unmapped httpx error at that point. The native body read now records a body cut and the host raises the same exception the active transport would Python wraps a connect, DNS, TLS or reset failure in a synthesized 500 AnthropicError, which maps to InternalServerError. The native path surfaced these as APIConnectionError with the bug-report footer, which also changed cooldown decisions The proxy bills the estimated prompt of a failed call only when the Timeout was caused by an httpx ReadTimeout. The native timeout now chains the same httpx timeout Python's client raises, so its failure row bills the prompt Python rebuilds an upstream error response from the decoded body without content-encoding and content-length. The native failure now drops those wire headers, and a failure raised without an HTTP response states no headers at all, like Python --- litellm-rust/crates/http/src/error.rs | 104 ++++++++++ litellm-rust/crates/http/src/lib.rs | 2 +- .../crates/inference-messages/src/handler.rs | 21 ++- litellm-rust/crates/inference/src/error.rs | 8 + litellm-rust/crates/llms/src/error.rs | 7 + .../crates/python-bridge/src/errors.rs | 5 + .../host/callbacks_legacy_python.py | 6 + litellm/rust_bridge/host/routes/messages.py | 133 ++++++++++++- .../rust_bridge/host/routes/test_messages.py | 178 ++++++++++++++++++ .../host/test_callbacks_legacy_python.py | 24 +++ 10 files changed, 475 insertions(+), 13 deletions(-) diff --git a/litellm-rust/crates/http/src/error.rs b/litellm-rust/crates/http/src/error.rs index 7078d4aba2a..9505aacf44b 100644 --- a/litellm-rust/crates/http/src/error.rs +++ b/litellm-rust/crates/http/src/error.rs @@ -16,6 +16,15 @@ pub enum TlsSource { #[error(transparent)] pub struct Error(Arc); +// A streamed body the connection ended early, after the response head arrived; `expected` is +// the declared `content-length`, `None` for a chunked body. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct BodyCut { + pub received: u64, + pub expected: Option, + pub reset: bool, +} + #[derive(Debug, thiserror::Error)] pub(crate) enum Repr { #[error("could not read {}", path.display())] @@ -56,6 +65,12 @@ pub(crate) enum Repr { Connect(#[source] reqwest::Error), #[error("upstream network error")] Network(#[source] reqwest::Error), + #[error("upstream network error")] + BodyCut { + #[source] + source: reqwest::Error, + cut: BodyCut, + }, #[error("upstream request timed out")] Timeout(#[source] reqwest::Error), /// No bytes arrived within the read timeout, which the caller enforces per read; `body` says @@ -201,6 +216,37 @@ impl Error { matches!(&*self.0, Repr::ReadTimeout { body: true, .. }) } + pub fn body_read(error: reqwest::Error, received: u64, expected: Option) -> Self { + let mut source: Option<&(dyn std::error::Error + 'static)> = Some(&error); + let kind = std::iter::from_fn(|| { + let current = source?; + source = current.source(); + Some(current) + }) + .find_map(|cause| cause.downcast_ref::().map(io::Error::kind)); + let reset = match kind { + Some(io::ErrorKind::UnexpectedEof) => false, + Some(io::ErrorKind::ConnectionReset) => true, + _ => return error.into(), + }; + Repr::BodyCut { + source: error.without_url(), + cut: BodyCut { + received, + expected, + reset, + }, + } + .into() + } + + pub fn body_cut(&self) -> Option { + match &*self.0 { + Repr::BodyCut { cut, .. } => Some(*cut), + _ => None, + } + } + pub fn computed_header(name: impl Into) -> Self { Repr::ComputedHeader(name.into()).into() } @@ -226,6 +272,7 @@ impl Error { pub fn dispatched(&self) -> bool { match &*self.0 { Repr::Network(_) + | Repr::BodyCut { .. } | Repr::Timeout(_) | Repr::ReadTimeout { .. } | Repr::Status { .. } @@ -285,6 +332,7 @@ impl Error { | Repr::Signature(_) | Repr::Connect(_) | Repr::Network(_) + | Repr::BodyCut { .. } | Repr::Timeout(_) | Repr::ReadTimeout { .. } | Repr::Status { .. } @@ -318,6 +366,7 @@ impl Error { | Repr::Signature(_) | Repr::Connect(_) | Repr::Network(_) + | Repr::BodyCut { .. } | Repr::Timeout(_) | Repr::ReadTimeout { .. } | Repr::Status { .. } @@ -355,6 +404,7 @@ impl Error { | Repr::Signature(_) | Repr::Connect(_) | Repr::Network(_) + | Repr::BodyCut { .. } | Repr::Timeout(_) | Repr::ReadTimeout { .. } | Repr::Status { .. } @@ -394,6 +444,7 @@ impl Error { | Repr::UnsignableHeader(_) | Repr::Signature(_) | Repr::Network(_) + | Repr::BodyCut { .. } | Repr::Timeout(_) | Repr::ReadTimeout { .. } | Repr::Status { .. } @@ -437,6 +488,7 @@ impl Error { | Repr::Signature(_) | Repr::Connect(_) | Repr::Network(_) + | Repr::BodyCut { .. } | Repr::Status { .. } | Repr::WebSocketUrl(_) | Repr::WebSocketScheme(_) @@ -473,6 +525,7 @@ impl Error { | Repr::Signature(_) | Repr::Connect(_) | Repr::Network(_) + | Repr::BodyCut { .. } | Repr::Timeout(_) | Repr::ReadTimeout { .. } | Repr::WebSocketUrl(_) @@ -513,6 +566,7 @@ impl Error { | Repr::Signature(_) | Repr::Connect(_) | Repr::Network(_) + | Repr::BodyCut { .. } | Repr::WebSocketUrl(_) | Repr::WebSocketScheme(_) | Repr::WebSocketSchemeRewrite(_) @@ -548,6 +602,7 @@ impl Error { | Repr::Signature(_) | Repr::Connect(_) | Repr::Network(_) + | Repr::BodyCut { .. } | Repr::Timeout(_) | Repr::ReadTimeout { .. } | Repr::Status { .. } @@ -742,4 +797,53 @@ mod tests { assert_eq!(error.upstream_response(), Some((429, body.as_str()))); assert_eq!(error.status(), Some(429)); } + + #[rstest] + #[case::chunked("transfer-encoding: chunked\r\n\r\n5\r\nhello\r\n", None)] + #[case::content_length("content-length: 100\r\n\r\nhello", Some(100))] + #[tokio::test] + async fn a_connection_that_ends_mid_body_is_a_body_cut( + #[case] tail: &'static str, + #[case] expected: Option, + ) { + use std::io::{Read as _, Write as _}; + + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind"); + let address = listener.local_addr().expect("address"); + let server = std::thread::spawn(move || { + let (mut stream, _) = listener.accept().expect("accept"); + let mut request = [0_u8; 1024]; + let _ = stream.read(&mut request).expect("request"); + stream + .write_all(format!("HTTP/1.1 200 OK\r\n{tail}").as_bytes()) + .expect("response"); + }); + let mut response = reqwest::Client::builder() + .no_proxy() + .build() + .expect("client") + .get(format!("http://{address}/")) + .send() + .await + .expect("head"); + server.join().expect("server"); + let mut received = 0_u64; + let error = loop { + match response.chunk().await { + Ok(Some(chunk)) => received += chunk.len() as u64, + Ok(None) => panic!("the body must end with an error"), + Err(error) => break crate::Error::body_read(error, received, expected), + } + }; + assert_eq!( + error.body_cut(), + Some(super::BodyCut { + received: 5, + expected, + reset: false, + }) + ); + assert!(error.dispatched()); + assert_eq!(error.to_string(), "upstream network error"); + } } diff --git a/litellm-rust/crates/http/src/lib.rs b/litellm-rust/crates/http/src/lib.rs index 02129c62083..1aa11843b46 100644 --- a/litellm-rust/crates/http/src/lib.rs +++ b/litellm-rust/crates/http/src/lib.rs @@ -18,7 +18,7 @@ pub mod websocket; pub use client::Client; pub use config::{ClientIdentity, HttpClientConfig, Resolution, Verify}; -pub use error::{Error, TlsSource}; +pub use error::{BodyCut, Error, TlsSource}; pub use pool::{ClientVariant, HttpClientPool}; pub use proxy::EnvironmentProxies; pub use settings::{HttpSettings, HttpSettingsLayer, SslVerify, TcpKeepalive}; diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index 5eb017ddfc5..6a8ed7fff8e 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -386,13 +386,20 @@ fn timed_chunks( response: reqwest::Response, timeout: Duration, ) -> impl futures_util::Stream> + Send + 'static { - futures_util::stream::try_unfold(response, move |mut response| async move { - let chunk = tokio::time::timeout(timeout, response.chunk()) - .await - .map_err(litellm_http::Error::body_read_timeout)? - .map_err(litellm_http::Error::from)?; - Ok(chunk.map(|chunk| (chunk, response))) - }) + let expected = response.content_length(); + futures_util::stream::try_unfold( + (response, 0_u64), + move |(mut response, received)| async move { + let chunk = tokio::time::timeout(timeout, response.chunk()) + .await + .map_err(litellm_http::Error::body_read_timeout)? + .map_err(|error| litellm_http::Error::body_read(error, received, expected))?; + Ok(chunk.map(|chunk| { + let received = received + chunk.len() as u64; + (chunk, (response, received)) + })) + }, + ) } /// The provider's error, its body read under `timeout`: a stream has no overall deadline once diff --git a/litellm-rust/crates/inference/src/error.rs b/litellm-rust/crates/inference/src/error.rs index 841e0c53799..4135424d6ef 100644 --- a/litellm-rust/crates/inference/src/error.rs +++ b/litellm-rust/crates/inference/src/error.rs @@ -164,6 +164,14 @@ impl Error { } } + pub fn body_cut(&self) -> Option { + match self { + Self::Llm(error) => error.body_cut(), + Self::Http(error) => error.body_cut(), + _ => None, + } + } + /// `Some(connect)` when the call timed out, where `connect` says the connection was never /// established; `None` for every other error. pub fn timeout(&self) -> Option { diff --git a/litellm-rust/crates/llms/src/error.rs b/litellm-rust/crates/llms/src/error.rs index 2844fbf3e01..2975ceee3cb 100644 --- a/litellm-rust/crates/llms/src/error.rs +++ b/litellm-rust/crates/llms/src/error.rs @@ -103,6 +103,13 @@ impl Error { matches!(self, Self::Http(error) if error.is_body_read_timeout()) } + pub fn body_cut(&self) -> Option { + match self { + Self::Http(error) => error.body_cut(), + _ => None, + } + } + pub fn is_timeout(&self) -> Option { match self { Self::Http(error) if error.is_timeout() => Some(error.is_connect_timeout()), diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index 426cd861729..db5696d52ce 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -29,6 +29,8 @@ const MESSAGES_UNSUPPORTED_PARAM: &str = "messages_unsupported_param"; /// The phase label (`Connection` or `Read`) of a timed-out call, for Python's timeout text, or /// `Body` for a streamed body read that timed out after the head arrived. const MESSAGES_TIMEOUT: &str = "messages_timeout"; +// `(received, expected, reset)` of a streamed body the connection ended early. +const MESSAGES_BODY_CUT: &str = "messages_body_cut"; /// No credential reached the provider, which Python raises as `AuthenticationError`. const MESSAGES_MISSING_CREDENTIAL: &str = "messages_missing_credential"; /// Python refuses the request only once its logger knows the resolved model and provider (the @@ -134,6 +136,9 @@ pub(crate) fn route_failure( }; value.setattr(MESSAGES_TIMEOUT, phase)?; } + if let Some(cut) = error.body_cut() { + value.setattr(MESSAGES_BODY_CUT, (cut.received, cut.expected, cut.reset))?; + } if error.is_missing_credential() { value.setattr(MESSAGES_MISSING_CREDENTIAL, true)?; } diff --git a/litellm/rust_bridge/host/callbacks_legacy_python.py b/litellm/rust_bridge/host/callbacks_legacy_python.py index 8c281c73842..5a368573f3d 100644 --- a/litellm/rust_bridge/host/callbacks_legacy_python.py +++ b/litellm/rust_bridge/host/callbacks_legacy_python.py @@ -510,9 +510,15 @@ def stream_failure( chunks: list[bytes], error: Exception, ) -> Coroutine[object, object, None]: + from litellm._logging import verbose_proxy_logger from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType + verbose_proxy_logger.error("Error in chunk_processor: %s", error) + # Python's stream fails before its success handler records `cache_hit`; only a cache hit + # stamps it earlier. + if logger.model_call_details.get("cache_hit") is not True: + logger.model_call_details.pop("cache_hit", None) return PassThroughStreamingHandler.schedule_stream_failure_logging( litellm_logging_obj=logger, endpoint_type=EndpointType(endpoint_type), diff --git a/litellm/rust_bridge/host/routes/messages.py b/litellm/rust_bridge/host/routes/messages.py index c8b53097f46..c50791f8a9c 100644 --- a/litellm/rust_bridge/host/routes/messages.py +++ b/litellm/rust_bridge/host/routes/messages.py @@ -1,11 +1,16 @@ from __future__ import annotations import datetime +import errno +import ipaddress +import os +import socket from collections.abc import Callable, Mapping, Sequence from dataclasses import asdict, dataclass from typing import Final, TypeAlias, cast # noqa: TID251 # narrows native payloads to public TypedDicts import httpx +from aiohttp.http_exceptions import ContentLengthError, TransferEncodingError from pydantic import TypeAdapter, ValidationError import litellm @@ -14,7 +19,7 @@ from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.reasoning_effort_capability import deployment_reasoning_effort_levels from litellm.litellm_core_utils.request_timeout_resolver import get_configured_request_timeout -from litellm.llms.anthropic.common_utils import AnthropicError +from litellm.llms.anthropic.common_utils import AnthropicError, anthropic_pass_through_api_base from litellm.llms.anthropic.messages.native import handler from litellm.llms.anthropic.messages.native.handler import ( _native_messages_config as native_messages_config, # pyright: ignore[reportPrivateUsage] # the handler's own config lookup @@ -23,8 +28,10 @@ from litellm.llms.anthropic.messages.native.handler import ( _with_deployment_output_effort as _deployment_output_effort, # pyright: ignore[reportPrivateUsage] # the handler's own effort rounding ) from litellm.llms.anthropic.messages.utils import is_reasoning_auto_summary_enabled +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.custom_httpx.http_handler import get_ssl_configuration from litellm.rust_bridge import failures -from litellm.rust_bridge.host import callbacks_legacy_python +from litellm.rust_bridge.host import callbacks_legacy_python, http_transport from litellm.rust_bridge.public_call import known_params, optional_str from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse @@ -120,6 +127,7 @@ def stream_hidden_params(headers: Sequence[tuple[str, str]]) -> Mapping[str, obj _UPSTREAM_ARGS: Final = TypeAdapter(tuple[int, str]) _UPSTREAM_HEADERS: Final = TypeAdapter(list[tuple[str, str]]) +_BODY_CUT: Final = TypeAdapter(tuple[int, int | None, bool]) _ERROR_CLASS: TypeAlias = Callable[..., Exception] _validate_metadata: Final = cast( # cast-ok: the handler's untyped metadata validator Callable[[dict[str, object]], object], handler.validate_anthropic_api_metadata @@ -128,6 +136,17 @@ _exception_type: Final = cast( # cast-ok: the legacy public exception mapper is Callable[..., Exception], exception_mapping_utils.exception_type ) _BODY_READ_TIMEOUT: Final = "Timeout on reading data from socket" # aiohttp's SocketTimeoutError text +_UNREACHABLE: Final = "could not reach the provider: " +_SEND_FAILED: Final = "upstream network error: error sending request: client error (SendRequest): " +_DNS_FAILURE: Final = "failed to lookup address information: " +_HTTPX_DNS_ERRNO: Final = { + "Name or service not known": socket.EAI_NONAME, + "Temporary failure in name resolution": socket.EAI_AGAIN, +} +_TRANSPORT_TIMEOUTS: Final[Mapping[object, type[httpx.TimeoutException]]] = { + "Connection": httpx.ConnectTimeout, + "Read": httpx.ReadTimeout, +} _MISSING_KEY: Final = ( "Missing Anthropic API Key - A call is being made to anthropic but no key is set " "either in the environment variables or via params. Please set `ANTHROPIC_API_KEY` " @@ -174,6 +193,82 @@ def _timeout_text(phase: str, request: Mapping[str, object]) -> str: ) +def _body_cut(received: int, expected: int | None, reset: bool) -> httpx.TransportError: + if not http_transport.uses_aiohttp_transport(): + if reset: + return httpx.ReadError("") + detail: Final = ( + "incomplete chunked read" if expected is None else f"received {received} bytes, expected {expected}" + ) + return httpx.RemoteProtocolError(f"peer closed connection without sending complete message body ({detail})") + payload: Final = ( + TransferEncodingError("Not enough data to satisfy transfer length header.") + if expected is None + else ContentLengthError( + f"Not enough data to satisfy content length header (received {received} of {expected} bytes)." + ) + ) + text: Final = f"Response payload is not completed: {payload!r}" + if not reset: + return httpx.ReadError(text) + return httpx.ReadError(f"{text}. {ConnectionResetError(errno.ECONNRESET, os.strerror(errno.ECONNRESET))!r}") + + +def _connect_failure(native: str, request: Mapping[str, object], request_provider: str) -> str | None: + api_base: Final = request.get("api_base") + if not isinstance(api_base, str) and request_provider != "anthropic": + return None + refused: Final = "tcp connect error: Connection refused" in native + dns_reason: Final = native.rsplit(_DNS_FAILURE, 1)[1] if _DNS_FAILURE in native else None + tls_eof: Final = native.endswith("tls handshake eof") + if not http_transport.uses_aiohttp_transport(): + if refused: + return "All connection attempts failed" + dns_errno: Final = None if dns_reason is None else _HTTPX_DNS_ERRNO.get(dns_reason) + if dns_errno is not None: + return f"[Errno {dns_errno}] {dns_reason}" + return "" if tls_eof else None + url: Final = httpx.URL(api_base if isinstance(api_base, str) else anthropic_pass_through_api_base()) + port: Final = url.port or (443 if url.scheme == "https" else 80) + if refused: + try: + address: Final = ipaddress.ip_address(url.host) + except ValueError: + return None + peer: Final = (str(address), port) if address.version == 4 else (str(address), port, 0, 0) + return _aiohttp_connect_failure(url, port, request, f"Connect call failed {peer}") + if dns_reason is not None: + return _aiohttp_connect_failure(url, port, request, dns_reason) + if tls_eof: + return _aiohttp_connect_failure(url, port, request, None) + return None + + +def _aiohttp_connect_failure(url: httpx.URL, port: int, request: Mapping[str, object], reason: str | None) -> str: + ssl_verify: Final = request.get("ssl_verify") + ssl_config: Final = get_ssl_configuration(ssl_verify if isinstance(ssl_verify, bool | str) else None) + return f"Cannot connect to host {url.host}:{port} ssl:{'default' if ssl_config is True else ssl_config} [{reason}]" + + +def _transport_failure(error: Exception, request: Mapping[str, object], request_provider: str) -> str | None: + """The text Python's HTTP client raises for a call that failed before any response arrived.""" + native: Final = str(error) + if native.startswith(_SEND_FAILED): + reset: Final = native.endswith(f"Connection reset by peer (os error {errno.ECONNRESET})") + closed: Final = native.endswith("connection closed before message completed") + if http_transport.uses_aiohttp_transport(): + if reset: + return str(ConnectionResetError(errno.ECONNRESET, os.strerror(errno.ECONNRESET))) + return "Server disconnected" if closed else native + if reset: + return "" + return "Server disconnected without sending a response." if closed else native + if not native.startswith(_UNREACHABLE): + return None + connect: Final = _connect_failure(native, request, request_provider) + return native if connect is None else connect + + def _upstream_failure( error: Exception, request: Mapping[str, object] ) -> tuple[int, str, Sequence[tuple[str, str]]] | None: @@ -183,14 +278,27 @@ def _upstream_failure( return 408, _timeout_text(timeout_phase, request), () try: status, text = _UPSTREAM_ARGS.validate_python(error.args) - return status, text, _UPSTREAM_HEADERS.validate_python(getattr(error, "headers", None)) + headers: Final = _UPSTREAM_HEADERS.validate_python(getattr(error, "headers", None)) except ValidationError: return None + if not headers: + return status, text, headers + # Python's `MaskedHTTPStatusError` rebuilds the response from its decoded body without the + # wire's content-encoding and content-length, so httpx states a length only for a non-empty body. + rebuilt: Final = httpx.Response( + status, + content=text.encode(), + headers=[ + (name, value) for name, value in headers if name.lower() not in ("content-encoding", "content-length") + ], + ) + return status, text, tuple(rebuilt.headers.items()) def _provider_error(error: Exception, request: Mapping[str, object], request_provider: str) -> Exception | None: """The provider exception Python's `_handle_error` builds for the same failure.""" - failure: Final = _upstream_failure(error, request) + transport: Final = _transport_failure(error, request, request_provider) + failure: Final = _upstream_failure(error, request) if transport is None else (500, transport, ()) if failure is None: return None status, text, headers = failure @@ -201,13 +309,17 @@ def _provider_error(error: Exception, request: Mapping[str, object], request_pro request.get("model_info"), api_base if isinstance(api_base, str) else None, ) - return ( + provider_error: Final = ( cast(_ERROR_CLASS, config.get_error_class)( # cast-ok: the config's untyped headers parameter error_message=text, status_code=status, headers=dict(headers) ) if config is not None else AnthropicError(status_code=status, message=text, headers=httpx.Headers(headers)) ) + if transport is not None and isinstance(provider_error, BaseLLMException): + # Python's `_handle_error` gives a failure without a response the status 500. + provider_error.status_code_is_synthesized = True + return provider_error def _log_resolved_call(request: Mapping[str, object], request_provider: str) -> None: @@ -232,6 +344,17 @@ def map_failure(error: Exception, request: Mapping[str, object], request_provide # Python reads a stream's body outside its HTTP client's error mapping, so the transport's # own timeout reaches the caller unmapped. return httpx.ReadTimeout(_BODY_READ_TIMEOUT) + body_cut: Final = getattr(error, "messages_body_cut", None) + if body_cut is not None: + # Python reads a stream's body outside its HTTP client's error mapping, so a body the + # connection cut short reaches the caller as its transport raised it. + return _body_cut(*_BODY_CUT.validate_python(body_cut)) + transport_timeout: Final = _TRANSPORT_TIMEOUTS.get(getattr(error, "messages_timeout", None)) + if transport_timeout is not None: + # Python's `litellm.Timeout` is raised while handling httpx's own timeout, which the proxy's + # failure spend reads to bill a read timeout's prompt. + chained: Final = transport_timeout(str(error)) + error.__context__ = chained # rebind-ok: Python chains the transport timeout as the exception context if getattr(error, "messages_late_rejection", False): _log_resolved_call(request, request_provider) if getattr(error, "messages_request_error", False): diff --git a/tests/unit/rust_bridge/host/routes/test_messages.py b/tests/unit/rust_bridge/host/routes/test_messages.py index 9d8b09dd1c4..704d52386a8 100644 --- a/tests/unit/rust_bridge/host/routes/test_messages.py +++ b/tests/unit/rust_bridge/host/routes/test_messages.py @@ -207,6 +207,16 @@ def test_a_native_read_timeout_maps_to_the_public_timeout_with_python_text() -> assert "Read timed out. Timeout passed=1.0" in str(mapped) +@pytest.mark.parametrize(("phase", "billed"), (("Read", True), ("Connection", False))) +def test_a_native_timeout_is_billed_on_the_failure_row_as_python_bills_it(phase: str, billed: bool) -> None: + from litellm.proxy.utils import _provider_billed_the_failure + + mapped: Final = messages_host.map_failure(_marked(phase), {"model": _ADAPTIVE_OPUS, "timeout": 1.0}, "anthropic") + + assert isinstance(mapped, litellm.Timeout) + assert _provider_billed_the_failure(mapped) is billed + + def test_an_upstream_status_maps_through_the_provider_error_class() -> None: upstream: Final = ValueError(429, "slow down") upstream.headers = [("retry-after", "3")] # pyright: ignore[reportAttributeAccessIssue] # set by the native host @@ -218,6 +228,44 @@ def test_an_upstream_status_maps_through_the_provider_error_class() -> None: assert mapped.__context__ is upstream +@pytest.mark.parametrize( + ("body", "expected"), + [ + ("", {"content-type": "text/plain"}), + ("gateway", {"content-type": "text/plain", "content-length": "7"}), + ], +) +def test_upstream_failure_headers_match_the_python_clients_rebuilt_response( + body: str, expected: dict[str, str] +) -> None: + from litellm.litellm_core_utils.exception_mapping_utils import _get_response_headers + + upstream: Final = ValueError(504, body) + upstream.headers = [ # pyright: ignore[reportAttributeAccessIssue] # set by the native host + ("content-type", "text/plain"), + ("content-encoding", "gzip"), + ("content-length", "42"), + ] + + mapped: Final = messages_host.map_failure(upstream, {"model": _ADAPTIVE_OPUS}, "anthropic") + + assert isinstance(mapped, litellm.Timeout) + assert dict(_get_response_headers(mapped) or {}) == expected + + +@pytest.mark.parametrize("status", [200, 500]) +def test_upstream_failure_without_wire_headers_states_no_headers(status: int) -> None: + from litellm.litellm_core_utils.exception_mapping_utils import _get_response_headers + + upstream: Final = ValueError(status, "not a messages response") + upstream.headers = [] # pyright: ignore[reportAttributeAccessIssue] # set by the native host + + mapped: Final = messages_host.map_failure(upstream, {"model": _ADAPTIVE_OPUS}, "anthropic") + + assert isinstance(mapped, Exception) + assert not dict(_get_response_headers(mapped) or {}) + + def test_a_late_rejection_tells_the_logger_the_resolved_call_first() -> None: import datetime @@ -245,3 +293,133 @@ def test_a_late_rejection_tells_the_logger_the_resolved_call_first() -> None: assert isinstance(mapped, litellm.BadRequestError) assert logger.model_call_details["custom_llm_provider"] == "anthropic" assert logger.model_call_details["optional_params"] == {"max_tokens": 8} + + +@pytest.mark.parametrize( + ("aiohttp", "cut", "error_class", "text"), + ( + ( + True, + (60, None, False), + "ReadError", + "Response payload is not completed: ", + ), + ( + True, + (19, 100, True), + "ReadError", + "Response payload is not completed: . " + "ConnectionResetError(104, 'Connection reset by peer')", + ), + ( + False, + (60, None, False), + "RemoteProtocolError", + "peer closed connection without sending complete message body (incomplete chunked read)", + ), + ( + False, + (19, 100, False), + "RemoteProtocolError", + "peer closed connection without sending complete message body (received 19 bytes, expected 100)", + ), + (False, (19, 100, True), "ReadError", ""), + ), +) +def test_a_stream_body_cut_reaches_the_caller_as_the_transport_raised_it( + monkeypatch: pytest.MonkeyPatch, + aiohttp: bool, + cut: tuple[int, int | None, bool], + error_class: str, + text: str, +) -> None: + import httpx + + from litellm.rust_bridge.host import http_transport + + monkeypatch.setattr(http_transport, "uses_aiohttp_transport", lambda: aiohttp) + native: Final = RuntimeError("upstream network error") + native.messages_body_cut = cut # pyright: ignore[reportAttributeAccessIssue] # marker the native host sets + + mapped: Final = messages_host.map_failure(native, {"model": _ADAPTIVE_OPUS}, "anthropic") + + assert type(mapped) is getattr(httpx, error_class) + assert str(mapped) == text + + +_TRANSPORT: Final = "error sending request: client error" + + +@pytest.mark.parametrize( + ("aiohttp", "native", "text"), + ( + ( + True, + f"could not reach the provider: {_TRANSPORT} (Connect): tcp connect error: Connection refused (os error 111)", + "Cannot connect to host 127.0.0.1:9 ssl: [Connect call failed ('127.0.0.1', 9)]", + ), + ( + True, + f"could not reach the provider: {_TRANSPORT} (Connect): dns error: error resolving DNS: " + "failed to lookup address information: Name or service not known", + "Cannot connect to host 127.0.0.1:9 ssl: [Name or service not known]", + ), + ( + True, + f"could not reach the provider: {_TRANSPORT} (Connect): tls handshake eof", + "Cannot connect to host 127.0.0.1:9 ssl: [None]", + ), + ( + True, + f"upstream network error: {_TRANSPORT} (SendRequest): connection error: Connection reset by peer (os error 104)", + "[Errno 104] Connection reset by peer", + ), + ( + True, + f"upstream network error: {_TRANSPORT} (SendRequest): connection closed before message completed", + "Server disconnected", + ), + ( + False, + f"could not reach the provider: {_TRANSPORT} (Connect): tcp connect error: Connection refused (os error 111)", + "All connection attempts failed", + ), + ( + False, + f"could not reach the provider: {_TRANSPORT} (Connect): dns error: error resolving DNS: " + "failed to lookup address information: Name or service not known", + "[Errno -2] Name or service not known", + ), + (False, f"could not reach the provider: {_TRANSPORT} (Connect): tls handshake eof", ""), + ( + False, + f"upstream network error: {_TRANSPORT} (SendRequest): connection error: Connection reset by peer (os error 104)", + "", + ), + ( + False, + f"upstream network error: {_TRANSPORT} (SendRequest): connection closed before message completed", + "Server disconnected without sending a response.", + ), + ), +) +def test_a_failure_before_any_response_maps_like_the_python_transport_error( + monkeypatch: pytest.MonkeyPatch, aiohttp: bool, native: str, text: str +) -> None: + from litellm.rust_bridge.host import http_transport + + monkeypatch.setattr(http_transport, "uses_aiohttp_transport", lambda: aiohttp) + monkeypatch.setattr(messages_host, "get_ssl_configuration", lambda _ssl_verify: "") + + mapped: Final = messages_host.map_failure( + RuntimeError(native), {"model": _ADAPTIVE_OPUS, "api_base": "http://127.0.0.1:9"}, "anthropic" + ) + + assert type(mapped) is litellm.InternalServerError + assert mapped.status_code == 500 + assert mapped.message == ( + f"litellm.InternalServerError: AnthropicException - {text}. Handle with `litellm.InternalServerError`." + ) + assert "This looks like a bug in LiteLLM" not in str(mapped) diff --git a/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py index 87b9e70612c..b00df6547d2 100644 --- a/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py +++ b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py @@ -338,3 +338,27 @@ def test_stream_head_is_stamped_in_place_with_the_call_hidden_params() -> None: assert head["additional_headers"] == {"request-id": "req_native", "llm_provider-request-id": "req_native"} assert head["litellm_call_id"] == "messages" assert "_response_ms" in head + + +@pytest.mark.asyncio +async def test_a_native_stream_failure_logs_like_the_python_chunk_processor(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm._logging import verbose_proxy_logger + from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler + + logged: Final[list[tuple[object, ...]]] = [] + scheduled: Final[list[Mapping[str, object]]] = [] + + async def schedule(**kwargs: object) -> None: + scheduled.append(MappingProxyType(dict(kwargs))) + + monkeypatch.setattr(verbose_proxy_logger, "error", lambda *args: logged.append(args)) + monkeypatch.setattr(PassThroughStreamingHandler, "schedule_stream_failure_logging", schedule) + logger: Final = _supplied_logger() + logger.model_call_details["cache_hit"] = False + cut: Final = httpx.ReadError("cut") + + await legacy.stream_failure(logger, "anthropic", {}, [b"event: ping\n\n"], cut) + + assert logged == [("Error in chunk_processor: %s", cut)] + assert "cache_hit" not in logger.model_call_details + assert scheduled[0]["exception"] is cut From 6a4aab41f032d9948801a05dafd0a53887d488f9 Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Mon, 9 Nov 2026 16:24:21 -0800 Subject: [PATCH 37/46] fix(messages): bill a native stream the client abandoned like python (#50425) When the client disconnects mid-stream, the native adapter now logs success for the chunks it delivered before teardown, so spend, SpendLogs and the key's budget and tpm reservations settle the same way they do on the Python path The proxy records a client disconnect, with its 499 error information, in the stream generator's cleanup. The native teardown billed the stream before that cleanup, so the row had empty error information. A proxy stream torn down by cancellation now bills from the task's done callback, after the disconnect is recorded --- .../python_contract.json | 3 +- .../callbacks-legacy-python/src/adapter.rs | 69 ++++++++++++++++++- .../src/test_support.rs | 2 +- .../host/callbacks_legacy_python.py | 17 ++++- .../messages/test_streaming.py | 20 ++++++ .../host/test_callbacks_legacy_python.py | 41 +++++++++++ 6 files changed, 147 insertions(+), 5 deletions(-) diff --git a/litellm-rust/crates/callbacks-legacy-python/python_contract.json b/litellm-rust/crates/callbacks-legacy-python/python_contract.json index c0fe1c7b52f..8a212223523 100644 --- a/litellm-rust/crates/callbacks-legacy-python/python_contract.json +++ b/litellm-rust/crates/callbacks-legacy-python/python_contract.json @@ -130,7 +130,8 @@ "chunks", "start", "end", - "first_chunk" + "first_chunk", + "stream_teardown" ], "stream_failure": [ "logger", diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index 28bd5f2cefb..f68541c5c7b 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -329,7 +329,12 @@ impl LegacyLogging { logger.sync_success_for_async_call(py, &self.response, &self.start, &self.end) } - fn stream_success(&self, py: Python<'_>, stream: &DeliveredStream) -> PyResult<()> { + fn stream_success( + &self, + py: Python<'_>, + stream: &DeliveredStream, + teardown: bool, + ) -> PyResult<()> { let logger = self.logger()?; let billing = self.stream_billing().ok_or(HostError::MissingState)?; let billed = Streaming::Success.call( @@ -346,6 +351,7 @@ impl LegacyLogging { &self.start, &self.end, &stream.first_chunk, + teardown, ), ); match billed { @@ -667,7 +673,7 @@ impl LegacyLogging { self.end = Some(datetime(py, timing.end_time)?); self.response = Some(response.clone_ref(py)); match &self.stream { - Some(stream) => self.stream_success(py, stream)?, + Some(stream) => self.stream_success(py, stream, false)?, None => self.dispatch_success(py)?, } Ok(HookStep::Ready(())) @@ -737,8 +743,28 @@ impl LegacyLogging { } } +impl LegacyLogging { + // A caller that goes away mid-stream cancels the call before it ends: Python's Messages + // stream iterator still bills the chunks it read, straight away rather than deferred. + fn bill_abandoned_stream(&mut self, py: Python<'_>) -> PyResult<()> { + if !self.asynchronous || self.end.is_some() { + return Ok(()); + } + let Some(stream) = self.stream.take() else { + return Ok(()); + }; + self.end = Some(datetime(py, epoch_seconds())?); + self.stream_success(py, &stream, true) + } +} + impl PythonOwned for LegacyLogging { fn close(&mut self, py: Python<'_>) { + if self.logger.is_some() + && let Err(error) = self.bill_abandoned_stream(py) + { + error.write_unraisable(py, None); + } if let Some(logger) = self.logger.take() && let Err(error) = logger.restore_context(py) { @@ -2186,6 +2212,45 @@ assert chunks[1] is last }); } + #[rstest] + #[case::cancelled_mid_stream(None, &["stamp_stream_head", "stream_opened", "stream_success", "restore"])] + #[case::after_success(Some(true), &["stamp_stream_head", "stream_opened", "stream_success", "restore"])] + #[case::after_failure(Some(false), &["stamp_stream_head", "stream_opened", "stream_failure", "restore"])] + fn closing_an_unfinished_stream_bills_its_delivered_chunks_once( + #[case] ended: Option, + #[case] expected: &[&str], + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c"first = b'first'\nresponse = None\nfailure = ValueError('provider')\nhead = {}", + ); + let mut logging = LegacyLogging { + operation: crate::LoggingOperation::Messages, + ..logged(py, &locals, true) + }; + logging + .on_stream_open(py, &local(&locals, "head").unbind()) + .unwrap(); + logging + .on_stream_chunk(py, &local(&locals, "first").unbind()) + .unwrap(); + match ended { + Some(true) => drop(succeed(py, &locals, &mut logging)), + Some(false) => drop(fail(py, &locals, &mut logging)), + None => {} + } + litellm_host_python::PythonOwned::close(&mut logging, py); + let names: Vec = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert_eq!(names, expected); + }); + } + #[rstest] #[case::synchronous(false, &["failure_handler"])] #[case::asynchronous(true, &[])] diff --git a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs index c88a6496423..ee16e0c0059 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs @@ -98,7 +98,7 @@ FAKES = { 'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type), 'stream_opened': lambda logger, head: logger.record('stream_opened', head), 'stamp_stream_head': lambda head, logger, kwargs, start, end: logger.record('stamp_stream_head', head), - 'stream_success': lambda logger, url_route, endpoint_type, request_body, chunks, start, end, first_chunk: logger.record( + 'stream_success': lambda logger, url_route, endpoint_type, request_body, chunks, start, end, first_chunk, stream_teardown: logger.record( 'stream_success', list(chunks) ), 'stream_failure': lambda logger, endpoint_type, request_body, chunks, error: logger.record('stream_failure', error), diff --git a/litellm/rust_bridge/host/callbacks_legacy_python.py b/litellm/rust_bridge/host/callbacks_legacy_python.py index 5a368573f3d..ba557f769eb 100644 --- a/litellm/rust_bridge/host/callbacks_legacy_python.py +++ b/litellm/rust_bridge/host/callbacks_legacy_python.py @@ -11,6 +11,7 @@ import contextvars import datetime import inspect import json +import sys import traceback import uuid import weakref @@ -457,6 +458,14 @@ def stream_opened(logger: Logging, head: Mapping[str, object]) -> None: ) +def _cancelling(task: asyncio.Task[object]) -> bool: + return task.cancelling() > 0 if sys.version_info >= (3, 11) else True + + +def _proxy_request(logger: LoggingSurface) -> object: + return logger.litellm_params.get("proxy_server_request") + + def stream_success( logger: Logging, url_route: str, @@ -466,6 +475,7 @@ def stream_success( start: datetime.datetime, end: datetime.datetime, first_chunk: datetime.datetime | None, + stream_teardown: bool, ) -> None: from litellm.llms.anthropic.messages.native.streaming_iterator import ( GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ, @@ -490,7 +500,7 @@ def stream_success( raw_bytes=chunks, end_time=end, ) - if getattr(logger, "_on_deferred_stream_complete", None) is not None: + if not stream_teardown and getattr(logger, "_on_deferred_stream_complete", None) is not None: logger._deferred_stream_complete_args = (coroutine,) # pyright: ignore[reportAttributeAccessIssue] # the proxy's deferred stream release reads this slot return try: @@ -500,6 +510,11 @@ def stream_success( executor.submit(contextvars.copy_context().run, asyncio.run, coroutine) return + cancelled: Final = asyncio.current_task() + if stream_teardown and cancelled is not None and _cancelling(cancelled) and _proxy_request(logger): + # the proxy's disconnect cleanup runs later in this task and stamps the 499 error_information this row carries + cancelled.add_done_callback(lambda _: enqueue_logging(coroutine)) + return enqueue_logging(coroutine) diff --git a/tests/test_litellm_rust/messages/test_streaming.py b/tests/test_litellm_rust/messages/test_streaming.py index 491d257b482..99b6101da4e 100644 --- a/tests/test_litellm_rust/messages/test_streaming.py +++ b/tests/test_litellm_rust/messages/test_streaming.py @@ -1,3 +1,4 @@ +import asyncio import json from collections.abc import AsyncIterator, Iterator from typing import Final @@ -73,6 +74,25 @@ async def test_stream_closed_early_logs_success_once(server: RecordingServer) -> await anext(stream) +@pytest.mark.asyncio +async def test_stream_cancelled_while_waiting_on_the_provider_logs_success_for_what_it_read( + server: RecordingServer, +) -> None: + server.enqueue(ResponseSpec(body=None, events=MESSAGES_EVENTS, chunk_delay=1)) + recorder: Final = RecordingLogger() + + stream: Final = await call(server, stream=True, callbacks=[recorder]) + assert isinstance(stream, AsyncIterator) + await anext(stream) + with pytest.raises(TimeoutError): + await asyncio.wait_for(anext(stream), timeout=0.2) + + success: Final = await recorder.wait_for_async("async_log_success_event") + assert len(success) == 1 + assert success[0].response.usage.prompt_tokens == MESSAGES_RESPONSE["usage"]["input_tokens"] + assert "log_failure_event" not in recorder.names + + async def _streamed_pre_call_input(server: RecordingServer) -> object: server.enqueue(STREAM) recorder: Final = RecordingLogger() diff --git a/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py index b00df6547d2..cb594767908 100644 --- a/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py +++ b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py @@ -1,3 +1,4 @@ +import asyncio import datetime import inspect from collections.abc import Mapping @@ -362,3 +363,43 @@ async def test_a_native_stream_failure_logs_like_the_python_chunk_processor(monk assert logged == [("Error in chunk_processor: %s", cut)] assert "cache_hit" not in logger.model_call_details assert scheduled[0]["exception"] is cut + + +@pytest.mark.asyncio +async def test_a_proxy_stream_abandoned_mid_cancel_bills_after_the_disconnect_cleanup( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler + + logger: Final = _messages_logger() + logger.update_environment_variables( + litellm_params={"metadata": {}, "proxy_server_request": {"url": "/v1/messages"}}, optional_params={} + ) + + async def billing() -> None: + return None + + events: Final[list[str]] = [] + monkeypatch.setattr(PassThroughStreamingHandler, "_route_streaming_logging_to_handler", lambda **_: billing()) + monkeypatch.setattr(legacy, "enqueue_logging", lambda coroutine: (events.append("billed"), coroutine.close())) + started: Final = asyncio.Event() + now: Final = datetime.datetime.now() + + async def relay() -> None: + try: + started.set() + await asyncio.Event().wait() + except asyncio.CancelledError: + legacy.stream_success(logger, "/v1/messages", "anthropic", {}, [b"event: ping\n\n"], now, now, None, True) + await asyncio.sleep(0) + events.append("disconnect recorded") + raise + + task: Final = asyncio.create_task(relay()) + await started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.sleep(0) + + assert events == ["disconnect recorded", "billed"] From 13c1c2a364b2eb04c9d7aea7bcfe302a51495657 Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Mon, 9 Nov 2026 16:24:21 -0800 Subject: [PATCH 38/46] fix(messages): write the native unary cache entry after the router stamps it like python (#50426) Python stores a unary messages response from a background cache write task that runs after the router stamps _hidden_params, so the entry carries model_id and the router headers and a later hit returns x-litellm-cache-key. The native path stored the entry inside the native call, before the router saw the response. The async unary store is now deferred until the route encodes its public response and is handed to Python's own cache write task --- .../crates/host-python/src/binding.rs | 5 + litellm-rust/crates/host-python/src/driver.rs | 12 ++- .../crates/python-bridge/src/cache/mod.rs | 2 + .../python-bridge/src/cache/python/AGENTS.md | 2 +- .../python-bridge/src/cache/python/host.rs | 27 +++++- .../python-bridge/src/cache/python/service.rs | 8 +- .../python-bridge/src/cache/selection.rs | 3 +- .../python-bridge/src/routes/messages/host.rs | 4 + litellm/rust_bridge/host/cache.py | 44 ++++++++- .../cache/test_python_cache.py | 96 +++++++++++++++++-- tests/unit/rust_bridge/host/test_cache.py | 34 ++++++- 11 files changed, 220 insertions(+), 17 deletions(-) diff --git a/litellm-rust/crates/host-python/src/binding.rs b/litellm-rust/crates/host-python/src/binding.rs index e36b7d70c28..fe6173dd100 100644 --- a/litellm-rust/crates/host-python/src/binding.rs +++ b/litellm-rust/crates/host-python/src/binding.rs @@ -28,6 +28,11 @@ pub trait PythonBinding: PythonOwned { ) -> PyResult>; /// What the stream carries at hand-off, as the caller's stream receives it. + /// The unary response after the deployment transform, before success hooks run. + fn deliver_response(&mut self, _py: Python<'_>, _response: &Bound<'_, PyAny>) -> PyResult<()> { + Ok(()) + } + fn encode_stream_head( &mut self, py: Python<'_>, diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index dbf843a5295..052156cbc5d 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -399,8 +399,16 @@ where Ok(step) => step, Err(error) => return self.failure(error, FailureOrigin::Host).await, }; - match self.settle(step).await { - Ok(response) => self.succeeded(response).await, + let response = match self.settle(step).await { + Ok(response) => response, + Err(error) => return self.failure(error, FailureOrigin::Host).await, + }; + let delivered = { + let mut parts = self.parts(); + Python::attach(|py| parts.binding.deliver_response(py, response.bind(py))) + }; + match delivered { + Ok(()) => self.succeeded(response).await, Err(error) => self.failure(error, FailureOrigin::Host).await, } } diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index 4e3847951d8..fe9e06b3d23 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -6,6 +6,8 @@ pub(crate) use selection::{Cached, Selection, configure, ensure_cache_disabled}; use pyo3::{exceptions::PyRuntimeError, prelude::*}; +const CACHE_HOST_MODULE: &str = "litellm.rust_bridge.host.cache"; + #[derive(Debug, thiserror::Error)] pub(crate) enum HandleError { #[error("this inference route does not support configured caches")] diff --git a/litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md b/litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md index 81a0997f02a..3b516251f43 100644 --- a/litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md +++ b/litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md @@ -4,6 +4,6 @@ This directory lets Rust inference use a selected Python cache. `service.rs` imp Keep the shared inference protocol wrapper and adapter selection in the parent module. Receive the configured cache and prepared arguments from the parent module. Do not discover global configuration, choose native backends, or move inference to Python. Core remains independent of Python objects and cache implementation details -Await asynchronous cache operations through the existing host driver in the caller's task. Do not create another asyncio task or event loop. Cancellation must prevent subsequent provider requests and cache writes. Preserve ordinary cache failure handling without swallowing cancellation or other Python base exceptions +Await asynchronous cache operations through the existing host driver in the caller's task. Do not create another asyncio task or event loop. The one exception is the asynchronous unary messages store: Python's `_async_set_cache` writes the returned response from `create_cache_write_task` after the router has stamped its `_hidden_params`, so `host.rs` defers that store until the deployment transform has returned the response the caller receives and hands that response to `litellm.rust_bridge.host.cache.store_response`, which schedules the same write task. Lookups, sync stores and stream stores stay in the caller's task. Cancellation must prevent subsequent provider requests and cache writes. Preserve ordinary cache failure handling without swallowing cancellation or other Python base exceptions Traverse retained Python references for GC and release pending replies when the call closes. Regression tests must require native inference with Python fallback disabled and assert observable hits, provider request counts, cache-key headers, task identity and cancellation 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 82b7d881df6..a8cbd0ce607 100644 --- a/litellm-rust/crates/python-bridge/src/cache/python/host.rs +++ b/litellm-rust/crates/python-bridge/src/cache/python/host.rs @@ -9,7 +9,7 @@ use pyo3::{ use serde_json::Value; use super::service::CacheCall; -use crate::cache::HandleError; +use crate::cache::{CACHE_HOST_MODULE, HandleError}; enum Pending { Lookup(Reply, Error>>), @@ -21,6 +21,7 @@ pub(crate) struct PythonCache { arguments: Option>, pending: Option, asynchronous: bool, + deferred: bool, } impl PythonCache { @@ -30,6 +31,7 @@ impl PythonCache { arguments: None, pending: None, asynchronous, + deferred: false, } } @@ -64,7 +66,14 @@ impl PythonCache { None, ) } - CacheCall::Store { value, reply } => { + CacheCall::Store { + unary: true, reply, .. + } if self.asynchronous => { + self.deferred = true; + reply.send(Ok(())); + return Ok(None); + } + CacheCall::Store { value, reply, .. } => { self.pending = Some(Pending::Store(reply)); ( if self.asynchronous { @@ -89,6 +98,19 @@ impl PythonCache { self.resume(py, result) } + pub fn store_response(&mut self, py: Python<'_>, response: &Bound<'_, PyAny>) -> PyResult<()> { + if !std::mem::take(&mut self.deferred) { + return Ok(()); + } + let (Some(cache), Some(arguments)) = (self.cache.as_ref(), self.arguments.as_ref()) else { + return Err(HandleError::Unconfigured.into()); + }; + py.import(CACHE_HOST_MODULE)? + .getattr("store_response")? + .call((cache, response), Some(&arguments.bind(py).copy()?))?; + Ok(()) + } + pub fn resume( &mut self, py: Python<'_>, @@ -124,6 +146,7 @@ impl PythonCache { pub fn close(&mut self) { self.pending = None; + self.deferred = false; self.cache = None; self.arguments = 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 38872b85d30..018c8647043 100644 --- a/litellm-rust/crates/python-bridge/src/cache/python/service.rs +++ b/litellm-rust/crates/python-bridge/src/cache/python/service.rs @@ -45,6 +45,7 @@ pub(crate) enum CacheCall { }, Store { value: Value, + unary: bool, reply: Reply>, }, } @@ -101,8 +102,13 @@ where ) -> Pin> + Send + 'a>> { Box::pin(async move { let value = to_python(value)?; + let unary = value.get(STREAM_EVENTS_KEY).is_none(); self.services - .call(|reply| CacheCall::Store { value, reply }) + .call(|reply| CacheCall::Store { + value, + unary, + reply, + }) .await .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 e96f079e5da..78c1606ea40 100644 --- a/litellm-rust/crates/python-bridge/src/cache/selection.rs +++ b/litellm-rust/crates/python-bridge/src/cache/selection.rs @@ -1,11 +1,10 @@ -use super::{HandleError, python}; +use super::{CACHE_HOST_MODULE, HandleError, python}; use litellm_cache_response::{CacheOptions, CachePolicy, CacheScope, ScopedCache}; use litellm_host::{ machine::{HostServices, MachineFault}, protocol::Protocol, }; use pyo3::{prelude::*, types::PyDict}; -const CACHE_HOST_MODULE: &str = "litellm.rust_bridge.host.cache"; pub(crate) struct Cached

(std::marker::PhantomData

); 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 e92daa67be4..2a5f79ca393 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -244,6 +244,10 @@ impl PythonBinding for MessagesPythonHost { public_response(py, ROUTE_HOST_MODULE, &response) } + fn deliver_response(&mut self, py: Python<'_>, response: &Bound<'_, PyAny>) -> PyResult<()> { + self.cache.store_response(py, response) + } + fn encode_stream_head( &mut self, py: Python<'_>, diff --git a/litellm/rust_bridge/host/cache.py b/litellm/rust_bridge/host/cache.py index cac4b75677f..bf2655f9f62 100644 --- a/litellm/rust_bridge/host/cache.py +++ b/litellm/rust_bridge/host/cache.py @@ -1,12 +1,54 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from collections.abc import Mapping +from typing import TYPE_CHECKING, Final + +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict if TYPE_CHECKING: from litellm.caching.caching import Cache +_RUST_RESPONSE_HEADER: Final = "x-litellm-rust" +_HEADERS: Final = TypeAdapter(dict[str, object]) + def configured_cache() -> Cache | None: import litellm return litellm.cache + + +def _python_entry(response: Mapping[str, object]) -> Mapping[str, object]: + hidden: Final = get_hidden_params_dict(response) + headers: Final = hidden.get("additional_headers") + if not isinstance(headers, Mapping) or _RUST_RESPONSE_HEADER not in headers: + return response + return { + **response, + "_hidden_params": { + **hidden, + "additional_headers": { + name: value + for name, value in _HEADERS.validate_python(headers).items() + if name != _RUST_RESPONSE_HEADER + }, + }, + } + + +def store_response( + cache: Cache, + response: Mapping[str, object], + /, + **kwargs: object, # kwargs-ok: the cache call's own arguments +) -> None: + from litellm.caching.caching_handler import create_cache_write_task + + create_cache_write_task( + lambda: cache.async_add_cache( # pyright: ignore[reportUnknownMemberType] # Cache.async_add_cache leaves its result untyped + _python_entry(response), + **kwargs, # pyright: ignore[reportArgumentType] # the cache call's own keyword arguments, as Python passes them + ) + ) diff --git a/tests/test_litellm_rust/cache/test_python_cache.py b/tests/test_litellm_rust/cache/test_python_cache.py index d26fe47d69d..d88feec7d5e 100644 --- a/tests/test_litellm_rust/cache/test_python_cache.py +++ b/tests/test_litellm_rust/cache/test_python_cache.py @@ -9,7 +9,11 @@ from litellm.caching.caching import Cache, CacheMode from litellm.caching.caching_handler import ( _PENDING_CACHE_WRITES, # pyright: ignore[reportPrivateUsage] # await the existing background cache writer before the next request ) -from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict +from litellm.litellm_core_utils.add_retry_fallback_headers import ( + _add_headers_to_response, # pyright: ignore[reportPrivateUsage] # stamp the response the way the router does + add_retry_headers_to_response, + get_hidden_params_dict, +) from litellm.proxy._types import Litellm_EntityType from litellm.proxy.hooks.model_max_budget_limiter import ( PROXY_VirtualKeyModelMaxBudgetLimiter, @@ -19,6 +23,8 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelReq from litellm.proxy.utils import InternalUsageCache from litellm.redis.stores.rate_limits import RateLimitScope from litellm.rust_bridge.catalog import Route +from litellm.types.llms.anthropic import AnthropicResponseContentBlockText +from litellm.types.utils import CallTypes, LLMResponseTypes from tests.test_litellm_rust.cache.support import cache_key, collect, invoke, payload from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging, request_body @@ -251,9 +257,11 @@ async def test_rust_responses_rejects_a_configured_python_cache( @pytest.mark.asyncio @pytest.mark.parametrize("cancel_lookup", (False, True)) -async def test_python_cache_operations_stay_in_the_rust_callers_task( +@pytest.mark.parametrize("stream", (False, True), ids=("response", "stream")) +async def test_python_cache_operations_stay_in_the_rust_callers_task_except_the_unary_store( recording_server: RecordingServer, cancel_lookup: bool, + stream: bool, ) -> None: from litellm.caching.base_cache import BaseCache from litellm.caching.in_memory_cache import InMemoryCache @@ -278,14 +286,19 @@ async def test_python_cache_operations_stay_in_the_rust_callers_task( return storage.get_cache(key, **kwargs) async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: - assert asyncio.current_task() is caller + writer: Final = asyncio.current_task() + if stream: + assert writer is caller + else: + assert writer is not caller + assert writer in _PENDING_CACHE_WRITES await asyncio.sleep(0) storage.set_cache(key, value, **kwargs) litellm.cache = Cache(_backend=CallerCache()) if cancel_lookup: recording_server.expected_requests = 0 - task: Final = asyncio.create_task(invoke("messages", recording_server, {})) + task: Final = asyncio.create_task(invoke("messages", recording_server, {"stream": stream})) await asyncio.wait_for(entered.wait(), timeout=5) task.cancel() with pytest.raises(asyncio.CancelledError): @@ -295,13 +308,82 @@ async def test_python_cache_operations_stay_in_the_rust_callers_task( assert len(recording_server.requests) == 0 assert storage.cache_dict == {} return - first: Final = await invoke("messages", recording_server, {}) - second: Final = await invoke("messages", recording_server, {}) - assert payload(first) == payload(second) + first: Final = await invoke("messages", recording_server, {"stream": stream}) + first_payload: Final = await collect(first) if stream else payload(first) + await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) + second: Final = await invoke("messages", recording_server, {"stream": stream}) + second_payload: Final = await collect(second) if stream else payload(second) + assert first_payload == second_payload assert cache_key(second) assert len(recording_server.requests) == 1 +@pytest.mark.asyncio +async def test_rust_unary_messages_store_keeps_the_hidden_params_stamped_after_the_call( + recording_server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.caching.in_memory_cache import InMemoryCache + + storage: Final = InMemoryCache() + monkeypatch.setattr(litellm, "cache", Cache(_backend=storage)) + first: Final = await invoke("messages", recording_server, {}) + assert isinstance(first, dict) + assert storage.cache_dict == {} + add_retry_headers_to_response( + _add_headers_to_response(first, {"x-litellm-model-group": "group"}), attempted_retries=0 + ) + get_hidden_params_dict(first)["model_id"] = "deployment" + await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) + (entry,) = storage.cache_dict.values() + stored: Final = TypeAdapter(dict[str, object]).validate_python( + TypeAdapter(dict[str, object]).validate_python(entry)["response"] + ) + assert stored["_hidden_params"] == { + "additional_headers": {"x-litellm-model-group": "group", "x-litellm-attempted-retries": 0}, + "model_id": "deployment", + } + assert get_hidden_params_dict(first)["additional_headers"] == { + "x-litellm-rust": "true", + "x-litellm-model-group": "group", + "x-litellm-attempted-retries": 0, + } + hit: Final = await invoke("messages", recording_server, {}) + assert get_hidden_params_dict(hit)["cache_key"] == cache_key(hit) + assert len(recording_server.requests) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) +async def test_messages_cache_stores_the_response_a_yielding_deployment_hook_returned( + recording_server: RecordingServer, native: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.caching.in_memory_cache import InMemoryCache + from litellm.integrations.custom_logger import CustomLogger + + blocks: Final = TypeAdapter(list[AnthropicResponseContentBlockText]) + + class Redact(CustomLogger): + async def async_post_call_success_deployment_hook( + self, request_data: dict[str, object], response: LLMResponseTypes, call_type: CallTypes | None + ) -> LLMResponseTypes | None: + await asyncio.sleep(0.05) + assert isinstance(response, dict) + return {**response, "content": [{"type": "text", "text": "redacted"}]} # pyright: ignore[reportReturnType] # messages hooks see and return the plain dict the route serves + + monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") + storage: Final = InMemoryCache() + monkeypatch.setattr(litellm, "cache", Cache(_backend=storage)) + monkeypatch.setattr(litellm, "callbacks", [Redact()]) + first: Final = await invoke("messages", recording_server, {}, native=native) + assert isinstance(first, dict) + assert [block.text for block in blocks.validate_python(first["content"])] == ["redacted"] + await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) + hit: Final = await invoke("messages", recording_server, {}, native=native) + assert isinstance(hit, dict) + assert [block.text for block in blocks.validate_python(hit["content"])] == ["redacted"] + assert len(recording_server.requests) == 1 + + @pytest.mark.asyncio @pytest.mark.parametrize("namespace_source", ("cache", "metadata")) async def test_rust_messages_legacy_cache_honors_request_namespaces( diff --git a/tests/unit/rust_bridge/host/test_cache.py b/tests/unit/rust_bridge/host/test_cache.py index 9d532631ef2..c2b87cca1dc 100644 --- a/tests/unit/rust_bridge/host/test_cache.py +++ b/tests/unit/rust_bridge/host/test_cache.py @@ -1,10 +1,12 @@ +import asyncio from typing import Final import pytest import litellm from litellm.caching.caching import Cache -from litellm.rust_bridge.host.cache import configured_cache +from litellm.caching.caching_handler import _PENDING_CACHE_WRITES +from litellm.rust_bridge.host.cache import configured_cache, store_response from litellm.types.caching import LiteLLMCacheType @@ -19,3 +21,33 @@ def test_configured_cache_preserves_identity_and_observes_runtime_replacement(mo monkeypatch.setattr(litellm, "cache", None) assert configured_cache() is None + + +_REQUEST: Final = {"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "hi"}]} + + +async def _stored(response: dict[str, object]) -> object: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + store_response(cache, response, **_REQUEST) + await asyncio.gather(*_PENDING_CACHE_WRITES) + return await cache.async_get_cache(**_REQUEST) + + +@pytest.mark.asyncio +async def test_store_response_drops_only_the_rust_marker_header_from_the_cached_entry() -> None: + stored: Final = await _stored( + { + "id": "msg_1", + "_hidden_params": {"additional_headers": {"x-litellm-rust": "true", "llm_provider-x": "1"}, "model": "m"}, + } + ) + assert stored == { + "id": "msg_1", + "_hidden_params": {"additional_headers": {"llm_provider-x": "1"}, "model": "m"}, + } + + +@pytest.mark.asyncio +async def test_store_response_keeps_an_entry_without_the_rust_marker_unchanged() -> None: + response: Final = {"id": "msg_2", "_hidden_params": {"additional_headers": {"llm_provider-x": "1"}}} + assert await _stored(response) == response From c3e7835d685c777f0d153c0de028ef0d7c00bfbd Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 10 Nov 2026 00:30:01 +0000 Subject: [PATCH 39/46] test(integration): assert response ids stay scoped to their minting key (#50703) Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_responses_compact_id_ownership.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/tests/integration/authorization/test_responses_compact_id_ownership.py b/tests/integration/authorization/test_responses_compact_id_ownership.py index b5a6215dd15..c6486a61c85 100644 --- a/tests/integration/authorization/test_responses_compact_id_ownership.py +++ b/tests/integration/authorization/test_responses_compact_id_ownership.py @@ -442,7 +442,7 @@ def test_provider_errors_on_an_owned_compact_reach_the_caller(gateway: Gateway, assert_compact_rows(actors.owner, frozenset({after})) -def test_response_minted_by_a_key_without_owner_is_addressable_by_any_key(gateway: Gateway) -> None: +def test_response_minted_by_a_key_without_user_or_team_is_addressable_only_by_that_key(gateway: Gateway) -> None: upstream: Final = CompactUpstream() with wire_server(upstream.respond) as wire, gateway.scenario() as scenario: actors: Final = tenants(scenario, wire) @@ -451,7 +451,18 @@ def test_response_minted_by_a_key_without_owner_is_addressable_by_any_key(gatewa minted: Final = mint(gateway, ownerless, actors.model, created) upstreamed(wire) text: Final = marker("compact_ownerless") - assert_compacted(sdk_compact(gateway, actors.other_user, "/v1", actors.model, text, minted), actors.model, text) + assert_refused( + compact_post( + gateway, + actors.other_user, + "/v1", + compact_body(actors.model, text, {"previous_response_id": minted}), + text, + ), + OTHER_USER_DETAIL, + ) + assert compacts(upstreamed(wire)) == () + assert_compacted(sdk_compact(gateway, ownerless, "/v1", actors.model, text, minted), actors.model, text) assert compacts(upstreamed(wire)) == ( upstream_compact_body(text, {"previous_response_id": minted_id(created)}), ) From a390f5362864f591f8eabac3fe860741a8d61378 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 10 Nov 2026 00:33:52 +0000 Subject: [PATCH 40/46] chore(cost-map): add github_copilot/claude-fable-5.2 from the Copilot models-and-pricing table (#50775) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 22 +++++++++++++++++++ model_prices_and_context_window.json | 22 +++++++++++++++++++ 2 files changed, 44 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 66d375b70ff..6c1f41b2545 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -88044,5 +88044,27 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "thinking_always_on": true + }, + "github_copilot/claude-fable-5.2": { + "cache_creation_input_token_cost": 1e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "github_copilot", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/messages" + ], + "supports_adaptive_thinking": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "thinking_always_on": true } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 66d375b70ff..6c1f41b2545 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -88044,5 +88044,27 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "thinking_always_on": true + }, + "github_copilot/claude-fable-5.2": { + "cache_creation_input_token_cost": 1e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "github_copilot", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/messages" + ], + "supports_adaptive_thinking": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "thinking_always_on": true } } From b161fbf8d2e099993a203cc2632ad12cfbdba389 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 10 Nov 2026 00:33:57 +0000 Subject: [PATCH 41/46] test(integration): assert team model visibility follows caller scope (#50772) Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/authorization/test_team_scoped_models.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/integration/authorization/test_team_scoped_models.py b/tests/integration/authorization/test_team_scoped_models.py index 8f8e12ae498..433a9e34e04 100644 --- a/tests/integration/authorization/test_team_scoped_models.py +++ b/tests/integration/authorization/test_team_scoped_models.py @@ -137,12 +137,13 @@ def _v2_team_public_names(gateway: Gateway, key: str, model: str) -> list[JsonVa return [entry["model_info"].get("team_public_model_name") for entry in response.json()["data"]] -def test_v2_model_info_reports_a_team_model_to_team_and_non_team_keys(gateway: Gateway) -> None: +def test_v2_model_info_reports_a_team_model_only_to_team_and_admin_keys(gateway: Gateway) -> None: with gateway.scenario() as scenario: team: Final = scenario.team(models=[]) model: Final = scenario.model(model_info={"team_id": team}) assert _v2_team_public_names(gateway, scenario.key(team_id=team), model) == [model] - assert _v2_team_public_names(gateway, scenario.key(), model) == [model] + assert _v2_team_public_names(gateway, scenario.key(), model) == [] + assert _v2_team_public_names(gateway, gateway.key, model) == [model] @pytest.mark.parametrize("set_on", ["team_new", "team_update"]) From a053bf2a5c06059eda6ff09a00dac5da7e9ca0b6 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 10 Nov 2026 00:34:19 +0000 Subject: [PATCH 42/46] fix(cost-map): price wandb DeepSeek-V4.1-Flash-1017 from the CoreWeave Forge listing (#50776) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 6 +++--- model_prices_and_context_window.json | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 6c1f41b2545..68bbf612e31 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -87346,13 +87346,13 @@ "supports_vision": true }, "wandb/deepseek-ai/DeepSeek-V4.1-Flash-1017": { - "cache_read_input_token_cost": 2e-09, - "input_cost_per_token": 1.3e-07, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, "litellm_provider": "wandb", "max_input_tokens": 1049000, "max_output_tokens": 1048576, "mode": "chat", - "output_cost_per_token": 6.5e-07, + "output_cost_per_token": 1.2e-06, "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", "supports_prompt_caching": true, "supports_reasoning": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 6c1f41b2545..68bbf612e31 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -87346,13 +87346,13 @@ "supports_vision": true }, "wandb/deepseek-ai/DeepSeek-V4.1-Flash-1017": { - "cache_read_input_token_cost": 2e-09, - "input_cost_per_token": 1.3e-07, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, "litellm_provider": "wandb", "max_input_tokens": 1049000, "max_output_tokens": 1048576, "mode": "chat", - "output_cost_per_token": 6.5e-07, + "output_cost_per_token": 1.2e-06, "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", "supports_prompt_caching": true, "supports_reasoning": true, From 13df479397bdd37bdbfd5a57c63edcebe9a074bf Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 9 Nov 2026 17:05:10 -0800 Subject: [PATCH 43/46] fix(proxy): write SpendLogs failure rows for config-defined pass-through endpoints (#41831) * fix(proxy): write SpendLogs failure rows for config-defined pass-through endpoints Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): simplify pass-through failure api_base lookup Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): exempt pass-through spend log regression from quality gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): reuse failure payload builder in pass-through exception path Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): cover config-defined pass-through failure spend rows Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): cover failure spend route eligibility Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): check configured routes for pass-through failure logs 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: kerry Co-authored-by: kenneth --- .../proxy/hooks/proxy_track_cost_callback.py | 13 ++-- .../pass_through_endpoints.py | 9 ++- .../spend/test_passthrough_failure.py | 53 ++++++++++++++ .../hooks/test_proxy_track_cost_callback.py | 71 ++++++++++++++++++- .../test_pass_through_endpoints.py | 34 +++++++++ 5 files changed, 173 insertions(+), 7 deletions(-) create mode 100644 tests/integration/spend/test_passthrough_failure.py diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 133731fc335..dd7a24ae879 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -479,15 +479,20 @@ class _ProxyDBLogger(CustomLogger): return request_route: Final = user_api_key_dict.request_route + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import InitPassThroughEndpointHelpers + + route_is_trackable: Final = ( + request_route is None + or RouteChecks.is_llm_api_route(route=request_route) + or RouteChecks.is_info_route(route=request_route) + or InitPassThroughEndpointHelpers.has_configured_route_ignoring_trailing_slash(route=request_route) + ) if ( ProxyDBLogger._should_track_errors_in_db() is False and not completed_cost or ProxyDBLogger._skips_pre_routing_rejections() and is_pre_routing_rejection(original_exception, _REQUEST_DATA.validate_python(request_data)) - or request_route is not None - and not ( - RouteChecks.is_llm_api_route(route=request_route) or RouteChecks.is_info_route(route=request_route) - ) + or not route_is_trackable or request_route is not None and RouteChecks.is_dashboard_session_info_route(request_route, user_api_key_dict.team_id) ): diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 3e25f7ce9af..4b84fe19559 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -964,11 +964,16 @@ def _build_passthrough_failure_request_payload( model and provider. A user-defined route never gets that, since its URL names an upstream model rather than a LiteLLM-managed one. """ - request_payload: Final[dict] = dict(cast(Mapping[str, object], parsed_body or {})) # cast-ok: json body + request_payload: Final = dict(cast(Mapping[str, object], parsed_body or {})) # cast-ok: json body if kwargs: - request_payload.update(kwargs) + request_payload.update(_json_object(_DECODED_JSON.validate_python(kwargs))) if logging_obj is not None: request_payload["litellm_logging_obj"] = logging_obj + logging_litellm_params: Final = _json_object(logging_obj.model_call_details.get("litellm_params")) + existing_litellm_params: Final = _json_object(request_payload.get("litellm_params")) + api_base: Final = logging_litellm_params.get("api_base") + if isinstance(api_base, str) and api_base and "api_base" not in existing_litellm_params: + request_payload["litellm_params"] = {**existing_litellm_params, "api_base": api_base} url_route: Final = str(url) if url is not None else "" google_provider: Final = ( _passthrough_google_provider(url_route, custom_llm_provider) if builtin_provider_route else None diff --git a/tests/integration/spend/test_passthrough_failure.py b/tests/integration/spend/test_passthrough_failure.py new file mode 100644 index 00000000000..cb102c406f7 --- /dev/null +++ b/tests/integration/spend/test_passthrough_failure.py @@ -0,0 +1,53 @@ +import json +from typing import Final +from pathlib import Path + +import pytest +from tests.integration._support.client import Gateway, eventually +from tests.integration._support.database import read_rows +from tests.integration._support.process import owned_proxy +from tests.integration._support.wire import Reply, Request, wire_server + + +@pytest.mark.timeout(180) +def test_config_pass_through_endpoint_failure_writes_a_spend_row(gateway: Gateway, tmp_path: Path) -> None: + def respond(request: Request) -> Reply: + assert request.target == "/" + assert json.loads(request.body) == {"search": "hello"} + return Reply(status=500, body=json.dumps({"error": "upstream search failure"}).encode()) + + with wire_server(respond) as wire: + config: Final = tmp_path / "proxy_config.yaml" + config.write_text( + "model_list: []\n" + "general_settings:\n" + " master_key: os.environ/LITELLM_MASTER_KEY\n" + " database_url: os.environ/DATABASE_URL\n" + " store_model_in_db: true\n" + " disable_spend_logs: false\n" + " proxy_batch_write_at: 1\n" + " pass_through_endpoints:\n" + " - path: /azure-search\n" + f" target: {wire.url}\n" + " headers:\n" + " Authorization: Bearer synthetic-search-key\n" + "router_settings:\n" + " disable_cooldowns: true\n" + ) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate: + response: Final = candidate.request("POST", "/azure-search", {"search": "hello"}) + assert response.status_code == 500, response.text + call_id: Final = response.headers["x-litellm-call-id"] + rows: Final = eventually( + lambda: read_rows( + 'SELECT status, call_type, api_base, spend FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (call_id,), + database_url=candidate.database_url, + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert rows[0]["status"] == "failure", rows + assert rows[0]["call_type"] == "pass_through_endpoint", rows + assert "127.0.0.1" in rows[0]["api_base"], rows + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/")] diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 0283c90f3d6..43f529faa6b 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -13,7 +13,7 @@ import pytest from litellm.redis.stores.connection import connect from litellm.redis.stores.spend_counters import SpendCounterStore -from fastapi import HTTPException +from fastapi import FastAPI, HTTPException import litellm from litellm._logging import verbose_proxy_logger @@ -35,6 +35,7 @@ from litellm.proxy.hooks.proxy_track_cost_callback import ( run_spend_event, ) from litellm.proxy.route_llm_request import ProxyModelNotFoundError +from litellm.proxy.pass_through_endpoints.pass_through_endpoints import InitPassThroughEndpointHelpers from litellm.proxy.spend_tracking.spend_event import SpendEventDecodeError, build_spend_event, decode_spend_event from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer, UnixAddress from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload @@ -907,6 +908,74 @@ async def test_async_post_call_failure_hook_non_llm_route(): mock_update_database.assert_not_called() +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("route", "call_type", "written"), + [ + ("/azure-search", CallTypes.pass_through.value, True), + ("/unregistered-search", CallTypes.acompletion.value, False), + ("/key/generate", CallTypes.pass_through.value, False), + ("/team/new", CallTypes.pass_through.value, False), + ("/unregistered-search", CallTypes.pass_through.value, False), + ("/azure-search/", CallTypes.pass_through.value, True), + ("/configured-search/indexes/docs", CallTypes.pass_through.value, True), + (None, CallTypes.acompletion.value, True), + ], +) +async def test_async_post_call_failure_hook_tracks_config_defined_pass_through_route( + route: str | None, call_type: str, written: bool +) -> None: + api_base: Final = "http://upstream.example/indexes/x/docs/search" + request_data: Final = { + "model": "azure-search", + "call_type": call_type, + "litellm_params": {"api_base": api_base, "metadata": {}}, + } + user_api_key_dict: Final = UserAPIKeyAuth( + api_key="test_api_key", + request_route=route, + ) + + spend_writer: Final = create_autospec(DBSpendUpdateWriter, instance=True) + with patch.dict( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes", clear=True + ): + InitPassThroughEndpointHelpers.add_exact_path_route( + app=FastAPI(), + path="/azure-search", + target=api_base, + custom_headers={}, + forward_headers=False, + merge_query_params=False, + dependencies=[], + cost_per_request=None, + endpoint_id="configured-exact", + ) + InitPassThroughEndpointHelpers.add_subpath_route( + app=FastAPI(), + path="/configured-search", + target=api_base, + custom_headers={}, + forward_headers=False, + merge_query_params=False, + dependencies=[], + cost_per_request=None, + endpoint_id="configured-subpath", + ) + await ProxyDBLogger(spend_writer=lambda: spend_writer).async_post_call_failure_hook( + request_data=request_data, + original_exception=Exception("Provider error"), + user_api_key_dict=user_api_key_dict, + ) + + assert spend_writer.update_database.await_count == (1 if written else 0) + if not written: + return + call_args: Final = spend_writer.update_database.call_args + assert call_args.kwargs["kwargs"]["litellm_params"]["api_base"] == api_base + assert call_args.kwargs["kwargs"]["litellm_params"]["metadata"]["status"] == "failure" + + @pytest.mark.asyncio @pytest.mark.parametrize( ("team_id", "route", "written"), diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 464c39d0ba4..31b7b499845 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -449,6 +449,39 @@ async def test_non_streaming_http_request_handler_multipart_with_non_empty_parse assert call_args["files"] == [("file", ("test.txt", file_content, "text/plain"))] +@pytest.mark.parametrize("existing_api_base", [None, "http://existing.example/search"]) +def test_build_passthrough_failure_request_payload_carries_masked_api_base(existing_api_base: str | None) -> None: + logging_obj: Final = LiteLLMLoggingObj( + model="passthrough", + messages=[], + stream=False, + call_type=CallTypes.pass_through.value, + start_time=datetime(2026, 1, 1), + litellm_call_id="test-call-id", + function_id="test-function-id", + ) + logging_obj.pre_call( + input={}, + api_key="", + additional_args={ + "complete_input_dict": {}, + "api_base": "http://upstream.example/search?api-version=1", + "headers": {}, + }, + ) + params: Final = {"metadata": {}, **({"api_base": existing_api_base} if existing_api_base else {})} + payload: Final = _build_passthrough_failure_request_payload( + parsed_body={}, + kwargs={"litellm_params": params}, + logging_obj=logging_obj, + custom_llm_provider=None, + ) + assert payload["litellm_params"]["api_base"] == ( + existing_api_base or logging_obj.model_call_details["litellm_params"]["api_base"] + ) + assert params == {"metadata": {}, **({"api_base": existing_api_base} if existing_api_base else {})} + + @pytest.mark.asyncio async def test_pass_through_request_failure_handler(): """ @@ -468,6 +501,7 @@ async def test_pass_through_request_failure_handler(): assert "Request failed" in (failure.traceback_str or "") assert failure.request_data["contents"] == "not a list" assert failure.request_data["model"] == "" + assert failure.request_data["litellm_params"]["api_base"] == "http://test.com" @pytest.mark.asyncio From 71c1a2aff6b682865fd1395d1ed371da62fed8e1 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 9 Nov 2026 17:08:35 -0800 Subject: [PATCH 44/46] test(integration): repair tests that merged changes left stale (#50782) Several integration tests on main still assert behavior that later changes replaced on purpose, so they fail on every scheduled run, and one of them stops the whole sdk group at import. The batch cost case expected cached tokens to bill at the Batch input price when a price row has no Batch cache price. Since #50062 they bill at the Standard cache price, so the case now expects 0.000235 input cost and is renamed to say so. The OTel v2 attribute test pinned langfuse_otel and langtrace as carrying no request metadata baggage. Since #50373 every OTel v2 destination carries it, so all seven callbacks share one expectation and the Langfuse baggage test that was deferred under LIT-9128 runs again. The clean-EOF outage test trips its database relay on the text of the spend log insert. The ledger writes that insert without the "public" schema prefix, so the relay never tripped and the test timed out before it checked anything. The trigger now matches the insert on both write paths. The usage routing tests expected cooldown, tpm and rpm keys in one MGET. Since #50170 the cooldown keys are read in their own MGET, at most once per DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS. Both the proxy and the SDK test now expect the two MGETs, and the proxy test waits out that interval after its /model/info poll, which reads the same keys. The concurrent deployment delete test held its writer inside Router._update_deployment_indices_after_removal, which went away with the Python router core in #50178, so the file failed at import and stopped the whole sdk group. It now holds the writer where the registry reads the model list again after router-core applied the change. --- .../cost_calculation/cases/batch_realtime.py | 10 +++---- ...test_arize_otel_v2_openinference_family.py | 24 +++++++-------- .../test_usage_based_routing_redis_reads.py | 29 ++++++++++--------- ...router_concurrent_deployment_delete_sdk.py | 19 ++++++++---- ...est_usage_based_routing_sdk_redis_reads.py | 14 +++++---- ..._stream_clean_eof_without_finish_reason.py | 2 +- 6 files changed, 54 insertions(+), 44 deletions(-) diff --git a/tests/integration/cost_calculation/cases/batch_realtime.py b/tests/integration/cost_calculation/cases/batch_realtime.py index 0284036f31f..52e2021dc7d 100644 --- a/tests/integration/cost_calculation/cases/batch_realtime.py +++ b/tests/integration/cost_calculation/cases/batch_realtime.py @@ -46,8 +46,8 @@ GPT_5_6_BATCH_CACHED_INPUT_HALVED: Final = BatchCostCase( ), ) -GPT_5_4_BATCH_EXPLICIT_BATCH_RATES_BILL_CACHED_AT_BATCH_INPUT_RATE: Final = BatchCostCase( - name="gpt-5.4-batch-explicit_batch_rates_bill_cached_at_batch_input_rate", +GPT_5_4_BATCH_NO_BATCH_CACHE_RATE_BILLS_CACHED_AT_STANDARD_CACHE_RATE: Final = BatchCostCase( + name="gpt-5.4-batch-no_batch_cache_rate_bills_cached_at_standard_cache_rate", covers="quota_management.spend_tracking.batch_costs.explicit_rates", model="gpt-5.4", litellm_model="openai/gpt-5.4", @@ -56,8 +56,8 @@ GPT_5_4_BATCH_EXPLICIT_BATCH_RATES_BILL_CACHED_AT_BATCH_INPUT_RATE: Final = Batc BatchOutputLine(status_code=200, prompt_tokens=120, completion_tokens=30), ), expected=ExactExpected( - spend=0.000875, - input_cost=0.000275, + spend=0.000835, + input_cost=0.000235, output_cost=0.0006, prompt_tokens=220, completion_tokens=80, @@ -188,7 +188,7 @@ GPT_REALTIME_MINI_2025_12_15_REALTIME_SESSION_WITHOUT_TURNS_ZERO_SPEND: Final = BATCH_CASES: Final[tuple[BatchCostCase, ...]] = ( GPT_5_6_BATCH_HALVED_RATES_WHEN_MAP_HAS_NO_BATCH_KEYS, GPT_5_6_BATCH_CACHED_INPUT_HALVED, - GPT_5_4_BATCH_EXPLICIT_BATCH_RATES_BILL_CACHED_AT_BATCH_INPUT_RATE, + GPT_5_4_BATCH_NO_BATCH_CACHE_RATE_BILLS_CACHED_AT_STANDARD_CACHE_RATE, GPT_5_6_BATCH_ALL_REQUESTS_FAILED_ZERO_SPEND, ) diff --git a/tests/integration/observability/test_arize_otel_v2_openinference_family.py b/tests/integration/observability/test_arize_otel_v2_openinference_family.py index f638a7360be..881a65c687e 100644 --- a/tests/integration/observability/test_arize_otel_v2_openinference_family.py +++ b/tests/integration/observability/test_arize_otel_v2_openinference_family.py @@ -23,7 +23,7 @@ from integration._support.client import Gateway from integration._support.wire import Reply, Request from pydantic import JsonValue -_GENAI_B3_KEYS_WITHOUT_BAGGAGE: Final = frozenset( +_GENAI_B3_KEYS: Final = frozenset( { "gen_ai.input.messages", "gen_ai.operation.name", @@ -56,6 +56,7 @@ _GENAI_B3_KEYS_WITHOUT_BAGGAGE: Final = frozenset( "litellm.cost.output", "litellm.cost.tool_usage", "litellm.cost.total", + "litellm.metadata.trace_marker", "litellm.provider.model", "litellm.request.route", "litellm.request.tools.declared", @@ -66,7 +67,6 @@ _GENAI_B3_KEYS_WITHOUT_BAGGAGE: Final = frozenset( "server.port", } ) -_GENAI_B3_KEYS_WITH_BAGGAGE: Final = _GENAI_B3_KEYS_WITHOUT_BAGGAGE | frozenset({"litellm.metadata.trace_marker"}) _LANGFUSE_B3_KEYS: Final = frozenset( { "gen_ai.input.messages", @@ -109,6 +109,7 @@ _LANGFUSE_B3_KEYS: Final = frozenset( "litellm.cost.output", "litellm.cost.tool_usage", "litellm.cost.total", + "litellm.metadata.trace_marker", "litellm.provider.model", "litellm.request.route", "litellm.request.tools.declared", @@ -122,15 +123,14 @@ _LANGFUSE_B3_KEYS: Final = frozenset( _B3_ATTRIBUTE_KEYS: Final[Mapping[str, frozenset[str]]] = MappingProxyType( { "langfuse_otel": _LANGFUSE_B3_KEYS, - "langtrace": _GENAI_B3_KEYS_WITHOUT_BAGGAGE, - "signoz": _GENAI_B3_KEYS_WITH_BAGGAGE, - "newrelic": _GENAI_B3_KEYS_WITH_BAGGAGE, - "levo": _GENAI_B3_KEYS_WITH_BAGGAGE, - "agentops": _GENAI_B3_KEYS_WITH_BAGGAGE, - "otel": _GENAI_B3_KEYS_WITH_BAGGAGE, + "langtrace": _GENAI_B3_KEYS, + "signoz": _GENAI_B3_KEYS, + "newrelic": _GENAI_B3_KEYS, + "levo": _GENAI_B3_KEYS, + "agentops": _GENAI_B3_KEYS, + "otel": _GENAI_B3_KEYS, } ) -_B3_BAGGAGE_CALLBACKS: Final = frozenset({"signoz", "newrelic", "levo", "agentops", "otel"}) _B4_ATTRIBUTE_KEYS: Final = frozenset( { "input.value", @@ -305,14 +305,10 @@ def test_arize_otel_v2_b3_non_openinference_callback_family(callback: str, gatew baggage: Final = tuple( sorted((key, value) for key, value in attributes.items() if key.startswith("litellm.metadata.")) ) - expected_baggage: Final = ( - (("litellm.metadata.trace_marker", marker),) if callback in _B3_BAGGAGE_CALLBACKS else () - ) - assert baggage == expected_baggage, attributes + assert baggage == (("litellm.metadata.trace_marker", marker),), attributes def test_arize_otel_v2_b3_langfuse_carries_metadata_baggage(gateway: Gateway, tmp_path: Path) -> None: - pytest.skip("BUG: LIT-9128 Langfuse OTel v2 preset omits request-metadata baggage") marker: Final = "b3-langfuse-" + uuid.uuid4().hex def upstream(request: Request) -> Reply: diff --git a/tests/integration/routing/test_usage_based_routing_redis_reads.py b/tests/integration/routing/test_usage_based_routing_redis_reads.py index 8a31eac7f91..edb24cc7930 100644 --- a/tests/integration/routing/test_usage_based_routing_redis_reads.py +++ b/tests/integration/routing/test_usage_based_routing_redis_reads.py @@ -3,6 +3,7 @@ from __future__ import annotations import json import shlex import threading +import time import uuid from collections.abc import Iterator from contextlib import contextmanager @@ -19,6 +20,7 @@ from integration._support.client import Gateway, eventually, object_value, strin from integration._support.process import owned_proxy from integration._support.redis_process import owned_redis from integration._support.wire import Reply, Request, wire_server +from litellm.constants import DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS from pydantic import JsonValue, TypeAdapter from redis import Redis from redis.exceptions import TimeoutError as RedisTimeoutError @@ -217,11 +219,12 @@ def test_proxy_usage_routing_reads_cooldown_tpm_then_rpm_from_redis( lambda registered: registered == frozenset(deployment_ids), seconds=15, ) + cooldown_read_throttle_ends: Final = time.monotonic() + DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS wire.drain() eventually( - lambda: datetime.now(UTC), - lambda current: current.second < 40, - seconds=65, + lambda: (datetime.now(UTC), time.monotonic()), + lambda clock: clock[0].second < 40 and clock[1] >= cooldown_read_throttle_ends, + seconds=65 + DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS, ) minute: Final = datetime.now(UTC).strftime("%H-%M") markers: Final = tuple(f"{run_id}-{index}" for index in range(3)) @@ -256,18 +259,18 @@ def test_proxy_usage_routing_reads_cooldown_tpm_then_rpm_from_redis( observed_bodies: Final = tuple(_request_object(request.body) for request in received) expected_bodies: Final = tuple(_expected_wire_body(endpoint, marker) for marker in markers[:2]) assert observed_bodies == expected_bodies, observed_bodies - expected_mget: Final = ( + cooldown_mget: Final = ( "MGET", - f"deployment:{deployment_ids[0]}:cooldown", - f"deployment:{deployment_ids[1]}:cooldown", - f"{deployment_ids[0]}:openai/{OPENAI_MODEL}:tpm:{minute}", - f"{deployment_ids[1]}:openai/{OPENAI_MODEL}:tpm:{minute}", - *( - f"{deployment_id}:openai/{OPENAI_MODEL}:rpm:{minute}" - for deployment_id in deployment_ids - ), + *(f"deployment:{deployment_id}:cooldown" for deployment_id in deployment_ids), + ) + usage_mget: Final = ( + "MGET", + *(f"{deployment_id}:openai/{OPENAI_MODEL}:tpm:{minute}" for deployment_id in deployment_ids), + *(f"{deployment_id}:openai/{OPENAI_MODEL}:rpm:{minute}" for deployment_id in deployment_ids), ) mgets: Final = _drain_mgets(commands) - assert any(arguments == expected_mget for _, arguments in mgets), mgets + mget_arguments: Final = tuple(arguments for _, arguments in mgets) + assert cooldown_mget in mget_arguments, mgets + assert usage_mget in mget_arguments, mgets raw_mgets: Final = tuple(line for line, _ in mgets) print(f"proxy {endpoint} MGETs: {raw_mgets}") diff --git a/tests/integration/sdk/test_router_concurrent_deployment_delete_sdk.py b/tests/integration/sdk/test_router_concurrent_deployment_delete_sdk.py index 3867808aa5f..9f8aa966335 100644 --- a/tests/integration/sdk/test_router_concurrent_deployment_delete_sdk.py +++ b/tests/integration/sdk/test_router_concurrent_deployment_delete_sdk.py @@ -19,6 +19,8 @@ from pydantic import JsonValue import litellm from litellm import Router +from litellm.router_rust.surface.model_list import HostedModelList +from litellm.router_rust.surface.registry import HostRegistry _DEPLOYMENTS: Final = 4 _TARGET: Final = _DEPLOYMENTS - 1 @@ -28,12 +30,13 @@ _BURST_THREADS: Final = 8 _BURST_CALLS: Final = 24 _CHURN_UPSERTS: Final = 120 _UNREPORTED_MODEL_ID: Final = frozenset({"aresponses-stream", "aanthropic-messages"}) -_REINDEX: Final = Router._update_deployment_indices_after_removal.__code__ +_LIST_REFRESH: Final = HostedModelList.refresh.__code__ +_HAND_CHANGES: Final = HostRegistry._hand_changes.__code__ def _router_source_files() -> frozenset[str]: facade_files: Final = (sys.modules[cls.__module__].__file__ for cls in Router.__mro__ if cls is not object) - return frozenset({_REINDEX.co_filename, *(path for path in facade_files if path is not None)}) + return frozenset({_HAND_CHANGES.co_filename, *(path for path in facade_files if path is not None)}) def _lock_lines(path: str) -> frozenset[int]: @@ -88,9 +91,9 @@ def _router(wire: Wire, team: bool = False) -> Router: def _read_during_held_write(read: Callable[[], T], write: Callable[[], object]) -> T: - """Run `read` on its own thread while `write` waits where the list has shrunk but the index maps have not caught - up. The write resumes once `read` returns or reaches the router's index lock, so a read that skips the lock meets - the half-updated list and a read that takes the lock waits for the write""" + """Run `read` on its own thread while `write` waits where router-core has applied the change but the Python model + list has not been read again. The write resumes once `read` returns or reaches the router's index lock, so a read + that skips the lock meets the half-updated list and a read that takes the lock waits for the write""" held: Final = threading.Event() release: Final = threading.Event() settled: Final = threading.Event() @@ -99,7 +102,11 @@ def _read_during_held_write(read: Callable[[], T], write: Callable[[], object]) failure: Final[list[BaseException]] = [] def hold_at_reindex(frame: FrameType, event: str, _arg: object) -> None: - if event == "call" and frame.f_code is _REINDEX and not held.is_set(): + caller: Final = frame.f_back + refreshes_after_a_write: Final = ( + frame.f_code is _LIST_REFRESH and caller is not None and caller.f_code is _HAND_CHANGES + ) + if event == "call" and refreshes_after_a_write and not held.is_set(): held.set() released_in_time.append(release.wait(timeout=_SYNC_SECONDS)) diff --git a/tests/integration/sdk/test_usage_based_routing_sdk_redis_reads.py b/tests/integration/sdk/test_usage_based_routing_sdk_redis_reads.py index da9a8309c8d..ff4761b1da3 100644 --- a/tests/integration/sdk/test_usage_based_routing_sdk_redis_reads.py +++ b/tests/integration/sdk/test_usage_based_routing_sdk_redis_reads.py @@ -140,7 +140,7 @@ async def _exercise_router(router: Router, model_name: str, markers: tuple[str, return _model_id(first), _model_id(second) -def test_sdk_usage_routing_reads_tpm_then_rpm_from_redis(tmp_path: Path) -> None: +def test_sdk_usage_routing_reads_cooldown_tpm_then_rpm_from_redis(tmp_path: Path) -> None: with owned_redis(tmp_path) as cache, wire_server(_reply) as wire: run_id: Final = uuid.uuid4().hex model_name: Final = f"usage-redis-{run_id}" @@ -176,15 +176,19 @@ def test_sdk_usage_routing_reads_tpm_then_rpm_from_redis(tmp_path: Path) -> None for marker in markers[:2] ) assert observed_bodies == expected_bodies, observed_bodies - expected_mget: Final = ( + cooldown_mget: Final = ( + "MGET", + *(f"deployment:{deployment_id}:cooldown" for deployment_id in deployment_ids), + ) + usage_mget: Final = ( "MGET", - f"deployment:{deployment_ids[0]}:cooldown", - f"deployment:{deployment_ids[1]}:cooldown", *(f"{deployment_id}:openai/{OPENAI_MODEL}:tpm:{minute}" for deployment_id in deployment_ids), *(f"{deployment_id}:openai/{OPENAI_MODEL}:rpm:{minute}" for deployment_id in deployment_ids), ) mgets: Final = _drain_mgets(commands) - assert any(arguments == expected_mget for _, arguments in mgets), mgets + mget_arguments: Final = tuple(arguments for _, arguments in mgets) + assert cooldown_mget in mget_arguments, mgets + assert usage_mget in mget_arguments, mgets raw_mgets: Final = tuple(line for line, _ in mgets) print(f"sdk MGETs: {raw_mgets}") finally: diff --git a/tests/integration/streaming/test_stream_clean_eof_without_finish_reason.py b/tests/integration/streaming/test_stream_clean_eof_without_finish_reason.py index f3b427d91c4..933c152eb26 100644 --- a/tests/integration/streaming/test_stream_clean_eof_without_finish_reason.py +++ b/tests/integration/streaming/test_stream_clean_eof_without_finish_reason.py @@ -464,7 +464,7 @@ def test_concurrent_burst_ends_each_stream_by_whether_its_provider_finished_it(g ) -_SPEND_LOG_INSERT: Final = b'INSERT INTO "public"."LiteLLM_SpendLogs"' +_SPEND_LOG_INSERT: Final = b'"LiteLLM_SpendLogs" (' def _spend_statuses(database_url: str | None) -> Counter[tuple[str, str]]: From b28005fdf66d7e5f810ac1f68d87113b2d9fc9fe Mon Sep 17 00:00:00 2001 From: paveln83 <349044532+paveln83@users.noreply.github.com> Date: Mon, 9 Nov 2026 17:23:52 -0800 Subject: [PATCH 45/46] fix(router): merge default metadata under the request's instead of replacing it (#50700) * fix(router): merge default metadata under the request's instead of replacing it * fix(router): join default metadata tags after the request's tags A router default_litellm_params.metadata.tags list replaced the request's tags, so the deployment and credential tags the router adds never reached the logged request tags. Default tags now follow the request's tags, skipping ones already present, like the proxy's key tag merge * fix(router): keep masking a router default turns on Merging default metadata under the request let a caller's mask_input or mask_output false switch off masking an admin set in default_litellm_params.metadata. A masking flag the default turns on now stays on, and a caller can still turn masking on. Also gives this PR's test suppressions a reason and adds a missing blank line * fix(router): keep the router's cache namespace and drop malformed caller tags and dicts A caller's metadata.redis_namespace no longer beats a router default, so it cannot pick another tenant's cache partition. A caller tags value that is not a list, or a non-dict value under a dict default, is treated as absent the way the proxy's merge_tags treats malformed tags, so the router defaults still apply. * test(router): check that mask_output stays on when the router default turns it on * fix(router): keep a router default guardrail whatever guardrails the request lists * fix(router): keep the admin's settings for a router default guardrail when the caller lists it * test(router): check that a router default guardrail entry keeps the admin's settings * fix(router): keep a router default private Langfuse trace private whatever the caller sends * fix(router): keep a router default function such as langfuse_masking_function over the caller's value --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../routing_common_utils.py | 71 +++++ litellm/router_rust/invoke/endpoints.py | 11 +- litellm/router_rust/invoke/shaping.py | 7 +- ...tex_labels_router_default_metadata_wire.py | 85 ++++++ ...st_endpoints_characterization_inference.py | 15 +- .../unit/test_router/core/test_gaps_l2800.py | 259 +++++++++++++++++- 6 files changed, 435 insertions(+), 13 deletions(-) create mode 100644 tests/integration/providers/test_vertex_labels_router_default_metadata_wire.py diff --git a/litellm/litellm_core_utils/routing_common_utils.py b/litellm/litellm_core_utils/routing_common_utils.py index d13ff4059aa..89d6729766f 100644 --- a/litellm/litellm_core_utils/routing_common_utils.py +++ b/litellm/litellm_core_utils/routing_common_utils.py @@ -1,6 +1,7 @@ import hashlib import json from collections.abc import Mapping, Sequence +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Final, TypeVar @@ -13,10 +14,12 @@ import litellm from litellm._logging import verbose_logger, verbose_router_logger from litellm.constants import ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS from litellm.exceptions import BadRequestError +from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure from litellm.llms.sagemaker.common_utils import deployment_inference_component +from litellm.secret_managers.main import str_to_bool from litellm.types.router import ( CredentialLiteLLMParams, Deployment, @@ -62,6 +65,74 @@ 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} +_ROUTER_OWNED_METADATA_KEYS: Final = frozenset({"caching_groups", "redis_namespace"}) +_PRIVATE_SETTING_BY_FLAG: Final = MappingProxyType({"mask_input": True, "mask_output": True, "trace_public": False}) +_JOINED_LIST_METADATA_KEYS: Final = frozenset({"tags", "guardrails"}) + + +def metadata_defaults_under_request( + defaults: Mapping[str, object], request_metadata: Mapping[str, object] +) -> Mapping[str, object]: + return { + key: _default_under_request(key=key, default=default, requested=request_metadata.get(key)) + for key, default in defaults.items() + } + + +def _default_under_request(key: str, default: object, requested: object) -> object: + if key in _ROUTER_OWNED_METADATA_KEYS or callable(default): + return default + if _PRIVATE_SETTING_BY_FLAG.get(key) is _is_flag_on(default): + return default + if key in _JOINED_LIST_METADATA_KEYS: + return _request_items_then_default_items(default=default, requested=requested) + return _request_value_over_default(default=default, requested=requested) + + +def _is_flag_on(value: object) -> bool: + parsed: Final = str_to_bool(value) if isinstance(value, str) else None + return bool(value) if parsed is None else parsed + + +def _request_items_then_default_items(default: object, requested: object) -> object: + default_items: Final = as_object_list(default) + if default_items is None: + return _request_value_over_default(default=default, requested=requested) + default_names: Final = frozenset(chain.from_iterable(_entry_names(item) for item in default_items)) + request_items: Final = tuple( + chain.from_iterable(_without_entries_named(item, default_names) for item in as_object_list(requested) or ()) + ) + return [*request_items, *(item for item in default_items if item not in request_items)] + + +def _entry_names(entry: object) -> frozenset[str]: + mapping: Final = as_str_mapping(entry) + if mapping is not None: + return frozenset(mapping) + return frozenset((entry,)) if isinstance(entry, str) else frozenset() + + +def _without_entries_named(entry: object, names: frozenset[str]) -> tuple[object, ...]: + mapping: Final = as_str_mapping(entry) + if mapping is None or names.isdisjoint(mapping): + return (entry,) + if names.issuperset(mapping): + return () + return ({key: value for key, value in mapping.items() if key not in names},) + + +def _request_value_over_default(default: object, requested: object) -> object: + if requested is None: + return default + default_mapping: Final = as_str_mapping(default) + if default_mapping is None: + return requested + requested_mapping: Final = as_str_mapping(requested) + if requested_mapping is None: + return default + return {**default_mapping, **requested_mapping} + + def as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]: if not isinstance(value, tuple): return () diff --git a/litellm/router_rust/invoke/endpoints.py b/litellm/router_rust/invoke/endpoints.py index 7a5b52761e2..70c83bed728 100644 --- a/litellm/router_rust/invoke/endpoints.py +++ b/litellm/router_rust/invoke/endpoints.py @@ -2118,7 +2118,10 @@ class RustEndpoints: def text_completion(self, request: Mapping[str, object]) -> object: """The sync `text_completion`: one pick, no retries, the group name as the model.""" - from litellm.litellm_core_utils.routing_common_utils import deployment_model_id_pin + from litellm.litellm_core_utils.routing_common_utils import ( + deployment_model_id_pin, + metadata_defaults_under_request, + ) kwargs: Final = without(request, "is_retry", "is_fallback", "is_async") model: Final = kwargs.get("model") @@ -2139,9 +2142,9 @@ class RustEndpoints: if key not in kwargs: kwargs[key] = value elif key == "metadata": - bucket = kwargs[key] - if isinstance(bucket, dict): - bucket.update(as_dict(value)) # pyright: ignore[reportUnknownMemberType] # the request's metadata dict + metadata_bucket(kwargs, key).update( + metadata_defaults_under_request(defaults=as_dict(value), request_metadata=as_dict(kwargs[key])) + ) pinned: Final = deployment_model_id_pin(data) if pinned is not None: kwargs["model_id"] = pinned diff --git a/litellm/router_rust/invoke/shaping.py b/litellm/router_rust/invoke/shaping.py index 258f919c9f1..bf6410d8242 100644 --- a/litellm/router_rust/invoke/shaping.py +++ b/litellm/router_rust/invoke/shaping.py @@ -32,7 +32,7 @@ from litellm.litellm_core_utils.core_helpers import without from litellm.litellm_core_utils.get_litellm_params import SUPPLEMENTAL_KWARGS_KEYS from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider, inferred_provider from litellm.litellm_core_utils.request_timeout_resolver import get_service_tier_timeout -from litellm.litellm_core_utils.routing_common_utils import deployment_model_id_pin +from litellm.litellm_core_utils.routing_common_utils import deployment_model_id_pin, metadata_defaults_under_request from litellm.litellm_core_utils.tool_fields import TOOL_FIELDS, drop_policy from litellm.responses.retrieve_stream_query import requests_streamed_retrieval, without_retrieve_stream_defaults from litellm.router_api import deployment_id as generate_deployment_id @@ -359,14 +359,15 @@ def update_kwargs_with_default_litellm_params( metadata_variable_name: str, ) -> None: """``_update_kwargs_with_default_litellm_params``: unset non-None defaults fill the request; - the default metadata is merged over the request's.""" + the default metadata fills the keys the request leaves unset, the router's caching groups excepted.""" defaults: Final = dict(without_retrieve_stream_defaults(shaping.default_litellm_params, kwargs)) metadata_defaults: Final = defaults.pop("metadata", {}) or {} # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values for key, value in defaults.items(): if value is None: continue kwargs.setdefault(key, value) - metadata_bucket(kwargs, metadata_variable_name).update(as_dict(metadata_defaults)) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + bucket: Final = metadata_bucket(kwargs, metadata_variable_name) + bucket.update(metadata_defaults_under_request(defaults=as_dict(metadata_defaults), request_metadata=bucket)) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values # --- clientside credentials ------------------------------------------------------------------- diff --git a/tests/integration/providers/test_vertex_labels_router_default_metadata_wire.py b/tests/integration/providers/test_vertex_labels_router_default_metadata_wire.py new file mode 100644 index 00000000000..89747a5bed8 --- /dev/null +++ b/tests/integration/providers/test_vertex_labels_router_default_metadata_wire.py @@ -0,0 +1,85 @@ +import json +from collections.abc import Mapping +from pathlib import Path +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_MODEL_NAME: Final = "labelled-gemini" +_PROJECT: Final = "scripted-project" +_LOCATION: Final = "us-central1" +_ROUTER_LABELS: Final = {"env": "prod"} +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_REPLY: Final = json.dumps( + { + "candidates": [{"content": {"role": "model", "parts": [{"text": "OK"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 1, "totalTokenCount": 4}, + } +).encode() + + +def _gemini_peer(request: Request) -> Reply: + if request.target.endswith(":generateContent"): + return Reply(body=_REPLY) + return Reply(status=404, body=json.dumps({"error": request.target}).encode()) + + +def _config(directory: Path, gateway: Gateway, api_base: str) -> Path: + config: Final = directory / "router-default-requester-metadata.yaml" + config.write_text( + json.dumps( + { + "model_list": [ + { + "model_name": _MODEL_NAME, + "litellm_params": { + "model": "vertex_ai/gemini-3.8-flash", + "api_base": api_base, + "vertex_project": _PROJECT, + "vertex_location": _LOCATION, + "vertex_credentials": service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")), + }, + } + ], + "router_settings": {"default_litellm_params": {"metadata": {"requester_metadata": _ROUTER_LABELS}}}, + } + ), + encoding="utf-8", + ) + return config + + +def _labels_sent(proxy: Gateway, wire: Wire, metadata: Mapping[str, JsonValue] | None) -> JsonValue: + response: Final = proxy.request( + "POST", + "/v1/chat/completions", + { + "model": _MODEL_NAME, + "messages": [{"role": "user", "content": "Return exactly OK."}], + **({"metadata": dict(metadata)} if metadata is not None else {}), + }, + ) + assert response.status_code == 200, response.text + generations: Final = tuple(request for request in wire.drain() if request.target.endswith(":generateContent")) + assert len(generations) == 1, generations + return _JSON_OBJECT.validate_json(generations[0].body).get("labels") + + +def test_router_default_requester_metadata_merges_under_each_callers_vertex_labels( + gateway: Gateway, tmp_path: Path +) -> None: + with ( + wire_server(_gemini_peer) as wire, + owned_proxy_process(gateway, tmp_path, {}, config=_config(tmp_path, gateway, wire.url)) as owned, + ): + proxy: Final = Gateway(owned.gateway.client, owned.gateway.key, gateway.upstream_url) + assert _labels_sent(proxy, wire, {"team": "ai-platform"}) == {**_ROUTER_LABELS, "team": "ai-platform"} + assert _labels_sent(proxy, wire, {"env": "staging"}) == {"env": "staging"} + assert _labels_sent(proxy, wire, None) == _ROUTER_LABELS diff --git a/tests/unit/test_router/core/test_endpoints_characterization_inference.py b/tests/unit/test_router/core/test_endpoints_characterization_inference.py index 8db0d6bea12..c975e45b1d3 100644 --- a/tests/unit/test_router/core/test_endpoints_characterization_inference.py +++ b/tests/unit/test_router/core/test_endpoints_characterization_inference.py @@ -205,6 +205,17 @@ async def test_sync_text_completion_merges_default_metadata_into_request_metadat assert (metadata["x"], metadata["y"], metadata["model_group"], sent["user"]) == (1, 2, "g", "u") +async def test_sync_text_completion_keeps_the_request_metadata_over_the_default_metadata( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core( + model_list=[deployment("g", "d1")], + default_litellm_params={"metadata": {"requester_metadata": {"env": "prod", "region": "us"}}}, + ) + core.text_completion(model="g", prompt="p", metadata={"requester_metadata": {"env": "dev", "team": "t1"}}) + assert _bucket(_sent(host))["requester_metadata"] == {"env": "dev", "region": "us", "team": "t1"} + + async def test_sync_text_completion_sends_the_model_group_name_not_the_deployment_model( make_core: MakeCore, host: EndpointHost ): @@ -329,12 +340,12 @@ async def test_default_litellm_params_metadata_merges_into_request_metadata(make assert (metadata["team"], metadata["trace"], metadata["model_group"]) == ("t", "x", "g") -async def test_default_litellm_params_metadata_overwrites_the_same_request_metadata_key( +async def test_the_request_metadata_key_wins_over_the_same_default_litellm_params_metadata_key( make_core: MakeCore, host: EndpointHost ): core: Final = make_core(model_list=[deployment("g", "d1")], default_litellm_params={"metadata": {"team": "t"}}) await core.acompletion(model="g", messages=MESSAGES, metadata={"team": "request"}) - assert _bucket(_sent(host))["team"] == "t" + assert _bucket(_sent(host))["team"] == "request" async def test_default_litellm_params_metadata_goes_to_litellm_metadata_for_generic_endpoints( diff --git a/tests/unit/test_router/core/test_gaps_l2800.py b/tests/unit/test_router/core/test_gaps_l2800.py index 78e237f0696..196ce20d4b3 100644 --- a/tests/unit/test_router/core/test_gaps_l2800.py +++ b/tests/unit/test_router/core/test_gaps_l2800.py @@ -246,7 +246,7 @@ async def test_responses_replace_kubernetes_pod_routing_in_the_request_metadata( """On the Responses path the routing goes to ``litellm_metadata``; a pod routing key the caller put in ``metadata`` is overwritten with the serving deployment's (none here).""" core: Final = make_core(model_list=[deployment("g", "d1")]) - await core.aresponses( # pyright: ignore[reportAttributeAccessIssue] + await core.aresponses( # pyright: ignore[reportAttributeAccessIssue] # endpoint method model="g", input="hi", metadata={KUBERNETES_POD_ROUTING_KEY: {"pod": "stale"}, "keep": 1} ) sent: Final = host.calls[0].kwargs @@ -332,7 +332,7 @@ async def test_the_process_request_timeout_outranks_the_router_timeout_once_a_ro process settings), and only when the router has a ``timeout`` of its own.""" monkeypatch.setattr(litellm, "request_timeout", 33.0) core: Final = make_core(model_list=[deployment("g", "d1")], timeout=42) - await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # endpoint method assert host.calls[0].kwargs["timeout"] == 33.0 @@ -353,7 +353,7 @@ async def test_default_litellm_params_timeout_applies_when_no_router_or_process_ ) -> None: host.set_dynamic(request_timeout=None) core: Final = make_core(model_list=[deployment("g", "d1")], default_litellm_params={"timeout": 9}) - await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # endpoint method assert host.calls[0].kwargs["timeout"] == 9 @@ -368,6 +368,257 @@ async def test_default_litellm_params_stream_timeout_is_the_last_stream_timeout_ assert host.calls[0].kwargs["stream_timeout"] == 4 +_PROXY_REQUESTER_METADATA: Final = {"metadata": {"requester_metadata": {"env": "prod"}}} + + +@pytest.mark.parametrize( + ("request_metadata", "sent_requester_metadata"), + [ + ({"requester_metadata": {"team": "t1"}}, {"env": "prod", "team": "t1"}), + ({"requester_metadata": {"env": "dev", "team": "t1"}}, {"env": "dev", "team": "t1"}), + ({"requester_metadata": None}, {"env": "prod"}), + ({"requester_metadata": {}}, {"env": "prod"}), + ({"requester_metadata": "x"}, {"env": "prod"}), + ({}, {"env": "prod"}), + ], +) +async def test_default_requester_metadata_merges_under_the_request_requester_metadata( + make_core: MakeCore, + host: ScriptedHost, + request_metadata: Mapping[str, object], + sent_requester_metadata: Mapping[str, str], +) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], default_litellm_params=_PROXY_REQUESTER_METADATA) + await core.acompletion(model="g", messages=MESSAGES, metadata=dict(request_metadata)) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert _metadata(host.calls[0].kwargs)["requester_metadata"] == sent_requester_metadata + + +async def test_responses_default_requester_metadata_merges_under_the_request_litellm_metadata( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], default_litellm_params=_PROXY_REQUESTER_METADATA) + await core.aresponses( # pyright: ignore[reportAttributeAccessIssue] # endpoint method + model="g", input="hi", litellm_metadata={"requester_metadata": {"team": "t1"}} + ) + sent: Final = _metadata(host.calls[0].kwargs, "litellm_metadata") + assert sent["requester_metadata"] == {"env": "prod", "team": "t1"} + + +async def test_default_requester_metadata_does_not_carry_one_request_labels_into_the_next( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], default_litellm_params=_PROXY_REQUESTER_METADATA) + await core.acompletion(model="g", messages=MESSAGES, metadata={"requester_metadata": {"team": "t1"}}) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert _metadata(host.calls[1].kwargs)["requester_metadata"] == {"env": "prod"} + + +@pytest.mark.parametrize( + ("request_metadata", "sent_regions"), + [({"regions": ["eu"]}, ["eu"]), ({"regions": None}, ["us"]), ({}, ["us"])], +) +async def test_a_default_metadata_list_applies_only_when_the_request_sets_none( + make_core: MakeCore, host: ScriptedHost, request_metadata: Mapping[str, object], sent_regions: list[str] +) -> None: + core: Final = make_core( + model_list=[deployment("g", "d1")], default_litellm_params={"metadata": {"regions": ["us"]}} + ) + await core.acompletion(model="g", messages=MESSAGES, metadata=dict(request_metadata)) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert _metadata(host.calls[0].kwargs)["regions"] == sent_regions + + +@pytest.mark.parametrize("mask_key", ["mask_input", "mask_output"]) +@pytest.mark.parametrize( + ("default_mask", "request_mask", "sent_mask"), + [ + (True, False, True), + (True, None, True), + ("true", False, "true"), + (False, True, True), + (False, None, False), + ("false", True, True), + ], +) +async def test_default_metadata_masking_stays_on_once_either_side_turns_it_on( + make_core: MakeCore, + host: ScriptedHost, + mask_key: str, + default_mask: object, + request_mask: object, + sent_mask: object, +) -> None: + core: Final = make_core( + model_list=[deployment("g", "d1")], default_litellm_params={"metadata": {mask_key: default_mask}} + ) + await core.acompletion(model="g", messages=MESSAGES, metadata={mask_key: request_mask}) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert _metadata(host.calls[0].kwargs)[mask_key] == sent_mask + + +@pytest.mark.parametrize( + ("default_public", "request_public", "sent_public"), + [ + (False, True, False), + ("false", True, "false"), + (False, None, False), + (True, False, False), + (True, None, True), + ("true", True, True), + ], +) +async def test_a_private_default_trace_stays_private_whatever_the_request_sends( + make_core: MakeCore, + host: ScriptedHost, + default_public: object, + request_public: object, + sent_public: object, +) -> None: + core: Final = make_core( + model_list=[deployment("g", "d1")], default_litellm_params={"metadata": {"trace_public": default_public}} + ) + await core.acompletion(model="g", messages=MESSAGES, metadata={"trace_public": request_public}) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert _metadata(host.calls[0].kwargs)["trace_public"] == sent_public + + +async def test_a_default_redis_namespace_beats_the_request_redis_namespace( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "d1")], default_litellm_params={"metadata": {"redis_namespace": "tenant-a"}} + ) + await core.acompletion(model="g", messages=MESSAGES, metadata={"redis_namespace": "tenant-b"}) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert _metadata(host.calls[0].kwargs)["redis_namespace"] == "tenant-a" + + +def _redact(_value: object) -> str: + return "redacted" + + +@pytest.mark.parametrize("request_value", ["", {"caller": "value"}]) +async def test_a_default_masking_function_beats_whatever_the_request_sends( + make_core: MakeCore, host: ScriptedHost, request_value: object +) -> None: + core: Final = make_core( + model_list=[deployment("g", "d1")], default_litellm_params={"metadata": {"langfuse_masking_function": _redact}} + ) + await core.acompletion(model="g", messages=MESSAGES, metadata={"langfuse_masking_function": request_value}) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert _metadata(host.calls[0].kwargs)["langfuse_masking_function"] is _redact + + +_PROXY_DEFAULT_TAGS: Final = {"metadata": {"tags": ["proxy-default"]}} + + +@pytest.mark.parametrize( + ("request_metadata", "sent_tags"), + [ + ({"tags": ["caller"]}, ["caller", "proxy-default"]), + ({"tags": ["proxy-default", "caller"]}, ["proxy-default", "caller"]), + ({"tags": None}, ["proxy-default"]), + ({}, ["proxy-default"]), + ({"tags": "caller"}, ["proxy-default"]), + ], +) +async def test_default_metadata_tags_follow_the_request_tags( + make_core: MakeCore, host: ScriptedHost, request_metadata: Mapping[str, object], sent_tags: object +) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], default_litellm_params=_PROXY_DEFAULT_TAGS) + await core.acompletion(model="g", messages=MESSAGES, metadata=dict(request_metadata)) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert _metadata(host.calls[0].kwargs)["tags"] == sent_tags + + +@pytest.mark.parametrize( + ("request_metadata", "sent_tags"), + [ + ({"tags": ["caller"]}, ["caller"]), + ({}, "proxy-default"), + ], +) +async def test_a_default_tags_that_is_not_a_list_yields_to_the_request_tags( + make_core: MakeCore, host: ScriptedHost, request_metadata: Mapping[str, object], sent_tags: object +) -> None: + core: Final = make_core( + model_list=[deployment("g", "d1")], default_litellm_params={"metadata": {"tags": "proxy-default"}} + ) + await core.acompletion(model="g", messages=MESSAGES, metadata=dict(request_metadata)) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert _metadata(host.calls[0].kwargs)["tags"] == sent_tags + + +@pytest.mark.parametrize( + ("request_metadata", "sent_tags"), + [ + ({}, ["t1", "Credential: cred", "proxy-default"]), + ({"tags": ["caller"]}, ["caller", "t1", "Credential: cred", "proxy-default"]), + ], +) +async def test_default_metadata_tags_join_the_deployment_and_credential_tags( + make_core: MakeCore, host: ScriptedHost, request_metadata: Mapping[str, object], sent_tags: list[str] +) -> None: + core: Final = make_core( + model_list=[deployment("g", "d1", tags=["t1"], litellm_credential_name="cred")], + default_litellm_params=_PROXY_DEFAULT_TAGS, + ) + await core.acompletion(model="g", messages=MESSAGES, metadata=dict(request_metadata)) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert _metadata(host.calls[0].kwargs)["tags"] == sent_tags + + +async def test_default_metadata_tags_do_not_carry_one_request_tags_into_the_next( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], default_litellm_params=_PROXY_DEFAULT_TAGS) + await core.acompletion(model="g", messages=MESSAGES, metadata={"tags": ["caller"]}) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert _metadata(host.calls[1].kwargs)["tags"] == ["proxy-default"] + + +@pytest.mark.parametrize( + ("request_metadata", "sent_guardrails"), + [ + ({"guardrails": []}, ["admin-guard"]), + ({"guardrails": ["caller-guard"]}, ["caller-guard", "admin-guard"]), + ({"guardrails": "caller-guard"}, ["admin-guard"]), + ({}, ["admin-guard"]), + ({"guardrails": [{"admin-guard": {"extra_body": {"threshold": 1}}}]}, ["admin-guard"]), + ( + { + "guardrails": [ + {"admin-guard": {"extra_body": {"threshold": 1}}, "caller-guard": {"extra_body": {"k": 2}}} + ] + }, + [{"caller-guard": {"extra_body": {"k": 2}}}, "admin-guard"], + ), + ], +) +async def test_a_default_guardrail_stays_on_with_its_own_settings_whatever_the_request_lists( + make_core: MakeCore, host: ScriptedHost, request_metadata: Mapping[str, object], sent_guardrails: list[object] +) -> None: + core: Final = make_core( + model_list=[deployment("g", "d1")], default_litellm_params={"metadata": {"guardrails": ["admin-guard"]}} + ) + await core.acompletion(model="g", messages=MESSAGES, metadata=dict(request_metadata)) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert _metadata(host.calls[0].kwargs)["guardrails"] == sent_guardrails + + +async def test_a_default_guardrail_entry_keeps_the_admins_settings_over_the_callers( + make_core: MakeCore, host: ScriptedHost +) -> None: + admin_entry: Final = {"admin-guard": {"extra_body": {"threshold": 0.9}}} + core: Final = make_core( + model_list=[deployment("g", "d1")], default_litellm_params={"metadata": {"guardrails": [admin_entry]}} + ) + await core.acompletion( # pyright: ignore[reportAttributeAccessIssue] # endpoint method + model="g", messages=MESSAGES, metadata={"guardrails": [{"admin-guard": {"extra_body": {"threshold": 0}}}]} + ) + assert _metadata(host.calls[0].kwargs)["guardrails"] == [admin_entry] + + +@pytest.mark.parametrize("router_caching_groups", [[("g", "other")], None]) +async def test_the_router_caching_groups_replace_the_request_caching_groups( + make_core: MakeCore, host: ScriptedHost, router_caching_groups: list[tuple[str, str]] | None +) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], caching_groups=router_caching_groups) + await core.acompletion(model="g", messages=MESSAGES, metadata={"caching_groups": [("g", "premium")]}) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert _metadata(host.calls[0].kwargs)["caching_groups"] == router_caching_groups + + # --- sync completion runs the routing pre-call checks for a model group only ------------------- @@ -412,7 +663,7 @@ async def test_acompletion_counts_total_success_and_failed_calls_per_deployment_ make_core: MakeCore, host: ScriptedHost ) -> None: core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0) - await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # endpoint method host.answer("openai/d1", rate_limit()) with pytest.raises(litellm.RateLimitError): await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] From cd9f98d09576d206526e286523b2b98364bcc300 Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Mon, 9 Nov 2026 17:35:35 -0800 Subject: [PATCH 46/46] fix(anthropic-messages): apply additional_drop_params before the unsupported param checks (#50753) * fix(anthropic-messages): apply additional_drop_params before the unsupported-param gates * fix(anthropic-messages): apply additional_drop_params before the unsupported param checks Share the tolerant drop list parser between the Python handler and the Rust host route, clear listed params before the Rust gate, and keep the prompt cache baseline body in step with the real request * fix(anthropic-messages): leave additional_drop_params keys out of the Rust host body * test(anthropic-messages): annotate the new test locals with Final --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../crates/inference-messages/src/handler.rs | 8 +- .../inference-messages/tests/messages/host.rs | 28 +++++ .../tests/messages/request.rs | 24 ++++ .../llms/anthropic/messages/native/handler.py | 12 +- .../llms/anthropic/messages/native/utils.py | 5 +- .../llms/anthropic/prompt_cache_prediction.py | 8 +- litellm/rust_bridge/host/routes/messages.py | 13 +- ...st_messages_additional_drop_params_wire.py | 119 ++++++++++++++++++ .../native/test_anthropic_messages_handler.py | 83 ++++++++++++ .../test_request_optional_param_utils.py | 36 ++++++ .../test_anthropic_prompt_cache_prediction.py | 20 +++ .../rust_bridge/host/routes/test_messages.py | 22 ++++ 12 files changed, 356 insertions(+), 22 deletions(-) create mode 100644 tests/integration/messages_endpoint/test_messages_additional_drop_params_wire.py diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index 6a8ed7fff8e..c7a7d93bf98 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -133,10 +133,10 @@ async fn provider_call( )?; let transform_context = MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params); - let filtered = requested_optional_params(sanitized, provider, &transform_context)?; - let requested_params = serde_json::to_value(&filtered.params).map_err(serialize_failure)?; - let trimmed = - without_additional_drop_params(filtered, &shaping.settings.additional_drop_params)?; + let without_listed = + without_additional_drop_params(sanitized, &shaping.settings.additional_drop_params)?; + let trimmed = requested_optional_params(without_listed, provider, &transform_context)?; + let requested_params = serde_json::to_value(&trimmed.params).map_err(serialize_failure)?; let scoped = get_provider_specific_headers(provider_specific_header.as_ref(), provider.as_str()); diff --git a/litellm-rust/crates/inference-messages/tests/messages/host.rs b/litellm-rust/crates/inference-messages/tests/messages/host.rs index 20f6a97aa37..063e9d18738 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/host.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/host.rs @@ -456,6 +456,34 @@ async fn the_request_context_carries_the_shaped_params_without_model_or_messages assert_eq!(optional_params, json!({"max_tokens": 16})); } +#[rstest] +#[tokio::test] +async fn the_request_context_leaves_out_a_listed_param_the_model_refuses(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + let host = RecordingHost::passthrough(authenticated( + MessagesCall { + shaping: MessagesShaping { + settings: MessagesSettings { + additional_drop_params: vec!["temperature".into()], + ..MessagesSettings::default() + }, + capabilities: AnthropicModelCapabilities { + supports_sampling_params: false, + ..AnthropicModelCapabilities::default() + }, + }, + ..with_fields(call, json!({"temperature": 0.2})) + }, + upstream.uri(), + )); + + run_through(&host).await.expect("messages call succeeds"); + + let [optional_params] = <[Value; 1]>::try_from(host.optional_params.into_inner().unwrap()) + .unwrap_or_else(|seen| panic!("before_provider_request runs once, saw {}", seen.len())); + assert_eq!(optional_params, json!({"max_tokens": 16})); +} + #[rstest] #[case::reasoning_effort_on_an_adaptive_model( AnthropicModelCapabilities { diff --git a/litellm-rust/crates/inference-messages/tests/messages/request.rs b/litellm-rust/crates/inference-messages/tests/messages/request.rs index 216aa96fbe6..3784eb6cc68 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/request.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/request.rs @@ -530,6 +530,30 @@ async fn unsupported_params_are_dropped_under_drop_params_and_rejected_without_i assert_eq!(sent["max_tokens"], 16); } +#[rstest] +#[tokio::test] +async fn additional_drop_params_clear_a_param_the_model_does_not_support(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + + run_message(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + shaping: MessagesShaping { + capabilities: sampling_removed(), + settings: MessagesSettings { + additional_drop_params: vec!["top_k".into()], + ..MessagesSettings::default() + }, + }, + ..with_fields(call, json!({"top_k": 5})) + }) + .await; + + let sent = only_request(&upstream).await.json(); + assert_eq!(sent.get("top_k"), None); + assert_eq!(sent["max_tokens"], 16); +} + #[rstest] #[case::adaptive_thinking(json!({"type": "adaptive"}), json!({"type": "adaptive", "display": "summarized"}))] #[case::disabled_thinking(json!({"type": "disabled"}), json!({"type": "disabled"}))] diff --git a/litellm/llms/anthropic/messages/native/handler.py b/litellm/llms/anthropic/messages/native/handler.py index 922de386fe4..59ffb6fef04 100644 --- a/litellm/llms/anthropic/messages/native/handler.py +++ b/litellm/llms/anthropic/messages/native/handler.py @@ -12,7 +12,7 @@ from dataclasses import dataclass from functools import partial from typing import Final, cast -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError import litellm from litellm.anthropic_beta_headers_manager import aload_beta_headers_config @@ -67,6 +67,15 @@ __all__ = ("anthropic_messages", "anthropic_messages_handler") _RESPONSES_API_PROVIDERS: Final = frozenset({"openai"}) _SAFEGUARD_KWARGS: Final = TypeAdapter(Mapping[str, object]) _SAFEGUARD_MESSAGES: Final = TypeAdapter(Sequence[Mapping[str, object]]) +_DROP_PATHS: Final = TypeAdapter(list[object]) + + +def additional_drop_paths(kwargs: Mapping[str, object]) -> tuple[str, ...]: + try: + configured: Final = _DROP_PATHS.validate_python(kwargs.get("additional_drop_params")) + except ValidationError: + return () + return tuple(path for path in configured if isinstance(path, str)) def _bridges_to_responses_api(model: str, custom_llm_provider: str, deployment_model_info: object) -> bool: @@ -837,6 +846,7 @@ def anthropic_messages_handler( model=model, drop_params=litellm_params.get("drop_params") is True, custom_llm_provider=custom_llm_provider, + additional_drop_params=additional_drop_paths(request_kwargs), ) ) resolved_api_base: Final = ( diff --git a/litellm/llms/anthropic/messages/native/utils.py b/litellm/llms/anthropic/messages/native/utils.py index ea3afdb9a8f..8f33d000c7c 100644 --- a/litellm/llms/anthropic/messages/native/utils.py +++ b/litellm/llms/anthropic/messages/native/utils.py @@ -299,6 +299,7 @@ class AnthropicMessagesRequestUtils: model: str | None = None, drop_params: bool = False, custom_llm_provider: str | None = None, + additional_drop_params: Iterable[str] = (), ) -> AnthropicMessagesRequestOptionalParams: """ Filter parameters to only include those defined in AnthropicMessagesRequestOptionalParams. @@ -308,11 +309,13 @@ class AnthropicMessagesRequestUtils: model: Resolved model id; when set, unsupported params may be dropped drop_params: Per-request drop_params flag (also respects litellm.drop_params) custom_llm_provider: Routed provider; fast mode is gated to direct Anthropic + additional_drop_params: Paths the deployment drops. A listed top-level key is left out + before the unsupported-param checks run, so it never raises Returns: AnthropicMessagesRequestOptionalParams instance with only the valid parameters """ - valid_keys: Final = anthropic_messages_optional_param_keys() + valid_keys: Final = anthropic_messages_optional_param_keys().difference(additional_drop_params) filtered_params: Final = {k: v for k, v in params.items() if k in valid_keys and v is not None} if model is not None: from litellm.llms.anthropic.chat.transformation import AnthropicConfig diff --git a/litellm/llms/anthropic/prompt_cache_prediction.py b/litellm/llms/anthropic/prompt_cache_prediction.py index 2da071ebdd2..28896c99eaf 100644 --- a/litellm/llms/anthropic/prompt_cache_prediction.py +++ b/litellm/llms/anthropic/prompt_cache_prediction.py @@ -654,17 +654,15 @@ def prepare_native_baseline_body(request: Mapping[str, object], model: str) -> M custom_llm_provider="anthropic", tools=_MESSAGES.validate_python(owned.get("tools") or []), ) + drop_paths: Final = TypeAdapter(tuple[str, ...]).validate_python(owned.get("additional_drop_params") or ()) options: Final = AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param( {**owned, "system": system}, model=resolved_model, custom_llm_provider="anthropic", drop_params=owned.get("drop_params") is True, + additional_drop_params=drop_paths, ) - filtered: Final = reduce( - delete_nested_value, - TypeAdapter(tuple[str, ...]).validate_python(owned.get("additional_drop_params") or ()), - dict(options), - ) + filtered: Final = reduce(delete_nested_value, drop_paths, dict(options)) config: Final = _BaselineAnthropicMessagesConfig() body: Final = config.transform_anthropic_messages_request( model=resolved_model, diff --git a/litellm/rust_bridge/host/routes/messages.py b/litellm/rust_bridge/host/routes/messages.py index c50791f8a9c..c6568264edb 100644 --- a/litellm/rust_bridge/host/routes/messages.py +++ b/litellm/rust_bridge/host/routes/messages.py @@ -35,7 +35,6 @@ from litellm.rust_bridge.host import callbacks_legacy_python, http_transport from litellm.rust_bridge.public_call import known_params, optional_str from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse -_DROP_PATHS: Final = TypeAdapter(list[object]) _METADATA: Final = TypeAdapter(dict[str, object]) _BODY_FIELDS: Final = frozenset( { @@ -92,7 +91,7 @@ def inferred_provider(model: str, api_base: str | None) -> str | None: def body_params(request: Mapping[str, object]) -> Mapping[str, object]: - params: Final = known_params(_BODY_FIELDS, request) + params: Final = known_params(_BODY_FIELDS.difference(handler.additional_drop_paths(request)), request) if deployment_reasoning_effort_levels(request.get("model_info")) is None: return params return _with_deployment_effort(params, request) @@ -386,20 +385,12 @@ def _drop_params(kwargs: Mapping[str, object]) -> bool: return bool(litellm.drop_params) or normalize_drop_params(kwargs.get("drop_params")) is True -def _additional_drop_params(kwargs: Mapping[str, object]) -> tuple[str, ...]: - try: - configured: Final = _DROP_PATHS.validate_python(kwargs.get("additional_drop_params")) - except ValidationError: - return () - return tuple(path for path in configured if isinstance(path, str)) - - def settings(kwargs: Mapping[str, object]) -> dict[str, object]: return asdict( MessagesSettings( drop_params=_drop_params(kwargs), reasoning_auto_summary=is_reasoning_auto_summary_enabled(), - additional_drop_params=_additional_drop_params(kwargs), + additional_drop_params=handler.additional_drop_paths(kwargs), request_timeout=get_configured_request_timeout(), ) ) diff --git a/tests/integration/messages_endpoint/test_messages_additional_drop_params_wire.py b/tests/integration/messages_endpoint/test_messages_additional_drop_params_wire.py new file mode 100644 index 00000000000..f19976bda74 --- /dev/null +++ b/tests/integration/messages_endpoint/test_messages_additional_drop_params_wire.py @@ -0,0 +1,119 @@ +import json +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final + +import pytest +from integration._support.client import Gateway, list_value, object_value +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +_JSON: Final = TypeAdapter(dict[str, JsonValue]) +_ANSWER: Final = "sampling params dropped" +_SAMPLING_SPELLINGS: Final = ("top_k", "top_p", "topK", "topP") +_AWS: Final[dict[str, JsonValue]] = { + "aws_access_key_id": "AKIASCRIPTEDPROVIDER", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": "us-east-1", +} + + +@dataclass(frozen=True, slots=True) +class Backend: + model: str + path_suffix: str + credentials: Mapping[str, JsonValue] + + +_BEDROCK_INVOKE: Final = Backend("bedrock/invoke/us.anthropic.claude-opus-4-8", "/invoke", _AWS) +_ANTHROPIC: Final = Backend("anthropic/claude-opus-4-8", "/v1/messages", {}) +_BEDROCK_CONVERSE: Final = Backend("bedrock/converse/us.anthropic.claude-opus-4-8", "/converse", _AWS) +_NATIVE_BACKENDS: Final = ( + pytest.param(_BEDROCK_INVOKE, id="bedrock_invoke"), + pytest.param(_ANTHROPIC, id="anthropic"), +) + + +def _answer(request: Request) -> Reply: + if request.target.endswith("/converse"): + return Reply( + body=json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": _ANSWER}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 5, "outputTokens": 3, "totalTokens": 8}, + "metrics": {"latencyMs": 1}, + } + ).encode() + ) + return Reply( + body=json.dumps( + { + "id": "msg_dropped_sampling_params", + "type": "message", + "role": "assistant", + "model": "claude-opus-4-8", + "content": [{"type": "text", "text": _ANSWER}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 3}, + } + ).encode() + ) + + +def _messages(gateway: Gateway, model: str, **sampling: JsonValue) -> tuple[int, str]: + response: Final = gateway.request( + "POST", + "/v1/messages", + {"model": model, "max_tokens": 32, "messages": [{"role": "user", "content": "hi"}], **sampling}, + ) + return response.status_code, response.text + + +def _answered(status: int, text: str) -> None: + assert status == 200, text + content: Final = list_value(_JSON.validate_json(text)["content"]) + assert object_value(content[0])["text"] == _ANSWER, text + + +def _sent_without_sampling_params(wire: Wire, backend: Backend) -> None: + (request,) = wire.drain() + assert request.target.endswith(backend.path_suffix), request.target + sent: Final = request.body.decode() + assert not any(spelling in sent for spelling in _SAMPLING_SPELLINGS), sent + + +@pytest.mark.parametrize( + "backend", (*_NATIVE_BACKENDS, pytest.param(_BEDROCK_CONVERSE, id="bedrock_converse_chat_bridge")) +) +def test_messages_drops_the_sampling_params_a_deployment_lists_on_a_model_without_them( + gateway: Gateway, backend: Backend +) -> None: + with wire_server(_answer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=backend.model, + api_base=wire.url, + additional_drop_params=["top_k", "top_p"], + **backend.credentials, + ) + _answered(*_messages(gateway, model, top_k=5, top_p=0.9)) + _sent_without_sampling_params(wire, backend) + + +@pytest.mark.parametrize("backend", _NATIVE_BACKENDS) +def test_messages_rejects_a_sampling_param_the_deployment_does_not_list_before_calling_the_provider( + gateway: Gateway, backend: Backend +) -> None: + with wire_server(_answer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=backend.model, + api_base=wire.url, + additional_drop_params=["top_k"], + **backend.credentials, + ) + rejected_status, rejected_text = _messages(gateway, model, top_k=5, top_p=0.9) + assert rejected_status == 400, rejected_text + assert "does not support top_p=0.9" in rejected_text, rejected_text + _answered(*_messages(gateway, model, top_k=5)) + _sent_without_sampling_params(wire, backend) diff --git a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py index 3296163c64e..a83e98afb89 100644 --- a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py +++ b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py @@ -2021,6 +2021,89 @@ async def test_anthropic_messages_forwards_safeguards_and_dangerous_tool_use_bet assert response["safeguard_results"] == safeguard_results +_OPUS_4_8_ON_BEDROCK_INVOKE: Final = "bedrock/invoke/us.anthropic.claude-opus-4-8" + + +def _bedrock_invoke_recording(bodies: list[dict[str, object]]) -> AsyncHTTPHandler: + def record_and_answer(request: httpx.Request) -> httpx.Response: + bodies.append(json.loads(request.content)) + return httpx.Response( + 200, + json={ + "id": "msg_drop_top_k", + "type": "message", + "role": "assistant", + "model": "claude-opus-4-8", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 3, "output_tokens": 1}, + }, + request=request, + ) + + upstream: Final = AsyncHTTPHandler() + upstream.client = httpx.AsyncClient(transport=httpx.MockTransport(record_and_answer)) + return upstream + + +async def _opus_4_8_messages_with_top_k(bodies: list[dict[str, object]], **deployment: object) -> object: + return await messages.acreate( + max_tokens=16, + messages=[{"role": "user", "content": "hi"}], + model=_OPUS_4_8_ON_BEDROCK_INVOKE, + top_k=5, + aws_access_key_id="test-access-key", + aws_secret_access_key="test-secret-key", + aws_region_name="us-east-1", + client=_bedrock_invoke_recording(bodies), + **deployment, + ) + + +@pytest.mark.asyncio +async def test_bedrock_invoke_messages_drops_top_k_the_deployment_lists_on_a_model_without_sampling_params( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + bodies: Final[list[dict[str, object]]] = [] + + response: Final = await _opus_4_8_messages_with_top_k(bodies, additional_drop_params=["top_k"]) + + assert response["content"] == [{"type": "text", "text": "ok"}] + (body,) = bodies + assert "top_k" not in body + assert body["max_tokens"] == 16 + + +@pytest.mark.asyncio +async def test_bedrock_invoke_messages_still_rejects_top_k_the_deployment_does_not_list( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + bodies: Final[list[dict[str, object]]] = [] + + with pytest.raises(litellm.UnsupportedParamsError, match="does not support top_k=5") as raised: + await _opus_4_8_messages_with_top_k(bodies, additional_drop_params=["top_p"]) + + assert raised.value.status_code == 400 + assert bodies == [] + + +@pytest.mark.asyncio +async def test_bedrock_invoke_messages_answers_the_top_k_400_for_a_drop_list_that_is_not_a_list( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + bodies: Final[list[dict[str, object]]] = [] + + with pytest.raises(litellm.UnsupportedParamsError, match="does not support top_k=5") as raised: + await _opus_4_8_messages_with_top_k(bodies, additional_drop_params="top_k") + + assert raised.value.status_code == 400 + assert bodies == [] + + @pytest.mark.asyncio @pytest.mark.parametrize("client_headers", _CLIENT_BETA_HEADERS) async def test_anthropic_messages_forwards_safeguards_and_dangerous_tool_use_beta_to_vertex( diff --git a/tests/unit/llms/anthropic/messages/native/test_request_optional_param_utils.py b/tests/unit/llms/anthropic/messages/native/test_request_optional_param_utils.py index d525ccfec37..8104c0cef56 100644 --- a/tests/unit/llms/anthropic/messages/native/test_request_optional_param_utils.py +++ b/tests/unit/llms/anthropic/messages/native/test_request_optional_param_utils.py @@ -6,6 +6,8 @@ Regression tests for the /v1/messages request-parse fast paths: while resolving the (static) type hints only once per process. """ +from typing import Final + import pytest import litellm @@ -133,3 +135,37 @@ def test_sampling_param_raises_clean_400_without_drop_params(monkeypatch): model="claude-opus-4-7", drop_params=False, ) + + +@pytest.mark.parametrize( + ("model", "param", "value"), + [ + ("claude-opus-4-8", "temperature", 0.3), + ("claude-opus-4-8", "top_p", 0.9), + ("claude-opus-4-8", "top_k", 5), + ("claude-sonnet-4-6", "speed", "fast"), + ], +) +def test_deployment_listed_param_is_dropped_before_the_unsupported_param_gate(monkeypatch, model, param, value): + monkeypatch.setattr(litellm, "drop_params", False) + result: Final = AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param( + params={param: value, "stream": True}, + model=model, + drop_params=False, + additional_drop_params=(param,), + ) + + assert result == {"stream": True} + + +def test_param_missing_from_the_drop_list_still_raises_the_unsupported_param_400(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.utils.UnsupportedParamsError, match="does not support top_p") as raised: + AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param( + params={"top_k": 5, "top_p": 0.9}, + model="claude-opus-4-8", + drop_params=False, + additional_drop_params=("top_k", "top_p.nested"), + ) + + assert raised.value.status_code == 400 diff --git a/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py b/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py index d95f831df3e..89ea412af0c 100644 --- a/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py +++ b/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py @@ -618,6 +618,26 @@ def test_native_baseline_body_drops_a_non_numeric_budget_without_logging_a_drop( assert [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] == [] +@pytest.mark.usefixtures("local_model_cost_map") +def test_native_baseline_body_leaves_out_a_listed_sampling_param_the_model_does_not_support( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + request: Final = { + "model": "claude-opus-4-8", + "max_tokens": 64, + "messages": [{"role": "user", "content": "hi"}], + "top_k": 5, + "additional_drop_params": ["top_k"], + } + + body: Final = prepare_native_baseline_body(request, "anthropic/claude-opus-4-8") + + assert body is not None + assert "top_k" not in body + assert body["max_tokens"] == 64 + + @pytest.mark.usefixtures("local_model_cost_map") def test_native_baseline_body_drops_a_non_string_reasoning_effort_without_logging_a_drop( caplog: pytest.LogCaptureFixture, diff --git a/tests/unit/rust_bridge/host/routes/test_messages.py b/tests/unit/rust_bridge/host/routes/test_messages.py index 704d52386a8..aa7dcc74841 100644 --- a/tests/unit/rust_bridge/host/routes/test_messages.py +++ b/tests/unit/rust_bridge/host/routes/test_messages.py @@ -184,6 +184,28 @@ def test_body_params_keep_the_requested_effort_without_deployment_levels() -> No assert messages_host.body_params(request)["output_config"] == {"effort": "max"} +@pytest.mark.parametrize( + ("additional_drop_params", "expected_keys"), + [ + (["thinking"], {"max_tokens"}), + (["top_k"], {"max_tokens", "thinking", "output_config"}), + ], +) +def test_body_params_derive_no_effort_from_a_param_the_deployment_lists( + additional_drop_params: list[str], expected_keys: set[str] +) -> None: + request: Final[dict[str, object]] = { + "model": _ADAPTIVE_OPUS, + "messages": [], + "max_tokens": 8192, + "thinking": {"type": "enabled", "budget_tokens": 4096}, + "additional_drop_params": additional_drop_params, + "model_info": {"id": "deployment-1", "reasoning_effort_levels": ["low", "medium", "high"]}, + } + + assert set(messages_host.body_params(request)).difference({"messages", "model"}) == expected_keys + + def _marked(phase: str | None = None, *args: object) -> ValueError: error: Final = ValueError(*args) if phase is not None: