scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

fix(count_tokens): narrow stable backport to the October 19 patch

Browse source 
Keep #43900, #45301, #46121, #46784, #46759, #46767, #46772, #46815, #47057 and #46777. Port only the Gemini prerequisites from #46430.

Remove the deferred Bedrock Mantle behavior and unrelated picks from the final tree. Restore the PDF test provider fixture and deployment, and retain required test imports.
This commit is contained in:
mateo-berri 2026-10-19 14:03:20 -07:00
parent 94828be8b0
commit 828ddf0c37
48 changed files with 971 additions and 5263 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
.github/workflows/auto_update_price_and_context_window.yml vendored
View file

@ -24,9 +24,12 @@ jobs:
- name: Update JSON Data
run: |
uv run --frozen --with 'aiohttp==3.13.3' python ".github/scripts/auto_update_price_and_context_window_file.py"
- name: Regenerate JSON Schema
run: |
uv run --frozen python ci_cd/generate_model_prices_schema.py
- name: Create Pull Request
run: |
git add model_prices_and_context_window.json
git add model_prices_and_context_window.json model_prices_and_context_window.schema.json
git commit -m "Update model_prices_and_context_window.json file: $(date +'%Y-%m-%d')"
gh pr create --title "Update model_prices_and_context_window.json file" \
--body "Automated update for model_prices_and_context_window.json" \

7
.github/workflows/sync-together-ai-models.yml vendored
View file

@ -41,6 +41,10 @@ jobs:
uv run --frozen python scripts/sync_together_ai_models.py --write --pr-body-file "$RUNNER_TEMP/pr_body.md"
env:
TOGETHER_API_KEY: ${{ secrets.TOGETHER_API_KEY }}
- name: Regenerate the JSON schema
if: steps.existing.outputs.open_pr == ''
run: |
uv run --frozen python ci_cd/generate_model_prices_schema.py
- name: Create a pull request when the registry changed
if: steps.existing.outputs.open_pr == ''
run: |
@ -53,7 +57,8 @@ jobs:
git config user.email "41898282+github-actions[bot]@users.noreply.github.com"
git checkout -b "$branch"
git add model_prices_and_context_window.json \
litellm/model_prices_and_context_window_backup.json
litellm/model_prices_and_context_window_backup.json \
model_prices_and_context_window.schema.json
git commit -m "feat(models): sync together_ai model registry $(date +'%Y-%m-%d')"
gh auth setup-git
git push origin "$branch"

45
ci_cd/cost_map_guard.py
View file

