diff --git a/.github/workflows/auto_update_price_and_context_window.yml b/.github/workflows/auto_update_price_and_context_window.yml
index 838ba35f831..7e40a860ee9 100644
--- a/.github/workflows/auto_update_price_and_context_window.yml
+++ b/.github/workflows/auto_update_price_and_context_window.yml
@@ -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" \
diff --git a/.github/workflows/sync-together-ai-models.yml b/.github/workflows/sync-together-ai-models.yml
index 5325510e45e..18d5e3de1eb 100644
--- a/.github/workflows/sync-together-ai-models.yml
+++ b/.github/workflows/sync-together-ai-models.yml
@@ -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"
diff --git a/ci_cd/cost_map_guard.py b/ci_cd/cost_map_guard.py
index 8e34e79a354..14eecf01461 100644
--- a/ci_cd/cost_map_guard.py
+++ b/ci_cd/cost_map_guard.py
@@ -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),
diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py
new file mode 100644
index 00000000000..359df8e6e32
--- /dev/null
+++ b/ci_cd/generate_model_prices_schema.py
@@ -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())
diff --git a/litellm-rust/crates/model-catalog/AGENTS.md b/litellm-rust/crates/model-catalog/AGENTS.md
index 0a380788c19..9fcbcd57a76 100644
--- a/litellm-rust/crates/model-catalog/AGENTS.md
+++ b/litellm-rust/crates/model-catalog/AGENTS.md
@@ -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
diff --git a/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py b/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py
deleted file mode 100644
index 3c437ca7fbb..00000000000
--- a/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py
+++ /dev/null
@@ -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)
diff --git a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py
index 37c32144bbf..b5e9afca86b 100644
--- a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py
+++ b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py
@@ -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] = {
diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py
index 5a76e85da75..c62233f49e2 100644
--- a/litellm/litellm_core_utils/token_counter.py
+++ b/litellm/litellm_core_utils/token_counter.py
@@ -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:
diff --git a/litellm/llms/anthropic/count_tokens/handler.py b/litellm/llms/anthropic/count_tokens/handler.py
index 0cecf62368c..dd2135f4918 100644
--- a/litellm/llms/anthropic/count_tokens/handler.py
+++ b/litellm/llms/anthropic/count_tokens/handler.py
@@ -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,
diff --git a/litellm/llms/anthropic/count_tokens/transformation.py b/litellm/llms/anthropic/count_tokens/transformation.py
index 8c47d1cf884..fb12747cec0 100644
--- a/litellm/llms/anthropic/count_tokens/transformation.py
+++ b/litellm/llms/anthropic/count_tokens/transformation.py
@@ -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}
diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py
index e1eb83e4ebf..36a227c03ee 100644
--- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py
+++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py
@@ -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,
diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py
index 7ff2fddbb57..57bbb5b6cb6 100644
--- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py
+++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py
@@ -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
diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py
index be8243e48ef..d7fb510f057 100644
--- a/litellm/llms/bedrock/count_tokens/handler.py
+++ b/litellm/llms/bedrock/count_tokens/handler.py
@@ -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,
diff --git a/litellm/llms/bedrock/count_tokens/mantle_handler.py b/litellm/llms/bedrock/count_tokens/mantle_handler.py
deleted file mode 100644
index 23c3dc54c95..00000000000
--- a/litellm/llms/bedrock/count_tokens/mantle_handler.py
+++ /dev/null
@@ -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}")
diff --git a/litellm/llms/bedrock_mantle/common_utils.py b/litellm/llms/bedrock_mantle/common_utils.py
index ca637af0110..9892ef6224e 100644
--- a/litellm/llms/bedrock_mantle/common_utils.py
+++ b/litellm/llms/bedrock_mantle/common_utils.py
@@ -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()
diff --git a/litellm/llms/openai/responses/count_tokens/transformation.py b/litellm/llms/openai/responses/count_tokens/transformation.py
index 3ea5cf5c4b7..a26aebd8526 100644
--- a/litellm/llms/openai/responses/count_tokens/transformation.py
+++ b/litellm/llms/openai/responses/count_tokens/transformation.py
@@ -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)
diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py
index eb57ef4e6c6..2c7e07616b5 100644
--- a/litellm/llms/vertex_ai/common_utils.py
+++ b/litellm/llms/vertex_ai/common_utils.py
@@ -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
diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py
index 7100c9efde3..285380f5fc7 100644
--- a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py
+++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py
@@ -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,
diff --git a/litellm/proxy/common_utils/openai_error_payload.py b/litellm/proxy/common_utils/openai_error_payload.py
index b24761ac039..202c61b620e 100644
--- a/litellm/proxy/common_utils/openai_error_payload.py
+++ b/litellm/proxy/common_utils/openai_error_payload.py
@@ -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
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 8369e79596c..4ec1cf36408 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -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)
diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py
index d33d8a83bfa..80652e63300 100644
--- a/litellm/proxy/response_api_endpoints/endpoints.py
+++ b/litellm/proxy/response_api_endpoints/endpoints.py
@@ -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,
diff --git a/tests/integration/_support/bedrock_runtime_peer.py b/tests/integration/_support/bedrock_runtime_peer.py
deleted file mode 100644
index 64115d9eaa0..00000000000
--- a/tests/integration/_support/bedrock_runtime_peer.py
+++ /dev/null
@@ -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"why marker-{marker} {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 ("why ", f"marker-{marker}", " 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()
diff --git a/tests/integration/_support/openai_wire.py b/tests/integration/_support/openai_wire.py
deleted file mode 100644
index bf9d5f194c1..00000000000
--- a/tests/integration/_support/openai_wire.py
+++ /dev/null
@@ -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())
diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py
index 1b39eb81b01..1b137c38898 100644
--- a/tests/integration/conftest.py
+++ b/tests/integration/conftest.py
@@ -15,6 +15,7 @@ from redis import Redis
from tests.integration._support.client import Gateway, eventually, gateway_from_environment
from tests.integration._support.generation import LIFECYCLE_SETTINGS
from tests.integration._support.manifest import OWNED_DIRECTORIES
+from tests.integration._support.provider import SharedProvider, shared_provider
from tests.integration._support.routing import RoutingPlugin
COLLECTED: Final = pytest.StashKey[tuple[str, ...]]()
@@ -126,6 +127,35 @@ def gateway() -> Iterator[Gateway]:
yield value
+@pytest.fixture(scope="session")
+def shared_provider_server() -> Iterator[SharedProvider]:
+ if os.environ.get("PYTEST_XDIST_WORKER") and os.environ.get("INTEGRATION_WORKER_ISOLATION") != "1":
+ pytest.fail("the shared fake provider needs tests to run one at a time; this group runs under pytest-xdist")
+ with shared_provider() as server:
+ yield server
+ late: Final = server.received()
+ assert late == (), (
+ f"the shared fake provider got {[item.target for item in late]} after {server.last_test} finished"
+ )
+
+
+@pytest.fixture
+def provider(shared_provider_server: SharedProvider, request: pytest.FixtureRequest) -> Iterator[SharedProvider]:
+ stray: Final = shared_provider_server.received()
+ shared_provider_server.replies.clear()
+ assert stray == (), (
+ f"the shared fake provider got {[item.target for item in stray]} "
+ f"after {shared_provider_server.last_test} finished"
+ )
+ yield shared_provider_server
+ shared_provider_server.last_test = request.node.nodeid
+ unused: Final = len(shared_provider_server.replies)
+ unread: Final = shared_provider_server.received()
+ shared_provider_server.replies.clear()
+ assert unused == 0, f"{unused} queued provider replies were never requested"
+ assert unread == (), f"the test never read the provider requests {[item.target for item in unread]}"
+
+
@pytest.fixture
def peer(gateway: Gateway) -> Iterator[Gateway]:
url: Final = os.environ["INTEGRATION_PEER_URL"]
diff --git a/tests/integration/providers/_count_tokens_system_lift.py b/tests/integration/providers/_count_tokens_system_lift.py
deleted file mode 100644
index 56619d1c963..00000000000
--- a/tests/integration/providers/_count_tokens_system_lift.py
+++ /dev/null
@@ -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),
- )
diff --git a/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py b/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py
deleted file mode 100644
index b34b0b6b526..00000000000
--- a/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py
+++ /dev/null
@@ -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
diff --git a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py
deleted file mode 100644
index e20a397a708..00000000000
--- a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py
+++ /dev/null
@@ -1,1568 +0,0 @@
-import asyncio
-import json
-import re
-import signal
-import socket
-import threading
-import uuid
-from collections.abc import Callable, Iterator, Mapping, Sequence
-from concurrent.futures import Future, ThreadPoolExecutor
-from contextlib import ExitStack
-from pathlib import Path
-from queue import SimpleQueue
-from types import MappingProxyType
-from typing import Final, cast
-
-import anthropic
-import httpx
-import psutil
-import pytest
-import yaml
-from anthropic.types import MessageCountTokensToolParam, MessageParam
-from integration._support.bedrock_runtime_peer import answer, marker_of, target_of
-from integration._support.bedrock_runtime_peer import respond as runtime_generation
-from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment, object_value
-from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
-from integration._support.wire import Reply, Request, Wire, wire_server
-from integration.providers._count_tokens_system_lift import (
- ASSISTANT,
- FOLLOW_UP,
- INSTRUCTION,
- LEADING,
- LIFT_CASES,
- LIFTED,
- MID_SYSTEM,
- REMINDER,
- STRING_CASE,
- USER,
- USER_TEXT,
- LiftCase,
- accepts_count_body,
- async_openai_client,
- count_request,
- expected_count_body,
- openai_client,
-)
-from pydantic import JsonValue, TypeAdapter
-
-pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
-
-_OPUS: Final = "global.anthropic.claude-opus-4-8"
-_OPUS_BASE: Final = "anthropic.claude-opus-4-8"
-_SONNET: Final = "anthropic.claude-sonnet-4-6"
-_NOVA: Final = "amazon.nova-lite-v1:0"
-_REGION: Final = "us-east-1"
-_OWNED_OPUS: Final = "mantle-opus"
-_OWNED_SONNET: Final = "mantle-sonnet"
-_OWNED_NOVA: Final = "mantle-nova"
-_OWNED_MODELS: Final = MappingProxyType({_OWNED_OPUS: _OPUS, _OWNED_SONNET: _SONNET, _OWNED_NOVA: _NOVA})
-_ACCESS_KEY: Final = "AKIAINTEGRATIONMANTLE"
-_SECRET_KEY: Final = "integration-mantle-secret"
-_SIGV4_SCOPE: Final = f"/{_REGION}/bedrock/aws4_request"
-_MANTLE_TARGET: Final = "/anthropic/v1/messages/count_tokens"
-_MANTLE_VERSION: Final = "2023-06-01"
-_MANTLE_COUNT: Final = 4242
-_RUNTIME_COUNT: Final = 1345
-_UNSUPPORTED: Final = "The provided model doesn't support counting tokens."
-_REJECTION: Final = "scripted mantle rejection"
-_COUNT_TARGET: Final = re.compile(r"^/model/(.+)/count-tokens$")
-_INVOKE_TARGET: Final = re.compile(r"^/model/(.+)/invoke$")
-_SCRIPTED_STATUS: Final = re.compile(r"status=(\d{3})")
-_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")
-_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
-_JSON_LIST: Final = TypeAdapter(list[JsonValue])
-
-_SDK_MESSAGES: Final[list[MessageParam]] = [{"role": "user", "content": "Count this message"}]
-_MESSAGES: Final = _JSON_LIST.validate_json(json.dumps(_SDK_MESSAGES))
-_SYSTEM: Final = "You are a terse assistant that answers in one sentence"
-_SYSTEM_BLOCKS: Final[list[JsonValue]] = [
- {"type": "text", "text": "You are a terse assistant"},
- {"type": "text", "text": "Answer in one sentence"},
-]
-_SDK_TOOLS: Final[list[MessageCountTokensToolParam]] = [
- {
- "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"],
- },
- }
-]
-_TOOLS: Final = _JSON_LIST.validate_json(json.dumps(_SDK_TOOLS))
-_GEMINI_BODY: Final[dict[str, JsonValue]] = {"contents": [{"role": "user", "parts": [{"text": "Count this"}]}]}
-_GEMINI_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": "Count this"}]
-
-
-def _mantle_body(**fields: JsonValue) -> dict[str, JsonValue]:
- return {"model": _OPUS_BASE, "messages": _MESSAGES, **fields}
-
-
-_MANTLE_BARE: Final = _mantle_body()
-_MANTLE_LIFTED: Final = _mantle_body(system=[LIFTED])
-_LIFT_SDK_MESSAGES: Final = cast(
- list[MessageParam], [LEADING, USER]
-) # cast-ok: the SDK types reject the role the proxy lifts
-_MANTLE_FULL: Final = _mantle_body(system=_SYSTEM, tools=_TOOLS)
-
-
-def _json_reply(status: int, payload: Mapping[str, JsonValue]) -> Reply:
- return Reply(status=status, body=json.dumps(payload).encode())
-
-
-def _runtime_count(request: Request, model: str) -> Reply:
- scripted: Final = _SCRIPTED_STATUS.search(request.body.decode(errors="replace"))
- if scripted is not None:
- return _json_reply(int(scripted.group(1)), {"message": f"scripted {scripted.group(1)}"})
- if "sonnet" in model:
- return _json_reply(200, {"inputTokens": _RUNTIME_COUNT})
- return _json_reply(400, {"message": _UNSUPPORTED})
-
-
-def _invoke_reply(marker: str) -> Reply:
- return _json_reply(
- 200,
- {
- "id": f"msg_{marker}",
- "type": "message",
- "role": "assistant",
- "model": _OPUS_BASE,
- "content": [{"type": "text", "text": answer(marker)}],
- "stop_reason": "end_turn",
- "stop_sequence": None,
- "usage": {"input_tokens": 5, "output_tokens": 3},
- },
- )
-
-
-def _runtime(request: Request) -> Reply:
- target: Final = target_of(request)
- counted: Final = _COUNT_TARGET.match(target)
- if counted is not None:
- return _runtime_count(request, counted.group(1))
- if _INVOKE_TARGET.match(target):
- return _invoke_reply(marker_of(request))
- return runtime_generation(request)
-
-
-def _mantle_counted(_request: Request) -> Reply:
- return _json_reply(200, {"input_tokens": _MANTLE_COUNT})
-
-
-def _rejected(status: int) -> Reply:
- return _json_reply(status, {"type": "error", "error": {"type": "invalid_request_error", "message": _REJECTION}})
-
-
-def _rejecting(status: int) -> Callable[[Request], Reply]:
- def count(_request: Request) -> Reply:
- return _rejected(status)
-
- return count
-
-
-def _strict(request: Request) -> Reply:
- accepted: Final = accepts_count_body(_JSON_OBJECT.validate_json(request.body))
- return _mantle_counted(request) if accepted else _rejected(400)
-
-
-def _mantle(count: Callable[[Request], Reply] = _mantle_counted) -> Callable[[Request], Reply]:
- def respond(request: Request) -> Reply:
- if request.target == _MANTLE_TARGET:
- return count(request)
- return _json_reply(404, {"error": f"unscripted mantle target {request.target}"})
-
- return respond
-
-
-def _mantle_environment(port: int) -> Mapping[str, str]:
- return {"BEDROCK_MANTLE_API_BASE": f"http://127.0.0.1:{port}"}
-
-
-_INHERITED_BEARER: Final = ("AWS_BEARER_TOKEN_BEDROCK",)
-
-
-def _reserved_port() -> int:
- with socket.socket() as reserve:
- reserve.bind(("127.0.0.1", 0))
- return reserve.getsockname()[1]
-
-
-def _closed_port_url() -> str:
- return f"http://127.0.0.1:{_reserved_port()}"
-
-
-@pytest.fixture(scope="module")
-def mantle_port() -> int:
- return _reserved_port()
-
-
-@pytest.fixture(scope="module")
-def counting_proxy(tmp_path_factory: pytest.TempPathFactory, mantle_port: int) -> Iterator[OwnedProxy]:
- with (
- gateway_from_environment() as gateway,
- owned_proxy_process(
- gateway,
- tmp_path_factory.mktemp("mantle-count"),
- _mantle_environment(mantle_port),
- workers=1,
- remove_environment=_INHERITED_BEARER,
- ) as owned,
- ):
- yield owned
-
-
-def _litellm_params(model: str, api_base: str) -> dict[str, JsonValue]:
- return {
- "model": f"bedrock/{model}",
- "api_base": api_base,
- "aws_access_key_id": _ACCESS_KEY,
- "aws_secret_access_key": _SECRET_KEY,
- "aws_region_name": _REGION,
- }
-
-
-def _owned_config(path: Path, runtime_url: str, settings: Mapping[str, JsonValue]) -> Path:
- config: Final = object_value(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
- litellm_settings: Final = object_value(config["litellm_settings"])
- path.write_text(
- yaml.safe_dump(
- {
- **config,
- "model_list": [
- {"model_name": name, "litellm_params": _litellm_params(model, runtime_url)}
- for name, model in _OWNED_MODELS.items()
- ],
- "litellm_settings": {**litellm_settings, **settings},
- }
- )
- )
- return path
-
-
-def _deployment(scenario: Scenario, api_base: str, model: str = _OPUS) -> str:
- return scenario.model(model_info=None, api_key=None, **_litellm_params(model, api_base))
-
-
-def _bare(model: str) -> dict[str, JsonValue]:
- return {"model": model, "messages": _MESSAGES}
-
-
-def _full(model: str) -> dict[str, JsonValue]:
- return {**_bare(model), "system": _SYSTEM, "tools": _TOOLS}
-
-
-def _count(gateway: Gateway, body: Mapping[str, JsonValue]) -> httpx.Response:
- return gateway.request("POST", "/v1/messages/count_tokens", body)
-
-
-def _payload(response: httpx.Response) -> dict[str, JsonValue]:
- assert response.status_code == 200, response.text
- return _JSON_OBJECT.validate_json(response.content)
-
-
-def _local_count(gateway: Gateway, body: Mapping[str, JsonValue]) -> int:
- response: Final = gateway.request("POST", "/utils/token_counter", body, params={"call_endpoint": "false"})
- payload: Final = _payload(response)
- total: Final = payload["total_tokens"]
- assert payload["tokenizer_type"] not in ("bedrock_api", "bedrock_mantle_api"), response.text
- assert isinstance(total, int) and total > 0 and total not in (_MANTLE_COUNT, _RUNTIME_COUNT), response.text
- return total
-
-
-def _assert_sigv4(request: Request) -> None:
- authorization: Final = request.headers.get("authorization", "")
- assert authorization.startswith(f"AWS4-HMAC-SHA256 Credential={_ACCESS_KEY}/"), request.headers
- assert _SIGV4_SCOPE in authorization, authorization
-
-
-def _mantle_requests(requests: Sequence[Request]) -> tuple[dict[str, JsonValue], ...]:
- for request in requests:
- assert (request.method, request.target) == ("POST", _MANTLE_TARGET), request.target
- assert request.headers["anthropic-version"] == _MANTLE_VERSION, request.headers
- assert request.headers["content-type"] == "application/json", request.headers
- _assert_sigv4(request)
- return tuple(_JSON_OBJECT.validate_json(request.body) for request in requests)
-
-
-def _mantle_bodies(wire: Wire) -> tuple[dict[str, JsonValue], ...]:
- return _mantle_requests(wire.drain())
-
-
-def _runtime_count_targets(wire: Wire) -> tuple[str, ...]:
- counts: Final = tuple(request for request in wire.drain() if _COUNT_TARGET.match(target_of(request)))
- for request in counts:
- assert request.method == "POST", request.method
- _assert_sigv4(request)
- return tuple(target_of(request) for request in counts)
-
-
-def _clients(stack: ExitStack, base_url: str, count: int, timeout: float = 30) -> tuple[httpx.Client, ...]:
- return tuple(
- stack.enter_context(httpx.Client(base_url=base_url, timeout=timeout, trust_env=False)) for _ in range(count)
- )
-
-
-def _counted_on(client: httpx.Client, key: str, body: Mapping[str, JsonValue]) -> tuple[int, JsonValue]:
- response: Final = client.post(
- "/v1/messages/count_tokens", json=dict(body), headers={"Authorization": f"Bearer {key}"}
- )
- return response.status_code, _JSON_OBJECT.validate_json(response.content).get("input_tokens")
-
-
-def _generated_then_counted(
- client: httpx.Client, key: str, model: str, body: Mapping[str, JsonValue]
-) -> tuple[int, int, JsonValue]:
- generated: Final = client.post(
- "/v1/messages",
- json={
- "model": model,
- "max_tokens": 16,
- "messages": [{"role": "user", "content": f"Generate before counting marker-{uuid.uuid4().hex}"}],
- },
- headers={"Authorization": f"Bearer {key}"},
- )
- return generated.status_code, *_counted_on(client, key, body)
-
-
-def _counted_or_dropped(client: httpx.Client, key: str, body: Mapping[str, JsonValue]) -> tuple[int, JsonValue] | None:
- try:
- return _counted_on(client, key, body)
- except httpx.TransportError:
- return None
-
-
-def _probed_then_counted_or_dropped(
- client: httpx.Client, key: str, body: Mapping[str, JsonValue], probed: SimpleQueue[int]
-) -> tuple[int, tuple[int, JsonValue] | None]:
- port: Final = _local_port(client)
- probed.put(port)
- return port, _counted_or_dropped(client, key, body)
-
-
-def _local_port(client: httpx.Client) -> int:
- with client.stream("GET", "/health/liveliness") as response:
- port: Final = int(response.extensions["network_stream"].get_extra_info("client_addr")[1])
- response.read()
- assert response.status_code == 200, response.text
- return port
-
-
-def _accepted_client_ports(pid: int, proxy_port: int) -> frozenset[int]:
- return frozenset(
- connection.raddr.port
- for connection in psutil.Process(pid).net_connections(kind="tcp")
- if connection.raddr and connection.laddr.port == proxy_port
- )
-
-
-def _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 _mantle_counted(request)
-
- return hold
-
-
-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)
-
-
-@pytest.mark.parametrize("system", [_SYSTEM, _SYSTEM_BLOCKS], ids=["string", "blocks"])
-def test_messages_count_tokens_counts_through_mantle_when_the_runtime_cannot(
- counting_proxy: OwnedProxy, mantle_port: int, system: JsonValue
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS})
- assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
- assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",)
- assert _mantle_bodies(mantle) == (_mantle_body(system=system, tools=_TOOLS),)
-
-
-def test_messages_count_tokens_without_system_or_tools_sends_a_bare_body_to_mantle(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- response: Final = _count(gateway, _bare(model))
- assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (_MANTLE_BARE,)
-
-
-def test_utils_token_counter_call_endpoint_reports_the_mantle_tokenizer(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- response: Final = gateway.request(
- "POST", "/utils/token_counter", _full(model), params={"call_endpoint": "true"}
- )
- payload: Final = _payload(response)
- assert (payload["total_tokens"], payload["tokenizer_type"]) == (_MANTLE_COUNT, "bedrock_mantle_api")
- assert payload["original_response"] == {"input_tokens": _MANTLE_COUNT}, response.text
- assert (payload["request_model"], payload["model_used"]) == (model, _OPUS), response.text
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (_MANTLE_FULL,)
-
-
-def test_gemini_count_tokens_route_reports_the_mantle_total(counting_proxy: OwnedProxy, mantle_port: int) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", _GEMINI_BODY)
- assert _payload(response) == {
- "totalTokens": _MANTLE_COUNT,
- "promptTokensDetails": [],
- "litellm_count_source": "provider",
- }, response.text
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == ({"model": _OPUS_BASE, "messages": _GEMINI_MESSAGES},)
-
-
-def test_gemini_count_tokens_route_reports_the_runtime_total_for_a_model_bedrock_counts(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url, _SONNET)
- response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", _GEMINI_BODY)
- assert _payload(response) == {
- "totalTokens": _RUNTIME_COUNT,
- "promptTokensDetails": [],
- "litellm_count_source": "provider",
- }, response.text
- assert _runtime_count_targets(runtime) == (f"/model/{_SONNET}/count-tokens",)
- assert mantle.drain() == ()
-
-
-def test_runtime_404_on_claude_falls_back_to_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None:
- gateway: Final = counting_proxy.gateway
- system: Final = f"{_SYSTEM} status=404"
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS})
- assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (_mantle_body(system=system, tools=_TOOLS),)
-
-
-def test_runtime_404_and_mantle_404_surface_provider_error(counting_proxy: OwnedProxy, mantle_port: int) -> None:
- gateway: Final = counting_proxy.gateway
- system: Final = f"{_SYSTEM} status=404"
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(_rejecting(404)), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS})
- assert response.status_code == 404, response.text
- assert _REJECTION in response.text
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (_mantle_body(system=system, tools=_TOOLS),)
-
-
-def test_runtime_404_on_nova_does_not_call_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url, _NOVA)
- body: Final = {**_bare(model), "messages": [{"role": "user", "content": "status=404"}]}
- response: Final = _count(gateway, body)
- assert response.status_code == 404, response.text
- assert len(_runtime_count_targets(runtime)) == 1
- assert mantle.drain() == ()
-
-
-def test_google_count_tokens_forwards_system_and_tools_to_mantle(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- system_text: Final = "Use the provided tool"
- tool_schema: Final = {"type": "object", "properties": {"location": {"type": "string"}}}
- body: Final[dict[str, JsonValue]] = {
- "contents": [{"role": "user", "parts": [{"text": "status=404"}]}],
- "systemInstruction": {"parts": [{"text": system_text}]},
- "tools": [
- {
- "functionDeclarations": [
- {"name": "get_weather", "description": "Look up weather", "parameters": tool_schema}
- ]
- }
- ],
- }
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", body)
- assert _payload(response) == {
- "totalTokens": _MANTLE_COUNT,
- "promptTokensDetails": [],
- "litellm_count_source": "provider",
- }, response.text
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (
- _mantle_body(
- messages=[{"role": "user", "content": "status=404"}],
- system=[{"type": "text", "text": system_text}],
- tools=[
- {
- "name": "get_weather",
- "description": "Look up weather",
- "input_schema": tool_schema,
- }
- ],
- ),
- )
-
-
-def test_google_count_tokens_normalizes_uppercase_parameters_for_mantle(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- body: Final[dict[str, JsonValue]] = {
- "contents": [{"role": "user", "parts": [{"text": "status=404"}]}],
- "tools": [
- {
- "functionDeclarations": [
- {
- "name": "get_weather",
- "parameters": {
- "type": "OBJECT",
- "properties": {"location": {"type": "STRING"}},
- "required": ["location"],
- },
- }
- ]
- }
- ],
- }
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", body)
- assert _payload(response) == {
- "totalTokens": _MANTLE_COUNT,
- "promptTokensDetails": [],
- "litellm_count_source": "provider",
- }, response.text
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (
- _mantle_body(
- messages=[{"role": "user", "content": "status=404"}],
- tools=[
- {
- "name": "get_weather",
- "input_schema": {
- "type": "object",
- "properties": {"location": {"type": "string"}},
- "required": ["location"],
- },
- }
- ],
- ),
- )
-
-
-def test_responses_input_tokens_counts_through_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- response: Final = gateway.request(
- "POST", "/v1/responses/input_tokens", {"model": model, "input": "Count this message"}
- )
- assert _payload(response) == {
- "object": "response.input_tokens",
- "input_tokens": _MANTLE_COUNT,
- "litellm_count_source": "provider",
- }, response.text
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (_MANTLE_BARE,)
-
-
-def test_anthropic_sdk_count_tokens_counts_through_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- counted: Final = _anthropic_client(gateway).messages.count_tokens(
- model=model, messages=_SDK_MESSAGES, system=_SYSTEM, tools=_SDK_TOOLS
- )
- assert counted.input_tokens == _MANTLE_COUNT, counted
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (_MANTLE_FULL,)
-
-
-def test_async_anthropic_sdk_count_tokens_counts_through_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- counted: Final = asyncio.run(
- _async_anthropic_client(gateway).messages.count_tokens(model=model, messages=_SDK_MESSAGES, system=_SYSTEM)
- )
- assert counted.input_tokens == _MANTLE_COUNT, counted
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (_mantle_body(system=_SYSTEM),)
-
-
-def test_model_the_runtime_counts_never_reaches_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url, _SONNET)
- assert _payload(_count(gateway, _full(model))) == {
- "input_tokens": _RUNTIME_COUNT,
- "litellm_count_source": "provider",
- }
- detailed: Final = gateway.request(
- "POST", "/utils/token_counter", _full(model), params={"call_endpoint": "true"}
- )
- payload: Final = _payload(detailed)
- assert (payload["total_tokens"], payload["tokenizer_type"]) == (_RUNTIME_COUNT, "bedrock_api"), detailed.text
- assert _runtime_count_targets(runtime) == (f"/model/{_SONNET}/count-tokens",) * 2
- assert mantle.drain() == ()
-
-
-def test_non_claude_model_the_runtime_cannot_count_surfaces_the_error_without_mantle(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url, _NOVA)
- response: Final = _count(gateway, _bare(model))
- assert response.status_code == 400, response.text
- assert _UNSUPPORTED in response.text and "input_tokens" not in response.text, response.text
- assert _runtime_count_targets(runtime) == (f"/model/{_NOVA}/count-tokens",)
- assert mantle.drain() == ()
-
-
-def test_runtime_403_on_a_claude_model_never_reaches_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- body: Final[dict[str, JsonValue]] = {
- "model": model,
- "messages": [{"role": "user", "content": "status=403 Count this message"}],
- }
- response: Final = _count(gateway, body)
- assert response.status_code == 403, response.text
- assert "input_tokens" not in response.text, response.text
- assert len(_runtime_count_targets(runtime)) == 1
- assert mantle.drain() == ()
-
-
-def test_unreachable_runtime_surfaces_the_error_without_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None:
- gateway: Final = counting_proxy.gateway
- with wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario:
- model: Final = _deployment(scenario, _closed_port_url())
- response: Final = _count(gateway, _full(model))
- assert 500 <= response.status_code < 600, response.text
- assert "input_tokens" not in response.text, response.text
- assert mantle.drain() == ()
-
-
-def test_bedrock_passthrough_count_tokens_still_answers_the_runtime_rejection(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- response: Final = gateway.request("POST", "/bedrock/v1/messages/count_tokens", _full(model))
- assert response.status_code == 400, response.text
- assert _UNSUPPORTED in response.text and "input_tokens" not in response.text, response.text
- assert len(_runtime_count_targets(runtime)) == 1
- assert mantle.drain() == ()
-
-
-def test_responses_input_tokens_with_instructions_counts_through_mantle(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(_strict), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- response: Final = gateway.request(
- "POST",
- "/v1/responses/input_tokens",
- {"model": model, "input": "Count this message", "instructions": "Be terse"},
- )
- assert _payload(response) == {
- "object": "response.input_tokens",
- "input_tokens": _MANTLE_COUNT,
- "litellm_count_source": "provider",
- }, response.text
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (_mantle_body(system=[{"type": "text", "text": "Be terse"}]),)
-
-
-@pytest.mark.parametrize("case", LIFT_CASES.values(), ids=LIFT_CASES.keys())
-def test_messages_count_tokens_lifts_the_leading_system_run_for_mantle(
- counting_proxy: OwnedProxy, mantle_port: int, case: LiftCase
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(_strict), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- response: Final = _count(gateway, count_request(model, case))
- assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (expected_count_body(_OPUS_BASE, case),)
-
-
-def test_anthropic_sdk_count_tokens_lifts_the_leading_system_through_mantle(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(_strict), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- counted: Final = _anthropic_client(gateway).messages.count_tokens(model=model, messages=_LIFT_SDK_MESSAGES)
- assert counted.input_tokens == _MANTLE_COUNT, counted
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,)
-
-
-def test_async_anthropic_sdk_count_tokens_lifts_the_leading_system_through_mantle(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(_strict), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- counted: Final = asyncio.run(
- _async_anthropic_client(gateway).messages.count_tokens(model=model, messages=_LIFT_SDK_MESSAGES)
- )
- assert counted.input_tokens == _MANTLE_COUNT, counted
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,)
-
-
-def test_utils_token_counter_call_endpoint_counts_a_leading_system_through_mantle(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(_strict), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- response: Final = gateway.request(
- "POST", "/utils/token_counter", count_request(model, STRING_CASE), params={"call_endpoint": "true"}
- )
- payload: Final = _payload(response)
- assert (payload["total_tokens"], payload["tokenizer_type"]) == (_MANTLE_COUNT, "bedrock_mantle_api")
- assert payload["original_response"] == {"input_tokens": _MANTLE_COUNT}, response.text
- assert (payload["request_model"], payload["model_used"]) == (model, _OPUS), response.text
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,)
-
-
-def test_openai_sdk_input_tokens_lifts_instructions_through_mantle(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(_strict), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- counted: Final = openai_client(gateway).responses.input_tokens.count(
- model=model, input=USER_TEXT, instructions=INSTRUCTION
- )
- assert counted.input_tokens == _MANTLE_COUNT, counted
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,)
-
-
-def test_async_openai_sdk_input_tokens_lifts_instructions_through_mantle(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(_strict), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- counted: Final = asyncio.run(
- async_openai_client(gateway).responses.input_tokens.count(
- model=model, input=USER_TEXT, instructions=INSTRUCTION
- )
- )
- assert counted.input_tokens == _MANTLE_COUNT, counted
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,)
-
-
-def test_responses_input_tokens_lifts_instructions_ahead_of_a_leading_system_item_through_mantle(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(_strict), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- response: Final = gateway.request(
- "POST",
- "/v1/responses/input_tokens",
- {"model": model, "input": [MID_SYSTEM, USER], "instructions": INSTRUCTION},
- )
- assert _payload(response) == {
- "object": "response.input_tokens",
- "input_tokens": _MANTLE_COUNT,
- "litellm_count_source": "provider",
- }, response.text
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (_mantle_body(system=[LIFTED, {"type": "text", "text": REMINDER}]),)
-
-
-def test_messages_count_tokens_forwards_tools_beside_the_lifted_system_to_mantle(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(_strict), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- response: Final = _count(gateway, {**count_request(model, STRING_CASE), "tools": _TOOLS})
- assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (_mantle_body(system=[LIFTED], tools=_TOOLS),)
-
-
-def test_messages_count_tokens_keeps_a_mid_conversation_system_in_place_for_mantle(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(_strict), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- body: Final[dict[str, JsonValue]] = {
- "model": model,
- "messages": [LEADING, USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP],
- }
- response: Final = _count(gateway, body)
- assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (
- _mantle_body(messages=[USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP], system=[LIFTED]),
- )
-
-
-def test_messages_count_tokens_surfaces_the_mantle_error_when_every_message_is_system(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(_strict), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- body: Final[dict[str, JsonValue]] = {"model": 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 len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (_mantle_body(messages=[], system=[LIFTED]),)
-
-
-def test_messages_count_tokens_surfaces_the_mantle_error_beside_a_non_text_system(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(_strict), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- response: Final = _count(gateway, {**count_request(model, STRING_CASE), "system": 5})
- assert response.status_code == 400, response.text
- assert "input_tokens" not in response.text, response.text
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (_mantle_body(messages=[LEADING, USER], system=5),)
-
-
-def test_messages_count_tokens_repeated_request_lifts_the_leading_system_each_time_for_mantle(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(_strict), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- answers: Final = tuple(_payload(_count(gateway, count_request(model, STRING_CASE))) for _ in range(2))
- assert answers == ({"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"},) * 2
- assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",) * 2
- assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) * 2
-
-
-def test_messages_count_tokens_surfaces_the_counting_error_for_a_leading_system_without_content(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(_strict), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- body: Final[dict[str, JsonValue]] = {"model": 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 mantle.drain() == ()
- follow_up: Final = _count(gateway, count_request(model, STRING_CASE))
- assert _payload(follow_up) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, (
- follow_up.text
- )
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,)
-
-
-def test_messages_count_tokens_duplicate_messages_key_lifts_the_last_value_for_mantle(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(_strict), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- first: Final = json.dumps([USER])
- last: Final = json.dumps([LEADING, USER])
- response: Final = gateway.client.post(
- "/v1/messages/count_tokens",
- content=f'{{"model": "{model}", "messages": {first}, "messages": {last}}}',
- headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"},
- )
- assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,)
-
-
-def test_disabled_token_counter_counts_a_leading_system_through_mantle(
- gateway: Gateway, mantle_port: int, tmp_path: Path
-) -> None:
- with ExitStack() as stack:
- runtime: Final = stack.enter_context(wire_server(_runtime))
- config: Final = _owned_config(
- tmp_path / "disabled-token-counter-lift.yaml", runtime.url, {"disable_token_counter": True}
- )
- owned: Final = stack.enter_context(
- owned_proxy_process(
- gateway,
- tmp_path,
- _mantle_environment(mantle_port),
- config=config,
- workers=2,
- remove_environment=_INHERITED_BEARER,
- )
- )
- with wire_server(_mantle(_strict), port=mantle_port) as strict:
- counted: Final = _count(owned.gateway, count_request(_OWNED_OPUS, STRING_CASE))
- assert _payload(counted) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, (
- counted.text
- )
- assert _mantle_bodies(strict) == (_MANTLE_LIFTED,)
- assert len(_runtime_count_targets(runtime)) == 1
-
-
-def test_mantle_outage_between_concurrent_lifted_waves_errors_then_recovers(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- with ExitStack() as stack:
- clients: Final = _clients(stack, str(gateway.client.base_url), 8)
- pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients)))
- runtime: Final = stack.enter_context(wire_server(_runtime))
- scenario: Final = stack.enter_context(gateway.scenario())
- model: Final = _deployment(scenario, runtime.url)
- body: Final = count_request(model, STRING_CASE)
-
- def count(client: httpx.Client) -> tuple[int, JsonValue]:
- return _counted_on(client, gateway.key, body)
-
- with wire_server(_mantle(_strict), port=mantle_port) as mantle:
- assert tuple(pool.map(count, clients)) == ((200, _MANTLE_COUNT),) * len(clients)
- assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) * len(clients)
- outage: Final = tuple(pool.map(count, clients))
- for status, tokens in outage:
- assert 500 <= status < 600 and tokens is None, outage
- with wire_server(_mantle(_strict), port=mantle_port) as revived:
- assert tuple(pool.map(count, clients)) == ((200, _MANTLE_COUNT),) * len(clients)
- assert _mantle_bodies(revived) == (_MANTLE_LIFTED,) * len(clients)
- assert len(_runtime_count_targets(runtime)) == 3 * len(clients)
-
-
-@pytest.mark.parametrize("status", [400, 403, 404, 500, 503])
-def test_messages_count_tokens_surfaces_the_mantle_status_when_mantle_errors(
- counting_proxy: OwnedProxy, mantle_port: int, status: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(_rejecting(status)), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- response: Final = _count(gateway, _full(model))
- assert response.status_code == status, response.text
- assert "input_tokens" not in response.text, response.text
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (_MANTLE_FULL,)
-
-
-@pytest.mark.parametrize(
- "body",
- [b'{"inputTokens": 7}', b"not json at all", b"{}"],
- ids=["runtime_key", "not_json", "empty_object"],
-)
-def test_messages_count_tokens_errors_when_mantle_answers_without_input_tokens(
- counting_proxy: OwnedProxy, mantle_port: int, body: bytes
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(lambda _request: Reply(body=body)), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- response: Final = _count(gateway, _bare(model))
- assert 500 <= response.status_code < 600, response.text
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (_MANTLE_BARE,)
-
-
-def test_messages_count_tokens_errors_when_mantle_is_unreachable(counting_proxy: OwnedProxy) -> None:
- gateway: Final = counting_proxy.gateway
- with wire_server(_runtime) as runtime, gateway.scenario() as scenario:
- model: Final = _deployment(scenario, runtime.url)
- response: Final = _count(gateway, _full(model))
- assert 500 <= response.status_code < 600, response.text
- assert "input_tokens" not in response.text, response.text
- assert len(_runtime_count_targets(runtime)) == 1
-
-
-def test_messages_count_tokens_errors_when_mantle_never_answers(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- held: Final[SimpleQueue[str]] = SimpleQueue()
- release: Final = threading.Event()
-
- def hold(request: Request) -> Reply:
- held.put(request.target)
- release.wait(timeout=120)
- return Reply(drop_connection=True)
-
- with ExitStack() as stack:
- (client,) = _clients(stack, str(gateway.client.base_url), 1, timeout=120)
- pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=1))
- runtime: Final = stack.enter_context(wire_server(_runtime))
- mantle: Final = stack.enter_context(wire_server(_mantle(hold), port=mantle_port))
- stack.callback(release.set)
- scenario: Final = stack.enter_context(gateway.scenario())
- model: Final = _deployment(scenario, runtime.url)
- body: Final = _full(model)
- future: Final = pool.submit(_counted_on, client, gateway.key, body)
- eventually(held.qsize, lambda size: size == 1, seconds=30)
- assert gateway.request("GET", "/health/liveliness").status_code == 200
- assert _local_count(gateway, body) > 0
- assert not future.done()
- status, tokens = future.result(timeout=120)
- assert 500 <= status < 600 and tokens is None, (status, tokens)
- assert len(_runtime_count_targets(runtime)) == 1
- assert len(mantle.drain()) == 1
-
-
-def test_messages_count_tokens_surfaces_the_mantle_error_for_a_non_text_system(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(_strict), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- response: Final = _count(gateway, {**_bare(model), "system": 5})
- assert response.status_code == 400, response.text
- assert "input_tokens" not in response.text, response.text
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (_mantle_body(system=5),)
-
-
-@pytest.mark.parametrize(
- "tools",
- [5, "", "x" * 5120, ["get_weather"]],
- ids=["int", "empty_string", "5kb_string", "list_of_strings"],
-)
-def test_messages_count_tokens_rejects_malformed_tools_without_calling_either_peer(
- counting_proxy: OwnedProxy, mantle_port: int, tools: JsonValue
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- refused: Final = _count(gateway, {**_bare(model), "tools": tools})
- assert 400 <= refused.status_code < 600, refused.text
- assert "input_tokens" not in refused.text, refused.text
- assert runtime.drain() == () and mantle.drain() == ()
- assert _generated_then_counted(gateway.client, gateway.key, model, _bare(model)) == (200, 200, _MANTLE_COUNT)
- assert _mantle_bodies(mantle) == (_MANTLE_BARE,)
-
-
-@pytest.mark.parametrize("fields", [{}, {"messages": []}], ids=["missing", "empty"])
-def test_messages_count_tokens_without_messages_is_rejected(
- counting_proxy: OwnedProxy, mantle_port: int, fields: dict[str, JsonValue]
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- response: Final = _count(gateway, {"model": model, "system": _SYSTEM, "tools": _TOOLS, **fields})
- assert response.status_code == 400, response.text
- assert "messages parameter is required" in response.text, response.text
- assert runtime.drain() == () and mantle.drain() == ()
-
-
-def test_messages_count_tokens_unauthenticated_request_never_reaches_either_peer(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- response: Final = gateway.request(
- "POST", "/v1/messages/count_tokens", _full(model), key="sk-not-a-key-this-proxy-issued"
- )
- assert response.status_code == 401, response.text
- assert "input_tokens" not in response.text, response.text
- assert runtime.drain() == () and mantle.drain() == ()
-
-
-def test_messages_count_tokens_forwards_a_5kb_system_verbatim_to_mantle(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- system: Final = "Answer in one sentence. " * 214
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- response: Final = _count(gateway, {**_bare(model), "system": system})
- assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (_mantle_body(system=system),)
-
-
-def test_messages_count_tokens_duplicate_system_and_tools_keys_forward_one_value_each_to_mantle(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- fields: Final = f'"system": {json.dumps(_SYSTEM)}, "tools": {json.dumps(_TOOLS)}'
- response: Final = gateway.client.post(
- "/v1/messages/count_tokens",
- content=f'{{"model": "{model}", "messages": {json.dumps(_MESSAGES)}, {fields}, {fields}}}',
- headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"},
- )
- assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
- assert len(_runtime_count_targets(runtime)) == 1
- (sent,) = mantle.drain()
- assert _mantle_requests((sent,)) == (_MANTLE_FULL,)
- assert (sent.body.count(b'"system"'), sent.body.count(b'"tools"')) == (1, 1), sent.body
-
-
-def test_messages_count_tokens_leaves_empty_tools_and_system_out_of_the_mantle_body(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- response: Final = _count(gateway, {**_bare(model), "tools": [], "system": ""})
- assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
- assert len(_runtime_count_targets(runtime)) == 1
- assert _mantle_bodies(mantle) == (_MANTLE_BARE,)
-
-
-def test_messages_count_tokens_repeated_request_reaches_both_peers_each_time(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- answers: Final = tuple(_payload(_count(gateway, _bare(model))) for _ in range(2))
- assert answers == ({"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"},) * 2
- assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",) * 2
- assert _mantle_bodies(mantle) == (_MANTLE_BARE,) * 2
-
-
-def test_disabled_token_counter_surfaces_the_mantle_error_instead_of_counting_locally(
- gateway: Gateway, mantle_port: int, tmp_path: Path
-) -> None:
- with ExitStack() as stack:
- runtime: Final = stack.enter_context(wire_server(_runtime))
- config: Final = _owned_config(
- tmp_path / "disabled-token-counter.yaml", runtime.url, {"disable_token_counter": True}
- )
- owned: Final = stack.enter_context(
- owned_proxy_process(
- gateway,
- tmp_path,
- _mantle_environment(mantle_port),
- config=config,
- workers=2,
- remove_environment=_INHERITED_BEARER,
- )
- )
- with wire_server(_mantle(_rejecting(403)), port=mantle_port) as refusing:
- refused: Final = _count(owned.gateway, _full(_OWNED_OPUS))
- assert refused.status_code == 403, refused.text
- assert _REJECTION in refused.text and "input_tokens" not in refused.text, refused.text
- assert _mantle_bodies(refusing) == (_MANTLE_FULL,)
- with wire_server(_mantle(), port=mantle_port) as counting:
- assert _payload(_count(owned.gateway, _full(_OWNED_OPUS))) == {
- "input_tokens": _MANTLE_COUNT,
- "litellm_count_source": "provider",
- }
- assert _mantle_bodies(counting) == (_MANTLE_FULL,)
- assert _payload(_count(owned.gateway, _full(_OWNED_SONNET))) == {
- "input_tokens": _RUNTIME_COUNT,
- "litellm_count_source": "provider",
- }
- rejected: Final = _count(owned.gateway, _full(_OWNED_NOVA))
- assert rejected.status_code == 400, rejected.text
- assert _UNSUPPORTED in rejected.text and "input_tokens" not in rejected.text, rejected.text
- assert counting.drain() == ()
- assert len(_runtime_count_targets(runtime)) == 4
-
-
-def test_mantle_outage_between_concurrent_waves_errors_then_recovers(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- with ExitStack() as stack:
- clients: Final = _clients(stack, str(gateway.client.base_url), 8)
- pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients)))
- runtime: Final = stack.enter_context(wire_server(_runtime))
- scenario: Final = stack.enter_context(gateway.scenario())
- model: Final = _deployment(scenario, runtime.url)
- body: Final = _full(model)
-
- def generate_then_count(client: httpx.Client) -> tuple[int, int, JsonValue]:
- return _generated_then_counted(client, gateway.key, model, body)
-
- def count_only(client: httpx.Client) -> tuple[int, JsonValue]:
- return _counted_on(client, gateway.key, body)
-
- with wire_server(_mantle(), port=mantle_port) as mantle:
- assert tuple(pool.map(generate_then_count, clients)) == ((200, 200, _MANTLE_COUNT),) * len(clients)
- assert _mantle_bodies(mantle) == (_MANTLE_FULL,) * len(clients)
- outage: Final = tuple(pool.map(count_only, clients))
- for status, tokens in outage:
- assert 500 <= status < 600 and tokens is None, outage
- with wire_server(_mantle(), port=mantle_port) as revived:
- assert tuple(pool.map(generate_then_count, clients)) == ((200, 200, _MANTLE_COUNT),) * len(clients)
- assert _mantle_bodies(revived) == (_MANTLE_FULL,) * len(clients)
- assert len(_runtime_count_targets(runtime)) == 3 * len(clients)
-
-
-def test_slow_mantle_holds_concurrent_counts_without_stalling_the_proxy(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- held: Final[SimpleQueue[str]] = SimpleQueue()
- release: Final = threading.Event()
- with ExitStack() as stack:
- clients: Final = _clients(stack, str(gateway.client.base_url), 6)
- runtime: Final = stack.enter_context(wire_server(_runtime))
- mantle: Final = stack.enter_context(wire_server(_mantle(_holding(held, release, 20)), port=mantle_port))
- scenario: Final = stack.enter_context(gateway.scenario())
- pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients)))
- stack.callback(release.set)
- model: Final = _deployment(scenario, runtime.url)
- futures: Final = tuple(
- pool.submit(_generated_then_counted, client, gateway.key, model, _full(model)) for client in clients
- )
- eventually(held.qsize, lambda size: size == len(clients), seconds=30)
- assert gateway.request("GET", "/health/liveliness").status_code == 200
- assert _local_count(gateway, _full(model)) > 0
- assert not any(future.done() for future in futures)
- release.set()
- assert tuple(future.result(timeout=30) for future in futures) == ((200, 200, _MANTLE_COUNT),) * len(clients)
- assert _mantle_bodies(mantle) == (_MANTLE_FULL,) * len(clients)
- assert len(_runtime_count_targets(runtime)) == len(clients)
-
-
-def test_worker_sigkill_mid_burst_leaves_the_sibling_counting(
- gateway: Gateway, mantle_port: int, tmp_path: Path
-) -> None:
- held: Final[SimpleQueue[str]] = SimpleQueue()
- release: Final = threading.Event()
- with ExitStack() as stack:
- runtime: Final = stack.enter_context(wire_server(_runtime))
- mantle: Final = stack.enter_context(wire_server(_mantle(_holding(held, release, 60)), port=mantle_port))
- config: Final = _owned_config(tmp_path / "worker-kill.yaml", runtime.url, {})
- owned: Final = stack.enter_context(
- owned_proxy_process(
- gateway,
- tmp_path,
- _mantle_environment(mantle_port),
- config=config,
- workers=2,
- remove_environment=_INHERITED_BEARER,
- )
- )
- body: Final = _full(_OWNED_OPUS)
- proxy_url: Final = owned.gateway.client.base_url
- workers: Final = eventually(
- lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())),
- lambda pids: len(pids) == 2,
- seconds=30,
- )
- clients: Final = _clients(stack, str(proxy_url), 12)
- pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients)))
- stack.callback(release.set)
- probed: Final[SimpleQueue[int]] = SimpleQueue()
- futures: Final[tuple[Future[tuple[int, tuple[int, JsonValue] | None]], ...]] = tuple(
- pool.submit(_probed_then_counted_or_dropped, client, owned.gateway.key, body, probed) for client in clients
- )
- eventually(held.qsize, lambda size: size == len(clients), seconds=30)
- ports: Final = frozenset(probed.get_nowait() for _ in clients)
- shares: Final = {pid: _accepted_client_ports(pid, proxy_url.port or 0) & ports for pid in workers}
- assert sum(map(len, shares.values())) == len(clients), shares
- victim: Final = min((pid for pid in workers if shares[pid]), key=lambda pid: len(shares[pid]))
- psutil.Process(victim).send_signal(signal.SIGKILL)
- release.set()
- for port, result in (future.result(timeout=60) for future in futures):
- assert result == (None if port in shares[victim] else (200, _MANTLE_COUNT)), (port, result, shares)
- second_wave: Final = _clients(stack, str(proxy_url), 6)
- assert tuple(_counted_on(client, owned.gateway.key, body) for client in second_wave) == (
- (200, _MANTLE_COUNT),
- ) * len(second_wave)
- assert len(_mantle_bodies(mantle)) == len(clients) + len(second_wave)
- eventually(lambda: len(_STARTED_WORKER.findall(owned.log.read_text())), lambda started: started >= 3, 120)
- assert owned.process.poll() is None
-
-
-def _streamed_text(text: str) -> str:
- events: Final = tuple(
- _JSON_OBJECT.validate_json(line.removeprefix("data: "))
- for line in text.splitlines()
- if line.startswith("data: {")
- )
- return "".join(_delta_content(event) for event in events)
-
-
-def _delta_content(event: dict[str, JsonValue]) -> str:
- choices: Final = event.get("choices")
- if not isinstance(choices, list) or not choices:
- return ""
- delta: Final = object_value(choices[0]).get("delta")
- if not isinstance(delta, dict):
- return ""
- content: Final = delta.get("content")
- return content if isinstance(content, str) else ""
-
-
-@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"])
-def test_chat_completions_on_the_same_deployment_still_generate(
- counting_proxy: OwnedProxy, mantle_port: int, stream: bool
-) -> None:
- gateway: Final = counting_proxy.gateway
- marker: Final = uuid.uuid4().hex
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- response: Final = gateway.request(
- "POST",
- "/v1/chat/completions",
- {
- "model": model,
- "messages": [{"role": "user", "content": f"chat control marker-{marker}"}],
- "stream": stream,
- "cache": {"no-cache": True},
- },
- )
- assert response.status_code == 200, response.text
- generated: Final = _streamed_text(response.text) if stream else response.text
- assert answer(marker) in generated, response.text
- assert not stream or response.text.rstrip().endswith("data: [DONE]"), response.text
- (sent,) = runtime.drain()
- assert target_of(sent) == f"/model/{_OPUS}/{'converse-stream' if stream else 'converse'}", sent.target
- assert mantle.drain() == ()
-
-
-def test_messages_endpoint_on_the_same_deployment_still_generates(counting_proxy: OwnedProxy, mantle_port: int) -> None:
- gateway: Final = counting_proxy.gateway
- marker: Final = uuid.uuid4().hex
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- response: Final = gateway.request(
- "POST",
- "/v1/messages",
- {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": f"marker-{marker}"}]},
- )
- assert response.status_code == 200, response.text
- assert answer(marker) in response.text, response.text
- (sent,) = runtime.drain()
- assert target_of(sent) == f"/model/{_OPUS}/invoke", sent.target
- assert mantle.drain() == ()
-
-
-def test_responses_endpoint_on_the_same_deployment_still_generates(
- counting_proxy: OwnedProxy, mantle_port: int
-) -> None:
- gateway: Final = counting_proxy.gateway
- marker: Final = uuid.uuid4().hex
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- response: Final = gateway.request("POST", "/v1/responses", {"model": model, "input": f"marker-{marker}"})
- assert response.status_code == 200, response.text
- assert answer(marker) in response.text, response.text
- (sent,) = runtime.drain()
- assert target_of(sent) == f"/model/{_OPUS}/converse", sent.target
- assert mantle.drain() == ()
-
-
-def test_utils_token_counter_local_mode_never_calls_either_peer(counting_proxy: OwnedProxy, mantle_port: int) -> None:
- gateway: Final = counting_proxy.gateway
- with (
- wire_server(_runtime) as runtime,
- wire_server(_mantle(), port=mantle_port) as mantle,
- gateway.scenario() as scenario,
- ):
- model: Final = _deployment(scenario, runtime.url)
- assert _local_count(gateway, _full(model)) > 0
- assert runtime.drain() == () and mantle.drain() == ()
diff --git a/tests/integration/providers/test_chat_template_tool_count_wire.py b/tests/integration/providers/test_chat_template_tool_count_wire.py
deleted file mode 100644
index d2b85e7168f..00000000000
--- a/tests/integration/providers/test_chat_template_tool_count_wire.py
+++ /dev/null
@@ -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"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"" in sent for index in range(_FILES)) == files_kept
diff --git a/tests/integration/providers/test_openai_count_tokens_tools_wire.py b/tests/integration/providers/test_openai_count_tokens_tools_wire.py
index 16a5b8685ad..2b23ccdf4f2 100644
--- a/tests/integration/providers/test_openai_count_tokens_tools_wire.py
+++ b/tests/integration/providers/test_openai_count_tokens_tools_wire.py
@@ -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"}]
diff --git a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py
index 8de93256f6c..5fe82d37e9e 100644
--- a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py
+++ b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py
@@ -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 == ()
diff --git a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py
index d5c841d9335..73c135944e0 100644
--- a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py
+++ b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py
@@ -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)
diff --git a/tests/integration/proxy_config.yaml b/tests/integration/proxy_config.yaml
index 8ed5237868a..58312faa1f9 100644
--- a/tests/integration/proxy_config.yaml
+++ b/tests/integration/proxy_config.yaml
@@ -1,4 +1,11 @@
model_list:
+ - model_name: anthropic/claude-opus-5-5
+ litellm_params:
+ model: anthropic/claude-opus-5-5
+ api_base: http://127.0.0.1:8191
+ api_key: synthetic-anthropic-key
+ model_info:
+ token_count_framing: 0
- model_name: perplexity/pplx-decider-v1-27b
litellm_params:
model: perplexity/pplx-decider-v1-27b
diff --git a/tests/test_litellm/proxy/google_endpoints/test_endpoints.py b/tests/test_litellm/proxy/google_endpoints/test_endpoints.py
index 17af52b5579..92001118e2c 100644
--- a/tests/test_litellm/proxy/google_endpoints/test_endpoints.py
+++ b/tests/test_litellm/proxy/google_endpoints/test_endpoints.py
@@ -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()
diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_utils.py b/tests/test_litellm/proxy/proxy_server/test_routes_utils.py
index ee165d7b204..dc4e929637e 100644
--- a/tests/test_litellm/proxy/proxy_server/test_routes_utils.py
+++ b/tests/test_litellm/proxy/proxy_server/test_routes_utils.py
@@ -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)
diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py
index d7648a08b6a..e684aa55b33 100644
--- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py
+++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py
@@ -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
diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py
deleted file mode 100644
index 51249e01d4d..00000000000
--- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py
+++ /dev/null
@@ -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
diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py
index 9e3704a9df0..d1e23a17747 100644
--- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py
+++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py
@@ -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"}},
- )
diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py
index 8d2884756cb..1285a291650 100644
--- a/tests/unit/litellm_core_utils/test_token_counter.py
+++ b/tests/unit/litellm_core_utils/test_token_counter.py
@@ -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"""
diff --git a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py
index fea65a9c821..ddac561f337 100644
--- a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py
+++ b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py
@@ -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?"}],
- }
diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py
deleted file mode 100644
index 71b451f834c..00000000000
--- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py
+++ /dev/null
@@ -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
diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py
deleted file mode 100644
index 8a21027872a..00000000000
--- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py
+++ /dev/null
@@ -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,
- }
diff --git a/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py b/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py
index ddcbfc61822..77b8e66c5a8 100644
--- a/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py
+++ b/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py
@@ -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():
diff --git a/tests/unit/llms/vertex_ai/count_tokens/__init__.py b/tests/unit/llms/vertex_ai/count_tokens/__init__.py
deleted file mode 100644
index e69de29bb2d..00000000000
diff --git a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py
index 71e6ecc1f81..f9561a0d710 100644
--- a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py
+++ b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py
@@ -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
diff --git a/tests/unit/proxy/test_custom_tokenizer_bug.py b/tests/unit/proxy/test_custom_tokenizer_bug.py
index 29625ea1a32..23b6b302202 100644
--- a/tests/unit/proxy/test_custom_tokenizer_bug.py
+++ b/tests/unit/proxy/test_custom_tokenizer_bug.py
@@ -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."
diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py
index 86556f1c7bd..38f2181d4e7 100644
--- a/tests/unit/proxy/test_proxy_token_counter.py
+++ b/tests/unit/proxy/test_proxy_token_counter.py
@@ -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
diff --git a/tests/unit/test_cost_map_guard.py b/tests/unit/test_cost_map_guard.py
index 410b3acfd42..e2c54d381d8 100644
--- a/tests/unit/test_cost_map_guard.py
+++ b/tests/unit/test_cost_map_guard.py
@@ -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
diff --git a/tests/unit/test_model_prices_and_context_window.py b/tests/unit/test_model_prices_schema.py
similarity index 67%
rename from tests/unit/test_model_prices_and_context_window.py
rename to tests/unit/test_model_prices_schema.py
index 49a6fb05b15..211a43d2574 100644
--- a/tests/unit/test_model_prices_and_context_window.py
+++ b/tests/unit/test_model_prices_schema.py
@@ -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] = []