From 8d1d740bb919bad6a87dbb79144c944949c2dc4a Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 2 Oct 2026 18:12:56 -0700 Subject: [PATCH 01/12] fix(proxy): stop queued registry read-throughs spending the resync budget (#44277) * fix(proxy): stop queued registry read-throughs spending the resync budget RegistryReadThrough.attempt serializes misses behind one lock, but every request that was queued behind the first one still spent a unit of the 20-per-5s resync budget and re-ran the DB resync, even though the first request had already loaded the object. A burst of more than 20 requests for a model created on another worker therefore exhausted the budget and the rest got 400 "Invalid model name". attempt now checks whether the key is already loaded once it holds the lock and returns early without touching the budget. Models check the router's model names and deployment ids, guardrails and agents reuse their existing registry lookups. * test(proxy): gate the queued read-through test on events and cover each registry's loaded check The queued-requests test now holds the first resync on an asyncio.Event instead of a timed sleep and records calls in a recorder with tuple and frozenset state. New tests show the model, guardrail and agent read-throughs each answer an object that is already loaded without reading the database, so rewiring any registry's loaded check now fails a test * test(proxy): keep the queued read-through recorder inside its test and type the agent registry fixture (cherry picked from commit b877a38e5fc41d797f67fd973d8ccad53b23fbac) --- .../common_utils/registry_read_through.py | 28 ++- .../test_registry_read_through.py | 166 ++++++++++++++++-- tests/unit/proxy/conftest.py | 4 +- 3 files changed, 184 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/common_utils/registry_read_through.py b/litellm/proxy/common_utils/registry_read_through.py index 8a82e253c5c..32beb0b7447 100644 --- a/litellm/proxy/common_utils/registry_read_through.py +++ b/litellm/proxy/common_utils/registry_read_through.py @@ -36,6 +36,7 @@ READ_THROUGH_MAX_RESYNCS_PER_WINDOW: Final = 20 class RegistryReadThrough: __slots__ = ( + "_is_loaded", "_lock", "_max_resyncs_per_window", "_miss_ttl_seconds", @@ -49,11 +50,13 @@ class RegistryReadThrough: def __init__( self, resync: Callable[[str], Awaitable[bool]], + is_loaded: Callable[[str], bool], miss_ttl_seconds: float = READ_THROUGH_MISS_TTL_SECONDS, max_resyncs_per_window: int = READ_THROUGH_MAX_RESYNCS_PER_WINDOW, resync_window_seconds: float = READ_THROUGH_RESYNC_WINDOW_SECONDS, ) -> None: self._resync = resync + self._is_loaded = is_loaded self._miss_ttl_seconds = miss_ttl_seconds self._max_resyncs_per_window = max_resyncs_per_window self._resync_window_seconds = resync_window_seconds @@ -78,6 +81,8 @@ class RegistryReadThrough: async with self._lock: if self._recent_misses.get_cache(key) is not None: return False + if self._is_loaded(key): + return True if not self._consume_resync_budget(): verbose_proxy_logger.warning( "registry read-through for %r skipped: resync budget of %s per %ss exhausted", @@ -190,9 +195,26 @@ async def _resync_agents(agent_id_or_name: str) -> bool: return True -model_registry_read_through: Final = RegistryReadThrough(resync=_resync_model_deployments) -guardrail_registry_read_through: Final = RegistryReadThrough(resync=_resync_guardrails) -agent_registry_read_through: Final = RegistryReadThrough(resync=_resync_agents) +def _model_is_loaded(model_name_or_id: str) -> bool: + from litellm.proxy import proxy_server + + router: Final = proxy_server.llm_router + if router is None: + return False + return model_name_or_id in router.model_names or router.has_model_id(model_name_or_id) + + +def _guardrail_is_loaded(guardrail_name: str) -> bool: + return _initialized_guardrail(guardrail_name) is not None + + +def _agent_is_loaded(agent_id_or_name: str) -> bool: + return _agent_from_registry(agent_id_or_name) is not None + + +model_registry_read_through: Final = RegistryReadThrough(resync=_resync_model_deployments, is_loaded=_model_is_loaded) +guardrail_registry_read_through: Final = RegistryReadThrough(resync=_resync_guardrails, is_loaded=_guardrail_is_loaded) +agent_registry_read_through: Final = RegistryReadThrough(resync=_resync_agents, is_loaded=_agent_is_loaded) def _agent_from_registry(agent_id_or_name: str) -> "AgentResponse | None": diff --git a/tests/unit/proxy/common_utils/test_registry_read_through.py b/tests/unit/proxy/common_utils/test_registry_read_through.py index 9e20386bf3d..5713f7dfaa7 100644 --- a/tests/unit/proxy/common_utils/test_registry_read_through.py +++ b/tests/unit/proxy/common_utils/test_registry_read_through.py @@ -1,10 +1,17 @@ import asyncio -from typing import Final +from typing import TYPE_CHECKING, Final import pytest from litellm.proxy.common_utils.registry_read_through import RegistryReadThrough +if TYPE_CHECKING: + from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry + + +def nothing_loaded(_key: str) -> bool: + return False + class ResyncSpy: def __init__(self, found: bool = True, error: Exception | None = None) -> None: @@ -22,7 +29,7 @@ class ResyncSpy: @pytest.mark.asyncio async def test_attempt_returns_true_when_resync_finds_object(): spy: Final = ResyncSpy(found=True) - read_through: Final = RegistryReadThrough(resync=spy) + read_through: Final = RegistryReadThrough(resync=spy, is_loaded=nothing_loaded) assert await read_through.attempt("new-model") is True assert spy.calls == ["new-model"] @@ -31,7 +38,7 @@ async def test_attempt_returns_true_when_resync_finds_object(): @pytest.mark.asyncio async def test_attempt_found_key_is_not_negative_cached(): spy: Final = ResyncSpy(found=True) - read_through: Final = RegistryReadThrough(resync=spy) + read_through: Final = RegistryReadThrough(resync=spy, is_loaded=nothing_loaded) assert await read_through.attempt("new-model") is True assert await read_through.attempt("new-model") is True @@ -41,7 +48,7 @@ async def test_attempt_found_key_is_not_negative_cached(): @pytest.mark.asyncio async def test_missing_key_is_negative_cached_within_ttl(): spy: Final = ResyncSpy(found=False) - read_through: Final = RegistryReadThrough(resync=spy, miss_ttl_seconds=60.0) + read_through: Final = RegistryReadThrough(resync=spy, is_loaded=nothing_loaded, miss_ttl_seconds=60.0) assert await read_through.attempt("ghost-model") is False assert await read_through.attempt("ghost-model") is False @@ -51,7 +58,7 @@ async def test_missing_key_is_negative_cached_within_ttl(): @pytest.mark.asyncio async def test_negative_cache_expires_and_resync_runs_again(): spy: Final = ResyncSpy(found=False) - read_through: Final = RegistryReadThrough(resync=spy, miss_ttl_seconds=0.05) + read_through: Final = RegistryReadThrough(resync=spy, is_loaded=nothing_loaded, miss_ttl_seconds=0.05) assert await read_through.attempt("ghost-model") is False await asyncio.sleep(0.1) @@ -62,7 +69,7 @@ async def test_negative_cache_expires_and_resync_runs_again(): @pytest.mark.asyncio async def test_resync_exception_returns_false_without_negative_caching(): spy: Final = ResyncSpy(error=RuntimeError("db down")) - read_through: Final = RegistryReadThrough(resync=spy) + read_through: Final = RegistryReadThrough(resync=spy, is_loaded=nothing_loaded) assert await read_through.attempt("new-model") is False assert await read_through.attempt("new-model") is False @@ -77,7 +84,7 @@ async def test_concurrent_attempts_for_missing_key_resync_once(): return await super().__call__(key) spy: Final = SlowResyncSpy(found=False) - read_through: Final = RegistryReadThrough(resync=spy, miss_ttl_seconds=60.0) + read_through: Final = RegistryReadThrough(resync=spy, is_loaded=nothing_loaded, miss_ttl_seconds=60.0) results: Final = await asyncio.gather(*(read_through.attempt("ghost-model") for _ in range(5))) assert results == [False] * 5 @@ -87,7 +94,7 @@ async def test_concurrent_attempts_for_missing_key_resync_once(): @pytest.mark.asyncio async def test_distinct_keys_do_not_share_negative_cache(): spy: Final = ResyncSpy(found=False) - read_through: Final = RegistryReadThrough(resync=spy, miss_ttl_seconds=60.0) + read_through: Final = RegistryReadThrough(resync=spy, is_loaded=nothing_loaded, miss_ttl_seconds=60.0) assert await read_through.attempt("ghost-a") is False assert await read_through.attempt("ghost-b") is False @@ -98,7 +105,11 @@ async def test_distinct_keys_do_not_share_negative_cache(): async def test_resync_budget_exhausted_blocks_resync_without_negative_caching(): spy: Final = ResyncSpy(found=False) read_through: Final = RegistryReadThrough( - resync=spy, miss_ttl_seconds=60.0, max_resyncs_per_window=2, resync_window_seconds=60.0 + resync=spy, + is_loaded=nothing_loaded, + miss_ttl_seconds=60.0, + max_resyncs_per_window=2, + resync_window_seconds=60.0, ) assert await read_through.attempt("ghost-a") is False @@ -108,10 +119,48 @@ async def test_resync_budget_exhausted_blocks_resync_without_negative_caching(): assert read_through._recent_misses.get_cache("ghost-c") is None +@pytest.mark.asyncio +async def test_requests_queued_behind_a_successful_resync_spend_no_budget(): + from unittest.mock import AsyncMock, call + + entered: Final = asyncio.Event() + release: Final = asyncio.Event() + new_model_loaded: Final = asyncio.Event() + + async def gated_load(key: str) -> bool: + entered.set() + await release.wait() + if key == "new-model": + new_model_loaded.set() + return True + + def is_loaded(key: str) -> bool: + return key == "new-model" and new_model_loaded.is_set() + + resync: Final = AsyncMock(side_effect=gated_load) + read_through: Final = RegistryReadThrough( + resync=resync, + is_loaded=is_loaded, + max_resyncs_per_window=2, + resync_window_seconds=60.0, + ) + + burst: Final = asyncio.gather(*(read_through.attempt("new-model") for _ in range(25))) + await entered.wait() + release.set() + + assert await burst == [True] * 25 + assert resync.await_args_list == [call("new-model")] + assert await read_through.attempt("other-model") is True + assert resync.await_args_list == [call("new-model"), call("other-model")] + + @pytest.mark.asyncio async def test_resync_budget_replenishes_after_window(): spy: Final = ResyncSpy(found=True) - read_through: Final = RegistryReadThrough(resync=spy, max_resyncs_per_window=1, resync_window_seconds=0.05) + read_through: Final = RegistryReadThrough( + resync=spy, is_loaded=nothing_loaded, max_resyncs_per_window=1, resync_window_seconds=0.05 + ) assert await read_through.attempt("model-a") is True assert await read_through.attempt("model-b") is False @@ -551,3 +600,100 @@ async def test_agent_read_through_hydrates_identity_binding(lookup, clean_agent_ assert agent.identity is not None assert agent.identity.model_dump(include=set(binding)) == binding assert clean_agent_registry.get_agent_by_id(agent_id="agent-id").identity == agent.identity + + +def test_model_is_loaded_matches_router_model_names_and_deployment_ids(monkeypatch: pytest.MonkeyPatch): + import litellm.proxy.proxy_server as proxy_server + from litellm import Router + from litellm.proxy.common_utils.registry_read_through import _model_is_loaded + + router: Final = Router( + model_list=[ + { + "model_name": "loaded-model", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, + "model_info": {"id": "loaded-deployment-id"}, + } + ] + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + + assert _model_is_loaded("loaded-model") is True + assert _model_is_loaded("loaded-deployment-id") is True + assert _model_is_loaded("model-created-on-a-sibling") is False + + monkeypatch.setattr(proxy_server, "llm_router", None) + assert _model_is_loaded("loaded-model") is False + + +@pytest.mark.asyncio +async def test_model_read_through_answers_a_loaded_model_without_reading_the_db(monkeypatch: pytest.MonkeyPatch): + from unittest.mock import AsyncMock, MagicMock + + import litellm.proxy.proxy_server as proxy_server + from litellm import Router + from litellm.proxy.common_utils.registry_read_through import model_registry_read_through + + prisma_client: Final = MagicMock() + prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(side_effect=AssertionError("db read")) + router: Final = Router( + model_list=[ + { + "model_name": "wired-loaded-model", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, + } + ] + ) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr(proxy_server, "llm_router", router) + + assert await model_registry_read_through.attempt("wired-loaded-model") is True + prisma_client.db.litellm_proxymodeltable.find_many.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_guardrail_read_through_answers_a_loaded_guardrail_without_reading_the_db( + monkeypatch: pytest.MonkeyPatch, +): + from unittest.mock import AsyncMock, MagicMock + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.registry_read_through import guardrail_registry_read_through + from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER + from litellm.types.guardrails import Guardrail + + guardrail_id: Final = "wired-loaded-guardrail-id" + guardrail_name: Final = "wired-loaded-guardrail" + prisma_client: Final = MagicMock() + prisma_client.db.litellm_guardrailstable.find_first = AsyncMock(side_effect=AssertionError("db read")) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + + IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db( + guardrail=Guardrail(**dict(FakeGuardrailRow(guardrail_id, guardrail_name))) + ) + try: + assert await guardrail_registry_read_through.attempt(guardrail_name) is True + prisma_client.db.litellm_guardrailstable.find_first.assert_not_awaited() + finally: + IN_MEMORY_GUARDRAIL_HANDLER.delete_in_memory_guardrail(guardrail_id) + + +@pytest.mark.asyncio +async def test_agent_read_through_answers_a_loaded_agent_without_reading_the_db( + clean_agent_registry: "AgentRegistry", monkeypatch: pytest.MonkeyPatch +): + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.registry_read_through import agent_registry_read_through + from litellm.types.agents import AgentResponse + + monkeypatch.setattr(proxy_server, "store_model_in_db", False) + clean_agent_registry.register_agent( + agent_config=AgentResponse.model_validate( + FakeAgentRow("wired-loaded-agent-id", "wired-loaded-agent").model_dump() + ) + ) + + assert await agent_registry_read_through.attempt("wired-loaded-agent-id") is True + assert await agent_registry_read_through.attempt("wired-loaded-agent") is True diff --git a/tests/unit/proxy/conftest.py b/tests/unit/proxy/conftest.py index 50c89387d80..6dec588b763 100644 --- a/tests/unit/proxy/conftest.py +++ b/tests/unit/proxy/conftest.py @@ -411,6 +411,8 @@ def create_proxy_test_client( def fresh_agent_read_through(monkeypatch): from litellm.proxy.common_utils import registry_read_through - read_through = registry_read_through.RegistryReadThrough(resync=registry_read_through._resync_agents) + read_through = registry_read_through.RegistryReadThrough( + resync=registry_read_through._resync_agents, is_loaded=registry_read_through._agent_is_loaded + ) monkeypatch.setattr(registry_read_through, "agent_registry_read_through", read_through) return read_through From bc05d543ca99e1398b81dccd3e4013342adf0189 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 3 Oct 2026 06:44:52 +0000 Subject: [PATCH 02/12] fix(proxy): stamp the client alias on a copy of each streamed chunk so pricing sees the deployment model (#44341) * fix(proxy): stamp the client alias on a copy of each streamed chunk so pricing sees the deployment model Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(spend): streamed alias matching a capability rule bills the deployment price Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(spend): assert every streamed chunk carries the client alias Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(logging): log the client alias on the priced streamed response, the same as non-streamed Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> (cherry picked from commit 7b432d78d2be711c864204fc948073bfd8d658a2) --- litellm/litellm_core_utils/litellm_logging.py | 21 +- litellm/proxy/proxy_server.py | 11 +- .../spend/test_stream_alias_billing.py | 253 ++++++++++++++++++ .../test_litellm_logging.py | 12 + .../proxy_server/test_streaming_helpers.py | 34 ++- 5 files changed, 320 insertions(+), 11 deletions(-) create mode 100644 tests/integration/spend/test_stream_alias_billing.py diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 9e25e7d9d2b..9756e122c5c 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -680,6 +680,9 @@ class Logging(LiteLLMLoggingBaseClass): self.truncated_messages_for_logging: str | list | dict | None = None # mutable-ok: logged messages shape ## TIME TO FIRST TOKEN LOGGING ## self.completion_start_time: datetime.datetime | None = None + # The model the proxy shows the client on streamed chunks. The logged streamed response carries it + # once that response is priced, the same way a non-streamed response is logged + self.client_facing_stream_model: str | None = None self.zero_cost_warned: bool = False self._llm_caching_handler: LLMCachingHandler | None = None @@ -2494,6 +2497,15 @@ class Logging(LiteLLMLoggingBaseClass): setattr(result, "usage", transformed_usage) return result + def _with_client_facing_stream_model( + self, + response: ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse, + ) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse: + model: Final = self.client_facing_stream_model + if model is None or getattr(response, "model", None) in (None, model): + return response + return response.model_copy(update={"model": model}) + def _success_handler_helper_fn( self, result=None, @@ -2811,9 +2823,11 @@ class Logging(LiteLLMLoggingBaseClass): result=complete_streaming_response ) self._merge_hidden_params_from_response_into_metadata(complete_streaming_response) + logged_streaming_response: Final = self._with_client_facing_stream_model(complete_streaming_response) + self.model_call_details["complete_streaming_response"] = logged_streaming_response ## STANDARDIZED LOGGING PAYLOAD self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( - complete_streaming_response, start_time, end_time + logged_streaming_response, start_time, end_time ) standard_logging_payload: Final[StandardLoggingPayload | None] = self.model_call_details.get( "standard_logging_object" @@ -3352,10 +3366,13 @@ class Logging(LiteLLMLoggingBaseClass): await self._prepare_baseline_cache_estimate(complete_streaming_response) + logged_streaming_response: Final = self._with_client_facing_stream_model(complete_streaming_response) + self.model_call_details["async_complete_streaming_response"] = logged_streaming_response + ## STANDARDIZED LOGGING PAYLOAD try: self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( - complete_streaming_response, start_time, end_time + logged_streaming_response, start_time, end_time ) except Exception: # noqa: BLE001 # payload build must never block later callbacks (slot release) verbose_logger.exception( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 26b9a5798bd..7f2cd5d10e6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -9454,12 +9454,17 @@ def _restamp_streaming_chunk_model( ) model_mismatch_logged = True + # The streaming wrapper keeps these same chunk objects to assemble the response it + # prices, so stamp a copy for the client and leave the provider's model for pricing. + # The logging object stamps the same model on the assembled response after pricing it. + logging_obj: Final = request_data.get("litellm_logging_obj") + if isinstance(logging_obj, LiteLLMLoggingObj): + logging_obj.client_facing_stream_model = target_model if isinstance(chunk, dict): - chunk["model"] = target_model - return chunk, model_mismatch_logged + return {**chunk, "model": target_model}, model_mismatch_logged try: - chunk.model = target_model + return chunk.model_copy(update={"model": target_model}), model_mismatch_logged except Exception as e: verbose_proxy_logger.error( "litellm_call_id=%s: failed to override chunk.model=%r on chunk_type=%s. error=%s", diff --git a/tests/integration/spend/test_stream_alias_billing.py b/tests/integration/spend/test_stream_alias_billing.py new file mode 100644 index 00000000000..c9f8dba615a --- /dev/null +++ b/tests/integration/spend/test_stream_alias_billing.py @@ -0,0 +1,253 @@ +"""A streamed alias never replaces the deployment's model for pricing (LIT-9065). + +The proxy shows the client's alias on every streamed chunk, but the chunks kept for end-of-stream cost calculation +keep the deployment's model. "claude-opus-4.8-" is no cost-map key and only matches the claude capability +rules, whose model info carries no prices, so a stream through that alias must bill exactly what the plain alias +"integration-" bills at the same deployment rates, and the client must still see the alias it asked for. +Logging callbacks see that alias as the response model on streamed requests, the same as on non-streamed ones +""" + +import json +from collections.abc import Callable, Iterator, Mapping +from hashlib import sha256 +from pathlib import Path +from typing import Final +from uuid import uuid4 + +import pytest +import yaml +from integration._support.client import ( + Gateway, + Scenario, + eventually, + gateway_from_environment, + object_value, + string_value, +) +from integration._support.database import read_rows +from integration._support.otlp_sink import owned_sinks, recorded_spans +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + + +def _sse_event(name: str, payload: dict[str, JsonValue]) -> bytes: + return f"event: {name}\ndata: {json.dumps(payload, separators=(',', ':'))}\n\n".encode() + + +def _anthropic_reply(request: Request) -> Reply: + assert request.target.endswith("/v1/messages"), request.target + body: Final = json.loads(request.body) + assert body["model"] == "claude-opus-4-8", body + if body.get("stream") is not True: + return Reply( + body=json.dumps( + { + "id": f"msg_{uuid4().hex[:12]}", + "type": "message", + "role": "assistant", + "model": "claude-opus-4-8", + "content": [{"type": "text", "text": "hi"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 30, "output_tokens": 40}, + } + ).encode() + ) + return Reply( + content_type="text/event-stream", + chunks=( + _sse_event( + "message_start", + { + "type": "message_start", + "message": { + "id": f"msg_{uuid4().hex[:12]}", + "type": "message", + "role": "assistant", + "model": "claude-opus-4-8", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 30, "output_tokens": 1}, + }, + }, + ), + _sse_event( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + _sse_event( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hi"}}, + ), + _sse_event("content_block_stop", {"type": "content_block_stop", "index": 0}), + _sse_event( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 40}, + }, + ), + _sse_event("message_stop", {"type": "message_stop"}), + ), + ) + + +def _deployment( + scenario: Scenario, + model_name: str, + litellm_params: dict[str, JsonValue], + model_info: dict[str, JsonValue] | None = None, +) -> str: + created: Final = scenario.gateway.post( + "/model/new", {"model_name": model_name, "litellm_params": litellm_params, "model_info": model_info or {}} + ) + identity: Final = string_value(object_value(created["model_info"])["id"]) + scenario.cleanups.callback(scenario.delete_model, identity) + return model_name + + +def _streamed_spend(gateway: Gateway, scenario: Scenario, model: str, content: str) -> dict[str, JsonValue]: + key: Final = scenario.key(models=[model]) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": content}], + "stream": True, + "stream_options": {"include_usage": True}, + }, + key=key, + ) + assert response.status_code == 200, response.text + chunks: Final = tuple( + json.loads(line.removeprefix("data: ")) + for line in response.text.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + assert chunks and {chunk["model"] for chunk in chunks} == {model}, response.text + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE api_key=%s', + (sha256(key.encode()).hexdigest(),), + ), + lambda values: len(values) == 1, + seconds=70, + ) + return rows[0] + + +def _listed_deployments(gateway: Gateway, model_name: str) -> tuple[dict[str, JsonValue], ...]: + entries: Final = gateway.get("/model/info")["data"] + assert isinstance(entries, list) + return tuple(object_value(entry) for entry in entries if object_value(entry)["model_name"] == model_name) + + +def _deployment_pricing(gateway: Gateway, model_name: str) -> dict[str, JsonValue]: + listed: Final = eventually(lambda: _listed_deployments(gateway, model_name), lambda found: len(found) == 1) + return object_value(listed[0]["model_info"]) + + +_BACKENDS: Final = ( + pytest.param( + lambda _: {"model": "vertex_ai/claude-opus-4-8@default", "mock_response": "hi"}, + id="vertex-mock-response", + ), + pytest.param( + lambda wire_url: { + "model": "anthropic/claude-opus-4-8", + "api_key": "integration-provider-key", + "api_base": wire_url, + }, + id="anthropic-upstream", + ), +) + + +@pytest.mark.parametrize("litellm_params", _BACKENDS) +@pytest.mark.timeout(180) +def test_streamed_alias_matching_a_capability_rule_bills_the_deployment_price( + gateway: Gateway, litellm_params: Callable[[str], dict[str, JsonValue]] +) -> None: + with wire_server(_anthropic_reply) as wire, gateway.scenario() as scenario: + content: Final = f"alias billing {uuid4().hex}" + plain_alias: Final = f"integration-{uuid4().hex}" + rule_alias: Final = f"claude-opus-4.8-{uuid4().int % 10**8:08d}" + exact_row: Final = _streamed_spend( + gateway, scenario, _deployment(scenario, plain_alias, litellm_params(wire.url)), content + ) + alias_row: Final = _streamed_spend( + gateway, scenario, _deployment(scenario, rule_alias, litellm_params(wire.url)), content + ) + + for model_name, row in ((plain_alias, exact_row), (rule_alias, alias_row)): + pricing: Final = _deployment_pricing(gateway, model_name) + input_rate: Final = float(str(pricing["input_cost_per_token"])) + output_rate: Final = float(str(pricing["output_cost_per_token"])) + uplift: Final = float(str(pricing["regional_endpoint_uplift_multiplier"] or 1)) + assert input_rate > 0 and output_rate > 0, pricing + assert float(str(row["spend"])) == pytest.approx( + uplift + * (float(str(row["prompt_tokens"])) * input_rate + float(str(row["completion_tokens"])) * output_rate) + ), (model_name, row, pricing) + + +@pytest.fixture(scope="module") +def otel_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, str]]: + directory: Final = tmp_path_factory.mktemp("stream-alias-otel") + with owned_sinks(directory / "sinks") as sinks, gateway_from_environment() as base: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["litellm_settings"] = {**config["litellm_settings"], "callbacks": ["otel"]} + config["callback_settings"] = { + "otel": {"exporter": "http/json", "endpoint": sinks.operator, "use_simple_processor": True} + } + path: Final = directory / "otel.yaml" + path.write_text(yaml.safe_dump(config)) + overrides: Final = {"OTEL_EXPORTER": "http/json", "OTEL_ENDPOINT": sinks.operator} + with owned_proxy(base, directory, overrides, config=path) as candidate: + yield candidate, sinks.operator + + +def _logged_response_models(sink: str, call_ids: Mapping[str, str]) -> dict[str, JsonValue]: + _, spans = recorded_spans(sink) + return { + label: span["attributes"]["gen_ai.response.model"] + for span in spans + for label, call_id in call_ids.items() + if span["attributes"].get("litellm.call_id") == call_id and "gen_ai.response.model" in span["attributes"] + } + + +@pytest.mark.parametrize("litellm_params", _BACKENDS) +@pytest.mark.timeout(240) +def test_logged_response_model_is_the_client_alias_whether_or_not_the_request_streams( + otel_proxy: tuple[Gateway, str], litellm_params: Callable[[str], dict[str, JsonValue]] +) -> None: + candidate, sink = otel_proxy + with wire_server(_anthropic_reply) as wire, candidate.scenario() as scenario: + alias: Final = f"claude-opus-4.8-{uuid4().int % 10**8:08d}" + key: Final = scenario.key(models=[_deployment(scenario, alias, litellm_params(wire.url))]) + call_ids: Final[dict[str, str]] = {} + for label, stream_fields in ( + ("non-streamed", {}), + ("streamed", {"stream": True, "stream_options": {"include_usage": True}}), + ): + response = candidate.request( + "POST", + "/v1/chat/completions", + {"model": alias, "messages": [{"role": "user", "content": f"logged alias {uuid4().hex}"}]} + | stream_fields, + key=key, + ) + assert response.status_code == 200, response.text + call_ids[label] = response.headers["x-litellm-call-id"] + logged: Final = eventually( + lambda: _logged_response_models(sink, call_ids), + lambda found: len(found) == 2, + seconds=60, + return_last_on_timeout=True, + ) + assert logged == {"non-streamed": alias, "streamed": alias}, call_ids diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 50d5786b418..afe60df5d1d 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -321,6 +321,18 @@ async def test_mcp_direct_content_edit_invalidates_stale_structured_data(logging assert "SECRET-1234" not in result.model_dump_json() +def test_with_client_facing_stream_model_stamps_a_copy_of_the_priced_response(logging_obj): + response = ModelResponse(model="claude-opus-4-6@default") + logging_obj.client_facing_stream_model = "claude-opus-4.6" + logged = logging_obj._with_client_facing_stream_model(response) + assert (logged.model, response.model) == ("claude-opus-4.6", "claude-opus-4-6@default") + + +def test_with_client_facing_stream_model_keeps_the_response_when_the_proxy_set_no_model(logging_obj): + response = ModelResponse(model="claude-opus-4-6@default") + assert logging_obj._with_client_facing_stream_model(response) is response + + def test_get_combined_callback_list_preserves_insertion_order(logging_obj): assert logging_obj.get_combined_callback_list( dynamic_success_callbacks=["prometheus", "langfuse", "datadog", "otel", "s3"], diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index 92de00a4a3f..69fa195e9d6 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -283,8 +283,9 @@ def test_restamp_streaming_chunk_model_overrides_model_on_basemodel(): "model": new_chunk.model, "logged": logged, "same_object": new_chunk is chunk, + "original_model": chunk.model, } - assert snapshot == {"model": "gpt-4", "logged": True, "same_object": True} + assert snapshot == {"model": "gpt-4", "logged": True, "same_object": False, "original_model": "openai/internal-x"} @pytest.mark.parametrize("return_raw_model_name", [False, True]) @@ -310,8 +311,7 @@ def test_restamp_streaming_chunk_model_overrides_model_on_dict(): request_data={}, model_mismatch_logged=True, ) - assert new_chunk["model"] == "gpt-4" - assert logged is True + assert (new_chunk["model"], chunk["model"], logged) == ("gpt-4", "internal", True) def test_restamp_streaming_chunk_model_uses_fallback_model_from_metadata(): @@ -443,7 +443,7 @@ def test_restamp_streaming_chunk_model_fastest_response_preserves_model(): assert logged is False -def test_restamp_streaming_chunk_model_setattr_exception_logs_and_returns(): +def test_restamp_streaming_chunk_model_restamps_a_frozen_chunk_through_a_copy(): from pydantic import ConfigDict class FrozenChunk(_simple_chunk().__class__): @@ -462,8 +462,30 @@ def test_restamp_streaming_chunk_model_setattr_exception_logs_and_returns(): request_data={"litellm_call_id": "test-id"}, model_mismatch_logged=False, ) - assert new_chunk.model == "openai/internal-x" - assert logged is True + assert (new_chunk.model, chunk.model, logged) == ("gpt-4", "openai/internal-x", True) + + +def test_restamp_streaming_chunk_model_records_the_client_model_on_the_logging_object(): + import time + + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj = Logging( + model="openai/internal-x", + messages=[], + stream=True, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="test-id", + function_id="test-id", + ) + _restamp_streaming_chunk_model( + chunk=_simple_chunk(model="openai/internal-x"), + requested_model_from_client="gpt-4", + request_data={"litellm_call_id": "test-id", "litellm_logging_obj": logging_obj}, + model_mismatch_logged=False, + ) + assert logging_obj.client_facing_stream_model == "gpt-4" def test_format_fallback_metadata_sse_event(): From b4a795a11c5937530ac0ddaa0eb7cca1f28244a9 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 3 Oct 2026 07:00:52 +0000 Subject: [PATCH 03/12] fix(responses): merge bridged tool calls into the same choice as the text (#44346) * revert(responses): revert "fix(responses): keep gpt-5.4/5.5 tool calls on chat and merge bridged tool calls into one choice" (#44295) This reverts commit ca1994e403aa0a04daf18a39a0d6bd8b229cfa90. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(responses): merge bridged tool calls into the same choice as the text Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> (cherry picked from commit 5724117116102665495b9bad0c998fd210cdee06) --- .../transformation.py | 29 +- .../providers/test_openai_chat_wire.py | 188 +++++++++ .../test_responses_bridge_incomplete.py | 378 ++++++++++++++++++ ...responses_transformation_transformation.py | 151 +++++++ 4 files changed, 744 insertions(+), 2 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 391dbc44eec..25ea4b05c62 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -187,7 +187,7 @@ def _reasoning_items_from_output_items(output_items: Sequence[object]) -> tuple[ def _as_chat_reasoning_items( - reasoning_items: Sequence[_BuiltReasoningItem], + reasoning_items: Sequence[_BuiltReasoningItem | ChatCompletionReasoningItem], ) -> list[ChatCompletionReasoningItem] | None: if not reasoning_items: return None @@ -784,7 +784,32 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): else: pass # don't fail request if item in list is not supported - # If we accumulated tool calls, create a single choice with all of them + if accumulated_tool_calls and choices: + last_choice: Final = choices[-1] + last_reasoning_content: Final = getattr(last_choice.message, "reasoning_content", None) + last_reasoning_items: Final = getattr(last_choice.message, "reasoning_items", None) + merged_reasoning_content: Final = ( + " ".join(value for value in (last_reasoning_content, reasoning_content) if value) or None + ) + merged_reasoning_items: Final = _as_chat_reasoning_items( + ( + *(last_reasoning_items or ()), + *(() if pending_reasoning_item is None else (pending_reasoning_item,)), + ) + ) + merged_message: Final = Message( + role=last_choice.message.role, + content=last_choice.message.content, + annotations=getattr(last_choice.message, "annotations", None), + tool_calls=accumulated_tool_calls, + reasoning_content=merged_reasoning_content, + reasoning_items=merged_reasoning_items, + ) + return [ + *choices[:-1], + Choices(message=merged_message, finish_reason="tool_calls", index=last_choice.index), + ] + if accumulated_tool_calls: msg = Message( content=None, diff --git a/tests/integration/providers/test_openai_chat_wire.py b/tests/integration/providers/test_openai_chat_wire.py index 24d7d83e519..4cab61db4d0 100644 --- a/tests/integration/providers/test_openai_chat_wire.py +++ b/tests/integration/providers/test_openai_chat_wire.py @@ -1,5 +1,6 @@ import json import uuid +from itertools import chain from typing import Final import pytest @@ -64,3 +65,190 @@ def test_openai_chat_tool_choice_without_tools_is_not_forwarded(gateway: Gateway } ] assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")] + + +def test_azure_gpt_6_bridged_stream_returns_text_and_tool_call_on_one_choice(gateway: Gateway) -> None: + identity: Final = f"azure-gpt-6-sol-stream-{uuid.uuid4().hex}" + expected_text: Final = "Let me check the weather." + events: Final = ( + { + "type": "response.created", + "response": { + "id": "resp_weather", + "object": "response", + "created_at": 1, + "status": "in_progress", + "model": "gpt-6-sol", + }, + }, + { + "type": "response.output_item.added", + "output_index": 0, + "item": { + "id": "msg_weather", + "type": "message", + "status": "in_progress", + "role": "assistant", + "content": [], + }, + }, + { + "type": "response.output_text.delta", + "item_id": "msg_weather", + "output_index": 0, + "content_index": 0, + "delta": "Let me check ", + }, + { + "type": "response.output_text.delta", + "item_id": "msg_weather", + "output_index": 0, + "content_index": 0, + "delta": "the weather.", + }, + { + "type": "response.output_item.done", + "output_index": 0, + "item": { + "id": "msg_weather", + "type": "message", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": expected_text, "annotations": []}], + }, + }, + { + "type": "response.output_item.added", + "output_index": 1, + "item": { + "id": "fc_1", + "type": "function_call", + "status": "in_progress", + "call_id": "call_1", + "name": "get_weather", + "arguments": "", + }, + }, + { + "type": "response.function_call_arguments.delta", + "item_id": "fc_1", + "output_index": 1, + "delta": '{"city":', + }, + { + "type": "response.function_call_arguments.delta", + "item_id": "fc_1", + "output_index": 1, + "delta": '"Paris"}', + }, + { + "type": "response.output_item.done", + "output_index": 1, + "item": { + "id": "fc_1", + "type": "function_call", + "status": "completed", + "call_id": "call_1", + "name": "get_weather", + "arguments": '{"city":"Paris"}', + }, + }, + { + "type": "response.completed", + "response": { + "id": "resp_weather", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-6-sol", + "output": [ + { + "id": "msg_weather", + "type": "message", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": expected_text, "annotations": []}], + }, + { + "id": "fc_1", + "type": "function_call", + "status": "completed", + "call_id": "call_1", + "name": "get_weather", + "arguments": '{"city":"Paris"}', + }, + ], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + }, + }, + ) + stream_chunks: Final = tuple(f"data: {json.dumps(event)}\n\n".encode() for event in events) + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/openai/responses?api-version=2025-04-01-preview" + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["model"] == "gpt-6-sol" + return Reply(content_type="text/event-stream", chunks=stream_chunks) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="azure/gpt-6-sol", + api_base=wire.url, + api_key=_API_KEY, + api_version="2025-04-01-preview", + ) + with gateway.client.stream( + "POST", + "/v1/chat/completions", + headers={"Authorization": f"Bearer {gateway.key}"}, + json={ + "model": model, + "messages": [{"role": "user", "content": f"What is the weather in Paris? {identity}"}], + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } + ], + "stream": True, + "cache": {"no-cache": True}, + }, + ) as response: + response_body: Final = response.read() + assert response.status_code == 200, response.text + chunks: Final = tuple( + _JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in response_body.decode().splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + choices: Final = tuple(chain.from_iterable(chunk["choices"] for chunk in chunks)) + assert choices, response.text + assert all(choice["index"] == 0 for choice in choices), response.text + assert "".join(str(choice["delta"].get("content") or "") for choice in choices) == expected_text, ( + response.text + ) + tool_call_chunks: Final = tuple( + chain.from_iterable(choice["delta"].get("tool_calls", []) for choice in choices) + ) + assert ( + "".join(str(tool_call["function"].get("name") or "") for tool_call in tool_call_chunks) == "get_weather" + ), response.text + assert ( + "".join(str(tool_call["function"].get("arguments") or "") for tool_call in tool_call_chunks) + == '{"city":"Paris"}' + ), response.text + assert tuple( + choice.get("finish_reason") for choice in choices if choice.get("finish_reason") is not None + ) == ("tool_calls",), response.text + assert [(request.method, request.target) for request in wire.drain()] == [ + ("POST", "/openai/responses?api-version=2025-04-01-preview") + ] diff --git a/tests/integration/providers/test_responses_bridge_incomplete.py b/tests/integration/providers/test_responses_bridge_incomplete.py index 2252f1634e0..5d3877c954e 100644 --- a/tests/integration/providers/test_responses_bridge_incomplete.py +++ b/tests/integration/providers/test_responses_bridge_incomplete.py @@ -194,3 +194,381 @@ def test_messages_over_responses_deployment_with_max_tokens_one_reaches_openai_a assert len(tuple(request for request in wire.drain() if request.method == "POST")) == 1 assert body["content"] == [{"type": "text", "text": "ok"}], response.text assert body["usage"]["input_tokens"] == 9 and body["usage"]["output_tokens"] == 1, response.text + + +def test_chat_over_responses_deployment_merges_message_and_function_call(gateway: Gateway) -> None: + identity: Final = "responses-bridge-" + uuid.uuid4().hex + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=b'{"object":"list","data":[]}') + assert request.method == "POST" and request.target == "/responses", request.target + return Reply( + body=json.dumps( + { + "id": "resp_weather", + "object": "response", + "created_at": 1789788253, + "status": "completed", + "model": "gpt-6-sol", + "output": [ + { + "type": "message", + "id": "msg_weather", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "Let me check the weather.", + "annotations": [], + } + ], + }, + { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "get_weather", + "arguments": '{"city":"Paris"}', + "status": "completed", + }, + ], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/responses/gpt-6-sol", api_base=wire.url, api_key="synthetic-openai-key" + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"What is the weather in Paris? {identity}"}], + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } + ], + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + body: Final = response.json() + assert body["choices"] == [ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "role": "assistant", + "content": "Let me check the weather.", + "tool_calls": [ + { + "id": "fc_1", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city":"Paris"}', + }, + "index": 0, + } + ], + }, + } + ], response.text + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")] + + +def test_chat_over_responses_deployment_keeps_reasoning_with_merged_tool_call(gateway: Gateway) -> None: + identity: Final = "responses-bridge-reasoning-" + uuid.uuid4().hex + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=b'{"object":"list","data":[]}') + assert request.method == "POST" and request.target == "/responses", request.target + return Reply( + body=json.dumps( + { + "id": "resp_weather_reasoning", + "object": "response", + "created_at": 1789788253, + "status": "completed", + "model": "gpt-6-sol", + "output": [ + { + "type": "message", + "id": "msg_weather_reasoning", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "Let me check the weather.", + "annotations": [], + } + ], + }, + { + "type": "reasoning", + "id": "rs_weather", + "summary": [{"type": "summary_text", "text": "Checking the forecast."}], + }, + { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "get_weather", + "arguments": '{"city":"Paris"}', + "status": "completed", + }, + ], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/responses/gpt-6-sol", api_base=wire.url, api_key="synthetic-openai-key" + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"What is the weather in Paris? {identity}"}], + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } + ], + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + body: Final = response.json() + assert body["choices"] == [ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "role": "assistant", + "content": "Let me check the weather.", + "reasoning_content": "Checking the forecast.", + "reasoning_items": [ + { + "type": "reasoning", + "id": "rs_weather", + "summary": [{"type": "summary_text", "text": "Checking the forecast."}], + } + ], + "tool_calls": [ + { + "id": "fc_1", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city":"Paris"}', + }, + "index": 0, + } + ], + }, + } + ], response.text + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")] + + +def test_chat_over_responses_deployment_returns_tool_call_only_reply_as_one_choice(gateway: Gateway) -> None: + identity: Final = "responses-bridge-tool-only-" + uuid.uuid4().hex + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=b'{"object":"list","data":[]}') + assert request.method == "POST" and request.target == "/responses", request.target + return Reply( + body=json.dumps( + { + "id": "resp_weather_tool_only", + "object": "response", + "created_at": 1789788253, + "status": "completed", + "model": "gpt-6-sol", + "output": [ + { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "get_weather", + "arguments": '{"city":"Paris"}', + "status": "completed", + } + ], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/responses/gpt-6-sol", api_base=wire.url, api_key="synthetic-openai-key" + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"What is the weather in Paris? {identity}"}], + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } + ], + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + body: Final = response.json() + assert body["choices"] == [ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "fc_1", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city":"Paris"}', + }, + "index": 0, + } + ], + }, + } + ], response.text + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")] + + +def test_chat_over_responses_deployment_merges_function_call_followed_by_message(gateway: Gateway) -> None: + identity: Final = "responses-bridge-tool-then-message-" + uuid.uuid4().hex + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=b'{"object":"list","data":[]}') + assert request.method == "POST" and request.target == "/responses", request.target + return Reply( + body=json.dumps( + { + "id": "resp_weather_tool_then_message", + "object": "response", + "created_at": 1789788253, + "status": "completed", + "model": "gpt-6-sol", + "output": [ + { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "get_weather", + "arguments": '{"city":"Paris"}', + "status": "completed", + }, + { + "type": "message", + "id": "msg_after_tool", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "After the tool.", "annotations": []}], + }, + ], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/responses/gpt-6-sol", api_base=wire.url, api_key="synthetic-openai-key" + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"What is the weather in Paris? {identity}"}], + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } + ], + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + body: Final = response.json() + assert body["choices"] == [ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "role": "assistant", + "content": "After the tool.", + "tool_calls": [ + { + "id": "fc_1", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city":"Paris"}', + }, + "index": 0, + } + ], + }, + } + ], response.text + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")] diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 282b84104a6..f5ec9eec1f5 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -7,6 +7,15 @@ from unittest.mock import ANY, MagicMock, Mock, patch import httpx import pytest +from openai.types.responses import ( + ResponseFunctionToolCall, + ResponseOutputMessage, + ResponseOutputText, +) +from openai.types.responses.response_reasoning_item import ( + ResponseReasoningItem, + Summary, +) import litellm from litellm.completion_extras.litellm_responses_transformation.transformation import ( @@ -3307,6 +3316,148 @@ def test_convert_response_output_generic_pydantic_message_item(): assert choices[0].finish_reason == "stop" +def test_convert_response_output_merges_message_reasoning_and_function_call() -> None: + message: Final = ResponseOutputMessage( + id="msg_weather", + content=[ + ResponseOutputText( + annotations=[ + { + "type": "url_citation", + "start_index": 0, + "end_index": 5, + "title": "Forecast", + "url": "https://example.com/forecast", + } + ], + text="Sunny.", + type="output_text", + logprobs=[], + ) + ], + role="assistant", + status="completed", + type="message", + ) + reasoning: Final = ResponseReasoningItem( + id="rs_before", + summary=[Summary(type="summary_text", text="Checking the forecast.")], + type="reasoning", + content=None, + encrypted_content=None, + status=None, + ) + pending_reasoning: Final = ResponseReasoningItem( + id="rs_after", + summary=[Summary(type="summary_text", text="The location is Paris.")], + type="reasoning", + content=None, + encrypted_content=None, + status=None, + ) + function_call: Final = ResponseFunctionToolCall( + id="fc_1", + type="function_call", + status="completed", + arguments='{"city":"Paris"}', + call_id="call_1", + name="get_weather", + ) + + message_and_call: Final = LiteLLMResponsesTransformationHandler._convert_response_output_to_choices( + (message, function_call) + ) + assert len(message_and_call) == 1 + assert message_and_call[0].index == 0 + assert message_and_call[0].finish_reason == "tool_calls" + assert message_and_call[0].message.role == "assistant" + assert message_and_call[0].message.content == "Sunny." + assert message_and_call[0].message.annotations == [ + { + "type": "url_citation", + "start_index": 0, + "end_index": 5, + "title": "Forecast", + "url": "https://example.com/forecast", + } + ] + function_calls: Final = message_and_call[0].message.tool_calls + assert function_calls is not None + assert len(function_calls) == 1 + assert function_calls[0].function.name == "get_weather" + assert function_calls[0].function.arguments == '{"city":"Paris"}' + + reasoning_before_message: Final = LiteLLMResponsesTransformationHandler._convert_response_output_to_choices( + (reasoning, message, function_call) + ) + assert len(reasoning_before_message) == 1 + assert reasoning_before_message[0].message.reasoning_content == "Checking the forecast." + reasoning_before_items: Final = reasoning_before_message[0].message.reasoning_items + assert reasoning_before_items is not None + assert reasoning_before_items[0]["id"] == "rs_before" + + reasoning_after_message: Final = LiteLLMResponsesTransformationHandler._convert_response_output_to_choices( + (message, pending_reasoning, function_call) + ) + assert len(reasoning_after_message) == 1 + assert reasoning_after_message[0].message.reasoning_content == "The location is Paris." + reasoning_after_items: Final = reasoning_after_message[0].message.reasoning_items + assert reasoning_after_items is not None + assert reasoning_after_items[0]["id"] == "rs_after" + + merged_reasoning: Final = LiteLLMResponsesTransformationHandler._convert_response_output_to_choices( + (reasoning, message, pending_reasoning, function_call) + ) + assert len(merged_reasoning) == 1 + assert merged_reasoning[0].message.reasoning_content == "Checking the forecast. The location is Paris." + merged_reasoning_items: Final = merged_reasoning[0].message.reasoning_items + assert merged_reasoning_items is not None + assert [item["id"] for item in merged_reasoning_items] == ["rs_before", "rs_after"] + + tool_only: Final = LiteLLMResponsesTransformationHandler._convert_response_output_to_choices((function_call,)) + assert len(tool_only) == 1 + assert tool_only[0].index == 0 + assert tool_only[0].finish_reason == "tool_calls" + assert tool_only[0].message.content is None + assert tool_only[0].message.tool_calls is not None + assert len(tool_only[0].message.tool_calls) == 1 + + message_only: Final = LiteLLMResponsesTransformationHandler._convert_response_output_to_choices((message,)) + assert len(message_only) == 1 + assert message_only[0].index == 0 + assert message_only[0].finish_reason == "stop" + assert message_only[0].message.content == "Sunny." + assert message_only[0].message.tool_calls is None + + +def test_convert_response_output_merges_raw_dict_message_and_function_call() -> None: + handler: Final = LiteLLMResponsesTransformationHandler() + raw_message: Final = { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Let me check.", "annotations": []}], + } + raw_function_call: Final = { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "get_weather", + "arguments": '{"city":"Paris"}', + } + choices: Final = LiteLLMResponsesTransformationHandler._convert_response_output_to_choices( + (raw_message, raw_function_call), + handle_raw_dict_callback=handler._handle_raw_dict_response_item, + ) + + assert len(choices) == 1 + assert choices[0].index == 0 + assert choices[0].finish_reason == "tool_calls" + assert choices[0].message.role == "assistant" + assert choices[0].message.content == "Let me check." + assert choices[0].message.tool_calls is not None + assert len(choices[0].message.tool_calls) == 1 + + def test_convert_tools_to_responses_format_flattens_nested_custom_tool(): from litellm.completion_extras.litellm_responses_transformation.transformation import ( LiteLLMResponsesTransformationHandler, From d442ef4cbbebb4ec8cab394613906e2da93d466c Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:33:50 -0700 Subject: [PATCH 04/12] fix(responses): fold bridged Responses output into one chat choice (#46099) * fix(responses): fold bridged Responses output into one chat choice A Chat Completions request bridged to /responses got one choice per Responses message item, so a reply shaped [message, message, function_call] put the tool call on choices[1] and clients reading only choices[0] never saw it. Fold every message text, annotation, reasoning item and tool call into choices[0] Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(responses): drop mutable annotations and the extra cast in the choice fold Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(responses): remove a cast-ok marker that suppressed nothing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(responses): join folded message texts with a newline Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * revert(responses): join folded message texts with no separator so stream and non-stream match 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> (cherry picked from commit 4b7b8efd4cf91c302ecb3196bc7ac3a4288ebd51) --- .../transformation.py | 111 +++++--- tests/integration/_support/client.py | 5 + .../test_responses_bridge_incomplete.py | 256 +++++++++++++++++- ...responses_transformation_transformation.py | 58 ++++ 4 files changed, 389 insertions(+), 41 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 25ea4b05c62..66ff7163e67 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -5,6 +5,7 @@ Handler for transforming /chat/completions api requests to litellm.responses req import json import os from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence +from itertools import accumulate, chain from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, TypeVar, Union, cast @@ -69,7 +70,7 @@ if TYPE_CHECKING: ChatCompletionToolReferenceObject, OpenAIMessageContentListBlock, ) - from litellm.types.utils import Choices + from litellm.types.utils import Choices, Message _CHAT_COMPLETION_FIELDS: Final = frozenset((*ModelResponse.model_fields, "usage")) @@ -194,6 +195,67 @@ def _as_chat_reasoning_items( return cast(list[ChatCompletionReasoningItem], list(reasoning_items)) +def _shift_annotation(annotation: ChatCompletionAnnotation, offset: int) -> ChatCompletionAnnotation: + shifted: Final = { + key: value + offset if key in ("start_index", "end_index") and isinstance(value, int) else value + for key, value in annotation.items() + } + return cast(ChatCompletionAnnotation, shifted) + + +def _message_annotations(message: "Message") -> Sequence[ChatCompletionAnnotation]: + annotations: Sequence[ChatCompletionAnnotation] | None = getattr(message, "annotations", None) + return annotations or () + + +def _message_reasoning_content(message: "Message") -> str | None: + reasoning_content: str | None = getattr(message, "reasoning_content", None) + return reasoning_content + + +def _message_reasoning_items(message: "Message") -> Sequence[ChatCompletionReasoningItem]: + reasoning_items: Sequence[ChatCompletionReasoningItem] | None = getattr(message, "reasoning_items", None) + return reasoning_items or () + + +def _shifted_annotations(messages: Sequence["Message"], offsets: Sequence[int]) -> Iterator[ChatCompletionAnnotation]: + for message, offset in zip(messages, offsets): + for annotation in _message_annotations(message): + yield _shift_annotation(annotation, offset) + + +def _fold_choices_into_one( + choices: Sequence["Choices"], + tool_calls: Sequence[Mapping[str, object]], + trailing_reasoning_content: str | None, + trailing_reasoning_item: _BuiltReasoningItem | None, +) -> "Choices": + """Clients that only read choices[0] must still see every text and tool call.""" + from litellm.types.utils import Choices, Message + + messages: Final = tuple(choice.message for choice in choices) + texts: Final = tuple(str(message.content or "") for message in messages) + offsets: Final = (0, *accumulate(len(text) for text in texts)) + annotations: Final = tuple(_shifted_annotations(messages, offsets)) + reasoning_contents: Final = (*(_message_reasoning_content(m) for m in messages), trailing_reasoning_content) + reasoning_items: Final = ( + *chain.from_iterable(_message_reasoning_items(m) for m in messages), + *(() if trailing_reasoning_item is None else (trailing_reasoning_item,)), + ) + return Choices( + message=Message( + role=messages[0].role if messages else "assistant", + content="".join(texts) if messages else None, + annotations=list(annotations) or None, + tool_calls=list(tool_calls) or None, + reasoning_content=" ".join(value for value in reasoning_contents if value) or None, + reasoning_items=_as_chat_reasoning_items(reasoning_items), + ), + finish_reason="tool_calls" if tool_calls else "stop", + index=0, + ) + + _ToolChoiceT = TypeVar("_ToolChoiceT") @@ -784,47 +846,16 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): else: pass # don't fail request if item in list is not supported - if accumulated_tool_calls and choices: - last_choice: Final = choices[-1] - last_reasoning_content: Final = getattr(last_choice.message, "reasoning_content", None) - last_reasoning_items: Final = getattr(last_choice.message, "reasoning_items", None) - merged_reasoning_content: Final = ( - " ".join(value for value in (last_reasoning_content, reasoning_content) if value) or None - ) - merged_reasoning_items: Final = _as_chat_reasoning_items( - ( - *(last_reasoning_items or ()), - *(() if pending_reasoning_item is None else (pending_reasoning_item,)), - ) - ) - merged_message: Final = Message( - role=last_choice.message.role, - content=last_choice.message.content, - annotations=getattr(last_choice.message, "annotations", None), + if not choices and not accumulated_tool_calls: + return choices + return [ + _fold_choices_into_one( + choices=choices, tool_calls=accumulated_tool_calls, - reasoning_content=merged_reasoning_content, - reasoning_items=merged_reasoning_items, + trailing_reasoning_content=reasoning_content, + trailing_reasoning_item=pending_reasoning_item, ) - return [ - *choices[:-1], - Choices(message=merged_message, finish_reason="tool_calls", index=last_choice.index), - ] - - if accumulated_tool_calls: - msg = Message( - content=None, - tool_calls=accumulated_tool_calls, - reasoning_content=reasoning_content, - reasoning_items=cast( - list[ChatCompletionReasoningItem] | None, - ([pending_reasoning_item] if pending_reasoning_item is not None else None), - ), - ) - choices.append(Choices(message=msg, finish_reason="tool_calls", index=index)) - reasoning_content = None - pending_reasoning_item = None - - return choices + ] @staticmethod def _build_empty_incomplete_choice( diff --git a/tests/integration/_support/client.py b/tests/integration/_support/client.py index a57792255ab..d2706fdeb38 100644 --- a/tests/integration/_support/client.py +++ b/tests/integration/_support/client.py @@ -23,6 +23,11 @@ def object_value(value: JsonValue) -> dict[str, JsonValue]: return JSON_OBJECT.validate_python(value) +def list_value(value: JsonValue) -> Sequence[JsonValue]: + assert isinstance(value, list), f"Expected a list, received {type(value).__name__}" + return value + + def string_value(value: JsonValue) -> str: assert isinstance(value, str), f"Expected a string, received {type(value).__name__}" return value diff --git a/tests/integration/providers/test_responses_bridge_incomplete.py b/tests/integration/providers/test_responses_bridge_incomplete.py index 5d3877c954e..6e005f928ef 100644 --- a/tests/integration/providers/test_responses_bridge_incomplete.py +++ b/tests/integration/providers/test_responses_bridge_incomplete.py @@ -3,8 +3,9 @@ import uuid from typing import Final import pytest -from integration._support.client import Gateway +from integration._support.client import Gateway, list_value, object_value from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue @pytest.mark.covers("other.provider_wire.responses_bridge.max_output_tokens_incomplete_maps_to_length") @@ -572,3 +573,256 @@ def test_chat_over_responses_deployment_merges_function_call_followed_by_message } ], response.text assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")] + + +_WEATHER_TOOL: Final[dict[str, JsonValue]] = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, +} + + +def _message_item(item_id: str, text: str) -> dict[str, JsonValue]: + return { + "type": "message", + "id": item_id, + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": text, "annotations": []}], + } + + +_WEATHER_CALL: Final[dict[str, JsonValue]] = { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "get_weather", + "arguments": '{"city":"Paris"}', + "status": "completed", +} + +_TWO_MESSAGES_THEN_CALL: Final = ( + _message_item("msg_first", "First message."), + _message_item("msg_second", "Second message."), + _WEATHER_CALL, +) + + +def test_chat_over_responses_deployment_folds_two_messages_and_function_call_into_one_choice( + gateway: Gateway, +) -> None: + identity: Final = "responses-bridge-two-messages-" + uuid.uuid4().hex + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=b'{"object":"list","data":[]}') + assert request.method == "POST" and request.target == "/responses", request.target + return Reply( + body=json.dumps( + { + "id": "resp_two_messages", + "object": "response", + "created_at": 1789788253, + "status": "completed", + "model": "gpt-6-sol", + "output": list(_TWO_MESSAGES_THEN_CALL), + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/responses/gpt-6-sol", api_base=wire.url, api_key="synthetic-openai-key" + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"What is the weather in Paris? {identity}"}], + "tools": [_WEATHER_TOOL], + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + body: Final = response.json() + assert body["choices"] == [ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "role": "assistant", + "content": "First message.Second message.", + "tool_calls": [ + { + "id": "fc_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"Paris"}'}, + "index": 0, + } + ], + }, + } + ], response.text + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")] + + +def _responses_frame(event: str, payload: dict[str, JsonValue]) -> bytes: + return f"event: {event}\ndata: {json.dumps(payload, separators=(',', ':'))}\n\n".encode() + + +def _message_item_frames(output_index: int, item: dict[str, JsonValue]) -> tuple[bytes, ...]: + item_id: Final = item["id"] + text: Final = object_value(list_value(item["content"])[0])["text"] + return ( + _responses_frame( + "response.output_item.added", + { + "type": "response.output_item.added", + "output_index": output_index, + "item": {**item, "status": "in_progress", "content": []}, + }, + ), + _responses_frame( + "response.output_text.delta", + { + "type": "response.output_text.delta", + "item_id": item_id, + "output_index": output_index, + "content_index": 0, + "delta": text, + }, + ), + _responses_frame( + "response.output_item.done", + {"type": "response.output_item.done", "output_index": output_index, "item": item}, + ), + ) + + +def _function_call_frames(output_index: int, item: dict[str, JsonValue]) -> tuple[bytes, ...]: + return ( + _responses_frame( + "response.output_item.added", + { + "type": "response.output_item.added", + "output_index": output_index, + "item": {**item, "status": "in_progress", "arguments": ""}, + }, + ), + _responses_frame( + "response.function_call_arguments.delta", + { + "type": "response.function_call_arguments.delta", + "item_id": item["id"], + "output_index": output_index, + "delta": item["arguments"], + }, + ), + _responses_frame( + "response.function_call_arguments.done", + { + "type": "response.function_call_arguments.done", + "item_id": item["id"], + "output_index": output_index, + "arguments": item["arguments"], + }, + ), + _responses_frame( + "response.output_item.done", + {"type": "response.output_item.done", "output_index": output_index, "item": item}, + ), + ) + + +def _sse_payloads(lines: list[str]) -> list[dict[str, JsonValue]]: + return [json.loads(line[6:]) for line in lines if line.startswith("data: ") and line != "data: [DONE]"] + + +def test_chat_stream_over_responses_deployment_folds_two_messages_and_function_call_into_one_choice( + gateway: Gateway, +) -> None: + identity: Final = "responses-bridge-two-messages-stream-" + uuid.uuid4().hex + response_payload: Final[dict[str, JsonValue]] = { + "id": "resp_two_messages_stream", + "object": "response", + "created_at": 1789788253, + "model": "gpt-6-sol", + "status": "in_progress", + "output": [], + } + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=b'{"object":"list","data":[]}') + assert request.method == "POST" and request.target == "/responses", request.target + assert json.loads(request.body)["stream"] is True, request.body + return Reply( + content_type="text/event-stream", + chunks=( + _responses_frame("response.created", {"type": "response.created", "response": response_payload}), + *_message_item_frames(0, _TWO_MESSAGES_THEN_CALL[0]), + *_message_item_frames(1, _TWO_MESSAGES_THEN_CALL[1]), + *_function_call_frames(2, _TWO_MESSAGES_THEN_CALL[2]), + _responses_frame( + "response.completed", + { + "type": "response.completed", + "response": { + **response_payload, + "status": "completed", + "output": list(_TWO_MESSAGES_THEN_CALL), + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + }, + }, + ), + ), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/responses/gpt-6-sol", api_base=wire.url, api_key="synthetic-openai-key" + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"What is the weather in Paris? {identity}"}], + "tools": [_WEATHER_TOOL], + "stream": True, + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + chunks: Final = _sse_payloads(list(response.iter_lines())) + choices: Final = [object_value(choice) for chunk in chunks for choice in list_value(chunk["choices"])] + assert [choice["index"] for choice in choices] == [0] * len(choices), chunks + deltas: Final = [object_value(choice["delta"]) for choice in choices] + assert ( + "".join(str(delta["content"]) for delta in deltas if delta.get("content")) + == "First message.Second message." + ) + tool_calls: Final = [ + object_value(call) + for delta in deltas + if delta.get("tool_calls") + for call in list_value(delta["tool_calls"]) + ] + assert [ + object_value(call["function"])["name"] for call in tool_calls if object_value(call["function"]).get("name") + ] == ["get_weather"], chunks + assert ( + "".join(str(object_value(call["function"]).get("arguments") or "") for call in tool_calls) + == '{"city":"Paris"}' + ), chunks + assert [choice["finish_reason"] for choice in choices if choice.get("finish_reason")] == ["tool_calls"], chunks + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")] diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index f5ec9eec1f5..860d1f67bb7 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -4539,3 +4539,61 @@ def test_every_bridged_chunk_after_response_created_carries_the_served_service_t relayed = [iterator.chunk_parser(event).model_dump().get("service_tier") for event in events] assert relayed == ["default"] * len(events), relayed + + +def test_convert_response_output_folds_two_messages_and_function_call_into_one_choice() -> None: + first: Final = ResponseOutputMessage.model_validate( + { + "id": "msg_first", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "First message.", "annotations": []}], + } + ) + second_annotation: Final = { + "type": "url_citation", + "start_index": 0, + "end_index": 6, + "title": "Second", + "url": "https://example.com/second", + } + second: Final = ResponseOutputMessage.model_validate( + { + "id": "msg_second", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "Second message.", "annotations": [second_annotation]}], + } + ) + function_call: Final = ResponseFunctionToolCall( + id="fc_1", + type="function_call", + status="completed", + arguments='{"city":"Paris"}', + call_id="call_1", + name="get_weather", + ) + + choices: Final = LiteLLMResponsesTransformationHandler._convert_response_output_to_choices( # pyright: ignore[reportPrivateUsage] # exercises the converter directly + (first, second, function_call) + ) + + assert len(choices) == 1 + assert choices[0].index == 0 + assert choices[0].finish_reason == "tool_calls" + assert choices[0].message.content == "First message.Second message." + assert choices[0].message.annotations == [ + {**second_annotation, "start_index": len("First message."), "end_index": len("First message.") + 6} + ] + tool_calls: Final = choices[0].message.tool_calls + assert tool_calls is not None + assert [call.model_dump(exclude_none=True) for call in tool_calls] == [ + { + "id": "fc_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"Paris"}'}, + "index": 0, + } + ] From c2e40e31e8d251570787b0d4361af5ea76fd0e26 Mon Sep 17 00:00:00 2001 From: Chloe Lu <3273594+clonylu@users.noreply.github.com> Date: Thu, 15 Oct 2026 14:30:03 +0800 Subject: [PATCH 05/12] fix(anthropic): keep signed empty thinking blocks on /v1/messages requests (#41362) The request-side sanitizer on /v1/messages dropped every thinking block whose text is empty, even with a valid signature. Claude Opus 4.8 / Fable 5 / Fable 5.1 return exactly that shape by default (thinking.display omitted): the signature is the preserved reasoning. Replaying the assistant turn as returned therefore lost the model's reasoning on every multi-turn request, silently (the provider answers 200 without it), and left preserved-thinking enforcement (thinking.block_binding) nothing to bind. Anthropic accepts an empty thinking block that carries a signature and binds it to its prefix; only an empty block with no signature (the LIT-6357 case, a turn a non-Anthropic model served through the bridge) is rejected with "each thinking block must contain thinking". Use the signature-aware predicate is_empty_unsigned_thinking_block, which #38809 already applied on the emit side, for the request side too. The unsigned case is still dropped. Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> (cherry picked from commit 897bfea38d22666c67be8e8e523ae316f4e42265) --- litellm/llms/anthropic/common_utils.py | 42 +++++++++++-------- .../anthropic/test_anthropic_common_utils.py | 40 +++++++++++++++--- 2 files changed, 59 insertions(+), 23 deletions(-) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 65c2fccceeb..fb2480cf0ac 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -1390,7 +1390,7 @@ def strip_empty_content_blocks_from_anthropic_messages( ) -> list[Any]: """ Return a new message list with empty or whitespace-only ``{"type": "text"}`` - and ``{"type": "thinking"}`` content blocks removed. + content blocks and empty *unsigned* ``{"type": "thinking"}`` blocks removed. Anthropic's API rejects requests containing such blocks with ``"messages: text content blocks must be non-empty"`` and @@ -1406,8 +1406,13 @@ def strip_empty_content_blocks_from_anthropic_messages( on the unified ``/v1/messages`` path. ``/v1/chat/completions`` already handles this in ``anthropic_messages_pt``; this helper provides the equivalent guarantee for the native Anthropic Messages path. - ``redacted_thinking`` blocks are never touched: they carry opaque - ``data`` instead of thinking text. + + A thinking block with empty text but a non-empty ``signature`` is kept: + it is the default response shape of Claude Opus 4.8 / Fable 5 / Fable 5.1 + (``thinking.display`` omitted), the signature is the model's preserved + reasoning, and Anthropic accepts it on input. ``redacted_thinking`` + blocks are never touched: they carry opaque ``data`` instead of thinking + text. Messages whose content is a list and becomes empty after stripping are omitted, matching :func:`strip_thinking_blocks_from_anthropic_messages`. @@ -1420,7 +1425,7 @@ def strip_empty_content_blocks_from_anthropic_messages( out.append(m) continue content = m["content"] - filtered = [b for b in content if not _is_empty_text_block(b) and not is_empty_thinking_block(b)] + filtered = [b for b in content if not _is_empty_text_block(b) and not is_empty_unsigned_thinking_block(b)] if len(filtered) == len(content): out.append(m) elif filtered: @@ -1438,11 +1443,11 @@ def _is_empty_text_block(block: object) -> bool: def is_empty_thinking_block(block: object) -> bool: """ True for a ``{"type": "thinking"}`` content block whose thinking text is - missing, not a string, or empty/whitespace-only after ``.strip()``. - Anthropic rejects such blocks with ``"each thinking block must contain - thinking"`` (whitespace-only included, verified live), regardless of any - signature they carry. ``redacted_thinking`` blocks are a different type - and always return False. + missing, not a string, or empty/whitespace-only after ``.strip()``, + whatever its signature. This is a shape check only; whether such a block + may be sent to Anthropic depends on the signature, see + :func:`is_empty_unsigned_thinking_block`. ``redacted_thinking`` blocks are + a different type and always return False. """ if not isinstance(block, dict) or block.get("type") != "thinking": return False @@ -1454,15 +1459,16 @@ def is_empty_unsigned_thinking_block(block: object) -> bool: """ True for an empty ``{"type": "thinking"}`` block carrying no signature. - The emit-side predicate: response paths drop a thinking block only when it - holds nothing the client could need. A signature-only block is a real - provider response (Bedrock Converse under adaptive thinking emits a - reasoning block with empty text and only a signature) and the client needs - the signature to replay reasoning across tool-use turns, so it must be - emitted. Request paths keep using :func:`is_empty_thinking_block`: - Anthropic rejects empty thinking blocks in request history regardless of - signature, and the inbound strip self-heals a replayed signature-only - block. + Used on both the emit and the request side: a thinking block is dropped + only when it holds nothing the provider or the client could use. A + signature-only block is a real provider response (Claude Opus 4.8 / + Fable 5 / Fable 5.1 return it by default, and Bedrock Converse under + adaptive thinking emits a reasoning block with empty text and only a + signature); the client needs the signature to replay reasoning across + tool-use turns, and Anthropic accepts the replayed block and binds it to + its prefix. Only an empty block with no signature (a turn a + non-Anthropic reasoning model served through the bridge) is rejected + with ``"each thinking block must contain thinking"`` and must go. """ if not isinstance(block, dict) or not is_empty_thinking_block(block): return False diff --git a/tests/unit/llms/anthropic/test_anthropic_common_utils.py b/tests/unit/llms/anthropic/test_anthropic_common_utils.py index 68e2e9650d5..e3cb138b508 100644 --- a/tests/unit/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/unit/llms/anthropic/test_anthropic_common_utils.py @@ -1409,11 +1409,11 @@ class TestAnthropicThinkingSignatureSelfHeal: assert is_empty_unsigned_thinking_block({"type": "redacted_thinking", "data": "opaque"}) is False assert is_empty_unsigned_thinking_block("not a dict") is False - def test_strip_empty_content_blocks_drops_empty_thinking_blocks(self): + def test_strip_empty_content_blocks_drops_empty_unsigned_thinking_blocks(self): """LIT-6357 ingestion half: an assistant tool-loop turn carrying an - empty (even signed) thinking block keeps its tool_use blocks and loses - the poison; whitespace-only counts as empty; a non-empty thinking block - and redacted_thinking are untouched.""" + empty, unsigned thinking block keeps its tool_use blocks and loses the + poison; whitespace-only counts as empty; a non-empty thinking block and + redacted_thinking are untouched.""" from litellm.llms.anthropic.common_utils import ( strip_empty_content_blocks_from_anthropic_messages, ) @@ -1423,7 +1423,7 @@ class TestAnthropicThinkingSignatureSelfHeal: {"role": "user", "content": "weather?"}, { "role": "assistant", - "content": [{"type": "thinking", "thinking": "", "signature": "sig_abc"}, tu], + "content": [{"type": "thinking", "thinking": "", "signature": ""}, tu], }, { "role": "assistant", @@ -1442,6 +1442,36 @@ class TestAnthropicThinkingSignatureSelfHeal: assert out[2]["content"][0]["thinking"] == "real plan" assert len(msgs[1]["content"]) == 2 + def test_strip_empty_content_blocks_keeps_signed_empty_thinking_blocks(self): + """A thinking block with empty text and a signature is the default + response shape of Claude Opus 4.8 / Fable 5 / Fable 5.1 (thinking.display + omitted): the signature IS the preserved reasoning, Anthropic accepts it + on input and binds it to its prefix. Dropping it silently discards the + model's reasoning on every replayed turn and leaves preserved-thinking + enforcement (thinking.block_binding) nothing to check. The message and + block order must survive untouched.""" + from litellm.llms.anthropic.common_utils import ( + strip_empty_content_blocks_from_anthropic_messages, + ) + + tu = {"type": "tool_use", "id": "toolu_01A", "name": "get_weather", "input": {"city": "Paris"}} + signed = {"type": "thinking", "thinking": "", "signature": "EqQBCkYIBxgCKkD"} + msgs = [ + {"role": "user", "content": "weather?"}, + {"role": "assistant", "content": [signed, tu]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01A", "content": "sunny"}]}, + { + "role": "assistant", + "content": [signed, {"type": "text", "text": ""}, {"type": "text", "text": "Sunny."}], + }, + ] + out = strip_empty_content_blocks_from_anthropic_messages(msgs) + assert len(out) == 4 + assert out[1] is msgs[1] + assert [b["type"] for b in out[3]["content"]] == ["thinking", "text"] + assert out[3]["content"][0] is signed + assert out[3]["content"][1]["text"] == "Sunny." + def test_strip_thinking_blocks_from_anthropic_messages(self): from litellm.llms.anthropic.common_utils import ( strip_thinking_blocks_from_anthropic_messages, From 05de300ed610e5bcdfe07f0d3dc351438145aedf Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 15 Oct 2026 11:49:53 -0700 Subject: [PATCH 06/12] fix(anthropic): keep signed empty thinking blocks and honor thinking.block_binding (#46342) * fix(anthropic): treat adaptive thinking as enabled when mapping response_format is_thinking_enabled only recognised thinking.type=enabled, so a request that paired thinking.type=adaptive with response_format took the synthetic JSON tool path with a forced tool_choice, which Anthropic rejects alongside thinking. The same check also gates the Bedrock Converse and Databricks forced tool_choice downgrades, so they now see adaptive thinking too. * fix(anthropic): keep signed thinking blocks with empty text on chat completions Models that default to display: omitted return signed thinking blocks whose text is empty, and Anthropic rejects the next turn when they are not replayed unchanged. is_unsignable_thinking_block dropped them along with unsigned blocks, and the inline content path required non-empty text, so a tool-use turn on Opus 4.7+ or the 5 families lost its thinking blocks and 400ed on the tool-result turn. Only a missing signature now drops a block, and the mid-conversation system placement mirrors that rule. The /v1/messages strip helper is left to #41362. * fix(anthropic): skip the thinking-strip retry when block_binding asks for an error The /v1/messages retry loop strips thinking blocks and resends on any invalid-thinking 400, including the binding mismatch a caller opted into with thinking.block_binding.prefix_mismatch_behavior: error. The retry decision now sees the request body and surfaces that 400 instead of overriding the caller's choice. #42888 covers the litellm_params side. * fix(bedrock): carry thinking.block_binding through the Invoke thinking rebuilds Bedrock Invoke rebuilt the thinking param from scratch when display: updates was requested, on both the /v1/messages and chat completions paths, and the legacy enabled-to-adaptive translation plus the clear_thinking injection did the same, so thinking.block_binding never reached Bedrock. Every rebuild now spreads the caller's thinking and only overrides the keys it owns, and AnthropicThinkingParam types block_binding. The thinking-binding-controls beta was already on the Bedrock allowlist. * fix(anthropic): keep forcing the response_format tool under adaptive thinking This reverts ad9e861ae7. Anthropic, Bedrock Invoke and Vertex all accept a forced tool_choice next to thinking.type=adaptive and only reject it with type=enabled, so counting adaptive as enabled fixed no live failure. It did drop the forced JSON tool, so Claude answered response_format requests in prose, and it downgraded a required tool_choice to auto on Bedrock Converse * test(e2e): drop the adaptive thinking tool round trip case With adaptive thinking the model often skips thinking on the weather prompt, so turn 1 carries no signed empty block and the case fails before the replay it is meant to cover. Anthropic also answers the replay with a 200 whether or not the block is forwarded, so the case could not tell the fix from main. The factory unit tests pin the outgoing shape instead * fix(anthropic): only skip the thinking-strip retry for a binding mismatch 400 prefix_mismatch_behavior: error asks Anthropic to reject a block bound to another conversation. It says nothing about unrelated bad blocks such as a null signature, so those still get the strip-and-retry self-heal. Also trims the empty-thinking docstrings to what the code guarantees and types the new test parameters * refactor(anthropic): narrow thinking fields with the mapping validator instead of typing.cast typing.cast is banned by the strict ruff gate (TID251), so the three new casts pushed that rule one over its merge-base count. The existing pydantic mapping validator gives the same typed view without an unchecked cast, and it now has a generic name since it is no longer only used for Claude Code payloads * refactor(anthropic): read thinking settings through the validated mapping so basedpyright stays at its merge-base count The Bedrock clear_thinking rewrite still narrowed with a cast and an untyped dict fallback, which added two reportUnknownVariableType errors. It now uses the same pydantic mapping validator, made public for that reuse, and the retry hook validates the request body itself so the handler passes a typed dict instead of an untyped one --------- Co-authored-by: yuneng-jiang (cherry picked from commit f470e418eaa73ef5fc9d001a8a06723364b55cdb) --- .../prompt_templates/common_utils.py | 11 +- .../prompt_templates/factory.py | 7 +- .../mid_conversation_system.py | 3 +- litellm/llms/anthropic/chat/transformation.py | 9 +- litellm/llms/anthropic/common_utils.py | 43 ++++-- .../pass_through/messages/transformation.py | 5 + .../anthropic_messages/transformation.py | 18 ++- .../anthropic_claude3_transformation.py | 12 +- .../anthropic_claude3_transformation.py | 24 ++- litellm/llms/custom_httpx/llm_http_handler.py | 4 +- litellm/types/llms/anthropic.py | 5 + ...llm_core_utils_prompt_templates_factory.py | 94 ++++++------ ...rompt_templates_mid_conversation_system.py | 2 + .../anthropic/test_anthropic_common_utils.py | 145 ++++++++++++++++++ ...ations_anthropic_claude3_transformation.py | 57 +++++++ .../test_anthropic_claude3_transformation.py | 37 +++++ .../custom_httpx/test_llm_http_handler.py | 65 ++++++++ 17 files changed, 442 insertions(+), 99 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 2f1a4147544..74d18e6f783 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -1993,10 +1993,8 @@ def is_unsignable_thinking_block(block: object) -> bool: Anthropic verifies the signature cryptographically, so a block with a null, empty, or missing signature (e.g. from an open-source reasoning model) is rejected with a 400, and so is a block whose signature or data carries - another provider's encrypted reasoning. It also rejects a `thinking` block - whose text is empty or whitespace-only ("each thinking block must contain - thinking"), regardless of signature, e.g. when a `thinking_blocks` history - item from a non-Anthropic reasoning provider is replayed through this path. + another provider's encrypted reasoning. A signed block is kept whatever its + text, since ``display: omitted`` returns signed blocks with empty ``thinking``. `redacted_thinking` blocks carry no signature and are always kept. """ if is_encrypted_reasoning_block(block): @@ -2007,10 +2005,7 @@ def is_unsignable_thinking_block(block: object) -> bool: if mapping.get("type") != "thinking": return False signature: Final = mapping.get("signature") - if not (isinstance(signature, str) and len(signature) > 0): - return True - thinking_text: Final = mapping.get("thinking") - return not (isinstance(thinking_text, str) and len(thinking_text.strip()) > 0) + return not (isinstance(signature, str) and len(signature) > 0) def strip_encrypted_reasoning_from_messages( diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 0c48b7c1c2a..84a3da288f3 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -2684,13 +2684,8 @@ def anthropic_messages_pt( if not isinstance(m, dict): continue # handle thinking blocks - thinking_block = cast(str, m.get("thinking", "")) text_block = cast(str, m.get("text", "")) - if ( - m.get("type", "") == "thinking" - and len(thinking_block) > 0 - and not is_unsignable_thinking_block(m) - ): # don't pass empty text blocks. anthropic api raises errors. + if m.get("type", "") == "thinking" and not is_unsignable_thinking_block(m): anthropic_message: ChatCompletionThinkingBlock | AnthropicMessagesTextParam = cast( ChatCompletionThinkingBlock, m ) diff --git a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py index 2169dfcad39..2a214f11a8b 100644 --- a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py +++ b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py @@ -257,8 +257,7 @@ def _assistant_part_renders(part: object) -> bool: """A text part always renders: the converter pads empty text with a placeholder.""" part_type: Final = message_field(part, "type") if part_type == "thinking": - thinking: Final = message_field(part, "thinking") - return isinstance(thinking, str) and bool(thinking) and _thinking_block_renders(part) + return _thinking_block_renders(part) return part_type in _RENDERED_ASSISTANT_PART_TYPES or ( isinstance(part_type, str) and part_type.endswith("_tool_result") ) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 490912d42eb..18fdbbba487 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -327,6 +327,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): def _resolved_provider(self) -> str: return self.custom_llm_provider or "anthropic" + @property + def _carries_thinking_block_binding(self) -> bool: + return self._resolved_provider in ("anthropic", "azure_ai") + @classmethod def get_config(cls, *, model: str | None = None): config: Final = super().get_config() @@ -1588,7 +1592,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): else: optional_params["thinking"] = value AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( - model=model, optional_params=optional_params, custom_llm_provider=self._resolved_provider + model=model, + optional_params=optional_params, + custom_llm_provider=self._resolved_provider, + carries_block_binding=self._carries_thinking_block_binding, ) elif param == "reasoning_effort": # Accept both string ("low") and dict ({"effort": "low", diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index fb2480cf0ac..4448c67a212 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -78,7 +78,7 @@ _INFERENCE_PROFILE_MINOR_RE: Final = re.compile(r":\d+$") _DATED_RELEASE_SUFFIX_RE: Final = re.compile(r"-\d{8}$") _DOTTED_VERSION_RE: Final = re.compile(r"(\d)\.(\d)") _CLAUDE_CODE_BILLING_HEADER_PREFIX: Final = "x-anthropic-billing-header:" -_CLAUDE_CODE_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object]) +_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object]) _CLAUDE_CODE_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) @@ -127,9 +127,9 @@ def is_claude_code_user_agent(user_agent: str) -> bool: return user_agent.startswith(_CLAUDE_CODE_USER_AGENT_PREFIXES) -def _validated_claude_code_mapping(value: object) -> dict[object, object] | None: +def validated_object_mapping(value: object) -> dict[object, object] | None: try: - return _CLAUDE_CODE_OBJECT_MAPPING_ADAPTER.validate_python(value) + return _OBJECT_MAPPING_ADAPTER.validate_python(value) except ValidationError: return None @@ -166,7 +166,7 @@ def _claude_code_billing_texts(system: object) -> tuple[str, ...] | None: blocks: Final = _validated_claude_code_list(system) if blocks is None: return None - block_mappings: Final = tuple(_validated_claude_code_mapping(block) for block in blocks) + block_mappings: Final = tuple(validated_object_mapping(block) for block in blocks) if any(block is None for block in block_mappings): return None text_values: Final = tuple( @@ -199,7 +199,7 @@ def is_claude_code_one_shot_subagent_request( tools: object, user_agent: str | None, ) -> bool: - only_message: Final = _validated_claude_code_mapping(messages[0]) if len(messages) == 1 else None + only_message: Final = validated_object_mapping(messages[0]) if len(messages) == 1 else None return ( user_agent is not None and is_claude_code_user_agent(user_agent) @@ -748,6 +748,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): model: str, optional_params: MutableMapping[str, object], # mutable-ok: in-place out-param like the sibling helpers custom_llm_provider: str, + carries_block_binding: bool = False, ) -> None: """Translate legacy ``thinking.type=enabled`` to adaptive for the adaptive-thinking models that reject it (4.7+ and the 5 families). @@ -760,21 +761,25 @@ class AnthropicModelInfo(BaseLLMModelInfo): return if AnthropicModelInfo._supports_legacy_thinking(model, custom_llm_provider): return - thinking: Final = optional_params.get("thinking") - if not isinstance(thinking, dict) or thinking.get("type") != "enabled": + thinking: Final = validated_object_mapping(optional_params.get("thinking")) + if thinking is None or thinking.get("type") != "enabled": return + budget_tokens: Final = thinking.get("budget_tokens") or 0 effort: Final = AnthropicModelInfo._legacy_budget_to_effort( model=model, - budget_tokens=int(thinking.get("budget_tokens") or 0), + budget_tokens=int(budget_tokens) if isinstance(budget_tokens, (int, float, str)) else 0, custom_llm_provider=custom_llm_provider, ) existing_output_config: Final = optional_params.get("output_config") display: Final = thinking.get("display") - if display in ("summarized", "omitted"): - optional_params["thinking"] = {"type": "adaptive", "display": display} - else: - optional_params["thinking"] = {"type": "adaptive"} + supported_display: Final = {"display": display} if display in ("summarized", "omitted") else {} + carried_block_binding: Final = ( + {"block_binding": thinking["block_binding"]} + if carries_block_binding and "block_binding" in thinking + else {} + ) + optional_params["thinking"] = {**carried_block_binding, "type": "adaptive", **supported_display} optional_params["output_config"] = { "effort": effort, **(existing_output_config if isinstance(existing_output_config, dict) else MappingProxyType({})), @@ -1320,6 +1325,20 @@ def is_anthropic_invalid_thinking_block_error(error_text: str) -> bool: return "must contain thinking" in lower +def is_anthropic_thinking_binding_mismatch_error(error_text: str) -> bool: + """messages.N.content.M: Invalid `signature` in `thinking` block. The block is bound to a different conversation.""" + return "bound to a different conversation" in error_text.lower() + + +def anthropic_request_errors_on_thinking_binding_mismatch(request_body: object) -> bool: + """``thinking.block_binding.prefix_mismatch_behavior: error`` is the caller asking + for the 400, so a retry that strips thinking must not override it.""" + body: Final = validated_object_mapping(request_body) + thinking: Final = validated_object_mapping(body.get("thinking")) if body is not None else None + binding: Final = validated_object_mapping(thinking.get("block_binding")) if thinking is not None else None + return binding is not None and binding.get("prefix_mismatch_behavior") == "error" + + def strip_thinking_blocks_from_anthropic_messages(messages: Sequence[object]) -> list[object]: """ Return a new message list with thinking / redacted_thinking content blocks removed diff --git a/litellm/llms/anthropic/pass_through/messages/transformation.py b/litellm/llms/anthropic/pass_through/messages/transformation.py index 2fbb51ec949..b3a83e60bcc 100644 --- a/litellm/llms/anthropic/pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/pass_through/messages/transformation.py @@ -63,6 +63,10 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): def _resolved_provider(self) -> str: return self.custom_llm_provider or "anthropic" + @property + def _carries_thinking_block_binding(self) -> bool: + return self._resolved_provider in ("anthropic", "azure_ai") + def get_supported_anthropic_messages_params(self, model: str) -> list: return [ "messages", @@ -514,6 +518,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): model=model, optional_params=anthropic_messages_optional_request_params, custom_llm_provider=self._resolved_provider, + carries_block_binding=self._carries_thinking_block_binding, ) self._translate_adaptive_effort_for_non_adaptive_model( diff --git a/litellm/llms/base_llm/anthropic_messages/transformation.py b/litellm/llms/base_llm/anthropic_messages/transformation.py index 101a5e6c58c..d87ccb50af4 100644 --- a/litellm/llms/base_llm/anthropic_messages/transformation.py +++ b/litellm/llms/base_llm/anthropic_messages/transformation.py @@ -1,6 +1,6 @@ from abc import ABC, abstractmethod from collections.abc import AsyncIterator -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Final import httpx @@ -156,16 +156,28 @@ class BaseAnthropicMessagesConfig(ABC): """ return 2 - def should_retry_anthropic_messages_on_http_error(self, e: httpx.HTTPStatusError, litellm_params: dict) -> bool: + def should_retry_anthropic_messages_on_http_error( + self, e: httpx.HTTPStatusError, litellm_params: dict, request_body: object = None + ) -> bool: """ When True, async_anthropic_messages_handler will transform the request body and issue one more attempt (bounded by max_retry_on_anthropic_messages_http_error). + A binding-mismatch 400 is not retried when the body's ``thinking.block_binding`` + asked for it, since the retry would strip the very blocks the caller pinned. """ from litellm.llms.anthropic.common_utils import ( + anthropic_request_errors_on_thinking_binding_mismatch, is_anthropic_invalid_thinking_block_error, + is_anthropic_thinking_binding_mismatch_error, ) - return e.response.status_code == 400 and is_anthropic_invalid_thinking_block_error(e.response.text) + if e.response.status_code != 400: + return False + error_text: Final[str] = e.response.text + caller_asked_for_this_error: Final = anthropic_request_errors_on_thinking_binding_mismatch( + request_body + ) and is_anthropic_thinking_binding_mismatch_error(error_text) + return not caller_asked_for_this_error and is_anthropic_invalid_thinking_block_error(error_text) def transform_anthropic_messages_request_on_http_error(self, e: httpx.HTTPStatusError, request_data: dict) -> dict: """ diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 0be087d2dd4..b3673eb1e6c 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -67,6 +67,10 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): def custom_llm_provider(self) -> str | None: return "bedrock" + @property + def _carries_thinking_block_binding(self) -> bool: + return True + def should_strip_billing_metadata(self) -> bool: return True @@ -115,7 +119,10 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): model = original_model AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( - model=original_model, optional_params=optional_params, custom_llm_provider="bedrock" + model=original_model, + optional_params=optional_params, + custom_llm_provider="bedrock", + carries_block_binding=self._carries_thinking_block_binding, ) translated_thinking: Final = optional_params.get("thinking") if ( @@ -123,8 +130,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): and isinstance(translated_thinking, dict) and translated_thinking.get("type") == "adaptive" ): - thinking_with_display: Final[AnthropicThinkingParam] = {"type": "adaptive", "display": "updates"} - optional_params["thinking"] = thinking_with_display + optional_params["thinking"] = {**translated_thinking, "display": "updates"} # The stub model hides the original model from the parent's forced-tool-use backstop response_format_tool_choice: Final = optional_params.get("tool_choice") diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 6234ca3a9c3..cf733f5605c 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -17,7 +17,7 @@ from litellm.llms.anthropic.chat.transformation import ( DROP_UNSUPPORTED_OUTPUT_CONFIG_WARNING, AnthropicConfig, ) -from litellm.llms.anthropic.common_utils import AnthropicModelInfo +from litellm.llms.anthropic.common_utils import AnthropicModelInfo, validated_object_mapping from litellm.llms.anthropic.pass_through.messages.transformation import ( AnthropicMessagesConfig, ) @@ -49,7 +49,6 @@ from litellm.types.llms.anthropic import ( ANTHROPIC_BETA_HEADER_VALUES, ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER, ANTHROPIC_TOOL_SEARCH_BETA_HEADER, - AnthropicThinkingParam, ) from litellm.types.llms.bedrock import BedrockInvokeAnthropicMessagesRequest from litellm.types.llms.openai import AllMessageValues @@ -82,6 +81,10 @@ class AmazonAnthropicClaudeMessagesConfig( def custom_llm_provider(self) -> str | None: return "bedrock" + @property + def _carries_thinking_block_binding(self) -> bool: + return True + @property def beta_headers_provider(self) -> str: return self.custom_llm_provider or "bedrock" @@ -349,18 +352,10 @@ class AmazonAnthropicClaudeMessagesConfig( if not isinstance(output_config, dict): output_config = {} output_config.setdefault("effort", self._effort_from_thinking_budget(budget_tokens)) - thinking: Final = anthropic_messages_request.get("thinking") - display: Final = thinking.get("display") if isinstance(thinking, dict) else None + thinking: Final = validated_object_mapping(anthropic_messages_request.get("thinking")) + carried: Final = {key: value for key, value in (thinking or {}).items() if key not in ("type", "budget_tokens")} anthropic_messages_request["output_config"] = output_config - if display is None: - adaptive_thinking: Final[AnthropicThinkingParam] = {"type": "adaptive"} - anthropic_messages_request["thinking"] = adaptive_thinking - else: - adaptive_thinking_with_display: Final[AnthropicThinkingParam] = { - "type": "adaptive", - "display": display, - } - anthropic_messages_request["thinking"] = adaptive_thinking_with_display + anthropic_messages_request["thinking"] = {**carried, "type": "adaptive"} verbose_logger.debug( "Bedrock clear_thinking_20251015: injected adaptive thinking with effort=%s for model=%s", output_config.get("effort"), @@ -701,8 +696,7 @@ class AmazonAnthropicClaudeMessagesConfig( and isinstance(translated_thinking, dict) and translated_thinking.get("type") == "adaptive" ): - thinking_with_display: Final[AnthropicThinkingParam] = {"type": "adaptive", "display": "updates"} - anthropic_messages_request["thinking"] = thinking_with_display + anthropic_messages_request["thinking"] = {**translated_thinking, "display": "updates"} self._normalize_system_role_messages(anthropic_messages_request, model=model) ######################################################### ############## BEDROCK Invoke SPECIFIC TRANSFORMATION ### diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 8383078b7ba..97aa7bbd7be 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2018,7 +2018,7 @@ class BaseLLMHTTPHandler: # str when the caller passes a pre-serialized (unsigned) body to avoid # re-dumping; bytes when a provider signed the request (e.g. Bedrock). signed_json_body: str | bytes | None, - request_body: dict, + request_body: dict[str, object], stream: bool, logging_obj: LiteLLMLoggingObj, provider_config: BaseAnthropicMessagesConfig, @@ -2045,7 +2045,7 @@ class BaseLLMHTTPHandler: except httpx.HTTPStatusError as e: hit_max_attempt = attempt_idx + 1 == max_attempts should_retry = provider_config.should_retry_anthropic_messages_on_http_error( - e=e, litellm_params=litellm_params_dict + e=e, litellm_params=litellm_params_dict, request_body=request_body ) if should_retry and not hit_max_attempt: if logging_obj.baseline_cache_context is not None: diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index ee357cd6581..4f147437b38 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -730,10 +730,15 @@ ANTHROPIC_API_ONLY_HEADERS: Final = { # fails if calling anthropic on vertex ai } +class AnthropicThinkingBlockBinding(TypedDict, total=False): + prefix_mismatch_behavior: ReadOnly[Literal["error", "drop_block"]] + + class AnthropicThinkingParam(TypedDict, total=False): type: ReadOnly[Literal["enabled", "adaptive", "disabled"]] budget_tokens: int display: ReadOnly[Literal["summarized", "omitted", "updates"]] + block_binding: ReadOnly[AnthropicThinkingBlockBinding] class ANTHROPIC_HOSTED_TOOLS(str, Enum): diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 8d2e6b9fd0c..a8adc33bf1f 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -27,7 +27,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( ollama_pt, sanitize_messages_for_tool_calling, ) -from litellm.types.llms.openai import ChatCompletionToolMessage +from litellm.types.llms.openai import ChatCompletionAssistantMessage, ChatCompletionToolMessage from litellm.utils import validate_and_fix_openai_messages @@ -3922,44 +3922,42 @@ def test_anthropic_messages_pt_drops_a_system_message_with_no_text(): assert [m["role"] for m in result] == ["user", "assistant"] -def test_anthropic_messages_pt_drops_empty_but_signed_thinking_block(): - """ - Anthropic rejects a `thinking` block whose `thinking` text is empty, even - when it carries a valid-looking signature, with: - 400 messages.N.content.M.thinking: each thinking block must contain thinking - This shape is reachable via cross-provider replay of a `thinking_blocks` - history item (see PR #36033), so `is_unsignable_thinking_block()` must - also check the thinking text, not just the signature. - """ +@pytest.mark.parametrize( + "assistant_message", + ( + { + "role": "assistant", + "content": "4", + "thinking_blocks": [{"type": "thinking", "thinking": "", "signature": "sig_abc123_looks_valid"}], + }, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "", "signature": "sig_abc123_looks_valid"}, + {"type": "text", "text": "4"}, + ], + }, + ), + ids=("thinking_blocks", "inline_content"), +) +def test_anthropic_messages_pt_keeps_empty_but_signed_thinking_block( + assistant_message: ChatCompletionAssistantMessage, +) -> None: + """``display: omitted`` returns signed thinking blocks with empty text, and they + must reach the next turn unchanged.""" from litellm.litellm_core_utils.prompt_templates.factory import ( anthropic_messages_pt, ) - messages = [ - {"role": "user", "content": "What's 2+2?"}, - { - "role": "assistant", - "content": "4", - "thinking_blocks": [ - { - "type": "thinking", - "thinking": "", - "signature": "sig_abc123_looks_valid", - } - ], - }, - ] - result = anthropic_messages_pt( - messages=messages, - model="claude-sonnet-4-5-20250929", + messages=[{"role": "user", "content": "What's 2+2?"}, assistant_message], + model="claude-opus-5-5", llm_provider="anthropic", ) assistant_msg = result[1] assert isinstance(assistant_msg["content"], list) - content_types = [block.get("type") for block in assistant_msg["content"]] - assert "thinking" not in content_types, "empty-text thinking block must be dropped even though it has a signature" + assert assistant_msg["content"][0] == {"type": "thinking", "thinking": "", "signature": "sig_abc123_looks_valid"} def test_anthropic_messages_pt_keeps_non_empty_signed_thinking_block(): @@ -4073,29 +4071,31 @@ def test_anthropic_messages_pt_drops_unsigned_thinking_block(): assert "thinking" not in content_types, "unsigned thinking block must still be dropped" -def test_is_unsignable_thinking_block_treats_whitespace_only_as_empty(): +def test_anthropic_messages_pt_drops_unsigned_empty_thinking_block(): """ - Edge case: a `thinking` field that is present but whitespace-only (e.g. - a single trailing newline forwarded from another provider's empty - reasoning summary) is functionally empty and Anthropic's API will still - reject it with "each thinking block must contain thinking". We treat it - the same as a fully empty string and drop the block. - - The check lives in the shared `is_unsignable_thinking_block` helper, which - `_drop_unsignable_thinking_blocks` calls standalone, so the whitespace-aware - test has to hold there rather than only at the factory call site. + An empty block without a signature is the shape Claude Code replays after a + non-Anthropic turn, and Anthropic rejects it, so only the missing signature + decides whether a block is dropped. """ - from litellm.litellm_core_utils.prompt_templates.common_utils import ( - is_unsignable_thinking_block, + from litellm.litellm_core_utils.prompt_templates.factory import ( + anthropic_messages_pt, ) - whitespace_only_block = { - "type": "thinking", - "thinking": " \n\t ", - "signature": "sig_abc123_looks_valid", - } + result = anthropic_messages_pt( + messages=[ + {"role": "user", "content": "What's 2+2?"}, + { + "role": "assistant", + "content": [{"type": "thinking", "thinking": "", "signature": ""}, {"type": "text", "text": "4"}], + }, + ], + model="claude-opus-5-5", + llm_provider="anthropic", + ) - assert is_unsignable_thinking_block(whitespace_only_block) is True + assistant_msg = result[1] + assert isinstance(assistant_msg["content"], list) + assert [block.get("type") for block in assistant_msg["content"]] == ["text"] _CONTENT_LESS_USER_MESSAGES: Final = ({"role": "user"}, {"role": "user", "content": None}) diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py index d1e23a17747..2d9c4c99c3e 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py @@ -336,6 +336,7 @@ def test_flagged_placement_keeps_a_system_whose_empty_assistant_follower_ends_th }, {"role": "assistant", "content": None, "thinking_blocks": [{"type": "redacted_thinking", "data": "x"}]}, {"role": "assistant", "content": [{"type": "thinking", "thinking": "hm", "signature": "s"}]}, + {"role": "assistant", "content": [{"type": "thinking", "thinking": "", "signature": "s"}]}, {"role": "assistant", "content": None, "function_call": {"name": "f", "arguments": "{}"}}, litellm.Message( role="assistant", @@ -348,6 +349,7 @@ def test_flagged_placement_keeps_a_system_whose_empty_assistant_follower_ends_th "signed-thinking-block", "redacted-thinking-block", "signed-thinking-part", + "signed-empty-thinking-part", "function-call", "pydantic-tool-calls", ], diff --git a/tests/unit/llms/anthropic/test_anthropic_common_utils.py b/tests/unit/llms/anthropic/test_anthropic_common_utils.py index e3cb138b508..677ef16fc64 100644 --- a/tests/unit/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/unit/llms/anthropic/test_anthropic_common_utils.py @@ -13,10 +13,12 @@ Verifies that: import json import os import sys +from collections.abc import Mapping from types import SimpleNamespace from typing import Final from unittest.mock import patch +import httpx import pytest sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../.."))) @@ -1319,6 +1321,12 @@ class TestPassthroughAuthToken: assert url == "https://custom.example.com/v1/messages" +_BOUND_THINKING_ERROR: Final = ( + "messages.1.content.0: Invalid `signature` in `thinking` block. The block is bound to a different conversation." +) +_NULL_SIGNATURE_ERROR: Final = "messages.1.content.0.thinking.signature.str: Input should be a valid string" + + class TestAnthropicThinkingSignatureSelfHeal: """Helpers for retrying after invalid thinking blocks in replayed history: invalid encrypted signatures, and blocks with empty thinking text.""" @@ -1984,6 +1992,49 @@ class TestAnthropicThinkingSignatureSelfHeal: assert "thinking" not in data assert data["messages"] == [] + @pytest.mark.parametrize( + ("thinking", "error_message", "expected"), + ( + ( + {"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": "error"}}, + _BOUND_THINKING_ERROR, + False, + ), + ( + {"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": "error"}}, + _NULL_SIGNATURE_ERROR, + True, + ), + ( + {"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": "drop_block"}}, + _BOUND_THINKING_ERROR, + True, + ), + ({"type": "adaptive"}, _BOUND_THINKING_ERROR, True), + ), + ids=("error", "error_on_unrelated_bad_block", "drop_block", "unbound"), + ) + def test_anthropic_messages_retry_honors_thinking_block_binding( + self, thinking: Mapping[str, object], error_message: str, expected: bool + ) -> None: + from litellm.llms.anthropic.pass_through.messages.transformation import ( + AnthropicMessagesConfig, + ) + + req: Final = httpx.Request("POST", "https://api.anthropic.com/v1/messages") + resp: Final = httpx.Response( + 400, + request=req, + text=json.dumps({"type": "error", "error": {"type": "invalid_request_error", "message": error_message}}), + ) + err: Final = httpx.HTTPStatusError("bad", request=req, response=resp) + + should_retry = AnthropicMessagesConfig().should_retry_anthropic_messages_on_http_error( + err, {}, request_body={"model": "claude-opus-5-5", "thinking": thinking} + ) + + assert should_retry is expected + class TestClaudeOpus48AdaptiveThinking: """Opus 4.8 requires adaptive thinking (``thinking.type='adaptive'`` + @@ -2482,6 +2533,100 @@ def test_shared_legacy_thinking_translation_preserves_supported_display( assert optional_params["thinking"] == expected_thinking +LEGACY_THINKING_WITH_BINDING: Final = { + "type": "enabled", + "budget_tokens": 2048, + "block_binding": {"prefix_mismatch_behavior": "error"}, + "future_key": 1, +} + + +def _rebuilt_thinking(carries_block_binding: bool) -> dict[str, object]: + return { + "type": "adaptive", + **({"block_binding": LEGACY_THINKING_WITH_BINDING["block_binding"]} if carries_block_binding else {}), + } + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize( + ("provider", "carries_block_binding"), + ( + ("anthropic", True), + ("azure_ai", True), + ("bedrock_invoke", True), + ("vertex_ai", False), + ("databricks", False), + ("bedrock_converse", False), + ), +) +def test_chat_legacy_thinking_rebuild_carries_only_block_binding_and_only_where_supported( + provider: str, carries_block_binding: bool +) -> None: + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + from litellm.llms.azure_ai.anthropic.transformation import AzureAnthropicConfig + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeConfig, + ) + from litellm.llms.databricks.chat.transformation import DatabricksConfig + from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import VertexAIAnthropicConfig + + config, model = { + "anthropic": (AnthropicConfig(), "claude-opus-4-8"), + "azure_ai": (AzureAnthropicConfig(), "claude-opus-4-8"), + "bedrock_invoke": (AmazonAnthropicClaudeConfig(), "us.anthropic.claude-opus-4-8"), + "vertex_ai": (VertexAIAnthropicConfig(), "claude-opus-4-8"), + "databricks": (DatabricksConfig(), "databricks-claude-opus-4-8"), + "bedrock_converse": (AmazonConverseConfig(), "us.anthropic.claude-opus-4-8"), + }[provider] + + mapped: Final = config.map_openai_params( + non_default_params={"max_tokens": 4096, "thinking": dict(LEGACY_THINKING_WITH_BINDING)}, + optional_params={}, + model=model, + drop_params=False, + ) + + assert mapped["thinking"] == _rebuilt_thinking(carries_block_binding) + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize( + ("provider", "carries_block_binding"), + (("anthropic", True), ("azure_ai", True), ("bedrock_invoke", True), ("vertex_ai", False)), +) +def test_messages_legacy_thinking_rebuild_carries_only_block_binding_and_only_where_supported( + provider: str, carries_block_binding: bool +) -> None: + from litellm.llms.anthropic.pass_through.messages.transformation import AnthropicMessagesConfig + from litellm.llms.azure_ai.anthropic.messages_transformation import AzureAnthropicMessagesConfig + from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeMessagesConfig, + ) + from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation import ( + VertexAIPartnerModelsAnthropicMessagesConfig, + ) + from litellm.types.router import GenericLiteLLMParams + + config, model = { + "anthropic": (AnthropicMessagesConfig(), "claude-opus-4-8"), + "azure_ai": (AzureAnthropicMessagesConfig(), "claude-opus-4-8"), + "bedrock_invoke": (AmazonAnthropicClaudeMessagesConfig(), "us.anthropic.claude-opus-4-8"), + "vertex_ai": (VertexAIPartnerModelsAnthropicMessagesConfig(), "claude-opus-4-8"), + }[provider] + + request: Final = config.transform_anthropic_messages_request( + model=model, + messages=[{"role": "user", "content": "Reply with OK"}], + anthropic_messages_optional_request_params={"max_tokens": 4096, "thinking": dict(LEGACY_THINKING_WITH_BINDING)}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert request["thinking"] == _rebuilt_thinking(carries_block_binding) + + @pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config") @pytest.mark.parametrize("action", (None, "tool_addition", "tool_removal")) @pytest.mark.parametrize("explicit_beta", (False, True)) diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 2c82ba1c5b8..ce5c5ad1b27 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -759,6 +759,30 @@ def test_bedrock_chat_invoke_tool_based_response_format_still_upgrades_legacy_th assert result["output_config"] == {"effort": "high"} +def test_bedrock_chat_invoke_adaptive_thinking_keeps_the_forced_response_format_tool(local_model_cost_map): + """Anthropic accepts a forced tool_choice next to adaptive thinking, and dropping + it lets the model answer in prose instead of the requested JSON.""" + result = AmazonAnthropicClaudeConfig().map_openai_params( + non_default_params={ + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "test_schema", + "schema": {"type": "object", "properties": {"result": {"type": "string"}}}, + }, + }, + "thinking": {"type": "adaptive"}, + "max_tokens": 8192, + }, + optional_params={}, + model="us.anthropic.claude-opus-5", + drop_params=False, + ) + + assert result["tool_choice"] == {"type": "tool", "name": result["tools"][0]["name"]} + assert result["thinking"] == {"type": "adaptive"} + + def test_bedrock_chat_invoke_response_format_stub_still_upgrades_legacy_thinking(local_model_cost_map): """Regression: the tool-based ``response_format`` path swaps in a Claude 3 stub model before the shared Anthropic mapping, which hid the adaptive-only model @@ -973,6 +997,39 @@ def test_bedrock_chat_invoke_preserves_display_when_translating_legacy_thinking( assert ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER in result.get("anthropic_beta", []) +BLOCK_BINDING: Final = {"prefix_mismatch_behavior": "error"} + + +@pytest.mark.parametrize( + "requested_thinking", + ( + {"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING}, + {"type": "enabled", "budget_tokens": 2048, "display": "updates", "block_binding": BLOCK_BINDING}, + ), + ids=("adaptive", "legacy"), +) +def test_bedrock_chat_invoke_keeps_block_binding_with_display_updates( + local_model_cost_map, local_beta_headers_config, requested_thinking +) -> None: + config: Final = AmazonAnthropicClaudeConfig() + model: Final = "us.anthropic.claude-opus-5" + optional_params: Final = config.map_openai_params( + non_default_params={"max_tokens": 512, "thinking": requested_thinking}, + optional_params={}, + model=model, + drop_params=False, + ) + result: Final = config.transform_request( + model=model, + messages=[{"role": "user", "content": "Reply with OK"}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert result.get("thinking") == {"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING} + + FINE_GRAINED_TOOL_STREAMING_BETA: Final = "fine-grained-tool-streaming-2025-05-14" EAGER_TOOL_SCHEMA: Final = {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]} diff --git a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index a269d556262..f303c899898 100644 --- a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -3529,6 +3529,43 @@ def test_bedrock_clear_thinking_preserves_display_updates() -> None: assert ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER in result.get("anthropic_beta", []) +BLOCK_BINDING: Final = {"prefix_mismatch_behavior": "error"} + + +@pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config") +@pytest.mark.parametrize( + ("model", "requested_thinking", "extra_params"), + ( + ("eu.anthropic.claude-opus-5", {"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING}, {}), + ( + "eu.anthropic.claude-opus-5", + {"type": "enabled", "budget_tokens": 24000, "display": "updates", "block_binding": BLOCK_BINDING}, + {}, + ), + ( + "us.anthropic.claude-opus-4-6", + {"type": "enabled", "budget_tokens": 2048, "display": "updates", "block_binding": BLOCK_BINDING}, + {"context_management": {"edits": [{"type": "clear_thinking_20251015"}]}}, + ), + ), + ids=("adaptive", "legacy", "legacy_clear_thinking"), +) +def test_bedrock_messages_keeps_block_binding_with_display_updates( + model: str, requested_thinking: dict[str, object], extra_params: dict[str, object] +) -> None: + from litellm.types.router import GenericLiteLLMParams + + result: Final = AmazonAnthropicClaudeMessagesConfig().transform_anthropic_messages_request( + model=model, + messages=[{"role": "user", "content": "Reply with OK"}], + anthropic_messages_optional_request_params={"max_tokens": 512, "thinking": requested_thinking, **extra_params}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result.get("thinking") == {"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING} + + @pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config") @pytest.mark.parametrize("action", (None, "tool_addition", "tool_removal")) @pytest.mark.parametrize("explicit_beta", (False, True)) diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index f3332cb513c..a6b6f5196d0 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -1636,6 +1636,71 @@ async def test_anthropic_post_retry_reserializes_mutated_body(): assert "MUTATED" in second_sent # ... the mutated body +_BOUND_THINKING_ERROR: Final = ( + '{"type":"error","error":{"type":"invalid_request_error","message":"messages.1.content.0: ' + 'Invalid `signature` in `thinking` block. The block is bound to a different conversation."}}' +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("prefix_mismatch_behavior", "expected_posts"), + (("error", 1), ("drop_block", 2)), +) +async def test_anthropic_post_skips_the_thinking_strip_retry_when_block_binding_errors( + prefix_mismatch_behavior: str, expected_posts: int +) -> None: + """A caller that pinned ``prefix_mismatch_behavior: error`` asked for the 400, + so the handler must surface it instead of stripping the bound blocks and retrying.""" + from litellm.llms.anthropic.pass_through.messages.transformation import AnthropicMessagesConfig + + handler = BaseLLMHTTPHandler() + thinking_block = {"type": "thinking", "thinking": "", "signature": "sig"} + request_body = { + "model": "claude-opus-5-5", + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": [thinking_block, {"type": "text", "text": "hello"}]}, + ], + "thinking": {"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": prefix_mismatch_behavior}}, + } + request = httpx.Request("POST", "https://api.anthropic.com/v1/messages") + http_client = Mock() + http_client.post = AsyncMock( + side_effect=[ + httpx.Response(400, request=request, text=_BOUND_THINKING_ERROR), + httpx.Response(200, request=request, json={"type": "message", "content": []}), + ] + ) + logging_obj: Final = Mock(baseline_cache_context=None) + logging_obj.model_call_details = {} + + async def _post() -> httpx.Response: + return await handler._async_post_anthropic_messages_with_http_error_retry( + async_httpx_client=http_client, + request_url="https://api.anthropic.com/v1/messages", + headers={}, + signed_json_body=None, + request_body=request_body, + stream=False, + logging_obj=logging_obj, + provider_config=AnthropicMessagesConfig(), + litellm_params=GenericLiteLLMParams(), + api_key="k", + model="claude-opus-5-5", + ) + + if expected_posts == 1: + with pytest.raises(BaseLLMException) as raised: + await _post() + assert raised.value.status_code == 400 + assert "bound to a different conversation" in raised.value.message + assert request_body["messages"][1]["content"][0] == thinking_block + else: + assert (await _post()).status_code == 200 + assert http_client.post.await_count == expected_posts + + def test_base_responses_config_sign_request_is_noop_by_default(): """Default responses sign_request must be a no-op: unchanged headers, no signed body. From 3ed027b5a684242bb6d7b2bead913081892978c6 Mon Sep 17 00:00:00 2001 From: Mrunmayee Daware <165297479+Mrun25@users.noreply.github.com> Date: Sun, 18 Oct 2026 12:48:11 +0530 Subject: [PATCH 07/12] fix(proxy): authorize the public name of a team-scoped model on batch retrieve/cancel (#47067) * fix(proxy): authorize the public name of a team-scoped model on batch retrieve/cancel A unified batch id resolves to the deployment's internal model_name (model_name_{team_id}_{uuid}), which is in no key or team allowlist, so authorize_model_for_key always denied and team keys got 403 retrieving or cancelling their own batches. Authorize against model_info team_public_model_name when the resolved deployment is team-scoped, falling back to the resolved name otherwise. Fixes #46958 * fix(proxy): authorize team-scoped batch models by their public name on both auth paths * fix(proxy): accept the owning team's public name for its own team deployment on every auth path Managed batch, file, fine-tuning, vector store, character and video ids resolve to the team deployment's internal name again, so teams that granted the internal name keep access. The model access check now also accepts the owning team's public name and its legacy team aliases for that internal name, and only when the caller's team owns the deployment. * fix(proxy): match a team deployment's public and alias names against the caller's own grants --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> (cherry picked from commit e979bd0ecf836976373e12624585962dd3356f3a) --- litellm/proxy/auth/auth_checks.py | 30 +++++ tests/unit/proxy/auth/test_auth_utils.py | 119 +++++++++++++++++ .../proxy/batches_endpoints/test_endpoints.py | 125 +++++++++++++++++- 3 files changed, 273 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index dbd6f28a183..2afe56da3fb 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4420,6 +4420,24 @@ def _check_model_access_helper( return True +def _owning_team_names_for_team_deployment( + model: str, + llm_router: Router | None, + team_id: str | None, + team_model_aliases: Mapping[str, str] | None, +) -> tuple[str, ...]: + if llm_router is None or team_id is None or not model.startswith(f"model_name_{team_id}_"): + return () + deployment: Final = llm_router.get_deployment_by_model_group_name(model_group_name=model) + public_name: Final = ( + deployment.model_info.team_public_model_name + if deployment is not None and deployment.model_info.team_id == team_id + else None + ) + legacy_aliases: Final = tuple(alias for alias, target in (team_model_aliases or {}).items() if target == model) + return (*((public_name,) if public_name else ()), *legacy_aliases) + + def _can_object_call_model( model: str | list[str], llm_router: Router | None, @@ -4491,6 +4509,13 @@ def _can_object_call_model( *((global_or_router_alias_target,) if global_or_router_alias_target else ()), ) ) + owning_team_names: Final = ( + () + if key_alias_applied + else _owning_team_names_for_team_deployment( + model=model, llm_router=llm_router, team_id=team_id, team_model_aliases=team_model_aliases + ) + ) ## check model access for alias + underlying model - allow if either is in allowed models for m in potential_models: @@ -4502,6 +4527,11 @@ def _can_object_call_model( team_id=team_id, ): return True + if any( + _check_model_access_helper(model=name, llm_router=llm_router, models=models, team_id=team_id) + for name in owning_team_names + ): + return True internal_message: Final = ( f"{object_type} not allowed to access model. This {object_type} can only access models={models}. " diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index aadb5881ce0..03817bc7819 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -30,6 +30,7 @@ from litellm.proxy.auth.auth_utils import ( log_once_if_budget_reservation_disabled, warn_once_if_custom_auth_skips_common_checks, ) +from litellm.router import Router class TestCustomAuthCommonChecksWarning: @@ -1167,6 +1168,124 @@ async def test_managed_batch_routes_pass_team_model_access_check(route, request_ ) +_TEAM_DEPLOYMENT_ID: Final = "team-a-deployment-id" +_TEAM_INTERNAL_MODEL_NAME: Final = "model_name_team-a_6d1f3c2e-0000-4000-8000-000000000000" +_TEAM_PUBLIC_MODEL_NAME: Final = "team-batch-model" +_TEAM_MANAGED_BATCH_ID: Final = _encode_managed_id( + f"litellm_proxy;model_id:{_TEAM_DEPLOYMENT_ID};llm_batch_id:provider-batch-789" +) + + +def _team_scoped_router(public_model_name: str | None = _TEAM_PUBLIC_MODEL_NAME) -> Router: + return Router( + model_list=[ + { + "model_name": _TEAM_INTERNAL_MODEL_NAME, + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"}, + "model_info": { + "id": _TEAM_DEPLOYMENT_ID, + "team_id": "team-a", + **({"team_public_model_name": public_model_name} if public_model_name else {}), + }, + }, + ] + ) + + +def _team_scoped_character_id() -> str: + from litellm.types.videos.utils import encode_character_id_with_provider + + return encode_character_id_with_provider(character_id="char_team", provider="openai", model_id=_TEAM_DEPLOYMENT_ID) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "route, request_data", + [ + pytest.param("/v1/batches/{batch_id}", {"batch_id": _TEAM_MANAGED_BATCH_ID}, id="batch-retrieve"), + pytest.param("/v1/batches/{batch_id}/cancel", {"batch_id": _TEAM_MANAGED_BATCH_ID}, id="batch-cancel"), + pytest.param( + "/v1/fine_tuning/jobs/{fine_tuning_job_id}", + { + "fine_tuning_job_id": _encode_managed_id( + f"litellm_proxy;model_id:{_TEAM_DEPLOYMENT_ID};generic_response_id:ftjob-1" + ) + }, + id="fine-tuning-job", + ), + pytest.param( + "/v1/vector_stores/{vector_store_id}", + { + "vector_store_id": _encode_managed_id( + "litellm_proxy:vector_store;unified_id,vs-unified;" + f"target_model_names,{_TEAM_PUBLIC_MODEL_NAME};resource_id,vs_1;model_id,{_TEAM_DEPLOYMENT_ID}" + ) + }, + id="vector-store", + ), + pytest.param( + "/v1/videos/characters/{character_id}", + {"character_id": _team_scoped_character_id()}, + id="character", + ), + ], +) +async def test_team_scoped_managed_ids_authorize_only_the_owning_team_by_public_name( + route: str, request_data: dict[str, str] +) -> None: + from litellm.proxy._types import ProxyException + from litellm.proxy.auth.auth_checks import can_key_call_resolved_model + + llm_router: Final = _team_scoped_router() + + async def check(team_id: str) -> None: + await can_key_call_resolved_model( + model=get_model_from_request( + request_data=request_data, route=route, llm_router=llm_router, team_id=team_id + ), + llm_model_list=None, + valid_token=UserAPIKeyAuth(api_key="sk-team", team_id=team_id, team_models=[_TEAM_PUBLIC_MODEL_NAME]), + llm_router=llm_router, + ) + + await check("team-a") + with pytest.raises(ProxyException) as exc_info: + await check("team-b") + assert exc_info.value.code == "403" + + +@pytest.mark.asyncio +async def test_legacy_team_alias_authorizes_a_team_batch_only_for_keys_granted_the_alias() -> None: + from litellm.proxy._types import ProxyException + from litellm.proxy.auth.auth_checks import can_key_call_resolved_model + + llm_router: Final = _team_scoped_router(public_model_name=None) + + async def check(key_models: list[str]) -> None: + await can_key_call_resolved_model( + model=get_model_from_request( + request_data={"batch_id": _TEAM_MANAGED_BATCH_ID}, + route="/v1/batches/{batch_id}", + llm_router=llm_router, + team_id="team-a", + ), + llm_model_list=None, + valid_token=UserAPIKeyAuth( + api_key="sk-team", + team_id="team-a", + models=key_models, + team_models=["legacy-team-model"], + team_model_aliases={"legacy-team-model": _TEAM_INTERNAL_MODEL_NAME}, + ), + llm_router=llm_router, + ) + + await check(["legacy-team-model"]) + with pytest.raises(ProxyException) as exc_info: + await check(["some-other-model"]) + assert exc_info.value.code == "403" + + def test_get_model_from_request_resolves_character_id_model_with_router(): from litellm.types.videos.utils import encode_character_id_with_provider diff --git a/tests/unit/proxy/batches_endpoints/test_endpoints.py b/tests/unit/proxy/batches_endpoints/test_endpoints.py index 3bf51f02d34..7efa7b22252 100644 --- a/tests/unit/proxy/batches_endpoints/test_endpoints.py +++ b/tests/unit/proxy/batches_endpoints/test_endpoints.py @@ -35,7 +35,7 @@ import logging from contextlib import ExitStack from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from typing import Any, Dict, Optional +from typing import Any, Dict, Final, Optional from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -3259,6 +3259,129 @@ async def test_cancel__unified_batch_id_rejects_key_without_model_grant(cancel_h cancel_harness.router_acancel.assert_not_called() +TEAM_INTERNAL_MODEL_NAME: Final = "model_name_team-a_00000000-0000-0000-0000-000000000000" +TEAM_PUBLIC_MODEL_NAME: Final = "team-gpt4o-mini" + + +def _team_scoped_deployment() -> Deployment: + return Deployment( + model_name=TEAM_INTERNAL_MODEL_NAME, + litellm_params=LiteLLM_Params(model="gpt-4o-mini"), + model_info=ModelInfo(id="dep-team-a-1", team_id="team-a", team_public_model_name=TEAM_PUBLIC_MODEL_NAME), + ) + + +def _route_unified_batch_to_team_deployment( + harness: RetrieveHarness | CancelHarness, deployment: Deployment | None = None +) -> None: + target: Final = deployment or _team_scoped_deployment() + harness.router.resolve_model_name_from_model_id = MagicMock(return_value=target.model_name) + harness.router.get_deployment_by_model_group_name = MagicMock(return_value=target) + + +@pytest.mark.asyncio +async def test_retrieve__team_scoped_deployment_authorizes_owning_team_by_public_name( + retrieve_harness: RetrieveHarness, +) -> None: + _route_unified_batch_to_team_deployment(retrieve_harness) + + resp: Final = await call_retrieve( + retrieve_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=_key_restricted_to(TEAM_PUBLIC_MODEL_NAME) + ) + + assert retrieve_harness.router_aretrieve.call_count == 1 + assert resp._hidden_params["model_id"] == "gpt-4o-mini" + + +@pytest.mark.asyncio +async def test_cancel__team_scoped_deployment_authorizes_owning_team_by_public_name( + cancel_harness: CancelHarness, +) -> None: + _route_unified_batch_to_team_deployment(cancel_harness) + + await call_cancel(cancel_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=_key_restricted_to(TEAM_PUBLIC_MODEL_NAME)) + + assert cancel_harness.router_acancel.call_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "deployment, user", + [ + pytest.param( + _team_scoped_deployment(), _key_restricted_to("some-other-model"), id="owning-team-key-without-the-model" + ), + pytest.param( + _team_scoped_deployment(), + UserAPIKeyAuth( + api_key="sk-team-b", + team_id="team-b", + team_models=[TEAM_PUBLIC_MODEL_NAME], + models=[TEAM_PUBLIC_MODEL_NAME], + ), + id="other-team-key-granted-the-same-public-name", + ), + pytest.param( + Deployment( + model_name="model_name_team-a_ml_00000000-0000-0000-0000-000000000000", + litellm_params=LiteLLM_Params(model="gpt-4o-mini"), + model_info=ModelInfo( + id="dep-team-a-ml-1", team_id="team-a_ml", team_public_model_name=TEAM_PUBLIC_MODEL_NAME + ), + ), + _key_restricted_to(TEAM_PUBLIC_MODEL_NAME), + id="team-whose-id-prefixes-the-owning-team-id", + ), + ], +) +async def test_retrieve__team_scoped_deployment_rejects_keys_outside_its_grant( + retrieve_harness: RetrieveHarness, deployment: Deployment, user: UserAPIKeyAuth +) -> None: + _route_unified_batch_to_team_deployment(retrieve_harness, deployment) + + with pytest.raises(ProxyException) as exc_info: + await call_retrieve(retrieve_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=user) + + assert exc_info.value.code == "403" + retrieve_harness.router_aretrieve.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "deployment, user", + [ + pytest.param( + _team_scoped_deployment(), + _key_restricted_to(TEAM_INTERNAL_MODEL_NAME), + id="owning-team-key-granted-the-internal-name", + ), + pytest.param( + Deployment( + model_name=TEAM_INTERNAL_MODEL_NAME, + litellm_params=LiteLLM_Params(model="gpt-4o-mini"), + model_info=ModelInfo(id="dep-team-a-1", team_id="team-a"), + ), + UserAPIKeyAuth( + api_key="sk-legacy", + team_id="team-a", + team_models=["legacy-team-model"], + models=["legacy-team-model"], + team_model_aliases={"legacy-team-model": TEAM_INTERNAL_MODEL_NAME}, + ), + id="legacy-team-model-reached-through-a-team-alias", + ), + ], +) +async def test_retrieve__team_scoped_deployment_authorizes_internal_name_and_legacy_alias_grants( + retrieve_harness: RetrieveHarness, deployment: Deployment, user: UserAPIKeyAuth +) -> None: + _route_unified_batch_to_team_deployment(retrieve_harness, deployment) + + await call_retrieve(retrieve_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=user) + + assert retrieve_harness.router_aretrieve.call_count == 1 + + @pytest.mark.asyncio async def test_cancel__executed_batch_rejects_key_without_model_grant(cancel_harness, executed_runner): runner, factory = executed_runner From de32a2a06e024866048fc2c08caef852d8de9567 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 19 Oct 2026 15:42:02 -0700 Subject: [PATCH 08/12] Merge pull request #47389 from BerriAI/litellm_backport_count_tokens_stable_1_104_x fix(count_tokens): backport October 19 fixes to stable/1.104.x (cherry picked from commit 5e4d160423016d171b8a3c5a8ff792e8c1db7c86) --- ci_cd/generate_model_prices_schema.py | 20 + .../gitlab/gitlab_prompt_manager.py | 2 +- litellm/litellm_core_utils/core_helpers.py | 26 +- .../prompt_templates/factory.py | 22 + .../anthropic/count_tokens/token_counter.py | 33 +- .../pass_through/adapters/transformation.py | 2 +- .../pass_through/convertible_blocks.py | 106 + .../anthropic/pass_through/messages/utils.py | 13 + .../responses_adapters/transformation.py | 17 +- .../anthropic/count_tokens/handler.py | 11 +- .../anthropic/count_tokens/token_counter.py | 39 +- .../anthropic/count_tokens/transformation.py | 9 +- litellm/llms/base_llm/base_utils.py | 6 +- .../count_tokens/bedrock_token_counter.py | 6 +- litellm/llms/bedrock/count_tokens/kimi_k3.py | 502 ++++ litellm/llms/gemini/common_utils.py | 44 +- litellm/llms/gemini/count_tokens/handler.py | 44 +- .../openai/responses/count_tokens/handler.py | 7 +- .../responses/count_tokens/token_counter.py | 142 +- .../responses/count_tokens/transformation.py | 88 +- litellm/llms/vertex_ai/common_utils.py | 86 +- .../llms/vertex_ai/count_tokens/context.py | 242 ++ .../llms/vertex_ai/count_tokens/handler.py | 32 +- .../llms/vertex_ai/gemini/transformation.py | 8 +- .../vertex_ai_partner_models/main.py | 6 + ...odel_prices_and_context_window_backup.json | 28 +- .../proxy/anthropic_endpoints/endpoints.py | 51 +- litellm/proxy/common_request_processing.py | 21 + .../common_utils/token_count_settings.py | 164 ++ litellm/proxy/google_endpoints/endpoints.py | 188 +- litellm/proxy/proxy_server.py | 495 +++- .../proxy/response_api_endpoints/endpoints.py | 5 +- litellm/types/llms/vertex_ai.py | 6 +- litellm/types/token_count.py | 16 + litellm/types/utils.py | 13 + litellm/utils.py | 3 + model_prices_and_context_window.json | 28 +- model_prices_and_context_window.schema.json | 12 + tests/integration/_support/provider.py | 49 + tests/integration/_support/vertex.py | 31 + tests/integration/conftest.py | 30 + .../providers/_cache_control_marks_support.py | 479 ++++ .../test_count_tokens_error_shapes_wire.py | 215 ++ .../test_openai_count_tokens_tools_wire.py | 393 ++++ .../test_vertex_gemini_count_tokens_wire.py | 596 +++++ .../test_vertex_partner_count_tokens_wire.py | 858 +++++++ tests/integration/proxy_config.yaml | 7 + .../test_vertex_partner_count_tokens_sdk.py | 90 + .../integration/spend/test_cache_and_quota.py | 8 +- .../litellm_core_utils/test_token_counter.py | 55 +- .../pass_through/test_convertible_blocks.py | 192 ++ .../llms/anthropic/test_count_tokens_oauth.py | 51 + .../unit/llms/azure_ai/anthropic/__init__.py | 0 .../anthropic/count_tokens/__init__.py | 0 .../count_tokens/test_token_counter.py | 217 ++ .../count_tokens/test_bedrock_kimi_k3.py | 524 +++++ .../llms/gemini/test_gemini_common_utils.py | 238 +- .../openai/responses/count_tokens/__init__.py | 0 .../count_tokens/test_token_counter.py | 436 ++++ ...test_openai_count_tokens_transformation.py | 109 + .../openai/test_openai_workload_identity.py | 53 + .../llms/vertex_ai/count_tokens/__init__.py | 0 .../vertex_ai/count_tokens/test_context.py | 507 ++++ .../vertex_ai/test_vertex_ai_common_utils.py | 382 ++- .../anthropic_endpoints/test_endpoints.py | 15 +- .../common_utils/test_token_count_settings.py | 190 ++ .../test_google_api_endpoints.py | 87 +- .../test_model_management_endpoints.py | 48 + .../proxy/proxy_server/test_routes_utils.py | 8 +- tests/unit/proxy/test_proxy_token_counter.py | 2064 ++++++++++++++--- tests/unit/test_count_tokens_public_api.py | 2 +- tests/unit/test_model_prices_schema.py | 46 + .../unit/test_router_model_cost_isolation.py | 34 + tests/unit/test_utils.py | 7 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 12 +- 75 files changed, 9892 insertions(+), 684 deletions(-) create mode 100644 litellm/llms/anthropic/pass_through/convertible_blocks.py create mode 100644 litellm/llms/bedrock/count_tokens/kimi_k3.py create mode 100644 litellm/llms/vertex_ai/count_tokens/context.py create mode 100644 litellm/proxy/common_utils/token_count_settings.py create mode 100644 litellm/types/token_count.py create mode 100644 tests/integration/_support/provider.py create mode 100644 tests/integration/_support/vertex.py create mode 100644 tests/integration/providers/_cache_control_marks_support.py create mode 100644 tests/integration/providers/test_count_tokens_error_shapes_wire.py create mode 100644 tests/integration/providers/test_openai_count_tokens_tools_wire.py create mode 100644 tests/integration/providers/test_vertex_gemini_count_tokens_wire.py create mode 100644 tests/integration/providers/test_vertex_partner_count_tokens_wire.py create mode 100644 tests/integration/sdk/test_vertex_partner_count_tokens_sdk.py create mode 100644 tests/unit/llms/anthropic/pass_through/test_convertible_blocks.py create mode 100644 tests/unit/llms/azure_ai/anthropic/__init__.py create mode 100644 tests/unit/llms/azure_ai/anthropic/count_tokens/__init__.py create mode 100644 tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py create mode 100644 tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py create mode 100644 tests/unit/llms/openai/responses/count_tokens/__init__.py create mode 100644 tests/unit/llms/openai/responses/count_tokens/test_token_counter.py create mode 100644 tests/unit/llms/vertex_ai/count_tokens/__init__.py create mode 100644 tests/unit/llms/vertex_ai/count_tokens/test_context.py create mode 100644 tests/unit/proxy/common_utils/test_token_count_settings.py diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py index 8eec07dadda..a5a4b7f6f3b 100644 --- a/ci_cd/generate_model_prices_schema.py +++ b/ci_cd/generate_model_prices_schema.py @@ -19,6 +19,7 @@ NONNEG_NUMBER: JsonSchema = {"type": "number", "minimum": 0} NONNEG_INTEGER: JsonSchema = {"type": "integer", "minimum": 0} BOOLEAN: JsonSchema = {"type": "boolean"} STRING: JsonSchema = {"type": "string"} +NON_BLANK_STRING: Final[JsonSchema] = {"type": "string", "pattern": r"\S"} TIME_WINDOW: Final[JsonSchema] = {"type": "string", "pattern": r"^([01]\d|2[0-3]):[0-5]\d-([01]\d|2[0-3]):[0-5]\d$"} WEEKDAY_PATTERN: Final = ( r"(?i)^(mon|monday|tue|tues|tuesday|wed|wednesday|thu|thur|thurs|thursday|fri|friday|sat|saturday|sun|sunday)$" @@ -103,6 +104,16 @@ OBJECT_KEYS: dict[str, JsonSchema] = { "type": "object", "description": "Provider-internal routing hints (e.g. bedrock_invocation_schema).", }, + "token_count_model": { + "type": "object", + "required": ["model"], + "additionalProperties": False, + "properties": { + "model": NON_BLANK_STRING, + "region": NON_BLANK_STRING, + }, + "description": "Model that counts tokens for this one, and the AWS region its count API runs in.", + }, } ARRAY_KEYS: dict[str, JsonSchema] = { @@ -189,6 +200,10 @@ INTEGER_KEYS: dict[str, JsonSchema] = { **NONNEG_INTEGER, "description": "Smallest prefix the provider will actually cache; absent means the provider default applies.", }, + "token_count_framing": { + **NONNEG_INTEGER, + "description": "Tokens added once per request to a local count estimate, for the provider's chat framing.", + }, "tpm": {**NONNEG_INTEGER, "description": "Provider default tokens-per-minute limit."}, "rpm": {**NONNEG_INTEGER, "description": "Provider default requests-per-minute limit."}, } @@ -244,6 +259,11 @@ def cost_schema(key: str) -> JsonSchema: def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]: return { + "token_count_tokenizer": { + "type": "string", + "description": "Tokenizer litellm counts this model's requests with when the provider cannot count them.", + "enum": ["kimi_k3"], + }, "litellm_provider": { "type": "string", "description": "LiteLLM provider slug; one of https://docs.litellm.ai/docs/providers.", diff --git a/litellm/integrations/gitlab/gitlab_prompt_manager.py b/litellm/integrations/gitlab/gitlab_prompt_manager.py index 817d280074f..ee1d4dc7043 100644 --- a/litellm/integrations/gitlab/gitlab_prompt_manager.py +++ b/litellm/integrations/gitlab/gitlab_prompt_manager.py @@ -215,7 +215,7 @@ class GitLabTemplateManager: elif value.replace(".", "").isdigit(): try: result[key] = float(value) - except Exception: + except ValueError: result[key] = value else: result[key] = value.strip("\"'") diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 39e95fbf687..aad4e80b287 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -3,9 +3,9 @@ import copy import logging import re -from collections.abc import Collection, Iterable, Mapping +from collections.abc import Collection, Iterable, Iterator, Mapping from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Protocol +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, runtime_checkable import httpx from pydantic import TypeAdapter, ValidationError @@ -51,6 +51,28 @@ def safe_divide_seconds(seconds: float, denominator: float, default: float | Non _DROP_PARAMS_BOOL: Final = TypeAdapter(bool) +def as_str_mapping(value: object) -> Mapping[str, object] | None: + return value if isinstance(value, Mapping) else None + + +@runtime_checkable +class _SupportsIter(Protocol): + def __iter__(self) -> Iterator[object]: ... + + +@runtime_checkable +class _SupportsGetItem(Protocol): + def __getitem__(self, index: int, /) -> object: ... + + +def iter_items(value: object) -> Iterable[object]: + if isinstance(value, _SupportsIter): + return value + if isinstance(value, _SupportsGetItem): + return iter(value) + raise TypeError(f"'{type(value).__name__}' object is not iterable") + + def normalize_drop_params(value: object) -> bool | None: if value is None or isinstance(value, bool): return value diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 84a3da288f3..6567f00f121 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -5666,3 +5666,25 @@ def resolve_structured_messages( if structured: return [msg if isinstance(msg, dict) else msg.model_dump() for msg in structured] return None + + +def is_flat_responses_tool(tool: Mapping[str, object]) -> bool: + tool_type: Final = tool.get("type") + return tool_type in ("function", "custom") and tool_type not in tool and "input_schema" not in tool + + +def chat_completion_tools_for(tool: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + if is_flat_responses_tool(tool): + chat_tools, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( + [dict(tool)] # pyright: ignore[reportArgumentType] # a JSON dict in a flat Responses tool shape + ) + return tuple(chat_tools) + anthropic_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai( + [dict(tool)] # pyright: ignore[reportArgumentType] # a JSON dict in an Anthropic or chat tool shape + ) + return tuple(anthropic_tools) diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index 8e8d10c961b..2154189a9c7 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -5,10 +5,13 @@ Anthropic Token Counter implementation using the CountTokens API. import os from typing import Any, Final +from pydantic import JsonValue + +import litellm from litellm._logging import verbose_logger from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler from litellm.llms.base_llm.base_utils import BaseTokenCounter -from litellm.types.utils import LlmProviders, TokenCountResponse +from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests anthropic_count_tokens_handler: Final = AnthropicCountTokensHandler() @@ -30,8 +33,10 @@ class AnthropicTokenCounter(BaseTokenCounter): contents: list[dict[str, Any]] | None, deployment: dict[str, Any] | None = None, request_model: str = "", - tools: list[dict[str, Any]] | None = None, - system: Any | None = None, + tools: list[dict[str, JsonValue]] | None = None, + system: JsonValue = None, + messages_format: TokenCountMessagesFormat = "anthropic", + output_format: JsonValue = None, ) -> TokenCountResponse | None: """ Count tokens using Anthropic's CountTokens API. @@ -54,14 +59,22 @@ class AnthropicTokenCounter(BaseTokenCounter): deployment = deployment or {} litellm_params: Final = deployment.get("litellm_params", {}) - # Get Anthropic API key from deployment config or environment - api_key = litellm_params.get("api_key") + api_key: Final = ( + litellm_params.get("api_key") or litellm.anthropic_key or litellm.api_key or os.getenv("ANTHROPIC_API_KEY") + ) if not api_key: - api_key = os.getenv("ANTHROPIC_API_KEY") - - if not api_key: - verbose_logger.warning("No Anthropic API key found for token counting") - return None + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="anthropic_api", + error=True, + error_message=( + f"No Anthropic API key is set for model {request_model or model_to_use}, so tokens can't be " + "counted. Set api_key on the deployment or ANTHROPIC_API_KEY in the proxy environment." + ), + status_code=401, + ) try: result: Final = await anthropic_count_tokens_handler.handle_count_tokens_request( diff --git a/litellm/llms/anthropic/pass_through/adapters/transformation.py b/litellm/llms/anthropic/pass_through/adapters/transformation.py index 022bc6337b5..84513933b21 100644 --- a/litellm/llms/anthropic/pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/pass_through/adapters/transformation.py @@ -460,7 +460,7 @@ class LiteLLMAnthropicMessagesAdapter: *, custom_llm_provider: str | None = None, preserve_midturn_system: bool = False, - ) -> list: + ) -> list[AllMessageValues]: new_messages: Final[list[AllMessageValues]] = [] replayable_messages: Final = strip_encrypted_reasoning_blocks_from_anthropic_messages(messages) leading_count: Final = next( diff --git a/litellm/llms/anthropic/pass_through/convertible_blocks.py b/litellm/llms/anthropic/pass_through/convertible_blocks.py new file mode 100644 index 00000000000..543a6efd4c4 --- /dev/null +++ b/litellm/llms/anthropic/pass_through/convertible_blocks.py @@ -0,0 +1,106 @@ +from collections.abc import Iterator, Mapping, Sequence +from itertools import chain +from types import MappingProxyType +from typing import Final + +from pydantic import JsonValue, TypeAdapter, ValidationError +from typing_extensions import assert_never + +from litellm.llms.anthropic.pass_through.messages.utils import normalize_native_messages +from litellm.types.utils import TokenCountMessagesFormat + +_CONVERTIBLE_BLOCK_TYPES: Final = MappingProxyType( + { + "user": frozenset({"text", "image", "document", "tool_result"}), + "assistant": frozenset({"text", "tool_use", "thinking", "redacted_thinking"}), + } +) +_SOURCED_BLOCK_TYPES: Final = frozenset({"image", "document"}) +_TOOL_RESULT_BLOCK_TYPES: Final = frozenset({"text", "image", "document"}) + + +_JSON_MESSAGES: Final = TypeAdapter(tuple[Mapping[str, JsonValue], ...]) + + +def _json_messages(messages: Sequence[Mapping[str, object]]) -> tuple[Mapping[str, JsonValue], ...]: + try: + return _JSON_MESSAGES.validate_python(messages) + except ValidationError: + return () + + +def _source_label(block_type: str, source: JsonValue) -> str | None: + if not isinstance(source, dict): + return f"{block_type} without a source" + match source.get("type"): + case "base64" if source.get("data"): + return None + case "url" if source.get("url"): + return None + case "base64" | "url" as source_type: + return f"{block_type} with an empty {source_type} source" + case str() as source_type: + return f"{block_type} with a {source_type} source" + case _: + return f"{block_type} with an untyped source" + + +def _block_label(block: JsonValue, convertible: frozenset[str]) -> str | None: + if not isinstance(block, dict): + return None + block_type: Final = block.get("type") + if not isinstance(block_type, str): + return "a block without a type" + if block_type not in convertible: + return block_type + return _source_label(block_type, block.get("source")) if block_type in _SOURCED_BLOCK_TYPES else None + + +def _tool_result_labels(block: JsonValue) -> Iterator[str]: + if not isinstance(block, dict) or block.get("type") != "tool_result": + return iter(()) + inner: Final = block.get("content") + if not isinstance(inner, list): + return iter(()) + return ( + f"{label} inside a tool_result" + for item in inner + if (label := _block_label(item, _TOOL_RESULT_BLOCK_TYPES)) is not None + ) + + +def _content_labels(block: JsonValue, convertible: frozenset[str]) -> Iterator[str]: + label: Final = _block_label(block, convertible) + return chain(() if label is None else (label,), _tool_result_labels(block)) + + +def _message_labels(message: Mapping[str, JsonValue]) -> Iterator[str]: + role: Final = message.get("role") + convertible: Final = _CONVERTIBLE_BLOCK_TYPES.get(role) if isinstance(role, str) else None + content: Final = message.get("content") + if convertible is None or not isinstance(content, list): + return iter(()) + return chain.from_iterable(_content_labels(block, convertible) for block in content) + + +def unconvertible_anthropic_blocks(messages: Sequence[Mapping[str, object]]) -> tuple[str, ...]: + normalized: Final = _json_messages(normalize_native_messages(messages)) + return tuple(sorted(set(chain.from_iterable(_message_labels(message) for message in normalized)))) + + +def unconvertible_blocks_error( + messages: Sequence[Mapping[str, object]] | None, messages_format: TokenCountMessagesFormat +) -> str | None: + match messages_format: + case "openai": + return None + case "anthropic": + labels: Final = unconvertible_anthropic_blocks(messages or ()) + return ( + "count_tokens can't convert these content blocks for this model, so the count would leave them out: " + + ", ".join(labels) + if labels + else None + ) + case _: + assert_never(messages_format) diff --git a/litellm/llms/anthropic/pass_through/messages/utils.py b/litellm/llms/anthropic/pass_through/messages/utils.py index fe8ac2cd7a2..32181b221f1 100644 --- a/litellm/llms/anthropic/pass_through/messages/utils.py +++ b/litellm/llms/anthropic/pass_through/messages/utils.py @@ -119,6 +119,19 @@ def anthropic_system_to_openai_message(system: object) -> ChatCompletionSystemMe return ChatCompletionSystemMessage(role="system", content=system) +def normalize_native_messages(messages: Sequence[object]) -> list[dict[str, object]]: + """Apply the block rewrites every /v1/messages request gets before it reaches a provider.""" + from litellm.llms.anthropic.common_utils import ( + flatten_unencrypted_web_search_results_in_anthropic_messages, + sanitize_tool_use_ids_in_anthropic_messages, + strip_empty_content_blocks_from_anthropic_messages, + ) + + return flatten_unencrypted_web_search_results_in_anthropic_messages( + sanitize_tool_use_ids_in_anthropic_messages(strip_empty_content_blocks_from_anthropic_messages(messages)) + ) + + @lru_cache(maxsize=1) def _anthropic_messages_optional_param_keys() -> frozenset[str]: """ diff --git a/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py index a3ebbbcb830..6056d434144 100644 --- a/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py @@ -495,6 +495,19 @@ class LiteLLMAnthropicToResponsesAPIAdapter: result["summary"] = "detailed" return result + @staticmethod + def system_blocks_to_instructions(system: object) -> str: + if not isinstance(system, list): + return "" + return "\n".join( + text + for block in system + if isinstance(block, dict) + and block.get("type") == "text" + and isinstance(text := block.get("text"), str) + and text + ) + def translate_request( self, anthropic_request: AnthropicMessagesRequest, @@ -544,9 +557,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: if isinstance(system, str): responses_kwargs["instructions"] = system elif isinstance(system, list): - responses_kwargs["instructions"] = "\n".join( - filter(None, (b.get("text", "") for b in system if isinstance(b, dict) and b.get("type") == "text")) - ) + responses_kwargs["instructions"] = self.system_blocks_to_instructions(system) # max_tokens -> max_output_tokens max_tokens: Final = anthropic_request.get("max_tokens") diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index 36d5a56db0d..36a227c03ee 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -28,7 +28,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): self, model: str, messages: list[dict[str, Any]], - api_key: str, + api_key: str | None, api_base: str, litellm_params: dict[str, Any] | None = None, timeout: float | httpx.Timeout | None = None, @@ -78,6 +78,15 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): api_key=api_key, litellm_params=litellm_params, ) + if not {"x-api-key", "Authorization"} & headers.keys(): + raise AnthropicError( + status_code=401, + message=( + f"No Azure AI API key or Entra ID credential is set for model {model}, so tokens can't be " + "counted. Set api_key or an Entra ID credential on the deployment, or AZURE_AI_API_KEY or " + "AZURE_API_KEY in the proxy environment." + ), + ) # Use LiteLLM's async httpx client async_client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.AZURE_AI) diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py index c31d2c427bb..5de5240ee47 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -5,17 +5,35 @@ Azure AI Anthropic Token Counter implementation using the CountTokens API. import os from typing import Any, Final +from pydantic import JsonValue + +import litellm from litellm._logging import verbose_logger from litellm.llms.azure_ai.anthropic.count_tokens.handler import ( AzureAIAnthropicCountTokensHandler, ) from litellm.llms.base_llm.base_utils import BaseTokenCounter -from litellm.types.utils import LlmProviders, TokenCountResponse +from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests azure_ai_anthropic_count_tokens_handler: Final = AzureAIAnthropicCountTokensHandler() +def _missing_api_base_response(request_model: str, model_to_use: str) -> TokenCountResponse: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="azure_ai_anthropic_api", + error=True, + error_message=( + f"No Azure AI API base is set for model {request_model or model_to_use}, so tokens can't be counted. " + "Set api_base on the deployment or AZURE_AI_API_BASE or AZURE_API_BASE in the proxy environment." + ), + status_code=500, + ) + + class AzureAIAnthropicTokenCounter(BaseTokenCounter): """Token counter implementation for Azure AI Anthropic provider using the CountTokens API.""" @@ -34,6 +52,8 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): request_model: str = "", tools: list[dict[str, object]] | None = None, system: object | None = None, + messages_format: TokenCountMessagesFormat = "anthropic", + output_format: JsonValue = None, ) -> TokenCountResponse | None: """ Count tokens using Azure AI Anthropic's CountTokens API. @@ -61,18 +81,15 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): if not api_key: api_key = os.getenv("AZURE_AI_API_KEY") - # Get API base from deployment config or environment - api_base = litellm_params.get("api_base") - if not api_base: - api_base = os.getenv("AZURE_AI_API_BASE") - - if not api_key: - verbose_logger.warning("No Azure AI API key found for token counting") - return None + api_base: Final = ( + litellm_params.get("api_base") + or os.getenv("AZURE_AI_API_BASE") + or os.getenv("AZURE_API_BASE") + or litellm.api_base + ) if not api_base: - verbose_logger.warning("No Azure AI API base found for token counting") - return None + return _missing_api_base_response(request_model, model_to_use) try: result: Final = await azure_ai_anthropic_count_tokens_handler.handle_count_tokens_request( diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py index e96a2f20bbb..c73192fe50a 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py @@ -24,7 +24,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): def get_required_headers( self, - api_key: str, + api_key: str | None, litellm_params: dict[str, object] | None = None, ) -> dict[str, str]: """ @@ -45,12 +45,12 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): "Content-Type": "application/json", "anthropic-version": "2023-06-01", "anthropic-beta": ANTHROPIC_TOKEN_COUNTING_BETA_VERSION, - "x-api-key": api_key, # Azure AI Anthropic requires this header + **({"x-api-key": api_key} if api_key else {}), } # Also set up Azure auth headers for flexibility litellm_params = litellm_params or {} - if "api_key" not in litellm_params: + if api_key and "api_key" not in litellm_params: litellm_params["api_key"] = api_key litellm_params_obj: Final = GenericLiteLLMParams.model_validate(litellm_params) @@ -58,8 +58,9 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): # Get Azure auth headers (api-key or Authorization) azure_headers = BaseAzureLLM._base_validate_azure_environment(headers={}, litellm_params=litellm_params_obj) - # Merge Azure auth headers headers.update(azure_headers) + if "api-key" in azure_headers: + headers.setdefault("x-api-key", azure_headers["api-key"]) return headers diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py index a446a721307..80b9c466dfb 100644 --- a/litellm/llms/base_llm/base_utils.py +++ b/litellm/llms/base_llm/base_utils.py @@ -8,11 +8,11 @@ from abc import ABC, abstractmethod from typing import Any, Final from openai.lib import _parsing, _pydantic -from pydantic import BaseModel +from pydantic import BaseModel, JsonValue from litellm._logging import verbose_logger from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk -from litellm.types.utils import Message, ProviderSpecificModelInfo, TokenCountResponse +from litellm.types.utils import Message, ProviderSpecificModelInfo, TokenCountMessagesFormat, TokenCountResponse class BaseTokenCounter(ABC): @@ -26,6 +26,8 @@ class BaseTokenCounter(ABC): request_model: str = "", tools: list[dict[str, Any]] | None = None, system: Any | None = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, ) -> TokenCountResponse | None: pass diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 2d02b152c61..57bbb5b6cb6 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -5,11 +5,13 @@ Bedrock Token Counter implementation using the CountTokens API. from collections.abc import Mapping, Sequence from typing import Any, Final +from pydantic import JsonValue + from litellm._logging import verbose_logger from litellm.llms.base_llm.base_utils import BaseTokenCounter from litellm.llms.bedrock.common_utils import BedrockError, get_bedrock_base_model from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler -from litellm.types.utils import LlmProviders, TokenCountResponse +from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse class BedrockTokenCounter(BaseTokenCounter): @@ -33,6 +35,8 @@ class BedrockTokenCounter(BaseTokenCounter): request_model: str = "", tools: Sequence[Mapping[str, object]] | None = None, system: object | None = None, + messages_format: TokenCountMessagesFormat = "anthropic", + output_format: JsonValue = None, ) -> TokenCountResponse | None: """ Count tokens using AWS Bedrock's CountTokens API. diff --git a/litellm/llms/bedrock/count_tokens/kimi_k3.py b/litellm/llms/bedrock/count_tokens/kimi_k3.py new file mode 100644 index 00000000000..5e3fd56e4ba --- /dev/null +++ b/litellm/llms/bedrock/count_tokens/kimi_k3.py @@ -0,0 +1,502 @@ +""" +Exact Bedrock token counts for Kimi K3, which has no CountTokens support. + +Renders the Converse body litellm sends through Moonshot's Kimi K3 chat template and counts it with the +official tokenizer, downloaded once from Hugging Face at a pinned revision and checked against its SHA-256. +""" + +import asyncio +import base64 +import copy +import hashlib +import json +import threading +import time +from collections.abc import Callable, Iterable, Mapping, Sequence +from concurrent.futures import Future +from dataclasses import dataclass +from enum import Enum +from itertools import chain, groupby +from pathlib import Path +from typing import Annotated, Final, Literal, TypeAlias + +import regex +import tiktoken +from pydantic import BaseModel, ConfigDict, Field, JsonValue, ValidationError, model_validator +from typing_extensions import Self + +import litellm +from litellm._logging import verbose_logger +from litellm.litellm_core_utils.prompt_templates.factory import ( + chat_completion_tools_for, + is_flat_responses_tool, +) +from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam + +KIMI_K3_REPO: Final = "moonshotai/Kimi-K3" +KIMI_K3_REVISION: Final = "f831ab66814297da540d832a5235f8e904f29d06" +KIMI_K3_RANKS_FILE: Final = "tiktoken.model" +KIMI_K3_RANKS_SHA256: Final = "b6c497a7469b33ced9c38afb1ad6e47f03f5e5dc05f15930799210ec050c5103" +_PATTERN: Final = "|".join( + ( + r"[\p{Han}]+", + r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?", + r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?", + r"\p{N}{1,3}", + r" ?[^\s\p{L}\p{N}]+[\r\n]*", + r"\s*[\r\n]+", + r"\s+(?!\S)", + r"\s+", + ) +) +# Ends chunks only where the pattern above always splits, so their counts sum to the whole text's count while BPE +# stays linear. Only a run of over 1024 characters with no such split, like 2000 spaces, is cut mid-piece. +_CHUNK: Final = regex.compile( + r"(?s).{1,1024}(?:(?<=\S)(?= \S)|(?<=\p{N})(?=\P{N})|(?<=[^\s\p{N}])(?=\p{N})|(?<=\p{L})(?=[^\p{L}\p{M}'])" + r"|(?<=[\r\n])(?=\S)|\Z)|.{1024}\p{N}*" +) + + +@dataclass(frozen=True, slots=True) +class KimiK3TokenizerSource: + ranks_path: Callable[[], str] + sha256: str + + +def _hugging_face_ranks_path() -> str: + import huggingface_hub + + return huggingface_hub.hf_hub_download( + repo_id=KIMI_K3_REPO, + filename=KIMI_K3_RANKS_FILE, + revision=KIMI_K3_REVISION, + local_files_only=litellm.disable_hf_tokenizer_download is True, + ) + + +HUGGING_FACE_SOURCE: Final = KimiK3TokenizerSource(ranks_path=_hugging_face_ranks_path, sha256=KIMI_K3_RANKS_SHA256) + + +def _loaded(source: KimiK3TokenizerSource) -> tiktoken.Encoding | str: + try: + data: Final = Path(source.ranks_path()).read_bytes() + except Exception as e: # noqa: BLE001 # a download or read failure of any kind falls back to the estimate + return type(e).__name__ + if hashlib.sha256(data).hexdigest() != source.sha256: + return "SHA-256 mismatch" + rows: Final = (line.split() for line in data.splitlines() if line) + return tiktoken.Encoding( + name="kimi_k3", + pat_str=_PATTERN, + mergeable_ranks={base64.b64decode(token): int(rank) for token, rank in rows}, + special_tokens={}, + ) + + +def load_kimi_k3_encoding(source: KimiK3TokenizerSource) -> tiktoken.Encoding | None: + loaded: Final = _loaded(source) + if isinstance(loaded, tiktoken.Encoding): + return loaded + verbose_logger.warning( + "Kimi K3 tokenizer unavailable (%s), so Kimi K3 token counts are estimated until a later count loads it. " + "Allow downloads of %s from Hugging Face repo %s at revision %s, or put it in the Hugging Face cache", + loaded, + KIMI_K3_RANKS_FILE, + KIMI_K3_REPO, + KIMI_K3_REVISION, + ) + return None + + +def _started_load(source: KimiK3TokenizerSource) -> "Future[tiktoken.Encoding | None]": + encoding: Final = Future[tiktoken.Encoding | None]() + threading.Thread( + target=lambda: encoding.set_result(load_kimi_k3_encoding(source)), name="litellm-kimi-k3-tokenizer", daemon=True + ).start() + return encoding + + +@dataclass(frozen=True, slots=True) +class _EncodingLoad: + encoding: "Future[tiktoken.Encoding | None]" + started_at: float + is_retry: bool + + def retry_due(self, now: float, retry_after_seconds: float) -> bool: + return self.encoding.done() and self.encoding.result() is None and now - self.started_at >= retry_after_seconds + + +@dataclass(slots=True) +class KimiK3Encodings: + source: KimiK3TokenizerSource + retry_after_seconds: float + clock: Callable[[], float] = time.monotonic + current: _EncodingLoad | None = None + + def _load(self) -> _EncodingLoad: + now: Final = self.clock() + current: Final = self.current + if current is not None and not current.retry_due(now, self.retry_after_seconds): + return current + load: Final = _EncodingLoad(encoding=_started_load(self.source), started_at=now, is_retry=current is not None) + self.current = load + return load + + async def encoding(self, wait_seconds: float) -> tiktoken.Encoding | None: + load: Final = self._load() + remaining: Final = 0 if load.is_retry else load.started_at + wait_seconds - self.clock() + if load.encoding.done() or remaining <= 0: + return load.encoding.result() if load.encoding.done() else None + await asyncio.wait((asyncio.wrap_future(load.encoding),), timeout=remaining) + if load.encoding.done(): + return load.encoding.result() + verbose_logger.warning( + "Kimi K3 tokenizer still loading after %s s, so Kimi K3 token counts are estimated until it loads. " + "Check that %s can be downloaded from Hugging Face repo %s at revision %s, or put it in the Hugging Face cache", + wait_seconds, + KIMI_K3_RANKS_FILE, + KIMI_K3_REPO, + KIMI_K3_REVISION, + ) + return None + + +class _Converse(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + +class _Text(_Converse): + text: str + + +class _CachePoint(_Converse): + cache_point: JsonValue = Field(alias="cachePoint") + + +class _ToolUse(_Converse): + tool_use_id: str = Field(alias="toolUseId") + name: str + arguments: Mapping[str, JsonValue] = Field(alias="input") + + +class _ToolUseBlock(_Converse): + tool_use: _ToolUse = Field(alias="toolUse") + + +class _ToolResult(_Converse): + tool_use_id: str = Field(alias="toolUseId") + content: tuple[_Text, ...] + status: str | None = None + + +class _ToolResultBlock(_Converse): + tool_result: _ToolResult = Field(alias="toolResult") + + +class _ReasoningText(_Converse): + text: str + signature: str | None = None + + +class _Reasoning(_Converse): + reasoning_text: _ReasoningText = Field(alias="reasoningText") + + +class _ReasoningBlock(_Converse): + reasoning_content: _Reasoning = Field(alias="reasoningContent") + + +class _UserMessage(_Converse): + role: Literal["user"] + content: tuple[_Text | _ToolResultBlock | _CachePoint, ...] + + +class _AssistantMessage(_Converse): + role: Literal["assistant"] + content: tuple[_Text | _ToolUseBlock | _ReasoningBlock | _CachePoint, ...] + + def calls(self) -> tuple[_ToolUse, ...]: + return tuple(block.tool_use for block in self.content if isinstance(block, _ToolUseBlock)) + + +class _InputSchema(_Converse): + definition: JsonValue = Field(alias="json") + + +class _ToolSpec(_Converse): + name: str + description: str | None = None + input_schema: _InputSchema = Field(alias="inputSchema") + + def declared(self) -> Mapping[str, JsonValue]: + described: Final[Mapping[str, JsonValue]] = ( + {} if self.description is None else {"description": self.description} + ) + return {"name": self.name, **described, "parameters": self.input_schema.definition} + + +class _ToolSpecBlock(_Converse): + tool_spec: _ToolSpec = Field(alias="toolSpec") + + +class _ToolConfig(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + tools: tuple[_ToolSpecBlock | _CachePoint, ...] + + +_Message: TypeAlias = Annotated[_UserMessage | _AssistantMessage, Field(discriminator="role")] + + +def _result_ids(message: _Message) -> frozenset[str]: + if isinstance(message, _AssistantMessage): + return frozenset() + return frozenset(block.tool_result.tool_use_id for block in message.content if isinstance(block, _ToolResultBlock)) + + +def _call_ids(message: _Message | None) -> frozenset[str]: + if not isinstance(message, _AssistantMessage): + return frozenset() + return frozenset(call.tool_use_id for call in message.calls()) + + +class _ConverseRequest(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + messages: tuple[_Message, ...] + system: tuple[_Text | _CachePoint, ...] = () + tool_config: _ToolConfig | None = Field(default=None, alias="toolConfig") + + @model_validator(mode="after") + def _results_answer_the_previous_calls(self) -> Self: + pairs: Final = zip((None, *self.messages), self.messages, strict=False) + if not all(_result_ids(message) <= _call_ids(previous) for previous, message in pairs): + raise ValueError("a toolResult does not answer a toolUse of the previous assistant message") + return self + + +class _Control(Enum): + OPEN = "<|open|>" + CLOSE = "<|close|>" + SEP = "<|sep|>" + END_OF_MSG = "<|end_of_msg|>" + + +_Segment: TypeAlias = str | _Control +_Segments: TypeAlias = tuple[_Segment, ...] + + +def _open(tag: str, attrs: Iterable[tuple[str, str]] = ()) -> _Segments: + escaped: Final = ( + (f" {key}", '="', value.replace("&", "&").replace('"', """), '"') for key, value in attrs + ) + return (_Control.OPEN, tag, *chain.from_iterable(escaped), _Control.SEP) + + +def _close(tag: str) -> _Segments: + return (_Control.CLOSE, tag, _Control.SEP) + + +def _message(attrs: Iterable[tuple[str, str]], body: Iterable[_Segment]) -> _Segments: + return (*_open("message", attrs), *body, *_close("message"), _Control.END_OF_MSG) + + +def _tool_declare(specs: Sequence[_ToolSpec]) -> _Segments: + tools: Final = json.dumps( + [{"type": "function", "function": spec.declared()} for spec in specs], + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ) + return _message( + (("role", "system"), ("type", "tool-declare")), + (f"# Tools\nHere are the available tools, described in JSONSchema.\n\n```json\n{tools}\n```",), + ) + + +def _text(blocks: Iterable[_Converse]) -> str: + return "".join(block.text for block in blocks if isinstance(block, _Text)) + + +def _json_type(value: JsonValue) -> str: + match value: + case bool(): + return "boolean" + case None: + return "null" + case int() | float(): + return "number" + case str(): + return "string" + case dict(): + return "object" + case _: + return "array" + + +def _argument(key: str, value: JsonValue) -> _Segments: + text: Final = value if isinstance(value, str) else json.dumps(value, ensure_ascii=False) + return (*_open("argument", (("key", key), ("type", _json_type(value)))), text, *_close("argument")) + + +def _call(index: int, call: _ToolUse) -> _Segments: + arguments: Final = chain.from_iterable(_argument(key, value) for key, value in call.arguments.items()) + return (*_open("call", (("tool", call.name), ("index", str(index)))), *arguments, *_close("call")) + + +def _tool_calls(calls: Sequence[_ToolUse]) -> _Segments: + if not calls: + return () + rendered: Final = chain.from_iterable(_call(index, call) for index, call in enumerate(calls, start=1)) + return (*_open("tools"), *rendered, *_close("tools")) + + +def _assistant_message(message: _AssistantMessage) -> _Segments: + reasoning: Final = "".join( + block.reasoning_content.reasoning_text.text for block in message.content if isinstance(block, _ReasoningBlock) + ) + return _message( + (("role", "assistant"),), + ( + *_open("think"), + reasoning, + *_close("think"), + *_open("response"), + _text(message.content), + *_close("response"), + *_tool_calls(message.calls()), + ), + ) + + +def _tool_messages(results: Sequence[_ToolResult], previous: _AssistantMessage) -> _Segments: + names: Final = {call.tool_use_id: call.name for call in previous.calls()} + return tuple( + chain.from_iterable( + _message( + (("role", "tool"), ("tool", names[result.tool_use_id]), ("index", str(index))), + (_text(result.content),), + ) + for index, result in enumerate(results, start=1) + ) + ) + + +def _user_run(blocks: Sequence[_Text | _ToolResultBlock], previous: _Message | None) -> _Segments: + results: Final = tuple(block.tool_result for block in blocks if isinstance(block, _ToolResultBlock)) + if results and isinstance(previous, _AssistantMessage): + return _tool_messages(results, previous) + return _message((("role", "user"),), (_text(blocks),)) + + +def _user_messages(message: _UserMessage, previous: _Message | None) -> _Segments: + blocks: Final = tuple(block for block in message.content if not isinstance(block, _CachePoint)) + runs: Final = groupby(blocks, key=lambda block: isinstance(block, _ToolResultBlock)) + return tuple(chain.from_iterable(_user_run(tuple(run), previous) for _, run in runs)) + + +def _turn(previous: _Message | None, message: _Message) -> _Segments: + if isinstance(message, _AssistantMessage): + return _assistant_message(message) + return _user_messages(message, previous) + + +def _rendered(request: _ConverseRequest) -> _Segments: + specs: Final = ( + () + if request.tool_config is None + else tuple(block.tool_spec for block in request.tool_config.tools if isinstance(block, _ToolSpecBlock)) + ) + system: Final = _text(request.system) + turns: Final = chain.from_iterable( + _turn(previous, message) for previous, message in zip((None, *request.messages), request.messages, strict=False) + ) + return ( + *(_tool_declare(specs) if specs else ()), + *(_message((("role", "system"),), (system,)) if system else ()), + *turns, + *_open("message", (("role", "assistant"),)), + *_open("think"), + ) + + +def _text_tokens(encoding: tiktoken.Encoding, text: str) -> int: + spans: Final = (match.span() for match in _CHUNK.finditer(text)) + return sum(len(encoding.encode_ordinary(text[start:end])) for start, end in spans) + + +class _TextPart(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + type: Literal["text"] + + +class _TextOnlyMessage(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + content: str | tuple[_TextPart, ...] | None = None + + +def _text_only(message: AllMessageValues) -> bool: + try: + _TextOnlyMessage.model_validate(message) + except ValidationError: + return False + return True + + +def _as_converse_input(tool: ChatCompletionToolParam) -> tuple[Mapping[str, object], ...]: + return chat_completion_tools_for(tool) if is_flat_responses_tool(tool) else (tool,) + + +def _converse_request( + model: str, messages: Sequence[AllMessageValues], tools: Sequence[ChatCompletionToolParam] | None +) -> _ConverseRequest | None: + if not all(_text_only(message) for message in messages): + return None + try: + chat_tools: Final = tuple( + chat_tool for tool in tools or () for chat_tool in _as_converse_input(tool) + ) # comprehension-ok: flattens each tool's translation + return _ConverseRequest.model_validate( + litellm.AmazonConverseConfig().transform_request( + model=model, + messages=copy.deepcopy(list(messages)), + optional_params={"tools": copy.deepcopy(list(chat_tools))} if chat_tools else {}, + litellm_params={}, + headers={}, + ) + ) + except Exception: # noqa: BLE001 # a body litellm cannot send to Bedrock, or cannot render, falls back to the estimate + return None + + +def kimi_k3_tokens( + encoding: tiktoken.Encoding, + model: str, + messages: Sequence[AllMessageValues], + tools: Sequence[ChatCompletionToolParam] | None, +) -> int | None: + request: Final = _converse_request(model, messages, tools) + if request is None: + return None + return sum( + 1 if isinstance(segment, _Control) else _text_tokens(encoding, segment) for segment in _rendered(request) + ) + + +def kimi_k3_anthropic_tokens(encoding: tiktoken.Encoding, model: str, fields: Mapping[str, object]) -> int | None: + body: Final = {**fields, "model": model, "max_tokens": 1} + try: + translated, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( + copy.deepcopy(body), # pyright: ignore[reportArgumentType] # raw /v1/messages fields + custom_llm_provider="bedrock", + ) + except Exception: # noqa: BLE001 # a request the adapter cannot translate falls back to the estimate + return None + match translated: + case {"tools": tools}: + return kimi_k3_tokens(encoding, model, translated["messages"], tools) + case _: + return kimi_k3_tokens(encoding, model, translated["messages"], None) + + +def kimi_k3_template_tokens(message_count: int, has_tools: bool) -> int: + return 14 + 13 * max(message_count - 1, 0) + (40 if has_tools else 0) diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 78e6e6aaf82..155366f476e 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -6,6 +6,7 @@ from collections.abc import Mapping, Sequence from typing import Any, Final import httpx +from pydantic import JsonValue import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH @@ -13,7 +14,7 @@ from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import TokenCountResponse +from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse GEMINI_IMAGE_ASPECT_RATIOS: Final[dict[str, float]] = { "1:1": 1 / 1, @@ -491,21 +492,54 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): request_model: str = "", tools: list[dict[str, object]] | None = None, system: object | None = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, ) -> TokenCountResponse | None: import copy from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter + from litellm.llms.vertex_ai.count_tokens.context import ( + VertexGeminiCountContextError, + count_context_error_response, + gemini_count_contents, + inference_request_view, + provider_error_response, + unconvertible_messages_error, + vertex_gemini_count_context, + ) + + request_view: Final = inference_request_view(messages or [], system, messages_format) + count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools) + if isinstance(count_context, VertexGeminiCountContextError): + return count_context_error_response(count_context, request_model, model_to_use) + blocks_error: Final = unconvertible_messages_error(request_view.messages, messages_format) + if blocks_error is not None: + return count_context_error_response(blocks_error, request_model, model_to_use) deployment = deployment or {} count_tokens_params_request: Final = copy.deepcopy(deployment.get("litellm_params", {})) + resolved_contents: Final = ( + contents + if contents is not None + else await gemini_count_contents( + request_view.messages, messages_format, model_to_use, "gemini", count_tokens_params_request + ) + ) + if isinstance(resolved_contents, VertexGeminiCountContextError): + return count_context_error_response(resolved_contents, request_model, model_to_use) count_tokens_params: Final = { "model": model_to_use, - "contents": contents, + "contents": resolved_contents, + "system_instruction": count_context.system_instruction, + "tools": count_context.tools, } count_tokens_params_request.update(count_tokens_params) - result: Final = await GoogleAIStudioTokenCounter().acount_tokens( - **count_tokens_params_request, - ) + try: + result: Final = await GoogleAIStudioTokenCounter().acount_tokens( + **count_tokens_params_request, + ) + except litellm.APIError as error: + return provider_error_response(error, request_model, model_to_use) if result is not None: return TokenCountResponse( diff --git a/litellm/llms/gemini/count_tokens/handler.py b/litellm/llms/gemini/count_tokens/handler.py index c2f0ef473ae..e11a401a5f9 100644 --- a/litellm/llms/gemini/count_tokens/handler.py +++ b/litellm/llms/gemini/count_tokens/handler.py @@ -1,9 +1,10 @@ +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx import litellm -from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client from litellm.types.utils import LlmProviders if TYPE_CHECKING: @@ -28,12 +29,12 @@ class GoogleAIStudioTokenCounter: """ import copy - from google.genai.types import FunctionResponse - # Handle None or empty contents if not contents: return contents + from google.genai.types import FunctionResponse + cleaned_contents: Final = copy.deepcopy(contents) for content in cleaned_contents: @@ -54,6 +55,28 @@ class GoogleAIStudioTokenCounter: base_url: Final = api_base or "https://generativelanguage.googleapis.com" return f"{base_url}/v1beta/models/{model}:countTokens" + def __init__(self, client: AsyncHTTPHandler | None = None) -> None: + super().__init__() + self._client: Final = client + + def _build_count_tokens_request_body( + self, + model: str, + contents: object, + system_instruction: Mapping[str, object] | None, + tools: Sequence[Mapping[str, object]] | None, + ) -> Mapping[str, object]: + if system_instruction is None and not tools: + return {"contents": contents} + + generate_content_request: Final = { + "model": model if model.startswith("models/") else f"models/{model}", + "contents": contents, + **({"systemInstruction": dict(system_instruction)} if system_instruction is not None else {}), + **({"tools": [dict(tool) for tool in tools]} if tools else {}), + } + return {"generateContentRequest": generate_content_request} + async def validate_environment( self, api_base: str | None = None, @@ -84,6 +107,8 @@ class GoogleAIStudioTokenCounter: api_key: str | None = None, api_base: str | None = None, timeout: float | httpx.Timeout | None = None, + system_instruction: Mapping[str, object] | None = None, + tools: Sequence[Mapping[str, object]] | None = None, **kwargs: object, ) -> dict[str, Any]: """ @@ -130,14 +155,19 @@ class GoogleAIStudioTokenCounter: # Prepare request body - clean up contents to remove unsupported fields cleaned_contents: Final = self._clean_contents_for_gemini_api(contents) - request_body: Final = {"contents": cleaned_contents} + request_body: Final = self._build_count_tokens_request_body( + model=model, + contents=cleaned_contents, + system_instruction=system_instruction, + tools=tools, + ) - async_httpx_client: Final = get_async_httpx_client( - llm_provider=LlmProviders.GEMINI, + async_httpx_client: Final = ( + self._client if self._client is not None else get_async_httpx_client(llm_provider=LlmProviders.GEMINI) ) try: - response: Final = await async_httpx_client.post(url=url, headers=headers, json=request_body) + response: Final = await async_httpx_client.post(url=url, headers=headers, json=dict(request_body)) # Check for HTTP errors response.raise_for_status() diff --git a/litellm/llms/openai/responses/count_tokens/handler.py b/litellm/llms/openai/responses/count_tokens/handler.py index 66782a83a19..0f4e4655575 100644 --- a/litellm/llms/openai/responses/count_tokens/handler.py +++ b/litellm/llms/openai/responses/count_tokens/handler.py @@ -5,6 +5,7 @@ Uses httpx for HTTP requests to OpenAI's /v1/responses/input_tokens endpoint. """ import json +from collections.abc import Mapping, Sequence from typing import Any, Final import httpx @@ -30,9 +31,10 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): api_key: str, api_base: str | None = None, timeout: float | httpx.Timeout | None = None, - tools: list[dict[str, Any]] | None = None, + tools: Sequence[Mapping[str, object]] | None = None, instructions: str | None = None, - ) -> dict[str, Any]: + text: Mapping[str, object] | None = None, + ) -> dict[str, object]: """ Handle a token counting request to OpenAI's Responses API. @@ -52,6 +54,7 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): input=input, tools=tools, instructions=instructions, + text=text, ) endpoint_url: Final = self.get_openai_count_tokens_endpoint(api_base) diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index c05d943b7cf..cd0c3c4ec86 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -2,24 +2,47 @@ OpenAI Token Counter implementation using the Responses API /input_tokens endpoint. """ -import os +from collections.abc import Mapping, Sequence from typing import Any, Final +import openai +from pydantic import JsonValue +from typing_extensions import assert_never + from litellm._logging import verbose_logger +from litellm.litellm_core_utils.asyncify import asyncify +from litellm.llms.anthropic.pass_through.convertible_blocks import unconvertible_blocks_error +from litellm.llms.anthropic.pass_through.responses_adapters.transformation import ( + LiteLLMAnthropicToResponsesAPIAdapter, +) from litellm.llms.base_llm.base_utils import BaseTokenCounter -from litellm.llms.openai.common_utils import OpenAIError +from litellm.llms.openai.common_utils import OpenAIError, get_openai_credentials from litellm.llms.openai.responses.count_tokens.handler import ( OpenAICountTokensHandler, ) from litellm.llms.openai.responses.count_tokens.transformation import ( OpenAICountTokensConfig, + ResponsesCountRequestParts, ) -from litellm.types.utils import LlmProviders, TokenCountResponse +from litellm.llms.openai.workload_identity import ( + OpenAIWorkloadIdentityConfig, + get_workload_identity_bearer_token, + resolve_openai_workload_identity_config, +) +from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests openai_count_tokens_handler: Final = OpenAICountTokensHandler() +def _workload_identity_config(litellm_params: Mapping[str, object]) -> OpenAIWorkloadIdentityConfig | None: + api_base: Final = litellm_params.get("api_base") + return resolve_openai_workload_identity_config( + api_key=None, + api_base=api_base if isinstance(api_base, str) else None, + ) + + class OpenAITokenCounter(BaseTokenCounter): """Token counter implementation for OpenAI provider using the Responses API.""" @@ -29,6 +52,37 @@ class OpenAITokenCounter(BaseTokenCounter): ) -> bool: return custom_llm_provider == LlmProviders.OPENAI.value + @staticmethod + def _system_instructions(system: object) -> str | None: + if system is None or isinstance(system, str): + return system + return LiteLLMAnthropicToResponsesAPIAdapter.system_blocks_to_instructions(system) or None + + @staticmethod + def _responses_request_parts( + model_to_use: str, + messages: Sequence[Mapping[str, object]], + tools: Sequence[Mapping[str, object]] | None, + system: object, + messages_format: TokenCountMessagesFormat, + output_format: JsonValue, + ) -> ResponsesCountRequestParts: + match messages_format: + case "anthropic": + return OpenAICountTokensConfig.anthropic_messages_to_responses_input( + model=model_to_use, messages=messages, system=system, tools=tools, output_format=output_format + ) + case "openai": + input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input( + [dict(message) for message in messages] + ) + fallback_instructions: Final = OpenAITokenCounter._system_instructions(system) + return OpenAICountTokensConfig.count_request_parts( + input_items, instructions if instructions is not None else fallback_instructions, tools + ) + case _: + assert_never(messages_format) + async def count_tokens( self, model_to_use: str, @@ -38,6 +92,8 @@ class OpenAITokenCounter(BaseTokenCounter): request_model: str = "", tools: list[dict[str, object]] | None = None, system: object = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, ) -> TokenCountResponse | None: """ Count tokens using OpenAI's Responses API /input_tokens endpoint. @@ -45,39 +101,76 @@ class OpenAITokenCounter(BaseTokenCounter): if not messages: return None + malformed_error: Final = ( + OpenAICountTokensConfig.malformed_anthropic_messages_error(messages) + if messages_format == "anthropic" + else None + ) + request_error: Final = malformed_error or unconvertible_blocks_error(messages, messages_format) + if request_error is not None: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="openai_api", + error=True, + error_message=request_error, + status_code=400, + rejected_before_provider_call=True, + ) + deployment = deployment or {} litellm_params: Final = deployment.get("litellm_params", {}) - # Get OpenAI API key from deployment config or environment - api_key = litellm_params.get("api_key") - if not api_key: - api_key = os.getenv("OPENAI_API_KEY") + credentials: Final = get_openai_credentials( + api_base=litellm_params.get("api_base"), api_key=litellm_params.get("api_key") + ) + credential: Final = credentials.api_key or _workload_identity_config(litellm_params) - if not api_key: - verbose_logger.warning("No OpenAI API key found for token counting") - return None + if credential is None: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="openai_api", + error=True, + error_message=( + f"No OpenAI API key is set for model {request_model or model_to_use}, so tokens can't be " + "counted. Set api_key or workload identity on the deployment, or OPENAI_API_KEY in the proxy " + "environment." + ), + status_code=401, + ) - api_base: Final = litellm_params.get("api_base") + api_base: Final = credentials.api_base - # Convert chat messages to Responses API input format - input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages) - - # Use system param if instructions not extracted from messages - if instructions is None and system is not None: - instructions = system if isinstance(system, str) else str(system) + input_items, instructions, counted_tools, text = self._responses_request_parts( + model_to_use=model_to_use, + messages=messages, + tools=tools, + system=system, + messages_format=messages_format, + output_format=output_format, + ) # If no input items were produced (e.g., system-only messages), fall back to local counting if not input_items: return None try: + api_key: Final = ( + credential + if isinstance(credential, str) + else await asyncify(get_workload_identity_bearer_token)(credential) + ) result: Final = await openai_count_tokens_handler.handle_count_tokens_request( model=model_to_use, - input=input_items if input_items is not None else [], + input=input_items, api_key=api_key, api_base=api_base, - tools=tools, + tools=counted_tools, instructions=instructions, + text=text, ) if result is not None: @@ -99,6 +192,17 @@ class OpenAITokenCounter(BaseTokenCounter): error_message=e.message, status_code=e.status_code, ) + except openai.AuthenticationError as e: + verbose_logger.warning("OpenAI workload identity token exchange failed: %s", e.message) + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="openai_api", + error=True, + error_message=e.message, + status_code=401, + ) except Exception as e: verbose_logger.warning("Error calling OpenAI CountTokens API: %s", e) return TokenCountResponse( diff --git a/litellm/llms/openai/responses/count_tokens/transformation.py b/litellm/llms/openai/responses/count_tokens/transformation.py index 9f596505f91..eda9777801c 100644 --- a/litellm/llms/openai/responses/count_tokens/transformation.py +++ b/litellm/llms/openai/responses/count_tokens/transformation.py @@ -5,10 +5,18 @@ This module handles the transformation of requests to OpenAI's /v1/responses/inp """ from collections.abc import Mapping, Sequence -from typing import Any, Final, Literal +from typing import Any, Final, Literal, TypeAlias +from pydantic import JsonValue, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict +from litellm.llms.anthropic.pass_through.messages.utils import normalize_native_messages +from litellm.llms.anthropic.pass_through.responses_adapters.transformation import ( + LiteLLMAnthropicToResponsesAPIAdapter, +) +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.llms.anthropic import AnthropicMessagesRequest, AnthropicOutputSchema + class ResponsesInputTextPart(TypedDict): type: ReadOnly[Literal["input_text"]] @@ -31,6 +39,30 @@ ResponsesInputPart = ResponsesInputTextPart | ResponsesInputImagePart | Response ResponsesContentRole = Literal["user", "assistant"] +ResponsesCountRequestParts: TypeAlias = tuple[ + tuple[Mapping[str, object], ...], + str | None, + tuple[Mapping[str, object], ...] | None, + Mapping[str, object] | None, +] +_RESPONSES_COUNT_REQUEST_PARTS: Final = TypeAdapter(ResponsesCountRequestParts) +_TRANSLATED_REQUEST: Final = TypeAdapter(Mapping[str, object]) + + +class _AnthropicCountMessage(TypedDict): + role: ReadOnly[str] + + +_ANTHROPIC_COUNT_MESSAGES: Final = TypeAdapter(tuple[_AnthropicCountMessage, ...]) +_ANTHROPIC_OUTPUT_SCHEMA: Final = TypeAdapter(AnthropicOutputSchema) + + +def _anthropic_output_schema(output_format: JsonValue) -> AnthropicOutputSchema | None: + try: + return _ANTHROPIC_OUTPUT_SCHEMA.validate_python(output_format) + except ValidationError: + return None + def _chat_image_block_to_responses_part(image_url: object) -> ResponsesInputImagePart | None: url: Final = image_url.get("url") if isinstance(image_url, Mapping) else image_url @@ -118,8 +150,9 @@ class OpenAICountTokensConfig: self, model: str, input: str | Sequence[object], - tools: list[dict[str, Any]] | None = None, + tools: Sequence[Mapping[str, object]] | None = None, instructions: str | None = None, + text: Mapping[str, object] | None = None, ) -> dict[str, object]: """ Transform request to OpenAI Responses API token counting format. @@ -135,7 +168,10 @@ class OpenAICountTokensConfig: request["instructions"] = instructions if tools is not None: - request["tools"] = self._transform_tools_for_responses_api(tools) + request["tools"] = self._transform_tools_for_responses_api([dict(tool) for tool in tools]) + + if text is not None: + request["text"] = text return request @@ -183,7 +219,7 @@ class OpenAICountTokensConfig: @staticmethod def messages_to_responses_input( messages: list[dict[str, Any]], - ) -> tuple: + ) -> tuple[list[dict[str, object]], str | None]: """ Convert standard chat messages format to OpenAI Responses API input format. @@ -246,3 +282,47 @@ class OpenAICountTokensConfig: instructions: Final = "\n".join(instructions_parts) if instructions_parts else None return input_items, instructions + + @staticmethod + def malformed_anthropic_messages_error(messages: Sequence[Mapping[str, object]]) -> str | None: + try: + _ANTHROPIC_COUNT_MESSAGES.validate_python(messages) + except ValidationError as e: + return "; ".join( + f"messages.{'.'.join(str(part) for part in error['loc'])}: {error['msg']}" for error in e.errors() + ) + return None + + @staticmethod + def anthropic_messages_to_responses_input( + model: str, + messages: Sequence[Mapping[str, object]], + system: object, + tools: Sequence[Mapping[str, object]] | None, + output_format: JsonValue = None, + ) -> ResponsesCountRequestParts: + """Convert with the adapter /v1/messages inference uses, so the count sees the same input it bills.""" + anthropic_request: Final = AnthropicMessagesRequest( + model=model, + messages=normalize_native_messages(messages), + system=system if isinstance(system, (str, list)) else None, + tools=[dict(tool) for tool in tools] if tools else None, + output_format=_anthropic_output_schema(output_format), + ) + translated: Final = _TRANSLATED_REQUEST.validate_python( + LiteLLMAnthropicToResponsesAPIAdapter().translate_request( + anthropic_request, include_encrypted_reasoning=False + ) + ) + return OpenAICountTokensConfig.count_request_parts( + ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input(translated.get("input")), # pyright: ignore[reportPrivateUsage] # the input restore inference runs + translated.get("instructions"), + translated.get("tools"), + translated.get("text"), + ) + + @staticmethod + def count_request_parts( + input_items: object, instructions: object, tools: object, text: object = None + ) -> ResponsesCountRequestParts: + return _RESPONSES_COUNT_REQUEST_PARTS.validate_python((input_items, instructions, tools, text)) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 5b5e1403c58..2c7e07616b5 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -7,7 +7,7 @@ from types import MappingProxyType from typing import Any, Final, Literal, cast, get_type_hints import httpx -from pydantic import TypeAdapter, ValidationError +from pydantic import JsonValue, TypeAdapter, ValidationError from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm @@ -22,7 +22,7 @@ from litellm.types.llms.vertex_ai import ( PartType, Schema, ) -from litellm.types.utils import TokenCountResponse +from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse from litellm.utils import supports_response_schema, supports_system_messages VERTEX_SELF_DEPLOYED_ENDPOINT_UNSUPPORTED_PARAMS: Final = frozenset( @@ -1237,9 +1237,14 @@ class VertexAITokenCounter(BaseTokenCounter): request_model: str = "", tools: list[dict[str, object]] | None = None, system: object | None = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, ) -> TokenCountResponse | None: import copy + from litellm.llms.vertex_ai.vertex_ai_partner_models.main import ( + VertexAIError as PartnerVertexAIError, + ) from litellm.llms.vertex_ai.vertex_ai_partner_models.main import ( VertexAIPartnerModels, ) @@ -1269,14 +1274,32 @@ class VertexAITokenCounter(BaseTokenCounter): "vertex_ai_credentials" ) - result = await partner_models_handler.count_tokens( - model=model_to_use, - messages=messages or [], - litellm_params=partner_litellm_params, - vertex_project=vertex_project, - vertex_location=vertex_location, - vertex_credentials=vertex_credentials, - ) + try: + result = await partner_models_handler.count_tokens( + model=model_to_use, + messages=messages or [], + litellm_params=partner_litellm_params, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_credentials=vertex_credentials, + system=system, + tools=tools, + ) + except (PartnerVertexAIError, httpx.HTTPStatusError) as e: + status_code: Final = e.response.status_code + error_message: Final = e.message if isinstance(e, PartnerVertexAIError) else e.response.text + verbose_logger.warning( + "Vertex AI partner CountTokens API error: status=%s, message=%s", status_code, error_message + ) + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="vertex_ai_partner_models", + error=True, + error_message=error_message, + status_code=status_code, + ) if result is not None: return TokenCountResponse( @@ -1287,30 +1310,55 @@ class VertexAITokenCounter(BaseTokenCounter): original_response=result, ) else: - from litellm.llms.vertex_ai.count_tokens.handler import VertexAITokenCounter - from litellm.llms.vertex_ai.gemini.transformation import ( - _gemini_convert_messages_with_history, # pyright: ignore[reportPrivateUsage] # shared helper already used by gemini/chat, context_caching, and vertex_and_google_ai_studio_gemini + from litellm.llms.vertex_ai.count_tokens.context import ( + VertexGeminiCountContextError, + count_context_error_response, + gemini_count_contents, + inference_request_view, + provider_error_response, + unconvertible_messages_error, + vertex_gemini_count_context, ) + from litellm.llms.vertex_ai.count_tokens.handler import VertexAITokenCounter + + request_view: Final = inference_request_view(messages or [], system, messages_format) + count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools) + if isinstance(count_context, VertexGeminiCountContextError): + return count_context_error_response(count_context, request_model, model_to_use) + blocks_error: Final = unconvertible_messages_error(request_view.messages, messages_format) + if blocks_error is not None: + return count_context_error_response(blocks_error, request_model, model_to_use) resolved_contents: Final = ( - contents if contents is not None else _gemini_convert_messages_with_history(messages=messages or []) + contents + if contents is not None + else await gemini_count_contents( + request_view.messages, messages_format, model_to_use, "vertex_ai", count_tokens_params_request + ) ) + if isinstance(resolved_contents, VertexGeminiCountContextError): + return count_context_error_response(resolved_contents, request_model, model_to_use) count_tokens_params: Final = { "model": model_to_use, "contents": resolved_contents, + "system_instruction": count_context.system_instruction, + "tools": count_context.tools, } count_tokens_params_request.update(count_tokens_params) - result = await VertexAITokenCounter().acount_tokens( - **count_tokens_params_request, - ) + try: + result = await VertexAITokenCounter().acount_tokens( + **count_tokens_params_request, + ) + except litellm.APIError as error: + return provider_error_response(error, request_model, model_to_use) - if result is not None and "totalTokens" in result: + if "totalTokens" in result: return TokenCountResponse( total_tokens=result["totalTokens"], request_model=request_model, model_used=model_to_use, - tokenizer_type=result.get("tokenizer_used", ""), + tokenizer_type=result.get("tokenizer_used", "gemini"), original_response=result, ) diff --git a/litellm/llms/vertex_ai/count_tokens/context.py b/litellm/llms/vertex_ai/count_tokens/context.py new file mode 100644 index 00000000000..27f2987ea50 --- /dev/null +++ b/litellm/llms/vertex_ai/count_tokens/context.py @@ -0,0 +1,242 @@ +import copy +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal + +from pydantic import JsonValue, TypeAdapter, ValidationError +from typing_extensions import assert_never + +from litellm.litellm_core_utils.prompt_templates.factory import chat_completion_tools_for +from litellm.types.llms.anthropic import AnthropicMessagesTextParam +from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.vertex_ai import ContentType +from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse + +if TYPE_CHECKING: + from litellm.exceptions import APIError + from litellm.llms.anthropic.pass_through.context_management.result import PolyfillResult + + +@dataclass(frozen=True, slots=True) +class VertexGeminiCountContext: + system_instruction: Mapping[str, object] | None + tools: tuple[Mapping[str, object], ...] | None + + +@dataclass(frozen=True, slots=True) +class VertexGeminiCountContextError: + error_message: str + + +_SYSTEM_BLOCKS: Final = TypeAdapter(tuple[AnthropicMessagesTextParam, ...]) +_TOOLS: Final = TypeAdapter(tuple[Mapping[str, JsonValue], ...]) +_VERTEX_TOOL_KEYS: Final = MappingProxyType( + { + "function_declarations": "functionDeclarations", + "url_context": "urlContext", + "code_execution": "codeExecution", + } +) + + +def _system_parts( + system: object | None, +) -> tuple[Mapping[str, JsonValue], ...] | VertexGeminiCountContextError: + if system is None: + return () + if isinstance(system, str): + return ({"text": system},) if system else () + try: + blocks: Final = _SYSTEM_BLOCKS.validate_python(system) + except ValidationError: + return VertexGeminiCountContextError(error_message="Gemini system instructions only support text blocks.") + return tuple({"text": block["text"]} for block in blocks) + + +_INVALID_TOOLS: Final = VertexGeminiCountContextError( + error_message="tools must be a list of objects, and each function tool needs a JSON object schema." +) + + +def _has_object_schema(tool: Mapping[str, JsonValue]) -> bool: + function: Final = tool.get("function") + schema: Final = ( + function.get("parameters") if isinstance(function, dict) else tool.get("input_schema", tool.get("parameters")) + ) + return schema is None or isinstance(schema, dict) + + +def _is_function_tool(tool: Mapping[str, JsonValue]) -> bool: + return isinstance(tool.get("name"), str) or "input_schema" in tool or "function" in tool + + +def _map_function_tools( + tools: object | None, +) -> tuple[Mapping[str, object], ...] | None | VertexGeminiCountContextError: + try: + tool_values: Final = _TOOLS.validate_python(tools or ()) + except ValidationError: + return _INVALID_TOOLS + function_tools: Final = tuple(tool for tool in tool_values if _is_function_tool(tool)) + if not function_tools: + return None + if not all(_has_object_schema(tool) for tool in function_tools): + return _INVALID_TOOLS + + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig + + chat_tools: Final = tuple( + chat_tool for tool in function_tools for chat_tool in chat_completion_tools_for(tool) + ) # comprehension-ok: flattens each tool's translation + try: + mapped_tools: Final = VertexGeminiConfig()._map_function( # pyright: ignore[reportPrivateUsage] # the tool mapping inference uses + value=[dict(tool) for tool in chat_tools], + optional_params={}, + ) + except (AttributeError, TypeError, ValueError): + return _INVALID_TOOLS + return tuple( + MappingProxyType({_VERTEX_TOOL_KEYS.get(key, key): value for key, value in tool.items()}) + for tool in mapped_tools + ) + + +def vertex_gemini_count_context( + system: object | None, + tools: object | None, +) -> VertexGeminiCountContext | VertexGeminiCountContextError: + system_parts: Final = _system_parts(system) + if isinstance(system_parts, VertexGeminiCountContextError): + return system_parts + function_tools: Final = _map_function_tools(tools) + if isinstance(function_tools, VertexGeminiCountContextError): + return function_tools + return VertexGeminiCountContext( + system_instruction=MappingProxyType({"parts": system_parts}) if system_parts else None, + tools=function_tools, + ) + + +def unconvertible_messages_error( + messages: Sequence[Mapping[str, object]] | None, messages_format: TokenCountMessagesFormat +) -> VertexGeminiCountContextError | None: + from litellm.llms.anthropic.pass_through.convertible_blocks import unconvertible_blocks_error + + error_message: Final = unconvertible_blocks_error(messages, messages_format) + return None if error_message is None else VertexGeminiCountContextError(error_message=error_message) + + +def count_context_error_response( + error: VertexGeminiCountContextError, request_model: str, model_used: str +) -> TokenCountResponse: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_used, + tokenizer_type="gemini", + error=True, + error_message=error.error_message, + status_code=400, + rejected_before_provider_call=True, + ) + + +def provider_error_response(error: "APIError", request_model: str, model_used: str) -> TokenCountResponse: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_used, + tokenizer_type="gemini", + error=True, + error_message=error.message, + status_code=error.status_code, + ) + + +@dataclass(frozen=True, slots=True) +class InferenceRequestView: + messages: Sequence[Mapping[str, object]] + system: object | None + + +def _client_compaction(messages: Sequence[Mapping[str, object]], system: object | None) -> "PolyfillResult | None": + from litellm.llms.anthropic.pass_through.context_management.editors.compact import ( + apply_client_compaction_block_history, + ) + + chat_messages: Final = [dict(message) for message in messages] + if system is None or isinstance(system, str): + return apply_client_compaction_block_history(messages=chat_messages, system=system) + try: + blocks: Final = _SYSTEM_BLOCKS.validate_python(system) + except ValidationError: + return None + return apply_client_compaction_block_history(messages=chat_messages, system=[dict(block) for block in blocks]) + + +def inference_request_view( + messages: Sequence[Mapping[str, object]], system: object | None, messages_format: TokenCountMessagesFormat +) -> InferenceRequestView: + """Apply the history rewrites /v1/messages runs before it converts a body for Gemini.""" + from litellm.llms.anthropic.pass_through.messages.utils import normalize_native_messages + + match messages_format: + case "openai": + return InferenceRequestView(messages=messages, system=system) + case "anthropic": + normalized: Final = normalize_native_messages(messages) + compacted: Final = _client_compaction(normalized, system) + if compacted is None: + return InferenceRequestView(messages=normalized, system=system) + return InferenceRequestView(messages=compacted.messages, system=compacted.system) + case _: + assert_never(messages_format) + + +def _chat_messages( + messages: Sequence[Mapping[str, object]], + messages_format: TokenCountMessagesFormat, + model: str, + custom_llm_provider: Literal["vertex_ai", "gemini"], +) -> Sequence[AllMessageValues]: + from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter + + match messages_format: + case "anthropic": + return LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( + messages=list(messages), # pyright: ignore[reportArgumentType] # raw Anthropic body, read the way /v1/messages inference reads it + model=model, + custom_llm_provider=custom_llm_provider, + ) + case "openai": + return copy.deepcopy(list(messages)) # pyright: ignore[reportReturnType] # chat messages from the request, copied because the AI Studio conversion rewrites them in place + case _: + assert_never(messages_format) + + +async def gemini_count_contents( + messages: Sequence[Mapping[str, object]], + messages_format: TokenCountMessagesFormat, + model: str, + custom_llm_provider: Literal["vertex_ai", "gemini"], + litellm_params: Mapping[str, object], +) -> Sequence[ContentType] | VertexGeminiCountContextError: + from litellm.exceptions import BadRequestError + from litellm.litellm_core_utils.asyncify import asyncify + from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media + from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig + from litellm.llms.vertex_ai.gemini.transformation import should_inline_remote_media + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig + + config: Final = GoogleAIStudioGeminiConfig() if custom_llm_provider == "gemini" else VertexGeminiConfig() + try: + chat_messages: Final = await async_inline_remote_media( + list(_chat_messages(messages, messages_format, model, custom_llm_provider)), + should_inline=should_inline_remote_media(custom_llm_provider), + ) + return await asyncify(config._transform_messages)( # pyright: ignore[reportPrivateUsage] # the message transform inference runs + messages=chat_messages, model=model, litellm_params=dict(litellm_params) + ) + except BadRequestError as error: + return VertexGeminiCountContextError(error_message=error.message) diff --git a/litellm/llms/vertex_ai/count_tokens/handler.py b/litellm/llms/vertex_ai/count_tokens/handler.py index 47a08ff054d..03970099f2e 100644 --- a/litellm/llms/vertex_ai/count_tokens/handler.py +++ b/litellm/llms/vertex_ai/count_tokens/handler.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping, Sequence from typing import Any, Final from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter @@ -5,6 +6,19 @@ from litellm.llms.vertex_ai.vertex_llm_base import VertexBase class VertexAITokenCounter(GoogleAIStudioTokenCounter, VertexBase): + def _build_count_tokens_request_body( + self, + model: str, + contents: object, + system_instruction: Mapping[str, object] | None, + tools: Sequence[Mapping[str, object]] | None, + ) -> Mapping[str, object]: + return { + "contents": contents, + **({"systemInstruction": dict(system_instruction)} if system_instruction else {}), + **({"tools": [dict(tool) for tool in tools]} if tools else {}), + } + async def validate_environment( self, api_base: str | None = None, @@ -16,29 +30,29 @@ class VertexAITokenCounter(GoogleAIStudioTokenCounter, VertexBase): """ Returns a Tuple of headers and url for the Vertex AI countTokens endpoint. """ - litellm_params = litellm_params or {} - vertex_credentials: Final = self.get_vertex_ai_credentials(litellm_params=litellm_params) - vertex_project = self.get_vertex_ai_project(litellm_params=litellm_params) - vertex_location: Final = self.get_vertex_ai_location(litellm_params=litellm_params) - _auth_header, vertex_project = await self._ensure_access_token_async( + resolved_litellm_params: Final = litellm_params or {} + vertex_credentials: Final = self.get_vertex_ai_credentials(litellm_params=resolved_litellm_params) + vertex_project: Final = self.get_vertex_ai_project(litellm_params=resolved_litellm_params) + vertex_location: Final = self.get_vertex_ai_location(litellm_params=resolved_litellm_params) + _auth_header, resolved_vertex_project = await self._ensure_access_token_async( credentials=vertex_credentials, project_id=vertex_project, custom_llm_provider="vertex_ai", ) - auth_header, api_base = self._get_token_and_url( + auth_header, url = self._get_token_and_url( model=model, gemini_api_key=None, auth_header=_auth_header, - vertex_project=vertex_project, + vertex_project=resolved_vertex_project, vertex_location=vertex_location, vertex_credentials=vertex_credentials, stream=False, custom_llm_provider="vertex_ai", - api_base=None, + api_base=api_base, mode="count_tokens", ) headers = { "Authorization": f"Bearer {auth_header}", } - return headers, api_base + return headers, url diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index e3cc3bbb2dc..83b796d5df2 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -7,7 +7,7 @@ Why separate file? Make it easy to see how transformation works import json import os import re -from collections.abc import Mapping +from collections.abc import Callable, Mapping from typing import TYPE_CHECKING, Any, Final, Literal, cast from urllib.parse import quote @@ -1327,6 +1327,10 @@ def _vertex_inlines(media: RemoteMedia) -> bool: ) +def should_inline_remote_media(custom_llm_provider: str) -> Callable[[RemoteMedia], bool]: + return _ai_studio_inlines if custom_llm_provider == "gemini" else _vertex_inlines + + async def async_transform_request_body( gemini_api_key: str | None, messages: list[AllMessageValues], @@ -1369,7 +1373,7 @@ async def async_transform_request_body( ) inlined_messages: Final = await async_inline_remote_media( - messages, should_inline=_ai_studio_inlines if custom_llm_provider == "gemini" else _vertex_inlines + messages, should_inline=should_inline_remote_media(custom_llm_provider) ) if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages): diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py index 40503edbb9e..c855a073648 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py @@ -2,6 +2,7 @@ ## API Handler for calling Vertex AI Partner Models from collections.abc import Callable from enum import Enum +from types import MappingProxyType from typing import Final import httpx @@ -263,6 +264,8 @@ class VertexAIPartnerModels(VertexBase): vertex_project=None, vertex_location=None, vertex_credentials=None, + system: object | None = None, + tools: list[dict[str, object]] | None = None, ): """ Count tokens for Vertex AI partner models (Anthropic Claude, Mistral, etc.) @@ -296,6 +299,9 @@ class VertexAIPartnerModels(VertexBase): request_data: Final = { "model": model, "messages": messages, + **MappingProxyType( + {key: value for key, value in (("system", system), ("tools", tools)) if value is not None} + ), } # Prepare litellm_params with credentials diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d42e66605d7..78f17ee3f07 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -19620,7 +19620,13 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": false, + "token_count_framing": 97 + }, + "databricks/system.ai.deepseek-v4-flash-0731": { + "litellm_provider": "databricks", + "mode": "chat", + "token_count_framing": 97 }, "databricks/databricks-deepseek-v4-pro-0813": { "cache_creation_input_token_cost": 1.31999e-06, @@ -20611,7 +20617,13 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 100 + }, + "databricks/system.ai.kimi-k3": { + "litellm_provider": "databricks", + "mode": "chat", + "token_count_framing": 100 }, "databricks/databricks-llama-4-maverick": { "cache_creation_input_token_cost": 5.0001e-07, @@ -76780,7 +76792,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3" }, "global.moonshotai.kimi-k3": { "cache_creation_input_token_cost": 3.75e-06, @@ -76800,7 +76814,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3" }, "us.moonshotai.kimi-k3": { "cache_creation_input_token_cost": 4.125e-06, @@ -76820,7 +76836,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3" }, "openrouter/xiaomi/mimo-v2.6-flash": { "cache_read_input_token_cost": 2.8e-09, diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 2911e7801f7..a5f3fc48811 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -2,10 +2,12 @@ Unified /v1/messages endpoint - (Anthropic Spec) """ +from collections.abc import Mapping from typing import Final from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse +from pydantic import JsonValue import litellm from litellm.anthropic_interface.exceptions import ( @@ -27,6 +29,7 @@ from litellm.proxy.common_request_processing import ( create_response, log_llm_api_exception, proxy_exception_from_http_exception, + proxy_exception_from_route_error, resolve_litellm_call_id, ) from litellm.proxy.common_utils.error_body_call_id import error_body_call_id @@ -38,7 +41,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( openai_error_type, with_litellm_call_id, ) -from litellm.types.utils import TokenCountResponse +from litellm.proxy.common_utils.token_count_settings import count_source_fields router: Final = APIRouter() @@ -286,6 +289,14 @@ async def anthropic_response( ) +def _requested_output_format(data: Mapping[str, JsonValue]) -> JsonValue: + output_format: Final = data.get("output_format") + output_config: Final = data.get("output_config") + if isinstance(output_format, dict) or not isinstance(output_config, dict): + return output_format + return output_config.get("format") + + @router.post( "/v1/messages/count_tokens", tags=["[beta] Anthropic Messages Token Counting"], @@ -315,7 +326,7 @@ async def count_tokens( Returns: {"input_tokens": } """ - from litellm.proxy.proxy_server import token_counter as internal_token_counter + from litellm.proxy.proxy_server import count_request_tokens litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) try: @@ -327,10 +338,10 @@ async def count_tokens( messages: Final = data.get("messages", []) if not model_name: - raise HTTPException(status_code=400, detail={"error": "model parameter is required"}) + raise HTTPException(status_code=400, detail="model parameter is required") if not messages: - raise HTTPException(status_code=400, detail={"error": "messages parameter is required"}) + raise HTTPException(status_code=400, detail="messages parameter is required") # Create TokenCountRequest for the internal endpoint from litellm.proxy._types import TokenCountRequest @@ -342,35 +353,21 @@ async def count_tokens( system=data.get("system"), ) - # Call the internal token counter function with direct request flag set to False - token_response: Final = await internal_token_counter( + token_response: Final = await count_request_tokens( request=token_request, call_endpoint=True, + messages_format="anthropic", + output_format=_requested_output_format(data), ) - _token_response_dict: dict = {} - if isinstance(token_response, TokenCountResponse): - _token_response_dict = token_response.model_dump() - elif isinstance(token_response, dict): - _token_response_dict = token_response - - # Convert the internal response to Anthropic API format - return {"input_tokens": _token_response_dict.get("total_tokens", 0)} - - except HTTPException: - raise + except HTTPException as e: + return _anthropic_error_json_response(proxy_exception_from_http_exception(e, {}), request) except ProxyException as e: - status_code: Final = int(e.code) if e.code and e.code.isdigit() else 500 - detail: Final = AnthropicExceptionMapping.transform_to_anthropic_error( - status_code=status_code, - raw_message=e.message, - ) - raise HTTPException( - status_code=status_code, - detail=detail, - ) + return _anthropic_error_json_response(e, request) except Exception as e: log_llm_api_exception(e, litellm_call_id) - raise HTTPException(status_code=500, detail={"error": f"Internal server error: {e}"}) + return _anthropic_error_json_response(proxy_exception_from_route_error(e), request) + else: + return {"input_tokens": token_response.total_tokens, **count_source_fields(token_response)} @router.post( diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 31e84fd4d9e..3698751f78f 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -614,6 +614,27 @@ def proxy_exception_from_http_exception(exc: HTTPException, headers: dict[str, s ) +def proxy_exception_from_route_error(exc: Exception) -> ProxyException: + if isinstance(exc, ProxyException): + return exc + if isinstance(exc, HTTPException): + return proxy_exception_from_http_exception(exc, {}) + error_status: Final = error_status_code(exc, status.HTTP_500_INTERNAL_SERVER_ERROR) + message: Final = attribute_of(exc, "message", str(exc)) + carried_code: Final = attribute_of(exc, "code") + provider_fields: Final = attribute_of(exc, "provider_specific_fields") + return ProxyException( + message=redact_internal_details_from_client_message( + strip_bug_report_notice(message) if isinstance(message, str) else str(exc) + ), + type=openai_error_type(exc, error_status), + param=openai_error_param(exc), + code=error_status, + openai_code=carried_code if isinstance(carried_code, str) else None, + provider_specific_fields=provider_fields if isinstance(provider_fields, dict) else None, + ) + + def _collect_response_file_search_vector_store_ids(data: Mapping[str, object]) -> set[str]: vector_store_ids: Final[set[str]] = set() tools: Final = data.get("tools") diff --git a/litellm/proxy/common_utils/token_count_settings.py b/litellm/proxy/common_utils/token_count_settings.py new file mode 100644 index 00000000000..6f805616820 --- /dev/null +++ b/litellm/proxy/common_utils/token_count_settings.py @@ -0,0 +1,164 @@ +from collections.abc import Callable, Mapping +from dataclasses import dataclass, replace +from types import MappingProxyType +from typing import Final, Literal + +from pydantic import BaseModel, ConfigDict, ValidationError + +import litellm +from litellm.anthropic_interface.exceptions.exception_mapping_utils import AnthropicExceptionMapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.proxy._types import ProxyException +from litellm.types.token_count import TokenCountSourceFields, TokenCountTokenizer +from litellm.types.utils import TokenCountResponse + +CANNOT_COUNT_MODEL_STATUSES: Final = frozenset({400, 404}) +TOKENIZER_LOAD_WAIT_SECONDS: Final = 5.0 +TOKENIZER_LOAD_REUSE_SECONDS: Final = 60.0 +TOKENIZER_LOAD_MAX_RUNNING: Final = 4 + + +class _UpstreamErrorDetail(BaseModel): + model_config = ConfigDict(frozen=True, extra="allow") + message: str + + +class _UpstreamErrorEnvelope(BaseModel): + model_config = ConfigDict(frozen=True, extra="allow") + type: Literal["error"] + error: _UpstreamErrorDetail + + +def _joined(message: str, hint: str) -> str: + return f"{message.rstrip().rstrip('.')}. {hint}" + + +def _with_hint(message: str, hint: str) -> str: + try: + envelope: Final = _UpstreamErrorEnvelope.model_validate_json(message) + except ValidationError: + return _joined(AnthropicExceptionMapping.extract_error_message(message), hint) + detail: Final = envelope.error.model_copy(update={"message": _joined(envelope.error.message, hint)}) + return envelope.model_copy(update={"error": detail}).model_dump_json() + + +@dataclass(frozen=True, slots=True) +class TokenCountModel: + model: str + region: str | None + + +@dataclass(frozen=True, slots=True) +class TokenCountSettings: + framing: int | None + count_model: TokenCountModel | None + tokenizer: TokenCountTokenizer | None = None + + +@dataclass(frozen=True, slots=True) +class ProviderCountFailure: + status_code: int + message: str + rejected_before_provider_call: bool = False + + @property + def cannot_count_model(self) -> bool: + return not self.rejected_before_provider_call and self.status_code in CANNOT_COUNT_MODEL_STATUSES + + def with_hint(self, hint: str) -> "ProviderCountFailure": + return replace(self, message=_with_hint(self.message, hint)) + + def to_proxy_exception(self) -> ProxyException: + return ProxyException( + message=self.message, + type="token_counting_error", + param="model", + code=self.status_code, + ) + + +def cost_map_entry(model: str, custom_llm_provider: str | None = None) -> Mapping[str, object]: + try: + return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) + except Exception: # noqa: BLE001 # get_model_info raises bare Exception for unmapped models + return MappingProxyType({}) + + +def _framing(value: object) -> int | None: + match value: + case bool(): + return None + case int() if value >= 0: + return value + case float() if value.is_integer() and value >= 0: + return int(value) + case _: + return None + + +def _non_blank(value: object) -> str | None: + return value if isinstance(value, str) and value.strip() else None + + +def _count_model(value: Mapping[str, object]) -> TokenCountModel | None: + model: Final = _non_blank(value.get("model")) + return None if model is None else TokenCountModel(model=model, region=_non_blank(value.get("region"))) + + +def _geography(region: str) -> str: + return region.rsplit("-", 2)[0] + + +def _within_deployment_geography( + count_model: TokenCountModel | None, deployment_region: Callable[[], str] +) -> TokenCountModel | None: + if count_model is None or count_model.region is None or count_model.model.startswith("global."): + return count_model + if _geography(deployment_region()) == _geography(count_model.region): + return count_model + return replace(count_model, region=None) + + +def _resolved_count_model( + configured: object, default: object, deployment_region: Callable[[], str] +) -> TokenCountModel | None: + configured_mapping: Final = as_str_mapping(configured) + if configured_mapping is not None: + return _count_model(configured_mapping) + default_mapping: Final = as_str_mapping(default) + return ( + None + if default_mapping is None + else _within_deployment_geography(_count_model(default_mapping), deployment_region) + ) + + +def _tokenizer(value: object) -> TokenCountTokenizer | None: + match value: + case "kimi_k3": + return "kimi_k3" + case _: + return None + + +def resolve_token_count_settings( + configured: Mapping[str, object], default: Mapping[str, object], deployment_region: Callable[[], str] +) -> TokenCountSettings: + framing: Final = _framing(configured.get("token_count_framing")) + return TokenCountSettings( + framing=framing if framing is not None else _framing(default.get("token_count_framing")), + count_model=_resolved_count_model( + configured.get("token_count_model"), default.get("token_count_model"), deployment_region + ), + tokenizer=_tokenizer(default.get("token_count_tokenizer")), + ) + + +def count_source_fields(response: TokenCountResponse) -> TokenCountSourceFields: + match response.count_source: + case None: + return {} + case "token_count_model": + return {"litellm_count_source": "token_count_model", "litellm_count_model": response.model_used} + case source: + return {"litellm_count_source": source} diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 9d39430cfa9..5fbe780379e 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -1,19 +1,134 @@ +from collections.abc import Mapping +from itertools import chain +from types import MappingProxyType from typing import Final -from fastapi import APIRouter, Depends, Request, Response -from fastapi.responses import ORJSONResponse +from fastapi import APIRouter, Depends, HTTPException, Request, Response +from fastapi.responses import JSONResponse, ORJSONResponse +from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError +from litellm.litellm_core_utils.json_validation_rule import normalize_json_schema_types from litellm.proxy._types import * +from litellm.proxy._types import ProxyException, TokenCountRequest from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth -from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + log_llm_api_exception, + proxy_exception_from_http_exception, + resolve_litellm_call_id, +) from litellm.proxy.common_utils.http_parsing_utils import _read_request_body +from litellm.proxy.common_utils.token_count_settings import count_source_fields from litellm.types.llms.vertex_ai import TokenCountDetailsResponse + +class _GoogleCountTokensTextPart(BaseModel): + model_config = ConfigDict(frozen=True) + text: str + + +class _GoogleCountTokensSystemInstruction(BaseModel): + model_config = ConfigDict(frozen=True) + parts: tuple[_GoogleCountTokensTextPart, ...] + + +class _GoogleCountTokensFunctionDeclaration(BaseModel): + model_config = ConfigDict(frozen=True, populate_by_name=True) + name: str + description: str | None = None + parameters: JsonValue | None = None + parameters_json_schema: JsonValue | None = Field(default=None, alias="parametersJsonSchema") + + +class _GoogleCountTokensTool(BaseModel): + model_config = ConfigDict(frozen=True, populate_by_name=True) + function_declarations: tuple[_GoogleCountTokensFunctionDeclaration, ...] = Field( + default=(), alias="functionDeclarations" + ) + + +class _GoogleCountTokensRequest(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore", populate_by_name=True) + contents: tuple[Mapping[str, JsonValue], ...] = () + system_instruction: _GoogleCountTokensSystemInstruction | None = Field(default=None, alias="systemInstruction") + tools: tuple[_GoogleCountTokensTool, ...] = () + + router: Final = APIRouter( tags=["google genai endpoints"], ) +def _input_schema(declaration: _GoogleCountTokensFunctionDeclaration) -> object: + if declaration.parameters_json_schema is not None: + return declaration.parameters_json_schema + if declaration.parameters is not None: + return normalize_json_schema_types(declaration.parameters) + return {"type": "object", "properties": {}} + + +_GOOGLE_RPC_STATUS: Final = MappingProxyType( + { + 400: "INVALID_ARGUMENT", + 401: "UNAUTHENTICATED", + 403: "PERMISSION_DENIED", + 404: "NOT_FOUND", + 409: "ABORTED", + 429: "RESOURCE_EXHAUSTED", + 499: "CANCELLED", + 501: "UNIMPLEMENTED", + 503: "UNAVAILABLE", + 504: "DEADLINE_EXCEEDED", + } +) + + +def _google_error_response(request: Request, error: ProxyException) -> JSONResponse: + from litellm.proxy.proxy_server import ( + _close_dangling_otel_server_span, # pyright: ignore[reportPrivateUsage] # proxy_server keeps the span-close helper private; error JSONResponses returned by the route must stamp the OTel server span like the global ProxyException handler does + ) + + status_code: Final = int(error.code) if error.code is not None and error.code.isdigit() else 500 + _close_dangling_otel_server_span(request, status_code, exc=error) + status: Final = _GOOGLE_RPC_STATUS.get(status_code, "INTERNAL" if status_code >= 500 else "FAILED_PRECONDITION") + return JSONResponse( + status_code=status_code, + content={"error": {"code": status_code, "message": error.message, "status": status}}, + headers=error.headers, + ) + + +def _google_count_tokens_request(model_name: str, data: object) -> TokenCountRequest: + from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter + + validated_data: Final = TypeAdapter(_GoogleCountTokensRequest).validate_python(data) + contents: Final = [dict(content) for content in validated_data.contents] + messages: Final = ( + GoogleGenAIAdapter().translate_generate_content_to_completion(model_name, contents).get("messages", []) + ) + system: Final = ( + [{"type": "text", "text": part.text} for part in validated_data.system_instruction.parts] + if validated_data.system_instruction is not None + else None + ) + declarations: Final = chain.from_iterable(tool.function_declarations for tool in validated_data.tools) + tools: Final = [ + { + "name": declaration.name, + **({"description": declaration.description} if declaration.description is not None else {}), + "input_schema": _input_schema(declaration), + } + for declaration in declarations + ] + return TokenCountRequest( + model=model_name, + contents=contents, + messages=messages, + system=system, + tools=tools or None, + ) + + @router.post( "/v1beta/models/{model_name:path}:generateContent", dependencies=[Depends(user_api_key_auth)], @@ -165,50 +280,39 @@ async def google_count_tokens(request: Request, model_name: str): } ``` """ - from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.proxy.proxy_server import token_counter as internal_token_counter data: Final = await _read_request_body(request=request) - contents: Final = data.get("contents", []) - # Create TokenCountRequest for the internal endpoint - from litellm.proxy._types import TokenCountRequest + try: + token_request: Final = _google_count_tokens_request(model_name, data) + except ValidationError: + return _google_error_response( + request, + ProxyException( + message="Invalid Google countTokens request body", type="invalid_request_error", param=None, code=400 + ), + ) - # Translate contents to openai format messages using the adapter - messages = GoogleGenAIAdapter().translate_generate_content_to_completion(model_name, contents).get("messages", []) - - token_request: Final = TokenCountRequest( - model=model_name, - contents=contents, - messages=messages, # compatibility when use openai-like endpoint - ) - - # Call the internal token counter function with direct request flag set to False - token_response: Final = await internal_token_counter( - request=token_request, - call_endpoint=True, - ) - if token_response is not None: - # cast the response to the well known format - original_response: Final[dict] = token_response.original_response or {} - if original_response: - return TokenCountDetailsResponse( - totalTokens=original_response.get("totalTokens", 0), - promptTokensDetails=original_response.get("promptTokensDetails", []), - ) - else: - return TokenCountDetailsResponse( - totalTokens=token_response.total_tokens or 0, - promptTokensDetails=[], - ) - - ######################################################### - # Return the response in the well known format - ######################################################### - return TokenCountDetailsResponse( - totalTokens=0, - promptTokensDetails=[], - ) + try: + token_response: Final = await internal_token_counter(request=token_request, call_endpoint=True) + except ProxyException as error: + return _google_error_response(request, error) + except HTTPException as error: + return _google_error_response(request, proxy_exception_from_http_exception(error, {})) + except Exception as error: + log_llm_api_exception(error, resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))) + return _google_error_response( + request, + ProxyException(message="Internal server error", type="internal_server_error", param=None, code=500), + ) + original_response: Final[dict] = token_response.original_response or {} + counted: Final[TokenCountDetailsResponse] = { + "totalTokens": original_response.get("totalTokens") or token_response.total_tokens or 0, + "promptTokensDetails": original_response.get("promptTokensDetails", []), + **count_source_fields(token_response), + } + return counted # ============================================================ diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7f2cd5d10e6..e995ab384db 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -48,6 +48,7 @@ from typing import ( ) import anyio +import anyio.lowlevel import websockets import websockets.exceptions from pydantic import BaseModel, Json, JsonValue, TypeAdapter, ValidationError @@ -88,6 +89,13 @@ from litellm.litellm_core_utils.litellm_logging import ( from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.token_counter import offload_token_count +from litellm.llms.bedrock.count_tokens.kimi_k3 import ( + HUGGING_FACE_SOURCE, + KimiK3Encodings, + kimi_k3_anthropic_tokens, + kimi_k3_template_tokens, + kimi_k3_tokens, +) from litellm.proxy._types import ( UI_TEAM_ID, CallbackDelete, @@ -135,6 +143,15 @@ from litellm.proxy.common_utils.callback_utils import ( strip_callback_config, ) from litellm.proxy.common_utils.realtime_utils import _realtime_request_body +from litellm.proxy.common_utils.token_count_settings import ( + TOKENIZER_LOAD_MAX_RUNNING, + TOKENIZER_LOAD_REUSE_SECONDS, + TOKENIZER_LOAD_WAIT_SECONDS, + ProviderCountFailure, + TokenCountSettings, + cost_map_entry, + resolve_token_count_settings, +) from litellm.proxy.management_helpers.auto_router_availability import AutoRouterCatalogEntry, build_auto_router_catalog from litellm.router_utils.access_windows import access_windows_config_error from litellm.router_utils.add_retry_fallback_headers import ( @@ -164,6 +181,7 @@ from litellm.types.utils import ( ModelResponseStream, StreamingChoices, TextCompletionResponse, + TokenCountMessagesFormat, TokenCountResponse, echoed_cost_map_pricing_fields, is_server_derived_pricing_key, @@ -287,6 +305,7 @@ from litellm.constants import ( CLI_SSO_SESSION_TTL_SECONDS, DAYS_IN_A_MONTH, DEFAULT_HEALTH_CHECK_INTERVAL, + DEFAULT_MAX_LRU_CACHE_SIZE, DEFAULT_MODEL_CREATED_AT_TIME, GLOBAL_PROXY_SPEND_CACHE_KEY, LITELLM_PROXY_ADMIN_NAME, @@ -317,6 +336,7 @@ from litellm.litellm_core_utils.agentic_loop_settings import ( from litellm.litellm_core_utils.audio_utils.utils import resolve_speech_media_type from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, + as_str_mapping, drop_params_flag, get_litellm_metadata_from_kwargs, ) @@ -901,7 +921,13 @@ from litellm.types.secret_managers.main import ( KeyManagementSettings, KeyManagementSystem, ) -from litellm.types.utils import CredentialItem, CustomHuggingfaceTokenizer, RawRequestTypedDict, StandardLoggingPayload +from litellm.types.utils import ( + CredentialItem, + CustomHuggingfaceTokenizer, + RawRequestTypedDict, + SelectTokenizerResponse, + StandardLoggingPayload, +) from litellm.types.utils import ModelInfo as ModelMapInfo from litellm.utils import _add_custom_logger_callback_to_specific_event @@ -13885,52 +13911,367 @@ async def _try_provider_token_count( provider_counter: "BaseTokenCounter", custom_llm_provider: str | None, model_to_use: str, + region: str | None, messages: list | None, contents: list | None, - deployment: dict[str, object] | None, + deployment: Mapping[str, object] | None, request_model: str, tools: list | None = None, system: str | None = None, -) -> Optional["TokenCountResponse"]: - """Attempt provider-specific token counting. Returns result on success, None to fall through to local counting.""" + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, +) -> Union["TokenCountResponse", ProviderCountFailure, None]: + """Attempt provider-specific token counting. Returns None to fall through to local counting.""" if not provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider): return None + deployment_in_region: Final = ( + None + if deployment is None + else dict(deployment) + if region is None + else { + **deployment, + "litellm_params": {**(as_str_mapping(deployment.get("litellm_params")) or {}), "aws_region_name": region}, + } + ) try: result: Final = await provider_counter.count_tokens( model_to_use=model_to_use or "", messages=messages, contents=contents, - deployment=deployment, + deployment=deployment_in_region, request_model=request_model, tools=tools, system=system, + messages_format=messages_format, + output_format=output_format, ) except httpx.HTTPStatusError as e: - error_message: Final = getattr(e, "message", None) or str(e) - status_code: Final = getattr(e, "status_code", None) or e.response.status_code - raise ProxyException( - message=error_message, - type="token_counting_error", - param="model", - code=status_code, + return ProviderCountFailure( + status_code=getattr(e, "status_code", None) or e.response.status_code, + message=getattr(e, "message", None) or str(e), ) if result is not None and result.error is True: - if litellm.disable_token_counter is True: - raise ProxyException( - message=result.error_message or "Token counting failed", - type="token_counting_error", - param="model", - code=result.status_code or 500, - ) - verbose_proxy_logger.warning( - "Provider token counting failed (%s): %s. Falling back to local tokenizer.", - result.status_code, - result.error_message, + return ProviderCountFailure( + status_code=result.status_code or 500, + message=result.error_message or "Token counting failed", + rejected_before_provider_call=result.rejected_before_provider_call, ) - return None return result +async def _provider_token_count( + count_with: Callable[[str, str | None], Awaitable[Union["TokenCountResponse", ProviderCountFailure, None]]], + model_to_use: str, + settings: TokenCountSettings, + local_estimate_allowed: bool, +) -> Union["TokenCountResponse", ProviderCountFailure, None]: + count_model: Final = settings.count_model + counted: Final = await ( + count_with(model_to_use, None) if count_model is None else count_with(count_model.model, count_model.region) + ) + if counted is None: + return None + if isinstance(counted, TokenCountResponse): + if count_model is None: + return counted.model_copy(update={"count_source": "provider"}) + return counted.model_copy(update={"count_source": "token_count_model", "model_used": count_model.model}) + failure: Final = ( + counted + if count_model is None or count_model.region is None + else counted.with_hint( + f"This model counts tokens with {count_model.model} in {count_model.region}. " + "If that region is not enabled for the AWS account, enable it or set token_count_model in model_info" + ) + ) + if local_estimate_allowed: + return failure + raise failure.to_proxy_exception() + + +@dataclass(frozen=True, slots=True) +class _LocalTokenEstimate: + total_tokens: int + tokenizer_type: str + + +def _is_fork_guard_error(error: BaseException) -> bool: + from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted, ProcessReservedForForking + + return isinstance(error, (ForkedAfterNativeRuntimeStarted, ProcessReservedForForking)) + + +def _is_tokenizer_panic(error: BaseException) -> bool: + return type(error).__module__ == "pyo3_runtime" and type(error).__name__ == "PanicException" + + +def _answers_provider_error(error: BaseException) -> bool: + return (isinstance(error, Exception) or _is_tokenizer_panic(error)) and not _is_fork_guard_error(error) + + +def _provider_error_after_failed_estimate( + provider_failure: ProviderCountFailure, error: BaseException +) -> ProxyException: + verbose_proxy_logger.warning( + "Local token estimate failed after the provider's count error %s, answering that error: %s", + provider_failure.status_code, + type(error).__name__, + ) + return provider_failure.to_proxy_exception() + + +async def _estimate_or_provider_error(estimate: Awaitable[int], provider_failure: ProviderCountFailure) -> int: + try: + return await estimate + except BaseException as e: + if not _answers_provider_error(e): + raise + raise _provider_error_after_failed_estimate(provider_failure, e) from e + + +@dataclass(frozen=True, slots=True) +class _CustomTokenizerKey: + identifier: str + revision: str + auth_token: str | None + + +class _TokenizerLoadsBusy(Exception): + pass + + +@dataclass(frozen=True, slots=True) +class _TokenizerLoadOutcome: + tokenizer: SelectTokenizerResponse | BaseException + finished_at: float + + +@dataclass(frozen=True, slots=True) +class _SharedTokenizerLoad: + outcome: asyncio.Task[_TokenizerLoadOutcome] + started_at: float + + def reusable(self, now: float) -> bool: + return not self.outcome.done() or now - self.outcome.result().finished_at < TOKENIZER_LOAD_REUSE_SECONDS + + +async def _loaded_tokenizer( + model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, threads: anyio.CapacityLimiter +) -> SelectTokenizerResponse | BaseException: + try: + return await asyncify(litellm.utils._select_tokenizer, limiter=threads)( + model=model, custom_tokenizer=custom_tokenizer + ) + except asyncio.CancelledError: + raise + except BaseException as e: # noqa: BLE001 # every count sharing this load decides what it answers, pyo3 panics included + return e + + +async def _tokenizer_load_outcome( + model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, threads: anyio.CapacityLimiter +) -> _TokenizerLoadOutcome: + tokenizer: Final = await _loaded_tokenizer(model, custom_tokenizer, threads) + return _TokenizerLoadOutcome(tokenizer=tokenizer, finished_at=asyncio.get_running_loop().time()) + + +@dataclass(slots=True) +class _TokenizerLoads: + threads: anyio.CapacityLimiter + running: int = 0 + + def start( + self, model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, loop: asyncio.AbstractEventLoop + ) -> _SharedTokenizerLoad | None: + if self.running >= self.threads.total_tokens: + return None + self.running += 1 + outcome: Final = loop.create_task(_tokenizer_load_outcome(model, custom_tokenizer, self.threads)) + outcome.add_done_callback(self._finished) + return _SharedTokenizerLoad(outcome=outcome, started_at=loop.time()) + + def _finished(self, _outcome: asyncio.Task[_TokenizerLoadOutcome]) -> None: + self.running -= 1 + + +_TOKENIZER_LOADS: Final = anyio.lowlevel.RunVar[_TokenizerLoads]("litellm_proxy_tokenizer_loads") + + +def _tokenizer_loads_for_this_loop() -> _TokenizerLoads: + existing: Final = _TOKENIZER_LOADS.get(None) + if existing is not None: + return existing + created: Final = _TokenizerLoads(threads=anyio.CapacityLimiter(TOKENIZER_LOAD_MAX_RUNNING)) + _TOKENIZER_LOADS.set(created) + return created + + +@dataclass(slots=True) +class _TokenizerLoadSlot: + model: str + custom_tokenizer: CustomHuggingfaceTokenizer | None + current: _SharedTokenizerLoad | None = None + + def load(self, loads: _TokenizerLoads, loop: asyncio.AbstractEventLoop) -> _SharedTokenizerLoad | None: + current: Final = self.current + if current is not None and current.reusable(loop.time()): + return current + self.current = loads.start(self.model, self.custom_tokenizer, loop) + return self.current + + +@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) +def _tokenizer_load_slot( + model: str, custom_tokenizer: _CustomTokenizerKey | None, loop: asyncio.AbstractEventLoop +) -> _TokenizerLoadSlot: + return _TokenizerLoadSlot( + model=model, + custom_tokenizer=None + if custom_tokenizer is None + else CustomHuggingfaceTokenizer( + identifier=custom_tokenizer.identifier, + revision=custom_tokenizer.revision, + auth_token=custom_tokenizer.auth_token, + ), + ) + + +async def _tokenizer_after_provider_error( + model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None +) -> SelectTokenizerResponse | BaseException: + loop: Final = asyncio.get_running_loop() + try: + load: Final = _tokenizer_load_slot( + model, + None + if custom_tokenizer is None + else _CustomTokenizerKey( + identifier=custom_tokenizer["identifier"], + revision=custom_tokenizer["revision"], + auth_token=custom_tokenizer["auth_token"], + ), + loop, + ).load(_tokenizer_loads_for_this_loop(), loop) + if load is None: + return _TokenizerLoadsBusy() + outcome: Final = await asyncio.wait_for( + asyncio.shield(load.outcome), + timeout=max(load.started_at + TOKENIZER_LOAD_WAIT_SECONDS - loop.time(), 0), + ) + except Exception as e: # noqa: BLE001 # a tokenizer that cannot be read, loaded in time, or loaded at all answers the provider's error + return e + return outcome.tokenizer + + +async def _tokenizer_or_provider_error( + model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, provider_failure: ProviderCountFailure +) -> SelectTokenizerResponse: + tokenizer: Final = await _tokenizer_after_provider_error(model, custom_tokenizer) + if not isinstance(tokenizer, BaseException): + return tokenizer + if not _answers_provider_error(tokenizer): + raise tokenizer.with_traceback(None) + raise _provider_error_after_failed_estimate(provider_failure, tokenizer) from tokenizer + + +async def _local_token_estimate( + model: str, + prompt: str | None, + messages: Sequence[AllMessageValues] | None, + tools: list[ChatCompletionToolParam] | None, + custom_tokenizer: CustomHuggingfaceTokenizer | None, + provider_failure: ProviderCountFailure | None, +) -> _LocalTokenEstimate: + tokenizer: Final = ( + await asyncify(litellm.utils._select_tokenizer)(model=model, custom_tokenizer=custom_tokenizer) + if provider_failure is None + else await _tokenizer_or_provider_error(model, custom_tokenizer, provider_failure) + ) + count: Final = offload_token_count(litellm.token_counter)( + model=model, + text=prompt, + messages=messages, + tools=tools, + custom_tokenizer=tokenizer, + ) + total_tokens: Final = await ( + count if provider_failure is None else _estimate_or_provider_error(count, provider_failure) + ) + return _LocalTokenEstimate(total_tokens=total_tokens, tokenizer_type=str(tokenizer["type"])) + + +_KIMI_K3_ENCODINGS: Final = KimiK3Encodings( + source=HUGGING_FACE_SOURCE, retry_after_seconds=TOKENIZER_LOAD_REUSE_SECONDS +) + + +async def _kimi_k3_count( + request: TokenCountRequest, + messages_format: TokenCountMessagesFormat, + model: str, + messages: Sequence[AllMessageValues], + tools: Sequence[ChatCompletionToolParam] | None, +) -> int | None: + encoding: Final = await _KIMI_K3_ENCODINGS.encoding(TOKENIZER_LOAD_WAIT_SECONDS) + if encoding is None: + return None + match messages_format: + case "anthropic": + return await offload_token_count(kimi_k3_anthropic_tokens)( + encoding, model, request.model_dump(include={"messages", "system", "tools"}, exclude_none=True) + ) + case "openai": + return await offload_token_count(kimi_k3_tokens)(encoding, model, messages, tools) + + +class _TokenCountDeployment(TypedDict, total=False): + litellm_params: ReadOnly[Mapping[str, object]] + model_info: ReadOnly[Mapping[str, object]] + + +_TOKEN_COUNT_DEPLOYMENT: Final = TypeAdapter(_TokenCountDeployment) + + +def _model_group_deployment(router: Router | None, model_group: str) -> Mapping[str, object] | None: + group: Final = () if router is None else router._get_all_deployments(model_name=model_group) + return group[0] if group else None + + +def _non_empty_str(value: object) -> str | None: + return value if isinstance(value, str) and value else None + + +def _deployment_aws_region(deployment: _TokenCountDeployment) -> str: + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + params: Final = deployment.get("litellm_params", {}) + return BaseAWSLLM()._get_aws_region_name( + {"aws_region_name": params.get("aws_region_name")}, model=_non_empty_str(params.get("model")) + ) + + +def _token_count_settings_model(deployment: _TokenCountDeployment, request_model: str) -> tuple[str, str | None]: + model_info: Final = deployment.get("model_info", {}) + params: Final = deployment.get("litellm_params", {}) + base_model: Final = _non_empty_str(model_info.get("base_model")) or _non_empty_str(params.get("base_model")) + if base_model is not None: + return base_model, None + model: Final = _non_empty_str(params.get("model")) + if model is None: + return request_model, None + return model, _non_empty_str(params.get("custom_llm_provider")) + + +_PROVIDERS_COUNTING_WITHOUT_MESSAGES: Final = frozenset( + {litellm.LlmProviders.GEMINI.value, litellm.LlmProviders.VERTEX_AI.value} +) + + +def _drops_params(deployment: Mapping[str, object] | None) -> bool: + litellm_params: Final = as_str_mapping(None if deployment is None else deployment.get("litellm_params")) + return litellm.drop_params is True or (litellm_params is not None and litellm_params.get("drop_params") is True) + + def _system_message(system: object) -> ChatCompletionSystemMessage | None: if not isinstance(system, (str, list)) or not system: return None @@ -13953,8 +14294,15 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) Returns: TokenCountResponse """ - global llm_router + return await count_request_tokens(request=request, call_endpoint=call_endpoint) + +async def count_request_tokens( + request: TokenCountRequest, + call_endpoint: bool = False, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, +) -> TokenCountResponse: prompt: Final = request.prompt messages: Final = request.messages contents: Final = request.contents @@ -13990,8 +14338,16 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) litellm_model_name = litellm_model_name.split("/", 1)[1] model_to_use: str = ( - litellm_model_name or request.model - ) # use litellm model name, if it's not avalable then fallback to request.model + litellm_model_name if isinstance(litellm_model_name, str) and litellm_model_name else request.model + ) + deployment_view: Final = _TOKEN_COUNT_DEPLOYMENT.validate_python( + deployment or _model_group_deployment(llm_router, request.model) or {} + ) + settings: Final = await asyncify(resolve_token_count_settings)( + configured=deployment_view.get("model_info", {}), + default=await asyncify(cost_map_entry)(*_token_count_settings_model(deployment_view, request.model)), + deployment_region=partial(_deployment_aws_region, deployment_view), + ) # Try provider-specific token counting first - only for non-direct requests (from provider endpoints) provider_counter: BaseTokenCounter | None = None @@ -14002,20 +14358,40 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) if _model is not None: model_to_use = _model - if provider_counter is not None: - result: Final = await _try_provider_token_count( - provider_counter=provider_counter, - custom_llm_provider=custom_llm_provider, - model_to_use=model_to_use, - messages=messages, - contents=contents, - deployment=deployment, - request_model=request.model, - tools=tools, - system=system, + conversation_empty: Final = not messages and not contents + if prompt and not conversation_empty and not _drops_params(deployment): + raise ProxyException( + message="prompt cannot be sent with messages or contents. Set drop_params to count them without the prompt", + type="invalid_request_error", + param="prompt", + code=400, ) - if result is not None: - return result + prompt_only: Final = bool(prompt) and conversation_empty + prompt_as_message: Final = prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES + ladder: Final = ( + None + if provider_counter is None + else await _provider_token_count( + count_with=partial( + _try_provider_token_count, + provider_counter, + custom_llm_provider, + messages=[{"role": "user", "content": prompt}] if prompt_as_message else messages, + contents=None if prompt_as_message else contents, + deployment=deployment, + request_model=request.model, + tools=tools, + system=system, + messages_format=messages_format, + output_format=output_format, + ), + model_to_use=model_to_use, + settings=settings, + local_estimate_allowed=litellm.disable_token_counter is not True, + ) + ) + if isinstance(ladder, TokenCountResponse): + return ladder # Check if token counter is disabled before fallback if litellm.disable_token_counter is True: @@ -14033,14 +14409,11 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) CustomHuggingfaceTokenizer | None, model_info.get("custom_tokenizer", None), ) - _tokenizer_used: Final = await asyncify(litellm.utils._select_tokenizer)( - model=model_to_use, custom_tokenizer=custom_tokenizer - ) - - tokenizer_used: Final = str(_tokenizer_used["type"]) system_message: Final = _system_message(system) + prompt_turn: Final = prompt_only and not (system_message is None and not tools) typed_messages: Final = cast( # cast-ok: request messages are raw chat-shaped dicts that token_counter normalizes - Sequence[AllMessageValues] | None, messages + Sequence[AllMessageValues] | None, + [{"role": "user", "content": prompt}] if prompt_turn else (None if prompt_only else messages), ) counted_messages: Final = ( typed_messages if typed_messages is None or system_message is None else (system_message, *typed_messages) @@ -14048,18 +14421,38 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) counted_tools: Final = cast( # cast-ok: raw OpenAI or Anthropic tool dicts, both of which token_counter formats list[ChatCompletionToolParam] | None, tools if counted_messages is not None else None ) - total_tokens: Final = await offload_token_count(litellm.token_counter)( + framing: Final = settings.framing if counted_messages is not None and settings.framing is not None else 0 + kimi_k3_messages: Final = counted_messages if settings.tokenizer == "kimi_k3" else None + exact: Final = ( + await _kimi_k3_count(request, messages_format, model_to_use, kimi_k3_messages, counted_tools) + if kimi_k3_messages is not None and custom_tokenizer is None + else None + ) + if exact is not None: + return TokenCountResponse( + total_tokens=exact + framing, + request_model=request.model, + model_used=model_to_use, + tokenizer_type="huggingface_tokenizer", + count_source="local_estimate", + ) + estimate: Final = await _local_token_estimate( model=model_to_use, - text=prompt, + prompt=prompt if counted_messages is None else None, messages=counted_messages, tools=counted_tools, - custom_tokenizer=_tokenizer_used, + custom_tokenizer=custom_tokenizer, + provider_failure=ladder, + ) + template: Final = ( + 0 if kimi_k3_messages is None else kimi_k3_template_tokens(len(kimi_k3_messages), bool(counted_tools)) ) return TokenCountResponse( - total_tokens=total_tokens, + total_tokens=estimate.total_tokens + template + framing, request_model=request.model, model_used=model_to_use, - tokenizer_type=tokenizer_used, + tokenizer_type=estimate.tokenizer_type, + count_source="local_estimate", ) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 7badf0e79bb..383a24dc6a2 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -34,6 +34,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_set_request_parsed_body, ) +from litellm.proxy.common_utils.token_count_settings import count_source_fields from litellm.proxy.route_llm_request import raise_if_required_body_param_missing from litellm.types.llms.openai import ( REASONING_EFFORT, @@ -41,6 +42,7 @@ from litellm.types.llms.openai import ( ResponsesAPIResponse, ) from litellm.types.responses.main import DeleteResponseResult +from litellm.types.token_count import TokenCountSourceFields from litellm.types.utils import TokenCountResponse if TYPE_CHECKING: @@ -1063,7 +1065,7 @@ class _ResponsesApiErrorBody(TypedDict): error: ReadOnly[_ResponsesApiErrorDetail] -class _ResponsesInputTokensResult(TypedDict): +class _ResponsesInputTokensResult(TokenCountSourceFields): object: ReadOnly[str] input_tokens: ReadOnly[int] @@ -1194,6 +1196,7 @@ async def responses_input_tokens( result: Final[_ResponsesInputTokensResult] = { "object": "response.input_tokens", "input_tokens": token_response.total_tokens, + **count_source_fields(token_response), } return result diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index ce51e46ef15..28fdf3acf67 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -7,6 +7,7 @@ from typing_extensions import ( ) from litellm.types.llms.openai import EmbeddingInput +from litellm.types.token_count import TokenCountSourceFields # Gemini supports nested-list inputs (e.g. [["text", "image"]]) as an explicit # opt-in for combined embeddings — a provider-specific extension of the @@ -308,12 +309,13 @@ class UsageMetadata(TypedDict, total=False): candidatesTokensDetails: list[PromptTokensDetails] # Alternative key name used in some responses -class TokenCountDetailsResponse(TypedDict): +class TokenCountDetailsResponse(TokenCountSourceFields): """ Response structure for token count details with modality breakdown. Example: - {'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}]} + {'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}], + 'litellm_count_source': 'provider'} """ totalTokens: int diff --git a/litellm/types/token_count.py b/litellm/types/token_count.py new file mode 100644 index 00000000000..3b21401066c --- /dev/null +++ b/litellm/types/token_count.py @@ -0,0 +1,16 @@ +from typing import Literal, TypeAlias + +from typing_extensions import NotRequired, ReadOnly, TypedDict + +TokenCountSource: TypeAlias = Literal["provider", "token_count_model", "local_estimate"] +TokenCountTokenizer: TypeAlias = Literal["kimi_k3"] + + +class TokenCountModelInfo(TypedDict): + model: ReadOnly[str] + region: ReadOnly[NotRequired[str]] + + +class TokenCountSourceFields(TypedDict, total=False): + litellm_count_source: ReadOnly[TokenCountSource] + litellm_count_model: ReadOnly[str] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index a673db3db86..30de80a460d 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -9,6 +9,7 @@ from typing import ( Any, Final, Literal, + TypeAlias, get_args, ) @@ -55,6 +56,7 @@ from litellm.types.llms.base import ( LiteLLMPydanticObjectBase, ) from litellm.types.mcp import MCPServerCostInfo +from litellm.types.token_count import TokenCountModelInfo, TokenCountSource, TokenCountTokenizer from ..litellm_core_utils.core_helpers import map_finish_reason, process_response_headers from . import litellm_params as _litellm_params @@ -211,6 +213,9 @@ class ProviderSpecificModelInfo(TypedDict, total=False): vertex_ai_audio_api: ReadOnly[Literal["lyria_predict", "lyria_interactions"] | None] bedrock_output_config_effort_ceiling: Literal["low", "medium", "high", "max", "xhigh"] | None bedrock_converse_supports_strict_tools: bool | None + token_count_framing: ReadOnly[int | None] + token_count_model: ReadOnly[TokenCountModelInfo | None] + token_count_tokenizer: ReadOnly[TokenCountTokenizer | None] class SearchContextCostPerQuery(TypedDict, total=False): @@ -3856,10 +3861,13 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): DEPLOYMENT_SCOPED_PRICING_FIELDS: Final[frozenset[str]] = frozenset({"off_peak_pricing"}) +DEPLOYMENT_SCOPED_TOKEN_COUNT_FIELDS: Final[frozenset[str]] = frozenset({"token_count_framing", "token_count_model"}) + SHARED_BACKEND_MODEL_INFO_FIELDS: Final[frozenset[str]] = ( frozenset(ModelInfoBase.__required_keys__ | ModelInfoBase.__optional_keys__) - frozenset(CustomPricingLiteLLMParams.model_fields) - DEPLOYMENT_SCOPED_PRICING_FIELDS + - DEPLOYMENT_SCOPED_TOKEN_COUNT_FIELDS ) @@ -4255,11 +4263,15 @@ class LiteLLMLoggingBaseClass: pass +TokenCountMessagesFormat: TypeAlias = Literal["openai", "anthropic"] + + class TokenCountResponse(LiteLLMPydanticObjectBase): total_tokens: int request_model: str model_used: str tokenizer_type: str + count_source: TokenCountSource | None = None original_response: dict | None = None """ Original Response from upstream API call - if an API call was made for token counting @@ -4270,6 +4282,7 @@ class TokenCountResponse(LiteLLMPydanticObjectBase): HTTP status code from the token counting API (e.g., 200 for success, 429 for rate limit, 400 for bad request) """ status_code: int | None = None + rejected_before_provider_call: bool = Field(default=False, exclude=True) class CustomHuggingfaceTokenizer(TypedDict): diff --git a/litellm/utils.py b/litellm/utils.py index 05969bba4fd..922f2436123 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6364,6 +6364,9 @@ def _get_model_info_helper( uses_embed_content=_model_info.get("uses_embed_content", None), supports_image_size=_model_info.get("supports_image_size", None), supported_audio_formats=_model_info.get("supported_audio_formats", None), + token_count_framing=_model_info.get("token_count_framing", None), + token_count_model=_model_info.get("token_count_model", None), + token_count_tokenizer=_model_info.get("token_count_tokenizer", None), vertex_ai_audio_api=_model_info.get("vertex_ai_audio_api", None), ) for cost_key, cost_value in _model_info.items(): diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d42e66605d7..78f17ee3f07 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -19620,7 +19620,13 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": false, + "token_count_framing": 97 + }, + "databricks/system.ai.deepseek-v4-flash-0731": { + "litellm_provider": "databricks", + "mode": "chat", + "token_count_framing": 97 }, "databricks/databricks-deepseek-v4-pro-0813": { "cache_creation_input_token_cost": 1.31999e-06, @@ -20611,7 +20617,13 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 100 + }, + "databricks/system.ai.kimi-k3": { + "litellm_provider": "databricks", + "mode": "chat", + "token_count_framing": 100 }, "databricks/databricks-llama-4-maverick": { "cache_creation_input_token_cost": 5.0001e-07, @@ -76780,7 +76792,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3" }, "global.moonshotai.kimi-k3": { "cache_creation_input_token_cost": 3.75e-06, @@ -76800,7 +76814,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3" }, "us.moonshotai.kimi-k3": { "cache_creation_input_token_cost": 4.125e-06, @@ -76820,7 +76836,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3" }, "openrouter/xiaomi/mimo-v2.6-flash": { "cache_read_input_token_cost": 2.8e-09, diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index d7126df23d5..a59b2bd342d 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -1161,6 +1161,18 @@ "additionalProperties": false } }, + "token_count_framing": { + "type": "integer", + "minimum": 0, + "description": "Tokens added once per request to a local count estimate, for the provider's chat framing." + }, + "token_count_tokenizer": { + "type": "string", + "description": "Tokenizer litellm counts this model's requests with when the provider cannot count them.", + "enum": [ + "kimi_k3" + ] + }, "tpm": { "type": "integer", "minimum": 0, diff --git a/tests/integration/_support/provider.py b/tests/integration/_support/provider.py new file mode 100644 index 00000000000..fa0bb7ad379 --- /dev/null +++ b/tests/integration/_support/provider.py @@ -0,0 +1,49 @@ +"""The fake provider shared by every integration test that takes the `provider` fixture. + +Deployments in `proxy_config.yaml` point at `PROVIDER_URL`, so one server answers for all of them. A test +queues the replies it expects with `expect` and reads what the proxy sent with `received`. Tests run one at a +time against it; the `provider` fixture checks nothing is left over between tests. +The proxy's model-info refresh is answered here. +""" + +from __future__ import annotations + +import os +from collections import deque +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import Final + +from tests.integration._support.wire import Reply, Request, Wire, wire_server + +PROVIDER_PORT: Final = int(os.environ.get("INTEGRATION_PROVIDER_PORT", "8191")) +PROVIDER_URL: Final = f"http://127.0.0.1:{PROVIDER_PORT}" +MODEL_DISCOVERY: Final = ("GET", "/v1/models") +_UNQUEUED: Final = Reply(status=500, body=b'{"error": "the shared fake provider has no reply queued for this request"}') + + +@dataclass(slots=True) +class SharedProvider: + wire: Wire + replies: deque[Reply] + last_test: str | None = field(default=None) + + def expect(self, *replies: Reply) -> None: + self.replies.extend(replies) + + def received(self) -> tuple[Request, ...]: + return tuple(request for request in self.wire.drain() if (request.method, request.target) != MODEL_DISCOVERY) + + +@contextmanager +def shared_provider() -> Iterator[SharedProvider]: + replies: Final[deque[Reply]] = deque() + + def respond(request: Request) -> Reply: + if (request.method, request.target) == MODEL_DISCOVERY: + return Reply(body=b'{"object":"list","data":[]}') + return replies.popleft() if replies else _UNQUEUED + + with wire_server(respond, port=PROVIDER_PORT) as wire: + yield SharedProvider(wire, replies) diff --git a/tests/integration/_support/vertex.py b/tests/integration/_support/vertex.py new file mode 100644 index 00000000000..af7178fa6e9 --- /dev/null +++ b/tests/integration/_support/vertex.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +import json +from typing import Final + +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa + + +def service_account_json(project: str, token_url: str) -> str: + private_key: Final = ( + rsa.generate_private_key(public_exponent=65537, key_size=2048) + .private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ) + .decode() + ) + return json.dumps( + { + "type": "service_account", + "project_id": project, + "private_key_id": "scripted", + "private_key": private_key, + "client_email": f"scripted@{project}.iam.gserviceaccount.com", + "client_id": "0", + "auth_uri": f"{token_url}/_oauth/authorize", + "token_uri": f"{token_url}/_oauth/token", + } + ) diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 1b39eb81b01..1b137c38898 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -15,6 +15,7 @@ from redis import Redis from tests.integration._support.client import Gateway, eventually, gateway_from_environment from tests.integration._support.generation import LIFECYCLE_SETTINGS from tests.integration._support.manifest import OWNED_DIRECTORIES +from tests.integration._support.provider import SharedProvider, shared_provider from tests.integration._support.routing import RoutingPlugin COLLECTED: Final = pytest.StashKey[tuple[str, ...]]() @@ -126,6 +127,35 @@ def gateway() -> Iterator[Gateway]: yield value +@pytest.fixture(scope="session") +def shared_provider_server() -> Iterator[SharedProvider]: + if os.environ.get("PYTEST_XDIST_WORKER") and os.environ.get("INTEGRATION_WORKER_ISOLATION") != "1": + pytest.fail("the shared fake provider needs tests to run one at a time; this group runs under pytest-xdist") + with shared_provider() as server: + yield server + late: Final = server.received() + assert late == (), ( + f"the shared fake provider got {[item.target for item in late]} after {server.last_test} finished" + ) + + +@pytest.fixture +def provider(shared_provider_server: SharedProvider, request: pytest.FixtureRequest) -> Iterator[SharedProvider]: + stray: Final = shared_provider_server.received() + shared_provider_server.replies.clear() + assert stray == (), ( + f"the shared fake provider got {[item.target for item in stray]} " + f"after {shared_provider_server.last_test} finished" + ) + yield shared_provider_server + shared_provider_server.last_test = request.node.nodeid + unused: Final = len(shared_provider_server.replies) + unread: Final = shared_provider_server.received() + shared_provider_server.replies.clear() + assert unused == 0, f"{unused} queued provider replies were never requested" + assert unread == (), f"the test never read the provider requests {[item.target for item in unread]}" + + @pytest.fixture def peer(gateway: Gateway) -> Iterator[Gateway]: url: Final = os.environ["INTEGRATION_PEER_URL"] diff --git a/tests/integration/providers/_cache_control_marks_support.py b/tests/integration/providers/_cache_control_marks_support.py new file mode 100644 index 00000000000..2f6071f9e45 --- /dev/null +++ b/tests/integration/providers/_cache_control_marks_support.py @@ -0,0 +1,479 @@ +import json +import re +import uuid +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import yaml +from integration._support.client import Gateway, eventually, object_value +from integration._support.database import read_rows +from integration._support.wire import Reply, Request +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter + +# https://platform.claude.com/docs/en/build-with-claude/prompt-caching (read 2026-09-28): at most 4 blocks with cache_control +ANTHROPIC_CACHE_CONTROL_CAP: Final = 4 +# https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-caching.html (read 2026-09-28): at most 4 cache checkpoints +BEDROCK_CACHE_CHECKPOINT_CAP: Final = 4 + +ANTHROPIC_MODEL: Final = "claude-opus-5-5" +BEDROCK_MODEL: Final = "anthropic.claude-opus-5-5" +UNMAPPED_BEDROCK_ARN: Final = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/scripted" +PROVIDER_KEY: Final = "synthetic-provider-key" +SYSTEM: Final = "Use the provided tool results to answer the user." +ASK: Final = "Look up the weather in London, Paris, and Tokyo." +CITIES: Final = ("London", "Paris", "Tokyo") +EPHEMERAL: Final[dict[str, JsonValue]] = {"type": "ephemeral"} +POINTS: Final[list[JsonValue]] = [ + {"location": "message", "role": "system"}, + {"location": "message", "index": -1}, +] +TOOL: Final[dict[str, JsonValue]] = { + "type": "function", + "function": { + "name": "lookup_weather", + "description": "Look up the weather in a city.", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, +} +SYSTEM_LABEL: Final = f"system:{SYSTEM}" +ASK_LABEL: Final = f"user:text:{ASK}" +_MARKER: Final = re.compile(rb"marker-([0-9a-f]{32})") +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +@dataclass(frozen=True, slots=True) +class Mark: + label: str + ttl: str | None + + +def new_marker() -> str: + return uuid.uuid4().hex + + +def final_text(marker: str) -> str: + return f"Summarize the results in one word. marker-{marker}" + + +def final_label(marker: str) -> str: + return f"user:text:{final_text(marker)}" + + +def call_id(city: str) -> str: + return f"call_weather_{city.lower()}" + + +def tool_use_label(city: str) -> str: + return f"assistant:tool_use:{call_id(city)}" + + +def tool_call(city: str, **fields: JsonValue) -> dict[str, JsonValue]: + return { + "id": call_id(city), + "type": "function", + "function": {"name": "lookup_weather", "arguments": json.dumps({"city": city})}, + **fields, + } + + +def marked_calls(mark: JsonValue = EPHEMERAL, cities: Sequence[str] = CITIES) -> list[JsonValue]: + return [tool_call(city, cache_control=mark) for city in cities] + + +def client_marked() -> list[str]: + return [ASK_LABEL, *(tool_use_label(city) for city in CITIES)] + + +def ask(*, marked: bool) -> dict[str, JsonValue]: + block: Final[dict[str, JsonValue]] = {"type": "text", "text": ASK} + return {"role": "user", "content": [{**block, "cache_control": EPHEMERAL} if marked else block]} + + +def tool_results(calls: Sequence[JsonValue]) -> list[JsonValue]: + return [ + {"role": "tool", "tool_call_id": call["id"], "content": "sunny."} + for call in calls + if isinstance(call, dict) and str(call.get("id", "")).startswith("call_") + ] + + +def conversation( + marker: str, + calls: Sequence[JsonValue], + *, + ask_marked: bool = True, + assistant: dict[str, JsonValue] | None = None, +) -> list[JsonValue]: + turn: Final[dict[str, JsonValue]] = { + "role": "assistant", + "content": "", + "tool_calls": list(calls), + **(assistant or {}), + } + return [ + {"role": "system", "content": SYSTEM}, + ask(marked=ask_marked), + turn, + *tool_results(calls), + {"role": "user", "content": final_text(marker)}, + ] + + +def chat_body(model: str, messages: Sequence[JsonValue], **fields: JsonValue) -> dict[str, JsonValue]: + return {"model": model, "messages": list(messages), "tools": [TOOL], "max_tokens": 64, **fields} + + +def messages_body(model: str, marker: str, *, stream: bool) -> dict[str, JsonValue]: + return { + "model": model, + "max_tokens": 64, + "stream": stream, + "system": [{"type": "text", "text": SYSTEM}], + "tools": [ + { + "name": "lookup_weather", + "description": "Look up the weather in a city.", + "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + } + ], + "messages": [ + {"role": "user", "content": [{"type": "text", "text": ASK, "cache_control": EPHEMERAL}]}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": call_id(city), + "name": "lookup_weather", + "input": {"city": city}, + "cache_control": EPHEMERAL, + } + for city in CITIES + ], + }, + { + "role": "user", + "content": [ + *({"type": "tool_result", "tool_use_id": call_id(city), "content": "sunny."} for city in CITIES), + {"type": "text", "text": final_text(marker)}, + ], + }, + ], + } + + +def responses_body(model: str, marker: str) -> dict[str, JsonValue]: + items: Final[list[JsonValue]] = [ + {"role": "user", "content": [{"type": "input_text", "text": ASK, "cache_control": EPHEMERAL}]}, + *( + { + "type": "function_call", + "call_id": call_id(city), + "name": "lookup_weather", + "arguments": json.dumps({"city": city}), + "cache_control": EPHEMERAL, + } + for city in CITIES + ), + *({"type": "function_call_output", "call_id": call_id(city), "output": "sunny."} for city in CITIES), + {"role": "user", "content": final_text(marker)}, + ] + return { + "model": model, + "instructions": SYSTEM, + "max_output_tokens": 64, + "tools": [{"type": "function", "name": "lookup_weather", "parameters": {"type": "object"}}], + "input": items, + } + + +def marker_of(request: Request) -> str: + found: Final = _MARKER.findall(request.body) + return found[-1].decode() if found else "unmarked" + + +class _AnthropicBlock(BaseModel): + model_config = ConfigDict(extra="allow", frozen=True) + type: str + text: str | None = None + id: str | None = None + tool_use_id: str | None = None + name: str | None = None + cache_control: JsonValue = None + + +class _AnthropicMessage(BaseModel): + model_config = ConfigDict(extra="allow", frozen=True) + role: str + content: str | tuple[_AnthropicBlock, ...] + + +class _AnthropicTool(BaseModel): + model_config = ConfigDict(extra="allow", frozen=True) + name: str = "" + cache_control: JsonValue = None + + +class _AnthropicBody(BaseModel): + model_config = ConfigDict(extra="allow", frozen=True) + system: str | tuple[_AnthropicBlock, ...] = () + messages: tuple[_AnthropicMessage, ...] = () + tools: tuple[_AnthropicTool, ...] = () + cache_control: JsonValue = None + stream: bool = False + + +def _ttl(cache_control: JsonValue) -> str | None: + if not isinstance(cache_control, dict): + return None + ttl: Final = cache_control.get("ttl") + return ttl if isinstance(ttl, str) else None + + +def _block_label(role: str, block: _AnthropicBlock) -> str: + detail: Final = block.text if block.type == "text" else block.id or block.tool_use_id or "" + return f"{role}:{block.type}:{detail}" + + +def _message_blocks(body: _AnthropicBody) -> Iterator[tuple[str, _AnthropicBlock]]: + for message in body.messages: + if isinstance(message.content, tuple): + yield from ((message.role, block) for block in message.content) + + +def _anthropic_body(request: Request) -> _AnthropicBody: + return _AnthropicBody.model_validate_json(request.body) + + +def anthropic_marks(request: Request) -> tuple[Mark, ...]: + body: Final = _anthropic_body(request) + system: Final = body.system if isinstance(body.system, tuple) else () + return ( + *(Mark(f"tool:{tool.name}", _ttl(tool.cache_control)) for tool in body.tools if tool.cache_control is not None), + *( + Mark(f"system:{block.text}", _ttl(block.cache_control)) + for block in system + if block.cache_control is not None + ), + *( + Mark(_block_label(role, block), _ttl(block.cache_control)) + for role, block in _message_blocks(body) + if block.cache_control is not None + ), + *((Mark("request", _ttl(body.cache_control)),) if body.cache_control is not None else ()), + ) + + +def anthropic_labels(request: Request) -> list[str]: + return [mark.label for mark in anthropic_marks(request)] + + +def _anthropic_error(message: str) -> Reply: + return Reply( + status=400, + body=json.dumps({"type": "error", "error": {"type": "invalid_request_error", "message": message}}).encode(), + ) + + +def _ttl_out_of_order(marks: Sequence[Mark]) -> bool: + first_short: Final = next((index for index, mark in enumerate(marks) if mark.ttl != "1h"), len(marks)) + return any(mark.ttl == "1h" for mark in marks[first_short:]) + + +def _usage() -> dict[str, JsonValue]: + return {"input_tokens": 12, "output_tokens": 1, "cache_creation_input_tokens": 12, "cache_read_input_tokens": 0} + + +def _anthropic_message(identity: str) -> dict[str, JsonValue]: + return { + "id": identity, + "type": "message", + "role": "assistant", + "model": ANTHROPIC_MODEL, + "content": [{"type": "text", "text": "sunny"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": _usage(), + } + + +def _event(name: str, data: dict[str, JsonValue]) -> bytes: + return f"event: {name}\ndata: {json.dumps(data)}\n\n".encode() + + +def _anthropic_stream(identity: str) -> tuple[bytes, ...]: + return ( + _event( + "message_start", + {"type": "message_start", "message": {**_anthropic_message(identity), "content": [], "stop_reason": None}}, + ), + _event( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + _event( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "sunny"}}, + ), + _event("content_block_stop", {"type": "content_block_stop", "index": 0}), + _event( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 1}, + }, + ), + _event("message_stop", {"type": "message_stop"}), + ) + + +def anthropic_peer(request: Request) -> Reply: + marks: Final = anthropic_marks(request) + if len(marks) > ANTHROPIC_CACHE_CONTROL_CAP: + return _anthropic_error( + f"A maximum of {ANTHROPIC_CACHE_CONTROL_CAP} blocks with cache_control may be provided. Found {len(marks)}." + ) + if _ttl_out_of_order(marks): + return _anthropic_error("a ttl='1h' cache_control block must not come after a ttl='5m' cache_control block") + identity: Final = f"msg_{marker_of(request)}" + if _anthropic_body(request).stream: + return Reply(content_type="text/event-stream", chunks=_anthropic_stream(identity)) + return Reply(body=json.dumps(_anthropic_message(identity)).encode()) + + +def _bedrock_label(role: str, block: JsonValue) -> str: + if not isinstance(block, dict): + return f"{role}:start" + if isinstance(block.get("text"), str): + return f"{role}:text:{block['text']}" + for kind in ("toolUse", "toolResult"): + inner = block.get(kind) + if isinstance(inner, dict): + return f"{role}:{kind}:{inner.get('toolUseId')}" + spec: Final = block.get("toolSpec") + return f"tool:{spec.get('name')}" if isinstance(spec, dict) else f"{role}:other" + + +def _cache_points(role: str, blocks: JsonValue) -> Iterator[str]: + listed: Final = blocks if isinstance(blocks, list) else [] + for previous, block in zip([None, *listed], listed): + if isinstance(block, dict) and "cachePoint" in block: + yield _bedrock_label(role, previous) + + +def _bedrock_sections(body: dict[str, JsonValue]) -> Iterator[tuple[str, JsonValue]]: + tool_config: Final = body.get("toolConfig") + yield ("tool", tool_config.get("tools") if isinstance(tool_config, dict) else None) + yield ("system", body.get("system")) + messages: Final = body.get("messages") + for message in messages if isinstance(messages, list) else []: + if isinstance(message, dict): + yield (str(message.get("role")), message.get("content")) + + +def bedrock_labels(request: Request) -> list[str]: + body: Final = _JSON_OBJECT.validate_json(request.body) + return [label for role, blocks in _bedrock_sections(body) for label in _cache_points(role, blocks)] + + +def bedrock_peer(request: Request) -> Reply: + found: Final = len(bedrock_labels(request)) + if found > BEDROCK_CACHE_CHECKPOINT_CAP: + return Reply( + status=400, + headers={"x-amzn-errortype": "ValidationException"}, + body=json.dumps( + { + "message": f"A maximum of {BEDROCK_CACHE_CHECKPOINT_CAP} cache checkpoints may be provided. Found {found}." + } + ).encode(), + ) + return Reply( + body=json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": "sunny"}]}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 12, + "outputTokens": 1, + "totalTokens": 13, + "cacheWriteInputTokens": 12, + "cacheReadInputTokens": 0, + }, + "metrics": {"latencyMs": 1}, + } + ).encode() + ) + + +def _without_cache_points(value: JsonValue) -> JsonValue: + if isinstance(value, dict): + return {key: _without_cache_points(inner) for key, inner in value.items()} + if isinstance(value, list): + return [ + _without_cache_points(inner) for inner in value if not (isinstance(inner, dict) and "cachePoint" in inner) + ] + return value + + +def sent_without_cache_points(requests: Sequence[Request], marker: str) -> list[JsonValue]: + return [ + _without_cache_points(_JSON_OBJECT.validate_json(request.body.replace(marker.encode(), b"marker"))) + for request in requests + ] + + +def gateway_injected(response_id: str) -> bool: + rows: Final = eventually( + lambda: read_rows( + """SELECT metadata->>'litellm_gateway_injected_cache' AS injected FROM "LiteLLM_SpendLogs" """ + "WHERE request_id=%s", + (response_id,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + return rows[0]["injected"] is not None + + +def post_chat(gateway: Gateway, body: dict[str, JsonValue], *, key: str | None = None) -> tuple[int, str, str]: + response: Final = gateway.request("POST", "/v1/chat/completions", body, key=key) + identity: Final = _JSON_OBJECT.validate_json(response.content).get("id") if response.status_code == 200 else None + return response.status_code, str(identity), response.text + + +def anthropic_deployment(name: str, api_base: str, **fields: JsonValue) -> dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": { + "model": f"anthropic/{ANTHROPIC_MODEL}", + "api_base": api_base, + "api_key": PROVIDER_KEY, + **fields, + }, + } + + +def owned_config( + directory: Path, + model_list: Sequence[JsonValue], + *, + litellm_settings: Mapping[str, JsonValue] = MappingProxyType({}), + router_settings: Mapping[str, JsonValue] = MappingProxyType({}), +) -> Path: + config: Final = _JSON_OBJECT.validate_python( + yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + ) + merged: Final = { + **config, + "model_list": list(model_list), + "litellm_settings": {**object_value(config["litellm_settings"]), **litellm_settings}, + "router_settings": {**object_value(config["router_settings"]), "num_retries": 0, **router_settings}, + } + path: Final = directory / f"cache-control-marks-{uuid.uuid4().hex}.yaml" + path.write_text(yaml.safe_dump(merged)) + return path diff --git a/tests/integration/providers/test_count_tokens_error_shapes_wire.py b/tests/integration/providers/test_count_tokens_error_shapes_wire.py new file mode 100644 index 00000000000..74581e44310 --- /dev/null +++ b/tests/integration/providers/test_count_tokens_error_shapes_wire.py @@ -0,0 +1,215 @@ +import json +from collections.abc import Iterator +from pathlib import Path +from typing import Final + +import anthropic +import openai +import pytest +from integration._support.client import Gateway, gateway_from_environment +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": "Count this synthetic request."}] +_PROVIDER_KEY_VARIABLES: Final = ( + "ANTHROPIC_API_KEY", + "ANTHROPIC_AUTH_TOKEN", + "AZURE_AI_API_KEY", + "AZURE_API_KEY", + "AZURE_OPENAI_API_KEY", + "AZURE_AD_TOKEN", + "AZURE_CLIENT_ID", + "AZURE_CLIENT_SECRET", + "AZURE_TENANT_ID", + "OPENAI_API_KEY", +) + + +@pytest.fixture(scope="module") +def provider_error_config(tmp_path_factory: pytest.TempPathFactory) -> Path: + config: Final = tmp_path_factory.mktemp("count-tokens-provider-errors") / "config.yaml" + base: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" + config.write_text(base.read_text() + "\nlitellm_settings:\n disable_token_counter: true\n") + return config + + +@pytest.fixture(scope="module") +def keyless_proxy(tmp_path_factory: pytest.TempPathFactory, provider_error_config: Path) -> Iterator[OwnedProxy]: + with ( + gateway_from_environment() as gateway, + owned_proxy_process( + gateway, + tmp_path_factory.mktemp("count-tokens-error-shapes"), + {}, + config=provider_error_config, + remove_environment=_PROVIDER_KEY_VARIABLES, + workers=1, + ) as owned, + ): + yield owned + + +@pytest.fixture(scope="module") +def azure_environment_key_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]: + with ( + gateway_from_environment() as gateway, + owned_proxy_process( + gateway, + tmp_path_factory.mktemp("count-tokens-azure-environment-key"), + {"AZURE_API_KEY": "synthetic-azure-env-key"}, + remove_environment=tuple(name for name in _PROVIDER_KEY_VARIABLES if name != "AZURE_API_KEY"), + workers=1, + ) as owned, + ): + yield owned + + +def _anthropic(gateway: Gateway) -> anthropic.Anthropic: + return anthropic.Anthropic( + api_key=gateway.key, base_url=str(gateway.client.base_url), max_retries=0, http_client=gateway.client + ) + + +def _never_called(request: Request) -> Reply: + raise AssertionError(f"count reached the provider without a key: {request.target}") + + +@pytest.mark.parametrize( + ("deployment", "missing"), + ( + ({"model": "anthropic/claude-sonnet-4-6"}, "No Anthropic API key is set for model "), + ({"model": "azure_ai/claude-sonnet-4-6"}, "No Azure AI API key or Entra ID credential is set for model "), + ), + ids=("anthropic", "azure_ai_anthropic"), +) +def test_anthropic_count_without_a_provider_credential_is_an_anthropic_401( + keyless_proxy: OwnedProxy, deployment: dict[str, JsonValue], missing: str +) -> None: + gateway: Final = keyless_proxy.gateway + with wire_server(_never_called) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(**deployment, api_base=wire.url, api_key=None, num_retries=0) + with pytest.raises(anthropic.AuthenticationError) as raised: + _anthropic(gateway).messages.count_tokens(model=model, messages=_MESSAGES) + assert raised.value.body == { + "type": "error", + "error": {"type": "authentication_error", "message": raised.value.body["error"]["message"]}, + } + assert missing in raised.value.message, raised.value.message + assert wire.drain() == () + + +def test_azure_ai_anthropic_count_sends_the_azure_environment_key_as_x_api_key( + azure_environment_key_proxy: OwnedProxy, +) -> None: + gateway: Final = azure_environment_key_proxy.gateway + + def count(request: Request) -> Reply: + assert request.headers.get("x-api-key") == "synthetic-azure-env-key", request.headers + return Reply(body=json.dumps({"input_tokens": 58}).encode()) + + with wire_server(count) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="azure_ai/claude-sonnet-4-6", + api_base=wire.url, + api_key=None, + num_retries=0, + ) + result: Final = _anthropic(gateway).messages.count_tokens(model=model, messages=_MESSAGES) + assert result.input_tokens == 58 + + +def test_azure_ai_anthropic_count_authenticates_with_an_entra_token(keyless_proxy: OwnedProxy) -> None: + gateway: Final = keyless_proxy.gateway + + def count(request: Request) -> Reply: + assert request.headers.get("authorization") == "Bearer synthetic-entra-token", request.headers + assert "x-api-key" not in request.headers, request.headers + return Reply(body=json.dumps({"input_tokens": 61}).encode()) + + with wire_server(count) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="azure_ai/claude-sonnet-4-6", + api_base=wire.url, + api_key=None, + azure_ad_token="synthetic-entra-token", + num_retries=0, + ) + result: Final = _anthropic(gateway).messages.count_tokens(model=model, messages=_MESSAGES) + assert result.input_tokens == 61 + assert len(wire.drain()) == 1 + + +def test_anthropic_count_request_errors_use_the_anthropic_envelope(keyless_proxy: OwnedProxy) -> None: + gateway: Final = keyless_proxy.gateway + upstream_error: Final = {"type": "error", "error": {"type": "invalid_request_error", "message": "scripted reject"}} + with ( + wire_server(lambda _request: Reply(status=400, body=json.dumps(upstream_error).encode())) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model="anthropic/claude-sonnet-4-6", api_base=wire.url, api_key="synthetic-key", num_retries=0 + ) + missing_messages: Final = gateway.request("POST", "/v1/messages/count_tokens", {"model": model}) + assert missing_messages.status_code == 400, missing_messages.text + assert _JSON_OBJECT.validate_json(missing_messages.content) == { + "type": "error", + "error": {"type": "invalid_request_error", "message": "messages parameter is required"}, + } + + rejected: Final = gateway.request( + "POST", "/v1/messages/count_tokens", {"model": model, "messages": _MESSAGES} + ) + assert rejected.status_code == 400, rejected.text + body: Final = _JSON_OBJECT.validate_json(rejected.content) + assert set(body) <= {"type", "error", "request_id"}, body + assert body["error"] == upstream_error["error"], body + assert len(wire.drain()) == 1 + + +def test_openai_input_tokens_without_a_provider_key_is_an_openai_401(keyless_proxy: OwnedProxy) -> None: + gateway: Final = keyless_proxy.gateway + with wire_server(_never_called) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/gpt-4o-mini", api_base=f"{wire.url}/v1", api_key=None, num_retries=0 + ) + client: Final = openai.OpenAI( + api_key=gateway.key, base_url=str(gateway.client.base_url), max_retries=0, http_client=gateway.client + ) + with pytest.raises(openai.AuthenticationError) as raised: + client.responses.input_tokens.count(model=model, input="Count this synthetic request.") + assert f"No OpenAI API key is set for model {model}" in raised.value.message + assert wire.drain() == () + + +def test_google_count_tokens_errors_use_the_google_envelope(keyless_proxy: OwnedProxy) -> None: + gateway: Final = keyless_proxy.gateway + upstream_error: Final = { + "error": {"code": 400, "message": "API key not valid. Please pass a valid API key.", "status": "INVALID_ARGUMENT"} + } + with ( + wire_server(lambda _request: Reply(status=400, body=json.dumps(upstream_error).encode())) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model="gemini/gemini-2.5-flash", api_base=wire.url, api_key="synthetic-key", num_retries=0 + ) + rejected: Final = gateway.request( + "POST", + f"/v1beta/models/{model}:countTokens", + {"contents": [{"role": "user", "parts": [{"text": "Count this synthetic request."}]}]}, + ) + assert rejected.status_code == 400, rejected.text + error: Final = _JSON_OBJECT.validate_json(rejected.content)["error"] + assert isinstance(error, dict) and error["code"] == 400 and error["status"] == "INVALID_ARGUMENT", error + assert "API key not valid" in str(error["message"]), error + assert len(wire.drain()) == 1 + + malformed: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", {"contents": "not-a-list"}) + assert malformed.status_code == 400, malformed.text + assert _JSON_OBJECT.validate_json(malformed.content) == { + "error": {"code": 400, "message": "Invalid Google countTokens request body", "status": "INVALID_ARGUMENT"} + } diff --git a/tests/integration/providers/test_openai_count_tokens_tools_wire.py b/tests/integration/providers/test_openai_count_tokens_tools_wire.py new file mode 100644 index 00000000000..2b23ccdf4f2 --- /dev/null +++ b/tests/integration/providers/test_openai_count_tokens_tools_wire.py @@ -0,0 +1,393 @@ +import json +from collections.abc import Iterator, Mapping +from typing import Final + +import pytest +from integration._support.client import Scenario, gateway_from_environment +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_COUNT: Final = 4242 +_TOOL_SCHEMA: Final[dict[str, JsonValue]] = {"type": "object", "properties": {}} +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +@pytest.fixture(scope="module") +def counting_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]: + with ( + gateway_from_environment() as gateway, + owned_proxy_process(gateway, tmp_path_factory.mktemp("openai-count-tools"), {}, workers=1) as owned, + ): + yield owned + + +def _deployment(scenario: Scenario, api_base: str) -> str: + return scenario.model( + model="openai/gpt-4o-mini", + api_base=f"{api_base}/v1", + api_key="sk-integration-openai", + ) + + +def _count_response(_request: Request) -> Reply: + return Reply(body=json.dumps({"input_tokens": _COUNT}).encode()) + + +_PDF_BASE64: Final = "JVBERi0xLjQK" +_PNG_BASE64: Final = "aW1hZ2U=" +_LOOKUP_TOOL: Final[dict[str, JsonValue]] = { + "name": "lookup", + "description": "Find a record", + "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}, +} +_RESPONSES_REPLY: Final[Mapping[str, JsonValue]] = { + "id": "resp_count_parity", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "type": "message", + "id": "msg_count_parity", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "OK", "annotations": []}], + } + ], + "usage": {"input_tokens": 10, "output_tokens": 1, "total_tokens": 11}, +} +_WRAPPED_REASONING: Final = "litellm_enc:bW9kZWxfaWQ6ZGVwbG95bWVudC0x;gAAAAoriginal" +_ANSWER_SCHEMA: Final[dict[str, JsonValue]] = {"type": "object", "properties": {"answer": {"type": "string"}}} +_ANTHROPIC_BODIES: Final[Mapping[str, Mapping[str, JsonValue]]] = { + "structured-output": { + "output_config": {"format": {"type": "json_schema", "schema": _ANSWER_SCHEMA}}, + "messages": [{"role": "user", "content": "Answer in JSON."}], + }, + "affinity-wrapped-reasoning": { + "messages": [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "redacted_thinking", "data": f"litellm_encrypted_reasoning:{_WRAPPED_REASONING}"}, + {"type": "text", "text": "OK"}, + ], + }, + {"role": "user", "content": "again"}, + ], + }, + "base64-pdf": { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the attached contract."}, + { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": _PDF_BASE64}, + }, + ], + } + ], + }, + "base64-image": { + "system": [{"type": "text", "text": "Describe images briefly."}], + "messages": [ + { + "role": "user", + "content": [ + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": _PNG_BASE64}}, + {"type": "text", "text": "What is this?"}, + ], + } + ], + }, + "replayed-web-search": { + "messages": [ + {"role": "user", "content": "What is the newest LiteLLM release?"}, + { + "role": "assistant", + "content": [ + { + "type": "server_tool_use", + "id": "srvtoolu_01", + "name": "web_search", + "input": {"query": "newest LiteLLM release"}, + }, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": "", + "snippet": "LiteLLM ships a weekly release candidate and a stable line.", + } + ], + }, + {"type": "text", "text": "The newest release is in the release notes."}, + ], + }, + {"role": "user", "content": "Thanks."}, + ], + }, + "tool-turns": { + "system": "Use the lookup tool.", + "tools": [_LOOKUP_TOOL], + "messages": [ + {"role": "user", "content": "Find record 14."}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Looking it up."}, + {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {"query": "14"}}, + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "14. Indemnification."}], + }, + ], + }, +} + + +def _responses_peer(request: Request) -> Reply: + if request.target == "/v1/responses/input_tokens": + return _count_response(request) + if request.target == "/v1/responses": + return Reply(body=json.dumps(_RESPONSES_REPLY).encode()) + return Reply(status=404, body=json.dumps({"error": request.target}).encode()) + + +def _body_for(wire_requests: tuple[Request, ...], target: str) -> dict[str, JsonValue]: + matching: Final = tuple(request for request in wire_requests if request.target == target) + assert len(matching) == 1, [request.target for request in wire_requests] + return _JSON_OBJECT.validate_json(matching[0].body) + + +def _counted_fields(body: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]: + return {key: body[key] for key in ("input", "instructions", "tools", "text") if body.get(key) is not None} + + +def _count_body(wire: Wire) -> dict[str, JsonValue]: + requests: Final = wire.drain() + assert len(requests) == 1, requests + request: Final = requests[0] + assert request.method == "POST" + assert request.target == "/v1/responses/input_tokens", request.target + return _JSON_OBJECT.validate_json(request.body) + + +def test_openai_count_tokens_uses_deployment_api_base(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_count_response) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + {"model": model, "messages": [{"role": "user", "content": "hello"}]}, + ) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["input_tokens"] == _COUNT + assert "tools" not in _count_body(wire) + + +def test_google_count_tokens_maps_anthropic_tools_to_openai(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_count_response) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + f"/v1beta/models/{model}:countTokens", + { + "contents": [{"role": "user", "parts": [{"text": "hello"}]}], + "tools": [ + { + "functionDeclarations": [ + { + "name": "lookup", + "description": "Find a record", + "parameters": _TOOL_SCHEMA, + } + ] + } + ], + }, + ) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["totalTokens"] == _COUNT + assert _count_body(wire)["tools"] == [ + { + "type": "function", + "name": "lookup", + "description": "Find a record", + "parameters": _TOOL_SCHEMA, + "strict": False, + } + ] + + +def test_messages_count_tokens_maps_anthropic_tools_to_openai(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_count_response) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + { + "model": model, + "messages": [{"role": "user", "content": "hello"}], + "tools": [ + { + "name": "lookup", + "description": "Find a record", + "input_schema": _TOOL_SCHEMA, + } + ], + }, + ) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["input_tokens"] == _COUNT + assert _count_body(wire)["tools"] == [ + { + "type": "function", + "name": "lookup", + "description": "Find a record", + "parameters": _TOOL_SCHEMA, + "strict": False, + } + ] + + +def test_utils_token_counter_keeps_chat_tool_conversion(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_count_response) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + "/utils/token_counter", + { + "model": model, + "messages": [{"role": "user", "content": "hello"}], + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "description": "Find a record", + "parameters": _TOOL_SCHEMA, + }, + } + ], + }, + params={"call_endpoint": "true"}, + ) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["total_tokens"] == _COUNT + assert _count_body(wire)["tools"] == [ + { + "type": "function", + "name": "lookup", + "description": "Find a record", + "parameters": _TOOL_SCHEMA, + } + ] + + +@pytest.mark.parametrize("case", sorted(_ANTHROPIC_BODIES)) +def test_messages_count_tokens_sends_the_input_v1_messages_inference_sends( + counting_proxy: OwnedProxy, case: str +) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_responses_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + body: Final = {"model": model, **_ANTHROPIC_BODIES[case]} + generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16}) + counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body) + assert generated.status_code == 200, generated.text + assert counted.status_code == 200, counted.text + assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == _COUNT + requests: Final = wire.drain() + inference_body: Final = _body_for(requests, "/v1/responses") + count_body: Final = _body_for(requests, "/v1/responses/input_tokens") + assert _counted_fields(count_body) == _counted_fields(inference_body) + assert count_body["model"] == inference_body["model"] + + +_UNCONVERTIBLE_BODY: Final[Mapping[str, JsonValue]] = { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the notes."}, + {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "notes"}}, + {"type": "search_result", "source": "https://example.com/notes", "title": "Notes", "content": []}, + ], + }, + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "search", "input": {}}]}, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_01", + "content": [ + { + "type": "search_result", + "source": "https://example.com/hits", + "title": "Hits", + "content": [{"type": "text", "text": "hit"}], + } + ], + } + ], + }, + ], +} + + +def test_messages_count_tokens_rejects_blocks_it_cannot_convert(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_responses_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request("POST", "/v1/messages/count_tokens", {"model": model, **_UNCONVERTIBLE_BODY}) + assert response.status_code == 400, response.text + assert "document with a text source, search_result, search_result inside a tool_result" in response.text + assert wire.drain() == () + + +def test_messages_count_tokens_without_a_role_answers_400_before_calling_openai(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_count_response) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + {"model": model, "messages": [{"role": "user", "content": "hello"}, {"content": "no role"}]}, + ) + assert response.status_code == 400, response.text + assert "messages.1.role: Field required" in response.text + assert wire.drain() == () + + +def test_google_count_tokens_sends_the_system_instruction_text(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_count_response) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + f"/v1beta/models/{model}:countTokens", + { + "systemInstruction": {"parts": [{"text": "Be brief."}, {"text": "Answer in English."}]}, + "contents": [{"role": "user", "parts": [{"text": "hello"}]}], + }, + ) + assert response.status_code == 200, response.text + assert _count_body(wire)["instructions"] == "Be brief.\nAnswer in English." diff --git a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py new file mode 100644 index 00000000000..5fe82d37e9e --- /dev/null +++ b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py @@ -0,0 +1,596 @@ +import asyncio +import json +import threading +from collections.abc import Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack +from types import MappingProxyType +from typing import Final +from uuid import uuid4 + +import anthropic +import httpx +import pytest +from integration._support.client import Gateway, Scenario +from integration._support.process import graceful_stop_seconds +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: Final = "gemini-4-argon" +_PROJECT: Final = "scripted-project" +_LOCATION: Final = "us-east5" +_COUNT: Final = 777 +_UNIT: Final = " red blue green yellow orange purple black white silver gold" +_USER: Final = f"{_UNIT * 100}\nReturn exactly OK." +_MESSAGES: Final[JsonValue] = [{"role": "user", "content": _USER}] +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_JSON_LIST: Final = TypeAdapter(list[JsonValue]) +_REPLY: Final[Mapping[str, JsonValue]] = { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "OK"}]}, + "finishReason": "STOP", + } + ], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 1, "totalTokenCount": 11}, +} + + +class _GeminiPeer: + def __init__(self) -> None: + self.status: int = 200 + self.requests: tuple[Request, ...] = () + self._lock: Final = threading.Lock() + + def set_status(self, status: int) -> None: + self.status = status + + def __call__(self, request: Request) -> Reply: + with self._lock: + self.requests = (*self.requests, request) + status: Final = self.status + if request.target.endswith(":countTokens"): + return Reply(status=status, body=json.dumps({"totalTokens": _COUNT}).encode()) + if request.target.endswith(":generateContent"): + return Reply(body=json.dumps(_REPLY).encode()) + return Reply(status=404, body=json.dumps({"error": request.target}).encode()) + + def drain(self) -> tuple[Request, ...]: + with self._lock: + requests: Final = self.requests + self.requests = () + return requests + + +def _deployment(gateway: Gateway, scenario: Scenario, api_base: str) -> str: + return scenario.model( + model=f"vertex_ai/{_MODEL}", + api_base=api_base, + api_key=None, + vertex_project=_PROJECT, + vertex_location=_LOCATION, + vertex_credentials=service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")), + ) + + +def _messages(model: str, system: JsonValue | None, tools: JsonValue | None) -> Mapping[str, JsonValue]: + return { + "model": model, + "messages": _MESSAGES, + **({"system": system} if system is not None else {}), + **({"tools": tools} if tools is not None else {}), + } + + +def _json_body(request: Request) -> Mapping[str, JsonValue]: + return _JSON_OBJECT.validate_json(request.body) + + +def _count_requests(wire: Wire) -> tuple[Request, ...]: + return tuple(request for request in wire.drain() if request.target.endswith(":countTokens")) + + +def _count_tokens_tools(generated_tools: JsonValue) -> JsonValue: + return _JSON_LIST.validate_python( + [ + { + { + "function_declarations": "functionDeclarations", + "url_context": "urlContext", + "code_execution": "codeExecution", + }.get(key, key): value + for key, value in tool.items() + } + for tool in TypeAdapter(list[Mapping[str, JsonValue]]).validate_python(generated_tools) + ] + ) + + +def _count_tokens_request_from_generate_content( + generated_body: Mapping[str, JsonValue], +) -> Mapping[str, JsonValue]: + system_instruction: Final = generated_body.get("systemInstruction", generated_body.get("system_instruction")) + generated_tools: Final = generated_body.get("tools") + return MappingProxyType( + { + "contents": generated_body["contents"], + **({"systemInstruction": system_instruction} if system_instruction is not None else {}), + **({"tools": _count_tokens_tools(generated_tools)} if generated_tools is not None else {}), + } + ) + + +def _tools(description: str) -> JsonValue: + return _JSON_LIST.validate_python( + [ + { + "name": "noop", + "description": description, + "input_schema": {"type": "object", "properties": {}}, + } + ] + ) + + +@pytest.mark.parametrize( + ("large_system", "large_tool"), + ((False, False), (True, False), (False, True), (True, True)), + ids=("baseline", "large-system", "large-tool", "large-system-and-tool"), +) +def test_messages_and_count_tokens_send_matching_gemini_context( + gateway: Gateway, large_system: bool, large_tool: bool +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + system: Final = ( + f"{uuid4().hex}\n" + "This is a synthetic context-budget diagnostic. Return exactly OK." + (_UNIT * 100 if large_system else "") + if large_system + else None + ) + tools: Final = _tools(_UNIT * 100) if large_tool else None + body: Final = _messages(model, system, tools) + generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16}) + counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body) + assert generated.status_code == 200, generated.text + assert counted.status_code == 200, counted.text + assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == _COUNT, counted.text + requests: Final = peer.drain() + generated_body: Final = _json_body( + next(request for request in requests if request.target.endswith(":generateContent")) + ) + count_body: Final = _json_body(next(request for request in requests if request.target.endswith(":countTokens"))) + assert count_body == _count_tokens_request_from_generate_content(generated_body) + + +_LOOKUP_TOOL: Final[JsonValue] = { + "name": "lookup", + "description": "Look up a value", + "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}, +} +_CONTENT_BLOCK_BODIES: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingProxyType( + { + "base64-pdf": { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the document."}, + { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"}, + }, + ], + } + ], + }, + "base64-image": { + "messages": [ + { + "role": "user", + "content": [ + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}}, + {"type": "text", "text": "What is this?"}, + ], + } + ], + }, + "tool-turns": { + "tools": [_LOOKUP_TOOL], + "messages": [ + {"role": "user", "content": "Look up the value."}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Looking it up."}, + {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {"query": "value"}}, + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "42"}], + }, + ], + }, + } +) + + +def _ai_studio_deployment(scenario: Scenario, api_base: str) -> str: + return scenario.model(model=f"gemini/{_MODEL}", api_base=api_base, api_key="scripted-gemini-key") + + +def _counted_contents(count_body: Mapping[str, JsonValue]) -> JsonValue: + if "contents" in count_body: + return count_body["contents"] + return _JSON_OBJECT.validate_python(count_body["generateContentRequest"])["contents"] + + +@pytest.mark.parametrize("provider", ("vertex_ai", "gemini")) +@pytest.mark.parametrize("case", sorted(_CONTENT_BLOCK_BODIES)) +def test_count_tokens_sends_the_contents_messages_inference_sends(gateway: Gateway, provider: str, case: str) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = ( + _deployment(gateway, scenario, wire.url) + if provider == "vertex_ai" + else _ai_studio_deployment(scenario, wire.url) + ) + body: Final = {"model": model, **_CONTENT_BLOCK_BODIES[case]} + generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16}) + counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body) + assert generated.status_code == 200, generated.text + assert counted.status_code == 200, counted.text + assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == _COUNT, counted.text + requests: Final = peer.drain() + generated_body: Final = _json_body( + next(request for request in requests if request.target.endswith(":generateContent")) + ) + count_body: Final = _json_body(next(request for request in requests if request.target.endswith(":countTokens"))) + assert _counted_contents(count_body) == generated_body["contents"] + + +_UNCONVERTIBLE_MESSAGES: Final[JsonValue] = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the notes."}, + {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "notes"}}, + {"type": "search_result", "source": "https://example.com/notes", "title": "Notes", "content": []}, + ], + }, + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "search", "input": {}}]}, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_01", + "content": [ + { + "type": "search_result", + "source": "https://example.com/hits", + "title": "Hits", + "content": [{"type": "text", "text": "hit"}], + } + ], + } + ], + }, +] + + +@pytest.mark.parametrize("provider", ("vertex_ai", "gemini")) +def test_count_tokens_rejects_blocks_it_cannot_convert(gateway: Gateway, provider: str) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = ( + _deployment(gateway, scenario, wire.url) + if provider == "vertex_ai" + else _ai_studio_deployment(scenario, wire.url) + ) + response: Final = gateway.request( + "POST", "/v1/messages/count_tokens", {"model": model, "messages": _UNCONVERTIBLE_MESSAGES} + ) + assert response.status_code == 400, response.text + assert "document with a text source, search_result, search_result inside a tool_result" in response.text + assert peer.drain() == () + + +def test_count_tokens_keeps_each_text_block_in_system_instruction(gateway: Gateway) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + system: Final[JsonValue] = [ + {"type": "text", "text": "First instruction"}, + {"type": "text", "text": "Second instruction"}, + ] + response: Final = gateway.request("POST", "/v1/messages/count_tokens", _messages(model, system, None)) + assert response.status_code == 200, response.text + requests: Final = _count_requests(wire) + assert len(requests) == 1 + assert _json_body(requests[0]) == { + "contents": [{"role": "user", "parts": [{"text": _USER}]}], + "systemInstruction": {"parts": [{"text": "First instruction"}, {"text": "Second instruction"}]}, + } + + +def test_count_tokens_without_system_or_tools_sends_contents_only(gateway: Gateway) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = gateway.request("POST", "/v1/messages/count_tokens", _messages(model, None, None)) + assert response.status_code == 200, response.text + requests: Final = _count_requests(wire) + assert len(requests) == 1 + assert _json_body(requests[0]) == {"contents": [{"role": "user", "parts": [{"text": _USER}]}]} + + +_SDK_COUNT_ARGS: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "messages": [{"role": "user", "content": _USER}], + "system": "Follow the system instruction.", + "tools": [ + { + "name": "noop", + "description": "Return no output", + "input_schema": {"type": "object", "properties": {}}, + } + ], + } +) + + +def _sync_sdk_count_tokens(gateway: Gateway, model: str) -> int: + client: Final = anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + return client.messages.count_tokens(model=model, **_SDK_COUNT_ARGS).input_tokens # pyright: ignore[reportArgumentType] # SDK TypedDicts reject a JsonValue mapping + + +async def _async_sdk_count_tokens(gateway: Gateway, model: str) -> int: + client: Final = anthropic.AsyncAnthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + return (await client.messages.count_tokens(model=model, **_SDK_COUNT_ARGS)).input_tokens # pyright: ignore[reportArgumentType] # SDK TypedDicts reject a JsonValue mapping + + +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync-sdk", "async-sdk")) +def test_anthropic_sdk_count_tokens_forwards_system_and_tools(gateway: Gateway, asynchronous: bool) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + input_tokens: Final = ( + asyncio.run(_async_sdk_count_tokens(gateway, model)) + if asynchronous + else _sync_sdk_count_tokens(gateway, model) + ) + assert input_tokens == _COUNT + requests: Final = _count_requests(wire) + assert len(requests) == 1 + request_body: Final = _json_body(requests[0]) + assert request_body["systemInstruction"] == {"parts": [{"text": "Follow the system instruction."}]} + assert request_body["tools"] == [ + { + "functionDeclarations": [ + {"name": "noop", "description": "Return no output", "parameters": {"type": "object"}} + ] + } + ] + + +def test_google_count_tokens_forwards_system_and_tools_to_vertex(gateway: Gateway) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + body: Final[Mapping[str, JsonValue]] = { + "contents": [ + {"role": "user", "parts": [{"text": _USER}]}, + {"role": "model", "parts": [{"functionCall": {"name": "noop", "args": {"query": "value"}}}]}, + { + "role": "user", + "parts": [{"functionResponse": {"name": "noop", "response": {"result": "ok"}}}], + }, + ], + "systemInstruction": {"parts": [{"text": "Follow this instruction."}]}, + "tools": [ + { + "functionDeclarations": [ + { + "name": "noop", + "description": "Return no output", + "parameters": {"type": "OBJECT", "properties": {}}, + } + ] + } + ], + } + response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", body) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["totalTokens"] == _COUNT, response.text + requests: Final = _count_requests(wire) + assert len(requests) == 1 + request_body: Final = _json_body(requests[0]) + assert request_body["contents"] == body["contents"] + assert request_body["systemInstruction"] == body["systemInstruction"] + assert request_body["tools"] == [ + { + "functionDeclarations": [ + { + "name": "noop", + "description": "Return no output", + "parameters": {"type": "object"}, + } + ] + } + ] + + +def test_responses_input_tokens_maps_flat_function_tools_to_vertex(gateway: Gateway) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + body: Final[Mapping[str, JsonValue]] = { + "model": model, + "instructions": "Follow the system instruction.", + "input": [ + { + "type": "function_call", + "call_id": "call_lookup", + "name": "lookup", + "arguments": '{"query":"value"}', + }, + { + "type": "function_call_output", + "call_id": "call_lookup", + "output": "result", + }, + ], + "tools": [ + { + "type": "function", + "name": "lookup", + "description": "Look up a value", + "parameters": {"type": "object", "properties": {}}, + } + ], + } + response: Final = gateway.request("POST", "/v1/responses/input_tokens", body) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["input_tokens"] == _COUNT, response.text + requests: Final = _count_requests(wire) + assert len(requests) == 1 + request_body: Final = _json_body(requests[0]) + assert request_body["tools"] == [ + { + "functionDeclarations": [ + { + "name": "lookup", + "description": "Look up a value", + "parameters": {"type": "object"}, + } + ] + } + ] + + +def test_utils_token_counter_maps_openai_function_tools_to_vertex(gateway: Gateway) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + tool_schema: Final[JsonValue] = {"type": "object", "properties": {}} + body: Final[Mapping[str, JsonValue]] = { + "model": model, + "messages": _MESSAGES, + "system": "Follow the system instruction.", + "tools": [ + { + "type": "function", + "function": { + "name": "noop", + "description": "Return no output", + "parameters": tool_schema, + }, + } + ], + } + response: Final = gateway.request("POST", "/utils/token_counter", body, params={"call_endpoint": "true"}) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["total_tokens"] == _COUNT, response.text + requests: Final = _count_requests(wire) + assert len(requests) == 1 + request_body: Final = _json_body(requests[0]) + assert request_body["systemInstruction"] == {"parts": [{"text": "Follow the system instruction."}]} + assert request_body["tools"] == [ + { + "functionDeclarations": [ + { + "name": "noop", + "description": "Return no output", + "parameters": {"type": "object"}, + } + ] + } + ] + + +def test_count_tokens_surfaces_peer_errors_and_rejects_non_text_system(gateway: Gateway) -> None: + peer: Final = _GeminiPeer() + peer.set_status(400) + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = gateway.request("POST", "/v1/messages/count_tokens", _messages(model, None, None)) + assert response.status_code == 400, response.text + assert len(_count_requests(wire)) == 1 + + invalid_peer: Final = _GeminiPeer() + with wire_server(invalid_peer) as invalid_wire, gateway.scenario() as scenario: + invalid_model: Final = _deployment(gateway, scenario, invalid_wire.url) + invalid_response: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + _messages( + invalid_model, + [{"type": "image", "source": {"type": "base64", "data": "AAAA"}}], + None, + ), + ) + assert invalid_response.status_code == 400, invalid_response.text + assert _count_requests(invalid_wire) == () + + +@pytest.mark.parametrize( + "tools", + ( + [{"name": "noop", "input_schema": {"type": "object", "properties": {"query": "not-a-schema"}}}], + [{"name": "noop", "input_schema": "not-a-schema"}], + ), + ids=("malformed-property", "string-schema"), +) +def test_count_tokens_rejects_unconvertible_tools_without_request(gateway: Gateway, tools: JsonValue) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = gateway.request("POST", "/v1/messages/count_tokens", _messages(model, None, tools)) + assert response.status_code == 400, response.text + assert "tools" in response.text + assert _count_requests(wire) == () + + +def _count_wave(clients: Sequence[httpx.Client], body: Mapping[str, JsonValue]) -> tuple[tuple[int, JsonValue], ...]: + def count(client: httpx.Client) -> tuple[int, JsonValue]: + response: Final = client.post("/v1/messages/count_tokens", json=dict(body)) + payload: Final = _JSON_OBJECT.validate_json(response.content) + return response.status_code, payload.get("input_tokens") + + with ThreadPoolExecutor(max_workers=len(clients)) as executor: + return tuple(executor.map(count, clients)) + + +def test_concurrent_count_waves_recover_after_vertex_count_outage(gateway: Gateway) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario, ExitStack() as stack: + model: Final = _deployment(gateway, scenario, wire.url) + clients: Final = tuple( + stack.enter_context( + httpx.Client( + base_url=str(gateway.client.base_url), + headers={"Authorization": f"Bearer {gateway.key}"}, + timeout=30, + trust_env=False, + ) + ) + for _ in range(3) + ) + body: Final = _messages(model, None, None) + assert _count_wave(clients, body) == ((200, _COUNT), (200, _COUNT), (200, _COUNT)) + peer.set_status(503) + assert _count_wave(clients, body) == ((503, None), (503, None), (503, None)) + peer.set_status(200) + assert _count_wave(clients, body) == ((200, _COUNT), (200, _COUNT), (200, _COUNT)) + requests: Final = _count_requests(wire) + assert len(requests) == 9 + assert ( + tuple(_json_body(request) for request in requests) + == ({"contents": [{"role": "user", "parts": [{"text": _USER}]}]},) * 9 + ) diff --git a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py new file mode 100644 index 00000000000..73c135944e0 --- /dev/null +++ b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py @@ -0,0 +1,858 @@ +import json +import re +import signal +import socket +import threading +import uuid +from collections.abc import Callable, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final + +import httpx +import psutil +import pytest +import yaml +from integration._support.client import Gateway, Scenario, eventually +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter +from tokenizers import Regex, Tokenizer, models, pre_tokenizers + +_BACKEND: Final = "claude-sonnet-4-6" +_PROJECT: Final = "scripted-project" +_LOCATION: Final = "us-east5" +_MODELS_PATH: Final = f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/anthropic/models" +_COUNT_TARGET: Final = f"{_MODELS_PATH}/count-tokens:rawPredict" +_MESSAGE_TARGET: Final = f"{_MODELS_PATH}/{_BACKEND}:rawPredict" +_STREAM_TARGET: Final = f"{_MODELS_PATH}/{_BACKEND}:streamRawPredict?alt=sse" +_PEER_COUNT: Final = 4242 +_REJECTION: Final = "scripted partner rejection" +_REJECT_TEXT: Final = "The peer must reject this message" +_REPLY_TEXT: Final = "scripted reply" +_OWNED_MODEL: Final = "partner-claude" +_OWNED_UNREACHABLE_MODEL: Final = "partner-claude-unreachable" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_GLM_BACKEND: Final = "zai-org/glm-5.2-maas" +_MISTRAL_BACKEND: Final = "mistral-medium-3" +_LLAMA_BACKEND: Final = "meta/llama-4-maverick-17b-128e-instruct-maas" +_LLAMA_COUNT_TARGET: Final = ( + f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/meta/models/count-tokens:rawPredict" +) +_COUNTED_TEXT: Final = "Count this message" +_TOKENIZER_REPO: Final = "scripted/char-tokenizer" +_TOKENIZER_TARGET: Final = f"/{_TOKENIZER_REPO}/resolve/main/tokenizer.json" +_TOKENIZER_HEADERS: Final = MappingProxyType({"X-Repo-Commit": "0" * 40, "ETag": '"scripted-char-tokenizer"'}) +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") + +_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": _COUNTED_TEXT}] +_SYSTEM: Final = "You are a terse assistant that answers in one sentence" +_SYSTEM_BLOCKS: Final[list[JsonValue]] = [ + {"type": "text", "text": "You are a terse assistant"}, + {"type": "text", "text": "Answer in one sentence"}, +] +_WEATHER_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"city": {"type": "string", "description": "City to look up"}}, + "required": ["city"], +} +_TOOLS: Final[list[JsonValue]] = [ + {"name": "get_weather", "description": "Look up the current weather for a city", "input_schema": _WEATHER_SCHEMA} +] +_OPENAI_TOOLS: Final[list[JsonValue]] = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Look up the current weather for a city", + "parameters": _WEATHER_SCHEMA, + }, + } +] +_RESPONSES_TOOLS: Final[list[JsonValue]] = [ + { + "type": "function", + "name": "get_weather", + "description": "Look up the current weather for a city", + "parameters": _WEATHER_SCHEMA, + } +] +_PEER_BARE: Final[dict[str, JsonValue]] = {"model": _BACKEND, "messages": _MESSAGES} +_PEER_FULL: Final[dict[str, JsonValue]] = {**_PEER_BARE, "system": _SYSTEM, "tools": _TOOLS} +_GEMINI_BODY: Final[dict[str, JsonValue]] = {"contents": [{"role": "user", "parts": [{"text": "Count this"}]}]} +_GEMINI_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": "Count this"}] + +_REPLY: Final[dict[str, JsonValue]] = { + "id": "msg_scripted", + "type": "message", + "role": "assistant", + "model": _BACKEND, + "content": [{"type": "text", "text": _REPLY_TEXT}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 3}, +} +_EVENTS: Final[tuple[tuple[str, dict[str, JsonValue]], ...]] = ( + ( + "message_start", + { + "type": "message_start", + "message": {**_REPLY, "content": [], "stop_reason": None, "usage": {"input_tokens": 5, "output_tokens": 1}}, + }, + ), + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + ( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": _REPLY_TEXT}}, + ), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 3}, + }, + ), + ("message_stop", {"type": "message_stop"}), +) +_SSE: Final = tuple(f"event: {name}\ndata: {json.dumps(data)}\n\n".encode() for name, data in _EVENTS) + + +def _counted(_request: Request) -> Reply: + return Reply(body=json.dumps({"input_tokens": _PEER_COUNT}).encode()) + + +def _rejected(status: int, error_type: str = "invalid_request_error") -> Reply: + return Reply( + status=status, + body=json.dumps({"type": "error", "error": {"type": error_type, "message": _REJECTION}}).encode(), + ) + + +def _rejecting(status: int, error_type: str = "invalid_request_error") -> Callable[[Request], Reply]: + def count(_request: Request) -> Reply: + return _rejected(status, error_type) + + return count + + +def _anthropic_message(message: JsonValue) -> bool: + return isinstance(message, dict) and message.get("role") in ("user", "assistant") + + +def _anthropic_tool(tool: JsonValue) -> bool: + return isinstance(tool, dict) and isinstance(tool.get("name"), str) and isinstance(tool.get("input_schema"), dict) + + +def _strict(request: Request) -> Reply: + body: Final = _JSON_OBJECT.validate_json(request.body) + messages: Final = body.get("messages") + tools: Final = body.get("tools", []) + accepted: Final = ( + isinstance(messages, list) + and all(map(_anthropic_message, messages)) + and isinstance(body.get("system", ""), (str, list)) + and isinstance(tools, list) + and all(map(_anthropic_tool, tools)) + ) + return _counted(request) if accepted else _rejected(400) + + +def _rejecting_marked_messages(request: Request) -> Reply: + return _rejected(400) if _REJECT_TEXT in request.body.decode() else _counted(request) + + +def _peer(count: Callable[[Request], Reply] = _counted) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.target.endswith("/count-tokens:rawPredict"): + return count(request) + if request.target.endswith(":streamRawPredict?alt=sse"): + return Reply(content_type="text/event-stream", chunks=_SSE) + if request.target.endswith(f"/{_BACKEND}:rawPredict"): + return Reply(body=json.dumps(_REPLY).encode()) + return Reply(status=404, body=json.dumps({"error": f"unscripted target {request.target}"}).encode()) + + return respond + + +def _count_requests(requests: Sequence[Request]) -> tuple[Request, ...]: + return tuple(request for request in requests if "count-tokens" in request.target) + + +def _count_bodies(requests: Sequence[Request], target: str = _COUNT_TARGET) -> tuple[dict[str, JsonValue], ...]: + counts: Final = _count_requests(requests) + for request in counts: + assert (request.method, request.target) == ("POST", target), request.target + assert request.headers["authorization"] == "Bearer scripted-token", request.headers + return tuple(_JSON_OBJECT.validate_json(request.body) for request in counts) + + +def _counted_bodies(wire: Wire) -> tuple[dict[str, JsonValue], ...]: + return _count_bodies(wire.drain()) + + +def _bare(model: str) -> dict[str, JsonValue]: + return {"model": model, "messages": _MESSAGES} + + +def _full(model: str) -> dict[str, JsonValue]: + return {**_bare(model), "system": _SYSTEM, "tools": _TOOLS} + + +def _deployment(gateway: Gateway, scenario: Scenario, api_base: str, **overrides: JsonValue) -> str: + return scenario.model( + **{ + "model": f"vertex_ai/{_BACKEND}", + "api_base": api_base, + "api_key": None, + "vertex_project": _PROJECT, + "vertex_location": _LOCATION, + "vertex_credentials": service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")), + **overrides, + } + ) + + +def _count(gateway: Gateway, body: Mapping[str, JsonValue]) -> httpx.Response: + return gateway.request("POST", "/v1/messages/count_tokens", body) + + +def _payload(response: httpx.Response) -> dict[str, JsonValue]: + assert response.status_code == 200, response.text + return _JSON_OBJECT.validate_json(response.content) + + +def _local_count(gateway: Gateway, body: Mapping[str, JsonValue]) -> int: + response: Final = gateway.request("POST", "/utils/token_counter", body, params={"call_endpoint": "false"}) + payload: Final = _payload(response) + total: Final = payload["total_tokens"] + assert payload["tokenizer_type"] != "vertex_ai_partner_models", response.text + assert isinstance(total, int) and 0 < total != _PEER_COUNT, response.text + return total + + +@dataclass(frozen=True, slots=True) +class _CountRoute: + path: Callable[[str], str] + body: Callable[[str], dict[str, JsonValue]] + rejection: Callable[[str], dict[str, JsonValue]] + count_field: str + params: Mapping[str, str] = MappingProxyType({}) + + def send(self, gateway: Gateway, model: str) -> httpx.Response: + return gateway.request("POST", self.path(model), self.body(model), params=self.params) + + +_MESSAGES_ROUTE: Final = _CountRoute( + path=lambda _model: "/v1/messages/count_tokens", + body=_bare, + rejection=lambda message: {"type": "error", "error": {"type": "invalid_request_error", "message": message}}, + count_field="input_tokens", +) +_RESPONSES_ROUTE: Final = _CountRoute( + path=lambda _model: "/v1/responses/input_tokens", + body=lambda model: {"model": model, "input": _COUNTED_TEXT}, + rejection=lambda message: { + "error": {"message": message, "type": "invalid_request_error", "param": None, "code": None} + }, + count_field="input_tokens", +) +_GOOGLE_ROUTE: Final = _CountRoute( + path=lambda model: f"/v1beta/models/{model}:countTokens", + body=lambda _model: {"contents": [{"role": "user", "parts": [{"text": _COUNTED_TEXT}]}]}, + rejection=lambda message: {"error": {"code": 400, "message": message, "status": "INVALID_ARGUMENT"}}, + count_field="totalTokens", +) +_UTILS_ROUTE: Final = _CountRoute( + path=lambda _model: "/utils/token_counter", + body=_bare, + rejection=lambda message: { + "error": {"message": message, "type": "token_counting_error", "param": "model", "code": "400"} + }, + count_field="total_tokens", + params=MappingProxyType({"call_endpoint": "true"}), +) +_COUNT_ROUTES: Final = (_MESSAGES_ROUTE, _RESPONSES_ROUTE, _GOOGLE_ROUTE, _UTILS_ROUTE) + + +def _error_message(payload: Mapping[str, JsonValue]) -> str: + error: Final = payload.get("error") + assert isinstance(error, dict), payload + message: Final = error.get("message") + assert isinstance(message, str), payload + return message + + +def _route_count(gateway: Gateway, route: _CountRoute, model: str) -> tuple[int, JsonValue]: + response: Final = route.send(gateway, model) + return response.status_code, _JSON_OBJECT.validate_json(response.content).get(route.count_field) + + +def _char_tokenizer() -> bytes: + tokenizer: Final = Tokenizer(models.WordLevel({"[UNK]": 0}, unk_token="[UNK]")) + tokenizer.pre_tokenizer = pre_tokenizers.Split(Regex("."), behavior="isolated") + return tokenizer.to_str().encode() + + +def _tokenizer_hub(tokenizer: bytes) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.target == _TOKENIZER_TARGET: + return Reply(body=tokenizer, headers=_TOKENIZER_HEADERS) + return Reply(status=404, body=json.dumps({"error": f"unscripted target {request.target}"}).encode()) + + return respond + + +def _closed_port_url() -> str: + with socket.socket() as reserve: + reserve.bind(("127.0.0.1", 0)) + return f"http://127.0.0.1:{reserve.getsockname()[1]}" + + +def _clients(stack: ExitStack, base_url: str, count: int) -> tuple[httpx.Client, ...]: + return tuple( + stack.enter_context(httpx.Client(base_url=base_url, timeout=30, trust_env=False)) for _ in range(count) + ) + + +def _counted_on(client: httpx.Client, key: str, body: Mapping[str, JsonValue]) -> tuple[int, JsonValue]: + response: Final = client.post( + "/v1/messages/count_tokens", json=dict(body), headers={"Authorization": f"Bearer {key}"} + ) + return response.status_code, _JSON_OBJECT.validate_json(response.content).get("input_tokens") + + +def _generated_then_counted( + client: httpx.Client, key: str, model: str, body: Mapping[str, JsonValue] +) -> tuple[int, int, JsonValue]: + generated: Final = client.post( + "/v1/messages", + json={ + "model": model, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"Generate before counting {uuid.uuid4().hex}"}], + }, + headers={"Authorization": f"Bearer {key}"}, + ) + return generated.status_code, *_counted_on(client, key, body) + + +def _local_port(client: httpx.Client) -> int: + with client.stream("GET", "/health/liveliness") as response: + port: Final = int(response.extensions["network_stream"].get_extra_info("client_addr")[1]) + response.read() + assert response.status_code == 200, response.text + return port + + +def _counted_or_dropped(client: httpx.Client, key: str, body: Mapping[str, JsonValue]) -> tuple[int, JsonValue] | None: + try: + return _counted_on(client, key, body) + except httpx.TransportError: + return None + + +def _accepted_client_ports(pid: int, proxy_port: int) -> frozenset[int]: + return frozenset( + connection.raddr.port + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.raddr and connection.laddr.port == proxy_port + ) + + +def _owned_config( + path: Path, gateway: Gateway, api_bases: Mapping[str, str], settings: Mapping[str, JsonValue] +) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + path.write_text( + yaml.safe_dump( + { + **config, + "model_list": [ + { + "model_name": name, + "litellm_params": { + "model": f"vertex_ai/{_BACKEND}", + "api_base": api_base, + "vertex_project": _PROJECT, + "vertex_location": _LOCATION, + "vertex_credentials": service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")), + }, + } + for name, api_base in api_bases.items() + ], + "litellm_settings": {**config["litellm_settings"], **settings}, + } + ) + ) + return path + + +@pytest.mark.parametrize("system", [_SYSTEM, _SYSTEM_BLOCKS], ids=["string", "blocks"]) +def test_messages_count_tokens_forwards_system_and_tools(gateway: Gateway, system: JsonValue) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS}) + assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text + assert _counted_bodies(wire) == ({**_PEER_BARE, "system": system, "tools": _TOOLS},) + + +def test_messages_count_tokens_without_system_or_tools_sends_bare_body(gateway: Gateway) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = _count(gateway, _bare(model)) + assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text + assert _counted_bodies(wire) == (_PEER_BARE,) + + +def test_utils_token_counter_call_endpoint_forwards_system_and_tools(gateway: Gateway) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = gateway.request( + "POST", "/utils/token_counter", _full(model), params={"call_endpoint": "true"} + ) + payload: Final = _payload(response) + assert (payload["total_tokens"], payload["tokenizer_type"]) == (_PEER_COUNT, "vertex_ai_partner_models") + assert (payload["request_model"], payload["model_used"]) == (model, _BACKEND), response.text + assert _counted_bodies(wire) == (_PEER_FULL,) + + +def test_utils_token_counter_local_mode_never_calls_the_peer(gateway: Gateway) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + assert _local_count(gateway, _full(model)) > 0 + assert wire.drain() == () + + +def test_utils_token_counter_surfaces_the_peer_error_for_openai_tools(gateway: Gateway) -> None: + with wire_server(_peer(_strict)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + body: Final = {**_bare(model), "tools": _OPENAI_TOOLS} + response: Final = gateway.request("POST", "/utils/token_counter", body, params={"call_endpoint": "true"}) + assert response.status_code == 400, response.text + assert "total_tokens" not in response.text, response.text + assert _counted_bodies(wire) == ({**_PEER_BARE, "tools": _OPENAI_TOOLS},) + + +def test_responses_input_tokens_counts_through_the_partner_peer(gateway: Gateway) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = gateway.request( + "POST", "/v1/responses/input_tokens", {"model": model, "input": "Count this message"} + ) + assert _payload(response) == { + "object": "response.input_tokens", + "input_tokens": _PEER_COUNT, + "litellm_count_source": "provider", + }, response.text + assert _counted_bodies(wire) == (_PEER_BARE,) + + +def test_responses_input_tokens_surfaces_the_peer_error(gateway: Gateway) -> None: + with wire_server(_peer(_strict)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = gateway.request( + "POST", + "/v1/responses/input_tokens", + {"model": model, "input": "Count this message", "instructions": "Be terse", "tools": _RESPONSES_TOOLS}, + ) + assert response.status_code == 400, response.text + assert response.json() == _RESPONSES_ROUTE.rejection(_REJECTION), response.text + (sent,) = _counted_bodies(wire) + assert sent["tools"] == _RESPONSES_TOOLS, sent + + + + + + + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_custom_tokenizer_partner_counts_locally_on_every_count_route(gateway: Gateway, tmp_path: Path) -> None: + environment: Final = {"HF_HOME": str(tmp_path / "hf-home")} + with ( + wire_server(_tokenizer_hub(_char_tokenizer())) as hub, + wire_server(_peer()) as vertex, + owned_proxy_process(gateway, tmp_path, {**environment, "HF_ENDPOINT": hub.url}) as owned, + owned.gateway.scenario() as scenario, + ): + model: Final = _deployment( + owned.gateway, + scenario, + vertex.url, + model=f"vertex_ai/{_GLM_BACKEND}", + model_info={"custom_tokenizer": {"identifier": _TOKENIZER_REPO, "revision": "main", "auth_token": None}}, + ) + local: Final = _payload( + owned.gateway.request("POST", "/utils/token_counter", _bare(model), params={"call_endpoint": "false"}) + ) + assert local["tokenizer_type"] == "huggingface_tokenizer", local + total: Final = local["total_tokens"] + assert isinstance(total, int) and total >= len(_COUNTED_TEXT), local + counts: Final = tuple(_route_count(owned.gateway, route, model) for route in _COUNT_ROUTES) + assert counts == ((200, total),) * len(_COUNT_ROUTES), counts + contents_only: Final = owned.gateway.request( + "POST", + "/utils/token_counter", + {"model": model, "contents": [{"role": "user", "parts": [{"text": _COUNTED_TEXT}]}]}, + params={"call_endpoint": "true"}, + ) + assert contents_only.status_code == 400, contents_only.text + assert _GLM_BACKEND in _error_message(_JSON_OBJECT.validate_json(contents_only.content)), contents_only.text + assert vertex.drain() == () + assert {request.target for request in hub.drain()} >= {_TOKENIZER_TARGET} + + +def test_gemini_count_tokens_route_reaches_the_partner_peer(gateway: Gateway) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", _GEMINI_BODY) + assert "totalTokens" in _payload(response), response.text + assert _counted_bodies(wire) == ({"model": _BACKEND, "messages": _GEMINI_MESSAGES},) + + +def test_gemini_count_tokens_route_surfaces_the_peer_error(gateway: Gateway) -> None: + with wire_server(_peer(_rejecting(400))) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", _GEMINI_BODY) + assert response.status_code == 400, response.text + assert "totalTokens" not in response.text, response.text + assert _counted_bodies(wire) == ({"model": _BACKEND, "messages": _GEMINI_MESSAGES},) + + +@pytest.mark.parametrize("status", [400, 500, 503]) +def test_messages_count_tokens_surfaces_the_peer_error_when_peer_errors(gateway: Gateway, status: int) -> None: + with wire_server(_peer(_rejecting(status))) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = _count(gateway, _full(model)) + assert response.status_code == status, response.text + assert "input_tokens" not in response.text, response.text + assert _counted_bodies(wire) == (_PEER_FULL,) + + +def test_messages_count_tokens_surfaces_the_error_when_token_endpoint_rejects(gateway: Gateway) -> None: + def respond(request: Request) -> Reply: + if request.target == "/_oauth/token": + return Reply( + status=400, + body=json.dumps({"error": "invalid_grant", "error_description": "scripted refusal"}).encode(), + ) + return _peer()(request) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _deployment( + gateway, scenario, wire.url, vertex_credentials=service_account_json(_PROJECT, wire.url) + ) + response: Final = _count(gateway, _full(model)) + assert 500 <= response.status_code < 600, response.text + assert "input_tokens" not in response.text, response.text + targets: Final = frozenset(request.target for request in wire.drain()) + assert targets == {"/_oauth/token"}, targets + + +def test_messages_count_tokens_errors_when_peer_is_unreachable(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, _closed_port_url()) + response: Final = _count(gateway, _full(model)) + assert 500 <= response.status_code < 600, response.text + assert "input_tokens" not in response.text, response.text + + +@pytest.mark.parametrize( + "tools", + [5, "", "x" * 5120, ["get_weather"]], + ids=["int", "empty_string", "5kb_string", "list_of_strings"], +) +def test_messages_count_tokens_rejects_malformed_tools_without_calling_the_peer( + gateway: Gateway, tools: JsonValue +) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + refused: Final = _count(gateway, {**_bare(model), "tools": tools}) + assert 400 <= refused.status_code < 600, refused.text + assert "input_tokens" not in refused.text, refused.text + assert wire.drain() == () + assert _generated_then_counted(gateway.client, gateway.key, model, _bare(model)) == (200, 200, _PEER_COUNT) + assert _counted_bodies(wire) == (_PEER_BARE,) + + +def test_messages_count_tokens_forwards_an_empty_tools_list(gateway: Gateway) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = _count(gateway, {**_bare(model), "tools": []}) + assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text + assert _counted_bodies(wire) == ({**_PEER_BARE, "tools": []},) + + +def test_messages_count_tokens_forwards_an_empty_system_string(gateway: Gateway) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = _count(gateway, {**_bare(model), "system": ""}) + assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text + assert _counted_bodies(wire) == ({**_PEER_BARE, "system": ""},) + + +def test_messages_count_tokens_leaves_null_system_and_tools_out(gateway: Gateway) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = _count(gateway, {**_bare(model), "system": None, "tools": None}) + assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text + assert _counted_bodies(wire) == (_PEER_BARE,) + + +def test_messages_count_tokens_surfaces_the_peer_error_for_a_non_text_system(gateway: Gateway) -> None: + with wire_server(_peer(_strict)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = _count(gateway, {**_bare(model), "system": 5}) + assert response.status_code == 400, response.text + assert "input_tokens" not in response.text, response.text + assert _counted_bodies(wire) == ({**_PEER_BARE, "system": 5},) + + +def test_messages_count_tokens_forwards_a_5kb_system_verbatim(gateway: Gateway) -> None: + system: Final = "Answer in one sentence. " * 214 + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = _count(gateway, {**_bare(model), "system": system}) + assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text + assert _counted_bodies(wire) == ({**_PEER_BARE, "system": system},) + + +def test_messages_count_tokens_duplicate_system_and_tools_keys_forward_one_value_each(gateway: Gateway) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + fields: Final = f'"system": {json.dumps(_SYSTEM)}, "tools": {json.dumps(_TOOLS)}' + response: Final = gateway.client.post( + "/v1/messages/count_tokens", + content=f'{{"model": "{model}", "messages": {json.dumps(_MESSAGES)}, {fields}, {fields}}}', + headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"}, + ) + assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text + (sent,) = _count_requests(wire.drain()) + assert _count_bodies((sent,)) == (_PEER_FULL,) + assert (sent.body.count(b'"system"'), sent.body.count(b'"tools"')) == (1, 1), sent.body + + +def test_messages_count_tokens_unauthenticated_request_never_reaches_the_peer(gateway: Gateway) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = gateway.request( + "POST", "/v1/messages/count_tokens", _full(model), key="sk-not-a-key-this-proxy-issued" + ) + assert response.status_code == 401, response.text + assert "input_tokens" not in response.text, response.text + assert wire.drain() == () + + +@pytest.mark.parametrize("fields", [{}, {"messages": []}], ids=["missing", "empty"]) +def test_messages_count_tokens_without_messages_is_rejected(gateway: Gateway, fields: dict[str, JsonValue]) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = _count(gateway, {"model": model, "system": _SYSTEM, "tools": _TOOLS, **fields}) + assert response.status_code == 400, response.text + assert "messages parameter is required" in response.text, response.text + assert wire.drain() == () + + +def test_count_tokens_location_override_targets_the_count_region(gateway: Gateway) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment( + gateway, scenario, wire.url, vertex_location="global", vertex_count_tokens_location="europe-west1" + ) + response: Final = _count(gateway, _bare(model)) + assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text + target: Final = _COUNT_TARGET.replace(f"/locations/{_LOCATION}/", "/locations/europe-west1/") + assert _count_bodies(wire.drain(), target) == (_PEER_BARE,) + + +def test_messages_count_tokens_repeated_request_reaches_the_peer_each_time(gateway: Gateway) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + answers: Final = tuple(_payload(_count(gateway, _bare(model))) for _ in range(2)) + assert answers == ({"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"},) * 2 + assert _counted_bodies(wire) == (_PEER_BARE,) * 2 + + +@pytest.mark.timeout(240) # boots an owned two-worker proxy with litellm_settings.disable_token_counter +def test_disabled_token_counter_surfaces_provider_failures_instead_of_counting_locally( + gateway: Gateway, tmp_path: Path +) -> None: + with wire_server(_peer(_rejecting_marked_messages)) as wire: + config: Final = _owned_config( + tmp_path / "disabled-token-counter.yaml", + gateway, + {_OWNED_MODEL: wire.url, _OWNED_UNREACHABLE_MODEL: _closed_port_url()}, + {"disable_token_counter": True}, + ) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: + counted: Final = _count(owned.gateway, _full(_OWNED_MODEL)) + assert _payload(counted) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, counted.text + rejected: Final = _count( + owned.gateway, {**_full(_OWNED_MODEL), "messages": [{"role": "user", "content": _REJECT_TEXT}]} + ) + assert rejected.status_code == 400, rejected.text + assert _REJECTION in rejected.text and "input_tokens" not in rejected.text, rejected.text + unreachable: Final = _count(owned.gateway, _full(_OWNED_UNREACHABLE_MODEL)) + assert 500 <= unreachable.status_code < 600, unreachable.text + assert "input_tokens" not in unreachable.text, unreachable.text + local: Final = owned.gateway.request( + "POST", "/utils/token_counter", _full(_OWNED_MODEL), params={"call_endpoint": "false"} + ) + assert local.status_code == 503, local.text + assert len(_counted_bodies(wire)) == 2 + + +def test_peer_outage_between_concurrent_waves_errors_then_recovers(gateway: Gateway) -> None: + with ExitStack() as stack: + clients: Final = _clients(stack, str(gateway.client.base_url), 8) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) + scenario: Final = stack.enter_context(gateway.scenario()) + with wire_server(_peer()) as wire: + port: Final = int(wire.url.rsplit(":", 1)[1]) + model: Final = _deployment(gateway, scenario, wire.url) + body: Final = _full(model) + + def generate_then_count(client: httpx.Client) -> tuple[int, int, JsonValue]: + return _generated_then_counted(client, gateway.key, model, body) + + assert tuple(pool.map(generate_then_count, clients)) == ((200, 200, _PEER_COUNT),) * len(clients) + assert _counted_bodies(wire) == (_PEER_FULL,) * len(clients) + outage: Final = tuple(pool.map(lambda client: _counted_on(client, gateway.key, body), clients)) + for status, tokens in outage: + assert 500 <= status < 600 and tokens is None, outage + with wire_server(_peer(), port=port) as revived: + assert tuple(pool.map(generate_then_count, clients)) == ((200, 200, _PEER_COUNT),) * len(clients) + assert _counted_bodies(revived) == (_PEER_FULL,) * len(clients) + + +def test_slow_peer_holds_concurrent_counts_without_stalling_the_proxy(gateway: Gateway) -> None: + held: Final[SimpleQueue[str]] = SimpleQueue() + release: Final = threading.Event() + + def hold(request: Request) -> Reply: + held.put(request.target) + assert release.wait(timeout=20), "Held count was never released" + return _counted(request) + + with ExitStack() as stack: + clients: Final = _clients(stack, str(gateway.client.base_url), 6) + wire: Final = stack.enter_context(wire_server(_peer(hold))) + scenario: Final = stack.enter_context(gateway.scenario()) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) + model: Final = _deployment(gateway, scenario, wire.url) + try: + futures: Final = tuple( + pool.submit(_generated_then_counted, client, gateway.key, model, _full(model)) for client in clients + ) + eventually(held.qsize, lambda size: size == len(clients), seconds=30) + assert gateway.request("GET", "/health/liveliness").status_code == 200 + assert _local_count(gateway, _full(model)) > 0 + assert not any(future.done() for future in futures) + finally: + release.set() + assert tuple(future.result(timeout=30) for future in futures) == ((200, 200, _PEER_COUNT),) * len(clients) + assert len(_counted_bodies(wire)) == len(clients) + + +@pytest.mark.timeout(300) # boots an owned two-worker proxy, kills one worker, and waits for its replacement +def test_worker_sigkill_mid_burst_leaves_the_sibling_counting(gateway: Gateway, tmp_path: Path) -> None: + held: Final[SimpleQueue[str]] = SimpleQueue() + release: Final = threading.Event() + + def hold(request: Request) -> Reply: + held.put(request.target) + assert release.wait(timeout=60), "Held count was never released" + return _counted(request) + + with ExitStack() as stack: + wire: Final = stack.enter_context(wire_server(_peer(hold))) + stack.callback(release.set) + config: Final = _owned_config(tmp_path / "worker-kill.yaml", gateway, {_OWNED_MODEL: wire.url}, {}) + owned: Final = stack.enter_context(owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2)) + proxy_url: Final = owned.gateway.client.base_url + workers: Final = eventually( + lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())), + lambda pids: len(pids) == 2, + seconds=30, + ) + clients: Final = _clients(stack, str(proxy_url), 12) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) + stack.callback(release.set) + ports: Final = tuple(_local_port(client) for client in clients) + futures: Final = tuple( + pool.submit(_counted_or_dropped, client, gateway.key, _full(_OWNED_MODEL)) for client in clients + ) + eventually(held.qsize, lambda size: size == len(clients), seconds=30) + shares: Final = {pid: _accepted_client_ports(pid, proxy_url.port or 0) & frozenset(ports) for pid in workers} + assert sum(map(len, shares.values())) == len(clients), shares + victim: Final = min((pid for pid in workers if shares[pid]), key=lambda pid: len(shares[pid])) + psutil.Process(victim).send_signal(signal.SIGKILL) + release.set() + results: Final = tuple(future.result(timeout=60) for future in futures) + for port, result in zip(ports, results, strict=True): + assert result == (None if port in shares[victim] else (200, _PEER_COUNT)), (port, result, shares) + second_wave: Final = _clients(stack, str(proxy_url), 6) + assert tuple(_counted_on(client, gateway.key, _full(_OWNED_MODEL)) for client in second_wave) == ( + (200, _PEER_COUNT), + ) * len(second_wave) + assert len(_counted_bodies(wire)) == len(clients) + len(second_wave) + eventually(lambda: len(_STARTED_WORKER.findall(owned.log.read_text())), lambda started: started >= 3, 120) + assert owned.process.poll() is None + + +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +def test_chat_completions_on_the_same_deployment_still_generate(gateway: Gateway, stream: bool) -> None: + marker: Final = f"chat control {uuid.uuid4().hex}" + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": marker}], + "stream": stream, + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + assert _REPLY_TEXT in response.text, response.text + assert not stream or response.text.rstrip().endswith("data: [DONE]"), response.text + (sent,) = wire.drain() + assert sent.target == (_STREAM_TARGET if stream else _MESSAGE_TARGET), sent.target + assert marker in sent.body.decode(), sent.body + + +def test_messages_endpoint_on_the_same_deployment_still_generates(gateway: Gateway) -> None: + marker: Final = f"messages control {uuid.uuid4().hex}" + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = gateway.request( + "POST", + "/v1/messages", + {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": marker}]}, + ) + assert response.status_code == 200, response.text + assert _REPLY_TEXT in response.text, response.text + (sent,) = wire.drain() + assert sent.target == _MESSAGE_TARGET, sent.target + assert marker in sent.body.decode(), sent.body + + +def test_responses_endpoint_on_the_same_deployment_still_generates(gateway: Gateway) -> None: + marker: Final = f"responses control {uuid.uuid4().hex}" + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = gateway.request("POST", "/v1/responses", {"model": model, "input": marker}) + assert response.status_code == 200, response.text + assert _REPLY_TEXT in response.text, response.text + (sent,) = wire.drain() + assert sent.target == _MESSAGE_TARGET, sent.target + assert marker in sent.body.decode(), sent.body diff --git a/tests/integration/proxy_config.yaml b/tests/integration/proxy_config.yaml index 8ed5237868a..58312faa1f9 100644 --- a/tests/integration/proxy_config.yaml +++ b/tests/integration/proxy_config.yaml @@ -1,4 +1,11 @@ model_list: + - model_name: anthropic/claude-opus-5-5 + litellm_params: + model: anthropic/claude-opus-5-5 + api_base: http://127.0.0.1:8191 + api_key: synthetic-anthropic-key + model_info: + token_count_framing: 0 - model_name: perplexity/pplx-decider-v1-27b litellm_params: model: perplexity/pplx-decider-v1-27b diff --git a/tests/integration/sdk/test_vertex_partner_count_tokens_sdk.py b/tests/integration/sdk/test_vertex_partner_count_tokens_sdk.py new file mode 100644 index 00000000000..e59d3f9264a --- /dev/null +++ b/tests/integration/sdk/test_vertex_partner_count_tokens_sdk.py @@ -0,0 +1,90 @@ +import json +from collections.abc import Callable +from typing import Final + +import litellm +import pytest +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_BACKEND: Final = "claude-sonnet-4-6" +_PROJECT: Final = "scripted-project" +_LOCATION: Final = "us-east5" +_COUNT_TARGET: Final = ( + f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/anthropic/models/count-tokens:rawPredict" +) +_PEER_COUNT: Final = 4242 +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MESSAGES: Final[list[dict[str, str]]] = [{"role": "user", "content": "Count this message"}] +_SYSTEM: Final = "You are a terse assistant that answers in one sentence" +_TOOLS: Final[list[dict[str, JsonValue]]] = [ + { + "name": "get_weather", + "description": "Look up the current weather for a city", + "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + } +] + + +def _peer(status: int) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.target == "/_oauth/token": + token: Final = {"access_token": "scripted-token", "token_type": "Bearer", "expires_in": 3600} + return Reply(body=json.dumps(token).encode()) + if status == 200: + return Reply(body=json.dumps({"input_tokens": _PEER_COUNT}).encode()) + rejection: Final = {"type": "error", "error": {"type": "invalid_request_error", "message": "scripted"}} + return Reply(status=status, body=json.dumps(rejection).encode()) + + return respond + + +def _count_requests(requests: tuple[Request, ...]) -> tuple[Request, ...]: + return tuple(request for request in requests if "count-tokens" in request.target) + + +@pytest.fixture +def vertex_environment(monkeypatch: pytest.MonkeyPatch) -> Callable[[str], None]: + def configure(token_url: str) -> None: + monkeypatch.setenv("VERTEXAI_PROJECT", _PROJECT) + monkeypatch.setenv("VERTEXAI_LOCATION", _LOCATION) + monkeypatch.setenv("VERTEXAI_CREDENTIALS", service_account_json(_PROJECT, token_url)) + + return configure + + +async def test_acount_tokens_forwards_system_and_tools_to_the_partner_peer( + vertex_environment: Callable[[str], None], +) -> None: + with wire_server(_peer(200)) as wire: + vertex_environment(wire.url) + counted: Final = await litellm.acount_tokens( + model=f"vertex_ai/{_BACKEND}", messages=_MESSAGES, tools=_TOOLS, system=_SYSTEM, api_base=wire.url + ) + (sent,) = _count_requests(wire.drain()) + assert (sent.method, sent.target, sent.headers["authorization"]) == ( + "POST", + _COUNT_TARGET, + "Bearer scripted-token", + ) + assert _JSON_OBJECT.validate_json(sent.body) == { + "model": _BACKEND, + "messages": _MESSAGES, + "system": _SYSTEM, + "tools": _TOOLS, + } + assert (counted.total_tokens, counted.tokenizer_type) == (_PEER_COUNT, "vertex_ai_partner_models"), counted + + +async def test_acount_tokens_falls_back_to_the_local_tokenizer_when_the_peer_rejects( + vertex_environment: Callable[[str], None], +) -> None: + with wire_server(_peer(400)) as wire: + vertex_environment(wire.url) + counted: Final = await litellm.acount_tokens( + model=f"vertex_ai/{_BACKEND}", messages=_MESSAGES, tools=_TOOLS, system=_SYSTEM, api_base=wire.url + ) + assert len(_count_requests(wire.drain())) == 1 + assert counted.tokenizer_type == "local_tokenizer", counted + assert counted.total_tokens > 0 and counted.total_tokens != _PEER_COUNT, counted diff --git a/tests/integration/spend/test_cache_and_quota.py b/tests/integration/spend/test_cache_and_quota.py index 1cc03838f8d..b5145ea49d4 100644 --- a/tests/integration/spend/test_cache_and_quota.py +++ b/tests/integration/spend/test_cache_and_quota.py @@ -399,9 +399,11 @@ def test_in_flight_count_tokens_does_not_reserve_key_budget_away_from_a_completi [{"role": "user", "content": prompt}] ] assert count.status_code == 200, count.text - assert count.json() == {"totalTokens": 12, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12}]}, ( - count.text - ) + assert count.json() == { + "totalTokens": 12, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12}], + "litellm_count_source": "provider", + }, count.text provider_calls: Final = wire.drain() assert [(call.method, call.target) for call in provider_calls] == [ ("POST", "/v1beta/models/gemini-3.8-flash:countTokens") diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index c71b1496bdd..18d379e5fdf 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -11,6 +11,7 @@ import threading import time from collections.abc import Mapping from concurrent.futures import Future, wait +from itertools import accumulate, chain from pathlib import Path from typing import Final from unittest.mock import MagicMock @@ -1434,14 +1435,51 @@ def _count_user_content(content: list[dict]) -> int: ) +def _pdf_base64(pages: tuple[str, ...], width: int = 612, height: int = 792) -> str: + def page_objects(index: int, text: str) -> tuple[bytes, bytes]: + escaped: Final = text.replace("\\", "\\\\").replace("(", "\\(").replace(")", "\\)") + stream: Final = f"BT /F1 12 Tf 72 720 Td ({escaped}) Tj ET".encode("latin-1") + content: Final = f"<< /Length {len(stream)} >>\nstream\n".encode() + stream + b"\nendstream" + page: Final = ( + f"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 {width} {height}] " + f"/Resources << /Font << /F1 3 0 R >> >> /Contents {4 + 2 * index} 0 R >>" + ).encode() + return content, page + + kids: Final = " ".join(f"{5 + 2 * index} 0 R" for index in range(len(pages))) + bodies: Final = ( + b"<< /Type /Catalog /Pages 2 0 R >>", + f"<< /Type /Pages /Count {len(pages)} /Kids [ {kids} ] >>".encode(), + b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>", + *chain.from_iterable(page_objects(index, text) for index, text in enumerate(pages)), + ) + header: Final = b"%PDF-1.4\n" + objects: Final = tuple( + f"{number} 0 obj\n".encode() + body + b"\nendobj\n" for number, body in enumerate(bodies, start=1) + ) + offsets: Final = accumulate((len(header), *(len(obj) for obj in objects[:-1]))) + xref: Final = f"xref\n0 {len(objects) + 1}\n0000000000 65535 f \n".encode() + b"".join( + f"{offset:010d} 00000 n \n".encode() for offset in offsets + ) + trailer: Final = ( + f"trailer\n<< /Size {len(objects) + 1} /Root 1 0 R >>\n" + f"startxref\n{len(header) + sum(len(obj) for obj in objects)}\n%%EOF\n" + ).encode() + return base64.b64encode(header + b"".join(objects) + xref + trailer).decode() + + +_READABLE_PDF_BASE64: Final = _pdf_base64(("Count this document without parsing its text locally.",)) + + @pytest.mark.parametrize( "source", [ {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"}, {"type": "url", "url": "https://example.com/report.pdf"}, {"type": "file", "file_id": "file-abc123"}, + {"type": "base64", "media_type": "application/pdf", "data": _READABLE_PDF_BASE64}, ], - ids=["base64", "url", "file"], + ids=["base64", "url", "file", "readable_pdf"], ) def test_anthropic_document_block_with_opaque_source_is_priced_like_an_image(source: dict[str, str]): """A `document` whose bytes can't be tokenized locally is priced like an `image`, not raised on.""" @@ -1483,7 +1521,8 @@ def test_anthropic_document_title_and_context_add_their_tokens(): ) -def test_openai_file_block_prices_like_the_equivalent_anthropic_document(): +@pytest.mark.parametrize("pdf_data", ("JVBERi0xLjQK", _READABLE_PDF_BASE64), ids=("opaque", "readable")) +def test_openai_file_block_prices_like_the_equivalent_anthropic_document(pdf_data: str) -> None: """An inline `file` is a `document` in the chat-completions dialect, so it must price identically, not raise. Before the fix `file` was missing from the content-block match even though `ChatCompletionFileObject` @@ -1493,17 +1532,21 @@ def test_openai_file_block_prices_like_the_equivalent_anthropic_document(): prompt = {"type": "text", "text": "Summarize this file."} inline_file = { "type": "file", - "file": {"filename": "report.pdf", "file_data": "data:application/pdf;base64,JVBERi0xLjQK"}, + "file": {"filename": "report.pdf", "file_data": "data:application/pdf;base64," + pdf_data}, } document = { "type": "document", "title": "report.pdf", - "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"}, + "source": {"type": "base64", "media_type": "application/pdf", "data": pdf_data}, } assert _count_user_content([prompt, inline_file]) == _count_user_content([prompt, document]) assert _count_user_content([prompt, inline_file]) > _count_user_content([prompt]) + assert _count_user_content([prompt, inline_file]) == _count_user_content( + [prompt, {"type": "text", "text": "report.pdf"}, {"type": "image", "source": document["source"]}] + ) + def test_openai_file_block_without_inline_bytes_counts_what_it_carries(): """A `file` block naming an uploaded file has no bytes to price, so it adds only the filename's tokens.""" @@ -1513,9 +1556,7 @@ def test_openai_file_block_without_inline_bytes_counts_what_it_carries(): assert _count_user_content([prompt, by_id]) == _count_user_content([prompt]) named = {"type": "file", "file": {"file_id": "file-abc123", "filename": "report.pdf"}} - assert _count_user_content([prompt, named]) == _count_user_content( - [prompt, {"type": "text", "text": "report.pdf"}] - ) + assert _count_user_content([prompt, named]) == _count_user_content([prompt, {"type": "text", "text": "report.pdf"}]) def _png_data_url(width: int, height: int) -> str: diff --git a/tests/unit/llms/anthropic/pass_through/test_convertible_blocks.py b/tests/unit/llms/anthropic/pass_through/test_convertible_blocks.py new file mode 100644 index 00000000000..b5301b413ff --- /dev/null +++ b/tests/unit/llms/anthropic/pass_through/test_convertible_blocks.py @@ -0,0 +1,192 @@ +import asyncio +import json +from collections.abc import Mapping +from typing import Final + +import pytest +from pydantic import JsonValue, TypeAdapter + +from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter +from litellm.llms.anthropic.pass_through.convertible_blocks import ( + unconvertible_anthropic_blocks, + unconvertible_blocks_error, +) +from litellm.llms.anthropic.pass_through.responses_adapters.transformation import ( + LiteLLMAnthropicToResponsesAPIAdapter, +) +from litellm.llms.vertex_ai.count_tokens.context import gemini_count_contents +from litellm.types.llms.anthropic import AllAnthropicPassThroughMessageValues + +_JSON_OBJECTS: Final = TypeAdapter(list[dict[str, JsonValue]]) + +_PDF: Final = {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"} +_PNG: Final = {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="} + +_CONVERTED_BLOCKS: Final[tuple[tuple[str, Mapping[str, object]], ...]] = ( + ("user", {"type": "text", "text": "hi"}), + ("user", {"type": "image", "source": _PNG}), + ("user", {"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}), + ("user", {"type": "document", "source": _PDF}), + ("user", {"type": "document", "source": {"type": "url", "url": "https://example.com/a.pdf"}}), + ("user", {"type": "tool_result", "tool_use_id": "toolu_01", "content": "42"}), + ("assistant", {"type": "text", "text": "hi"}), + ("assistant", {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}), + ("assistant", {"type": "thinking", "thinking": "Let me think.", "signature": "sig"}), +) + +_DROPPED_BLOCKS: Final[tuple[tuple[str, Mapping[str, object], str], ...]] = ( + ("user", {"type": "search_result", "source": "https://example.com", "title": "t", "content": []}, "search_result"), + ("user", {"type": "document", "source": {"type": "text", "data": "notes"}}, "document with a text source"), + ("user", {"type": "image", "source": {"type": "file", "file_id": "file_01"}}, "image with a file source"), + ("user", {"type": "image", "source": {"type": "base64", "data": ""}}, "image with an empty base64 source"), + ("user", {"type": "document"}, "document without a source"), + ("user", {"text": "untyped"}, "a block without a type"), + ("user", {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}, "tool_use"), + ( + "assistant", + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {}}, + "server_tool_use", + ), +) + + +def _adapter_messages(role: str, block: Mapping[str, object]) -> list[AllAnthropicPassThroughMessageValues]: + return [{"role": role, "content": [dict(block)]}] # pyright: ignore[reportReturnType] # sends block shapes the TypedDicts leave out on purpose + + +def _converted_by_both_adapters(role: str, block: Mapping[str, object]) -> tuple[bool, bool]: + messages: Final = _adapter_messages(role, block) + chat: Final = _JSON_OBJECTS.validate_python( + LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai(messages=messages) + ) + responses: Final = LiteLLMAnthropicToResponsesAPIAdapter().translate_messages_to_responses_input(messages) + chat_kept: Final = any( + message.get("content") or message.get("tool_calls") or message.get("thinking_blocks") for message in chat + ) + return chat_kept, bool(responses) + + +@pytest.mark.parametrize(("role", "block"), _CONVERTED_BLOCKS) +def test_blocks_both_conversions_keep_are_accepted(role: str, block: Mapping[str, object]) -> None: + assert _converted_by_both_adapters(role, block) == (True, True) + assert unconvertible_anthropic_blocks([{"role": role, "content": [block]}]) == () + + +@pytest.mark.parametrize(("role", "block", "label"), _DROPPED_BLOCKS) +def test_blocks_both_conversions_drop_are_named(role: str, block: Mapping[str, object], label: str) -> None: + assert _converted_by_both_adapters(role, block) == (False, False) + assert unconvertible_anthropic_blocks([{"role": role, "content": [block]}]) == (label,) + + +def test_labels_are_listed_once_in_order_and_string_content_is_accepted() -> None: + messages: Final = [ + {"role": "user", "content": "plain text"}, + {"role": "user", "content": [{"type": "search_result"}, {"type": "container_upload"}]}, + {"role": "assistant", "content": [{"type": "search_result"}]}, + ] + assert unconvertible_anthropic_blocks(messages) == ("container_upload", "search_result") + + +def test_redacted_thinking_is_accepted_like_other_replayed_thinking() -> None: + messages: Final = [{"role": "assistant", "content": [{"type": "redacted_thinking", "data": "opaque"}]}] + assert unconvertible_anthropic_blocks(messages) == () + + +def test_messages_that_are_not_json_are_left_to_the_count() -> None: + messages: Final = [{"role": "user", "content": [{"type": "search_result", "source": object()}]}] + assert unconvertible_anthropic_blocks(messages) == () + + +def _tool_turn(inner: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + return ( + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": [dict(inner)]}]}, + ) + + +def _sent_by_each_counter(inner: Mapping[str, object]) -> tuple[str, str]: + messages: Final = _tool_turn(inner) + gemini: Final = asyncio.run(gemini_count_contents(messages, "anthropic", "gemini-2.5-flash", "gemini", {})) + openai: Final = LiteLLMAnthropicToResponsesAPIAdapter().translate_messages_to_responses_input( + [dict(message) for message in messages] # pyright: ignore[reportArgumentType] # sends block shapes the TypedDicts leave out on purpose + ) + return json.dumps(gemini, default=str), json.dumps(openai, default=str) + + +_KEPT_INSIDE_TOOL_RESULT: Final[tuple[tuple[Mapping[str, object], str], ...]] = ( + ({"type": "text", "text": "forty-two"}, "forty-two"), + ({"type": "image", "source": _PNG}, _PNG["data"]), + ({"type": "document", "source": _PDF}, _PDF["data"]), +) + +_DROPPED_INSIDE_TOOL_RESULT: Final[tuple[tuple[Mapping[str, object], str, str], ...]] = ( + ( + { + "type": "search_result", + "source": "https://example.com", + "title": "t", + "content": [{"type": "text", "text": "search-snippet"}], + }, + "search-snippet", + "search_result inside a tool_result", + ), + ( + {"type": "document", "source": {"type": "text", "data": "text-notes"}}, + "text-notes", + "document with a text source inside a tool_result", + ), + ({"type": "tool_reference", "tool_name": "find_docs"}, "find_docs", "tool_reference inside a tool_result"), +) + + +@pytest.mark.parametrize(("inner", "marker"), _KEPT_INSIDE_TOOL_RESULT) +def test_tool_result_content_both_counters_send_is_accepted(inner: Mapping[str, object], marker: str) -> None: + gemini, openai = _sent_by_each_counter(inner) + + assert marker in gemini + assert marker in openai + assert unconvertible_anthropic_blocks(_tool_turn(inner)) == () + + +@pytest.mark.parametrize(("inner", "marker", "label"), _DROPPED_INSIDE_TOOL_RESULT) +def test_tool_result_content_both_counters_drop_is_named(inner: Mapping[str, object], marker: str, label: str) -> None: + gemini, openai = _sent_by_each_counter(inner) + + assert marker not in gemini + assert marker not in openai + assert unconvertible_anthropic_blocks(_tool_turn(inner)) == (label,) + + +def _replayed_web_search(encrypted_content: str) -> tuple[Mapping[str, object], ...]: + return ( + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": encrypted_content, + "snippet": "LiteLLM ships a weekly release candidate.", + } + ], + }, + ], + }, + ) + + +def test_a_replayed_web_search_inference_sends_as_text_is_accepted() -> None: + assert unconvertible_blocks_error(_replayed_web_search(""), "anthropic") is None + + +def test_an_encrypted_web_search_inference_cannot_send_is_named() -> None: + error: Final = unconvertible_blocks_error(_replayed_web_search("opaque"), "anthropic") + assert error is not None + assert error.endswith("server_tool_use, web_search_tool_result") diff --git a/tests/unit/llms/anthropic/test_count_tokens_oauth.py b/tests/unit/llms/anthropic/test_count_tokens_oauth.py index bcfc56577eb..169097ee47d 100644 --- a/tests/unit/llms/anthropic/test_count_tokens_oauth.py +++ b/tests/unit/llms/anthropic/test_count_tokens_oauth.py @@ -14,6 +14,12 @@ sys.path.insert( 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")) ) +import httpx +import pytest +import respx + +import litellm +from litellm.llms.anthropic.count_tokens.token_counter import AnthropicTokenCounter from litellm.llms.anthropic.count_tokens.transformation import ( AnthropicCountTokensConfig, ) @@ -84,3 +90,48 @@ class TestCountTokensOAuthHeaders: assert ( "oauth-2025-04-20" in beta_value ), f"oauth beta missing from OAuth headers: {beta_value}" + + +@pytest.mark.asyncio +async def test_a_global_anthropic_key_counts_like_chat_would_authenticate(monkeypatch): + for name in ("ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "anthropic_key", FAKE_REGULAR_KEY) + + with respx.mock(assert_all_called=True) as router: + route = router.post("https://api.anthropic.com/v1/messages/count_tokens").mock( + return_value=httpx.Response(200, json={"input_tokens": 13}) + ) + result = await AnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"model": "anthropic/claude-sonnet-4-5"}}, + request_model="claude-sonnet-4-5", + ) + + assert result is not None + assert result.total_tokens == 13 + assert route.calls.last.request.headers["x-api-key"] == FAKE_REGULAR_KEY + + +@pytest.mark.asyncio +async def test_no_anthropic_key_returns_a_401_count_error(monkeypatch): + for name in ("ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(litellm, "anthropic_key", None) + monkeypatch.setattr(litellm, "api_key", None) + + result = await AnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"model": "anthropic/claude-sonnet-4-5"}}, + request_model="claude-sonnet-4-5", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 401 + assert result.total_tokens == 0 diff --git a/tests/unit/llms/azure_ai/anthropic/__init__.py b/tests/unit/llms/azure_ai/anthropic/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/anthropic/count_tokens/__init__.py b/tests/unit/llms/azure_ai/anthropic/count_tokens/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py new file mode 100644 index 00000000000..4913d2b0d1d --- /dev/null +++ b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py @@ -0,0 +1,217 @@ +import json +import operator +from typing import Final + +import httpx +import pytest +import respx +from pydantic import ValidationError + +import litellm +from litellm.llms.azure_ai.anthropic.count_tokens.token_counter import AzureAIAnthropicTokenCounter +from litellm.types.utils import TokenCountResponse + +API_BASE: Final = "https://my-resource.services.ai.azure.com" +COUNT_URL: Final = f"{API_BASE}/anthropic/v1/messages/count_tokens" + + +async def _count(monkeypatch: pytest.MonkeyPatch, response: httpx.Response) -> TokenCountResponse | None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock: + respx.post(COUNT_URL).mock(return_value=response) + return await AzureAIAnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"api_key": "azure-key", "api_base": API_BASE}}, + request_model="azure_ai/claude-sonnet-4-5", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("payload", "expected_total"), + [ + ({"input_tokens": 15}, 15), + ({"input_tokens": "15"}, 15), + ({"input_tokens": 15.0}, 15), + ({"input_tokens": True}, 1), + ({"input_tokens": -3}, -3), + ({}, 0), + ({"type": "error", "error": {"message": "429 rate limit"}}, 0), + ({"input_tokens": 7, "context_management": {"original_input_tokens": 9}, "nested": {"a": [1, None]}}, 7), + ], +) +async def test_count_reports_upstream_input_tokens_and_keeps_the_raw_payload( + monkeypatch: pytest.MonkeyPatch, payload: dict[str, object], expected_total: int +) -> None: + result: Final = await _count(monkeypatch, httpx.Response(200, json=payload)) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": expected_total, + "request_model": "azure_ai/claude-sonnet-4-5", + "model_used": "claude-sonnet-4-5", + "tokenizer_type": "azure_ai_anthropic_api", + "count_source": None, + "original_response": payload, + "error": False, + "error_message": None, + "status_code": None, + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "rejected_count", + [ + None, + 15.5, + "Request timed out", + "429 Too Many Requests", + [15], + [*range(403), "Request timed out"], + {"value": 429}, + ], +) +async def test_non_integer_input_tokens_becomes_an_error_response_carrying_the_validation_message( + monkeypatch: pytest.MonkeyPatch, rejected_count: object +) -> None: + with pytest.raises(ValidationError) as rejection: + TokenCountResponse.model_validate( + { + "total_tokens": rejected_count, + "request_model": "azure_ai/claude-sonnet-4-5", + "model_used": "claude-sonnet-4-5", + "tokenizer_type": "azure_ai_anthropic_api", + } + ) + + result: Final = await _count(monkeypatch, httpx.Response(200, json={"input_tokens": rejected_count})) + + assert result is not None + assert [error["loc"] for error in rejection.value.errors()] == [("total_tokens",)] + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "azure_ai/claude-sonnet-4-5", + "model_used": "claude-sonnet-4-5", + "tokenizer_type": "azure_ai_anthropic_api", + "count_source": None, + "original_response": None, + "error": True, + "error_message": str(rejection.value), + "status_code": 500, + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "payload", + [[*({"input_tokens": position} for position in range(403)), "429"], "input_tokens", 15, 0, False], +) +async def test_count_payload_that_is_not_an_object_becomes_an_error_response_naming_the_failed_lookup( + monkeypatch: pytest.MonkeyPatch, payload: object +) -> None: + with pytest.raises(AttributeError) as failed_lookup: + operator.attrgetter("get")(payload) + + result: Final = await _count(monkeypatch, httpx.Response(200, json=payload)) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "azure_ai/claude-sonnet-4-5", + "model_used": "claude-sonnet-4-5", + "tokenizer_type": "azure_ai_anthropic_api", + "count_source": None, + "original_response": None, + "error": True, + "error_message": str(failed_lookup.value), + "status_code": 500, + } + ) + + +@pytest.mark.asyncio +async def test_null_count_payload_falls_back_to_local_counting(monkeypatch: pytest.MonkeyPatch) -> None: + assert await _count(monkeypatch, httpx.Response(200, content=b"null")) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("upstream_status", "upstream_text"), + [(429, "rate limit exceeded"), (403, "forbidden"), (408, "Request timed out")], +) +async def test_upstream_error_status_becomes_an_error_response_with_that_status_and_text( + monkeypatch: pytest.MonkeyPatch, upstream_status: int, upstream_text: str +) -> None: + result: Final = await _count(monkeypatch, httpx.Response(upstream_status, text=upstream_text)) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "azure_ai/claude-sonnet-4-5", + "model_used": "claude-sonnet-4-5", + "tokenizer_type": "azure_ai_anthropic_api", + "count_source": None, + "original_response": None, + "error": True, + "error_message": upstream_text, + "status_code": upstream_status, + } + ) + + +@pytest.mark.asyncio +async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_decoder_message( + monkeypatch: pytest.MonkeyPatch, +) -> None: + with pytest.raises(json.JSONDecodeError) as undecodable: + json.loads("not json") + + result: Final = await _count(monkeypatch, httpx.Response(200, text="not json")) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "azure_ai/claude-sonnet-4-5", + "model_used": "claude-sonnet-4-5", + "tokenizer_type": "azure_ai_anthropic_api", + "count_source": None, + "original_response": None, + "error": True, + "error_message": f"CountTokens processing error: {undecodable.value}", + "status_code": 500, + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("set_base", ("env", "global"), ids=("AZURE_API_BASE", "litellm.api_base")) +async def test_count_uses_the_shared_api_base_chat_would_use(monkeypatch: pytest.MonkeyPatch, set_base: str) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("AZURE_AI_API_BASE", raising=False) + monkeypatch.delenv("AZURE_API_BASE", raising=False) + if set_base == "env": + monkeypatch.setenv("AZURE_API_BASE", API_BASE) + else: + monkeypatch.setattr(litellm, "api_base", API_BASE) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 6})) + result: Final = await AzureAIAnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"api_key": "azure-key"}}, + request_model="azure_ai/claude-sonnet-4-5", + ) + assert result is not None + assert result.total_tokens == 6 + assert route.called diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py new file mode 100644 index 00000000000..ffddb09cca7 --- /dev/null +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py @@ -0,0 +1,524 @@ +import asyncio +import base64 +import hashlib +import json +import logging +import re +import threading +from collections.abc import Iterable +from itertools import product +from pathlib import Path +from typing import Final + +import httpx +import pytest +import respx +import tiktoken + +import litellm +from litellm._logging import verbose_logger +from litellm.llms.bedrock.count_tokens.kimi_k3 import ( + KimiK3Encodings, + KimiK3TokenizerSource, + kimi_k3_anthropic_tokens, + kimi_k3_tokens, + load_kimi_k3_encoding, +) +from litellm.responses.litellm_completion_transformation.transformation import LiteLLMCompletionResponsesConfig + +MODEL: Final = "us.moonshotai.kimi-k3" +_CONTROL: Final = re.compile(r"<\|(?:open|close|sep|end_of_msg)\|>") +_WEATHER: Final = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather in °C", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}, "days": {"type": "integer"}}, + "required": ["city"], + }, + }, +} + + +def _call(call_id: str, arguments: dict[str, object]) -> dict[str, object]: + return {"id": call_id, "type": "function", "function": {"name": "get_weather", "arguments": json.dumps(arguments)}} + + +def _written_ranks(path: Path, tokens: Iterable[bytes]) -> KimiK3TokenizerSource: + path.write_bytes(b"\n".join(base64.b64encode(token) + b" %d" % rank for rank, token in enumerate(tokens))) + return KimiK3TokenizerSource(ranks_path=lambda: str(path), sha256=hashlib.sha256(path.read_bytes()).hexdigest()) + + +def _single_bytes() -> tuple[bytes, ...]: + return tuple(bytes((value,)) for value in range(256)) + + +@pytest.fixture(scope="module") +def byte_source(tmp_path_factory: pytest.TempPathFactory) -> KimiK3TokenizerSource: + return _written_ranks(tmp_path_factory.mktemp("bytes") / "tiktoken.model", _single_bytes()) + + +@pytest.fixture(scope="module") +def byte_encoding(byte_source: KimiK3TokenizerSource) -> tiktoken.Encoding: + loaded: Final = load_kimi_k3_encoding(byte_source) + assert loaded is not None + return loaded + + +@pytest.fixture(scope="module") +def pair_encoding(tmp_path_factory: pytest.TempPathFactory) -> tiktoken.Encoding: + pairs: Final = (bytes(pair) for pair in product(range(256), repeat=2)) + source: Final = _written_ranks(tmp_path_factory.mktemp("pairs") / "tiktoken.model", (*_single_bytes(), *pairs)) + loaded: Final = load_kimi_k3_encoding(source) + assert loaded is not None + return loaded + + +def _template_count(encoding: tiktoken.Encoding, rendered: str) -> int: + texts: Final = _CONTROL.split(rendered) + return len(_CONTROL.findall(rendered)) + sum(len(encoding.encode_ordinary(text)) for text in texts) + + +_GENERATION_PROMPT: Final = '<|open|>message role="assistant"<|sep|><|open|>think<|sep|>' + + +# Expected renders follow the Kimi K3 chat template in moonshotai/Kimi-K3 at revision f831ab6, which matched +# Bedrock's prompt_tokens for us.moonshotai.kimi-k3 in 19 live calls on 2026-10-16 +@pytest.mark.parametrize( + ("messages", "tools", "rendered"), + [ + pytest.param( + [{"role": "user", "content": "Hello there"}], + None, + '<|open|>message role="user"<|sep|>Hello there<|close|>message<|sep|><|end_of_msg|>' + _GENERATION_PROMPT, + id="user", + ), + pytest.param( + [ + { + "role": "system", + "content": [{"type": "text", "text": "Be terse."}, {"type": "text", "text": "Answer in French."}], + }, + {"role": "user", "content": [{"type": "text", "text": "Hi"}, {"type": "text", "text": " there"}]}, + ], + None, + '<|open|>message role="system"<|sep|>Be terse.Answer in French.<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="user"<|sep|>Hi there<|close|>message<|sep|><|end_of_msg|>' + _GENERATION_PROMPT, + id="text_parts_join", + ), + pytest.param( + [ + {"role": "user", "content": "What is 2+2?"}, + {"role": "assistant", "content": "4"}, + {"role": "user", "content": "And 3+3?"}, + ], + None, + '<|open|>message role="user"<|sep|>What is 2+2?<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="assistant"<|sep|><|open|>think<|sep|><|close|>think<|sep|>' + "<|open|>response<|sep|>4<|close|>response<|sep|><|close|>message<|sep|><|end_of_msg|>" + '<|open|>message role="user"<|sep|>And 3+3?<|close|>message<|sep|><|end_of_msg|>' + _GENERATION_PROMPT, + id="assistant_turn", + ), + pytest.param( + [ + {"role": "user", "content": "Weather in Paris and Oslo?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [_call("c1", {"city": "Paris", "days": 3}), _call("c2", {"city": "Oslo"})], + }, + {"role": "tool", "tool_call_id": "c2", "content": "Snow"}, + {"role": "tool", "tool_call_id": "c1", "content": "Sun"}, + ], + [_WEATHER], + '<|open|>message role="system" type="tool-declare"<|sep|># Tools\n' + "Here are the available tools, described in JSONSchema.\n\n```json\n" + '[{"function":{"description":"Get the weather in °C","name":"get_weather","parameters":{"properties":' + '{"city":{"type":"string"},"days":{"type":"integer"}},"required":["city"],"type":"object"}},' + '"type":"function"}]\n```<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="user"<|sep|>Weather in Paris and Oslo?<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="assistant"<|sep|><|open|>think<|sep|><|close|>think<|sep|>' + "<|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|>" + '<|open|>call tool="get_weather" index="1"<|sep|>' + '<|open|>argument key="city" type="string"<|sep|>Paris<|close|>argument<|sep|>' + '<|open|>argument key="days" type="number"<|sep|>3<|close|>argument<|sep|><|close|>call<|sep|>' + '<|open|>call tool="get_weather" index="2"<|sep|>' + '<|open|>argument key="city" type="string"<|sep|>Oslo<|close|>argument<|sep|><|close|>call<|sep|>' + "<|close|>tools<|sep|><|close|>message<|sep|><|end_of_msg|>" + '<|open|>message role="tool" tool="get_weather" index="1"<|sep|>Sun<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="tool" tool="get_weather" index="2"<|sep|>Snow<|close|>message<|sep|><|end_of_msg|>' + + _GENERATION_PROMPT, + id="tool_results_out_of_call_order", + ), + pytest.param( + [ + {"role": "user", "content": "Weather?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + _call( + "c1", {"a&b": 'x"y', "flag": True, "none": None, "list": [1, "zé"], "object": {"k": True}} + ) + ], + }, + {"role": "tool", "tool_call_id": "c1", "content": "Rain"}, + ], + [_WEATHER], + '<|open|>message role="system" type="tool-declare"<|sep|># Tools\n' + "Here are the available tools, described in JSONSchema.\n\n```json\n" + '[{"function":{"description":"Get the weather in °C","name":"get_weather","parameters":{"properties":' + '{"city":{"type":"string"},"days":{"type":"integer"}},"required":["city"],"type":"object"}},' + '"type":"function"}]\n```<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="user"<|sep|>Weather?<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="assistant"<|sep|><|open|>think<|sep|><|close|>think<|sep|>' + "<|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|>" + '<|open|>call tool="get_weather" index="1"<|sep|>' + '<|open|>argument key="a&b" type="string"<|sep|>x"y<|close|>argument<|sep|>' + '<|open|>argument key="flag" type="boolean"<|sep|>true<|close|>argument<|sep|>' + '<|open|>argument key="none" type="null"<|sep|>null<|close|>argument<|sep|>' + '<|open|>argument key="list" type="array"<|sep|>[1, "zé"]<|close|>argument<|sep|>' + '<|open|>argument key="object" type="object"<|sep|>{"k": true}<|close|>argument<|sep|>' + "<|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|><|end_of_msg|>" + '<|open|>message role="tool" tool="get_weather" index="1"<|sep|>Rain<|close|>message<|sep|><|end_of_msg|>' + + _GENERATION_PROMPT, + id="argument_types_and_escaping", + ), + ], +) +@pytest.mark.parametrize("encoding_fixture", ["byte_encoding", "pair_encoding"]) +def test_count_is_the_kimi_k3_chat_template_render_of_the_bedrock_request( + request: pytest.FixtureRequest, + encoding_fixture: str, + messages: list[dict[str, object]], + tools: list[dict[str, object]] | None, + rendered: str, +): + encoding: Final[tiktoken.Encoding] = request.getfixturevalue(encoding_fixture) + + assert kimi_k3_tokens(encoding, MODEL, messages, tools) == _template_count(encoding, rendered) + + +def _parallel_calls(count: int) -> list[dict[str, object]]: + ids: Final = tuple(f"c{number}" for number in range(1, count + 1)) + return [ + {"role": "user", "content": "Weather?"}, + {"role": "assistant", "content": None, "tool_calls": [_call(call_id, {"city": "Paris"}) for call_id in ids]}, + *({"role": "tool", "tool_call_id": call_id, "content": "Sun"} for call_id in ids), + ] + + +def test_tenth_parallel_call_and_its_result_are_numbered_ten(byte_encoding: tiktoken.Encoding): + nine: Final = kimi_k3_tokens(byte_encoding, MODEL, _parallel_calls(9), [_WEATHER]) + ten: Final = kimi_k3_tokens(byte_encoding, MODEL, _parallel_calls(10), [_WEATHER]) + assert nine is not None and ten is not None + + assert ten - nine == _template_count( + byte_encoding, + '<|open|>call tool="get_weather" index="10"<|sep|>' + '<|open|>argument key="city" type="string"<|sep|>Paris<|close|>argument<|sep|><|close|>call<|sep|>' + '<|open|>message role="tool" tool="get_weather" index="10"<|sep|>Sun<|close|>message<|sep|><|end_of_msg|>', + ) + + +@pytest.mark.parametrize( + "text", + [ + pytest.param("The quick brown fox, jumps over 12345 lazy dogs. " * 60, id="prose"), + pytest.param("a" + "1234567890" * 250 + "1", id="long_digit_run"), + pytest.param("--- " * 600, id="symbols"), + pytest.param("def f(x):\n return x # it's 42\n\n" * 80, id="code"), + pytest.param("word next\t\ttab\n \n" * 120, id="whitespace"), + pytest.param("naïve café, “quotes” & 中文汉字，日本語のテキスト 😀🎉\r\n" * 40, id="unicode"), + pytest.param("We're sure they'd say it's DON'T or I'M " * 60, id="contractions"), + pytest.param("abc123def4567ghi89 " * 120, id="letters_and_digits"), + ], +) +def test_long_text_counts_what_the_tokenizer_counts_for_the_whole_text(pair_encoding: tiktoken.Encoding, text: str): + def counted(content: str) -> int | None: + return kimi_k3_tokens(pair_encoding, MODEL, [{"role": "user", "content": content}], None) + + short: Final = counted("x") + assert short is not None + + assert counted(text) == short - len(pair_encoding.encode_ordinary("x")) + len(pair_encoding.encode_ordinary(text)) + + +@pytest.mark.parametrize( + ("messages", "tools"), + [ + pytest.param( + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}}, + ], + } + ], + None, + id="image", + ), + pytest.param( + [{"role": "user", "content": "Weather?"}, {"role": "tool", "tool_call_id": "c1", "content": "Sun"}], + [_WEATHER], + id="tool_result_without_its_call", + ), + ], +) +def test_request_the_template_cannot_render_is_not_counted( + pair_encoding: tiktoken.Encoding, messages: list[dict[str, object]], tools: list[dict[str, object]] | None +): + assert kimi_k3_tokens(pair_encoding, MODEL, messages, tools) is None + + +@pytest.mark.parametrize( + "tool", + [ + pytest.param( + { + "type": "function", + "name": "get_weather", + "description": "Get the weather in °C", + "parameters": _WEATHER["function"]["parameters"], + }, + id="responses_api", + ), + pytest.param( + { + "name": "get_weather", + "description": "Get the weather in °C", + "input_schema": _WEATHER["function"]["parameters"], + }, + id="anthropic_messages", + ), + ], +) +def test_tool_in_another_api_shape_counts_like_the_same_chat_completions_tool( + pair_encoding: tiktoken.Encoding, tool: dict[str, object] +): + messages: Final = [{"role": "user", "content": "Weather in Paris?"}] + + counted: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [tool]) + + assert counted is not None + assert counted == kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER]) + + +@pytest.mark.parametrize( + "tool", + [ + pytest.param({"type": "web_search_20250305", "name": "web_search"}, id="anthropic_web_search"), + pytest.param({"type": "web_search"}, id="responses_web_search"), + pytest.param( + { + "type": "mcp", + "server_label": "docs", + "server_url": "https://mcp.example/mcp", + "require_approval": "never", + }, + id="mcp", + ), + ], +) +def test_tool_bedrock_inference_drops_counts_like_the_function_tools_alone( + pair_encoding: tiktoken.Encoding, tool: dict[str, object] +): + messages: Final = [{"role": "user", "content": "Weather in Paris?"}] + + counted: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER, tool]) + + assert counted is not None + assert counted == kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER]) + + +def test_responses_custom_tool_counts_like_the_function_tool_inference_sends(pair_encoding: tiktoken.Encoding): + messages: Final = [{"role": "user", "content": "Weather in Paris?"}] + custom_tool: Final = {"type": "custom", "name": "grammar", "description": "Free-form SQL"} + sent_to_inference, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( + [dict(custom_tool)] + ) + + without_it: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER]) + + counted: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER, custom_tool]) + + assert counted is not None and without_it is not None + assert counted > without_it + assert counted == kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER, *sent_to_inference]) + + +def test_counting_a_request_with_an_image_url_never_downloads_the_image( + pair_encoding: tiktoken.Encoding, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "user_url_validation", False) + image: Final = respx_mock.get("https://images.example/cat.png").mock(return_value=httpx.Response(200)) + messages: Final = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + {"type": "image_url", "image_url": {"url": "https://images.example/cat.png"}}, + ], + } + ] + + assert kimi_k3_tokens(pair_encoding, MODEL, messages, None) is None + assert not image.called + + +def test_anthropic_request_counts_like_the_same_openai_request(pair_encoding: tiktoken.Encoding): + anthropic: Final = kimi_k3_anthropic_tokens( + pair_encoding, + MODEL, + { + "system": "Be terse.", + "messages": [ + {"role": "user", "content": "Weather in Paris?"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "c1", "name": "get_weather", "input": {"city": "Paris"}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "c1", "content": "Sun"}]}, + ], + "tools": [ + { + "name": "get_weather", + "description": "Get the weather in °C", + "input_schema": _WEATHER["function"]["parameters"], + } + ], + }, + ) + openai: Final = kimi_k3_tokens( + pair_encoding, + MODEL, + [ + {"role": "system", "content": "Be terse."}, + {"role": "user", "content": "Weather in Paris?"}, + {"role": "assistant", "content": None, "tool_calls": [_call("c1", {"city": "Paris"})]}, + {"role": "tool", "tool_call_id": "c1", "content": "Sun"}, + ], + [_WEATHER], + ) + + assert openai is not None + assert anthropic == openai + + +@pytest.mark.parametrize("failure", ["checksum", "missing_file"]) +def test_tokenizer_that_fails_its_checksum_or_cannot_be_read_is_not_loaded( + byte_source: KimiK3TokenizerSource, tmp_path: Path, caplog: pytest.LogCaptureFixture, failure: str +): + source: Final = ( + KimiK3TokenizerSource(ranks_path=byte_source.ranks_path, sha256=hashlib.sha256(b"other").hexdigest()) + if failure == "checksum" + else KimiK3TokenizerSource(ranks_path=lambda: str(tmp_path / "missing"), sha256=byte_source.sha256) + ) + verbose_logger.addHandler(caplog.handler) + try: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + loaded: Final = load_kimi_k3_encoding(source) + finally: + verbose_logger.removeHandler(caplog.handler) + + assert loaded is None + assert any("Kimi K3 tokenizer unavailable" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_slow_tokenizer_load_is_not_waited_on_past_the_wait_and_is_shared_once_done( + byte_source: KimiK3TokenizerSource, +): + released: Final = threading.Event() + + def slow_ranks_path() -> str: + released.wait(10) + return byte_source.ranks_path() + + encodings: Final = KimiK3Encodings( + source=KimiK3TokenizerSource(ranks_path=slow_ranks_path, sha256=byte_source.sha256), + retry_after_seconds=60, + clock=lambda: 0.0, + ) + + assert await encodings.encoding(wait_seconds=0) is None + released.set() + loaded: Final = await encodings.encoding(wait_seconds=5) + assert loaded is not None + assert await encodings.encoding(wait_seconds=0) is loaded + + +class _Clock: + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now + + +async def _loaded_eventually(encodings: KimiK3Encodings) -> tiktoken.Encoding | None: + for _ in range(500): + if (loaded := await encodings.encoding(wait_seconds=0)) is not None: + return loaded + await asyncio.sleep(0.01) + return None + + +@pytest.mark.asyncio +async def test_failed_tokenizer_load_is_retried_in_the_background_only_after_the_retry_window( + byte_source: KimiK3TokenizerSource, tmp_path: Path +): + clock: Final = _Clock() + reads: Final[list[float]] = [] + released: Final = threading.Event() + + def ranks_path() -> str: + reads.append(clock.now) + if len(reads) == 1: + return str(tmp_path / "missing") + released.wait(10) + return byte_source.ranks_path() + + encodings: Final = KimiK3Encodings( + source=KimiK3TokenizerSource(ranks_path=ranks_path, sha256=byte_source.sha256), + retry_after_seconds=60, + clock=clock, + ) + + assert await encodings.encoding(wait_seconds=5) is None + clock.now = 59.0 + assert await encodings.encoding(wait_seconds=5) is None + clock.now = 60.0 + assert await asyncio.wait_for(encodings.encoding(wait_seconds=5), timeout=1) is None + released.set() + assert await _loaded_eventually(encodings) is not None + assert reads == [0.0, 60.0] + + +@pytest.mark.asyncio +async def test_tokenizer_load_still_running_after_the_wait_is_logged( + byte_source: KimiK3TokenizerSource, caplog: pytest.LogCaptureFixture +): + released: Final = threading.Event() + + def hung_ranks_path() -> str: + released.wait(10) + return byte_source.ranks_path() + + encodings: Final = KimiK3Encodings( + source=KimiK3TokenizerSource(ranks_path=hung_ranks_path, sha256=byte_source.sha256), retry_after_seconds=60 + ) + verbose_logger.addHandler(caplog.handler) + try: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + loaded: Final = await encodings.encoding(wait_seconds=0.05) + finally: + verbose_logger.removeHandler(caplog.handler) + released.set() + + assert loaded is None + assert any("Kimi K3 tokenizer still loading" in record.getMessage() for record in caplog.records) diff --git a/tests/unit/llms/gemini/test_gemini_common_utils.py b/tests/unit/llms/gemini/test_gemini_common_utils.py index 70946e10590..e54023c6596 100644 --- a/tests/unit/llms/gemini/test_gemini_common_utils.py +++ b/tests/unit/llms/gemini/test_gemini_common_utils.py @@ -1,8 +1,200 @@ +from typing import Final from unittest.mock import AsyncMock, patch +import httpx import pytest +from pydantic import JsonValue, TypeAdapter +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.gemini.common_utils import GeminiModelInfo, GoogleAIStudioTokenCounter +from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler +from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter as GeminiCountTokensHandler + +_JSON_BODY: Final = TypeAdapter(dict[str, JsonValue]) + + +class _GeminiCountTokensPeer: + def __init__(self) -> None: + self.request_body: dict[str, JsonValue] | None = None + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.request_body = _JSON_BODY.validate_json(request.content) + return httpx.Response(200, json={"totalTokens": 11}, request=request) + + +@pytest.mark.asyncio +async def test_google_ai_studio_count_tokens_forwards_system_and_tools() -> None: + peer: Final = _GeminiCountTokensPeer() + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + result: Final = await GoogleAIStudioTokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=None, + contents=[{"role": "user", "parts": [{"text": "hello"}]}], + deployment={ + "litellm_params": { + "api_key": "test-key", + "api_base": "https://count-tokens.test", + "system_instruction": {"parts": [{"text": "deployment-level"}]}, + } + }, + request_model="gemini/gemini-2.5-flash", + tools=[ + { + "name": "lookup", + "description": "Find a record", + "input_schema": {"type": "object", "properties": {}}, + } + ], + system="You are terse.", + ) + + assert result is not None + assert result.error is False + assert result.total_tokens == 11 + assert peer.request_body == { + "generateContentRequest": { + "model": "models/gemini-2.5-flash", + "contents": [{"role": "user", "parts": [{"text": "hello"}]}], + "systemInstruction": {"parts": [{"text": "You are terse."}]}, + "tools": [ + { + "functionDeclarations": [ + { + "name": "lookup", + "description": "Find a record", + "parameters": {"type": "object"}, + } + ] + } + ], + } + } + + +@pytest.mark.asyncio +async def test_google_ai_studio_count_tokens_rejects_image_system_without_request() -> None: + peer: Final = _GeminiCountTokensPeer() + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + result: Final = await GoogleAIStudioTokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=None, + contents=[{"role": "user", "parts": [{"text": "hello"}]}], + deployment={"litellm_params": {"api_key": "test-key", "api_base": "https://count-tokens.test"}}, + request_model="gemini/gemini-2.5-flash", + system=[ + {"type": "text", "text": "You are terse."}, + { + "type": "image", + "source": {"type": "base64", "media_type": "image/png", "data": "iVBORw0KGgo="}, + }, + ], + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert peer.request_body is None + + +@pytest.mark.asyncio +async def test_google_ai_studio_count_tokens_rejects_an_image_url_it_refuses_to_fetch_without_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", True) + peer: Final = _GeminiCountTokensPeer() + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + result: Final = await GoogleAIStudioTokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=[ + { + "role": "user", + "content": [{"type": "image", "source": {"type": "url", "url": "http://127.0.0.1:9/image.png"}}], + } + ], + contents=None, + deployment={"litellm_params": {"api_key": "test-key", "api_base": "https://count-tokens.test"}}, + request_model="gemini/gemini-2.5-flash", + messages_format="anthropic", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert "http://127.0.0.1:9/image.png" in (result.error_message or "") + assert peer.request_body is None + + +@pytest.mark.asyncio +async def test_google_ai_studio_count_tokens_answers_with_the_status_google_rejected_it_with() -> None: + rejection: Final = '{"error": {"code": 400, "message": "Invalid function name", "status": "INVALID_ARGUMENT"}}' + client: Final = AsyncHTTPHandler( + transport=httpx.MockTransport(lambda request: httpx.Response(400, text=rejection, request=request)) + ) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + result: Final = await GoogleAIStudioTokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=None, + contents=[{"role": "user", "parts": [{"text": "hello"}]}], + deployment={"litellm_params": {"api_key": "test-key", "api_base": "https://count-tokens.test"}}, + request_model="gemini/gemini-2.5-flash", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert "Invalid function name" in (result.error_message or "") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("system_instruction", "tools", "expected_body"), + ( + ( + {"parts": [{"text": "Follow the system rule."}]}, + [{"functionDeclarations": [{"name": "lookup", "parameters": {"type": "OBJECT"}}]}], + { + "generateContentRequest": { + "model": "models/gemini-4-argon", + "contents": [{"role": "user", "parts": [{"text": "hello"}]}], + "systemInstruction": {"parts": [{"text": "Follow the system rule."}]}, + "tools": [{"functionDeclarations": [{"name": "lookup", "parameters": {"type": "OBJECT"}}]}], + } + }, + ), + ( + None, + [], + {"contents": [{"role": "user", "parts": [{"text": "hello"}]}]}, + ), + ), +) +async def test_gemini_count_tokens_sends_optional_request_fields( + system_instruction: dict[str, JsonValue] | None, + tools: list[dict[str, JsonValue]], + expected_body: dict[str, JsonValue], +) -> None: + peer: Final = _GeminiCountTokensPeer() + token_counter: Final = GeminiCountTokensHandler(client=AsyncHTTPHandler(transport=httpx.MockTransport(peer))) + + result: Final = await token_counter.acount_tokens( + contents=[{"role": "user", "parts": [{"text": "hello"}]}], + model="gemini-4-argon", + api_key="test-key", + api_base="https://count-tokens.test", + system_instruction=system_instruction, + tools=tools, + ) + + assert result == {"totalTokens": 11} + assert peer.request_body == expected_body class TestGeminiModelInfo: @@ -36,18 +228,10 @@ class TestGeminiModelInfo: # Test edge cases where model names end with characters from "models/" # These would be incorrectly processed if using strip("models/") instead of replace("models/", "") models = [ - { - "name": "models/gemini-1.5-pro" - }, # ends with 'o' - would become "gemini-1.5-pr" with strip() - { - "name": "models/test-model" - }, # ends with 'l' - would become "gemini/test-mode" with strip() - { - "name": "models/custom-models" - }, # ends with 's' - would become "gemini/custom-model" with strip() - { - "name": "models/demo" - }, # ends with 'o' - would become "gemini/dem" with strip() + {"name": "models/gemini-1.5-pro"}, # ends with 'o' - would become "gemini-1.5-pr" with strip() + {"name": "models/test-model"}, # ends with 'l' - would become "gemini/test-mode" with strip() + {"name": "models/custom-models"}, # ends with 's' - would become "gemini/custom-model" with strip() + {"name": "models/demo"}, # ends with 'o' - would become "gemini/dem" with strip() ] result = gemini_model_info.process_model_name(models) @@ -98,16 +282,10 @@ class TestGoogleAIStudioTokenCounter: token_counter = GoogleAIStudioTokenCounter() # Test with gemini provider - should return True - assert ( - token_counter.should_use_token_counting_api(LlmProviders.GEMINI.value) - is True - ) + assert token_counter.should_use_token_counting_api(LlmProviders.GEMINI.value) is True # Test with other providers - should return False - assert ( - token_counter.should_use_token_counting_api(LlmProviders.OPENAI.value) - is False - ) + assert token_counter.should_use_token_counting_api(LlmProviders.OPENAI.value) is False assert token_counter.should_use_token_counting_api("anthropic") is False assert token_counter.should_use_token_counting_api("vertex_ai") is False @@ -158,7 +336,10 @@ class TestGoogleAIStudioTokenCounter: # Verify the mock was called correctly mock_acount_tokens.assert_called_once_with( - model=model_to_use, contents=contents + model=model_to_use, + contents=contents, + system_instruction=None, + tools=None, ) def test_clean_contents_for_gemini_api_removes_id_field(self): @@ -176,9 +357,7 @@ class TestGoogleAIStudioTokenCounter: "functionResponse": { "id": "read_many_files-1757526647518-730a691aac11c", # This should be removed "name": "read_many_files", - "response": { - "output": "No files matching the criteria were found or all were skipped." - }, + "response": {"output": "No files matching the criteria were found or all were skipped."}, } } ], @@ -187,9 +366,7 @@ class TestGoogleAIStudioTokenCounter: ] # Clean the contents - cleaned_contents = token_counter._clean_contents_for_gemini_api( - contents_with_id - ) + cleaned_contents = token_counter._clean_contents_for_gemini_api(contents_with_id) # Verify the 'id' field was removed function_response = cleaned_contents[1]["parts"][0]["functionResponse"] @@ -198,8 +375,7 @@ class TestGoogleAIStudioTokenCounter: assert "response" in function_response assert function_response["name"] == "read_many_files" assert ( - function_response["response"]["output"] - == "No files matching the criteria were found or all were skipped." + function_response["response"]["output"] == "No files matching the criteria were found or all were skipped." ) def test_clean_contents_for_gemini_api_preserves_other_fields(self): @@ -215,9 +391,7 @@ class TestGoogleAIStudioTokenCounter: ] # Clean the contents - cleaned_contents = token_counter._clean_contents_for_gemini_api( - contents_without_function_response - ) + cleaned_contents = token_counter._clean_contents_for_gemini_api(contents_without_function_response) # Verify the contents are unchanged assert cleaned_contents == contents_without_function_response diff --git a/tests/unit/llms/openai/responses/count_tokens/__init__.py b/tests/unit/llms/openai/responses/count_tokens/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py new file mode 100644 index 00000000000..a45773b93fc --- /dev/null +++ b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py @@ -0,0 +1,436 @@ +import json +import operator +from typing import Final + +import httpx +import pytest +import respx +from pydantic import JsonValue, ValidationError + +import litellm +from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse + +COUNT_URL: Final = "https://api.openai.com/v1/responses/input_tokens" + + +async def _count(monkeypatch: pytest.MonkeyPatch, response: httpx.Response) -> TokenCountResponse | None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock: + respx.post(COUNT_URL).mock(return_value=response) + return await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"api_key": "sk-test"}}, + request_model="openai/gpt-4o", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("payload", "expected_total"), + [ + ({"input_tokens": 15}, 15), + ({"input_tokens": "15"}, 15), + ({"input_tokens": 15.0}, 15), + ({"input_tokens": True}, 1), + ({"input_tokens": -3}, -3), + ({}, 0), + ({"object": "response.input_tokens", "error": {"message": "429 rate limit"}}, 0), + ({"input_tokens": 7, "object": "response.input_tokens", "nested": {"a": [1, None]}}, 7), + ], +) +async def test_count_reports_upstream_input_tokens_and_keeps_the_raw_payload( + monkeypatch: pytest.MonkeyPatch, payload: dict[str, object], expected_total: int +) -> None: + result: Final = await _count(monkeypatch, httpx.Response(200, json=payload)) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": expected_total, + "request_model": "openai/gpt-4o", + "model_used": "gpt-4o", + "tokenizer_type": "openai_api", + "count_source": None, + "original_response": payload, + "error": False, + "error_message": None, + "status_code": None, + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "rejected_count", + [ + None, + 15.5, + "Request timed out", + "429 Too Many Requests", + [15], + [*range(403), "Request timed out"], + {"value": 429}, + ], +) +async def test_non_integer_input_tokens_becomes_an_error_response_carrying_the_validation_message( + monkeypatch: pytest.MonkeyPatch, rejected_count: object +) -> None: + with pytest.raises(ValidationError) as rejection: + TokenCountResponse.model_validate( + { + "total_tokens": rejected_count, + "request_model": "openai/gpt-4o", + "model_used": "gpt-4o", + "tokenizer_type": "openai_api", + } + ) + + result: Final = await _count(monkeypatch, httpx.Response(200, json={"input_tokens": rejected_count})) + + assert result is not None + assert [error["loc"] for error in rejection.value.errors()] == [("total_tokens",)] + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "openai/gpt-4o", + "model_used": "gpt-4o", + "tokenizer_type": "openai_api", + "count_source": None, + "original_response": None, + "error": True, + "error_message": str(rejection.value), + "status_code": 500, + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "payload", + [[*({"input_tokens": position} for position in range(403)), "429"], "input_tokens", 15, 0, False], +) +async def test_count_payload_that_is_not_an_object_becomes_an_error_response_naming_the_failed_lookup( + monkeypatch: pytest.MonkeyPatch, payload: object +) -> None: + with pytest.raises(AttributeError) as failed_lookup: + operator.attrgetter("get")(payload) + + result: Final = await _count(monkeypatch, httpx.Response(200, json=payload)) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "openai/gpt-4o", + "model_used": "gpt-4o", + "tokenizer_type": "openai_api", + "count_source": None, + "original_response": None, + "error": True, + "error_message": str(failed_lookup.value), + "status_code": 500, + } + ) + + +@pytest.mark.asyncio +async def test_null_count_payload_falls_back_to_local_counting(monkeypatch: pytest.MonkeyPatch) -> None: + assert await _count(monkeypatch, httpx.Response(200, content=b"null")) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("upstream_status", "upstream_text"), + [(429, "rate limit exceeded"), (403, "forbidden"), (408, "Request timed out")], +) +async def test_upstream_error_status_becomes_an_error_response_with_that_status_and_text( + monkeypatch: pytest.MonkeyPatch, upstream_status: int, upstream_text: str +) -> None: + result: Final = await _count(monkeypatch, httpx.Response(upstream_status, text=upstream_text)) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "openai/gpt-4o", + "model_used": "gpt-4o", + "tokenizer_type": "openai_api", + "count_source": None, + "original_response": None, + "error": True, + "error_message": upstream_text, + "status_code": upstream_status, + } + ) + + +async def _sent_count_body( + monkeypatch: pytest.MonkeyPatch, + messages: list[dict[str, object]], + system: object, + messages_format: TokenCountMessagesFormat, + output_format: JsonValue = None, +) -> object: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 1})) + await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=messages, + contents=None, + deployment={"litellm_params": {"api_key": "sk-test"}}, + request_model="openai/gpt-4o", + system=system, + messages_format=messages_format, + output_format=output_format, + ) + return json.loads(route.calls.last.request.content) + + +@pytest.mark.asyncio +async def test_anthropic_format_sends_the_document_image_and_tool_turns_to_the_count_endpoint( + monkeypatch: pytest.MonkeyPatch, +) -> None: + messages: Final[list[dict[str, object]]] = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize."}, + {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBE"}}, + {"type": "image", "source": {"type": "url", "url": "https://example.com/page.png"}}, + ], + }, + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "found"}]}, + ] + + sent: Final = await _sent_count_body( + monkeypatch, messages, [{"type": "text", "text": "Be brief."}], messages_format="anthropic" + ) + + assert sent == { + "model": "gpt-4o", + "input": [ + { + "type": "message", + "role": "user", + "content": [ + {"type": "input_text", "text": "Summarize."}, + {"type": "input_file", "filename": "document.pdf", "file_data": "data:application/pdf;base64,JVBE"}, + {"type": "input_image", "image_url": "https://example.com/page.png"}, + ], + }, + {"type": "function_call", "call_id": "toolu_01", "name": "lookup", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "toolu_01", "output": "found"}, + ], + "instructions": "Be brief.", + } + + +@pytest.mark.asyncio +async def test_openai_format_keeps_the_chat_conversion(monkeypatch: pytest.MonkeyPatch) -> None: + messages: Final[list[dict[str, object]]] = [ + {"role": "system", "content": "Be brief."}, + {"role": "user", "content": [{"type": "text", "text": "Summarize."}]}, + {"role": "tool", "tool_call_id": "call_1", "content": "found"}, + ] + + sent: Final = await _sent_count_body(monkeypatch, messages, system=None, messages_format="openai") + + assert sent == { + "model": "gpt-4o", + "input": [ + {"role": "user", "content": "Summarize."}, + {"type": "function_call_output", "call_id": "call_1", "output": "found"}, + ], + "instructions": "Be brief.", + } + + +@pytest.mark.asyncio +async def test_anthropic_format_sends_affinity_wrapped_reasoning_unwrapped_as_inference_does( + monkeypatch: pytest.MonkeyPatch, +) -> None: + wrapped: Final = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAoriginal", "deployment-1") # pyright: ignore[reportPrivateUsage] # the wrap inference applies + messages: Final[list[dict[str, object]]] = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "redacted_thinking", "data": f"litellm_encrypted_reasoning:{wrapped}"}, + {"type": "text", "text": "OK"}, + ], + }, + {"role": "user", "content": "again"}, + ] + + sent: Final = await _sent_count_body(monkeypatch, messages, system=None, messages_format="anthropic") + + assert sent == { + "model": "gpt-4o", + "input": [ + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]}, + {"type": "reasoning", "summary": [], "encrypted_content": "gAAAAoriginal"}, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "OK"}]}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "again"}]}, + ], + } + + +@pytest.mark.asyncio +async def test_anthropic_format_counts_the_structured_output_schema(monkeypatch: pytest.MonkeyPatch) -> None: + schema: Final[dict[str, JsonValue]] = {"type": "object", "properties": {"answer": {"type": "string"}}} + + sent: Final = await _sent_count_body( + monkeypatch, + [{"role": "user", "content": "Answer in JSON."}], + system=None, + messages_format="anthropic", + output_format={"type": "json_schema", "schema": schema}, + ) + + assert sent == { + "model": "gpt-4o", + "input": [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Answer in JSON."}]}], + "text": {"format": {"type": "json_schema", "name": "structured_output", "schema": schema, "strict": False}}, + } + + +@pytest.mark.asyncio +async def test_anthropic_format_sends_a_replayed_web_search_result_as_its_text(monkeypatch: pytest.MonkeyPatch) -> None: + snippet: Final = "LiteLLM ships a weekly release candidate and a stable line." + messages: Final[list[dict[str, object]]] = [ + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": "", + "snippet": snippet, + } + ], + }, + ], + }, + {"role": "user", "content": "Thanks."}, + ] + + sent: Final = json.dumps(await _sent_count_body(monkeypatch, messages, system=None, messages_format="anthropic")) + + assert snippet in sent + assert "https://docs.litellm.ai/release_notes" in sent + + +@pytest.mark.asyncio +async def test_anthropic_format_message_without_a_role_answers_400_without_calling_openai( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 1})) + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}, {"content": "no role"}], + contents=None, + deployment={"litellm_params": {"api_key": "sk-test"}}, + request_model="openai/gpt-4o", + messages_format="anthropic", + ) + + assert route.called is False + assert result is not None + assert (result.error, result.status_code, result.error_message) == (True, 400, "messages.1.role: Field required") + + +@pytest.mark.asyncio +async def test_openai_format_sends_a_list_system_as_its_text(monkeypatch: pytest.MonkeyPatch) -> None: + sent: Final = await _sent_count_body( + monkeypatch, + [{"role": "user", "content": "hi"}], + system=[{"type": "text", "text": "Be brief."}, {"type": "text", "text": "Answer in English."}], + messages_format="openai", + ) + + assert sent == { + "model": "gpt-4o", + "input": [{"role": "user", "content": "hi"}], + "instructions": "Be brief.\nAnswer in English.", + } + + +@pytest.mark.asyncio +async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_decoder_message( + monkeypatch: pytest.MonkeyPatch, +) -> None: + with pytest.raises(json.JSONDecodeError) as undecodable: + json.loads("not json") + + result: Final = await _count(monkeypatch, httpx.Response(200, text="not json")) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "openai/gpt-4o", + "model_used": "gpt-4o", + "tokenizer_type": "openai_api", + "count_source": None, + "original_response": None, + "error": True, + "error_message": f"CountTokens processing error: {undecodable.value}", + "status_code": 500, + } + ) + + +@pytest.mark.asyncio +async def test_count_uses_the_global_openai_key_chat_would_use(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "openai_key", "sk-from-litellm-settings") + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 9})) + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="openai/gpt-4o", + ) + assert result is not None + assert result.total_tokens == 9 + assert route.calls.last.request.headers["authorization"] == "Bearer sk-from-litellm-settings" + + +@pytest.mark.asyncio +async def test_count_sends_the_global_key_only_to_the_global_api_base(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "openai_key", "sk-private-gateway") + monkeypatch.setattr(litellm, "api_base", "https://gateway.internal.example/v1") + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + with respx.mock: + route: Final = respx.post("https://gateway.internal.example/v1/responses/input_tokens").mock( + return_value=httpx.Response(200, json={"input_tokens": 4}) + ) + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="openai/gpt-4o", + ) + assert result is not None + assert result.total_tokens == 4 + assert route.calls.last.request.headers["authorization"] == "Bearer sk-private-gateway" diff --git a/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py b/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py index c2efc1acdb9..8850fd13d10 100644 --- a/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py +++ b/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py @@ -1,6 +1,10 @@ +from typing import Final import pytest +from litellm.llms.anthropic.pass_through.responses_adapters.transformation import ( + LiteLLMAnthropicToResponsesAPIAdapter, +) from litellm.llms.openai.responses.count_tokens.transformation import ( OpenAICountTokensConfig, ) @@ -446,3 +450,108 @@ def test_get_required_headers(): assert headers["Authorization"] == "Bearer sk-test-key" assert headers["Content-Type"] == "application/json" + + +PDF_BASE64: Final = "JVBERi0xLjcK" +PNG_BASE64: Final = "iVBORw0KGgo=" + +ANTHROPIC_AGENT_TURN: Final[list[dict[str, object]]] = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the indemnity terms."}, + { + "type": "document", + "title": "contract.pdf", + "source": {"type": "base64", "media_type": "application/pdf", "data": PDF_BASE64}, + }, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": PNG_BASE64}}, + ], + }, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Let me pull clause 14."}, + {"type": "tool_use", "id": "toolu_01", "name": "lookup_clause", "input": {"id": "14"}}, + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "14. Indemnification."}], + }, +] + +ANTHROPIC_TOOLS: Final[list[dict[str, object]]] = [ + { + "name": "lookup_clause", + "description": "Fetch a clause by id", + "input_schema": {"type": "object", "properties": {"id": {"type": "string"}}}, + } +] + + +@pytest.mark.parametrize( + "system", + [None, "You review contracts.", [{"type": "text", "text": "You review contracts."}]], +) +def test_anthropic_messages_convert_exactly_as_v1_messages_inference_converts_them(system: object) -> None: + inference_request: Final[dict[str, object]] = { + "model": "gpt-4o", + "messages": ANTHROPIC_AGENT_TURN, + "tools": ANTHROPIC_TOOLS, + **({"system": system} if system is not None else {}), + } + inference: Final = LiteLLMAnthropicToResponsesAPIAdapter().translate_request( + inference_request, include_encrypted_reasoning=False + ) + + counted: Final = OpenAICountTokensConfig.anthropic_messages_to_responses_input( + model="gpt-4o", messages=ANTHROPIC_AGENT_TURN, system=system, tools=ANTHROPIC_TOOLS + ) + + assert counted == ( + tuple(inference["input"]), + inference.get("instructions"), + tuple(inference["tools"]), + inference.get("text"), + ) + + +def test_anthropic_messages_keep_every_block_the_chat_converter_drops() -> None: + input_items, _, _, _ = OpenAICountTokensConfig.anthropic_messages_to_responses_input( + model="gpt-4o", messages=ANTHROPIC_AGENT_TURN, system=None, tools=None + ) + + assert input_items == ( + { + "type": "message", + "role": "user", + "content": [ + {"type": "input_text", "text": "Summarize the indemnity terms."}, + { + "type": "input_file", + "filename": "contract.pdf", + "file_data": f"data:application/pdf;base64,{PDF_BASE64}", + }, + {"type": "input_image", "image_url": f"data:image/png;base64,{PNG_BASE64}"}, + ], + }, + {"type": "function_call", "call_id": "toolu_01", "name": "lookup_clause", "arguments": '{"id": "14"}'}, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Let me pull clause 14."}], + }, + {"type": "function_call_output", "call_id": "toolu_01", "output": "14. Indemnification."}, + ) + + +def test_anthropic_list_system_becomes_joined_instruction_text() -> None: + _, instructions, _, _ = OpenAICountTokensConfig.anthropic_messages_to_responses_input( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + system=[{"type": "text", "text": "First rule."}, {"type": "text", "text": "Second rule."}], + tools=None, + ) + + assert instructions == "First rule.\nSecond rule." diff --git a/tests/unit/llms/openai/test_openai_workload_identity.py b/tests/unit/llms/openai/test_openai_workload_identity.py index db107e00df0..aa0cad178f3 100644 --- a/tests/unit/llms/openai/test_openai_workload_identity.py +++ b/tests/unit/llms/openai/test_openai_workload_identity.py @@ -279,3 +279,56 @@ class TestResponsesValidateEnvironment: headers={}, model="gpt-4o-mini", litellm_params=GenericLiteLLMParams() ) assert headers["Authorization"] == "Bearer None" + + +class TestCountTokens: + @pytest.mark.asyncio + @respx.mock + async def test_count_authenticates_with_the_workload_identity_bearer( + self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + mock_token_exchange() + count_route: Final = respx.post("https://api.openai.com/v1/responses/input_tokens").mock( + return_value=httpx.Response(200, json={"input_tokens": 12}) + ) + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="openai/gpt-4o", + ) + assert result is not None + assert result.error is False + assert result.total_tokens == 12 + assert count_route.calls.last.request.headers["authorization"] == "Bearer exchanged-bearer-token" + + @pytest.mark.asyncio + @respx.mock + async def test_failed_mint_becomes_a_count_error_response( + self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + respx.post(TOKEN_EXCHANGE_URL).mock(return_value=httpx.Response(400, json={"error": "invalid_grant"})) + count_route: Final = respx.post("https://api.openai.com/v1/responses/input_tokens") + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="openai/gpt-4o", + ) + assert result is not None + assert result.error is True + assert result.status_code == 401 + assert result.total_tokens == 0 + assert not count_route.called diff --git a/tests/unit/llms/vertex_ai/count_tokens/__init__.py b/tests/unit/llms/vertex_ai/count_tokens/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/vertex_ai/count_tokens/test_context.py b/tests/unit/llms/vertex_ai/count_tokens/test_context.py new file mode 100644 index 00000000000..235c3710ea5 --- /dev/null +++ b/tests/unit/llms/vertex_ai/count_tokens/test_context.py @@ -0,0 +1,507 @@ +import asyncio +import base64 +import time +import uuid +from collections.abc import Generator +from contextlib import contextmanager +from copy import deepcopy +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from threading import Thread +from typing import Final, Literal +from unittest.mock import patch + +import pytest +from pydantic import JsonValue, TypeAdapter + +import litellm +from litellm.llms.anthropic.pass_through.adapters.transformation import ( + LiteLLMAnthropicMessagesAdapter, + truncate_tool_name, +) +from litellm.llms.vertex_ai.count_tokens.context import ( + VertexGeminiCountContext, + VertexGeminiCountContextError, + gemini_count_contents, + inference_request_view, + vertex_gemini_count_context, +) +from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig + +_SCHEMA: Final[JsonValue] = { + "type": "object", + "properties": {}, + "additionalProperties": False, +} +_FUNCTION_DECLARATION: Final = { + "functionDeclarations": [ + { + "name": "lookup", + "description": "Look up a value", + "parameters": {"type": "object"}, + } + ] +} + + +@pytest.mark.parametrize( + "tool", + ( + { + "name": "lookup", + "description": "Look up a value", + "input_schema": _SCHEMA, + }, + { + "type": "function", + "function": { + "name": "lookup", + "description": "Look up a value", + "parameters": _SCHEMA, + }, + }, + { + "type": "function", + "name": "lookup", + "description": "Look up a value", + "parameters": _SCHEMA, + }, + ), + ids=("anthropic", "openai-chat", "responses"), +) +def test_function_tool_shapes_map_to_vertex_declarations(tool: JsonValue) -> None: + context: Final = vertex_gemini_count_context(system=None, tools=[tool]) + + assert isinstance(context, VertexGeminiCountContext) + assert context.tools == (_FUNCTION_DECLARATION,) + assert _SCHEMA["additionalProperties"] is False + + +def test_long_anthropic_tool_name_is_truncated_for_vertex() -> None: + name: Final = "mcp__docs__" + "a" * 70 + context: Final = vertex_gemini_count_context( + system=None, + tools=[{"name": name, "description": "Look up a value", "input_schema": _SCHEMA}], + ) + + assert isinstance(context, VertexGeminiCountContext) + assert context.tools == ( + { + "functionDeclarations": [ + { + "name": truncate_tool_name(name), + "description": "Look up a value", + "parameters": {"type": "object"}, + } + ] + }, + ) + assert len(truncate_tool_name(name)) == 64 + + +def test_64_character_anthropic_tool_name_is_unchanged_for_vertex() -> None: + name: Final = "mcp__docs__" + "a" * 53 + context: Final = vertex_gemini_count_context( + system=None, + tools=[{"name": name, "input_schema": _SCHEMA}], + ) + + assert isinstance(context, VertexGeminiCountContext) + assert context.tools == ({"functionDeclarations": [{"name": name, "parameters": {"type": "object"}}]},) + + +def test_long_openai_chat_tool_name_is_unchanged_for_vertex() -> None: + name: Final = "mcp__docs__" + "a" * 70 + context: Final = vertex_gemini_count_context( + system=None, + tools=[{"type": "function", "function": {"name": name, "parameters": _SCHEMA}}], + ) + + assert isinstance(context, VertexGeminiCountContext) + assert context.tools == ({"functionDeclarations": [{"name": name, "parameters": {"type": "object"}}]},) + + +def test_unknown_tools_are_omitted() -> None: + context: Final = vertex_gemini_count_context( + system=None, + tools=[{"type": "web_search_preview"}, {"type": "computer_use"}], + ) + + assert isinstance(context, VertexGeminiCountContext) + assert context.tools is None + + +def test_text_block_system_instruction_keeps_one_part_per_block() -> None: + context: Final = vertex_gemini_count_context( + system=[{"type": "text", "text": "First instruction."}, {"type": "text", "text": "Second instruction."}], + tools=None, + ) + + assert isinstance(context, VertexGeminiCountContext) + assert context.system_instruction == {"parts": ({"text": "First instruction."}, {"text": "Second instruction."})} + + +def test_string_system_instruction_maps_to_one_part() -> None: + context: Final = vertex_gemini_count_context(system="Follow this instruction.", tools=None) + + assert isinstance(context, VertexGeminiCountContext) + assert context.system_instruction == {"parts": ({"text": "Follow this instruction."},)} + + +def test_non_text_system_block_returns_an_error_value() -> None: + context: Final = vertex_gemini_count_context( + system=[{"type": "image", "source": {"type": "base64", "data": "AAAA"}}], + tools=None, + ) + + assert isinstance(context, VertexGeminiCountContextError) + assert "text blocks" in context.error_message + + +@pytest.mark.parametrize(("system", "tools"), ((None, None), ("", None), ([], []))) +def test_empty_context_is_omitted(system: JsonValue | None, tools: JsonValue | None) -> None: + context: Final = vertex_gemini_count_context(system=system, tools=tools) + + assert isinstance(context, VertexGeminiCountContext) + assert context.system_instruction is None + assert context.tools is None + + +@pytest.mark.parametrize( + "tool", + ( + {"description": "Look up a value", "input_schema": _SCHEMA}, + {"type": "bash_20250124", "name": "bash"}, + ), + ids=("unnamed", "anthropic-hosted"), +) +def test_anthropic_tools_count_the_same_declarations_inference_sends(tool: dict[str, JsonValue]) -> None: + translated_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai([deepcopy(tool)]) # pyright: ignore[reportArgumentType] # raw Anthropic tool dict + inference_tools: Final = VertexGeminiConfig()._map_function( # pyright: ignore[reportPrivateUsage] # the tool mapping inference uses + value=[dict(translated) for translated in translated_tools], + optional_params={}, + ) + + context: Final = vertex_gemini_count_context(system=None, tools=[tool]) + + assert isinstance(context, VertexGeminiCountContext) + assert context.tools is not None + assert [tool_value["functionDeclarations"] for tool_value in context.tools] == [ + inference_tool["function_declarations"] for inference_tool in inference_tools + ] + + +@pytest.mark.parametrize( + "tools", + ( + "lookup", + [1], + [{"name": "lookup", "input_schema": "not-a-schema"}], + [{"type": "function", "function": {"name": "lookup", "parameters": [1]}}], + [{"name": "lookup", "input_schema": {"type": "object", "properties": {"query": "not-a-schema"}}}], + ), + ids=("not-a-list", "non-object-tool", "string-schema", "list-parameters", "malformed-property"), +) +def test_invalid_tools_return_an_error_value(tools: JsonValue) -> None: + context: Final = vertex_gemini_count_context(system=None, tools=tools) + + assert isinstance(context, VertexGeminiCountContextError) + assert "tools" in context.error_message + + +_ANTHROPIC_CONVERSATION: Final = ( + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize both"}, + {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0="}}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}}, + ], + }, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {"query": "x"}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "42"}]}, +) + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +async def test_anthropic_documents_images_and_tool_turns_become_gemini_parts( + custom_llm_provider: Literal["vertex_ai", "gemini"], +) -> None: + contents: Final = await gemini_count_contents( + _ANTHROPIC_CONVERSATION, "anthropic", "gemini-2.5-flash", custom_llm_provider, {} + ) + + assert contents == [ + { + "role": "user", + "parts": [ + {"text": "Summarize both"}, + {"inline_data": {"data": "JVBERi0=", "mime_type": "application/pdf"}}, + {"inline_data": {"data": "aW1hZ2U=", "mime_type": "image/png"}}, + ], + }, + {"role": "model", "parts": [{"function_call": {"name": "lookup", "args": {"query": "x"}}}]}, + {"role": "user", "parts": [{"function_response": {"name": "lookup", "response": {"content": "42"}}}]}, + ] + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +async def test_openai_messages_keep_their_chat_conversion(custom_llm_provider: Literal["vertex_ai", "gemini"]) -> None: + contents: Final = await gemini_count_contents( + ({"role": "user", "content": "hello"}, {"role": "assistant", "content": "hi"}), + "openai", + "gemini-2.5-flash", + custom_llm_provider, + {}, + ) + + assert contents == [ + {"role": "user", "parts": [{"text": "hello"}]}, + {"role": "model", "parts": [{"text": "hi"}]}, + ] + + +_PNG_BYTES: Final = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" +) +_DOWNLOAD_SECONDS: Final = 0.5 + + +class _SlowPngHandler(BaseHTTPRequestHandler): + def do_GET(self) -> None: + time.sleep(_DOWNLOAD_SECONDS) + self.send_response(200) + self.send_header("Content-Type", "image/png") + self.send_header("Content-Length", str(len(_PNG_BYTES))) + self.end_headers() + self.wfile.write(_PNG_BYTES) + + def log_message(self, format: str, *args: object) -> None: + return None + + +@contextmanager +def _slow_png_server() -> Generator[str, None, None]: + with ThreadingHTTPServer(("127.0.0.1", 0), _SlowPngHandler) as server: + thread: Final = Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}" + finally: + server.shutdown() + + +async def _event_loop_wake_delay(nap_seconds: float) -> float: + started: Final = time.monotonic() + await asyncio.sleep(nap_seconds) + return time.monotonic() - started + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +async def test_remote_image_downloads_without_blocking_the_event_loop( + custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", False) + with _slow_png_server() as base_url: + image_url: Final = f"{base_url}/{uuid.uuid4().hex}.png" + wake_delay, contents = await asyncio.gather( + _event_loop_wake_delay(0.05), + gemini_count_contents( + ( + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe it"}, + {"type": "image", "source": {"type": "url", "url": image_url}}, + ], + }, + ), + "anthropic", + "gemini-2.5-flash", + custom_llm_provider, + {}, + ), + ) + + assert contents == [ + { + "role": "user", + "parts": [ + {"text": "Describe it"}, + {"inline_data": {"data": base64.b64encode(_PNG_BYTES).decode(), "mime_type": "image/png"}}, + ], + } + ] + assert wake_delay < _DOWNLOAD_SECONDS / 2 + + +class _MissingImageHandler(BaseHTTPRequestHandler): + def do_GET(self) -> None: + self.send_response(404) + self.send_header("Content-Length", "0") + self.end_headers() + + def log_message(self, format: str, *args: object) -> None: + return None + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +async def test_image_that_cannot_be_fetched_returns_an_error_value( + custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", False) + with ThreadingHTTPServer(("127.0.0.1", 0), _MissingImageHandler) as server: + Thread(target=server.serve_forever, daemon=True).start() + try: + contents = await gemini_count_contents( + ( + { + "role": "user", + "content": [ + { + "type": "image", + "source": { + "type": "url", + "url": f"http://127.0.0.1:{server.server_address[1]}/{uuid.uuid4().hex}.png", + }, + } + ], + }, + ), + "anthropic", + "gemini-2.5-flash", + custom_llm_provider, + {}, + ) + finally: + server.shutdown() + + assert isinstance(contents, VertexGeminiCountContextError) + + +@pytest.mark.parametrize( + "system", ("Be brief.", [{"type": "text", "text": "Be brief."}]), ids=("string-system", "text-block-system") +) +def test_anthropic_view_keeps_only_the_turns_after_a_client_compaction_block(system: JsonValue) -> None: + view: Final = inference_request_view( + ( + {"role": "user", "content": "first question " * 50}, + {"role": "assistant", "content": [{"type": "compaction", "content": "Summary of the earlier turns."}]}, + {"role": "user", "content": "latest question"}, + ), + system, + "anthropic", + ) + + assert view.messages == [{"role": "user", "content": "latest question"}] + assert "Summary of the earlier turns." in str(view.system) + assert "Be brief." in str(view.system) + + +def test_anthropic_view_flattens_a_replayed_web_search_result_into_text() -> None: + snippet: Final = "LiteLLM ships a weekly release candidate and a stable line." + view: Final = inference_request_view( + ( + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": "", + "snippet": snippet, + } + ], + }, + ], + }, + ), + None, + "anthropic", + ) + + assistant_blocks: Final = TypeAdapter(list[dict[str, str]]).validate_python(view.messages[1]["content"]) + assert [block["type"] for block in assistant_blocks] == ["text"] + assert snippet in assistant_blocks[0]["text"] + + +def test_openai_view_is_the_request_unchanged() -> None: + messages: Final = ({"role": "user", "content": [{"type": "compaction", "content": "kept"}]},) + + view: Final = inference_request_view(messages, "Be brief.", "openai") + + assert view.messages == messages + assert view.system == "Be brief." + + +async def test_gcs_media_type_lookup_uses_the_deployment_vertex_project() -> None: + with patch( + "litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type", + return_value="image/png", + ) as lookup: + contents: Final = await gemini_count_contents( + ( + { + "role": "user", + "content": [ + {"type": "image", "source": {"type": "url", "url": "gs://bucket/image-without-extension"}} + ], + }, + ), + "anthropic", + "gemini-2.5-flash", + "vertex_ai", + {"vertex_project": "deployment-project"}, + ) + + assert lookup.call_args.kwargs.get("vertex_project") == "deployment-project" + assert not isinstance(contents, VertexGeminiCountContextError) + assert contents[0]["parts"][0] == { + "file_data": {"mime_type": "image/png", "file_uri": "gs://bucket/image-without-extension"} + } + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +def test_anthropic_replayed_web_search_result_is_counted_as_its_text( + custom_llm_provider: Literal["vertex_ai", "gemini"], +) -> None: + snippet: Final = "LiteLLM ships a weekly release candidate and a stable line." + messages: Final = ( + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": "", + "snippet": snippet, + } + ], + }, + ], + }, + {"role": "user", "content": "Thanks."}, + ) + + view: Final = inference_request_view(messages, None, "anthropic") + contents: Final = asyncio.run( + gemini_count_contents(view.messages, "anthropic", "gemini-2.5-flash", custom_llm_provider, {}) + ) + + assert snippet in str(contents) diff --git a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py index 04a7ee451c4..f9561a0d710 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -1,11 +1,18 @@ -from unittest.mock import patch +from collections.abc import Iterator +from contextlib import contextmanager +from typing import Final +from unittest.mock import AsyncMock, patch +import httpx import pytest +from pydantic import JsonValue, TypeAdapter +import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH - - +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler from litellm.llms.vertex_ai.common_utils import ( + VertexAITokenCounter, _get_vertex_url, convert_anyof_null_to_nullable, get_vertex_location_from_url, @@ -16,6 +23,110 @@ from litellm.llms.vertex_ai.common_utils import ( validate_vertex_location, vertex_request_labels_from_litellm_params, ) +from litellm.llms.vertex_ai.count_tokens import handler as vertex_count_tokens_handler +from litellm.types.utils import TokenCountMessagesFormat + + +@pytest.mark.asyncio +@pytest.mark.parametrize("messages_format", ("openai", "anthropic")) +async def test_vertex_gemini_rejects_non_text_system_blocks_before_request( + messages_format: TokenCountMessagesFormat, +) -> None: + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="gemini-4-argon", + messages=[{"role": "user", "content": "Count this"}], + contents=None, + request_model="vertex_ai/gemini-4-argon", + system=[{"type": "image", "source": {"type": "base64", "data": "AAAA"}}], + messages_format=messages_format, + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert "text blocks" in (result.error_message or "") + + +class _VertexCountTokensPeer: + def __init__(self, status_code: int) -> None: + self.status_code: Final = status_code + self.requests: tuple[httpx.Request, ...] = () + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.requests = (*self.requests, request) + if self.status_code == 200: + return httpx.Response(200, json={"totalTokens": 37}, request=request) + return httpx.Response( + self.status_code, + json={"error": {"code": self.status_code, "message": "Invalid function name", "status": "INVALID_ARGUMENT"}}, + request=request, + ) + + +@contextmanager +def _vertex_count_tokens_peer(status_code: int) -> Iterator[_VertexCountTokensPeer]: + peer: Final = _VertexCountTokensPeer(status_code) + with ( + patch.object( + gemini_count_tokens_handler, + "get_async_httpx_client", + return_value=AsyncHTTPHandler(transport=httpx.MockTransport(peer)), + ), + patch.object( + vertex_count_tokens_handler.VertexAITokenCounter, + "_ensure_access_token_async", + new=AsyncMock(return_value=("fake-token", "test-project")), + ), + ): + yield peer + + +_VERTEX_DEPLOYMENT: Final = {"litellm_params": {"vertex_project": "test-project", "vertex_location": "us-central1"}} + + +@pytest.mark.asyncio +async def test_vertex_gemini_count_answers_with_the_status_google_rejected_it_with() -> None: + with _vertex_count_tokens_peer(400) as peer: + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=[{"role": "user", "content": "Hello"}], + contents=None, + deployment=_VERTEX_DEPLOYMENT, + request_model="vertex_ai/gemini-2.5-flash", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert "Invalid function name" in (result.error_message or "") + assert len(peer.requests) == 1 + + +@pytest.mark.asyncio +async def test_vertex_gemini_count_rejects_an_image_url_it_refuses_to_fetch_without_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", True) + with _vertex_count_tokens_peer(200) as peer: + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=[ + { + "role": "user", + "content": [{"type": "image", "source": {"type": "url", "url": "http://127.0.0.1:9/image.png"}}], + } + ], + contents=None, + deployment=_VERTEX_DEPLOYMENT, + request_model="vertex_ai/gemini-2.5-flash", + messages_format="anthropic", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert "http://127.0.0.1:9/image.png" in (result.error_message or "") + assert peer.requests == () @pytest.mark.parametrize("location", ["us", "eu", "us-central1", "europe-west1", "global"]) @@ -149,15 +260,11 @@ async def test_get_supports_system_message(): from litellm.llms.vertex_ai.common_utils import get_supports_system_message # fine-tuned vertex gemini models will specifiy they are in the /gemini spec format - result = get_supports_system_message( - model="gemini/1234567890", custom_llm_provider="vertex_ai" - ) + result = get_supports_system_message(model="gemini/1234567890", custom_llm_provider="vertex_ai") assert result == True # non-fine-tuned vertex gemini models will not specifiy they are in the /gemini spec format - result = get_supports_system_message( - model="random-model-name", custom_llm_provider="vertex_ai" - ) + result = get_supports_system_message(model="random-model-name", custom_llm_provider="vertex_ai") assert result == False @@ -230,13 +337,9 @@ def test_build_vertex_schema(): "properties": { "tags": {"items": {"type": "string"}, "type": "array"}, "metadata": {"type": "object"}, - "callbacks": { - "anyOf": [{"items": {}, "type": "array"}, {}, {"type": "null"}] - }, + "callbacks": {"anyOf": [{"items": {}, "type": "array"}, {}, {"type": "null"}]}, "run_name": {"type": "string"}, - "max_concurrency": { - "anyOf": [{"type": "integer"}, {"type": "null"}] - }, + "max_concurrency": {"anyOf": [{"type": "integer"}, {"type": "null"}]}, "recursion_limit": {"type": "integer"}, "configurable": {"type": "object"}, "run_id": { @@ -280,9 +383,7 @@ def test_build_vertex_schema(): ] }, "run_name": {"type": "string"}, - "max_concurrency": { - "anyOf": [{"type": "integer", "nullable": True}] - }, + "max_concurrency": {"anyOf": [{"type": "integer", "nullable": True}]}, "recursion_limit": {"type": "integer"}, "configurable": {"type": "object"}, "run_id": {"anyOf": [{"type": "string", "nullable": True}]}, @@ -689,9 +790,7 @@ def test_get_vertex_url_global_region(stream, expected_endpoint_suffix): ({}, "us-east1", "us-east1"), ], ) -def test_get_vertex_region_global_only_model( - model_cost_entry, vertex_region, expected_region -): +def test_get_vertex_region_global_only_model(model_cost_entry, vertex_region, expected_region): """Test get_vertex_region resolves region from model_cost supported_regions""" import litellm from litellm.llms.vertex_ai.vertex_llm_base import VertexBase @@ -703,9 +802,7 @@ def test_get_vertex_region_global_only_model( {"vertex_ai/test-model": model_cost_entry}, clear=False, ): - result = vertex_base.get_vertex_region( - vertex_region=vertex_region, model="test-model" - ) + result = vertex_base.get_vertex_region(vertex_region=vertex_region, model="test-model") assert result == expected_region @@ -823,9 +920,7 @@ def test_convert_schema_types_type_array_conversion(): assert anyof_types[1]["type"] == "number" # 4. Other properties preserved - assert ( - input_schema["properties"]["studio"]["description"] == "The studio ID or name" - ) + assert input_schema["properties"]["studio"]["description"] == "The studio ID or name" assert input_schema["required"] == ["studio"] @@ -992,7 +1087,9 @@ def test_construct_target_url_with_version_prefix(): ), ], ) -def test_construct_target_url_versionless_project_route_gets_api_version(requested_route: str, expected_url: str) -> None: +def test_construct_target_url_versionless_project_route_gets_api_version( + requested_route: str, expected_url: str +) -> None: from litellm.llms.vertex_ai.common_utils import construct_target_url target_url = construct_target_url( @@ -1125,10 +1222,7 @@ def test_fix_enum_types(): # 2. Non-string enums are removed assert "enum" not in input_schema["properties"]["maxLength"] assert "enum" not in input_schema["properties"]["enabled"] - assert ( - "enum" - not in input_schema["properties"]["nested"]["properties"]["innerNonStringEnum"] - ) + assert "enum" not in input_schema["properties"]["nested"]["properties"]["innerNonStringEnum"] # 3. anyOf with string type keeps enum, non-string removes it assert "enum" in input_schema["properties"]["anyOfField"]["anyOf"][0] @@ -1182,8 +1276,6 @@ def test_get_token_url(): assert "v1beta1" not in url assert "/v1/" in url - pass - @pytest.mark.asyncio async def test_vertex_ai_token_counter_routes_partner_models(): @@ -1229,6 +1321,149 @@ async def test_vertex_ai_token_counter_routes_partner_models(): assert result.tokenizer_type == "vertex_ai_partner_models" +@pytest.mark.asyncio +async def test_vertex_ai_token_counter_forwards_system_and_tools_to_partner_request(): + from typing import Final + from unittest.mock import AsyncMock, patch + + from litellm.llms.vertex_ai.common_utils import VertexAITokenCounter + from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens import handler + + class FakeResponse: + status_code = 200 + + def json(self) -> dict[str, int]: + return {"input_tokens": 37} + + class FakeHttpClient: + posted_bodies: tuple[dict[str, object], ...] = () + + async def post( + self, + url: str, + headers: dict[str, str], + json: dict[str, object], + timeout: float, + ) -> FakeResponse: + self.posted_bodies = (*self.posted_bodies, json) + return FakeResponse() + + fake_http_client: Final = FakeHttpClient() + counter: Final = VertexAITokenCounter() + model: Final = "claude-opus-5-5" + messages: Final = [{"role": "user", "content": "Hello"}] + system: Final = "Follow the system instructions" + tools: Final = [ + { + "name": "lookup", + "description": "Look up a value", + "input_schema": {"type": "object", "properties": {}}, + } + ] + deployment: Final = { + "litellm_params": { + "vertex_project": "test-project", + "vertex_location": "us-east5", + } + } + + with ( + patch.object(handler, "get_async_httpx_client", return_value=fake_http_client), + patch.object( + handler.VertexAIPartnerModelsTokenCounter, + "_ensure_access_token_async", + new=AsyncMock(return_value=("fake-token", "test-project")), + ), + ): + with_optional_fields: Final = await counter.count_tokens( + model_to_use=model, + messages=messages, + contents=None, + deployment=deployment, + system=system, + tools=tools, + ) + without_optional_fields: Final = await counter.count_tokens( + model_to_use=model, + messages=messages, + contents=None, + deployment=deployment, + ) + + assert fake_http_client.posted_bodies == ( + {"model": model, "messages": messages, "system": system, "tools": tools}, + {"model": model, "messages": messages}, + ) + assert with_optional_fields is not None + assert with_optional_fields.total_tokens == 37 + assert without_optional_fields is not None + assert without_optional_fields.total_tokens == 37 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("provider_failure", "expected_status", "expected_message"), + [ + ("http_400", 400, 'tools.0: Input tag "function" does not match any of the expected tags'), + ("credentials", 500, "could not resolve credentials"), + ], +) +async def test_vertex_ai_token_counter_returns_partner_provider_error_as_value( + provider_failure: str, expected_status: int, expected_message: str +): + from typing import Final + from unittest.mock import AsyncMock, patch + + import httpx + + from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError + from litellm.llms.vertex_ai.common_utils import VertexAITokenCounter + from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens import handler + + class RejectingHttpClient: + async def post( + self, + url: str, + headers: dict[str, str], + json: dict[str, object], + timeout: float, + ) -> None: + request: Final = httpx.Request("POST", url) + response: Final = httpx.Response(400, text=expected_message, request=request) + raise MaskedHTTPStatusError( + httpx.HTTPStatusError("Client error '400 Bad Request'", request=request, response=response), + message=expected_message, + text=expected_message, + ) + + access_token: Final = ( + AsyncMock(side_effect=ValueError(expected_message)) + if provider_failure == "credentials" + else AsyncMock(return_value=("fake-token", "test-project")) + ) + with ( + patch.object(handler, "get_async_httpx_client", return_value=RejectingHttpClient()), + patch.object(handler.VertexAIPartnerModelsTokenCounter, "_ensure_access_token_async", new=access_token), + ): + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="claude-opus-5-5", + messages=[{"role": "user", "content": "Hello"}], + contents=None, + deployment={"litellm_params": {"vertex_project": "test-project", "vertex_location": "us-east5"}}, + request_model="vertex-claude", + tools=[{"type": "function", "function": {"name": "lookup", "parameters": {}}}], + ) + + assert result is not None + assert result.error is True + assert result.status_code == expected_status + assert result.error_message is not None + assert expected_message in result.error_message + assert result.total_tokens == 0 + assert result.request_model == "vertex-claude" + assert result.tokenizer_type == "vertex_ai_partner_models" + + @pytest.mark.asyncio async def test_vertex_ai_token_counter_uses_count_tokens_location(): """ @@ -1318,6 +1553,49 @@ async def test_vertex_ai_token_counter_routes_gemini_models(): assert result.total_tokens == 50 +@pytest.mark.asyncio +async def test_vertex_gemini_count_tokens_uses_request_system_over_deployment_system_instruction() -> None: + class VertexCountTokensPeer: + def __init__(self) -> None: + self.request_body: dict[str, JsonValue] | None = None + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.request_body = TypeAdapter(dict[str, JsonValue]).validate_json(request.content) + return httpx.Response(200, json={"totalTokens": 37}, request=request) + + peer: Final = VertexCountTokensPeer() + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer)) + + with ( + patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client), + patch.object( + vertex_count_tokens_handler.VertexAITokenCounter, + "_ensure_access_token_async", + new=AsyncMock(return_value=("fake-token", "test-project")), + ), + ): + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=[{"role": "user", "content": "Hello"}], + contents=None, + deployment={ + "litellm_params": { + "vertex_project": "test-project", + "vertex_location": "us-central1", + "system_instruction": {"parts": [{"text": "deployment-level"}]}, + } + }, + request_model="vertex_ai/gemini-2.5-flash", + system="You are terse.", + ) + + assert result is not None + assert result.error is False + assert result.total_tokens == 37 + assert peer.request_body is not None + assert peer.request_body["systemInstruction"] == {"parts": [{"text": "You are terse."}]} + + @pytest.mark.asyncio async def test_vertex_ai_token_counter_converts_messages_to_contents_for_gemini(): """ @@ -1332,9 +1610,7 @@ async def test_vertex_ai_token_counter_converts_messages_to_contents_for_gemini( token_counter = VertexAITokenCounter() - with patch( - "litellm.llms.vertex_ai.count_tokens.handler.VertexAITokenCounter.acount_tokens" - ) as mock_acount_tokens: + with patch("litellm.llms.vertex_ai.count_tokens.handler.VertexAITokenCounter.acount_tokens") as mock_acount_tokens: mock_acount_tokens.return_value = { "totalTokens": 42, "tokenizer_used": "gemini", @@ -1376,9 +1652,7 @@ async def test_vertex_ai_token_counter_returns_none_when_api_omits_total_tokens( token_counter = VertexAITokenCounter() - with patch( - "litellm.llms.vertex_ai.count_tokens.handler.VertexAITokenCounter.acount_tokens" - ) as mock_acount_tokens: + with patch("litellm.llms.vertex_ai.count_tokens.handler.VertexAITokenCounter.acount_tokens") as mock_acount_tokens: mock_acount_tokens.return_value = {"tokenizer_used": "gemini"} result = await token_counter.count_tokens( @@ -1421,9 +1695,7 @@ async def test_vertex_ai_partner_model_detection(): # Test Minimax models assert VertexAIPartnerModels.is_vertex_partner_model("minimaxai/minimax-m2-maas") # Test Moonshot models - assert VertexAIPartnerModels.is_vertex_partner_model( - "moonshotai/kimi-k2-thinking-maas" - ) + assert VertexAIPartnerModels.is_vertex_partner_model("moonshotai/kimi-k2-thinking-maas") # Test Gemini models (should NOT be detected as partner model) assert not VertexAIPartnerModels.is_vertex_partner_model("gemini-1.5-pro") @@ -1454,9 +1726,7 @@ def test_vertex_ai_moonshot_uses_openai_handler(): VertexAIPartnerModels, ) - assert VertexAIPartnerModels.should_use_openai_handler( - "moonshotai/kimi-k2-thinking-maas" - ) + assert VertexAIPartnerModels.should_use_openai_handler("moonshotai/kimi-k2-thinking-maas") def test_vertex_ai_zai_uses_openai_handler(): @@ -1491,9 +1761,7 @@ def test_vertex_ai_gemma_maas_is_partner_model(): VertexAIPartnerModels, ) - assert VertexAIPartnerModels.is_vertex_partner_model( - "google/gemma-4-26b-a4b-it-maas" - ) + assert VertexAIPartnerModels.is_vertex_partner_model("google/gemma-4-26b-a4b-it-maas") def test_vertex_ai_gemma_maas_uses_openai_handler(): @@ -1504,9 +1772,7 @@ def test_vertex_ai_gemma_maas_uses_openai_handler(): VertexAIPartnerModels, ) - assert VertexAIPartnerModels.should_use_openai_handler( - "google/gemma-4-26b-a4b-it-maas" - ) + assert VertexAIPartnerModels.should_use_openai_handler("google/gemma-4-26b-a4b-it-maas") def test_vertex_ai_gemma_maas_routes_to_partner_models(): @@ -1708,12 +1974,8 @@ def test_pop_vertex_request_labels_prefers_explicit_labels_then_metadata(): def test_pop_vertex_request_labels_uses_litellm_metadata_when_metadata_absent(): optional: dict = {} - litellm_params = { - "litellm_metadata": {"requester_metadata": {"team": "from_litellm_meta"}} - } - assert pop_vertex_request_labels(optional, litellm_params) == { - "team": "from_litellm_meta" - } + litellm_params = {"litellm_metadata": {"requester_metadata": {"team": "from_litellm_meta"}}} + assert pop_vertex_request_labels(optional, litellm_params) == {"team": "from_litellm_meta"} def test_vertex_text_embedding_request_includes_labels_from_metadata(): @@ -1723,9 +1985,7 @@ def test_vertex_text_embedding_request_includes_labels_from_metadata(): input="hi", optional_params={}, model="text-embedding-004", - litellm_params={ - "metadata": {"requester_metadata": {"project_id": "cost-center-1"}} - }, + litellm_params={"metadata": {"requester_metadata": {"project_id": "cost-center-1"}}}, ) assert req.get("labels") == {"project_id": "cost-center-1"} diff --git a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py index 801f61aa498..1351680d617 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py @@ -284,7 +284,9 @@ class TestHttpExceptionDictDetail: "base_process_llm_request", new=AsyncMock(side_effect=exc), ), - patch.object(proxy_server, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global imported at call time; no injection seam + patch.object( + proxy_server, "proxy_logging_obj" + ) as mock_logging, # test-quality-ok: module global imported at call time; no injection seam ): mock_logging.post_call_failure_hook = AsyncMock() response = await ep.anthropic_response( @@ -424,8 +426,6 @@ class TestErrorLogCarriesCallId: @pytest.mark.asyncio async def test_count_tokens_failure_log_carries_callers_call_id(self, caplog: pytest.LogCaptureFixture): - from fastapi import HTTPException - import litellm.proxy.anthropic_endpoints.endpoints as ep import litellm.proxy.proxy_server as proxy_server from litellm.proxy._types import UserAPIKeyAuth @@ -440,13 +440,14 @@ class TestErrorLogCarriesCallId: "_read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet", "messages": [{"role": "user", "content": "hi"}]}), ), - patch.object(proxy_server, "token_counter", new=AsyncMock(side_effect=RuntimeError("tokenizer down"))), # test-quality-ok: module global imported at call time; the test targets the endpoint's except block + patch.object( + proxy_server, "count_request_tokens", new=AsyncMock(side_effect=RuntimeError("tokenizer down")) + ), # test-quality-ok: module global imported at call time; the test targets the endpoint's except block caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), - pytest.raises(HTTPException) as raised, ): - await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) + response = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) - assert raised.value.status_code == 500 + assert response.status_code == 500 record = self._error_record(caplog) assert record.litellm_call_id == call_id assert call_id in record.getMessage() diff --git a/tests/unit/proxy/common_utils/test_token_count_settings.py b/tests/unit/proxy/common_utils/test_token_count_settings.py new file mode 100644 index 00000000000..84a32a02a5c --- /dev/null +++ b/tests/unit/proxy/common_utils/test_token_count_settings.py @@ -0,0 +1,190 @@ +import json +from typing import Final + +import pytest + +from litellm.proxy.common_utils.token_count_settings import ( + ProviderCountFailure, + TokenCountModel, + TokenCountSettings, + count_source_fields, + resolve_token_count_settings, +) +from litellm.types.utils import TokenCountResponse + +_OPUS_US_EAST_1: Final = TokenCountModel(model="anthropic.claude-opus-5-5", region="us-east-1") +_DEFAULT: Final = { + "token_count_framing": 97, + "token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-east-1"}, +} + + +def _in_us_west_2() -> str: + return "us-west-2" + + +@pytest.mark.parametrize( + ("configured", "expected"), + [ + ({}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)), + ( + {"token_count_framing": 5, "token_count_model": {"model": "custom-model", "region": "ap-southeast-4"}}, + TokenCountSettings(framing=5, count_model=TokenCountModel(model="custom-model", region="ap-southeast-4")), + ), + ( + {"token_count_model": {"model": "custom-model"}}, + TokenCountSettings(framing=97, count_model=TokenCountModel(model="custom-model", region=None)), + ), + ( + {"token_count_model": {"model": "custom-model", "region": " "}}, + TokenCountSettings(framing=97, count_model=TokenCountModel(model="custom-model", region=None)), + ), + ({"token_count_framing": 0}, TokenCountSettings(framing=0, count_model=_OPUS_US_EAST_1)), + ( + {"token_count_framing": None, "token_count_model": None}, + TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1), + ), + ( + {"token_count_framing": True, "token_count_model": "custom-model"}, + TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1), + ), + ({"token_count_framing": 78.0}, TokenCountSettings(framing=78, count_model=_OPUS_US_EAST_1)), + ({"token_count_framing": 1.5}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)), + ({"token_count_framing": -1}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)), + ({"token_count_framing": "78"}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)), + ({"token_count_model": {}}, TokenCountSettings(framing=97, count_model=None)), + ({"token_count_model": ""}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)), + ( + {"token_count_model": {"model": " ", "region": "us-east-1"}}, + TokenCountSettings(framing=97, count_model=None), + ), + ], +) +def test_config_overrides_the_cost_map_default_key_by_key( + configured: dict[str, object], expected: TokenCountSettings +) -> None: + assert ( + resolve_token_count_settings(configured=configured, default=_DEFAULT, deployment_region=_in_us_west_2) + == expected + ) + + +@pytest.mark.parametrize( + ("deployment_region", "configured", "expected_region"), + [ + ("us-west-2", {}, "us-east-1"), + ("eu-central-1", {}, None), + ("ca-central-1", {}, None), + ("us-gov-west-1", {}, None), + ( + "eu-central-1", + {"token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-east-1"}}, + "us-east-1", + ), + ], +) +def test_shipped_count_region_never_takes_the_prompt_out_of_the_deployment_geography( + deployment_region: str, configured: dict[str, object], expected_region: str | None +) -> None: + settings: Final = resolve_token_count_settings( + configured=configured, default=_DEFAULT, deployment_region=lambda: deployment_region + ) + assert settings.count_model == TokenCountModel(model="anthropic.claude-opus-5-5", region=expected_region) + + +def test_shipped_global_count_model_keeps_its_region_for_a_deployment_in_any_geography() -> None: + settings: Final = resolve_token_count_settings( + configured={}, + default={"token_count_model": {"model": "global.anthropic.claude-opus-5-5", "region": "us-east-1"}}, + deployment_region=lambda: "eu-west-1", + ) + assert settings.count_model == TokenCountModel(model="global.anthropic.claude-opus-5-5", region="us-east-1") + + +def test_shipped_australian_count_region_stays_with_an_australian_deployment() -> None: + settings: Final = resolve_token_count_settings( + configured={}, + default={"token_count_model": {"model": "au.anthropic.claude-opus-5-5", "region": "ap-southeast-4"}}, + deployment_region=lambda: "ap-southeast-2", + ) + assert settings.count_model == TokenCountModel(model="au.anthropic.claude-opus-5-5", region="ap-southeast-4") + + +def test_no_setting_anywhere_means_the_model_counts_itself_without_framing() -> None: + assert resolve_token_count_settings( + configured={}, default={}, deployment_region=_in_us_west_2 + ) == TokenCountSettings(framing=None, count_model=None) + + +@pytest.mark.parametrize( + ("value", "expected"), + [("kimi_k3", "kimi_k3"), ("cl100k_base", None), (None, None), (True, None)], +) +def test_only_a_tokenizer_litellm_ships_is_taken_from_the_cost_map(value: object, expected: str | None) -> None: + settings: Final = resolve_token_count_settings( + configured={}, default={"token_count_tokenizer": value}, deployment_region=_in_us_west_2 + ) + assert settings.tokenizer == expected + + +def _counted(source: str | None, model_used: str = "us.anthropic.claude-opus-5-5") -> TokenCountResponse: + return TokenCountResponse( + total_tokens=7, + request_model="sonnet", + model_used=model_used, + tokenizer_type="bedrock_api", + count_source=source, + ) + + +@pytest.mark.parametrize( + ("source", "expected"), + [ + (None, {}), + ("provider", {"litellm_count_source": "provider"}), + ("local_estimate", {"litellm_count_source": "local_estimate"}), + ( + "token_count_model", + {"litellm_count_source": "token_count_model", "litellm_count_model": "us.anthropic.claude-opus-5-5"}, + ), + ], +) +def test_count_source_fields_name_the_counting_model_only_when_token_count_model_counted( + source: str | None, expected: dict[str, str] +) -> None: + assert count_source_fields(_counted(source)) == expected + + +_UPSTREAM_404: Final = { + "type": "error", + "request_id": "req_upstream", + "error": {"type": "not_found_error", "message": "The model 'anthropic.claude-opus-5-5' does not exist"}, +} + + +@pytest.mark.parametrize( + ("message", "expected"), + [ + ( + json.dumps(_UPSTREAM_404), + { + **_UPSTREAM_404, + "error": { + "type": "not_found_error", + "message": "The model 'anthropic.claude-opus-5-5' does not exist. Enable eu-west-1", + }, + }, + ), + (json.dumps({"message": "Too many tokens"}), "Too many tokens. Enable eu-west-1"), + ("scripted 404", "scripted 404. Enable eu-west-1"), + ( + "Credential should be scoped to a valid region. ", + "Credential should be scoped to a valid region. Enable eu-west-1", + ), + ], +) +def test_hint_lands_in_the_upstream_error_message_and_keeps_its_envelope(message: str, expected: object) -> None: + hinted: Final = ProviderCountFailure(status_code=404, message=message).with_hint("Enable eu-west-1") + + assert (json.loads(hinted.message) if isinstance(expected, dict) else hinted.message) == expected + assert hinted.status_code == 404 diff --git a/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py b/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py index e4cd7d9dfa8..2777d18a399 100644 --- a/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py +++ b/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py @@ -3,10 +3,12 @@ Test to verify the Google GenAI proxy API endpoints """ -from unittest.mock import AsyncMock, MagicMock, patch +from typing import Final +from unittest.mock import AsyncMock, patch import pytest +from litellm.types.utils import TokenCountResponse def _build_test_client(): @@ -87,9 +89,7 @@ def test_google_stream_generate_content_endpoint(): assert init_kwargs["data"]["_litellm_raw_sse_stream"] is True assert init_kwargs["data"]["_litellm_skip_openai_stream_done"] is True assert init_kwargs["data"]["model"] == "test-model" - assert init_kwargs["data"]["contents"] == [ - {"role": "user", "parts": [{"text": "Hello"}]} - ] + assert init_kwargs["data"]["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}] def test_google_generate_content_data_flows_through_processor(): @@ -122,9 +122,7 @@ def test_google_generate_content_data_flows_through_processor(): data = mock_init.call_args.kwargs["data"] assert data["model"] == "test-model" assert data["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}] - assert data["systemInstruction"] == { - "parts": [{"text": "Your name is Doodle."}] - } + assert data["systemInstruction"] == {"parts": [{"text": "Your name is Doodle."}]} # generationConfig arrives intact here; the rename to `config` is # done downstream in route_request (see test_route_llm_request). assert data["generationConfig"]["responseModalities"] == ["TEXT", "IMAGE"] @@ -159,12 +157,17 @@ def test_google_count_tokens_unchanged(): except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - fake_response = MagicMock() - fake_response.original_response = { - "totalTokens": 7, - "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 7}], - } - fake_response.total_tokens = 7 + fake_response = TokenCountResponse( + total_tokens=7, + request_model="test-model", + model_used="test-model", + tokenizer_type="gemini_api", + original_response={ + "totalTokens": 7, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 7}], + }, + count_source="provider", + ) with patch( "litellm.proxy.proxy_server.token_counter", @@ -179,3 +182,61 @@ def test_google_count_tokens_unchanged(): assert response.status_code == 200 body = response.json() assert body["totalTokens"] == 7 + assert body["litellm_count_source"] == "provider" + + +def test_google_count_tokens_unexpected_failure_uses_the_google_error_envelope(): + try: + client = _build_test_client() + except ImportError as e: + pytest.skip(f"Skipping test due to missing dependency: {e}") + + with patch( + "litellm.proxy.proxy_server.token_counter", + new_callable=AsyncMock, + side_effect=RuntimeError("tokenizer unavailable"), + ): + response = client.post( + "/v1beta/models/test-model:countTokens", + json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, + ) + + assert response.status_code == 500 + error = response.json()["error"] + assert error["code"] == 500 + assert error["status"] == "INTERNAL" + assert "tokenizer unavailable" not in error["message"] + + +def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs(): + """A provider that counts in Anthropic shape (Bedrock, Anthropic) carries no totalTokens key, so the normalized total must win over the raw response's missing field.""" + try: + client: Final = _build_test_client() + except ImportError as e: + pytest.skip(f"Skipping test due to missing dependency: {e}") + + fake_response: Final = TokenCountResponse( + total_tokens=2167, + request_model="claude-opus-4-8", + model_used="us.anthropic.claude-opus-5-5", + tokenizer_type="bedrock_api", + original_response={"input_tokens": 2167}, + count_source="token_count_model", + ) + + with patch( + "litellm.proxy.proxy_server.token_counter", + new_callable=AsyncMock, + return_value=fake_response, + ): + response: Final = client.post( + "/v1beta/models/claude-opus-4-8:countTokens", + json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, + ) + + assert response.status_code == 200 + body: Final = response.json() + assert body["totalTokens"] == 2167 + assert body["promptTokensDetails"] == [] + assert body["litellm_count_source"] == "token_count_model" + assert body["litellm_count_model"] == "us.anthropic.claude-opus-5-5" diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index 5f7807650e1..b7f6865ab7e 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -4239,6 +4239,54 @@ class TestModelInfoCostMapEchoFilter: assert "litellm_provider" not in info assert "supported_openai_params" not in info + @pytest.fixture + def bedrock_sonnet_5_5_mapped_as_on_main(self) -> Iterator[None]: + import litellm + + entry: Final = { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": "us-east-1"}, + } + with pytest.MonkeyPatch.context() as patched: + patched.setattr(litellm, "model_cost", {**litellm.model_cost, "us.anthropic.claude-sonnet-5-5": entry}) + litellm.utils._invalidate_model_cost_lowercase_map() + yield + litellm.utils._invalidate_model_cost_lowercase_map() + + @pytest.mark.parametrize( + ("model", "field", "edited"), + [ + ("bedrock/us.anthropic.claude-sonnet-5-5", "token_count_model", {}), + ("bedrock/us.moonshotai.kimi-k3", "token_count_framing", 0), + ], + ) + def test_token_count_setting_is_stored_only_when_the_operator_changed_it( + self, model: str, field: str, edited: object, bedrock_sonnet_5_5_mapped_as_on_main: None + ) -> None: + import litellm + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + entry: Final = litellm.get_model_info(model) + assert entry[field] not in (None, edited), entry + db_model: Final = Deployment( + model_name="count-settings", + litellm_params=LiteLLM_Params(model=model), + model_info=ModelInfo(id="dep-echo-count"), + ) + + echoed: Final = update_db_model( + db_model=db_model, + updated_patch=updateDeployment(model_info=ModelInfo(**{**entry, "id": "dep-echo-count"})), + ) + changed: Final = update_db_model( + db_model=db_model, + updated_patch=updateDeployment(model_info=ModelInfo(**{**entry, "id": "dep-echo-count", field: edited})), + ) + + assert field not in json.loads(echoed["model_info"]) + assert json.loads(changed["model_info"])[field] == edited + def test_metadata_without_a_cost_map_key_is_persisted(self): import litellm diff --git a/tests/unit/proxy/proxy_server/test_routes_utils.py b/tests/unit/proxy/proxy_server/test_routes_utils.py index 91329d122ee..77a12cf0d43 100644 --- a/tests/unit/proxy/proxy_server/test_routes_utils.py +++ b/tests/unit/proxy/proxy_server/test_routes_utils.py @@ -54,6 +54,7 @@ def test_token_counter_happy_path(client, auth_as, patched_token_counter): "request_model": "gpt-4", "model_used": "gpt-4", "tokenizer_type": "openai_tokenizer", + "count_source": "local_estimate", "original_response": None, "error": False, "error_message": None, @@ -390,8 +391,7 @@ def test_token_counter_fallback_counts_tools_system_and_anthropic_blocks(client, assert full > bare -def test_token_counter_fallback_prompt_with_tools_does_not_500(client, auth_as, monkeypatch): - """Regression: a ``prompt`` request carrying ``tools`` but no ``messages`` still counts, because the fallback attaches tools only when counting messages (``token_counter`` rejects tools on the text path).""" +def test_token_counter_fallback_counts_a_prompt_with_tools_like_its_messages_form(client, auth_as, monkeypatch): monkeypatch.setattr(proxy_server, "llm_router", None) monkeypatch.setattr(litellm, "disable_token_counter", False, raising=False) prompt = "count the tokens in this sentence please" @@ -412,4 +412,6 @@ def test_token_counter_fallback_prompt_with_tools_does_not_500(client, auth_as, ) assert response.status_code == 200, response.text - assert response.json()["total_tokens"] == litellm.token_counter(model="claude-fable-5", text=prompt) + assert response.json()["total_tokens"] == litellm.token_counter( + model="claude-fable-5", messages=[{"role": "user", "content": prompt}], tools=tools + ) diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index fcd72e3777c..995089105c4 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -2,32 +2,120 @@ # 1. Generate a Key, and use it to make a call +import asyncio +import base64 +import hashlib +import itertools +import json +import logging +import threading +import time +import traceback +from collections.abc import Callable, Coroutine, Iterator, Mapping +from functools import partial +from pathlib import Path +from types import MappingProxyType +from typing import Final, NoReturn from unittest.mock import AsyncMock, MagicMock, patch +import anyio import httpx import pytest +import respx +import tiktoken from dotenv import load_dotenv +from pydantic import BaseModel, ConfigDict, JsonValue +from tokenizers import Tokenizer +from tokenizers.models import WordLevel +from tokenizers.pre_tokenizers import Whitespace +from tokenizers.processors import TemplateProcessing load_dotenv() # this file is to test litellm/proxy -from fastapi import HTTPException, Request +from fastapi import Request +from fastapi.responses import JSONResponse import litellm from litellm import Router +from litellm._logging import verbose_proxy_logger +from litellm.caching.llm_caching_handler import LLMClientCache +from litellm.llms.anthropic.count_tokens import handler as anthropic_count_tokens_handler +from litellm.llms.base_llm.base_utils import BaseTokenCounter from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler +from litellm.llms.bedrock.count_tokens.kimi_k3 import ( + KIMI_K3_RANKS_SHA256, + KimiK3Encodings, + KimiK3TokenizerSource, + kimi_k3_template_tokens, + kimi_k3_tokens, + load_kimi_k3_encoding, +) +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler +from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import VertexAIPartnerModelsTokenCounter from litellm.proxy._types import ProxyException, TokenCountRequest from litellm.proxy.anthropic_endpoints.endpoints import ( count_tokens as anthropic_count_tokens, ) -from litellm.proxy.proxy_server import token_counter -from litellm.types.utils import TokenCountResponse +from litellm.proxy.common_utils.token_count_settings import ( + ProviderCountFailure, + TokenCountModel, + TokenCountSettings, + cost_map_entry, +) +from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request +from litellm.proxy.proxy_server import ( + _estimate_or_provider_error, + _provider_token_count, + _try_provider_token_count, + count_request_tokens, + token_counter, +) +from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted +from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem +from litellm.types.utils import SelectTokenizerResponse, TokenCountMessagesFormat, TokenCountResponse +def _hugging_face_offline() -> NoReturn: + raise OSError("unit tests never download the Kimi K3 tokenizer") + + +@pytest.fixture(autouse=True) +def _kimi_k3_tokenizer_unavailable(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + litellm.proxy.proxy_server, + "_KIMI_K3_ENCODINGS", + KimiK3Encodings( + source=KimiK3TokenizerSource(ranks_path=_hugging_face_offline, sha256=KIMI_K3_RANKS_SHA256), + retry_after_seconds=60, + ), + ) + + +@pytest.fixture +def bedrock_sonnet_5_5_mapped_as_on_main() -> Iterator[None]: + entries: Final = { + name: { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "token_count_model": {"model": count_model, "region": "us-east-1"}, + } + for name, count_model in ( + ("anthropic.claude-sonnet-5-5", "anthropic.claude-opus-5-5"), + ("us.anthropic.claude-sonnet-5-5", "us.anthropic.claude-opus-5-5"), + ) + } + with pytest.MonkeyPatch.context() as patched: + patched.setattr(litellm, "model_cost", {**litellm.model_cost, **entries}) + litellm.utils._invalidate_model_cost_lowercase_map() + yield + litellm.utils._invalidate_model_cost_lowercase_map() + @pytest.mark.asyncio async def test_vLLM_token_counting(): @@ -61,9 +149,7 @@ async def test_vLLM_token_counting(): print("response: ", response) - assert ( - response.tokenizer_type == "openai_tokenizer" - ) # SHOULD use the default tokenizer + assert response.tokenizer_type == "openai_tokenizer" # SHOULD use the default tokenizer assert response.model_used == "wolfram/miquliz-120b-v2.0" @@ -96,9 +182,7 @@ async def test_token_counting_model_not_in_model_list(): print("response: ", response) - assert ( - response.tokenizer_type == "openai_tokenizer" - ) # SHOULD use the OpenAI tokenizer + assert response.tokenizer_type == "openai_tokenizer" # SHOULD use the OpenAI tokenizer assert response.model_used == "special-alias" @@ -131,9 +215,7 @@ async def test_gpt_token_counting(): print("response: ", response) - assert ( - response.tokenizer_type == "openai_tokenizer" - ) # SHOULD use the OpenAI tokenizer + assert response.tokenizer_type == "openai_tokenizer" # SHOULD use the OpenAI tokenizer assert response.request_model == "gpt-4" @@ -171,10 +253,8 @@ async def test_anthropic_messages_count_tokens_endpoint(): anthropic_endpoints._read_request_body = mock_read_request_body # Mock the internal token_counter function to return a controlled response - async def mock_token_counter(request, call_endpoint=False): - assert ( - call_endpoint == True - ), "Should be called with call_endpoint=True for Anthropic endpoint" + async def mock_token_counter(request, call_endpoint=False, messages_format="openai", output_format=None): + assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" assert request.model == "claude-3-sonnet-20240229" assert request.messages == [{"role": "user", "content": "Hello Claude!"}] @@ -190,8 +270,8 @@ async def test_anthropic_messages_count_tokens_endpoint(): # Patch the imported token_counter function from proxy_server import litellm.proxy.proxy_server as proxy_server - original_token_counter = proxy_server.token_counter - proxy_server.token_counter = mock_token_counter + original_token_counter = proxy_server.count_request_tokens + proxy_server.count_request_tokens = mock_token_counter try: # Call the endpoint @@ -208,7 +288,7 @@ async def test_anthropic_messages_count_tokens_endpoint(): finally: # Restore original functions anthropic_endpoints._read_request_body = original_read_request_body - proxy_server.token_counter = original_token_counter + proxy_server.count_request_tokens = original_token_counter @pytest.mark.asyncio @@ -245,10 +325,8 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): anthropic_endpoints._read_request_body = mock_read_request_body # Mock the internal token_counter function to return a controlled response - async def mock_token_counter(request, call_endpoint=True): - assert ( - call_endpoint == True - ), "Should be called with call_endpoint=True for Anthropic endpoint" + async def mock_token_counter(request, call_endpoint=True, messages_format="openai", output_format=None): + assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" assert request.model == "gpt-4" assert request.messages == [{"role": "user", "content": "Hello GPT!"}] @@ -264,8 +342,8 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): # Patch the imported token_counter function from proxy_server import litellm.proxy.proxy_server as proxy_server - original_token_counter = proxy_server.token_counter - proxy_server.token_counter = mock_token_counter + original_token_counter = proxy_server.count_request_tokens + proxy_server.count_request_tokens = mock_token_counter try: # Call the endpoint @@ -282,77 +360,79 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): finally: # Restore original functions anthropic_endpoints._read_request_body = original_read_request_body - proxy_server.token_counter = original_token_counter + proxy_server.count_request_tokens = original_token_counter @pytest.mark.asyncio -async def test_internal_token_counter_anthropic_provider_detection(): - """ - Test that the internal token_counter correctly detects Anthropic providers - and handles the from_anthropic_endpoint flag appropriately - """ - - # Test with Anthropic provider - llm_router = Router( - model_list=[ - { - "model_name": "claude-test", - "litellm_params": { - "model": "anthropic/claude-3-sonnet-20240229", - "api_key": "test-key", - }, - } - ] +@pytest.mark.parametrize("status_code", (400, 401, 403, 404, 429, 500)) +@pytest.mark.parametrize( + ("model", "count_url"), + ( + ("anthropic/claude-sonnet-5-5", "https://api.anthropic.com/v1/messages/count_tokens"), + ("openai/gpt-6.1-sol", "https://api.openai.com/v1/responses/input_tokens"), + ), +) +async def test_provider_errors_return_labelled_local_estimates( + monkeypatch: pytest.MonkeyPatch, status_code: int, model: str, count_url: str +) -> None: + messages: Final = [ + {"role": "user", "content": "Count the whole request locally when the count API is unavailable."} + ] + router: Final = Router( + model_list=[{"model_name": "counted", "litellm_params": {"model": model, "api_key": "test-count-key"}}] ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + monkeypatch.setattr(litellm, "disable_token_counter", False) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") - setattr(litellm.proxy.proxy_server, "llm_router", llm_router) + with respx.mock(assert_all_called=True) as upstream: + upstream.post(count_url).mock( + return_value=httpx.Response( + status_code, json={"type": "error", "error": {"type": "api_error", "message": "Count unavailable"}} + ) + ) + counted: Final = await token_counter( + request=TokenCountRequest(model="counted", messages=messages), call_endpoint=True + ) - # Test with is_direct_request=False (simulating call from Anthropic endpoint) - response = await token_counter( - request=TokenCountRequest( - model="claude-test", - messages=[{"role": "user", "content": "hello"}], - ), - call_endpoint=True, + assert counted.total_tokens == litellm.token_counter(model=model.split("/", 1)[1], messages=messages) + assert counted.count_source == "local_estimate" + assert counted.error is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "part", + ( + {"fileData": {"mimeType": "image/png", "fileUri": "gs://bucket/image.png"}}, + {"inlineData": {"mimeType": "image/png", "data": "aW1hZ2U="}}, + {"text": "Reasoning about the request", "thought": True}, + ), +) +async def test_google_provider_errors_preserve_legacy_translated_message_estimates( + monkeypatch: pytest.MonkeyPatch, part: dict[str, JsonValue] +) -> None: + model: Final = "openai/gpt-6.1-sol" + request: Final = _google_count_tokens_request( + "counted", {"contents": [{"role": "user", "parts": [{"text": "Describe this"}, part]}]} ) - - print("Anthropic provider test response:", response) - - # Verify response structure - assert response.request_model == "claude-test" - assert response.model_used == "claude-3-sonnet-20240229" - assert response.total_tokens > 0 - - # Test with non-Anthropic provider - llm_router = Router( - model_list=[ - { - "model_name": "gpt-test", - "litellm_params": { - "model": "gpt-4", - }, - } - ] + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[{"model_name": "counted", "litellm_params": {"model": model, "api_key": "test-key"}}]), ) - - setattr(litellm.proxy.proxy_server, "llm_router", llm_router) - - # Test with is_direct_request=False but non-Anthropic provider - response = await token_counter( - request=TokenCountRequest( - model="gpt-test", - messages=[{"role": "user", "content": "hello"}], - ), - call_endpoint=True, - ) - - print("Non-Anthropic provider test response:", response) - - # Verify response structure - assert response.request_model == "gpt-test" - assert response.model_used == "gpt-4" - assert response.total_tokens > 0 - assert response.tokenizer_type == "openai_tokenizer" # Should use LiteLLM tokenizer + monkeypatch.setattr(litellm, "disable_token_counter", False) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + with respx.mock(assert_all_called=True) as upstream: + upstream.post("https://api.openai.com/v1/responses/input_tokens").mock( + return_value=httpx.Response(400, json={"error": {"message": "Count unavailable"}}) + ) + counted: Final = await token_counter(request=request, call_endpoint=True) + assert counted.total_tokens == litellm.token_counter(model="gpt-6.1-sol", messages=request.messages) + assert counted.count_source == "local_estimate" + assert counted.error is False @pytest.mark.asyncio @@ -362,12 +442,14 @@ async def test_anthropic_endpoint_error_handling(): """ from unittest.mock import MagicMock - from fastapi import HTTPException, Request + from fastapi import Request from litellm.proxy.anthropic_endpoints.endpoints import count_tokens # Mock request object mock_request = MagicMock(spec=Request) + mock_request.headers = {} + mock_request.scope = {} mock_user_api_key_dict = MagicMock() # Test missing model parameter @@ -385,12 +467,13 @@ async def test_anthropic_endpoint_error_handling(): anthropic_endpoints._read_request_body = mock_read_request_body try: - # Should raise HTTPException for missing model - with pytest.raises(HTTPException) as exc_info: - await count_tokens(mock_request, mock_user_api_key_dict) + response = await count_tokens(mock_request, mock_user_api_key_dict) - assert exc_info.value.status_code == 400 - assert "model parameter is required" in str(exc_info.value.detail) + assert response.status_code == 400 + assert json.loads(response.body) == { + "type": "error", + "error": {"type": "invalid_request_error", "message": "model parameter is required"}, + } print("✅ Error handling test passed!") @@ -409,9 +492,7 @@ async def test_factory_anthropic_endpoint_calls_anthropic_counter(): # Mock the global handler instance in token_counter module mock_handler = MagicMock() - mock_handler.handle_count_tokens_request = AsyncMock( - return_value={"input_tokens": 42} - ) + mock_handler.handle_count_tokens_request = AsyncMock(return_value={"input_tokens": 42}) with patch( "litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler", @@ -468,9 +549,7 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): # Mock the global handler instance in token_counter module mock_handler = MagicMock() - mock_handler.handle_count_tokens_request = AsyncMock( - return_value={"input_tokens": 42} - ) + mock_handler.handle_count_tokens_request = AsyncMock(return_value={"input_tokens": 42}) with patch( "litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler", @@ -485,7 +564,7 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): mock_router.model_list = [ { "model_name": "gpt-4", - "litellm_params": {"model": "openai/gpt-4"}, + "litellm_params": {"model": "databricks/gpt-4"}, "model_info": {}, } ] @@ -494,7 +573,7 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): mock_router.async_get_available_deployment = AsyncMock( return_value={ "model_name": "gpt-4", - "litellm_params": {"model": "openai/gpt-4"}, + "litellm_params": {"model": "databricks/gpt-4"}, "model_info": {}, } ) @@ -529,9 +608,7 @@ async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic(): # Mock the global handler instance in token_counter module mock_handler = MagicMock() - mock_handler.handle_count_tokens_request = AsyncMock( - return_value={"input_tokens": 42} - ) + mock_handler.handle_count_tokens_request = AsyncMock(return_value={"input_tokens": 42}) with patch( "litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler", @@ -546,9 +623,7 @@ async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic(): mock_router.model_list = [ { "model_name": "claude-3-5-sonnet", - "litellm_params": { - "model": "anthropic/claude-3-5-sonnet-20241022" - }, + "litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"}, "model_info": {}, } ] @@ -557,9 +632,7 @@ async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic(): mock_router.async_get_available_deployment = AsyncMock( return_value={ "model_name": "claude-3-5-sonnet", - "litellm_params": { - "model": "anthropic/claude-3-5-sonnet-20241022" - }, + "litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"}, "model_info": {}, } ) @@ -594,9 +667,7 @@ async def test_factory_registration(): assert counter is not None # Create test deployments - anthropic_deployment = { - "litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"} - } + anthropic_deployment = {"litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"}} non_anthropic_deployment = {"litellm_params": {"model": "openai/gpt-4"}} @@ -632,9 +703,7 @@ async def test_bedrock_count_tokens_endpoint(): model_list=[ { "model_name": "claude-bedrock", - "litellm_params": { - "model": "bedrock/anthropic.claude-3-sonnet-20240229-v1:0" - }, + "litellm_params": {"model": "bedrock/anthropic.claude-3-sonnet-20240229-v1:0"}, } ] ) @@ -648,9 +717,7 @@ async def test_bedrock_count_tokens_endpoint(): } # Test the mock handler directly to verify correct parameter extraction - await mock_count_tokens_handler( - request_data, {}, "anthropic.claude-3-sonnet-20240229-v1:0" - ) + await mock_count_tokens_handler(request_data, {}, "anthropic.claude-3-sonnet-20240229-v1:0") @pytest.mark.asyncio @@ -712,10 +779,7 @@ async def test_vertex_ai_anthropic_token_counting(): assert call_args is not None assert call_args.kwargs["model"] == "claude-3-5-sonnet-20241022" assert "messages" in call_args.kwargs["request_data"] - assert ( - call_args.kwargs["request_data"]["messages"][0]["content"] - == "Hello Claude on Vertex AI! How are you?" - ) + assert call_args.kwargs["request_data"]["messages"][0]["content"] == "Hello Claude on Vertex AI! How are you?" # Validate response structure assert response.model_used == "claude-3-5-sonnet-20241022" @@ -748,9 +812,7 @@ def test_vertex_ai_partner_models_token_counting_endpoint(vertex_location): if vertex_location == "global": assert endpoint.startswith("https://aiplatform.googleapis.com") else: - assert endpoint.startswith( - f"https://{vertex_location}-aiplatform.googleapis.com" - ) + assert endpoint.startswith(f"https://{vertex_location}-aiplatform.googleapis.com") @pytest.mark.asyncio @@ -840,20 +902,14 @@ async def test_bedrock_handler_httpx_error_status_code_propagation(): 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, "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: + 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 @@ -862,9 +918,7 @@ async def test_bedrock_handler_httpx_error_status_code_propagation(): await handler.handle_count_tokens_request( request_data={ "model": "test", - "messages": [ - {"role": "user", "content": "hello"} - ], + "messages": [{"role": "user", "content": "hello"}], }, litellm_params={}, resolved_model="anthropic.claude-3-sonnet", @@ -872,81 +926,9 @@ async def test_bedrock_handler_httpx_error_status_code_propagation(): 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.message == "Forbidden - Invalid credentials" -@pytest.mark.asyncio -async def test_token_counter_httpx_status_error_raises_proxy_exception(): - """ - When provider_counter.count_tokens() raises httpx.HTTPStatusError, - the token_counter endpoint should catch it and raise a ProxyException - with the upstream status code and error message. - """ - - upstream_status = 429 - upstream_message = "Rate limit exceeded" - response = httpx.Response( - status_code=upstream_status, - request=httpx.Request("POST", "https://provider.example.com/count"), - ) - http_error = httpx.HTTPStatusError( - message=upstream_message, - request=response.request, - response=response, - ) - - mock_counter = MagicMock() - mock_counter.should_use_token_counting_api.return_value = True - mock_counter.count_tokens = AsyncMock(side_effect=http_error) - - # Save originals - original_get_provider_token_counter = ( - litellm.proxy.proxy_server._get_provider_token_counter - ) - original_router = litellm.proxy.proxy_server.llm_router - - try: - - def mock_get_provider_token_counter(deployment, model_to_use): - return (mock_counter, "claude-4-6-sonnet", "vertex_ai") - - litellm.proxy.proxy_server._get_provider_token_counter = ( - mock_get_provider_token_counter - ) - - mock_router = MagicMock() - mock_router.async_get_available_deployment = AsyncMock( - return_value={ - "litellm_params": { - "model": "vertex_ai/claude-4-6-sonnet", - "api_key": "fake-key", - }, - "model_info": {}, - } - ) - litellm.proxy.proxy_server.llm_router = mock_router - - with pytest.raises(ProxyException) as exc_info: - await token_counter( - request=TokenCountRequest( - model="claude-4-6-sonnet", - messages=[{"role": "user", "content": "hello"}], - ), - call_endpoint=True, - ) - - assert exc_info.value.code == str(upstream_status) - assert upstream_message in exc_info.value.message - assert exc_info.value.type == "token_counting_error" - assert exc_info.value.param == "model" - finally: - litellm.proxy.proxy_server._get_provider_token_counter = ( - original_get_provider_token_counter - ) - litellm.proxy.proxy_server.llm_router = original_router @pytest.mark.asyncio @@ -981,9 +963,7 @@ async def test_proxy_token_counter_error_raises_exception_when_disabled(): # Save original value and function original_disable = litellm.disable_token_counter - original_get_provider_token_counter = ( - litellm.proxy.proxy_server._get_provider_token_counter - ) + original_get_provider_token_counter = litellm.proxy.proxy_server._get_provider_token_counter try: litellm.disable_token_counter = True @@ -997,9 +977,7 @@ async def test_proxy_token_counter_error_raises_exception_when_disabled(): def mock_get_provider_token_counter(deployment, model_to_use): return (mock_counter, "anthropic.claude-3-sonnet", "bedrock") - litellm.proxy.proxy_server._get_provider_token_counter = ( - mock_get_provider_token_counter - ) + litellm.proxy.proxy_server._get_provider_token_counter = mock_get_provider_token_counter with pytest.raises(ProxyException) as exc_info: await token_counter( @@ -1014,81 +992,9 @@ async def test_proxy_token_counter_error_raises_exception_when_disabled(): assert "Rate limit exceeded" in exc_info.value.message finally: litellm.disable_token_counter = original_disable - litellm.proxy.proxy_server._get_provider_token_counter = ( - original_get_provider_token_counter - ) + litellm.proxy.proxy_server._get_provider_token_counter = original_get_provider_token_counter -@pytest.mark.asyncio -async def test_proxy_token_counter_error_falls_back_when_enabled(): - """ - Test that proxy token_counter falls back to local tokenizer when disable_token_counter=False - and provider returns an error response. - """ - # Create error response - error_response = TokenCountResponse( - total_tokens=0, - request_model="bedrock/anthropic.claude-3-sonnet", - model_used="anthropic.claude-3-sonnet", - tokenizer_type="bedrock_api", - error=True, - error_message="Rate limit exceeded", - status_code=429, - ) - - # Create mock router that returns a deployment - mock_deployment = { - "litellm_params": { - "model": "bedrock/anthropic.claude-3-sonnet", - }, - "model_info": {}, - } - - mock_router = MagicMock() - mock_router.async_get_available_deployment = AsyncMock(return_value=mock_deployment) - - setattr(litellm.proxy.proxy_server, "llm_router", mock_router) - - # Save original value and function - original_disable = litellm.disable_token_counter - original_get_provider_token_counter = ( - litellm.proxy.proxy_server._get_provider_token_counter - ) - - try: - litellm.disable_token_counter = False - - # Create a mock counter that returns an error response - mock_counter = MagicMock(spec=BedrockTokenCounter) - mock_counter.should_use_token_counting_api.return_value = True - mock_counter.count_tokens = AsyncMock(return_value=error_response) - - # Replace the function directly - def mock_get_provider_token_counter(deployment, model_to_use): - return (mock_counter, "anthropic.claude-3-sonnet", "bedrock") - - litellm.proxy.proxy_server._get_provider_token_counter = ( - mock_get_provider_token_counter - ) - - # Should not raise, should fall back to local tokenizer - result = await token_counter( - request=TokenCountRequest( - model="claude-bedrock", - messages=[{"role": "user", "content": "hello"}], - ), - call_endpoint=True, - ) - - # Should have used the fallback tokenizer - assert result.error is False - assert result.total_tokens > 0 - assert result.tokenizer_type != "bedrock_api" - finally: - litellm.disable_token_counter = original_disable - litellm.proxy.proxy_server._get_provider_token_counter = ( - original_get_provider_token_counter - ) @pytest.mark.asyncio @@ -1101,6 +1007,7 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): # Mock request object mock_request = MagicMock(spec=Request) + mock_request.headers = {} mock_request_data = { "model": "claude-bedrock", "messages": [{"role": "user", "content": "Hello!"}], @@ -1114,10 +1021,10 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): original_read_request_body = anthropic_endpoints._read_request_body anthropic_endpoints._read_request_body = mock_read_request_body - original_token_counter = proxy_server.token_counter + original_token_counter = proxy_server.count_request_tokens # Mock token_counter to raise ProxyException with Bedrock-style error - async def mock_token_counter_error(request, call_endpoint=False): + async def mock_token_counter_error(request, call_endpoint=False, messages_format="openai", output_format=None): raise ProxyException( message='{"detail":{"message":"Input is too long for requested model."}}', type="token_counting_error", @@ -1125,23 +1032,20 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): code=400, ) - proxy_server.token_counter = mock_token_counter_error + proxy_server.count_request_tokens = mock_token_counter_error try: - with pytest.raises(HTTPException) as exc_info: - await anthropic_count_tokens(mock_request, mock_user_api_key_dict) + response: Final[object] = await anthropic_count_tokens(mock_request, mock_user_api_key_dict) - # Verify HTTP status code is correct - assert exc_info.value.status_code == 400 - - # Verify error is in Anthropic format - detail = exc_info.value.detail + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + detail = json.loads(bytes(response.body)) assert detail["type"] == "error" assert detail["error"]["type"] == "invalid_request_error" assert detail["error"]["message"] == "Input is too long for requested model." finally: anthropic_endpoints._read_request_body = original_read_request_body - proxy_server.token_counter = original_token_counter + proxy_server.count_request_tokens = original_token_counter @pytest.mark.asyncio @@ -1153,6 +1057,7 @@ async def test_anthropic_endpoint_403_permission_error_format(): import litellm.proxy.proxy_server as proxy_server mock_request = MagicMock(spec=Request) + mock_request.headers = {} mock_request_data = { "model": "claude-bedrock", "messages": [{"role": "user", "content": "Hello!"}], @@ -1166,10 +1071,10 @@ async def test_anthropic_endpoint_403_permission_error_format(): original_read_request_body = anthropic_endpoints._read_request_body anthropic_endpoints._read_request_body = mock_read_request_body - original_token_counter = proxy_server.token_counter + original_token_counter = proxy_server.count_request_tokens # Mock token_counter to raise ProxyException with 403 error - async def mock_token_counter_error(request, call_endpoint=False): + async def mock_token_counter_error(request, call_endpoint=False, messages_format="openai", output_format=None): raise ProxyException( message='{"Message":"Bearer Token has expired"}', type="token_counting_error", @@ -1177,21 +1082,20 @@ async def test_anthropic_endpoint_403_permission_error_format(): code=403, ) - proxy_server.token_counter = mock_token_counter_error + proxy_server.count_request_tokens = mock_token_counter_error try: - with pytest.raises(HTTPException) as exc_info: - await anthropic_count_tokens(mock_request, mock_user_api_key_dict) + response: Final[object] = await anthropic_count_tokens(mock_request, mock_user_api_key_dict) - assert exc_info.value.status_code == 403 - - detail = exc_info.value.detail + assert isinstance(response, JSONResponse) + assert response.status_code == 403 + detail = json.loads(bytes(response.body)) assert detail["type"] == "error" assert detail["error"]["type"] == "permission_error" assert detail["error"]["message"] == "Bearer Token has expired" finally: anthropic_endpoints._read_request_body = original_read_request_body - proxy_server.token_counter = original_token_counter + proxy_server.count_request_tokens = original_token_counter @pytest.mark.asyncio @@ -1203,6 +1107,7 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): import litellm.proxy.proxy_server as proxy_server mock_request = MagicMock(spec=Request) + mock_request.headers = {} mock_request_data = { "model": "claude-bedrock", "messages": [{"role": "user", "content": "Hello!"}], @@ -1216,10 +1121,10 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): original_read_request_body = anthropic_endpoints._read_request_body anthropic_endpoints._read_request_body = mock_read_request_body - original_token_counter = proxy_server.token_counter + original_token_counter = proxy_server.count_request_tokens # Mock token_counter to raise ProxyException with 429 error - async def mock_token_counter_error(request, call_endpoint=False): + async def mock_token_counter_error(request, call_endpoint=False, messages_format="openai", output_format=None): raise ProxyException( message="Rate limit exceeded", type="token_counting_error", @@ -1227,18 +1132,1537 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): code=429, ) - proxy_server.token_counter = mock_token_counter_error + proxy_server.count_request_tokens = mock_token_counter_error try: - with pytest.raises(HTTPException) as exc_info: - await anthropic_count_tokens(mock_request, mock_user_api_key_dict) + response: Final[object] = await anthropic_count_tokens(mock_request, mock_user_api_key_dict) - assert exc_info.value.status_code == 429 - - detail = exc_info.value.detail + assert isinstance(response, JSONResponse) + assert response.status_code == 429 + detail = json.loads(bytes(response.body)) assert detail["type"] == "error" assert detail["error"]["type"] == "rate_limit_error" assert detail["error"]["message"] == "Rate limit exceeded" finally: anthropic_endpoints._read_request_body = original_read_request_body - proxy_server.token_counter = original_token_counter + proxy_server.count_request_tokens = original_token_counter + + +class _ScriptedCounter(BaseTokenCounter): + def __init__(self, replies: dict[str, TokenCountResponse], burst: int = 1) -> None: + self._replies = replies + self._burst = burst + self._burst_arrived = asyncio.Event() + self.calls: list[tuple[str, object, object, object]] = [] + self.regions: list[object] = [] + + def should_use_token_counting_api(self, custom_llm_provider: str | None = None) -> bool: + return True + + async def count_tokens( + self, + model_to_use: str, + messages: object, + contents: object, + deployment: object = None, + request_model: str = "", + tools: object = None, + system: object = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, + context_management: object = None, + ) -> TokenCountResponse | None: + self.calls.append((model_to_use, messages, tools, system)) + self.regions.append(None if deployment is None else deployment["litellm_params"].get("aws_region_name")) + if len(self.calls) >= self._burst: + self._burst_arrived.set() + await self._burst_arrived.wait() + return self._replies[model_to_use] + + +def _provider_reply( + model: str, total_tokens: int = 0, status_code: int | None = None, rejected_before_provider_call: bool = False +) -> TokenCountResponse: + return TokenCountResponse( + total_tokens=total_tokens, + request_model="sonnet-5-5", + model_used=model, + tokenizer_type="bedrock_api", + error=status_code is not None, + error_message=None if status_code is None else f"scripted {status_code} for {model}", + status_code=status_code, + rejected_before_provider_call=rejected_before_provider_call, + ) + + +_COUNT_REQUEST = TokenCountRequest( + model="sonnet-5-5", + messages=[{"role": "user", "content": "hello"}], + tools=[{"name": "get_quote", "input_schema": {"type": "object"}}], + system="be brief", +) + + +_SONNET_DEPLOYMENT: Final = MappingProxyType( + {"litellm_params": {"model": "bedrock/us.anthropic.claude-sonnet-5-5", "aws_region_name": "us-west-2"}} +) +_OPUS_IN_US_EAST_1: Final = TokenCountModel(model="us.anthropic.claude-opus-5-5", region="us-east-1") + + +async def _ladder( + counter: _ScriptedCounter, settings: TokenCountSettings, local_estimate_allowed: bool = True +) -> TokenCountResponse | ProviderCountFailure | None: + return await _provider_token_count( + count_with=partial( + _try_provider_token_count, + counter, + "bedrock", + messages=_COUNT_REQUEST.messages, + contents=_COUNT_REQUEST.contents, + deployment=_SONNET_DEPLOYMENT, + request_model=_COUNT_REQUEST.model, + tools=_COUNT_REQUEST.tools, + system=_COUNT_REQUEST.system, + ), + model_to_use="us.anthropic.claude-sonnet-5-5", + settings=settings, + local_estimate_allowed=local_estimate_allowed, + ) + + +@pytest.mark.asyncio +async def test_model_without_token_count_model_counts_itself_in_the_deployment_region(): + counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", total_tokens=447)}) + + result: Final = await _ladder(counter, TokenCountSettings(framing=10, count_model=None)) + + assert result is not None + assert (result.total_tokens, result.count_source) == (447, "provider") + assert ([call[0] for call in counter.calls], counter.regions) == (["us.anthropic.claude-sonnet-5-5"], ["us-west-2"]) + + +@pytest.mark.parametrize( + ("count_model", "region"), + [ + (_OPUS_IN_US_EAST_1, "us-east-1"), + (TokenCountModel(model="us.anthropic.claude-opus-5-5", region=None), "us-west-2"), + ], +) +@pytest.mark.asyncio +async def test_token_count_model_counts_the_same_request_in_its_own_region_without_trying_the_model( + count_model: TokenCountModel, region: str +): + counter: Final = _ScriptedCounter({"us.anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)}) + + result: Final = await _ladder(counter, TokenCountSettings(framing=None, count_model=count_model)) + + assert result is not None + assert (result.total_tokens, result.count_source, result.model_used) == ( + 447, + "token_count_model", + "us.anthropic.claude-opus-5-5", + ) + assert ([call[0] for call in counter.calls], counter.regions) == (["us.anthropic.claude-opus-5-5"], [region]) + assert counter.calls[0][1:] == (_COUNT_REQUEST.messages, _COUNT_REQUEST.tools, _COUNT_REQUEST.system) + assert _SONNET_DEPLOYMENT["litellm_params"]["aws_region_name"] == "us-west-2" + + +@pytest.mark.parametrize("status_code", (400, 401, 403, 404, 429, 500)) +@pytest.mark.parametrize("framing", (None, 0, 10)) +@pytest.mark.asyncio +async def test_provider_count_errors_preserve_legacy_local_fallback(status_code: int, framing: int | None): + counter: Final = _ScriptedCounter( + {"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=status_code)} + ) + + result: Final = await _ladder(counter, TokenCountSettings(framing=framing, count_model=None)) + + assert isinstance(result, ProviderCountFailure) + assert (result.status_code, result.message) == (status_code, f"scripted {status_code} for sonnet") + + +@pytest.mark.parametrize("status_code", [400, 403, 404, 500]) +@pytest.mark.asyncio +async def test_token_count_model_failure_names_its_region_so_the_user_can_enable_it(status_code: int): + counter: Final = _ScriptedCounter( + {"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=status_code)} + ) + + result: Final = await _ladder(counter, TokenCountSettings(framing=None, count_model=_OPUS_IN_US_EAST_1)) + + assert isinstance(result, ProviderCountFailure) + assert result.status_code == status_code + assert result.message.startswith(f"scripted {status_code} for opus") + assert "us.anthropic.claude-opus-5-5 in us-east-1" in result.message + assert len(counter.calls) == 1 + + +@pytest.mark.asyncio +async def test_token_count_model_without_a_region_preserves_its_error_for_a_failed_local_estimate(): + counter: Final = _ScriptedCounter({"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=404)}) + + result: Final = await _ladder( + counter, + TokenCountSettings(framing=None, count_model=TokenCountModel(model="us.anthropic.claude-opus-5-5", region=None)), + ) + + assert isinstance(result, ProviderCountFailure) + assert (result.status_code, result.message) == (404, "scripted 404 for opus") + + +@pytest.mark.asyncio +async def test_uncountable_model_without_framing_preserves_legacy_local_fallback(): + counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=404)}) + + result: Final = await _ladder(counter, TokenCountSettings(framing=None, count_model=None)) + + assert isinstance(result, ProviderCountFailure) + assert (result.status_code, result.message) == (404, "scripted 404 for sonnet") + + +@pytest.mark.asyncio +async def test_uncountable_model_with_framing_falls_through_to_the_local_estimate(): + counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=400)}) + + result: Final = await _ladder(counter, TokenCountSettings(framing=0, count_model=None)) + + assert isinstance(result, ProviderCountFailure) + assert len(counter.calls) == 1 + + +@pytest.mark.parametrize("status_code", [400, 404]) +@pytest.mark.asyncio +async def test_token_count_model_failure_allows_a_framed_local_estimate(status_code: int): + counter: Final = _ScriptedCounter( + {"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=status_code)} + ) + + result: Final = await _ladder(counter, TokenCountSettings(framing=78, count_model=_OPUS_IN_US_EAST_1)) + + assert isinstance(result, ProviderCountFailure) + assert result.status_code == status_code + assert "us.anthropic.claude-opus-5-5 in us-east-1" in result.message + + +@pytest.mark.asyncio +async def test_uncountable_model_with_framing_raises_when_local_counting_is_disabled(): + counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=400)}) + + with pytest.raises(ProxyException) as raised: + await _ladder(counter, TokenCountSettings(framing=10, count_model=None), local_estimate_allowed=False) + + assert raised.value.code == "400" + + +@pytest.mark.asyncio +async def test_rejected_provider_conversion_still_allows_a_local_estimate(): + counter: Final = _ScriptedCounter( + { + "us.anthropic.claude-sonnet-5-5": _provider_reply( + "sonnet", status_code=400, rejected_before_provider_call=True + ) + } + ) + + result: Final = await _ladder(counter, TokenCountSettings(framing=100, count_model=None)) + + assert isinstance(result, ProviderCountFailure) + assert (result.status_code, result.message) == (400, "scripted 400 for sonnet") + + +@pytest.mark.parametrize("model", ["openai/gpt-5.4-mini", "gemini/gemini-3.8-flash"]) +@pytest.mark.asyncio +async def test_framing_keeps_the_400_for_blocks_the_provider_counter_cannot_convert( + monkeypatch: pytest.MonkeyPatch, model: str +): + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "framed", + "litellm_params": {"model": model, "api_key": "sk-unused-framed"}, + "model_info": {"token_count_framing": 100}, + } + ] + ), + ) + search_result: Final = {"type": "search_result", "source": "https://example.com", "title": "Notes", "content": []} + + with pytest.raises(ProxyException) as raised: + await count_request_tokens( + request=TokenCountRequest( + model="framed", + messages=[{"role": "user", "content": [{"type": "text", "text": "Summarize."}, search_result]}], + ), + call_endpoint=True, + messages_format="anthropic", + ) + + assert raised.value.code == "400" + assert "search_result" in raised.value.message + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "unreadable_body", + [ + pytest.param( + { + "messages": [ + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + { + "type": "server_tool_use", + "id": "srvtoolu_01", + "name": "web_search", + "input": {"query": "new"}, + }, + {"type": "web_search_tool_result", "tool_use_id": "srvtoolu_01", "content": []}, + ], + }, + ] + }, + id="web_search_blocks", + ), + pytest.param( + { + "messages": [{"role": "user", "content": "What is new?"}], + "tools": [{"name": "search", "input_schema": {"type": "object", "properties": {"query": True}}}], + }, + id="boolean_subschema_tool", + ), + pytest.param( + { + "messages": [{"role": "user", "content": "What is new?"}], + "tools": [ + {"name": "search", "input_schema": {"type": "object", "properties": {"tags": {"type": "array"}}}} + ], + }, + id="array_schema_without_items", + ), + ], +) +async def test_framed_estimate_that_cannot_read_the_body_answers_the_provider_error( + monkeypatch: pytest.MonkeyPatch, unreadable_body: dict[str, list[dict[str, object]]] +): + counter: Final = _ScriptedCounter({"us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[{"model_name": "kimi", "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}}]), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + with pytest.raises(ProxyException) as raised: + await token_counter(request=TokenCountRequest(model="kimi", **unreadable_body), call_endpoint=True) + readable: Final = await token_counter( + request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]), + call_endpoint=True, + ) + + assert (raised.value.code, raised.value.message) == ("400", "scripted 400 for kimi") + assert readable.count_source == "local_estimate" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "contents", + [ + pytest.param([{"role": "user", "parts": [{"text": "What is new in the release notes?"}]}], id="text"), + pytest.param( + [ + {"role": "user", "parts": [{"text": "Weather in Paris?"}]}, + { + "role": "model", + "parts": [ + {"functionCall": {"name": "get_weather", "args": {"city": "Paris"}}, "thoughtSignature": "c2ln"} + ], + }, + { + "role": "user", + "parts": [{"functionResponse": {"name": "get_weather", "response": {"forecast": "Sun"}}}], + }, + ], + id="sdk_agent_loop", + ), + pytest.param( + [ + {"role": "user", "parts": [{"text": "Hi"}]}, + {"role": "model", "parts": [{"text": "Hello!", "thoughtSignature": "c2ln"}]}, + {"role": "user", "parts": [{"text": "What is new?"}]}, + ], + id="sdk_chat_with_a_thought_signature", + ), + pytest.param([], id="empty"), + ], +) +async def test_framed_estimate_counts_google_contents_like_the_same_conversation_sent_as_messages( + monkeypatch: pytest.MonkeyPatch, contents: list[dict[str, object]] +): + _framed_kimi(monkeypatch) + google_request: Final = _google_count_tokens_request("kimi", {"contents": contents}) + + from_google: Final = await token_counter(request=google_request, call_endpoint=True) + from_messages: Final = await token_counter( + request=TokenCountRequest(model="kimi", messages=google_request.messages), call_endpoint=True + ) + + assert from_google.count_source == "local_estimate" + assert from_google.total_tokens == from_messages.total_tokens + + +@pytest.mark.asyncio +async def test_framed_deployment_whose_custom_tokenizer_cannot_load_answers_the_provider_error( + monkeypatch: pytest.MonkeyPatch, +): + counter: Final = _ScriptedCounter({"us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "kimi", + "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}, + "model_info": {"custom_tokenizer": {"identifier": "moonshotai/kimi-k3", "revision": "main"}}, + } + ] + ), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + with pytest.raises(ProxyException) as raised: + await token_counter( + request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]), + call_endpoint=True, + ) + + assert (raised.value.code, raised.value.message) == ("400", "scripted 400 for kimi") + + +def _framed_kimis_with_custom_tokenizers( + monkeypatch: pytest.MonkeyPatch, tokenizers: Mapping[str, tuple[str, str]], burst: int = 1 +) -> None: + counter: Final = _ScriptedCounter( + { + "us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400), + "us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", total_tokens=12), + }, + burst=burst, + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "sonnet", + "litellm_params": {"model": "bedrock/us.anthropic.claude-sonnet-5-5"}, + "model_info": {"token_count_model": {}}, + }, + *( + { + "model_name": name, + "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}, + "model_info": { + "custom_tokenizer": {"identifier": identifier, "revision": revision, "auth_token": None} + }, + } + for name, (identifier, revision) in tokenizers.items() + ), + ] + ), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + +def _framed_kimi_with_custom_tokenizer(monkeypatch: pytest.MonkeyPatch, identifier: str) -> None: + _framed_kimis_with_custom_tokenizers(monkeypatch, {"kimi": (identifier, "main")}) + + +async def _count_status(model: str = "kimi") -> str: + try: + response: Final = await token_counter( + request=TokenCountRequest(model=model, messages=[{"role": "user", "content": "What is new?"}]), + call_endpoint=True, + ) + except ProxyException as e: + return e.code + return response.count_source or "" + + +class _SkippableClockLoop(asyncio.SelectorEventLoop): + skipped_seconds: float = 0.0 + + def time(self) -> float: + return super().time() + self.skipped_seconds + + def skip(self, seconds: float) -> None: + self.skipped_seconds += seconds + + +def _run_on_skippable_clock( + scenario: Callable[[_SkippableClockLoop], Coroutine[object, object, tuple[str, ...]]], +) -> tuple[str, ...]: + loop: Final = _SkippableClockLoop() + try: + return loop.run_until_complete(scenario(loop)) + finally: + loop.run_until_complete(loop.shutdown_default_executor()) + loop.close() + + +async def _until_finished(tasks: tuple[asyncio.Future[str], ...], count: int) -> None: + for finished in itertools.islice(asyncio.as_completed(tasks), count): + await finished + + +def _tokenizer_that_panics_while_encoding() -> SelectTokenizerResponse: + template: Final = Tokenizer(WordLevel({"[UNK]": 0, "": 1}, unk_token="[UNK]")) + template.pre_tokenizer = Whitespace() + template.post_processor = TemplateProcessing(single=" $A", special_tokens=[("", 1)]) + spec: Final = json.loads(template.to_str()) + return litellm.utils.create_tokenizer( + json.dumps({**spec, "post_processor": {**spec["post_processor"], "special_tokens": {}}}) + ) + + +def test_hanging_tokenizer_download_answers_the_provider_error_after_the_wait_with_one_download_while_it_runs( + monkeypatch: pytest.MonkeyPatch, +): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/hanging-tokenizer") + download_started: Final = threading.Event() + hub_released: Final = threading.Event() + downloads: Final[list[str]] = [] + + def hanging_download(model: str, custom_tokenizer: object = None) -> NoReturn: + downloads.append(model) + download_started.set() + hub_released.wait(10) + raise OSError("hub unreachable") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", hanging_download) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_WAIT_SECONDS", 30) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", 1e-9) + + async def burst_past_the_wait_then_one_more(loop: _SkippableClockLoop) -> tuple[str, ...]: + try: + burst: Final = asyncio.gather(*(_count_status() for _ in range(60))) + await asyncio.to_thread(download_started.wait, 10) + loop.skip(31) + answered: Final = await asyncio.wait_for(burst, 8) + later: Final = await asyncio.wait_for(_count_status(), 8) + finally: + hub_released.set() + return (*answered, later) + + assert set(_run_on_skippable_clock(burst_past_the_wait_then_one_more)) == {"400"} + assert len(downloads) == 1 + + +@pytest.mark.asyncio +async def test_framed_count_waits_for_a_tokenizer_download_that_finishes_within_the_wait( + monkeypatch: pytest.MonkeyPatch, +): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/slow-tokenizer") + real_select_tokenizer: Final = litellm.utils._select_tokenizer + download_started: Final = threading.Event() + hub_released: Final = threading.Event() + + def slow_download(model: str, custom_tokenizer: object = None): + download_started.set() + hub_released.wait(10) + return real_select_tokenizer(model="gpt-5.4-mini") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", slow_download) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_WAIT_SECONDS", 30) + count: Final = asyncio.create_task(_count_status()) + try: + await asyncio.to_thread(download_started.wait, 10) + finally: + hub_released.set() + + assert await asyncio.wait_for(count, 10) == "local_estimate" + + +@pytest.mark.asyncio +async def test_failed_tokenizer_download_is_reused_until_the_retry_window_passes(monkeypatch: pytest.MonkeyPatch): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/flaky-tokenizer") + real_select_tokenizer: Final = litellm.utils._select_tokenizer + downloads: Final[list[str]] = [] + + def hub_down_once(model: str, custom_tokenizer: object = None): + downloads.append(model) + if len(downloads) == 1: + raise OSError("hub unreachable") + return real_select_tokenizer(model="gpt-5.4-mini") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", hub_down_once) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", 1e9) + within_window: Final = [await _count_status() for _ in range(3)] + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", 1e-9) + after_window: Final = await _count_status() + + assert within_window == ["400", "400", "400"] + assert after_window == "local_estimate" + assert len(downloads) == 2 + + +def test_failed_tokenizer_load_is_reused_for_the_window_after_it_fails(monkeypatch: pytest.MonkeyPatch): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/slow-failing-tokenizer") + download_started: Final = threading.Event() + hub_released: Final = threading.Event() + downloads: Final[list[str]] = [] + + def slow_failing_download(model: str, custom_tokenizer: object = None) -> NoReturn: + downloads.append(model) + download_started.set() + hub_released.wait(10) + raise OSError("hub unreachable") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", slow_failing_download) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_WAIT_SECONDS", 1000) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", 60) + + async def fail_after_longer_than_the_window_then_retry(loop: _SkippableClockLoop) -> tuple[str, ...]: + first: Final = asyncio.ensure_future(_count_status()) + try: + await asyncio.to_thread(download_started.wait, 10) + loop.skip(120) + finally: + hub_released.set() + return await asyncio.wait_for(first, 8), await asyncio.wait_for(_count_status(), 8) + + assert _run_on_skippable_clock(fail_after_longer_than_the_window_then_retry) == ("400", "400") + assert len(downloads) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "second_tokenizer", + [("lit-tests/tokenizer-b", "main"), ("lit-tests/tokenizer-a", "v2")], + ids=["identifier", "revision"], +) +async def test_framed_deployments_with_different_custom_tokenizers_each_load_their_own( + monkeypatch: pytest.MonkeyPatch, second_tokenizer: tuple[str, str] +): + _framed_kimis_with_custom_tokenizers( + monkeypatch, {"kimi-a": ("lit-tests/tokenizer-a", "main"), "kimi-b": second_tokenizer} + ) + real_select_tokenizer: Final = litellm.utils._select_tokenizer + loaded: Final[list[object]] = [] + + def load(model: str, custom_tokenizer: object = None) -> SelectTokenizerResponse: + loaded.append(custom_tokenizer) + return real_select_tokenizer(model="gpt-5.4-mini") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", load) + + statuses: Final = [await _count_status(name) for name in ("kimi-a", "kimi-b")] + + assert statuses == ["local_estimate", "local_estimate"] + assert loaded == [ + {"identifier": "lit-tests/tokenizer-a", "revision": "main", "auth_token": None}, + {"identifier": second_tokenizer[0], "revision": second_tokenizer[1], "auth_token": None}, + ] + + +@pytest.mark.asyncio +async def test_tokenizer_loads_past_the_cap_answer_the_provider_error_without_stalling_other_counts( + monkeypatch: pytest.MonkeyPatch, +): + tokenizers: Final = {f"kimi-{n}": (f"lit-tests/unreachable-tokenizer-{n}", "main") for n in range(5)} + _framed_kimis_with_custom_tokenizers(monkeypatch, tokenizers, burst=len(tokenizers)) + real_select_tokenizer: Final = litellm.utils._select_tokenizer + hub_released: Final = threading.Event() + downloads: Final[list[object]] = [] + + def held_download(model: str, custom_tokenizer: object = None) -> SelectTokenizerResponse: + downloads.append(custom_tokenizer) + hub_released.wait(10) + return real_select_tokenizer(model="gpt-5.4-mini") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", held_download) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_MAX_RUNNING", 2) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_WAIT_SECONDS", 30) + monkeypatch.setattr(anyio.to_thread.current_default_thread_limiter(), "total_tokens", 2) + framed: Final = tuple(asyncio.ensure_future(_count_status(name)) for name in tokenizers) + try: + await asyncio.wait_for(_until_finished(framed, 3), 8) + unrelated: Final = await asyncio.wait_for(_count_status("sonnet"), 5) + finally: + hub_released.set() + statuses: Final = await asyncio.wait_for(asyncio.gather(*framed), 8) + + assert sorted(statuses) == ["400"] * 3 + ["local_estimate"] * 2 + assert unrelated == "provider" + assert len(downloads) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("panics_while", ["loading", "counting"]) +async def test_framed_count_answers_the_provider_error_when_the_tokenizer_panics( + monkeypatch: pytest.MonkeyPatch, panics_while: str +): + _framed_kimi_with_custom_tokenizer(monkeypatch, f"lit-tests/panicking-tokenizer-{panics_while}") + panicking: Final = _tokenizer_that_panics_while_encoding() + + def load(model: str, custom_tokenizer: object = None) -> SelectTokenizerResponse: + if panics_while == "loading": + panicking["tokenizer"].encode("hello world") + return panicking + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", load) + + assert await _count_status() == "400" + + +def _framed_kimi(monkeypatch: pytest.MonkeyPatch) -> None: + counter: Final = _ScriptedCounter({"us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[{"model_name": "kimi", "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}}]), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + +@pytest.mark.asyncio +async def test_local_count_loads_the_tokenizer_again_after_a_failed_load(monkeypatch: pytest.MonkeyPatch): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/local-route-tokenizer") + real_select_tokenizer: Final = litellm.utils._select_tokenizer + downloads: Final[list[str]] = [] + + def hub_down_once(model: str, custom_tokenizer: object = None): + downloads.append(model) + if len(downloads) == 1: + raise OSError("hub unreachable") + return real_select_tokenizer(model="gpt-5.4-mini") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", hub_down_once) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", 1e9) + request: Final = TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]) + + with pytest.raises(OSError, match="hub unreachable"): + await token_counter(request=request) + retried: Final = await token_counter(request=request) + + assert retried.count_source == "local_estimate" and retried.total_tokens > 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("reuse_seconds", "tokenizer_types"), + [(1e9, ("openai_tokenizer", "openai_tokenizer")), (1e-9, ("openai_tokenizer", "huggingface_tokenizer"))], + ids=["within_the_window", "after_the_window"], +) +async def test_framed_count_loads_the_tokenizer_again_only_after_the_window_of_a_load_that_fell_back( + monkeypatch: pytest.MonkeyPatch, reuse_seconds: float, tokenizer_types: tuple[str, str] +): + _framed_kimi(monkeypatch) + real_select_tokenizer: Final = litellm.utils._select_tokenizer + hub_tokenizer: Final = litellm.utils.create_tokenizer( + Tokenizer(WordLevel({"[UNK]": 0}, unk_token="[UNK]")).to_str() + ) + loads: Final[list[str]] = [] + + def hub_up_on_the_second_load(model: str, custom_tokenizer: object = None): + loads.append(model) + return real_select_tokenizer(model="gpt-5.4-mini") if len(loads) == 1 else hub_tokenizer + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", hub_up_on_the_second_load) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", reuse_seconds) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_MAX_RUNNING", 1) + request: Final = TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]) + + first: Final = await token_counter(request=request, call_endpoint=True) + second: Final = await token_counter(request=request, call_endpoint=True) + + assert (first.tokenizer_type, second.tokenizer_type) == tokenizer_types + + +async def _fork_guard_traceback_frames() -> int: + with pytest.raises(ForkedAfterNativeRuntimeStarted) as raised: + await token_counter( + request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]), + call_endpoint=True, + ) + return len(traceback.extract_tb(raised.value.__traceback__)) + + +@pytest.mark.asyncio +async def test_fork_guard_error_while_loading_the_tokenizer_is_raised_to_every_count_sharing_the_load( + monkeypatch: pytest.MonkeyPatch, +): + _framed_kimi(monkeypatch) + + def forked_load(model: str, custom_tokenizer: object = None) -> NoReturn: + raise ForkedAfterNativeRuntimeStarted("native runtime started before fork") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", forked_load) + + frames: Final = [await _fork_guard_traceback_frames() for _ in range(3)] + + assert len(set(frames)) == 1 + + +@pytest.mark.asyncio +async def test_swallowed_estimate_error_is_logged_by_type_without_the_request_content( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +): + counter: Final = _ScriptedCounter({"us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[{"model_name": "kimi", "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}}]), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + verbose_proxy_logger.addHandler(caplog.handler) + try: + with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name), pytest.raises(ProxyException): + await token_counter( + request=TokenCountRequest( + model="kimi", + messages=[ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "ok", "tool_calls": "private-marker-in-the-prompt"}, + ], + ), + call_endpoint=True, + ) + finally: + verbose_proxy_logger.removeHandler(caplog.handler) + + warnings: Final = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] + assert any("count error 400" in message and "ValueError" in message for message in warnings) + assert not any("private-marker-in-the-prompt" in message for message in warnings) + + +@pytest.mark.asyncio +async def test_fork_guard_error_in_the_estimate_is_never_answered_as_the_provider_error(): + async def forked_estimate() -> NoReturn: + raise ForkedAfterNativeRuntimeStarted("native runtime started before fork") + + with pytest.raises(ForkedAfterNativeRuntimeStarted): + await _estimate_or_provider_error( + forked_estimate(), provider_failure=ProviderCountFailure(status_code=400, message="scripted 400") + ) + + +@pytest.mark.asyncio +async def test_local_estimate_adds_the_configured_framing_once_per_request(monkeypatch: pytest.MonkeyPatch): + def deployment(name: str, model_info: dict[str, object]) -> dict[str, object]: + return { + "model_name": name, + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": f"sk-unused-{name}"}, + "model_info": model_info, + } + + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[deployment("framed", {"token_count_framing": 100}), deployment("plain", {})]), + ) + messages: Final = [{"role": "system", "content": "be brief"}, {"role": "user", "content": "hello there"}] + + plain: Final = await token_counter(request=TokenCountRequest(model="plain", messages=messages)) + framed: Final = await token_counter(request=TokenCountRequest(model="framed", messages=messages)) + framed_prompt: Final = await token_counter(request=TokenCountRequest(model="framed", prompt="hello there")) + + assert plain.total_tokens == litellm.token_counter(model="gpt-5.4-mini", messages=messages) + assert framed.total_tokens == plain.total_tokens + 100 + assert (plain.count_source, framed.count_source) == ("local_estimate", "local_estimate") + assert framed_prompt.total_tokens == litellm.token_counter(model="gpt-5.4-mini", text="hello there") + + +@pytest.mark.asyncio +async def test_framing_comes_from_the_row_the_deployment_provider_or_base_model_names(monkeypatch: pytest.MonkeyPatch): + def deployment(name: str, litellm_params: dict[str, object], model_info: dict[str, object]) -> dict[str, object]: + return { + "model_name": name, + "litellm_params": {**litellm_params, "api_key": "unused", "api_base": "https://unused.invalid"}, + "model_info": model_info, + } + + kimi: Final = "databricks/system.ai.kimi-k3" + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + deployment("prefixed", {"model": kimi}, {}), + deployment("split", {"model": "system.ai.kimi-k3", "custom_llm_provider": "databricks"}, {}), + deployment("aliased", {"model": "gpt-5.4-mini", "custom_llm_provider": "openai"}, {"base_model": kimi}), + ] + ), + ) + messages: Final = [{"role": "user", "content": "hello there"}] + framing: Final = cost_map_entry(kimi)["token_count_framing"] + + prefixed: Final = await token_counter(request=TokenCountRequest(model="prefixed", messages=messages)) + split: Final = await token_counter(request=TokenCountRequest(model="split", messages=messages)) + aliased: Final = await token_counter(request=TokenCountRequest(model="aliased", messages=messages)) + + assert prefixed.total_tokens == litellm.token_counter(model="system.ai.kimi-k3", messages=messages) + framing + assert split.total_tokens == prefixed.total_tokens + assert aliased.total_tokens == litellm.token_counter(model="gpt-5.4-mini", messages=messages) + framing + + +_KIMI_WEATHER: Final = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, +} + + +def _byte_level_kimi_k3(path: Path) -> KimiK3TokenizerSource: + path.write_bytes(b"\n".join(base64.b64encode(bytes((value,))) + b" %d" % value for value in range(256))) + return KimiK3TokenizerSource(ranks_path=lambda: str(path), sha256=hashlib.sha256(path.read_bytes()).hexdigest()) + + +def _available_kimi_k3(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> tiktoken.Encoding: + source: Final = _byte_level_kimi_k3(tmp_path / "tiktoken.model") + monkeypatch.setattr( + litellm.proxy.proxy_server, "_KIMI_K3_ENCODINGS", KimiK3Encodings(source=source, retry_after_seconds=60) + ) + encoding: Final = load_kimi_k3_encoding(source) + assert encoding is not None + return encoding + + +def _shipped_kimi_k3_framing() -> int: + framing: Final = cost_map_entry("bedrock/us.moonshotai.kimi-k3")["token_count_framing"] + assert isinstance(framing, int) + return framing + + +@pytest.mark.asyncio +async def test_bedrock_kimi_k3_is_counted_with_its_tokenizer_plus_the_shipped_framing_in_both_request_formats( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +): + _framed_kimi(monkeypatch) + encoding: Final = _available_kimi_k3(monkeypatch, tmp_path) + question: Final = {"role": "user", "content": "Weather in Paris?"} + openai_messages: Final = [ + {"role": "system", "content": "Be terse."}, + question, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "c1", "type": "function", "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}} + ], + }, + {"role": "tool", "tool_call_id": "c1", "content": "Sun"}, + ] + + openai: Final = await count_request_tokens( + request=TokenCountRequest(model="kimi", messages=openai_messages, tools=[_KIMI_WEATHER]), call_endpoint=True + ) + anthropic: Final = await count_request_tokens( + request=TokenCountRequest( + model="kimi", + system="Be terse.", + messages=[ + question, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "c1", "name": "get_weather", "input": {"city": "Paris"}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "c1", "content": "Sun"}]}, + ], + tools=[ + { + "name": "get_weather", + "description": "Get the weather", + "input_schema": _KIMI_WEATHER["function"]["parameters"], + } + ], + ), + call_endpoint=True, + messages_format="anthropic", + ) + exact: Final = kimi_k3_tokens(encoding, "us.moonshotai.kimi-k3", openai_messages, [_KIMI_WEATHER]) + + assert exact is not None + assert (openai.total_tokens, openai.tokenizer_type, openai.count_source) == ( + exact + _shipped_kimi_k3_framing(), + "huggingface_tokenizer", + "local_estimate", + ) + assert (anthropic.total_tokens, anthropic.tokenizer_type) == (openai.total_tokens, "huggingface_tokenizer") + + +@pytest.mark.asyncio +async def test_bedrock_kimi_k3_deployment_with_a_custom_tokenizer_keeps_counting_with_it( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/kimi-custom-tokenizer") + _available_kimi_k3(monkeypatch, tmp_path) + custom: Final = litellm.utils._select_tokenizer(model="gpt-5.4-mini") + monkeypatch.setattr(litellm.utils, "_select_tokenizer", lambda model, custom_tokenizer=None: custom) + messages: Final = [{"role": "user", "content": "What is new?"}] + + counted: Final = await token_counter(request=TokenCountRequest(model="kimi", messages=messages), call_endpoint=True) + + assert counted.total_tokens == ( + litellm.token_counter(model=counted.model_used, messages=messages, custom_tokenizer=custom) + + kimi_k3_template_tokens(len(messages), False) + + _shipped_kimi_k3_framing() + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("messages", "tools", "template_tokens"), + [ + pytest.param([{"role": "user", "content": "What is new?"}], None, 14, id="one_message"), + pytest.param( + [ + {"role": "system", "content": "Be terse."}, + {"role": "user", "content": "Weather in Paris?"}, + {"role": "assistant", "content": "Sunny."}, + {"role": "user", "content": "And Oslo?"}, + ], + [_KIMI_WEATHER], + 14 + 3 * 13 + 40, + id="turns_and_tools", + ), + ], +) +async def test_bedrock_kimi_k3_without_its_tokenizer_adds_the_template_offsets_to_the_local_estimate( + monkeypatch: pytest.MonkeyPatch, + messages: list[dict[str, object]], + tools: list[dict[str, object]] | None, + template_tokens: int, +): + _framed_kimi(monkeypatch) + + counted: Final = await token_counter( + request=TokenCountRequest(model="kimi", messages=messages, tools=tools), call_endpoint=True + ) + estimate: Final = litellm.token_counter( + model=counted.model_used, + messages=messages, + tools=tools, + custom_tokenizer=litellm.utils._select_tokenizer(model=counted.model_used), + ) + + assert counted.total_tokens == estimate + template_tokens + _shipped_kimi_k3_framing() + + +@pytest.mark.asyncio +async def test_slow_cost_map_lookup_does_not_stall_other_requests_on_the_worker(monkeypatch: pytest.MonkeyPatch): + real_get_model_info: Final = litellm.get_model_info + lookup_started: Final = threading.Event() + other_request_ran: Final = threading.Event() + released_by_the_other_request: Final[list[bool]] = [] + + def slow_get_model_info(model: str, *args: object, **kwargs: object): + if model == "slow-lookup-model": + lookup_started.set() + released_by_the_other_request.append(other_request_ran.wait(5)) + return real_get_model_info(model, *args, **kwargs) + + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", None) + monkeypatch.setattr(litellm, "get_model_info", slow_get_model_info) + counting: Final = asyncio.create_task( + token_counter( + request=TokenCountRequest(model="slow-lookup-model", messages=[{"role": "user", "content": "hi"}]) + ) + ) + await asyncio.to_thread(lookup_started.wait, 5) + other_request_ran.set() + await counting + + assert set(released_by_the_other_request) == {True} + + +@pytest.mark.asyncio +async def test_uncountable_model_keeps_the_provider_error_when_local_counting_is_disabled( + monkeypatch: pytest.MonkeyPatch, +): + counter: Final = _ScriptedCounter({"us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "kimi", + "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}, + "model_info": {"token_count_framing": 78}, + } + ] + ), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + monkeypatch.setattr(litellm, "disable_token_counter", True) + + with pytest.raises(ProxyException) as raised: + await token_counter( + request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True + ) + + assert raised.value.code == "400" + assert "scripted 400 for kimi" in raised.value.message + + +@pytest.mark.parametrize( + ("model_info", "counted_with"), + [ + ({}, ("us.anthropic.claude-opus-5-5", "us-east-1", 447)), + ( + {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": "ap-southeast-4"}}, + ("us.anthropic.claude-opus-5-5", "ap-southeast-4", 447), + ), + ({"token_count_model": {}}, ("us.anthropic.claude-sonnet-5-5", "us-west-2", 300)), + ], +) +@pytest.mark.usefixtures("bedrock_sonnet_5_5_mapped_as_on_main") +@pytest.mark.asyncio +async def test_configured_token_count_model_overrides_the_shipped_one_and_an_empty_one_counts_the_model_itself( + monkeypatch: pytest.MonkeyPatch, model_info: dict[str, object], counted_with: tuple[str, str, int] +): + counter: Final = _ScriptedCounter( + { + "us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", total_tokens=300), + "us.anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447), + } + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "sonnet", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-sonnet-5-5", + "aws_region_name": "us-west-2", + }, + "model_info": model_info, + } + ] + ), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + counted: Final = await token_counter( + request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True + ) + + assert (counter.calls[0][0], counter.regions[0], counted.total_tokens) == counted_with + assert len(counter.calls) == 1 + + +@pytest.mark.parametrize( + ("aws_region_name", "counted_in"), + [("eu-central-1", "eu-central-1"), ("us-west-2", "us-east-1")], +) +@pytest.mark.usefixtures("bedrock_sonnet_5_5_mapped_as_on_main") +@pytest.mark.asyncio +async def test_shipped_token_count_model_keeps_the_prompt_in_the_deployment_geography( + monkeypatch: pytest.MonkeyPatch, aws_region_name: str, counted_in: str +): + counter: Final = _ScriptedCounter({"anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "sonnet", + "litellm_params": { + "model": "bedrock/anthropic.claude-sonnet-5-5", + "aws_region_name": aws_region_name, + }, + } + ] + ), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + await token_counter( + request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True + ) + + assert (counter.calls[0][0], counter.regions[0]) == ("anthropic.claude-opus-5-5", counted_in) + + +class _SlowVault: + def sync_read_secret(self, secret_name: str) -> None: + time.sleep(0.5) + + +async def _longest_stall(done: asyncio.Event) -> float: + if done.is_set(): + return 0.0 + started: Final = time.monotonic() + await asyncio.sleep(0.01) + return max(time.monotonic() - started, await _longest_stall(done)) + + +@pytest.mark.usefixtures("bedrock_sonnet_5_5_mapped_as_on_main") +@pytest.mark.asyncio +async def test_region_lookup_for_a_shipped_count_model_leaves_the_event_loop_free(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy import proxy_server + + counter: Final = _ScriptedCounter({"anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)}) + + def bedrock_counter(deployment: object, model: str) -> tuple[_ScriptedCounter, str, str]: + return counter, model, "bedrock" + + monkeypatch.setattr( + proxy_server, + "llm_router", + Router( + model_list=[{"model_name": "sonnet", "litellm_params": {"model": "bedrock/anthropic.claude-sonnet-5-5"}}] + ), + ) + monkeypatch.setattr(proxy_server, "_get_provider_token_counter", bedrock_counter) + monkeypatch.setattr(litellm, "secret_manager_client", _SlowVault()) + monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.HASHICORP_VAULT) + monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only")) + done: Final = asyncio.Event() + + async def count() -> TokenCountResponse: + try: + return await token_counter( + request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]), + call_endpoint=True, + ) + finally: + done.set() + + started: Final = time.monotonic() + stall, response = await asyncio.gather(_longest_stall(done), count()) + + assert time.monotonic() - started >= 1.0 + assert stall < 0.4 + assert response.total_tokens == 447 + + +@pytest.mark.asyncio +async def test_framing_still_applies_while_the_deployment_is_cooling_down(monkeypatch: pytest.MonkeyPatch): + router: Final = Router( + model_list=[ + { + "model_name": "framed", + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "sk-unused-framed"}, + "model_info": {"id": "framed-deployment", "token_count_framing": 100}, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + messages: Final = [{"role": "user", "content": "hello there"}] + warm: Final = await token_counter(request=TokenCountRequest(model="framed", messages=messages)) + + router.cooldown_cache.add_deployment_to_cooldown( + model_id="framed-deployment", + original_exception=Exception("rate limited"), + exception_status=429, + cooldown_time=60, + ) + cooling: Final = await token_counter(request=TokenCountRequest(model="framed", messages=messages)) + + assert warm.total_tokens == litellm.token_counter(model="gpt-5.4-mini", messages=messages) + 100 + assert cooling.total_tokens - 100 == litellm.token_counter(model="framed", messages=messages) + assert cooling.count_source == "local_estimate" + + +class _GeminiPart(BaseModel): + model_config = ConfigDict(frozen=True) + text: str = "" + + +class _GeminiContent(BaseModel): + model_config = ConfigDict(frozen=True) + parts: tuple[_GeminiPart, ...] = () + + +class _GeminiCountTokensBody(BaseModel): + model_config = ConfigDict(frozen=True) + contents: tuple[_GeminiContent, ...] = () + + +def _word_count(content: _GeminiContent) -> int: + return sum(len(part.text.split()) for part in content.parts) + + +def _gemini_word_counting_peer(request: httpx.Request) -> httpx.Response: + body: Final = _GeminiCountTokensBody.model_validate_json(request.content) + total: Final = sum(_word_count(content) for content in body.contents) + return httpx.Response(200, json={"totalTokens": total}, request=request) + + +def _unavailable_peer(request: httpx.Request) -> httpx.Response: + return httpx.Response(503, json={"error": {"message": "unavailable"}}, request=request) + + +def _gemini_router(drop_params: bool = False) -> Router: + return Router( + model_list=[ + { + "model_name": "gemini-counted", + "litellm_params": { + "model": "gemini/gemini-3.8-flash", + "api_key": "test-key", + "api_base": "https://count-tokens.test", + "drop_params": drop_params, + }, + } + ] + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "text", ["hello there", "The quick brown fox jumps over the lazy dog. " * 50], ids=["short", "long"] +) +async def test_call_endpoint_counts_a_prompt_like_the_same_text_sent_as_a_user_message( + text: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + from_prompt: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", prompt=text), call_endpoint=True + ) + from_message: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", messages=[{"role": "user", "content": text}]), + call_endpoint=True, + ) + + assert from_prompt.error is False + assert from_prompt.total_tokens == from_message.total_tokens == len(text.split()) + + +class _PartnerMessage(BaseModel): + model_config = ConfigDict(frozen=True) + content: str = "" + + +class _PartnerCountRequest(BaseModel): + model_config = ConfigDict(frozen=True) + messages: tuple[_PartnerMessage, ...] = () + + +async def _vertex_partner_word_count( + model: str, request_data: JsonValue, litellm_params: JsonValue +) -> dict[str, JsonValue]: + body: Final = _PartnerCountRequest.model_validate(request_data) + return {"input_tokens": sum(len(message.content.split()) for message in body.messages)} + + +@pytest.mark.asyncio +async def test_call_endpoint_counts_a_prompt_on_a_vertex_ai_deployment(monkeypatch: pytest.MonkeyPatch) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "vertex-counted", + "litellm_params": { + "model": "vertex_ai/claude-sonnet-5-5", + "vertex_project": "test-project", + "vertex_location": "us-east5", + }, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + text: Final = "The quick brown fox jumps over the lazy dog." + + with patch.object( + VertexAIPartnerModelsTokenCounter, "handle_count_tokens_request", side_effect=_vertex_partner_word_count + ): + from_prompt: Final = await token_counter( + request=TokenCountRequest(model="vertex-counted", prompt=text), call_endpoint=True + ) + + assert (from_prompt.error, from_prompt.total_tokens) == (False, 9) + + +_FOX: Final = "The quick brown fox jumps over the lazy dog." + + +_CONVERSATIONS: Final = ( + {"messages": [{"role": "user", "content": _FOX}]}, + {"contents": [{"role": "user", "parts": [{"text": _FOX}]}]}, +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("conversation", _CONVERSATIONS, ids=["messages", "contents"]) +@pytest.mark.parametrize("call_endpoint", [True, False], ids=["provider", "local"]) +async def test_prompt_next_to_a_conversation_answers_400_without_drop_params( + conversation: dict[str, JsonValue], call_endpoint: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer)) + + with ( + patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client), + pytest.raises(ProxyException) as rejected, + ): + await token_counter( + request=TokenCountRequest(model="gemini-counted", prompt="hi", **conversation), call_endpoint=call_endpoint + ) + + assert (rejected.value.code, rejected.value.param) == ("400", "prompt") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("conversation", "call_endpoint", "global_drop_params", "deployment_drop_params"), + [ + (_CONVERSATIONS[0], True, False, True), + (_CONVERSATIONS[1], True, False, True), + (_CONVERSATIONS[0], False, False, True), + (_CONVERSATIONS[0], False, True, False), + ], + ids=[ + "provider_messages", + "provider_contents", + "local_messages", + "local_messages_global_setting", + ], +) +async def test_drop_params_counts_a_conversation_without_the_prompt( + conversation: dict[str, JsonValue], + call_endpoint: bool, + global_drop_params: bool, + deployment_drop_params: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router(drop_params=deployment_drop_params)) + monkeypatch.setattr(litellm, "drop_params", global_drop_params) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + with_prompt: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", prompt="hi", **conversation), call_endpoint=call_endpoint + ) + conversation_only: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", **conversation), call_endpoint=call_endpoint + ) + + assert with_prompt.error is False + assert with_prompt.total_tokens == conversation_only.total_tokens > 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("empty_conversation", [{"messages": []}, {"contents": []}], ids=["messages", "contents"]) +async def test_call_endpoint_prompt_next_to_an_empty_conversation_counts_the_prompt( + empty_conversation: dict[str, JsonValue], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + counted: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", prompt=_FOX, **empty_conversation), call_endpoint=True + ) + + assert (counted.error, counted.total_tokens) == (False, 9) + + +@pytest.mark.asyncio +async def test_local_count_of_a_prompt_next_to_empty_messages_counts_the_prompt( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) + + with_empty: Final = await token_counter(request=TokenCountRequest(model="gemini-counted", prompt=_FOX, messages=[])) + alone: Final = await token_counter(request=TokenCountRequest(model="gemini-counted", prompt=_FOX)) + + assert with_empty.error is False + assert with_empty.total_tokens == alone.total_tokens > 0 + + +_WEATHER_TOOL: Final = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a city", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "chat_fields", + [{"system": "You answer in one short sentence. " * 20}, {"tools": [_WEATHER_TOOL]}], + ids=["system", "tools"], +) +@pytest.mark.parametrize("empty_conversation", [{}, {"messages": []}], ids=["no_messages", "empty_messages"]) +async def test_local_count_of_a_prompt_with_system_or_tools_counts_them_like_the_message_form( + chat_fields: dict[str, JsonValue], empty_conversation: dict[str, JsonValue], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) + + from_prompt: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", prompt=_FOX, **empty_conversation, **chat_fields) + ) + from_message: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", messages=[{"role": "user", "content": _FOX}], **chat_fields) + ) + prompt_alone: Final = await token_counter(request=TokenCountRequest(model="gemini-counted", prompt=_FOX)) + + assert from_prompt.error is False + assert from_prompt.total_tokens == from_message.total_tokens > prompt_alone.total_tokens + + +@pytest.mark.asyncio +async def test_call_endpoint_prompt_keeps_the_local_count_on_a_counter_that_needs_messages( + monkeypatch: pytest.MonkeyPatch, +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "anthropic-counted", + "litellm_params": { + "model": "anthropic/claude-sonnet-5-5", + "api_key": "test-key", + "api_base": "https://count-tokens.test", + }, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_unavailable_peer)) + request: Final = TokenCountRequest(model="anthropic-counted", prompt="The quick brown fox jumps over the lazy dog.") + + with patch.object(anthropic_count_tokens_handler, "get_async_httpx_client", return_value=client): + with_endpoint: Final = await token_counter(request=request, call_endpoint=True) + local: Final = await token_counter(request=request, call_endpoint=False) + + assert (with_endpoint.error, with_endpoint.total_tokens, with_endpoint.tokenizer_type) == ( + False, + local.total_tokens, + local.tokenizer_type, + ) diff --git a/tests/unit/test_count_tokens_public_api.py b/tests/unit/test_count_tokens_public_api.py index 84121dd63e4..e33bc1e5907 100644 --- a/tests/unit/test_count_tokens_public_api.py +++ b/tests/unit/test_count_tokens_public_api.py @@ -97,7 +97,7 @@ def test_acount_tokens_with_tools(): assert result.total_tokens == 30 mock_handler.assert_called_once() call_kwargs = mock_handler.call_args - assert call_kwargs.kwargs.get("tools") == tools + assert list(call_kwargs.kwargs.get("tools") or ()) == tools def test_acount_tokens_with_system(): diff --git a/tests/unit/test_model_prices_schema.py b/tests/unit/test_model_prices_schema.py index a05c345b5cd..211a43d2574 100644 --- a/tests/unit/test_model_prices_schema.py +++ b/tests/unit/test_model_prices_schema.py @@ -176,6 +176,52 @@ def test_generated_off_peak_schema_rejects_malformed_blocks(block: dict): assert not validator.is_valid({"some-model": {**OFF_PEAK_ENTRY, "off_peak_pricing": block}}) +TOKEN_COUNT_ENTRY: Final = MappingProxyType( + { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "token_count_framing": 78, + "token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": "us-east-1"}, + } +) + + +@pytest.mark.parametrize( + "setting", + [ + {"token_count_framing": 0}, + {"token_count_framing": 78}, + {"token_count_model": {"model": "us.anthropic.claude-opus-5-5"}}, + ], +) +def test_generator_classifies_token_count_settings(setting: dict): + generator = load_generator() + schema = json.loads(generator.render(generator.build_schema({"some-model": dict(TOKEN_COUNT_ENTRY)}))) + assert build_validator(schema).is_valid({"some-model": {**TOKEN_COUNT_ENTRY, **setting}}) + + +@pytest.mark.parametrize( + "setting", + [ + {"token_count_framing": -1}, + {"token_count_framing": "78"}, + {"token_count_framing": 1.5}, + {"token_count_model": {"model": ""}}, + {"token_count_model": {"model": " "}}, + {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": ""}}, + {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": " "}}, + {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": 5}}, + {"token_count_model": {"region": "us-east-1"}}, + {"token_count_model": "us.anthropic.claude-opus-5-5"}, + {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "zone": "us-east-1"}}, + ], +) +def test_generated_token_count_schema_rejects_malformed_settings(setting: dict): + generator = load_generator() + schema = json.loads(generator.render(generator.build_schema({"some-model": dict(TOKEN_COUNT_ENTRY)}))) + assert not build_validator(schema).is_valid({"some-model": {**TOKEN_COUNT_ENTRY, **setting}}) + + def find_duplicate_keys(path: Path) -> list[str]: duplicates: list[str] = [] diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index 74839831ca1..60b19cd742b 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -514,6 +514,40 @@ def test_should_not_pollute_shared_key_with_custom_nonzero_pricing(): ) +def test_token_count_settings_stay_on_the_deployment_that_set_them() -> None: + """A deployment's ``token_count_model`` and ``token_count_framing`` must not land on + the shared ``{provider}/{model}`` key, or a sibling deployment of that model would count with them.""" + from litellm.proxy.common_utils.token_count_settings import cost_map_entry + + backend_model: Final = "bedrock/us.moonshotai.kimi-k3" + deploy_id: Final = "kimi-deploy-keep-token-count" + configured: Final = { + "token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-west-2"}, + "token_count_framing": 7, + } + builtin: Final = {field: cost_map_entry(backend_model).get(field) for field in configured} + assert all(builtin[field] != value for field, value in configured.items()), builtin + model_keys: Final = { + key: copy.deepcopy(litellm.model_cost.get(key)) + for key in (deploy_id, backend_model, backend_model.removeprefix("bedrock/")) + } + try: + Router( + model_list=[ + { + "model_name": "kimi-keep-token-count", + "litellm_params": {"model": backend_model}, + "model_info": {"id": deploy_id, **configured}, + } + ], + ) + + assert {field: litellm.model_cost[deploy_id][field] for field in configured} == configured + assert {field: cost_map_entry(backend_model).get(field) for field in configured} == builtin + finally: + _restore_model_cost_entries(model_keys) + + def test_should_store_full_pricing_under_deployment_model_id(): """ Per-deployment pricing (including zero) should be stored and diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index da770fcf695..d01daa6d49d 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -956,6 +956,13 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_web_search": {"type": "boolean"}, "supports_bedrock_runtime_chat_completions_tools_with_reasoning": {"type": "boolean"}, "supports_bedrock_runtime_chat_completions_response_format": {"type": "boolean"}, + "token_count_framing": {"type": "integer"}, + "token_count_model": { + "type": "object", + "properties": {"model": {"type": "string"}, "region": {"type": "string"}}, + "required": ["model"], + }, + "token_count_tokenizer": {"type": "string", "enum": ["kimi_k3"]}, "supports_url_context": {"type": "boolean"}, "supports_multimodal": {"type": "boolean"}, "uses_embed_content": {"type": "boolean"}, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 30c3b471559..ab9c3cff95f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -46043,9 +46043,17 @@ export interface components { * @description Response structure for token count details with modality breakdown. * * Example: - * {'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}]} + * {'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}], + * 'litellm_count_source': 'provider'} */ TokenCountDetailsResponse: { + /** Litellm Count Model */ + litellm_count_model?: string; + /** + * Litellm Count Source + * @enum {string} + */ + litellm_count_source?: "provider" | "token_count_model" | "local_estimate"; /** Prompttokensdetails */ promptTokensDetails: components["schemas"]["PromptTokensDetails"][]; /** Totaltokens */ @@ -46074,6 +46082,8 @@ export interface components { }; /** TokenCountResponse */ TokenCountResponse: { + /** Count Source */ + count_source?: ("provider" | "token_count_model" | "local_estimate") | null; /** * Error * @default false From a251c0d6b4a7cae9c0cf3db49d5b2ca4cabbae66 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 00:29:29 +0000 Subject: [PATCH 09/12] Merge pull request #47446 from BerriAI/litellm_backport_google_count_shapes_stable_1_104_x fix(count_tokens): restore Google request compatibility on stable/1.104.x (cherry picked from commit a63df167da30438b08fb1568b340ab787a4e182a) --- litellm/proxy/google_endpoints/endpoints.py | 11 ++++-- .../test_google_gemini_proxy_request.py | 39 ++++++++++++++++++- 2 files changed, 46 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 5fbe780379e..5dbcbcba8e1 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -5,7 +5,7 @@ from typing import Final from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse, ORJSONResponse -from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError, field_validator from litellm.litellm_core_utils.json_validation_rule import normalize_json_schema_types from litellm.proxy._types import * @@ -31,6 +31,11 @@ class _GoogleCountTokensSystemInstruction(BaseModel): model_config = ConfigDict(frozen=True) parts: tuple[_GoogleCountTokensTextPart, ...] + @field_validator("parts", mode="before") + @classmethod + def normalize_parts(cls, parts: object) -> object: + return (parts,) if isinstance(parts, Mapping) else parts + class _GoogleCountTokensFunctionDeclaration(BaseModel): model_config = ConfigDict(frozen=True, populate_by_name=True) @@ -51,7 +56,7 @@ class _GoogleCountTokensRequest(BaseModel): model_config = ConfigDict(frozen=True, extra="ignore", populate_by_name=True) contents: tuple[Mapping[str, JsonValue], ...] = () system_instruction: _GoogleCountTokensSystemInstruction | None = Field(default=None, alias="systemInstruction") - tools: tuple[_GoogleCountTokensTool, ...] = () + tools: tuple[_GoogleCountTokensTool, ...] | None = () router: Final = APIRouter( @@ -111,7 +116,7 @@ def _google_count_tokens_request(model_name: str, data: object) -> TokenCountReq if validated_data.system_instruction is not None else None ) - declarations: Final = chain.from_iterable(tool.function_declarations for tool in validated_data.tools) + declarations: Final = chain.from_iterable(tool.function_declarations for tool in validated_data.tools or ()) tools: Final = [ { "name": declaration.name, diff --git a/tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py b/tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py index 6f8f90efc73..c465b6167e6 100644 --- a/tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py +++ b/tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py @@ -9,11 +9,12 @@ The request payload is correctly processed and forwarded to the httpx client. import json import unittest.mock -from typing import Optional +from typing import Final, Optional from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +from pydantic import JsonValue, ValidationError # Add the parent directory to the system path @@ -493,3 +494,39 @@ async def test_gemini_proxy_config_with_custom_api_base(): if __name__ == "__main__": # Run the tests pytest.main([__file__, "-v"]) + + +@pytest.mark.parametrize("system_field", ("systemInstruction", "system_instruction")) +@pytest.mark.parametrize("parts", ({"text": "Follow these rules"}, [{"text": "Follow these rules"}])) +@pytest.mark.parametrize("tools", (None, [])) +def test_google_count_tokens_accepts_equivalent_system_parts_and_no_tools( + system_field: str, parts: JsonValue, tools: JsonValue +) -> None: + from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request + + contents: Final = [{"role": "user", "parts": [{"text": "Count this request"}]}] + request: Final = _google_count_tokens_request( + "counted", {"contents": contents, system_field: {"parts": parts}, "tools": tools} + ) + + assert request.messages == [{"role": "user", "content": "Count this request"}] + assert request.system == [{"type": "text", "text": "Follow these rules"}] + assert request.tools is None + assert request.contents == contents + + +@pytest.mark.parametrize( + "parts", + ( + {"inlineData": {"mimeType": "image/png", "data": "aGVsbG8="}}, + [{"inlineData": {"mimeType": "image/png", "data": "aGVsbG8="}}], + ), +) +def test_google_count_tokens_rejects_media_only_system_parts(parts: JsonValue) -> None: + from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request + + with pytest.raises(ValidationError): + _google_count_tokens_request( + "counted", + {"contents": [{"role": "user", "parts": [{"text": "Hello"}]}], "systemInstruction": {"parts": parts}}, + ) From 422c2802148b290c0bac7014678daadf9a4d330e Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Mon, 19 Oct 2026 17:59:47 -0700 Subject: [PATCH 10/12] Merge pull request #47454 from BerriAI/litellm_restore_token_count_input_precedence fix(proxy): restore legacy token count input precedence (cherry picked from commit 5346a40c37a25b6e264c7b3783c5afb7d99ba456) --- litellm/proxy/proxy_server.py | 12 ------ .../test_vertex_gemini_count_tokens_wire.py | 38 +++++++++++++++- tests/unit/proxy/test_proxy_token_counter.py | 43 ++++--------------- 3 files changed, 45 insertions(+), 48 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e995ab384db..b4768df010c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -14267,11 +14267,6 @@ _PROVIDERS_COUNTING_WITHOUT_MESSAGES: Final = frozenset( ) -def _drops_params(deployment: Mapping[str, object] | None) -> bool: - litellm_params: Final = as_str_mapping(None if deployment is None else deployment.get("litellm_params")) - return litellm.drop_params is True or (litellm_params is not None and litellm_params.get("drop_params") is True) - - def _system_message(system: object) -> ChatCompletionSystemMessage | None: if not isinstance(system, (str, list)) or not system: return None @@ -14359,13 +14354,6 @@ async def count_request_tokens( model_to_use = _model conversation_empty: Final = not messages and not contents - if prompt and not conversation_empty and not _drops_params(deployment): - raise ProxyException( - message="prompt cannot be sent with messages or contents. Set drop_params to count them without the prompt", - type="invalid_request_error", - param="prompt", - code=400, - ) prompt_only: Final = bool(prompt) and conversation_empty prompt_as_message: Final = prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES ladder: Final = ( diff --git a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py index 5fe82d37e9e..e54fb1a4a13 100644 --- a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py @@ -220,8 +220,8 @@ _CONTENT_BLOCK_BODIES: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingPro ) -def _ai_studio_deployment(scenario: Scenario, api_base: str) -> str: - return scenario.model(model=f"gemini/{_MODEL}", api_base=api_base, api_key="scripted-gemini-key") +def _ai_studio_deployment(scenario: Scenario, api_base: str, **parameters: JsonValue) -> str: + return scenario.model(model=f"gemini/{_MODEL}", api_base=api_base, api_key="scripted-gemini-key", **parameters) def _counted_contents(count_body: Mapping[str, JsonValue]) -> JsonValue: @@ -594,3 +594,37 @@ def test_concurrent_count_waves_recover_after_vertex_count_outage(gateway: Gatew tuple(_json_body(request) for request in requests) == ({"contents": [{"role": "user", "parts": [{"text": _USER}]}]},) * 9 ) + + +@pytest.mark.parametrize( + ("conversation", "call_endpoint"), + ( + ({"messages": _MESSAGES}, True), + ({"contents": [{"role": "user", "parts": [{"text": _USER}]}]}, True), + ({"messages": _MESSAGES}, False), + ), + ids=("provider-count-messages", "provider-count-contents", "local-count-messages"), +) +@pytest.mark.parametrize("drop_params", (False, True), ids=("no-drop", "drop")) +def test_utils_token_counter_counts_the_conversation_regardless_of_prompt_and_drop_params( + gateway: Gateway, conversation: Mapping[str, JsonValue], call_endpoint: bool, drop_params: bool +) -> None: + peer: Final = _GeminiPeer() + params: Final = {"call_endpoint": "true" if call_endpoint else "false"} + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _ai_studio_deployment(scenario, wire.url, drop_params=drop_params) + mixed: Final = gateway.request( + "POST", "/utils/token_counter", {"model": model, "prompt": "hi", **conversation}, params=params + ) + mixed_requests: Final = _count_requests(wire) + alone: Final = gateway.request("POST", "/utils/token_counter", {"model": model, **conversation}, params=params) + alone_requests: Final = _count_requests(wire) + assert mixed.status_code == 200, mixed.text + assert alone.status_code == 200, alone.text + mixed_tokens: Final = TypeAdapter(int).validate_python(_JSON_OBJECT.validate_json(mixed.content)["total_tokens"]) + alone_tokens: Final = TypeAdapter(int).validate_python(_JSON_OBJECT.validate_json(alone.content)["total_tokens"]) + assert mixed_tokens == alone_tokens > 0 + assert len(mixed_requests) == len(alone_requests) == int(call_endpoint) + assert [_counted_contents(_json_body(request)) for request in mixed_requests] == [ + _counted_contents(_json_body(request)) for request in alone_requests + ] diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 995089105c4..9703e8ad1dc 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -2515,43 +2515,18 @@ _CONVERSATIONS: Final = ( ) -@pytest.mark.asyncio -@pytest.mark.parametrize("conversation", _CONVERSATIONS, ids=["messages", "contents"]) -@pytest.mark.parametrize("call_endpoint", [True, False], ids=["provider", "local"]) -async def test_prompt_next_to_a_conversation_answers_400_without_drop_params( - conversation: dict[str, JsonValue], call_endpoint: bool, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) - client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer)) - - with ( - patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client), - pytest.raises(ProxyException) as rejected, - ): - await token_counter( - request=TokenCountRequest(model="gemini-counted", prompt="hi", **conversation), call_endpoint=call_endpoint - ) - - assert (rejected.value.code, rejected.value.param) == ("400", "prompt") - - @pytest.mark.asyncio @pytest.mark.parametrize( - ("conversation", "call_endpoint", "global_drop_params", "deployment_drop_params"), - [ - (_CONVERSATIONS[0], True, False, True), - (_CONVERSATIONS[1], True, False, True), - (_CONVERSATIONS[0], False, False, True), - (_CONVERSATIONS[0], False, True, False), - ], - ids=[ - "provider_messages", - "provider_contents", - "local_messages", - "local_messages_global_setting", - ], + ("conversation", "call_endpoint"), + ((_CONVERSATIONS[0], True), (_CONVERSATIONS[1], True), (_CONVERSATIONS[0], False)), + ids=("provider_messages", "provider_contents", "local_messages"), ) -async def test_drop_params_counts_a_conversation_without_the_prompt( +@pytest.mark.parametrize( + ("global_drop_params", "deployment_drop_params"), + ((False, False), (False, True), (True, False), (True, True)), + ids=("no_drop", "deployment_drop", "global_drop", "both_drop"), +) +async def test_conversation_takes_precedence_over_prompt_regardless_of_drop_params( conversation: dict[str, JsonValue], call_endpoint: bool, global_drop_params: bool, From 60876f2b2aa5a70512b4b8cc23a1370e07e3bc3e Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 19 Oct 2026 22:57:30 -0700 Subject: [PATCH 11/12] chore: bump litellm 1.105.1 -> 1.105.2 litellm 1.105.1 already shipped, so the next release off this line is 1.105.2 --- pyproject.toml | 4 ++-- uv.lock | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 0d3692db5e8..3127fab91d0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.105.1" +version = "1.105.2" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.15" @@ -355,7 +355,7 @@ litellm-enterprise = { workspace = true } members = ["enterprise", "litellm-proxy-extras"] [tool.commitizen] -version = "1.105.1" +version = "1.105.2" version_files = [ "pyproject.toml:^version", ] diff --git a/uv.lock b/uv.lock index d98caabf059..8093175ed86 100644 --- a/uv.lock +++ b/uv.lock @@ -4500,7 +4500,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.105.1" +version = "1.105.2" source = { editable = "." } dependencies = [ { name = "aiohttp" }, From 9a4884cd121ae4a33516669431cc741fe8c67a8d Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 20 Oct 2026 03:11:23 -0700 Subject: [PATCH 12/12] fix(anthropic): carry thinking.block_binding only on the exact Anthropic, Azure AI and Bedrock Invoke configs The #46342 backport carried block_binding through every subclass of the Anthropic and Bedrock Invoke configs, so Bedrock Mantle, Claude Platform and OpenAI-compatible Anthropic endpoints started sending it. Each of the six chat and messages configs now opts in by exact class, and the Bedrock Invoke display updates and clear_thinking rebuilds carry block_binding only when the config opts in. Every other config sends the same thinking as 1.105.1 --- litellm/llms/anthropic/chat/transformation.py | 4 +- litellm/llms/anthropic/common_utils.py | 13 ++-- .../pass_through/messages/transformation.py | 4 +- .../anthropic/messages_transformation.py | 4 ++ .../llms/azure_ai/anthropic/transformation.py | 4 ++ .../anthropic_claude3_transformation.py | 12 ++-- .../anthropic_claude3_transformation.py | 22 +++++-- .../anthropic/test_anthropic_common_utils.py | 22 ++++++- ...ations_anthropic_claude3_transformation.py | 28 ++++++++- .../test_anthropic_claude3_transformation.py | 59 ++++++++++++++++++- 10 files changed, 147 insertions(+), 25 deletions(-) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 18fdbbba487..d3824897226 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -329,7 +329,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): @property def _carries_thinking_block_binding(self) -> bool: - return self._resolved_provider in ("anthropic", "azure_ai") + return type(self) is AnthropicConfig @classmethod def get_config(cls, *, model: str | None = None): @@ -1595,7 +1595,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): model=model, optional_params=optional_params, custom_llm_provider=self._resolved_provider, - carries_block_binding=self._carries_thinking_block_binding, + carry_block_binding=self._carries_thinking_block_binding, ) elif param == "reasoning_effort": # Accept both string ("low") and dict ({"effort": "low", diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 4448c67a212..234ffd2cc32 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -134,6 +134,11 @@ def validated_object_mapping(value: object) -> dict[object, object] | None: return None +def thinking_block_binding(thinking: Mapping[object, object] | None) -> dict[str, object]: + block_binding: Final = thinking.get("block_binding") if thinking is not None else None + return {} if block_binding is None else {"block_binding": block_binding} + + def _validated_claude_code_list(value: object) -> list[object] | None: try: return _CLAUDE_CODE_OBJECT_LIST_ADAPTER.validate_python(value) @@ -748,7 +753,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): model: str, optional_params: MutableMapping[str, object], # mutable-ok: in-place out-param like the sibling helpers custom_llm_provider: str, - carries_block_binding: bool = False, + carry_block_binding: bool = False, ) -> None: """Translate legacy ``thinking.type=enabled`` to adaptive for the adaptive-thinking models that reject it (4.7+ and the 5 families). @@ -774,11 +779,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): existing_output_config: Final = optional_params.get("output_config") display: Final = thinking.get("display") supported_display: Final = {"display": display} if display in ("summarized", "omitted") else {} - carried_block_binding: Final = ( - {"block_binding": thinking["block_binding"]} - if carries_block_binding and "block_binding" in thinking - else {} - ) + carried_block_binding: Final = thinking_block_binding(thinking) if carry_block_binding else {} optional_params["thinking"] = {**carried_block_binding, "type": "adaptive", **supported_display} optional_params["output_config"] = { "effort": effort, diff --git a/litellm/llms/anthropic/pass_through/messages/transformation.py b/litellm/llms/anthropic/pass_through/messages/transformation.py index b3a83e60bcc..1569c8f8380 100644 --- a/litellm/llms/anthropic/pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/pass_through/messages/transformation.py @@ -65,7 +65,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): @property def _carries_thinking_block_binding(self) -> bool: - return self._resolved_provider in ("anthropic", "azure_ai") + return type(self) is AnthropicMessagesConfig def get_supported_anthropic_messages_params(self, model: str) -> list: return [ @@ -518,7 +518,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): model=model, optional_params=anthropic_messages_optional_request_params, custom_llm_provider=self._resolved_provider, - carries_block_binding=self._carries_thinking_block_binding, + carry_block_binding=self._carries_thinking_block_binding, ) self._translate_adaptive_effort_for_non_adaptive_model( diff --git a/litellm/llms/azure_ai/anthropic/messages_transformation.py b/litellm/llms/azure_ai/anthropic/messages_transformation.py index 3d8b574e8c0..71cfe1c9d63 100644 --- a/litellm/llms/azure_ai/anthropic/messages_transformation.py +++ b/litellm/llms/azure_ai/anthropic/messages_transformation.py @@ -22,6 +22,10 @@ class AzureAnthropicMessagesConfig(AnthropicMessagesConfig): def custom_llm_provider(self) -> str | None: return "azure_ai" + @property + def _carries_thinking_block_binding(self) -> bool: + return type(self) is AzureAnthropicMessagesConfig + def should_strip_billing_metadata(self) -> bool: return True diff --git a/litellm/llms/azure_ai/anthropic/transformation.py b/litellm/llms/azure_ai/anthropic/transformation.py index 864d2134a84..79af7f0a049 100644 --- a/litellm/llms/azure_ai/anthropic/transformation.py +++ b/litellm/llms/azure_ai/anthropic/transformation.py @@ -39,6 +39,10 @@ class AzureAnthropicConfig(AnthropicConfig): def custom_llm_provider(self) -> str | None: return "azure_ai" + @property + def _carries_thinking_block_binding(self) -> bool: + return type(self) is AzureAnthropicConfig + def should_strip_billing_metadata(self) -> bool: return True diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index b3673eb1e6c..7f99b647025 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -12,7 +12,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import ( convert_url_to_base64, ) from litellm.llms.anthropic.chat.transformation import AnthropicConfig -from litellm.llms.anthropic.common_utils import AnthropicModelInfo +from litellm.llms.anthropic.common_utils import AnthropicModelInfo, thinking_block_binding from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( AmazonInvokeConfig, ) @@ -69,7 +69,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): @property def _carries_thinking_block_binding(self) -> bool: - return True + return type(self) is AmazonAnthropicClaudeConfig def should_strip_billing_metadata(self) -> bool: return True @@ -122,7 +122,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): model=original_model, optional_params=optional_params, custom_llm_provider="bedrock", - carries_block_binding=self._carries_thinking_block_binding, + carry_block_binding=self._carries_thinking_block_binding, ) translated_thinking: Final = optional_params.get("thinking") if ( @@ -130,7 +130,11 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): and isinstance(translated_thinking, dict) and translated_thinking.get("type") == "adaptive" ): - optional_params["thinking"] = {**translated_thinking, "display": "updates"} + optional_params["thinking"] = { + **(thinking_block_binding(translated_thinking) if self._carries_thinking_block_binding else {}), + "type": "adaptive", + "display": "updates", + } # The stub model hides the original model from the parent's forced-tool-use backstop response_format_tool_choice: Final = optional_params.get("tool_choice") diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index cf733f5605c..43582cfc252 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -17,7 +17,11 @@ from litellm.llms.anthropic.chat.transformation import ( DROP_UNSUPPORTED_OUTPUT_CONFIG_WARNING, AnthropicConfig, ) -from litellm.llms.anthropic.common_utils import AnthropicModelInfo, validated_object_mapping +from litellm.llms.anthropic.common_utils import ( + AnthropicModelInfo, + thinking_block_binding, + validated_object_mapping, +) from litellm.llms.anthropic.pass_through.messages.transformation import ( AnthropicMessagesConfig, ) @@ -83,7 +87,7 @@ class AmazonAnthropicClaudeMessagesConfig( @property def _carries_thinking_block_binding(self) -> bool: - return True + return type(self) is AmazonAnthropicClaudeMessagesConfig @property def beta_headers_provider(self) -> str: @@ -353,9 +357,13 @@ class AmazonAnthropicClaudeMessagesConfig( output_config = {} output_config.setdefault("effort", self._effort_from_thinking_budget(budget_tokens)) thinking: Final = validated_object_mapping(anthropic_messages_request.get("thinking")) - carried: Final = {key: value for key, value in (thinking or {}).items() if key not in ("type", "budget_tokens")} + display: Final = thinking.get("display") if thinking is not None else None anthropic_messages_request["output_config"] = output_config - anthropic_messages_request["thinking"] = {**carried, "type": "adaptive"} + anthropic_messages_request["thinking"] = { + **(thinking_block_binding(thinking) if self._carries_thinking_block_binding else {}), + "type": "adaptive", + **({} if display is None else {"display": display}), + } verbose_logger.debug( "Bedrock clear_thinking_20251015: injected adaptive thinking with effort=%s for model=%s", output_config.get("effort"), @@ -696,7 +704,11 @@ class AmazonAnthropicClaudeMessagesConfig( and isinstance(translated_thinking, dict) and translated_thinking.get("type") == "adaptive" ): - anthropic_messages_request["thinking"] = {**translated_thinking, "display": "updates"} + anthropic_messages_request["thinking"] = { + **(thinking_block_binding(translated_thinking) if self._carries_thinking_block_binding else {}), + "type": "adaptive", + "display": "updates", + } self._normalize_system_role_messages(anthropic_messages_request, model=model) ######################################################### ############## BEDROCK Invoke SPECIFIC TRANSFORMATION ### diff --git a/tests/unit/llms/anthropic/test_anthropic_common_utils.py b/tests/unit/llms/anthropic/test_anthropic_common_utils.py index 677ef16fc64..c6536b88e72 100644 --- a/tests/unit/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/unit/llms/anthropic/test_anthropic_common_utils.py @@ -2558,6 +2558,7 @@ def _rebuilt_thinking(carries_block_binding: bool) -> dict[str, object]: ("vertex_ai", False), ("databricks", False), ("bedrock_converse", False), + ("bedrock_mantle", False), ), ) def test_chat_legacy_thinking_rebuild_carries_only_block_binding_and_only_where_supported( @@ -2569,6 +2570,7 @@ def test_chat_legacy_thinking_rebuild_carries_only_block_binding_and_only_where_ from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( AmazonAnthropicClaudeConfig, ) + from litellm.llms.bedrock.chat.mantle.transformation import AmazonMantleConfig from litellm.llms.databricks.chat.transformation import DatabricksConfig from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import VertexAIAnthropicConfig @@ -2579,6 +2581,7 @@ def test_chat_legacy_thinking_rebuild_carries_only_block_binding_and_only_where_ "vertex_ai": (VertexAIAnthropicConfig(), "claude-opus-4-8"), "databricks": (DatabricksConfig(), "databricks-claude-opus-4-8"), "bedrock_converse": (AmazonConverseConfig(), "us.anthropic.claude-opus-4-8"), + "bedrock_mantle": (AmazonMantleConfig(), "anthropic.claude-opus-4-8"), }[provider] mapped: Final = config.map_openai_params( @@ -2594,16 +2597,29 @@ def test_chat_legacy_thinking_rebuild_carries_only_block_binding_and_only_where_ @pytest.mark.usefixtures("local_model_cost_map") @pytest.mark.parametrize( ("provider", "carries_block_binding"), - (("anthropic", True), ("azure_ai", True), ("bedrock_invoke", True), ("vertex_ai", False)), + ( + ("anthropic", True), + ("azure_ai", True), + ("bedrock_invoke", True), + ("vertex_ai", False), + ("bedrock_mantle", False), + ("bedrock_mantle_provider", False), + ("bedrock_claude_platform", False), + ("openai_like", False), + ), ) def test_messages_legacy_thinking_rebuild_carries_only_block_binding_and_only_where_supported( provider: str, carries_block_binding: bool ) -> None: from litellm.llms.anthropic.pass_through.messages.transformation import AnthropicMessagesConfig from litellm.llms.azure_ai.anthropic.messages_transformation import AzureAnthropicMessagesConfig + from litellm.llms.bedrock.claude_platform.messages_transformation import BedrockClaudePlatformMessagesConfig from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( AmazonAnthropicClaudeMessagesConfig, ) + from litellm.llms.bedrock.messages.mantle_transformation import AmazonMantleMessagesConfig + from litellm.llms.bedrock_mantle.messages.transformation import BedrockMantleAnthropicMessagesConfig + from litellm.llms.openai_like.messages.transformation import OpenAILikeAnthropicMessagesConfig from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation import ( VertexAIPartnerModelsAnthropicMessagesConfig, ) @@ -2614,6 +2630,10 @@ def test_messages_legacy_thinking_rebuild_carries_only_block_binding_and_only_wh "azure_ai": (AzureAnthropicMessagesConfig(), "claude-opus-4-8"), "bedrock_invoke": (AmazonAnthropicClaudeMessagesConfig(), "us.anthropic.claude-opus-4-8"), "vertex_ai": (VertexAIPartnerModelsAnthropicMessagesConfig(), "claude-opus-4-8"), + "bedrock_mantle": (AmazonMantleMessagesConfig(), "anthropic.claude-opus-4-8"), + "bedrock_mantle_provider": (BedrockMantleAnthropicMessagesConfig(), "anthropic.claude-opus-4-8"), + "bedrock_claude_platform": (BedrockClaudePlatformMessagesConfig(), "claude-opus-4-8"), + "openai_like": (OpenAILikeAnthropicMessagesConfig(), "claude-opus-4-8"), }[provider] request: Final = config.transform_anthropic_messages_request( diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index ce5c5ad1b27..97f83b48885 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -1003,8 +1003,14 @@ BLOCK_BINDING: Final = {"prefix_mismatch_behavior": "error"} @pytest.mark.parametrize( "requested_thinking", ( - {"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING}, - {"type": "enabled", "budget_tokens": 2048, "display": "updates", "block_binding": BLOCK_BINDING}, + {"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING, "future_key": 1}, + { + "type": "enabled", + "budget_tokens": 2048, + "display": "updates", + "block_binding": BLOCK_BINDING, + "future_key": 1, + }, ), ids=("adaptive", "legacy"), ) @@ -1030,6 +1036,24 @@ def test_bedrock_chat_invoke_keeps_block_binding_with_display_updates( assert result.get("thinking") == {"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING} +def test_bedrock_chat_mantle_rebuilds_display_updates_without_block_binding( + local_model_cost_map, local_beta_headers_config +) -> None: + from litellm.llms.bedrock.chat.mantle.transformation import AmazonMantleConfig + + optional_params: Final = AmazonMantleConfig().map_openai_params( + non_default_params={ + "max_tokens": 512, + "thinking": {"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING, "future_key": 1}, + }, + optional_params={}, + model="anthropic.claude-opus-4-8", + drop_params=False, + ) + + assert optional_params.get("thinking") == {"type": "adaptive", "display": "updates"} + + FINE_GRAINED_TOOL_STREAMING_BETA: Final = "fine-grained-tool-streaming-2025-05-14" EAGER_TOOL_SCHEMA: Final = {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]} diff --git a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index f303c899898..00277c98e80 100644 --- a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -36,6 +36,8 @@ from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_tran AmazonAnthropicClaudeMessagesConfig, AmazonAnthropicClaudeMessagesStreamDecoder, ) +from litellm.llms.bedrock.messages.mantle_transformation import AmazonMantleMessagesConfig +from litellm.llms.bedrock_mantle.messages.transformation import BedrockMantleAnthropicMessagesConfig @pytest.mark.asyncio @@ -3536,15 +3538,31 @@ BLOCK_BINDING: Final = {"prefix_mismatch_behavior": "error"} @pytest.mark.parametrize( ("model", "requested_thinking", "extra_params"), ( - ("eu.anthropic.claude-opus-5", {"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING}, {}), ( "eu.anthropic.claude-opus-5", - {"type": "enabled", "budget_tokens": 24000, "display": "updates", "block_binding": BLOCK_BINDING}, + {"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING, "future_key": 1}, + {}, + ), + ( + "eu.anthropic.claude-opus-5", + { + "type": "enabled", + "budget_tokens": 24000, + "display": "updates", + "block_binding": BLOCK_BINDING, + "future_key": 1, + }, {}, ), ( "us.anthropic.claude-opus-4-6", - {"type": "enabled", "budget_tokens": 2048, "display": "updates", "block_binding": BLOCK_BINDING}, + { + "type": "enabled", + "budget_tokens": 2048, + "display": "updates", + "block_binding": BLOCK_BINDING, + "future_key": 1, + }, {"context_management": {"edits": [{"type": "clear_thinking_20251015"}]}}, ), ), @@ -3566,6 +3584,41 @@ def test_bedrock_messages_keeps_block_binding_with_display_updates( assert result.get("thinking") == {"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING} +@pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config") +@pytest.mark.parametrize( + "config_class", + (AmazonMantleMessagesConfig, BedrockMantleAnthropicMessagesConfig), + ids=("bedrock_mantle_route", "bedrock_mantle_provider"), +) +@pytest.mark.parametrize( + ("requested_thinking", "extra_params"), + ( + ({"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING, "future_key": 1}, {}), + ( + {"type": "enabled", "budget_tokens": 2048, "display": "updates", "block_binding": BLOCK_BINDING}, + {"context_management": {"edits": [{"type": "clear_thinking_20251015"}]}}, + ), + ), + ids=("adaptive", "legacy_clear_thinking"), +) +def test_bedrock_mantle_messages_rebuilds_display_updates_without_block_binding( + config_class: type[AmazonMantleMessagesConfig], + requested_thinking: dict[str, object], + extra_params: dict[str, object], +) -> None: + from litellm.types.router import GenericLiteLLMParams + + result: Final = config_class().transform_anthropic_messages_request( + model="anthropic.claude-opus-4-8", + messages=[{"role": "user", "content": "Reply with OK"}], + anthropic_messages_optional_request_params={"max_tokens": 512, "thinking": requested_thinking, **extra_params}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result.get("thinking") == {"type": "adaptive", "display": "updates"} + + @pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config") @pytest.mark.parametrize("action", (None, "tool_addition", "tool_removal")) @pytest.mark.parametrize("explicit_beta", (False, True))