@ -1,11 +1,11 @@
"""Guard the cost map on pull requests.
Every pull request whose diff against its merge base touches one of the two cost map files gets the file
checks: the files parse and the backup copy matches the root file. A pull request that leaves both untouched
skips them, since merging it keeps the base branch's copies and its head tree only carries whatever state the
branch was cut from. Pull requests from the cost map sync bot (branches named litellm_cost_map_sync_*) always
get the file checks and additionally may only touch those two files or the legacy schema
path and may only add or update models.
Every pull request whose diff against its merge base touches one of the three cost map files gets the file
checks: the files parse, the backup copy matches the root file, and the JSON schema is in sync and validates the
map. A pull request that leaves all three untouched skips them, since merging it keeps the base branch's copies
and its head tree only carries whatever state the branch was cut from. Pull requests from the cost map sync bot
(branches named litellm_cost_map_sync_*) always get the file checks and additionally may only touch those three
files and may only add or update models.
"""
from __future__ import annotations
@ -19,12 +19,13 @@ from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
from generate_model_prices_schema import SPECIAL_ROOT_KEYS, build_schema, render, validation_errors
COST_MAP_PATH: Final = "model_prices_and_context_window.json"
BACKUP_PATH: Final = "litellm/model_prices_and_context_window_backup.json"
LEGACY_SCHEMA_PATH: Final = "model_prices_and_context_window.schema.json"
GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH)
SCHEMA_PATH: Final = "model_prices_and_context_window.schema.json"
GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH, SCHEMA_PATH)
BOT_BRANCH_PREFIX: Final = "litellm_cost_map_sync_"
SPECIAL_ROOT_KEYS: Final = frozenset({"sample_spec", "fallback_generalizations"})
EVALUATION_OUTPUT_PRICE_SOURCES: Final[Mapping[str, str]] = MappingProxyType({})
CostMap = dict[str, object]
@ -34,6 +35,7 @@ CostMap = dict[str, object]
class Snapshot:
cost_map: str
backup: str
schema: str
def _parse_object(text: str, path: str) -> CostMap | str:
@ -44,6 +46,13 @@ def _parse_object(text: str, path: str) -> CostMap | str:
return parsed if isinstance(parsed, dict) else f"{path} must be a JSON object at the root"
def _rendered_schema(cost_map: CostMap) -> str:
try:
return render(build_schema(cost_map))
except SystemExit as error:
return str(error)
def _evaluation_output_price_failure(
key: str, entry: object, evaluation_output_price_sources: Mapping[str, str]
) -> str | None:
@ -69,13 +78,29 @@ def _file_failures(
*,
evaluation_output_price_sources: Mapping[str, str] = EVALUATION_OUTPUT_PRICE_SOURCES,
) -> tuple[str, ...]:
schema_text: Final = _rendered_schema(head_map)
if not schema_text.startswith("{"):
return (schema_text,)
backup_failure: Final = (
()
if head.backup == head.cost_map
else (f"{BACKUP_PATH} differs from {COST_MAP_PATH}; copy the root file over it",)
)
schema_failure: Final = (
()
if head.schema == schema_text
else (
f"{SCHEMA_PATH} is out of sync with {COST_MAP_PATH}; "
"run `python ci_cd/generate_model_prices_schema.py` and commit the result",
)
)
return (
*backup_failure,
*schema_failure,
*(
f"{COST_MAP_PATH} does not validate against its schema: {error}"
for error in validation_errors(head_map, json.loads(schema_text))[:20]
),
*(
failure
for key, entry in head_map.items()
@ -110,7 +135,7 @@ def _bot_failures(base: Snapshot, head_map: CostMap, changed_files: Sequence[str
*(
f"bot PRs may only change the cost map files, not {path}"
for path in changed_files
if path not in GUARDED_PATHS and path != LEGACY_SCHEMA_PATH
if path not in GUARDED_PATHS
),
*(f"bot PRs may not remove models: {key}" for key in base_map if key not in head_map),
*(f"bot PRs may not remove fields: {ref}" for ref in removed_fields),

424
ci_cd/generate_model_prices_schema.py Normal file
View file

@ -0,0 +1,424 @@
from __future__ import annotations
import json
import sys
from pathlib import Path
from typing import Final, Optional
import jsonschema
REPO_ROOT = Path(__file__).parent.parent
PRICES_PATH = REPO_ROOT / "model_prices_and_context_window.json"
SCHEMA_PATH = REPO_ROOT / "model_prices_and_context_window.schema.json"
SPECIAL_ROOT_KEYS = frozenset({"sample_spec", "fallback_generalizations"})
JsonSchema = dict
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)$"
)
EXTRA_BOOLEAN_KEYS = frozenset(
{
"gemini_native_audio",
"gemini_audio_only_live",
"uses_embed_content",
"use_openai_responses_path",
"bedrock_converse_supports_strict_tools",
"thinking_always_on",
}
)
HOURS_UTC: Final[JsonSchema] = {
"description": 'UTC "HH:MM-HH:MM" window, or a list of them; a window may wrap past midnight.',
"oneOf": [TIME_WINDOW, {"type": "array", "items": TIME_WINDOW, "minItems": 1}],
}
OFF_PEAK_WINDOW: Final[JsonSchema] = {
"type": "object",
"properties": {
"hours_utc": HOURS_UTC,
"weekdays": {
"type": "array",
"description": "ISO-8601 weekday numbers (1 = Monday .. 7 = Sunday) or English day names the window applies on.",
"items": {
"oneOf": [
{"type": "integer", "minimum": 1, "maximum": 7},
{"type": "string", "pattern": WEEKDAY_PATTERN},
]
},
"minItems": 1,
},
},
"required": ["hours_utc"],
"additionalProperties": False,
}
OBJECT_KEYS: dict[str, JsonSchema] = {
"off_peak_pricing": {
"type": "object",
"description": "Rates that replace the same-named base fields while the request falls inside the stated UTC windows.",
"properties": {
"hours_utc": HOURS_UTC,
"windows": {"type": "array", "items": OFF_PEAK_WINDOW, "minItems": 1},
"weekday_timezone": {
"type": "string",
"description": "IANA zone the weekdays of each window are read on; defaults to UTC.",
},
"input_cost_per_token": NONNEG_NUMBER,
"output_cost_per_token": NONNEG_NUMBER,
"output_cost_per_reasoning_token": NONNEG_NUMBER,
"cache_read_input_token_cost": NONNEG_NUMBER,
"cache_creation_input_token_cost": NONNEG_NUMBER,
},
"anyOf": [{"required": ["hours_utc"]}, {"required": ["windows"]}],
"additionalProperties": False,
},
"search_context_cost_per_query": {
"type": "object",
"description": "USD cost per web search query, keyed by search context size.",
"properties": {
"search_context_size_low": NONNEG_NUMBER,
"search_context_size_medium": NONNEG_NUMBER,
"search_context_size_high": NONNEG_NUMBER,
},
"additionalProperties": False,
},
"guardrail_cost_per_unit": {
"type": "object",
"description": "USD cost per billable guardrail unit, keyed by the provider's usage counter name (e.g. Bedrock's contentPolicyUnits).",
"additionalProperties": NONNEG_NUMBER,
},
"metadata": {
"type": "object",
"description": "Free-form notes about the entry (e.g. pricing derivation).",
},
"provider_specific_entry": {
"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] = {
"supported_audio_formats": {
"type": "array",
"description": "Audio container formats the model can return.",
"items": {"type": "string", "enum": ["mp3", "wav"]},
},
"supported_endpoints": {
"type": "array",
"description": "OpenAI-style API routes this model can be called through, e.g. /v1/chat/completions.",
"items": STRING,
},
"supported_modalities": {
"type": "array",
"description": "Input modalities the model accepts.",
"items": {"type": "string", "enum": ["text", "image", "audio", "video"]},
},
"supported_output_modalities": {
"type": "array",
"description": "Output modalities the model can produce.",
"items": {"type": "string", "enum": ["text", "image", "audio", "video", "code"]},
},
"reasoning_effort_levels": {
"type": "array",
"description": "Exact reasoning_effort levels this deployment accepts; wins over supports_* flags.",
"items": {"type": "string", "enum": ["none", "minimal", "low", "medium", "high", "xhigh", "max"]},
},
"supported_regions": {
"type": "array",
"description": "Cloud regions the model is available in ('global' or region ids).",
"items": STRING,
},
"tiered_pricing": {
"type": "array",
"description": "Context-length or result-count tiered rates; each tier's costs apply within its range.",
"items": {
"type": "object",
"properties": {
"range": {
"type": "array",
"description": "[min, max] prompt-token span this tier applies to.",
"items": NONNEG_NUMBER,
"minItems": 2,
"maxItems": 2,
},
"max_results_range": {
"type": "array",
"description": "[min, max] result-count span this tier applies to (search models).",
"items": NONNEG_NUMBER,
"minItems": 2,
"maxItems": 2,
},
"input_cost_per_token": NONNEG_NUMBER,
"output_cost_per_token": NONNEG_NUMBER,
"output_cost_per_reasoning_token": NONNEG_NUMBER,
"cache_read_input_token_cost": NONNEG_NUMBER,
"cache_creation_input_token_cost": NONNEG_NUMBER,
"input_cost_per_query": NONNEG_NUMBER,
},
"additionalProperties": False,
},
},
}
INTEGER_KEYS: dict[str, JsonSchema] = {
"max_tokens": {
**NONNEG_INTEGER,
"description": "Legacy field: max output tokens if the provider specifies it, else max input tokens.",
},
"max_input_tokens": {
**NONNEG_INTEGER,
"description": "Maximum prompt/context tokens the model accepts.",
},
"max_output_tokens": {
**NONNEG_INTEGER,
"description": "Maximum tokens the model can generate in one response.",
},
"output_vector_size": {
**NONNEG_INTEGER,
"description": "Embedding dimension for embedding models.",
},
"prompt_cache_min_tokens": {
**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."},
}
NUMBER_KEYS: dict[str, JsonSchema] = {
"regional_processing_uplift_multiplier_eu": {
"type": "number",
"minimum": 1,
"description": "Multiplier applied to all token costs for EU data residency (e.g. 1.10 = +10%).",
},
"regional_processing_uplift_multiplier_us": {
"type": "number",
"minimum": 1,
"description": "Multiplier applied to all token costs for US data residency (e.g. 1.10 = +10%).",
},
"regional_endpoint_uplift_multiplier": {
"type": "number",
"minimum": 1,
"description": "Multiplier applied to all token costs when served from a non-global Vertex AI endpoint (e.g. 1.10 = +10%).",
},
}
COST_DESCRIPTIONS: dict[str, str] = {
"input_cost_per_token": "USD per prompt token.",
"output_cost_per_token": "USD per generated token.",
"output_cost_per_reasoning_token": "USD per reasoning/thinking token, when billed separately.",
"google_maps_grounding_cost_per_query": (
"USD per Grounding with Google Maps request; billed per query or per prompt per web_search_billing_unit."
),
"cache_creation_input_token_cost": "USD per token written to the provider's prompt cache.",
"cache_read_input_token_cost": "USD per prompt token served from the provider's prompt cache.",
"input_cost_per_token_batches": "USD per prompt token via the provider's batch API.",
"output_cost_per_token_batches": "USD per generated token via the provider's batch API.",
}
def cost_description(key: str) -> Optional[str]:
if key in COST_DESCRIPTIONS:
return COST_DESCRIPTIONS[key]
if key.endswith("_flex"):
return "Flex service-tier rate for the same-named base field."
if key.endswith("_priority"):
return "Priority service-tier rate for the same-named base field."
if "_above_" in key:
return "Rate applied once the prompt exceeds the token threshold in the field name."
return None
def cost_schema(key: str) -> JsonSchema:
description = cost_description(key)
return {**NONNEG_NUMBER, "description": description} if description else dict(NONNEG_NUMBER)
def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]:
return {
"litellm_provider": {
"type": "string",
"description": "LiteLLM provider slug; one of https://docs.litellm.ai/docs/providers.",
},
"mode": {
"type": "string",
"description": "Primary API surface / task type of the model.",
"enum": list(modes),
},
"source": {
"type": "string",
"description": "URL of the provider pricing/model page this entry was taken from.",
},
"deprecation_date": {
"type": "string",
"description": "Date the provider deprecates the model, YYYY-MM-DD.",
"format": "date",
"pattern": "^\\d{4}-(0[1-9]|1[0-2])-(0[1-9]|[12]\\d|3[01])$",
},
"web_search_billing_unit": {
"type": "string",
"description": "Whether web search is billed per query or per prompt.",
"enum": ["per_query", "per_prompt"],
},
"bedrock_output_config_effort_ceiling": {
"type": "string",
"description": "Highest reasoning effort the Bedrock output_config accepts for this model.",
"enum": ["low", "medium", "high", "max", "xhigh"],
},
"default_reasoning_effort": {
"type": "string",
"description": (
"Reasoning effort the provider applies when the request omits reasoning_effort. "
"Gates whether a non-default temperature or the top_p/logprobs sampling params are "
"accepted, which hold only when the effort resolves to 'none'."
),
"enum": ["none", "minimal", "low", "medium", "high", "xhigh"],
},
"comment": STRING,
"audio_transcription_config": STRING,
"vertex_ai_audio_api": {
"type": "string",
"enum": ["lyria_predict", "lyria_interactions"],
},
}
def classify(key: str, modes: tuple) -> Optional[JsonSchema]:
curated = {**OBJECT_KEYS, **ARRAY_KEYS, **string_key_schemas(modes), **INTEGER_KEYS, **NUMBER_KEYS}
if key in curated:
return curated[key]
if key.startswith("supports_") or key in EXTRA_BOOLEAN_KEYS:
return BOOLEAN
if "cost" in key:
return cost_schema(key)
return None
def build_schema(prices: dict) -> JsonSchema:
entries = {name: entry for name, entry in prices.items() if name not in SPECIAL_ROOT_KEYS}
all_keys = tuple(sorted({key for entry in entries.values() for key in entry}))
modes = tuple(sorted({entry["mode"] for entry in entries.values() if "mode" in entry}))
unclassified = tuple(key for key in all_keys if classify(key, modes) is None)
if unclassified:
raise SystemExit(
f"Unclassified keys in {PRICES_PATH.name}: {', '.join(unclassified)}. "
f"Add them to the key tables in {Path(__file__).name} and rerun it."
)
entry_properties = {key: classify(key, modes) for key in all_keys}
return {
"$schema": "https://json-schema.org/draft/2020-12/schema",
"title": "LiteLLM model_prices_and_context_window.json",
"description": (
"Schema for LiteLLM's model price and context window registry "
"(https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json). "
"Every top-level key except 'sample_spec' and 'fallback_generalizations' is a model id, "
"optionally prefixed with its provider (e.g. 'azure/gpt-5.4'), mapping to a model entry. "
"All costs are USD per unit. New optional fields are added regularly, so consumers should "
"ignore unknown fields rather than reject them."
),
"type": "object",
"properties": {
"sample_spec": {
"type": "object",
"description": (
"Documentation placeholder illustrating the entry shape; not a real model and not "
"schema-conformant (several values are prose)."
),
},
"fallback_generalizations": {
"type": "object",
"description": "Regex rules that generalize unknown model ids to known families; not a model entry.",
"properties": {
"rules": {
"type": "array",
"items": {
"type": "object",
"properties": {
"name": STRING,
"pattern": STRING,
"description": STRING,
},
"required": ["name", "pattern"],
"additionalProperties": True,
},
}
},
"additionalProperties": False,
},
},
"additionalProperties": {"$ref": "#/$defs/modelEntry"},
"$defs": {
"modelEntry": {
"type": "object",
"description": (
"Pricing, limits, and capability flags for one model. Fields other than litellm_provider "
"are optional; boolean capability flags are simply omitted when unknown or false."
),
"required": ["litellm_provider"],
"properties": entry_properties,
"additionalProperties": True,
}
},
}
def render(schema: JsonSchema) -> str:
return json.dumps(schema, indent=2) + "\n"
def validation_errors(prices: dict, schema: JsonSchema) -> tuple:
validator = jsonschema.Draft202012Validator(schema, format_checker=jsonschema.Draft202012Validator.FORMAT_CHECKER)
return tuple(
f"{'.'.join(str(part) for part in error.absolute_path)}: {error.message}"
for error in validator.iter_errors(prices)
)
def main() -> int:
check = "--check" in sys.argv[1:]
prices = json.loads(PRICES_PATH.read_text())
rendered = render(build_schema(prices))
errors = validation_errors(prices, json.loads(rendered))
if errors:
print(f"{PRICES_PATH.name} does not validate against the generated schema:")
print("\n".join(errors[:20]))
return 1
if not check:
SCHEMA_PATH.write_text(rendered)
print(f"wrote {SCHEMA_PATH}")
return 0
if not SCHEMA_PATH.exists() or SCHEMA_PATH.read_text() != rendered:
print(
f"{SCHEMA_PATH.name} is out of sync with {PRICES_PATH.name}. "
f"Run `python {Path(__file__).relative_to(REPO_ROOT)}` and commit the result."
)
return 1
print(f"{SCHEMA_PATH.name} is in sync and {PRICES_PATH.name} validates against it")
return 0
if __name__ == "__main__":
sys.exit(main())

1
litellm-rust/crates/model-catalog/AGENTS.md
View file

@ -2,4 +2,5 @@
For `model_prices_and_context_window.json` validation, we should eventually:
- Remove any schema file like `model_prices_and_context_window.schema.json`
- Stop skipping this crate's tests

37
litellm/litellm_core_utils/prompt_templates/chat_template_tools.py
View file

@ -1,37 +0,0 @@
import json
import re
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from typing import Final
@dataclass(frozen=True, slots=True)
class ChatTemplateToolSection:
preamble_tokens: int
implicit_system_turn_tokens: int
render_tool_schemas: Callable[[Sequence[Mapping[str, object]]], str]
def _deepseek_v4_tool_schemas(functions: Sequence[Mapping[str, object]]) -> str:
return "\n".join(json.dumps(function, ensure_ascii=False) for function in functions)
def _qwen3_5_tool_schemas(functions: Sequence[Mapping[str, object]]) -> str:
return "".join(
"\n" + json.dumps({"type": "function", "function": function}, ensure_ascii=False) for function in functions
)
_TOOL_SECTIONS: Final = (
(
re.compile(r"deepseek-v4(?![.\d]|p\d)", re.IGNORECASE),
ChatTemplateToolSection(222, 0, _deepseek_v4_tool_schemas),
),
(re.compile(r"qwen3[._]?5(?!\d)", re.IGNORECASE), ChatTemplateToolSection(196, 5, _qwen3_5_tool_schemas)),
)
def chat_template_tool_section(model: str | None) -> ChatTemplateToolSection | None:
if not model:
return None
return next((section for pattern, section in _TOOL_SECTIONS if pattern.search(model)), None)

16
litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py
View file

@ -31,7 +31,7 @@ Anthropic wire shape is built later by ``anthropic_messages_pt``.
from collections.abc import Iterator, Mapping, Sequence
from itertools import chain, groupby
from typing import Final, Literal, TypeAlias, TypeVar
from typing import Final, Literal, TypeAlias
from litellm.types.llms.anthropic import AnthropicMessagesSystemMessageParam, AnthropicSystemMessageContent
from litellm.types.llms.openai import (
@ -55,7 +55,6 @@ _RENDERED_ASSISTANT_PART_TYPES: Final = frozenset({"text", "server_tool_use"})
_THINKING_BLOCK_TYPES: Final = frozenset({"thinking", "redacted_thinking"})
_MessageKind: TypeAlias = Literal["system", "tool", "user", "other"]
_Message: Final = TypeVar("_Message")
_TextPart: TypeAlias = tuple[str, ChatCompletionCachedContent | None]
@ -97,8 +96,8 @@ def _kind(message: object) -> _MessageKind:
def split_leading_system_run(
messages: Sequence[_Message],
) -> tuple[tuple[_Message, ...], tuple[_Message, ...]]:
messages: Sequence[AllMessageValues],
) -> tuple[tuple[AllMessageValues, ...], tuple[AllMessageValues, ...]]:
"""Split ``messages`` into the leading run of system messages and everything after it."""
leading_count: Final = next(
(index for index, message in enumerate(messages) if not is_system_message(message)),
@ -161,16 +160,9 @@ def _anthropic_text_block(part: _TextPart) -> AnthropicSystemMessageContent:
return cached
def anthropic_system_blocks(run: Sequence[object]) -> tuple[AnthropicSystemMessageContent, ...]:
"""The top-level ``system`` blocks for a run of system messages: every non-empty text part in order,
each keeping its ``cache_control``, which is the shape the chat path sends for the leading run."""
parts: Final = chain.from_iterable(_text_parts(message) for message in run)
return tuple(_anthropic_text_block(part) for part in parts)
def anthropic_system_messages(message: object) -> tuple[AnthropicMessagesSystemMessageParam, ...]:
"""The Anthropic wire message for a system message, or nothing when it carries no text."""
blocks: Final = anthropic_system_blocks((message,))
blocks: Final = tuple(_anthropic_text_block(part) for part in _text_parts(message))
if not blocks:
return ()
wire: Final[AnthropicMessagesSystemMessageParam] = {

180
litellm/litellm_core_utils/token_counter.py
View file

@ -34,10 +34,6 @@ from litellm.constants import (
TOKEN_COUNTER_MAX_EXACT_CHARS,
)
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.litellm_core_utils.prompt_templates.chat_template_tools import (
ChatTemplateToolSection,
chat_template_tool_section,
)
from litellm.litellm_core_utils.tokenizer import Encoding, HuggingFace, HuggingFaceTokenizer, OpenAIEncoding
from litellm.litellm_core_utils.url_utils import safe_get
from litellm.llms.custom_httpx.http_handler import _get_httpx_client
@ -347,9 +343,6 @@ EXTRAPOLATION_SAMPLES: Final = 16
T_ParamSpec: Final = ParamSpec("T_ParamSpec")
T_Retval = TypeVar("T_Retval")
_COUNT_OFFLOAD_LIMITER: Final = anyio.lowlevel.RunVar[anyio.CapacityLimiter]("litellm_count_offload_limiter")
_ANTHROPIC_TOOL_KEYS_MAPPED_BY_BRIDGE: Final = frozenset(
("name", "input_schema", "description", "cache_control", "strict", "type", "eager_input_streaming")
)
def _count_offload_limiter_for_this_loop() -> anyio.CapacityLimiter:
@ -407,7 +400,6 @@ class _MessageCountParams:
self.tokens_per_message = 3
self.tokens_per_name = 1
self.count_function = _get_count_function(model, custom_tokenizer)
self.tool_section = chat_template_tool_section(model)
def token_counter(
@ -471,9 +463,7 @@ def token_counter(
num_tokens = _count_messages(params, new_messages, use_default_image_token_count, default_token_count)
if count_response_tokens is False:
includes_system_message: Final = any(message.get("role", None) == "system" for message in new_messages)
num_tokens += _count_extra(
params.count_function, tools, tool_choice, includes_system_message, params.tool_section
)
num_tokens += _count_extra(params.count_function, tools, tool_choice, includes_system_message)
else:
raise ValueError("Either text or messages must be provided")
@ -567,7 +557,6 @@ def _count_extra(
tools: list[ChatCompletionToolParam] | None,
tool_choice: ChatCompletionNamedToolChoiceParam | None,
includes_system_message: bool,
tool_section: ChatTemplateToolSection | None,
) -> int:
"""Count extra tokens for function definitions and tool choices.
Args:
@ -575,56 +564,26 @@ def _count_extra(
tools (Optional[List[ChatCompletionToolParam]]): The available tools.
tool_choice (Optional[ChatCompletionNamedToolChoiceParam]): The tool choice.
includes_system_message (bool): Whether the messages include a system message.
tool_section (Optional[ChatTemplateToolSection]): How the model's chat template renders tools, if known.
"""
reply_primer_tokens: Final = 3
tool_choice_tokens: Final = 0 if tool_section is not None else _count_tool_choice(count_function, tool_choice)
return (
reply_primer_tokens
+ _count_tool_definitions(count_function, tools, includes_system_message, tool_section)
+ tool_choice_tokens
)
num_tokens = 3 # every reply is primed with <|start|>assistant<|message|>
def _count_tool_definitions(
count_function: TokenCounterFunction,
tools: Sequence[object] | None,
includes_system_message: bool,
tool_section: ChatTemplateToolSection | None,
) -> int:
if not tools:
return 0
if tool_section is not None:
return _count_chat_template_tools(count_function, tools, includes_system_message, tool_section)
system_message_discount: Final = 4 if includes_system_message else 0
return count_function(_format_function_definitions(tools)) + 9 - system_message_discount
def _count_tool_choice(
count_function: TokenCounterFunction,
tool_choice: ChatCompletionNamedToolChoiceParam | None,
) -> int:
if tools:
num_tokens += count_function(_format_function_definitions(tools))
num_tokens += 9 # Additional tokens for function definition of tools
# If there's a system message and tools are present, subtract four tokens
if tools and includes_system_message:
num_tokens -= 4
# If tool_choice is 'none', add one token.
# If it's an object, add 4 + the number of tokens in the function name.
# If it's undefined or 'auto', don't add anything.
if tool_choice == "none":
return 1
if isinstance(tool_choice, dict):
return 7 + count_function(str(tool_choice["function"]["name"]))
return 0
num_tokens += 1
elif isinstance(tool_choice, dict):
num_tokens += 7
num_tokens += count_function(str(tool_choice["function"]["name"]))
def _count_chat_template_tools(
count_function: TokenCounterFunction,
tools: Sequence[object],
includes_system_message: bool,
tool_section: ChatTemplateToolSection,
) -> int:
functions: Final = tuple(
function for function in _function_definitions(_tools_sent_as_functions(tools)) if function.get("name")
)
if not functions:
return 0
system_turn_tokens: Final = 0 if includes_system_message else tool_section.implicit_system_turn_tokens
schemas: Final = tool_section.render_tool_schemas(functions)
return tool_section.preamble_tokens + system_turn_tokens + count_function(schemas)
return num_tokens
def _get_extrapolating_count_function(
@ -1034,90 +993,49 @@ def _count_content_list(
)
def _format_function_definitions(tools: Sequence[object]) -> str:
def _format_function_definitions(tools):
"""Formats tool definitions in the format that OpenAI appears to use.
Based on https://github.com/forestwanglin/openai-java/blob/main/jtokkit/src/main/java/xyz/felh/openai/jtokkit/utils/TikTokenUtils.java
"""
lines: Final = []
lines.append("namespace functions {")
lines.append("")
for function in _function_definitions(tools):
lines.extend(_format_single_function_definition(function))
for tool in tools:
if not isinstance(tool, dict):
continue
function = tool.get("function")
if not isinstance(function, dict):
# Anthropic tool shape → OpenAI function dict for token counting.
params = tool.get("input_schema") or tool.get("parameters") or {}
if not isinstance(params, dict):
params = {}
function = {
"name": tool.get("name"),
"description": tool.get("description"),
"parameters": params,
}
function_name = function.get("name")
if not function_name:
# Skip malformed tools missing a name to avoid emitting
# ``type None = ...`` which would produce inaccurate token counts.
continue
if function_description := function.get("description"):
lines.append(f"// {function_description}")
parameters = function.get("parameters") or {}
if not isinstance(parameters, dict):
parameters = {}
properties = parameters.get("properties")
if properties and properties.keys():
lines.append(f"type {function_name} = (_: {{")
lines.append(_format_object_parameters(parameters, 0))
lines.append("}) => any;")
else:
lines.append(f"type {function_name} = () => any;")
lines.append("")
lines.append("} // namespace functions")
return "\n".join(lines)
def _tools_sent_as_functions(tools: Sequence[object]) -> tuple[Mapping[str, object], ...]:
return tuple(tool for tool in _tool_mappings(tools) if not _is_anthropic_web_search_tool(tool))
def _is_anthropic_web_search_tool(tool: Mapping[str, object]) -> bool:
return tool.get("name") == "web_search" and "input_schema" not in tool and tool.get("type") != "function"
def _tool_mappings(tools: Sequence[object]) -> Iterable[Mapping[str, object]]:
for tool in tools:
if isinstance(tool, Mapping):
yield cast(Mapping[str, object], tool)
def _function_definitions(tools: Sequence[object]) -> Iterable[Mapping[str, object]]:
for tool in _tool_mappings(tools):
yield from _function_definitions_for_tool(tool)
def _function_definitions_for_tool(tool: Mapping[str, object]) -> Iterable[Mapping[str, object]]:
function: Final = tool.get("function")
if isinstance(function, Mapping):
yield function
return
declarations: Final = tool.get("function_declarations") or tool.get("functionDeclarations")
if isinstance(declarations, list):
for declaration in declarations:
if isinstance(declaration, Mapping):
yield declaration
return
input_schema: Final = tool.get("input_schema")
parameters: Final = input_schema or tool.get("parameters") or {}
normalized_parameters: Final = parameters if isinstance(parameters, Mapping) else {}
unmapped_anthropic_keys: Final = (
{key: value for key, value in tool.items() if key not in _ANTHROPIC_TOOL_KEYS_MAPPED_BY_BRIDGE}
if input_schema is not None
else {}
)
description: Final = tool.get("description")
strict: Final = tool.get("strict")
yield {
"name": tool.get("name"),
"parameters": {**normalized_parameters, **unmapped_anthropic_keys},
**({"description": description} if description is not None else {}),
**({"strict": bool(strict)} if strict is not None else {}),
}
def _format_single_function_definition(function: Mapping[str, object]) -> tuple[str, ...]:
function_name: Final = function.get("name")
if not function_name:
return ()
function_description: Final = function.get("description")
parameters_value: Final = function.get("parameters") or {}
parameters: Final = parameters_value if isinstance(parameters_value, Mapping) else {}
properties: Final = parameters.get("properties")
if isinstance(properties, Mapping) and properties:
return (
*((f"// {function_description}",) if function_description else ()),
f"type {function_name} = (_: {{",
_format_object_parameters(parameters, 0),
"}) => any;",
"",
)
return (
*((f"// {function_description}",) if function_description else ()),
f"type {function_name} = () => any;",
"",
)
def _format_object_parameters(parameters, indent):
properties: Final = parameters.get("properties")
if not properties:

3
litellm/llms/anthropic/count_tokens/handler.py
View file

@ -12,7 +12,6 @@ from pydantic import JsonValue, TypeAdapter
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.llms.anthropic.common_utils import AnthropicError
from litellm.llms.anthropic.count_tokens.transformation import (
AnthropicCountTokensConfig,
@ -63,7 +62,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig):
verbose_logger.debug("Processing Anthropic CountTokens request for model: %s", model)
# Transform request to Anthropic format
request_body: Final = await asyncify(self.transform_request_to_count_tokens)(
request_body: Final = self.transform_request_to_count_tokens(
model=model,
messages=messages,
tools=tools,

70
litellm/llms/anthropic/count_tokens/transformation.py
View file

@ -11,67 +11,11 @@ from typing import Final
from pydantic import JsonValue, TypeAdapter
from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION
from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import (
anthropic_system_blocks,
split_leading_system_run,
)
from litellm.types.llms.openai import ChatCompletionImageObject
_COUNT_REQUEST: Final = TypeAdapter(dict[str, JsonValue])
_SYSTEM_BLOCKS: Final = TypeAdapter(list[JsonValue])
_IMAGE_BLOCK: Final = TypeAdapter(ChatCompletionImageObject)
COUNT_TOKEN_OPTION_NAMES: Final = ("thinking", "tool_choice", "output_config")
def _count_image(block: JsonValue) -> JsonValue:
if not isinstance(block, dict) or block.get("type") != "image_url":
return block
from litellm.litellm_core_utils.prompt_templates.factory import convert_to_anthropic_image_obj
image_block: Final = _IMAGE_BLOCK.validate_python(block)
image_url: Final = image_block["image_url"]
source: Final = convert_to_anthropic_image_obj(
openai_image_url=image_url if isinstance(image_url, str) else image_url["url"],
format=image_url.get("format") if isinstance(image_url, dict) else None,
)
image: Final = _COUNT_REQUEST.validate_python({"type": "image", "source": source})
return {**{key: value for key, value in block.items() if key not in {"type", "image_url"}}, **image}
def _count_block(block: JsonValue) -> JsonValue:
if not isinstance(block, dict) or block.get("type") != "tool_result":
return _count_image(block)
content: Final = block.get("content")
if not isinstance(content, list):
return block
return {**block, "content": [_count_image(part) for part in content]}
def _count_content(content: JsonValue) -> JsonValue:
return [_count_block(block) for block in content] if isinstance(content, list) else content
def _lift_leading_system(
messages: Sequence[Mapping[str, JsonValue]], system: JsonValue
) -> tuple[tuple[Mapping[str, JsonValue], ...], JsonValue]:
"""Move the leading run of system-role messages into the top-level ``system`` parameter.
count_tokens only takes the initial system prompt there and answers 400 on ``role: "system"``
at the head of ``messages``; the chat path sends the same run as ``system``. A caller's own
``system`` keeps its place ahead of the lifted blocks, and a ``system`` that is neither text
nor a block list is left as sent, messages included, for the provider to judge.
"""
leading, conversation = split_leading_system_run(messages)
if not leading or not (system is None or isinstance(system, (str, list))):
return tuple(messages), system
lifted: Final = _SYSTEM_BLOCKS.validate_python(list(anthropic_system_blocks(leading)))
if isinstance(system, list):
return conversation, [*system, *lifted]
if isinstance(system, str) and system:
return conversation, [{"type": "text", "text": system}, *lifted]
return conversation, lifted or system
class AnthropicCountTokensConfig:
"""
Configuration and transformation logic for Anthropic CountTokens API.
@ -102,24 +46,16 @@ class AnthropicCountTokensConfig:
"""
Transform request to Anthropic CountTokens format.
Includes optional system and tools fields for accurate token counting; a leading run of
system-role messages is counted through ``system``, the only place count_tokens accepts it.
Includes optional system and tools fields for accurate token counting.
"""
options: Final[Mapping[str, JsonValue]] = optional_params or MappingProxyType({})
counted_messages, counted_system = _lift_leading_system(messages, system)
return _COUNT_REQUEST.validate_python(
MappingProxyType(
{
"model": model,
"messages": [
{**message, "content": _count_content(message["content"])} for message in counted_messages
],
"messages": messages,
**MappingProxyType(
{
key: value
for key, value in (("system", counted_system), ("tools", tools))
if value is not None
}
{key: value for key, value in (("system", system), ("tools", tools)) if value is not None}
),
**MappingProxyType(
{key: value for key, value in options.items() if key in COUNT_TOKEN_OPTION_NAMES}

3
litellm/llms/azure_ai/anthropic/count_tokens/handler.py
View file

@ -10,7 +10,6 @@ import httpx
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.llms.anthropic.common_utils import AnthropicError
from litellm.llms.azure_ai.anthropic.count_tokens.transformation import (
AzureAIAnthropicCountTokensConfig,
@ -60,7 +59,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig):
verbose_logger.debug("Processing Azure AI Anthropic CountTokens request for model: %s", model)
# Transform request to Anthropic format
request_body: Final = await asyncify(self.transform_request_to_count_tokens)(
request_body: Final = self.transform_request_to_count_tokens(
model=model,
messages=messages,
tools=tools,

181
litellm/llms/bedrock/count_tokens/bedrock_token_counter.py
View file

@ -3,101 +3,29 @@ Bedrock Token Counter implementation using the CountTokens API.
"""
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Any, Final
from pydantic import JsonValue
from typing_extensions import assert_never
from litellm._logging import verbose_logger
from litellm.llms.base_llm.base_utils import BaseTokenCounter
from litellm.llms.bedrock.common_utils import BedrockError, get_bedrock_base_model
from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler
from litellm.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, TokenCountMessagesFormat, TokenCountResponse
RUNTIME_TOKENIZER_TYPE: Final = "bedrock_api"
MANTLE_TOKENIZER_TYPE: Final = "bedrock_mantle_api"
_RUNTIME_FALLBACK_STATUSES: Final = frozenset({400, 404})
class _CountTokensReply(LiteLLMPydanticObjectBase):
input_tokens: int
@dataclass(frozen=True, slots=True)
class CountedTokens:
input_tokens: int
original_response: Mapping[str, JsonValue]
tokenizer_type: str
@dataclass(frozen=True, slots=True)
class CountTokensFailure:
status_code: int
message: str
tokenizer_type: str
CountTokensOutcome = CountedTokens | CountTokensFailure
def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model: str) -> bool:
return (
isinstance(outcome, CountTokensFailure)
and outcome.status_code in _RUNTIME_FALLBACK_STATUSES
and is_mantle_claude_model(resolved_model)
)
from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse
class BedrockTokenCounter(BaseTokenCounter):
"""Token counter for AWS Bedrock: bedrock-runtime CountTokens first, and Anthropic's count_tokens
on bedrock-mantle for Claude models bedrock-runtime answers 400 or 404 for"""
def __init__(
self,
runtime_handler: BedrockCountTokensHandler | None = None,
mantle_handler: BedrockMantleCountTokensHandler | None = None,
client: AsyncHTTPHandler | None = None,
) -> None:
self._runtime_handler: Final = runtime_handler or BedrockCountTokensHandler()
self._mantle_handler: Final = mantle_handler or BedrockMantleCountTokensHandler()
self._client: Final = client
"""Token counter implementation for AWS Bedrock provider using the CountTokens API."""
def should_use_token_counting_api(
self,
custom_llm_provider: str | None = None,
) -> bool:
"""
Returns True if we should use the Bedrock CountTokens API for token counting.
"""
return custom_llm_provider == LlmProviders.BEDROCK.value
async def _count_with(
self,
handler: BedrockCountTokensHandler | BedrockMantleCountTokensHandler,
tokenizer_type: str,
request_data: dict[str, object],
litellm_params: dict[str, object],
resolved_model: str,
) -> CountTokensOutcome:
try:
result: Final = await handler.handle_count_tokens_request(
request_data=request_data,
litellm_params=litellm_params,
resolved_model=resolved_model,
client=self._client,
)
reply: Final = _CountTokensReply.model_validate(result)
except BedrockError as e:
verbose_logger.debug(
"%s CountTokens API error: status=%s, message=%s", tokenizer_type, e.status_code, e.message
)
return CountTokensFailure(status_code=e.status_code, message=e.message, tokenizer_type=tokenizer_type)
except Exception as e:
verbose_logger.debug("Error calling %s CountTokens API: %s", tokenizer_type, e)
return CountTokensFailure(status_code=500, message=str(e), tokenizer_type=tokenizer_type)
return CountedTokens(input_tokens=reply.input_tokens, original_response=result, tokenizer_type=tokenizer_type)
async def count_tokens(
self,
model_to_use: str,
@ -110,52 +38,81 @@ class BedrockTokenCounter(BaseTokenCounter):
messages_format: TokenCountMessagesFormat = "anthropic",
output_format: JsonValue = None,
) -> TokenCountResponse | None:
"""
Count tokens using AWS Bedrock's CountTokens API.
This method calls the existing BedrockCountTokensHandler to make an API call
to Bedrock's token counting endpoint, bypassing the local tiktoken-based counting.
Args:
model_to_use: The model identifier
messages: The messages to count tokens for
contents: Alternative content format (not used for Bedrock)
deployment: Deployment configuration containing litellm_params
request_model: The original request model name
Returns:
TokenCountResponse with token count, or None if counting fails
"""
if not messages:
return None
litellm_params: Final = (deployment or {}).get("litellm_params", {})
deployment = deployment or {}
litellm_params: Final = deployment.get("litellm_params", {})
# Build request data in the format expected by BedrockCountTokensHandler
request_data: Final[dict[str, object]] = {
"model": model_to_use,
"messages": messages,
**({"tools": tools} if tools else {}),
**({"system": system} if system else {}),
}
if tools:
request_data["tools"] = tools
if system:
request_data["system"] = system
# Get the resolved model (strip prefixes like bedrock/, converse/, etc.)
resolved_model: Final = get_bedrock_base_model(model_to_use)
runtime: Final = await self._count_with(
self._runtime_handler, RUNTIME_TOKENIZER_TYPE, request_data, litellm_params, resolved_model
)
outcome: Final = (
await self._count_with(
self._mantle_handler, MANTLE_TOKENIZER_TYPE, request_data, litellm_params, resolved_model
try:
handler: Final = BedrockCountTokensHandler()
result: Final = await handler.handle_count_tokens_request(
request_data=request_data,
litellm_params=litellm_params,
resolved_model=resolved_model,
)
if _runtime_rejected_claude_model(runtime, resolved_model)
else runtime
)
match outcome:
case CountedTokens():
# Transform response to TokenCountResponse
if result is not None:
return TokenCountResponse(
total_tokens=outcome.input_tokens,
total_tokens=result.get("input_tokens", 0),
request_model=request_model,
model_used=model_to_use,
tokenizer_type=outcome.tokenizer_type,
original_response=dict(outcome.original_response),
tokenizer_type="bedrock_api",
original_response=result,
)
case CountTokensFailure():
verbose_logger.warning(
"%s CountTokens API error: status=%s, message=%s",
outcome.tokenizer_type,
outcome.status_code,
outcome.message,
)
return TokenCountResponse(
total_tokens=0,
request_model=request_model,
model_used=model_to_use,
tokenizer_type=outcome.tokenizer_type,
error=True,
error_message=outcome.message,
status_code=outcome.status_code,
)
case _:
assert_never(outcome)
except BedrockError as e:
verbose_logger.warning("Bedrock CountTokens API error: status=%s, message=%s", e.status_code, e.message)
return TokenCountResponse(
total_tokens=0,
request_model=request_model,
model_used=model_to_use,
tokenizer_type="bedrock_api",
error=True,
error_message=e.message,
status_code=e.status_code,
)
except Exception as e:
verbose_logger.warning("Error calling Bedrock CountTokens API: %s", e)
return TokenCountResponse(
total_tokens=0,
request_model=request_model,
model_used=model_to_use,
tokenizer_type="bedrock_api",
error=True,
error_message=str(e),
status_code=500,
)
return None

2
litellm/llms/bedrock/count_tokens/handler.py
View file

@ -125,7 +125,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
raise
except httpx.HTTPStatusError as e:
# HTTP errors - preserve the actual status code
verbose_logger.debug("HTTP error in CountTokens handler: %s", e)
verbose_logger.error("HTTP error in CountTokens handler: %s", e)
raise BedrockError(
status_code=e.response.status_code,
message=e.response.text,

99
litellm/llms/bedrock/count_tokens/mantle_handler.py
View file

@ -1,99 +0,0 @@
from typing import Final
import httpx
from pydantic import JsonValue, TypeAdapter
from typing_extensions import NotRequired, ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_logger
from litellm.llms.anthropic.count_tokens.transformation import AnthropicCountTokensConfig
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing
from litellm.llms.bedrock.common_utils import BedrockError, build_mantle_messages_url
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client
MANTLE_COUNT_TOKENS_SUFFIX: Final = "/count_tokens"
MANTLE_ANTHROPIC_VERSION: Final = "2023-06-01"
class MantleCountTokensRequest(TypedDict):
messages: ReadOnly[list[dict[str, JsonValue]]]
system: ReadOnly[NotRequired[JsonValue]]
tools: ReadOnly[NotRequired[list[dict[str, JsonValue]]]]
_COUNT_REQUEST: Final = TypeAdapter(MantleCountTokensRequest)
_COUNT_RESPONSE: Final = TypeAdapter(dict[str, JsonValue])
class BedrockMantleCountTokensHandler(BaseAWSLLM):
"""Counts tokens through Anthropic's count_tokens on the bedrock-mantle endpoint.
Claude models that Bedrock offers only through cross-region inference answer 400 on
bedrock-runtime's CountTokens; AWS documents Mantle's /anthropic/v1/messages/count_tokens
as the way to count them, with the base model id and the deployment's AWS credentials
"""
def __init__(self, anthropic_config: AnthropicCountTokensConfig | None = None) -> None:
super().__init__()
self._anthropic_config: Final = anthropic_config or AnthropicCountTokensConfig()
def get_mantle_count_tokens_endpoint(self, aws_region_name: str) -> str:
messages_url: Final = build_mantle_messages_url(
api_base=None, aws_bedrock_runtime_endpoint=None, region=aws_region_name
)
return f"{messages_url}{MANTLE_COUNT_TOKENS_SUFFIX}"
async def handle_count_tokens_request(
self,
request_data: dict[str, object],
litellm_params: dict[str, object],
resolved_model: str,
client: AsyncHTTPHandler | None = None,
) -> dict[str, JsonValue]:
try:
request: Final = _COUNT_REQUEST.validate_python(request_data)
aws_region_name: Final = self._get_aws_region_name(
optional_params=litellm_params, model=resolved_model, model_id=None
)
endpoint_url: Final = self.get_mantle_count_tokens_endpoint(aws_region_name)
body: Final = self._anthropic_config.transform_request_to_count_tokens(
model=resolved_model,
messages=request["messages"],
tools=request.get("tools"),
system=request.get("system"),
)
verbose_logger.debug("Making bedrock-mantle count_tokens request to: %s", endpoint_url)
api_key: Final = litellm_params.get("api_key")
signed_headers, signed_body = await run_aws_signing(
self._sign_request,
service_name="bedrock",
headers={"Content-Type": "application/json", "anthropic-version": MANTLE_ANTHROPIC_VERSION},
optional_params=litellm_params,
request_data=body,
api_base=endpoint_url,
model=resolved_model,
api_key=api_key if isinstance(api_key, str) else None,
)
async_client: Final = client or get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK)
response: Final = await async_client.post(
endpoint_url, headers=signed_headers, data=signed_body, timeout=30.0
)
if response.status_code != 200:
raise BedrockError(
status_code=response.status_code,
message=response.text,
headers=response.headers,
response=response,
)
return _COUNT_RESPONSE.validate_json(response.content)
except BedrockError:
raise
except httpx.HTTPStatusError as e:
raise BedrockError(
status_code=e.response.status_code,
message=e.response.text,
headers=e.response.headers,
response=e.response,
)
except Exception as e:
raise BedrockError(status_code=500, message=f"bedrock-mantle count_tokens processing error: {e}")

4
litellm/llms/bedrock_mantle/common_utils.py
View file

@ -158,7 +158,3 @@ def mantle_base_segment(model: str | None, model_cost: dict) -> str:
"""
entry: Final = model_cost.get(f"bedrock_mantle/{split_mantle_region_prefix(model)[1]}", {}) if model else {}
return "openai/v1" if entry.get("use_openai_responses_path") is True else "v1"
def is_mantle_claude_model(model: str) -> bool:
return "claude" in model.lower()

13
litellm/llms/openai/responses/count_tokens/transformation.py
View file

@ -193,10 +193,9 @@ class OpenAICountTokensConfig:
tools: list[dict[str, Any]],
) -> list[dict[str, object]]:
"""
Transform OpenAI chat and Anthropic tools to Responses API tools format.
Transform OpenAI chat tools format to Responses API tools format.
Chat format: {"type": "function", "function": {"name": "...", "parameters": {...}}}
Anthropic format: {"name": "...", "input_schema": {...}}
Responses format: {"type": "function", "name": "...", "parameters": {...}}
"""
transformed: Final[list[dict[str, object]]] = []
@ -212,16 +211,6 @@ class OpenAICountTokensConfig:
if "strict" in func:
item["strict"] = func["strict"]
transformed.append(item)
elif "input_schema" in tool and "name" in tool and tool.get("type") in (None, "custom"):
transformed.append(
{
"type": "function",
"name": tool["name"],
"description": tool.get("description", ""),
"parameters": tool["input_schema"],
"strict": bool(tool.get("strict")),
}
)
else:
# Pass through non-function tools (e.g., web_search, file_search)
transformed.append(tool)

18
litellm/llms/vertex_ai/common_utils.py
View file

@ -1242,7 +1242,6 @@ class VertexAITokenCounter(BaseTokenCounter):
) -> TokenCountResponse | None:
import copy
from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import count_tokens_publisher
from litellm.llms.vertex_ai.vertex_ai_partner_models.main import (
VertexAIError as PartnerVertexAIError,
)
@ -1253,23 +1252,6 @@ class VertexAITokenCounter(BaseTokenCounter):
deployment = deployment or {}
count_tokens_params_request: Final = copy.deepcopy(deployment.get("litellm_params", {}))
without_count_route: Final = (
VertexAIPartnerModels.is_vertex_partner_model(model_to_use) and count_tokens_publisher(model_to_use) is None
)
model_info: Final[Mapping[str, object]] = deployment.get("model_info") or {}
if without_count_route and messages and model_info.get("custom_tokenizer") is not None:
return None
if without_count_route:
return TokenCountResponse(
total_tokens=0,
request_model=request_model,
model_used=model_to_use,
tokenizer_type="vertex_ai_partner_models",
error=True,
error_message=f"LiteLLM does not support counting tokens for Vertex AI model {model_to_use}.",
status_code=400,
)
# Check if this is a partner model (Claude, Mistral, etc.)
if VertexAIPartnerModels.is_vertex_partner_model(model_to_use):
# Use partner models token counter

36
litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py
View file

@ -1,5 +1,5 @@
"""
Token counter for Vertex AI Partner Models (Anthropic Claude, Meta Llama)
Token counter for Vertex AI Partner Models (Anthropic Claude, Mistral, etc.)
This handler provides token counting for partner models hosted on Vertex AI.
Unlike Gemini models which use Google's token counting API, partner models use
@ -12,30 +12,36 @@ from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.llms.vertex_ai.common_utils import get_vertex_base_url
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
_COUNT_TOKENS_PUBLISHERS: Final = (
("claude", "anthropic"),
("llama", "meta"),
("meta/", "meta"),
)
def count_tokens_publisher(model: str) -> str | None:
return next((publisher for marker, publisher in _COUNT_TOKENS_PUBLISHERS if marker in model), None)
class VertexAIPartnerModelsTokenCounter(VertexBase):
"""
Token counter for Vertex AI Partner Models.
Handles token counting for models like Claude (Anthropic) and Llama (Meta)
Handles token counting for models like Claude (Anthropic), Mistral, etc.
that are available through Vertex AI's partner model program.
"""
def _get_publisher_for_model(self, model: str) -> str:
publisher: Final = count_tokens_publisher(model)
if publisher is None:
"""
Determine the publisher name for the given model.
Args:
model: The model name (e.g., "claude-3-5-sonnet-20241022")
Returns:
Publisher name to use in the Vertex AI endpoint URL
Raises:
ValueError: If the model is not a recognized partner model
"""
if "claude" in model:
return "anthropic"
elif "mistral" in model or "codestral" in model:
return "mistralai"
elif "llama" in model or "meta/" in model:
return "meta"
else:
raise ValueError(f"Unknown partner model: {model}")
return publisher
def _build_count_tokens_endpoint(
self,

4
litellm/proxy/common_utils/openai_error_payload.py
View file

@ -42,10 +42,6 @@ def openai_error_type(exc: object, status_code: int) -> str:
carried: Final = attribute_of(exc, "type")
if isinstance(carried, str) and carried != STRINGIFIED_NONE:
return carried
return openai_error_type_for_status(status_code)
def openai_error_type_for_status(status_code: int) -> str:
mapped: Final = _OPENAI_ERROR_TYPE_BY_STATUS.get(status_code)
if mapped is not None:
return mapped

4
litellm/proxy/proxy_server.py
View file

@ -14153,10 +14153,8 @@ async def count_request_tokens(
model_info.get("custom_tokenizer", None),
)
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,
[{"role": "user", "content": prompt}] if prompt_turn else (None if prompt_only else messages),
Sequence[AllMessageValues] | None, messages
)
counted_messages: Final = (
typed_messages if typed_messages is None or system_message is None else (system_message, *typed_messages)

97
litellm/proxy/response_api_endpoints/endpoints.py
View file

@ -14,12 +14,11 @@ from fastapi import APIRouter, Depends, HTTPException, Request, Response
from fastapi.responses import JSONResponse
from openai.types.responses import ResponseItemList
from openai.types.responses.response_create_params import ResponseInputParam
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError, with_config
from pydantic import BaseModel, ConfigDict, ValidationError
from starlette.websockets import WebSocket, WebSocketDisconnect
from typing_extensions import NotRequired, ReadOnly, TypedDict
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.anthropic_interface.exceptions.exception_mapping_utils import AnthropicExceptionMapping
from litellm.constants import EMPTY_MAPPING
from litellm.integrations.custom_guardrail import ModifyResponseException
from litellm.llms.base_llm.guardrail_translation.utils import (
@ -32,16 +31,10 @@ from litellm.proxy.auth.user_api_key_auth import (
user_api_key_auth_websocket,
)
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.common_utils.error_body_call_id import error_body_call_id, with_call_id
from litellm.proxy.common_utils.http_parsing_utils import (
_read_request_body,
_safe_set_request_parsed_body,
)
from litellm.proxy.common_utils.openai_error_payload import (
LITELLM_CALL_ID_HEADER,
error_status_code,
openai_error_type_for_status,
)
from litellm.proxy.common_utils.token_count_settings import count_source_fields
from litellm.proxy.route_llm_request import raise_if_required_body_param_missing
from litellm.types.llms.openai import (
@ -1112,85 +1105,6 @@ def _responses_invalid_request_response(message: str, param: str | None, code: s
return JSONResponse(status_code=400, content=body)
class _ProviderErrorMessage(TypedDict):
message: ReadOnly[str]
class _NestedProviderError(TypedDict):
error: ReadOnly[_ProviderErrorMessage]
@with_config(ConfigDict(extra="ignore"))
class _ProviderOpenAIErrorDetail(TypedDict):
message: ReadOnly[str]
type: ReadOnly[str]
param: NotRequired[ReadOnly[str | None]]
code: NotRequired[ReadOnly[str | None]]
@with_config(ConfigDict(extra="forbid"))
class _ProviderOpenAIError(TypedDict):
error: ReadOnly[_ProviderOpenAIErrorDetail]
_PROVIDER_COUNT_ERROR_TYPE: Final = "token_counting_error"
_PROVIDER_OPENAI_ERROR: Final = TypeAdapter(_ProviderOpenAIError)
_NESTED_PROVIDER_ERROR: Final = TypeAdapter(_NestedProviderError)
def _provider_error_json(raw_message: str) -> str:
start: Final = raw_message.find("{")
return raw_message if start == -1 else raw_message[start:]
def _provider_error_message(raw_message: str) -> str:
error_json: Final = _provider_error_json(raw_message)
try:
return _NESTED_PROVIDER_ERROR.validate_json(error_json)["error"]["message"]
except ValidationError:
flat_message: Final = AnthropicExceptionMapping.extract_error_message(error_json)
return raw_message if flat_message == error_json else flat_message
def _provider_count_error_body(raw_message: str, status_code: int) -> _ResponsesApiErrorBody:
try:
provider_error: Final = _PROVIDER_OPENAI_ERROR.validate_json(_provider_error_json(raw_message))["error"]
except ValidationError:
return {
"error": {
"message": _provider_error_message(raw_message),
"type": openai_error_type_for_status(status_code),
"param": None,
"code": None,
}
}
return {
"error": {
"message": provider_error["message"],
"type": provider_error["type"],
"param": provider_error.get("param"),
"code": provider_error.get("code"),
}
}
def _provider_count_error_response(exc: ProxyException, request: Request) -> JSONResponse:
from litellm.proxy.proxy_server import (
_close_dangling_otel_server_span, # pyright: ignore[reportPrivateUsage] # proxy_server keeps the span-close helper private; a route that answers an error itself must stamp the OTel server span like the global ProxyException handler does
general_settings_view,
)
status_code: Final = error_status_code(exc, 500)
_close_dangling_otel_server_span(request, status_code, exc=exc)
body: Final = _provider_count_error_body(exc.message, status_code)
call_id: Final = error_body_call_id(general_settings_view(), exc.headers.get(LITELLM_CALL_ID_HEADER))
return JSONResponse(
status_code=status_code,
content={"error": with_call_id({**body["error"]}, call_id)},
headers=exc.headers,
)
def _missing_responses_param_response(param: str) -> JSONResponse:
return _responses_invalid_request_response(
message=f"Missing required parameter: '{param}'.",
@ -1282,12 +1196,7 @@ async def responses_input_tokens(
message=f"Invalid request for token counting: {e}", param=None, code=None
)
try:
token_response: Final = await token_counter(request=token_request, call_endpoint=True)
except ProxyException as e:
if e.type != _PROVIDER_COUNT_ERROR_TYPE:
raise
return _provider_count_error_response(e, request)
token_response: Final = await token_counter(request=token_request, call_endpoint=True)
result: Final[_ResponsesInputTokensResult] = {
"object": "response.input_tokens",
"input_tokens": token_response.total_tokens,

278
tests/integration/_support/bedrock_runtime_peer.py
View file

@ -1,278 +0,0 @@
from __future__ import annotations
import json
import re
import threading
from collections.abc import Mapping
from multiprocessing.sharedctypes import Synchronized
from types import MappingProxyType
from typing import Final
from urllib.parse import unquote
from integration._support.upstream import _aws_event_frame
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
MARKER: Final = re.compile(r"marker-([0-9a-f]{32})")
EVENT_STREAM: Final = "application/vnd.amazon.eventstream"
REASONING_EFFORTS: Final = frozenset(("none", "minimal", "low", "medium", "high", "xhigh"))
NATIVE_CHAT: Final = "/openai/v1/chat/completions"
NATIVE_RESPONSES: Final = "/openai/v1/responses"
PNG_1X1: Final = bytes.fromhex(
"89504e470d0a1a0a0000000d49484452000000010000000108060000001f15c489"
"0000000d49444154789c63f8cfc0f01f00050001ff89993d1d0000000049454e44ae426082"
)
USAGE: Final[Mapping[str, JsonValue]] = MappingProxyType(
{
"prompt_tokens": 9,
"completion_tokens": 5,
"total_tokens": 14,
"completion_tokens_details": {"reasoning_tokens": 3},
}
)
_STATUS: Final = re.compile(r"status=(\d{3})")
_CONVERSE: Final = re.compile(r"^/model/(.+)/converse$")
_CONVERSE_STREAM: Final = re.compile(r"^/model/(.+)/converse-stream$")
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_NO_MARKER: Final = "0" * 32
def marker_of(request: Request) -> str:
found: Final = MARKER.search(request.body.decode(errors="replace"))
return _NO_MARKER if found is None else found.group(1)
def body_of(request: Request) -> Mapping[str, JsonValue]:
try:
return _JSON_OBJECT.validate_json(request.body)
except ValueError:
return {}
def target_of(request: Request) -> str:
return unquote(request.target)
def answer(marker: str) -> str:
return f"answer marker-{marker}"
def reasoning_answer(marker: str) -> str:
return f"<reasoning>why marker-{marker}</reasoning> {answer(marker)}"
def _headers(marker: str) -> Mapping[str, str]:
return MappingProxyType({"x-amzn-requestid": marker})
def _json_reply(status: int, payload: Mapping[str, JsonValue], marker: str) -> Reply:
return Reply(status=status, body=json.dumps(payload).encode(), headers=_headers(marker))
def _error(status: int, message: str, marker: str) -> Reply:
return _json_reply(status, {"message": message}, marker)
def _effort_of(target: str, body: Mapping[str, JsonValue]) -> JsonValue:
if not _CONVERSE.match(target) and not _CONVERSE_STREAM.match(target):
return body.get("reasoning_effort")
fields: Final = body.get("additionalModelRequestFields")
reasoning: Final = fields.get("reasoning") if isinstance(fields, Mapping) else None
return reasoning.get("effort") if isinstance(reasoning, Mapping) else None
def forwarded_effort(request: Request) -> JsonValue:
return _effort_of(target_of(request), body_of(request))
def _sse(frames: tuple[Mapping[str, JsonValue], ...], pause: float) -> Reply:
return Reply(
content_type="text/event-stream",
chunks=(*(b"data: " + json.dumps(frame).encode() + b"\n\n" for frame in frames), b"data: [DONE]\n\n"),
pause_between_chunks=pause,
)
def _with_headers(reply: Reply, marker: str) -> Reply:
return Reply(
status=reply.status,
body=reply.body,
content_type=reply.content_type,
chunks=reply.chunks,
abort_after=reply.abort_after,
gate_after_first=reply.gate_after_first,
pause_between_chunks=reply.pause_between_chunks,
headers=_headers(marker),
)
def _content_deltas(model: str, marker: str) -> tuple[str, ...]:
if "gpt-oss" in model:
return ("<reason", "ing>why ", f"marker-{marker}", "</reas", "oning> answer ", f"marker-{marker}")
return ("answer ", f"marker-{marker}")
def _chat_text(model: str, marker: str) -> str:
return reasoning_answer(marker) if "gpt-oss" in model else answer(marker)
def _chat_reply(model: str, marker: str, stream: bool, pause: float) -> Reply:
identity: Final = f"chatcmpl-{marker}"
if not stream:
return _json_reply(
200,
{
"id": identity,
"object": "chat.completion",
"created": 1,
"model": model,
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": _chat_text(model, marker)},
"finish_reason": "stop",
}
],
"usage": dict(USAGE),
},
marker,
)
deltas: Final = _content_deltas(model, marker)
frames: Final = tuple(
{
"id": identity,
"object": "chat.completion.chunk",
"created": 1,
"model": model,
"choices": [{"index": 0, "delta": {"role": "assistant", "content": delta}, "finish_reason": None}],
}
for delta in deltas
)
finish: Final[Mapping[str, JsonValue]] = {
"id": identity,
"object": "chat.completion.chunk",
"created": 1,
"model": model,
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
"usage": dict(USAGE),
}
return _with_headers(_sse((*frames, finish), pause), marker)
def _responses_reply(model: str, marker: str, stream: bool, pause: float) -> Reply:
identity: Final = f"resp_upstream_{marker}"
item_id: Final = f"msg_{marker}"
response: Final[Mapping[str, JsonValue]] = {
"id": identity,
"object": "response",
"created_at": 1,
"status": "completed",
"model": model,
"output": [
{
"type": "message",
"id": item_id,
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": answer(marker), "annotations": []}],
}
],
"usage": {"input_tokens": 30, "output_tokens": 5, "total_tokens": 35},
}
if not stream:
return _json_reply(200, response, marker)
events: Final[tuple[Mapping[str, JsonValue], ...]] = (
{
"type": "response.created",
"sequence_number": 0,
"response": {**response, "status": "in_progress", "output": []},
},
{
"type": "response.output_text.delta",
"sequence_number": 1,
"item_id": item_id,
"output_index": 0,
"content_index": 0,
"delta": answer(marker),
},
{"type": "response.completed", "sequence_number": 2, "response": response},
)
return Reply(
content_type="text/event-stream",
chunks=tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events),
pause_between_chunks=pause,
headers=_headers(marker),
)
def _converse_reply(marker: str) -> Reply:
return _json_reply(
200,
{
"output": {"message": {"role": "assistant", "content": [{"text": answer(marker)}]}},
"stopReason": "end_turn",
"usage": {"inputTokens": 9, "outputTokens": 5, "totalTokens": 14},
"metrics": {"latencyMs": 1},
},
marker,
)
def _converse_stream_reply(marker: str, pause: float) -> Reply:
events: Final[tuple[tuple[str, Mapping[str, JsonValue]], ...]] = (
("messageStart", {"role": "assistant"}),
("contentBlockDelta", {"delta": {"text": "answer "}, "contentBlockIndex": 0}),
("contentBlockDelta", {"delta": {"text": f"marker-{marker}"}, "contentBlockIndex": 0}),
("contentBlockStop", {"contentBlockIndex": 0}),
("messageStop", {"stopReason": "end_turn"}),
("metadata", {"usage": {"inputTokens": 9, "outputTokens": 5, "totalTokens": 14}, "metrics": {"latencyMs": 1}}),
)
return Reply(
content_type=EVENT_STREAM,
chunks=tuple(_aws_event_frame(kind, payload, "sc", marker) for kind, payload in events),
pause_between_chunks=pause,
headers=_headers(marker),
)
def respond(request: Request, *, pause: float = 0.0) -> Reply:
target: Final = target_of(request)
marker: Final = marker_of(request)
if request.method == "GET":
if target == "/image.png":
return Reply(body=PNG_1X1, content_type="image/png", headers=_headers(marker))
return _error(404, f"no scripted object at {target}", marker)
scripted_status: Final = _STATUS.search(request.body.decode(errors="replace"))
if scripted_status is not None:
status: Final = int(scripted_status.group(1))
return _error(status, f"scripted {status}", marker)
body: Final = body_of(request)
effort: Final = _effort_of(target, body)
if effort is not None and (not isinstance(effort, str) or effort not in REASONING_EFFORTS):
return _error(400, f"Invalid reasoning effort: {json.dumps(effort)}", marker)
model: Final = str(body.get("model", ""))
stream: Final = body.get("stream") is True
if request.method == "POST" and target == NATIVE_CHAT:
return _chat_reply(model, marker, stream, pause)
if request.method == "POST" and target == NATIVE_RESPONSES:
return _responses_reply(model, marker, stream, pause)
if request.method == "POST" and _CONVERSE.match(target):
return _converse_reply(marker)
if request.method == "POST" and _CONVERSE_STREAM.match(target):
return _converse_stream_reply(marker, pause)
return _error(404, f"unknown bedrock route {request.method} {target}", marker)
def serve_peer(port: int, received: Synchronized[int], answer_first: int) -> None:
held: Final = threading.Event()
def respond_or_hold(request: Request) -> Reply:
with received.get_lock():
received.value += 1
ordinal: Final = received.value
if ordinal > answer_first:
held.wait()
return respond(request)
with wire_server(respond_or_hold, port=port):
threading.Event().wait()

126
tests/integration/_support/openai_wire.py
View file

@ -1,126 +0,0 @@
from __future__ import annotations
import json
from collections.abc import Callable
from typing import Final
from integration._support.wire import Reply, Request, Wire
from pydantic import JsonValue
from tests.integration._support.provider import MODEL_DISCOVERY
_USAGE: Final = {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}
def answering_model_discovery(respond: Callable[[Request], Reply]) -> Callable[[Request], Reply]:
def guarded(request: Request) -> Reply:
if (request.method, request.target) == MODEL_DISCOVERY:
return Reply(body=json.dumps({"object": "list", "data": []}).encode())
return respond(request)
return guarded
def posted_targets(wire: Wire) -> tuple[str, ...]:
return tuple(request.target for request in wire.drain() if request.method == "POST")
def openai_error(status: int) -> Reply:
return Reply(
status=status,
body=json.dumps({"error": {"message": f"scripted {status}", "type": "server_error", "code": None}}).encode(),
)
def _data_frame(frame: dict[str, JsonValue]) -> bytes:
return b"data: " + json.dumps(frame).encode() + b"\n\n"
def _typed_frame(event: dict[str, JsonValue]) -> bytes:
return f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode()
def chat_stream(identity: str, model: str, text: str) -> tuple[bytes, bytes, bytes]:
chunk: Final = {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": model}
role_only: Final = _data_frame({**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}}]})
content: Final = _data_frame({**chunk, "choices": [{"index": 0, "delta": {"content": text}}]})
finish: Final = _data_frame(
{**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], "usage": _USAGE}
)
return (role_only, content, finish + b"data: [DONE]\n\n")
def chat_reply(identity: str, model: str, text: str, *, stream: bool) -> Reply:
if stream:
return Reply(content_type="text/event-stream", chunks=chat_stream(identity, model, text))
return Reply(
body=json.dumps(
{
"id": identity,
"object": "chat.completion",
"created": 1,
"model": model,
"choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}],
"usage": _USAGE,
}
).encode()
)
def _response_object(identity: str, model: str, text: str) -> dict[str, JsonValue]:
return {
"id": identity,
"object": "response",
"created_at": 1,
"status": "completed",
"model": model,
"output": [_message_item(identity, text, "completed")],
"usage": {"input_tokens": 5, "output_tokens": 3, "total_tokens": 8},
}
def _message_item(identity: str, text: str, status: str) -> dict[str, JsonValue]:
return {
"id": f"msg_{identity}",
"type": "message",
"role": "assistant",
"status": status,
"content": [{"type": "output_text", "text": text, "annotations": []}] if status == "completed" else [],
}
def responses_stream(identity: str, model: str, text: str) -> tuple[bytes, bytes, bytes]:
response: Final = _response_object(identity, model, text)
item: Final = _message_item(identity, text, "in_progress")
opened: Final = _typed_frame(
{
"type": "response.created",
"sequence_number": 0,
"response": {**response, "status": "in_progress", "output": []},
}
) + _typed_frame({"type": "response.output_item.added", "sequence_number": 1, "output_index": 0, "item": item})
delta: Final = _typed_frame(
{
"type": "response.output_text.delta",
"sequence_number": 2,
"item_id": f"msg_{identity}",
"output_index": 0,
"content_index": 0,
"delta": text,
}
)
closed: Final = _typed_frame(
{
"type": "response.output_item.done",
"sequence_number": 3,
"output_index": 0,
"item": _message_item(identity, text, "completed"),
}
) + _typed_frame({"type": "response.completed", "sequence_number": 4, "response": response})
return (opened, delta, closed)
def responses_reply(identity: str, model: str, text: str, *, stream: bool) -> Reply:
if stream:
return Reply(content_type="text/event-stream", chunks=responses_stream(identity, model, text))
return Reply(body=json.dumps(_response_object(identity, model, text)).encode())

30
tests/integration/conftest.py
View file

@ -15,6 +15,7 @@ from redis import Redis
from tests.integration._support.client import Gateway, eventually, gateway_from_environment
from tests.integration._support.generation import LIFECYCLE_SETTINGS
from tests.integration._support.manifest import OWNED_DIRECTORIES
from tests.integration._support.provider import SharedProvider, shared_provider
from tests.integration._support.routing import RoutingPlugin
COLLECTED: Final = pytest.StashKey[tuple[str, ...]]()
@ -126,6 +127,35 @@ def gateway() -> Iterator[Gateway]:
yield value
@pytest.fixture(scope="session")
def shared_provider_server() -> Iterator[SharedProvider]:
if os.environ.get("PYTEST_XDIST_WORKER") and os.environ.get("INTEGRATION_WORKER_ISOLATION") != "1":
pytest.fail("the shared fake provider needs tests to run one at a time; this group runs under pytest-xdist")
with shared_provider() as server:
yield server
late: Final = server.received()
assert late == (), (
f"the shared fake provider got {[item.target for item in late]} after {server.last_test} finished"
)
@pytest.fixture
def provider(shared_provider_server: SharedProvider, request: pytest.FixtureRequest) -> Iterator[SharedProvider]:
stray: Final = shared_provider_server.received()
shared_provider_server.replies.clear()
assert stray == (), (
f"the shared fake provider got {[item.target for item in stray]} "
f"after {shared_provider_server.last_test} finished"
)
yield shared_provider_server
shared_provider_server.last_test = request.node.nodeid
unused: Final = len(shared_provider_server.replies)
unread: Final = shared_provider_server.received()
shared_provider_server.replies.clear()
assert unused == 0, f"{unused} queued provider replies were never requested"
assert unread == (), f"the test never read the provider requests {[item.target for item in unread]}"
@pytest.fixture
def peer(gateway: Gateway) -> Iterator[Gateway]:
url: Final = os.environ["INTEGRATION_PEER_URL"]

159
tests/integration/providers/_count_tokens_system_lift.py
View file

@ -1,159 +0,0 @@
from collections.abc import Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
import anthropic
import httpx
import openai
from integration._support.client import Gateway
from pydantic import JsonValue, TypeAdapter
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
TOKEN_COUNTING_BETA: Final = "token-counting-2024-11-01"
ANTHROPIC_VERSION: Final = "2023-06-01"
REJECTION: Final = 'messages.0: Unexpected role "system". The Messages API accepts a top-level `system` parameter'
USER_TEXT: Final = "Count this message"
USER: Final[dict[str, JsonValue]] = {"role": "user", "content": USER_TEXT}
ASSISTANT: Final[dict[str, JsonValue]] = {"role": "assistant", "content": "One."}
FOLLOW_UP: Final[dict[str, JsonValue]] = {"role": "user", "content": "Again"}
INSTRUCTION: Final = "You are a terse assistant"
REMINDER: Final = "Answer in one sentence"
LEADING: Final[dict[str, JsonValue]] = {"role": "system", "content": INSTRUCTION}
LIFTED: Final[dict[str, JsonValue]] = {"type": "text", "text": INSTRUCTION}
MID_SYSTEM: Final[dict[str, JsonValue]] = {"role": "system", "content": REMINDER}
EPHEMERAL: Final[dict[str, JsonValue]] = {"type": "ephemeral"}
ONE_HOUR: Final[dict[str, JsonValue]] = {"type": "ephemeral", "ttl": "1h"}
IMAGE_PART: Final[dict[str, JsonValue]] = {
"type": "image_url",
"image_url": {
"url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg=="
},
}
FIVE_KB: Final = "Answer in one sentence. " * 214
CALLER_SYSTEM: Final = "Prefer metric units"
CALLER_BLOCKS: Final[list[JsonValue]] = [
{"type": "text", "text": CALLER_SYSTEM},
{"type": "text", "text": "Never guess", "cache_control": EPHEMERAL},
]
TOOLS: Final[list[JsonValue]] = [
{
"name": "get_weather",
"description": "Look up the current weather for a city",
"input_schema": {
"type": "object",
"properties": {"city": {"type": "string", "description": "City to look up"}},
"required": ["city"],
},
}
]
@dataclass(frozen=True, slots=True)
class LiftCase:
messages: tuple[dict[str, JsonValue], ...]
system: JsonValue | None
lifted: list[JsonValue] | None
LIFT_CASES: Final[Mapping[str, LiftCase]] = MappingProxyType(
{
"string": LiftCase((LEADING, USER), None, [LIFTED]),
"run_with_cache_control": LiftCase(
(
{"role": "system", "content": INSTRUCTION, "cache_control": ONE_HOUR},
{
"role": "system",
"content": [
{"type": "text", "text": REMINDER, "cache_control": EPHEMERAL},
{"type": "text", "text": ""},
IMAGE_PART,
],
},
USER,
),
None,
[
{"type": "text", "text": INSTRUCTION, "cache_control": ONE_HOUR},
{"type": "text", "text": REMINDER, "cache_control": EPHEMERAL},
],
),
"caller_system_string_first": LiftCase(
(LEADING, USER), CALLER_SYSTEM, [{"type": "text", "text": CALLER_SYSTEM}, LIFTED]
),
"caller_system_blocks_first": LiftCase((LEADING, USER), CALLER_BLOCKS, [*CALLER_BLOCKS, LIFTED]),
"caller_empty_system_dropped": LiftCase((LEADING, USER), "", [LIFTED]),
"empty_content_dropped": LiftCase(({"role": "system", "content": ""}, USER), None, None),
"image_only_content_dropped": LiftCase(({"role": "system", "content": [IMAGE_PART]}, USER), None, None),
"integer_content_dropped": LiftCase(({"role": "system", "content": 5}, USER), None, None),
"integer_text_part_dropped": LiftCase(
({"role": "system", "content": [{"type": "text", "text": 7}]}, USER), None, None
),
"five_kb": LiftCase(({"role": "system", "content": FIVE_KB}, USER), None, [{"type": "text", "text": FIVE_KB}]),
}
)
STRING_CASE: Final = LIFT_CASES["string"]
def count_request(model: str, case: LiftCase) -> dict[str, JsonValue]:
return {"model": model, "messages": list(case.messages), **({} if case.system is None else {"system": case.system})}
def expected_count_body(model: str, case: LiftCase) -> dict[str, JsonValue]:
conversation: Final[list[JsonValue]] = [message for message in case.messages if message["role"] != "system"]
return {"model": model, "messages": conversation, **({} if case.lifted is None else {"system": case.lifted})}
def _opening_message(message: JsonValue) -> bool:
return isinstance(message, dict) and message.get("role") in ("user", "assistant")
def _conversation_message(message: JsonValue) -> bool:
return isinstance(message, dict) and message.get("role") in ("user", "assistant", "system")
def _anthropic_tool(tool: JsonValue) -> bool:
return isinstance(tool, dict) and isinstance(tool.get("name"), str) and isinstance(tool.get("input_schema"), dict)
def accepts_count_body(body: Mapping[str, JsonValue]) -> bool:
# Anthropic, Azure AI Foundry and Bedrock Mantle count_tokens verdicts observed live on 2026-10-09: an empty
# messages list and a system role at messages[0] answer 400 invalid_request_error, a later system role counts
messages: Final = body.get("messages")
tools: Final = body.get("tools", [])
return (
isinstance(messages, list)
and len(messages) > 0
and _opening_message(messages[0])
and all(map(_conversation_message, messages))
and isinstance(body.get("system", ""), (str, list))
and isinstance(tools, list)
and all(map(_anthropic_tool, tools))
)
def anthropic_client(gateway: Gateway) -> anthropic.Anthropic:
return anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0)
def async_anthropic_client(gateway: Gateway) -> anthropic.AsyncAnthropic:
return anthropic.AsyncAnthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0)
def openai_client(gateway: Gateway) -> openai.OpenAI:
return openai.OpenAI(
base_url=f"{gateway.client.base_url}/v1",
api_key=gateway.key,
max_retries=0,
http_client=httpx.Client(trust_env=False),
)
def async_openai_client(gateway: Gateway) -> openai.AsyncOpenAI:
return openai.AsyncOpenAI(
base_url=f"{gateway.client.base_url}/v1",
api_key=gateway.key,
max_retries=0,
http_client=httpx.AsyncClient(trust_env=False),
)

581
tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py
View file

@ -1,581 +0,0 @@
import asyncio
import json
import socket
import threading
import uuid
from collections.abc import Callable, Iterator, Mapping
from concurrent.futures import ThreadPoolExecutor
from contextlib import AbstractContextManager, ExitStack
from dataclasses import dataclass
from queue import SimpleQueue
from types import MappingProxyType
from typing import Final, cast
import httpcore
import httpx
import psutil
import pytest
from anthropic.types import MessageParam
from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment
from integration._support.wire import Reply, Request, Wire, wire_server
from integration.providers._count_tokens_system_lift import (
ANTHROPIC_VERSION,
ASSISTANT,
EPHEMERAL,
FOLLOW_UP,
IMAGE_PART,
INSTRUCTION,
JSON_OBJECT,
LEADING,
LIFT_CASES,
LIFTED,
MID_SYSTEM,
REJECTION,
REMINDER,
STRING_CASE,
TOKEN_COUNTING_BETA,
TOOLS,
USER,
USER_TEXT,
LiftCase,
accepts_count_body,
anthropic_client,
async_anthropic_client,
async_openai_client,
count_request,
expected_count_body,
openai_client,
)
from pydantic import JsonValue, TypeAdapter
_MODEL: Final = "claude-opus-5-5"
_API_KEY: Final = "synthetic-count-tokens-key"
_COUNT: Final = 3131
_PROVIDER_TOKENIZERS: Final = frozenset({"anthropic_api", "azure_ai_anthropic_api"})
_SDK_MESSAGES: Final = cast(
list[MessageParam], [LEADING, USER]
) # cast-ok: the SDK types reject the role the proxy lifts
_STRING_BODY: Final = expected_count_body(_MODEL, STRING_CASE)
_PROXY_MODULE: Final = "integration._support.proxy"
_PROBES_PER_ROUND: Final = 8
_CLIENT_ADDRESS: Final = TypeAdapter(tuple[str, int])
_ARGUMENTS: Final = TypeAdapter(tuple[str, ...])
_NAME: Final = TypeAdapter(str)
@dataclass(frozen=True, slots=True)
class _Provider:
prefix: str
target: str
tokenizer: str
azure: bool
@dataclass(frozen=True, slots=True)
class _Deployment:
provider: _Provider
port: int
model: str
_ANTHROPIC: Final = _Provider("anthropic", "/v1/messages/count_tokens", "anthropic_api", False)
_AZURE: Final = _Provider("azure_ai", "/anthropic/v1/messages/count_tokens", "azure_ai_anthropic_api", True)
_PROVIDERS: Final = MappingProxyType({"anthropic": _ANTHROPIC, "azure_ai": _AZURE})
@pytest.fixture(params=_PROVIDERS.keys())
def provider(request: pytest.FixtureRequest) -> _Provider:
name: Final[object] = request.param # pyright: ignore[reportAny] # pytest types the fixture param as Any
return _PROVIDERS[_NAME.validate_python(name)]
def _json_reply(status: int, payload: Mapping[str, JsonValue]) -> Reply:
return Reply(status=status, body=json.dumps(payload).encode())
def _rejected(status: int) -> Reply:
return _json_reply(status, {"type": "error", "error": {"type": "invalid_request_error", "message": REJECTION}})
def _counted(request: Request) -> Reply:
accepted: Final = accepts_count_body(JSON_OBJECT.validate_json(request.body))
return _json_reply(200, {"input_tokens": _COUNT}) if accepted else _rejected(400)
def _rejecting(status: int) -> Callable[[Request], Reply]:
def count(_request: Request) -> Reply:
return _rejected(status)
return count
def _holding(held: SimpleQueue[str], release: threading.Event, seconds: float) -> Callable[[Request], Reply]:
def hold(request: Request) -> Reply:
held.put(request.target)
assert release.wait(timeout=seconds), "Held count was never released"
return _counted(request)
return hold
def _peer(provider: _Provider, count: Callable[[Request], Reply] = _counted) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
if request.target == provider.target:
return count(request)
return _json_reply(404, {"error": f"unscripted target {request.target}"})
return respond
def _listening(deployment: _Deployment, count: Callable[[Request], Reply] = _counted) -> AbstractContextManager[Wire]:
return wire_server(_peer(deployment.provider, count), port=deployment.port)
def _reserved_port() -> int:
with socket.socket() as reserve:
reserve.bind(("127.0.0.1", 0))
return _CLIENT_ADDRESS.validate_python(reserve.getsockname())[1]
def _cmdline(process: psutil.Process) -> tuple[str, ...]:
try:
return _ARGUMENTS.validate_python(process.cmdline())
except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
return ()
def _serves(cmdline: tuple[str, ...], proxy_port: int) -> bool:
if _PROXY_MODULE not in cmdline or "--port" not in cmdline:
return False
return cmdline[cmdline.index("--port") + 1] == str(proxy_port)
def _listens(process: psutil.Process, proxy_port: int) -> bool:
return any(
connection.status == psutil.CONN_LISTEN and connection.laddr and connection.laddr.port == proxy_port
for connection in process.net_connections(kind="tcp")
)
def _proxy_workers(proxy_port: int) -> frozenset[int]:
(master,) = tuple(process for process in psutil.process_iter() if _serves(_cmdline(process), proxy_port))
spawned: Final = frozenset(child.pid for child in master.children() if _listens(child, proxy_port))
return spawned or frozenset({master.pid})
def _holder(workers: frozenset[int], client_port: int) -> int | None:
def holds(pid: int) -> bool:
return any(
connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == client_port
for connection in psutil.Process(pid).net_connections(kind="tcp")
)
return next((pid for pid in sorted(workers) if holds(pid)), None)
def _client_port(response: httpx.Response) -> int:
stream: Final[object] = response.extensions["network_stream"] # pyright: ignore[reportAny] # httpx types extensions as Any
assert isinstance(stream, httpcore.NetworkStream), stream
return _CLIENT_ADDRESS.validate_python(stream.get_extra_info("client_addr"))[1]
def _probe(gateway: Gateway, body: Mapping[str, JsonValue], workers: frozenset[int]) -> tuple[int | None, JsonValue]:
with httpx.Client(base_url=str(gateway.client.base_url), timeout=30, trust_env=False) as client:
response: Final = client.post(
"/v1/messages/count_tokens", json=dict(body), headers={"Authorization": f"Bearer {gateway.key}"}
)
return _holder(workers, _client_port(response)), JSON_OBJECT.validate_json(response.content).get("input_tokens")
def _round(
gateway: Gateway, body: Mapping[str, JsonValue], workers: frozenset[int]
) -> tuple[tuple[int | None, JsonValue], ...]:
def probe(_index: int) -> tuple[int | None, JsonValue]:
return _probe(gateway, body, workers)
with ThreadPoolExecutor(max_workers=_PROBES_PER_ROUND) as pool:
return tuple(pool.map(probe, range(_PROBES_PER_ROUND)))
def _settled_on_every_worker(gateway: Gateway, body: Mapping[str, JsonValue]) -> None:
proxy_port: Final = gateway.client.base_url.port
assert proxy_port is not None, gateway.client.base_url
workers: Final = _proxy_workers(proxy_port)
eventually(
lambda: _round(gateway, body, workers),
lambda observed: (
frozenset(pid for pid, _ in observed) == workers and all(count == _COUNT for _, count in observed)
),
seconds=60,
)
def _deploy(gateway: Gateway, scenario: Scenario, provider: _Provider) -> _Deployment:
port: Final = _reserved_port()
model: Final = scenario.model(
model_info=None, model=f"{provider.prefix}/{_MODEL}", api_base=f"http://127.0.0.1:{port}", api_key=_API_KEY
)
with wire_server(_peer(provider), port=port):
_settled_on_every_worker(gateway, {"model": model, "messages": [USER]})
return _Deployment(provider, port, model)
@pytest.fixture(scope="module")
def deployments() -> Iterator[Mapping[str, _Deployment]]:
with gateway_from_environment() as gateway, gateway.scenario() as scenario:
yield MappingProxyType({chosen.prefix: _deploy(gateway, scenario, chosen) for chosen in (_ANTHROPIC, _AZURE)})
@pytest.fixture
def deployment(provider: _Provider, deployments: Mapping[str, _Deployment]) -> _Deployment:
return deployments[provider.prefix]
def _count(gateway: Gateway, body: Mapping[str, JsonValue], key: str | None = None) -> httpx.Response:
return gateway.request("POST", "/v1/messages/count_tokens", body, key=key)
def _payload(response: httpx.Response) -> dict[str, JsonValue]:
assert response.status_code == 200, response.text
return JSON_OBJECT.validate_json(response.content)
def _local_count(gateway: Gateway, body: Mapping[str, JsonValue]) -> int:
response: Final = gateway.request("POST", "/utils/token_counter", body, params={"call_endpoint": "false"})
payload: Final = _payload(response)
total: Final = payload["total_tokens"]
assert payload["tokenizer_type"] not in _PROVIDER_TOKENIZERS, response.text
assert isinstance(total, int) and total > 0 and total != _COUNT, response.text
return total
def _bodies(wire: Wire, provider: _Provider) -> tuple[dict[str, JsonValue], ...]:
received: Final = wire.drain()
for request in received:
assert (request.method, request.target) == ("POST", provider.target), request.target
assert request.headers["anthropic-version"] == ANTHROPIC_VERSION, request.headers
assert TOKEN_COUNTING_BETA in request.headers["anthropic-beta"], request.headers
assert request.headers["content-type"] == "application/json", request.headers
assert request.headers["x-api-key"] == _API_KEY, request.headers
assert (request.headers.get("api-key") == _API_KEY) is provider.azure, request.headers
return tuple(JSON_OBJECT.validate_json(request.body) for request in received)
def _clients(stack: ExitStack, base_url: str, count: int) -> tuple[httpx.Client, ...]:
return tuple(
stack.enter_context(httpx.Client(base_url=base_url, timeout=30, trust_env=False)) for _ in range(count)
)
def _counted_on(client: httpx.Client, key: str, body: Mapping[str, JsonValue]) -> tuple[int, JsonValue]:
response: Final = client.post(
"/v1/messages/count_tokens", json=dict(body), headers={"Authorization": f"Bearer {key}"}
)
return response.status_code, JSON_OBJECT.validate_json(response.content).get("input_tokens")
@pytest.mark.parametrize("case", LIFT_CASES.values(), ids=LIFT_CASES.keys())
def test_messages_count_tokens_lifts_the_leading_system_run(
deployment: _Deployment, gateway: Gateway, case: LiftCase
) -> None:
with _listening(deployment) as peer:
response: Final = _count(gateway, count_request(deployment.model, case))
assert _payload(response) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, response.text
assert _bodies(peer, deployment.provider) == (expected_count_body(_MODEL, case),)
def test_anthropic_sdk_count_tokens_lifts_the_leading_system_message(deployment: _Deployment, gateway: Gateway) -> None:
with _listening(deployment) as peer:
counted: Final = anthropic_client(gateway).messages.count_tokens(model=deployment.model, messages=_SDK_MESSAGES)
assert counted.input_tokens == _COUNT, counted
assert _bodies(peer, deployment.provider) == (_STRING_BODY,)
def test_async_anthropic_sdk_count_tokens_lifts_the_leading_system_message(
deployment: _Deployment, gateway: Gateway
) -> None:
with _listening(deployment) as peer:
counted: Final = asyncio.run(
async_anthropic_client(gateway).messages.count_tokens(model=deployment.model, messages=_SDK_MESSAGES)
)
assert counted.input_tokens == _COUNT, counted
assert _bodies(peer, deployment.provider) == (_STRING_BODY,)
def test_utils_token_counter_call_endpoint_counts_a_leading_system_through_the_provider(
deployment: _Deployment, gateway: Gateway
) -> None:
with _listening(deployment) as peer:
response: Final = gateway.request(
"POST",
"/utils/token_counter",
count_request(deployment.model, STRING_CASE),
params={"call_endpoint": "true"},
)
payload: Final = _payload(response)
expected_tokenizer: Final = deployment.provider.tokenizer
assert (payload["total_tokens"], payload["tokenizer_type"]) == (_COUNT, expected_tokenizer), response.text
assert payload["original_response"] == {"input_tokens": _COUNT}, response.text
assert (payload["request_model"], payload["model_used"]) == (deployment.model, _MODEL), response.text
assert _bodies(peer, deployment.provider) == (_STRING_BODY,)
def test_utils_token_counter_local_mode_never_calls_the_peer_for_a_leading_system(
deployment: _Deployment, gateway: Gateway
) -> None:
with _listening(deployment) as peer:
assert _local_count(gateway, count_request(deployment.model, STRING_CASE)) > 0
assert peer.drain() == ()
def test_responses_input_tokens_lifts_instructions(deployment: _Deployment, gateway: Gateway) -> None:
with _listening(deployment) as peer:
response: Final = gateway.request(
"POST",
"/v1/responses/input_tokens",
{"model": deployment.model, "input": USER_TEXT, "instructions": INSTRUCTION},
)
assert _payload(response) == {
"object": "response.input_tokens",
"input_tokens": _COUNT,
"litellm_count_source": "provider",
}, response.text
assert _bodies(peer, deployment.provider) == (_STRING_BODY,)
def test_responses_input_tokens_lifts_instructions_ahead_of_a_leading_system_item(
deployment: _Deployment, gateway: Gateway
) -> None:
with _listening(deployment) as peer:
response: Final = gateway.request(
"POST",
"/v1/responses/input_tokens",
{"model": deployment.model, "input": [MID_SYSTEM, USER], "instructions": INSTRUCTION},
)
assert _payload(response) == {
"object": "response.input_tokens",
"input_tokens": _COUNT,
"litellm_count_source": "provider",
}, response.text
assert _bodies(peer, deployment.provider) == (
{"model": _MODEL, "messages": [USER], "system": [LIFTED, {"type": "text", "text": REMINDER}]},
)
def test_openai_sdk_input_tokens_lifts_instructions(deployment: _Deployment, gateway: Gateway) -> None:
with _listening(deployment) as peer:
counted: Final = openai_client(gateway).responses.input_tokens.count(
model=deployment.model, input=USER_TEXT, instructions=INSTRUCTION
)
assert counted.input_tokens == _COUNT, counted
assert _bodies(peer, deployment.provider) == (_STRING_BODY,)
def test_async_openai_sdk_input_tokens_lifts_instructions(deployment: _Deployment, gateway: Gateway) -> None:
with _listening(deployment) as peer:
counted: Final = asyncio.run(
async_openai_client(gateway).responses.input_tokens.count(
model=deployment.model, input=USER_TEXT, instructions=INSTRUCTION
)
)
assert counted.input_tokens == _COUNT, counted
assert _bodies(peer, deployment.provider) == (_STRING_BODY,)
def test_messages_count_tokens_forwards_tools_beside_the_lifted_system(
deployment: _Deployment, gateway: Gateway
) -> None:
with _listening(deployment) as peer:
response: Final = _count(gateway, {**count_request(deployment.model, STRING_CASE), "tools": TOOLS})
assert _payload(response) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, response.text
assert _bodies(peer, deployment.provider) == ({**_STRING_BODY, "tools": TOOLS},)
def test_messages_count_tokens_keeps_a_mid_conversation_system_in_place(
deployment: _Deployment, gateway: Gateway
) -> None:
with _listening(deployment) as peer:
body: Final[dict[str, JsonValue]] = {
"model": deployment.model,
"messages": [LEADING, USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP],
}
response: Final = _count(gateway, body)
assert _payload(response) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, response.text
assert _bodies(peer, deployment.provider) == (
{"model": _MODEL, "messages": [USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP], "system": [LIFTED]},
)
def test_messages_count_tokens_surfaces_the_peer_error_when_every_message_is_system(
deployment: _Deployment, gateway: Gateway
) -> None:
with _listening(deployment) as peer:
body: Final[dict[str, JsonValue]] = {"model": deployment.model, "messages": [LEADING]}
response: Final = _count(gateway, body)
assert response.status_code == 400, response.text
assert "input_tokens" not in response.text, response.text
assert _bodies(peer, deployment.provider) == ({"model": _MODEL, "messages": [], "system": [LIFTED]},)
def test_messages_count_tokens_surfaces_the_peer_error_for_a_non_text_system(
deployment: _Deployment, gateway: Gateway
) -> None:
with _listening(deployment) as peer:
response: Final = _count(gateway, {**count_request(deployment.model, STRING_CASE), "system": 5})
assert response.status_code == 400, response.text
assert "input_tokens" not in response.text, response.text
assert _bodies(peer, deployment.provider) == ({"model": _MODEL, "messages": [LEADING, USER], "system": 5},)
def test_messages_count_tokens_repeated_request_lifts_each_time(deployment: _Deployment, gateway: Gateway) -> None:
with _listening(deployment) as peer:
answers: Final = tuple(
_payload(_count(gateway, count_request(deployment.model, STRING_CASE))) for _ in range(2)
)
assert answers == ({"input_tokens": _COUNT, "litellm_count_source": "provider"},) * 2
assert _bodies(peer, deployment.provider) == (_STRING_BODY,) * 2
def test_messages_count_tokens_surfaces_the_counting_error_for_a_leading_system_without_content(
deployment: _Deployment, gateway: Gateway
) -> None:
with _listening(deployment) as peer:
body: Final[dict[str, JsonValue]] = {"model": deployment.model, "messages": [{"role": "system"}, USER]}
response: Final = _count(gateway, body)
assert 500 <= response.status_code < 600, response.text
assert "input_tokens" not in response.text, response.text
assert peer.drain() == ()
follow_up: Final = _count(gateway, count_request(deployment.model, STRING_CASE))
assert _payload(follow_up) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, follow_up.text
assert _bodies(peer, deployment.provider) == (_STRING_BODY,)
def test_messages_count_tokens_duplicate_messages_key_lifts_the_last_value(
deployment: _Deployment, gateway: Gateway
) -> None:
with _listening(deployment) as peer:
first: Final = json.dumps([USER])
last: Final = json.dumps([LEADING, USER])
response: Final = gateway.client.post(
"/v1/messages/count_tokens",
content=f'{{"model": "{deployment.model}", "messages": {first}, "messages": {last}}}',
headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"},
)
assert _payload(response) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, response.text
assert _bodies(peer, deployment.provider) == (_STRING_BODY,)
@pytest.mark.parametrize("status", [400, 403, 404, 500, 503])
def test_messages_count_tokens_surfaces_the_peer_status_when_the_peer_rejects_the_lifted_body(
deployment: _Deployment, gateway: Gateway, status: int
) -> None:
with _listening(deployment, _rejecting(status)) as peer:
body: Final = count_request(deployment.model, STRING_CASE)
response: Final = _count(gateway, body)
assert response.status_code == status, response.text
assert "input_tokens" not in response.text, response.text
assert _bodies(peer, deployment.provider) == (_STRING_BODY,)
def test_messages_count_tokens_unauthenticated_request_never_reaches_the_peer(
deployment: _Deployment, gateway: Gateway
) -> None:
with _listening(deployment) as peer:
response: Final = _count(gateway, count_request(deployment.model, STRING_CASE), key="sk-not-a-key")
assert response.status_code == 401, response.text
assert peer.drain() == ()
def test_peer_outage_between_concurrent_waves_errors_then_recovers(
deployment: _Deployment, gateway: Gateway
) -> None:
with ExitStack() as stack:
clients: Final = _clients(stack, str(gateway.client.base_url), 8)
pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients)))
body: Final = count_request(deployment.model, STRING_CASE)
def count(client: httpx.Client) -> tuple[int, JsonValue]:
return _counted_on(client, gateway.key, body)
with _listening(deployment) as peer:
assert tuple(pool.map(count, clients)) == ((200, _COUNT),) * len(clients)
assert _bodies(peer, deployment.provider) == (_STRING_BODY,) * len(clients)
outage: Final = tuple(pool.map(count, clients))
for status, tokens in outage:
assert 500 <= status < 600 and tokens is None, outage
with _listening(deployment) as revived:
assert tuple(pool.map(count, clients)) == ((200, _COUNT),) * len(clients)
assert _bodies(revived, deployment.provider) == (_STRING_BODY,) * len(clients)
def test_slow_peer_holds_concurrent_lifted_counts_without_stalling_the_proxy(
deployment: _Deployment, gateway: Gateway
) -> None:
held: Final[SimpleQueue[str]] = SimpleQueue()
release: Final = threading.Event()
with ExitStack() as stack:
clients: Final = _clients(stack, str(gateway.client.base_url), 6)
peer: Final = stack.enter_context(_listening(deployment, _holding(held, release, 20)))
pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients)))
stack.callback(release.set)
body: Final = count_request(deployment.model, STRING_CASE)
futures: Final = tuple(pool.submit(_counted_on, client, gateway.key, body) for client in clients)
eventually(held.qsize, lambda size: size == len(clients), seconds=30)
assert gateway.request("GET", "/health/liveliness").status_code == 200
assert _local_count(gateway, body) > 0
assert not any(future.done() for future in futures)
release.set()
assert tuple(future.result(timeout=30) for future in futures) == ((200, _COUNT),) * len(clients)
assert _bodies(peer, deployment.provider) == (_STRING_BODY,) * len(clients)
def test_chat_completions_on_the_same_deployment_keeps_a_mid_conversation_system_in_place(
deployments: Mapping[str, _Deployment], gateway: Gateway
) -> None:
identity: Final = f"msg_{uuid.uuid4().hex}"
anthropic_deployment: Final = deployments[_ANTHROPIC.prefix]
def respond(request: Request) -> Reply:
assert (request.method, request.target) == ("POST", "/v1/messages"), request.target
return _json_reply(
200,
{
"id": identity,
"type": "message",
"role": "assistant",
"model": _MODEL,
"content": [{"type": "text", "text": "done"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 12, "output_tokens": 3},
},
)
with wire_server(respond, port=anthropic_deployment.port) as peer:
reminder: Final[dict[str, JsonValue]] = {
"role": "system",
"content": [
{"type": "text", "text": REMINDER, "cache_control": EPHEMERAL},
{"type": "text", "text": ""},
IMAGE_PART,
],
}
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": anthropic_deployment.model,
"max_tokens": 16,
"messages": [LEADING, USER, reminder, ASSISTANT, FOLLOW_UP],
},
)
assert response.status_code == 200, response.text
(sent,) = peer.drain()
body: Final = JSON_OBJECT.validate_json(sent.body)
assert body["system"] == [LIFTED], body
assert body["messages"] == [
{"role": "user", "content": [{"type": "text", "text": USER_TEXT}]},
{"role": "system", "content": [{"type": "text", "text": REMINDER, "cache_control": EPHEMERAL}]},
{"role": "assistant", "content": [{"type": "text", "text": "One."}]},
{"role": "user", "content": [{"type": "text", "text": "Again"}]},
], body

1568
tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py
View file

File diff suppressed because it is too large Load diff

383
tests/integration/providers/test_chat_template_tool_count_wire.py
View file

@ -1,383 +0,0 @@
import json
import re
import uuid
from collections.abc import Iterator, Mapping, Sequence
from contextlib import ExitStack
from dataclasses import dataclass
from itertools import chain
from pathlib import Path
from typing import Final
import httpx
import pytest
import yaml
from integration._support.client import (
JSON_OBJECT,
Gateway,
Scenario,
gateway_from_environment,
object_value,
string_value,
)
from integration._support.openai_wire import chat_reply
from integration._support.process import graceful_stop_seconds, owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
from litellm.litellm_core_utils.default_encoding import encoding as cl100k_base
from pydantic import JsonValue
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
_API_KEY: Final = "scripted-provider-key"
_DEEPSEEK_V4: Final = "azure_ai/DeepSeek-V4-Flash-0731"
_QWEN3_5: Final = "databricks/databricks-qwen35-122b-a10b"
_GROK: Final = "azure_ai/grok-4-1"
_DEEPSEEK_V4_TOOL_PREAMBLE: Final = 222
_QWEN3_5_TOOL_PREAMBLE: Final = 196
_QWEN3_5_OPENED_SYSTEM_TURN: Final = 5
_OPENAI_TOOL_OVERHEAD: Final = 9
_SCHEMA: Final[dict[str, JsonValue]] = {
"type": "object",
"properties": {"query": {"type": "string"}},
"required": ["query"],
}
_FUNCTION: Final[dict[str, JsonValue]] = {"name": "lookup", "parameters": _SCHEMA}
_OPENAI_NAMESPACE: Final = "namespace functions {\n\ntype lookup = (_: {\nquery: string,\n}) => any;\n\n} // namespace functions"
_ANTHROPIC_TOOL: Final[dict[str, JsonValue]] = {"name": "lookup", "input_schema": _SCHEMA}
_OPENAI_TOOL: Final[dict[str, JsonValue]] = {"type": "function", "function": _FUNCTION}
_RESPONSES_TOOL: Final[dict[str, JsonValue]] = {"type": "function", **_FUNCTION}
_NAMED_TOOL_CHOICE: Final[dict[str, JsonValue]] = {"type": "function", "function": {"name": "lookup"}}
_WEB_SEARCH_TOOL: Final[dict[str, JsonValue]] = {"type": "web_search_20250305", "name": "web_search", "max_uses": 3}
_PROMPT: Final = "Find record 14."
_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": _PROMPT}]
_INPUT_PRICE: Final = 0.001
_OUTPUT_PRICE: Final = 0.002
_TINY_BUDGET: Final = 0.000000001
_ESTIMATED_COST: Final = re.compile(r"Estimated request cost: ([^\s,]+)")
def _tokens(text: str) -> int:
return len(cl100k_base.encode(text))
def _deepseek_v4_tool_tokens(function: Mapping[str, JsonValue]) -> int:
return _DEEPSEEK_V4_TOOL_PREAMBLE + _tokens(json.dumps(function, ensure_ascii=False))
def _qwen3_5_tool_tokens(function: Mapping[str, JsonValue], *, with_system: bool) -> int:
opened_system_turn: Final = 0 if with_system else _QWEN3_5_OPENED_SYSTEM_TURN
schema: Final = json.dumps({"type": "function", "function": function}, ensure_ascii=False)
return _QWEN3_5_TOOL_PREAMBLE + opened_system_turn + _tokens("\n" + schema)
def _completion(_request: Request) -> Reply:
return chat_reply(f"chatcmpl-{uuid.uuid4().hex}", "scripted-model", "OK", stream=False)
def _deployment(scenario: Scenario, group: str, model: str, api_base: str) -> str:
name: Final = f"{group}-{uuid.uuid4().hex}"
created: Final = scenario.gateway.post(
"/model/new",
{
"model_name": name,
"litellm_params": {"model": model, "api_base": api_base, "api_key": _API_KEY},
"model_info": {},
},
)
scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"]))
return name
@dataclass(frozen=True, slots=True)
class _CountRoute:
path: str
field: str
tool: dict[str, JsonValue]
def body(self, model: str, tools: Sequence[JsonValue]) -> dict[str, JsonValue]:
prompt: Final[dict[str, JsonValue]] = (
{"input": _PROMPT} if self.path == "/v1/responses/input_tokens" else {"messages": _MESSAGES}
)
return {"model": model, **prompt, **({"tools": list(tools)} if tools else {})}
_COUNT_ROUTES: Final[Mapping[str, _CountRoute]] = {
"messages-count-tokens": _CountRoute("/v1/messages/count_tokens", "input_tokens", _ANTHROPIC_TOOL),
"utils-token-counter": _CountRoute("/utils/token_counter", "total_tokens", _OPENAI_TOOL),
"responses-input-tokens": _CountRoute("/v1/responses/input_tokens", "input_tokens", _RESPONSES_TOOL),
}
_MESSAGES_COUNT: Final = _COUNT_ROUTES["messages-count-tokens"]
def _count(gateway: Gateway, route: _CountRoute, body: Mapping[str, JsonValue]) -> int:
response: Final = gateway.request("POST", route.path, body)
assert response.status_code == 200, response.text
counted: Final = JSON_OBJECT.validate_json(response.content)[route.field]
assert isinstance(counted, int), response.text
return counted
@pytest.mark.parametrize("route", sorted(_COUNT_ROUTES))
def test_count_routes_count_deepseek_v4_tools_the_way_its_chat_template_renders_them(
gateway: Gateway, route: str
) -> None:
selected: Final = _COUNT_ROUTES[route]
with wire_server(_completion) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, "count", _DEEPSEEK_V4, wire.url)
without_tools: Final = _count(gateway, selected, selected.body(model, ()))
with_tool: Final = _count(gateway, selected, selected.body(model, (selected.tool,)))
upstream: Final = wire.drain()
assert upstream == ()
assert with_tool - without_tools == _deepseek_v4_tool_tokens(_FUNCTION)
@pytest.mark.parametrize("with_system", (False, True), ids=("no-system", "system"))
def test_messages_count_tokens_counts_qwen3_5_tools_the_way_its_chat_template_renders_them(
gateway: Gateway, with_system: bool
) -> None:
system: Final[dict[str, JsonValue]] = {"system": "Be brief."} if with_system else {}
with wire_server(_completion) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, "count", _QWEN3_5, wire.url)
without_tools: Final = _count(gateway, _MESSAGES_COUNT, {**_MESSAGES_COUNT.body(model, ()), **system})
with_tool: Final = _count(
gateway, _MESSAGES_COUNT, {**_MESSAGES_COUNT.body(model, (_ANTHROPIC_TOOL,)), **system}
)
upstream: Final = wire.drain()
assert upstream == ()
assert with_tool - without_tools == _qwen3_5_tool_tokens(_FUNCTION, with_system=with_system)
def test_messages_count_tokens_skips_web_search_and_counts_unmapped_tool_keys_for_deepseek_v4(
gateway: Gateway,
) -> None:
deferred_function: Final[dict[str, JsonValue]] = {"name": "lookup", "parameters": {**_SCHEMA, "defer_loading": True}}
with wire_server(_completion) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, "count", _DEEPSEEK_V4, wire.url)
plain: Final = _count(gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, (_ANTHROPIC_TOOL,)))
with_web_search: Final = _count(
gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, (_ANTHROPIC_TOOL, _WEB_SEARCH_TOOL))
)
deferred: Final = _count(
gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, ({**_ANTHROPIC_TOOL, "defer_loading": True},))
)
upstream: Final = wire.drain()
assert upstream == ()
assert with_web_search == plain
assert deferred - plain == _deepseek_v4_tool_tokens(deferred_function) - _deepseek_v4_tool_tokens(_FUNCTION)
def test_messages_count_tokens_keeps_openai_tool_framing_for_other_models(gateway: Gateway) -> None:
with wire_server(_completion) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, "count", _GROK, wire.url)
without_tools: Final = _count(gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, ()))
plain: Final = _count(gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, (_ANTHROPIC_TOOL,)))
deferred: Final = _count(
gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, ({**_ANTHROPIC_TOOL, "defer_loading": True},))
)
upstream: Final = wire.drain()
assert upstream == ()
assert plain - without_tools == _OPENAI_TOOL_OVERHEAD + _tokens(_OPENAI_NAMESPACE)
assert deferred == plain
@dataclass(frozen=True, slots=True)
class _Endpoint:
path: str
body: dict[str, JsonValue]
tools: dict[str, JsonValue]
_ENDPOINTS: Final[Mapping[str, _Endpoint]] = {
"chat-completions": _Endpoint(
"/v1/chat/completions",
{"messages": _MESSAGES, "max_tokens": 1},
{"tools": [_OPENAI_TOOL], "tool_choice": _NAMED_TOOL_CHOICE},
),
"messages": _Endpoint(
"/v1/messages",
{"messages": _MESSAGES, "max_tokens": 1},
{"tools": [_ANTHROPIC_TOOL], "tool_choice": {"type": "auto"}},
),
"responses": _Endpoint(
"/v1/responses",
{"input": _PROMPT, "max_output_tokens": 16},
{"tools": [_RESPONSES_TOOL]},
),
}
def _send(gateway: Gateway, endpoint: _Endpoint, group: str, key: str, *, tools: bool) -> httpx.Response:
return gateway.request(
"POST",
endpoint.path,
{"model": group, **endpoint.body, **(endpoint.tools if tools else {})},
key=key,
)
@dataclass(frozen=True, slots=True)
class _BudgetRig:
gateway: Gateway
wire: Wire
group: str
def _strict_budget_config(directory: Path, group: str, api_base: str) -> Path:
base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
config: Final = {
**base,
"model_list": [
{
"model_name": group,
"litellm_params": {
"model": _DEEPSEEK_V4,
"api_base": api_base,
"api_key": _API_KEY,
"input_cost_per_token": _INPUT_PRICE,
"output_cost_per_token": _OUTPUT_PRICE,
},
}
],
"general_settings": {**object_value(base["general_settings"]), "fail_closed_budget_enforcement": True},
}
path: Final = directory / f"chat-template-tool-budget-{uuid.uuid4().hex[:8]}.yaml"
path.write_text(yaml.safe_dump(config))
return path
@pytest.fixture(scope="module")
def strict_budget(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_BudgetRig]:
directory: Final = tmp_path_factory.mktemp("chat-template-tool-budget")
group: Final = f"deepseek-v4-flash-{uuid.uuid4().hex}"
with ExitStack() as stack:
rig_gateway: Final = stack.enter_context(gateway_from_environment())
wire: Final = stack.enter_context(wire_server(_completion))
owned: Final = stack.enter_context(
owned_proxy_process(rig_gateway, directory, {}, config=_strict_budget_config(directory, group, wire.url))
)
yield _BudgetRig(owned.gateway, wire, group)
def _estimated_cost(response: httpx.Response) -> float:
assert response.status_code in (400, 422, 429), f"{response.status_code} {response.text}"
error: Final = object_value(JSON_OBJECT.validate_json(response.content)["error"])
estimate: Final = _ESTIMATED_COST.search(string_value(error["message"]))
assert estimate is not None, response.text
return float(estimate[1])
@pytest.mark.parametrize("endpoint", ("chat-completions", "messages"))
def test_budget_reservation_estimates_deepseek_v4_tools_the_way_its_chat_template_renders_them(
strict_budget: _BudgetRig, endpoint: str
) -> None:
selected: Final = _ENDPOINTS[endpoint]
with strict_budget.gateway.scenario() as scenario:
key: Final = scenario.key(models=[strict_budget.group], max_budget=_TINY_BUDGET)
without_tools: Final = _estimated_cost(_send(strict_budget.gateway, selected, strict_budget.group, key, tools=False))
with_tools: Final = _estimated_cost(_send(strict_budget.gateway, selected, strict_budget.group, key, tools=True))
assert strict_budget.wire.drain() == ()
assert with_tools - without_tools == pytest.approx(_deepseek_v4_tool_tokens(_FUNCTION) * _INPUT_PRICE)
def _project_key(scenario: Scenario, group: str) -> str:
team: Final = scenario.team()
project: Final = scenario.project(team, model_itpm_limit={group: _deepseek_v4_tool_tokens(_FUNCTION)})
return scenario.key(team_id=team, project_id=project)
@pytest.mark.parametrize("endpoint", sorted(_ENDPOINTS))
def test_project_itpm_reserves_deepseek_v4_tools_the_way_its_chat_template_renders_them(
gateway: Gateway, endpoint: str
) -> None:
with wire_server(_completion) as wire, gateway.scenario() as scenario:
group: Final = _deployment(scenario, "deepseek-v4-flash", _DEEPSEEK_V4, wire.url)
response: Final = _send(gateway, _ENDPOINTS[endpoint], group, _project_key(scenario, group), tools=True)
upstream: Final = wire.drain()
assert response.status_code == 429, response.text
assert "model_per_project_itpm" in response.text, response.text
assert upstream == ()
def test_project_itpm_keeps_openai_tool_framing_for_other_models(gateway: Gateway) -> None:
chat: Final = _ENDPOINTS["chat-completions"]
with wire_server(_completion) as wire, gateway.scenario() as scenario:
group: Final = _deployment(scenario, "grok-4-1", _GROK, wire.url)
response: Final = gateway.request(
"POST",
chat.path,
{"model": group, **chat.body, "tools": [_OPENAI_TOOL]},
key=_project_key(scenario, group),
)
upstream: Final = wire.drain()
assert response.status_code == 200, response.text
assert len(upstream) == 1, upstream
_FILES: Final = 3
_READ_FILE_FUNCTION: Final[dict[str, JsonValue]] = {
"name": "read_file",
"parameters": {"type": "object", "properties": {"path": {"type": "string"}}},
}
_READ_FILE_TOOL: Final[dict[str, JsonValue]] = {
"name": "read_file",
"input_schema": _READ_FILE_FUNCTION["parameters"],
}
def _read_file_round(marker: str, index: int) -> tuple[JsonValue, ...]:
return (
{"role": "user", "content": f"Read file {index}."},
{
"role": "assistant",
"content": [
{"type": "tool_use", "id": f"toolu_{index}", "name": "read_file", "input": {"path": f"{index}.txt"}}
],
},
{
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": f"toolu_{index}",
"content": f"<file {marker}-{index}>short file body.",
}
],
},
)
def _agent_body(group: str, marker: str) -> dict[str, JsonValue]:
rounds: Final = tuple(chain.from_iterable(_read_file_round(marker, index) for index in range(_FILES)))
return {
"model": group,
"max_tokens": 16,
"tools": [_READ_FILE_TOOL],
"messages": [*rounds, {"role": "user", "content": "Summarize what you read."}],
"context_management": {
"edits": [
{
"type": "clear_tool_uses_20250919",
"trigger": {"type": "input_tokens", "value": _deepseek_v4_tool_tokens(_READ_FILE_FUNCTION)},
"keep": {"type": "tool_uses", "value": 1},
}
]
},
}
@pytest.mark.parametrize(
("model", "files_kept"),
((_DEEPSEEK_V4, (False, False, True)), (_GROK, (True, True, True))),
ids=("deepseek-v4", "grok-4-1"),
)
def test_messages_context_management_trigger_counts_tools_the_way_the_models_chat_template_renders_them(
gateway: Gateway, model: str, files_kept: tuple[bool, ...]
) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(_completion) as wire, gateway.scenario() as scenario:
group: Final = _deployment(scenario, "context", model, wire.url)
response: Final = gateway.request("POST", "/v1/messages", _agent_body(group, marker))
upstream: Final = wire.drain()
assert response.status_code == 200, response.text
assert len(upstream) == 1, upstream
sent: Final = upstream[0].body.decode()
assert tuple(f"<file {marker}-{index}>" in sent for index in range(_FILES)) == files_kept

