Merge pull request #47517 from BerriAI/patch_4fe6df3e
Some checks are pending
LiteLLM Rust / rust-wheel (push) Waiting to run
LiteLLM Rust / rust-lint (push) Waiting to run
LiteLLM Rust / rust-test (push) Waiting to run

chore(release): v1.106.1 patch
This commit is contained in:
yuneng-jiang 2026-10-20 09:11:00 -07:00 • committed by GitHub
commit 978b4d1e4e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
107 changed files with 10309 additions and 861 deletions

View file

@ -670,6 +670,7 @@ jobs:
- shard: proxy-extras
test-path: >-
tests/unit/litellm_proxy_extras
tests/unit/db_scripts
python-version: "3.12"
workers: 2
reruns: 0

View file

@ -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.",

View file

@ -0,0 +1,201 @@
import argparse
import os
import sys
from collections.abc import Iterator, Sequence
from dataclasses import dataclass
from itertools import accumulate, repeat, takewhile
from typing import Final, Literal, TextIO
import httpx
from pydantic import BaseModel, ConfigDict, JsonValue, ValidationError
from litellm.vector_stores.utils import is_litellm_written_metadata_key
DEFAULT_API_BASE: Final = "https://api.openai.com/v1"
LIST_PAGE_SIZE: Final = "100"
class VectorStore(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
id: str
metadata: dict[str, JsonValue] | None = None
class VectorStorePage(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
data: tuple[VectorStore, ...]
has_more: bool = False
class OpenAIErrorBody(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
message: str = ""
class OpenAIErrorResponse(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
error: OpenAIErrorBody
@dataclass(frozen=True, slots=True)
class Options:
api_key: str
api_base: str
vector_store_ids: tuple[str, ...]
extra_keys: frozenset[str]
apply: bool
@dataclass(frozen=True, slots=True)
class StoreResult:
store_id: str
matched_keys: tuple[str, ...]
status: Literal["dry_run", "updated", "failed"]
error: str = ""
def is_proxy_written_key(key: str, extra_keys: frozenset[str]) -> bool:
return is_litellm_written_metadata_key(key) or key in extra_keys
def matching_keys(store: VectorStore, extra_keys: frozenset[str]) -> tuple[str, ...]:
return tuple(sorted(key for key in (store.metadata or {}) if is_proxy_written_key(key, extra_keys)))
def fetch_page(client: httpx.Client, after: str | None) -> VectorStorePage:
params: Final = {"limit": LIST_PAGE_SIZE} if after is None else {"limit": LIST_PAGE_SIZE, "after": after}
return VectorStorePage.model_validate_json(client.get("/vector_stores", params=params).raise_for_status().content)
def fetch_next_page(client: httpx.Client, page: VectorStorePage | None) -> VectorStorePage | None:
if page is None or not page.has_more or not page.data:
return None
return fetch_page(client, after=page.data[-1].id)
def list_all_stores(client: httpx.Client) -> Iterator[VectorStore]:
pages: Final = accumulate(
repeat(None), lambda page, _: fetch_next_page(client, page), initial=fetch_page(client, after=None)
)
for page in takewhile(lambda page: page is not None, pages):
yield from page.data
def retrieve_stores(client: httpx.Client, store_ids: tuple[str, ...]) -> Iterator[VectorStore]:
for store_id in store_ids:
yield VectorStore.model_validate_json(client.get(f"/vector_stores/{store_id}").raise_for_status().content)
def describe_http_error(response: httpx.Response) -> str:
try:
message: Final = OpenAIErrorResponse.model_validate_json(response.content).error.message
except ValidationError:
return f"HTTP {response.status_code}"
return f"HTTP {response.status_code}: {message}"
def delete_keys(client: httpx.Client, store_id: str, keys: tuple[str, ...]) -> StoreResult:
try:
response: Final = client.post(f"/vector_stores/{store_id}", json={"metadata": dict.fromkeys(keys)})
except httpx.TransportError as error:
return StoreResult(store_id=store_id, matched_keys=keys, status="failed", error=type(error).__name__)
if response.is_success:
return StoreResult(store_id=store_id, matched_keys=keys, status="updated")
return StoreResult(store_id=store_id, matched_keys=keys, status="failed", error=describe_http_error(response))
def clean_store(client: httpx.Client, store: VectorStore, options: Options) -> StoreResult | None:
keys: Final = matching_keys(store, options.extra_keys)
if not keys:
return None
if not options.apply:
return StoreResult(store_id=store.id, matched_keys=keys, status="dry_run")
return delete_keys(client, store.id, keys)
def format_result(result: StoreResult) -> str:
key_list: Final = ", ".join(result.matched_keys)
count: Final = len(result.matched_keys)
match result.status:
case "dry_run":
return f"{result.store_id}: would delete {count} keys: {key_list}"
case "updated":
return f"{result.store_id}: deleted {count} keys: {key_list}"
case "failed":
return f"{result.store_id}: failed to delete {count} keys ({result.error}): {key_list}"
def clean_and_report(
client: httpx.Client, stores: Iterator[VectorStore], options: Options, out: TextIO
) -> Iterator[StoreResult | None]:
for store in stores:
result: Final = clean_store(client, store, options)
if result is not None:
print(format_result(result), file=out)
yield result
def run(options: Options, client: httpx.Client, out: TextIO) -> int:
stores: Final = (
retrieve_stores(client, options.vector_store_ids) if options.vector_store_ids else list_all_stores(client)
)
scanned: Final = tuple(clean_and_report(client, stores, options, out))
results: Final = tuple(result for result in scanned if result is not None)
updated: Final = sum(1 for result in results if result.status == "updated")
failed: Final = sum(1 for result in results if result.status == "failed")
mode: Final = "applied" if options.apply else "dry run, pass --apply to delete"
print(
f"scanned={len(scanned)} matched={len(results)} updated={updated} failed={failed} ({mode})",
file=out,
)
return 1 if failed else 0
def parse_options(argv: Sequence[str] | None) -> Options:
parser: Final = argparse.ArgumentParser(
description="Remove LiteLLM proxy-written keys from OpenAI vector store metadata. Dry run by default."
)
parser.add_argument("--api-key", default=os.environ.get("OPENAI_API_KEY"), help="defaults to $OPENAI_API_KEY")
parser.add_argument("--api-base", default=DEFAULT_API_BASE)
parser.add_argument("--vector-store-id", action="append", default=[], help="repeatable, defaults to every store")
parser.add_argument("--extra-key", action="append", default=[], help="repeatable, another key name to delete")
parser.add_argument("--apply", action="store_true", help="delete the matching keys instead of listing them")
args: Final = parser.parse_args(argv)
if not args.api_key:
parser.error("--api-key or OPENAI_API_KEY is required")
return Options(
api_key=args.api_key,
api_base=args.api_base.rstrip("/"),
vector_store_ids=tuple(args.vector_store_id),
extra_keys=frozenset(args.extra_key),
apply=args.apply,
)
def main(argv: Sequence[str] | None = None, transport: httpx.BaseTransport | None = None) -> int:
options: Final = parse_options(argv)
with httpx.Client(
base_url=options.api_base,
headers={"Authorization": f"Bearer {options.api_key}"},
timeout=60.0,
transport=transport,
) as client:
try:
return run(options, client, sys.stdout)
except httpx.HTTPStatusError as error:
sys.stderr.write(f"could not read vector stores: {describe_http_error(error.response)}\n")
return 2
except httpx.TransportError as error:
sys.stderr.write(f"could not reach {options.api_base}: {type(error).__name__}\n")
return 2
except ValidationError:
sys.stderr.write(f"unexpected vector store response shape from {options.api_base}\n")
return 2
if __name__ == "__main__":
sys.exit(main())

View file

@ -74,7 +74,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"))
@ -308,6 +308,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")
@ -951,47 +1012,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(

View file

@ -1588,6 +1588,7 @@ OUTPUT_TOKEN_CEILING_PARAMS: Final = frozenset({"max_tokens", "max_completion_to
CLIENT_OUTPUT_CEILING_METADATA_KEY: Final = "_client_output_ceiling"
CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags"
ROUTING_REQUEST_TAGS_METADATA_KEY: Final = "_routing_request_tags"
VECTOR_STORE_CALLER_METADATA_KWARG: Final = "litellm_vector_store_caller_metadata"
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: Final = "_litellm_router_usage_counted_tokens"
INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated"

View file

@ -219,7 +219,7 @@ class GitLabTemplateManager:
elif value.replace(".", "").isdigit():
try:
result[key] = float(value)
except Exception:
except ValueError:
result[key] = value
else:
result[key] = value.strip("\"'")

View file

@ -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, cast
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, runtime_checkable
import httpx
from pydantic import TypeAdapter, ValidationError
@ -52,6 +52,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

View file

@ -2100,10 +2100,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):
@ -2114,10 +2112,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(

View file

@ -2774,13 +2774,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
)
@ -5781,3 +5776,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)

View file

@ -265,8 +265,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")
)

View file

@ -329,6 +329,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 type(self) is AnthropicConfig
@classmethod
def get_config(cls, *, model: str | None = None):
config: Final = super().get_config()
@ -1656,7 +1660,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,
carry_block_binding=self._carries_thinking_block_binding,
)
elif param == "reasoning_effort":
# Accept both string ("low") and dict ({"effort": "low",

View file

@ -106,7 +106,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])
@ -155,13 +155,18 @@ 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
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)
@ -194,7 +199,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(
@ -227,7 +232,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)
@ -914,6 +919,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,
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).
@ -926,21 +932,21 @@ 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 = 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,
**(existing_output_config if isinstance(existing_output_config, dict) else MappingProxyType({})),
@ -1608,6 +1614,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
@ -1678,7 +1698,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
@ -1694,8 +1714,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`.
@ -1708,7 +1733,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:
@ -1726,11 +1751,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
@ -1742,15 +1767,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
@ -1884,7 +1910,7 @@ def _replayed_server_tool_use(block: object) -> _ReplayedServerToolUse | None:
def tool_call_is_rebuilt_as_server_tool_use(tool_call_id: object, provider_specific_fields: object) -> bool:
fields: Final = _validated_claude_code_mapping(provider_specific_fields)
fields: Final = validated_object_mapping(provider_specific_fields)
if not isinstance(tool_call_id, str) or fields is None:
return False
return (

View file

@ -4,11 +4,14 @@ Anthropic Token Counter implementation using the CountTokens API.
from typing import Any, Final
from pydantic import JsonValue
import litellm
from litellm._logging import verbose_logger
from litellm.exceptions import AuthenticationError
from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler
from litellm.llms.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, object]] | 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.
@ -57,14 +62,24 @@ class AnthropicTokenCounter(BaseTokenCounter):
try:
auth_header: Final = await AnthropicModelInfo.aget_auth_header(
api_key=litellm_params.get("api_key"),
api_key=litellm_params.get("api_key") or litellm.anthropic_key or litellm.api_key,
api_base=api_base,
litellm_params=litellm_params,
allow_workload_identity=True,
)
if auth_header is None:
verbose_logger.warning("No Anthropic credential 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,
)
result: Final = await anthropic_count_tokens_handler.handle_count_tokens_request(
model=model_to_use,

View file

@ -461,7 +461,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(

View file

@ -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)

View file

@ -80,6 +80,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 type(self) is AnthropicMessagesConfig
def get_supported_anthropic_messages_params(self, model: str) -> list:
return [
"messages",
@ -607,6 +611,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
model=model,
optional_params=anthropic_messages_optional_request_params,
custom_llm_provider=self._resolved_provider,
carry_block_binding=self._carries_thinking_block_binding,
)
self._translate_adaptive_effort_for_non_adaptive_model(

View file

@ -160,6 +160,13 @@ def prepare_native_messages(
)
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."""
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]:
"""

View file

@ -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")

View file

@ -29,7 +29,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, object] | None = None,
timeout: float | httpx.Timeout | None = None,
@ -79,6 +79,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)

View file

@ -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(

View file

@ -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

View file

@ -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

View file

@ -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

View file

@ -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
@ -179,16 +179,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:
"""

View file

@ -9,11 +9,11 @@ from collections.abc import Mapping
from typing import Any, Final
from openai.lib import _parsing, _pydantic # pyright: ignore[reportPrivateUsage] # SDK parser internals
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):
@ -27,6 +27,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

View file

@ -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,
)
@ -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 type(self) is AmazonAnthropicClaudeConfig
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",
carry_block_binding=self._carries_thinking_block_binding,
)
translated_thinking: Final = optional_params.get("thinking")
if (
@ -123,8 +130,11 @@ 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"] = {
**(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")

View file

@ -16,7 +16,7 @@ from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler
from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler
from litellm.llms.bedrock_mantle.common_utils import is_mantle_claude_model
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.types.utils import LiteLLMPydanticObjectBase, LlmProviders, TokenCountResponse
from litellm.types.utils import LiteLLMPydanticObjectBase, LlmProviders, TokenCountMessagesFormat, TokenCountResponse
RUNTIME_TOKENIZER_TYPE: Final = "bedrock_api"
MANTLE_TOKENIZER_TYPE: Final = "bedrock_mantle_api"
@ -106,6 +106,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:
if not messages:
return None

View file

@ -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("&", "&amp;").replace('"', "&quot;"), '"') 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)

View file

@ -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
from litellm.llms.anthropic.common_utils import (
AnthropicModelInfo,
thinking_block_binding,
validated_object_mapping,
)
from litellm.llms.anthropic.pass_through.messages.transformation import (
AnthropicMessagesConfig,
)
@ -49,7 +53,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 +85,10 @@ class AmazonAnthropicClaudeMessagesConfig(
def custom_llm_provider(self) -> str | None:
return "bedrock"
@property
def _carries_thinking_block_binding(self) -> bool:
return type(self) is AmazonAnthropicClaudeMessagesConfig
@property
def beta_headers_provider(self) -> str:
return self.custom_llm_provider or "bedrock"
@ -349,18 +356,14 @@ 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"))
display: Final = thinking.get("display") if thinking is not None else None
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"] = {
**(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"),
@ -701,8 +704,11 @@ 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"] = {
**(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 ###

View file

@ -2091,7 +2091,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,
@ -2118,7 +2118,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:

View file

@ -6,7 +6,7 @@ from collections.abc import Mapping, Sequence
from typing import Any, Final
import httpx
from pydantic import TypeAdapter
from pydantic import JsonValue, TypeAdapter
import litellm
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
@ -14,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,
@ -495,21 +495,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(

View file

@ -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()

View file

@ -5,8 +5,8 @@ Uses httpx for HTTP requests to OpenAI's /v1/responses/input_tokens endpoint.
"""
import json
from collections.abc import Sequence
from typing import Any, Final
from collections.abc import Mapping, Sequence
from typing import Final
import httpx
@ -31,9 +31,10 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig):
api_key: str,
api_base: str | None = None,
timeout: float | httpx.Timeout | None = None,
tools: list[dict[str, object]] | 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.
@ -53,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)

View file

@ -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(

View file

@ -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))

View file

@ -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
GEMINI_VIDEO_METADATA_KEYS: Final = MappingProxyType(
@ -1271,6 +1271,8 @@ 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
@ -1342,30 +1344,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,
)

View file

@ -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)

View file

@ -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

View file

@ -6,7 +6,7 @@ Why separate file? Make it easy to see how transformation works
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
@ -1311,6 +1311,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],
@ -1353,7 +1357,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):

View file

@ -20077,7 +20077,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,
@ -21068,7 +21074,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,
@ -77877,7 +77889,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": {
"supports_regex_lookaround": false,
@ -77907,7 +77921,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": {
"supports_regex_lookaround": false,
@ -77937,7 +77953,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,

View file

@ -2450,6 +2450,20 @@ async def fire_mcp_tool_call_failure_logging(
)
def _single_scoped_server(
mcp_servers: Sequence[str] | None, allowed_mcp_servers: Sequence[MCPServer]
) -> MCPServer | None:
"""The one server a ``/{server_name}/mcp`` path or ``x-mcp-servers`` header scoped this call to.
A bare tool name on such a request resolves against that server, the way a REST ``server_id``
does, instead of against the per-process tool-name map, which stays empty for a server whose
upstream only lists tools with the caller's own credentials until this process has listed it.
"""
if mcp_servers is None or len(mcp_servers) != 1 or len(allowed_mcp_servers) != 1:
return None
return allowed_mcp_servers[0]
@client
@catalog_operation(lambda: global_mcp_server_manager)
async def call_mcp_tool(
@ -2511,6 +2525,7 @@ async def call_mcp_tool(
detail="User not allowed to call this tool.",
)
scoped_server: Final = _single_scoped_server(mcp_servers, allowed_mcp_servers)
# Delegate to execute_mcp_tool for execution
response = await execute_mcp_tool(
name=name,
@ -2527,6 +2542,7 @@ async def call_mcp_tool(
input_responses=input_responses,
request_state=request_state,
incoming_bearer_token=incoming_bearer_token,
requested_server_id=scoped_server.server_id if scoped_server is not None else None,
**kwargs,
)
except Exception as e:

View file

@ -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
@ -41,7 +44,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()
@ -289,6 +292,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"],
@ -318,7 +329,7 @@ async def count_tokens(
Returns: {"input_tokens": <number>}
"""
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:
@ -330,10 +341,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
@ -345,35 +356,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(

View file

@ -4597,6 +4597,24 @@ def check_model_access_helper(
_check_model_access_helper: Final = check_model_access_helper
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,
@ -4668,6 +4686,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:
@ -4679,6 +4704,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}. "

View file

@ -637,6 +637,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")

View file

@ -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}

View file

@ -1,22 +1,142 @@
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, field_validator
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 ( # noqa: F401 # legacy module exports
_read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
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, ...]
@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)
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, ...] | None = ()
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 or ())
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)],
@ -168,36 +288,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 None:
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 {}
return TokenCountDetailsResponse(
totalTokens=original_response.get("totalTokens") or token_response.total_tokens or 0,
promptTokensDetails=original_response.get("promptTokensDetails", []),
)
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
# ============================================================

View file

@ -31,6 +31,7 @@ from litellm.constants import (
SESSION_ID_GENERATED_METADATA_KEY,
SESSION_ID_OMITTED_METADATA_KEY,
TYPESAFE_DANGER_ZONE_HEADER_PREFIX,
VECTOR_STORE_CALLER_METADATA_KWARG,
X_LITELLM_DISABLE_CALLBACKS,
)
from litellm.integrations.custom_guardrail import without_server_streaming_classification
@ -2099,7 +2100,7 @@ def refresh_proxy_server_request_body_snapshot(
if not isinstance(proxy_server_request, dict):
return
_body_snapshot_exclude: Final = (
frozenset({"secret_fields", "proxy_server_request", "litellm_logging_obj"})
frozenset({"secret_fields", "proxy_server_request", "litellm_logging_obj", VECTOR_STORE_CALLER_METADATA_KWARG})
| _TRANSPORT_ONLY_CREDENTIAL_KEYS
| _CALLBACK_CREDENTIAL_KEYS
)

View file

@ -47,6 +47,7 @@ from typing import (
)
import anyio
import anyio.lowlevel
import websockets
import websockets.exceptions
from pydantic import BaseModel, ConfigDict, Json, JsonValue, TypeAdapter, ValidationError, with_config
@ -87,6 +88,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,
@ -136,6 +144,15 @@ from litellm.proxy.common_utils.realtime_utils import ( # noqa: F401, RUF100 #
_realtime_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
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 (
@ -165,6 +182,7 @@ from litellm.types.utils import (
ModelResponseStream,
StreamingChoices,
TextCompletionResponse,
TokenCountMessagesFormat,
TokenCountResponse,
echoed_cost_map_pricing_fields,
is_server_derived_pricing_key,
@ -289,6 +307,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,
@ -318,6 +337,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 (
as_str_mapping,
drop_params_flag,
get_litellm_metadata_from_kwargs,
get_parent_otel_span_from_kwargs,
@ -960,7 +980,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
@ -14206,52 +14232,362 @@ 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 _system_message(system: object) -> ChatCompletionSystemMessage | None:
if not isinstance(system, (str, list)) or not system:
return None
@ -14274,8 +14610,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
@ -14311,8 +14654,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
@ -14323,20 +14674,33 @@ 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,
conversation_empty: Final = not messages and not contents
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,
messages=messages,
contents=contents,
deployment=deployment,
request_model=request.model,
tools=tools,
system=system,
settings=settings,
local_estimate_allowed=litellm.disable_token_counter is not True,
)
if result is not None:
return result
)
if isinstance(ladder, TokenCountResponse):
return ladder
# Check if token counter is disabled before fallback
if litellm.disable_token_counter is True:
@ -14354,14 +14718,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)
@ -14369,18 +14730,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",
)

View file

@ -36,6 +36,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg
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.base import LiteLLMBaseModel
from litellm.types.llms.openai import (
@ -44,6 +45,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:
@ -1067,7 +1069,7 @@ class _ResponsesApiErrorBody(TypedDict):
error: ReadOnly[_ResponsesApiErrorDetail]
class _ResponsesInputTokensResult(TypedDict):
class _ResponsesInputTokensResult(TokenCountSourceFields):
object: ReadOnly[str]
input_tokens: ReadOnly[int]
@ -1198,6 +1200,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

View file

@ -9,12 +9,14 @@ from typing import (
from fastapi import APIRouter, Depends, HTTPException, Request, Response
from litellm.constants import VECTOR_STORE_CALLER_METADATA_KWARG
from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import (
LiteLLM_ManagedVectorStore,
)
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.common_utils.http_parsing_utils import read_raw_json_body
from litellm.proxy.utils import jsonify_object
from litellm.proxy.vector_store_endpoints.utils import (
assert_proxy_admin_for_vector_store_index_management,
@ -23,6 +25,7 @@ from litellm.proxy.vector_store_endpoints.utils import (
)
from litellm.repositories.table_repositories import ManagedVectorStoreIndexRepository
from litellm.types.vector_stores import IndexCreateRequest, IndexListResponse
from litellm.vector_stores.utils import VectorStoreWriteBody, callers_metadata_snapshot
from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry
router: Final = APIRouter()
@ -213,6 +216,9 @@ async def vector_store_create(
)
data: Final = await read_request_body(request=request)
data[VECTOR_STORE_CALLER_METADATA_KWARG] = callers_metadata_snapshot(
await read_raw_json_body(request), VectorStoreWriteBody.model_validate(data)
)
# Check for target_model_names parameter
target_model_names: Final = data.pop("target_model_names", None)
@ -448,6 +454,9 @@ async def vector_store_update(
data = await read_request_body(request=request) # rebind-ok: pre-existing rebinding on a rename-only line
if "vector_store_id" not in data:
data["vector_store_id"] = vector_store_id
data[VECTOR_STORE_CALLER_METADATA_KWARG] = callers_metadata_snapshot(
await read_raw_json_body(request), VectorStoreWriteBody.model_validate(data)
)
data = await _update_request_data_with_litellm_managed_vector_store_registry(
data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict

View file

@ -732,10 +732,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):

View file

@ -8,6 +8,7 @@ from typing_extensions import (
)
from litellm.types.llms.openai import ChatCompletionFileObject, 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
@ -319,12 +320,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

View file

@ -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]

View file

@ -9,6 +9,7 @@ from typing import (
Any,
Final,
Literal,
TypeAlias,
get_args,
)
@ -56,6 +57,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.types.workload_identity import (
ANTHROPIC_WIF_KWARGS_KEYS,
OPENAI_WIF_KWARGS_KEYS,
@ -227,6 +229,9 @@ class ProviderSpecificModelInfo(TypedDict, total=False):
bedrock_output_config_effort_ceiling: Literal["low", "medium", "high", "max", "xhigh"] | None
bedrock_converse_supports_strict_tools: bool | None
supports_regex_lookaround: ReadOnly[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):
@ -3958,11 +3963,14 @@ DEPLOYMENT_SCOPED_PRICING_FIELDS: Final[frozenset[str]] = frozenset({"off_peak_p
DEPLOYMENT_SCOPED_CAPABILITY_FIELDS: Final[frozenset[str]] = frozenset({"supports_regex_lookaround"})
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_CAPABILITY_FIELDS
- DEPLOYMENT_SCOPED_TOKEN_COUNT_FIELDS
)
@ -4364,11 +4372,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
@ -4379,6 +4391,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):

View file

@ -6429,6 +6429,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():

View file

@ -1,11 +1,92 @@
import copy
from collections.abc import Container, Mapping
from typing import Any, Final, cast, get_type_hints
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from litellm.constants import ROUTING_REQUEST_TAGS_METADATA_KEY, VECTOR_STORE_CALLER_METADATA_KWARG
from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
from litellm.types.vector_stores import (
VectorStoreCreateOptionalRequestParams,
VectorStoreSearchOptionalRequestParams,
)
LITELLM_WRITTEN_METADATA_KEYS: Final = frozenset(
{
"requester_ip_address",
"user_agent",
"endpoint",
"litellm_api_version",
"deployment",
"deployment_model_name",
"api_base",
"agent_id",
"actor_agent_id",
"target_agent_id",
"billing_agent_id",
"agent_execution_mode",
"verified_human_user_id",
"litellm_caller_session_id",
}
)
def is_litellm_written_metadata_key(key: str) -> bool:
return key == "user_api_key" or key.startswith("user_api_key_") or key in LITELLM_WRITTEN_METADATA_KEYS
def _is_callers_key(key: object) -> bool:
return not isinstance(key, str) or not is_litellm_written_metadata_key(key)
def _sent_by_proxy_or_router(params: Mapping[str, object], metadata: Container[object]) -> bool:
return params.get("proxy_server_request") is not None or ROUTING_REQUEST_TAGS_METADATA_KEY in metadata
def _callers_metadata(metadata: Mapping[object, object]) -> Mapping[object, object]:
return {key: value for key, value in metadata.items() if _is_callers_key(key)}
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object])
class VectorStoreWriteBody(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
metadata: object = None
def _json_object(raw: str | bytes) -> Mapping[str, object] | None:
try:
return _JSON_OBJECT.validate_json(raw)
except ValidationError:
return None
def _decoded_metadata(metadata: object) -> object:
if not isinstance(metadata, str):
return metadata
decoded: Final = _json_object(metadata)
return metadata if decoded is None else decoded
def callers_metadata_snapshot(raw_json_body: bytes | None, parsed_body: VectorStoreWriteBody) -> object:
raw_body: Final = None if raw_json_body is None else _json_object(raw_json_body)
metadata: Final = copy.deepcopy(parsed_body.metadata) if raw_body is None else raw_body.get("metadata")
return _decoded_metadata(metadata)
def _metadata_for_provider(params: Mapping[str, object]) -> object:
if VECTOR_STORE_CALLER_METADATA_KWARG in params:
return params[VECTOR_STORE_CALLER_METADATA_KWARG]
metadata: Final = params.get("metadata")
if not isinstance(metadata, Mapping):
return metadata
request_metadata: Final = cast(Mapping[object, object], metadata) # cast-ok: request metadata is a JSON object
if not _sent_by_proxy_or_router(params, request_metadata):
return request_metadata
return _callers_metadata(request_metadata)
class VectorStoreRequestUtils:
"""Helper utils for constructing Vector Store search requests"""
@ -48,7 +129,11 @@ class VectorStoreRequestUtils:
Returns:
VectorStoreCreateOptionalRequestParams instance with only the valid parameters
"""
valid_keys: Final = get_type_hints(VectorStoreCreateOptionalRequestParams).keys()
valid_keys: Final = get_type_hints(VectorStoreCreateOptionalRequestParams).keys() - {"metadata"}
filtered_params: Final = {k: v for k, v in params.items() if k in valid_keys and v is not None}
return cast(VectorStoreCreateOptionalRequestParams, filtered_params)
metadata: Final = _metadata_for_provider(params)
if metadata is None:
return cast(VectorStoreCreateOptionalRequestParams, filtered_params)
return cast( # cast-ok: same params as above plus the metadata the provider should store
VectorStoreCreateOptionalRequestParams, {**filtered_params, "metadata": metadata}
)

View file

@ -20077,7 +20077,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,
@ -21068,7 +21074,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,
@ -77877,7 +77889,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": {
"supports_regex_lookaround": false,
@ -77907,7 +77921,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": {
"supports_regex_lookaround": false,
@ -77937,7 +77953,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,

View file

@ -1219,6 +1219,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,

View file

@ -1,6 +1,6 @@
[project]
name = "litellm"
version = "1.106.0"
version = "1.106.1"
description = "Library to easily interface with LLM API providers"
readme = "README.md"
requires-python = ">=3.10, <3.15"
@ -373,7 +373,7 @@ litellm-enterprise = { workspace = true }
members = ["enterprise", "litellm-proxy-extras"]
[tool.commitizen]
version = "1.106.0"
version = "1.106.1"
version_files = [
"pyproject.toml:^version",
]

View file

@ -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"}
}

View file

@ -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."

View file

@ -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")]

View file

@ -0,0 +1,630 @@
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"<filler>{_UNIT * 100}</filler>\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"<cache-buster>{uuid4().hex}</cache-buster>\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, **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:
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
)
@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
]

View file

@ -24,6 +24,8 @@ model_list:
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: anthropic/claude-sonnet-4-6
litellm_params:
model: anthropic/claude-sonnet-4-6

View file

@ -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")

View file

@ -5019,4 +5019,62 @@ def test_transform_request_drop_params_in_litellm_params_gates_the_prompt_cache_
litellm_logging_obj=Mock(),
)
assert "prompt_cache_breakpoint" not in result["input"][0]["content"][0]
assert "prompt_cache_breakpoint" not in result["input"][0]["content"][0]
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,
}
]

View file

View file

@ -0,0 +1,198 @@
import importlib.util
import json
import sys
from collections.abc import Mapping
from pathlib import Path
from typing import Final
import httpx
import pytest
_SCRIPT: Final = Path(__file__).resolve().parents[3] / "db_scripts" / "strip_vector_store_internal_metadata.py"
_spec: Final = importlib.util.spec_from_file_location("strip_vector_store_internal_metadata", _SCRIPT)
strip_script: Final = importlib.util.module_from_spec(_spec)
sys.modules[_spec.name] = strip_script
_spec.loader.exec_module(strip_script)
POLLUTED: Final = {
"id": "vs_1",
"metadata": {
"user_api_key": "hash-value-1",
"user_api_key_hash": "hash-value-1",
"user_api_key_user_email": "member@example.com",
"endpoint": "http://proxy/v1/vector_stores",
"deployment": "openai/gpt-5",
"team": "llmproxy",
"trace_id": "trace-123",
},
}
CLEAN: Final = {"id": "vs_2", "metadata": {"team": "llmproxy"}}
NO_METADATA: Final = {"id": "vs_3", "metadata": None}
def _fake_openai(
stores: tuple[Mapping[str, object], ...], page_size: int = 100, post_status: int = 200
) -> tuple[httpx.MockTransport, list[httpx.Request]]:
seen: Final[list[httpx.Request]] = []
by_id: Final = {store["id"]: store for store in stores}
def handle(request: httpx.Request) -> httpx.Response:
seen.append(request)
path: Final = request.url.path.removeprefix("/v1")
if request.method == "POST":
if post_status != 200:
return httpx.Response(post_status, json={"error": {"message": "upstream rejected the update"}})
return httpx.Response(200, json=by_id[path.removeprefix("/vector_stores/")])
if path == "/vector_stores":
after: Final = request.url.params.get("after")
start: Final = 0 if after is None else [store["id"] for store in stores].index(after) + 1
page: Final = stores[start : start + page_size]
return httpx.Response(200, json={"data": list(page), "has_more": start + page_size < len(stores)})
return httpx.Response(200, json=by_id[path.removeprefix("/vector_stores/")])
return httpx.MockTransport(handle), seen
def _run(args: list[str], transport: httpx.MockTransport) -> int:
return strip_script.main(["--api-key", "sk-test", *args], transport=transport)
def _posts(seen: list[httpx.Request]) -> list[tuple[str, dict[str, object]]]:
return [(request.url.path, json.loads(request.content)) for request in seen if request.method == "POST"]
def test_dry_run_lists_matching_key_names_without_values_and_sends_no_update(capsys):
transport, seen = _fake_openai((POLLUTED, CLEAN))
exit_code: Final = _run([], transport)
out: Final = capsys.readouterr().out
assert exit_code == 0
assert (
"vs_1: would delete 5 keys: deployment, endpoint, user_api_key, user_api_key_hash, user_api_key_user_email"
in out
)
assert "scanned=2 matched=1 updated=0 failed=0" in out
assert "hash-value-1" not in out
assert "member@example.com" not in out
assert _posts(seen) == []
assert {request.headers["authorization"] for request in seen} == {"Bearer sk-test"}
def test_apply_nulls_exactly_the_proxy_written_keys_and_leaves_caller_keys_alone(capsys):
transport, seen = _fake_openai((POLLUTED, CLEAN, NO_METADATA))
exit_code: Final = _run(["--apply"], transport)
assert exit_code == 0
assert _posts(seen) == [
(
"/v1/vector_stores/vs_1",
{
"metadata": {
"deployment": None,
"endpoint": None,
"user_api_key": None,
"user_api_key_hash": None,
"user_api_key_user_email": None,
}
},
)
]
assert "scanned=3 matched=1 updated=1 failed=0" in capsys.readouterr().out
def test_extra_key_is_deleted_alongside_the_built_in_keys():
transport, seen = _fake_openai((POLLUTED,))
_run(["--apply", "--extra-key", "trace_id"], transport)
assert _posts(seen)[0][1]["metadata"]["trace_id"] is None
assert "team" not in _posts(seen)[0][1]["metadata"]
def test_listing_follows_has_more_with_the_last_store_id(capsys):
transport, seen = _fake_openai((CLEAN, POLLUTED, NO_METADATA), page_size=2)
_run([], transport)
assert [request.url.params.get("after") for request in seen] == [None, "vs_1"]
assert "scanned=3 matched=1" in capsys.readouterr().out
def test_listing_walks_more_pages_than_the_interpreter_recursion_limit(capsys):
stores: Final = tuple({"id": f"vs_{index}", "metadata": None} for index in range(sys.getrecursionlimit() + 50))
transport, _ = _fake_openai(stores, page_size=1)
assert _run([], transport) == 0
assert f"scanned={len(stores)} matched=0" in capsys.readouterr().out
def test_non_string_metadata_values_are_scanned_without_printing_them(capsys):
transport, _ = _fake_openai(({"id": "vs_9", "metadata": {"user_api_key": "secret-value", "count": 7}},))
exit_code: Final = _run([], transport)
out: Final = capsys.readouterr().out
assert exit_code == 0
assert "vs_9: would delete 1 keys: user_api_key" in out
assert "secret-value" not in out
def test_an_unexpected_response_shape_exits_non_zero_without_echoing_metadata(capsys):
def malformed(request: httpx.Request) -> httpx.Response:
return httpx.Response(200, json={"data": [{"metadata": {"user_api_key": "secret-value"}}], "has_more": False})
exit_code: Final = _run([], httpx.MockTransport(malformed))
captured: Final = capsys.readouterr()
assert exit_code == 2
assert "unexpected vector store response shape" in captured.err
assert "secret-value" not in captured.out + captured.err
def test_vector_store_id_retrieves_only_the_named_stores_without_listing(capsys):
transport, seen = _fake_openai((POLLUTED, CLEAN, NO_METADATA))
_run(["--vector-store-id", "vs_2", "--vector-store-id", "vs_1"], transport)
assert [request.url.path for request in seen] == ["/v1/vector_stores/vs_2", "/v1/vector_stores/vs_1"]
assert "scanned=2 matched=1" in capsys.readouterr().out
def test_a_failed_update_is_reported_and_exits_non_zero(capsys):
transport, _ = _fake_openai((POLLUTED,), post_status=400)
exit_code: Final = _run(["--apply"], transport)
out: Final = capsys.readouterr().out
assert exit_code == 1
assert "vs_1: failed to delete 5 keys (HTTP 400: upstream rejected the update)" in out
assert "updated=0 failed=1" in out
def test_a_listing_error_exits_non_zero_without_updating_anything():
seen: Final[list[httpx.Request]] = []
def reject(request: httpx.Request) -> httpx.Response:
seen.append(request)
return httpx.Response(401, json={"error": {"message": "bad key"}})
assert _run(["--apply"], httpx.MockTransport(reject)) == 2
assert [request.method for request in seen] == ["GET"]
def test_an_unreachable_api_exits_non_zero():
def unreachable(request: httpx.Request) -> httpx.Response:
raise httpx.ConnectError("connection refused", request=request)
assert _run([], httpx.MockTransport(unreachable)) == 2
def test_an_api_key_is_required(monkeypatch):
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
with pytest.raises(SystemExit) as exit_info:
strip_script.main([])
assert exit_info.value.code == 2

View file

@ -32,7 +32,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
sanitize_messages_for_tool_calling,
THOUGHT_SIGNATURE_SEPARATOR,
)
from litellm.types.llms.openai import ChatCompletionToolMessage
from litellm.types.llms.openai import ChatCompletionAssistantMessage, ChatCompletionToolMessage
from litellm.utils import (
_invalidate_model_cost_lowercase_map,
function_setup,
@ -4130,44 +4130,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():
@ -4281,29 +4279,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})

View file

@ -337,6 +337,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",
@ -349,6 +350,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",
],

View file

@ -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")

View file

@ -14,8 +14,9 @@ import json
import os
import sys
import threading
from collections.abc import Mapping
from pathlib import Path
from types import SimpleNamespace
from types import MappingProxyType, SimpleNamespace
from typing import Final
from unittest.mock import patch
@ -24,6 +25,9 @@ import pytest
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")))
from litellm.llms.base_llm.anthropic_messages.transformation import ( # noqa: E402 # sys.path must be patched before importing litellm
BaseAnthropicMessagesConfig,
)
from litellm.proxy._types import SpecialHeaders # noqa: E402 # sys.path must be patched before importing litellm
# Fake tokens for testing (not real secrets)
@ -1332,6 +1336,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."""
@ -1422,11 +1432,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,
)
@ -1436,7 +1446,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",
@ -1455,6 +1465,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,
@ -1967,6 +2007,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'`` +
@ -4092,6 +4175,117 @@ def test_shared_legacy_thinking_translation_preserves_supported_display(
assert optional_params["thinking"] == expected_thinking
LEGACY_THINKING_WITH_BLOCK_BINDING: Final = MappingProxyType(
{
"type": "enabled",
"budget_tokens": 2048,
"block_binding": {"prefix_mismatch_behavior": "error"},
"future_key": 1,
}
)
ADAPTIVE_THINKING_WITH_BLOCK_BINDING: Final = MappingProxyType(
{"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": "error"}}
)
@pytest.mark.usefixtures("local_model_cost_map")
@pytest.mark.parametrize(
("model", "custom_llm_provider", "expected_thinking"),
(
("claude-opus-4-8", "anthropic", ADAPTIVE_THINKING_WITH_BLOCK_BINDING),
("claude-opus-4-8", "azure_ai", ADAPTIVE_THINKING_WITH_BLOCK_BINDING),
("invoke/us.anthropic.claude-opus-4-8", "bedrock", ADAPTIVE_THINKING_WITH_BLOCK_BINDING),
("converse/us.anthropic.claude-opus-4-8", "bedrock", {"type": "adaptive"}),
("mantle/anthropic.claude-opus-4-8", "bedrock", {"type": "adaptive"}),
("claude_platform/claude-opus-4-8", "bedrock", {"type": "adaptive"}),
("claude-opus-4-8", "vertex_ai", {"type": "adaptive"}),
("databricks-claude-opus-4-8", "databricks", {"type": "adaptive"}),
),
ids=(
"anthropic",
"azure_ai",
"bedrock_invoke",
"bedrock_converse",
"bedrock_mantle",
"bedrock_claude_platform",
"vertex_ai",
"databricks",
),
)
def test_chat_legacy_thinking_translation_carries_only_block_binding_where_supported(
model: str, custom_llm_provider: str, expected_thinking: Mapping[str, object]
) -> None:
from litellm.utils import get_optional_params
optional_params: Final = get_optional_params(
model=model,
custom_llm_provider=custom_llm_provider,
thinking=dict(LEGACY_THINKING_WITH_BLOCK_BINDING),
)
assert optional_params["thinking"] == expected_thinking
@pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config")
@pytest.mark.parametrize(
("model", "custom_llm_provider", "expected_thinking"),
(
("claude-opus-4-8", "anthropic", ADAPTIVE_THINKING_WITH_BLOCK_BINDING),
("claude-opus-4-8", "azure_ai", ADAPTIVE_THINKING_WITH_BLOCK_BINDING),
("us.anthropic.claude-opus-4-8", "bedrock", ADAPTIVE_THINKING_WITH_BLOCK_BINDING),
("mantle/anthropic.claude-opus-4-8", "bedrock", {"type": "adaptive"}),
("anthropic.claude-opus-4-8", "bedrock_mantle", {"type": "adaptive"}),
("claude_platform/claude-opus-4-8", "bedrock", {"type": "adaptive"}),
("claude-opus-4-8", "vertex_ai", {"type": "adaptive"}),
),
ids=(
"anthropic",
"azure_ai",
"bedrock_invoke",
"bedrock_mantle",
"bedrock_mantle_provider",
"bedrock_claude_platform",
"vertex_ai",
),
)
def test_messages_legacy_thinking_translation_carries_only_block_binding_where_supported(
model: str, custom_llm_provider: str, expected_thinking: Mapping[str, object]
) -> None:
from litellm.types.utils import LlmProviders
from litellm.utils import ProviderConfigManager
config: Final = ProviderConfigManager.get_provider_anthropic_messages_config(
model=model, provider=LlmProviders(custom_llm_provider)
)
assert config is not None
assert _thinking_sent_for_legacy_messages_request(config, model) == expected_thinking
@pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config")
def test_openai_compatible_messages_passthrough_keeps_legacy_thinking_without_block_binding() -> None:
from litellm.llms.openai_like.messages.transformation import OpenAILikeAnthropicMessagesConfig
assert _thinking_sent_for_legacy_messages_request(OpenAILikeAnthropicMessagesConfig(), "claude-opus-4-8") == {
"type": "adaptive"
}
def _thinking_sent_for_legacy_messages_request(config: BaseAnthropicMessagesConfig, model: str) -> object:
from litellm.types.router import GenericLiteLLMParams
return config.transform_anthropic_messages_request(
model=model,
messages=[{"role": "user", "content": "Reply with OK"}],
anthropic_messages_optional_request_params={
"max_tokens": 512,
"thinking": dict(LEGACY_THINKING_WITH_BLOCK_BINDING),
},
litellm_params=GenericLiteLLMParams(),
headers={},
).get("thinking")
@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))

View file

@ -352,3 +352,48 @@ def test_count_token_counters_select_their_own_provider(
assert counter.should_use_token_counting_api(custom_llm_provider=provider) is True
assert counter.should_use_token_counting_api(custom_llm_provider="unknown") is False
@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

View file

@ -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

View file

@ -747,7 +747,7 @@ def test_bedrock_chat_invoke_tool_based_response_format_still_upgrades_legacy_th
"schema": {"type": "object", "properties": {"result": {"type": "string"}}},
},
},
"thinking": {"type": "enabled", "budget_tokens": 4096},
"thinking": {"type": "enabled", "budget_tokens": 4096, "block_binding": BLOCK_BINDING},
"max_tokens": 8192,
},
optional_params={},
@ -756,10 +756,34 @@ def test_bedrock_chat_invoke_tool_based_response_format_still_upgrades_legacy_th
)
assert "tools" in result
assert result["thinking"] == {"type": "adaptive"}
assert result["thinking"] == {"type": "adaptive", "block_binding": BLOCK_BINDING}
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
@ -974,6 +998,71 @@ 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, "future_key": 1},
{
"type": "enabled",
"budget_tokens": 2048,
"display": "updates",
"block_binding": BLOCK_BINDING,
"future_key": 1,
},
),
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}
@pytest.mark.parametrize(
"config_path",
(
"litellm.llms.bedrock.chat.mantle.transformation.AmazonMantleConfig",
"litellm.llms.bedrock_mantle.chat.claude_transformation.BedrockMantleClaudeChatConfig",
),
ids=("bedrock_mantle_route", "bedrock_mantle_provider"),
)
def test_bedrock_chat_mantle_rebuilds_display_updates_without_block_binding(
local_model_cost_map, local_beta_headers_config, config_path: str
) -> None:
module_name, _, class_name = config_path.rpartition(".")
config: Final = getattr(importlib.import_module(module_name), class_name)()
optional_params: Final = config.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"]}

View file

@ -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&amp;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” & <tags> 中文汉字,日本語のテキスト 😀🎉\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)

View file

@ -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
@ -3529,6 +3531,94 @@ 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, "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,
"future_key": 1,
},
{"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(
"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))

View file

@ -1804,6 +1804,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.

View file

@ -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

View file

@ -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"

View file

@ -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."

View file

@ -630,3 +630,56 @@ class TestClientsideBaseOverride:
)
is 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

View file

@ -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)

View file

@ -1,9 +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,
convert_anyof_null_to_nullable,
get_vertex_location_from_url,
get_vertex_project_id_from_url,
@ -14,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"])
@ -147,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
@ -228,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": {
@ -278,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}]},
@ -685,9 +788,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
@ -699,9 +800,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
@ -819,9 +918,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"]
@ -988,7 +1085,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(
@ -1121,10 +1220,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]
@ -1178,8 +1274,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():
@ -1457,6 +1551,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():
"""
@ -1471,9 +1608,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",
@ -1515,9 +1650,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(
@ -1560,9 +1693,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")
@ -1593,9 +1724,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():
@ -1630,9 +1759,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():
@ -1643,9 +1770,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():
@ -1835,12 +1960,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():
@ -1850,9 +1971,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"}

View file

@ -2994,11 +2994,7 @@ async def test_call_mcp_tool_uses_manager_permission_lookup():
"get_mcp_server_by_id",
return_value=mock_server,
),
patch.object(
global_mcp_server_manager,
"get_mcp_server_from_tool_name",
return_value=mock_server,
) as mock_get_server,
patch.object(global_mcp_server_manager, "server_exposes_tool", return_value=True),
patch(
"litellm.proxy._experimental.mcp_server.operations.global_mcp_tool_registry"
) as mock_tool_registry,
@ -3023,24 +3019,18 @@ async def test_call_mcp_tool_uses_manager_permission_lookup():
assert result == expected_response
mock_get_allowed.assert_awaited_once()
# We call `_get_mcp_server_from_tool_name` multiple times:
# - for logging/metadata
# - for resolving the server during dispatch
# - and inside `call_tool` for guardrails/hooks
# The exact count isn't important, only that it is used.
assert mock_get_server.call_count >= 2
# First call should use the prefixed tool name
assert (
mock_get_server.call_args_list[0][0][0]
== f"{mock_server.name}/gmail_send_email"
)
# The request names exactly one server, so that server is the dispatch target
# and the tool name reaches it untouched (the "/" is not a LiteLLM prefix).
mock_handle_managed.assert_awaited_once()
assert mock_handle_managed.await_args.kwargs["server_name"] == mock_server.name
assert mock_handle_managed.await_args.kwargs["name"] == f"{mock_server.name}/gmail_send_email"
@pytest.mark.asyncio
async def test_call_mcp_tool_resolves_unprefixed_tool_name_and_checks_permissions():
"""
Ensure `call_mcp_tool` correctly resolves the MCP server for an unprefixed tool
name and enforces server-level permissions using that resolved server.
Ensure `call_mcp_tool` resolves an unprefixed tool name on a request scoped to
one server against that server and enforces server-level permissions with it.
"""
from litellm.proxy._experimental.mcp_server.server import (
call_mcp_tool,
@ -3070,11 +3060,7 @@ async def test_call_mcp_tool_resolves_unprefixed_tool_name_and_checks_permission
"get_mcp_server_by_id",
return_value=mock_server,
),
patch.object(
global_mcp_server_manager,
"get_mcp_server_from_tool_name",
return_value=mock_server,
) as mock_get_server,
patch.object(global_mcp_server_manager, "server_exposes_tool", return_value=True),
patch(
"litellm.proxy._experimental.mcp_server.operations.global_mcp_tool_registry"
) as mock_tool_registry,
@ -3091,7 +3077,7 @@ async def test_call_mcp_tool_resolves_unprefixed_tool_name_and_checks_permission
mock_tool_registry.get_tool.return_value = None
mock_handle_managed.return_value = expected_response
# Call with UNPREFIXED tool name; server should be resolved via mapping
# Call with UNPREFIXED tool name; the single server the request names is the target
result = await call_mcp_tool(
name="gmail_send_email",
arguments={"body": "hello"},
@ -3100,8 +3086,9 @@ async def test_call_mcp_tool_resolves_unprefixed_tool_name_and_checks_permission
assert result == expected_response
mock_get_allowed.assert_awaited_once()
# We should resolve the server at least once using the unprefixed name
assert mock_get_server.call_count >= 1
assert mock_get_server.call_args_list[0][0][0] == "gmail_send_email"
mock_handle_managed.assert_awaited_once()
assert mock_handle_managed.await_args.kwargs["server_name"] == mock_server.name
assert mock_handle_managed.await_args.kwargs["name"] == "gmail_send_email"
# Permissions check should be invoked with the resolved server name
mock_is_allowed.assert_called_once()
assert mock_is_allowed.call_args.kwargs["server_name"] == mock_server.name

View file

@ -565,6 +565,88 @@ async def test_tools_call_warmup_does_not_consume_mcp_server_rpm() -> None:
assert [tool.name for tool in listing.tools] == ["echo"]
def _scoped_rig(
monkeypatch: pytest.MonkeyPatch,
) -> tuple[MCPServer, MCPServer, AsyncMock, AsyncMock]:
"""Two config servers that both expose a bare ``list_rules`` on a manager that has listed neither."""
from mcp.types import CallToolResult, TextContent
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager
manager: Final = MCPServerManager()
memory: Final = MCPServer(
server_id="id-memory", name="memory_server", server_name="memory_server", transport=MCPTransport.http
)
other: Final = MCPServer(
server_id="id-other", name="other_server", server_name="other_server", transport=MCPTransport.http
)
manager.config_mcp_servers = {server.server_id: server for server in (memory, other)}
monkeypatch.setattr(proxy_server, "prisma_client", None)
monkeypatch.setattr(operations, "global_mcp_server_manager", manager)
monkeypatch.setattr(
manager, "get_allowed_mcp_servers", AsyncMock(return_value=[memory.server_id, other.server_id])
)
monkeypatch.setattr(global_mcp_tool_registry, "published_tools", {})
listing: Final = AsyncMock()
monkeypatch.setattr(operations, "_get_tools_from_mcp_servers", listing)
handled: Final = AsyncMock(
return_value=CallToolResult(content=[TextContent(type="text", text="ok")], isError=False)
)
monkeypatch.setattr(operations, "_handle_managed_mcp_tool", handled)
return memory, other, listing, handled
@pytest.mark.asyncio
@pytest.mark.parametrize("scope", ["memory_server", "other_server"])
async def test_bare_tool_name_on_a_single_server_scope_routes_to_that_server_cold(
monkeypatch: pytest.MonkeyPatch, scope: str
) -> None:
"""``/{server_name}/mcp`` (or ``x-mcp-servers`` naming one server) with the upstream's bare tool name
must list and call the scoped server on a process whose tool-name map has never seen it, and when two
servers share the bare name the scope decides (BerriAI/litellm#44831)."""
memory, other, listing, handled = _scoped_rig(monkeypatch)
expected: Final = memory if scope == "memory_server" else other
async with operations.global_mcp_server_manager.catalog.operation():
result: Final = await operations.call_mcp_tool(
name="list_rules",
arguments={},
user_api_key_auth=UserAPIKeyAuth(api_key="sk-scoped"),
mcp_servers=[scope],
)
assert result.is_error is False
assert listing.await_args.kwargs["mcp_servers"] == [expected.server_id]
assert handled.await_args.kwargs["server_name"] == expected.name
assert handled.await_args.kwargs["name"] == "list_rules"
@pytest.mark.asyncio
@pytest.mark.parametrize("mcp_servers", [None, ["memory_server", "other_server"]])
async def test_bare_tool_name_without_a_single_server_scope_is_still_unresolved_cold(
monkeypatch: pytest.MonkeyPatch, mcp_servers: list[str] | None
) -> None:
"""The scope shortcut only applies when the request names exactly one server; an ambiguous bare name
keeps today's answer rather than being guessed onto the first allowed server."""
from fastapi import HTTPException
_memory, _other, listing, handled = _scoped_rig(monkeypatch)
async with operations.global_mcp_server_manager.catalog.operation():
with pytest.raises(HTTPException) as unresolved:
await operations.call_mcp_tool(
name="list_rules",
arguments={},
user_api_key_auth=UserAPIKeyAuth(api_key="sk-scoped"),
mcp_servers=mcp_servers,
)
assert unresolved.value.status_code == 404
listing.assert_not_awaited()
handled.assert_not_awaited()
@pytest.mark.asyncio
async def test_tools_call_pre_call_check_enforces_mcp_server_rpm() -> None:
from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager

View file

@ -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()

View file

@ -1354,6 +1354,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

View file

@ -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

View file

@ -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

View file

@ -4,10 +4,11 @@ Test to verify the Google GenAI proxy API endpoints
"""
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
from unittest.mock import AsyncMock, patch
import pytest
from litellm.types.utils import TokenCountResponse
def _build_test_client():
@ -88,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():
@ -123,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"]
@ -160,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",
@ -180,6 +182,30 @@ 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():
@ -189,9 +215,14 @@ def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs()
except ImportError as e:
pytest.skip(f"Skipping test due to missing dependency: {e}")
fake_response: Final = MagicMock()
fake_response.original_response = {"input_tokens": 2167}
fake_response.total_tokens = 2167
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",
@ -207,3 +238,5 @@ def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs()
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"

View file

@ -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}},
)

View file

@ -4147,6 +4147,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

View file

@ -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
)

File diff suppressed because it is too large Load diff

View file

@ -6,6 +6,7 @@ import pytest
from fastapi import HTTPException, Request
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import (
LiteLLM_ManagedVectorStore,
)
@ -3256,3 +3257,159 @@ class TestConfigOwnedVectorStores:
assert response["status"] == "success", response
prisma.db.litellm_managedvectorstorestable.delete.assert_awaited_once_with(where={"vector_store_id": self.DB_ID})
assert registry.get_litellm_managed_vector_store_from_registry(self.DB_ID) is None
CALLERS_METADATA = {"team": "llmproxy", "endpoint": "ingest", "trace_id": "callers-trace"}
class _PreCallBodies(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.logged_bodies = []
def log_pre_api_call(self, model, messages, kwargs):
self.logged_bodies.append(kwargs["litellm_params"]["proxy_server_request"]["body"])
@pytest.mark.parametrize("operation", ["create", "update"])
@pytest.mark.parametrize("model", [None, "gpt-5"])
@pytest.mark.parametrize(
("metadata_field", "expected_metadata"),
[
({"metadata": CALLERS_METADATA}, CALLERS_METADATA),
({"metadata": {}}, {}),
({"metadata": json.dumps(CALLERS_METADATA)}, CALLERS_METADATA),
({"metadata": None}, None),
({}, None),
],
)
def test_proxy_vector_store_writes_send_the_provider_exactly_the_callers_metadata(
monkeypatch, operation, model, metadata_field, expected_metadata
):
import httpx
import respx
from fastapi.testclient import TestClient
from litellm.constants import VECTOR_STORE_CALLER_METADATA_KWARG
from litellm.proxy import proxy_server
from litellm.proxy.proxy_server import app
monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
litellm.in_memory_llm_clients_cache.flush_cache()
monkeypatch.setattr(proxy_server, "master_key", "sk-master-test")
monkeypatch.setattr(
proxy_server,
"llm_router",
litellm.Router(
model_list=[{"model_name": "gpt-5", "litellm_params": {"model": "openai/gpt-5", "api_key": "sk-test"}}]
),
)
pre_call_bodies = _PreCallBodies()
monkeypatch.setattr(litellm, "callbacks", [pre_call_bodies])
sent_bodies = []
def echo(request: httpx.Request) -> httpx.Response:
body = json.loads(request.content)
sent_bodies.append(body)
return httpx.Response(
200,
json={
"id": "vs_1",
"object": "vector_store",
"created_at": 0,
"name": "x",
"usage_bytes": 0,
"file_counts": {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0},
"status": "completed",
"metadata": body.get("metadata"),
},
)
path = "/v1/vector_stores" if operation == "create" else "/v1/vector_stores/vs_1"
routing = {} if model is None else {"model": model}
with respx.mock(assert_all_called=False) as upstream:
upstream.post(url__regex=r"https://api\.openai\.com/v1/vector_stores.*").mock(side_effect=echo)
response = TestClient(app).post(
path,
json={"name": "x", **routing, **metadata_field},
headers={
"Authorization": "Bearer sk-master-test",
"x-litellm-tags": "header-tag",
"x-litellm-trace-id": "header-trace",
"x-litellm-session-id": "header-session",
},
)
assert response.status_code == 200, response.text
assert len(sent_bodies) == 1
assert sent_bodies[0].get("metadata") == expected_metadata
assert len(pre_call_bodies.logged_bodies) == 1
assert "name" in pre_call_bodies.logged_bodies[0]
assert VECTOR_STORE_CALLER_METADATA_KWARG not in pre_call_bodies.logged_bodies[0]
@pytest.mark.parametrize("operation", ["create", "update"])
@pytest.mark.parametrize(
"request_body",
[
{"data": {"name": "x", "metadata": json.dumps(CALLERS_METADATA)}},
{
"content": json.dumps({"name": "x", "metadata": CALLERS_METADATA})[:-1]
+ ', "nested": '
+ "[" * 300
+ "]" * 300
+ "}",
"headers": {"Content-Type": "application/json"},
},
],
ids=["form", "json-too-deep-to-reparse"],
)
def test_proxy_vector_store_writes_send_only_the_callers_metadata_when_the_raw_body_cannot_be_reparsed(
monkeypatch, operation, request_body
):
import httpx
import respx
from fastapi.testclient import TestClient
from litellm.proxy import proxy_server
from litellm.proxy.proxy_server import app
monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
litellm.in_memory_llm_clients_cache.flush_cache()
monkeypatch.setattr(proxy_server, "master_key", "sk-master-test")
monkeypatch.setattr(proxy_server, "llm_router", litellm.Router(model_list=[]))
sent_bodies = []
def echo(request: httpx.Request) -> httpx.Response:
body = json.loads(request.content)
sent_bodies.append(body)
return httpx.Response(
200,
json={
"id": "vs_1",
"object": "vector_store",
"created_at": 0,
"name": "x",
"usage_bytes": 0,
"file_counts": {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0},
"status": "completed",
"metadata": body.get("metadata"),
},
)
path = "/v1/vector_stores" if operation == "create" else "/v1/vector_stores/vs_1"
headers = {
"Authorization": "Bearer sk-master-test",
"x-litellm-trace-id": "header-trace",
**request_body.get("headers", {}),
}
body_kwargs = {key: value for key, value in request_body.items() if key != "headers"}
with respx.mock(assert_all_called=False) as upstream:
upstream.post(url__regex=r"https://api\.openai\.com/v1/vector_stores.*").mock(side_effect=echo)
response = TestClient(app).post(path, headers=headers, **body_kwargs)
assert response.status_code == 200, response.text
assert len(sent_bodies) == 1
assert sent_bodies[0].get("metadata") == CALLERS_METADATA

View file

@ -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():

Some files were not shown because too many files have changed in this diff Show more