scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-05 07:37:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 233

fix(rust): reconcile restacked base and body regression

Browse source
This commit is contained in:
Devin AI 2026-10-28 14:41:47 +00:00
commit 3a1a2017cc
110 changed files with 9964 additions and 648 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
README.md
View file

@ -660,6 +660,8 @@ This covers:
We welcome contributions to LiteLLM! Whether you're fixing bugs, adding features, or improving documentation, we appreciate your help.
If you want to work closely with our team, check out [LiteLLM OSS Champions](https://www.litellm.ai/champions), our program for trusted OSS contributors with their own Slack channel and a voice in the roadmap
## Quick Start for Contributors
This requires uv to be installed.

2
cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md
View file

@ -40,6 +40,8 @@ env_http_headers = { "x-litellm-api-key" = "LITELLM_API_KEY" }
Replace `https://litellm.example.com` with your proxy URL. A local proxy normally uses `http://localhost:4000/v1`. Launch Codex, sign into ChatGPT if needed, and send a prompt with a model available to your subscription. The selected model must match the name exposed by the gateway.
The `lite` CLI can set up this provider each time it starts Codex, with no `config.toml` edit. Export `LITELLM_PROXY_URL` and `LITELLM_PROXY_API_KEY`, then run `lite codex --subscription`. Run `lite config set codex_subscription true` to make it the default.
For a managed workspace rollout, use OpenAI's [gateway configuration guide](https://learn.chatgpt.com/docs/enterprise/sign-in-with-chatgpt-through-a-gateway).
## Native compatibility and Pi

111
enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py
View file

@ -15,7 +15,11 @@ from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Final, cast
from fastapi import APIRouter, Depends, HTTPException, Request
from pydantic import TypeAdapter
from litellm_enterprise.proxy.management_endpoints.project_model_quota import (
describe_missing_quota,
missing_project_model_quota,
)
from pydantic import TypeAdapter, ValidationError
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
@ -27,6 +31,15 @@ from litellm.proxy.management.teams.dependencies import get_team_access
from litellm.proxy.management.users.service import org_admin_org_ids
from litellm.proxy.management_endpoints.common_utils import _set_object_metadata_field
from litellm.proxy.management_endpoints.team_admin_field_permissions import team_admin_may_manage_projects
from litellm.proxy.management_helpers.model_quota_policy import (
ENFORCE_PROJECT_MODEL_QUOTA_SETTING,
REQUIRE_BOTH_LIMITS,
ModelQuotaPolicy,
RouteOf,
no_route,
parse_model_quota_setting,
router_model_routes,
)
from litellm.proxy.management_helpers.utils import (
management_endpoint_wrapper,
)
@ -276,26 +289,6 @@ def _check_team_project_limits(
)
def _project_models_missing_positive_quota(
models: list[str] | None,
rpm_limits: Mapping[str, object] | None,
tpm_limits: Mapping[str, object] | None,
) -> list[str]:
"""Return the models that lack a positive `rpm` AND `tpm` quota.
A valid quota is a positive integer; null, zero, and negative are rejected
because downstream rate limiters treat a non-positive limit as immediately
exhausted (every request blocked).
"""
def _is_positive(value: object) -> bool:
return isinstance(value, int) and not isinstance(value, bool) and value > 0
rpm = rpm_limits or {}
tpm = tpm_limits or {}
return [model for model in (models or []) if not _is_positive(rpm.get(model)) or not _is_positive(tpm.get(model))]
def _router_access_group_names(llm_router: "Router | None") -> frozenset[str]:
return frozenset(llm_router.get_model_access_groups()) if llm_router is not None else frozenset()
@ -322,64 +315,79 @@ def _raise_on_project_models_expanding_at_request_time(
raise HTTPException(
status_code=400,
detail={
"error": f"models {list(expanding)} expand to multiple models at request time, so a per-model rpm/tpm quota cannot be enforced for them while 'enforce_project_model_quota' is enabled. List concrete model names instead."
"error": f"models {list(expanding)} expand to multiple models at request time, so a per-model rpm/tpm quota cannot be enforced for them while '{ENFORCE_PROJECT_MODEL_QUOTA_SETTING}' is enabled. List concrete model names instead."
},
)
def _project_model_quota_policy(general_settings: Mapping[str, object]) -> ModelQuotaPolicy | None:
try:
return parse_model_quota_setting(general_settings.get(ENFORCE_PROJECT_MODEL_QUOTA_SETTING))
except ValidationError as e:
raise HTTPException(
status_code=500,
detail={"error": f"general_settings.{ENFORCE_PROJECT_MODEL_QUOTA_SETTING} is invalid: {e}"},
) from e
def _raise_on_missing_project_model_quota(
data: NewProjectRequest | UpdateProjectRequest, access_group_names: frozenset[str] = frozenset()
data: NewProjectRequest | UpdateProjectRequest,
access_group_names: frozenset[str] = frozenset(),
policy: ModelQuotaPolicy = REQUIRE_BOTH_LIMITS,
route_of: RouteOf = no_route,
) -> None:
"""Require a positive `rpm`/`tpm` quota for every model on project CREATE.
"""Require the positive `rpm`/`tpm` quotas the policy asks of each model on project CREATE.
`model_rpm_limit`/`model_tpm_limit` are relocated into `metadata` by the request
model's `set_model_info` validator, so they are read from there.
Only invoked when `general_settings.enforce_project_model_quota` is enabled
(default off), so it is opt-in and does not change behavior for existing users.
"""
_raise_on_project_models_expanding_at_request_time(data.models, access_group_names)
metadata = data.metadata or {}
missing = _project_models_missing_positive_quota(
data.models, metadata.get("model_rpm_limit"), metadata.get("model_tpm_limit")
missing: Final = missing_project_model_quota(
data.models or (), metadata.get("model_rpm_limit"), metadata.get("model_tpm_limit"), policy, route_of
)
if not missing:
return
raise HTTPException(
status_code=400,
detail={
"error": f"models {missing} added to project without a positive rpm/tpm quota. Set a positive model_rpm_limit and model_tpm_limit for each model."
"error": f"models added to project without a positive rpm/tpm quota: {describe_missing_quota(missing)}. Set a positive value for each listed limit."
},
)
def _raise_on_missing_project_model_quota_on_update(
data: UpdateProjectRequest, existing_project: object, access_group_names: frozenset[str] = frozenset()
data: UpdateProjectRequest,
existing_project: object,
access_group_names: frozenset[str] = frozenset(),
policy: ModelQuotaPolicy = REQUIRE_BOTH_LIMITS,
route_of: RouteOf = no_route,
) -> None:
"""Require a positive `rpm`/`tpm` quota over the RESULTING state on project UPDATE.
"""Require the positive `rpm`/`tpm` quotas the policy asks of each model over the RESULTING state on project UPDATE.
`/project/update` replaces `models` and `metadata` when they are provided, so the
check runs on what the project WILL look like: a partial update that doesn't touch
models/quota keeps the existing values, while one that adds a model or clears a
model's quota must leave every resulting model with a positive limit.
Only invoked when `general_settings.enforce_project_model_quota` is enabled
(default off), so it is opt-in and does not change behavior for existing users.
model's quota must leave every resulting model with the limits it needs.
"""
resulting_models = data.models if data.models is not None else (getattr(existing_project, "models", None) or [])
resulting_metadata = (
data.metadata if data.metadata is not None else (getattr(existing_project, "metadata", None) or {})
)
_raise_on_project_models_expanding_at_request_time(resulting_models, access_group_names)
missing = _project_models_missing_positive_quota(
resulting_models, resulting_metadata.get("model_rpm_limit"), resulting_metadata.get("model_tpm_limit")
missing: Final = missing_project_model_quota(
resulting_models,
resulting_metadata.get("model_rpm_limit"),
resulting_metadata.get("model_tpm_limit"),
policy,
route_of,
)
if not missing:
return
raise HTTPException(
status_code=400,
detail={
"error": f"models {missing} would be left on the project without a positive rpm/tpm quota. Set a positive model_rpm_limit and model_tpm_limit for each model."
"error": f"models would be left on the project without a positive rpm/tpm quota: {describe_missing_quota(missing)}. Set a positive value for each listed limit."
},
)
@ -530,6 +538,7 @@ async def new_project(
"""
from litellm.proxy.proxy_server import (
general_settings,
general_settings_view,
litellm_proxy_admin_name,
llm_router,
premium_user,
@ -578,10 +587,6 @@ async def new_project(
data=data,
)
# Opt-in (default off): require rpm/tpm for every model added to the project.
if general_settings.get("enforce_project_model_quota", False):
_raise_on_missing_project_model_quota(data, _router_access_group_names(llm_router))
# Check if user has permission to create projects for this team
# only team admins can create projects for their team
has_permission = await _check_user_permission_for_project(
@ -600,6 +605,15 @@ async def new_project(
},
)
quota_policy: Final = _project_model_quota_policy(general_settings_view())
if quota_policy is not None:
_raise_on_missing_project_model_quota(
data,
_router_access_group_names(llm_router),
quota_policy,
router_model_routes(llm_router, data.team_id),
)
# Generate project_id if not provided
if data.project_id is None:
data.project_id = str(uuid.uuid4())
@ -721,6 +735,7 @@ async def update_project(
"""
from litellm.proxy.proxy_server import (
general_settings,
general_settings_view,
litellm_proxy_admin_name,
llm_router,
premium_user,
@ -828,10 +843,14 @@ async def update_project(
data=data,
)
# Opt-in (default off): require rpm/tpm for every model the update would leave on the project.
if general_settings.get("enforce_project_model_quota", False):
quota_policy: Final = _project_model_quota_policy(general_settings_view())
if quota_policy is not None:
_raise_on_missing_project_model_quota_on_update(
data, existing_project, _router_access_group_names(llm_router)
data,
existing_project,
_router_access_group_names(llm_router),
quota_policy,
router_model_routes(llm_router, target_team_id),
)
# Prepare update data

49
enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py Normal file
View file

@ -0,0 +1,49 @@
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
from litellm.proxy.management_helpers.model_quota_policy import (
ModelQuotaPolicy,
QuotaLimit,
RouteOf,
required_limits,
unmet_limits,
)
_NO_LIMITS_SET: Final[Mapping[str, object]] = MappingProxyType({})
@dataclass(frozen=True, slots=True)
class MissingModelQuota:
model: str
limits: tuple[QuotaLimit, ...]
def _limits_set_for(
model: str, rpm_set: Mapping[str, object], tpm_set: Mapping[str, object]
) -> Mapping[QuotaLimit, object]:
return MappingProxyType({"rpm": rpm_set.get(model), "tpm": tpm_set.get(model)})
def missing_project_model_quota(
models: Sequence[str],
rpm_limits: Mapping[str, object] | None,
tpm_limits: Mapping[str, object] | None,
policy: ModelQuotaPolicy,
route_of: RouteOf,
) -> tuple[MissingModelQuota, ...]:
rpm_set: Final = rpm_limits or _NO_LIMITS_SET
tpm_set: Final = tpm_limits or _NO_LIMITS_SET
gaps: Final = tuple(
MissingModelQuota(
model,
unmet_limits(required_limits(policy, model, route_of(model)), _limits_set_for(model, rpm_set, tpm_set)),
)
for model in models
)
return tuple(gap for gap in gaps if gap.limits)
def describe_missing_quota(missing: Sequence[MissingModelQuota]) -> str:
return ", ".join(f"{gap.model} ({' and '.join(f'model_{limit}_limit' for limit in gap.limits)})" for gap in missing)

19
litellm-rust/crates/inference-messages/tests/messages/response.rs
View file

@ -173,9 +173,13 @@ async fn a_json_error_envelope_is_kept_verbatim(call: MessagesCall) {
}
#[rstest]
#[case::ascii(format!("{}private-tail", "x".repeat(600)))]
#[case::unicode(format!("{}private-tail", "é".repeat(600)))]
#[tokio::test]
async fn a_long_error_body_is_truncated_at_the_documented_cap(call: MessagesCall) {
let long = "x".repeat(600);
async fn a_long_error_body_is_preserved_with_bounded_diagnostics(
call: MessagesCall,
#[case] long: String,
) {
let upstream = upstream([ResponseTemplate::new(500).set_body_string(long.clone())]).await;
let error = run(MessagesCall {
@ -186,10 +190,13 @@ async fn a_long_error_body_is_truncated_at_the_documented_cap(call: MessagesCall
.await
.expect_err("upstream error propagates");
assert_eq!(
error.upstream_response(),
Some((500, format!("{}... (truncated)", &long[..256]).as_str()))
);
assert_eq!(error.upstream_response(), Some((500, long.as_str())));
let diagnostic_body: String = long.chars().take(256).collect();
let report = litellm_error::Report::new(&error).to_string();
assert!(report.ends_with(&format!("{diagnostic_body}... (truncated)")));
assert_eq!(report.matches(&diagnostic_body).count(), 1, "{report}");
assert!(!report.contains("private-tail"), "{report}");
assert!(!format!("{error:?}").contains("private-tail"));
}
#[rstest]

5
litellm/__init__.py
View file

@ -710,6 +710,7 @@ transcribe_models: Set = set()
gigachat_models: Set = set()
llamagate_models: Set = set()
coralbricks_models: Set[str] = set()
aiand_models: Set[str] = set()
reducto_models: Set = set()
bedrock_mantle_models: Set = set()
@ -1022,6 +1023,8 @@ def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, obje
llamagate_models.add(key)
elif value.get("litellm_provider") == "coralbricks":
coralbricks_models.add(key)
elif value.get("litellm_provider") == "aiand":
aiand_models.add(key)
elif value.get("litellm_provider") == "reducto":
reducto_models.add(key)
elif value.get("litellm_provider") == "bedrock_mantle":
@ -1159,6 +1162,7 @@ model_list = list(
| reducto_models
| bedrock_mantle_models
| coralbricks_models
| aiand_models
| set(clarifai_models)
)
@ -1283,6 +1287,7 @@ def _build_models_by_provider() -> dict:
"gigachat": gigachat_models,
"llamagate": llamagate_models,
"coralbricks": coralbricks_models,
"aiand": aiand_models,
"reducto": reducto_models,
"bedrock_mantle": bedrock_mantle_models,
}

2
litellm/constants.py
View file

@ -1035,6 +1035,7 @@ openai_compatible_endpoints: Final[list[str]] = [
"https://api-inference.modelscope.cn/v1",
"https://api.moonshot.ai/v1",
"https://api.publicai.co/v1",
"https://api.aiand.com/v1",
"https://api.synthetic.new/openai/v1",
"https://serverless.tensormesh.ai/v1",
"https://api.stima.tech/v1",
@ -1100,6 +1101,7 @@ openai_compatible_providers: Final[list[str]] = [
"novita",
"meta_llama",
"publicai", # PublicAI - JSON-configured provider
"aiand", # ai& - JSON-configured provider
"synthetic", # Synthetic - JSON-configured provider
"tensormesh", # Tensormesh - JSON-configured provider
"apertis", # Apertis - JSON-configured provider

117
litellm/litellm_core_utils/fallback_generalizations.py
View file

@ -19,17 +19,11 @@ ones on key conflicts, and the caller backfills ``litellm_provider`` with the
provider it requested. If no capability rule matches, model-info resolution misses
as if no rules existed.
LEGACY-SCHEMA SHIM (temporary, until the new-schema JSON reaches main): released
proxies fetch this JSON remotely from main, whose block still ships the old schema
where a rule mixes ``litellm_provider`` with capability keys and may inherit a
parent's ``model_info`` via ``extends``. Such a legacy rule is tolerated rather
than skipped: ``extends`` is resolved once at install time (single level, against
raw parents), and the resolved rule acts as BOTH kinds, a routing rule (its
``litellm_provider`` participates in first-hit inference) and a capability rule
(its full ``model_info``, provider included, participates in the union). New-schema
rules never mix the two and never use ``extends``. A rule whose
``litellm_provider`` is not a string is invalid and is warned about and skipped
(a warning rather than a crash, for the same remote-fetch reason).
A rule whose ``model_info`` carries ``litellm_provider`` alongside any other key,
or carries it as a non-string, is neither kind: it is warned about and skipped
rather than raising, since released proxies fetch this JSON remotely. A rule
carrying ``extends`` is skipped the same way. Mixed rules and ``extends`` are the
old rule format, so their warning points the operator at a current cost map.
Rules are only consulted after exact and case-insensitive lookups miss, so an
exact cost-map entry always takes precedence over any rule.
@ -62,37 +56,8 @@ NAME_FIELD: Final = "name"
PATTERN_FIELD: Final = "pattern"
MODEL_INFO_FIELD: Final = "model_info"
PROVIDER_KEY: Final = "litellm_provider"
LEGACY_EXTENDS_FIELD: Final = "extends"
FILL_MISSING_FOR_PROVIDERS_FIELD: Final = "fill_missing_for_providers"
def _resolve_legacy_extends(rules: list) -> list:
"""Expand legacy ``extends`` inheritance so each rule's ``model_info`` is self-contained.
Compatibility shim for the old remote schema: single level, resolved against each
parent's raw ``model_info``, with the child's own keys winning on conflict. Non-dict
rules and dangling parents pass through unchanged; new-schema rules carry no
``extends`` and are untouched.
"""
base_by_name: Final = {
rule[NAME_FIELD]: rule[MODEL_INFO_FIELD]
for rule in rules
if isinstance(rule, dict)
and isinstance(rule.get(NAME_FIELD), str)
and isinstance(rule.get(MODEL_INFO_FIELD), dict)
}
def resolved(rule: object) -> object:
if not isinstance(rule, dict):
return rule
parent_name: Final = rule.get(LEGACY_EXTENDS_FIELD)
own_info: Final = rule.get(MODEL_INFO_FIELD)
parent_info: Final = base_by_name.get(parent_name) if isinstance(parent_name, str) else None
if parent_info is None or not isinstance(own_info, dict):
return rule
return {**rule, MODEL_INFO_FIELD: {**parent_info, **own_info}}
return [resolved(rule) for rule in rules]
EXTENDS_FIELD: Final = "extends"
@dataclass(frozen=True, slots=True)
@ -127,19 +92,33 @@ def _parse_fill_missing_for_providers(rule: Mapping[str, object], pattern_label:
return frozenset(raw_fill_missing_for_providers)
def _compile_rule(rule: object) -> tuple[_CompiledRule, ...]:
def _warn_old_format_rule(rule_label: str, reason: str) -> None:
verbose_logger.warning(
"LiteLLM: skipping fallback generalization rule %s, written in the old cost map rule format (%s). "
"Point LITELLM_MODEL_COST_MAP_URL at a current cost map, or with LITELLM_LOCAL_MODEL_COST_MAP set, "
"replace the bundled backup with a current one.",
rule_label,
reason,
)
def _compile_rule(rule: object) -> _CompiledRule | None:
if not isinstance(rule, dict):
return ()
return None
pattern: Final = rule.get(PATTERN_FIELD)
model_info: Final = rule.get(MODEL_INFO_FIELD)
rule_label: Final = f"{rule.get(NAME_FIELD, pattern)}"
if not isinstance(pattern, str) or not isinstance(model_info, dict):
verbose_logger.warning(
"LiteLLM: skipping malformed fallback generalization rule %s (needs string '%s' and dict '%s').",
rule.get(NAME_FIELD, pattern),
rule_label,
PATTERN_FIELD,
MODEL_INFO_FIELD,
)
return ()
return None
if EXTENDS_FIELD in rule:
_warn_old_format_rule(rule_label, f"'{EXTENDS_FIELD}' is no longer supported")
return None
try:
compiled: Final = re.compile(pattern, re.IGNORECASE)
except re.error as e:
@ -148,37 +127,28 @@ def _compile_rule(rule: object) -> tuple[_CompiledRule, ...]:
pattern,
e,
)
return ()
return None
fill_missing_for_providers: Final = _parse_fill_missing_for_providers(rule, pattern)
if fill_missing_for_providers is None:
return ()
return None
if PROVIDER_KEY not in model_info:
return (
_CapabilityRule(
pattern=compiled,
model_info=model_info,
fill_missing_for_providers=fill_missing_for_providers,
),
return _CapabilityRule(
pattern=compiled,
model_info=model_info,
fill_missing_for_providers=fill_missing_for_providers,
)
provider: Final = model_info[PROVIDER_KEY]
if not isinstance(provider, str):
verbose_logger.warning(
"LiteLLM: skipping invalid fallback generalization rule %s: '%s' in '%s' must be a string.",
rule.get(NAME_FIELD, pattern),
"LiteLLM: skipping malformed fallback generalization rule %s ('%s' must be a string).",
rule_label,
PROVIDER_KEY,
MODEL_INFO_FIELD,
)
return ()
if len(model_info) == 1:
return (_RoutingRule(pattern=compiled, provider=provider),)
return (
_RoutingRule(pattern=compiled, provider=provider),
_CapabilityRule(
pattern=compiled,
model_info=model_info,
fill_missing_for_providers=fill_missing_for_providers,
),
)
return None
if len(model_info) != 1:
_warn_old_format_rule(rule_label, f"'{PROVIDER_KEY}' must be the only key in '{MODEL_INFO_FIELD}'")
return None
return _RoutingRule(pattern=compiled, provider=provider)
class _FallbackGeneralizations:
@ -192,7 +162,7 @@ class _FallbackGeneralizations:
def set_rules(self, rules: list | None) -> None:
installed: Final = rules if isinstance(rules, list) else []
compiled: Final = tuple(kind for rule in _resolve_legacy_extends(installed) for kind in _compile_rule(rule))
compiled: Final = tuple(kind for kind in map(_compile_rule, installed) if kind is not None)
self.rules = installed
self.routing_rules = tuple(rule for rule in compiled if isinstance(rule, _RoutingRule))
self.capability_rules = tuple(rule for rule in compiled if isinstance(rule, _CapabilityRule))
@ -224,10 +194,7 @@ class _FallbackGeneralizations:
)
if not matched:
return None
fill_missing: Final[Mapping[str, object]] = {
key: value for model_info in matched for key, value in model_info.items() if key != PROVIDER_KEY
}
return fill_missing or None
return {key: value for model_info in matched for key, value in model_info.items()} or None
_registry: Final = _FallbackGeneralizations()
@ -236,10 +203,8 @@ _registry: Final = _FallbackGeneralizations()
def set_fallback_generalizations(rules: list | None) -> None:
"""Install the active rule list, compiling and classifying each rule.
Legacy ``extends`` inheritance is resolved here, once, before classification;
a legacy rule mixing ``litellm_provider`` with capability keys installs as both
kinds. Malformed and invalid-regex rules are warned about and skipped. Called
once when the model cost map is loaded (and again on any reload).
Malformed, mixed-kind, and invalid-regex rules are warned about and skipped.
Called once when the model cost map is loaded (and again on any reload).
"""
_registry.set_rules(rules)

10
litellm/llms/anthropic/count_tokens/token_counter.py
View file

@ -27,6 +27,9 @@ anthropic_count_tokens_handler: Final = AnthropicCountTokensHandler()
class AnthropicTokenCounter(BaseTokenCounter):
"""Token counter implementation for Anthropic provider using the CountTokens API."""
def __init__(self, client_oauth_token: str | None = None) -> None:
self.client_oauth_token: Final = client_oauth_token
def should_use_token_counting_api(
self,
custom_llm_provider: str | None = None,
@ -95,12 +98,17 @@ class AnthropicTokenCounter(BaseTokenCounter):
api_base: Final = litellm_params.get("api_base")
try:
auth_header: Final = await AnthropicModelInfo.aget_auth_header(
configured_auth_header: Final = await AnthropicModelInfo.aget_auth_header(
api_key=litellm_params.get("api_key") or litellm.anthropic_key or litellm.api_key,
api_base=api_base,
litellm_params=litellm_params,
allow_workload_identity=True,
)
auth_header: Final = (
configured_auth_header
if configured_auth_header is not None or self.client_oauth_token is None
else AnthropicModelInfo.get_auth_header(api_key=self.client_oauth_token, api_base=api_base)
)
if auth_header is None:
return TokenCountResponse(
total_tokens=0,

4
litellm/llms/base_llm/base_model_iterator.py
View file

@ -3,6 +3,7 @@ from abc import abstractmethod
from typing import TYPE_CHECKING, Final, cast
import litellm
from litellm.constants import STREAM_SSE_DONE_STRING
if TYPE_CHECKING:
import httpx
@ -142,7 +143,8 @@ class BaseModelResponseIterator:
def _handle_string_chunk(self, str_line: str) -> GenericStreamingChunk | ModelResponseStream:
# chunk is a str at this point
stripped_json_chunk: Final = BaseModelResponseIterator._string_to_dict_parser(str_line=str_line)
if "[DONE]" in str_line:
stripped_sse_chunk: Final = litellm.CustomStreamWrapper.strip_sse_data_from_chunk(str_line)
if stripped_sse_chunk is not None and stripped_sse_chunk.strip() == STREAM_SSE_DONE_STRING:
return GenericStreamingChunk(
text="",
is_finished=True,

48
litellm/llms/bedrock/chat/native_output_fallback.py
View file

@ -1,11 +1,14 @@
from collections.abc import Callable, Coroutine
from functools import partial
from typing import Final, TypeAlias
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.bedrock.common_utils import (
NativeOutputFallback,
is_bedrock_native_output_rejection,
is_bedrock_strict_tools_rejection,
warn_bedrock_native_output_rejected,
warn_bedrock_strict_tools_dropped,
)
from litellm.types.utils import ModelResponse
from litellm.utils import CustomStreamWrapper
@ -14,12 +17,13 @@ NativeOutputResponse: TypeAlias = ModelResponse | CustomStreamWrapper
NativeOutputAttemptResult: TypeAlias = NativeOutputResponse | Coroutine[object, object, NativeOutputResponse]
NativeOutputAttempt: TypeAlias = Callable[[NativeOutputFallback | None], NativeOutputAttemptResult]
NativeOutputFallbackFactory: TypeAlias = Callable[[], NativeOutputFallback | None]
_RetryFor: TypeAlias = Callable[[BaseLLMException], NativeOutputFallback | None]
def _fallback_for(
error: BaseLLMException, fallback: NativeOutputFallbackFactory, model: str
def _native_output_retry(
error: BaseLLMException, fallback: NativeOutputFallbackFactory | None, model: str
) -> NativeOutputFallback | None:
if not is_bedrock_native_output_rejection(error.status_code, error.message):
if fallback is None or not is_bedrock_native_output_rejection(error.status_code, error.message):
return None
retry: Final = fallback()
if retry is not None:
@ -27,16 +31,38 @@ def _fallback_for(
return retry
def _strict_tools_retry(
error: BaseLLMException, fallback: NativeOutputFallbackFactory | None, model: str
) -> NativeOutputFallback | None:
if fallback is None or not is_bedrock_strict_tools_rejection(error.status_code, error.message):
return None
retry: Final = fallback()
if retry is not None:
warn_bedrock_strict_tools_dropped(model)
return retry
def _fallback_for(
error: BaseLLMException,
fallback: NativeOutputFallbackFactory | None,
strict_tools_fallback: NativeOutputFallbackFactory | None,
model: str,
) -> NativeOutputFallback | None:
native_output_retry: Final = _native_output_retry(error, fallback, model)
if native_output_retry is not None:
return native_output_retry
return _strict_tools_retry(error, strict_tools_fallback, model)
async def _await_with_native_output_fallback(
first_attempt: Coroutine[object, object, NativeOutputResponse],
attempt: NativeOutputAttempt,
fallback: NativeOutputFallbackFactory,
model: str,
retry_for: _RetryFor,
) -> NativeOutputResponse:
try:
return await first_attempt
except BaseLLMException as error:
retry: Final = _fallback_for(error, fallback, model)
retry: Final = retry_for(error)
if retry is None:
raise
second_attempt: Final = attempt(retry)
@ -47,16 +73,20 @@ def call_with_native_output_fallback(
attempt: NativeOutputAttempt,
fallback: NativeOutputFallbackFactory | None,
model: str,
strict_tools_fallback: NativeOutputFallbackFactory | None = None,
) -> NativeOutputAttemptResult:
if fallback is None:
if fallback is None and strict_tools_fallback is None:
return attempt(None)
retry_for: Final = partial(
_fallback_for, fallback=fallback, strict_tools_fallback=strict_tools_fallback, model=model
)
try:
first_attempt: Final = attempt(None)
except BaseLLMException as error:
retry: Final = _fallback_for(error, fallback, model)
retry: Final = retry_for(error)
if retry is None:
raise
return attempt(retry)
if isinstance(first_attempt, Coroutine):
return _await_with_native_output_fallback(first_attempt, attempt, fallback, model)
return _await_with_native_output_fallback(first_attempt, attempt, retry_for)
return first_attempt

54
litellm/llms/bedrock/common_utils.py
View file

@ -362,6 +362,60 @@ class NativeOutputFallback:
litellm_params: Mapping[str, object]
_BEDROCK_STRICT_TOOLS_REJECTION: Final = re.compile(
r"Too many strict tools \(\d+\)"
r"|Schemas contains too many (?:optional parameters|parameters with union types)"
r"|[Cc]ompiled grammar (?:is too large|size)"
r"|tools\.\d+\.custom(?:\.strict)?: "
r"|Grammar compilation timed out"
)
_TOOLS_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, object], ...])
_TOOL_SPEC_ADAPTER: Final = TypeAdapter(Mapping[str, object])
def is_bedrock_strict_tools_rejection(status_code: int, error_text: str) -> bool:
return status_code == 400 and _BEDROCK_STRICT_TOOLS_REJECTION.search(error_text) is not None
def warn_bedrock_strict_tools_dropped(model: str) -> None:
verbose_logger.warning(
"Bedrock rejected the strict tool schemas for model=%s, so litellm sent the request again without `strict` "
"because drop_params is set.",
model,
)
def _without_strict(tool: Mapping[str, object]) -> Mapping[str, object]:
spec_key: Final = next((key for key in ("function", "toolSpec") if key in tool), None)
if spec_key is None:
return tool
spec: Final = _TOOL_SPEC_ADAPTER.validate_python(tool[spec_key])
return {**tool, spec_key: {key: value for key, value in spec.items() if key != "strict"}}
def _sends_strict(tool: Mapping[str, object], supports_strict_tools: bool) -> bool:
if "toolSpec" in tool:
return "strict" in _TOOL_SPEC_ADAPTER.validate_python(tool["toolSpec"])
if "function" not in tool:
return False
return supports_strict_tools and bool(_TOOL_SPEC_ADAPTER.validate_python(tool["function"]).get("strict"))
def bedrock_strict_tools_fallback(
model: str, optional_params: Mapping[str, object], litellm_params: Mapping[str, object]
) -> NativeOutputFallback | None:
tools: Final = _TOOLS_ADAPTER.validate_python(optional_params.get("tools", ()))
supports_strict_tools: Final = bedrock_converse_supports_strict_tools(
bedrock_capability_model(model, litellm_params.get("base_model"))
)
if not any(_sends_strict(tool, supports_strict_tools) for tool in tools):
return None
return NativeOutputFallback(
optional_params={**optional_params, "tools": [_without_strict(tool) for tool in tools]},
litellm_params=litellm_params,
)
def requests_output_config_format(optional_params: Mapping[str, object]) -> bool:
match optional_params:
case {"output_config": {"format": object()}}:

6
litellm/llms/openai_like/providers.json
View file

@ -107,6 +107,12 @@
"api_key_env": "AIHUBMIX_API_KEY",
"api_base_env": "AIHUBMIX_API_BASE"
},
"aiand": {
"base_url": "https://api.aiand.com/v1",
"api_key_env": "AIAND_API_KEY",
"api_base_env": "AIAND_API_BASE",
"supported_endpoints": ["/v1/chat/completions", "/v1/responses", "/v1/messages"]
},
"crusoe": {
"base_url": "https://api.inference.crusoecloud.com/v1",
"api_key_env": "CRUSOE_API_KEY",

13
litellm/main.py
View file

@ -123,6 +123,7 @@ from litellm.llms.bedrock.common_utils import (
BedrockModelInfo,
NativeOutputFallback,
bedrock_route_for_request,
bedrock_strict_tools_fallback,
deployment_model_info,
without_bedrock_native_output_opt_in,
without_bedrock_route_prefix,
@ -206,7 +207,12 @@ from litellm.utils import (
from ._logging import verbose_logger
from .caching.caching import disable_cache, enable_cache, update_cache
from .litellm_core_utils.core_helpers import as_str_mapping, normalize_drop_params, safe_deep_copy
from .litellm_core_utils.core_helpers import (
as_str_mapping,
drops_unsupported_params,
normalize_drop_params,
safe_deep_copy,
)
from .litellm_core_utils.fallback_utils import (
async_completion_with_fallbacks,
completion_with_fallbacks,
@ -4378,6 +4384,11 @@ def _complete_bedrock(ctx: CompletionDispatchContext) -> _CompletionDispatchResu
else None
),
converse_model,
strict_tools_fallback=(
partial(bedrock_strict_tools_fallback, converse_model, converse_params, dict(litellm_params))
if drops_unsupported_params(litellm_params.get("drop_params"))
else None
),
)
elif bedrock_route == "converse_like":
model = model.replace("converse_like/", "")

199
litellm/model_prices_and_context_window_backup.json
View file

@ -15764,7 +15764,6 @@
"us": 1.1
},
"supports_output_config": true,
"supports_speed": true,
"prompt_cache_min_tokens": 2048,
"source": "https://platform.claude.com/docs/en/about-claude/pricing"
},
@ -15807,7 +15806,6 @@
"us": 1.1
},
"supports_output_config": true,
"supports_speed": true,
"prompt_cache_min_tokens": 2048
},
"claude-fable-5": {
@ -37504,7 +37502,6 @@
"input_cost_per_token": 5e-08,
"input_cost_per_token_batches": 2.5e-08,
"input_cost_per_token_flex": 2.5e-08,
"input_cost_per_token_priority": 2.5e-06,
"litellm_provider": "openai",
"max_input_tokens": 272000,
"max_output_tokens": 128000,
@ -37552,7 +37549,6 @@
"deprecation_date": "2026-12-11",
"input_cost_per_token": 5e-08,
"input_cost_per_token_batches": 2.5e-08,
"input_cost_per_token_priority": 2.5e-06,
"input_cost_per_token_flex": 2.5e-08,
"litellm_provider": "openai",
"max_input_tokens": 272000,
@ -69827,6 +69823,201 @@
"supports_reasoning": true,
"supports_vision": true
},
"aiand/deepseek-ai/deepseek-v4-flash": {
"cache_read_input_token_cost": 8e-08,
"input_cost_per_token": 1.5e-07,
"litellm_provider": "aiand",
"max_input_tokens": 1048576,
"max_output_tokens": 384000,
"mode": "chat",
"output_cost_per_token": 2.5e-07,
"source": "https://api.aiand.com/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": false
},
"aiand/deepseek-ai/deepseek-v4-pro": {
"cache_read_input_token_cost": 2.5e-07,
"input_cost_per_token": 1e-06,
"litellm_provider": "aiand",
"max_input_tokens": 1048576,
"max_output_tokens": 384000,
"mode": "chat",
"output_cost_per_token": 2.5e-06,
"source": "https://api.aiand.com/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": false
},
"aiand/deepseek-ai/deepseek-v4.1-flash": {
"cache_read_input_token_cost": 2e-08,
"input_cost_per_token": 3e-07,
"litellm_provider": "aiand",
"max_input_tokens": 1048576,
"max_output_tokens": 384000,
"mode": "chat",
"output_cost_per_token": 6e-07,
"source": "https://api.aiand.com/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"aiand/google/gemma-4-31b-it": {
"cache_read_input_token_cost": 5e-08,
"input_cost_per_token": 2e-07,
"litellm_provider": "aiand",
"max_input_tokens": 262144,
"max_output_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 5e-07,
"source": "https://api.aiand.com/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"aiand/moonshotai/kimi-k2.7-code": {
"cache_read_input_token_cost": 2e-07,
"input_cost_per_token": 7.5e-07,
"litellm_provider": "aiand",
"max_input_tokens": 262144,
"max_output_tokens": 262144,
"mode": "chat",
"output_cost_per_token": 3.5e-06,
"source": "https://api.aiand.com/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"aiand/moonshotai/kimi-k3": {
"cache_read_input_token_cost": 5e-07,
"input_cost_per_token": 3e-06,
"litellm_provider": "aiand",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 1.25e-05,
"source": "https://api.aiand.com/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"aiand/motif-technologies/motif-3": {
"cache_read_input_token_cost": 2e-07,
"input_cost_per_token": 5e-07,
"litellm_provider": "aiand",
"max_input_tokens": 262144,
"max_output_tokens": 262144,
"mode": "chat",
"output_cost_per_token": 2e-06,
"source": "https://api.aiand.com/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": false
},
"aiand/openai/gpt-oss-120b": {
"cache_read_input_token_cost": 8e-08,
"input_cost_per_token": 1.5e-07,
"litellm_provider": "aiand",
"max_input_tokens": 131072,
"max_output_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 6e-07,
"source": "https://api.aiand.com/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": false
},
"aiand/qwen/qwen3.6-27b": {
"cache_read_input_token_cost": 2e-07,
"input_cost_per_token": 3.2e-07,
"litellm_provider": "aiand",
"max_input_tokens": 262144,
"max_output_tokens": 65536,
"mode": "chat",
"output_cost_per_token": 3.2e-06,
"source": "https://api.aiand.com/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"aiand/qwen/qwen3.8-27b": {
"cache_read_input_token_cost": 2e-07,
"input_cost_per_token": 4e-07,
"litellm_provider": "aiand",
"max_input_tokens": 262144,
"max_output_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 3e-06,
"source": "https://api.aiand.com/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"aiand/zai-org/glm-5.2": {
"cache_read_input_token_cost": 3e-07,
"input_cost_per_token": 1e-06,
"litellm_provider": "aiand",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 4e-06,
"source": "https://api.aiand.com/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": false
},
"aiand/zai-org/glm-5.3": {
"cache_read_input_token_cost": 3e-07,
"input_cost_per_token": 1e-06,
"litellm_provider": "aiand",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 4e-06,
"source": "https://api.aiand.com/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": false
},
"aiand/zai-org/glm-5.3-flash": {
"cache_read_input_token_cost": 3e-08,
"input_cost_per_token": 1.5e-07,
"litellm_provider": "aiand",
"max_input_tokens": 1048550,
"max_output_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 5e-07,
"source": "https://api.aiand.com/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"bespoke/nimble-latest": {
"input_cost_per_token": 0.0,
"litellm_provider": "bespoke",

26
litellm/provider_endpoints_support_backup.json
View file

@ -189,6 +189,32 @@
"generateContent": false
}
},
"aiand": {
"display_name": "ai& (`aiand`)",
"url": "https://docs.aiand.com/api/chat-completions/",
"endpoints": {
"chat_completions": true,
"responses": true,
"text_completion": true,
"embeddings": false,
"image_generations": false,
"image_edits": false,
"audio_transcriptions": false,
"audio_speech": false,
"rerank": false,
"ocr": false,
"search": false,
"moderations": false,
"realtime": false,
"video_generations": false,
"batches": false,
"messages": true,
"interactions": true,
"systemone": false,
"a2a": true,
"generateContent": true
}
},
"amazon_nova": {
"display_name": "Amazon Nova (`amazon_nova`)",
"url": "https://docs.litellm.ai/docs/providers/amazon_nova",

6
litellm/proxy/anthropic_endpoints/endpoints.py
View file

@ -50,7 +50,10 @@ from litellm.proxy.common_utils.openai_error_payload import (
with_litellm_call_id,
)
from litellm.proxy.common_utils.token_count_settings import count_source_fields
from litellm.proxy.litellm_pre_call_utils import should_auto_drop_params_for_agentic_cli
from litellm.proxy.litellm_pre_call_utils import (
forwarded_anthropic_oauth_token,
should_auto_drop_params_for_agentic_cli,
)
from litellm.router import Router
router: Final = APIRouter()
@ -412,6 +415,7 @@ async def count_tokens(
request=token_request,
call_endpoint=True,
messages_format="anthropic",
client_oauth_token=forwarded_anthropic_oauth_token(request.headers),
**_count_tokens_body_options(data, auto_drop_params),
)
except UnsupportedContentTypeError as e:

87
litellm/proxy/client/cli/commands/agents.py
View file

@ -34,7 +34,7 @@ from .claude_settings import (
load_json_or_empty,
)
from .cmd_quoting import quote_for_cmd
from .config import claude_subscription_default
from .config import subscription_default
from .pi import (
LITELLM_PROXY_API_KEY_ENV,
PI_PROVIDER_NAME,
@ -83,9 +83,11 @@ _INSTALL_DOCS: Final[dict[str, str]] = {
}
_HIDDEN_AGENTS: Final = frozenset({"pi"})
_SUBSCRIPTION_AGENTS: Final = frozenset({"claude"})
_SUBSCRIPTION_LOGINS: Final = MappingProxyType({"claude": "claude.ai", "codex": "ChatGPT"})
_SUBSCRIPTION_AGENTS: Final = frozenset(_SUBSCRIPTION_LOGINS)
CODEX_PROXY_PROVIDER: Final = "litellm"
CODEX_SUBSCRIPTION_PROVIDER: Final = "lite_subscription"
CODEX_HOME_ENV: Final = "CODEX_HOME"
CODEX_MODEL_CATALOG_FILENAME: Final = "litellm-models.json"
_CODEX_PREFLIGHT_TIMEOUT_SECONDS: Final = 10.0
@ -143,6 +145,10 @@ def build_agent_env(
travels in the x-litellm-api-key header through ANTHROPIC_CUSTOM_HEADERS, and no
ANTHROPIC_AUTH_TOKEN or ANTHROPIC_API_KEY is left in the environment, since either one makes
Claude Code stop using the subscription (and its connectors).
With subscription=True the OpenAI profile keeps Codex on its ChatGPT login: the key goes to
LITELLM_PROXY_API_KEY, which Codex sends in the x-litellm-api-key header, and OPENAI_API_KEY and
OPENAI_BASE_URL are left exactly as the caller had them, so the key is never offered to Codex as an
OpenAI credential and a child process never pairs the caller's own OpenAI key with the proxy URL.
"""
env: Final = dict(base_env)
root: Final = base_url.rstrip("/")
@ -161,21 +167,24 @@ def build_agent_env(
if ENABLE_GATEWAY_MODEL_DISCOVERY_ENV not in env:
env[ENABLE_GATEWAY_MODEL_DISCOVERY_ENV] = ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE
if PROFILE_OPENAI in profiles:
env[OPENAI_BASE_URL_ENV] = root + "/v1"
env[OPENAI_API_KEY_ENV] = api_key
if subscription:
env[LITELLM_PROXY_API_KEY_ENV] = api_key
else:
env[OPENAI_BASE_URL_ENV] = root + "/v1"
env[OPENAI_API_KEY_ENV] = api_key
if PROFILE_LITELLM in profiles:
env[LITELLM_PROXY_API_KEY_ENV] = api_key
return env
def codex_proxy_provider(base_url: str) -> Mapping[str, str | bool]:
def codex_proxy_provider(base_url: str, *, subscription: bool = False) -> Mapping[str, str | bool]:
return MappingProxyType(
{
"name": "LiteLLM proxy",
"base_url": base_url.rstrip("/") + "/v1",
"wire_api": "responses",
"supports_websockets": False,
"requires_openai_auth": False,
"requires_openai_auth": subscription,
}
)
@ -304,7 +313,7 @@ def subscription_launch_args(
return (_SETTINGS_FLAG, str(overlay)), rest
def _codex_proxy_args(base_url: str) -> list[str]:
def _codex_proxy_args(base_url: str, subscription: bool) -> list[str]:
"""Codex `-c` overrides that point it at the proxy.
Codex ignores OPENAI_BASE_URL (it always dials api.openai.com), so the env
@ -312,24 +321,35 @@ def _codex_proxy_args(base_url: str) -> list[str]:
inline; supports_websockets=false forces the HTTP/SSE Responses transport
because the proxy does not speak the Responses WebSocket protocol. The key is
read from OPENAI_API_KEY, which build_agent_env already exports.
With subscription=True the provider is `lite_subscription` with requires_openai_auth, so Codex
signs in with its own ChatGPT login, and the key travels in the x-litellm-api-key header from
LITELLM_PROXY_API_KEY. It sets no env_key, which would win over the login. A `-c` override cannot
remove keys, so the id is one no config.toml defines: `lite configure codex` writes `litellm`, and
the cookbook tells users to write `litellm_subscription`, whose env_key or headers would otherwise apply.
"""
provider: Final = f"model_providers.{CODEX_PROXY_PROVIDER}"
provider_id: Final = CODEX_SUBSCRIPTION_PROVIDER if subscription else CODEX_PROXY_PROVIDER
provider: Final = f"model_providers.{provider_id}"
credential: Final = (
f"{provider}.env_http_headers={{{json.dumps(LITELLM_KEY_HEADER)}={json.dumps(LITELLM_PROXY_API_KEY_ENV)}}}"
if subscription
else f'{provider}.env_key="{OPENAI_API_KEY_ENV}"'
)
return [
"-c",
f'model_provider="{CODEX_PROXY_PROVIDER}"',
f'model_provider="{provider_id}"',
*(
argument
for key, value in codex_proxy_provider(base_url).items()
for key, value in codex_proxy_provider(base_url, subscription=subscription).items()
for argument in ("-c", f"{provider}.{key}={json.dumps(value)}")
),
"-c",
f'{provider}.env_key="{OPENAI_API_KEY_ENV}"',
credential,
"-c",
f"{provider}.http_headers={{}}",
]
_PROXY_ARGS: Final[dict[str, Callable[[str], list[str]]]] = {
_PROXY_ARGS: Final[dict[str, Callable[[str, bool], list[str]]]] = {
"codex": _codex_proxy_args,
}
@ -410,14 +430,14 @@ _Preparer: TypeAlias = Callable[[str, str, Mapping[str, str]], Sequence[str]]
_PREPARERS: Final[Mapping[str, _Preparer]] = MappingProxyType({"pi": prepare_pi, "codex": prepare_codex})
def agent_launch_args(command: str, base_url: str) -> list[str]:
def agent_launch_args(command: str, base_url: str, *, subscription: bool = False) -> list[str]:
"""Extra CLI args an agent needs to actually honor the proxy.
Claude Code and OpenCode respect the exported env vars, so they get nothing
here; Codex needs its provider pointed via config overrides.
"""
builder: Final = _PROXY_ARGS.get(os.path.basename(command))
return builder(base_url) if builder else []
return builder(base_url, subscription) if builder else []
class ListedModel(LiteLLMBaseModel):
@ -889,11 +909,13 @@ def run_agent(
(verify_subscription if subscription else verify)(base_url, api_key)
env_before_sync: Final = base_env if base_env is not None else os.environ
synced: Final = sync_models(binary, env_before_sync, base_url, api_key, skip_verify)
synced: Final = (
_NO_EXTRA_ENV if subscription else sync_models(binary, env_before_sync, base_url, api_key, skip_verify)
)
if isinstance(synced, ModelSyncSkipped):
warn(f"litellm: not syncing {display_name} models from the proxy: {synced.reason}")
prepare: Final = preparers.get(os.path.basename(command[0]))
prepare: Final = None if subscription else preparers.get(os.path.basename(command[0]))
prepared_args: Final = tuple(prepare(base_url, api_key, env_before_sync)) if prepare is not None else ()
env: Final = MappingProxyType(
@ -903,7 +925,11 @@ def run_agent(
}
)
synced_args: Final = synced.args if isinstance(synced, ModelSyncArgs) else ()
extra_args: Final = (*agent_launch_args(command[0], base_url), *synced_args, *prepared_args)
extra_args: Final = (
*agent_launch_args(command[0], base_url, subscription=subscription),
*synced_args,
*prepared_args,
)
settings_args, forwarded = (
subscription_launch_args(
command,
@ -914,7 +940,7 @@ def run_agent(
),
write_overlay,
)
if subscription and os.path.basename(command[0]) in _SUBSCRIPTION_AGENTS
if subscription and os.path.basename(command[0]) == "claude"
else ((), tuple(command[1:]))
)
if reattach_terminal is not None:
@ -948,10 +974,19 @@ def resolve_api_key(ctx: click.Context) -> str:
_SKIP_VERIFY_HELP: Final = "Skip the pre-launch key check against the proxy."
_SUBSCRIPTION_HELP: Final = (
"Keep using your claude.ai login (and its connectors). The LiteLLM key is sent in the "
"x-litellm-api-key header instead of replacing Claude Code's own credential. "
"Make it the default with `lite config set claude_subscription true`."
_SUBSCRIPTION_HELP: Final = MappingProxyType(
{
"claude": (
"Keep using your claude.ai login (and its connectors). The LiteLLM key is sent in the "
"x-litellm-api-key header instead of replacing Claude Code's own credential. "
"Make it the default with `lite config set claude_subscription true`."
),
"codex": (
"Keep using your ChatGPT login. The LiteLLM key is sent in the "
"x-litellm-api-key header instead of replacing Codex's own credential. "
"Make it the default with `lite config set codex_subscription true`."
),
}
)
@ -964,10 +999,8 @@ def _launch(
api_key: Final = resolve_api_key(ctx)
display_name, _profiles = agent_profile(binary)
use_subscription: Final = (
subscription if subscription is not None else binary in _SUBSCRIPTION_AGENTS and claude_subscription_default()
)
login_note: Final = " (keeping your claude.ai login)" if use_subscription else ""
use_subscription: Final = subscription if subscription is not None else subscription_default(binary)
login_note: Final = f" (keeping your {_SUBSCRIPTION_LOGINS[binary]} login)" if use_subscription else ""
click.echo(f"litellm: routing {display_name} through proxy at {base_url.rstrip('/')}{login_note}")
try:
@ -989,7 +1022,7 @@ _Decorated: Final = TypeVar("_Decorated", bound=Callable[..., object] | click.Co
def _subscription_option(binary: str) -> Callable[[_Decorated], _Decorated]:
if binary not in _SUBSCRIPTION_AGENTS:
return lambda func: func
return click.option("--subscription/--no-subscription", default=None, help=_SUBSCRIPTION_HELP)
return click.option("--subscription/--no-subscription", default=None, help=_SUBSCRIPTION_HELP[binary])
def _make_agent_command(binary: str, display_name: str) -> click.Command:

14
litellm/proxy/client/cli/commands/config.py
View file

@ -14,6 +14,7 @@ from litellm.litellm_core_utils.private_json import ensure_private_dir, write_pr
HIDDEN_COMMANDS_KEY: Final = "hidden_commands"
CLAUDE_SUBSCRIPTION_KEY: Final = "claude_subscription"
CODEX_SUBSCRIPTION_KEY: Final = "codex_subscription"
PKCE_KEY: Final = "pkce"
_config_adapter: Final[TypeAdapter[Mapping[str, str]]] = TypeAdapter(Mapping[str, str])
@ -114,9 +115,15 @@ def _normalize_boolean(key: str) -> Callable[[str], str]:
return normalize
def claude_subscription_default() -> bool:
"""Whether `lite claude` keeps the claude.ai login unless told otherwise."""
return get_config_value(CLAUDE_SUBSCRIPTION_KEY) == "true"
_SUBSCRIPTION_KEYS: Final[Mapping[str, str]] = MappingProxyType(
{"claude": CLAUDE_SUBSCRIPTION_KEY, "codex": CODEX_SUBSCRIPTION_KEY}
)
def subscription_default(agent: str) -> bool:
"""Whether `lite <agent>` keeps the agent's own subscription login unless told otherwise."""
key: Final = _SUBSCRIPTION_KEYS.get(agent)
return key is not None and get_config_value(key) == "true"
def pkce_default() -> bool:
@ -129,6 +136,7 @@ _NORMALIZERS: Final[Mapping[str, Callable[[str], str]]] = MappingProxyType(
"base_url": normalize_base_url,
HIDDEN_COMMANDS_KEY: _normalize_hidden_commands,
CLAUDE_SUBSCRIPTION_KEY: _normalize_boolean(CLAUDE_SUBSCRIPTION_KEY),
CODEX_SUBSCRIPTION_KEY: _normalize_boolean(CODEX_SUBSCRIPTION_KEY),
PKCE_KEY: _normalize_boolean(PKCE_KEY),
}
)

42
litellm/proxy/litellm_pre_call_utils.py
View file

@ -363,6 +363,8 @@ _UNTRUSTED_ROOT_CONTROL_FIELDS: Final = (
"fallback_depth",
"_target_order",
"attempted_targets",
"_excluded_deployment_ids",
"_retry_skipped_deployment_ids",
"proxy_server_request",
"standard_logging_object",
"secret_fields",
@ -444,6 +446,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = (
"_pipeline_managed_guardrails",
"client_disconnected",
"error_information",
"_failover_excluded_ids",
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
)
@ -1412,12 +1415,34 @@ def _tenant_service_name(user_api_key_dict: UserAPIKeyAuth) -> str | None:
)
def litellm_auth_header_name(headers: Headers) -> str:
if "x-litellm-api-key" in headers:
return "x-litellm-api-key"
if "authorization" in headers:
return "authorization"
return "x-api-key"
def forwarded_anthropic_oauth_token(headers: Headers) -> str | None:
from litellm.llms.anthropic.common_utils import is_anthropic_oauth_key
forwarded: Final = clean_headers(headers, authenticated_with_header=litellm_auth_header_name(headers))
return next(
(
value.removeprefix("Bearer ")
for header, value in forwarded.items()
if header.lower() == "authorization" and is_anthropic_oauth_key(value)
),
None,
)
def clean_headers(
headers: Headers,
litellm_key_header_name: str | None = None,
forward_llm_provider_auth_headers: bool = False,
authenticated_with_header: str | None = None,
) -> dict:
) -> dict[str, str]:
"""
Removes litellm api key from headers
@ -2510,26 +2535,13 @@ async def add_litellm_data_to_request(
forward_llm_auth = general_settings.get("forward_llm_provider_auth_headers", False)
if not forward_llm_auth:
forward_llm_auth = getattr(litellm, "forward_llm_provider_auth_headers", False)
# Determine which header was used for authentication
# This enables forwarding provider keys (e.g., x-api-key) when they weren't used for LiteLLM auth
authenticated_with_header = None
if "x-litellm-api-key" in request.headers:
# If x-litellm-api-key is present, it was used for auth
authenticated_with_header = "x-litellm-api-key"
elif "authorization" in request.headers:
# Authorization header was used for auth
authenticated_with_header = "authorization"
else:
# x-api-key or another header was used for auth
authenticated_with_header = "x-api-key"
_headers: Final[dict[str, str]] = clean_headers(
request.headers,
litellm_key_header_name=(
general_settings.get("litellm_key_header_name") if general_settings is not None else None
),
forward_llm_provider_auth_headers=forward_llm_auth,
authenticated_with_header=authenticated_with_header,
authenticated_with_header=litellm_auth_header_name(request.headers),
)
_strip_untrusted_request_header_controls(
_headers,

296
litellm/proxy/management_endpoints/model_management_endpoints.py
View file

@ -12,6 +12,7 @@ model/{model_id}/update - PATCH endpoint for model update.
import asyncio
import datetime
import itertools
import json
from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence
from contextlib import AbstractAsyncContextManager, asynccontextmanager, suppress
@ -42,6 +43,7 @@ from litellm.litellm_core_utils.ptu_pricing import (
ptu_config_error,
)
from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload
from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping
from litellm.proxy._types import (
BlockModelRequest,
CommonProxyErrors,
@ -102,6 +104,18 @@ from litellm.proxy.management_helpers.auto_router_permissions import (
authorize_member_auto_router_write,
)
from litellm.proxy.management_helpers.model_allowlist_rename_sync import sync_model_allowlists_for_renamed_model
from litellm.proxy.management_helpers.model_quota_policy import (
ENFORCE_PROJECT_MODEL_QUOTA_SETTING,
ModelQuotaPolicy,
QuotaLimit,
deployment_route,
exempts_any_route,
limits_added_by_team_write,
parse_model_quota_setting,
required_limits,
team_write_serves,
unmet_limits,
)
from litellm.proxy.spend_tracking.ptu_feature_flag import (
PTU_COST_ATTRIBUTION_ENV_VAR,
is_ptu_cost_attribution_enabled,
@ -111,6 +125,7 @@ from litellm.repositories.chunked_in import find_many_in
from litellm.repositories.credentials_repository import CredentialsRepository
from litellm.repositories.model_repository import ModelRepository
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.project_repository import ProjectRepository
from litellm.repositories.table_repositories import ModelTableRepository
from litellm.repositories.team_repository import TeamRepository
from litellm.router import Router
@ -152,6 +167,7 @@ from litellm.types.router import (
SPECIAL_MODEL_INFO_PARAMS,
Deployment,
GenericLiteLLMParams,
LiteLLM_Params,
ModelInfo,
updateDeployment,
updateLiteLLMParams,
@ -762,26 +778,93 @@ async def _auto_router_capability_slot(
ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING: Final = "enforce_rpm_tpm_on_model_add"
def _raise_if_rate_limits_required_but_missing(*, litellm_params: GenericLiteLLMParams, enforced: bool) -> None:
"""Require both rpm and tpm (each a positive value) when the operator opts in via config.yaml.
_ROUTING_KEYS: Final = ("model", "custom_llm_provider", "api_base")
_RATE_LIMIT_KEYS: Final = ("rpm", "tpm")
_NOTHING_SENT: Final[Mapping[str, object]] = MappingProxyType({})
Off by default, so deployments keep adding models without limits. When
``enforce_rpm_tpm_on_model_add: true`` is set under general_settings, a model added
without both rpm and tpm set to a positive value is rejected rather than stored
unbounded (or effectively excluded from routing by a zero/negative limit).
"""
if not enforced:
def _merged_litellm_params(
incoming: GenericLiteLLMParams | None, existing: GenericLiteLLMParams | None, keys: tuple[str, ...]
) -> Mapping[str, object]:
sent: Final = _NOTHING_SENT if incoming is None else _dump_without_none(incoming)
stored: Final = _NOTHING_SENT if existing is None else _decrypted_litellm_params(existing)
return MappingProxyType({key: sent.get(key, stored.get(key)) for key in keys})
def _written_routing_params(
incoming: GenericLiteLLMParams | None, existing: GenericLiteLLMParams | None
) -> Mapping[str, object]:
return _merged_litellm_params(incoming, existing, _ROUTING_KEYS)
def _quota_policy(general_settings: Mapping[str, object], setting: str) -> ModelQuotaPolicy | None:
try:
return parse_model_quota_setting(general_settings.get(setting))
except ValidationError as e:
raise ProxyException(
message=f"general_settings.{setting} is invalid: {e}",
type=ProxyErrorTypes.internal_server_error.value,
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
param=None,
) from e
def _raise_if_rate_limits_required_but_missing(
*, model_name: str, litellm_params: LiteLLM_Params, team_model: bool, policy: ModelQuotaPolicy | None
) -> None:
"""Off by default, so deployments keep adding models without limits. Once the operator opts in under
general_settings, the policy decides per model name and provider which of rpm and tpm a new model needs,
and a model missing one, or carrying a zero or negative one, is rejected rather than stored unbounded
(or effectively excluded from routing)."""
if policy is None:
return
missing: Final = tuple(
field
for field, value in (("rpm", litellm_params.rpm), ("tpm", litellm_params.tpm))
if value is None or value <= 0
_raise_on_missing_rate_limits(
model_name,
unmet_limits(
required_limits(
policy, model_name, deployment_route(_written_routing_params(litellm_params, None), team_model)
),
MappingProxyType({"rpm": litellm_params.rpm, "tpm": litellm_params.tpm}),
),
)
def _raise_if_model_update_adds_missing_rate_limits(
*,
existing: Deployment,
written_name: str,
written_team_model: bool,
incoming_params: GenericLiteLLMParams | None,
remaining_params: GenericLiteLLMParams,
policy: ModelQuotaPolicy | None,
) -> None:
if policy is None:
return
before: Final = required_limits(
policy,
existing.model_info.team_public_model_name or existing.model_name,
deployment_route(
_written_routing_params(None, existing.litellm_params), existing.model_info.team_id is not None
),
)
after: Final = required_limits(
policy,
written_name,
deployment_route(_written_routing_params(incoming_params, remaining_params), written_team_model),
)
written_limits: Final = _merged_litellm_params(incoming_params, remaining_params, _RATE_LIMIT_KEYS)
checked: Final = after - before
_raise_on_missing_rate_limits(
written_name, unmet_limits(checked, MappingProxyType({limit: written_limits.get(limit) for limit in checked}))
)
def _raise_on_missing_rate_limits(model_name: str, missing: tuple[QuotaLimit, ...]) -> None:
if not missing:
return
raise ProxyException(
message=(
f"{' and '.join(missing)} must be set to a positive value when "
f"{' and '.join(missing)} must be set to a positive value for '{model_name}' when "
f"'{ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING}' is enabled in general_settings"
),
type=ProxyErrorTypes.validation_error.value,
@ -790,6 +873,113 @@ def _raise_if_rate_limits_required_but_missing(*, litellm_params: GenericLiteLLM
)
def _project_limit(metadata: object, limits_key: str, model: str) -> object:
limits: Final = metadata.get(limits_key) if is_str_keyed_mapping(metadata) else None
return limits.get(model) if is_str_keyed_mapping(limits) else None
def _projects_short_on(
projects: Sequence["prisma_models.LiteLLM_ProjectTable"], model: str, added: frozenset[QuotaLimit]
) -> tuple[str, ...]:
return tuple(
project.project_alias or project.project_id
for project in projects
if model in project.models
and unmet_limits(
added,
MappingProxyType(
{
"rpm": _project_limit(project.metadata, "model_rpm_limit", model),
"tpm": _project_limit(project.metadata, "model_tpm_limit", model),
}
),
)
)
async def _raise_if_team_write_leaves_projects_short(
*,
prisma_client: PrismaClient,
general_settings: Mapping[str, object],
llm_router: Router | None,
team_id: str | None,
public_name: str,
replaced_id: str | None,
written_params: Mapping[str, object] | None,
) -> None:
if team_id is None:
return
policy: Final = _quota_policy(general_settings, ENFORCE_PROJECT_MODEL_QUOTA_SETTING)
if policy is None or not exempts_any_route(policy):
return
projects: Final = await ProjectRepository(prisma_client).table.find_many(where={"team_id": team_id})
listed: Final = frozenset(itertools.chain.from_iterable(project.models for project in projects))
served: Final = sorted(model for model in listed if team_write_serves(llm_router, team_id, public_name, model))
shortfall: Final = next(
(
(model, added, short)
for model in served
if (
added := limits_added_by_team_write(
policy, llm_router, team_id, public_name, replaced_id, written_params, model
)
)
and (short := _projects_short_on(projects, model, added))
),
None,
)
if shortfall is None:
return
short_model, short_limits, short_projects = shortfall
raise ProxyException(
message=(
f"projects {', '.join(short_projects)} list '{short_model}' without a positive "
f"{' and '.join(sorted(short_limits))} "
f"quota, which '{ENFORCE_PROJECT_MODEL_QUOTA_SETTING}' requires once this team model change applies. "
"Set those limits on the projects first"
),
type=ProxyErrorTypes.validation_error.value,
code=status.HTTP_400_BAD_REQUEST,
param=None,
)
async def _raise_if_team_update_leaves_projects_short(
*,
prisma_client: PrismaClient,
general_settings: Mapping[str, object],
llm_router: Router | None,
existing: Deployment,
written_team_id: str | None,
written_name: str,
incoming_params: GenericLiteLLMParams | None,
remaining_params: GenericLiteLLMParams,
) -> None:
if written_team_id is None and existing.model_info.team_id is None:
return
stored_name: Final = existing.model_info.team_public_model_name or existing.model_name
await _raise_if_team_write_leaves_projects_short(
prisma_client=prisma_client,
general_settings=general_settings,
llm_router=llm_router,
team_id=written_team_id,
public_name=written_name,
replaced_id=existing.model_info.id,
written_params=_written_routing_params(incoming_params, remaining_params),
)
if (written_team_id, written_name) == (existing.model_info.team_id, stored_name):
return
await _raise_if_team_write_leaves_projects_short(
prisma_client=prisma_client,
general_settings=general_settings,
llm_router=llm_router,
team_id=existing.model_info.team_id,
public_name=stored_name,
replaced_id=existing.model_info.id,
written_params=None,
)
_PTU_PRICED_PAIR: Final = frozenset({"ptu_count", "cost_per_ptu_per_hour"})
@ -1308,6 +1498,7 @@ async def patch_model(
ProxyException: For various error conditions including authentication and database errors
"""
from litellm.proxy.proxy_server import (
general_settings_view,
litellm_proxy_admin_name,
llm_router,
premium_user,
@ -1434,6 +1625,29 @@ async def patch_model(
existing_params=remaining_params,
)
written_team_id: Final = (
patch_data.model_info.team_id if patch_data.model_info else None
) or db_model.model_info.team_id
written_name: Final = _get_public_model_name(effective_patch, db_model)
_raise_if_model_update_adds_missing_rate_limits(
existing=db_model,
written_name=written_name,
written_team_model=written_team_id is not None,
incoming_params=incoming_params,
remaining_params=remaining_params,
policy=_quota_policy(general_settings_view(), ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING),
)
await _raise_if_team_update_leaves_projects_short(
prisma_client=prisma_client,
general_settings=general_settings_view(),
llm_router=llm_router,
existing=db_model,
written_team_id=written_team_id,
written_name=written_name,
incoming_params=incoming_params,
remaining_params=remaining_params,
)
effective_params: Final = _effective_complexity_router_params(incoming_params, remaining_params)
requested_model_name: Final = patch_data.model_name
stored_model_name: str | None = None
@ -2509,6 +2723,7 @@ async def delete_model(
from litellm.proxy.proxy_server import (
MODEL_RECONCILE_LOCK,
general_settings_view,
llm_router,
premium_user,
prisma_client,
@ -2542,6 +2757,15 @@ async def delete_model(
incoming_params=None,
allow_missing_team=True,
)
await _raise_if_team_write_leaves_projects_short(
prisma_client=prisma_client,
general_settings=general_settings_view(),
llm_router=llm_router,
team_id=model_params.model_info.team_id,
public_name=model_params.model_info.team_public_model_name or model_params.model_name,
replaced_id=model_params.model_info.id,
written_params=None,
)
# update DB
if store_model_in_db is True:
@ -2708,6 +2932,8 @@ async def add_new_model(
"""
from litellm.proxy.proxy_server import (
general_settings,
general_settings_view,
llm_router,
premium_user,
prisma_client,
proxy_config,
@ -2766,8 +2992,19 @@ async def add_new_model(
)
_raise_if_rate_limits_required_but_missing(
model_name=model_params.model_name,
litellm_params=model_params.litellm_params,
enforced=bool(general_settings.get(ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING, False)),
team_model=model_params.model_info.team_id is not None,
policy=_quota_policy(general_settings_view(), ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING),
)
await _raise_if_team_write_leaves_projects_short(
prisma_client=prisma_client,
general_settings=general_settings_view(),
llm_router=llm_router,
team_id=model_params.model_info.team_id,
public_name=model_params.model_name,
replaced_id=None,
written_params=_written_routing_params(model_params.litellm_params, None),
)
clean_model_info: Final = ModelInfo.model_validate(
@ -2906,6 +3143,7 @@ async def update_model(
"""
from litellm.proxy.proxy_server import (
LITELLM_PROXY_ADMIN_NAME,
general_settings_view,
llm_router,
premium_user,
prisma_client,
@ -2970,6 +3208,30 @@ async def update_model(
incoming_params=model_params.litellm_params,
existing_params=deployment.litellm_params,
)
renamed_to: Final = (
model_params.model_name
if model_params.model_name not in (None, deployment.model_name) and deployment.model_info.team_id is None
else None
)
written_name: Final = renamed_to or deployment.model_info.team_public_model_name or deployment.model_name
_raise_if_model_update_adds_missing_rate_limits(
existing=deployment,
written_name=written_name,
written_team_model=deployment.model_info.team_id is not None,
incoming_params=model_params.litellm_params,
remaining_params=deployment.litellm_params,
policy=_quota_policy(general_settings_view(), ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING),
)
await _raise_if_team_update_leaves_projects_short(
prisma_client=prisma_client,
general_settings=general_settings_view(),
llm_router=llm_router,
existing=deployment,
written_team_id=deployment.model_info.team_id,
written_name=written_name,
incoming_params=model_params.litellm_params,
remaining_params=deployment.litellm_params,
)
effective_params: Final = _effective_complexity_router_params(
model_params.litellm_params, deployment.litellm_params
)
@ -3009,12 +3271,6 @@ async def update_model(
if value is not None or _existing_litellm_params_dict.get(key) is not None
}
renamed_to: Final = (
model_params.model_name
if model_params.model_name not in (None, deployment.model_name)
and deployment.model_info.team_id is None
else None
)
base_update: Final[PrismaCompatibleUpdateDBModel] = {
"litellm_params": json.dumps(merged_dictionary),
"updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME,

327
litellm/proxy/management_helpers/model_quota_policy.py Normal file
View file

@ -0,0 +1,327 @@
import re
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, TypeAlias
from pydantic import BaseModel, ConfigDict, Field, field_validator
import litellm
from litellm.exceptions import BadRequestError
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider, is_openai_compatible_endpoint
from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping
from litellm.llms.openai_like.json_loader import JSONProviderRegistry
from litellm.router_utils.common_utils import provider_for_generic_call
from litellm.router_utils.pattern_match_deployments import PatternMatchRouter, PatternUtils
from litellm.secret_managers.main import get_secret_str
if TYPE_CHECKING:
from litellm import Router
QuotaLimit: TypeAlias = Literal["rpm", "tpm"]
ENFORCE_PROJECT_MODEL_QUOTA_SETTING: Final = "enforce_project_model_quota"
_LIMITS: Final[tuple[QuotaLimit, ...]] = ("rpm", "tpm")
_ROUTER_PATTERNS: Final = PatternMatchRouter()
def _is_known_provider(provider: str) -> bool:
return provider in litellm.provider_list or JSONProviderRegistry.exists(provider)
class ModelQuotaPolicy(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
default: frozenset[QuotaLimit] = frozenset(_LIMITS)
providers: Mapping[str, frozenset[QuotaLimit]] = Field(default_factory=dict)
models: Mapping[str, frozenset[QuotaLimit]] = Field(default_factory=dict)
@field_validator("providers")
@classmethod
def _known_providers(cls, providers: Mapping[str, frozenset[QuotaLimit]]) -> Mapping[str, frozenset[QuotaLimit]]:
unknown: Final = sorted(provider for provider in providers if not _is_known_provider(provider))
if unknown:
raise ValueError(f"unknown provider ids {unknown}")
return providers
REQUIRE_BOTH_LIMITS: Final = ModelQuotaPolicy()
_POLICY_KEYS: Final = frozenset(ModelQuotaPolicy.model_fields)
@dataclass(frozen=True, slots=True)
class TeamDeployment:
provider: str | None
called_model: str | None
@dataclass(frozen=True, slots=True)
class ModelRoute:
providers: frozenset[str] = frozenset()
team_deployments: frozenset[TeamDeployment] = frozenset()
RouteOf: TypeAlias = Callable[[str], ModelRoute]
NO_ROUTE: Final = ModelRoute()
def parse_model_quota_setting(setting: object) -> ModelQuotaPolicy | None:
"""Any falsy value is off and any other value requires both limits, as the boolean setting always
behaved, so only a mapping that names a policy key is read as a per-provider and per-model policy."""
if not setting:
return None
if is_str_keyed_mapping(setting) and not _POLICY_KEYS.isdisjoint(setting):
return ModelQuotaPolicy.model_validate(setting)
return REQUIRE_BOTH_LIMITS
def _wildcard_matches(pattern: str, model: str) -> bool:
return re.match(_ROUTER_PATTERNS.pattern_to_regex(pattern), model) is not None
def _router_specificity(pattern: str) -> tuple[int, int]:
return PatternUtils.calculate_pattern_specificity(_ROUTER_PATTERNS.pattern_to_regex(pattern))
def _most_specific_model_pattern(model: str, entries: Mapping[str, frozenset[QuotaLimit]]) -> str | None:
matching: Final = tuple(pattern for pattern in entries if "*" in pattern and _wildcard_matches(pattern, model))
return max(matching, key=_router_specificity, default=None)
def _model_entry(entries: Mapping[str, frozenset[QuotaLimit]], model: str) -> frozenset[QuotaLimit] | None:
exact: Final = entries.get(model)
if exact is not None:
return exact
pattern: Final = _most_specific_model_pattern(model, entries)
return None if pattern is None else entries[pattern]
def _provider_limits(policy: ModelQuotaPolicy, provider: str | None) -> frozenset[QuotaLimit]:
return policy.default if provider is None else policy.providers.get(provider, policy.default)
def _team_deployment_limits(policy: ModelQuotaPolicy, deployment: TeamDeployment) -> frozenset[QuotaLimit]:
entry: Final = None if deployment.called_model is None else _model_entry(policy.models, deployment.called_model)
return _provider_limits(policy, deployment.provider) if entry is None else entry
def required_limits(policy: ModelQuotaPolicy, model: str, route: ModelRoute) -> frozenset[QuotaLimit]:
"""A model entry (exact name, then the most specific `*` pattern) wins over the providers of the
model's deployments, which win over the default. A model served by several providers needs every
limit any of them requires. A team picks its model's public name, so a team deployment matches
model entries by the model it calls instead."""
entry: Final = None if route.team_deployments else _model_entry(policy.models, model)
if entry is not None:
return entry
if not route.providers and not route.team_deployments:
return policy.default
return frozenset[QuotaLimit]().union(
*(_team_deployment_limits(policy, deployment) for deployment in route.team_deployments),
*(_provider_limits(policy, provider) for provider in route.providers),
)
def _is_positive(value: object) -> bool:
return isinstance(value, int) and not isinstance(value, bool) and value > 0
def unmet_limits(required: frozenset[QuotaLimit], values: Mapping[QuotaLimit, object]) -> tuple[QuotaLimit, ...]:
return tuple(
limit
for limit in _LIMITS
if (limit in required or values.get(limit) is not None) and not _is_positive(values.get(limit))
)
def no_route(_model: str) -> ModelRoute:
return NO_ROUTE
def _resolved_api_base(api_base: object) -> object:
if isinstance(api_base, str) and api_base.startswith("os.environ/"):
return get_secret_str(api_base)
return api_base
def _provider_behind_api_base(litellm_params: Mapping[str, object]) -> str | None:
model: Final = litellm_params.get("model")
api_base: Final = _resolved_api_base(litellm_params.get("api_base"))
if litellm_params.get("custom_llm_provider") or not isinstance(model, str) or not isinstance(api_base, str):
return None
if model.split("/", 1)[0] in litellm.provider_list or not is_openai_compatible_endpoint(api_base):
return None
try:
_, provider, _, _ = get_llm_provider(model=model, api_base=api_base)
except BadRequestError:
return None
return provider
def deployment_provider(litellm_params: object) -> str | None:
if not is_str_keyed_mapping(litellm_params):
return None
return _provider_behind_api_base(litellm_params) or provider_for_generic_call(litellm_params)
def _called_model(litellm_params: Mapping[str, object]) -> str | None:
model: Final = litellm_params.get("model")
if not isinstance(model, str):
return None
prefix, _, rest = model.partition("/")
return rest if rest and _is_known_provider(prefix) else model
def team_deployment(litellm_params: object) -> TeamDeployment:
if not is_str_keyed_mapping(litellm_params):
return TeamDeployment(provider=None, called_model=None)
return TeamDeployment(provider=deployment_provider(litellm_params), called_model=_called_model(litellm_params))
def deployment_route(litellm_params: Mapping[str, object], team_model: bool) -> ModelRoute:
if team_model:
return ModelRoute(team_deployments=frozenset((team_deployment(litellm_params),)))
provider: Final = deployment_provider(litellm_params)
return ModelRoute(providers=frozenset() if provider is None else frozenset((provider,)))
def _named_deployments(llm_router: "Router", model: str, team_id: str | None) -> tuple[Mapping[str, object], ...]:
return (
*llm_router.get_all_deployments(model_name=model, team_id=team_id),
*llm_router.get_model_list_from_routing_groups(model_name=model),
)
def _routed_deployments(llm_router: "Router", model: str, team_id: str | None) -> Sequence[Mapping[str, object]]:
target: Final = llm_router.routable_model_group(model)
return _named_deployments(llm_router, target, team_id) or llm_router.wildcard_route_deployments(target, team_id)
def _is_team_deployment(deployment: Mapping[str, object]) -> bool:
model_info: Final = deployment.get("model_info")
return is_str_keyed_mapping(model_info) and model_info.get("team_id") is not None
def router_model_routes(llm_router: "Router | None", team_id: str | None) -> RouteOf:
if llm_router is None:
return no_route
def route_of(model: str) -> ModelRoute:
deployments: Final = _routed_deployments(llm_router, model, team_id)
shared: Final = tuple(deployment for deployment in deployments if not _is_team_deployment(deployment))
return ModelRoute(
providers=frozenset(
provider
for deployment in shared
if (provider := deployment_provider(deployment.get("litellm_params"))) is not None
),
team_deployments=frozenset(
team_deployment(deployment.get("litellm_params"))
for deployment in deployments
if _is_team_deployment(deployment)
),
)
return route_of
def _deployment_id(deployment: Mapping[str, object]) -> object:
model_info: Final = deployment.get("model_info")
return model_info.get("id") if is_str_keyed_mapping(model_info) else None
def _routable_name(llm_router: "Router | None", model: str) -> str:
return model if llm_router is None else llm_router.routable_model_group(model)
def _patterns_tried_before(patterns: Sequence[str], regex: str) -> Sequence[str]:
if regex in patterns:
return patterns[: patterns.index(regex)]
specificity: Final = PatternUtils.calculate_pattern_specificity(regex)
return tuple(rival for rival in patterns if PatternUtils.calculate_pattern_specificity(rival) >= specificity)
def _team_pattern_serves(llm_router: "Router", team_id: str, pattern: str, model: str) -> bool:
regex: Final = _ROUTER_PATTERNS.pattern_to_regex(pattern)
if re.match(regex, model) is None or _named_deployments(llm_router, model, team_id):
return False
team_router: Final = llm_router.team_pattern_routers.get(team_id)
patterns: Final = () if team_router is None else tuple(team_router.patterns)
return not any(re.match(rival, model) is not None for rival in _patterns_tried_before(patterns, regex))
def team_write_serves(llm_router: "Router | None", team_id: str, public_name: str, model: str) -> bool:
"""Whether a request from the team for ``model`` reaches the team route named ``public_name``: the name
itself or an alias of it, or, for a `*` name, a model no named deployment shadows and no team pattern the
router tries first takes."""
target: Final = _routable_name(llm_router, model)
if "*" not in public_name or llm_router is None:
return target == public_name
return _team_pattern_serves(llm_router, team_id, public_name, target)
def _served_params(litellm_params: object, public_name: str, model: str) -> object:
if "*" not in public_name or not is_str_keyed_mapping(litellm_params):
return litellm_params
matched: Final = re.match(_ROUTER_PATTERNS.pattern_to_regex(public_name), model)
called: Final = litellm_params.get("model")
if matched is None or not isinstance(called, str):
return litellm_params
return MappingProxyType(
{
**litellm_params,
"model": PatternMatchRouter.set_deployment_model_name(
matched_pattern=matched, litellm_deployment_litellm_model=called
),
}
)
def team_route_after_write(
llm_router: "Router | None",
team_id: str,
public_name: str,
replaced_id: str | None,
written_params: Mapping[str, object] | None,
model: str,
) -> ModelRoute | None:
deployments: Final = () if llm_router is None else llm_router.get_all_deployments(public_name, team_id=team_id)
kept: Final = tuple(
deployment.get("litellm_params")
for deployment in deployments
if _is_team_deployment(deployment) and _deployment_id(deployment) != replaced_id
)
team_params: Final = kept if written_params is None else (*kept, written_params)
target: Final = _routable_name(llm_router, model)
if team_params:
return ModelRoute(
team_deployments=frozenset(
team_deployment(_served_params(litellm_params, public_name, target)) for litellm_params in team_params
)
)
if llm_router is None or not _routed_deployments(llm_router, model, None):
return None
return router_model_routes(llm_router, None)(model)
def limits_added_by_team_write(
policy: ModelQuotaPolicy,
llm_router: "Router | None",
team_id: str,
public_name: str,
replaced_id: str | None,
written_params: Mapping[str, object] | None,
model: str,
) -> frozenset[QuotaLimit]:
after: Final = team_route_after_write(llm_router, team_id, public_name, replaced_id, written_params, model)
if after is None:
return frozenset()
now: Final = router_model_routes(llm_router, team_id)(model)
return required_limits(policy, model, after) - required_limits(policy, model, now)
def exempts_any_route(policy: ModelQuotaPolicy) -> bool:
return any(
limits != REQUIRE_BOTH_LIMITS.default
for limits in (policy.default, *policy.providers.values(), *policy.models.values())
)

13
litellm/proxy/proxy_server.py
View file

@ -14628,6 +14628,7 @@ async def run_thread(
# )
# async def get_available_routes(user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth)):
from litellm.litellm_core_utils.tool_fields import DropPolicy, drop_policy
from litellm.llms.anthropic.count_tokens.token_counter import AnthropicTokenCounter
from litellm.llms.base_llm.base_utils import BaseTokenCounter
from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter
from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient
@ -14649,6 +14650,14 @@ from litellm.responses.litellm_completion_transformation.count_request import (
)
def _with_client_oauth_token(
counter: BaseTokenCounter | None, client_oauth_token: str | None
) -> BaseTokenCounter | None:
if client_oauth_token is None or not isinstance(counter, AnthropicTokenCounter):
return counter
return AnthropicTokenCounter(client_oauth_token=client_oauth_token)
def _get_provider_token_counter(
deployment: dict, model_to_use: str
) -> tuple[BaseTokenCounter | None, str | None, str | None]:
@ -15264,6 +15273,7 @@ async def count_request_tokens(
context_management: JsonValue = None,
request_litellm_params: Mapping[str, JsonValue] = MappingProxyType({}),
responses_request: Mapping[str, object] | None = None,
client_oauth_token: str | None = None,
) -> TokenCountResponse:
prompt: Final = request.prompt
messages: Final = request.messages
@ -15321,7 +15331,8 @@ async def count_request_tokens(
custom_llm_provider: str | None = None
if call_endpoint is True and deployment is not None:
# Auto-route to the correct provider based on model
provider_counter, _model, custom_llm_provider = _get_provider_token_counter(deployment, model_to_use)
routed_counter, _model, custom_llm_provider = _get_provider_token_counter(deployment, model_to_use)
provider_counter = _with_client_oauth_token(routed_counter, client_oauth_token)
if _model is not None:
model_to_use = _model

28
litellm/proxy/public_endpoints/provider_create_fields.json
View file

@ -1,4 +1,32 @@
[
{
"provider": "AIAND",
"provider_display_name": "ai&",
"litellm_provider": "aiand",
"credential_fields": [
{
"key": "api_base",
"label": "API Base",
"placeholder": "https://api.aiand.com/v1",
"tooltip": null,
"required": false,
"field_type": "text",
"options": null,
"default_value": null
},
{
"key": "api_key",
"label": "API Key",
"placeholder": null,
"tooltip": null,
"required": true,
"field_type": "password",
"options": null,
"default_value": null
}
],
"default_model_placeholder": "aiand/openai/gpt-oss-120b"
},
{
"provider": "SCALEDOWN",
"provider_display_name": "ScaleDown",

9
litellm/responses/litellm_completion_transformation/session_handler.py
View file

@ -54,8 +54,15 @@ def _normalize_redacted_tool_call_arguments(message: Message) -> None:
_MCP_TOOL_CALL_TYPES: Final = frozenset((CallTypes.list_mcp_tools.value, CallTypes.call_mcp_tool.value))
_FAILED_TURN_STATUS: Final = "failure"
def _is_a_failed_turn(spend_log: "SpendLogsPayload") -> bool:
return "status" in spend_log and spend_log["status"] == _FAILED_TURN_STATUS
def _is_a_conversation_turn(spend_log: "SpendLogsPayload") -> bool:
return spend_log.get("call_type") not in _MCP_TOOL_CALL_TYPES
return spend_log.get("call_type") not in _MCP_TOOL_CALL_TYPES and not _is_a_failed_turn(spend_log)
_MCP_AUTO_EXECUTE_OUTPUT_TYPES: Final = frozenset(("mcp_tools_fetched", "tool_execution_results"))

1
litellm/responses/litellm_completion_transformation/streaming_iterator.py
View file

@ -423,6 +423,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
if self._is_tool_search_call(call_id, fn_name):
self._queue_tool_search_call_added(call_id, output_index, fn_args_delta)
continue
self._started_tool_call = True
if call_id not in self._tool_args_by_call_id:
self._tool_args_by_call_id[call_id] = ""

71
litellm/responses/main.py
View file

@ -6,7 +6,7 @@ from contextlib import contextmanager
from dataclasses import dataclass
from functools import partial
from types import MappingProxyType
from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, NoReturn, Optional, TypeAlias, cast
from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, NoReturn, Optional, TypeAlias, cast, get_type_hints
import httpx
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
@ -22,6 +22,7 @@ from litellm.completion_extras.litellm_responses_transformation.transformation i
LiteLLMResponsesTransformationHandler,
)
from litellm.constants import DEFAULT_CHAT_COMPLETION_PARAM_VALUES, request_timeout
from litellm.exceptions import UnsupportedParamsError
from litellm.integrations.anthropic_cache_control_hook import (
CARRY_UNMATCHED_MESSAGE_POINTS,
AnthropicCacheControlHook,
@ -29,7 +30,7 @@ from litellm.integrations.anthropic_cache_control_hook import (
)
from litellm.litellm_core_utils.asyncify import run_async_function
from litellm.litellm_core_utils.completion_timeout import CompletionTimeout, StreamKind
from litellm.litellm_core_utils.core_helpers import normalize_drop_params
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, normalize_drop_params
from litellm.litellm_core_utils.get_llm_provider_logic import get_provider_affinity_header
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
@ -616,16 +617,16 @@ def previous_response_client_ids(previous_response_id: str | None, kwargs: Mappi
)
def _bridge_kwargs(
kwargs: Mapping[str, object],
responses_api_provider_config: BaseResponsesAPIConfig | None,
allowed_openai_params: Sequence[str] | None,
) -> Mapping[str, object]:
if responses_api_provider_config is None:
return MappingProxyType(
{key: value for key, value in kwargs.items() if key not in _RESPONSES_ONLY_REQUEST_FIELDS_NEVER_BRIDGED}
)
forwarded_keys: Final = frozenset(
_RESPONSES_REQUEST_FIELDS: Final = frozenset(get_type_hints(ResponsesAPIOptionalRequestParams))
@dataclass(frozen=True, slots=True)
class _UntranslatableResponsesFields:
fields: tuple[str, ...]
def _chat_completion_request_fields(allowed_openai_params: Sequence[str] | None) -> frozenset[str]:
return frozenset(
(
*litellm.OPENAI_CHAT_COMPLETION_PARAMS,
*DEFAULT_CHAT_COMPLETION_PARAM_VALUES,
@ -634,7 +635,49 @@ def _bridge_kwargs(
*(allowed_openai_params or ()),
)
).difference(_RESPONSES_ONLY_REQUEST_FIELDS_NEVER_BRIDGED)
return MappingProxyType({key: value for key, value in kwargs.items() if key in forwarded_keys})
_ADDITIONAL_DROP_PARAMS: Final = TypeAdapter(list[str])
def _explicitly_dropped_fields(kwargs: Mapping[str, object]) -> frozenset[str]:
try:
return frozenset(_ADDITIONAL_DROP_PARAMS.validate_python(kwargs.get("additional_drop_params")))
except ValidationError:
return frozenset()
def _bridge_kwargs(
kwargs: Mapping[str, object],
responses_api_provider_config: BaseResponsesAPIConfig | None,
allowed_openai_params: Sequence[str] | None,
) -> Mapping[str, object] | _UntranslatableResponsesFields:
forwarded_keys: Final = _chat_completion_request_fields(allowed_openai_params)
if responses_api_provider_config is not None:
return MappingProxyType({key: value for key, value in kwargs.items() if key in forwarded_keys})
candidate_keys: Final = _RESPONSES_REQUEST_FIELDS.difference(forwarded_keys, _explicitly_dropped_fields(kwargs))
untranslatable: Final = tuple(sorted(key for key in candidate_keys if kwargs.get(key) is not None))
if untranslatable and not drops_unsupported_params(kwargs.get("drop_params")):
return _UntranslatableResponsesFields(untranslatable)
dropped_keys: Final = _RESPONSES_ONLY_REQUEST_FIELDS_NEVER_BRIDGED.union(untranslatable)
return MappingProxyType({key: value for key, value in kwargs.items() if key not in dropped_keys})
def _raise_untranslatable_responses_fields(
failure: _UntranslatableResponsesFields, model: str, custom_llm_provider: str | None
) -> NoReturn:
fields: Final = list(failure.fields)
raise UnsupportedParamsError(
status_code=400,
model=model,
llm_provider=custom_llm_provider,
message=(
f"{custom_llm_provider} does not support parameters: {fields}, for model={model}. "
"The Responses API to chat completions bridge cannot translate them. "
"To drop these, set `litellm.drop_params=True` or for proxy:\n\n`litellm_settings:\n drop_params: true`\n. \n "
f"If you want to use these params dynamically send allowed_openai_params={fields} in your request."
),
)
_ResponsesCompatibilityFailure: TypeAlias = Literal["encrypted_task_unsupported"]
@ -1709,6 +1752,8 @@ def responses(
if _bridges_to_chat_completions(responses_api_provider_config, use_chat_completions_api):
bridge_kwargs: Final = _bridge_kwargs(kwargs, responses_api_provider_config, allowed_openai_params)
if isinstance(bridge_kwargs, _UntranslatableResponsesFields):
_raise_untranslatable_responses_fields(bridge_kwargs, model, custom_llm_provider)
return litellm_completion_transformation_handler.response_api_handler(
previous_response_client_ids(previous_response_id, kwargs),
model=model,

7
litellm/responses/mcp/mcp_streaming_iterator.py
View file

@ -20,6 +20,7 @@ from litellm.types.llms.openai import (
MCPListToolsCompletedEvent,
MCPListToolsFailedEvent,
MCPListToolsInProgressEvent,
ResponseInputParam,
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
ResponsesAPIStreamingResponse,
@ -290,6 +291,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
self.mcp_tools_with_litellm_proxy = mcp_tools_with_litellm_proxy or []
self.user_api_key_auth = user_api_key_auth
self.original_request_params = original_request_params or {}
self._follow_up_history: str | ResponseInputParam | None = self.original_request_params.get("input")
configured_rounds: Final = validated_max_agentic_loops(
self.original_request_params.get("max_agentic_loops"), field="litellm_params.max_agentic_loops"
)
@ -1003,9 +1005,12 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
follow_up_input: Final = LiteLLM_Proxy_MCP_Handler.create_follow_up_input(
response=self.collected_response,
tool_results=self.tool_results,
original_input=self.original_request_params.get("input"),
original_input=self._follow_up_history,
preserve_reasoning=persistence_disabled,
)
self._follow_up_history = cast( # cast-ok: helper returns valid input items
"ResponseInputParam", follow_up_input
)
# Make follow-up call with streaming
follow_up_params: Final = self.original_request_params.copy()

27
litellm/responses/streaming_iterator.py
View file

@ -365,6 +365,7 @@ class BaseResponsesAPIStreamingIterator:
"""
_pending_logging_tasks: tuple[asyncio.Task[object], ...] = ()
_started_tool_call: bool = False
_earlier_rounds_spend: BilledRoundsSpend | None = None
def __init__(
@ -488,12 +489,16 @@ class BaseResponsesAPIStreamingIterator:
_args_delta: Final = getattr(openai_responses_api_chunk, "delta", None)
if isinstance(_args_delta, str):
self._generated_tool_arguments += _args_delta
if _event_type in _CLIENT_RUN_TOOL_ARGUMENTS_DELTA_EVENTS:
self._started_tool_call = True
_stream_model_id: Final = _model_id_from_metadata(self.litellm_metadata)
if _event_type in (
ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
):
_item: Final[object] = getattr(openai_responses_api_chunk, "item", None)
if getattr(_item, "type", None) in _CLIENT_RUN_TOOL_CALL_ITEM_TYPES:
self._started_tool_call = True
if _item is not None:
ResponsesAPIRequestUtils.encode_container_id_on_output_item(
item=_item,
@ -1044,6 +1049,11 @@ class BaseResponsesAPIStreamingIterator:
if pending:
await asyncio.wait(pending)
@property
def started_tool_call(self) -> bool:
"""Whether the stream began a tool call, which a text continuation could never finish."""
return self._started_tool_call
def partial_response_on_disconnect(self) -> ResponseIncompleteEvent | None:
"""
The response.incomplete event to bill when the client went away mid-stream,
@ -1690,13 +1700,26 @@ def _usage_as_model(usage: object) -> ResponseAPIUsage | None:
return None
_TOOL_ARGUMENTS_DELTA_EVENTS: Final = frozenset(
_CLIENT_RUN_TOOL_CALL_ITEM_TYPES: Final = frozenset(
{
"function_call",
"custom_tool_call",
"mcp_approval_request",
"computer_call",
"local_shell_call",
"shell_call",
"apply_patch_call",
}
)
_CLIENT_RUN_TOOL_ARGUMENTS_DELTA_EVENTS: Final = frozenset(
{
ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA,
ResponsesAPIStreamEvents.CUSTOM_TOOL_CALL_INPUT_DELTA,
ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DELTA,
}
)
_TOOL_ARGUMENTS_DELTA_EVENTS: Final = _CLIENT_RUN_TOOL_ARGUMENTS_DELTA_EVENTS | {
ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DELTA
}
def _estimate_usage_from_text(

162
litellm/router.py
View file

@ -212,6 +212,7 @@ from litellm.router_utils.clientside_credential_handler import (
)
from litellm.router_utils.common_utils import (
MID_STREAM_CONTINUATION_KWARG,
as_retry_skipped_deployment_ids,
deployments_the_caller_may_use,
filter_team_based_models,
filter_web_search_deployments,
@ -238,11 +239,12 @@ from litellm.router_utils.cooldown_handlers import (
async_get_cooldown_deployments_with_debug_info,
get_cooldown_deployment_info,
get_cooldown_deployments,
is_advisor_orchestration_failure,
is_background_response_poll_not_found,
is_caller_scoped_auth_failure,
is_caller_timeout_408,
is_failure_outside_selected_deployment,
is_video_lookup_not_found,
mark_deployment_selection_failure,
set_cooldown_deployments,
)
from litellm.router_utils.deployment_pricing import deployment_pricing_models
@ -303,14 +305,20 @@ from litellm.router_utils.kubernetes_pod_discovery import (
)
from litellm.router_utils.pre_call_checks.continuation_prefill_check import (
MID_STREAM_CONTINUATION_MARKER,
RESPONSES_CONTINUATION_KWARG,
ContinuationOutput,
ContinuationPrefillDeploymentCheck,
ContinuedPartial,
ResponsesContinuation,
ResponsesContinuationSeam,
constrains_continuation,
constrains_responses_continuation,
continuation_deployments,
continuation_output_ceilings,
continued_text,
deployment_can_continue,
responses_continuation_input,
responses_continuation_input_for,
)
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
DeploymentAffinityCheck,
@ -455,6 +463,7 @@ from litellm.utils import (
provider_rejectable_params,
resolve_api_key_secret_reference,
set_live_deployment_replay,
token_counter,
)
from .router_utils.pattern_match_deployments import PatternMatchRouter
@ -486,7 +495,6 @@ if TYPE_CHECKING:
from litellm.types.llms.openai import (
ResponseAPIUsage,
ResponseIncompleteEvent,
ResponseInputParam,
ResponsesAPIResponse,
)
@ -684,10 +692,6 @@ def _configured_model_info(deployment: Mapping[str, object]) -> Mapping[str, obj
return _MODEL_INFO_ADAPTER.validate_python(deployment.get("model_info") or {})
def _as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]:
return tuple(item for item in value if isinstance(item, str)) if isinstance(value, tuple) else ()
def _silent_experiment_targets(silent_model: object) -> tuple[str, ...]:
if silent_model is None:
return ()
@ -1519,6 +1523,7 @@ class Router:
self.allowed_fails = allowed_fails
else:
self.allowed_fails = litellm.allowed_fails
self.explicit_allowed_fails = allowed_fails
self.cooldown_time = cooldown_time or DEFAULT_COOLDOWN_TIME_SECONDS
self.cooldown_cache = CooldownCache(cache=self.cache, default_cooldown_time=self.cooldown_time)
self.disable_cooldowns = disable_cooldowns
@ -3866,7 +3871,8 @@ class Router:
) -> Mapping[str, object]:
"""The request as the deployment received it, merged in the order _update_kwargs_with_deployment applies."""
deployment_params: Final = deployment.litellm_params.model_dump(exclude_none=True) if deployment else {}
defaults: Final = {key: value for key, value in self.default_litellm_params.items() if value is not None}
default_params: Final = _REQUEST_KWARGS_ADAPTER.validate_python(self.default_litellm_params)
defaults: Final = {key: value for key, value in default_params.items() if value is not None}
deployment_tool_choice: Final = (
{"tool_choice": deployment_params["tool_choice"]}
if deployment_params.get("tool_choice") is not None
@ -3945,67 +3951,6 @@ class Router:
extended: Final = _extend_trailing_text_block(blocks, new_text)
return [*messages[:-1], {**last, "content": extended, "prefix": True}]
@staticmethod
def _build_responses_continuation_input(
input_val: Union[str, "ResponseInputParam"] | None,
generated_content: str,
) -> "ResponseInputParam":
"""
Convert Responses-API input + partial assistant output into a
continuation input that asks the fallback model to pick up where the
prior assistant message stopped.
Best effort across providers. The chat-completions path uses
Anthropic's `prefix: True` prefill trick on the assistant message;
the Responses-API input schema has no direct equivalent, so we
append an instruction (developer role) plus a prior assistant
message containing the partial output. Providers without prefill
semantics (OpenAI, Vertex) treat this as conversational context
and may regenerate — same trade-off as the chat-completions path
for non-Anthropic fallbacks.
"""
# base/continuation are List[Any] because ResponseInputParam items
# are a wide Union of TypedDicts (EasyInputMessageParam, Message,
# ResponseOutputMessageParam, ...) — annotating as List[Dict[str, Any]]
# rejects the list() spread of input_val. We cast the combined list to
# ResponseInputParam at the return.
base: list[object]
if isinstance(input_val, str):
base = [
{
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": input_val}],
}
]
elif isinstance(input_val, list):
base = list(input_val)
else:
base = []
continuation: Final[list[object]] = [
{
"type": "message",
"role": "developer",
"content": [
{
"type": "input_text",
"text": (
"The previous assistant response was interrupted "
"mid-stream. Continue exactly where it stopped — "
"do not repeat any of its content. Your response "
"must read as a seamless continuation."
),
}
],
},
{
"type": "message",
"role": "assistant",
"content": [{"type": "output_text", "text": generated_content}],
},
]
return cast("ResponseInputParam", base + continuation)
async def _aresponses_streaming_iterator(
self,
response: "BaseResponsesAPIStreamingIterator",
@ -4027,9 +3972,12 @@ class Router:
Full parity with the chat-completions path:
- Pre-first-chunk: retry with the original input unchanged.
- Partial content: inject a developer instruction + prior
assistant message carrying the generated text so the fallback
model continues rather than restarts.
- Partial content: continue from the generated text, as an
assistant prefill on a fallback that takes one and as assistant
history plus a user request to continue on any other. A request
whose output is constrained (a JSON schema, a forced tool), or a
stream that began a tool call, is not continued, and the primary's
error surfaces instead.
- Usage combining: merge partial-stream usage onto the fallback's
response.completed event so accounting reflects both attempts.
- Stream cleanup: shielded aclose() on both source and fallback
@ -4218,6 +4166,42 @@ class Router:
wrapper: Final = FallbackResponsesStreamWrapper(stream_with_fallbacks())
return wrapper
def _responses_continuation(
self,
e: "MidStreamFallbackError",
source_iterator: "BaseResponsesAPIStreamingIterator",
request_kwargs: Mapping[str, object],
) -> ResponsesContinuation | None:
"""What a fallback continues from once a Responses stream broke after text, raising the primary's error
when the request constrains its output or the stream began a tool call, since a continued JSON document or
tool call would not parse, and when a content filter cut it, so no other model finishes filtered text."""
if not e.generated_content or e.is_pre_first_chunk:
return None
if source_iterator.started_tool_call or isinstance(e.original_exception, litellm.ContentPolicyViolationError):
Router._raise_original_mid_stream_error(e)
cut_deployment_id: Final = (get_hidden_params(source_iterator) or {}).get("model_id")
cut_deployment: Final = (
self.get_deployment(model_id=cut_deployment_id) if isinstance(cut_deployment_id, str) else None
)
if constrains_responses_continuation(self._params_sent_to_deployment(request_kwargs, cut_deployment)):
Router._raise_original_mid_stream_error(e)
earlier: Final = request_kwargs.get(RESPONSES_CONTINUATION_KWARG)
if isinstance(earlier, ResponsesContinuation):
return earlier.extended(e.generated_content)
return ResponsesContinuation(request_kwargs.get("input"), e.generated_content)
@staticmethod
def _responses_continuation_ceilings(
e: "MidStreamFallbackError", request_kwargs: Mapping[str, object]
) -> Mapping[str, object]:
"""The caller's output ceilings less what the cut stream already emitted, raising the primary's error once
nothing is left."""
emitted_tokens: Final = token_counter(model=e.model or "", text=e.generated_content)
reduced: Final = continuation_output_ceilings(request_kwargs, emitted_tokens)
if reduced is None:
Router._raise_original_mid_stream_error(e)
return reduced
async def _aresponses_fallback_attempt(
self,
e: "MidStreamFallbackError",
@ -4236,6 +4220,11 @@ class Router:
"""
from litellm.exceptions import MidStreamFallbackError
continuation: Final = self._responses_continuation(e, source_iterator, initial_kwargs)
reduced_ceilings: Final[Mapping[str, object]] = (
Router._responses_continuation_ceilings(e, initial_kwargs) if continuation is not None else {}
)
seam: Final = ResponsesContinuationSeam(continuation.emitted_text if continuation is not None else "")
partial_usage: Final = Router._extract_partial_responses_usage(source_iterator)
fallback_response = None # rebind-ok: pre-init so finally can close it if a fallback was actually attempted
fallback_yielded = False # rebind-ok: flipped on the first fallback item so a fallback that dies before its first event still replays the primary's held announcement
@ -4253,11 +4242,10 @@ class Router:
initial_kwargs["original_function"] = ( # rebind-ok: the fallback chain re-enters on the same kwargs
self._ageneric_api_call_with_fallbacks_responses_attempt
)
if e.generated_content and not e.is_pre_first_chunk:
initial_kwargs["input"] = Router._build_responses_continuation_input( # rebind-ok: fallback hop input
initial_kwargs.get("input"),
e.generated_content,
)
if continuation is not None:
initial_kwargs.update(reduced_ceilings)
initial_kwargs[RESPONSES_CONTINUATION_KWARG] = continuation # rebind-ok: shaped per fallback deployment
initial_kwargs["input"] = responses_continuation_input(continuation) # rebind-ok: fallback hop input
# The Responses-API path stores observability metadata
# under "litellm_metadata" (not the default "metadata") —
# see _ageneric_api_call_with_fallbacks. Mirroring that
@ -4286,7 +4274,7 @@ class Router:
)
prepared_fallback_hidden_params: Final = adopt_headers(fallback_response)
if hasattr(fallback_response, "__aiter__"):
async for fallback_item in fallback_response:
async for fallback_item in seam.splice(fallback_response):
Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params)
if partial_usage is not None:
Router._combine_responses_fallback_usage(fallback_item, partial_usage)
@ -4809,11 +4797,12 @@ class Router:
splits_stream_timeout: bool = False,
) -> None:
"""
4 jobs:
5 jobs:
- Adds selected deployment, model_info and api_base to kwargs["metadata"] (used for logging)
- Adds default litellm params to kwargs, if set.
- Merges tools from deployment with request (proxy-configured tools + request tools).
- On a fallback hop, drops the encrypted reasoning this deployment cannot decrypt, keeping its summary.
- On a Responses continuation, sends the partial response as prefill where this deployment takes one.
"""
for key in self._forwarded_alias_marker_keys_the_deployment_sets(
deployment=deployment, forwarded_keys=kwargs.pop(_ALIAS_MARKER_FORWARDED_PARAMS_KWARG, ())
@ -5017,6 +5006,11 @@ class Router:
(_FALLBACK_HOP_ADAPTER.validate_python(deployment),),
unmarked_origin=deployment_that_just_failed,
)
if RESPONSES_CONTINUATION_KWARG in hop_kwargs:
del kwargs[RESPONSES_CONTINUATION_KWARG] # rebind-ok: this method prepares the call's kwargs in place
kwargs["input"] = responses_continuation_input_for( # rebind-ok: the deployment's continuation input
hop_kwargs, _FALLBACK_HOP_ADAPTER.validate_python(deployment)
)
def _get_async_openai_model_client(self, deployment: dict, kwargs: dict):
"""
@ -8940,7 +8934,7 @@ class Router:
return ()
if litellm.should_retry(status_code):
return ()
already_skipped_ids: Final = _as_retry_skipped_deployment_ids(already_skipped)
already_skipped_ids: Final = as_retry_skipped_deployment_ids(already_skipped)
skipped: Final = tuple(sorted(frozenset((*already_skipped_ids, failed_deployment_id))))
verbose_router_logger.debug(
"Retry skips deployments that already answered %s to this request: %s", status_code, skipped
@ -9579,10 +9573,10 @@ class Router:
try:
exception: Final = kwargs.get("exception", None)
if is_advisor_orchestration_failure(exception):
if is_failure_outside_selected_deployment(exception):
verbose_router_logger.debug(
"Router: Exiting 'deployment_callback_on_failure' without cooldown. "
"Failure originated from advisor orchestration, not the selected deployment."
"Failure came from advisor orchestration or deployment selection, not the selected deployment."
)
return False
@ -15184,7 +15178,7 @@ class Router:
## RETRY SKIP ## -> drop deployments that already refused this request with a
## non-retryable status, unless that leaves nothing, so the caller still gets
## the provider's own error instead of a no-deployments error.
_retry_skipped_deployment_ids: Final = _as_retry_skipped_deployment_ids(
_retry_skipped_deployment_ids: Final = as_retry_skipped_deployment_ids(
request_kwargs.pop("_retry_skipped_deployment_ids", None) if request_kwargs else None
)
healthy_deployments = (
@ -15473,6 +15467,7 @@ class Router:
phase_event(_DEPLOYMENT_SELECTED_EVENT, pick_attributes)
return deployment
except Exception as e:
mark_deployment_selection_failure(e)
traceback_exception: Final = traceback.format_exc()
# if router rejects call -> log to langfuse/otel/etc.
if request_kwargs is not None:
@ -15641,6 +15636,7 @@ class Router:
return deployment
except Exception as e:
mark_deployment_selection_failure(e)
traceback_exception: Final = traceback.format_exc()
if request_kwargs is not None:
logging_obj: Final = request_kwargs.get("litellm_logging_obj", None)
@ -16414,7 +16410,7 @@ class Router:
)
## RETRY SKIP ## -> see async counterpart in async_get_healthy_deployments.
_retry_skipped_deployment_ids: Final = _as_retry_skipped_deployment_ids(
_retry_skipped_deployment_ids: Final = as_retry_skipped_deployment_ids(
request_kwargs.pop("_retry_skipped_deployment_ids", None) if request_kwargs else None
)
healthy_deployments = (

28
litellm/router_utils/common_utils.py
View file

@ -4,7 +4,7 @@ from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, TypeVar
from pydantic import TypeAdapter
from pydantic import TypeAdapter, ValidationError
if TYPE_CHECKING:
from litellm.types.llms.openai import OpenAIFileObject
@ -38,6 +38,7 @@ _OBJECT_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[object, object]]] = TypeAdapt
_DEPLOYMENT_ROWS_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter(
tuple[Mapping[str, object], ...]
)
_DEPLOYMENT_IDS_ADAPTER: Final = TypeAdapter(tuple[str, ...])
def with_admin_api_key_secret_reference(deployment: Deployment) -> Deployment:
@ -57,6 +58,31 @@ def without_router_only_kwargs(kwargs: Mapping[str, _V]) -> dict[str, _V]:
return {key: value for key, value in kwargs.items() if key not in ROUTER_ONLY_CALL_KWARGS}
def as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]:
if not isinstance(value, tuple):
return ()
items: Final[tuple[object, ...]] = value
return tuple(item for item in items if isinstance(item, str))
def as_excluded_deployment_ids(value: object) -> tuple[str, ...]:
if value is None:
return ()
try:
return _DEPLOYMENT_IDS_ADAPTER.validate_python(value)
except ValidationError:
return ()
def deployment_ids_failed_in_request(request_kwargs: Mapping[str, object]) -> frozenset[str]:
return frozenset(
(
*as_excluded_deployment_ids(request_kwargs.get("_excluded_deployment_ids")),
*as_retry_skipped_deployment_ids(request_kwargs.get("_retry_skipped_deployment_ids")),
)
)
def request_user_api_key_auth(request_kwargs: Mapping[str, object]) -> object:
"""The caller's auth object, from whichever metadata bucket this surface writes to."""
metadata_value: Final = request_kwargs.get("metadata")

20
litellm/router_utils/cooldown_handlers.py
View file

@ -57,6 +57,7 @@ else:
Span = Any
_ADVISOR_ORCHESTRATION_FAILURE_ATTR: Final = "_litellm_advisor_orchestration_failure"
_DEPLOYMENT_SELECTION_FAILURE_ATTR: Final = "_litellm_deployment_selection_failure"
_CREDENTIAL_VALUES_ADAPTER: Final = TypeAdapter(Mapping[str, object])
@ -79,6 +80,16 @@ def is_advisor_orchestration_failure(exception: BaseException | None) -> bool:
return bool(getattr(exception, _ADVISOR_ORCHESTRATION_FAILURE_ATTR, False))
def mark_deployment_selection_failure(exception: BaseException) -> None:
setattr(exception, _DEPLOYMENT_SELECTION_FAILURE_ATTR, True)
def is_failure_outside_selected_deployment(exception: BaseException | None) -> bool:
return is_advisor_orchestration_failure(exception) or bool(
getattr(exception, _DEPLOYMENT_SELECTION_FAILURE_ATTR, False)
)
_BACKGROUND_RESPONSE_POLL_CALL_ORIGINS: Final[frozenset[InternalCallOrigin]] = frozenset(
{BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN}
)
@ -781,14 +792,17 @@ def _is_allowed_fails_set_on_router(
litellm_router_instance: LitellmRouter,
) -> bool:
"""
Check if Router.allowed_fails is set or is Non-default Value
Check if Router.allowed_fails was passed to the Router or is a Non-default Value
Returns:
- True if Router.allowed_fails is set or is Non-default Value
- False if Router.allowed_fails is None or is Default Value
- True if allowed_fails was passed to the Router, even when it equals the default
- True if Router.allowed_fails is a Non-default Value
- False if Router.allowed_fails is None, or is the Default Value and was not passed
"""
if litellm_router_instance.allowed_fails is None:
return False
if litellm_router_instance.explicit_allowed_fails is not None:
return True
if litellm_router_instance.allowed_fails != litellm.allowed_fails:
return True
return False

149
litellm/router_utils/pre_call_checks/continuation_prefill_check.py
View file

@ -1,12 +1,13 @@
"""Keeps a mid-stream continuation on deployments whose model supports assistant prefill and whose own params leave
the output unconstrained."""
the output unconstrained, and shapes a Responses API continuation for the deployment that serves it."""
from collections.abc import AsyncIterable, AsyncIterator, Mapping, Sequence
import json
from collections.abc import AsyncIterable, AsyncIterator, Callable, Mapping, Sequence
from dataclasses import dataclass
from itertools import accumulate
from os.path import commonprefix
from types import MappingProxyType
from typing import Final
from typing import Final, TypeVar
from pydantic import TypeAdapter, ValidationError
@ -14,7 +15,7 @@ import litellm
from litellm.integrations.custom_logger import CustomLogger, Span
from litellm.litellm_core_utils.rules import Rules
from litellm.router_utils.common_utils import MID_STREAM_CONTINUATION_KWARG
from litellm.types.llms.openai import AllMessageValues
from litellm.types.llms.openai import AllMessageValues, OutputTextDeltaEvent
from litellm.types.utils import ModelResponse, ModelResponseStream, Usage
from litellm.utils import supports_assistant_prefill
@ -28,7 +29,9 @@ MID_STREAM_CONTINUATION_MARKER: Final = _ContinuationMarker()
_STR_KEYED_DICT_ADAPTER: Final = TypeAdapter(dict[str, object])
_NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({})
_UNCONSTRAINED_RESPONSE_FORMAT: Final[Mapping[str, str]] = MappingProxyType({"type": "text"})
_OUTPUT_CEILING_KEYS: Final = ("max_tokens", "max_completion_tokens")
_OUTPUT_CEILING_KEYS: Final = ("max_tokens", "max_completion_tokens", "max_output_tokens")
_QUOTED_ENDING_CHARS: Final = 100
RESPONSES_CONTINUATION_KWARG: Final = "_responses_continuation"
def _str_keyed(value: object) -> Mapping[str, object]:
@ -38,32 +41,55 @@ def _str_keyed(value: object) -> Mapping[str, object]:
return _NO_FIELDS
def deployment_supports_prefill(deployment: object) -> bool:
def _underlying_model(deployment: object) -> str | None:
deployment_map: Final = _str_keyed(deployment)
model_info: Final = _str_keyed(deployment_map.get("model_info"))
declared: Final = model_info.get("supports_assistant_prefill")
if isinstance(declared, bool):
return declared
litellm_params: Final = _str_keyed(deployment_map.get("litellm_params"))
candidates: Final = (model_info.get("base_model"), litellm_params.get("base_model"), litellm_params.get("model"))
model: Final = next((candidate for candidate in candidates if isinstance(candidate, str) and candidate), None)
return next((candidate for candidate in candidates if isinstance(candidate, str) and candidate), None)
def deployment_supports_prefill(deployment: object) -> bool:
declared: Final = _str_keyed(_str_keyed(deployment).get("model_info")).get("supports_assistant_prefill")
if isinstance(declared, bool):
return declared
model: Final = _underlying_model(deployment)
return model is not None and supports_assistant_prefill(model=model)
def constrains_continuation(params: Mapping[str, object]) -> bool:
def _forces_tool_call(tool_choice: object) -> bool:
return tool_choice == "required" or isinstance(tool_choice, Mapping)
def _forces_responses_tool_call(tool_choice: object) -> bool:
choice: Final = _str_keyed(tool_choice)
leaves_text_open: Final = choice.get("type") == "allowed_tools" and choice.get("mode") == "auto"
return _forces_tool_call(tool_choice) and not leaves_text_open
def _constrains(params: Mapping[str, object], forces_tool_call: Callable[[object], bool]) -> bool:
sent: Final = {**params, **_str_keyed(params.get("extra_body"))}
choice_count: Final = sent.get("n")
response_format: Final = sent.get("response_format")
tool_choice: Final = sent.get("tool_choice")
text_format: Final = _str_keyed(sent.get("text")).get("format")
return (
(isinstance(choice_count, int) and choice_count > 1)
or (response_format is not None and response_format != _UNCONSTRAINED_RESPONSE_FORMAT)
or tool_choice == "required"
or isinstance(tool_choice, Mapping)
or (text_format is not None and text_format != _UNCONSTRAINED_RESPONSE_FORMAT)
or sent.get("text_format") is not None
or forces_tool_call(sent.get("tool_choice"))
or sent.get("merge_reasoning_content_in_choices") is True
)
def constrains_continuation(params: Mapping[str, object]) -> bool:
return _constrains(params, _forces_tool_call)
def constrains_responses_continuation(params: Mapping[str, object]) -> bool:
return _constrains(params, _forces_responses_tool_call)
def deployment_can_continue(deployment: object) -> bool:
litellm_params: Final = _str_keyed(_str_keyed(deployment).get("litellm_params"))
return deployment_supports_prefill(deployment) and not constrains_continuation(litellm_params)
@ -99,6 +125,7 @@ def _delta_text(chunk: ModelResponseStream) -> str:
_CONTENT_BLOCKS_ADAPTER: Final = TypeAdapter(list[dict[str, object]])
_INPUT_ITEMS_ADAPTER: Final = TypeAdapter(list[object])
def _prefill_text(prefill_message: Mapping[str, object]) -> str:
@ -114,7 +141,7 @@ def _prefill_text(prefill_message: Mapping[str, object]) -> str:
)
def _resendable_whitespace(prefill: str) -> str:
def resendable_whitespace(prefill: str) -> str:
return prefill[len(prefill.rstrip()) :] if prefill.strip() else ""
@ -124,7 +151,7 @@ def continued_text(prefill_message: Mapping[str, object], generated_content: str
return None
echo: Final = len(prefill) if prefill and generated_content.startswith(prefill) else 0
after_echo: Final = generated_content[echo:]
return after_echo[len(commonprefix((after_echo, _resendable_whitespace(prefill)))) :]
return after_echo[len(commonprefix((after_echo, resendable_whitespace(prefill)))) :]
def _reduced_ceilings(params: Mapping[str, object], emitted_tokens: int) -> Mapping[str, int]:
@ -211,7 +238,7 @@ class ContinuationOutput:
self._rules: Final = Rules()
self._held: tuple[ModelResponseStream, ...] = ()
self._echo_settled = not self._prefill_text
self._seam = _resendable_whitespace(self._prefill_text)
self._seam = resendable_whitespace(self._prefill_text)
async def splice(self, stream: AsyncIterable[object]) -> AsyncIterator[object]:
async for item in stream:
@ -262,3 +289,91 @@ class ContinuationOutput:
self._text += _delta_text(chunk)
self._rules.post_call_rules(input=self._text, model=self._model)
return chunks
@dataclass(frozen=True, slots=True)
class ResponsesContinuation:
prior_input: object
emitted_text: str
def extended(self, generated_content: str) -> "ResponsesContinuation":
return ResponsesContinuation(self.prior_input, f"{self.emitted_text}{generated_content}")
def deployment_takes_responses_prefill(deployment: object) -> bool:
"""Claude takes a trailing assistant turn as prefill with no `prefix` flag, which a Responses input item cannot
carry."""
model: Final = _underlying_model(deployment)
return model is not None and "claude" in model.lower() and deployment_supports_prefill(deployment)
def _continuation_request(emitted_text: str) -> Mapping[str, object]:
seam: Final = (
"It already ends with whitespace, so do not add another space before your first word."
if emitted_text[-1:].isspace()
else "It ends without trailing whitespace, so begin with a space if your first word starts a new word."
)
ending: Final = json.dumps(emitted_text[-_QUOTED_ENDING_CHARS:], ensure_ascii=False)
instruction: Final = (
f"The previous assistant response was interrupted mid-stream. It ended with {ending}. Continue exactly where "
"it stopped, so your first characters are the ones that would have come right after it, and do not repeat "
f"any of its content. {seam} Your response must read as a seamless continuation."
)
return {"type": "message", "role": "user", "content": [{"type": "input_text", "text": instruction}]}
def _input_items(request_input: object) -> tuple[object, ...]:
if isinstance(request_input, str):
return ({"type": "message", "role": "user", "content": [{"type": "input_text", "text": request_input}]},)
try:
return tuple(_INPUT_ITEMS_ADAPTER.validate_python(request_input))
except ValidationError:
return ()
def responses_continuation_input(continuation: ResponsesContinuation) -> Sequence[object]:
"""The prior input, the partial response as assistant history, and a user request to continue it."""
partial: Final = {
"type": "message",
"role": "assistant",
"content": [{"type": "output_text", "text": continuation.emitted_text}],
}
return [*_input_items(continuation.prior_input), partial, _continuation_request(continuation.emitted_text)]
def responses_continuation_input_for(request_kwargs: Mapping[str, object], deployment: Mapping[str, object]) -> object:
"""The continuation input as the deployment receives it, without the request to continue where the deployment
takes the partial response as prefill. It reads the input the router already stripped of reasoning this
deployment cannot decrypt, never the continuation's prior input."""
request_input: Final = request_kwargs.get("input")
continuation: Final = request_kwargs.get(RESPONSES_CONTINUATION_KWARG)
if not isinstance(continuation, ResponsesContinuation):
return request_input
items: Final = _input_items(request_input)
if items[-1:] != (_continuation_request(continuation.emitted_text),) or not deployment_takes_responses_prefill(
deployment
):
return request_input
return list(items[:-1])
_StreamItem = TypeVar("_StreamItem")
class ResponsesContinuationSeam:
"""Drops whitespace a fallback sends again at the start of its text when the client already has it, since
Anthropic strips a prefill's trailing whitespace and the model then writes it back."""
def __init__(self, emitted_text: str) -> None:
self._seam = resendable_whitespace(emitted_text)
async def splice(self, stream: AsyncIterable[_StreamItem]) -> AsyncIterator[_StreamItem]:
async for item in stream:
yield self._past(item)
def _past(self, item: _StreamItem) -> _StreamItem:
if not self._seam or not isinstance(item, OutputTextDeltaEvent):
return item
matched: Final = len(commonprefix((item.delta, self._seam)))
self._seam = self._seam[matched:] if matched == len(item.delta) else ""
return item.model_copy(update={"delta": item.delta[matched:]}) if matched else item

8
litellm/router_utils/pre_call_checks/deployment_affinity_check.py
View file

@ -31,6 +31,7 @@ from litellm.caching.dual_cache import DualCache
from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY
from litellm.integrations.custom_logger import CustomLogger, Span
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.router_utils.common_utils import deployment_ids_failed_in_request
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import CallTypes
@ -429,6 +430,9 @@ class DeploymentAffinityCheck(CustomLogger):
session_affinity_active: Final = (
enable_session_id or self._get_marker_session_affinity_ttl(request_kwargs=request_kwargs) is not None
)
failed_deployment_ids: Final = deployment_ids_failed_in_request(
cast(Mapping[str, object], request_kwargs) # cast-ok: request kwargs keys are param names
)
user_key: Final = (
self.get_user_key_from_request_kwargs(request_kwargs=request_kwargs)
if (session_affinity_active or enable_user_key)
@ -458,7 +462,7 @@ class DeploymentAffinityCheck(CustomLogger):
elif isinstance(session_cache_result, str):
session_model_id = session_cache_result
if session_model_id:
if session_model_id and session_model_id not in failed_deployment_ids:
session_deployment: Final = self._find_deployment_by_model_id(
healthy_deployments=typed_healthy_deployments,
model_id=session_model_id,
@ -493,7 +497,7 @@ class DeploymentAffinityCheck(CustomLogger):
# Backwards / safety: allow raw string values.
model_id = cache_result
if not model_id:
if not model_id or model_id in failed_deployment_ids:
return typed_healthy_deployments
deployment = self._find_deployment_by_model_id(

39
litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py
View file

@ -51,6 +51,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
strip_encrypted_reasoning_from_messages,
)
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.router_utils.common_utils import as_excluded_deployment_ids, as_retry_skipped_deployment_ids
from litellm.types.llms.openai import AllMessageValues
from litellm.types.router import Deployment
from litellm.utils import get_order_filtered_deployments
@ -182,6 +183,14 @@ class EncryptedContentAffinityCheck(CustomLogger):
None,
)
@staticmethod
def _deployment_model_id(deployment: Mapping[str, object]) -> str | None:
model_info: Final = deployment.get("model_info")
if not isinstance(model_info, Mapping):
return None
deployment_id: Final = cast(Mapping[str, object], model_info).get("id") # cast-ok: narrowed by isinstance
return None if deployment_id is None else str(deployment_id)
@staticmethod
def _find_deployment_by_model_id(healthy_deployments: list[dict], model_id: str) -> dict | None:
for deployment in healthy_deployments:
@ -345,7 +354,8 @@ class EncryptedContentAffinityCheck(CustomLogger):
typed_healthy_deployments: Final = cast(list[dict[str, object]], healthy_deployments)
if not self._is_enabled_for_model_group(model):
return typed_healthy_deployments
target_order: Final = _REQUEST_KWARGS_ADAPTER.validate_python(request_kwargs).get("_target_order")
typed_request_kwargs: Final = _REQUEST_KWARGS_ADAPTER.validate_python(request_kwargs)
target_order: Final = typed_request_kwargs.get("_target_order")
candidates: Final = (
get_order_filtered_deployments(typed_healthy_deployments, target_order=target_order)
if isinstance(target_order, int)
@ -375,8 +385,31 @@ class EncryptedContentAffinityCheck(CustomLogger):
model_id,
)
previous_response_id: Final = typed_request_kwargs.get("previous_response_id")
previous_response_origin: Final = ResponsesAPIRequestUtils.get_model_id_from_response_id(
None if previous_response_id is None else str(previous_response_id)
)
previous_response_binds_origin: Final = previous_response_origin == model_id
excluded_ids: Final[frozenset[str]] = (
frozenset()
if previous_response_binds_origin
else frozenset(as_excluded_deployment_ids(typed_request_kwargs.get("_excluded_deployment_ids")))
)
retry_skipped_ids: Final[frozenset[str]] = (
frozenset()
if previous_response_binds_origin
else frozenset(as_retry_skipped_deployment_ids(typed_request_kwargs.get("_retry_skipped_deployment_ids")))
)
not_excluded_candidates: Final = [
candidate for candidate in candidates if self._deployment_model_id(candidate) not in excluded_ids
]
untried_candidates: Final = [
candidate
for candidate in not_excluded_candidates
if self._deployment_model_id(candidate) not in retry_skipped_ids
] or not_excluded_candidates
deployment: Final = self._find_deployment_by_model_id(
healthy_deployments=candidates,
healthy_deployments=untried_candidates,
model_id=model_id,
)
if deployment is not None:
@ -392,7 +425,7 @@ class EncryptedContentAffinityCheck(CustomLogger):
# Follow-up switched model_name (LIT-2531): pin by Azure resource instead.
boundary_matches, _originating = self._find_deployments_on_same_encryption_boundary(
healthy_deployments=candidates,
healthy_deployments=untried_candidates,
model_id=model_id,
)
if boundary_matches:

5
litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py
View file

@ -19,6 +19,7 @@ from litellm.integrations.anthropic_cache_control_hook import (
from litellm.integrations.custom_logger import CustomLogger, Span
from litellm.litellm_core_utils.token_counter import offload_token_count
from litellm.llms.bedrock.common_utils import bedrock_cache_point_ttl_supported, bedrock_capability_model
from litellm.router_utils.common_utils import deployment_ids_failed_in_request
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
from litellm.types.utils import CallTypes, StandardLoggingPayload
from litellm.utils import (
@ -152,7 +153,9 @@ class PromptCachingDeploymentCheck(CustomLogger):
request_tools,
),
)
if model_id_dict is not None:
if model_id_dict is not None and model_id_dict["model_id"] not in deployment_ids_failed_in_request(
request_data or {}
):
model_id: Final = model_id_dict["model_id"]
for deployment in healthy_deployments:
if deployment["model_info"]["id"] == model_id:

1
litellm/types/utils.py
View file

@ -4573,6 +4573,7 @@ class LlmProviders(str, Enum):
POE = "poe"
CHUTES = "chutes"
NEOSANTARA = "neosantara"
AIAND = "aiand"
PARASAIL = "parasail"
XIAOMI_MIMO = "xiaomi_mimo"
TENSORMESH = "tensormesh"

199
model_prices_and_context_window.json
View file

@ -15764,7 +15764,6 @@
"us": 1.1
},
"supports_output_config": true,
"supports_speed": true,
"prompt_cache_min_tokens": 2048,
"source": "https://platform.claude.com/docs/en/about-claude/pricing"
},
@ -15807,7 +15806,6 @@
"us": 1.1
},
"supports_output_config": true,
"supports_speed": true,
"prompt_cache_min_tokens": 2048
},
"claude-fable-5": {
@ -37504,7 +37502,6 @@
"input_cost_per_token": 5e-08,
"input_cost_per_token_batches": 2.5e-08,
"input_cost_per_token_flex": 2.5e-08,
"input_cost_per_token_priority": 2.5e-06,
"litellm_provider": "openai",
"max_input_tokens": 272000,
"max_output_tokens": 128000,
@ -37552,7 +37549,6 @@
"deprecation_date": "2026-12-11",
"input_cost_per_token": 5e-08,
"input_cost_per_token_batches": 2.5e-08,
"input_cost_per_token_priority": 2.5e-06,
"input_cost_per_token_flex": 2.5e-08,
"litellm_provider": "openai",
"max_input_tokens": 272000,
@ -69827,6 +69823,201 @@
"supports_reasoning": true,
"supports_vision": true
},
"aiand/deepseek-ai/deepseek-v4-flash": {
"cache_read_input_token_cost": 8e-08,
"input_cost_per_token": 1.5e-07,
"litellm_provider": "aiand",
"max_input_tokens": 1048576,
"max_output_tokens": 384000,
"mode": "chat",
"output_cost_per_token": 2.5e-07,
"source": "https://api.aiand.com/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": false
},
"aiand/deepseek-ai/deepseek-v4-pro": {
"cache_read_input_token_cost": 2.5e-07,
"input_cost_per_token": 1e-06,
"litellm_provider": "aiand",
"max_input_tokens": 1048576,
"max_output_tokens": 384000,
"mode": "chat",
"output_cost_per_token": 2.5e-06,
"source": "https://api.aiand.com/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": false
},
"aiand/deepseek-ai/deepseek-v4.1-flash": {
"cache_read_input_token_cost": 2e-08,
"input_cost_per_token": 3e-07,
"litellm_provider": "aiand",
"max_input_tokens": 1048576,
"max_output_tokens": 384000,
"mode": "chat",
"output_cost_per_token": 6e-07,
"source": "https://api.aiand.com/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"aiand/google/gemma-4-31b-it": {
"cache_read_input_token_cost": 5e-08,
"input_cost_per_token": 2e-07,
"litellm_provider": "aiand",
"max_input_tokens": 262144,
"max_output_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 5e-07,
"source": "https://api.aiand.com/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"aiand/moonshotai/kimi-k2.7-code": {
"cache_read_input_token_cost": 2e-07,
"input_cost_per_token": 7.5e-07,
"litellm_provider": "aiand",
"max_input_tokens": 262144,
"max_output_tokens": 262144,
"mode": "chat",
"output_cost_per_token": 3.5e-06,
"source": "https://api.aiand.com/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"aiand/moonshotai/kimi-k3": {
"cache_read_input_token_cost": 5e-07,
"input_cost_per_token": 3e-06,
"litellm_provider": "aiand",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 1.25e-05,
"source": "https://api.aiand.com/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"aiand/motif-technologies/motif-3": {
"cache_read_input_token_cost": 2e-07,
"input_cost_per_token": 5e-07,
"litellm_provider": "aiand",
"max_input_tokens": 262144,
"max_output_tokens": 262144,
"mode": "chat",
"output_cost_per_token": 2e-06,
"source": "https://api.aiand.com/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": false
},
"aiand/openai/gpt-oss-120b": {
"cache_read_input_token_cost": 8e-08,
"input_cost_per_token": 1.5e-07,
"litellm_provider": "aiand",
"max_input_tokens": 131072,
"max_output_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 6e-07,
"source": "https://api.aiand.com/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": false
},
"aiand/qwen/qwen3.6-27b": {
"cache_read_input_token_cost": 2e-07,
"input_cost_per_token": 3.2e-07,
"litellm_provider": "aiand",
"max_input_tokens": 262144,
"max_output_tokens": 65536,
"mode": "chat",
"output_cost_per_token": 3.2e-06,
"source": "https://api.aiand.com/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"aiand/qwen/qwen3.8-27b": {
"cache_read_input_token_cost": 2e-07,
"input_cost_per_token": 4e-07,
"litellm_provider": "aiand",
"max_input_tokens": 262144,
"max_output_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 3e-06,
"source": "https://api.aiand.com/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"aiand/zai-org/glm-5.2": {
"cache_read_input_token_cost": 3e-07,
"input_cost_per_token": 1e-06,
"litellm_provider": "aiand",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 4e-06,
"source": "https://api.aiand.com/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": false
},
"aiand/zai-org/glm-5.3": {
"cache_read_input_token_cost": 3e-07,
"input_cost_per_token": 1e-06,
"litellm_provider": "aiand",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 4e-06,
"source": "https://api.aiand.com/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": false
},
"aiand/zai-org/glm-5.3-flash": {
"cache_read_input_token_cost": 3e-08,
"input_cost_per_token": 1.5e-07,
"litellm_provider": "aiand",
"max_input_tokens": 1048550,
"max_output_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 5e-07,
"source": "https://api.aiand.com/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"bespoke/nimble-latest": {
"input_cost_per_token": 0.0,
"litellm_provider": "bespoke",

26
provider_endpoints_support.json
View file

@ -189,6 +189,32 @@
"generateContent": false
}
},
"aiand": {
"display_name": "ai& (`aiand`)",
"url": "https://docs.aiand.com/api/chat-completions/",
"endpoints": {
"chat_completions": true,
"responses": true,
"text_completion": true,
"embeddings": false,
"image_generations": false,
"image_edits": false,
"audio_transcriptions": false,
"audio_speech": false,
"rerank": false,
"ocr": false,
"search": false,
"moderations": false,
"realtime": false,
"video_generations": false,
"batches": false,
"messages": true,
"interactions": true,
"systemone": false,
"a2a": true,
"generateContent": true
}
},
"amazon_nova": {
"display_name": "Amazon Nova (`amazon_nova`)",
"url": "https://docs.litellm.ai/docs/providers/amazon_nova",

2
tests/code_coverage_tests/router_code_coverage.py
View file

@ -171,6 +171,8 @@ ignored_function_names = [
"_deployment_field", # Read by is_batch_capable, tested through the mixed-group tests in test_router.py
"_deployment_provider", # Read by without_deployments_that_cannot_batch, tested through the mixed-group tests in test_router.py
"_cannot_run_batches", # Tested through the Azure and OpenAI mixed-group acreate_batch tests in test_router.py
"_responses_continuation", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py
"_responses_continuation_ceilings", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py
"_team_deployments_named", # Tested through the team lookups in the concurrent delete tests in test_router.py
]

143
tests/integration/_support/riva.py Normal file
View file

@ -0,0 +1,143 @@
from __future__ import annotations
import io
import threading
import wave
from array import array
from collections.abc import Iterable, Iterator
from concurrent import futures
from contextlib import contextmanager
from dataclasses import dataclass
from enum import Enum
from queue import SimpleQueue
from typing import Final
import grpc
from riva.client.proto import riva_asr_pb2, riva_asr_pb2_grpc
SAMPLE_RATE_HZ: Final = 16000
OFFLINE_ONLY_REFUSAL: Final = (
"Error: Unavailable model requested given these parameters: language_code=en-US; sample_rate=16000; type=online;"
)
_STALL_LIMIT_SECONDS: Final = 120
class Rpc(Enum):
RECOGNIZE = "Recognize"
STREAMING = "StreamingRecognize"
@dataclass(frozen=True, slots=True)
class Word:
word: str
start_ms: int
end_ms: int
@dataclass(frozen=True, slots=True)
class RivaCall:
rpc: Rpc
audio: bytes
cancelled: bool
@dataclass(frozen=True, slots=True)
class RivaServer:
address: str
calls: SimpleQueue[RivaCall]
def drain(self) -> tuple[RivaCall, ...]:
return tuple(self.calls.get_nowait() for _ in range(self.calls.qsize()))
def _clip_samples(seconds: float) -> array[int]:
return array("h", (((index * 37) % 2001 - 1000) * 16 for index in range(int(seconds * SAMPLE_RATE_HZ))))
def pcm_wav(seconds: float) -> bytes:
buffer: Final = io.BytesIO()
with wave.open(buffer, "wb") as writer:
writer.setnchannels(1)
writer.setsampwidth(2)
writer.setframerate(SAMPLE_RATE_HZ)
writer.writeframes(_clip_samples(seconds).tobytes())
return buffer.getvalue()
def pcm_bytes(seconds: float) -> int:
return int(seconds * SAMPLE_RATE_HZ) * 2
def carries_clip(audio: bytes, seconds: float) -> bool:
sent: Final = _clip_samples(seconds)
if len(audio) != len(sent) * 2:
return False
received: Final = array("h", audio)
return all(abs(left - right) <= 1 for left, right in zip(sent, received, strict=True))
class _Asr(riva_asr_pb2_grpc.RivaSpeechRecognitionServicer):
def __init__(self, words: tuple[Word, ...], offline_only: bool, stall: bool, calls: SimpleQueue[RivaCall]) -> None:
self._words: Final = words
self._offline_only: Final = offline_only
self._stall: Final = stall
self._calls: Final = calls
def _alternative(self, config: riva_asr_pb2.RecognitionConfig) -> riva_asr_pb2.SpeechRecognitionAlternative:
timed: Final = self._words if config.enable_word_time_offsets else ()
return riva_asr_pb2.SpeechRecognitionAlternative(
transcript=" ".join(word.word for word in self._words),
confidence=0.9,
words=[
riva_asr_pb2.WordInfo(word=word.word, start_time=word.start_ms, end_time=word.end_ms) for word in timed
],
)
def Recognize(
self, request: riva_asr_pb2.RecognizeRequest, context: grpc.ServicerContext
) -> riva_asr_pb2.RecognizeResponse:
if self._stall:
released: Final = threading.Event()
context.add_callback(released.set)
released.wait(timeout=_STALL_LIMIT_SECONDS)
self._calls.put(RivaCall(Rpc.RECOGNIZE, request.audio, cancelled=not context.is_active()))
return riva_asr_pb2.RecognizeResponse()
self._calls.put(RivaCall(Rpc.RECOGNIZE, request.audio, cancelled=False))
return riva_asr_pb2.RecognizeResponse(
results=[riva_asr_pb2.SpeechRecognitionResult(alternatives=[self._alternative(request.config)])]
)
def StreamingRecognize(
self, request_iterator: Iterable[riva_asr_pb2.StreamingRecognizeRequest], context: grpc.ServicerContext
) -> Iterator[riva_asr_pb2.StreamingRecognizeResponse]:
if self._offline_only:
self._calls.put(RivaCall(Rpc.STREAMING, b"", cancelled=False))
context.abort(grpc.StatusCode.INVALID_ARGUMENT, OFFLINE_ONLY_REFUSAL)
requests: Final = tuple(request_iterator)
config: Final = requests[0].streaming_config.config if requests else riva_asr_pb2.RecognitionConfig()
self._calls.put(
RivaCall(Rpc.STREAMING, b"".join(request.audio_content for request in requests), cancelled=False)
)
yield riva_asr_pb2.StreamingRecognizeResponse(
results=[riva_asr_pb2.StreamingRecognitionResult(alternatives=[self._alternative(config)], is_final=True)]
)
@contextmanager
def riva_server(
words: tuple[Word, ...],
*,
offline_only: bool = False,
stall: bool = False,
max_receive_bytes: int | None = None,
) -> Iterator[RivaServer]:
calls: Final[SimpleQueue[RivaCall]] = SimpleQueue()
options: Final = () if max_receive_bytes is None else (("grpc.max_receive_message_length", max_receive_bytes),)
server: Final = grpc.server(futures.ThreadPoolExecutor(max_workers=8), options=options)
riva_asr_pb2_grpc.add_RivaSpeechRecognitionServicer_to_server(_Asr(words, offline_only, stall, calls), server)
port: Final = server.add_insecure_port("127.0.0.1:0")
server.start()
try:
yield RivaServer(f"127.0.0.1:{port}", calls)
finally:
server.stop(grace=None).wait()

392
tests/integration/compatibility/test_compressed_request_bodies.py Normal file
View file

@ -0,0 +1,392 @@
from __future__ import annotations
import gzip
import json
import os
import secrets
import threading
import uuid
from collections.abc import Callable, Iterator, Mapping
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from functools import partial
from pathlib import Path
from types import MappingProxyType
from typing import Final
import httpx
import pytest
import yaml
from integration._support.budget import case_timeout_seconds
from integration._support.client import (
JSON_OBJECT,
Gateway,
eventually,
gateway_from_environment,
list_value,
object_value,
string_value,
)
from integration._support.process import graceful_stop_seconds, owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import JsonValue
MIB: Final = 1024 * 1024
CHAT: Final = "/v1/chat/completions"
CHAT_MODEL: Final = "compressed-chat"
OTHER_MODEL: Final = "compressed-chat-other"
ANTHROPIC_MODEL: Final = "claude-haiku-5-5"
GATED_PREFIX: Final = "gated-"
ORIGIN: Final = "http://browser.example"
MIDDLEWARE_HEADERS: Final = (
"access-control-allow-origin",
"x-frame-options",
"content-security-policy",
"x-content-type-options",
)
ZSTD_MAGIC: Final = b"\x28\xb5\x2f\xfd"
ZSTD_SINGLE_SEGMENT_FOUR_BYTE_SIZE: Final = b"\xa0"
ZSTD_MAX_BLOCK: Final = 128 * 1024
NO_HEADERS: Final[Mapping[str, str]] = MappingProxyType({})
pytestmark: Final = [
pytest.mark.timeout(2 * graceful_stop_seconds() + 120),
pytest.mark.xdist_group("compressed-request-bodies"),
]
def zstd_frame(data: bytes) -> bytes:
blocks: Final = tuple(data[start : start + ZSTD_MAX_BLOCK] for start in range(0, len(data), ZSTD_MAX_BLOCK))
framed: Final = blocks or (b"",)
return (
ZSTD_MAGIC
+ ZSTD_SINGLE_SEGMENT_FOUR_BYTE_SIZE
+ len(data).to_bytes(4, "little")
+ b"".join(
(len(block) << 3 | (index == len(framed) - 1)).to_bytes(3, "little") + block
for index, block in enumerate(framed)
)
)
def gzip_body(data: bytes) -> bytes:
return gzip.compress(data, mtime=0)
def as_sent(data: bytes) -> bytes:
return data
@dataclass(frozen=True, slots=True)
class Rig:
gateway: Gateway
upstream: Wire
gate: threading.Event
limited_key: str
def calls(self) -> tuple[Request, ...]:
return tuple(request for request in self.upstream.drain() if request.method == "POST")
def _chat_completion() -> bytes:
return json.dumps(
{
"id": f"chatcmpl-{uuid.uuid4().hex}",
"object": "chat.completion",
"created": 1,
"model": "gpt-4o-mini",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "scripted"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5},
}
).encode()
def _stream_frame(identity: str, delta: Mapping[str, JsonValue], finish_reason: str | None) -> bytes:
chunk: Final = {
"id": identity,
"object": "chat.completion.chunk",
"created": 1,
"model": "gpt-4o-mini",
"choices": [{"index": 0, "delta": dict(delta), "finish_reason": finish_reason}],
}
return b"data: " + json.dumps(chunk).encode() + b"\n\n"
def _chat_stream(gate: threading.Event | None) -> Reply:
identity: Final = f"chatcmpl-{uuid.uuid4().hex}"
return Reply(
content_type="text/event-stream",
chunks=(
_stream_frame(identity, {"role": "assistant", "content": "streamed "}, None),
_stream_frame(identity, {"content": "answer"}, "stop"),
b"data: [DONE]\n\n",
),
gate_after_first=gate,
)
def _anthropic_message() -> bytes:
return json.dumps(
{
"id": f"msg_{uuid.uuid4().hex}",
"type": "message",
"role": "assistant",
"model": ANTHROPIC_MODEL,
"content": [{"type": "text", "text": "scripted"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 3, "output_tokens": 2},
}
).encode()
def _respond(gate: threading.Event, request: Request) -> Reply:
if request.method == "GET":
return Reply(body=b'{"object":"list","data":[]}')
if request.target == "/v1/messages":
return Reply(body=_anthropic_message())
body: Final = JSON_OBJECT.validate_json(request.body)
if body.get("stream") is not True:
return Reply(body=_chat_completion())
text: Final = string_value(object_value(list_value(body["messages"])[0])["content"])
return _chat_stream(gate if text.startswith(GATED_PREFIX) else None)
def _deployment(name: str, upstream: Wire) -> Mapping[str, JsonValue]:
return {
"model_name": name,
"litellm_params": {"model": "openai/gpt-4o-mini", "api_base": f"{upstream.url}/v1", "api_key": "sk-fixture"},
}
def _config(directory: Path, upstream: Wire) -> Path:
base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
configuration: Final = {
**base,
"model_list": [_deployment(CHAT_MODEL, upstream), _deployment(OTHER_MODEL, upstream)],
"general_settings": {
**object_value(base.get("general_settings", {})),
"max_request_size_mb": 1,
"enable_claude_code_gateway": True,
},
"router_settings": {**object_value(base.get("router_settings", {})), "num_retries": 0},
}
path: Final = directory / "compressed-request-bodies.yaml"
path.write_text(yaml.safe_dump(configuration))
return path
@pytest.fixture(scope="module")
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
assert os.environ.get("LITELLM_LICENSE"), (
"LITELLM_LICENSE must be set: max_request_size_mb is an enterprise feature"
)
gate: Final = threading.Event()
directory: Final = tmp_path_factory.mktemp("compressed-request-bodies")
with gateway_from_environment() as root, wire_server(partial(_respond, gate)) as upstream:
overrides: Final = {
"ANTHROPIC_API_BASE": upstream.url,
"ANTHROPIC_API_KEY": "sk-ant-synthetic",
"MAX_DECOMPRESSED_REQUEST_BODIES_IN_FLIGHT_MB": "1",
}
with (
owned_proxy_process(root, directory, overrides, config=_config(directory, upstream), workers=1) as owned,
owned.gateway.scenario() as scenario,
):
yield Rig(owned.gateway, upstream, gate, scenario.key(models=[CHAT_MODEL]))
@pytest.fixture(autouse=True)
def empty_upstream_log(rig: Rig) -> None:
rig.upstream.drain()
def _headers(key: str, coding: str | None, extra: Mapping[str, str] = NO_HEADERS) -> Mapping[str, str]:
return {
"Authorization": f"Bearer {key}",
"content-type": "application/json",
**({"content-encoding": coding} if coding is not None else {}),
**extra,
}
def _send(
rig: Rig,
data: bytes,
coding: str | None,
*,
key: str | None = None,
extra: Mapping[str, str] = NO_HEADERS,
) -> httpx.Response:
return rig.gateway.client.post(
CHAT,
content=data,
headers=_headers(rig.gateway.key if key is None else key, coding, extra),
timeout=case_timeout_seconds(),
)
def _chat_body(model: str, text: str, *, stream: bool = False) -> bytes:
return json.dumps(
{
"model": model,
"messages": [{"role": "user", "content": text}],
"stream": stream,
"cache": {"no-cache": True},
}
).encode()
def _error_type(response: httpx.Response) -> JsonValue:
return object_value(JSON_OBJECT.validate_json(response.content)["error"])["type"]
def _answer(response: httpx.Response) -> JsonValue:
choice: Final = object_value(list_value(JSON_OBJECT.validate_json(response.content)["choices"])[0])
return object_value(choice["message"])["content"]
def _stream(gateway: Gateway, data: bytes, first_text: threading.Event) -> tuple[int, str]:
with (
httpx.Client(base_url=gateway.client.base_url, timeout=case_timeout_seconds(), trust_env=False) as client,
client.stream("POST", CHAT, content=data, headers=_headers(gateway.key, "gzip")) as response,
):
pieces: Final = response.iter_text()
first: Final = next(pieces, "")
first_text.set()
return response.status_code, first + "".join(pieces)
CODINGS: Final = (
pytest.param("zstd", zstd_frame, id="zstd"),
pytest.param("gzip", gzip_body, id="gzip"),
pytest.param("gzip", as_sent, id="plain-labeled-gzip"),
)
@pytest.mark.parametrize(("coding", "encode"), CODINGS)
def test_a_compressed_chat_body_reaches_the_provider_as_its_plain_twin(
rig: Rig, coding: str, encode: Callable[[bytes], bytes]
) -> None:
body: Final = _chat_body(CHAT_MODEL, f"twin {uuid.uuid4().hex}")
plain: Final = _send(rig, body, None)
compressed: Final = _send(rig, encode(body), coding)
assert plain.status_code == 200, plain.text
assert compressed.status_code == 200, compressed.text
assert _answer(compressed) == _answer(plain)
assert compressed.headers["x-litellm-response-cost"] == plain.headers["x-litellm-response-cost"]
calls: Final = rig.calls()
assert len(calls) == 2, [call.target for call in calls]
assert JSON_OBJECT.validate_json(calls[1].body) == JSON_OBJECT.validate_json(calls[0].body)
assert "content-encoding" not in calls[1].headers, calls[1].headers
def test_auth_reads_the_model_from_a_compressed_body(rig: Rig) -> None:
refused: Final = _send(rig, zstd_frame(_chat_body(OTHER_MODEL, "outside the key")), "zstd", key=rig.limited_key)
assert refused.status_code == 403, refused.text
assert _error_type(refused) == "key_model_access_denied", refused.text
assert rig.calls() == ()
served: Final = _send(rig, zstd_frame(_chat_body(CHAT_MODEL, "inside the key")), "zstd", key=rig.limited_key)
assert served.status_code == 200, served.text
assert len(rig.calls()) == 1
@pytest.mark.parametrize(
("coding", "encode"),
(pytest.param("zstd", zstd_frame, id="zstd"), pytest.param("gzip", as_sent, id="plain-labeled-gzip")),
)
def test_a_compressed_pass_through_body_reaches_the_vendor_decoded(
rig: Rig, coding: str, encode: Callable[[bytes], bytes]
) -> None:
body: Final = json.dumps(
{
"model": ANTHROPIC_MODEL,
"max_tokens": 8,
"messages": [{"role": "user", "content": f"pass through {uuid.uuid4().hex}"}],
}
).encode()
response: Final = rig.gateway.client.post(
"/anthropic/v1/messages",
content=encode(body),
headers={
"x-litellm-api-key": rig.gateway.key,
"anthropic-version": "2023-06-01",
"content-type": "application/json",
"content-encoding": coding,
},
timeout=case_timeout_seconds(),
)
assert response.status_code == 200, response.text
calls: Final = rig.calls()
assert [call.target for call in calls] == ["/v1/messages"]
assert JSON_OBJECT.validate_json(calls[0].body) == JSON_OBJECT.validate_json(body)
assert "content-encoding" not in calls[0].headers, calls[0].headers
def test_a_gzip_bomb_gets_a_413_with_the_headers_other_errors_carry(rig: Rig) -> None:
origin: Final = {"origin": ORIGIN}
rejected: Final = _send(rig, gzip_body(b"\0" * 800_000), "gzip", extra=origin)
unauthorized: Final = _send(
rig, _chat_body(CHAT_MODEL, "unknown key"), None, key=f"sk-unknown-{uuid.uuid4().hex}", extra=origin
)
assert rejected.status_code == 413, rejected.text
assert unauthorized.status_code == 401, unauthorized.text
assert unauthorized.headers.get("access-control-allow-origin") is not None, dict(unauthorized.headers)
assert {name: rejected.headers.get(name) for name in MIDDLEWARE_HEADERS} == {
name: unauthorized.headers.get(name) for name in MIDDLEWARE_HEADERS
}
assert rig.calls() == ()
def test_the_operator_size_limit_holds_a_compressed_body_to_its_plain_twins_limit(rig: Rig) -> None:
over: Final = _chat_body(CHAT_MODEL, secrets.token_hex(600_000))
compressed_over: Final = gzip_body(over)
assert len(compressed_over) < MIB < len(over)
plain: Final = _send(rig, over, None)
rejected: Final = _send(rig, compressed_over, "gzip")
served: Final = _send(rig, gzip_body(_chat_body(CHAT_MODEL, secrets.token_hex(300_000))), "gzip")
assert plain.status_code == 413, plain.text
assert rejected.status_code == plain.status_code, rejected.text
assert served.status_code == 200, served.text
assert len(rig.calls()) == 1
def test_the_decoded_body_budget_turns_away_a_second_body_until_the_stream_ends(rig: Rig) -> None:
held: Final = gzip_body(_chat_body(CHAT_MODEL, GATED_PREFIX + secrets.token_hex(300_000), stream=True))
second: Final = gzip_body(_chat_body(CHAT_MODEL, secrets.token_hex(300_000)))
keyless: Final = _send(rig, second, "gzip", key=f"sk-unknown-{uuid.uuid4().hex}")
assert keyless.status_code == 401, keyless.text
first_text: Final = threading.Event()
with ThreadPoolExecutor(1) as pool:
streaming: Final = pool.submit(_stream, rig.gateway, held, first_text)
try:
eventually(lambda: first_text.is_set() or streaming.done(), bool, seconds=case_timeout_seconds())
reached: Final = rig.calls()
turned_away: Final = _send(rig, second, "gzip")
finally:
rig.gate.set()
status, streamed = streaming.result(timeout=case_timeout_seconds())
assert len(reached) == 1, streamed
assert turned_away.status_code == 503, turned_away.text
assert _error_type(turned_away) == "overloaded_error", turned_away.text
assert turned_away.headers.get("retry-after") == "1", dict(turned_away.headers)
assert status == 200, streamed
assert "streamed " in streamed and "answer" in streamed and "data: [DONE]" in streamed, streamed
assert rig.calls() == ()
after_the_stream: Final = _send(rig, second, "gzip")
assert after_the_stream.status_code == 200, after_the_stream.text
assert len(rig.calls()) == 1
def test_a_claude_code_otlp_export_keeps_a_two_member_gzip_body(rig: Rig) -> None:
response: Final = rig.gateway.client.post(
"/claude_code_gateway/v1/metrics",
content=gzip_body(b"first member") + gzip_body(b"second member"),
headers={
"Authorization": f"Bearer {rig.gateway.key}",
"content-type": "application/x-protobuf",
"content-encoding": "gzip",
},
timeout=case_timeout_seconds(),
)
assert response.status_code == 200, response.text

165
tests/integration/configuration/test_openapi_docs_route.py Normal file
View file

@ -0,0 +1,165 @@
import time
from collections.abc import Iterator
from concurrent.futures import Future, ThreadPoolExecutor
from pathlib import Path
from typing import Final
import httpx
import pytest
from integration._support.client import Gateway, list_value, object_value
from integration._support.process import graceful_stop_seconds, owned_proxy_process
from pydantic import JsonValue, TypeAdapter
from litellm.proxy._lazy_features import LAZY_FEATURES
DOCS_SECONDS: Final = 120
COLD_DOCS: Final = 4
CHATS: Final = 8
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
OWNED_PROXY_TIMEOUT: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
TWO_OWNED_PROXIES_TIMEOUT: Final = pytest.mark.timeout(2 * (2 * graceful_stop_seconds() + 120))
LAZY_ROUTES_FLAG: Final = "LITELLM_DISABLE_LAZY_ROUTES"
WARMED_FEATURE: Final = next(feature for feature in LAZY_FEATURES if feature.name == "policies")
def _timed_get(candidate: Gateway, path: str) -> tuple[httpx.Response, float]:
started: Final = time.perf_counter()
response: Final = candidate.request("GET", path)
return response, time.perf_counter() - started
def _liveliness_until(done: Future[tuple[httpx.Response, float]], candidate: Gateway) -> Iterator[tuple[int, float]]:
while not done.done():
response, seconds = _timed_get(candidate, "/health/liveliness")
yield response.status_code, seconds
def _schema(response: httpx.Response) -> dict[str, JsonValue]:
assert response.status_code == 200, response.text[:500]
return JSON_OBJECT.validate_json(response.content)
def _servers(response: httpx.Response) -> list[JsonValue]:
return list(list_value(_schema(response).get("servers", [])))
def _chat(candidate: Gateway, model: str, key: str, content: str) -> httpx.Response:
return candidate.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": content}], "cache": {"no-cache": True}},
key=key,
)
def _warmed_feature_operations(response: httpx.Response) -> dict[str, JsonValue]:
paths: Final = object_value(_schema(response)["paths"])
return {path: operations for path, operations in paths.items() if WARMED_FEATURE.matches(path)}
@OWNED_PROXY_TIMEOUT
def test_liveliness_answers_while_a_cold_openapi_schema_builds(gateway: Gateway, tmp_path: Path) -> None:
with (
owned_proxy_process(gateway, tmp_path, {}, client_timeout=DOCS_SECONDS) as owned,
ThreadPoolExecutor(1) as pool,
):
docs: Final = pool.submit(_timed_get, owned.gateway, "/openapi.json")
probes: Final = tuple(_liveliness_until(docs, owned.gateway))
docs_response, docs_seconds = docs.result()
assert "/v1/chat/completions" in object_value(_schema(docs_response)["paths"])
assert probes, "no liveliness probe ran while the docs request was pending"
assert {status for status, _ in probes} == {200}, probes
slowest: Final = max(seconds for _, seconds in probes)
assert slowest < docs_seconds / 3, (
f"a liveliness probe waited {slowest:.2f}s of the {docs_seconds:.2f}s cold docs request"
)
@OWNED_PROXY_TIMEOUT
def test_concurrent_cold_openapi_requests_share_the_schema_later_requests_get(gateway: Gateway, tmp_path: Path) -> None:
with (
owned_proxy_process(gateway, tmp_path, {}, client_timeout=DOCS_SECONDS) as owned,
ThreadPoolExecutor(COLD_DOCS) as pool,
):
cold: Final = tuple(pool.map(lambda _: owned.gateway.request("GET", "/openapi.json"), range(COLD_DOCS)))
warm: Final = owned.gateway.request("GET", "/openapi.json")
assert [response.status_code for response in (*cold, warm)] == [200] * (COLD_DOCS + 1)
assert {response.content for response in cold} == {warm.content}
assert "/v1/chat/completions" in object_value(_schema(warm)["paths"])
@OWNED_PROXY_TIMEOUT
def test_chat_completions_overlapping_a_cold_openapi_build_all_succeed(gateway: Gateway, tmp_path: Path) -> None:
with (
owned_proxy_process(gateway, tmp_path, {}, client_timeout=DOCS_SECONDS) as owned,
owned.gateway.scenario() as scenario,
ThreadPoolExecutor(CHATS + 1) as pool,
):
model: Final = scenario.model()
key: Final = scenario.key(models=[model])
first_chat: Final = _chat(owned.gateway, model, key, "before the docs build")
docs: Final = pool.submit(owned.gateway.request, "GET", "/openapi.json")
chats: Final = tuple(
pool.map(lambda index: _chat(owned.gateway, model, key, f"overlapping docs build {index}"), range(CHATS))
)
chats_returned_during_build: Final = not docs.done()
docs_response: Final = docs.result()
assert first_chat.status_code == 200, first_chat.text[:300]
assert chats_returned_during_build, "the chat calls returned only after the cold docs build finished"
assert [(response.status_code, response.text[:300]) for response in chats if response.status_code != 200] == []
replies: Final = tuple(JSON_OBJECT.validate_json(response.content) for response in chats)
assert {reply["object"] for reply in replies} == {"chat.completion"}
assert all(list_value(reply["choices"]) for reply in replies), replies
assert "/v1/chat/completions" in object_value(_schema(docs_response)["paths"])
@OWNED_PROXY_TIMEOUT
def test_docs_filtered_mode_serves_only_llm_routes(gateway: Gateway, tmp_path: Path) -> None:
with owned_proxy_process(gateway, tmp_path, {"DOCS_FILTERED": "True"}, client_timeout=DOCS_SECONDS) as owned:
paths: Final = object_value(_schema(owned.gateway.request("GET", "/openapi.json"))["paths"])
assert "/v1/chat/completions" in paths
assert {"/key/generate", "/user/new", "/team/new"}.isdisjoint(paths), sorted(paths)
@OWNED_PROXY_TIMEOUT
def test_each_root_path_prefix_lists_itself_first_in_servers(gateway: Gateway, tmp_path: Path) -> None:
prefixes: Final = ("/tenant-a", "/tenant-b", "/tenant-a")
with owned_proxy_process(
gateway, tmp_path, {"SERVER_ROOT_PATHS": "/tenant-a,/tenant-b"}, client_timeout=DOCS_SECONDS
) as owned:
served: Final = tuple(
(prefix, _servers(owned.gateway.request("GET", f"{prefix}/openapi.json"))) for prefix in prefixes
)
assert served == tuple((prefix, [{"url": prefix}]) for prefix in prefixes)
@OWNED_PROXY_TIMEOUT
def test_a_scalar_server_root_path_is_listed_once_in_servers(gateway: Gateway, tmp_path: Path) -> None:
overrides: Final = {"SERVER_ROOT_PATH": "/gateway", "LITELLM_UI_PATH": str(tmp_path / "ui")}
with owned_proxy_process(gateway, tmp_path, overrides, client_timeout=DOCS_SECONDS) as owned:
servers: Final = _servers(owned.gateway.request("GET", "/gateway/openapi.json"))
assert servers == [{"url": "/gateway"}]
@TWO_OWNED_PROXIES_TIMEOUT
def test_a_lazy_feature_loaded_during_a_cold_openapi_build_shows_its_live_routes_in_the_next_schema(
gateway: Gateway, tmp_path: Path
) -> None:
with owned_proxy_process(gateway, tmp_path, {LAZY_ROUTES_FLAG: "true"}, client_timeout=DOCS_SECONDS) as eager:
live: Final = _warmed_feature_operations(eager.gateway.request("GET", "/openapi.json"))
with (
owned_proxy_process(
gateway, tmp_path, {}, remove_environment=(LAZY_ROUTES_FLAG,), client_timeout=DOCS_SECONDS
) as lazy,
ThreadPoolExecutor(1) as pool,
):
docs: Final = pool.submit(lazy.gateway.request, "GET", "/openapi.json")
warmed: Final = lazy.gateway.request("POST", f"/lazy/warm/{WARMED_FEATURE.name}")
warmed_during_build: Final = not docs.done()
during: Final = _warmed_feature_operations(docs.result())
after: Final = _warmed_feature_operations(lazy.gateway.request("GET", "/openapi.json"))
assert warmed.status_code == 200, warmed.text
assert warmed_during_build, "the feature finished loading only after the cold docs build finished"
assert live, "the eager proxy published no route of the warmed feature"
assert during != live, "the cold docs build started after the feature had already loaded"
assert after == live

32
tests/integration/cost_calculation/cases/decisions.py
View file

@ -162,6 +162,37 @@ TYPESAFE_JEV_LATEST_SYSTEMONE: Final = CostTrackingTestCase(
completion_tokens=22,
),
)
TYPESAFE_JEV_PREVIEW_ANSWERED_AS_PRICED_VERSION_SYSTEMONE: Final = CostTrackingTestCase(
name="typesafe/jev-preview-answered-as-priced-version-systemone",
covers="quota_management.spend_tracking.systemone.logs_cost",
model="typesafe/jev-preview",
endpoint="/v1/systemone",
request={
"model": "$MODEL",
"state": "The replacement part arrived but the wrong size.",
"questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}},
},
upstream_body={
"model": "jev-preview",
"state": "The replacement part arrived but the wrong size.",
"questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}},
},
response=JsonResponse(
content_type="application/json",
body={
"model": "jev-latest",
"answers": {"is_refund": {"type": "noul", "noul": 0.31}},
"usage": {"input_tokens": 311, "output_tokens": 9},
},
),
expected=ExactExpected(
spend=311 * 4.2e-08,
input_cost=311 * 4.2e-08,
output_cost=0.0,
prompt_tokens=311,
completion_tokens=9,
),
)
# perplexity/pplx-decider-v1-27b
PERPLEXITY_PPLX_DECIDER_V1_27B_DECISIONS: Final = CostTrackingTestCase(
@ -514,6 +545,7 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = (
OPENAI_GPT_6_LUNA_SYSTEMONE,
TYPESAFE_JEV_LATEST_DECISIONS,
TYPESAFE_JEV_LATEST_SYSTEMONE,
TYPESAFE_JEV_PREVIEW_ANSWERED_AS_PRICED_VERSION_SYSTEMONE,
PERPLEXITY_PPLX_DECIDER_V1_27B_DECISIONS,
PERPLEXITY_PPLX_DECIDER_V1_27B_SYSTEMONE,
CLOUDFLARE_CLEF_FLASH_DECISIONS,

6
tests/integration/cost_calculation/cost_map.py
View file

@ -121,6 +121,12 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = {
"input_cost_per_token": 4.2e-08,
"output_cost_per_token": 0.0,
},
"typesafe/jev-preview": {
"litellm_provider": "typesafe",
"mode": "evaluation",
"input_cost_per_token": 8.4e-08,
"output_cost_per_token": 0.0,
},
"fireworks_ai/accounts/fireworks/routers/rune": {
"litellm_provider": "fireworks_ai",
"mode": "evaluation",

3
tests/integration/management/test_scim_admin_group.py
View file

@ -691,9 +691,6 @@ def test_scim_group_put_returns_the_updated_members(gateway: Gateway, tmp_path:
@pytest.mark.timeout(300)
def test_scim_admin_group_delete_revokes_the_members_admin_key(gateway: Gateway, tmp_path: Path) -> None:
pytest.skip(
"BUG: DELETE /scim/v2/Groups/{id} demotes the admin group's members in the DB but their existing keys keep proxy_admin access on /key/generate"
)
group_name: Final = f"scim-admin-delete-{uuid.uuid4().hex}"
config: Final = _owned_scim_config(tmp_path / "scim-admin-group-delete.yaml", {"scim_admin_group": group_name})
with owned_proxy_process(gateway, tmp_path, {}, config=config) as owned, owned.gateway.scenario() as scenario:

106
tests/integration/management/test_scim_group_delete.py
View file

@ -50,7 +50,7 @@ from integration._support.scim import (
)
from pydantic import JsonValue, TypeAdapter
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS, TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS
pytestmark: Final = pytest.mark.shares_worker_state
@ -70,6 +70,17 @@ FROM "LiteLLM_AuditLog"
WHERE object_id IN (%s, %s) AND action = 'deleted'
ORDER BY table_name
"""
AUDIT_ROWS_NAMING_SQL: Final = """
SELECT table_name, object_id, changed_by_api_key
FROM "LiteLLM_AuditLog"
WHERE object_id = ANY(%s::text[]) AND action = 'deleted'
ORDER BY table_name, object_id
"""
OLDEST_TRANSACTION_BLOCKED_BY_SQL: Final = """
SELECT coalesce(max(extract(epoch FROM clock_timestamp() - xact_start)), 0)::float AS age
FROM pg_stat_activity
WHERE %s::int = ANY(pg_blocking_pids(pid))
"""
_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")
_PIDS: Final = TypeAdapter(tuple[int, ...])
@ -115,15 +126,28 @@ def _deleted_audit_rows(team: str, hashed_key: str) -> list[dict[str, JsonValue]
return read_rows(DELETED_AUDIT_ROWS_SQL, (team, hashed_key))
def _audit_rows_naming(object_ids: tuple[str, ...]) -> list[dict[str, JsonValue]]:
return read_rows(AUDIT_ROWS_NAMING_SQL, ("{" + ",".join(object_ids) + "}",))
def _oldest_transaction_blocked_by(holder: int) -> float:
age: Final = read_rows(OLDEST_TRANSACTION_BLOCKED_BY_SQL, (str(holder),))[0]["age"]
assert isinstance(age, float), age
return age
def _worker_startups(log: Path) -> tuple[tuple[int, ...], int]:
text: Final = log.read_text()
return _PIDS.validate_python(_STARTED_WORKER.findall(text)), text.count("Application startup complete.")
@contextmanager
def _fresh_connection(candidate: Gateway) -> Iterator[Gateway]:
def _fresh_connection(candidate: Gateway, timeout: httpx.Timeout | float | None = None) -> Iterator[Gateway]:
with httpx.Client(
base_url=candidate.client.base_url, timeout=candidate.client.timeout, trust_env=False, limits=GATEWAY_LIMITS
base_url=candidate.client.base_url,
timeout=candidate.client.timeout if timeout is None else timeout,
trust_env=False,
limits=GATEWAY_LIMITS,
) as client:
yield Gateway(client, candidate.key, candidate.upstream_url)
@ -294,6 +318,82 @@ def test_a_group_delete_revokes_member_keys_and_a_re_created_group_gets_its_serv
assert _models_status(gateway, member_key) == 401
@pytest.mark.timeout(OWNED_PROXY_TIMEOUT)
def test_a_group_delete_sent_with_the_idps_scim_token_revokes_member_keys_and_audits_them_under_that_token(
gateway: Gateway, tmp_path: Path
) -> None:
with (
owned_proxy_process(
gateway,
tmp_path,
{"DATABASE_URL": os.environ["DATABASE_URL"], "LITELLM_STORE_AUDIT_LOGS": "true"},
remove_environment=("DATABASE_URL_READ_REPLICA",),
) as owned,
owned.gateway.scenario() as scenario,
):
candidate: Final = owned.gateway
scim_token: Final = scenario.key(allowed_routes=["/scim/*"], models=[])
users: Final = seed_users(scenario, 3)
team: Final = created_team(scenario, create_group(candidate, users))
member_keys: Final = (team_key(candidate, team, users[0]), team_key(candidate, team, users[1]))
service_key: Final = service_account_key(candidate, scenario, team)
deleted: Final = delete_group(candidate, team, key=scim_token)
assert deleted.status_code == 204, deleted.text
assert keys_of(team) == frozenset({_hashed(service_key)})
for key in member_keys:
_assert_revoked_everywhere((candidate,), key)
audited: Final = (team, *(_hashed(key) for key in member_keys), _hashed(service_key))
rows: Final = eventually(lambda: _audit_rows_naming(audited), lambda found: len(found) >= 3, seconds=30)
token: Final = _hashed(scim_token)
assert rows == [
{"table_name": "LiteLLM_TeamTable", "object_id": team, "changed_by_api_key": token},
*(
{"table_name": "LiteLLM_VerificationToken", "object_id": hashed, "changed_by_api_key": token}
for hashed in sorted(_hashed(key) for key in member_keys)
),
]
@pytest.mark.timeout(TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS + 120)
def test_a_group_delete_that_outlives_its_transaction_on_a_locked_member_row_changes_nothing_and_a_retry_lands(
gateway: Gateway,
) -> None:
patient_timeout: Final = TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS + 90
with (
gateway.scenario() as scenario,
_fresh_connection(gateway, timeout=patient_timeout) as patient,
ThreadPoolExecutor(max_workers=1) as pool,
):
users: Final = seed_users(scenario, 3)
team: Final = created_team(scenario, create_group(gateway, users))
key: Final = team_key(gateway, team, users[0])
with held_user_row(users[0]) as holder:
delete: Final = pool.submit(delete_group, patient, team)
eventually(
lambda: (
delete.done()
or _oldest_transaction_blocked_by(holder) > TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS + 1
),
lambda outlived: outlived,
seconds=TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS + 30,
)
failed: Final = delete.result(timeout=patient_timeout)
assert failed.status_code >= 500, failed.text
assert_landed(gateway, team, users)
assert keys_of(team) == frozenset({_hashed(key)})
assert read_rows(DELETED_TEAM_RECORDS_SQL, (team,)) == []
assert read_rows(DELETED_KEY_RECORDS_SQL, (_hashed(key),)) == []
assert _models_status(gateway, key) == 200
retried: Final = delete_group(gateway, team)
assert retried.status_code == 204, retried.text
assert_gone(gateway, team, users, key)
def test_deleting_an_unknown_group_is_404(gateway: Gateway) -> None:
missing: Final = f"missing-{uuid.uuid4().hex}"
response: Final = delete_group(gateway, missing)

44
tests/integration/management/test_team_delete_large_membership.py
View file

@ -33,7 +33,7 @@ from tests.integration._support.client import (
string_value,
)
from tests.integration._support.database import read_rows, write_rows
from tests.integration._support.process import owned_proxy_process
from tests.integration._support.process import graceful_stop_seconds, owned_proxy_process
LARGE_ROSTER: Final = 250
POOL_LIMIT: Final = 5
@ -131,6 +131,13 @@ def _roster_user_ids(roster: JsonValue) -> list[str]:
return sorted(string_value(object_value(member)["user_id"]) for member in roster)
def _deleted_audit_rows(object_id: str, table_name: str) -> list[dict[str, JsonValue]]:
return read_rows(
'SELECT before_value FROM "LiteLLM_AuditLog" WHERE object_id = %s AND table_name = %s AND action = %s',
(object_id, table_name, "deleted"),
)
def _waiters_on_lock_held_by(backend_pid: int) -> int:
rows: Final = read_rows(WAITERS_ON_HELD_LOCK_SQL, ("%pg_advisory_xact_lock%", str(backend_pid)))
waiting: Final = rows[0]["waiting"]
@ -551,6 +558,41 @@ def test_delete_writes_one_audit_row_for_the_team_and_one_per_key(gateway: Gatew
]
@pytest.mark.timeout(int(2 * graceful_stop_seconds() + 120))
def test_two_deletes_racing_on_one_team_write_one_tombstone_and_one_audit_row(gateway: Gateway, tmp_path: Path) -> None:
with (
owned_proxy_process(
gateway,
tmp_path,
{"DATABASE_URL": os.environ["DATABASE_URL"], "LITELLM_STORE_AUDIT_LOGS": "true"},
remove_environment=("DATABASE_URL_READ_REPLICA",),
) as owned,
owned.gateway.scenario() as scenario,
ThreadPoolExecutor(max_workers=2) as pool,
):
user: Final = scenario.user()
team: Final = _create_team(scenario)
_bulk_member_add(owned.gateway, team, [user])
with psycopg.connect(os.environ["DATABASE_URL"]) as holder:
holder.execute(TAKE_TEAM_LOCK_SQL, (team,))
deletes: Final = tuple(
_post_in_background(pool, owned.gateway, "/team/delete", {"team_ids": [team]}) for _ in range(2)
)
eventually(
lambda: _waiters_on_lock_held_by(holder.info.backend_pid), lambda waiting: waiting == 2, seconds=30
)
assert [delete.result(timeout=60).status_code for delete in deletes] == [200, 200]
tombstones: Final = _tombstones(team)
assert [row["team_id"] for row in tombstones] == [team]
assert _roster_user_ids(tombstones[0]["members_with_roles"]) == sorted([PROXY_ADMIN, user])
team_audit_rows: Final = _deleted_audit_rows(team, "LiteLLM_TeamTable")
assert len(team_audit_rows) == 1, team_audit_rows
assert object_value(team_audit_rows[0]["before_value"])["team_id"] == team
def test_second_delete_of_the_same_team_is_404_with_one_tombstone(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
user: Final = scenario.user()

69
tests/integration/mcp/test_mcp_llm_endpoints.py
View file

@ -702,6 +702,36 @@ def test_streaming_responses_opens_the_gateway_tool_call_empty_and_streams_its_a
assert len(rig.upstream_tools()) == 2
STREAM_END_EVENT: Final[Mapping[Surface, str]] = {
"chat": "data: [DONE]",
"responses": "response.completed",
"messages": "message_stop",
"messages_bridge": "message_stop",
}
@pytest.mark.parametrize("surface", SURFACES)
def test_streamed_auto_execute_gives_back_the_keys_only_parallel_slot_once_the_stream_ends(
gateway: Gateway, surface: Surface
) -> None:
with _rig(gateway, surface) as rig:
key: Final = rig.scenario.key(object_permission={"mcp_servers": [rig.server_id]}, max_parallel_requests=1)
streamed: Final = rig.send(key, [AUTO], stream=True)
assert streamed.status_code == 200, streamed.text
assert streamed.headers["content-type"].startswith("text/event-stream"), streamed.headers
assert STREAM_END_EVENT[surface] in streamed.text, streamed.text
assert ANSWER in streamed.text, streamed.text
assert [call["body"]["params"]["name"] for call in _peer_add_calls(rig.peer)] == ["add"]
follow_up: Final = eventually(
lambda: rig.send(key, [], prompt=f"{DIRECT} {rig.alias}"),
lambda response: response.status_code == 200,
seconds=20,
return_last_on_timeout=True,
)
assert follow_up.status_code == 200, follow_up.text
assert rig.final_text(follow_up.json()) == ANSWER
def test_streaming_chat_through_a_toolset_gateway_url_serves_a_team_key_without_its_own_grant(
gateway: Gateway,
) -> None:
@ -920,6 +950,45 @@ def test_nonstream_failed_follow_up_bills_round_one_on_the_failure_row(
)
NO_CACHE: Final = {"no-cache": True}
def _upstream_rounds(rig: Rig) -> tuple[int, int]:
answered: Final = tuple(
_has_tool_result(JSON_OBJECT.validate_json(request.body))
for request in rig.wire.drain()
if request.method == "POST"
)
return answered.count(False), answered.count(True)
def _billed_rows(call_id: str, call_type: str, billed: float) -> list[dict[str, JsonValue]]:
return eventually(
lambda: read_rows(
'SELECT status, spend FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s AND call_type=%s',
(call_id, call_type),
),
lambda found: sum(float(str(row["spend"])) for row in found) >= billed * (1 - 1e-6),
seconds=70,
)
@pytest.mark.parametrize(("surface", "call_type"), (("responses", "aresponses"), ("chat", "acompletion")))
def test_nonstream_follow_up_failing_on_every_retry_bills_each_attempts_first_round_on_one_failure_row(
gateway: Gateway, surface: Surface, call_type: str
) -> None:
billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_fails=True)
with _rig(gateway, surface, billing=billing, prices={**PRICES, "num_retries": 2}) as rig:
response: Final = rig.send(_granted_key(rig), [AUTO], cache=NO_CACHE)
assert response.status_code == 500 and FOLLOW_UP_OUTAGE in response.text, response.text
first_rounds, follow_ups = _upstream_rounds(rig)
assert first_rounds == follow_ups > 1, (first_rounds, follow_ups)
paid: Final = _price(*(FIRST_ROUND,) * first_rounds)
rows: Final = _billed_rows(response.headers["x-litellm-call-id"], call_type, paid)
assert [row["status"] for row in rows] == ["failure"], rows
assert float(str(rows[0]["spend"])) == pytest.approx(paid), rows
def _tool_rows(call_id: str, expected: int) -> list[dict[str, JsonValue]]:
return eventually(
lambda: read_rows(

174
tests/integration/mcp/test_responses_mcp_streaming_round_history.py Normal file
View file

@ -0,0 +1,174 @@
from __future__ import annotations
import json
import uuid
from collections.abc import Mapping
from typing import Final
from integration._support.client import Gateway
from integration._support.mcp import mcp_peer, register_mcp, tool_calls
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
TOOL_ARGUMENTS: Final = '{"a":2,"b":3}'
ANSWER: Final = "All three tool rounds are done."
def _event_frame(event: Mapping[str, JsonValue]) -> bytes:
return f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode()
def _streamed_response(round_number: int, model: str, tool_name: str) -> Reply:
response_id: Final = f"resp_mcp_round_{round_number}"
call_id: Final = f"call_{round_number}"
is_tool_call: Final = round_number <= 3
output_item: Final[dict[str, JsonValue]] = (
{
"type": "function_call",
"id": f"fc_{round_number}",
"call_id": call_id,
"name": tool_name,
"arguments": TOOL_ARGUMENTS,
"status": "completed",
}
if is_tool_call
else {
"type": "message",
"id": f"msg_{round_number}",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": ANSWER, "annotations": []}],
}
)
added_item: Final[dict[str, JsonValue]] = (
{**output_item, "arguments": "", "status": "in_progress"}
if is_tool_call
else {**output_item, "content": [], "status": "in_progress"}
)
item_events: Final[tuple[dict[str, JsonValue], ...]] = (
(
{"type": "response.output_item.added", "output_index": 0, "item": added_item},
{
"type": "response.function_call_arguments.delta",
"item_id": f"fc_{round_number}",
"output_index": 0,
"delta": TOOL_ARGUMENTS,
},
{
"type": "response.function_call_arguments.done",
"item_id": f"fc_{round_number}",
"output_index": 0,
"arguments": TOOL_ARGUMENTS,
},
{"type": "response.output_item.done", "output_index": 0, "item": output_item},
)
if is_tool_call
else (
{"type": "response.output_item.added", "output_index": 0, "item": added_item},
{
"type": "response.output_text.delta",
"item_id": f"msg_{round_number}",
"output_index": 0,
"content_index": 0,
"delta": ANSWER,
},
{"type": "response.output_item.done", "output_index": 0, "item": output_item},
)
)
response: Final[dict[str, JsonValue]] = {
"id": response_id,
"object": "response",
"created_at": 1,
"status": "completed",
"model": model,
"output": [output_item],
"usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
}
events: Final[tuple[dict[str, JsonValue], ...]] = (
{
"type": "response.created",
"response": {**response, "status": "in_progress", "output": []},
},
{
"type": "response.in_progress",
"response": {**response, "status": "in_progress", "output": []},
},
*item_events,
{"type": "response.completed", "response": response},
)
chunks: Final = tuple(
_event_frame({**event, "sequence_number": sequence_number}) for sequence_number, event in enumerate(events)
)
return Reply(content_type="text/event-stream", chunks=chunks)
def _tool_call_history(body: Mapping[str, JsonValue]) -> tuple[tuple[str, str], ...]:
input_items: Final = body["input"]
assert isinstance(input_items, list)
return tuple(
(str(item["type"]), str(item["call_id"]))
for item in input_items
if isinstance(item, dict) and item.get("type") in ("function_call", "function_call_output")
)
def test_streaming_responses_mcp_follow_up_keeps_every_tool_round(gateway: Gateway) -> None:
alias: Final = "hist" + uuid.uuid4().hex[:8]
upstream_requests: Final[list[dict[str, JsonValue]]] = []
def respond(request: Request) -> Reply:
if request.method == "GET" and request.target.endswith("/models"):
return Reply(body=b'{"object":"list","data":[]}')
assert request.target.endswith("/responses"), request.target
request_body: Final = JSON_OBJECT.validate_json(request.body)
upstream_requests.append(request_body)
return _streamed_response(len(upstream_requests), str(request_body["model"]), f"{alias}-add")
with mcp_peer() as peer, wire_server(respond) as wire, gateway.scenario() as scenario:
server_id: Final = register_mcp(scenario, peer, alias)
model: Final = scenario.model(model="openai/responses/gpt-4o-mini", api_base=wire.url + "/v1")
key: Final = scenario.key(object_permission={"mcp_servers": [server_id]})
peer.drain()
response: Final = gateway.client.post(
"/v1/responses",
headers={"Authorization": f"Bearer {key}"},
json={
"model": model,
"input": "run each tool round in order",
"stream": True,
"tools": [
{
"type": "mcp",
"server_url": "litellm_proxy",
"server_label": "litellm",
"require_approval": "never",
}
],
},
timeout=90,
)
assert response.status_code == 200, response.text
assert "response.completed" in response.text and ANSWER in response.text, response.text
assert len(upstream_requests) == 4, upstream_requests
assert _tool_call_history(upstream_requests[2]) == (
("function_call", "call_1"),
("function_call_output", "call_1"),
("function_call", "call_2"),
("function_call_output", "call_2"),
)
assert _tool_call_history(upstream_requests[3]) == (
("function_call", "call_1"),
("function_call_output", "call_1"),
("function_call", "call_2"),
("function_call_output", "call_2"),
("function_call", "call_3"),
("function_call_output", "call_3"),
)
third_input: Final = upstream_requests[2]["input"]
fourth_input: Final = upstream_requests[3]["input"]
assert isinstance(third_input, list) and isinstance(fourth_input, list)
assert sum(1 for item in third_input if isinstance(item, dict) and item.get("role") == "user") == 1
assert sum(1 for item in fourth_input if isinstance(item, dict) and item.get("role") == "user") == 1
assert len(tool_calls(peer.drain())) == 3

9
tests/integration/providers/_api_base_url_join_support.py
View file

@ -13,6 +13,7 @@ from pydantic import JsonValue
COUNTED_TOKENS: Final = 7
RELEVANCE: Final = 0.75
AZURE_DEPLOYMENT: Final = "gpt-url-join"
_USAGE: Final = {"input_tokens": 3, "output_tokens": 2}
@ -146,12 +147,20 @@ def _cohere_rerank(request: Request) -> Reply:
)
def _titan_embedding(_: Request) -> Reply:
return _json({"embedding": list(VECTOR), "inputTextTokenCount": 2})
_BY_PREFIX: Final[Mapping[str, Callable[[Request], Reply]]] = MappingProxyType(
{
"an": _anthropic,
"bi": lambda request: _json(anthropic_message(marker_of(request))),
"bc": _bedrock_converse,
"be": _titan_embedding,
"gm": _gemini,
"vx": _gemini,
"fd": _anthropic,
"az": _openai_like(AZURE_DEPLOYMENT),
"mi": _openai_like("mistral-small-latest"),
"db": _openai_like("databricks-meta-llama-3-3-70b-instruct"),
"co": _cohere_rerank,

37
tests/integration/providers/test_api_base_url_join_host_owned_proxy.py
View file

@ -1,6 +1,7 @@
import os
from collections.abc import Callable, Iterator
from pathlib import Path
from types import MappingProxyType
from typing import Final
import pytest
@ -8,6 +9,7 @@ import yaml
from integration._support.client import JSON_OBJECT, Gateway, gateway_from_environment
from integration._support.forward_proxy import ForwardProxy, refusing_forward_proxy
from integration._support.process import graceful_stop_seconds, owned_proxy
from integration._support.vertex import service_account_json
from integration.providers._api_base_trailing_slash_support import KEY, new_marker
from pydantic import JsonValue
@ -17,6 +19,12 @@ _RERANK: Final = "host-rerank"
_AZURE_CLAUDE: Final = "host-azure-claude"
_AZURE_COUNT: Final = "host-azure-count"
_EMPTY_BASE: Final = "host-empty-base"
_VERTEX_HOST: Final = "vertex-gateway"
_VERTEX_PROJECT: Final = "url-join-project"
_VERTEX_LOCATION: Final = "us-east5"
_VERTEX_CLAUDE_BASES: Final = MappingProxyType(
{"host-vertex-claude": f"http://{_VERTEX_HOST}", "host-vertex-claude-slash": f"http://{_VERTEX_HOST}/"}
)
_SCRUBBED_PREFIXES: Final = ("ANTHROPIC_", "AZURE_AI_", "INFINITY_")
_PROXY_VARIABLES: Final = frozenset({"HTTPS_PROXY", "HTTP_PROXY", "ALL_PROXY", "NO_PROXY"})
_load_yaml: Final[Callable[[str], object]] = yaml.safe_load
@ -28,7 +36,20 @@ def _deployment(name: str, model: str, api_base: str) -> dict[str, JsonValue]:
return {"model_name": name, "litellm_params": {"model": model, "api_base": api_base, "api_key": KEY}}
def _config(directory: Path) -> Path:
def _vertex_claude_deployment(name: str, api_base: str, token_url: str) -> dict[str, JsonValue]:
return {
"model_name": name,
"litellm_params": {
"model": "vertex_ai/claude-haiku-5-5",
"api_base": api_base,
"vertex_project": _VERTEX_PROJECT,
"vertex_location": _VERTEX_LOCATION,
"vertex_credentials": service_account_json(_VERTEX_PROJECT, token_url),
},
}
def _config(directory: Path, token_url: str) -> Path:
base: Final = JSON_OBJECT.validate_python(_load_yaml(Path("tests/integration/proxy_config.yaml").read_text()))
configuration: Final = {
**base,
@ -38,6 +59,7 @@ def _config(directory: Path) -> Path:
_deployment(_AZURE_CLAUDE, "azure_ai/claude-haiku-5-5", f"http://{_AZURE_HOST}/anthropic"),
_deployment(_AZURE_COUNT, "azure_ai/claude-haiku-5-5", f"http://{_AZURE_HOST}"),
_deployment(_EMPTY_BASE, "anthropic/claude-haiku-5-5", ""),
*(_vertex_claude_deployment(name, base, token_url) for name, base in _VERTEX_CLAUDE_BASES.items()),
],
"router_settings": {"disable_cooldowns": True, "num_retries": 0},
}
@ -72,7 +94,7 @@ def isolated(forward: ForwardProxy, tmp_path_factory: pytest.TempPathFactory) ->
"NO_PROXY": "127.0.0.1,localhost",
"PROXY_BATCH_POLLING_ENABLED": "false",
},
config=_config(directory),
config=_config(directory, shared.upstream_url.rstrip("/")),
remove_environment=_scrubbed(),
) as candidate:
yield candidate
@ -138,3 +160,14 @@ def test_azure_claude_count_tokens_on_a_bare_host_starting_with_anthropic_adds_t
def test_empty_api_base_falls_back_to_the_provider_default_host(isolated: Gateway, forward: ForwardProxy) -> None:
targets: Final = _targets_after(forward, lambda: _chat(isolated, _EMPTY_BASE))
assert targets == {"api.anthropic.com:443"}
@pytest.mark.parametrize("model", tuple(_VERTEX_CLAUDE_BASES), ids=("no-slash", "slash"))
def test_vertex_partner_claude_on_a_bare_host_reaches_one_single_slash_url(
isolated: Gateway, forward: ForwardProxy, model: str
) -> None:
targets: Final = _targets_after(forward, lambda: _chat(isolated, model))
assert targets == {
f"http://{_VERTEX_HOST}/v1/projects/{_VERTEX_PROJECT}/locations/{_VERTEX_LOCATION}"
"/publishers/anthropic/models/claude-haiku-5-5:rawPredict"
}

99
tests/integration/providers/test_api_base_url_join_wire.py
View file

@ -7,6 +7,7 @@ import openai
import pytest
from integration._support.client import Gateway, Scenario, eventually, list_value, object_value, string_value
from integration._support.database import read_rows
from integration._support.vertex import service_account_json
from integration._support.wire import Reply, Request, wire_server
from integration.providers._api_base_trailing_slash_support import (
KEY,
@ -22,7 +23,7 @@ from integration.providers._api_base_trailing_slash_support import (
reply_text,
request_json,
)
from integration.providers._api_base_url_join_support import COUNTED_TOKENS, RELEVANCE, join_peer
from integration.providers._api_base_url_join_support import AZURE_DEPLOYMENT, COUNTED_TOKENS, RELEVANCE, join_peer
from pydantic import JsonValue
_HAIKU: Final = "anthropic.claude-haiku-5-5"
@ -49,10 +50,20 @@ _RESPONSES_ROWS: Final = (_ANTHROPIC, _GEMINI)
_EMBEDDING_ROWS: Final = (
Row("mistral/mistral-embed", "/mi/v1", "/embeddings"),
Row("gemini/gemini-embedding-2", "/gm/v1beta", "/models/gemini-embedding-2:batchEmbedContents"),
Row("bedrock/amazon.titan-embed-text-v2:0", "/be", "/model/amazon.titan-embed-text-v2%3A0/invoke"),
)
_TRANSCRIPTION: Final = Row("mistral/voxtral-mini-latest", "/mi/v1", "/audio/transcriptions")
_FULL_ROUTE_SUFFIXES: Final = ("", "/", "?beta=true")
_FULL_ROUTE_IDS: Final = ("exact", "slash", "query")
_FULL_ROUTE_SUFFIXES: Final = (("", ""), ("/", ""), ("?beta=true", "?beta=true"), ("/?beta=true", "?beta=true"))
_FULL_ROUTE_IDS: Final = ("exact", "slash", "query", "slash-query")
_VERTEX_PROJECT: Final = "url-join-project"
_VERTEX_LOCATION: Final = "us-east5"
_VERTEX_GEMINI_PATH: Final = (
f"/vx/v1/projects/{_VERTEX_PROJECT}/locations/{_VERTEX_LOCATION}/publishers/google/models/gemini-3.8-flash"
)
_AZURE_V1_BASES: Final = ("/az", "/az/", "/az/openai/v1", "/az/openai/v1/")
_AZURE_V1_IDS: Final = ("host", "host-slash", "route", "route-slash")
_AZURE_ANTHROPIC_BASES: Final = ("/fd/anthropic", "/fd/anthropic/v1/messages")
_AZURE_ANTHROPIC_IDS: Final = ("anthropic-root", "messages-route")
def _ids(rows: tuple[Row, ...]) -> tuple[str, ...]:
@ -65,6 +76,17 @@ def _deploy(scenario: Scenario, row: Row, api_base: str) -> str:
return scenario.model(model=row.model, api_base=api_base, api_key=KEY)
def _deploy_vertex(gateway: Gateway, scenario: Scenario, model: str, api_base: str) -> str:
return scenario.model(
model=model,
api_base=api_base,
api_key=None,
vertex_project=_VERTEX_PROJECT,
vertex_location=_VERTEX_LOCATION,
vertex_credentials=service_account_json(_VERTEX_PROJECT, gateway.upstream_url.rstrip("/")),
)
def _openai(gateway: Gateway) -> openai.OpenAI:
return openai.OpenAI(base_url=str(gateway.client.base_url) + "/v1", api_key=gateway.key, max_retries=0)
@ -80,6 +102,13 @@ def _chat(gateway: Gateway, model: str, marker: str) -> str | None:
return completion.choices[0].message.content
def _streamed_chat(gateway: Gateway, model: str, marker: str) -> str:
stream: Final = _openai(gateway).chat.completions.create(
model=model, messages=[{"role": "user", "content": marker}], stream=True
)
return "".join(chunk.choices[0].delta.content or "" for chunk in stream if chunk.choices)
def _message_text(gateway: Gateway, model: str, marker: str) -> list[str]:
message: Final = _anthropic(gateway).messages.create(
model=model, max_tokens=16, messages=[{"role": "user", "content": marker}]
@ -270,22 +299,74 @@ def test_cache_hit_on_a_slash_base_skips_the_upstream_and_logs_its_own_row(gatew
assert "//" not in urlsplit(string_value(hit["api_base"])).path, hit
@pytest.mark.parametrize("suffix", _FULL_ROUTE_SUFFIXES, ids=_FULL_ROUTE_IDS)
def test_chat_base_that_already_names_the_route_keeps_its_path_and_query(gateway: Gateway, suffix: str) -> None:
@pytest.mark.parametrize(("suffix", "sent"), _FULL_ROUTE_SUFFIXES, ids=_FULL_ROUTE_IDS)
def test_chat_base_that_already_names_the_route_keeps_its_path_and_query(
gateway: Gateway, suffix: str, sent: str
) -> None:
marker: Final = new_marker()
with wire_server(join_peer) as wire, gateway.scenario() as scenario:
model: Final = _deploy(scenario, _ANTHROPIC, f"{wire.url}/an/v1/messages{suffix}")
assert _chat(gateway, model, marker) == reply_text(marker)
assert_one_request(wire.drain(), f"/an/v1/messages{suffix.rstrip('/')}", marker)
assert_one_request(wire.drain(), f"/an/v1/messages{sent}", marker)
@pytest.mark.parametrize("suffix", _FULL_ROUTE_SUFFIXES, ids=_FULL_ROUTE_IDS)
def test_messages_base_that_already_names_the_route_keeps_its_path_and_query(gateway: Gateway, suffix: str) -> None:
@pytest.mark.parametrize(("suffix", "sent"), _FULL_ROUTE_SUFFIXES, ids=_FULL_ROUTE_IDS)
def test_messages_base_that_already_names_the_route_keeps_its_path_and_query(
gateway: Gateway, suffix: str, sent: str
) -> None:
marker: Final = new_marker()
with wire_server(join_peer) as wire, gateway.scenario() as scenario:
model: Final = _deploy(scenario, _ANTHROPIC, f"{wire.url}/an/v1/messages{suffix}")
assert _message_text(gateway, model, marker) == [reply_text(marker)]
assert_one_request(wire.drain(), f"/an/v1/messages{suffix.rstrip('/')}", marker)
assert_one_request(wire.drain(), f"/an/v1/messages{sent}", marker)
@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS)
def test_vertex_model_url_base_calls_its_method_with_no_slash_before_the_colon(gateway: Gateway, trailing: str) -> None:
marker: Final = new_marker()
with wire_server(join_peer) as wire, gateway.scenario() as scenario:
model: Final = _deploy_vertex(
gateway, scenario, "vertex_ai/gemini-3.8-flash", f"{wire.url}{_VERTEX_GEMINI_PATH}{trailing}"
)
assert _chat(gateway, model, marker) == reply_text(marker)
assert_one_request(wire.drain(), f"{_VERTEX_GEMINI_PATH}:generateContent", marker, keyed=False)
@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS)
def test_vertex_model_url_base_streams_with_no_slash_before_the_colon(gateway: Gateway, trailing: str) -> None:
marker: Final = new_marker()
with wire_server(join_peer) as wire, gateway.scenario() as scenario:
model: Final = _deploy_vertex(
gateway, scenario, "vertex_ai/gemini-3.8-flash", f"{wire.url}{_VERTEX_GEMINI_PATH}{trailing}"
)
assert _streamed_chat(gateway, model, marker) == reply_text(marker)
assert_one_request(wire.drain(), f"{_VERTEX_GEMINI_PATH}:streamGenerateContent?alt=sse", marker, keyed=False)
@pytest.mark.parametrize("base", _AZURE_V1_BASES, ids=_AZURE_V1_IDS)
def test_azure_v1_chat_reaches_one_openai_v1_route(gateway: Gateway, base: str) -> None:
marker: Final = new_marker()
with wire_server(join_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=f"azure/{AZURE_DEPLOYMENT}", api_base=f"{wire.url}{base}", api_key=KEY, api_version="v1"
)
assert _chat(gateway, model, marker) == reply_text(marker)
assert_one_request(wire.drain(), "/az/openai/v1/chat/completions", marker)
@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS)
@pytest.mark.parametrize("base", _AZURE_ANTHROPIC_BASES, ids=_AZURE_ANTHROPIC_IDS)
def test_azure_ai_count_tokens_reaches_one_count_tokens_route(gateway: Gateway, base: str, trailing: str) -> None:
marker: Final = new_marker()
with wire_server(join_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model="azure_ai/claude-haiku-5-5", api_base=f"{wire.url}{base}{trailing}", api_key=KEY
)
response: Final = gateway.request(
"POST", "/v1/messages/count_tokens", {"model": model, "messages": [{"role": "user", "content": marker}]}
)
assert body_of(response)["input_tokens"] == COUNTED_TOKENS
assert_one_request(wire.drain(), "/fd/anthropic/v1/messages/count_tokens", marker)
@pytest.mark.parametrize("suffix", ("", "/", "?tier=a"), ids=("exact", "slash", "query"))

91
tests/integration/providers/test_bedrock_converse_redacted_thinking_wire.py
View file

@ -6,7 +6,7 @@ from typing import Final
from urllib.parse import unquote
import pytest
from integration._support.client import Gateway
from integration._support.client import Gateway, Scenario
from integration._support.upstream import _aws_event_frame
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import BaseModel, ConfigDict, JsonValue
@ -414,3 +414,92 @@ def test_converse_redacted_thinking_replays_to_the_anthropic_api_as_claude_nativ
{"type": "redacted_thinking", "data": NATIVE_REDACTED},
{"type": "text", "text": FIRST_ANSWER},
)
_SIGNED_THINKING: Final[dict[str, JsonValue]] = {
"type": "thinking",
"thinking": "Checking the forecast.",
"signature": "scripted-signature-" + "s" * 32,
}
_CONVERSE_REDACTED_PART: Final[dict[str, JsonValue]] = {"reasoningContent": {"redactedContent": CONVERSE_REDACTED}}
_CONVERSE_SIGNED_PART: Final[dict[str, JsonValue]] = {
"reasoningContent": {
"reasoningText": {"text": _SIGNED_THINKING["thinking"], "signature": _SIGNED_THINKING["signature"]}
}
}
def _replayed_to_converse(
gateway: Gateway, wire: Wire, model: str, assistants: list[JsonValue], **fields: JsonValue
) -> tuple[dict[str, JsonValue], ...]:
reply: Final = _chat(
gateway,
{
"model": model,
"thinking": THINKING,
"max_tokens": 2048,
"messages": [{"role": "user", "content": QUESTION}, *assistants, {"role": "user", "content": FOLLOW_UP}],
**fields,
},
)
assert reply.content == FINAL_ANSWER
return _replayed_assistant_turn(wire).content
def _claude_on_converse(scenario: Scenario, wire: Wire) -> str:
return scenario.model(
model=f"bedrock/converse/{CLAUDE_MODEL_ID}",
api_key=TOKEN,
aws_region_name="us-east-1",
aws_bedrock_runtime_endpoint=wire.url,
)
@pytest.mark.parametrize(
"thinking_blocks", [None, [_redacted(NATIVE_REDACTED)]], ids=["content", "content-and-thinking-blocks"]
)
def test_claude_redacted_thinking_replayed_in_content_reaches_converse_once(
gateway: Gateway, thinking_blocks: list[JsonValue] | None
) -> None:
assistant: Final[dict[str, JsonValue]] = {
"role": "assistant",
"content": [_redacted(NATIVE_REDACTED), {"type": "text", "text": FIRST_ANSWER}],
**({} if thinking_blocks is None else {"thinking_blocks": thinking_blocks}),
}
with wire_server(_bedrock_peer()) as wire, gateway.scenario() as scenario:
sent: Final = _replayed_to_converse(gateway, wire, _claude_on_converse(scenario, wire), [assistant])
assert sent == (_CONVERSE_REDACTED_PART, {"text": FIRST_ANSWER})
@pytest.mark.parametrize(
"block, converse_part",
[(_redacted(NATIVE_REDACTED), _CONVERSE_REDACTED_PART), (_SIGNED_THINKING, _CONVERSE_SIGNED_PART)],
ids=["redacted", "signed"],
)
def test_block_split_across_merged_assistant_messages_reaches_converse_once(
gateway: Gateway, block: dict[str, JsonValue], converse_part: dict[str, JsonValue]
) -> None:
assistants: Final[list[JsonValue]] = [
{"role": "assistant", "content": [block, {"type": "text", "text": FIRST_ANSWER}]},
{"role": "assistant", "content": "Checking again.", "thinking_blocks": [block]},
]
with wire_server(_bedrock_peer()) as wire, gateway.scenario() as scenario:
sent: Final = _replayed_to_converse(gateway, wire, _claude_on_converse(scenario, wire), assistants)
assert sent == (converse_part, {"text": FIRST_ANSWER}, {"text": "Checking again."})
@pytest.mark.parametrize("drop", [True, False], ids=["drop-params", "no-drop-params"])
def test_converse_cache_point_after_replayed_thinking_follows_drop_params(gateway: Gateway, drop: bool) -> None:
assistant: Final[dict[str, JsonValue]] = {
"role": "assistant",
"content": [
{**_redacted(NATIVE_REDACTED), "cache_control": {"type": "ephemeral"}},
{"type": "text", "text": FIRST_ANSWER},
],
}
with wire_server(_bedrock_peer()) as wire, gateway.scenario() as scenario:
sent: Final = _replayed_to_converse(
gateway, wire, _claude_on_converse(scenario, wire), [assistant], drop_params=drop
)
replayed: Final = (_CONVERSE_REDACTED_PART, {"text": FIRST_ANSWER})
assert sent == (replayed if drop else (*replayed, {"cachePoint": {"type": "default"}}))

195
tests/integration/providers/test_claude_thinking_content_replay_wire.py Normal file
View file

@ -0,0 +1,195 @@
import uuid
from collections.abc import Mapping
from typing import Final, Literal
import pytest
from integration._support.anthropic_thinking import (
BEDROCK_MODEL,
MODEL,
NO_CACHE,
SIGNATURE,
THINKING,
answer,
marker_of,
prompt,
standard_peer,
thinking_block,
)
from integration._support.client import Gateway, Scenario
from integration._support.vertex import service_account_json
from integration._support.wire import Wire, wire_server
from pydantic import BaseModel, ConfigDict, JsonValue
Route = Literal["anthropic", "bedrock_invoke", "claude_platform", "vertex_ai", "azure_ai"]
Endpoint = Literal["chat", "messages"]
_VERTEX_PROJECT: Final = "scripted-project"
_VERTEX_LOCATION: Final = "us-east5"
_ENABLED: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": 1024}
_EPHEMERAL: Final[dict[str, JsonValue]] = {"type": "ephemeral"}
_SIGNED: Final = thinking_block(THINKING, SIGNATURE)
_REDACTED: Final[dict[str, JsonValue]] = {"type": "redacted_thinking", "data": "scripted-redacted-" + "r" * 32}
_BLANK_REDACTED: Final[dict[str, JsonValue]] = {"type": "redacted_thinking", "data": " "}
_ROUTES: Final[tuple[Route, ...]] = ("anthropic", "bedrock_invoke", "claude_platform", "vertex_ai", "azure_ai")
class _Frozen(BaseModel):
model_config = ConfigDict(frozen=True)
class _SentMessage(_Frozen):
role: str
content: str | tuple[dict[str, JsonValue], ...]
class _SentRequest(_Frozen):
messages: tuple[_SentMessage, ...]
def _deployment(scenario: Scenario, route: Route, wire_url: str, upstream_url: str) -> str:
match route:
case "anthropic":
return scenario.model(model=f"anthropic/{MODEL}", api_base=wire_url, api_key="scripted-key")
case "bedrock_invoke":
return scenario.model(
model=f"bedrock/invoke/{BEDROCK_MODEL}",
api_base=wire_url,
aws_access_key_id="AKIASCRIPTEDPROVIDER",
aws_secret_access_key="scripted-secret",
aws_region_name="us-east-1",
aws_bedrock_runtime_endpoint=wire_url,
)
case "claude_platform":
return scenario.model(
model=f"bedrock/claude_platform/{MODEL}",
api_base=wire_url,
api_key="scripted-platform-key",
aws_region_name="us-east-1",
workspace_id="scripted-workspace",
)
case "vertex_ai":
return scenario.model(
model=f"vertex_ai/{MODEL}",
api_base=(
f"{wire_url}/v1/projects/{_VERTEX_PROJECT}/locations/{_VERTEX_LOCATION}"
f"/publishers/anthropic/models/{MODEL}"
),
api_key=None,
vertex_project=_VERTEX_PROJECT,
vertex_location=_VERTEX_LOCATION,
vertex_credentials=service_account_json(_VERTEX_PROJECT, upstream_url.rstrip("/")),
)
case "azure_ai":
return scenario.model(model=f"azure_ai/{MODEL}", api_base=wire_url, api_key="scripted-azure-key")
def _text(text: str) -> dict[str, JsonValue]:
return {"type": "text", "text": text}
def _replayed_turn(
gateway: Gateway,
wire: Wire,
*,
endpoint: Endpoint = "chat",
model: str,
marker: str,
assistant: Mapping[str, JsonValue],
**fields: JsonValue,
) -> tuple[dict[str, JsonValue], ...]:
follow_up: Final = uuid.uuid4().hex
body: Final[dict[str, JsonValue]] = {
"model": model,
"max_tokens": 2048,
"thinking": _ENABLED,
"messages": [
{"role": "user", "content": prompt(marker)},
{"role": "assistant", **assistant},
{"role": "user", "content": prompt(follow_up)},
],
**NO_CACHE,
**fields,
}
path: Final = "/v1/chat/completions" if endpoint == "chat" else "/v1/messages"
response: Final = gateway.request("POST", path, body)
assert response.status_code == 200, response.text
assert answer(follow_up) in response.text, response.text
received: Final = wire.drain()
assert len(received) == 1, [request.body for request in received]
assert marker_of(received[0]) == follow_up, received[0].body
sent: Final = _SentRequest.model_validate_json(received[0].body).messages[1]
assert sent.role == "assistant", received[0].body
assert not isinstance(sent.content, str), received[0].body
return sent.content
@pytest.mark.parametrize("route", _ROUTES)
def test_redacted_thinking_replayed_in_assistant_content_reaches_claude_once(gateway: Gateway, route: Route) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(standard_peer) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, route, wire.url, gateway.upstream_url)
sent: Final = _replayed_turn(
gateway, wire, model=model, marker=marker, assistant={"content": [_REDACTED, _text(answer(marker))]}
)
assert sent == (_REDACTED, _text(answer(marker)))
def test_block_in_content_and_thinking_blocks_reaches_claude_once_in_order(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(standard_peer) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, "anthropic", wire.url, gateway.upstream_url)
sent: Final = _replayed_turn(
gateway,
wire,
model=model,
marker=marker,
assistant={"content": [_REDACTED, _text(answer(marker))], "thinking_blocks": [_SIGNED, _REDACTED]},
)
assert sent == (_SIGNED, _REDACTED, _text(answer(marker)))
def test_streamed_thinking_replayed_in_content_reaches_claude_without_its_index(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(standard_peer) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, "anthropic", wire.url, gateway.upstream_url)
sent: Final = _replayed_turn(
gateway,
wire,
model=model,
marker=marker,
assistant={"content": [{**_SIGNED, "index": 0}, _text(answer(marker))]},
)
assert sent == (_SIGNED, _text(answer(marker)))
@pytest.mark.parametrize("drop", [True, False], ids=["drop-params", "no-drop-params"])
def test_replayed_thinking_cache_control_and_blank_data_follow_drop_params(gateway: Gateway, drop: bool) -> None:
marker: Final = uuid.uuid4().hex
content: Final = ({**_SIGNED, "cache_control": _EPHEMERAL}, _BLANK_REDACTED, _REDACTED, _text(answer(marker)))
with wire_server(standard_peer) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, "anthropic", wire.url, gateway.upstream_url)
sent: Final = _replayed_turn(
gateway, wire, model=model, marker=marker, assistant={"content": list(content)}, drop_params=drop
)
assert sent == ((_SIGNED, _REDACTED, _text(answer(marker))) if drop else content)
@pytest.mark.parametrize("endpoint", ["chat", "messages"])
def test_assistant_cache_injection_point_never_marks_a_thinking_part(gateway: Gateway, endpoint: Endpoint) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(standard_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=f"anthropic/{MODEL}",
api_base=wire.url,
api_key="scripted-key",
cache_control_injection_points=[{"location": "message", "role": "assistant"}],
)
sent: Final = _replayed_turn(
gateway,
wire,
endpoint=endpoint,
model=model,
marker=marker,
assistant={"content": [_text(answer(marker)), _SIGNED]},
)
assert sent == ({**_text(answer(marker)), "cache_control": _EPHEMERAL}, _SIGNED)

143
tests/integration/providers/test_count_tokens_error_shapes_wire.py
View file

@ -1,5 +1,6 @@
import json
from collections.abc import Iterator
from collections.abc import Callable, Iterator
from types import MappingProxyType
from typing import Final
import anthropic
@ -8,6 +9,7 @@ import pytest
from integration._support.client import Gateway, gateway_from_environment
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
from integration._support.wire import Reply, Request, wire_server
from litellm.types.llms.anthropic import ANTHROPIC_OAUTH_BETA_HEADER
from pydantic import JsonValue, TypeAdapter
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
@ -26,6 +28,14 @@ _PROVIDER_KEY_VARIABLES: Final = (
"AZURE_TENANT_ID",
"OPENAI_API_KEY",
)
_ENVIRONMENT_KEYS: Final = MappingProxyType(
{
"ANTHROPIC_API_KEY": "synthetic-anthropic-env-key",
"AZURE_API_KEY": "synthetic-azure-env-key",
"OPENAI_API_KEY": "synthetic-openai-env-key",
}
)
_SUBSCRIPTION_TOKEN: Final = "sk-ant-oat01-synthetic-subscription-token"
@pytest.fixture(scope="module")
@ -44,14 +54,14 @@ def keyless_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedPro
@pytest.fixture(scope="module")
def azure_environment_key_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]:
def environment_key_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]:
with (
gateway_from_environment() as gateway,
owned_proxy_process(
gateway,
tmp_path_factory.mktemp("count-tokens-azure-environment-key"),
{"AZURE_API_KEY": "synthetic-azure-env-key"},
remove_environment=tuple(name for name in _PROVIDER_KEY_VARIABLES if name != "AZURE_API_KEY"),
tmp_path_factory.mktemp("count-tokens-environment-keys"),
_ENVIRONMENT_KEYS,
remove_environment=tuple(name for name in _PROVIDER_KEY_VARIABLES if name not in _ENVIRONMENT_KEYS),
workers=1,
) as owned,
):
@ -147,9 +157,9 @@ def test_anthropic_count_without_a_provider_credential_is_an_anthropic_401(
def test_azure_ai_anthropic_count_sends_the_azure_environment_key_as_x_api_key(
azure_environment_key_proxy: OwnedProxy,
environment_key_proxy: OwnedProxy,
) -> None:
gateway: Final = azure_environment_key_proxy.gateway
gateway: Final = environment_key_proxy.gateway
def count(request: Request) -> Reply:
assert request.headers.get("x-api-key") == "synthetic-azure-env-key", request.headers
@ -257,3 +267,122 @@ def test_google_count_tokens_errors_use_the_google_envelope(keyless_proxy: Owned
assert _JSON_OBJECT.validate_json(malformed.content) == {
"error": {"code": 400, "message": "Invalid Google countTokens request body", "status": "INVALID_ARGUMENT"}
}
def _counted(tokens: int) -> Reply:
return Reply(body=json.dumps({"object": "response.input_tokens", "input_tokens": tokens}).encode())
def _anthropic_count(model: str) -> tuple[str, dict[str, JsonValue]]:
return "/v1/messages/count_tokens", {"model": model, "messages": _MESSAGES}
def _responses_count(model: str) -> tuple[str, dict[str, JsonValue]]:
return "/v1/responses/input_tokens", {"model": model, "input": "Count this synthetic request."}
@pytest.mark.parametrize(
("provider_model", "base_path", "count_request", "header", "credential"),
(
("anthropic/claude-sonnet-5-5", "", _anthropic_count, "x-api-key", "synthetic-anthropic-env-key"),
("openai/gpt-4o-mini", "/v1", _responses_count, "authorization", "Bearer synthetic-openai-env-key"),
),
ids=("anthropic", "openai"),
)
def test_count_uses_a_provider_key_set_only_in_the_proxy_environment(
environment_key_proxy: OwnedProxy,
provider_model: str,
base_path: str,
count_request: Callable[[str], tuple[str, dict[str, JsonValue]]],
header: str,
credential: str,
) -> None:
gateway: Final = environment_key_proxy.gateway
with wire_server(lambda _request: _counted(47)) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=provider_model, api_base=f"{wire.url}{base_path}", api_key=None, num_retries=0
)
path, body = count_request(model)
counted: Final = gateway.request("POST", path, body)
assert counted.status_code == 200, counted.text
assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == 47, counted.text
received: Final = wire.drain()
assert [request.headers.get(header) for request in received] == [credential], received
def _google_count(model: str) -> tuple[str, dict[str, JsonValue]]:
return f"/v1beta/models/{model}:countTokens", {
"contents": [{"role": "user", "parts": [{"text": "Count this synthetic request."}]}]
}
def _utils_count(model: str) -> tuple[str, dict[str, JsonValue]]:
return "/utils/token_counter?call_endpoint=true", {"model": model, "messages": _MESSAGES}
@pytest.mark.parametrize(
("count_request", "google_status"),
((_google_count, "UNAUTHENTICATED"), (_utils_count, None)),
ids=("google", "utils"),
)
def test_count_entrypoints_report_a_missing_provider_key_in_their_own_shape(
keyless_proxy: OwnedProxy,
count_request: Callable[[str], tuple[str, dict[str, JsonValue]]],
google_status: str | None,
) -> None:
gateway: Final = keyless_proxy.gateway
with wire_server(_never_called) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model="anthropic/claude-sonnet-5-5", api_base=wire.url, api_key=None, num_retries=0
)
path, body = count_request(model)
rejected: Final = gateway.request("POST", path, body)
assert rejected.status_code == 401, rejected.text
error: Final = _JSON_OBJECT.validate_json(rejected.content)["error"]
assert isinstance(error, dict), rejected.text
assert f"No Anthropic API key is set for model {model}" in str(error["message"]), rejected.text
if google_status is not None:
assert (error["code"], error["status"]) == (401, google_status), rejected.text
assert wire.drain() == ()
def test_anthropic_count_uses_the_callers_subscription_token_when_the_deployment_has_no_key(
keyless_proxy: OwnedProxy,
) -> None:
gateway: Final = keyless_proxy.gateway
with wire_server(lambda _request: _counted(41)) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model="anthropic/claude-sonnet-5-5", api_base=wire.url, api_key=None, num_retries=0
)
counted: Final = gateway.request(
"POST",
"/v1/messages/count_tokens",
{"model": model, "messages": _MESSAGES},
headers={"Authorization": f"Bearer {_SUBSCRIPTION_TOKEN}", "x-litellm-api-key": f"Bearer {gateway.key}"},
)
assert counted.status_code == 200, counted.text
assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == 41, counted.text
(received,) = wire.drain()
assert received.headers.get("authorization") == f"Bearer {_SUBSCRIPTION_TOKEN}", received.headers
assert "x-api-key" not in received.headers, received.headers
assert ANTHROPIC_OAUTH_BETA_HEADER in received.headers.get("anthropic-beta", ""), received.headers
def test_anthropic_count_keeps_the_deployments_key_over_the_callers_subscription_token(
keyless_proxy: OwnedProxy,
) -> None:
gateway: Final = keyless_proxy.gateway
with wire_server(lambda _request: _counted(43)) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model="anthropic/claude-sonnet-5-5", api_base=wire.url, api_key="synthetic-deployment-key", num_retries=0
)
counted: Final = gateway.request(
"POST",
"/v1/messages/count_tokens",
{"model": model, "messages": _MESSAGES},
headers={"Authorization": f"Bearer {_SUBSCRIPTION_TOKEN}", "x-litellm-api-key": f"Bearer {gateway.key}"},
)
assert counted.status_code == 200, counted.text
(received,) = wire.drain()
assert received.headers.get("x-api-key") == "synthetic-deployment-key", received.headers
assert _SUBSCRIPTION_TOKEN not in json.dumps(dict(received.headers)), received.headers

23
tests/integration/providers/test_decisions_wire.py
View file

@ -137,7 +137,7 @@ _PROVIDERS: Final = (
"databricks-openjev-qwen35-4b",
_API_KEY,
False,
None,
"databricks/databricks-openjev-qwen35-4b",
),
)
_PERPLEXITY: Final = _PROVIDERS[0]
@ -576,13 +576,26 @@ def test_upstream_errors_keep_their_status_and_log_an_unbilled_failure(gateway:
)
def test_upstream_success_without_answers_is_a_gateway_side_server_error(gateway: Gateway) -> None:
@pytest.mark.parametrize(
("reply", "field"),
(
pytest.param({"model": _PERPLEXITY.body_model, "usage": _USAGE}, "answers", id="missing_answers"),
pytest.param(
{"model": _PERPLEXITY.body_model, "answers": {"defect": {"type": "Request timed out"}}, "usage": _USAGE},
"answers.defect",
id="answer_type_timeout_text",
),
),
)
def test_a_malformed_upstream_success_is_a_server_error_naming_the_field(
gateway: Gateway, reply: dict[str, JsonValue], field: str
) -> None:
with gateway.scenario() as scenario:
handle: Final = _register(scenario, {"model": _PERPLEXITY.body_model, "usage": _USAGE})
handle: Final = _register(scenario, reply)
model: Final = _deployment(scenario, handle, _PERPLEXITY)
response: Final = _decide(gateway, model, num_retries=0)
assert 500 <= response.status_code < 600, response.text
assert "answers" in response.text
assert response.status_code == 500, response.text
assert field in response.text
assert _spend_row(response.headers["x-litellm-call-id"])["status"] == "failure"

207
tests/integration/providers/test_gemini_default_output_budget_wire.py Normal file
View file

@ -0,0 +1,207 @@
from __future__ import annotations
import itertools
import uuid
from collections.abc import Iterator, Mapping
from typing import Final
import pytest
from _pytest.mark.structures import ParameterSet
from integration._support.client import JSON_OBJECT, Gateway, Scenario
from integration._support.process import graceful_stop_seconds, owned_proxy
from integration._support.wire import Wire, wire_server
from integration.providers._gemini_long_decode_support import (
API_KEY,
BACKEND,
CAP,
CUT_TOKEN,
FIRST_HALF,
HOP_TOKEN,
INPUT_COST,
NO_CACHE,
OUTPUT_COST,
PLAIN,
PROVIDERS,
RESUMED,
SECOND_HALF,
Endpoint,
Provider,
answered,
answers,
at,
complete,
post,
received,
register,
request_body,
upstream,
)
from pydantic import JsonValue
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
_BUDGET_ENV: Final = "LITELLM_GEMINI_DEFAULT_MAX_OUTPUT_TOKENS"
_SMALL_BACKEND: Final = "gemini-2.5-flash"
_LIMIT_KEYS: Final = ("max_output_tokens", "maxOutputTokens")
_JOINED: Final = FIRST_HALF.text + SECOND_HALF.text
def _row_limit(key: str) -> int:
limit: Final = JSON_OBJECT.validate_python(GetModelCostMap.load_local_model_cost_map()[key]).get(
"max_output_tokens"
)
assert isinstance(limit, int), (key, limit)
return limit
_SMALL_LIMIT: Final = _row_limit(f"gemini/{_SMALL_BACKEND}")
_LARGE_LIMIT: Final = min(_row_limit(f"{provider}/{BACKEND}") for provider in PROVIDERS)
_BUDGET: Final = (_SMALL_LIMIT + _LARGE_LIMIT) // 2
@pytest.fixture(scope="module")
def budgeted(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]:
assert _SMALL_LIMIT < _BUDGET < _LARGE_LIMIT, (_SMALL_LIMIT, _BUDGET, _LARGE_LIMIT)
directory: Final = tmp_path_factory.mktemp("gemini-default-output-budget")
with owned_proxy(session_gateway, directory, {_BUDGET_ENV: str(_BUDGET)}, workers=2) as owned:
yield owned
def _limits_sent(body: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
config: Final = body.get("generationConfig")
assert isinstance(config, dict), body
return {key: config[key] for key in _LIMIT_KEYS if key in config}
def _mode(stream: bool) -> str:
return "stream" if stream else "sync"
_REACHED: Final[tuple[ParameterSet, ...]] = tuple(
pytest.param(provider, endpoint, stream, id=f"{provider}-{endpoint}-{_mode(stream)}")
for provider, endpoint, stream in (
("gemini", "chat", False),
("gemini", "chat", True),
("vertex_ai", "chat", False),
("vertex_ai", "chat", True),
("gemini", "responses", False),
("vertex_ai", "responses", True),
)
)
@pytest.mark.parametrize(("provider", "endpoint", "stream"), _REACHED)
def test_request_with_no_output_limit_reaches_google_with_the_env_budget(
budgeted: Gateway, provider: Provider, endpoint: Endpoint, stream: bool
) -> None:
with wire_server(upstream(answers(f"plain-{uuid.uuid4().hex}", PLAIN))) as wire, budgeted.scenario() as scenario:
model: Final = register(budgeted, scenario, provider, wire)
response: Final = post(budgeted, endpoint, request_body(endpoint, model, stream, cap=None))
observed: Final = answered(endpoint, stream, response)
assert (observed.answer, observed.finishes) == (PLAIN.text, (complete(endpoint),)), observed
(only,) = received(wire, provider, stream, (False,))
assert _limits_sent(only) == {"max_output_tokens": _BUDGET}, only
@pytest.mark.parametrize(
("provider", "stream"),
[
pytest.param(provider, stream, id=f"{provider}-{_mode(stream)}")
for provider, stream in itertools.product(PROVIDERS, (False, True))
],
)
def test_request_with_no_output_limit_is_continued_into_one_joined_answer(
budgeted: Gateway, provider: Provider, stream: bool
) -> None:
respond: Final = upstream(
answers(f"first-{uuid.uuid4().hex}", FIRST_HALF), answers(f"second-{uuid.uuid4().hex}", SECOND_HALF)
)
with wire_server(respond) as wire, budgeted.scenario() as scenario:
model: Final = register(budgeted, scenario, provider, wire)
response: Final = post(budgeted, "chat", request_body("chat", model, stream, cap=None))
observed: Final = answered("chat", stream, response)
assert (observed.answer, observed.streamed, observed.finishes) == (_JOINED, _JOINED, ("stop",)), observed
assert observed.tokens == () and HOP_TOKEN not in observed.raw, observed.raw
opening, continuing = received(wire, provider, stream, (False, True))
assert _limits_sent(opening) == {"max_output_tokens": _BUDGET}, opening
assert continuing == {**opening, "continuationToken": HOP_TOKEN}, (opening, continuing)
@pytest.mark.parametrize(
("endpoint", "limit", "sent"),
[
pytest.param("chat", {"max_tokens": CAP}, {"max_output_tokens": CAP}, id="chat-max_tokens"),
pytest.param(
"chat", {"max_completion_tokens": CAP}, {"max_output_tokens": CAP}, id="chat-max_completion_tokens"
),
pytest.param(
"chat",
{"extra_body": {"generationConfig": {"maxOutputTokens": CAP}}},
{"maxOutputTokens": CAP},
id="chat-extra_body-generationConfig",
),
pytest.param(
"responses", {"max_output_tokens": CAP}, {"max_output_tokens": CAP}, id="responses-max_output_tokens"
),
],
)
def test_callers_own_output_limit_wins_over_the_env_budget(
budgeted: Gateway, endpoint: Endpoint, limit: Mapping[str, JsonValue], sent: Mapping[str, JsonValue]
) -> None:
with wire_server(upstream(answers(f"plain-{uuid.uuid4().hex}", PLAIN))) as wire, budgeted.scenario() as scenario:
model: Final = register(budgeted, scenario, "gemini", wire)
body: Final = request_body(endpoint, model, False, {**NO_CACHE, **limit}, cap=None)
observed: Final = answered(endpoint, False, post(budgeted, endpoint, body))
assert observed.answer == PLAIN.text, observed
(only,) = received(wire, "gemini", False, (False,))
assert _limits_sent(only) == sent, only
def _deploy(scenario: Scenario, wire: Wire, backend: str, model_info: Mapping[str, JsonValue] | None) -> str:
return scenario.model(
model=f"gemini/{backend}",
api_base=wire.url,
api_key=API_KEY,
input_cost_per_token=INPUT_COST,
output_cost_per_token=OUTPUT_COST,
model_info=model_info,
)
def test_env_budget_is_capped_per_deployment_and_never_at_a_siblings_limit(budgeted: Gateway) -> None:
lowered: Final = _BUDGET // 4
with wire_server(upstream(answers(f"plain-{uuid.uuid4().hex}", PLAIN))) as wire, budgeted.scenario() as scenario:
deployments: Final = {
"row-below-budget": _deploy(scenario, wire, _SMALL_BACKEND, None),
"model-info-above-row": _deploy(scenario, wire, _SMALL_BACKEND, {"max_output_tokens": 2 * _SMALL_LIMIT}),
"model-info-below-budget": _deploy(scenario, wire, BACKEND, {"max_output_tokens": lowered}),
"sibling-without-model-info": _deploy(scenario, wire, BACKEND, None),
}
def limit_sent(model: str) -> dict[str, JsonValue]:
_ = answered("chat", False, post(budgeted, "chat", request_body("chat", model, False, cap=None)))
(only,) = wire.drain()
return _limits_sent(JSON_OBJECT.validate_json(only.body))
assert {name: limit_sent(model) for name, model in deployments.items()} == {
"row-below-budget": {"max_output_tokens": _SMALL_LIMIT},
"model-info-above-row": {"max_output_tokens": _SMALL_LIMIT},
"model-info-below-budget": {"max_output_tokens": lowered},
"sibling-without-model-info": {"max_output_tokens": _BUDGET},
}
def test_resume_with_a_continuation_token_gets_no_env_budget(budgeted: Gateway) -> None:
respond: Final = upstream(
answers(f"restarted-{uuid.uuid4().hex}", PLAIN), answers(f"resumed-{uuid.uuid4().hex}", RESUMED)
)
with wire_server(respond) as wire, budgeted.scenario() as scenario:
model: Final = register(budgeted, scenario, "gemini", wire)
body: Final = request_body("chat", model, False, {**NO_CACHE, "continuation_token": CUT_TOKEN}, cap=None)
observed: Final = answered("chat", False, post(budgeted, "chat", body))
assert (observed.answer, observed.finishes) == (RESUMED.text, ("stop",)), observed
(only,) = received(wire, "gemini", False, (True,))
assert only["continuationToken"] == CUT_TOKEN, only
assert [at(only, "generationConfig", key) for key in _LIMIT_KEYS] == [None, None], only

17
tests/integration/providers/test_gemini_long_decode_wire.py
View file

@ -1163,6 +1163,23 @@ def test_request_for_two_candidates_is_never_continued_and_each_choice_keeps_its
assert at(only, "generationConfig", "candidate_count") == 2, only
@pytest.mark.parametrize("provider", PROVIDERS)
def test_request_for_two_candidates_in_camel_case_is_never_continued(gateway: Gateway, provider: Provider) -> None:
with wire_server(_two_drafts) as wire, gateway.scenario() as scenario:
model: Final = register(gateway, scenario, provider, wire)
two_candidates: Final = {"extra_body": {"generationConfig": {"candidateCount": 2}}}
body: Final = request_body("chat", model, False, {**NO_CACHE, **two_candidates})
response: Final = post(gateway, "chat", body)
assert response.status_code == 200, response.text
choices: Final = items_of(JSON_OBJECT.validate_json(response.content).get("choices"))
assert [
(at(choice, "message", "content"), at(choice, "message", "provider_specific_fields", "continuation_token"))
for choice in choices
] == list(_DRAFTS), response.text
(only,) = received(wire, provider, False, (False,))
assert at(only, "generationConfig", "candidateCount") == 2, only
def _uncapped_cells() -> tuple[ParameterSet, ...]:
reached: Final = (("gemini", "chat"), ("vertex_ai", "chat"), ("gemini", "responses"))
return tuple(

201
tests/integration/providers/test_huggingface_tokenizer_downloads_wire.py Normal file
View file

@ -0,0 +1,201 @@
import json
import threading
import uuid
from collections.abc import Callable, Iterator, Mapping
from pathlib import Path
from types import MappingProxyType
from typing import Final
import pytest
from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
from tokenizers import Regex, Tokenizer, models, pre_tokenizers
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
_LLAMA_3_BACKEND: Final = "databricks-meta-llama-3-3-70b-instruct"
_DEEPSEEK_V4: Final = "azure_ai/DeepSeek-V4-Flash-0731"
_API_KEY: Final = "synthetic-databricks-key"
_PROMPT: Final = "Count this message before the tokenizer comes back"
_COUNTS_AFTER_FAILURE: Final = 4
_TOKENIZER_FILE: Final = "/resolve/main/tokenizer.json"
_TOKENIZER_HEADERS: Final = MappingProxyType({"X-Repo-Commit": "0" * 40, "ETag": '"scripted-char-tokenizer"'})
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
class _Usage(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
prompt_tokens: int
class _Chunk(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
id: str
usage: _Usage | None = None
def _hf_environment(directory: Path) -> dict[str, str]:
return {
"HF_HOME": str(directory / "hf-home"),
"HF_HUB_CACHE": str(directory / "hf-cache"),
"HF_HUB_OFFLINE": "0",
}
def _missing(request: Request) -> Reply:
return Reply(status=404, body=json.dumps({"error": f"unscripted target {request.target}"}).encode())
@pytest.fixture(scope="module")
def unreachable_hub_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[OwnedProxy, Wire]]:
directory: Final = tmp_path_factory.mktemp("hf-tokenizer-unreachable")
with (
gateway_from_environment() as gateway,
wire_server(_missing) as hub,
owned_proxy_process(gateway, directory, {**_hf_environment(directory), "HF_ENDPOINT": hub.url}) as owned,
):
yield owned, hub
def _llama_3(scenario: Scenario, api_base: str) -> str:
return scenario.model(model=f"databricks/{_LLAMA_3_BACKEND}", api_base=api_base, api_key=_API_KEY)
def _local_count(gateway: Gateway, model: str) -> tuple[JsonValue, JsonValue]:
response: Final = gateway.request(
"POST",
"/utils/token_counter",
{"model": model, "messages": [{"role": "user", "content": _PROMPT}]},
params={"call_endpoint": "false"},
)
assert response.status_code == 200, response.text
payload: Final = _JSON_OBJECT.validate_json(response.content)
return payload["tokenizer_type"], payload["total_tokens"]
def _frame(identity: str, delta: Mapping[str, JsonValue], finish_reason: str | None) -> bytes:
chunk: Final = {
"id": identity,
"object": "chat.completion.chunk",
"created": 1,
"model": _LLAMA_3_BACKEND,
"choices": [{"index": 0, "delta": dict(delta), "finish_reason": finish_reason}],
}
return b"data: " + json.dumps(chunk).encode() + b"\n\n"
def _stream_without_usage(identity: str) -> Callable[[Request], Reply]:
frames: Final = (
_frame(identity, {"role": "assistant", "content": "Counted "}, None),
_frame(identity, {"content": "locally."}, None),
_frame(identity, {}, "stop"),
b"data: [DONE]\n\n",
)
def respond(_request: Request) -> Reply:
return Reply(content_type="text/event-stream", chunks=frames)
return respond
def test_counts_after_a_failed_llama_3_tokenizer_download_never_reach_the_hub_again(
unreachable_hub_proxy: tuple[OwnedProxy, Wire],
) -> None:
owned, hub = unreachable_hub_proxy
with owned.gateway.scenario() as scenario:
model: Final = _llama_3(scenario, owned.gateway.upstream_url)
first: Final = _local_count(owned.gateway, model)
attempted: Final = hub.received.qsize()
later: Final = tuple(_local_count(owned.gateway, model) for _ in range(_COUNTS_AFTER_FAILURE))
retried: Final = hub.received.qsize() - attempted
assert attempted > 0, "the proxy never asked the scripted hub for the llama-3 tokenizer"
assert first[0] == "openai_tokenizer", first
assert later == (first,) * _COUNTS_AFTER_FAILURE, later
assert retried == 0, f"{retried} hub requests after the download already failed"
def test_streamed_llama_3_chat_after_a_failed_download_counts_usage_without_reaching_the_hub(
unreachable_hub_proxy: tuple[OwnedProxy, Wire],
) -> None:
owned, hub = unreachable_hub_proxy
identity: Final = f"chatcmpl-{uuid.uuid4().hex}"
with wire_server(_stream_without_usage(identity)) as upstream, owned.gateway.scenario() as scenario:
model: Final = _llama_3(scenario, upstream.url)
counted: Final = _local_count(owned.gateway, model)
attempted: Final = hub.received.qsize()
with owned.gateway.client.stream(
"POST",
"/v1/chat/completions",
json={
"model": model,
"messages": [{"role": "user", "content": _PROMPT}],
"stream": True,
"stream_options": {"include_usage": True},
},
headers={"Authorization": f"Bearer {owned.gateway.key}"},
) as response:
assert response.status_code == 200, response.read()
lines: Final = tuple(line for line in response.iter_lines() if line.startswith("data: "))
retried: Final = hub.received.qsize() - attempted
sent: Final = tuple((request.method, request.target) for request in upstream.drain())
assert counted[0] == "openai_tokenizer", counted
assert lines[-1] == "data: [DONE]", lines
chunks: Final = tuple(_Chunk.model_validate_json(line.removeprefix("data: ")) for line in lines[:-1])
assert {chunk.id for chunk in chunks} == {identity}, lines
usages: Final = tuple(chunk.usage for chunk in chunks if chunk.usage is not None)
assert tuple(usage.prompt_tokens for usage in usages) == (counted[1],), (counted, lines)
assert sent == (("POST", "/chat/completions"),), sent
assert retried == 0, f"{retried} hub requests while streaming after the download already failed"
def test_a_failed_llama_3_download_leaves_the_deepseek_tokenizer_counting(
unreachable_hub_proxy: tuple[OwnedProxy, Wire],
) -> None:
owned, _hub = unreachable_hub_proxy
with owned.gateway.scenario() as scenario:
llama: Final = _llama_3(scenario, owned.gateway.upstream_url)
deepseek: Final = scenario.model(model=_DEEPSEEK_V4, api_base=owned.gateway.upstream_url, api_key=_API_KEY)
assert _local_count(owned.gateway, llama)[0] == "openai_tokenizer"
counted: Final = _local_count(owned.gateway, deepseek)
assert counted[0] == "huggingface_tokenizer", counted
def _char_tokenizer() -> bytes:
tokenizer: Final = Tokenizer(models.WordLevel({"[UNK]": 0}, unk_token="[UNK]"))
tokenizer.pre_tokenizer = pre_tokenizers.Split(Regex("."), behavior="isolated")
return tokenizer.to_str().encode()
def _recovering_hub(recovered: threading.Event, tokenizer: bytes) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
if recovered.is_set() and request.target.endswith(_TOKENIZER_FILE):
return Reply(body=tokenizer, headers=_TOKENIZER_HEADERS)
return _missing(request)
return respond
def test_a_failed_llama_3_download_is_retried_once_the_retry_window_passes(gateway: Gateway, tmp_path: Path) -> None:
recovered: Final = threading.Event()
with (
wire_server(_recovering_hub(recovered, _char_tokenizer())) as hub,
owned_proxy_process(
gateway,
tmp_path,
{**_hf_environment(tmp_path), "HF_ENDPOINT": hub.url, "HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS": "1"},
) as owned,
owned.gateway.scenario() as scenario,
):
model: Final = _llama_3(scenario, owned.gateway.upstream_url)
failed: Final = _local_count(owned.gateway, model)
hub.drain()
recovered.set()
counted: Final = eventually(
lambda: _local_count(owned.gateway, model), lambda count: count[0] == "huggingface_tokenizer", seconds=30
)
downloads: Final = tuple(request.target for request in hub.drain() if request.target.endswith(_TOKENIZER_FILE))
assert failed[0] == "openai_tokenizer", failed
assert counted[1] != failed[1], (failed, counted)
assert downloads, "the tokenizer came back without a new download"

137
tests/integration/providers/test_nvidia_riva_wire.py Normal file
View file

@ -0,0 +1,137 @@
from collections.abc import Mapping
from typing import Final
import httpx
import openai
import pytest
from integration._support.client import Gateway, Scenario, eventually
from integration._support.openai_wire import chat_reply
from integration._support.riva import RivaServer, Rpc, Word, carries_clip, pcm_bytes, pcm_wav, riva_server
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue
_MODEL: Final = "nvidia_riva/nvidia/parakeet-0.6b-tdt"
_WORDS: Final = (Word("riva", 0, 400), Word("offline", 400, 900), Word("works", 900, 1400))
_TRANSCRIPT: Final = "riva offline works"
_CLIP_SECONDS: Final = 0.5
_REFUSED_FIELD: Final = "riva_offline is not allowed in request body"
def _riva_model(scenario: Scenario, server: RivaServer, **parameters: JsonValue) -> str:
return scenario.model(
model=_MODEL,
api_base=server.address,
model_info={"mode": "audio_transcription"},
**parameters,
)
def _transcribe(
gateway: Gateway, key: str, model: str, seconds: float = _CLIP_SECONDS, **fields: str
) -> httpx.Response:
return gateway.request_multipart(
"/v1/audio/transcriptions",
{"model": model, **fields},
{"file": ("clip.wav", pcm_wav(seconds), "audio/wav")},
key=key,
)
def _sdk(gateway: Gateway, key: str) -> openai.OpenAI:
return openai.OpenAI(base_url=f"{str(gateway.client.base_url).rstrip('/')}/v1", api_key=key, max_retries=0)
def test_a_form_riva_offline_on_transcriptions_is_refused_before_riva_is_called(gateway: Gateway) -> None:
with riva_server(_WORDS) as server, gateway.scenario() as scenario:
model: Final = _riva_model(scenario, server)
refused: Final = _transcribe(gateway, scenario.key(), model, riva_offline="true")
assert refused.status_code == 401, refused.text
assert _REFUSED_FIELD in refused.json()["error"]["message"], refused.text
assert server.drain() == ()
def test_a_json_riva_offline_on_chat_is_refused_before_the_provider_is_called(gateway: Gateway) -> None:
def respond(_: Request) -> Reply:
return chat_reply("chatcmpl-riva-audit", "gpt-4o-mini", "served", stream=False)
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(api_base=f"{wire.url}/v1")
refused: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": "riva flag on chat"}], "riva_offline": True},
key=scenario.key(),
)
assert refused.status_code == 401, refused.text
assert _REFUSED_FIELD in refused.json()["error"]["message"], refused.text
assert wire.drain() == ()
def test_an_offline_deployment_sends_the_whole_clip_in_one_recognize_call(gateway: Gateway) -> None:
with riva_server(_WORDS, offline_only=True) as server, gateway.scenario() as scenario:
model: Final = _riva_model(scenario, server, riva_offline=True)
transcription: Final = _sdk(gateway, scenario.key()).audio.transcriptions.create(
model=model,
file=("clip.wav", pcm_wav(_CLIP_SECONDS), "audio/wav"),
response_format="verbose_json",
timestamp_granularities=["word"],
)
assert transcription.text == _TRANSCRIPT
assert [(word.word, word.start, word.end) for word in transcription.words or ()] == [
(word.word, word.start_ms / 1000, word.end_ms / 1000) for word in _WORDS
]
assert [(call.rpc, carries_clip(call.audio, _CLIP_SECONDS)) for call in server.drain()] == [
(Rpc.RECOGNIZE, True)
]
def test_a_deployment_without_riva_offline_still_streams_the_clip(gateway: Gateway) -> None:
with riva_server(_WORDS) as server, gateway.scenario() as scenario:
model: Final = _riva_model(scenario, server, timeout=30)
served: Final = _transcribe(gateway, scenario.key(), model)
assert served.status_code == 200, served.text
assert served.json()["text"] == _TRANSCRIPT, served.text
assert [(call.rpc, carries_clip(call.audio, _CLIP_SECONDS)) for call in server.drain()] == [
(Rpc.STREAMING, True)
]
@pytest.mark.parametrize(
("sent", "deployment", "status", "rpcs"),
(
pytest.param("true", {}, 200, (Rpc.RECOGNIZE,), id="true-on-a-streaming-deployment"),
pytest.param("false", {"riva_offline": True}, 200, (Rpc.STREAMING,), id="false-on-an-offline-deployment"),
pytest.param("yes", {}, 400, (), id="malformed"),
pytest.param("yes", {"drop_params": True}, 200, (Rpc.STREAMING,), id="malformed-under-drop-params"),
),
)
def test_an_opted_in_caller_riva_offline_picks_the_rpc(
gateway: Gateway, sent: str, deployment: Mapping[str, JsonValue], status: int, rpcs: tuple[Rpc, ...]
) -> None:
with riva_server(_WORDS) as server, gateway.scenario() as scenario:
model: Final = _riva_model(scenario, server, configurable_clientside_auth_params=["riva_offline"], **deployment)
answered: Final = _transcribe(gateway, scenario.key(), model, riva_offline=sent)
assert answered.status_code == status, answered.text
assert tuple(call.rpc for call in server.drain()) == rpcs
def test_offline_audio_over_the_servers_message_limit_answers_413(gateway: Gateway) -> None:
seconds: Final = 2.0
limit: Final = pcm_bytes(seconds) // 2
with (
riva_server(_WORDS, offline_only=True, max_receive_bytes=limit) as server,
gateway.scenario() as scenario,
):
model: Final = _riva_model(scenario, server, riva_offline=True)
refused: Final = _transcribe(gateway, scenario.key(), model, seconds)
assert refused.status_code == 413, refused.text
assert server.drain() == ()
def test_a_stalled_offline_recognize_answers_408_at_the_deployment_timeout(gateway: Gateway) -> None:
with riva_server(_WORDS, offline_only=True, stall=True) as server, gateway.scenario() as scenario:
model: Final = _riva_model(scenario, server, riva_offline=True, timeout=2)
timed_out: Final = _transcribe(gateway, scenario.key(), model)
assert timed_out.status_code == 408, timed_out.text
calls: Final = eventually(server.drain, lambda observed: len(observed) > 0, seconds=30)
assert all(call.rpc is Rpc.RECOGNIZE and call.cancelled for call in calls), calls

122
tests/integration/providers/test_sagemaker_chat_wire.py
View file

@ -7,6 +7,7 @@ from dataclasses import dataclass
from types import MappingProxyType
from typing import Final, Literal, TypeAlias
import anthropic
import httpx
import openai
import pytest
@ -703,3 +704,124 @@ async def test_sagemaker_chat_container_outage_mid_burst_answers_every_call_once
assert statuses == {
served.call_id: "failure" if served.call.marker in down else "success" for served in answered
}, rows
_MID_STREAM_ERROR: Final = "synthetic mid-stream container failure"
@dataclass(frozen=True, slots=True)
class _CutStream:
call_id: str
text: str
errors: tuple[str, ...]
finished: bool
async def _chat_chunks_until_error(
stream: openai.AsyncStream[openai.types.chat.ChatCompletionChunk],
) -> AsyncIterator[openai.types.chat.ChatCompletionChunk | openai.APIError]:
try:
async for chunk in stream:
yield chunk
except openai.APIError as error:
yield error
async def _chat_cut(gateway: Gateway, model: str, prompt: str) -> _CutStream:
async with async_openai_client(gateway) as client:
stream: Final = await client.chat.completions.create(
model=model,
messages=[{"role": "user", "content": prompt}],
max_tokens=16,
stream=True,
extra_body=dict(NO_CACHE),
)
items: Final = tuple([item async for item in _chat_chunks_until_error(stream)])
chunks: Final = tuple(item for item in items if not isinstance(item, openai.APIError))
choices: Final = tuple(itertools.chain.from_iterable(chunk.choices for chunk in chunks))
return _CutStream(
call_id=stream.response.headers["x-litellm-call-id"],
text="".join(choice.delta.content or "" for choice in choices),
errors=tuple(item.message for item in items if isinstance(item, openai.APIError)),
finished=any(choice.finish_reason for choice in choices),
)
async def _anthropic_events_until_error(
stream: anthropic.AsyncStream[anthropic.types.RawMessageStreamEvent],
) -> AsyncIterator[anthropic.types.RawMessageStreamEvent | anthropic.APIError]:
try:
async for event in stream:
yield event
except anthropic.APIError as error:
yield error
async def _messages_cut(gateway: Gateway, model: str, prompt: str) -> _CutStream:
async with async_anthropic_client(gateway) as client:
stream: Final = await client.messages.create(
model=model,
max_tokens=16,
messages=[{"role": "user", "content": prompt}],
stream=True,
extra_body=dict(NO_CACHE),
)
items: Final = tuple([item async for item in _anthropic_events_until_error(stream)])
events: Final = tuple(item for item in items if not isinstance(item, anthropic.APIError))
deltas: Final = tuple(event.delta for event in events if event.type == "content_block_delta")
return _CutStream(
call_id=stream.response.headers["x-litellm-call-id"],
text="".join(delta.text for delta in deltas if delta.type == "text_delta"),
errors=tuple(str(item) for item in items if isinstance(item, anthropic.APIError)),
finished=any(event.type == "message_stop" for event in events),
)
def _is_error_event(event: Mapping[str, JsonValue]) -> bool:
return "error" in event or event.get("type") in ("error", "response.failed")
def _responses_cut(gateway: Gateway, model: str, prompt: str) -> _CutStream:
with gateway.client.stream(
"POST",
"/v1/responses",
json={"model": model, "input": prompt, "max_output_tokens": 16, "stream": True, **NO_CACHE},
headers={"Authorization": f"Bearer {gateway.key}"},
) as response:
lines: Final = tuple(response.iter_lines())
events: Final = tuple(
JSON_OBJECT.validate_json(line.removeprefix("data: ")) for line in lines if line.startswith("data: {")
)
return _CutStream(
call_id=response.headers["x-litellm-call-id"],
text="".join(
string_value(event["delta"]) for event in events if event.get("type") == "response.output_text.delta"
),
errors=tuple(json.dumps(event) for event in events if _is_error_event(event)),
finished=any(event.get("type") == "response.completed" for event in events),
)
_CUT_DRIVERS: Final[Mapping[Endpoint, Callable[[Gateway, str, str], _CutStream]]] = MappingProxyType(
{
"chat": lambda gateway, model, prompt: asyncio.run(_chat_cut(gateway, model, prompt)),
"messages": lambda gateway, model, prompt: asyncio.run(_messages_cut(gateway, model, prompt)),
"responses": _responses_cut,
}
)
@pytest.mark.parametrize("endpoint", _ENDPOINTS)
def test_sagemaker_chat_stream_failure_after_the_first_chunks_reaches_every_endpoint_as_an_error_and_a_failed_spend_row(
gateway: Gateway, endpoint: Endpoint
) -> None:
events: Final = vllm_stream_events(f"sagemaker-chat-{uuid.uuid4().hex}")
reply: Final = event_stream_reply((*events[:2], _sse({"error": {"message": _MID_STREAM_ERROR, "code": 500}})))
with wire_server(lambda _request: reply) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire, {})
cut: Final = _CUT_DRIVERS[endpoint](gateway, model, marked_prompt())
only_request(wire)
assert (cut.text, cut.finished) == ("".join(CONTENT), False), cut
assert any(_MID_STREAM_ERROR in error for error in cut.errors), cut
assert _spend_row(cut.call_id)["status"] == "failure", cut

188
tests/integration/providers/test_vertex_anthropic_structured_output_wire.py
View file

@ -60,6 +60,8 @@ _OWNED_MODEL: Final = "partner-claude"
_DEPRECATED_OUTPUT_FORMAT: Final = "output_format: this field is deprecated, use output_config.format"
_BAD_OUTPUT_CONFIG: Final = "output_config: Input should be a valid dictionary"
_BAD_FORMAT: Final = "output_config.format: Input should be a json_schema object"
_BAD_OUTPUT_FORMAT: Final = "output_format: Input should be a valid dictionary"
_MISTYPED_OUTPUT_FORMAT: Final = "json_schema"
_SCHEMA: Final[dict[str, JsonValue]] = {
"type": "object",
"properties": {"tier": {"type": "integer"}, "summary": {"type": "string"}},
@ -94,6 +96,11 @@ _VALUES_RESPONSE_FORMAT: Final[dict[str, JsonValue]] = {
"type": "json_schema",
"json_schema": {"name": "valued", "schema": _VALUES_SCHEMA},
}
_REQUIRED_ONLY_SCHEMA: Final[dict[str, JsonValue]] = {
"type": "object",
"properties": {"tier": {"type": "integer"}},
"required": ["payload"],
}
_ADDITIONAL_OBJECT_SCHEMA: Final[dict[str, JsonValue]] = {
**_OPTIONAL_SCHEMA,
"additionalProperties": {"type": "object"},
@ -146,6 +153,7 @@ _STREAM_SHAPES: Final = (
"number_first",
"bool_first",
"null_first",
"named_wrapper",
)
_HELD_STREAM_SHAPES: Final = frozenset(
{
@ -166,12 +174,15 @@ _HELD_STREAM_SHAPES: Final = frozenset(
"named_encoded",
"undeclared_encoded",
"encoded_deep",
"named_wrapper",
}
)
_AS_WRITTEN_SHAPES: Final = ("label_non_ascii", "label_two_objects_non_ascii")
_TEXT_ANSWER_SHAPES: Final = frozenset({"undeclared_wrapper"})
_LABEL: Final = "RiskAssessment"
_NON_ASCII_LABEL: Final = "Évaluation des risques 风险"
_NON_ASCII_SUMMARY: Final = "Hôtel 東京"
_LONE_SURROGATE: Final = "\ud800"
_ENCODED_DEPTH: Final = 200_000
_FAMILY: Final = (
"anthropic_unmapped",
@ -181,6 +192,9 @@ _FAMILY: Final = (
"claude_platform_unmapped",
"vertex_haiku",
)
_CUT_FAMILY: Final = ("anthropic_unmapped", "bedrock_invoke")
_CUT_EVENTS: Final = 3
_OUTPUT_FORMAT_ROUTES: Final = ("anthropic_chat", "anthropic_messages", "vertex_chat", "vertex_messages")
_FRAGMENT_SIZE: Final = 20
_TIMING_BYTES: Final = 512 * 1024
_TIMING_FRAGMENT: Final = 4096
@ -228,9 +242,9 @@ def _nest(payload: JsonValue, depth: int) -> JsonValue:
def _tool_input(shape: str, marker_hex: str) -> JsonValue:
payload: Final = _payload(marker_hex)
match shape:
case "plain":
case "plain" | "cut_plain":
return payload
case "wrapped":
case "wrapped" | "cut_wrapped":
return {"tool_name": _JSON_TOOL, "parameters": payload}
case "named":
return {"name": _JSON_TOOL, "parameters": payload}
@ -256,6 +270,14 @@ def _tool_input(shape: str, marker_hex: str) -> JsonValue:
return {"tool_name": "other_tool", "parameters": payload}
case "named_input":
return {"name": _JSON_TOOL, "input": payload}
case "named_wrapper":
return {"name": _JSON_TOOL, "arguments": {"input": payload}}
case "required_only":
return {"payload": payload}
case "wrapped_non_ascii":
return {"tool_name": _JSON_TOOL, "parameters": _payload(f"{_NON_ASCII_SUMMARY} {marker_hex}")}
case "wrapped_lone_surrogate":
return {"tool_name": _JSON_TOOL, "parameters": _payload(marker_hex + _LONE_SURROGATE)}
case "named_two_objects":
return {"tool_name": _JSON_TOOL, "parameters": payload, "extra": {"note": marker_hex}}
case "named_string":
@ -348,8 +370,13 @@ def _expected(shape: str, marker_hex: str) -> JsonValue:
| "parameter_label"
| "function_label"
| "named_encoded"
| "named_wrapper"
| "cut_plain"
| "cut_wrapped"
):
return _payload(marker_hex)
case "wrapped_non_ascii" | "wrapped_lone_surrogate":
return object_value(_tool_input(shape, marker_hex))["parameters"]
case "large":
return {**_payload(marker_hex), "summary": marker_hex + "x" * 5000}
case "named_partial" | "inner_incomplete":
@ -423,9 +450,14 @@ def _aws_chunk(data: Mapping[str, JsonValue]) -> bytes:
def _stream_reply(
message: dict[str, JsonValue], block: dict[str, JsonValue], deltas: tuple[dict[str, JsonValue], ...], framing: str
message: dict[str, JsonValue],
block: dict[str, JsonValue],
deltas: tuple[dict[str, JsonValue], ...],
framing: str,
cut: bool = False,
) -> Reply:
events: Final = _events(message, block, deltas)
complete: Final = _events(message, block, deltas)
events: Final = complete[:-_CUT_EVENTS] if cut else complete
if framing == "aws":
return Reply(content_type=_EVENT_STREAM, chunks=tuple(_aws_chunk(data) for _, data in events))
return Reply(
@ -445,6 +477,8 @@ def _error(status: int, message: str) -> Reply:
def _format_fault(body: Mapping[str, JsonValue], vertex: bool) -> str | None:
if vertex and "output_format" in body:
return _DEPRECATED_OUTPUT_FORMAT
if body.get("output_format") is not None and not isinstance(body["output_format"], dict):
return _BAD_OUTPUT_FORMAT
output_config: Final = body.get("output_config")
if output_config is None:
return None
@ -486,7 +520,7 @@ def _tool_reply(message_id: str, shape: str, marker_hex: str, user_tool: bool, f
deltas: Final = tuple(
{"type": "input_json_delta", "partial_json": piece} for piece in _fragments(partial, _fragment_size(shape))
)
return _stream_reply(message, {**json_block, "input": {}}, deltas, framing)
return _stream_reply(message, {**json_block, "input": {}}, deltas, framing, cut=shape.startswith("cut_"))
def _text_reply(message_id: str, text: str, framing: str | None) -> Reply:
@ -748,7 +782,9 @@ def _messages_stream_text(response: httpx.Response) -> JsonValue:
return json.loads("".join(string_value(delta["text"]) for delta in deltas if delta["type"] == "text_delta"))
def _vertex_deployment(gateway: Gateway, scenario: Scenario, api_base: str, backend: str) -> str:
def _vertex_deployment(
gateway: Gateway, scenario: Scenario, api_base: str, backend: str, **litellm_params: JsonValue
) -> str:
return scenario.model(
model=f"vertex_ai/{backend}",
api_base=api_base,
@ -756,6 +792,7 @@ def _vertex_deployment(gateway: Gateway, scenario: Scenario, api_base: str, back
vertex_project=_PROJECT,
vertex_location=_LOCATION,
vertex_credentials=service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")),
**litellm_params,
)
@ -814,6 +851,33 @@ def _family_deployment(gateway: Gateway, scenario: Scenario, api_base: str, memb
raise AssertionError(member)
def _output_format_deployment(
gateway: Gateway, scenario: Scenario, api_base: str, route: str, **litellm_params: JsonValue
) -> str:
if route.startswith("vertex_"):
return _vertex_deployment(gateway, scenario, api_base, _NATIVE_MODEL, **litellm_params)
return scenario.model(
model=f"anthropic/{_NATIVE_MODEL}", api_base=api_base, api_key=_ANTHROPIC_KEY, **litellm_params
)
def _output_format_call(
gateway: Gateway, model: str, route: str, marker: str, output_format: JsonValue
) -> httpx.Response:
if route.endswith("_chat"):
return gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": marker}], "output_format": output_format},
)
return gateway.request(
"POST",
"/v1/messages",
_messages_body(model, marker, output_format=output_format),
headers={"anthropic-beta": _STRUCTURED_OUTPUTS_BETA},
)
def _owned_config(path: Path, gateway: Gateway, api_base: str) -> Path:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
path.write_text(
@ -1188,13 +1252,18 @@ def test_vertex_chat_json_schema_without_required_keys_unwraps_only_the_tool_nam
@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"])
@pytest.mark.parametrize(
("shape", "schema"),
[("values", _VALUES_SCHEMA), ("required_only", _REQUIRED_ONLY_SCHEMA)],
ids=["declared_property", "required_only"],
)
def test_vertex_chat_json_schema_that_declares_the_wrapper_key_keeps_the_answer_whole(
gateway: Gateway, stream: bool
gateway: Gateway, shape: str, schema: dict[str, JsonValue], stream: bool
) -> None:
marker, marker_hex = _marker("values")
marker, marker_hex = _marker(shape)
response_format: Final[dict[str, JsonValue]] = {
"type": "json_schema",
"json_schema": {"name": "valued", "schema": _VALUES_SCHEMA},
"json_schema": {"name": "valued", "schema": schema},
}
with wire_server(_peer) as wire, gateway.scenario() as scenario:
model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL)
@ -1202,9 +1271,9 @@ def test_vertex_chat_json_schema_that_declares_the_wrapper_key_keeps_the_answer_
"POST", "/v1/chat/completions", _chat_body(model, marker, stream=stream, response_format=response_format)
)
content: Final = _stream_content(response) if stream else _chat_content(response)
assert content == _tool_input("values", marker_hex), response.text
assert content == _tool_input(shape, marker_hex), response.text
if stream:
assert len(_content_frames(response)) >= _fragment_count("values", marker_hex), response.text
assert len(_content_frames(response)) >= _fragment_count(shape, marker_hex), response.text
_one_delivery(wire, marker_hex)
@ -1265,6 +1334,22 @@ def test_vertex_chat_json_schema_unwraps_a_tool_named_envelope_only_around_one_o
_one_delivery(wire, marker_hex)
@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"])
@pytest.mark.parametrize("shape", ["wrapped_non_ascii", "wrapped_lone_surrogate"])
def test_vertex_chat_json_schema_answer_keeps_non_ascii_text_and_escapes_a_lone_surrogate(
gateway: Gateway, shape: str, stream: bool
) -> None:
marker, marker_hex = _marker(shape)
with wire_server(_peer) as wire, gateway.scenario() as scenario:
model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL)
response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker, stream=stream))
text: Final = _stream_text(response) if stream else string_value(_chat_message(response)["content"])
assert json.loads(text) == _expected(shape, marker_hex), response.text
if shape == "wrapped_non_ascii":
assert _NON_ASCII_SUMMARY in text, text
_one_delivery(wire, marker_hex)
@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"])
@pytest.mark.parametrize("shape", ["name_property", "name_value_plain"])
def test_vertex_chat_json_schema_with_a_name_property_unwraps_only_the_tool_named_wrapper(
@ -1350,6 +1435,19 @@ def test_tool_path_family_stream_unwraps_the_envelope(gateway: Gateway, member:
assert _json_tool_schema(sent)["properties"] == _SCHEMA["properties"], sent.body
@pytest.mark.parametrize("shape", ["plain", "cut_plain", "cut_wrapped"])
@pytest.mark.parametrize("member", _CUT_FAMILY)
def test_tool_path_stream_cut_before_the_tool_block_closes_still_delivers_the_answer(
gateway: Gateway, member: str, shape: str
) -> None:
marker, marker_hex = _marker(shape)
with wire_server(_peer) as wire, gateway.scenario() as scenario:
model: Final = _family_deployment(gateway, scenario, wire.url, member)
response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker, stream=True))
assert _stream_content(response) == _payload(marker_hex), response.text
_one_delivery(wire, marker_hex)
@pytest.mark.parametrize("accepts_forced_tools", [True, False], ids=["accepts_forced_tools", "rejects_forced_tools"])
def test_vertex_chat_json_schema_forces_the_json_tool_only_when_the_model_accepts_forced_tools(
gateway: Gateway, accepts_forced_tools: bool
@ -1865,6 +1963,74 @@ def test_vertex_messages_duplicate_output_format_keys_fold_the_last_value(gatewa
assert _folded_format(_one_delivery(wire, marker_hex)) == _OUTPUT_FORMAT
@pytest.mark.parametrize("route", _OUTPUT_FORMAT_ROUTES)
def test_string_output_format_under_deployment_drop_params_is_dropped_before_the_peer(
gateway: Gateway, route: str
) -> None:
marker, marker_hex = _marker("plain_text")
with wire_server(_peer) as wire, gateway.scenario() as scenario:
model: Final = _output_format_deployment(gateway, scenario, wire.url, route, drop_params=True)
response: Final = _output_format_call(gateway, model, route, marker, _MISTYPED_OUTPUT_FORMAT)
assert response.status_code == 200, response.text
assert f"plain text {marker_hex}" in response.text, response.text
sent: Final = _one_delivery(wire, marker_hex)
assert not {"output_format", "output_config"} & sent.body.keys(), sent.body
@pytest.mark.parametrize("route", ["anthropic_chat", "anthropic_messages"])
def test_object_output_format_under_deployment_drop_params_still_reaches_anthropic(
gateway: Gateway, route: str
) -> None:
marker, marker_hex = _marker("native")
with wire_server(_peer) as wire, gateway.scenario() as scenario:
model: Final = _output_format_deployment(gateway, scenario, wire.url, route, drop_params=True)
response: Final = _output_format_call(gateway, model, route, marker, _OUTPUT_FORMAT)
assert response.status_code == 200, response.text
assert _one_delivery(wire, marker_hex).body["output_format"] == _OUTPUT_FORMAT
def test_vertex_messages_string_output_format_under_drop_params_keeps_the_callers_output_config_format(
gateway: Gateway,
) -> None:
marker, marker_hex = _marker("native")
with wire_server(_peer) as wire, gateway.scenario() as scenario:
model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL, drop_params=True)
response: Final = gateway.request(
"POST",
"/v1/messages",
_messages_body(
model, marker, output_format=_MISTYPED_OUTPUT_FORMAT, output_config={"format": _OUTPUT_FORMAT}
),
)
assert _messages_text(response) == _payload(marker_hex), response.text
assert _folded_format(_one_delivery(wire, marker_hex)) == _OUTPUT_FORMAT
@pytest.mark.parametrize("route", ["anthropic_chat", "anthropic_messages"])
def test_string_output_format_without_drop_params_reaches_anthropic_and_returns_its_400(
gateway: Gateway, route: str
) -> None:
marker, marker_hex = _marker("plain_text")
with wire_server(_peer) as wire, gateway.scenario() as scenario:
model: Final = _output_format_deployment(gateway, scenario, wire.url, route)
response: Final = _output_format_call(gateway, model, route, marker, _MISTYPED_OUTPUT_FORMAT)
assert response.status_code == 400, response.text
assert _BAD_OUTPUT_FORMAT in response.text, response.text
assert _one_delivery(wire, marker_hex).body["output_format"] == _MISTYPED_OUTPUT_FORMAT
def test_vertex_chat_string_output_format_without_drop_params_is_refused_without_reaching_the_peer(
gateway: Gateway,
) -> None:
marker, _ = _marker("plain_text")
with wire_server(_peer) as wire, gateway.scenario() as scenario:
model: Final = _output_format_deployment(gateway, scenario, wire.url, "vertex_chat")
response: Final = _output_format_call(gateway, model, "vertex_chat", marker, _MISTYPED_OUTPUT_FORMAT)
assert response.status_code == 400, response.text
assert "needs `output_format` to be a JSON object" in response.text, response.text
assert _delivered(wire) == ()
@pytest.mark.parametrize(
"schema",
[

426
tests/integration/routing/test_batch_mode_deployments_wire.py Normal file
View file

@ -0,0 +1,426 @@
import json
import uuid
from collections.abc import Iterator, Mapping, Sequence
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from typing import Final, Literal
from urllib.parse import urlsplit
import pytest
from integration._support.client import (
DEPLOYMENT_IDS,
Gateway,
Scenario,
eventually,
gateway_from_environment,
list_value,
object_value,
string_value,
)
from integration._support.openai_wire import chat_reply
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._cache_control_marks_support import owned_config
from pydantic import JsonValue, TypeAdapter
pytestmark: Final = pytest.mark.shares_worker_state
_AZURE_API_VERSION: Final = "2024-10-21"
_BACKEND: Final = "gpt-5.4-mini"
_BATCH_ENDPOINT: Final = "/v1/batch"
_CHAT_ENDPOINT: Final = "/v1/chat/completions"
_RESPONSES_ENDPOINT: Final = "/v1/responses"
_TEAM_GROUP: Final = "batch-team-mix"
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_JSON_LIST: Final = TypeAdapter(list[JsonValue])
_ENDPOINTS: Final = TypeAdapter(list[str] | None)
_OWNED_PROXY_TIMEOUT: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
Provider = Literal["azure", "openai"]
Mode = Literal["chat", "batch"]
@dataclass(frozen=True, slots=True)
class _Sku:
provider: Provider
mode: Mode
key: str
weight: int | None
endpoints: JsonValue
team_id: str | None
@dataclass(frozen=True, slots=True)
class _Create:
status: int
body: str
@dataclass(frozen=True, slots=True)
class _TeamRig:
owned: OwnedProxy
wire: Wire
team_key: str
other_team_key: str
no_team_key: str
chat: _Sku
batch: _Sku
def _sku(
provider: Provider,
mode: Mode,
*,
weight: int | None = None,
endpoints: JsonValue = None,
team_id: str | None = None,
) -> _Sku:
return _Sku(provider, mode, f"{mode}-sku-{uuid.uuid4().hex}", weight, endpoints, team_id)
def _serves_batches(key: str) -> bool:
return key.startswith("batch-sku-")
def _deployment(name: str, wire: Wire, sku: _Sku) -> dict[str, JsonValue]:
endpoint: Final[dict[str, JsonValue]] = (
{"model": f"azure/{_BACKEND}", "api_base": wire.url, "api_version": _AZURE_API_VERSION}
if sku.provider == "azure"
else {"model": f"openai/{_BACKEND}", "api_base": wire.url + "/v1"}
)
weight: Final[dict[str, JsonValue]] = {} if sku.weight is None else {"weight": sku.weight}
endpoints: Final[dict[str, JsonValue]] = {} if sku.endpoints is None else {"supported_endpoints": sku.endpoints}
team: Final[dict[str, JsonValue]] = {} if sku.team_id is None else {"team_id": sku.team_id}
return {
"model_name": name,
"litellm_params": {**endpoint, "api_key": sku.key, **weight},
"model_info": {"mode": sku.mode, **endpoints, **team},
}
def _json_reply(body: Mapping[str, JsonValue], status: int = 200) -> Reply:
return Reply(status=status, body=json.dumps(body).encode())
def _provider_key(request: Request) -> str:
return request.headers.get("api-key") or request.headers.get("authorization", "").removeprefix("Bearer ")
def _is_azure(request: Request) -> bool:
return urlsplit(request.target).path.startswith("/openai/")
def _posted_to(request: Request, suffix: str) -> bool:
return request.method == "POST" and urlsplit(request.target).path.endswith(suffix)
def _vendor(request: Request) -> Reply:
key: Final = _provider_key(request)
if _posted_to(request, "/files"):
return _json_reply(
{
"id": f"file-{key}",
"object": "file",
"bytes": len(request.body),
"created_at": 1700000000,
"filename": "in.jsonl",
"purpose": "batch",
"status": "processed",
}
)
if _posted_to(request, "/batches") and _is_azure(request) and not _serves_batches(key):
return _json_reply(
{"error": {"code": "invalid_deployment_type", "message": "The deployment does not support batch jobs."}},
400,
)
if _posted_to(request, "/batches"):
return _json_reply(
{
"id": f"batch-{uuid.uuid4().hex}",
"object": "batch",
"endpoint": _CHAT_ENDPOINT,
"input_file_id": f"file-{key}",
"status": "validating",
"created_at": 1700000000,
"completion_window": "24h",
}
)
if _posted_to(request, "/chat/completions") and _is_azure(request) and _serves_batches(key):
return _json_reply(
{"error": {"code": "OperationNotSupported", "message": "Batch deployments do not serve chat."}}, 400
)
if _posted_to(request, "/chat/completions"):
return chat_reply(f"chatcmpl-{uuid.uuid4().hex}", _BACKEND, "served", stream=False)
return _json_reply({"data": [], "object": "list"})
def _deployment_ids_on_a_worker(gateway: Gateway, name: str) -> tuple[str, frozenset[str]]:
response: Final = gateway.request(
"GET",
"/health/integration-deployment-ids",
params={"model_name": name},
headers={"connection": "close"},
)
assert response.status_code == 200, response.text
return response.headers["x-integration-worker-pid"], frozenset(
DEPLOYMENT_IDS.validate_json(response.content)["model_ids"]
)
def _wait_until_every_worker_serves(gateway: Gateway, name: str, identities: frozenset[str]) -> None:
from litellm.constants import PROXY_CONFIG_RELOAD_INTERVAL_SECONDS
def answers() -> tuple[tuple[str, frozenset[str]], ...]:
with ThreadPoolExecutor(max_workers=8) as pool:
return tuple(pool.map(lambda _: _deployment_ids_on_a_worker(gateway, name), range(16)))
eventually(
answers,
lambda seen: (
len({pid for pid, _ in seen}) >= gateway.workers and all(identities <= served for _, served in seen)
),
seconds=2 * PROXY_CONFIG_RELOAD_INTERVAL_SECONDS,
)
def _add_deployment(scenario: Scenario, deployment: Mapping[str, JsonValue]) -> str:
created: Final = scenario.gateway.post("/model/new", deployment)
identity: Final = string_value(object_value(created["model_info"])["id"])
scenario.cleanups.callback(scenario.delete_model, identity)
return identity
def _group(scenario: Scenario, wire: Wire, skus: Sequence[_Sku]) -> str:
name: Final = f"integration-batch-{uuid.uuid4().hex}"
identities: Final = frozenset(_add_deployment(scenario, _deployment(name, wire, sku)) for sku in skus)
_wait_until_every_worker_serves(scenario.gateway, name, identities)
return name
def _managed_upload(gateway: Gateway, group: str, key: str) -> str:
line: Final = {"custom_id": "r1", "method": "POST", "url": _CHAT_ENDPOINT, "body": {"model": group}}
uploaded: Final = gateway.request_multipart(
"/v1/files",
{"purpose": "batch", "target_model_names": group},
{"file": ("in.jsonl", json.dumps(line).encode() + b"\n", "application/jsonl")},
key=key,
)
assert uploaded.status_code == 200, uploaded.text
return string_value(_JSON_OBJECT.validate_json(uploaded.content)["id"])
def _create(gateway: Gateway, managed_file: str, key: str) -> _Create:
response: Final = gateway.request(
"POST",
"/v1/batches",
{"input_file_id": managed_file, "endpoint": _CHAT_ENDPOINT, "completion_window": "24h"},
key=key,
)
return _Create(response.status_code, response.text)
def _chat(gateway: Gateway, group: str) -> int:
return gateway.request(
"POST",
"/v1/chat/completions",
{"model": group, "messages": [{"role": "user", "content": "hello"}], "cache": {"no-cache": True}},
).status_code
def _keys_posted_to(requests: Sequence[Request], suffix: str) -> tuple[str, ...]:
return tuple(_provider_key(request) for request in requests if _posted_to(request, suffix))
def _assert_creates_reach_only(wire: Wire, creates: Sequence[_Create], sku: _Sku) -> None:
assert [create.status for create in creates] == [200] * len(creates), creates
assert _keys_posted_to(wire.drain(), "/batches") == (sku.key,) * len(creates)
@pytest.mark.parametrize(
("provider", "creates_reach"),
[("azure", "batch"), ("openai", "chat")],
ids=["azure-skips-chat", "openai-keeps-chat"],
)
def test_a_batch_create_on_a_mixed_group_reaches_the_deployment_its_provider_allows(
gateway: Gateway, provider: Provider, creates_reach: Mode
) -> None:
chat: Final = _sku(provider, "chat", weight=1)
batch: Final = _sku(provider, "batch", weight=0)
with wire_server(_vendor) as wire, gateway.scenario() as scenario:
group: Final = _group(scenario, wire, (chat, batch))
managed_file: Final = _managed_upload(gateway, group, gateway.key)
assert set(_keys_posted_to(wire.drain(), "/files")) == {chat.key, batch.key}
creates: Final = tuple(_create(gateway, managed_file, gateway.key) for _ in range(3))
_assert_creates_reach_only(wire, creates, batch if creates_reach == "batch" else chat)
def test_concurrent_batch_creates_and_chats_each_reach_their_own_deployment(gateway: Gateway) -> None:
chat: Final = _sku("azure", "chat", weight=1)
batch: Final = _sku("azure", "batch", weight=0)
with wire_server(_vendor) as wire, gateway.scenario() as scenario:
group: Final = _group(scenario, wire, (chat, batch))
managed_file: Final = _managed_upload(gateway, group, gateway.key)
wire.drain()
with ThreadPoolExecutor(max_workers=16) as pool:
burst: Final = tuple(
pool.map(
lambda turn: _create(gateway, managed_file, gateway.key) if turn % 2 else _chat(gateway, group),
range(24),
)
)
after: Final = tuple(_chat(gateway, group) for _ in range(6))
received: Final = wire.drain()
chat_statuses: Final = tuple(result for result in burst if isinstance(result, int))
creates: Final = tuple(result for result in burst if isinstance(result, _Create))
assert chat_statuses + after == (200,) * 18, (burst, after)
assert [create.status for create in creates] == [200] * 12, creates
assert _keys_posted_to(received, "/chat/completions") == (chat.key,) * 18
assert _keys_posted_to(received, "/batches") == (batch.key,) * 12
def _group_endpoints(gateway: Gateway, group: str, key: str | None = None) -> list[str] | None:
response: Final = gateway.request("GET", "/model_group/info", params={"model_group": group}, key=key)
assert response.status_code == 200, response.text
(entry,) = list_value(_JSON_OBJECT.validate_json(response.content)["data"])
return _ENDPOINTS.validate_python(object_value(entry).get("supported_endpoints"))
@pytest.mark.parametrize(
("shape", "expected"),
[
pytest.param(
(
("azure", "chat", [_CHAT_ENDPOINT, _RESPONSES_ENDPOINT, _BATCH_ENDPOINT]),
("azure", "batch", [_CHAT_ENDPOINT]),
),
[_BATCH_ENDPOINT, _CHAT_ENDPOINT, _RESPONSES_ENDPOINT],
id="azure-chat-and-batch",
),
pytest.param(
(("azure", "chat", [_CHAT_ENDPOINT, _BATCH_ENDPOINT]),),
[_CHAT_ENDPOINT],
id="azure-chat-only",
),
pytest.param(
(("azure", "batch", [_CHAT_ENDPOINT, _RESPONSES_ENDPOINT]),),
[_BATCH_ENDPOINT],
id="azure-batch-only",
),
pytest.param(
(
("openai", "chat", [_CHAT_ENDPOINT, _BATCH_ENDPOINT]),
("openai", "batch", [_RESPONSES_ENDPOINT]),
),
[_BATCH_ENDPOINT, _CHAT_ENDPOINT, _RESPONSES_ENDPOINT],
id="openai-chat-and-batch",
),
],
)
def test_model_group_info_lists_the_endpoints_its_deployments_serve(
gateway: Gateway, shape: tuple[tuple[Provider, Mode, list[str]], ...], expected: list[str]
) -> None:
skus: Final = tuple(_sku(provider, mode, endpoints=endpoints) for provider, mode, endpoints in shape)
with wire_server(_vendor) as wire, gateway.scenario() as scenario:
group: Final = _group(scenario, wire, skus)
listed: Final = _group_endpoints(gateway, group)
assert listed is not None and sorted(listed) == expected, listed
@pytest.mark.parametrize("malformed", ["/v1/batch", [7], {"batch": True}], ids=["string", "number-list", "object"])
def test_a_malformed_endpoint_list_keeps_group_info_and_chat_answering(gateway: Gateway, malformed: JsonValue) -> None:
with wire_server(_vendor) as wire, gateway.scenario() as scenario:
group: Final = _group(scenario, wire, (_sku("openai", "chat", endpoints=malformed),))
_group_endpoints(gateway, group)
assert _chat(gateway, group) == 200
@pytest.fixture(scope="module")
def team_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_TeamRig]:
directory: Final = tmp_path_factory.mktemp("batch-team")
team_id: Final = f"integration-batch-team-{uuid.uuid4().hex}"
chat: Final = _sku("azure", "chat", endpoints=[_CHAT_ENDPOINT])
batch: Final = _sku("azure", "batch", team_id=team_id)
with gateway_from_environment() as gateway, wire_server(_vendor) as wire:
config: Final = owned_config(
directory,
[_deployment(_TEAM_GROUP, wire, chat), _deployment(_TEAM_GROUP, wire, batch)],
litellm_settings={"public_model_groups": [_TEAM_GROUP]},
router_settings={"routing_strategy": "usage-based-routing"},
)
with owned_proxy_process(gateway, directory, {}, config=config) as owned, owned.gateway.scenario() as scenario:
scenario.team(team_id=team_id)
yield _TeamRig(
owned=owned,
wire=wire,
team_key=scenario.key(team_id=team_id),
other_team_key=scenario.key(team_id=scenario.team()),
no_team_key=scenario.key(),
chat=chat,
batch=batch,
)
def _caller_key(rig: _TeamRig, caller: str) -> str:
return {"team": rig.team_key, "other-team": rig.other_team_key, "no-team": rig.no_team_key}[caller]
@_OWNED_PROXY_TIMEOUT
@pytest.mark.parametrize(
("caller", "lists_batch"),
[("team", True), ("other-team", False), ("no-team", False)],
ids=["team", "other-team", "no-team"],
)
def test_model_group_info_lists_a_team_owned_batch_endpoint_only_to_that_team(
team_rig: _TeamRig, caller: str, lists_batch: bool
) -> None:
listed: Final = _group_endpoints(team_rig.owned.gateway, _TEAM_GROUP, _caller_key(team_rig, caller))
assert listed is not None and _CHAT_ENDPOINT in listed, listed
assert (_BATCH_ENDPOINT in listed) is lists_batch, listed
@_OWNED_PROXY_TIMEOUT
def test_public_model_hub_never_lists_a_team_owned_batch_endpoint(team_rig: _TeamRig) -> None:
response: Final = team_rig.owned.gateway.client.get("/public/model_hub")
assert response.status_code == 200, response.text
(entry,) = [
object_value(row)
for row in _JSON_LIST.validate_json(response.content)
if object_value(row)["model_group"] == _TEAM_GROUP
]
assert _BATCH_ENDPOINT not in (_ENDPOINTS.validate_python(entry.get("supported_endpoints")) or []), entry
@_OWNED_PROXY_TIMEOUT
def test_a_legacy_strategy_sends_a_team_batch_create_to_its_batch_deployment(team_rig: _TeamRig) -> None:
gateway: Final = team_rig.owned.gateway
team_rig.wire.drain()
managed_file: Final = _managed_upload(gateway, _TEAM_GROUP, team_rig.team_key)
assert team_rig.batch.key in _keys_posted_to(team_rig.wire.drain(), "/files")
creates: Final = tuple(_create(gateway, managed_file, team_rig.team_key) for _ in range(3))
_assert_creates_reach_only(team_rig.wire, creates, team_rig.batch)
@_OWNED_PROXY_TIMEOUT
@pytest.mark.parametrize("caller", ["other-team", "no-team"])
def test_a_legacy_strategy_keeps_other_callers_off_a_team_batch_deployment(team_rig: _TeamRig, caller: str) -> None:
gateway: Final = team_rig.owned.gateway
key: Final = _caller_key(team_rig, caller)
team_rig.wire.drain()
managed_file: Final = _managed_upload(gateway, _TEAM_GROUP, key)
team_rig.wire.drain()
tuple(_create(gateway, managed_file, key) for _ in range(3))
assert _keys_posted_to(team_rig.wire.drain(), "/batches") == (team_rig.chat.key,) * 3

61
tests/integration/routing/test_priority_scheduler_queue_cleanup.py
View file

@ -512,12 +512,55 @@ def test_other_prioritized_endpoints_are_served_once_and_billed(
marker: Final = new_marker()
answer: Final = call(gateway, endpoint, model, marker)
identity: Final = assert_served(answer, marker)
if endpoint == "queue":
if endpoint in ("queue", "queue-stream"):
assert answer.headers.get("x-litellm-priority") == "1", answer.headers
assert_once(received_markers(wire), (marker,))
spend_row_lands(identity)
@pytest.mark.parametrize("field", ({}, {"priority": "high", "drop_params": True}), ids=("missing", "dropped"))
def test_queue_route_refuses_a_request_it_cannot_queue_before_the_upstream(
gateway: Gateway, rig_model: tuple[str, Wire], field: Mapping[str, JsonValue]
) -> None:
model, wire = rig_model
marker: Final = new_marker()
answer: Final = post(gateway, "/queue/chat/completions", chat_body(model, marker, **field))
assert answer.status == 400, (answer.status, answer.text)
error: Final = JSON_OBJECT.validate_json(answer.text)["error"]
assert isinstance(error, dict), answer.text
assert (error["type"], error["param"]) == ("invalid_request_error", "priority"), answer.text
assert_never(received_markers(wire), (marker,))
def anthropic_error(status: int) -> Reply:
body: Final = {"type": "error", "error": {"type": "invalid_request_error", "message": f"scripted {status}"}}
return Reply(status=status, body=json.dumps(body).encode())
@pytest.mark.parametrize(
("status", "provider_model", "base_path", "respond"),
(
(400, "openai/gpt-4o-mini", "/v1", answering_model_discovery(lambda _: openai_error(400))),
(402, "anthropic/claude-sonnet-5-5", "", lambda _: anthropic_error(402)),
),
ids=("openai-400", "anthropic-402"),
)
def test_queue_route_answers_a_provider_bad_request_with_its_own_status(
gateway: Gateway, status: int, provider_model: str, base_path: str, respond: Callable[[Request], Reply]
) -> None:
marker: Final = new_marker()
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=provider_model, api_base=wire.url + base_path)
answer: Final = post(gateway, "/queue/chat/completions", chat_body(model, marker, priority=1))
assert marker in received_markers(wire)
assert answer.status == status, (answer.status, answer.text)
error: Final = JSON_OBJECT.validate_json(answer.text)["error"]
assert isinstance(error, dict), answer.text
assert error["type"] == "invalid_request_error", answer.text
assert f"scripted {status}" in str(error["message"]), answer.text
assert "Authentication Error" not in str(error["message"]), answer.text
def upstream_body_for(wire: Wire, marker: str) -> dict[str, JsonValue]:
(upstream_request,) = tuple(request for request in wire.drain() if marker.encode() in request.body)
return JSON_OBJECT.validate_json(upstream_request.body)
@ -751,6 +794,19 @@ def test_waiter_proceeds_once_the_entry_ahead_is_cleared(pair: Pair) -> None:
assert_drained(pair.cache, group)
@pytest.mark.timeout(PAIR_CELL_TIMEOUT)
def test_request_pinned_to_a_deployment_id_is_served_ahead_of_an_earlier_queue_entry(pair: Pair) -> None:
deployment: Final = f"{PAIR_GROUPS[0]}-deployment"
write_queue(pair.cache, deployment, [GHOST_ENTRY])
marker: Final = new_marker()
answer: Final = post(pair.second, "/v1/chat/completions", chat_body(deployment, marker, priority=2))
assert_served(answer, marker)
assert_prioritized(answer)
assert answer.seconds < PROMPT_SECONDS, answer.seconds
assert queue_entries(pair.cache, deployment) == [GHOST_ENTRY]
assert_once(received_markers(pair.wire), (marker,))
def pinned_connection(gateway: Gateway) -> http.client.HTTPConnection:
address: Final = urlsplit(str(gateway.client.base_url))
assert address.hostname is not None and address.port is not None, address
@ -1313,8 +1369,7 @@ def test_queue_route_without_priority_takes_the_default(defaulted: Defaulted, en
marker: Final = new_marker()
answer: Final = call(defaulted.pair.first, endpoint, group, marker, priority=None)
assert_served(answer, marker)
if endpoint == "queue":
assert answer.headers.get("x-litellm-priority") == str(DEFAULT_PRIORITY), answer.headers
assert answer.headers.get("x-litellm-priority") == str(DEFAULT_PRIORITY), answer.headers
assert "priority" not in upstream_body_for(defaulted.pair.wire, marker)

195
tests/integration/routing/test_silent_model_shadow_wire.py Normal file
View file

@ -0,0 +1,195 @@
from __future__ import annotations
import json
import re
import uuid
from collections.abc import Iterator, Mapping
from dataclasses import dataclass
from itertools import chain
from types import MappingProxyType
from typing import Final, Literal
from urllib.parse import unquote
import httpx
import pytest
from integration._support.anthropic_sse import parse_sse
from integration._support.client import Gateway, eventually, gateway_from_environment
from integration._support.openai_wire import chat_reply
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import JsonValue, TypeAdapter
Shape = Literal["single", "pair"]
_MODEL: Final = "gpt-5.6"
_API_KEY: Final = "synthetic-shadow-key"
_ROUTER_ONLY_KEYS: Final = ("silent_model", "include_fallback_errors")
_ANSWER: Final = "answered by"
_NONCE: Final = re.compile(r"nonce=([0-9a-f]{32})")
_JSON: Final = TypeAdapter(dict[str, JsonValue])
_CHOICES: Final = TypeAdapter(list[dict[str, JsonValue]])
_FRESH_CONNECTION: Final = MappingProxyType({"Connection": "close"})
_SPEND_ROW_SECONDS: Final = 70
_TARGETS: Final[Mapping[Shape, tuple[str, ...]]] = MappingProxyType(
{"single": ("shadow-a",), "pair": ("shadow-a", "shadow-b")}
)
_CASES: Final = (
pytest.param("single", False, id="non-stream-single"),
pytest.param("single", True, id="stream-single"),
pytest.param("pair", True, id="stream-pair"),
)
def _prompt(nonce: str) -> str:
return f"which deployment answers this? nonce={nonce}"
def _nonce_of(raw: str) -> str:
found: Final = _NONCE.search(raw)
assert found is not None, raw
return found.group(1)
def _identity(segment: str, nonce: str) -> str:
return f"chatcmpl-{segment}-{nonce}"
def _request(model: str, nonce: str, *, stream: bool) -> dict[str, JsonValue]:
body: Final[dict[str, JsonValue]] = {
"model": model,
"stream": stream,
"messages": [{"role": "user", "content": _prompt(nonce)}],
}
return {**body, "stream_options": {"include_usage": True}} if stream else body
def _segment_of(request: Request) -> str:
return unquote(request.target).lstrip("/").partition("/")[0]
def _peer(request: Request) -> Reply:
if request.method == "GET":
return Reply(body=json.dumps({"object": "list", "data": []}).encode())
segment: Final = _segment_of(request)
body: Final = _JSON.validate_json(request.body)
return chat_reply(
_identity(segment, _nonce_of(request.body.decode())),
_MODEL,
f"{_ANSWER} {segment}",
stream=body.get("stream") is True,
)
@dataclass(frozen=True, slots=True)
class _Posted:
segment: str
streamed: bool
leaked: tuple[str, ...]
def _posted(wire: Wire, nonce: str) -> tuple[_Posted, ...]:
return tuple(
_Posted(
_segment_of(request),
_JSON.validate_json(request.body).get("stream") is True,
tuple(key for key in _ROUTER_ONLY_KEYS if key in request.body.decode()),
)
for request in wire.drain()
if request.method == "POST" and nonce in request.body.decode()
)
def _delta_text(data: Mapping[str, JsonValue]) -> str:
deltas: Final = (
_JSON.validate_python(choice.get("delta") or {}) for choice in _CHOICES.validate_python(data["choices"])
)
return "".join(str(delta.get("content") or "") for delta in deltas)
def _served(response: httpx.Response, *, stream: bool) -> tuple[frozenset[str], str]:
if not stream:
body: Final = _JSON.validate_json(response.content)
message: Final = _JSON.validate_python(_CHOICES.validate_python(body["choices"])[0]["message"])
return frozenset({str(body["id"])}), str(message["content"])
events: Final = tuple(event.data for event in parse_sse(response.text))
return frozenset(str(data["id"]) for data in events), "".join(_delta_text(data) for data in events)
@dataclass(frozen=True, slots=True)
class _SpendRow:
request_id: str
model_group: str
spend: float
def _spend_row(gateway: Gateway, request_id: str) -> tuple[_SpendRow, ...]:
rows: Final = gateway.rows(
'SELECT model_group, spend FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (request_id,)
)
return tuple(_SpendRow(request_id, str(row["model_group"]), float(str(row["spend"]))) for row in rows)
def _spend_rows(gateway: Gateway, request_ids: tuple[str, ...]) -> tuple[_SpendRow, ...]:
return tuple(chain.from_iterable(_spend_row(gateway, request_id) for request_id in request_ids))
@dataclass(frozen=True, slots=True)
class _Mirrors:
gateway: Gateway
wire: Wire
groups: Mapping[str, str]
@pytest.fixture(scope="module")
def mirrors() -> Iterator[_Mirrors]:
with gateway_from_environment() as gateway, wire_server(_peer) as wire, gateway.scenario() as scenario:
def register(segment: str, silent_model: JsonValue) -> str:
return scenario.model(
model=f"openai/{_MODEL}", api_key=_API_KEY, api_base=f"{wire.url}/{segment}", silent_model=silent_model
)
shadow_a: Final = register("shadow-a", None)
shadow_b: Final = register("shadow-b", None)
groups: Final = MappingProxyType(
{
"shadow-a": shadow_a,
"shadow-b": shadow_b,
"primary-single": register("primary-single", shadow_a),
"primary-pair": register("primary-pair", [shadow_a, shadow_b]),
}
)
wire.drain()
yield _Mirrors(gateway, wire, groups)
@pytest.mark.parametrize(("shape", "stream"), _CASES)
def test_each_shadow_target_gets_one_request_in_the_callers_stream_mode(
mirrors: _Mirrors, shape: Shape, stream: bool
) -> None:
nonce: Final = uuid.uuid4().hex
primary: Final = f"primary-{shape}"
targets: Final = _TARGETS[shape]
response: Final = mirrors.gateway.request(
"POST",
"/v1/chat/completions",
_request(mirrors.groups[primary], nonce, stream=stream),
headers=_FRESH_CONNECTION,
)
assert response.status_code == 200, response.text
assert _served(response, stream=stream) == (frozenset({_identity(primary, nonce)}), f"{_ANSWER} {primary}")
expected_groups: Final = MappingProxyType(
{_identity(segment, nonce): mirrors.groups[segment] for segment in (primary, *targets)}
)
rows: Final = eventually(
lambda: _spend_rows(mirrors.gateway, tuple(expected_groups)),
lambda found: len(found) == len(expected_groups),
seconds=_SPEND_ROW_SECONDS,
)
assert {row.request_id: row.model_group for row in rows} == expected_groups, rows
assert all(row.spend > 0 for row in rows), rows
posted: Final = _posted(mirrors.wire, nonce)
assert sorted(item.segment for item in posted) == sorted((primary, *targets)), posted
assert all(item.streamed is stream for item in posted), posted
assert all(item.leaked == () for item in posted), posted

146
tests/integration/sdk/test_anthropic_json_tool_cut_stream_sync_sdk.py Normal file
View file

@ -0,0 +1,146 @@
import base64
import json
import uuid
from collections.abc import Callable, Mapping
from typing import Final
import pytest
from integration._support.upstream import aws_event_stream_frame
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue
import litellm
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
_JSON_TOOL: Final = "json_tool_call"
_FRAGMENT_SIZE: Final = 20
_CUT_EVENTS: Final = 3
_SCHEMA: Final[dict[str, JsonValue]] = {
"type": "object",
"properties": {"tier": {"type": "integer"}, "summary": {"type": "string"}},
"required": ["tier", "summary"],
"additionalProperties": False,
}
_RESPONSE_FORMAT: Final[dict[str, JsonValue]] = {
"type": "json_schema",
"json_schema": {"name": "risk", "schema": _SCHEMA},
}
_ANTHROPIC_MODEL: Final = "anthropic/claude-scripted-9"
_BEDROCK_INVOKE_MODEL: Final = "bedrock/invoke/anthropic.claude-sonnet-5"
def _payload(marker_hex: str) -> dict[str, JsonValue]:
return {"tier": 2, "summary": marker_hex}
def _tool_input(shape: str, marker_hex: str) -> dict[str, JsonValue]:
if shape == "wrapped":
return {"tool_name": _JSON_TOOL, "parameters": _payload(marker_hex)}
return _payload(marker_hex)
def _events(marker_hex: str, shape: str, cut: bool) -> tuple[tuple[str, dict[str, JsonValue]], ...]:
text: Final = json.dumps(_tool_input(shape, marker_hex))
deltas: Final = tuple(
(
"content_block_delta",
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "input_json_delta", "partial_json": text[start : start + _FRAGMENT_SIZE]},
},
)
for start in range(0, len(text), _FRAGMENT_SIZE)
)
message: Final[dict[str, JsonValue]] = {
"id": f"msg_{marker_hex}",
"type": "message",
"role": "assistant",
"model": "claude-scripted-9",
"content": [],
"stop_reason": None,
"stop_sequence": None,
"usage": {"input_tokens": 5, "output_tokens": 1},
}
complete: Final = (
("message_start", {"type": "message_start", "message": message}),
(
"content_block_start",
{
"type": "content_block_start",
"index": 0,
"content_block": {"type": "tool_use", "id": f"toolu_{marker_hex}", "name": _JSON_TOOL, "input": {}},
},
),
*deltas,
("content_block_stop", {"type": "content_block_stop", "index": 0}),
(
"message_delta",
{
"type": "message_delta",
"delta": {"stop_reason": "tool_use", "stop_sequence": None},
"usage": {"output_tokens": 3},
},
),
("message_stop", {"type": "message_stop"}),
)
return complete[:-_CUT_EVENTS] if cut else complete
def _aws_chunk(data: Mapping[str, JsonValue]) -> bytes:
payload: Final = json.dumps({"bytes": base64.b64encode(json.dumps(data).encode()).decode()}).encode()
return aws_event_stream_frame(
{":event-type": "chunk", ":content-type": "application/json", ":message-type": "event"}, payload
)
def _peer(marker_hex: str, shape: str, cut: bool) -> Callable[[Request], Reply]:
events: Final = _events(marker_hex, shape, cut)
def respond(request: Request) -> Reply:
assert marker_hex in request.body.decode(), request.body
if request.target.endswith("/invoke-with-response-stream"):
return Reply(
content_type="application/vnd.amazon.eventstream",
chunks=tuple(_aws_chunk(data) for _, data in events),
)
assert request.target.endswith("/v1/messages"), request.target
return Reply(
content_type="text/event-stream",
chunks=tuple(f"event: {name}\ndata: {json.dumps(data)}\n\n".encode() for name, data in events),
)
return respond
def _credentials(model: str) -> dict[str, str]:
if model == _ANTHROPIC_MODEL:
return {"api_key": "synthetic-anthropic-key"}
return {"api_key": "synthetic-bedrock-bearer", "aws_region_name": "us-east-1"}
@pytest.mark.parametrize("model", [_ANTHROPIC_MODEL, _BEDROCK_INVOKE_MODEL], ids=["anthropic", "bedrock_invoke"])
@pytest.mark.parametrize(
("shape", "cut"),
[("plain", False), ("plain", True), ("wrapped", True)],
ids=["plain_complete", "plain_cut", "wrapped_cut"],
)
def test_sync_stream_cut_before_the_json_tool_block_closes_still_yields_the_whole_answer(
model: str, shape: str, cut: bool
) -> None:
marker_hex: Final = uuid.uuid4().hex
with wire_server(_peer(marker_hex, shape, cut)) as wire:
response: Final = litellm.completion(
model=model,
messages=[{"role": "user", "content": f"Rate the risk {marker_hex}"}],
response_format=_RESPONSE_FORMAT,
max_tokens=64,
stream=True,
api_base=wire.url,
num_retries=0,
**_credentials(model),
)
assert isinstance(response, CustomStreamWrapper), type(response)
text: Final = "".join(str(chunk.choices[0].delta.content or "") for chunk in tuple(response))
assert len(wire.drain()) == 1
assert json.loads(text) == _payload(marker_hex), text

39
tests/integration/sdk/test_bedrock_converse_redacted_thinking_router_sdk.py
View file

@ -88,3 +88,42 @@ def test_sync_router_reads_base_model_to_replay_redacted_thinking_to_an_applicat
assistant: Final = next(message for message in received[-1].messages if message.role == "assistant")
redacted: Final = ({"reasoningContent": {"redactedContent": REDACTED}},) if replays_redacted else ()
assert assistant.content == (*redacted, {"text": FIRST_ANSWER})
def test_sync_router_replays_a_block_in_content_and_thinking_blocks_once() -> None:
redacted: Final[dict[str, JsonValue]] = {"type": "redacted_thinking", "data": REDACTED}
with wire_server(_respond) as wire:
router: Final = Router(
model_list=[
{
"model_name": "claude-profile",
"litellm_params": {
"model": f"bedrock/converse/{APPLICATION_PROFILE_ARN}",
"api_base": wire.url,
"api_key": "synthetic-bedrock-bearer",
"aws_region_name": "us-east-1",
},
"model_info": {"base_model": "bedrock/us.anthropic.claude-sonnet-4-6"},
}
],
num_retries=0,
)
response: Final = router.completion(
model="claude-profile",
max_tokens=256,
messages=[
{"role": "user", "content": "What's the weather in Paris?"},
{
"role": "assistant",
"content": [redacted, {"type": "text", "text": FIRST_ANSWER}],
"thinking_blocks": [redacted],
},
{"role": "user", "content": "Reply with the word ok."},
],
)
assert response.choices[0].message.content == FINAL_ANSWER
received: Final = tuple(_ConverseRequest.model_validate_json(request.body) for request in wire.drain())
assistant: Final = next(message for message in received[-1].messages if message.role == "assistant")
assert assistant.content == ({"reasoningContent": {"redactedContent": REDACTED}}, {"text": FIRST_ANSWER})

33
tests/integration/sdk/test_local_sandbox_sdk.py
View file

@ -1,16 +1,25 @@
import asyncio
import os
import sys
from collections.abc import AsyncIterator
from pathlib import Path
from typing import Final
import psutil
import pytest
import pytest_asyncio
from integration._support.client import eventually
from integration._support.process import kill_stray_process
from litellm import sandbox
from litellm.harness.sandbox import LocalSandbox
WRITE_AT_PWD: Final = "import os;open(os.path.join(os.environ['PWD'],'hello.txt'),'w').write('hi')"
SLEEPER_SPAWNED_ONCE_STDIN_CLOSES: Final = (
"import subprocess,sys;sys.stdin.read();"
"p=subprocess.Popen([sys.executable,'-c','import time;time.sleep(60)']);"
"open('grandchild.pid','w').write(str(p.pid));p.wait()"
)
@pytest_asyncio.fixture(loop_scope="function")
@ -37,3 +46,27 @@ async def test_child_trusting_pwd_writes_inside_its_cwd(
assert os.listdir(elsewhere) == []
assert await box.read("hello.txt") == b"hi"
assert await box.read("sub/hello.txt") == b"hi"
def running(pid: int) -> bool:
try:
return psutil.Process(pid).status() != psutil.STATUS_ZOMBIE
except psutil.NoSuchProcess:
return False
@pytest.mark.asyncio(loop_scope="function")
async def test_cancelled_run_stops_the_whole_command(box: LocalSandbox) -> None:
run: Final = asyncio.ensure_future(box.run([sys.executable, "-I", "-c", SLEEPER_SPAWNED_ONCE_STDIN_CLOSES]))
pid_file: Final = Path(box.workdir) / "grandchild.pid"
grandchild: Final = int(
await asyncio.to_thread(eventually, lambda: pid_file.read_text() if pid_file.exists() else "", bool)
)
try:
run.cancel()
with pytest.raises(asyncio.CancelledError):
await run
assert not eventually(lambda: running(grandchild), lambda alive: not alive, return_last_on_timeout=True)
finally:
kill_stray_process(grandchild)

129
tests/integration/sdk/test_router_silent_model_shadow_sdk.py Normal file
View file

@ -0,0 +1,129 @@
from __future__ import annotations
import re
import uuid
from collections.abc import Mapping
from dataclasses import dataclass
from queue import SimpleQueue
from typing import Final
import pytest
from integration._support.client import eventually
from integration._support.openai_wire import chat_reply
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import JsonValue, TypeAdapter
import litellm
from litellm import CustomStreamWrapper, Router
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.utils import ModelResponseStream
_MODEL: Final = "gpt-5.6"
_API_KEY: Final = "synthetic-shadow-key"
_ANSWER: Final = "answered by"
_NONCE: Final = re.compile(r"nonce=([0-9a-f]{32})")
_JSON: Final = TypeAdapter(dict[str, JsonValue])
_OBJECT: Final = TypeAdapter(dict[str, object])
_SHADOW_SECONDS: Final = 20
_TARGETS: Final = (
pytest.param(("shadow-a",), id="single"),
pytest.param(("shadow-a", "shadow-b"), id="pair"),
)
def _prompt(nonce: str) -> str:
return f"which deployment answers this? nonce={nonce}"
def _peer(request: Request) -> Reply:
deployment: Final = request.target.lstrip("/").partition("/")[0]
raw: Final = request.body.decode()
found: Final = _NONCE.search(raw)
assert found is not None, raw
return chat_reply(
f"chatcmpl-{deployment}-{found.group(1)}",
_MODEL,
f"{_ANSWER} {deployment}",
stream=_JSON.validate_json(request.body).get("stream") is True,
)
def _deployment(name: str, wire: Wire, **extra: JsonValue) -> dict[str, JsonValue]:
return {
"model_name": name,
"litellm_params": {"model": f"openai/{_MODEL}", "api_base": f"{wire.url}/{name}", "api_key": _API_KEY, **extra},
}
@dataclass(frozen=True, slots=True)
class _Posted:
deployment: str
streamed: bool
leaked: bool
def _posted(wire: Wire) -> tuple[_Posted, ...]:
return tuple(
_Posted(
request.target.lstrip("/").partition("/")[0],
_JSON.validate_json(request.body).get("stream") is True,
"silent_model" in request.body.decode(),
)
for request in wire.drain()
if request.method == "POST"
)
class _ShadowRecorder(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.shadow_groups: SimpleQueue[str] = SimpleQueue()
async def async_log_success_event(
self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object
) -> None:
metadata: Final = _OBJECT.validate_python(_OBJECT.validate_python(kwargs["litellm_params"])["metadata"] or {})
if metadata.get("is_silent_experiment") is True:
self.shadow_groups.put(str(metadata.get("model_group")))
def drained(self) -> tuple[str, ...]:
return tuple(self.shadow_groups.get_nowait() for _ in range(self.shadow_groups.qsize()))
def _delta(chunk: object) -> str:
assert isinstance(chunk, ModelResponseStream), chunk
return "".join(str(choice.delta.content or "") for choice in chunk.choices)
@pytest.mark.parametrize("targets", _TARGETS)
def test_sync_streamed_completion_streams_and_logs_each_shadow_target(
targets: tuple[str, ...], monkeypatch: pytest.MonkeyPatch
) -> None:
recorder: Final = _ShadowRecorder()
monkeypatch.setattr(litellm, "callbacks", [recorder])
with wire_server(_peer) as wire:
silent_model: Final[JsonValue] = targets[0] if len(targets) == 1 else [*targets]
router: Final = Router(
model_list=[
_deployment("mirrored", wire, silent_model=silent_model),
*(_deployment(target, wire) for target in targets),
],
num_retries=0,
disable_cooldowns=True,
)
response: Final = router.completion(
model="mirrored",
messages=[{"role": "user", "content": _prompt(uuid.uuid4().hex)}],
stream=True,
)
assert isinstance(response, CustomStreamWrapper), response
text: Final = "".join(_delta(chunk) for chunk in response)
eventually(wire.received.qsize, lambda count: count >= 1 + len(targets), seconds=_SHADOW_SECONDS)
posted: Final = _posted(wire)
eventually(recorder.shadow_groups.qsize, lambda count: count >= len(targets), seconds=_SHADOW_SECONDS)
logged: Final = recorder.drained()
assert text == f"{_ANSWER} mirrored"
assert sorted(item.deployment for item in posted) == sorted(("mirrored", *targets)), posted
assert all(item.streamed for item in posted), posted
assert not any(item.leaked for item in posted), posted
assert sorted(logged) == sorted(targets), logged

126
tests/integration/security/test_files_batches_access.py
View file

@ -48,6 +48,7 @@ VERTEX_ACCESS_TOKEN: Final = "ya29.scripted-delete-token"
AWS_ACCESS_KEY_ID: Final = "AKIAIOSFODNN7EXAMPLE"
AWS_SECRET_ACCESS_KEY: Final = "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"
AWS_REGION: Final = "us-east-1"
S3_OBJECT_BODY: Final = b'{"recordId":"r1","modelOutput":{"stop_reason":"end_turn"}}\n'
_ForbiddenSurface = Literal[
"upload_sdk_extra_body",
"upload_form",
@ -708,6 +709,86 @@ def test_non_admin_cannot_delete_a_raw_gcs_uri_before_vertex_auth_or_storage(gat
assert token_requests == (), "non-admin raw GCS delete reached the Vertex token endpoint"
def test_non_admin_still_reads_an_ordinary_provider_file_id(gateway: Gateway) -> None:
with wire_server(_openai_backend()) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="provider-key-" + uuid.uuid4().hex[:8]
)
_wait_until_every_worker_serves(gateway, model)
member: Final = _member(scenario, model)
file_id: Final = "file-" + uuid.uuid4().hex
read: Final = gateway.request("GET", f"/v1/files/{file_id}/content", key=member.key, params={"model": model})
assert read.status_code == 200, read.text
assert read.content == b"file-content\n", read.content
assert [
(request.method, urlsplit(request.target).path) for request in _drained_other_than_model_list_probes(wire)
] == [("GET", f"/v1/files/{file_id}/content")]
@pytest.mark.parametrize(
("method", "route_suffix", "verb"),
[("GET", "/content", "retrieved"), ("DELETE", "", "deleted")],
ids=["content", "delete"],
)
def test_non_admin_cannot_reach_vertex_with_a_twice_encoded_raw_gcs_uri(
gateway: Gateway, method: Literal["GET", "DELETE"], route_suffix: str, verb: str
) -> None:
with (
wire_server(_vertex_token_backend) as token_wire,
gateway.scenario() as scenario,
):
model: Final = scenario.model(
model=VERTEX_MODEL,
api_key=None,
vertex_project=VERTEX_PROJECT,
vertex_location=VERTEX_LOCATION,
vertex_credentials=_vertex_credentials(token_wire.url),
gcs_bucket_name=BUCKET,
)
_wait_until_every_worker_serves(gateway, model)
caller_user: Final = scenario.user(user_role="internal_user")
caller_key: Final = scenario.key(user_id=caller_user, models=[model])
object_uri: Final = f"gs://{BUCKET}/litellm-vertex-files/litellm-vertex-files-obj.jsonl"
twice_encoded: Final = quote(quote(object_uri, safe=""), safe="")
token_wire.drain()
denied: Final = gateway.request(
method, f"/v1/files/{twice_encoded}{route_suffix}", key=caller_key, params={"model": model}
)
token_requests: Final = token_wire.drain()
assert denied.status_code == 403, f"{denied.text}; token requests={token_requests}"
assert _error_message(denied, 403) == (
f"Raw cloud storage file ids can only be {verb} by a proxy admin key. "
"Use the LiteLLM managed file id returned when the file was created."
), denied.text
assert token_requests == (), f"non-admin twice-encoded raw GCS {method} reached the Vertex token endpoint"
def _assert_signed_empty_body_s3_request(request: Request, method: Literal["GET", "DELETE"], wire_url: str) -> None:
assert request.method == method, request.method
request_date: Final = request.headers["x-amz-date"]
content_hash: Final = hashlib.sha256(b"").hexdigest()
assert request.headers["x-amz-content-sha256"] == content_hash, request.headers
aws_request: Final = AWSRequest(
method=method,
url=wire_url + request.target,
headers={
"host": request.headers["host"],
"x-amz-date": request_date,
"x-amz-content-sha256": content_hash,
},
)
aws_request.context["timestamp"] = request_date
auth: Final = S3SigV4Auth(Credentials(AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY), "s3", AWS_REGION)
signature: Final = auth.signature(auth.string_to_sign(aws_request, auth.canonical_request(aws_request)), aws_request)
expected_authorization: Final = (
f"AWS4-HMAC-SHA256 Credential={AWS_ACCESS_KEY_ID}/{request_date[:8]}/{AWS_REGION}/s3/aws4_request, "
f"SignedHeaders={auth.signed_headers(auth.headers_to_sign(aws_request))}, Signature={signature}"
)
assert request.headers["authorization"] == expected_authorization, request.headers["authorization"]
def test_proxy_admin_deletes_a_raw_s3_uri_with_one_signed_delete(gateway: Gateway) -> None:
with wire_server(_s3_backend()) as wire, gateway.scenario() as scenario:
model: Final = _bedrock_model(scenario, wire)
@ -721,29 +802,30 @@ def test_proxy_admin_deletes_a_raw_s3_uri_with_one_signed_delete(gateway: Gatewa
assert parsed.deleted is True and parsed.id == object_uri, deleted.text
(request,) = _drained_other_than_model_list_probes(wire)
assert request.target == f"/{BUCKET}/litellm-bedrock-files-obj.jsonl", request.target
request_date: Final = request.headers["x-amz-date"]
content_hash: Final = hashlib.sha256(b"").hexdigest()
assert request.headers["x-amz-content-sha256"] == content_hash, request.headers
aws_request: Final = AWSRequest(
method="DELETE",
url=wire.url + request.target,
headers={
"host": request.headers["host"],
"x-amz-date": request_date,
"x-amz-content-sha256": content_hash,
},
_assert_signed_empty_body_s3_request(request, "DELETE", wire.url)
def _s3_object_backend(request: Request) -> Reply:
if request.method == "GET":
return Reply(body=S3_OBJECT_BODY, content_type="application/octet-stream")
return _json_reply({"error": {"message": f"unscripted {request.method} {request.target}"}}, 404)
def test_proxy_admin_reads_a_raw_s3_uri_with_one_signed_get(gateway: Gateway) -> None:
with wire_server(_s3_object_backend) as wire, gateway.scenario() as scenario:
model: Final = _bedrock_model(scenario, wire)
_wait_until_every_worker_serves(gateway, model)
object_key: Final = f"litellm-batch-outputs/{uuid.uuid4().hex}.jsonl.out"
object_uri: Final = f"s3://{BUCKET}/{object_key}"
read: Final = gateway.request(
"GET", "/v1/files/" + quote(object_uri, safe="") + "/content", params={"model": model}
)
aws_request.context["timestamp"] = request_date
auth: Final = S3SigV4Auth(Credentials(AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY), "s3", AWS_REGION)
signature: Final = auth.signature(
auth.string_to_sign(aws_request, auth.canonical_request(aws_request)), aws_request
)
expected_authorization: Final = (
f"AWS4-HMAC-SHA256 Credential={AWS_ACCESS_KEY_ID}/{request_date[:8]}/{AWS_REGION}/s3/aws4_request, "
f"SignedHeaders={auth.signed_headers(auth.headers_to_sign(aws_request))}, Signature={signature}"
)
assert request.headers["authorization"] == expected_authorization, request.headers["authorization"]
assert request.method == "DELETE", request.method
storage_requests: Final = _drained_other_than_model_list_probes(wire)
assert read.status_code == 200, f"{read.text}; storage requests={storage_requests}"
assert read.content == S3_OBJECT_BODY, read.content
assert [request.target for request in storage_requests] == [f"/{BUCKET}/{object_key}"], storage_requests
_assert_signed_empty_body_s3_request(storage_requests[0], "GET", wire.url)
def _granting_vertex_token_backend(request: Request) -> Reply:

31
tests/integration/spend/test_background_response_billed_once.py
View file

@ -642,6 +642,15 @@ def _key_budget(scenario: Scenario, budget_owner: str, model_budget: JsonValue)
return {}
def _assert_the_next_create_exceeds_the_model_budget(scenario: Scenario, key: str, model: str) -> None:
refused: Final = eventually(
lambda: _next_create(scenario, key, model),
lambda response: response.status_code != 200,
seconds=SPEND_LANDING_SECONDS,
)
assert f"exceeded budget for model={model}" in refused.text, refused.text
@pytest.mark.timeout(240)
@pytest.mark.parametrize("budget_owner", ["key", "team", "user", "budget_tier"])
def test_a_billed_background_response_counts_toward_its_creators_project_and_model_budget(
@ -663,12 +672,24 @@ def test_a_billed_background_response_counts_toward_its_creators_project_and_mod
lambda spends: spends[0] > 0 and costs(spends[0])(spends[1]),
seconds=SPEND_LANDING_SECONDS,
)
refused: Final = eventually(
lambda: _next_create(scenario, key, upstream.model),
lambda response: response.status_code != 200,
seconds=SPEND_LANDING_SECONDS,
_assert_the_next_create_exceeds_the_model_budget(scenario, key, upstream.model)
@pytest.mark.timeout(240)
def test_a_billed_background_response_on_a_team_model_counts_toward_the_budget_of_the_model_name_it_was_created_with(
gateway: Gateway,
) -> None:
with gateway.scenario() as scenario:
team_id: Final = scenario.team()
upstream: Final = scripted_upstream(scenario, background_routes("completed"), model_info={"team_id": team_id})
key: Final = scenario.key(
team_id=team_id, model_max_budget={upstream.model: {"max_budget": PRICE / 2, "budget_duration": "1d"}}
)
assert f"exceeded budget for model={upstream.model}" in refused.text, refused.text
response_id: Final = create_background_response(scenario, key, upstream.model)
read: Final = read_response(gateway, key, response_id)
assert (read,) in (settling_reads("completed", 1), free_reads("completed", 1))
eventually(lambda: key_spend(key), lambda spend: spend > 0, seconds=SPEND_LANDING_SECONDS)
_assert_the_next_create_exceeds_the_model_budget(scenario, key, upstream.model)
@pytest.mark.timeout(240)

229
tests/integration/spend/test_baggage_tag_mappings.py
View file

@ -2,8 +2,9 @@ import itertools
import json
import uuid
from collections.abc import Generator, Iterator, Mapping, Sequence
from concurrent.futures import ThreadPoolExecutor
from contextlib import ExitStack, contextmanager
from dataclasses import dataclass
from dataclasses import dataclass, replace
from hashlib import sha256
from pathlib import Path
from types import MappingProxyType
@ -15,14 +16,15 @@ import yaml
from integration._support.client import (
JSON_OBJECT,
Gateway,
Scenario,
eventually,
gateway_from_environment,
list_value,
object_value,
)
from integration._support.database import scratch_database
from integration._support.database import read_rows, scratch_database, write_rows
from integration._support.database_relay import dropped_connection_relay, plaintext_database_url
from integration._support.process import graceful_stop_seconds, owned_proxy
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy, owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import JsonValue, TypeAdapter
from websockets.exceptions import InvalidStatus
@ -33,6 +35,8 @@ pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml"
PRICED: Final[dict[str, JsonValue]] = {"input_cost_per_token": 0.01, "output_cost_per_token": 0.01}
WORKER_PROBES: Final = 6
VLLM_UPSTREAM_MODEL: Final = "baggage-model"
PINNED_CONNECTION: Final = httpx.Limits(max_connections=1)
JSON_LIST: Final = TypeAdapter(list[JsonValue])
GUARDRAIL_VERDICT: Final = b'{"action":"GUARDRAIL_INTERVENED","texts":["masked baggage prompt"]}'
GUARDRAIL_PATH: Final = "/beta/litellm_basic_guardrail_api"
@ -40,6 +44,7 @@ NO_FIELDS: Final[Mapping[str, JsonValue]] = MappingProxyType({})
Route = Literal["chat", "chat_litellm_metadata", "messages", "responses"]
CallerTag = Literal["none", "header", "metadata_json_with_header", "litellm_metadata"]
PassThrough = Literal["custom", "anthropic", "vllm"]
RejectedClientTag = Literal["none", "header", "body", "metadata_json", "litellm_metadata"]
@ -120,11 +125,31 @@ def _responses_reply() -> Reply:
)
def _anthropic_reply() -> Reply:
marker: Final = uuid.uuid4().hex
return Reply(
body=json.dumps(
{
"id": f"msg_{marker}",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-5",
"content": [{"type": "text", "text": marker}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 5, "output_tokens": 3},
}
).encode()
)
def _provider(request: Request) -> Reply:
if request.target == GUARDRAIL_PATH:
return Reply(body=b'{"action":"NONE"}')
if request.target == "/v1/responses":
return _responses_reply()
if request.target == "/v1/messages":
return _anthropic_reply()
return _chat_reply()
@ -183,7 +208,15 @@ def _mapped_models(marker: str, provider: Wire, upstream_url: str) -> list[JsonV
}
guarded: Final = _deployment(f"guarded-{marker}", provider, tags=[f"guarded-{marker}:on"])
pre_guarded: Final = _deployment(f"guarded-pre-{marker}", provider, tags=[f"guarded-{marker}:pre"])
return [_deployment(f"chat-{marker}", provider), guarded, pre_guarded, *routed, realtime]
vllm: Final[dict[str, JsonValue]] = {
"model_name": f"vllm-{marker}",
"litellm_params": {
"model": f"hosted_vllm/{VLLM_UPSTREAM_MODEL}",
"api_base": provider.url,
"api_key": "synthetic-vllm",
},
}
return [_deployment(f"chat-{marker}", provider), guarded, pre_guarded, *routed, realtime, vllm]
def _mapped_rules(marker: str, *, tenant_required: bool) -> list[JsonValue]:
@ -239,9 +272,12 @@ def _rig(
{"callbacks": ["generic_api"], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1},
[_deployment_tag_mode_guardrail(marker, provider)] if tag_mode_guardrail else (),
)
with owned_proxy(
base, directory, {"GENERIC_LOGGER_ENDPOINT": callbacks.url}, config=config, workers=workers
) as proxy:
environment: Final = {
"GENERIC_LOGGER_ENDPOINT": callbacks.url,
"ANTHROPIC_API_BASE": provider.url,
"ANTHROPIC_API_KEY": "synthetic-anthropic-key",
}
with owned_proxy(base, directory, environment, config=config, workers=workers) as proxy:
yield Rig(proxy, provider, callbacks, marker)
@ -676,6 +712,185 @@ def test_custom_passthrough_spend_row_carries_the_mapped_tags(mapped: Rig) -> No
assert {mapped.exact_tag, mapped.tenant_tag(tenant)} <= _spend_tags(mapped.proxy, key)
def _passthrough_path(rig: Rig, scenario: Scenario, entrypoint: PassThrough) -> str:
match entrypoint:
case "custom":
path: Final = f"/baggage-passthrough-{uuid.uuid4().hex}"
created: Final = rig.proxy.post(
"/config/pass_through_endpoint", {"path": path, "target": f"{rig.provider.url}/echo", "auth": True}
)
endpoint_id: Final = str(object_value(list_value(created["endpoints"])[0])["id"])
scenario.cleanups.callback(
lambda: rig.proxy.request(
"DELETE", "/config/pass_through_endpoint", params={"endpoint_id": endpoint_id}
)
)
return path
case "anthropic":
return "/anthropic/v1/messages"
case "vllm":
return "/vllm/v1/chat/completions"
def _passthrough_body(rig: Rig, entrypoint: PassThrough, text: str) -> dict[str, JsonValue]:
model: Final = {"custom": rig.chat_model, "anthropic": "claude-sonnet-4-5", "vllm": f"vllm-{rig.marker}"}
body: Final[dict[str, JsonValue]] = {
"model": model[entrypoint],
"max_tokens": 16,
"messages": [{"role": "user", "content": text}],
}
match entrypoint:
case "vllm":
return {**body, "metadata": {"user_id": f"user-{text}"}}
case "custom" | "anthropic":
return body
def _upstream_body(entrypoint: PassThrough, body: Mapping[str, JsonValue]) -> JsonValue:
match entrypoint:
case "vllm":
return {**body, "model": VLLM_UPSTREAM_MODEL}
case "custom" | "anthropic":
return dict(body)
def _forwarded_bodies(rig: Rig) -> tuple[JsonValue, ...]:
return tuple(
JSON_OBJECT.validate_json(request.body) for request in rig.provider.drain() if request.target != GUARDRAIL_PATH
)
@pytest.mark.parametrize("entrypoint", ["custom", "anthropic", "vllm"])
def test_passthrough_forwards_the_client_body_unchanged_while_its_spend_row_carries_the_mapped_tags(
optional: Rig, entrypoint: PassThrough
) -> None:
tenant: Final = uuid.uuid4().hex
with optional.proxy.scenario() as scenario:
path: Final = _passthrough_path(optional, scenario, entrypoint)
key_fields: Final[dict[str, JsonValue]] = (
{"allowed_passthrough_routes": [path]} if entrypoint == "custom" else {}
)
plain_key: Final = scenario.key(**key_fields)
key: Final = scenario.key(**key_fields)
body: Final = _passthrough_body(optional, entrypoint, tenant)
optional.provider.drain()
plain: Final = optional.proxy.request("POST", path, body, key=plain_key)
plain_forwarded: Final = _forwarded_bodies(optional)
response: Final = optional.proxy.request(
"POST", path, body, key=key, headers={"baggage": f"entrypoint=a,tenant={tenant}"}
)
assert (plain.status_code, response.status_code) == (200, 200), (plain.text, response.text)
assert plain_forwarded == (_upstream_body(entrypoint, body),)
assert _forwarded_bodies(optional) == plain_forwarded
mapped_tags: Final = {optional.exact_tag, optional.tenant_tag(tenant)}
assert _spend_tags(optional.proxy, key) == _spend_tags(optional.proxy, plain_key) | mapped_tags
def _general_settings_mappings(database_url: str) -> JsonValue:
rows: Final = read_rows(
'SELECT param_value FROM "LiteLLM_Config" WHERE param_name = %s',
("general_settings",),
database_url=database_url,
)
return object_value(rows[0]["param_value"]).get("request_tag_mappings") if rows else None
def _serving_pid(proxy: Gateway) -> int:
response: Final = proxy.request("GET", "/health/integration-readiness")
return int(response.headers["x-integration-worker-pid"])
def _status_on_one_worker(owned: OwnedProxy, rig: Rig, key: str) -> tuple[int, int] | None:
with httpx.Client(
base_url=str(owned.gateway.client.base_url), timeout=15, trust_env=False, limits=PINNED_CONNECTION
) as client:
pinned: Final = replace(rig, proxy=replace(owned.gateway, client=client))
before: Final = _serving_pid(pinned.proxy)
status: Final = _chat(pinned, key, None).status_code
return (before, status) if _serving_pid(pinned.proxy) == before else None
def _worker_statuses(owned: OwnedProxy, rig: Rig, key: str) -> frozenset[tuple[int, int]]:
def answers() -> frozenset[tuple[int, int]]:
with ThreadPoolExecutor(max_workers=8) as pool:
found: Final = tuple(pool.map(lambda _: _status_on_one_worker(owned, rig, key), range(8)))
return frozenset(answer for answer in found if answer is not None)
return eventually(answers, lambda found: frozenset(pid for pid, _ in found) >= owned.worker_pids(), seconds=30)
def _mappings_on_a_worker(proxy: Gateway) -> tuple[int, JsonValue]:
response: Final = proxy.request(
"GET",
"/health/integration-general-setting",
params={"name": "request_tag_mappings"},
headers={"connection": "close"},
)
assert response.status_code == 200, response.text
return int(response.headers["x-integration-worker-pid"]), JSON_OBJECT.validate_json(response.content)["value"]
def _mappings_settled_on_every_worker(owned: OwnedProxy, mappings: JsonValue) -> None:
def answers() -> tuple[tuple[int, JsonValue], ...]:
with ThreadPoolExecutor(max_workers=8) as pool:
return tuple(pool.map(lambda _: _mappings_on_a_worker(owned.gateway), range(16)))
eventually(
answers,
lambda found: (
frozenset(pid for pid, _ in found) >= owned.worker_pids() and all(value == mappings for _, value in found)
),
seconds=70,
)
def test_mappings_saved_through_the_settings_endpoint_apply_on_every_worker_and_an_invalid_stored_value_is_ignored(
gateway: Gateway, tmp_path: Path
) -> None:
marker: Final = uuid.uuid4().hex[:12]
rule: Final[dict[str, JsonValue]] = {"key": "tenant", "tag_prefix": f"tenant-{marker}:", "required": True}
with ExitStack() as stack:
provider: Final = stack.enter_context(wire_server(_provider))
config: Final = _write_config(tmp_path, [_deployment(f"chat-{marker}", provider)], {}, {})
owned: Final = stack.enter_context(owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2))
proxy: Final = owned.gateway
database_url: Final = proxy.database_url
assert database_url is not None
scenario: Final = stack.enter_context(proxy.scenario())
rig: Final = Rig(proxy, provider, provider, marker)
key: Final = scenario.key()
def save(rules: list[JsonValue]) -> httpx.Response:
return proxy.request(
"POST",
"/config/field/update",
{"field_name": "request_tag_mappings", "field_value": rules, "config_type": "general_settings"},
)
duplicate: Final = save([rule, {**rule, "required": False}])
assert duplicate.status_code == 400, duplicate.text
assert _general_settings_mappings(database_url) is None
saved: Final = save([rule])
assert saved.status_code == 200, saved.text
_mappings_settled_on_every_worker(owned, [rule])
assert _worker_statuses(owned, rig, key) == frozenset((pid, 400) for pid in owned.worker_pids())
tagged_key: Final = scenario.key()
provider.drain()
admitted: Final = _chat(rig, tagged_key, f"tenant={marker}")
assert admitted.status_code == 200, admitted.text
assert len(provider.drain()) == 1
assert _spend_tags(proxy, tagged_key) == {f"tenant-{marker}:{marker}"}
write_rows(
'UPDATE "LiteLLM_Config" SET param_value = jsonb_set(param_value::jsonb, %s, %s::jsonb) '
"WHERE param_name = %s",
("{request_tag_mappings}", json.dumps([{"key": "tenant", "required": True}]), "general_settings"),
database_url=database_url,
)
_mappings_settled_on_every_worker(owned, None)
assert _worker_statuses(owned, rig, key) == frozenset((pid, 200) for pid in owned.worker_pids())
def _rejected_client_tag(client_tag: RejectedClientTag, tag: str) -> tuple[Mapping[str, JsonValue], Mapping[str, str]]:
match client_tag:
case "none":

48
tests/integration/spend/test_failed_dispatch_tokens.py
View file

@ -167,3 +167,51 @@ def test_provider_500_failure_rows_count_deepseek_v4_prompt_tokens_with_deepseek
short: Final = _failure_row_prompt_tokens(gateway, model, key, SHORT_TEXT)
assert wire.drain()
assert_counted_with_deepseek(chinese - short, chinese_minus_short)
COMPACTION_INSTRUCTIONS: Final[JsonValue] = {"role": "system", "content": "answer briefly"}
COMPACTED_HISTORY: Final[tuple[JsonValue, ...]] = (
{"role": "user", "content": "a long turn " * 200},
{"role": "assistant", "content": "a long reply " * 200},
)
COMPACTION_ITEM: Final[JsonValue] = {
"type": "compaction",
"encrypted_content": json.dumps({"type": "compaction", "content": "the cat is named Biscuit"}),
}
TURN_AFTER_COMPACTION: Final[JsonValue] = {"role": "user", "content": "what is the cat's name?"}
def _failed_responses_prompt_tokens(gateway: Gateway, model: str, key: str, items: list[JsonValue]) -> int:
failed: Final = gateway.request("POST", "/v1/responses", {"model": model, "input": items}, key=key)
assert failed.status_code == 500 and "synthetic provider outage" in failed.text, failed.text
rows: Final = eventually(
lambda: read_rows(
'SELECT status, prompt_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
(failed.headers["x-litellm-call-id"],),
),
lambda values: len(values) == 1,
seconds=70,
)
assert rows[0]["status"] == "failure", rows
prompt_tokens: Final = rows[0]["prompt_tokens"]
assert isinstance(prompt_tokens, int) and prompt_tokens > 0, rows
return prompt_tokens
def test_a_failed_responses_request_counts_its_input_from_the_newest_compaction_item(gateway: Gateway) -> None:
with wire_server(_provider_500) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(api_base=wire.url + "/v1", num_retries=0)
key: Final = scenario.key(models=[model])
resent: Final = _failed_responses_prompt_tokens(
gateway, model, key, [COMPACTION_INSTRUCTIONS, *COMPACTED_HISTORY, COMPACTION_ITEM, TURN_AFTER_COMPACTION]
)
compacted: Final = _failed_responses_prompt_tokens(
gateway, model, key, [COMPACTION_INSTRUCTIONS, COMPACTION_ITEM, TURN_AFTER_COMPACTION]
)
uncompacted: Final = _failed_responses_prompt_tokens(
gateway, model, key, [COMPACTION_INSTRUCTIONS, *COMPACTED_HISTORY, TURN_AFTER_COMPACTION]
)
assert len(wire.drain()) == 3
assert resent == compacted
assert resent < uncompacted

73
tests/integration/spend/test_streamed_usage_estimate.py Normal file
View file

@ -0,0 +1,73 @@
import json
import uuid
from hashlib import sha256
from typing import Final
from integration._support.client import Gateway, eventually
from integration._support.database import read_rows
from integration._support.deepseek_tokens import (
CHINESE_TEXT,
SHORT_TEXT,
assert_counted_with_deepseek,
chinese_minus_short,
)
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue
_DEEPSEEK_V4: Final = "DeepSeek-V4-Flash-0731"
def _data_frame(frame: dict[str, JsonValue]) -> bytes:
return b"data: " + json.dumps(frame).encode() + b"\n\n"
def _stream_without_usage(request: Request) -> Reply:
assert request.method == "POST"
chunk: Final[dict[str, JsonValue]] = {
"id": f"chatcmpl-{uuid.uuid4().hex}",
"object": "chat.completion.chunk",
"created": 1,
"model": _DEEPSEEK_V4,
}
return Reply(
content_type="text/event-stream",
chunks=(
_data_frame({**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "OK"}}]}),
_data_frame({**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}),
b"data: [DONE]\n\n",
),
)
def _streamed_prompt_tokens(gateway: Gateway, key: str, model: str, text: str) -> int:
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": text}], "stream": True},
key=key,
)
assert response.status_code == 200, response.text
assert response.text.rstrip().endswith("data: [DONE]"), response.text
rows: Final = eventually(
lambda: read_rows(
'SELECT status, prompt_tokens FROM "LiteLLM_SpendLogs" WHERE api_key=%s',
(sha256(key.encode()).hexdigest(),),
),
lambda values: len(values) == 1,
seconds=70,
)
assert rows[0]["status"] == "success", rows
prompt_tokens: Final = rows[0]["prompt_tokens"]
assert isinstance(prompt_tokens, int), rows
return prompt_tokens
def test_a_deepseek_v4_stream_without_usage_bills_prompt_tokens_counted_with_deepseeks_tokenizer(
gateway: Gateway,
) -> None:
with wire_server(_stream_without_usage) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=f"azure_ai/{_DEEPSEEK_V4}", api_base=wire.url, num_retries=0)
chinese: Final = _streamed_prompt_tokens(gateway, scenario.key(models=[model]), model, CHINESE_TEXT)
short: Final = _streamed_prompt_tokens(gateway, scenario.key(models=[model]), model, SHORT_TEXT)
assert len(wire.drain()) == 2
assert_counted_with_deepseek(chinese - short, chinese_minus_short)

562
tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py
View file

@ -1,7 +1,7 @@
import asyncio
import os
import traceback
from collections.abc import Mapping
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Final
@ -30,6 +30,7 @@ from litellm_enterprise.proxy.management_endpoints.project_endpoints import (
delete_project,
project_info,
)
from litellm.proxy.management_helpers.model_quota_policy import ModelRoute, TeamDeployment, router_model_routes
from litellm.proxy.proxy_server import (
LitellmUserRoles,
)
@ -1205,7 +1206,9 @@ def test_update_quota_existing_quotaless_model_rejected():
_raise_on_missing_project_model_quota_on_update(data, existing)
def _enforced_new_project_mocks(monkeypatch, team_models: list[str], llm_router: mock.MagicMock | None) -> None:
def _enforced_new_project_mocks(
monkeypatch: pytest.MonkeyPatch, team_models: list[str], llm_router: object, quota_setting: object = True
) -> None:
from litellm.proxy._types import LiteLLM_TeamTable
from litellm_enterprise.proxy.management_endpoints import project_endpoints as pe
@ -1213,7 +1216,7 @@ def _enforced_new_project_mocks(monkeypatch, team_models: list[str], llm_router:
monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", mock.MagicMock())
monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True)
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", llm_router)
monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {"enforce_project_model_quota": True})
monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {"enforce_project_model_quota": quota_setting})
monkeypatch.setattr(pe, "_validate_team_exists", mock.AsyncMock(return_value=team))
monkeypatch.setattr(pe, "_check_user_permission_for_project", mock.AsyncMock(return_value=True))
@ -1604,6 +1607,559 @@ async def test_new_project_flag_on_access_group_model_returns_400(monkeypatch):
assert "expand to multiple models at request time" in str(exc_info.value)
_CUSTOMER_QUOTA_POLICY: Final = {
"providers": {"vertex_ai": ["rpm", "tpm"], "runwayml": []},
"models": {"veo-*": ["rpm"], "jev-latest": []},
}
_CUSTOMER_MODEL_PROVIDERS: Final = {
"gemini-3.8-flash": frozenset({"vertex_ai"}),
"veo-3.1-lite-generate-001": frozenset({"vertex_ai"}),
"runway-gen4-turbo": frozenset({"runwayml"}),
}
def _customer_route_of(model: str) -> ModelRoute:
return ModelRoute(providers=_CUSTOMER_MODEL_PROVIDERS.get(model, frozenset()))
def _quota_error(
models: list[str],
rpm: dict[str, int] | None = None,
tpm: dict[str, int] | None = None,
setting: object = _CUSTOMER_QUOTA_POLICY,
route_of: Callable[[str], ModelRoute] = _customer_route_of,
) -> str | None:
from litellm_enterprise.proxy.management_endpoints.project_endpoints import (
_raise_on_missing_project_model_quota,
)
from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting
policy: Final = parse_model_quota_setting(setting)
assert policy is not None
data: Final = NewProjectRequest(team_id="test-team", models=models, model_rpm_limit=rpm, model_tpm_limit=tpm)
try:
_raise_on_missing_project_model_quota(data, frozenset(), policy, route_of)
except HTTPException as e:
rejection: Final = e
else:
return None
assert rejection.status_code == 400
return str(rejection.detail)
def test_quota_policy_exempts_provider_and_passthrough_name_but_keeps_vertex_strict():
assert _quota_error(["runway-gen4-turbo", "jev-latest"]) is None
error: Final = _quota_error(["gemini-3.8-flash"], rpm={"gemini-3.8-flash": 10})
assert error is not None
assert "gemini-3.8-flash (model_tpm_limit)" in error
def test_quota_policy_model_pattern_overrides_its_provider_per_limit():
assert _quota_error(["veo-3.1-lite-generate-001"], rpm={"veo-3.1-lite-generate-001": 5}) is None
error: Final = _quota_error(["veo-3.1-lite-generate-001"], tpm={"veo-3.1-lite-generate-001": 5000})
assert error is not None
assert "veo-3.1-lite-generate-001 (model_rpm_limit)" in error
assert "model_tpm_limit" not in error
@pytest.mark.parametrize(
("model", "rpm", "tpm", "expected_missing"),
[
("runway-gen4-turbo", {"runway-gen4-turbo": 0}, None, "runway-gen4-turbo (model_rpm_limit)"),
("jev-latest", None, {"jev-latest": -5}, "jev-latest (model_tpm_limit)"),
(
"veo-3.1-lite-generate-001",
{"veo-3.1-lite-generate-001": 5},
{"veo-3.1-lite-generate-001": 0},
"veo-3.1-lite-generate-001 (model_tpm_limit)",
),
],
)
def test_quota_policy_rejects_a_non_positive_limit_set_on_an_exempt_model(
model: str, rpm: dict[str, int] | None, tpm: dict[str, int] | None, expected_missing: str
):
error: Final = _quota_error([model], rpm=rpm, tpm=tpm)
assert error is not None
assert expected_missing in error
def test_quota_policy_error_names_each_model_with_the_limits_it_misses():
error: Final = _quota_error(["gemini-3.8-flash", "veo-3.1-lite-generate-001", "runway-gen4-turbo"])
assert error is not None
assert "gemini-3.8-flash (model_rpm_limit and model_tpm_limit)" in error
assert "veo-3.1-lite-generate-001 (model_rpm_limit)" in error
assert "runway-gen4-turbo" not in error
@pytest.mark.parametrize(
"setting, model, expected_missing",
[
({"default": ["tpm"]}, "unrouted-name", "unrouted-name (model_tpm_limit)"),
({"default": [], "providers": {"runwayml": ["rpm"]}}, "runway-gen4-turbo", "(model_rpm_limit)"),
(
{"providers": {"runwayml": []}, "models": {"runway-gen4-turbo": ["tpm"]}},
"runway-gen4-turbo",
"(model_tpm_limit)",
),
({"models": {"runway-*": [], "runway-gen4-turbo": ["rpm"]}}, "runway-gen4-turbo", "(model_rpm_limit)"),
({"models": {"veo-*": [], "veo-3.1-lite-*": ["rpm"]}}, "veo-3.1-lite-generate-001", "(model_rpm_limit)"),
({"models": {"gpt-5-mini*": [], "gpt-*-*-*": ["rpm"]}}, "gpt-5-mini-fast", "(model_rpm_limit)"),
],
)
def test_quota_policy_precedence_is_model_then_pattern_then_provider_then_default(
setting: object, model: str, expected_missing: str
):
error: Final = _quota_error([model], setting=setting)
assert error is not None
assert expected_missing in error
def test_quota_policy_model_entry_matches_names_literally_apart_from_star():
setting: Final = {"models": {"claude-opus-5-5[1m]": [], "gpt.5": []}}
assert _quota_error(["claude-opus-5-5[1m]"], setting=setting) is None
assert _quota_error(["claude-opus-5-51"], setting=setting) is not None
assert _quota_error(["gptx5"], setting=setting) is not None
def test_quota_policy_model_on_several_providers_needs_every_required_limit():
def mixed_providers(_model: str) -> ModelRoute:
return ModelRoute(providers=frozenset({"runwayml", "vertex_ai"}))
error: Final = _quota_error(["video-pool"], rpm={"video-pool": 10}, route_of=mixed_providers)
assert error is not None
assert "video-pool (model_tpm_limit)" in error
def test_quota_policy_model_pattern_matches_from_the_start_like_the_router():
setting: Final = {"providers": {"vertex_ai": ["rpm", "tpm"]}, "models": {"*-lite": []}}
assert _quota_error(["veo-3.1-lite-generate-001"], setting=setting) is None
assert _quota_error(["gemini-3.8-flash"], setting=setting) is not None
def _team_vertex_route(called_model: str) -> Callable[[str], ModelRoute]:
def route_of(_model: str) -> ModelRoute:
return ModelRoute(team_deployments=frozenset({TeamDeployment(provider="vertex_ai", called_model=called_model)}))
return route_of
def test_quota_policy_matches_a_team_model_by_the_model_it_calls():
borrowed_name: Final = _quota_error(["jev-latest"], route_of=_team_vertex_route("gemini-3.8-flash"))
veo_call: Final = _quota_error(
["video-team"], rpm={"video-team": 5}, route_of=_team_vertex_route("veo-3.1-generate-001")
)
assert borrowed_name is not None and "jev-latest (model_rpm_limit and model_tpm_limit)" in borrowed_name
assert veo_call is None
def test_quota_policy_team_model_beside_a_shared_one_needs_both_routes_limits():
def team_vertex_beside_shared_runway(_model: str) -> ModelRoute:
return ModelRoute(
providers=frozenset({"runwayml"}),
team_deployments=frozenset({TeamDeployment(provider="vertex_ai", called_model="gemini-3.8-flash")}),
)
error: Final = _quota_error(["media-latest"], rpm={"media-latest": 5}, route_of=team_vertex_beside_shared_runway)
assert error is not None and "media-latest (model_tpm_limit)" in error
@pytest.mark.parametrize("setting", [True, "true", "false", 1, ["rpm"], {"enabled": True}, {"provider": {"runwayml": []}}])
def test_quota_setting_truthy_non_policy_value_still_requires_both_limits(setting: object):
error: Final = _quota_error(["runway-gen4-turbo"], setting=setting)
assert error is not None
assert "runway-gen4-turbo (model_rpm_limit and model_tpm_limit)" in error
@pytest.mark.parametrize("setting", [False, None, {}, "", 0, []])
def test_quota_setting_off_enforces_nothing(setting: object):
from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting
assert parse_model_quota_setting(setting) is None
def test_update_quota_policy_resolves_providers_for_the_resulting_models():
import types
from litellm_enterprise.proxy.management_endpoints.project_endpoints import (
_raise_on_missing_project_model_quota_on_update,
)
from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting
policy: Final = parse_model_quota_setting(_CUSTOMER_QUOTA_POLICY)
assert policy is not None
existing: Final = types.SimpleNamespace(models=["runway-gen4-turbo"], metadata={})
assert (
_raise_on_missing_project_model_quota_on_update(
UpdateProjectRequest(project_id="p", description="unrelated change"),
existing,
frozenset(),
policy,
_customer_route_of,
)
is None
)
with pytest.raises(HTTPException) as exc_info:
_raise_on_missing_project_model_quota_on_update(
UpdateProjectRequest(project_id="p", models=["runway-gen4-turbo", "gemini-3.8-flash"]),
existing,
frozenset(),
policy,
_customer_route_of,
)
assert "gemini-3.8-flash (model_rpm_limit and model_tpm_limit)" in str(exc_info.value.detail)
def _customer_router() -> "litellm.Router":
return litellm.Router(
model_list=[
{
"model_name": "gemini-3.8-flash",
"litellm_params": {"model": "vertex_ai/gemini-3.8-flash", "vertex_project": "qa-project"},
},
{
"model_name": "veo-3.1-lite-generate-001",
"litellm_params": {"model": "vertex_ai/veo-3.1-lite-generate-001", "vertex_project": "qa-project"},
},
{
"model_name": "runway-gen4-turbo",
"litellm_params": {"model": "runwayml/gen4_turbo", "api_key": "qa-runway-key"},
},
]
)
def test_quota_providers_come_from_the_one_wildcard_route_requests_use():
router: Final = litellm.Router(
model_list=[
{"model_name": "gen4*", "litellm_params": {"model": "runwayml/gen4*", "api_key": "qa-runway-key"}},
{
"model_name": "media-team-gen",
"litellm_params": {"model": "vertex_ai/*", "vertex_project": "qa-project"},
"model_info": {"team_id": "media-team", "team_public_model_name": "gen*"},
},
]
)
assert router_model_routes(router, "media-team")("gen4_turbo") == ModelRoute(
team_deployments=frozenset({TeamDeployment(provider="vertex_ai", called_model="4_turbo")})
)
assert router_model_routes(router, None)("gen4_turbo") == ModelRoute(providers=frozenset({"runwayml"}))
def test_quota_providers_read_the_provider_behind_a_known_api_base():
groq_api_base: Final = "https://api.groq.com/openai/v1"
router: Final = litellm.Router(
model_list=[
{
"model_name": "qa-llama",
"litellm_params": {"model": "qa-llama", "api_base": groq_api_base, "api_key": "qa-groq-key"},
},
{
"model_name": "qa-runway",
"litellm_params": {"model": "gen4_turbo", "custom_llm_provider": "runwayml", "api_base": groq_api_base},
},
]
)
route_of: Final = router_model_routes(router, None)
assert route_of("qa-llama").providers == frozenset({"groq"})
assert route_of("qa-runway").providers == frozenset({"runwayml"})
def test_quota_policy_rejects_an_unknown_provider_id():
from pydantic import ValidationError
from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting
with pytest.raises(ValidationError, match=r"unknown provider ids \['runway'\]"):
parse_model_quota_setting({"providers": {"vertex_ai": ["rpm"], "runway": []}})
def test_quota_policy_accepts_a_provider_registered_at_runtime(monkeypatch: pytest.MonkeyPatch):
from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting
monkeypatch.setattr(litellm, "provider_list", [*litellm.provider_list, "qa-custom-llm"])
policy: Final = parse_model_quota_setting({"providers": {"qa-custom-llm": []}})
assert policy is not None
assert policy.providers["qa-custom-llm"] == frozenset()
_TEAM_WRITE_POLICY: Final = {"providers": {"vertex_ai": ["rpm", "tpm"], "runwayml": []}, "models": {"jev-latest": []}}
_VERTEX_PARAMS: Final = {"model": "vertex_ai/gemini-3.8-flash", "vertex_project": "qa-project"}
_RUNWAY_PARAMS: Final = {"model": "gen4_turbo", "custom_llm_provider": "runwayml"}
def _team_deployment(public_name: str, model_id: str, litellm_params: Mapping[str, object]) -> dict[str, object]:
return {
"model_name": f"model_name_media-team_{model_id}",
"litellm_params": dict(litellm_params),
"model_info": {"id": model_id, "team_id": "media-team", "team_public_model_name": public_name},
}
@pytest.mark.parametrize(
("setting", "public_name", "proxy_params", "team_routes", "replaced_id", "written_params", "added"),
[
pytest.param(
_TEAM_WRITE_POLICY,
"jev-latest",
{"model": "openai/jev-latest"},
(),
None,
_VERTEX_PARAMS,
{"rpm", "tpm"},
id="team-add-on-vertex-behind-an-exempt-proxy-name",
),
pytest.param(
_TEAM_WRITE_POLICY,
"jev-latest",
{"model": "openai/jev-latest"},
(),
None,
_RUNWAY_PARAMS,
set(),
id="team-add-on-an-exempt-provider",
),
pytest.param(
_TEAM_WRITE_POLICY,
"jev-latest",
{"model": "openai/jev-latest"},
(),
None,
{"model": "aihubmix/jev-latest"},
set(),
id="team-add-through-a-json-provider-matches-the-model-it-calls",
),
pytest.param(
_TEAM_WRITE_POLICY,
"media-latest",
_VERTEX_PARAMS,
(("t1", _RUNWAY_PARAMS),),
"t1",
None,
{"rpm", "tpm"},
id="team-delete-falls-back-to-the-proxy-vertex-route",
),
pytest.param(
_TEAM_WRITE_POLICY,
"media-latest",
_VERTEX_PARAMS,
(("t1", _RUNWAY_PARAMS), ("t2", _RUNWAY_PARAMS)),
"t1",
None,
set(),
id="team-delete-keeps-a-sibling-team-route",
),
pytest.param(
_TEAM_WRITE_POLICY,
"media-latest",
_VERTEX_PARAMS,
(("t1", _RUNWAY_PARAMS),),
"t1",
_VERTEX_PARAMS,
{"rpm", "tpm"},
id="team-update-moves-the-route-to-vertex",
),
pytest.param(
True,
"jev-latest",
{"model": "openai/jev-latest"},
(),
None,
_VERTEX_PARAMS,
set(),
id="true-already-requires-both-limits",
),
],
)
def test_team_model_write_adds_only_the_limits_its_new_route_requires(
setting: object,
public_name: str,
proxy_params: Mapping[str, object],
team_routes: tuple[tuple[str, Mapping[str, object]], ...],
replaced_id: str | None,
written_params: Mapping[str, object] | None,
added: set[str],
):
from litellm.proxy.management_helpers.model_quota_policy import (
limits_added_by_team_write,
parse_model_quota_setting,
)
policy: Final = parse_model_quota_setting(setting)
router: Final = litellm.Router(
model_list=[
{"model_name": public_name, "litellm_params": dict(proxy_params)},
*(_team_deployment(public_name, model_id, params) for model_id, params in team_routes),
]
)
assert policy is not None
assert limits_added_by_team_write(
policy, router, "media-team", public_name, replaced_id, written_params, public_name
) == frozenset(added)
def test_quota_route_of_an_alias_reads_the_team_deployments_its_target_routes_to():
router: Final = litellm.Router(
model_list=[
{"model_name": "media-latest", "litellm_params": dict(_RUNWAY_PARAMS)},
_team_deployment("media-latest", "t1", _VERTEX_PARAMS),
],
model_group_alias={"media": "media-latest"},
)
route_of: Final = router_model_routes(router, "media-team")
assert route_of("media") == route_of("media-latest")
assert TeamDeployment(provider="vertex_ai", called_model="gemini-3.8-flash") in route_of("media").team_deployments
def _team_pattern_router() -> "litellm.Router":
return litellm.Router(
model_list=[
{"model_name": "media-latest", "litellm_params": dict(_VERTEX_PARAMS)},
{"model_name": "video-named", "litellm_params": dict(_VERTEX_PARAMS)},
_team_deployment("video-*", "t1", {"model": "runwayml/*"}),
_team_deployment("video-veo-*", "t2", {"model": "vertex_ai/veo-*", "vertex_project": "qa-project"}),
_team_deployment("clip-*x", "t3", {"model": "runwayml/*"}),
_team_deployment("clip-x*", "t4", {"model": "vertex_ai/*", "vertex_project": "qa-project"}),
],
model_group_alias={"media": "media-latest"},
)
@pytest.mark.parametrize(
("public_name", "model", "served"),
[
pytest.param("media-latest", "media-latest", True, id="the-name-itself"),
pytest.param("media-latest", "media", True, id="an-alias-of-the-name"),
pytest.param("media-latest", "video-named", False, id="another-name"),
pytest.param("video-*", "video-gen-1", True, id="a-name-the-pattern-matches"),
pytest.param("video-*", "video-veo-3", False, id="a-more-specific-team-pattern-takes-it"),
pytest.param("video-veo-*", "video-veo-3", True, id="the-most-specific-team-pattern"),
pytest.param("video-*", "video-named", False, id="a-named-deployment-shadows-the-pattern"),
pytest.param("video-*", "image-gen-1", False, id="a-name-the-pattern-misses"),
pytest.param("audio-*", "audio-gen-1", True, id="a-pattern-the-write-adds"),
pytest.param("clip-*x", "clip-xx", True, id="the-first-registered-of-two-tied-team-patterns"),
pytest.param("clip-x*", "clip-xx", False, id="a-tied-team-pattern-registered-first-takes-it"),
pytest.param("cli*-xx", "clip-xx", False, id="a-new-pattern-tied-with-a-registered-one"),
],
)
def test_team_write_serves_the_project_names_a_team_request_routes_through_it(
public_name: str, model: str, served: bool
):
from litellm.proxy.management_helpers.model_quota_policy import team_write_serves
assert team_write_serves(_team_pattern_router(), "media-team", public_name, model) is served
@pytest.mark.parametrize(
("model", "added"),
[
pytest.param("video-gen-1", {"rpm", "tpm"}, id="vertex-requires-both"),
pytest.param("video-veo-3", {"rpm"}, id="the-called-veo-model-entry-wins"),
],
)
def test_team_pattern_write_reads_each_served_name_by_the_model_it_calls(model: str, added: set[str]):
from litellm.proxy.management_helpers.model_quota_policy import (
limits_added_by_team_write,
parse_model_quota_setting,
)
policy: Final = parse_model_quota_setting(
{"default": [], "providers": {"vertex_ai": ["rpm", "tpm"]}, "models": {"veo-*": ["rpm"]}}
)
assert policy is not None
assert limits_added_by_team_write(
policy,
litellm.Router(model_list=[]),
"media-team",
"video-*",
None,
{"model": "vertex_ai/*", "vertex_project": "qa-project"},
model,
) == frozenset(added)
@pytest.mark.asyncio
async def test_new_project_quota_policy_reads_providers_from_router_deployments(monkeypatch: pytest.MonkeyPatch):
team_models: Final = ["gemini-3.8-flash", "veo-3.1-lite-generate-001", "runway-gen4-turbo", "jev-latest"]
_enforced_new_project_mocks(
monkeypatch, team_models=team_models, llm_router=_customer_router(), quota_setting=_CUSTOMER_QUOTA_POLICY
)
with pytest.raises(ProxyException) as exc_info:
await _run_new_project(NewProjectRequest(team_id="test-team", models=team_models))
assert exc_info.value.code == "400"
assert (
"gemini-3.8-flash (model_rpm_limit and model_tpm_limit), veo-3.1-lite-generate-001 (model_rpm_limit)."
in str(exc_info.value.message)
)
assert "runway-gen4-turbo" not in str(exc_info.value.message)
assert "jev-latest" not in str(exc_info.value.message)
@pytest.mark.asyncio
async def test_new_project_rejects_malformed_quota_setting_instead_of_skipping_it(monkeypatch: pytest.MonkeyPatch):
_enforced_new_project_mocks(
monkeypatch,
team_models=["runway-gen4-turbo"],
llm_router=_customer_router(),
quota_setting={"providers": {"runwayml": []}, "model": {"jev-latest": []}},
)
with pytest.raises(ProxyException) as exc_info:
await _run_new_project(NewProjectRequest(team_id="test-team", models=["runway-gen4-turbo"]))
assert exc_info.value.code == "500"
assert "general_settings.enforce_project_model_quota is invalid" in str(exc_info.value.message)
@pytest.mark.asyncio
async def test_new_project_checks_permission_before_reading_the_quota_setting(monkeypatch: pytest.MonkeyPatch):
from litellm_enterprise.proxy.management_endpoints import project_endpoints as pe
_enforced_new_project_mocks(
monkeypatch,
team_models=["runway-gen4-turbo"],
llm_router=_customer_router(),
quota_setting={"providers": {"runwayml": []}, "model": {"jev-latest": []}},
)
monkeypatch.setattr(pe, "_check_user_permission_for_project", mock.AsyncMock(return_value=False))
with pytest.raises(ProxyException) as exc_info:
await _run_new_project(NewProjectRequest(team_id="test-team", models=["runway-gen4-turbo"]))
assert exc_info.value.code == "403"
assert "enforce_project_model_quota" not in str(exc_info.value.message)
@pytest.mark.asyncio
async def test_new_project_quota_policy_ignores_model_entries_for_a_team_model_name(monkeypatch: pytest.MonkeyPatch):
router: Final = litellm.Router(
model_list=[
{
"model_name": "model_name_test-team_qa",
"litellm_params": {"model": "vertex_ai/gemini-3.8-flash", "vertex_project": "qa-project"},
"model_info": {"team_id": "test-team", "team_public_model_name": "jev-latest"},
},
]
)
_enforced_new_project_mocks(
monkeypatch, team_models=["jev-latest"], llm_router=router, quota_setting=_CUSTOMER_QUOTA_POLICY
)
with pytest.raises(ProxyException) as exc_info:
await _run_new_project(NewProjectRequest(team_id="test-team", models=["jev-latest"]))
assert exc_info.value.code == "400"
assert "jev-latest (model_rpm_limit and model_tpm_limit)" in str(exc_info.value.message)
_ACCESS_NOW: Final = datetime.now(timezone.utc)

138
tests/unit/litellm_core_utils/test_fallback_generalizations.py
View file

@ -7,7 +7,11 @@ its end-to-end wiring into provider routing (get_llm_provider) and model-info
resolution (get_model_info) including the shipped rules in the bundled cost map.
"""
import json
import logging
from collections.abc import Callable
from pathlib import Path
from typing import Final
import pytest
@ -21,6 +25,8 @@ from litellm.litellm_core_utils.fallback_generalizations import (
set_fallback_generalizations,
)
REPO_ROOT: Final = Path(__file__).parents[3]
@pytest.fixture
def restore_generalizations():
@ -148,7 +154,7 @@ def test_fill_missing_requires_per_rule_opt_in(restore_generalizations):
}
]
)
assert match_fill_missing_generalizations("acme-1", "openai") == {"supports_vision": True}
assert match_fill_missing_generalizations("acme-1", "openai") is None
restore_generalizations(
[
@ -203,89 +209,13 @@ def test_reinstalling_rules_replaces_compiled_rules(restore_generalizations):
# --------------------------------------------------------------------------- #
# Engine: install-time validation and legacy-schema shim
# Engine: install-time validation
# --------------------------------------------------------------------------- #
def test_legacy_mixed_rule_acts_as_both_kinds(restore_generalizations):
"""A legacy rule mixing ``litellm_provider`` with capability keys routes AND
contributes its full model_info (provider included) to the capability union."""
restore_generalizations(
[
{
"name": "legacy-mixed",
"pattern": r"^acme-",
"model_info": {"litellm_provider": "anthropic", "supports_vision": True},
},
{"name": "new-caps", "pattern": r"^acme-pro-", "model_info": {"supports_reasoning": True}},
]
)
assert match_routing_generalization("acme-pro-1") == "anthropic"
assert match_capability_generalizations("acme-pro-1") == {
"litellm_provider": "anthropic",
"supports_vision": True,
"supports_reasoning": True,
}
LEGACY_MAIN_RULES = [
{
"name": "anthropic-claude-adaptive-thinking",
"pattern": "(?:opus|sonnet|haiku)[-._](?:4[-._](?:[6-9]|[1-9]\\d)(?!\\d)|(?:[5-9]|[1-9]\\d{1,})[-._]\\d{1,2}(?!\\d))",
"description": "Claude opus/sonnet/haiku at version 4.6 or higher: 4.6 through 4.99, then any 5.x, 6.x or later major. The minor is capped at two digits so an 8-digit date suffix such as claude-opus-4-20250514 is never read as a >= 4.6 minor. Turns on adaptive thinking for new families with no code change.",
"extends": "anthropic-claude",
"model_info": {"supports_adaptive_thinking": True},
},
{
"name": "anthropic-claude",
"pattern": "^claude-[a-z]+-\\d+[-.]\\d+(?:-\\d{8})?$",
"description": "Any Claude family-major-minor id, optionally with an 8-digit date suffix, anchored to the whole name. Version-neutral fallback that gives an unmapped Claude provider routing and baseline capabilities; it carries no pricing, so cost stays on the standard unpriced behavior rather than a guessed number.",
"model_info": {
"litellm_provider": "anthropic",
"mode": "chat",
"max_input_tokens": 200000,
"max_output_tokens": 64000,
"max_tokens": 64000,
"supports_function_calling": True,
"supports_parallel_function_calling": True,
"supports_vision": True,
"supports_tool_choice": True,
"supports_assistant_prefill": True,
"supports_prompt_caching": True,
"supports_response_schema": True,
"supports_reasoning": True,
"supports_pdf_input": True,
"supports_system_messages": True,
},
},
]
def test_legacy_main_schema_keeps_unmapped_claude_working(restore_generalizations):
"""Pins the remote-map transition window: a released proxy running this engine
against main's old-schema block (mixed provider+capability rule plus ``extends``,
copied verbatim above) must keep unmapped-Claude inference and info resolution
working until the new-schema JSON reaches main."""
restore_generalizations([dict(rule) for rule in LEGACY_MAIN_RULES])
litellm.get_model_info.cache_clear()
_, provider, _, _ = litellm.get_llm_provider(model="claude-opus-9-9")
assert provider == "anthropic"
info = litellm.get_model_info("claude-opus-9-9")
assert info["litellm_provider"] == "anthropic"
assert info["supports_adaptive_thinking"] is True
assert info["supports_function_calling"] is True
assert info["max_input_tokens"] == 200000
assert not info.get("input_cost_per_token")
low = litellm.get_model_info("claude-opus-4-0")
assert low["litellm_provider"] == "anthropic"
assert low["supports_function_calling"] is True
assert low.get("supports_adaptive_thinking") is None
def test_non_string_provider_rule_warns_and_is_skipped(restore_generalizations, warning_messages):
def test_non_string_provider_rule_warns_and_is_skipped(
restore_generalizations: Callable[[list | None], None], warning_messages: list[str]
) -> None:
restore_generalizations([{"name": "bad-provider", "pattern": r"^acme-", "model_info": {"litellm_provider": 42}}])
assert any("bad-provider" in message for message in warning_messages)
assert match_routing_generalization("acme-1") is None
@ -306,6 +236,40 @@ def test_malformed_rules_are_skipped_not_fatal(restore_generalizations):
assert match_capability_generalizations("claude-opus-9-9") == {"good": True}
@pytest.mark.parametrize(
"rules",
[
[
{
"name": "stale-rule",
"pattern": r"^acme-",
"model_info": {"litellm_provider": "anthropic", "supports_vision": True},
}
],
[
{"name": "parent", "pattern": r"^zeta-only$", "model_info": {"supports_vision": True}},
{
"name": "stale-rule",
"pattern": r"^acme-",
"extends": "parent",
"model_info": {"supports_reasoning": True},
},
],
],
ids=["provider-mixed-with-capabilities", "extends"],
)
def test_old_format_rule_is_skipped_with_a_warning_pointing_at_a_current_cost_map(
restore_generalizations: Callable[[list | None], None], warning_messages: list[str], rules: list[dict]
) -> None:
restore_generalizations(rules)
assert any(
"stale-rule" in m and "LITELLM_MODEL_COST_MAP_URL" in m and "LITELLM_LOCAL_MODEL_COST_MAP " in m
for m in warning_messages
)
assert match_routing_generalization("acme-1") is None
assert match_capability_generalizations("acme-1") is None
# --------------------------------------------------------------------------- #
# End-to-end: provider routing + model-info resolution
# --------------------------------------------------------------------------- #
@ -435,6 +399,18 @@ def test_exact_entries_fill_only_missing_fields(restore_generalizations, monkeyp
# --------------------------------------------------------------------------- #
@pytest.mark.parametrize(
"cost_map_path",
["model_prices_and_context_window.json", "litellm/model_prices_and_context_window_backup.json"],
)
def test_every_shipped_rule_installs_without_a_skip_warning(
restore_generalizations: Callable[[list | None], None], warning_messages: list[str], cost_map_path: str
) -> None:
rules: Final = json.loads((REPO_ROOT / cost_map_path).read_text())["fallback_generalizations"]["rules"]
restore_generalizations(rules)
assert warning_messages == []
@pytest.fixture
def shipped_cost_map(monkeypatch):
"""Activate the bundled cost map so the shipped rules are installed."""

13
tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py
View file

@ -2378,10 +2378,8 @@ def test_anthropic_model_supports_effort_param_rejects_non_supporting_models(mod
"model",
[
"claude-opus-4-6",
"claude-opus-4-7",
"claude-opus-4-8",
"claude-opus-4-6-20260205",
"claude-opus-4-7-20260416",
],
)
def test_anthropic_model_supports_speed_param_recognizes_supporting_models(model):
@ -2397,6 +2395,10 @@ def test_anthropic_model_supports_speed_param_recognizes_supporting_models(model
"vertex_ai/claude-opus-4-8",
"azure_ai/claude-opus-4-8",
"anthropic.claude-opus-4-8",
# Anthropic removed Opus 4.7 fast mode on 2026-07-24 and answers 400 for `speed` on it
# (code.claude.com/docs/en/fast-mode, platform.claude.com/docs/en/build-with-claude/fast-mode, read 2026-10-28)
"claude-opus-4-7",
"claude-opus-4-7-20260416",
],
)
def test_anthropic_model_supports_speed_param_rejects_non_supporting_models(model):
@ -4677,8 +4679,9 @@ def test_fast_mode_parameter_mapping():
assert result["speed"] == "fast"
def test_anthropic_drop_params_strips_speed_for_unsupported_models():
"""``drop_params=True`` strips unsupported ``speed`` for non-Opus models."""
@pytest.mark.parametrize("model", ["claude-sonnet-4-6", "claude-opus-4-7"])
def test_anthropic_drop_params_strips_speed_for_unsupported_models(model: str):
"""``drop_params=True`` strips ``speed`` for models without fast mode."""
config = AnthropicConfig()
messages = [{"role": "user", "content": "Hello"}]
@ -4686,7 +4689,7 @@ def test_anthropic_drop_params_strips_speed_for_unsupported_models():
litellm.drop_params = True
try:
result = config.transform_request(
model="claude-sonnet-4-6",
model=model,
messages=messages,
optional_params={"speed": "fast", "max_tokens": 1024},
litellm_params={},

39
tests/unit/llms/anthropic/test_count_tokens_oauth.py
View file

@ -221,6 +221,45 @@ class TestCountTokensUsesWorkloadIdentity:
assert result.total_tokens == 13
assert route.calls.last.request.headers["x-api-key"] == FAKE_REGULAR_KEY
@pytest.mark.asyncio
@pytest.mark.parametrize(
("deployment_key", "expected_credentials"),
[
(None, {"authorization": f"Bearer {FAKE_OAUTH_TOKEN}", "x-api-key": None}),
(FAKE_REGULAR_KEY, {"authorization": None, "x-api-key": FAKE_REGULAR_KEY}),
],
ids=["keyless-deployment-uses-the-callers-token", "deployment-key-wins-over-the-callers-token"],
)
async def test_a_callers_subscription_token_counts_only_when_no_key_is_configured(
self,
monkeypatch: pytest.MonkeyPatch,
httpx_transport_clients: None,
deployment_key: str | None,
expected_credentials: dict[str, str | None],
) -> None:
for name in ("ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"):
monkeypatch.delenv(name, raising=False)
monkeypatch.setattr(litellm, "anthropic_key", None)
monkeypatch.setattr(litellm, "api_key", None)
with respx.mock(assert_all_called=True) as router:
route = router.post("https://api.anthropic.com/v1/messages/count_tokens").mock(
return_value=httpx.Response(200, json={"input_tokens": 17})
)
result = await AnthropicTokenCounter(client_oauth_token=FAKE_OAUTH_TOKEN).count_tokens(
model_to_use="claude-sonnet-4-5",
messages=[{"role": "user", "content": "hi"}],
contents=None,
deployment={"litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": deployment_key}},
request_model="claude-sonnet-4-5",
)
assert result is not None
assert result.total_tokens == 17
sent = route.calls.last.request.headers
assert {name: sent.get(name) for name in expected_credentials} == expected_credentials
assert (ANTHROPIC_OAUTH_BETA_HEADER in sent["anthropic-beta"].split(",")) is (deployment_key is None)
@pytest.mark.asyncio
async def test_a_failed_mint_degrades_like_an_anthropic_error(self, monkeypatch):
from litellm.llms.anthropic.count_tokens import token_counter as token_counter_module

605
tests/unit/llms/base_llm/responses/test_transformation.py
View file

@ -1,14 +1,17 @@
"""The shared Responses API config contract."""
import asyncio, importlib, json
from collections.abc import Mapping
import httpx
import pytest
import respx
import litellm
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
from litellm.types.router import GenericLiteLLMParams
from litellm import Router
from pydantic import BaseModel, ConfigDict
from litellm.constants import STREAM_SSE_DONE_STRING
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
@ -879,38 +882,6 @@ def test_combine_responses_fallback_usage_passthrough_for_unknown_event():
partial = ResponseAPIUsage(input_tokens=1, output_tokens=1, total_tokens=2)
Router._combine_responses_fallback_usage(other, partial)
@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit")
def test_build_responses_continuation_input_from_string():
out = Router._build_responses_continuation_input("Hello world", "partial assistant text")
assert len(out) == 3
assert out[0]["role"] == "user"
assert out[0]["content"][0]["text"] == "Hello world"
assert out[1]["role"] == "developer"
assert out[2]["role"] == "assistant"
assert out[2]["content"][0]["text"] == "partial assistant text"
@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit")
def test_build_responses_continuation_input_from_list_preserves_items():
existing: List[Any] = [
{
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "msg1"}],
}
]
out = Router._build_responses_continuation_input(existing, "partial")
assert len(out) == 3
assert out[0]["content"][0]["text"] == "msg1"
assert out[1]["role"] == "developer"
assert out[2]["role"] == "assistant"
@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit")
def test_build_responses_continuation_input_from_none():
out = Router._build_responses_continuation_input(None, "partial")
assert len(out) == 2
assert out[0]["role"] == "developer"
assert out[1]["role"] == "assistant"
@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit")
@pytest.mark.asyncio
async def test_aresponses_streaming_iterator_passthrough():
@ -1355,9 +1326,573 @@ async def test_aresponses_fallback_uses_continuation_input_after_partial_content
continuation = mock_fallback.await_args.kwargs["kwargs"]["input"]
assert isinstance(continuation, list)
assert continuation[0]["content"][0]["text"] == "original question"
assert continuation[-2]["role"] == "developer"
assert continuation[-1]["role"] == "assistant"
assert continuation[-1]["content"][0]["text"] == "partial answer"
assert continuation[-2]["role"] == "assistant"
assert continuation[-2]["content"][0]["text"] == "partial answer"
assert continuation[-1]["role"] == "user"
def _sse(*events: Mapping[str, object]) -> httpx.Response:
return httpx.Response(
200,
headers={"content-type": "text/event-stream"},
content="".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events),
)
def _openai_stream_cut_after(partial: str) -> httpx.Response:
return _sse(
*(
({
"type": "response.output_text.delta", "sequence_number": 0,
"item_id": "msg_primary", "output_index": 0, "content_index": 0, "delta": partial,
},)
if partial else ()
),
{
"type": "error", "sequence_number": 1,
"error": {"type": "server_error", "code": "internal_error", "message": "interrupted"},
},
)
def _anthropic_stream(text: str) -> httpx.Response:
return _sse(
{"type": "message_start", "message": {
"id": "msg_fallback", "type": "message", "role": "assistant",
"model": "claude-fallback", "content": [], "usage": {"input_tokens": 10, "output_tokens": 0},
}},
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}},
{"type": "content_block_stop", "index": 0},
{"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 5}},
{"type": "message_stop"},
)
def _primary_with_fallback(fallback_model: str) -> Router:
return Router(
model_list=[
{"model_name": "primary", "litellm_params": {"model": "openai/gpt-6-astra", "api_key": "synthetic"}},
{"model_name": "fallback", "litellm_params": {"model": fallback_model, "api_key": "synthetic"}},
],
fallbacks=[{"primary": ["fallback"]}],
num_retries=0,
)
def _streamed_text(events: tuple[Any, ...]) -> str:
return "".join(event.delta for event in events if event.type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA)
@pytest.fixture
def claude_with_prefill() -> str:
model: Final = "anthropic/claude-prefill-capable-test-model"
litellm.register_model(
{model: {"litellm_provider": "anthropic", "mode": "chat", "supports_assistant_prefill": True}},
persist_across_reloads=False,
)
litellm.get_model_info.cache_clear()
return model
@pytest.mark.asyncio
@pytest.mark.parametrize("partial", ["", "1, 2, "])
async def test_responses_fallback_can_continue_on_a_provider_without_prefill(
partial: str, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
prompt: Final = "Print the integers 1 through 5, separated by commas."
primary: Final = respx_mock.post("https://api.openai.com/v1/responses").mock(
return_value=_openai_stream_cut_after(partial)
)
def no_prefill_provider(request: httpx.Request) -> httpx.Response:
body: Final = json.loads(request.content)
if body["messages"][-1]["role"] != "user":
return httpx.Response(
400,
json={"type": "error", "error": {
"type": "invalid_request_error", "message": "The conversation must end with a user message",
}},
)
return _anthropic_stream("3, 4, 5" if partial else "1, 2, 3, 4, 5")
fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(side_effect=no_prefill_provider)
router: Final = _primary_with_fallback("anthropic/claude-opus-5-5")
try:
stream: Final = await router.aresponses(model="primary", input=prompt, stream=True, max_output_tokens=100)
events: Final = tuple([event async for event in stream])
assert _streamed_text(events) == "1, 2, 3, 4, 5"
assert events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED
assert primary.call_count == fallback.call_count == 1
messages: Final = json.loads(fallback.calls.last.request.content)["messages"]
assert messages[0] == {"role": "user", "content": [{"type": "text", "text": prompt}]}
if partial:
assert [message["role"] for message in messages] == ["user", "assistant", "user"]
assert messages[-2] == {"role": "assistant", "content": [{"type": "text", "text": partial}]}
assert json.dumps(partial) in messages[-1]["content"][0]["text"]
else:
assert len(messages) == 1
finally:
router.discard()
@pytest.mark.asyncio
async def test_responses_fallback_continues_a_claude_that_takes_prefill_from_its_own_partial(
claude_with_prefill: str, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
respx_mock.post("https://api.openai.com/v1/responses").mock(return_value=_openai_stream_cut_after("1, 2, "))
fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(
return_value=_anthropic_stream(" 3, 4, 5")
)
router: Final = _primary_with_fallback(claude_with_prefill)
try:
stream: Final = await router.aresponses(
model="primary", input="Print the integers 1 through 5.", stream=True, max_output_tokens=100
)
events: Final = tuple([event async for event in stream])
assert _streamed_text(events) == "1, 2, 3, 4, 5"
assert events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED
messages: Final = json.loads(fallback.calls.last.request.content)["messages"]
assert [message["role"] for message in messages] == ["user", "assistant"]
assert messages[-1]["content"][0]["text"].startswith("1, 2,")
finally:
router.discard()
@pytest.mark.asyncio
@pytest.mark.parametrize("fallback_takes_prefill", [True, False], ids=["prefill", "request_to_continue"])
async def test_responses_fallback_continuation_drops_reasoning_the_fallback_cannot_decrypt(
fallback_takes_prefill: bool,
claude_with_prefill: str,
respx_mock: respx.MockRouter,
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
respx_mock.post("https://api.openai.com/v1/responses").mock(return_value=_openai_stream_cut_after("1, 2, "))
fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(
return_value=_anthropic_stream(" 3, 4, 5")
)
router: Final = _primary_with_fallback(claude_with_prefill if fallback_takes_prefill else "anthropic/claude-opus-5-5")
try:
stream: Final = await router.aresponses(
model="primary",
input=[
{"role": "user", "content": "Print the integers 1 through 5."},
{"type": "reasoning", "id": "rs_1", "summary": [], "encrypted_content": "primary-only-ciphertext"},
],
stream=True,
max_output_tokens=100,
)
_ = [event async for event in stream]
assert fallback.call_count == 1
assert "primary-only-ciphertext" not in fallback.calls.last.request.content.decode()
finally:
router.discard()
class _Numbers(BaseModel):
model_config = ConfigDict(frozen=True)
numbers: tuple[int, ...]
@pytest.mark.asyncio
@pytest.mark.parametrize(
"constraint",
[
{"text": {"format": {"type": "json_schema", "name": "numbers", "strict": True, "schema": {"type": "object"}}}},
{"text": {"format": {"type": "json_object"}}},
{"tool_choice": "required", "tools": [{"type": "function", "name": "lookup", "parameters": {"type": "object"}}]},
{"text_format": _Numbers},
],
ids=["json_schema", "json_object", "required_tool", "sdk_text_format"],
)
async def test_responses_fallback_surfaces_the_primary_error_once_constrained_output_was_streamed(
constraint: dict[str, object],
claude_with_prefill: str,
respx_mock: respx.MockRouter,
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
respx_mock.post("https://api.openai.com/v1/responses").mock(
return_value=_openai_stream_cut_after('{"numbers": [1, 2, ')
)
fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(
return_value=_anthropic_stream("3]}")
)
router: Final = _primary_with_fallback(claude_with_prefill)
try:
stream: Final = await router.aresponses(
model="primary", input="Numbers 1 to 3 as JSON.", stream=True, max_output_tokens=100, **constraint
)
with pytest.raises(litellm.InternalServerError, match="interrupted"):
[event async for event in stream]
assert fallback.call_count == 0
finally:
router.discard()
@pytest.mark.asyncio
async def test_responses_fallback_continues_within_the_callers_output_ceiling(
claude_with_prefill: str, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
respx_mock.post("https://api.openai.com/v1/responses").mock(return_value=_openai_stream_cut_after("1, 2, "))
fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(
return_value=_anthropic_stream("3, 4, 5")
)
router: Final = _primary_with_fallback(claude_with_prefill)
try:
stream: Final = await router.aresponses(model="primary", input="Count to 5.", stream=True, max_output_tokens=100)
events: Final = tuple([event async for event in stream])
assert _streamed_text(events) == "1, 2, 3, 4, 5"
assert json.loads(fallback.calls.last.request.content)["max_tokens"] == 100 - litellm.token_counter(
text="1, 2, "
)
finally:
router.discard()
@pytest.mark.asyncio
async def test_responses_fallback_surfaces_the_primary_error_once_the_cut_stream_used_the_output_ceiling(
claude_with_prefill: str, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
respx_mock.post("https://api.openai.com/v1/responses").mock(return_value=_openai_stream_cut_after("1, 2, "))
fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(
return_value=_anthropic_stream("3, 4, 5")
)
router: Final = _primary_with_fallback(claude_with_prefill)
try:
stream: Final = await router.aresponses(
model="primary", input="Count to 5.", stream=True, max_output_tokens=litellm.token_counter(text="1, 2, ")
)
with pytest.raises(litellm.InternalServerError, match="interrupted"):
[event async for event in stream]
assert fallback.call_count == 0
finally:
router.discard()
def _openai_stream_cut_mid_tool_call(preamble: str, *, announce_item: bool = True) -> httpx.Response:
text: Final = {
"type": "response.output_text.delta", "sequence_number": 0,
"item_id": "msg_primary", "output_index": 0, "content_index": 0, "delta": preamble,
}
item_added: Final = {
"type": "response.output_item.added", "sequence_number": 1, "output_index": 1,
"item": {
"type": "function_call", "id": "fc_primary", "call_id": "call_primary",
"name": "lookup", "arguments": "", "status": "in_progress",
},
}
arguments: Final = {
"type": "response.function_call_arguments.delta", "sequence_number": 2,
"item_id": "fc_primary", "output_index": 1, "delta": '{"city": "Par',
}
error: Final = {
"type": "error", "sequence_number": 3,
"error": {"type": "server_error", "code": "internal_error", "message": "interrupted"},
}
return _sse(*((text, item_added, arguments, error) if announce_item else (text, arguments, error)))
def _anthropic_stream_cut_mid_tool_call(preamble: str, tool_block: Mapping[str, object] | None = None) -> httpx.Response:
return _sse(
{"type": "message_start", "message": {
"id": "msg_primary", "type": "message", "role": "assistant",
"model": "claude-primary", "content": [], "usage": {"input_tokens": 10, "output_tokens": 0},
}},
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": preamble}},
{"type": "content_block_stop", "index": 0},
{"type": "content_block_start", "index": 1, "content_block": tool_block or {
"type": "tool_use", "id": "toolu_primary", "name": "lookup", "input": {},
}},
{"type": "content_block_delta", "index": 1, "delta": {"type": "input_json_delta", "partial_json": '{"city": "Par'}},
{"type": "error", "error": {"type": "overloaded_error", "message": "interrupted"}},
)
@pytest.mark.asyncio
@pytest.mark.parametrize("announce_item", [True, False], ids=["announced_item", "arguments_only"])
async def test_responses_fallback_surfaces_the_primary_error_once_the_cut_stream_began_a_tool_call(
announce_item: bool, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
respx_mock.post("https://api.openai.com/v1/responses").mock(
return_value=_openai_stream_cut_mid_tool_call("Let me look that up. ", announce_item=announce_item)
)
fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(
return_value=_anthropic_stream("It is sunny in Paris.")
)
router: Final = _primary_with_fallback("anthropic/claude-opus-5-5")
try:
stream: Final = await router.aresponses(
model="primary",
input="What is the weather in Paris?",
stream=True,
tools=[{"type": "function", "name": "lookup", "parameters": {"type": "object"}}],
)
with pytest.raises(litellm.InternalServerError, match="interrupted"):
[event async for event in stream]
assert fallback.call_count == 0
finally:
router.discard()
def _openai_stream_cut_after_event(preamble: str, event: Mapping[str, object]) -> httpx.Response:
return _sse(
{
"type": "response.output_text.delta", "sequence_number": 0,
"item_id": "msg_primary", "output_index": 0, "content_index": 0, "delta": preamble,
},
{**event, "sequence_number": 1},
{
"type": "error", "sequence_number": 2,
"error": {"type": "server_error", "code": "internal_error", "message": "interrupted"},
},
)
def _openai_stream_cut_after_output_item(preamble: str, item: Mapping[str, object]) -> httpx.Response:
return _openai_stream_cut_after_event(
preamble, {"type": "response.output_item.added", "output_index": 1, "item": item}
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"item",
[
{"type": "shell_call", "id": "sh_1", "call_id": "call_1", "action": {"commands": ["ls"]}, "status": "in_progress"},
{
"type": "apply_patch_call", "id": "apc_1", "call_id": "call_1", "status": "in_progress",
"operation": {"type": "update_file", "path": "a.py", "diff": "@@"},
},
{
"type": "local_shell_call", "id": "lsh_1", "call_id": "call_1", "status": "in_progress",
"action": {"type": "exec", "command": ["ls"], "env": {}},
},
{
"type": "computer_call", "id": "cu_1", "call_id": "call_1", "status": "in_progress",
"action": {"type": "screenshot"}, "pending_safety_checks": [],
},
{"type": "mcp_approval_request", "id": "mcpr_1", "name": "lookup", "arguments": "{}", "server_label": "docs"},
],
ids=lambda item: item["type"],
)
async def test_responses_fallback_surfaces_the_primary_error_once_the_cut_stream_began_a_client_run_tool_call(
item: Mapping[str, object], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
respx_mock.post("https://api.openai.com/v1/responses").mock(
return_value=_openai_stream_cut_after_output_item("Let me check the repo. ", item)
)
fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(
return_value=_anthropic_stream("The repo has two files.")
)
router: Final = _primary_with_fallback("anthropic/claude-opus-5-5")
try:
stream: Final = await router.aresponses(model="primary", input="What is in the repo?", stream=True)
with pytest.raises(litellm.InternalServerError, match="interrupted"):
[event async for event in stream]
assert fallback.call_count == 0
finally:
router.discard()
@pytest.mark.asyncio
@pytest.mark.parametrize(
"server_run_event",
[
{
"type": "response.output_item.added", "output_index": 1,
"item": {
"type": "mcp_call", "id": "mcp_1", "name": "ask_question", "arguments": "{}",
"server_label": "docs", "status": "in_progress",
},
},
{"type": "response.mcp_call_arguments.delta", "output_index": 1, "item_id": "mcp_1", "delta": '{"question": "set'},
],
ids=["mcp_call_item", "mcp_call_arguments"],
)
async def test_responses_fallback_continues_after_a_server_run_mcp_call(
server_run_event: Mapping[str, object], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
respx_mock.post("https://api.openai.com/v1/responses").mock(
return_value=_openai_stream_cut_after_event("Let me check the docs. ", server_run_event)
)
fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(
return_value=_anthropic_stream("The docs cover setup.")
)
router: Final = _primary_with_fallback("anthropic/claude-opus-5-5")
try:
stream: Final = await router.aresponses(model="primary", input="What do the docs cover?", stream=True)
events: Final = tuple([event async for event in stream])
assert _streamed_text(events).endswith("The docs cover setup.")
assert fallback.call_count == 1
finally:
router.discard()
@pytest.mark.asyncio
async def test_responses_fallback_continues_when_allowed_tools_in_auto_mode_leaves_the_reply_open(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
openai: Final = respx_mock.post("https://api.openai.com/v1/responses").mock(
side_effect=[
_openai_stream_cut_after("1, 2, "),
_sse(
{
"type": "response.output_text.delta", "sequence_number": 0,
"item_id": "msg_fallback", "output_index": 0, "content_index": 0, "delta": "3, 4, 5",
},
{
"type": "response.completed", "sequence_number": 1,
"response": {
"id": "resp_fallback", "object": "response", "created_at": 0, "status": "completed",
"model": "gpt-6.1-terra", "output": [],
"usage": {"input_tokens": 5, "output_tokens": 5, "total_tokens": 10},
},
},
),
]
)
router: Final = _primary_with_fallback("openai/gpt-6.1-terra")
try:
stream: Final = await router.aresponses(
model="primary",
input="Print the integers 1 through 5.",
stream=True,
tools=[{"type": "function", "name": "lookup", "parameters": {"type": "object"}}],
tool_choice={"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "lookup"}]},
)
events: Final = tuple([event async for event in stream])
assert _streamed_text(events) == "1, 2, 3, 4, 5"
assert openai.call_count == 2
continued_input: Final = json.loads(openai.calls.last.request.content)["input"]
assert [item["role"] for item in continued_input[-2:]] == ["assistant", "user"]
finally:
router.discard()
@pytest.mark.asyncio
async def test_responses_fallback_surfaces_a_content_filter_that_cut_the_stream_mid_answer(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
respx_mock.post("https://api.openai.com/v1/responses").mock(
return_value=_sse(
{
"type": "response.output_text.delta", "sequence_number": 0,
"item_id": "msg_primary", "output_index": 0, "content_index": 0, "delta": "Step one is to obtain the",
},
{
"type": "error", "sequence_number": 1,
"error": {
"type": "invalid_request_error", "code": "content_policy_violation",
"message": "The response was halted mid-stream.", "param": None,
},
},
)
)
fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(
return_value=_anthropic_stream(" reagents.")
)
router: Final = Router(
model_list=[
{"model_name": "primary", "litellm_params": {"model": "openai/gpt-6-astra", "api_key": "synthetic"}},
{"model_name": "fallback", "litellm_params": {"model": "anthropic/claude-opus-5-5", "api_key": "synthetic"}},
],
fallbacks=[{"primary": ["fallback"]}],
content_policy_fallbacks=[{"primary": ["fallback"]}],
num_retries=0,
)
try:
stream: Final = await router.aresponses(model="primary", input="How do I make it?", stream=True)
with pytest.raises(litellm.ContentPolicyViolationError, match="halted mid-stream"):
[event async for event in stream]
assert fallback.call_count == 0
finally:
router.discard()
@pytest.mark.asyncio
async def test_bridged_responses_fallback_surfaces_the_primary_error_once_the_cut_stream_began_a_tool_call(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
provider: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(
side_effect=[
_anthropic_stream_cut_mid_tool_call("Let me look that up. "),
_anthropic_stream("It is sunny in Paris."),
]
)
router: Final = Router(
model_list=[
{"model_name": "primary", "litellm_params": {"model": "anthropic/claude-primary", "api_key": "synthetic"}},
{"model_name": "fallback", "litellm_params": {"model": "anthropic/claude-opus-5-5", "api_key": "synthetic"}},
],
fallbacks=[{"primary": ["fallback"]}],
num_retries=0,
)
try:
stream: Final = await router.aresponses(
model="primary",
input="What is the weather in Paris?",
stream=True,
tools=[{"type": "function", "name": "lookup", "parameters": {"type": "object"}}],
)
with pytest.raises(litellm.InternalServerError, match="interrupted"):
[event async for event in stream]
assert provider.call_count == 1
finally:
router.discard()
@pytest.mark.asyncio
async def test_bridged_responses_fallback_continues_after_a_server_run_tool_search(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
provider: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(
side_effect=[
_anthropic_stream_cut_mid_tool_call(
"Let me find the right tool. ",
{"type": "server_tool_use", "id": "srvtoolu_primary", "name": "tool_search_tool_regex", "input": {}},
),
_anthropic_stream("It is sunny in Paris."),
]
)
router: Final = Router(
model_list=[
{"model_name": "primary", "litellm_params": {"model": "anthropic/claude-primary", "api_key": "synthetic"}},
{"model_name": "fallback", "litellm_params": {"model": "anthropic/claude-opus-5-5", "api_key": "synthetic"}},
],
fallbacks=[{"primary": ["fallback"]}],
num_retries=0,
)
try:
stream: Final = await router.aresponses(
model="primary",
input="What is the weather in Paris?",
stream=True,
tools=[
{"type": "tool_search", "execution": "server"},
{"type": "function", "name": "lookup", "parameters": {"type": "object"}, "defer_loading": True},
],
)
events: Final = tuple([event async for event in stream])
assert _streamed_text(events).endswith("It is sunny in Paris.")
assert provider.call_count == 2
finally:
router.discard()
@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit")
@pytest.mark.asyncio

101
tests/unit/llms/base_llm/test_base_model_iterator.py
View file

@ -4,15 +4,19 @@ and non-string objects (e.g. Pydantic BaseModel events from the Responses API) p
"""
import json
from collections.abc import AsyncIterator
from collections.abc import AsyncIterator, Mapping
from itertools import chain
from typing import Final
import httpx
import pytest
from pydantic import BaseModel
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator, convert_model_response_to_streaming
from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler
from litellm.llms.openai.common_utils import OpenAIError
from litellm.types.utils import Choices, GenericStreamingChunk, Message, ModelResponse
@ -336,3 +340,98 @@ def test_fake_stream_chunk_indexes_thinking_blocks_by_position_without_touching_
{"type": "thinking", "thinking": "Step 2.", "signature": "sig-b", "index": 2},
]
assert [block.get("index") for block in response.choices[0].message.thinking_blocks] == [None, None, None]
def _sse_chunk(delta: Mapping[str, object], finish_reason: str | None = None) -> str:
return "data: " + json.dumps(
{
"id": "chatcmpl-done",
"object": "chat.completion.chunk",
"created": 1700000000,
"model": "mistral-small-latest",
"choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}],
}
)
async def _aiter_lines(lines: tuple[str, ...]) -> AsyncIterator[str]:
for line in lines:
yield line
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_done_inside_streamed_json_is_content_not_the_end_of_stream(sync_mode: bool) -> None:
lines: Final = (
_sse_chunk({"role": "assistant", "reasoning_content": "Mark it [DONE] once tests pass"}),
_sse_chunk({"content": "Checklist: "}),
_sse_chunk({"content": "- [DONE] write tests"}),
_sse_chunk({"content": " - [ ] ship it"}),
_sse_chunk(
{
"tool_calls": [
{
"index": 0,
"id": "call_1",
"type": "function",
"function": {"name": "set_status", "arguments": ""},
}
]
}
),
_sse_chunk({"tool_calls": [{"index": 0, "function": {"arguments": '{"status": '}}]}),
_sse_chunk({"tool_calls": [{"index": 0, "function": {"arguments": '"[DONE]"'}}]}),
_sse_chunk({"tool_calls": [{"index": 0, "function": {"arguments": ', "n": 1}'}}]}),
_sse_chunk({}, "tool_calls"),
"data: [DONE]",
)
wrapper: Final = CustomStreamWrapper(
completion_stream=OpenAIChatCompletionStreamingHandler(
streaming_response=iter(lines) if sync_mode else _aiter_lines(lines), sync_stream=sync_mode
),
model="mistral/mistral-small-latest",
custom_llm_provider="mistral",
logging_obj=Logging(
model="mistral/mistral-small-latest",
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type="completion",
start_time=1_790_000_000.0,
litellm_call_id="call-id",
function_id="function-id",
),
)
chunks: Final = list(wrapper) if sync_mode else [chunk async for chunk in wrapper]
deltas: Final = [chunk.choices[0].delta for chunk in chunks if chunk.choices]
tool_calls: Final = tuple(chain.from_iterable(delta.tool_calls or () for delta in deltas))
assert "".join(getattr(delta, "reasoning_content", None) or "" for delta in deltas) == (
"Mark it [DONE] once tests pass"
)
assert "".join(delta.content or "" for delta in deltas) == "Checklist: - [DONE] write tests - [ ] ship it"
assert json.loads("".join(tool_call.function.arguments or "" for tool_call in tool_calls)) == {
"status": "[DONE]",
"n": 1,
}
assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices and chunk.choices[0].finish_reason] == [
"tool_calls"
]
@pytest.mark.parametrize("line", ["data: [DONE]", "data:[DONE]", "[DONE]", "data: [DONE]", "data: [DONE] "])
def test_done_marker_ends_the_stream(line: str) -> None:
chunks: Final = list(BaseModelResponseIterator(streaming_response=iter((line,)), sync_stream=True))
assert chunks == [
GenericStreamingChunk(text="", is_finished=True, finish_reason="stop", usage=None, index=0, tool_use=None)
]
def test_error_frame_mentioning_done_raises_instead_of_ending_the_stream() -> None:
line: Final = 'data: {"error": {"message": "prompt too long [DONE]", "code": 400}}'
with pytest.raises(OpenAIError, match=r"prompt too long \[DONE\]") as error:
list(OpenAIChatCompletionStreamingHandler(streaming_response=iter((line,)), sync_stream=True))
assert error.value.status_code == 400

102
tests/unit/llms/bedrock/chat/test_native_output_fallback.py
View file

@ -121,3 +121,105 @@ def test_a_rejection_logs_the_retry_and_advises_opting_out_only_when_native_outp
assert _MODEL in caplog.text
assert ("supports_bedrock_native_structured_output: false" in caplog.text) is advises_opting_out
_STRICT_FALLBACK: Final = NativeOutputFallback(optional_params={"tools": "without_strict"}, litellm_params={})
# Bedrock's 400s for strict tools, seen live on us.anthropic.claude-sonnet-4-6 from 2026-10-20 to 2026-10-28 and quoted
# in issue #34388, each answered 200 once the request dropped strict
_STRICT_TOOLS_REJECTIONS: Final = tuple(
BedrockError(status_code=400, message=f'{{"message":"The model returned the following errors: {text}"}}')
for text in (
"Too many strict tools (21). The maximum number of strict tools supported is 20.",
"Schemas contains too many optional parameters (27), which would make grammar compilation inefficient.",
"Schemas contains too many parameters with union types (18 parameters with type arrays or anyOf). "
"This causes exponential compilation cost.",
"The compiled grammar is too large, which would cause performance issues.",
"tools.0.custom: Invalid schema: Schema type is missing for schema: {'description': 'tag'}",
"tools.3.custom: For 'integer' type, property 'minimum' is not supported",
"Compiled grammar size (410.4MB) exceeds maximum allowed size (300MB)",
"tools.1.custom.strict: Extra inputs are not permitted",
"Grammar compilation timed out.",
)
)
_TOO_MANY_STRICT_TOOLS: Final = _STRICT_TOOLS_REJECTIONS[0]
# Bedrock's 400s for requests with a strict tool that fail for other reasons, seen live on
# us.anthropic.claude-sonnet-4-6 on 2026-10-21, which echo tool names and tool call ids
_OTHER_400S_WITH_A_STRICT_TOOL: Final = tuple(
BedrockError(status_code=400, message=f'{{"message":"{text}"}}')
for text in (
"The tool check_grammar is already defined at toolConfig.tools.1.",
"Expected toolResult blocks at messages.0.content for the following Ids: call_grammar_9",
"tool_use ids were found without tool_result blocks immediately after: call_grammar_1",
)
)
def _call_with_strict_tools_fallback(
bedrock: ScriptedBedrock, fallback: NativeOutputFallback | None, strict_tools_fallback: NativeOutputFallback | None
) -> object:
return _resolve(
call_with_native_output_fallback(
bedrock,
None if fallback is None else lambda: fallback,
_MODEL,
strict_tools_fallback=None if strict_tools_fallback is None else lambda: strict_tools_fallback,
)
)
@_MODES
@pytest.mark.parametrize("rejection", _STRICT_TOOLS_REJECTIONS, ids=lambda error: error.message[60:100])
def test_a_strict_tools_rejection_is_sent_again_once_with_the_strict_tools_fallback(
asynchronous: bool, rejection: BedrockError
) -> None:
bedrock: Final = ScriptedBedrock(errors=(rejection, None), asynchronous=asynchronous)
response: Final = _call_with_strict_tools_fallback(bedrock, None, _STRICT_FALLBACK)
assert isinstance(response, ModelResponse) and response.id == "answer-2"
assert bedrock.sent == (None, _STRICT_FALLBACK)
@_MODES
@pytest.mark.parametrize(
("error", "strict_tools_fallback"),
[
(_OTHER_400, _STRICT_FALLBACK),
*((error, _STRICT_FALLBACK) for error in _OTHER_400S_WITH_A_STRICT_TOOL),
(BedrockError(status_code=500, message=_TOO_MANY_STRICT_TOOLS.message), _STRICT_FALLBACK),
(_TOO_MANY_STRICT_TOOLS, None),
],
ids=[
"other_400",
"duplicate_tool_name",
"missing_tool_result",
"orphan_tool_use",
"server_error",
"no_strict_tools_fallback",
],
)
def test_no_second_request_without_both_a_strict_tools_rejection_and_a_strict_tools_fallback(
asynchronous: bool, error: BedrockError, strict_tools_fallback: NativeOutputFallback | None
) -> None:
bedrock: Final = ScriptedBedrock(errors=(error,), asynchronous=asynchronous)
with pytest.raises(BedrockError) as raised:
_call_with_strict_tools_fallback(bedrock, None, strict_tools_fallback)
assert raised.value is error
assert bedrock.sent == (None,)
@pytest.mark.parametrize(
("error", "retried_with"),
[(_REJECTION, _FALLBACK), (_TOO_MANY_STRICT_TOOLS, _STRICT_FALLBACK)],
ids=["native_output_rejection", "strict_tools_rejection"],
)
def test_each_rejection_is_sent_again_with_its_own_fallback_when_both_are_set(
error: BedrockError, retried_with: NativeOutputFallback
) -> None:
bedrock: Final = ScriptedBedrock(errors=(error, None), asynchronous=False)
_call_with_strict_tools_fallback(bedrock, _FALLBACK, _STRICT_FALLBACK)
assert bedrock.sent == (None, retried_with)

79
tests/unit/llms/bedrock/test_bedrock_common_utils.py
View file

@ -6,7 +6,9 @@ from typing import Final
from litellm.llms.bedrock.common_utils import(
AmazonBedrockGlobalConfig,
BedrockModelInfo,
NativeOutputFallback,
bedrock_native_output_is_opt_in,
bedrock_strict_tools_fallback,
deployment_model_info,
extract_model_name_from_bedrock_arn,
get_bedrock_base_model,
@ -2150,3 +2152,80 @@ def test_bedrock_response_stream_shape_is_unavailable_when_botocore_describes_th
assert shape is None
(logged,) = warning.call_args_list
assert "only-in-the-description" not in logged.args[0] % logged.args[1:]
def test_bedrock_strict_tools_fallback_drops_only_strict_and_keeps_everything_else():
strict_tool: Final = {
"type": "function",
"cache_control": {"type": "ephemeral"},
"function": {"name": "lookup", "description": "Look up", "strict": True, "parameters": {"type": "object"}},
}
plain_tool: Final = {"type": "function", "function": {"name": "plain", "parameters": {"type": "object"}}}
anthropic_tool: Final = {"name": "native", "input_schema": {"type": "object"}}
bedrock_tool: Final = {"toolSpec": {"name": "spec", "strict": True, "inputSchema": {"json": {"type": "object"}}}}
litellm_params: Final = {"drop_params": True}
fallback: Final = bedrock_strict_tools_fallback(
"us.anthropic.claude-sonnet-4-6",
{"tools": [strict_tool, plain_tool, anthropic_tool, bedrock_tool], "maxTokens": 16, "stream": True},
litellm_params,
)
assert fallback == NativeOutputFallback(
optional_params={
"tools": [
{**strict_tool, "function": {"name": "lookup", "description": "Look up", "parameters": {"type": "object"}}},
plain_tool,
anthropic_tool,
{"toolSpec": {"name": "spec", "inputSchema": {"json": {"type": "object"}}}},
],
"maxTokens": 16,
"stream": True,
},
litellm_params=litellm_params,
)
assert strict_tool["function"]["strict"] is True
assert bedrock_tool["toolSpec"]["strict"] is True
@pytest.mark.parametrize(
"optional_params",
[
{"maxTokens": 16},
{"tools": [{"type": "function", "function": {"name": "plain", "parameters": {}}}]},
{"tools": [{"toolSpec": {"name": "spec", "inputSchema": {"json": {}}}}]},
],
ids=["no_tools", "no_strict_tool", "no_strict_tool_spec"],
)
def test_bedrock_strict_tools_fallback_is_none_without_a_strict_tool(optional_params: dict[str, object]) -> None:
assert bedrock_strict_tools_fallback("us.anthropic.claude-sonnet-4-6", optional_params, {}) is None
_CLAUDE_PROFILE_ARN: Final = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/claude"
@pytest.mark.parametrize(
("model", "litellm_params"),
[
("us.anthropic.claude-sonnet-4-6", {}),
("us.amazon.nova-pro-v1:0", {}),
(_CLAUDE_PROFILE_ARN, {"base_model": "bedrock/us.anthropic.claude-sonnet-4-6"}),
(_CLAUDE_PROFILE_ARN, {}),
],
ids=["claude", "nova", "claude_profile_with_base_model", "profile_without_base_model"],
)
@pytest.mark.parametrize("strict", [True, False])
def test_bedrock_strict_tools_fallback_retries_only_when_the_request_sent_strict(
model: str, litellm_params: dict[str, object], strict: bool
) -> None:
from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig
tools: Final = [{"type": "function", "function": {"name": "lookup", "strict": strict, "parameters": {}}}]
request: Final = AmazonConverseConfig().transform_request(
model, [{"role": "user", "content": "hi"}], {"tools": tools}, {**litellm_params}, {}
)
sent_strict: Final = any("strict" in tool["toolSpec"] for tool in request["toolConfig"]["tools"])
fallback: Final = bedrock_strict_tools_fallback(model, {"tools": tools}, litellm_params)
assert (fallback is not None) == sent_strict

145
tests/unit/llms/openai_like/test_aiand_provider.py Normal file
View file

@ -0,0 +1,145 @@
import json
from typing import Final
import httpx
import pytest
import respx
import litellm
from litellm.caching.llm_caching_handler import LLMClientCache
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.llms.openai_like.json_loader import JSONProviderRegistry
from litellm.proxy.auth.model_checks import get_known_models_from_wildcard
from litellm.utils import ProviderConfigManager
AIAND_API_BASE = "https://api.aiand.com/v1"
def test_aiand_chat_provider(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("AIAND_API_KEY", "test-key")
provider_config = JSONProviderRegistry.get("aiand")
assert provider_config is not None
assert litellm.LlmProviders.AIAND.value == "aiand"
assert "aiand" in litellm.openai_compatible_providers
assert provider_config.base_url == AIAND_API_BASE
assert provider_config.api_key_env == "AIAND_API_KEY"
assert provider_config.api_base_env == "AIAND_API_BASE"
model, provider, api_key, api_base = get_llm_provider("aiand/openai/gpt-oss-120b")
assert model == "openai/gpt-oss-120b"
assert provider == "aiand"
assert api_key == "test-key"
assert api_base == AIAND_API_BASE
def test_aiand_completion_request(respx_mock: respx.MockRouter):
messages = [{"role": "user", "content": "Hello"}]
route = respx_mock.post(f"{AIAND_API_BASE}/chat/completions").mock(
return_value=httpx.Response(
200,
json={
"id": "chatcmpl-test",
"object": "chat.completion",
"created": 1,
"model": "openai/gpt-oss-120b",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "Hello"},
"finish_reason": "stop",
}
],
"usage": {
"prompt_tokens": 1,
"completion_tokens": 1,
"total_tokens": 2,
},
},
)
)
litellm.completion(
model="aiand/openai/gpt-oss-120b",
api_key="test-key",
messages=messages,
)
assert route.called
request = route.calls.last.request
assert str(request.url) == f"{AIAND_API_BASE}/chat/completions"
assert request.headers["Authorization"] == "Bearer test-key"
request_body = json.loads(request.read())
assert request_body["model"] == "openai/gpt-oss-120b"
assert request_body["messages"] == messages
def test_aiand_responses_provider():
config = ProviderConfigManager.get_provider_responses_api_config(
provider="aiand",
model="aiand/openai/gpt-oss-120b",
)
assert config is not None
assert config.custom_llm_provider == "aiand"
assert config.get_complete_url(api_base=None, litellm_params={}) == f"{AIAND_API_BASE}/responses"
def _aiand_catalog() -> frozenset[str]:
return frozenset(name for name, info in litellm.model_cost.items() if info.get("litellm_provider") == "aiand")
def test_aiand_wildcard_lists_its_priced_models():
catalog: Final = _aiand_catalog()
assert catalog
assert frozenset(get_known_models_from_wildcard("aiand/*")) == catalog
def test_aiand_calls_are_billed_from_their_cost_rows():
catalog: Final = _aiand_catalog()
assert catalog
for model in sorted(catalog):
row: Final = litellm.model_cost[model]
prompt_cost, completion_cost = litellm.cost_per_token(model=model, prompt_tokens=1000, completion_tokens=500)
assert prompt_cost == pytest.approx(1000 * row["input_cost_per_token"])
assert completion_cost == pytest.approx(500 * row["output_cost_per_token"])
assert prompt_cost > 0 and completion_cost > 0
@pytest.mark.asyncio
async def test_aiand_anthropic_messages_reach_ai_and_with_the_thinking_budget(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache())
thinking: Final = {"type": "enabled", "budget_tokens": 1024}
with respx.mock() as upstream:
route: Final = upstream.post(f"{AIAND_API_BASE}/messages").respond(
200,
json={
"id": "msg_aiand",
"type": "message",
"role": "assistant",
"model": "zai-org/glm-5.3",
"content": [{"type": "thinking", "thinking": "A greeting", "signature": ""}, {"type": "text", "text": "Hello"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 4, "output_tokens": 3},
},
)
response: Final = await litellm.anthropic.messages.acreate(
model="aiand/zai-org/glm-5.3",
messages=[{"role": "user", "content": "Say hello"}],
max_tokens=2048,
thinking=thinking,
api_key="aiand-test-key",
)
request: Final = route.calls.last.request
body: Final = json.loads(request.content)
assert route.call_count == 1
assert body["model"] == "zai-org/glm-5.3"
assert body["thinking"] == thinking
assert "reasoning_effort" not in body
assert [block["type"] for block in response["content"]] == ["thinking", "text"]

270
tests/unit/proxy/client/cli/test_agents.py
View file

@ -1,10 +1,13 @@
import functools
import inspect
import json
import os
import stat
import subprocess
import sys
from collections.abc import Mapping, Sequence
from pathlib import Path
from typing import Final
from unittest.mock import patch
import click
@ -31,8 +34,14 @@ from litellm.proxy.client.cli.commands.agents import (
verify_subscription_key,
)
from litellm.proxy.client.cli.commands.claude_settings import ClaudeSettingsError
from litellm.proxy.client.cli.commands.config import save_config
from litellm.proxy.client.cli.main import cli as lite
if sys.version_info >= (3, 11):
import tomllib
else:
import tomli as tomllib
AGENTS_MODULE = "litellm.proxy.client.cli.commands.agents"
CLAUDE_SETTINGS_MODULE = "litellm.proxy.client.cli.commands.claude_settings"
@ -189,6 +198,30 @@ def _overlay(args):
return json.loads(Path(args[2]).read_text())
def _codex_config(args):
assert set(args[0::2]) == {"-c"}
return tomllib.loads("\n".join(args[1::2]))
def _set_codex_key(config: Mapping[str, object], keys: Sequence[str], value: object) -> Mapping[str, object]:
head, *rest = keys
if not rest:
return {**config, head: value}
nested: Final = config.get(head, {})
assert isinstance(nested, Mapping)
return {**config, head: _set_codex_key(nested, rest, value)}
def _with_codex_override(config: Mapping[str, object], override: str) -> Mapping[str, object]:
path, value = override.split("=", 1)
return _set_codex_key(config, path.split("."), tomllib.loads(f"value = {value}")["value"])
def _codex_config_over(user_config_toml: str, args: Sequence[str]) -> Mapping[str, object]:
assert set(args[0::2]) == {"-c"}
return functools.reduce(_with_codex_override, args[1::2], tomllib.loads(user_config_toml))
class TestAgentProfile:
def test_claude_is_anthropic(self):
name, profiles = agent_profile("claude")
@ -254,6 +287,7 @@ class TestBuildAgentEnv:
env = build_agent_env({}, "http://localhost:4000/", "sk-key", frozenset({"openai"}))
assert env["OPENAI_BASE_URL"] == "http://localhost:4000/v1"
assert env["OPENAI_API_KEY"] == "sk-key"
assert "LITELLM_PROXY_API_KEY" not in env
assert "ANTHROPIC_BASE_URL" not in env
assert "ENABLE_TOOL_SEARCH" not in env
assert "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" not in env
@ -317,11 +351,21 @@ class TestBuildAgentEnv:
)
assert env["ANTHROPIC_CUSTOM_HEADERS"].splitlines() == ["x-team: blue", "x-litellm-api-key: Bearer sk-new"]
def test_subscription_leaves_openai_agents_on_the_bearer_key(self):
env = build_agent_env({}, "http://localhost:4000", "sk-key", frozenset({"openai"}), subscription=True)
assert env["OPENAI_API_KEY"] == "sk-key"
def test_subscription_routes_openai_agents_without_offering_the_key_as_an_openai_credential(self):
env = build_agent_env({}, "http://localhost:4000/", "sk-key", frozenset({"openai"}), subscription=True)
assert "OPENAI_BASE_URL" not in env
assert env["LITELLM_PROXY_API_KEY"] == "sk-key"
assert "OPENAI_API_KEY" not in env
assert "ANTHROPIC_CUSTOM_HEADERS" not in env
def test_subscription_leaves_the_users_own_openai_key_and_base_url_as_they_were(self):
base = {"OPENAI_API_KEY": "sk-users-openai", "OPENAI_BASE_URL": "https://api.openai.com/v1"}
env = build_agent_env(base, "http://localhost:4000", "sk-key", frozenset({"openai"}), subscription=True)
assert env["OPENAI_API_KEY"] == "sk-users-openai"
assert env["OPENAI_BASE_URL"] == "https://api.openai.com/v1"
assert env["LITELLM_PROXY_API_KEY"] == "sk-key"
assert base == {"OPENAI_API_KEY": "sk-users-openai", "OPENAI_BASE_URL": "https://api.openai.com/v1"}
def test_default_mode_does_not_touch_custom_headers(self):
env = build_agent_env(
{"ANTHROPIC_CUSTOM_HEADERS": "x-team: blue"}, "http://localhost:4000", "sk-key", frozenset({"anthropic"})
@ -354,6 +398,67 @@ class TestAgentLaunchArgs:
agent_launch_args("codex", "http://localhost:4000")
)
def test_codex_default_provider_reads_the_key_from_openai_api_key(self):
assert _codex_config(agent_launch_args("codex", "http://localhost:4000/")) == {
"model_provider": "litellm",
"model_providers": {
"litellm": {
"name": "LiteLLM proxy",
"base_url": "http://localhost:4000/v1",
"wire_api": "responses",
"supports_websockets": False,
"requires_openai_auth": False,
"env_key": "OPENAI_API_KEY",
"http_headers": {},
}
},
}
def test_codex_subscription_provider_keeps_the_chatgpt_login_and_sends_the_key_in_the_litellm_header(self):
assert _codex_config(
agent_launch_args("/usr/local/bin/codex", "http://localhost:4000/", subscription=True)
) == {
"model_provider": "lite_subscription",
"model_providers": {
"lite_subscription": {
"name": "LiteLLM proxy",
"base_url": "http://localhost:4000/v1",
"wire_api": "responses",
"supports_websockets": False,
"requires_openai_auth": True,
"env_http_headers": {"x-litellm-api-key": "LITELLM_PROXY_API_KEY"},
"http_headers": {},
}
},
}
def test_codex_subscription_leaves_the_cookbook_provider_table_out_of_the_launch(self):
user_config = """
model_provider = "litellm_subscription"
[model_providers.litellm_subscription]
name = "LiteLLM subscription"
base_url = "https://litellm.example.com/v1"
wire_api = "responses"
requires_openai_auth = true
env_key = "OPENAI_API_KEY"
http_headers = { "x-team" = "qa" }
"""
config = _codex_config_over(
user_config, agent_launch_args("codex", "http://localhost:4000/", subscription=True)
)
launched = config["model_providers"][config["model_provider"]]
assert "env_key" not in launched
assert launched["requires_openai_auth"] is True
assert launched["base_url"] == "http://localhost:4000/v1"
assert config["model_providers"]["litellm_subscription"] == tomllib.loads(user_config)["model_providers"][
"litellm_subscription"
]
@pytest.mark.parametrize("agent", ["claude", "opencode", "pi", "mytool"])
def test_subscription_adds_no_args_for_agents_without_a_proxy_provider(self, agent):
assert agent_launch_args(agent, "http://localhost:4000", subscription=True) == []
def test_pi_gets_no_static_args(self):
assert agent_launch_args("pi", "http://localhost:4000") == []
@ -1286,6 +1391,25 @@ class TestRunAgent:
with pytest.raises(AgentRunError):
run_agent("http://localhost:4000", "sk-key", [])
@pytest.mark.parametrize("subscription", [True, False])
def test_codex_stop_hook_is_added_only_when_codex_sends_the_key_as_its_openai_key(self, subscription):
launched = {}
run_agent(
"http://localhost:4000",
"sk-key",
["codex"],
subscription=subscription,
base_env={"OPENAI_API_KEY": "sk-users-own-openai-key"},
which=lambda name: "/usr/local/bin/codex",
verify=lambda *a: None,
verify_subscription=lambda *a: None,
sync_models=lambda *a: ModelSyncSkipped("not needed"),
warn=lambda message: None,
launcher=lambda p, a, e: launched.update(args=tuple(a)),
preparers={"codex": lambda *a: ("-c", "hooks.Stop=lite")},
)
assert ("hooks.Stop=lite" in launched["args"]) is not subscription
def test_subscription_launch_probes_with_the_header_and_exports_it(self):
probed = []
launched = {}
@ -1574,6 +1698,56 @@ class TestRunAgent:
read_settings=lambda path: {},
)
def test_codex_subscription_launch_keeps_the_chatgpt_login_and_needs_no_settings_overlay(self):
probed = []
overlays = []
launched = {}
run_agent(
"http://localhost:4000",
"sk-key",
["codex", "exec", "--settings", "kept"],
subscription=True,
base_env={"PATH": "/usr/bin"},
which=lambda name: "/usr/local/bin/codex",
verify=lambda *a: probed.append(("bearer", a)),
verify_subscription=lambda *a: probed.append(("header", a)),
sync_models=lambda *a: {},
preparers={},
write_overlay=lambda text, environ: overlays.append(text),
launcher=lambda p, a, e: launched.update(path=p, args=tuple(a), env=dict(e)),
)
assert probed == [("header", ("http://localhost:4000", "sk-key"))]
assert overlays == []
assert launched["path"] == "/usr/local/bin/codex"
assert launched["args"][0] == "codex"
assert launched["args"][-3:] == ("exec", "--settings", "kept")
assert _codex_config(launched["args"][1:-3])["model_provider"] == "lite_subscription"
assert "OPENAI_BASE_URL" not in launched["env"]
assert launched["env"]["LITELLM_PROXY_API_KEY"] == "sk-key"
assert "OPENAI_API_KEY" not in launched["env"]
@pytest.mark.parametrize(
"sync_result",
[ModelSyncArgs(("-c", 'model_catalog_json="/tmp/c.json"')), ModelSyncSkipped("lists no chat models")],
)
def test_codex_subscription_launch_leaves_the_model_list_to_codex(self, sync_result):
warnings = []
launched = {}
run_agent(
"http://localhost:4000",
"sk-key",
["codex"],
subscription=True,
base_env={},
which=lambda name: "/usr/local/bin/codex",
verify_subscription=lambda *a: None,
sync_models=lambda *a: sync_result,
warn=warnings.append,
launcher=lambda p, a, e: launched.update(args=tuple(a)),
)
assert warnings == []
assert "model_catalog_json" not in json.dumps(launched["args"])
def test_default_launch_still_probes_with_the_bearer(self):
probed = []
run_agent(
@ -1964,7 +2138,11 @@ class TestAgentCommands:
assert captured["command"] == ["opencode"]
assert "routing OpenCode through proxy at http://localhost:4000" in result.output
def test_claude_subscription_flag_reaches_run_agent_and_is_not_forwarded(self):
@pytest.mark.parametrize(
("agent", "user_args", "login"),
[("claude", ["--resume"], "claude.ai"), ("codex", ["exec", "hi"], "ChatGPT")],
)
def test_subscription_flag_reaches_run_agent_and_is_not_forwarded(self, agent, user_args, login):
captured = {}
def fake_run_agent(base_url, api_key, command, **kwargs):
@ -1973,61 +2151,85 @@ class TestAgentCommands:
with patch(f"{AGENTS_MODULE}.run_agent", side_effect=fake_run_agent):
result = self.runner.invoke(
_agent_command("claude"),
["--subscription", "--resume"],
_agent_command(agent),
["--subscription", *user_args],
obj={"base_url": "http://localhost:4000", "api_key": "sk-key"},
)
assert result.exit_code == 0, result.output
assert captured["subscription"] is True
assert captured["command"] == ["claude", "--resume"]
assert "keeping your claude.ai login" in result.output
assert captured["command"] == [agent, *user_args]
assert f"keeping your {login} login" in result.output
@pytest.mark.parametrize(("agent", "login"), [("claude", "claude.ai"), ("codex", "ChatGPT")])
@pytest.mark.parametrize(
("configured", "args", "expected"),
("saved", "args", "expected"),
[
(False, [], False),
(True, [], True),
(True, ["--no-subscription"], False),
(False, ["--subscription"], True),
(None, [], False),
("false", [], False),
("true", [], True),
("true", ["--no-subscription"], False),
("false", ["--subscription"], True),
],
)
def test_subscription_comes_from_the_flag_else_the_saved_setting(self, configured, args, expected):
def test_subscription_comes_from_the_flag_else_the_saved_setting(self, agent, login, saved, args, expected):
if saved is not None:
save_config({f"{agent}_subscription": saved})
captured = {}
with (
patch(f"{AGENTS_MODULE}.claude_subscription_default", return_value=configured),
patch(
f"{AGENTS_MODULE}.run_agent",
side_effect=lambda b, k, c, **kw: captured.update(sub=kw.get("subscription"), command=list(c)),
),
with patch(
f"{AGENTS_MODULE}.run_agent",
side_effect=lambda b, k, c, **kw: captured.update(sub=kw.get("subscription"), command=list(c)),
):
result = self.runner.invoke(
_agent_command("claude"), args, obj={"base_url": "http://localhost:4000", "api_key": "sk-key"}
_agent_command(agent), args, obj={"base_url": "http://localhost:4000", "api_key": "sk-key"}
)
assert result.exit_code == 0, result.output
assert captured["sub"] is expected
assert captured["command"] == ["claude"]
assert ("keeping your claude.ai login" in result.output) is expected
assert captured["command"] == [agent]
assert (f"keeping your {login} login" in result.output) is expected
@pytest.mark.parametrize("agent", ["codex", "opencode"])
def test_the_saved_claude_setting_does_not_change_other_agents(self, agent):
@pytest.mark.parametrize(
("saved_key", "agent"),
[
("claude_subscription", "codex"),
("claude_subscription", "opencode"),
("codex_subscription", "claude"),
("codex_subscription", "opencode"),
],
)
def test_a_saved_subscription_setting_only_changes_its_own_agent(self, saved_key, agent):
save_config({saved_key: "true"})
captured = {}
with (
patch(f"{AGENTS_MODULE}.claude_subscription_default", return_value=True),
patch(
f"{AGENTS_MODULE}.run_agent",
side_effect=lambda b, k, c, **kw: captured.update(sub=kw.get("subscription")),
),
with patch(
f"{AGENTS_MODULE}.run_agent",
side_effect=lambda b, k, c, **kw: captured.update(sub=kw.get("subscription")),
):
result = self.runner.invoke(
_agent_command(agent), [], obj={"base_url": "http://localhost:4000", "api_key": "sk-key"}
)
assert result.exit_code == 0, result.output
assert captured["sub"] is False
assert "keeping your" not in result.output
@pytest.mark.parametrize("agent", ["codex", "opencode", "pi"])
def test_only_claude_offers_the_subscription_flag(self, agent):
assert "--subscription" not in {opt for param in _agent_command(agent).params for opt in param.opts}
def test_only_claude_and_codex_offer_the_subscription_flag(self):
offering = {
command.name
for command in agent_commands()
if "--subscription" in {opt for param in command.params for opt in param.opts}
}
assert offering == {"claude", "codex"}
@pytest.mark.parametrize(
("agent", "login", "own_credential"),
[("claude", "claude.ai", "Claude Code's"), ("codex", "ChatGPT", "Codex's")],
)
def test_subscription_help_names_the_agents_own_login_and_its_setting(self, agent, login, own_credential):
result = self.runner.invoke(lite, [agent, "--help"], prog_name="lite")
assert result.exit_code == 0, result.output
text = " ".join(result.output.split())
assert f"Keep using your {login} login" in text
assert f"instead of replacing {own_credential} own credential" in text
assert f"lite config set {agent}_subscription true" in text
def test_skip_verify_is_consumed_not_forwarded(self):
captured = {}

49
tests/unit/proxy/client/cli/test_config_commands.py
View file

@ -9,12 +9,12 @@ from click.testing import CliRunner
from litellm.litellm_core_utils.private_json import write_private_json
from litellm.proxy.client.cli import cli
from litellm.proxy.client.cli.commands.config import (
claude_subscription_default,
get_config_file_path,
get_config_value,
load_config,
pkce_default,
save_config,
subscription_default,
)
from litellm.proxy.client.cli.interface import show_commands
@ -362,33 +362,48 @@ class TestWritePrivateJson:
assert list(tmp_path.glob(".tmp-*")) == []
class TestClaudeSubscriptionSetting:
@pytest.mark.parametrize(("agent", "key"), [("claude", "claude_subscription"), ("codex", "codex_subscription")])
class TestSubscriptionSetting:
@pytest.mark.parametrize("word", ["true", "TRUE", "on", "1", "yes"])
def test_truthy_words_are_stored_as_true(self, cli_runner, isolated_home, word):
result = cli_runner.invoke(cli, ["config", "set", "claude_subscription", word])
def test_truthy_words_are_stored_as_true(self, cli_runner, isolated_home, agent, key, word):
result = cli_runner.invoke(cli, ["config", "set", key, word])
assert result.exit_code == 0, result.output
assert load_config()["claude_subscription"] == "true"
assert claude_subscription_default() is True
assert load_config()[key] == "true"
assert subscription_default(agent) is True
@pytest.mark.parametrize("word", ["false", "Off", "0", "no"])
def test_falsy_words_are_stored_as_false(self, cli_runner, isolated_home, word):
cli_runner.invoke(cli, ["config", "set", "claude_subscription", "true"])
result = cli_runner.invoke(cli, ["config", "set", "claude_subscription", word])
def test_falsy_words_are_stored_as_false(self, cli_runner, isolated_home, agent, key, word):
cli_runner.invoke(cli, ["config", "set", key, "true"])
result = cli_runner.invoke(cli, ["config", "set", key, word])
assert result.exit_code == 0, result.output
assert load_config()["claude_subscription"] == "false"
assert claude_subscription_default() is False
assert load_config()[key] == "false"
assert subscription_default(agent) is False
def test_anything_else_is_rejected_and_nothing_is_saved(self, cli_runner, isolated_home):
result = cli_runner.invoke(cli, ["config", "set", "claude_subscription", "maybe"])
def test_anything_else_is_rejected_and_nothing_is_saved(self, cli_runner, isolated_home, agent, key):
result = cli_runner.invoke(cli, ["config", "set", key, "maybe"])
assert result.exit_code != 0
assert "true or false" in result.output
assert "claude_subscription" not in load_config()
assert f"{key} must be true or false" in result.output
assert key not in load_config()
def test_unset_means_off(self, isolated_home):
assert claude_subscription_default() is False
def test_unset_means_off(self, cli_runner, isolated_home, agent, key):
cli_runner.invoke(cli, ["config", "set", key, "true"])
result = cli_runner.invoke(cli, ["config", "unset", key])
assert result.exit_code == 0, result.output
assert subscription_default(agent) is False
def test_turning_it_on_leaves_every_other_agent_off(self, cli_runner, isolated_home, agent, key):
cli_runner.invoke(cli, ["config", "set", key, "true"])
assert {name: subscription_default(name) for name in ("claude", "codex", "opencode", "pi")} == {
"claude": agent == "claude",
"codex": agent == "codex",
"opencode": False,
"pi": False,
}
class TestPkceSetting:

91
tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py
View file

@ -4332,6 +4332,97 @@ async def test_first_round_logged_after_tool_call_reuse_releases_only_at_request
assert await in_flight() == 0
@pytest.mark.asyncio
async def test_completed_mcp_replay_stream_frees_slot_its_logged_first_round_kept() -> None:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.anthropic.messages.native.fake_stream_iterator import (
FakeAnthropicMessagesStreamIterator,
)
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
reused_at: Final = datetime(2026, 1, 5, 12, 0, 0)
api_key: Final = hash_token("sk-mcp-replayed-stream")
local_cache: Final = DualCache()
handler: Final = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache),
time_provider=lambda: reused_at,
)
proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=DualCache())
proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = handler
user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1)
parallel_key: Final = f"{{api_key:{api_key}}}:max_parallel_requests"
owner_call_id: Final = "mcp-replayed-stream-owner"
messages: Final = [{"role": "user", "content": "add 2 and 3"}]
async def in_flight() -> int:
return handler._gauge_in_flight_from_cache_value(
await local_cache.async_get_cache(key=parallel_key)
)
async def pre_call(data: dict[str, object], call_type: str) -> None:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data=data,
call_type=call_type,
)
await pre_call(
{"model": "claude-haiku-5-5", "messages": messages, "stream": True, "litellm_call_id": owner_call_id},
CallTypes.anthropic_messages.value,
)
await pre_call(
{"model": "MCP: add", "messages": [{"role": "user", "content": "Tool: add"}], "mcp_server_name": "math"},
CallTypes.call_mcp_tool.value,
)
logging_obj: Final = LiteLLMLoggingObj(
model="claude-haiku-5-5",
messages=messages,
stream=True,
call_type=CallTypes.anthropic_messages.value,
start_time=None,
litellm_call_id=owner_call_id,
function_id="mcp-replayed-stream",
)
assert logging_obj.claim_async_success_log()
await handler.async_log_success_event(
kwargs={"litellm_call_id": owner_call_id, "stream": False},
response_obj=None,
start_time=None,
end_time=reused_at - timedelta(seconds=1),
)
assert await in_flight() == 1
replay: Final = FakeAnthropicMessagesStreamIterator(
response={
"id": "msg_replayed",
"type": "message",
"role": "assistant",
"model": "claude-haiku-5-5",
"content": [{"type": "text", "text": "2 + 3 = 5"}],
"stop_reason": "end_turn",
"usage": {"input_tokens": 20, "output_tokens": 8},
}
)
with _override_litellm_callbacks([]):
streamed: Final = [
chunk
async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator(
response=replay,
user_api_key_dict=user_api_key_dict,
request_data={
"model": "claude-haiku-5-5",
"litellm_call_id": owner_call_id,
"litellm_logging_obj": logging_obj,
},
proxy_logging_obj=proxy_logging_obj,
)
]
assert b"event: message_stop" in streamed[-1]
assert await in_flight() == 0
@pytest.mark.asyncio
async def test_mcp_tool_call_reusing_request_slot_keeps_request_max_parallel_headers() -> None:
api_key: Final = hash_token("sk-mcp-reuse-headers")

631
tests/unit/proxy/management_endpoints/test_model_management_endpoints.py
View file

@ -15,6 +15,7 @@ from fastapi.testclient import TestClient
from litellm._uuid import uuid
from litellm.models.credentials import CredentialItem
from litellm.models.project import LiteLLM_ProjectTable
from litellm.proxy._types import (
LiteLLM_ModelTable,
@ -31,6 +32,7 @@ from litellm.proxy.management_endpoints.model_management_endpoints import (
KubernetesPodsResponse,
ModelManagementAuthChecks,
_get_team_deployments,
_raise_if_model_update_adds_missing_rate_limits,
_raise_if_rate_limits_required_but_missing,
clear_cache,
delete_team_models,
@ -38,6 +40,7 @@ from litellm.proxy.management_endpoints.model_management_endpoints import (
patch_model,
update_model,
)
from litellm.proxy.management_helpers.model_quota_policy import REQUIRE_BOTH_LIMITS, ModelQuotaPolicy
from litellm.proxy.utils import PrismaClient
from litellm.router import Router
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo, updateDeployment, updateLiteLLMParams
@ -7369,12 +7372,19 @@ class TestAutoRouterClassifierDefaultPrompt:
assert response.system_prompt == classification_system_prompt(5)
_CUSTOMER_MODEL_ADD_POLICY: Final = ModelQuotaPolicy.model_validate(
{
"providers": {"vertex_ai": ["rpm", "tpm"], "runwayml": []},
"models": {"veo-*": ["rpm"], "jev-latest": []},
}
)
class TestEnforceRpmTpmOnModelAdd:
def test_passes_when_disabled_even_without_limits(self):
assert (
_raise_if_rate_limits_required_but_missing(
litellm_params=LiteLLM_Params(model="azure/gpt-5.2"),
enforced=False,
model_name="gpt-5.2", team_model=False, litellm_params=LiteLLM_Params(model="azure/gpt-5.2"), policy=None
)
is None
)
@ -7382,8 +7392,9 @@ class TestEnforceRpmTpmOnModelAdd:
def test_passes_when_enabled_and_both_set(self):
assert (
_raise_if_rate_limits_required_but_missing(
model_name="gpt-5.2", team_model=False,
litellm_params=LiteLLM_Params(model="azure/gpt-5.2", rpm=10, tpm=1000),
enforced=True,
policy=REQUIRE_BOTH_LIMITS,
)
is None
)
@ -7398,14 +7409,618 @@ class TestEnforceRpmTpmOnModelAdd:
(LiteLLM_Params(model="azure/gpt-5.2", rpm=10, tpm=-1), "tpm"),
],
)
def test_raises_when_enabled_and_missing(self, params, expected_missing):
from litellm.proxy._types import ProxyException
def test_raises_when_enabled_and_missing(self, params: LiteLLM_Params, expected_missing: str):
with pytest.raises(ProxyException) as exc_info:
_raise_if_rate_limits_required_but_missing(litellm_params=params, enforced=True)
assert expected_missing in str(exc_info.value.message)
_raise_if_rate_limits_required_but_missing(
model_name="gpt-5.2", team_model=False, litellm_params=params, policy=REQUIRE_BOTH_LIMITS
)
assert exc_info.value.message.startswith(f"{expected_missing} must be set")
assert exc_info.value.code == "400"
@pytest.mark.parametrize(
("model_name", "params"),
[
("runway-gen4", LiteLLM_Params(model="runwayml/gen4_turbo")),
("my-runway", LiteLLM_Params(model="gen4_turbo", custom_llm_provider="runwayml")),
("jev-latest", LiteLLM_Params(model="openai/jev-latest")),
("veo-3.1-generate-001", LiteLLM_Params(model="vertex_ai/veo-3.1-generate-001", rpm=10)),
("gemini-3-pro", LiteLLM_Params(model="vertex_ai/gemini-3-pro", rpm=10, tpm=1000)),
],
)
def test_policy_asks_each_model_only_for_the_limits_its_entry_requires(
self, model_name: str, params: LiteLLM_Params
):
assert (
_raise_if_rate_limits_required_but_missing(
model_name=model_name, team_model=False, litellm_params=params, policy=_CUSTOMER_MODEL_ADD_POLICY
)
is None
)
@pytest.mark.parametrize(
("model_name", "params", "expected_missing"),
[
("veo-3.1-generate-001", LiteLLM_Params(model="vertex_ai/veo-3.1-generate-001", tpm=1000), "rpm"),
("gemini-3-pro", LiteLLM_Params(model="vertex_ai/gemini-3-pro", rpm=10), "tpm"),
("openai-gpt", LiteLLM_Params(model="openai/gpt-5.2"), "rpm and tpm"),
("runway-gen4", LiteLLM_Params(model="runwayml/gen4_turbo", rpm=0), "rpm"),
],
)
def test_policy_rejects_a_model_missing_a_required_or_carrying_a_non_positive_limit(
self, model_name: str, params: LiteLLM_Params, expected_missing: str
):
with pytest.raises(ProxyException) as exc_info:
_raise_if_rate_limits_required_but_missing(
model_name=model_name, team_model=False, litellm_params=params, policy=_CUSTOMER_MODEL_ADD_POLICY
)
assert exc_info.value.message.startswith(f"{expected_missing} must be set to a positive value for '{model_name}'")
assert exc_info.value.code == "400"
@pytest.mark.parametrize("api_base", ["https://api.groq.com/openai/v1", "os.environ/QA_GROQ_API_BASE"])
def test_policy_reads_the_provider_behind_a_known_api_base(self, monkeypatch: pytest.MonkeyPatch, api_base: str):
monkeypatch.setenv("QA_GROQ_API_BASE", "https://api.groq.com/openai/v1")
policy: Final = ModelQuotaPolicy.model_validate({"default": [], "providers": {"groq": ["rpm", "tpm"]}})
with pytest.raises(ProxyException) as exc_info:
_raise_if_rate_limits_required_but_missing(
model_name="qa-llama", team_model=False,
litellm_params=LiteLLM_Params(model="qa-llama", api_base=api_base),
policy=policy,
)
assert exc_info.value.message.startswith("rpm and tpm must be set to a positive value for 'qa-llama'")
@pytest.mark.parametrize(
("model_name", "params", "shared_missing", "team_missing"),
[
pytest.param(
"jev-latest", LiteLLM_Params(model="vertex_ai/gemini-3.8-flash"), None, "rpm and tpm",
id="team-cannot-borrow-an-exempt-public-name",
),
pytest.param(
"video-team", LiteLLM_Params(model="vertex_ai/veo-3.1-generate-001", rpm=5), "tpm", None,
id="team-matches-the-model-it-calls",
),
],
)
def test_policy_matches_a_team_model_by_the_model_it_calls(
self, model_name: str, params: LiteLLM_Params, shared_missing: str | None, team_missing: str | None
):
assert _model_add_error(model_name, params, team_model=False) == _missing_limits_message(shared_missing, model_name)
assert _model_add_error(model_name, params, team_model=True) == _missing_limits_message(team_missing, model_name)
def _missing_limits_message(missing: str | None, model_name: str) -> str | None:
if missing is None:
return None
return (
f"{missing} must be set to a positive value for '{model_name}' when "
"'enforce_rpm_tpm_on_model_add' is enabled in general_settings"
)
def _model_add_error(model_name: str, params: LiteLLM_Params, team_model: bool) -> str | None:
try:
_raise_if_rate_limits_required_but_missing(
model_name=model_name, team_model=team_model, litellm_params=params, policy=_CUSTOMER_MODEL_ADD_POLICY
)
except ProxyException as error:
return error.message
return None
_STORED_RUNWAY: Final = {"model": "runwayml/gen4_turbo"}
def _stored_model(litellm_params: Mapping[str, object], team_id: str | None) -> Deployment:
return Deployment(
model_name="media-latest" if team_id is None else "model_name_media-team_team-row",
litellm_params=dict(litellm_params),
model_info={
"id": "team-row",
"team_id": team_id,
"team_public_model_name": None if team_id is None else "media-latest",
},
)
@pytest.mark.parametrize(
("stored", "team_id", "incoming", "missing"),
[
pytest.param(
_STORED_RUNWAY, "media-team", updateLiteLLMParams(model="vertex_ai/gemini-3.8-flash"), "rpm and tpm",
id="team-model-moves-to-a-provider-that-needs-limits",
),
pytest.param(
_STORED_RUNWAY, None, updateLiteLLMParams(model="vertex_ai/gemini-3.8-flash", rpm=10), "tpm",
id="proxy-model-moves-with-only-one-limit",
),
pytest.param(
{**_STORED_RUNWAY, "tpm": 1000}, None, updateLiteLLMParams(model="vertex_ai/gemini-3.8-flash", rpm=10),
None, id="stored-limit-counts-toward-the-move",
),
pytest.param(
_STORED_RUNWAY, "media-team", updateLiteLLMParams(model="vertex_ai/veo-3.1-generate-001", rpm=5), None,
id="team-model-moves-to-a-model-entry-it-meets",
),
pytest.param(
{"model": "vertex_ai/gemini-3.8-flash"}, None, updateLiteLLMParams(api_key="qa-rotated-key"), None,
id="unchanged-route-adds-no-limits",
),
pytest.param(
{"model": "vertex_ai/gemini-3.8-flash", "rpm": 10, "tpm": 1000}, None, updateLiteLLMParams(rpm=0), None,
id="unchanged-route-can-zero-a-limit-as-on-main",
),
pytest.param(
{"model": "vertex_ai/gemini-3.8-flash"}, None, updateLiteLLMParams(rpm=10), None,
id="legacy-model-without-tpm-stays-editable",
),
],
)
def test_model_update_requires_only_the_limits_its_new_route_adds(
stored: Mapping[str, object], team_id: str | None, incoming: updateLiteLLMParams, missing: str | None
) -> None:
assert _model_update_error(_stored_model(stored, team_id), incoming) == _missing_limits_message(
missing, "media-latest"
)
def _model_update_error(
existing: Deployment,
incoming: updateLiteLLMParams,
policy: ModelQuotaPolicy = _CUSTOMER_MODEL_ADD_POLICY,
remaining: LiteLLM_Params | None = None,
) -> str | None:
try:
_raise_if_model_update_adds_missing_rate_limits(
existing=existing,
written_name="media-latest",
written_team_model=existing.model_info.team_id is not None,
incoming_params=incoming,
remaining_params=existing.litellm_params if remaining is None else remaining,
policy=policy,
)
except ProxyException as error:
return error.message
return None
@pytest.mark.asyncio
@pytest.mark.parametrize(
("setting", "expected_code", "expected_message"),
[
({"default": [], "models": {"veo-*": ["rpm"]}}, "400", "rpm must be set to a positive value for 'veo-3.1'"),
("yes", "400", "rpm and tpm must be set to a positive value for 'veo-3.1'"),
({"enabled": True}, "400", "rpm and tpm must be set to a positive value for 'veo-3.1'"),
(
{"providers": {"vertex_ai": ["qps"]}},
"500",
"general_settings.enforce_rpm_tpm_on_model_add is invalid",
),
],
)
async def test_model_add_resolves_the_quota_policy_for_the_public_model_name(
monkeypatch: pytest.MonkeyPatch, setting: object, expected_code: str, expected_message: str
) -> None:
from litellm.proxy import proxy_server
from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model
prisma: Final = MagicMock()
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
monkeypatch.setattr(proxy_server, "general_settings", {"enforce_rpm_tpm_on_model_add": setting})
with pytest.raises(ProxyException) as error:
await add_new_model(
Deployment(model_name="veo-3.1", litellm_params={"model": "vertex_ai/veo-3.1"}),
UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert error.value.code == expected_code
assert error.value.message.startswith(expected_message)
prisma.db.litellm_proxymodeltable.create.assert_not_called()
@pytest.mark.asyncio
async def test_team_model_add_cannot_borrow_a_model_entry_by_its_public_name(monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.proxy import proxy_server
from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model
prisma: Final = MagicMock()
prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=LiteLLM_TeamTable(team_id="media-team"))
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
monkeypatch.setattr(proxy_server, "premium_user", True)
monkeypatch.setattr(
proxy_server,
"general_settings",
{"enforce_rpm_tpm_on_model_add": {"providers": {"vertex_ai": ["rpm", "tpm"]}, "models": {"jev-latest": []}}},
)
with pytest.raises(ProxyException) as error:
await add_new_model(
Deployment(
model_name="jev-latest",
litellm_params={"model": "vertex_ai/gemini-3.8-flash"},
model_info={"team_id": "media-team"},
),
UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert error.value.code == "400"
assert error.value.message.startswith("rpm and tpm must be set to a positive value for 'jev-latest'")
prisma.db.litellm_proxymodeltable.create.assert_not_called()
_PROJECT_QUOTA_SETTINGS: Final = {
"enforce_project_model_quota": {
"providers": {"vertex_ai": ["rpm", "tpm"], "runwayml": []},
"models": {"jev-latest": []},
}
}
_TEAM_VERTEX: Final = {"model": "vertex_ai/gemini-3.8-flash", "vertex_project": "qa-project"}
_TEAM_RUNWAY: Final = {"model": "gen4_turbo", "custom_llm_provider": "runwayml"}
def _media_team_db(public_name: str, project_metadata: Mapping[str, object]) -> MagicMock:
prisma: Final = MagicMock()
prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=LiteLLM_TeamTable(team_id="media-team"))
prisma.db.litellm_projecttable.find_many = AsyncMock(
return_value=[
LiteLLM_ProjectTable(
project_id="p-1",
project_alias="video-gen",
team_id="media-team",
models=[public_name],
metadata=dict(project_metadata),
),
LiteLLM_ProjectTable(
project_id="p-2", project_alias="chat-only", team_id="media-team", models=["chat-only-model"]
),
]
)
return prisma
def _team_model_row(public_name: str, litellm_params: Mapping[str, object]) -> LiteLLM_ProxyModelTable:
return LiteLLM_ProxyModelTable(
model_id="team-row",
model_name="model_name_media-team_team-row",
litellm_params=dict(litellm_params),
model_info={"id": "team-row", "team_id": "media-team", "team_public_model_name": public_name},
created_by="admin",
updated_by="admin",
)
def _quota_router(public_name: str, proxy_params: Mapping[str, object], *team_params: Mapping[str, object]) -> Router:
return Router(
model_list=[
{"model_name": public_name, "litellm_params": dict(proxy_params)},
*(
{
"model_name": "model_name_media-team_team-row",
"litellm_params": dict(params),
"model_info": {"id": "team-row", "team_id": "media-team", "team_public_model_name": public_name},
}
for params in team_params
),
]
)
def _use_quota_proxy(monkeypatch: pytest.MonkeyPatch, prisma: MagicMock, router: Router) -> None:
from litellm.proxy import proxy_server
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
monkeypatch.setattr(proxy_server, "general_settings", _PROJECT_QUOTA_SETTINGS)
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(proxy_server, "store_model_in_db", False)
monkeypatch.setattr(proxy_server, "premium_user", True)
_ADMIN: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
@pytest.mark.asyncio
async def test_team_model_add_rejects_a_route_its_projects_lack_the_quota_for(monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model
prisma: Final = _media_team_db("jev-latest", {})
_use_quota_proxy(monkeypatch, prisma, _quota_router("jev-latest", {"model": "openai/jev-latest"}))
with pytest.raises(ProxyException) as error:
await add_new_model(
Deployment(
model_name="jev-latest",
litellm_params={**_TEAM_VERTEX, "rpm": 10, "tpm": 1000},
model_info={"team_id": "media-team"},
),
_ADMIN,
)
assert error.value.code == "400"
assert error.value.message.startswith(
"projects video-gen list 'jev-latest' without a positive rpm and tpm quota, "
"which 'enforce_project_model_quota' requires once this team model change applies"
)
prisma.db.litellm_projecttable.find_many.assert_awaited_once_with(where={"team_id": "media-team"})
@pytest.mark.parametrize(
("litellm_params", "project_metadata"),
[
pytest.param(
_TEAM_VERTEX,
{"model_rpm_limit": {"jev-latest": 5}, "model_tpm_limit": {"jev-latest": 500}},
id="projects-carry-the-added-limits",
),
pytest.param(_TEAM_RUNWAY, {}, id="new-route-is-exempt"),
],
)
@pytest.mark.asyncio
async def test_team_model_add_goes_through_when_its_projects_meet_the_new_route(
monkeypatch: pytest.MonkeyPatch, litellm_params: Mapping[str, object], project_metadata: Mapping[str, object]
) -> None:
from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model
_use_quota_proxy(
monkeypatch, _media_team_db("jev-latest", project_metadata), _quota_router("jev-latest", {"model": "openai/jev-latest"})
)
with pytest.raises(ProxyException) as error:
await add_new_model(
Deployment(model_name="jev-latest", litellm_params=dict(litellm_params), model_info={"team_id": "media-team"}),
_ADMIN,
)
assert "STORE_MODEL_IN_DB" in error.value.message
@pytest.mark.parametrize("endpoint", ["patch", "put"])
@pytest.mark.asyncio
async def test_team_model_update_rejects_moving_its_route_where_its_projects_lack_quota(
monkeypatch: pytest.MonkeyPatch, endpoint: str
) -> None:
from litellm.proxy import proxy_server
runway_by_prefix: Final = {"model": "runwayml/gen4_turbo"}
prisma: Final = _media_team_db("jev-latest", {})
prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(
return_value=_team_model_row("jev-latest", runway_by_prefix)
)
_use_quota_proxy(monkeypatch, prisma, _quota_router("jev-latest", {"model": "openai/jev-latest"}, runway_by_prefix))
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
moved_to_vertex: Final = updateLiteLLMParams(model="vertex_ai/gemini-3.8-flash", vertex_project="qa-project")
write: Final = (
patch_model(
model_id="team-row", patch_data=updateDeployment(litellm_params=moved_to_vertex), user_api_key_dict=_ADMIN
)
if endpoint == "patch"
else update_model(
model_params=updateDeployment(litellm_params=moved_to_vertex, model_info=ModelInfo(id="team-row")),
user_api_key_dict=_ADMIN,
)
)
with pytest.raises(ProxyException) as error:
await write
assert error.value.code == "400"
assert error.value.message.startswith("projects video-gen list 'jev-latest' without a positive rpm and tpm quota")
def _team_route_router(public_name: str, team_params: Mapping[str, object]) -> Router:
return Router(
model_list=[
{"model_name": "jev-latest", "litellm_params": {"model": "openai/jev-latest"}},
{
"model_name": "model_name_media-team_team-row",
"litellm_params": dict(team_params),
"model_info": {"id": "team-row", "team_id": "media-team", "team_public_model_name": public_name},
},
],
model_group_alias={"media": "jev-latest"},
)
@pytest.mark.parametrize(
("public_name", "project_model", "team_params", "moved_to"),
[
pytest.param(
"video-*",
"video-gen-1",
{"model": "runwayml/*"},
updateLiteLLMParams(model="vertex_ai/*", vertex_project="qa-project"),
id="a-name-the-team-wildcard-serves",
),
pytest.param(
"jev-latest",
"media",
{"model": "runwayml/gen4_turbo"},
updateLiteLLMParams(model="vertex_ai/gemini-3.8-flash", vertex_project="qa-project"),
id="an-alias-of-the-team-name",
),
],
)
@pytest.mark.asyncio
async def test_team_model_patch_rechecks_projects_listing_any_name_it_serves(
monkeypatch: pytest.MonkeyPatch,
public_name: str,
project_model: str,
team_params: Mapping[str, object],
moved_to: updateLiteLLMParams,
) -> None:
from litellm.proxy import proxy_server
prisma: Final = _media_team_db(project_model, {})
prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=_team_model_row(public_name, team_params))
_use_quota_proxy(monkeypatch, prisma, _team_route_router(public_name, team_params))
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
with pytest.raises(ProxyException) as error:
await patch_model(
model_id="team-row", patch_data=updateDeployment(litellm_params=moved_to), user_api_key_dict=_ADMIN
)
assert error.value.code == "400"
assert error.value.message.startswith(
f"projects video-gen list '{project_model}' without a positive rpm and tpm quota"
)
prisma.db.litellm_proxymodeltable.update.assert_not_called()
@pytest.mark.asyncio
async def test_model_update_checks_the_limits_of_the_name_it_renames_to(monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.proxy import proxy_server
prisma: Final = MagicMock()
prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(
return_value=LiteLLM_ProxyModelTable(
model_id="vertex-row",
model_name="jev-latest",
litellm_params={"model": "vertex_ai/gemini-3.8-flash"},
model_info={"id": "vertex-row"},
created_by="admin",
updated_by="admin",
)
)
_use_quota_proxy(monkeypatch, prisma, Router(model_list=[]))
monkeypatch.setattr(
proxy_server,
"general_settings",
{"enforce_rpm_tpm_on_model_add": {"providers": {"vertex_ai": ["rpm", "tpm"]}, "models": {"jev-latest": []}}},
)
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
with pytest.raises(ProxyException) as error:
await update_model(
model_params=updateDeployment(
model_name="gemini-3.8-flash", litellm_params=updateLiteLLMParams(), model_info=ModelInfo(id="vertex-row")
),
user_api_key_dict=_ADMIN,
)
assert error.value.code == "400"
assert error.value.message.startswith("rpm and tpm must be set to a positive value for 'gemini-3.8-flash'")
prisma.db.litellm_proxymodeltable.update.assert_not_called()
@pytest.mark.parametrize(
("incoming", "remaining"),
[
pytest.param(updateLiteLLMParams(rpm=0), {"model": "vertex_ai/gemini-3.8-flash", "tpm": 1000}, id="zeroes-rpm"),
pytest.param(updateLiteLLMParams(), {"model": "vertex_ai/gemini-3.8-flash", "tpm": 1000}, id="removes-rpm"),
pytest.param(updateLiteLLMParams(model="runwayml/gen4_turbo"), {"model": "runwayml/gen4_turbo"}, id="moves"),
],
)
def test_model_update_under_the_boolean_setting_never_rejects_as_on_main(
incoming: updateLiteLLMParams, remaining: Mapping[str, object]
) -> None:
from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting
stored: Final = _stored_model({"model": "vertex_ai/gemini-3.8-flash", "rpm": 10, "tpm": 1000}, None)
policy: Final = parse_model_quota_setting(True)
assert policy is not None
assert _model_update_error(stored, incoming, policy, LiteLLM_Params(**remaining)) is None
@pytest.mark.asyncio
async def test_team_model_patch_rejects_a_removal_that_moves_its_route_where_its_projects_lack_quota(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.proxy import proxy_server
from litellm.types.proxy.management_endpoints.model_management_endpoints import PatchModelRequest
prisma: Final = _media_team_db("jev-latest", {})
prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=_team_model_row("jev-latest", _TEAM_RUNWAY))
_use_quota_proxy(monkeypatch, prisma, _quota_router("jev-latest", {"model": "openai/jev-latest"}, _TEAM_RUNWAY))
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
with pytest.raises(ProxyException) as error:
await patch_model(
model_id="team-row",
patch_data=PatchModelRequest(litellm_params_to_remove=("custom_llm_provider",)),
user_api_key_dict=_ADMIN,
)
assert error.value.code == "400"
assert error.value.message.startswith("projects video-gen list 'jev-latest' without a positive rpm and tpm quota")
@pytest.mark.asyncio
async def test_team_model_delete_rejects_falling_back_to_a_route_its_projects_lack_quota_for(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.proxy.management_endpoints.model_management_endpoints import ModelInfoDelete, delete_model
prisma: Final = _media_team_db("media-latest", {})
prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=_team_model_row("media-latest", _TEAM_RUNWAY))
_use_quota_proxy(monkeypatch, prisma, _quota_router("media-latest", _TEAM_VERTEX, _TEAM_RUNWAY))
with pytest.raises(ProxyException) as error:
await delete_model(model_info=ModelInfoDelete(id="team-row"), user_api_key_dict=_ADMIN)
assert error.value.code == "400"
assert error.value.message.startswith("projects video-gen list 'media-latest' without a positive rpm and tpm quota")
prisma.db.litellm_proxymodeltable.delete.assert_not_called()
@pytest.mark.asyncio
async def test_team_model_delete_goes_through_when_nothing_else_serves_its_name(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.proxy.management_endpoints.model_management_endpoints import ModelInfoDelete, delete_model
prisma: Final = _media_team_db("media-latest", {})
prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=_team_model_row("media-latest", _TEAM_RUNWAY))
only_the_team_model: Final = Router(
model_list=[
{
"model_name": "model_name_media-team_team-row",
"litellm_params": dict(_TEAM_RUNWAY),
"model_info": {"id": "team-row", "team_id": "media-team", "team_public_model_name": "media-latest"},
}
]
)
_use_quota_proxy(monkeypatch, prisma, only_the_team_model)
with pytest.raises(ProxyException) as error:
await delete_model(model_info=ModelInfoDelete(id="team-row"), user_api_key_dict=_ADMIN)
assert "STORE_MODEL_IN_DB" in str(error.value.message)
@pytest.mark.parametrize("endpoint", ["patch", "put"])
@pytest.mark.asyncio
async def test_team_model_update_rejects_moving_to_a_provider_that_needs_limits_it_lacks(
monkeypatch: pytest.MonkeyPatch, endpoint: str
) -> None:
from litellm.proxy import proxy_server
prisma: Final = _media_team_db("media-latest", {})
prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=_team_model_row("media-latest", _TEAM_RUNWAY))
_use_quota_proxy(monkeypatch, prisma, _quota_router("media-latest", _TEAM_RUNWAY, _TEAM_RUNWAY))
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
monkeypatch.setattr(
proxy_server,
"general_settings",
{"enforce_rpm_tpm_on_model_add": {"providers": {"vertex_ai": ["rpm", "tpm"], "runwayml": []}}},
)
moved_to_vertex: Final = updateLiteLLMParams(model="vertex_ai/gemini-3.8-flash", custom_llm_provider="vertex_ai")
write: Final = (
patch_model(
model_id="team-row", patch_data=updateDeployment(litellm_params=moved_to_vertex), user_api_key_dict=_ADMIN
)
if endpoint == "patch"
else update_model(
model_params=updateDeployment(litellm_params=moved_to_vertex, model_info=ModelInfo(id="team-row")),
user_api_key_dict=_ADMIN,
)
)
with pytest.raises(ProxyException) as error:
await write
assert error.value.code == "400"
assert error.value.message == _missing_limits_message("rpm and tpm", "media-latest")
prisma.db.litellm_proxymodeltable.update.assert_not_called()
class TestBlockModelResponseSerialization:
@pytest.mark.parametrize(

24
tests/unit/proxy/test_litellm_pre_call_utils.py
View file

@ -53,6 +53,7 @@ from litellm.proxy.litellm_pre_call_utils import (
add_provider_specific_headers_to_request,
check_if_token_is_service_account,
clean_headers,
forwarded_anthropic_oauth_token,
get_dynamic_logging_metadata,
get_metadata_variable_name,
move_guardrails_to_metadata,
@ -1393,6 +1394,8 @@ async def test_add_litellm_data_to_request_strips_string_encoded_admin_injection
("fallback_depth", True),
("_target_order", 2),
("attempted_targets", ["forged"]),
("_excluded_deployment_ids", ["pinned-deployment"]),
("_retry_skipped_deployment_ids", ["pinned-deployment"]),
],
)
async def test_add_litellm_data_to_request_strips_forged_fallback_hop_state(
@ -1477,6 +1480,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields():
"_managed_background_response": True,
"_guardrail_pipelines": [{"name": "spoofed"}],
"_pipeline_managed_guardrails": ["evaded"],
"_failover_excluded_ids": ["pinned-deployment"],
"safe_user_metadata": "kept",
}
data = {
@ -1537,6 +1541,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields():
"_managed_background_response",
"_guardrail_pipelines",
"_pipeline_managed_guardrails",
"_failover_excluded_ids",
}
assert "litellm_metadata" not in updated
for stripped_key in stripped_keys:
@ -5962,6 +5967,25 @@ def test_clean_headers_strips_x_api_key_when_byok_disabled():
assert "x-api-key" not in result
@pytest.mark.parametrize(
("headers", "expected_token"),
[
(
{"authorization": "Bearer sk-ant-oat01-caller", "x-litellm-api-key": "Bearer sk-proxy"},
"sk-ant-oat01-caller",
),
({"authorization": "Bearer sk-ant-oat01-caller"}, None),
({"authorization": "Bearer sk-ant-api03-caller", "x-litellm-api-key": "Bearer sk-proxy"}, None),
({"x-litellm-api-key": "Bearer sk-proxy"}, None),
],
ids=["oauth-beside-proxy-key", "oauth-used-as-proxy-key", "non-oauth-authorization", "no-authorization"],
)
def test_forwarded_anthropic_oauth_token_follows_chats_forwarding_rule(
headers: dict[str, str], expected_token: str | None
) -> None:
assert forwarded_anthropic_oauth_token(Headers(headers)) == expected_token
def test_clean_headers_strips_x_api_key_when_byok_enabled_but_x_api_key_was_auth_header():
"""
Anti-replay regression: even when forward_llm_provider_auth_headers=True,

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

@ -364,6 +364,7 @@ async def test_anthropic_messages_count_tokens_endpoint():
output_format=None,
context_management=None,
request_litellm_params=None,
client_oauth_token=None,
):
assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint"
assert request.model == "claude-3-sonnet-20240229"
@ -444,6 +445,7 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model():
output_format=None,
context_management=None,
request_litellm_params=None,
client_oauth_token=None,
):
assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint"
assert request.model == "gpt-4"
@ -1262,6 +1264,7 @@ async def test_anthropic_endpoint_returns_anthropic_error_format():
output_format=None,
context_management=None,
request_litellm_params=None,
client_oauth_token=None,
):
raise ProxyException(
message='{"detail":{"message":"Input is too long for requested model."}}',
@ -1319,6 +1322,7 @@ async def test_anthropic_endpoint_403_permission_error_format():
output_format=None,
context_management=None,
request_litellm_params=None,
client_oauth_token=None,
):
raise ProxyException(
message='{"Message":"Bearer Token has expired"}',
@ -1376,6 +1380,7 @@ async def test_anthropic_endpoint_429_rate_limit_error_format():
output_format=None,
context_management=None,
request_litellm_params=None,
client_oauth_token=None,
):
raise ProxyException(
message="Rate limit exceeded",

27
tests/unit/responses/litellm_completion_transformation/test_session_handler.py
View file

@ -569,6 +569,33 @@ async def test_message_history_reconstructs_every_turn_of_the_session_in_order()
assert result["litellm_session_id"] == session_id
@pytest.mark.asyncio
async def test_message_history_skips_the_turns_the_provider_rejected():
session_id: Final = "5c5f9a3e-1c86-4c0e-9d7c-0a54b8a0f2f1"
answered_request_id: Final = "chatcmpl-1111"
rejected_request_id: Final = "chatcmpl-2222"
answered_turn: Final = {
**_spend_log(answered_request_id, session_id, "What is the weather in Paris?", "Checking."),
"status": "success",
}
rejected_turn: Final = {
**_spend_log(rejected_request_id, session_id, "Never mind. Just say hello", ""),
"status": "failure",
"response": {},
}
fake_prisma_client: Final = _FakePrismaClient(results=[[answered_turn, rejected_turn]])
with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client):
result: Final = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(
answered_request_id
)
assert [(message.get("role"), message.get("content")) for message in result["messages"]] == [
("user", "What is the weather in Paris?"),
("assistant", "Checking."),
]
@pytest.mark.asyncio
async def test_session_lookup_stops_retrying_once_the_budget_is_spent(
instant_session_lookup_retries: None,

53
tests/unit/responses/mcp/test_mcp_streaming_iterator.py
View file

@ -2,7 +2,7 @@ import sys
import types
from contextlib import nullcontext
from types import SimpleNamespace
from typing import Final
from typing import Final, cast
from unittest.mock import AsyncMock, MagicMock
import pytest
@ -66,6 +66,14 @@ def _item_type(item: dict[str, object] | BaseLiteLLMOpenAIResponseObject) -> str
return str(item["type"]) if isinstance(item, dict) else str(item.type)
def _follow_up_tool_history(input_items: list[object]) -> tuple[tuple[str, str], ...]:
return tuple(
(str(item.get("type")), str(item.get("call_id")))
for item in input_items
if isinstance(item, dict) and item.get("type") in ("function_call", "function_call_output")
)
def _tool_call_stream(call_id: str, tool_name: str, response_id: str = "resp-1") -> _FakeAsyncStream:
return _FakeAsyncStream([_completed_chunk([_function_call(call_id, tool_name)], response_id=response_id)])
@ -157,6 +165,49 @@ async def test_second_round_tool_call_is_executed_and_reaches_final_text(monkeyp
assert final_output[-1]["content"][0]["text"] == "Here's what I found after retrying."
@pytest.mark.asyncio
async def test_follow_up_input_keeps_every_previous_tool_round(monkeypatch: pytest.MonkeyPatch) -> None:
_mock_mcp_environment(monkeypatch)
aresponses_mock: Final = AsyncMock(
side_effect=[
_tool_call_stream("call_2", "read_wiki_contents"),
_tool_call_stream("call_3", "read_wiki_contents"),
_text_only_stream("All three tool rounds are done."),
]
)
monkeypatch.setattr(responses_main_module, "aresponses", aresponses_mock)
iterator: Final = _make_iterator(
[
_output_item_added_chunk(),
_completed_chunk([_function_call("call_1", "read_wiki_contents")]),
]
)
_ = [chunk async for chunk in iterator]
second_follow_up_input: Final = cast("list[object]", aresponses_mock.call_args_list[1].kwargs["input"])
third_follow_up_input: Final = cast("list[object]", aresponses_mock.call_args_list[2].kwargs["input"])
assert _follow_up_tool_history(second_follow_up_input) == (
("function_call", "call_1"),
("function_call_output", "call_1"),
("function_call", "call_2"),
("function_call_output", "call_2"),
)
assert _follow_up_tool_history(third_follow_up_input) == (
("function_call", "call_1"),
("function_call_output", "call_1"),
("function_call", "call_2"),
("function_call_output", "call_2"),
("function_call", "call_3"),
("function_call_output", "call_3"),
)
user_messages: Final = [
item for item in third_follow_up_input if isinstance(item, dict) and item.get("role") == "user"
]
assert user_messages == [{"type": "message", "role": "user", "content": "what is berriai/litellm?"}]
@pytest.mark.asyncio
@pytest.mark.parametrize(
("max_agentic_loops", "expected_cap"),

161
tests/unit/responses/test_responses_api_bridge_flag.py
View file

@ -423,6 +423,149 @@ class TestUseResponsesApiBridgeFlag:
assert request.headers["authorization"] == "Bearer fake-azure-ad-token"
assert "chat_template_kwargs" not in json.loads(request.read())
def test_bridge_rejects_a_responses_field_it_cannot_translate_when_drop_params_is_off(
self, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
):
monkeypatch.setattr(litellm, "drop_params", False)
upstream: Final = _databricks_chat_upstream(respx_mock)
with pytest.raises(litellm.UnsupportedParamsError) as raised:
litellm.responses(
model="databricks/my-custom-model",
input="Hello",
max_tool_calls=5,
api_base="https://example.databricks.test/serving-endpoints",
api_key="fake-databricks-api-key",
num_retries=0,
)
assert upstream.call_count == 0, "the provider must not be called for a field the bridge cannot translate"
assert raised.value.status_code == 400
assert "max_tool_calls" in str(raised.value)
assert "drop_params" in str(raised.value)
@pytest.mark.parametrize("drop_params_source", ["request", "global"])
def test_bridge_drops_a_responses_field_it_cannot_translate_when_drop_params_is_on(
self, drop_params_source: str, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
):
monkeypatch.setattr(litellm, "drop_params", drop_params_source == "global")
upstream: Final = _databricks_chat_upstream(respx_mock)
response: Final = litellm.responses(
model="databricks/my-custom-model",
input="Hello",
max_tool_calls=5,
chat_template_kwargs={"thinking": True},
api_base="https://example.databricks.test/serving-endpoints",
api_key="fake-databricks-api-key",
num_retries=0,
**({"drop_params": True} if drop_params_source == "request" else {}),
)
assert upstream.call_count == 1
assert json.loads(upstream.calls[0].request.read()) == {
"model": "my-custom-model",
"messages": [{"role": "user", "content": "Hello"}],
"chat_template_kwargs": {"thinking": True},
}
assert response.output[0].content[0].text == "Answer"
def test_bridge_drops_a_responses_field_it_cannot_translate_when_the_deployment_sets_drop_params(
self, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
):
monkeypatch.setattr(litellm, "drop_params", False)
upstream: Final = _databricks_chat_upstream(respx_mock)
router: Final = Router(
model_list=[
{
"model_name": "bridged",
"litellm_params": {
"model": "databricks/my-custom-model",
"api_base": "https://example.databricks.test/serving-endpoints",
"api_key": "fake-databricks-api-key",
"drop_params": True,
},
}
],
num_retries=0,
)
response: Final = router.responses(model="bridged", input="Hello", max_tool_calls=5)
assert upstream.call_count == 1
request_body: Final = json.loads(upstream.calls[0].request.read())
assert "max_tool_calls" not in request_body
assert request_body["messages"] == [{"role": "user", "content": "Hello"}]
assert response.output[0].content[0].text == "Answer"
def test_allowed_openai_params_sends_a_responses_field_through_the_bridge(
self, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
):
monkeypatch.setattr(litellm, "drop_params", False)
upstream: Final = _databricks_chat_upstream(respx_mock)
litellm.responses(
model="databricks/my-custom-model",
input="Hello",
max_tool_calls=5,
allowed_openai_params=["max_tool_calls"],
api_base="https://example.databricks.test/serving-endpoints",
api_key="fake-databricks-api-key",
num_retries=0,
)
assert upstream.call_count == 1
assert json.loads(upstream.calls[0].request.read())["max_tool_calls"] == 5
def test_bridge_still_honors_additional_drop_params_for_a_responses_field_it_cannot_translate(
self, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
):
monkeypatch.setattr(litellm, "drop_params", False)
upstream: Final = _databricks_chat_upstream(respx_mock)
response: Final = litellm.responses(
model="databricks/my-custom-model",
input="Hello",
max_tool_calls=5,
additional_drop_params=["max_tool_calls"],
api_base="https://example.databricks.test/serving-endpoints",
api_key="fake-databricks-api-key",
num_retries=0,
)
assert upstream.call_count == 1
assert json.loads(upstream.calls[0].request.read()) == {
"model": "my-custom-model",
"messages": [{"role": "user", "content": "Hello"}],
}
assert response.output[0].content[0].text == "Answer"
def test_bridge_still_quietly_drops_responses_fields_it_already_handles(
self, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
):
monkeypatch.setattr(litellm, "drop_params", False)
upstream: Final = _databricks_chat_upstream(respx_mock)
litellm.responses(
model="databricks/my-custom-model",
input="Hello",
store=False,
include=["reasoning.encrypted_content"],
truncation="auto",
metadata={"k": "v"},
background=False,
text={"verbosity": "low"},
api_base="https://example.databricks.test/serving-endpoints",
api_key="fake-databricks-api-key",
num_retries=0,
)
assert upstream.call_count == 1
assert json.loads(upstream.calls[0].request.read()) == {
"model": "my-custom-model",
"messages": [{"role": "user", "content": "Hello"}],
}
@patch.object(import_module("litellm.responses.file_search.emulated_handler"), "_call_aresponses")
@patch.object(import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config")
async def test_bridge_flag_forwarded_to_file_search_emulation(self, mock_get_config, mock_call_aresponses):
@ -576,6 +719,24 @@ _CHAT: Final[Mapping[str, JsonValue]] = {"id": "dep-1", "mode": "chat"}
_RESPONSES: Final[Mapping[str, JsonValue]] = {"id": "dep-1", "mode": "responses"}
def _databricks_chat_upstream(respx_mock: respx.MockRouter) -> respx.Route:
return respx_mock.post("https://example.databricks.test/serving-endpoints/chat/completions").mock(
return_value=httpx.Response(
status_code=200,
json={
"id": "chatcmpl-123",
"object": "chat.completion",
"created": 1677652288,
"model": "my-custom-model",
"choices": [
{"index": 0, "message": {"role": "assistant", "content": "Answer"}, "finish_reason": "stop"}
],
"usage": {"prompt_tokens": 9, "completion_tokens": 1, "total_tokens": 10},
},
)
)
def _pin_shared_mode(monkeypatch: pytest.MonkeyPatch, model: str, mode: str) -> None:
cost_map: Final = cast(dict[str, dict[str, object]], litellm.model_cost)
monkeypatch.setitem(cost_map, model, {**cost_map.get(model, {}), "mode": mode})

250
tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py
View file

@ -4,13 +4,24 @@ import litellm
from litellm.router_utils.pre_call_checks.continuation_prefill_check import (
MID_STREAM_CONTINUATION_KWARG,
MID_STREAM_CONTINUATION_MARKER,
RESPONSES_CONTINUATION_KWARG,
ContinuationOutput,
ContinuationPrefillDeploymentCheck,
ResponsesContinuation,
ResponsesContinuationSeam,
constrains_continuation,
constrains_responses_continuation,
continuation_output_ceilings,
deployment_supports_prefill,
deployment_takes_responses_prefill,
responses_continuation_input,
responses_continuation_input_for,
)
from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents
PREFILL_MODEL = "anthropic/prefill-capable-test-model"
CLAUDE_PREFILL_MODEL = "anthropic/claude-prefill-capable-test-model"
CLAUDE_NON_PREFILL_MODEL = "anthropic/claude-prefill-rejecting-test-model"
NON_PREFILL_MODEL = "openai/prefill-unknown-test-model"
UNMAPPED_MODEL = "openai/unmapped-test-model"
@ -21,6 +32,12 @@ def _synthetic_cost_map_entries() -> None:
{
PREFILL_MODEL: {"litellm_provider": "anthropic", "mode": "chat", "supports_assistant_prefill": True},
NON_PREFILL_MODEL: {"litellm_provider": "openai", "mode": "chat"},
CLAUDE_PREFILL_MODEL: {"litellm_provider": "anthropic", "mode": "chat", "supports_assistant_prefill": True},
CLAUDE_NON_PREFILL_MODEL: {
"litellm_provider": "anthropic",
"mode": "chat",
"supports_assistant_prefill": False,
},
},
persist_across_reloads=False,
)
@ -191,7 +208,11 @@ def _text_chunk(text: str) -> litellm.ModelResponseStream:
("Hello", (" world",), (" world",)),
("Hello", ("Hel", "p me"), ("Hel", "p me")),
("Hello", ("Hel",), ("Hel",)),
([{"type": "text", "text": "Sure: "}, {"type": "text", "text": "Hello"}], ("Sure: Hello", " world"), ("", " world")),
(
[{"type": "text", "text": "Sure: "}, {"type": "text", "text": "Hello"}],
("Sure: Hello", " world"),
("", " world"),
),
],
ids=[
"echo_split_across_chunks_is_dropped",
@ -206,9 +227,7 @@ def _text_chunk(text: str) -> litellm.ModelResponseStream:
async def test_continuation_output_drops_only_a_full_echo_of_the_prefill(
prefill_content, fallback_texts, expected_texts
):
output = ContinuationOutput(
"Hello", {"role": "assistant", "content": prefill_content, "prefix": True}, "group", ()
)
output = ContinuationOutput("Hello", {"role": "assistant", "content": prefill_content, "prefix": True}, "group", ())
async def fallback_stream():
for text in fallback_texts:
@ -266,7 +285,14 @@ async def test_continuation_output_drops_only_whitespace_the_client_already_has(
({"extra_body": {"max_completion_tokens": 45}}, {"extra_body": {"max_completion_tokens": 15}}),
({"extra_body": {"top_k": 3}}, {}),
],
ids=["max_tokens", "both_ceilings", "no_ceiling", "extra_body_max_tokens", "extra_body_only", "extra_body_no_ceiling"],
ids=[
"max_tokens",
"both_ceilings",
"no_ceiling",
"extra_body_max_tokens",
"extra_body_only",
"extra_body_no_ceiling",
],
)
def test_continuation_output_ceilings_take_the_emitted_tokens_off_every_ceiling_the_request_sends(
request_kwargs, expected
@ -286,3 +312,217 @@ def test_continuation_output_ceilings_take_the_emitted_tokens_off_every_ceiling_
)
def test_continuation_output_ceilings_decline_once_a_ceiling_has_no_tokens_left(request_kwargs):
assert continuation_output_ceilings(request_kwargs, 30) is None
@pytest.mark.parametrize(
"sent,expected",
[
({"text": {"format": {"type": "json_schema", "name": "n", "schema": {"type": "object"}}}}, True),
({"text": {"format": {"type": "json_object"}}}, True),
({"text": {"format": {"type": "text"}}}, False),
({"text": {"verbosity": "low"}}, False),
({"tool_choice": "required"}, True),
({"tool_choice": "auto"}, False),
({"tool_choice": {"type": "function", "name": "lookup"}}, True),
({"tool_choice": {"type": "allowed_tools", "mode": "required", "tools": [{"type": "function", "name": "f"}]}}, True),
({"tool_choice": {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "f"}]}}, True),
({"tool_choice": {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": []}}}, True),
({"tool_choice": {"type": "web_search_preview"}}, True),
],
ids=[
"json_schema",
"json_object",
"text_format",
"verbosity_only",
"required_tool",
"auto_tool",
"forced_function",
"allowed_tools_required",
"allowed_tools_auto",
"chat_allowed_tools_auto",
"hosted_tool",
],
)
def test_constrains_continuation_reads_the_responses_text_format(sent, expected):
assert constrains_continuation(sent) is expected
@pytest.mark.parametrize(
"sent,expected",
[
({"tool_choice": {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "f"}]}}, False),
({"tool_choice": {"type": "allowed_tools", "mode": "required", "tools": [{"type": "function", "name": "f"}]}}, True),
({"tool_choice": {"type": "function", "name": "lookup"}}, True),
({"tool_choice": {"type": "web_search_preview"}}, True),
({"tool_choice": "required"}, True),
({"tool_choice": "auto"}, False),
(
{
"tool_choice": {"type": "allowed_tools", "mode": "auto", "tools": []},
"text": {"format": {"type": "json_object"}},
},
True,
),
],
ids=[
"allowed_tools_auto",
"allowed_tools_required",
"forced_function",
"hosted_tool",
"required_tool",
"auto_tool",
"allowed_tools_auto_with_json",
],
)
def test_constrains_responses_continuation_lets_allowed_tools_in_auto_mode_answer_in_text(sent, expected):
assert constrains_responses_continuation(sent) is expected
@pytest.mark.parametrize(
"deployment,expected",
[
(_deployment(CLAUDE_PREFILL_MODEL, "a"), True),
(_deployment(CLAUDE_NON_PREFILL_MODEL, "b"), False),
(_deployment(PREFILL_MODEL, "c"), False),
(
{"litellm_params": {"model": "anthropic/support-model-v1"}, "model_info": {"id": "d", "base_model": CLAUDE_PREFILL_MODEL}},
True,
),
],
ids=["claude_with_prefill", "claude_without_prefill", "prefill_capable_but_not_claude", "claude_named_by_base_model"],
)
def test_deployment_takes_responses_prefill_only_on_claude_with_prefill(deployment, expected):
assert deployment_takes_responses_prefill(deployment) is expected
def _continuation_kwargs() -> dict[str, object]:
continuation = ResponsesContinuation("Count to 5.", "1, 2, ")
return {RESPONSES_CONTINUATION_KWARG: continuation, "input": responses_continuation_input(continuation)}
@pytest.mark.parametrize(
"prior_input,expected_prior",
[
(
"Count to 5.",
[{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Count to 5."}]}],
),
(
[
{"role": "user", "content": "Count to 5."},
{"type": "function_call_output", "call_id": "c", "output": "x"},
],
[
{"role": "user", "content": "Count to 5."},
{"type": "function_call_output", "call_id": "c", "output": "x"},
],
),
(None, []),
],
ids=["string_input", "item_list_input", "no_input"],
)
def test_responses_continuation_input_keeps_the_prior_input_then_the_partial_then_a_request_to_continue(
prior_input, expected_prior
):
built = responses_continuation_input(ResponsesContinuation(prior_input, "1, 2, "))
partial = {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "1, 2, "}]}
assert built[:-1] == [*expected_prior, partial]
assert built[-1]["role"] == "user"
assert "already ends with whitespace" in built[-1]["content"][0]["text"]
def test_responses_continuation_request_asks_for_a_leading_space_only_after_a_word():
built = responses_continuation_input(ResponsesContinuation("Hi", "Transformers use"))
assert "begin with a space" in built[-1]["content"][0]["text"]
def test_responses_continuation_accumulates_across_hops():
continuation = ResponsesContinuation("Count to 5.", "1, ").extended("2, ")
assert continuation == ResponsesContinuation("Count to 5.", "1, 2, ")
@pytest.mark.parametrize(
"deployment,expected_roles",
[
(_deployment(CLAUDE_PREFILL_MODEL, "a"), ["user", "assistant"]),
(_deployment(CLAUDE_NON_PREFILL_MODEL, "b"), ["user", "assistant", "user"]),
(_deployment(NON_PREFILL_MODEL, "c"), ["user", "assistant", "user"]),
],
ids=["claude_prefill", "claude_without_prefill", "openai"],
)
def test_responses_continuation_input_for_shapes_the_input_for_the_deployment_that_serves_it(
deployment, expected_roles
):
shaped = responses_continuation_input_for(_continuation_kwargs(), deployment)
assert [item["role"] for item in shaped] == expected_roles
def test_responses_continuation_input_for_keeps_what_the_router_stripped_from_the_input():
continuation = ResponsesContinuation(
[
{"type": "reasoning", "id": "rs_1", "summary": [], "encrypted_content": "opaque"},
{"role": "user", "content": "Count to 5."},
],
"1, 2, ",
)
stripped_input = responses_continuation_input(continuation)[1:]
shaped = responses_continuation_input_for(
{RESPONSES_CONTINUATION_KWARG: continuation, "input": stripped_input}, _deployment(CLAUDE_PREFILL_MODEL, "a")
)
assert shaped == stripped_input[:-1]
def test_responses_continuation_input_for_leaves_an_input_without_its_request_to_continue_alone():
kwargs = {**_continuation_kwargs(), "input": [{"role": "user", "content": "Count to 5."}]}
assert responses_continuation_input_for(kwargs, _deployment(CLAUDE_PREFILL_MODEL, "a")) == kwargs["input"]
def test_responses_continuation_input_for_leaves_a_request_without_a_continuation_alone():
kwargs = {"input": "Count to 5."}
assert responses_continuation_input_for(kwargs, _deployment(CLAUDE_PREFILL_MODEL, "a")) == "Count to 5."
def _delta(text: str) -> OutputTextDeltaEvent:
return OutputTextDeltaEvent(
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, item_id="m", output_index=0, content_index=0, delta=text
)
@pytest.mark.parametrize(
"emitted,fallback_texts,expected_texts",
[
("1, 2, ", (" 3, 4", ", 5"), ("3, 4", ", 5")),
("1, 2, ", (" ", " 3"), ("", " 3")),
("1, 2,", (" 3",), (" 3",)),
("1, 2, ", ("3",), ("3",)),
("A.\n\n", ("\n Hi",), (" Hi",)),
("", (" Hi",), (" Hi",)),
],
ids=[
"re_sent_space_is_dropped",
"seam_split_across_events",
"no_trailing_whitespace_keeps_the_space",
"no_re_sent_space_passes_through",
"seam_ends_at_the_first_divergence",
"fresh_start_passes_through",
],
)
@pytest.mark.asyncio
async def test_responses_continuation_seam_drops_only_whitespace_the_client_already_has(
emitted, fallback_texts, expected_texts
):
async def fallback_stream():
for text in fallback_texts:
yield _delta(text)
released = [event async for event in ResponsesContinuationSeam(emitted).splice(fallback_stream())]
assert tuple(event.delta for event in released) == expected_texts

90
tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py
View file

@ -9,6 +9,7 @@ import pytest
import litellm
from litellm.caching.dual_cache import DualCache
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
DeploymentAffinityCheck,
)
@ -1209,3 +1210,92 @@ async def test_virtual_key_hash_wins_over_user_id_for_affinity():
parent_otel_span=None,
)
assert other_key_same_user == healthy_deployments
class _ForbiddenError(Exception):
status_code: Final = 403
@pytest.mark.asyncio
@pytest.mark.parametrize(
("pre_call_check", "metadata"),
[
("session_affinity", {"session_id": "sess-retry"}),
("deployment_affinity", {"user_api_key_hash": "hashed-key-retry"}),
],
)
async def test_retry_reaches_healthy_sibling_when_affinity_pins_deployment_that_refused(
pre_call_check: str, metadata: dict[str, str]
):
router: Final = litellm.Router(
model_list=[
{
"model_name": "test-model",
"litellm_params": {
"model": "gpt-4o",
"api_key": "bad",
"mock_response": _ForbiddenError("A refuses"),
"weight": 1,
},
"model_info": {"id": "A"},
},
{
"model_name": "test-model",
"litellm_params": {
"model": "gpt-4o",
"api_key": "good",
"mock_response": "ok from B",
"weight": 0,
},
"model_info": {"id": "B"},
},
],
routing_strategy="simple-shuffle",
num_retries=1,
retry_after=0,
disable_cooldowns=True,
optional_pre_call_checks=[pre_call_check],
)
responses: Final = [
await router.acompletion(
model="test-model",
messages=[{"role": "user", "content": "hi"}],
metadata=dict(metadata),
)
for _ in range(2)
]
assert [response._hidden_params["model_id"] for response in responses] == ["B", "B"]
assert [response._hidden_params["additional_headers"]["x-litellm-attempted-retries"] for response in responses] == [
1,
1,
]
@pytest.mark.asyncio
async def test_previous_response_id_pin_holds_even_after_its_deployment_failed():
callback: Final = DeploymentAffinityCheck(
cache=DualCache(),
ttl_seconds=60,
enable_user_key_affinity=True,
enable_responses_api_affinity=True,
)
healthy_deployments: Final = [
{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}, "model_info": {"id": "A"}},
{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}, "model_info": {"id": "B"}},
]
assert await callback.async_filter_deployments(
model="gpt-4o",
healthy_deployments=healthy_deployments,
messages=None,
request_kwargs={
"previous_response_id": ResponsesAPIRequestUtils._build_responses_api_response_id(
custom_llm_provider="openai", model_id="A", response_id="resp_upstream"
),
"_excluded_deployment_ids": ["A"],
"metadata": {"user_api_key_hash": "hashed-key"},
},
parent_otel_span=None,
) == [healthy_deployments[0]]

109
tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py
View file

@ -16,6 +16,7 @@ The mechanism works without any cache and supports two encoding strategies:
"""
import time
from typing import Final
from unittest.mock import AsyncMock, patch
import pytest
@ -2606,3 +2607,111 @@ async def test_affinity_strips_for_team_and_pattern_routes():
assert result is sibling_pool
assert not request_kwargs["input"][1].get("encrypted_content")
def _region_deployment(region: str) -> dict[str, object]:
return {
"model_name": "gpt-5.4",
"litellm_params": {
"model": "openai/gpt-5.4",
"api_base": f"https://{region}.example.com/v1",
"api_key": f"key-{region}",
},
"model_info": {"id": region},
}
def _reasoning_input_minted_by(origin: str) -> list[dict[str, object]]:
return [
{
"type": "reasoning",
"id": ResponsesAPIRequestUtils._build_encrypted_item_id(origin, "rs_test"),
"encrypted_content": ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id("gAAAAA-blob", origin),
"summary": [{"type": "summary_text", "text": "scattering"}],
},
{"role": "user", "content": "and sunsets?"},
]
def _previous_response_id_from(origin: str) -> str:
return ResponsesAPIRequestUtils._build_responses_api_response_id("openai", origin, "resp_test")
_ORIGIN_FAILED_THIS_REQUEST: Final = pytest.mark.parametrize(
"failed_in_request",
[{"_excluded_deployment_ids": ["region-a"]}, {"_retry_skipped_deployment_ids": ("region-a",)}],
)
@pytest.mark.asyncio
@_ORIGIN_FAILED_THIS_REQUEST
@pytest.mark.parametrize("continuing", [{}, {"previous_response_id": _previous_response_id_from("region-b")}])
async def test_affinity_strips_and_dispatches_when_origin_already_failed_this_request(
failed_in_request: dict[str, object], continuing: dict[str, object]
):
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
EncryptedContentAffinityCheck,
)
router: Final = litellm.Router(model_list=[_region_deployment("region-a"), _region_deployment("region-b")])
group: Final = router.model_list
reasoning_input: Final = _reasoning_input_minted_by("region-a")
result: Final = await EncryptedContentAffinityCheck(router=router).async_filter_deployments(
model="gpt-5.4",
healthy_deployments=group,
messages=None,
request_kwargs={"litellm_metadata": {}, "input": reasoning_input, **failed_in_request, **continuing},
)
assert result is group
assert reasoning_input[0] == {"type": "reasoning", "summary": [{"type": "summary_text", "text": "scattering"}]}
@pytest.mark.asyncio
@_ORIGIN_FAILED_THIS_REQUEST
async def test_affinity_keeps_origin_and_reasoning_when_previous_response_id_binds_the_failed_origin(
failed_in_request: dict[str, object],
):
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
EncryptedContentAffinityCheck,
)
pool: Final = (_region_deployment("region-a"), _region_deployment("region-b"))
reasoning_input: Final = _reasoning_input_minted_by("region-a")
original_reasoning: Final = dict(reasoning_input[0])
assert await EncryptedContentAffinityCheck(router=litellm.Router(model_list=list(pool))).async_filter_deployments(
model="gpt-5.4",
healthy_deployments=list(pool),
messages=None,
request_kwargs={
"litellm_metadata": {},
"input": reasoning_input,
"previous_response_id": _previous_response_id_from("region-a"),
**failed_in_request,
},
) == [pool[0]]
assert reasoning_input[0] == original_reasoning
@pytest.mark.asyncio
@pytest.mark.parametrize("regions", [("region-a",), ("region-a", "region-b")])
async def test_affinity_keeps_reasoning_when_every_deployment_already_refused_the_request(regions: tuple[str, ...]):
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
EncryptedContentAffinityCheck,
)
router: Final = litellm.Router(model_list=[_region_deployment(region) for region in regions])
reasoning_input: Final = _reasoning_input_minted_by("region-a")
original_reasoning: Final = dict(reasoning_input[0])
result: Final = await EncryptedContentAffinityCheck(router=router).async_filter_deployments(
model="gpt-5.4",
healthy_deployments=router.model_list,
messages=None,
request_kwargs={"litellm_metadata": {}, "input": reasoning_input, "_retry_skipped_deployment_ids": regions},
)
assert [deployment["model_info"]["id"] for deployment in result] == ["region-a"]
assert reasoning_input[0] == original_reasoning

21
tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py
View file

@ -313,6 +313,27 @@ async def test_async_filter_deployments_does_not_pin_when_target_order_is_set():
assert filtered == deployments
@pytest.mark.asyncio
@pytest.mark.parametrize(
"failed_in_request",
[{"_excluded_deployment_ids": ["dep-2"]}, {"_retry_skipped_deployment_ids": ("dep-2",)}],
)
async def test_async_filter_deployments_does_not_pin_to_a_deployment_that_failed_this_request(
failed_in_request: dict[str, object],
):
cache: Final = DualCache()
check: Final = PromptCachingDeploymentCheck(cache=cache)
messages: Final = _messages(word_count=5000)
await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=messages, tools=None)
assert await check.async_filter_deployments(
model=MODEL_GROUP_ALIAS,
healthy_deployments=_deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6"),
messages=messages,
request_kwargs=failed_in_request,
) == _deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6")
@pytest.mark.asyncio
async def test_async_filter_deployments_narrows_for_group_whose_model_minimum_is_lower():
"""

141
tests/unit/router_utils/test_cooldown_handlers.py
View file

@ -1,10 +1,12 @@
import asyncio
import copy
import datetime
import importlib
import itertools
import random
import time
from collections.abc import Mapping
import uuid
from collections.abc import Mapping, Sequence
from typing import Final
from unittest.mock import MagicMock, patch
@ -682,6 +684,45 @@ class TestDeploymentLevelAllowedFails:
)
def _gpt_4_group_with_a_sibling() -> Sequence[Mapping[str, object]]:
return (
{"model_name": "gpt-4", "litellm_params": {"model": "openai/gpt-4"}, "model_info": {"id": "primary"}},
{"model_name": "gpt-4", "litellm_params": {"model": "openai/gpt-4"}, "model_info": {"id": "secondary"}},
)
def _primary_cooldown_decisions_on_repeated_503s(router: Router) -> tuple[bool, ...]:
exception: Final = litellm.ServiceUnavailableError("down", "openai", "gpt-4")
return tuple(
_should_cooldown_deployment(
litellm_router_instance=router,
deployment="primary",
exception_status=503,
original_exception=exception,
)
for _ in range(litellm.allowed_fails + 1)
)
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
class TestRouterLevelAllowedFailsEqualToDefault:
def test_allowed_fails_equal_to_the_default_cools_down_after_that_many_fails(self):
router: Final = _make_router(model_list=list(_gpt_4_group_with_a_sibling()), allowed_fails=litellm.allowed_fails)
assert _primary_cooldown_decisions_on_repeated_503s(router) == (False,) * litellm.allowed_fails + (True,)
def test_unset_allowed_fails_keeps_the_rate_based_default(self):
router: Final = _make_router(model_list=list(_gpt_4_group_with_a_sibling()))
assert _primary_cooldown_decisions_on_repeated_503s(router) == (False,) * (litellm.allowed_fails + 1)
def test_default_allowed_fails_saved_through_update_settings_keeps_the_rate_based_default(self):
router: Final = _make_router(model_list=list(_gpt_4_group_with_a_sibling()))
router.update_settings(allowed_fails=litellm.allowed_fails)
assert _primary_cooldown_decisions_on_repeated_503s(router) == (False,) * (litellm.allowed_fails + 1)
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
class TestDeploymentLevelAllowedFailsPolicyByExceptionType:
def test_rate_limit_error_triggers_cooldown_with_zero_threshold(self):
@ -2931,6 +2972,104 @@ async def test_bad_request_error_does_not_cool_down_a_deployment(
assert tuple(cooldown_deployments) == ()
async def _complete_like_the_proxy(router: Router, model: str) -> str:
logging_obj, request = litellm.utils.function_setup(
original_function="acompletion",
rules_obj=litellm.utils.Rules(),
start_time=datetime.datetime(2026, 1, 1),
model=model,
messages=[{"role": "user", "content": "hi"}],
litellm_call_id=str(uuid.uuid4()),
)
response: Final = await router.acompletion(**request, litellm_logging_obj=logging_obj)
model_id: Final = response._hidden_params["model_id"]
assert isinstance(model_id, str)
return model_id
_CHAT_COMPLETION_OK: Final = {
"id": "chatcmpl-1",
"object": "chat.completion",
"created": 1,
"model": "gpt-4o",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
}
def _openai_deployment(model_name: str, host: str, deployment_id: str) -> Mapping[str, object]:
return {
"model_name": model_name,
"litellm_params": {"model": "openai/gpt-4o", "api_key": "k", "api_base": f"https://{host}/v1"},
"model_info": {"id": deployment_id},
}
@pytest.mark.asyncio
async def test_fallback_deployment_stays_up_while_the_primary_is_in_cooldown(
respx_mock: respx.MockRouter,
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
respx_mock.post("https://primary.test/v1/chat/completions").mock(
return_value=httpx.Response(503, json={"error": {"message": "down", "type": "server_error"}})
)
backup_route: Final = respx_mock.post("https://backup.test/v1/chat/completions").mock(
return_value=httpx.Response(200, json=_CHAT_COMPLETION_OK)
)
router: Final = litellm.Router(
model_list=[
_openai_deployment("primary", "primary.test", "primary-dep"),
_openai_deployment("backup", "backup.test", "backup-dep"),
],
fallbacks=[{"primary": ["backup"]}],
allowed_fails=litellm.allowed_fails,
num_retries=0,
)
served_by: Final = [await _complete_like_the_proxy(router, "primary") for _ in range(20)]
cooldown_deployments: Final = await async_get_cooldown_deployments(
litellm_router_instance=router,
parent_otel_span=None,
)
assert served_by == ["backup-dep"] * 20
assert len(backup_route.calls) == 20
assert tuple(cooldown_deployments) == ("primary-dep",)
@pytest.mark.asyncio
async def test_fallback_group_stays_up_after_a_bad_key_cools_the_primary_under_the_rate_based_default(
respx_mock: respx.MockRouter,
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
respx_mock.post("https://primary.test/v1/chat/completions").mock(
return_value=httpx.Response(401, json={"error": {"message": "bad key", "type": "invalid_request_error"}})
)
respx_mock.post("https://backup.test/v1/chat/completions").mock(
return_value=httpx.Response(200, json=_CHAT_COMPLETION_OK)
)
router: Final = litellm.Router(
model_list=[
_openai_deployment("primary", "primary.test", "primary-dep"),
_openai_deployment("backup", "backup.test", "backup-1"),
_openai_deployment("backup", "backup.test", "backup-2"),
],
fallbacks=[{"primary": ["backup"]}],
num_retries=0,
)
served_by: Final = [await _complete_like_the_proxy(router, "primary") for _ in range(12)]
cooldown_deployments: Final = await async_get_cooldown_deployments(
litellm_router_instance=router,
parent_otel_span=None,
)
assert set(served_by) <= {"backup-1", "backup-2"}
assert tuple(cooldown_deployments) == ("primary-dep",)
TIMED_OUT_DEPLOYMENT: Final = {
"model_name": "test-model",
"litellm_params": {"model": "openai/test-model", "api_key": "test-key"},

128
tests/unit/test_main.py
View file

@ -14,7 +14,7 @@ from dataclasses import dataclass
from datetime import datetime
from importlib import import_module
from pathlib import Path
from typing import Final
from typing import Final, Literal
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
@ -44,12 +44,13 @@ from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kw
from litellm.litellm_core_utils.get_litellm_params import stored_control_options
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.types.litellm_params import ControlOptions
from litellm.types.llms.openai import AllMessageValues, HttpxBinaryResponseContent
from litellm.types.prompts.init_prompts import PromptSpec
from litellm.types.utils import Delta, ModelResponseStream, StandardCallbackDynamicParams, StreamingChoices, Usage
from tests.unit.llms.bedrock.slow_upstream import (
CONVERSE_ANSWER,
STREAM_TIMEOUT_SECONDS,
UPSTREAM_FIRST_BYTE_SECONDS,
slow_upstream_sync_client,
@ -9950,3 +9951,126 @@ async def test_converse_fallback_prices_the_current_attempt(
else:
client.close()
configuration.reset_rust_configuration()
_LOOKUP_FUNCTION: Final = {
"name": "lookup_record",
"description": "Look up a record",
"parameters": {
"type": "object",
"properties": {"query": {"type": "string"}},
"required": ["query"],
"additionalProperties": False,
},
}
_STRICT_LOOKUP_TOOL: Final = {"type": "function", "function": {**_LOOKUP_FUNCTION, "strict": True}}
# Bedrock's 400 for a request past Claude's strict tool limits, seen live on us.anthropic.claude-sonnet-4-6 on 2026-10-20
_TOO_MANY_STRICT_TOOLS: Final = (
"The model returned the following errors: Too many strict tools (21). The maximum number of strict tools "
"supported is 20. Try reducing the number of tools marked as strict."
)
@dataclass(slots=True)
class _StrictToolsLimitRejectionFirst:
stream: Literal["off", "fake", "real"]
sent: tuple[tuple[str, dict[str, object]], ...] = ()
def __call__(self, request: httpx.Request) -> httpx.Response:
self.sent = (*self.sent, (request.url.path.rsplit("/", 1)[-1], json.loads(request.content)))
if len(self.sent) == 1:
return httpx.Response(400, json={"message": _TOO_MANY_STRICT_TOOLS})
return httpx.Response(200, content=b"" if self.stream == "real" else CONVERSE_ANSWER)
async def _converse_with_a_strict_tools_limit_rejection_first(
asynchronous: bool,
stream: Literal["off", "fake", "real"],
drop_params: bool | None,
tool: Mapping[str, object] = _STRICT_LOOKUP_TOOL,
) -> tuple[tuple[tuple[str, dict[str, object]], ...], litellm.BadRequestError | None]:
respond: Final = _StrictToolsLimitRejectionFirst(stream)
kwargs: Final = {
"model": "bedrock/us.anthropic.claude-sonnet-4-6",
"messages": [{"role": "user", "content": "Find invoice 42"}],
"tools": [tool],
"aws_access_key_id": "AKIASTRICTTOOLS",
"aws_secret_access_key": "strict-tools-secret",
"aws_region_name": "us-east-1",
"num_retries": 0,
"drop_params": drop_params,
"stream": stream != "off",
"fake_stream": stream == "fake",
}
try:
chunks: Final = (
await _acompletion_chunks(AsyncHTTPHandler(transport=httpx.MockTransport(respond)), kwargs, stream)
if asynchronous
else _completion_chunks(
HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))), kwargs, stream
)
)
except litellm.BadRequestError as error:
return respond.sent, error
assert chunks
return respond.sent, None
async def _acompletion_chunks(
client: AsyncHTTPHandler, kwargs: Mapping[str, object], stream: Literal["off", "fake", "real"]
) -> tuple[object, ...]:
response: Final = await litellm.acompletion(client=client, **kwargs)
return tuple([chunk async for chunk in response]) if stream == "fake" else (response,)
def _completion_chunks(
client: HTTPHandler, kwargs: Mapping[str, object], stream: Literal["off", "fake", "real"]
) -> tuple[object, ...]:
response: Final = litellm.completion(client=client, **kwargs)
return tuple(response) if stream == "fake" else (response,)
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
@pytest.mark.parametrize("stream", ["off", "fake", "real"], ids=["completion", "fake_stream", "stream"])
async def test_bedrock_strict_tools_limit_rejection_is_sent_again_without_strict_under_drop_params(
asynchronous: bool, stream: Literal["off", "fake", "real"], local_model_cost_map: None
) -> None:
sent, error = await _converse_with_a_strict_tools_limit_rejection_first(asynchronous, stream, drop_params=True)
assert error is None
(first_path, first_body), (second_path, second_body) = sent
first_tool_spec: Final = first_body["toolConfig"]["tools"][0]["toolSpec"]
assert first_tool_spec["strict"] is True
assert first_path == second_path == ("converse-stream" if stream == "real" else "converse")
assert second_body == {
**first_body,
"toolConfig": {
**first_body["toolConfig"],
"tools": [{"toolSpec": {key: value for key, value in first_tool_spec.items() if key != "strict"}}],
},
}
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
async def test_bedrock_strict_tools_limit_rejection_reaches_the_caller_without_drop_params(
asynchronous: bool, local_model_cost_map: None
) -> None:
sent, error = await _converse_with_a_strict_tools_limit_rejection_first(asynchronous, stream="off", drop_params=None)
assert error is not None and "Too many strict tools" in str(error)
assert len(sent) == 1
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
async def test_bedrock_strict_tools_limit_rejection_is_not_sent_again_when_no_tool_sent_strict(
asynchronous: bool, local_model_cost_map: None
) -> None:
sent, error = await _converse_with_a_strict_tools_limit_rejection_first(
asynchronous,
stream="off",
drop_params=True,
tool={"type": "function", "function": {**_LOOKUP_FUNCTION, "strict": False}},
)
assert error is not None and "Too many strict tools" in str(error)
assert len(sent) == 1

46
tests/unit/test_model_prices_and_context_window.py
View file

@ -2,7 +2,7 @@ from __future__ import annotations
import json
import re
from collections.abc import Mapping
from collections.abc import Iterator, Mapping
from pathlib import Path
from types import MappingProxyType
from typing import Final
@ -224,6 +224,50 @@ def test_dated_variants_carry_base_alias_service_tier_pricing(prices: dict):
)
PRIORITY_SUFFIX: Final = "_priority"
MIN_PRIORITY_MARKUP: Final = 1.0
MAX_PRIORITY_MARKUP: Final = 5.0
def _rate(value: JsonValue | None) -> float | None:
return float(value) if isinstance(value, (int, float)) and not isinstance(value, bool) else None
def _priority_markup(entry: Mapping[str, JsonValue], key: str) -> float | None:
if not key.endswith(PRIORITY_SUFFIX):
return None
priority: Final = _rate(entry[key])
standard: Final = _rate(entry.get(key.removesuffix(PRIORITY_SUFFIX)))
if priority is None or standard is None or standard <= 0:
return None
return priority / standard
def priority_markups(prices: Mapping[str, JsonValue]) -> Iterator[tuple[str, str, float]]:
for name, entry in prices.items():
if not isinstance(entry, dict):
continue
for key in entry:
if (markup := _priority_markup(entry, key)) is not None:
yield name, key, markup
@pytest.mark.parametrize("path", (PRICES_PATH, BACKUP_PRICES_PATH), ids=("main", "backup"))
def test_priority_rates_stay_within_a_premium_of_their_standard_rate(path: Path) -> None:
prices: Final = TypeAdapter(dict[str, JsonValue]).validate_json(path.read_bytes())
markups: Final = tuple(priority_markups(prices))
assert markups
overpriced: Final = tuple(
f"{name}.{key} is {markup:.1f}x its standard rate"
for name, key, markup in markups
if not MIN_PRIORITY_MARKUP <= markup <= MAX_PRIORITY_MARKUP
)
assert overpriced == (), (
"a priority rate is a premium on the standard rate, so one below it or far above it was likely copied from "
"another model; drop the key when the provider sells no priority tier for the model:\n" + "\n".join(overpriced)
)
# Azure Foundry GPT-6.1 Sol announcement, 2026-09-29: https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/introducing-gpt-6-1-sol-in-microsoft-foundry-advanced-intelligence-optimized-for/4560811
AZURE_GPT_6_1_SOL_REGIONAL_PREMIUM: Final = MappingProxyType({"us": 1.1, "eu": 1.2, "apac": 1.2})

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