25
tests/integration/providers/test_openai_count_tokens_tools_wire.py
View file

@ -391,28 +391,3 @@ def test_google_count_tokens_sends_the_system_instruction_text(counting_proxy: O
)
assert response.status_code == 200, response.text
assert _count_body(wire)["instructions"] == "Be brief.\nAnswer in English."
@pytest.mark.parametrize(
"error",
[
pytest.param(
{"message": "Invalid input", "type": "invalid_request_error", "param": "input", "code": "invalid_value"},
id="full",
),
pytest.param({"message": "Invalid input", "type": "invalid_request_error"}, id="bare"),
],
)
def test_responses_input_tokens_passes_an_openai_count_error_through(
counting_proxy: OwnedProxy, error: dict[str, JsonValue]
) -> None:
def rejecting(_request: Request) -> Reply:
return Reply(status=400, body=json.dumps({"error": error}).encode())
gateway: Final = counting_proxy.gateway
with wire_server(rejecting) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire.url)
response: Final = gateway.request("POST", "/v1/responses/input_tokens", {"model": model, "input": "hello"})
assert response.status_code == 400, response.text
assert response.json() == {"error": {"param": None, "code": None, **error}}, response.text
assert _count_body(wire)["input"] == [{"role": "user", "content": "hello"}]

114
tests/integration/providers/test_vertex_gemini_count_tokens_wire.py
View file

@ -65,7 +65,7 @@ class _GeminiPeer:
return requests
def _deployment(gateway: Gateway, scenario: Scenario, api_base: str, **parameters: JsonValue) -> str:
def _deployment(gateway: Gateway, scenario: Scenario, api_base: str) -> str:
return scenario.model(
model=f"vertex_ai/{_MODEL}",
api_base=api_base,
@ -73,7 +73,6 @@ def _deployment(gateway: Gateway, scenario: Scenario, api_base: str, **parameter
vertex_project=_PROJECT,
vertex_location=_LOCATION,
vertex_credentials=service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")),
**parameters,
)
@ -221,8 +220,8 @@ _CONTENT_BLOCK_BODIES: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingPro
)
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 _ai_studio_deployment(scenario: Scenario, api_base: str) -> str:
return scenario.model(model=f"gemini/{_MODEL}", api_base=api_base, api_key="scripted-gemini-key")
def _counted_contents(count_body: Mapping[str, JsonValue]) -> JsonValue:
@ -595,110 +594,3 @@ def test_concurrent_count_waves_recover_after_vertex_count_outage(gateway: Gatew
tuple(_json_body(request) for request in requests)
== ({"contents": [{"role": "user", "parts": [{"text": _USER}]}]},) * 9
)
_CONTENTS: Final[JsonValue] = [{"role": "user", "parts": [{"text": _USER}]}]
_OPENAI_TOOLS: Final[JsonValue] = [
{
"type": "function",
"function": {"name": "noop", "description": "Return no output", "parameters": {"type": "object"}},
}
]
def _provider_deployment(
gateway: Gateway, scenario: Scenario, provider: str, api_base: str, **parameters: JsonValue
) -> str:
if provider == "vertex_ai":
return _deployment(gateway, scenario, api_base, **parameters)
return _ai_studio_deployment(scenario, api_base, **parameters)
def _count(gateway: Gateway, body: Mapping[str, JsonValue], call_endpoint: bool = True) -> httpx.Response:
return gateway.request(
"POST", "/utils/token_counter", body, params={"call_endpoint": "true" if call_endpoint else "false"}
)
def _total_tokens(response: httpx.Response) -> JsonValue:
return _JSON_OBJECT.validate_json(response.content)["total_tokens"]
@pytest.mark.parametrize(
("provider", "empty_conversation"),
(("vertex_ai", {}), ("gemini", {}), ("gemini", {"messages": []}), ("gemini", {"contents": []})),
ids=("vertex-prompt-only", "ai-studio-prompt-only", "ai-studio-empty-messages", "ai-studio-empty-contents"),
)
def test_utils_token_counter_sends_a_prompt_as_the_user_turn_of_its_messages_form(
gateway: Gateway, provider: str, empty_conversation: Mapping[str, JsonValue]
) -> None:
peer: Final = _GeminiPeer()
with wire_server(peer) as wire, gateway.scenario() as scenario:
model: Final = _provider_deployment(gateway, scenario, provider, wire.url)
prompted: Final = _count(gateway, {"model": model, "prompt": _USER, **empty_conversation})
prompted_requests: Final = _count_requests(wire)
as_message: Final = _count(gateway, {"model": model, "messages": _MESSAGES})
message_requests: Final = _count_requests(wire)
assert prompted.status_code == 200, prompted.text
assert as_message.status_code == 200, as_message.text
assert _total_tokens(prompted) == _COUNT, prompted.text
assert len(prompted_requests) == 1
assert len(message_requests) == 1
assert _counted_contents(_json_body(prompted_requests[0])) == _counted_contents(_json_body(message_requests[0]))
@pytest.mark.parametrize(
("conversation", "call_endpoint"),
(({"messages": _MESSAGES}, True), ({"contents": _CONTENTS}, True), ({"messages": _MESSAGES}, False)),
ids=("provider-count-messages", "provider-count-contents", "local-count-messages"),
)
def test_utils_token_counter_rejects_a_prompt_next_to_a_conversation_without_calling_the_provider(
gateway: Gateway, conversation: Mapping[str, JsonValue], call_endpoint: bool
) -> None:
peer: Final = _GeminiPeer()
with wire_server(peer) as wire, gateway.scenario() as scenario:
model: Final = _ai_studio_deployment(scenario, wire.url)
response: Final = _count(gateway, {"model": model, "prompt": "hi", **conversation}, call_endpoint)
requests: Final = _count_requests(wire)
assert response.status_code == 400, response.text
assert _JSON_OBJECT.validate_python(_JSON_OBJECT.validate_json(response.content)["error"])["param"] == "prompt"
assert requests == ()
@pytest.mark.parametrize("call_endpoint", (True, False), ids=("provider-count", "local-count"))
def test_utils_token_counter_counts_only_the_conversation_when_drop_params_drops_the_prompt(
gateway: Gateway, call_endpoint: bool
) -> None:
peer: Final = _GeminiPeer()
with wire_server(peer) as wire, gateway.scenario() as scenario:
model: Final = _ai_studio_deployment(scenario, wire.url, drop_params=True)
mixed: Final = _count(gateway, {"model": model, "prompt": "hi", "messages": _MESSAGES}, call_endpoint)
mixed_requests: Final = _count_requests(wire)
alone: Final = _count(gateway, {"model": model, "messages": _MESSAGES}, call_endpoint)
alone_requests: Final = _count_requests(wire)
assert mixed.status_code == 200, mixed.text
assert alone.status_code == 200, alone.text
assert _total_tokens(mixed) == _total_tokens(alone)
assert [_counted_contents(_json_body(request)) for request in mixed_requests] == [
_counted_contents(_json_body(request)) for request in alone_requests
]
@pytest.mark.parametrize(
"extra",
({"system": "Follow the system instruction."}, {"tools": _OPENAI_TOOLS}),
ids=("system", "tools"),
)
def test_utils_token_counter_counts_a_local_prompt_with_system_or_tools_like_its_messages_form(
gateway: Gateway, extra: Mapping[str, JsonValue]
) -> None:
peer: Final = _GeminiPeer()
with wire_server(peer) as wire, gateway.scenario() as scenario:
model: Final = _ai_studio_deployment(scenario, wire.url)
prompted: Final = _count(gateway, {"model": model, "prompt": _USER, **extra}, call_endpoint=False)
as_message: Final = _count(gateway, {"model": model, "messages": _MESSAGES, **extra}, call_endpoint=False)
requests: Final = _count_requests(wire)
assert prompted.status_code == 200, prompted.text
assert as_message.status_code == 200, as_message.text
assert _total_tokens(prompted) == _total_tokens(as_message)
assert requests == ()

36
tests/integration/providers/test_vertex_partner_count_tokens_wire.py
View file

@ -468,46 +468,10 @@ def test_responses_input_tokens_surfaces_the_peer_error(gateway: Gateway) -> Non
assert sent["tools"] == _RESPONSES_TOOLS, sent
def test_responses_input_tokens_types_a_peer_error_by_its_status(gateway: Gateway) -> None:
with wire_server(_peer(_rejecting(404, "not_found_error"))) as wire, gateway.scenario() as scenario:
model: Final = _deployment(gateway, scenario, wire.url)
response: Final = _RESPONSES_ROUTE.send(gateway, model)
assert response.status_code == 404, response.text
assert response.json() == _RESPONSES_ROUTE.rejection(_REJECTION), response.text
assert _counted_bodies(wire) == (_PEER_BARE,)
@pytest.mark.parametrize(
("backend", "route"),
[
pytest.param(_GLM_BACKEND, _MESSAGES_ROUTE, id="glm-messages"),
pytest.param(_GLM_BACKEND, _RESPONSES_ROUTE, id="glm-responses"),
pytest.param(_GLM_BACKEND, _GOOGLE_ROUTE, id="glm-google"),
pytest.param(_GLM_BACKEND, _UTILS_ROUTE, id="glm-utils"),
pytest.param(_MISTRAL_BACKEND, _MESSAGES_ROUTE, id="mistral-messages"),
],
)
def test_partner_without_a_count_route_answers_400_without_calling_vertex(
gateway: Gateway, backend: str, route: _CountRoute
) -> None:
with wire_server(_peer()) as wire, gateway.scenario() as scenario:
model: Final = _deployment(gateway, scenario, wire.url, model=f"vertex_ai/{backend}")
response: Final = route.send(gateway, model)
assert response.status_code == 400, response.text
payload: Final = _JSON_OBJECT.validate_json(response.content)
message: Final = _error_message(payload)
assert backend in message, message
assert payload == route.rejection(message), response.text
assert wire.drain() == ()
def test_llama_count_still_reaches_the_meta_count_route(gateway: Gateway) -> None:
with wire_server(_peer()) as wire, gateway.scenario() as scenario:
model: Final = _deployment(gateway, scenario, wire.url, model=f"vertex_ai/{_LLAMA_BACKEND}")
response: Final = _count(gateway, _bare(model))
assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text
(sent,) = _count_bodies(wire.drain(), _LLAMA_COUNT_TARGET)
assert sent["messages"] == _MESSAGES, sent
@pytest.mark.timeout(2 * graceful_stop_seconds() + 120)

7
tests/integration/proxy_config.yaml
View file

@ -1,4 +1,11 @@
model_list:
- model_name: anthropic/claude-opus-5-5
litellm_params:
model: anthropic/claude-opus-5-5
api_base: http://127.0.0.1:8191
api_key: synthetic-anthropic-key
model_info:
token_count_framing: 0
- model_name: perplexity/pplx-decider-v1-27b
litellm_params:
model: perplexity/pplx-decider-v1-27b

78
tests/test_litellm/proxy/google_endpoints/test_endpoints.py
View file

@ -2,84 +2,14 @@
Test for google_endpoints/endpoints.py
"""
from typing import Final
import pytest
import sys, os
from dotenv import load_dotenv
from pydantic import JsonValue
from starlette.requests import Request
from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request, google_count_tokens
from litellm.proxy.google_endpoints.endpoints import google_count_tokens
from litellm.types.llms.vertex_ai import TokenCountDetailsResponse
@pytest.mark.parametrize(
("system_key", "schema_key"),
(("systemInstruction", "parameters"), ("system_instruction", "parametersJsonSchema")),
)
def test_google_count_tokens_translates_system_and_function_declarations(
system_key: str, schema_key: str
) -> None:
schema: Final = {"type": "object", "properties": {"query": {"type": "string"}}}
request_body: Final[dict[str, JsonValue]] = {
"contents": [{"role": "user", "parts": [{"text": "Look this up"}]}],
system_key: {"parts": [{"text": "Use the provided tool."}]},
"tools": [
{
"functionDeclarations": [
{
"name": "lookup",
"description": "Find a record",
schema_key: schema,
}
]
}
],
}
request: Final = _google_count_tokens_request("claude-sonnet", request_body)
assert request.system == [{"type": "text", "text": "Use the provided tool."}]
assert request.tools == [
{
"name": "lookup",
"description": "Find a record",
"input_schema": schema,
}
]
assert request.messages == [{"role": "user", "content": "Look this up"}]
@pytest.mark.parametrize(
("schema_key", "expected_schema"),
(
("parameters", {"type": "object", "properties": {"query": {"type": "string"}}}),
("parametersJsonSchema", {"type": "OBJECT", "properties": {"query": {"type": "STRING"}}}),
),
)
def test_google_count_tokens_normalizes_parameters_only(
schema_key: str, expected_schema: dict[str, JsonValue]
) -> None:
request: Final = _google_count_tokens_request(
"claude-sonnet",
{
"contents": [{"role": "user", "parts": [{"text": "Look this up"}]}],
"tools": [
{
"functionDeclarations": [
{
"name": "lookup",
schema_key: {
"type": "OBJECT",
"properties": {"query": {"type": "STRING"}},
},
}
]
}
],
},
)
assert request.tools == [{"name": "lookup", "input_schema": expected_schema}]
from starlette.requests import Request
load_dotenv()

7
tests/test_litellm/proxy/proxy_server/test_routes_utils.py
View file

@ -383,7 +383,8 @@ def test_token_counter_fallback_counts_tools_system_and_anthropic_blocks(client,
assert full > bare
def test_token_counter_fallback_counts_a_prompt_with_tools_like_its_messages_form(client, auth_as, monkeypatch):
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)."""
monkeypatch.setattr(proxy_server, "llm_router", None)
monkeypatch.setattr(litellm, "disable_token_counter", False, raising=False)
prompt = "count the tokens in this sentence please"
@ -404,6 +405,4 @@ def test_token_counter_fallback_counts_a_prompt_with_tools_like_its_messages_for
)
assert response.status_code == 200, response.text
assert response.json()["total_tokens"] == litellm.token_counter(
model="claude-fable-5", messages=[{"role": "user", "content": prompt}], tools=tools
)
assert response.json()["total_tokens"] == litellm.token_counter(model="claude-fable-5", text=prompt)

214
tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py
View file

@ -3,7 +3,6 @@ Test for response_api_endpoints/endpoints.py
"""
import unittest
from collections.abc import Mapping
from typing import Any, Final, Literal
from unittest.mock import AsyncMock, MagicMock, patch
@ -2429,219 +2428,6 @@ class TestResponsesInputTokens:
assert response.status_code == 429, response.text
assert response.json()["error"]["message"] == "rate limited"
@pytest.mark.parametrize(
"provider_body, status_code, expected_error",
[
pytest.param(
{
"error": {
"message": "Missing required parameter: 'tools[0].name'.",
"type": "invalid_request_error",
"param": "tools[0].name",
"code": "missing_required_parameter",
}
},
400,
{
"message": "Missing required parameter: 'tools[0].name'.",
"type": "invalid_request_error",
"param": "tools[0].name",
"code": "missing_required_parameter",
},
id="openai_error_passes_through",
),
pytest.param(
{
"error": {
"message": "Missing required parameter: 'tools[0].name'.",
"type": "invalid_request_error",
"param": "tools[0].name",
}
},
400,
{
"message": "Missing required parameter: 'tools[0].name'.",
"type": "invalid_request_error",
"param": "tools[0].name",
"code": None,
},
id="openai_error_without_code_keeps_type_and_param",
),
pytest.param(
{
"type": "error",
"error": {"type": "invalid_request_error", "message": "tools.0: Input tag 'function' is invalid"},
"request_id": "req_vrtx_1",
},
400,
{
"message": "tools.0: Input tag 'function' is invalid",
"type": "invalid_request_error",
"param": None,
"code": None,
},
id="anthropic_error_message_unwrapped",
),
pytest.param(
{
"type": "error",
"error": {"type": "not_found_error", "message": "model: claude-typo"},
"request_id": "req_vrtx_2",
},
404,
{"message": "model: claude-typo", "type": "invalid_request_error", "param": None, "code": None},
id="anthropic_error_type_follows_the_status",
),
pytest.param(
{"error": {"code": 404, "message": "Publisher model not found.", "status": "NOT_FOUND"}},
404,
{"message": "Publisher model not found.", "type": "invalid_request_error", "param": None, "code": None},
id="google_error_message_unwrapped",
),
pytest.param(
"litellm.APIError: Google Gen AI Studio API error: 404 - "
+ json.dumps(
{"error": {"code": 404, "message": "models/gemini-typo is not found.", "status": "NOT_FOUND"}}
),
404,
{
"message": "models/gemini-typo is not found.",
"type": "invalid_request_error",
"param": None,
"code": None,
},
id="prefixed_gemini_error_message_unwrapped",
),
pytest.param(
"litellm.BadRequestError: AzureException - "
+ json.dumps(
{
"error": {
"message": "Unrecognized request argument supplied: foo",
"type": "invalid_request_error",
"param": "foo",
"code": "unknown_parameter",
}
}
),
400,
{
"message": "Unrecognized request argument supplied: foo",
"type": "invalid_request_error",
"param": "foo",
"code": "unknown_parameter",
},
id="prefixed_openai_error_passes_through",
),
pytest.param(
{"message": "The provided model doesn't support counting tokens."},
400,
{
"message": "The provided model doesn't support counting tokens.",
"type": "invalid_request_error",
"param": None,
"code": None,
},
id="bedrock_error_message_unwrapped",
),
pytest.param(
"BedrockException - " + json.dumps({"message": "The provided model doesn't support counting tokens."}),
400,
{
"message": "The provided model doesn't support counting tokens.",
"type": "invalid_request_error",
"param": None,
"code": None,
},
id="prefixed_bedrock_error_message_unwrapped",
),
pytest.param(
"tool schema {unclosed is invalid",
400,
{
"message": "tool schema {unclosed is invalid",
"type": "invalid_request_error",
"param": None,
"code": None,
},
id="plain_text_with_a_brace_kept_whole",
),
pytest.param(
"upstream count service unavailable",
503,
{
"message": "upstream count service unavailable",
"type": "internal_server_error",
"param": None,
"code": None,
},
id="plain_text_error",
),
],
)
def test_provider_count_error_answers_in_openai_error_shape(
self,
provider_body: str | Mapping[str, object],
status_code: int,
expected_error: Mapping[str, str | None],
) -> None:
from litellm.proxy._types import ProxyException
failing_counter: Final = AsyncMock(
side_effect=ProxyException(
message=provider_body if isinstance(provider_body, str) else json.dumps(provider_body),
type="token_counting_error",
param="model",
code=str(status_code),
)
)
response, _ = self._post_input_tokens({"model": "gpt-4o", "input": "hi"}, counter=failing_counter)
assert response.status_code == status_code, response.text
assert response.json() == {"error": expected_error}
def test_non_provider_count_error_keeps_its_own_type(self) -> None:
from litellm.proxy._types import ProxyException
failing_counter: Final = AsyncMock(
side_effect=ProxyException(
message="Token counting is disabled and no provider API result available",
type="token_counting_disabled",
param="model",
code=503,
)
)
response, _ = self._post_input_tokens({"model": "gpt-4o", "input": "hi"}, counter=failing_counter)
assert response.status_code == 503, response.text
assert response.json()["error"]["type"] == "token_counting_disabled"
def test_provider_count_error_names_the_call_id_when_opted_in(self) -> None:
from litellm.proxy._types import ProxyException
failing_counter: Final = AsyncMock(
side_effect=ProxyException(
message="LiteLLM does not support counting tokens for Vertex AI model zai-org/glm-5.2-maas.",
type="token_counting_error",
param="model",
code="400",
headers={"x-litellm-call-id": "call-9564"},
)
)
with patch("litellm.proxy.proxy_server.general_settings", {"include_call_id_in_error_body": True}):
response, _ = self._post_input_tokens({"model": "gpt-4o", "input": "hi"}, counter=failing_counter)
assert response.status_code == 400, response.text
assert response.json() == {
"error": {
"message": "LiteLLM does not support counting tokens for Vertex AI model zai-org/glm-5.2-maas.",
"type": "invalid_request_error",
"param": None,
"code": None,
"litellm_call_id": "call-9564",
}
}
def test_responses_routes_document_response_models_in_openapi_schema():
from typing import cast

31
tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py
View file

@ -1,31 +0,0 @@
from typing import Final
import pytest
from litellm.litellm_core_utils.prompt_templates.chat_template_tools import chat_template_tool_section
_DESCRIBED_FUNCTION: Final = {"name": "f", "description": "é", "parameters": {}}
_BARE_FUNCTION: Final = {"name": "g", "parameters": {}}
# Tool schemas as rendered by render_tools in encoding/encoding_dsv4.py (deepseek-ai/DeepSeek-V4-Flash) and by
# chat_template.jinja (Qwen/Qwen3.5-122B-A10B), both read on Hugging Face on 2026-10-17
@pytest.mark.parametrize(
("model", "rendered"),
[
(
"DeepSeek-V4-Flash-0731",
'{"name": "f", "description": "é", "parameters": {}}\n{"name": "g", "parameters": {}}',
),
(
"databricks-qwen35-122b-a10b",
'\n{"type": "function", "function": {"name": "f", "description": "é", "parameters": {}}}'
'\n{"type": "function", "function": {"name": "g", "parameters": {}}}',
),
],
)
def test_chat_template_tool_section_renders_schemas_like_the_models_template(model: str, rendered: str):
section: Final = chat_template_tool_section(model)
assert section is not None
assert section.render_tool_schemas((_DESCRIBED_FUNCTION, _BARE_FUNCTION)) == rendered

21
tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py
View file

@ -11,7 +11,6 @@ import litellm
from litellm.litellm_core_utils.prompt_templates.common_utils import encrypted_reasoning_signature
from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import (
CONVERTED_SYSTEM_NOTE,
anthropic_system_blocks,
place_mid_conversation_system,
split_leading_system_run,
)
@ -391,23 +390,3 @@ def test_flagged_placement_keeps_a_system_before_an_assistant_turn_whose_empty_t
)
assert _roles(placed) == ["user", "system", "assistant", "user"]
def test_anthropic_system_blocks_keeps_text_parts_with_their_cache_control_and_drops_the_rest():
run = [
{"role": "system", "content": "one", "cache_control": {"type": "ephemeral", "ttl": "1h"}},
{
"role": "system",
"content": [
{"type": "text", "text": "two", "cache_control": {"type": "ephemeral"}},
{"type": "text", "text": ""},
{"type": "image_url", "image_url": {"url": "data:image/png;base64,aW1hZ2U="}},
],
},
{"role": "system", "content": ""},
]
assert anthropic_system_blocks(run) == (
{"type": "text", "text": "one", "cache_control": {"type": "ephemeral", "ttl": "1h"}},
{"type": "text", "text": "two", "cache_control": {"type": "ephemeral"}},
)

253
tests/unit/litellm_core_utils/test_token_counter.py
View file

@ -9,9 +9,8 @@ import subprocess
import sys
import threading
import time
from collections.abc import Callable, Mapping
from collections.abc import Mapping
from concurrent.futures import Future, wait
from copy import deepcopy
from itertools import accumulate, chain
from pathlib import Path
from typing import Final
@ -40,7 +39,6 @@ from litellm.litellm_core_utils.token_counter import (
offload_token_count,
)
from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new
from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter
from tests.large_text import text
from tests.unit.litellm_core_utils.event_loop_lag import (
assert_loop_stayed_free,
@ -445,255 +443,6 @@ def test_token_counter_with_tools(message_count_pair):
), f"Expected {expected_tokens} tokens, got {counted_tokens}."
def test_token_counter_counts_gemini_function_declarations():
openai_tools: Final = [
{
"type": "function",
"function": {
"name": "lookup_weather",
"description": "Find current weather conditions for a location",
"parameters": {
"type": "object",
"properties": {
"location": {"type": "string", "description": "City and region"},
"units": {"type": "string", "enum": ["celsius", "fahrenheit"]},
},
"required": ["location"],
},
},
}
]
gemini_tools: Final = litellm.utils.get_optional_params(
model="gemini-2.5-pro",
custom_llm_provider="gemini",
tools=openai_tools,
)["tools"]
camel_case_tools: Final = [{"functionDeclarations": gemini_tools[0]["function_declarations"]}]
openai_tokens: Final = token_counter_new(
model="gemini-2.5-pro",
messages=[{"role": "user", "content": "What's the weather?"}],
tools=openai_tools,
)
gemini_tokens: Final = token_counter_new(
model="gemini-2.5-pro",
messages=[{"role": "user", "content": "What's the weather?"}],
tools=gemini_tools,
)
camel_case_tokens: Final = token_counter_new(
model="gemini-2.5-pro",
messages=[{"role": "user", "content": "What's the weather?"}],
tools=camel_case_tools,
)
assert openai_tokens == gemini_tokens == camel_case_tokens
def test_token_counter_skips_non_mapping_tools():
openai_tool: Final = {
"type": "function",
"function": {
"name": "lookup_weather",
"description": "Find current weather conditions for a location",
"parameters": {
"type": "object",
"properties": {"location": {"type": "string", "description": "City and region"}},
"required": ["location"],
},
},
}
messages: Final = [{"role": "user", "content": "What's the weather?"}]
valid_tokens: Final = token_counter_new(
model="gemini-2.5-pro",
messages=messages,
tools=[openai_tool],
)
mixed_tokens: Final = token_counter_new(
model="gemini-2.5-pro",
messages=messages,
tools=["bad", None, openai_tool],
)
assert mixed_tokens == valid_tokens
_WEATHER_FUNCTION: Final = {
"name": "get_weather",
"description": "Get the current weather for a city",
"parameters": {
"type": "object",
"properties": {"location": {"type": "string", "description": "City name"}},
"required": ["location"],
},
}
_TIME_FUNCTION: Final = {
"name": "get_time",
"description": "Get the local time for a city",
"parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
}
_WEB_SEARCH_FUNCTION: Final = {
"name": "web_search",
"description": "Search the web",
"parameters": {"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]},
}
def _openai_tool(function: Mapping[str, object]) -> dict[str, object]:
return {"type": "function", "function": function}
def _anthropic_tool(function: Mapping[str, object]) -> dict[str, object]:
return {"name": function["name"], "description": function["description"], "input_schema": function["parameters"]}
_USER_MESSAGES: Final = ({"role": "user", "content": "What is the weather in Paris today?"},)
_SYSTEM_AND_USER_MESSAGES: Final = ({"role": "system", "content": "You are a helpful assistant."}, *_USER_MESSAGES)
@pytest.mark.parametrize("tool_shape", [_openai_tool, _anthropic_tool])
@pytest.mark.parametrize(
("model", "messages", "functions", "billed_prompt_tokens"),
[
# usage.prompt_tokens billed for this request by Azure AI Foundry DeepSeek-V4-Flash-0731,
# Databricks databricks-qwen35-122b-a10b and DashScope qwen3.5-plus on 2026-10-17
("DeepSeek-V4-Flash-0731", _USER_MESSAGES, (_WEATHER_FUNCTION,), 295),
("DeepSeek-V4-Flash-0731", _USER_MESSAGES, (_WEATHER_FUNCTION, _TIME_FUNCTION), 348),
("DeepSeek-V4-Flash-0731", _SYSTEM_AND_USER_MESSAGES, (_WEATHER_FUNCTION,), 300),
("databricks-qwen35-122b-a10b", _USER_MESSAGES, (_WEATHER_FUNCTION,), 289),
("databricks-qwen35-122b-a10b", _USER_MESSAGES, (_WEATHER_FUNCTION, _TIME_FUNCTION), 352),
("databricks-qwen35-122b-a10b", _SYSTEM_AND_USER_MESSAGES, (_WEATHER_FUNCTION,), 296),
("qwen3.5-plus", _USER_MESSAGES, (_WEATHER_FUNCTION,), 284),
("qwen3.5-plus", _USER_MESSAGES, (_WEATHER_FUNCTION, _TIME_FUNCTION), 342),
],
)
def test_token_counter_tools_match_billed_prompt_for_json_tool_chat_templates(
tool_shape: Callable[[Mapping[str, object]], dict[str, object]],
model: str,
messages: tuple[dict[str, str], ...],
functions: tuple[Mapping[str, object], ...],
billed_prompt_tokens: int,
):
counted: Final = token_counter_new(
model=model,
messages=list(messages),
tools=[tool_shape(function) for function in functions],
)
assert abs(counted - billed_prompt_tokens) <= billed_prompt_tokens * 0.05
@pytest.mark.parametrize(
("model", "counted_like"),
[
("DeepSeek-V4-Pro", "DeepSeek-V4-Flash-0731"),
("accounts/fireworks/models/deepseek-v4-flash", "DeepSeek-V4-Flash-0731"),
("DeepSeek-V4.1-Flash", "deepseek-v3.2"),
("accounts/fireworks/models/deepseek-v4p1-flash", "deepseek-v3.2"),
("Qwen/Qwen3.5-122B-A10B", "databricks-qwen35-122b-a10b"),
("qwen3-235b-a22b", "deepseek-v3.2"),
],
)
def test_token_counter_counts_tools_by_the_models_chat_template_family(model: str, counted_like: str):
tools: Final = [_openai_tool(_WEATHER_FUNCTION), _openai_tool(_TIME_FUNCTION)]
assert token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=tools) == token_counter_new(
model=counted_like, messages=list(_USER_MESSAGES), tools=tools
)
@pytest.mark.parametrize("model", ["DeepSeek-V4-Flash-0731", "databricks-qwen35-122b-a10b"])
@pytest.mark.parametrize(
"builtin_tool",
[
{"type": "web_search_preview"},
{"type": "web_search_20250305", "name": "web_search", "max_uses": 5},
{"name": "web_search"},
],
)
def test_token_counter_skips_tools_not_sent_as_functions_for_chat_template_families(
model: str, builtin_tool: dict[str, object]
):
weather_tool: Final = _openai_tool(_WEATHER_FUNCTION)
assert token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=[builtin_tool]) == token_counter_new(
model=model, messages=list(_USER_MESSAGES)
)
assert token_counter_new(
model=model, messages=list(_USER_MESSAGES), tools=[builtin_tool, weather_tool]
) == token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=[weather_tool])
def _responses_tool(function: Mapping[str, object]) -> dict[str, object]:
return {"type": "function", **function}
def _gemini_tool(function: Mapping[str, object]) -> dict[str, object]:
return {"function_declarations": [function]}
@pytest.mark.parametrize("model", ["DeepSeek-V4-Flash-0731", "databricks-qwen35-122b-a10b"])
@pytest.mark.parametrize("function", [_WEATHER_FUNCTION, _WEB_SEARCH_FUNCTION])
@pytest.mark.parametrize("tool_shape", [_responses_tool, _gemini_tool])
def test_token_counter_counts_function_tools_like_chat_completions_tools_for_chat_template_families(
model: str, function: Mapping[str, object], tool_shape: Callable[[Mapping[str, object]], dict[str, object]]
):
assert token_counter_new(
model=model, messages=list(_USER_MESSAGES), tools=[tool_shape(function)]
) == token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=[_openai_tool(function)])
# Azure AI Foundry DeepSeek-V4-Flash-0731 billed 295 prompt tokens for the weather tool with each of
# these tool_choice values on 2026-10-17, the same as without one
@pytest.mark.parametrize("tool_choice", ["none", "required", {"type": "function", "function": {"name": "get_weather"}}])
def test_token_counter_tool_choice_adds_nothing_for_deepseek_v4(tool_choice: object):
tools: Final = [_openai_tool(_WEATHER_FUNCTION)]
assert token_counter_new(
model="DeepSeek-V4-Flash-0731", messages=list(_USER_MESSAGES), tools=tools, tool_choice=tool_choice
) == token_counter_new(model="DeepSeek-V4-Flash-0731", messages=list(_USER_MESSAGES), tools=tools)
@pytest.mark.parametrize("model", ["DeepSeek-V4-Flash-0731", "databricks-qwen35-122b-a10b"])
@pytest.mark.parametrize(
"anthropic_tool",
[
{"name": "get_weather", "input_schema": _WEATHER_FUNCTION["parameters"]},
_anthropic_tool(_WEATHER_FUNCTION),
{**_anthropic_tool(_WEATHER_FUNCTION), "strict": True},
{**_anthropic_tool(_WEATHER_FUNCTION), "defer_loading": True, "input_examples": [{"location": "Paris"}]},
_anthropic_tool(_WEB_SEARCH_FUNCTION),
],
)
def test_token_counter_counts_anthropic_tools_as_sent_to_chat_template_families(
model: str, anthropic_tool: dict[str, object]
):
sent_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai([deepcopy(anthropic_tool)]) # pyright: ignore[reportArgumentType] # raw Anthropic tool dict
assert token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=[anthropic_tool]) == token_counter_new(
model=model, messages=list(_USER_MESSAGES), tools=sent_tools
)
# Qwen3.5's chat_template.jinja opens a system turn for the tool schemas when the request has none, 5 tokens that a
# request with a system message already pays, while DeepSeek V4's encoding_dsv4.py appends them to the first turn
# either way (both read on Hugging Face on 2026-10-17)
@pytest.mark.parametrize(
("model", "implicit_system_turn_tokens"), [("DeepSeek-V4-Flash-0731", 0), ("databricks-qwen35-122b-a10b", 5)]
)
def test_token_counter_counts_the_tools_system_turn_only_without_a_system_message(
model: str, implicit_system_turn_tokens: int
):
tools: Final = [_openai_tool(_WEATHER_FUNCTION)]
system_cost_with_tools: Final = token_counter_new(
model=model, messages=list(_SYSTEM_AND_USER_MESSAGES), tools=tools
) - token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=tools)
system_cost_without_tools: Final = token_counter_new(
model=model, messages=list(_SYSTEM_AND_USER_MESSAGES)
) - token_counter_new(model=model, messages=list(_USER_MESSAGES))
assert system_cost_with_tools == system_cost_without_tools - implicit_system_turn_tokens
class NeedsToleranceUpdateError(Exception):
"""Custom exception to mark tests that have improved"""

291
tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py
View file

@ -1,113 +1,7 @@
import asyncio
from base64 import b64encode
from copy import deepcopy
from threading import get_ident
from typing import Final
import httpx
import pytest
import respx
from pydantic import JsonValue, TypeAdapter
import litellm
from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler
from litellm.llms.anthropic.count_tokens.transformation import (
AnthropicCountTokensConfig,
)
from litellm.llms.azure_ai.anthropic.count_tokens.handler import (
AzureAIAnthropicCountTokensHandler,
)
from litellm.llms.azure_ai.anthropic.count_tokens.transformation import (
AzureAIAnthropicCountTokensConfig,
)
@pytest.mark.parametrize(
"config_type", (AnthropicCountTokensConfig, AzureAIAnthropicCountTokensConfig)
)
@pytest.mark.parametrize(
("image_url", "source"),
(
("data:image/png;base64,aW1hZ2U=", {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}),
({"url": "data:image/png;base64,aW1hZ2U="}, {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}),
(
{"url": "data:image/png;base64,aW1hZ2U=", "format": "image/jpeg", "detail": "high"},
{"type": "base64", "media_type": "image/jpeg", "data": "aW1hZ2U="},
),
),
)
def test_count_translates_openai_images_without_mutating_input(
config_type: type[AnthropicCountTokensConfig],
image_url: str | dict[str, JsonValue],
source: dict[str, JsonValue],
) -> None:
cache_control: Final[dict[str, JsonValue]] = {"type": "ephemeral"}
messages: Final[list[dict[str, JsonValue]]] = [{
"role": "user", "content": [
{"type": "text", "text": "Count this image"},
{"type": "image_url", "image_url": image_url, "cache_control": cache_control},
],
}]
original: Final = deepcopy(messages)
result: Final = config_type().transform_request_to_count_tokens(
model="claude-opus-5-5", messages=messages
)
assert result == {
"model": "claude-opus-5-5", "messages": [{
"role": "user", "content": [
{"type": "text", "text": "Count this image"},
{"type": "image", "source": source, "cache_control": cache_control},
],
}],
}
assert messages == original
@pytest.mark.parametrize(
"config_type", (AnthropicCountTokensConfig, AzureAIAnthropicCountTokensConfig)
)
def test_count_normalizes_nested_tool_images_and_preserves_native_fields(
config_type: type[AnthropicCountTokensConfig],
) -> None:
openai_image: Final[dict[str, JsonValue]] = {
"type": "image_url", "image_url": {"url": "data:image/png;base64,aW1hZ2U="}
}
native_image: Final[dict[str, JsonValue]] = {
"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}
}
assistant: Final[dict[str, JsonValue]] = {"role": "assistant", "content": [
{"type": "thinking", "thinking": "inspect screenshot", "signature": "fixture-signature"},
{"type": "tool_use", "id": "read-1", "name": "Read", "input": {"content": [openai_image]}},
]}
tool_result: Final[dict[str, JsonValue]] = {
"type": "tool_result", "tool_use_id": "read-1", "is_error": False,
"content": [{"type": "text", "text": "Screenshot"}, native_image, openai_image],
"cache_control": {"type": "ephemeral"},
}
text_result: Final[dict[str, JsonValue]] = {"type": "tool_result", "tool_use_id": "read-2", "content": "done"}
messages: Final[list[dict[str, JsonValue]]] = [
assistant, {"role": "user", "content": [native_image, tool_result, text_result]}
]
tools: Final[list[dict[str, JsonValue]]] = [{
"name": "Read", "input_schema": {"type": "object", "examples": [openai_image]}
}]
system: Final[JsonValue] = [{"type": "text", "text": "policy", "cache_control": {"type": "ephemeral"}}]
options: Final[dict[str, JsonValue]] = {
"thinking": {"type": "adaptive"}, "tool_choice": {"type": "auto"}, "output_config": {"effort": "high"}
}
original: Final = deepcopy((messages, tools, system, options))
result: Final = config_type().transform_request_to_count_tokens(
model="claude-opus-5-5", messages=messages, tools=tools, system=system, optional_params=options
)
assert result == {
"model": "claude-opus-5-5", "system": system, "tools": tools, **options,
"messages": [assistant, {"role": "user", "content": [native_image, {
**tool_result, "content": [{"type": "text", "text": "Screenshot"}, native_image, native_image]
}, text_result]}],
}
assert (messages, tools, system, options) == original
def test_transform_basic_request():
@ -194,188 +88,3 @@ def test_transform_no_system_no_tools():
assert "system" not in result
assert "tools" not in result
@pytest.fixture
def httpx_transport_clients(monkeypatch):
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
client_cache = getattr(litellm, "in_memory_llm_clients_cache", None)
if client_cache is not None:
client_cache.flush_cache()
yield
if client_cache is not None:
client_cache.flush_cache()
@pytest.mark.asyncio
@pytest.mark.usefixtures("httpx_transport_clients")
@pytest.mark.parametrize(
"handler_type", (AnthropicCountTokensHandler, AzureAIAnthropicCountTokensHandler)
)
@pytest.mark.parametrize("scheme", ("http", "https"))
@pytest.mark.parametrize("dict_url", (False, True))
async def test_remote_image_fetch_keeps_counting_handler_event_loop_responsive(
handler_type: type[AnthropicCountTokensHandler] | type[AzureAIAnthropicCountTokensHandler],
scheme: str,
dict_url: bool,
) -> None:
loop: Final = asyncio.get_running_loop()
loop_thread: Final = get_ident()
witness: Final = asyncio.Event()
image_bytes: Final = b"\x89PNG\r\n\x1a\ncount-image"
image_url: Final = f"{scheme}://1.1.1.1/{handler_type.__name__}-{dict_url}.png"
model: Final = "claude-opus-5-5"
api_base: Final = "https://gateway.example/anthropic"
image: Final[dict[str, JsonValue]] = {
"type": "image_url", "image_url": {"url": image_url} if dict_url else image_url,
"cache_control": {"type": "ephemeral"},
}
messages: Final[list[dict[str, JsonValue]]] = [{
"role": "user", "content": [image, {"type": "tool_result", "tool_use_id": "read-1", "content": [image]}]
}]
original: Final = deepcopy(messages)
async def run_witness() -> None:
witness.set()
def image_response(_request: httpx.Request) -> httpx.Response:
assert get_ident() != loop_thread, "image fetch blocked the counting handler's event loop"
asyncio.run_coroutine_threadsafe(run_witness(), loop).result(timeout=5)
return httpx.Response(200, content=image_bytes, headers={"Content-Type": "image/png"})
with respx.mock:
image_route: Final = respx.get(image_url).mock(side_effect=image_response)
count_route: Final = respx.post(url__startswith=api_base).mock(
return_value=httpx.Response(200, json={"input_tokens": 7})
)
handler: Final = handler_type()
result: Final = await handler.handle_count_tokens_request(
model=model, messages=messages, api_base=api_base, api_key="test-key"
)
assert witness.is_set()
assert image_route.call_count == count_route.call_count == 1
assert result == {"input_tokens": 7}
native_image: Final[dict[str, JsonValue]] = {
"type": "image", "source": {
"type": "base64", "media_type": "image/png", "data": b64encode(image_bytes).decode()
}, "cache_control": {"type": "ephemeral"},
}
assert TypeAdapter(dict[str, JsonValue]).validate_json(count_route.calls.last.request.content) == {
"model": model, "messages": [{"role": "user", "content": [
native_image, {"type": "tool_result", "tool_use_id": "read-1", "content": [native_image]}
]}],
}
assert messages == original
@pytest.mark.parametrize(
"config_type", (AnthropicCountTokensConfig, AzureAIAnthropicCountTokensConfig)
)
def test_count_lifts_the_leading_system_run_into_system(
config_type: type[AnthropicCountTokensConfig],
) -> None:
"""A Responses ``instructions`` arrives as a leading system-role message. count_tokens answers 400
on that role at the head of ``messages`` and only takes the initial prompt in ``system``, so the
leading run moves there with its cache_control, empty text dropped, and a later reminder stays."""
cache_control: Final[dict[str, JsonValue]] = {"type": "ephemeral"}
messages: Final[list[dict[str, JsonValue]]] = [
{"role": "system", "content": "Be terse", "cache_control": cache_control},
{"role": "system", "content": [{"type": "text", "text": "Answer in French"}, {"type": "text", "text": ""}]},
{"role": "user", "content": "Hello, how are you?"},
{"role": "system", "content": "later reminder"},
{"role": "assistant", "content": "Bonjour."},
]
original: Final = deepcopy(messages)
result: Final = config_type().transform_request_to_count_tokens(model="claude-opus-5-5", messages=messages)
assert result == {
"model": "claude-opus-5-5",
"system": [
{"type": "text", "text": "Be terse", "cache_control": cache_control},
{"type": "text", "text": "Answer in French"},
],
"messages": [
{"role": "user", "content": "Hello, how are you?"},
{"role": "system", "content": "later reminder"},
{"role": "assistant", "content": "Bonjour."},
],
}
assert messages == original
@pytest.mark.parametrize(
("system", "expected_system"),
(
(None, [{"type": "text", "text": "Be terse"}]),
("", [{"type": "text", "text": "Be terse"}]),
("Answer in French", [{"type": "text", "text": "Answer in French"}, {"type": "text", "text": "Be terse"}]),
(
[{"type": "text", "text": "Answer in French", "cache_control": {"type": "ephemeral"}}],
[
{"type": "text", "text": "Answer in French", "cache_control": {"type": "ephemeral"}},
{"type": "text", "text": "Be terse"},
],
),
),
ids=["absent", "empty", "string", "blocks"],
)
def test_count_keeps_the_callers_system_ahead_of_the_lifted_run(
system: JsonValue, expected_system: list[dict[str, JsonValue]]
) -> None:
result: Final = AnthropicCountTokensConfig().transform_request_to_count_tokens(
model="claude-opus-5-5",
messages=[{"role": "system", "content": "Be terse"}, {"role": "user", "content": "hi"}],
system=system,
)
assert result == {
"model": "claude-opus-5-5",
"system": expected_system,
"messages": [{"role": "user", "content": "hi"}],
}
def test_count_leaves_a_non_text_system_and_its_messages_as_sent() -> None:
"""A malformed ``system`` is the provider's to reject, so nothing is rearranged around it."""
messages: Final[list[dict[str, JsonValue]]] = [
{"role": "system", "content": "Be terse"},
{"role": "user", "content": "hi"},
]
result: Final = AnthropicCountTokensConfig().transform_request_to_count_tokens(
model="claude-opus-5-5", messages=messages, system=5
)
assert result == {"model": "claude-opus-5-5", "system": 5, "messages": messages}
def test_count_drops_a_leading_system_message_without_text() -> None:
result: Final = AnthropicCountTokensConfig().transform_request_to_count_tokens(
model="claude-opus-5-5",
messages=[{"role": "system", "content": ""}, {"role": "user", "content": "hi"}],
)
assert result == {"model": "claude-opus-5-5", "messages": [{"role": "user", "content": "hi"}]}
@pytest.mark.asyncio
async def test_handler_sends_the_leading_system_run_as_system_not_as_a_message(httpx_transport_clients):
"""The wire body is what the provider judges: ``system`` carries the prompt and no message has
``role: "system"``, so a Responses ``instructions`` is counted by Anthropic instead of 400ing."""
with respx.mock:
route = respx.post(url__startswith="https://gateway.example").mock(
return_value=httpx.Response(200, json={"input_tokens": 21})
)
result = await AnthropicCountTokensHandler().handle_count_tokens_request(
model="claude-opus-5-5",
messages=[{"role": "system", "content": "Be terse"}, {"role": "user", "content": "Hello, how are you?"}],
api_key="sk-ant-api03-test-key",
api_base="https://gateway.example",
)
assert result == {"input_tokens": 21}
assert TypeAdapter(dict[str, JsonValue]).validate_json(route.calls.last.request.content) == {
"model": "claude-opus-5-5",
"system": [{"type": "text", "text": "Be terse"}],
"messages": [{"role": "user", "content": "Hello, how are you?"}],
}

109
tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py
View file

@ -1,109 +0,0 @@
import json
from collections.abc import Mapping
from typing import Final
import httpx
import pytest
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
MANTLE_COUNT_URL: Final = "https://bedrock-mantle.us-east-1.api.aws/anthropic/v1/messages/count_tokens"
LITELLM_PARAMS: Final = {
"aws_access_key_id": "AKIATESTACCESSKEY",
"aws_secret_access_key": "test-secret",
"aws_region_name": "us-east-1",
}
REQUEST: Final[dict[str, object]] = {
"model": "global.anthropic.claude-opus-4-8",
"messages": [{"role": "user", "content": "The quick brown fox"}],
"system": "You are a terse assistant.",
"tools": [{"name": "get_weather", "input_schema": {"type": "object", "properties": {}}}],
}
class _MantleEndpoint:
def __init__(self, status_code: int, body: Mapping[str, object]) -> None:
self.status_code: Final = status_code
self.body: Final = body
self.requests: tuple[httpx.Request, ...] = ()
def __call__(self, request: httpx.Request) -> httpx.Response:
self.requests = (*self.requests, request)
return httpx.Response(self.status_code, json=dict(self.body), request=request)
def only_request(self) -> httpx.Request:
assert len(self.requests) == 1, self.requests
return self.requests[0]
def _client(endpoint: _MantleEndpoint) -> AsyncHTTPHandler:
return AsyncHTTPHandler(transport=httpx.MockTransport(endpoint))
@pytest.fixture(autouse=True)
def _sigv4_only(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
@pytest.mark.asyncio
async def test_counts_on_mantle_with_the_base_model_and_the_deployment_credentials() -> None:
mantle: Final = _MantleEndpoint(200, {"input_tokens": 2177})
result: Final = await BedrockMantleCountTokensHandler().handle_count_tokens_request(
request_data=dict(REQUEST),
litellm_params=dict(LITELLM_PARAMS),
resolved_model="anthropic.claude-opus-4-8",
client=_client(mantle),
)
assert result == {"input_tokens": 2177}
posted: Final = mantle.only_request()
assert str(posted.url) == MANTLE_COUNT_URL
assert posted.headers["Authorization"].startswith("AWS4-HMAC-SHA256")
assert "us-east-1/bedrock/aws4_request" in posted.headers["Authorization"]
assert posted.headers["anthropic-version"] == "2023-06-01"
assert json.loads(posted.content) == {
"model": "anthropic.claude-opus-4-8",
"messages": REQUEST["messages"],
"system": REQUEST["system"],
"tools": REQUEST["tools"],
}
@pytest.mark.asyncio
async def test_bedrock_mantle_api_base_env_names_the_host(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", "https://vpce-abc.bedrock-mantle.us-east-1.vpce.api.aws")
mantle: Final = _MantleEndpoint(200, {"input_tokens": 3})
await BedrockMantleCountTokensHandler().handle_count_tokens_request(
request_data=dict(REQUEST),
litellm_params=dict(LITELLM_PARAMS),
resolved_model="anthropic.claude-opus-4-8",
client=_client(mantle),
)
assert (
str(mantle.only_request().url)
== "https://vpce-abc.bedrock-mantle.us-east-1.vpce.api.aws/anthropic/v1/messages/count_tokens"
)
@pytest.mark.asyncio
async def test_non_200_answers_raise_bedrock_error_with_mantle_status() -> None:
mantle: Final = _MantleEndpoint(
404, {"type": "error", "error": {"type": "not_found_error", "message": "does not exist"}}
)
with pytest.raises(BedrockError) as raised:
await BedrockMantleCountTokensHandler().handle_count_tokens_request(
request_data=dict(REQUEST),
litellm_params=dict(LITELLM_PARAMS),
resolved_model="anthropic.claude-sonnet-5-5",
client=_client(mantle),
)
assert raised.value.status_code == 404
assert "does not exist" in raised.value.message

176
tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py
View file

@ -1,176 +0,0 @@
from collections.abc import Mapping
from typing import Final
import httpx
import pytest
from pydantic import JsonValue, TypeAdapter
from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
RUNTIME_HOST: Final = "bedrock-runtime.us-east-1.amazonaws.com"
MANTLE_HOST: Final = "bedrock-mantle.us-east-1.api.aws"
UNSUPPORTED: Final = {"message": "The provided model doesn't support counting tokens."}
DEPLOYMENT: Final = {
"litellm_params": {
"aws_access_key_id": "AKIATESTACCESSKEY",
"aws_secret_access_key": "test-secret",
"aws_region_name": "us-east-1",
}
}
MESSAGES: Final = [{"role": "user", "content": "The quick brown fox jumps over the lazy dog."}]
_JSON_BODY: Final = TypeAdapter(dict[str, JsonValue])
class _Bedrock:
def __init__(self, runtime: tuple[int, Mapping[str, object]], mantle: tuple[int, Mapping[str, object]]) -> None:
self.runtime: Final = runtime
self.mantle: Final = mantle
self.requests: tuple[httpx.Request, ...] = ()
def __call__(self, request: httpx.Request) -> httpx.Response:
self.requests = (*self.requests, request)
status_code, body = self.runtime if request.url.host == RUNTIME_HOST else self.mantle
return httpx.Response(status_code, json=dict(body), request=request)
def posted_hosts(self) -> tuple[str, ...]:
return tuple(request.url.host for request in self.requests)
def _counter(bedrock: _Bedrock) -> BedrockTokenCounter:
return BedrockTokenCounter(client=AsyncHTTPHandler(transport=httpx.MockTransport(bedrock)))
@pytest.fixture(autouse=True)
def _sigv4_only(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
@pytest.mark.asyncio
async def test_claude_model_bedrock_runtime_cannot_count_is_counted_on_mantle() -> None:
bedrock: Final = _Bedrock(runtime=(400, UNSUPPORTED), mantle=(200, {"input_tokens": 2177}))
result: Final = await _counter(bedrock).count_tokens(
model_to_use="global.anthropic.claude-opus-4-8",
messages=MESSAGES,
contents=None,
deployment=DEPLOYMENT,
request_model="claude-opus-4-8",
system="You are a terse assistant.",
)
assert result is not None
assert result.error is False
assert result.total_tokens == 2177
assert result.tokenizer_type == "bedrock_mantle_api"
assert result.original_response == {"input_tokens": 2177}
assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST)
mantle_body: Final = _JSON_BODY.validate_json(bedrock.requests[1].content)
assert mantle_body["model"] == "anthropic.claude-opus-4-8"
assert mantle_body["system"] == "You are a terse assistant."
@pytest.mark.asyncio
async def test_bedrock_runtime_count_is_kept_when_it_answers() -> None:
bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 1353}), mantle=(200, {"input_tokens": 1336}))
result: Final = await _counter(bedrock).count_tokens(
model_to_use="global.anthropic.claude-sonnet-4-6",
messages=MESSAGES,
contents=None,
deployment=DEPLOYMENT,
request_model="claude-sonnet-4-6",
)
assert result is not None
assert result.total_tokens == 1353
assert result.tokenizer_type == "bedrock_api"
assert bedrock.posted_hosts() == (RUNTIME_HOST,)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("model_to_use", "runtime"),
(
("global.anthropic.claude-opus-4-8", (401, {"Message": "not authorized to perform: bedrock:CountTokens"})),
("global.anthropic.claude-opus-4-8", (403, {"Message": "not authorized to perform: bedrock:CountTokens"})),
("amazon.nova-pro-v1:0", (400, UNSUPPORTED)),
("amazon.nova-pro-v1:0", (404, UNSUPPORTED)),
),
)
async def test_other_bedrock_runtime_failures_are_not_retried_on_mantle(
model_to_use: str, runtime: tuple[int, Mapping[str, object]]
) -> None:
bedrock: Final = _Bedrock(runtime=runtime, mantle=(200, {"input_tokens": 2177}))
result: Final = await _counter(bedrock).count_tokens(
model_to_use=model_to_use,
messages=MESSAGES,
contents=None,
deployment=DEPLOYMENT,
request_model=model_to_use,
)
assert result is not None
assert result.error is True
assert result.status_code == runtime[0]
assert result.tokenizer_type == "bedrock_api"
assert bedrock.posted_hosts() == (RUNTIME_HOST,)
@pytest.mark.asyncio
async def test_mantle_failure_is_reported_with_its_status() -> None:
bedrock: Final = _Bedrock(
runtime=(404, UNSUPPORTED),
mantle=(404, {"type": "error", "error": {"type": "not_found_error", "message": "does not exist"}}),
)
result: Final = await _counter(bedrock).count_tokens(
model_to_use="global.anthropic.claude-sonnet-5-5",
messages=MESSAGES,
contents=None,
deployment=DEPLOYMENT,
request_model="claude-sonnet-5-5",
)
assert result is not None
assert result.error is True
assert result.status_code == 404
assert result.tokenizer_type == "bedrock_mantle_api"
assert "does not exist" in (result.error_message or "")
assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST)
@pytest.mark.asyncio
async def test_claude_runtime_404_retries_with_system_and_tools_on_mantle() -> None:
bedrock: Final = _Bedrock(runtime=(404, UNSUPPORTED), mantle=(200, {"input_tokens": 2177}))
tools: Final = [
{
"name": "lookup",
"description": "Find a record",
"input_schema": {"type": "object", "properties": {"query": {"type": "string"}}},
}
]
result: Final = await _counter(bedrock).count_tokens(
model_to_use="global.anthropic.claude-opus-4-8",
messages=MESSAGES,
contents=None,
deployment=DEPLOYMENT,
request_model="claude-opus-4-8",
system="You are a terse assistant.",
tools=tools,
)
assert result is not None
assert result.error is False
assert result.total_tokens == 2177
assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST)
mantle_body: Final = _JSON_BODY.validate_json(bedrock.requests[1].content)
assert mantle_body == {
"model": "anthropic.claude-opus-4-8",
"messages": MESSAGES,
"system": "You are a terse assistant.",
"tools": tools,
}

69
tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py
View file

@ -83,54 +83,6 @@ def test_transform_includes_tools():
assert result["tools"] == tools
_SCHEMA: Final[dict[str, object]] = {"type": "object", "properties": {}} # mutable-ok: matches the tools param type
@pytest.mark.parametrize(
("tool", "expected"),
(
(
{"name": "lookup", "input_schema": _SCHEMA},
{"type": "function", "name": "lookup", "description": "", "parameters": _SCHEMA, "strict": False},
),
(
{"name": "lookup", "input_schema": _SCHEMA, "strict": True},
{"type": "function", "name": "lookup", "description": "", "parameters": _SCHEMA, "strict": True},
),
(
{"type": "custom", "name": "lookup", "description": "Find a record", "input_schema": _SCHEMA},
{
"type": "function",
"name": "lookup",
"description": "Find a record",
"parameters": _SCHEMA,
"strict": False,
},
),
),
ids=("anthropic", "anthropic-strict", "anthropic-custom-type"),
)
def test_transform_anthropic_tool_shapes(tool: dict[str, object], expected: dict[str, object]) -> None:
result: Final = OpenAICountTokensConfig().transform_request_to_count_tokens(
model="gpt-4o", input="Find a record", tools=[tool]
)
assert result["tools"] == [expected]
def test_transform_builtin_tool_passes_through():
config: Final = OpenAICountTokensConfig()
tools: Final = [{"type": "web_search", "search_context_size": "low"}]
result: Final = config.transform_request_to_count_tokens(
model="gpt-4o",
input="Find a record",
tools=tools,
)
assert result["tools"] == tools
def test_transform_no_instructions_no_tools():
"""Test that None values are not included."""
config = OpenAICountTokensConfig()
@ -154,7 +106,9 @@ def test_messages_to_responses_input_basic():
{"role": "user", "content": "How are you?"},
]
input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages)
input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(
messages
)
assert len(input_items) == 3
assert input_items[0] == {"role": "user", "content": "Hello"}
@ -170,7 +124,9 @@ def test_messages_to_responses_input_with_system():
{"role": "user", "content": "Hello"},
]
input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages)
input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(
messages
)
assert len(input_items) == 1
assert input_items[0] == {"role": "user", "content": "Hello"}
@ -184,7 +140,9 @@ def test_messages_to_responses_input_with_developer():
{"role": "user", "content": "Hello"},
]
input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages)
input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(
messages
)
assert len(input_items) == 1
assert instructions == "Be concise."
@ -197,7 +155,9 @@ def test_messages_to_responses_input_with_tool():
{"role": "tool", "content": "72°F", "tool_call_id": "call_123"},
]
input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages)
input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(
messages
)
assert len(input_items) == 2
assert input_items[1] == {
@ -468,7 +428,10 @@ def test_validate_request_missing_input():
def test_get_endpoint_default():
"""Test default endpoint URL."""
config = OpenAICountTokensConfig()
assert config.get_openai_count_tokens_endpoint() == "https://api.openai.com/v1/responses/input_tokens"
assert (
config.get_openai_count_tokens_endpoint()
== "https://api.openai.com/v1/responses/input_tokens"
)
def test_get_endpoint_custom_base():

0
tests/unit/llms/vertex_ai/count_tokens/__init__.py
View file

29
tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py
View file

@ -47,35 +47,6 @@ async def test_vertex_gemini_rejects_non_text_system_blocks_before_request(
assert "text blocks" in (result.error_message or "")
@pytest.mark.asyncio
@pytest.mark.parametrize(
"model",
(
"zai-org/glm-5.2-maas",
"deepseek-ai/deepseek-v3.2-maas",
"qwen/qwen3-coder-480b-a35b-instruct-maas",
"mistral-medium-3",
"codestral-2501",
"openai/gpt-oss-120b-maas",
"google/gemma-4-26b-a4b-it-maas",
),
)
async def test_vertex_partner_model_without_an_inference_matching_count_answers_400_naming_the_model(model: str) -> None:
result: Final = await VertexAITokenCounter().count_tokens(
model_to_use=model,
messages=[{"role": "user", "content": "Count this"}],
contents=None,
deployment={"litellm_params": {"vertex_project": "test-project", "vertex_location": "global"}},
request_model=f"vertex_ai/{model}",
messages_format="anthropic",
)
assert result is not None
assert result.error is True
assert result.status_code == 400
assert result.error_message == f"LiteLLM does not support counting tokens for Vertex AI model {model}."
class _VertexCountTokensPeer:
def __init__(self, status_code: int) -> None:
self.status_code: Final = status_code

60
tests/unit/proxy/test_custom_tokenizer_bug.py
View file

@ -9,7 +9,6 @@ The HuggingFace download boundary (Tokenizer.from_pretrained) is mocked so these
stay hermetic unit tests; the proxy's extraction-and-selection path runs for real.
"""
from typing import Final
from unittest.mock import MagicMock, patch
import pytest
@ -18,7 +17,7 @@ import litellm
import litellm.proxy.proxy_server
import litellm.utils
from litellm import Router
from litellm.proxy._types import ProxyException, TokenCountRequest
from litellm.proxy._types import TokenCountRequest
from litellm.proxy.proxy_server import token_counter
@ -103,60 +102,3 @@ async def test_model_without_custom_tokenizer_uses_default(monkeypatch):
assert response.tokenizer_type == "openai_tokenizer"
assert response.model_used == "gpt-4"
assert response.total_tokens > 0
def _vertex_glm_router_with_its_own_tokenizer() -> Router:
return Router(
model_list=[
{
"model_name": "vertex-glm",
"litellm_params": {
"model": "vertex_ai/zai-org/glm-5.2-maas",
"vertex_project": "test-project",
"vertex_location": "global",
},
"model_info": {
"custom_tokenizer": {"identifier": "zai-org/GLM-5.2", "revision": "main", "auth_token": None},
},
}
]
)
@pytest.mark.asyncio
async def test_custom_tokenizer_counts_a_vertex_partner_model_without_a_count_route(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _vertex_glm_router_with_its_own_tokenizer())
with patch.object(litellm.utils, "tokenizer_dispatch") as mock_tokenizer_cls:
mock_tokenizer_cls.from_pretrained.return_value = _fake_hf_tokenizer(9)
response: Final = await token_counter(
request=TokenCountRequest(model="vertex-glm", messages=[{"role": "user", "content": "Return exactly OK."}]),
call_endpoint=True,
)
mock_tokenizer_cls.from_pretrained.assert_called_once_with("zai-org/GLM-5.2", revision="main", token=None)
assert response.error is not True
assert response.tokenizer_type == "huggingface_tokenizer"
assert response.total_tokens >= 9
@pytest.mark.asyncio
async def test_contents_only_count_of_a_vertex_partner_model_without_a_count_route_answers_400(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _vertex_glm_router_with_its_own_tokenizer())
with pytest.raises(ProxyException) as raised:
await token_counter(
request=TokenCountRequest(
model="vertex-glm", contents=[{"role": "user", "parts": [{"text": "Return exactly OK."}]}]
),
call_endpoint=True,
)
assert raised.value.code == "400"
assert raised.value.message == "LiteLLM does not support counting tokens for Vertex AI model zai-org/glm-5.2-maas."

90
tests/unit/proxy/test_proxy_token_counter.py
View file

@ -801,45 +801,41 @@ def test_vertex_ai_partner_models_token_counting_endpoint(vertex_location):
assert endpoint.startswith(f"https://{vertex_location}-aiplatform.googleapis.com")
class _FailingRuntimeHandler(BedrockCountTokensHandler):
def __init__(self, error: Exception) -> None:
super().__init__()
self._error = error
async def handle_count_tokens_request(
self,
request_data: dict[str, object],
litellm_params: dict[str, object],
resolved_model: str,
client: AsyncHTTPHandler | None = None,
) -> dict[str, JsonValue]:
raise self._error
@pytest.mark.asyncio
async def test_bedrock_token_counter_error_propagation_bedrock_error():
"""
Test that BedrockTokenCounter properly returns error response when BedrockError is raised.
Verifies that the status code and error message are preserved.
"""
counter = BedrockTokenCounter(
runtime_handler=_FailingRuntimeHandler(BedrockError(status_code=429, message="Rate limit exceeded"))
)
counter = BedrockTokenCounter()
result = await counter.count_tokens(
model_to_use="anthropic.claude-3-sonnet",
messages=[{"role": "user", "content": "hello"}],
contents=None,
deployment={"litellm_params": {}},
request_model="bedrock/anthropic.claude-3-sonnet",
)
# Mock the handler to raise BedrockError with specific status code
with patch.object(
counter, "count_tokens", wraps=counter.count_tokens
) as mock_count:
# We need to patch at the handler level
with patch(
"litellm.llms.bedrock.count_tokens.bedrock_token_counter.BedrockCountTokensHandler"
) as MockHandler:
mock_handler_instance = MockHandler.return_value
mock_handler_instance.handle_count_tokens_request = AsyncMock(
side_effect=BedrockError(status_code=429, message="Rate limit exceeded")
)
assert result is not None
assert result.error is True
assert result.status_code == 429
assert "Rate limit exceeded" in result.error_message
assert result.tokenizer_type == "bedrock_api"
assert result.total_tokens == 0
result = await counter.count_tokens(
model_to_use="anthropic.claude-3-sonnet",
messages=[{"role": "user", "content": "hello"}],
contents=None,
deployment={"litellm_params": {}},
request_model="bedrock/anthropic.claude-3-sonnet",
)
assert result is not None
assert result.error is True
assert result.status_code == 429
assert "Rate limit exceeded" in result.error_message
assert result.tokenizer_type == "bedrock_api"
assert result.total_tokens == 0
@pytest.mark.asyncio
@ -847,20 +843,28 @@ async def test_bedrock_token_counter_error_propagation_generic_exception():
"""
Test that BedrockTokenCounter returns error response with 500 status for generic exceptions.
"""
counter = BedrockTokenCounter(runtime_handler=_FailingRuntimeHandler(Exception("Unexpected error")))
counter = BedrockTokenCounter()
result = await counter.count_tokens(
model_to_use="anthropic.claude-3-sonnet",
messages=[{"role": "user", "content": "hello"}],
contents=None,
deployment={"litellm_params": {}},
request_model="bedrock/anthropic.claude-3-sonnet",
)
with patch(
"litellm.llms.bedrock.count_tokens.bedrock_token_counter.BedrockCountTokensHandler"
) as MockHandler:
mock_handler_instance = MockHandler.return_value
mock_handler_instance.handle_count_tokens_request = AsyncMock(
side_effect=Exception("Unexpected error")
)
assert result is not None
assert result.error is True
assert result.status_code == 500
assert "Unexpected error" in result.error_message
result = await counter.count_tokens(
model_to_use="anthropic.claude-3-sonnet",
messages=[{"role": "user", "content": "hello"}],
contents=None,
deployment={"litellm_params": {}},
request_model="bedrock/anthropic.claude-3-sonnet",
)
assert result is not None
assert result.error is True
assert result.status_code == 500
assert "Unexpected error" in result.error_message
@pytest.mark.asyncio

52
tests/unit/test_cost_map_guard.py
View file

@ -22,6 +22,7 @@ def _load(name: str) -> ModuleType:
return module
schema_module: Final = _load("generate_model_prices_schema")
guard: Final = _load("cost_map_guard")
MAP_FILES: Final = (guard.COST_MAP_PATH,)
@ -51,9 +52,12 @@ def _serialize(cost_map: dict[str, object]) -> str:
return json.dumps(cost_map, indent=4, ensure_ascii=False) + "\n"
def _snapshot(cost_map: dict[str, object], backup: str | None = None) -> object:
def _snapshot(cost_map: dict[str, object], backup: str | None = None, schema: str | None = None) -> object:
text = _serialize(cost_map)
return guard.Snapshot(cost_map=text, backup=text if backup is None else backup)
rendered = schema_module.render(schema_module.build_schema(cost_map))
return guard.Snapshot(
cost_map=text, backup=text if backup is None else backup, schema=rendered if schema is None else schema
)
BASE: Final = _snapshot(BASE_MAP)
@ -131,14 +135,14 @@ def test_bot_may_add_and_reprice_models() -> None:
def test_broken_json_is_reported() -> None:
head = guard.Snapshot(cost_map="{not json", backup="{not json")
head = guard.Snapshot(cost_map="{not json", backup="{not json", schema="{}")
assert _failures(head, bot=False) == (
f"{guard.COST_MAP_PATH} is not valid JSON: Expecting property name enclosed in double quotes: line 1 column 2 (char 1)",
)
def test_non_object_root_is_reported() -> None:
head = guard.Snapshot(cost_map="[]", backup="[]")
head = guard.Snapshot(cost_map="[]", backup="[]", schema="{}")
assert _failures(head, bot=False) == (f"{guard.COST_MAP_PATH} must be a JSON object at the root",)
@ -147,7 +151,28 @@ def test_backup_drift_is_reported() -> None:
assert [failure for failure in _failures(head, bot=False) if failure.startswith(guard.BACKUP_PATH)]
STALE_HEAD: Final = _snapshot(BASE_MAP, backup=_serialize({**BASE_MAP, "openrouter/b": _entry(3e-06)}))
def test_schema_out_of_sync_is_reported() -> None:
head = _snapshot({**BASE_MAP, "openrouter/c": _entry(supports_audio_input=True)}, schema=BASE.schema)
assert [failure for failure in _failures(head, bot=False) if failure.startswith(guard.SCHEMA_PATH)]
def test_schema_validation_errors_are_reported() -> None:
head = _snapshot({**BASE_MAP, "openrouter/c": _entry(-1e-06)})
prefix = f"{guard.COST_MAP_PATH} does not validate against its schema: openrouter/c."
assert [failure.removeprefix(prefix).split(":")[0] for failure in _failures(head, bot=False)] == [
"input_cost_per_token",
"output_cost_per_token",
]
def test_unclassified_entry_key_is_reported() -> None:
text = _serialize({**BASE_MAP, "openrouter/c": _entry(weird_thing=1)})
head = guard.Snapshot(cost_map=text, backup=text, schema=BASE.schema)
(failure,) = _failures(head, bot=False)
assert "Unclassified keys" in failure and "weird_thing" in failure
STALE_HEAD: Final = _snapshot(BASE_MAP, backup=_serialize({**BASE_MAP, "openrouter/b": _entry(3e-06)}), schema="{}")
CODE_ONLY: Final = (
"litellm/utils.py",
"tests/test_litellm/test_utils.py",
@ -156,7 +181,7 @@ CODE_ONLY: Final = (
def test_human_pr_that_leaves_the_cost_map_alone_skips_the_file_checks() -> None:
unparseable: Final = guard.Snapshot(cost_map="{not json", backup="")
unparseable: Final = guard.Snapshot(cost_map="{not json", backup="", schema="")
assert _failures(STALE_HEAD, changed_files=CODE_ONLY, bot=False) == ()
assert _failures(STALE_HEAD, changed_files=(), bot=False) == ()
assert _failures(unparseable, changed_files=CODE_ONLY, bot=False) == ()
@ -166,16 +191,18 @@ def test_human_pr_that_leaves_the_cost_map_alone_skips_the_file_checks() -> None
def test_touching_any_cost_map_file_keeps_the_file_checks(guarded_path: str) -> None:
failures: Final = _failures(STALE_HEAD, changed_files=(*CODE_ONLY, guarded_path), bot=False)
assert [failure for failure in failures if failure.startswith(guard.BACKUP_PATH)]
assert [failure for failure in failures if failure.startswith(guard.SCHEMA_PATH)]
def test_bot_pr_always_gets_the_file_checks() -> None:
failures: Final = _failures(STALE_HEAD, changed_files=CODE_ONLY, bot=True)
assert [failure for failure in failures if failure.startswith(guard.BACKUP_PATH)]
assert [failure for failure in failures if failure.startswith(guard.SCHEMA_PATH)]
def test_contract_names_the_skip() -> None:
assert guard.contract_for(False, CODE_ONLY) == "human PR, cost map untouched"
assert guard.contract_for(False, (*CODE_ONLY, guard.COST_MAP_PATH)) == "human PR, file checks only"
assert guard.contract_for(False, (*CODE_ONLY, guard.SCHEMA_PATH)) == "human PR, file checks only"
assert guard.contract_for(True, CODE_ONLY) == "bot contract enforced"
@ -188,13 +215,6 @@ def test_bot_may_only_touch_the_cost_map_files() -> None:
)
def test_bot_may_update_cost_map_and_legacy_schema() -> None:
head: Final = _snapshot({**BASE_MAP, "openrouter/c": _entry()})
changed: Final = (*guard.GUARDED_PATHS, guard.LEGACY_SCHEMA_PATH)
assert _failures(head, changed_files=changed) == ()
def test_bot_may_not_remove_models() -> None:
head = _snapshot({key: value for key, value in BASE_MAP.items() if key != "openrouter/b"})
assert _failures(head, bot=False) == ()
@ -218,6 +238,7 @@ def _commit(repo: Path, cost_map: dict[str, object], message: str) -> str:
(repo / guard.COST_MAP_PATH).write_text(text)
(repo / guard.BACKUP_PATH).parent.mkdir(exist_ok=True)
(repo / guard.BACKUP_PATH).write_text(text)
(repo / guard.SCHEMA_PATH).write_text(schema_module.render(schema_module.build_schema(cost_map)))
return _git_commit(repo, message)
@ -271,7 +292,8 @@ def test_main_skips_the_file_checks_on_a_stale_base_the_pr_never_touched(tmp_pat
subprocess.run(("git", "init", "-q", str(tmp_path)), check=True)
_commit(tmp_path, BASE_MAP, "base")
(tmp_path / guard.BACKUP_PATH).write_text(_serialize({**BASE_MAP, "openrouter/b": _entry(3e-06)}))
stale_base: Final = _commit_code_only(tmp_path, "stale base with drifted backup")
(tmp_path / guard.SCHEMA_PATH).write_text("{}")
stale_base: Final = _commit_code_only(tmp_path, "stale base with drifted backup and schema")
head: Final = _commit_code_only(tmp_path, "code change on the stale base")
human: Final = _run_guard(tmp_path, stale_base, head, "litellm_fix_pricing")
assert human.returncode == 0, human.stdout + human.stderr

198
tests/unit/test_model_prices_and_context_window.py → tests/unit/test_model_prices_schema.py
View file

@ -1,5 +1,6 @@
from __future__ import annotations
import importlib.util
import json
import re
from collections.abc import Mapping
@ -7,6 +8,7 @@ from pathlib import Path
from types import MappingProxyType
from typing import Final
import jsonschema
import pytest
import litellm
@ -15,8 +17,27 @@ from litellm.llms.openai_like.json_loader import JSONProviderRegistry
from litellm.router_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts
REPO_ROOT = Path(__file__).parents[2]
GENERATOR_PATH = REPO_ROOT / "ci_cd" / "generate_model_prices_schema.py"
PRICES_PATH = REPO_ROOT / "model_prices_and_context_window.json"
BACKUP_PRICES_PATH = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json"
SCHEMA_PATH = REPO_ROOT / "model_prices_and_context_window.schema.json"
def build_validator(schema: dict) -> jsonschema.Draft202012Validator:
return jsonschema.Draft202012Validator(schema, format_checker=jsonschema.Draft202012Validator.FORMAT_CHECKER)
def load_generator():
spec = importlib.util.spec_from_file_location("generate_model_prices_schema", GENERATOR_PATH)
assert spec is not None and spec.loader is not None
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
@pytest.fixture(scope="module")
def committed_schema() -> dict:
return json.loads(SCHEMA_PATH.read_text())
@pytest.fixture(scope="module")
@ -24,6 +45,183 @@ def prices() -> dict:
return json.loads(PRICES_PATH.read_text())
def test_committed_schema_matches_generator_output(prices: dict, committed_schema: dict):
generator = load_generator()
regenerated = json.loads(generator.render(generator.build_schema(prices)))
assert regenerated == committed_schema, (
"model_prices_and_context_window.schema.json is out of sync; "
"run `python ci_cd/generate_model_prices_schema.py` and commit the result"
)
def test_prices_file_validates_against_committed_schema(prices: dict, committed_schema: dict):
validator = build_validator(committed_schema)
errors = [
f"{'.'.join(str(part) for part in error.absolute_path)}: {error.message}"
for error in validator.iter_errors(prices)
]
assert errors == []
@pytest.mark.parametrize(
"entry",
[
{"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": "0.01"},
{"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": -1},
{"litellm_provider": "openai", "mode": "not_a_real_mode"},
{"mode": "chat"},
{"litellm_provider": "openai", "deprecation_date": "June 2026"},
{"litellm_provider": "openai", "deprecation_date": "2026-99-99"},
{"litellm_provider": "openai", "deprecation_date": "2026-13-01"},
{"litellm_provider": "openai", "deprecation_date": "2026-01-32"},
{"litellm_provider": "openai", "deprecation_date": "2026-01-00"},
{"litellm_provider": "openai", "deprecation_date": "2026-02-31"},
{"litellm_provider": "openai", "supported_modalities": ["smell"]},
{"litellm_provider": "openai", "supports_vision": "yes"},
{"litellm_provider": "openai", "max_tokens": 8191.5},
{"litellm_provider": "openai", "tiered_pricing": [{"unknown_tier_field": 1}]},
],
ids=[
"cost_as_string",
"negative_cost",
"unknown_mode",
"missing_provider",
"non_iso_deprecation_date",
"impossible_month_and_day",
"month_out_of_range",
"day_out_of_range",
"day_zero",
"calendar_impossible_day",
"unknown_modality",
"boolean_flag_as_string",
"fractional_max_tokens",
"unknown_tiered_pricing_field",
],
)
def test_schema_rejects_malformed_entries(committed_schema: dict, entry: dict):
validator = build_validator(committed_schema)
assert not validator.is_valid({"some-model": entry})
def test_schema_accepts_minimal_and_unknown_optional_fields(committed_schema: dict):
validator = build_validator(committed_schema)
assert validator.is_valid({"some-model": {"litellm_provider": "openai"}})
assert validator.is_valid({"some-model": {"litellm_provider": "openai", "brand_new_field": {"nested": True}}})
def test_schema_accepts_cache_creation_cost_inside_a_pricing_tier(committed_schema: dict):
validator = build_validator(committed_schema)
entry = {
"litellm_provider": "dashscope",
"mode": "chat",
"tiered_pricing": [
{
"range": [0, 256000],
"input_cost_per_token": 3.25e-07,
"output_cost_per_token": 1.95e-06,
"cache_creation_input_token_cost": 4.063e-07,
"cache_read_input_token_cost": 3.25e-08,
}
],
}
assert validator.is_valid({"some-model": entry})
OFF_PEAK_ENTRY: Final = MappingProxyType(
{
"litellm_provider": "openrouter",
"mode": "chat",
"input_cost_per_token": 2e-6,
"output_cost_per_token": 8e-6,
"off_peak_pricing": {
"hours_utc": "16:30-00:30",
"windows": [{"hours_utc": ["00:30-02:00"], "weekdays": [6, "Sunday", "mon", "THURS"]}],
"weekday_timezone": "Asia/Shanghai",
"input_cost_per_token": 1e-6,
"output_cost_per_token": 4e-6,
"cache_read_input_token_cost": 1e-7,
},
}
)
def test_generator_classifies_off_peak_pricing_as_a_windowed_rate_block():
generator = load_generator()
schema = json.loads(generator.render(generator.build_schema({"some-model": dict(OFF_PEAK_ENTRY)})))
validator = build_validator(schema)
assert validator.is_valid({"some-model": dict(OFF_PEAK_ENTRY)})
@pytest.mark.parametrize(
"block",
[
{"hours_utc": "16:30-00:30", "input_cost_per_token": "1e-6"},
{"hours_utc": "16:30-00:30", "input_cost_per_token": -1e-6},
{"hours_utc": 1630, "input_cost_per_token": 1e-6},
{"hours_utc": "16:30-00:30", "discount": 0.5},
{"windows": [{"weekdays": [6]}], "input_cost_per_token": 1e-6},
{"windows": [{"hours_utc": "00:30-02:00", "weekdays": [0]}], "input_cost_per_token": 1e-6},
{"windows": [], "input_cost_per_token": 1e-6},
{"input_cost_per_token": 1e-6},
{"hours_utc": "16:30", "input_cost_per_token": 1e-6},
{"hours_utc": "25:00-01:00", "input_cost_per_token": 1e-6},
{"hours_utc": ["16:30-00:30", "4pm-midnight"], "input_cost_per_token": 1e-6},
{"windows": [{"hours_utc": "00:30-02:00", "weekdays": ["Funday"]}], "input_cost_per_token": 1e-6},
],
)
def test_generated_off_peak_schema_rejects_malformed_blocks(block: dict):
generator = load_generator()
schema = json.loads(generator.render(generator.build_schema({"some-model": dict(OFF_PEAK_ENTRY)})))
validator = build_validator(schema)
assert not validator.is_valid({"some-model": {**OFF_PEAK_ENTRY, "off_peak_pricing": block}})
TOKEN_COUNT_ENTRY: Final = MappingProxyType(
{
"litellm_provider": "bedrock_converse",
"mode": "chat",
"token_count_framing": 78,
"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": "us-east-1"},
}
)
@pytest.mark.parametrize(
"setting",
[
{"token_count_framing": 0},
{"token_count_framing": 78},
{"token_count_model": {"model": "us.anthropic.claude-opus-5-5"}},
],
)
def test_generator_classifies_token_count_settings(setting: dict):
generator = load_generator()
schema = json.loads(generator.render(generator.build_schema({"some-model": dict(TOKEN_COUNT_ENTRY)})))
assert build_validator(schema).is_valid({"some-model": {**TOKEN_COUNT_ENTRY, **setting}})
@pytest.mark.parametrize(
"setting",
[
{"token_count_framing": -1},
{"token_count_framing": "78"},
{"token_count_framing": 1.5},
{"token_count_model": {"model": ""}},
{"token_count_model": {"model": " "}},
{"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": ""}},
{"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": " "}},
{"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": 5}},
{"token_count_model": {"region": "us-east-1"}},
{"token_count_model": "us.anthropic.claude-opus-5-5"},
{"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "zone": "us-east-1"}},
],
)
def test_generated_token_count_schema_rejects_malformed_settings(setting: dict):
generator = load_generator()
schema = json.loads(generator.render(generator.build_schema({"some-model": dict(TOKEN_COUNT_ENTRY)})))
assert not build_validator(schema).is_valid({"some-model": {**TOKEN_COUNT_ENTRY, **setting}})
def find_duplicate_keys(path: Path) -> list[str]:
duplicates: list[str] = []