diff --git a/README.md b/README.md index 5dfa8d01fc4..66b6acb7a71 100644 --- a/README.md +++ b/README.md @@ -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. diff --git a/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md b/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md index c79632e9195..e1f606a2eb7 100644 --- a/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md +++ b/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md @@ -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 diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py index b91b4a970e8..c25da7cc187 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py @@ -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 diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py new file mode 100644 index 00000000000..e36ff9b53dc --- /dev/null +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py @@ -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) diff --git a/litellm-rust/crates/inference-messages/tests/messages/response.rs b/litellm-rust/crates/inference-messages/tests/messages/response.rs index a0969860822..7aff1286858 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/response.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/response.rs @@ -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] diff --git a/litellm/__init__.py b/litellm/__init__.py index 41e3fa34323..2cff21bf29a 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -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, } diff --git a/litellm/constants.py b/litellm/constants.py index 8cc51d0ab5e..0d6dfb5c0d4 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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 diff --git a/litellm/litellm_core_utils/fallback_generalizations.py b/litellm/litellm_core_utils/fallback_generalizations.py index be432d1fdd4..eb4cf421a19 100644 --- a/litellm/litellm_core_utils/fallback_generalizations.py +++ b/litellm/litellm_core_utils/fallback_generalizations.py @@ -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) diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index 0764fb41941..368bb929eeb 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -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, diff --git a/litellm/llms/base_llm/base_model_iterator.py b/litellm/llms/base_llm/base_model_iterator.py index 611051bb2dd..032951be74a 100644 --- a/litellm/llms/base_llm/base_model_iterator.py +++ b/litellm/llms/base_llm/base_model_iterator.py @@ -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, diff --git a/litellm/llms/bedrock/chat/native_output_fallback.py b/litellm/llms/bedrock/chat/native_output_fallback.py index 1b8ce5fcda4..6892cd37beb 100644 --- a/litellm/llms/bedrock/chat/native_output_fallback.py +++ b/litellm/llms/bedrock/chat/native_output_fallback.py @@ -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 diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 28888f1a23a..fde995c7907 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -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()}}: diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index 115d914d3ac..0034ed2b9d4 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -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", diff --git a/litellm/main.py b/litellm/main.py index 9f412e1f24b..39a0a0cd3c1 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -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/", "") diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 5774e699ad7..9c3bd9dacbb 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -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", diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index 9ba6ed342ea..1ab03fb3fdc 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -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", diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 13a1844a6de..1d2045eb743 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -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: diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index 52e36964bc1..4d09e4d1e0f 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -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: diff --git a/litellm/proxy/client/cli/commands/config.py b/litellm/proxy/client/cli/commands/config.py index 539d88dce4a..1689a08c3de 100644 --- a/litellm/proxy/client/cli/commands/config.py +++ b/litellm/proxy/client/cli/commands/config.py @@ -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 ` 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), } ) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 1e7cdc44c70..314f14b6279 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -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, diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index e625ea6f9f2..1130fb73e9a 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -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, diff --git a/litellm/proxy/management_helpers/model_quota_policy.py b/litellm/proxy/management_helpers/model_quota_policy.py new file mode 100644 index 00000000000..752c78df862 --- /dev/null +++ b/litellm/proxy/management_helpers/model_quota_policy.py @@ -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()) + ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3b36ea5afce..04f84f1f7b5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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 diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index d96f2976508..0dd96c856f8 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -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", diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index d50e9a61edc..1ab6293b432 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -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")) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index b5ef9e95741..219944de521 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -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] = "" diff --git a/litellm/responses/main.py b/litellm/responses/main.py index d993b05f5b7..8451834eab2 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -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, diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index cc79d1c88bb..189366feab4 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -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() diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 50ab3e4f75a..c1f422c7d7b 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -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( diff --git a/litellm/router.py b/litellm/router.py index c8153b596d6..f0bb6c3a7a9 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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 = ( diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index a00d84babf3..91b9e5cb114 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -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") diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index 330500789a4..e1baf66b201 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -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 diff --git a/litellm/router_utils/pre_call_checks/continuation_prefill_check.py b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py index 8f211ef67db..d2a46fc316d 100644 --- a/litellm/router_utils/pre_call_checks/continuation_prefill_check.py +++ b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py @@ -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 diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index 8693eaa5151..daf7d6bb345 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -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( diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index 5d47bf63a38..806c3db7601 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -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: diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index bb3475cdcbe..c04bab4d360 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -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: diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 2869f9b2577..055c18f2e36 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4573,6 +4573,7 @@ class LlmProviders(str, Enum): POE = "poe" CHUTES = "chutes" NEOSANTARA = "neosantara" + AIAND = "aiand" PARASAIL = "parasail" XIAOMI_MIMO = "xiaomi_mimo" TENSORMESH = "tensormesh" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 5774e699ad7..9c3bd9dacbb 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -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", diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 9ba6ed342ea..1ab03fb3fdc 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -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", diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 853b10ad97b..f4a295e697d 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -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 ] diff --git a/tests/integration/_support/riva.py b/tests/integration/_support/riva.py new file mode 100644 index 00000000000..8552bc2197d --- /dev/null +++ b/tests/integration/_support/riva.py @@ -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() diff --git a/tests/integration/compatibility/test_compressed_request_bodies.py b/tests/integration/compatibility/test_compressed_request_bodies.py new file mode 100644 index 00000000000..ef2490add40 --- /dev/null +++ b/tests/integration/compatibility/test_compressed_request_bodies.py @@ -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 diff --git a/tests/integration/configuration/test_openapi_docs_route.py b/tests/integration/configuration/test_openapi_docs_route.py new file mode 100644 index 00000000000..ccb30c2f784 --- /dev/null +++ b/tests/integration/configuration/test_openapi_docs_route.py @@ -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 diff --git a/tests/integration/cost_calculation/cases/decisions.py b/tests/integration/cost_calculation/cases/decisions.py index 1cc57afe8ed..6594510b232 100644 --- a/tests/integration/cost_calculation/cases/decisions.py +++ b/tests/integration/cost_calculation/cases/decisions.py @@ -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, diff --git a/tests/integration/cost_calculation/cost_map.py b/tests/integration/cost_calculation/cost_map.py index 16107def2ad..32722b37c9b 100644 --- a/tests/integration/cost_calculation/cost_map.py +++ b/tests/integration/cost_calculation/cost_map.py @@ -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", diff --git a/tests/integration/management/test_scim_admin_group.py b/tests/integration/management/test_scim_admin_group.py index 3aeed45d186..ffbec93d4a8 100644 --- a/tests/integration/management/test_scim_admin_group.py +++ b/tests/integration/management/test_scim_admin_group.py @@ -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: diff --git a/tests/integration/management/test_scim_group_delete.py b/tests/integration/management/test_scim_group_delete.py index 4a4a93f8bd2..53af34c12e0 100644 --- a/tests/integration/management/test_scim_group_delete.py +++ b/tests/integration/management/test_scim_group_delete.py @@ -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) diff --git a/tests/integration/management/test_team_delete_large_membership.py b/tests/integration/management/test_team_delete_large_membership.py index e752c835e4f..1a8140373b0 100644 --- a/tests/integration/management/test_team_delete_large_membership.py +++ b/tests/integration/management/test_team_delete_large_membership.py @@ -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() diff --git a/tests/integration/mcp/test_mcp_llm_endpoints.py b/tests/integration/mcp/test_mcp_llm_endpoints.py index c5ffb88506a..394221708bf 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/tests/integration/mcp/test_mcp_llm_endpoints.py @@ -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( diff --git a/tests/integration/mcp/test_responses_mcp_streaming_round_history.py b/tests/integration/mcp/test_responses_mcp_streaming_round_history.py new file mode 100644 index 00000000000..8313d96e23c --- /dev/null +++ b/tests/integration/mcp/test_responses_mcp_streaming_round_history.py @@ -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 diff --git a/tests/integration/providers/_api_base_url_join_support.py b/tests/integration/providers/_api_base_url_join_support.py index 44bd6ec4dca..5711cc77309 100644 --- a/tests/integration/providers/_api_base_url_join_support.py +++ b/tests/integration/providers/_api_base_url_join_support.py @@ -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, diff --git a/tests/integration/providers/test_api_base_url_join_host_owned_proxy.py b/tests/integration/providers/test_api_base_url_join_host_owned_proxy.py index 4a075138305..687c15283f4 100644 --- a/tests/integration/providers/test_api_base_url_join_host_owned_proxy.py +++ b/tests/integration/providers/test_api_base_url_join_host_owned_proxy.py @@ -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" + } diff --git a/tests/integration/providers/test_api_base_url_join_wire.py b/tests/integration/providers/test_api_base_url_join_wire.py index 42e53ff9f95..74f8dde3cf7 100644 --- a/tests/integration/providers/test_api_base_url_join_wire.py +++ b/tests/integration/providers/test_api_base_url_join_wire.py @@ -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")) diff --git a/tests/integration/providers/test_bedrock_converse_redacted_thinking_wire.py b/tests/integration/providers/test_bedrock_converse_redacted_thinking_wire.py index 79ec0545c44..efb0609d560 100644 --- a/tests/integration/providers/test_bedrock_converse_redacted_thinking_wire.py +++ b/tests/integration/providers/test_bedrock_converse_redacted_thinking_wire.py @@ -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"}})) diff --git a/tests/integration/providers/test_claude_thinking_content_replay_wire.py b/tests/integration/providers/test_claude_thinking_content_replay_wire.py new file mode 100644 index 00000000000..4f2f3cc66c4 --- /dev/null +++ b/tests/integration/providers/test_claude_thinking_content_replay_wire.py @@ -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) diff --git a/tests/integration/providers/test_count_tokens_error_shapes_wire.py b/tests/integration/providers/test_count_tokens_error_shapes_wire.py index a6c3a6dc62f..59114e4aced 100644 --- a/tests/integration/providers/test_count_tokens_error_shapes_wire.py +++ b/tests/integration/providers/test_count_tokens_error_shapes_wire.py @@ -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 diff --git a/tests/integration/providers/test_decisions_wire.py b/tests/integration/providers/test_decisions_wire.py index 89ca674ef31..04b78985b20 100644 --- a/tests/integration/providers/test_decisions_wire.py +++ b/tests/integration/providers/test_decisions_wire.py @@ -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" diff --git a/tests/integration/providers/test_gemini_default_output_budget_wire.py b/tests/integration/providers/test_gemini_default_output_budget_wire.py new file mode 100644 index 00000000000..3eab626690d --- /dev/null +++ b/tests/integration/providers/test_gemini_default_output_budget_wire.py @@ -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 diff --git a/tests/integration/providers/test_gemini_long_decode_wire.py b/tests/integration/providers/test_gemini_long_decode_wire.py index 1270b864cb7..5ecb5e2a942 100644 --- a/tests/integration/providers/test_gemini_long_decode_wire.py +++ b/tests/integration/providers/test_gemini_long_decode_wire.py @@ -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( diff --git a/tests/integration/providers/test_huggingface_tokenizer_downloads_wire.py b/tests/integration/providers/test_huggingface_tokenizer_downloads_wire.py new file mode 100644 index 00000000000..ca1bfb5f78a --- /dev/null +++ b/tests/integration/providers/test_huggingface_tokenizer_downloads_wire.py @@ -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" diff --git a/tests/integration/providers/test_nvidia_riva_wire.py b/tests/integration/providers/test_nvidia_riva_wire.py new file mode 100644 index 00000000000..ef36c449761 --- /dev/null +++ b/tests/integration/providers/test_nvidia_riva_wire.py @@ -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 diff --git a/tests/integration/providers/test_sagemaker_chat_wire.py b/tests/integration/providers/test_sagemaker_chat_wire.py index 261787fec38..fdf0d94f7c7 100644 --- a/tests/integration/providers/test_sagemaker_chat_wire.py +++ b/tests/integration/providers/test_sagemaker_chat_wire.py @@ -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 diff --git a/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py index 2918f364c56..4e31122440d 100644 --- a/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py +++ b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py @@ -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", [ diff --git a/tests/integration/routing/test_batch_mode_deployments_wire.py b/tests/integration/routing/test_batch_mode_deployments_wire.py new file mode 100644 index 00000000000..7d702f499f5 --- /dev/null +++ b/tests/integration/routing/test_batch_mode_deployments_wire.py @@ -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 diff --git a/tests/integration/routing/test_priority_scheduler_queue_cleanup.py b/tests/integration/routing/test_priority_scheduler_queue_cleanup.py index 9fc21843b04..11798542088 100644 --- a/tests/integration/routing/test_priority_scheduler_queue_cleanup.py +++ b/tests/integration/routing/test_priority_scheduler_queue_cleanup.py @@ -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) diff --git a/tests/integration/routing/test_silent_model_shadow_wire.py b/tests/integration/routing/test_silent_model_shadow_wire.py new file mode 100644 index 00000000000..b00fbca645e --- /dev/null +++ b/tests/integration/routing/test_silent_model_shadow_wire.py @@ -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 diff --git a/tests/integration/sdk/test_anthropic_json_tool_cut_stream_sync_sdk.py b/tests/integration/sdk/test_anthropic_json_tool_cut_stream_sync_sdk.py new file mode 100644 index 00000000000..da091b1c12e --- /dev/null +++ b/tests/integration/sdk/test_anthropic_json_tool_cut_stream_sync_sdk.py @@ -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 diff --git a/tests/integration/sdk/test_bedrock_converse_redacted_thinking_router_sdk.py b/tests/integration/sdk/test_bedrock_converse_redacted_thinking_router_sdk.py index 6b2f4cb8eea..2988f9b5f53 100644 --- a/tests/integration/sdk/test_bedrock_converse_redacted_thinking_router_sdk.py +++ b/tests/integration/sdk/test_bedrock_converse_redacted_thinking_router_sdk.py @@ -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}) diff --git a/tests/integration/sdk/test_local_sandbox_sdk.py b/tests/integration/sdk/test_local_sandbox_sdk.py index b30010ac3f8..b1b9c89adfd 100644 --- a/tests/integration/sdk/test_local_sandbox_sdk.py +++ b/tests/integration/sdk/test_local_sandbox_sdk.py @@ -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) diff --git a/tests/integration/sdk/test_router_silent_model_shadow_sdk.py b/tests/integration/sdk/test_router_silent_model_shadow_sdk.py new file mode 100644 index 00000000000..821e4ddd0b7 --- /dev/null +++ b/tests/integration/sdk/test_router_silent_model_shadow_sdk.py @@ -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 diff --git a/tests/integration/security/test_files_batches_access.py b/tests/integration/security/test_files_batches_access.py index 7b6fd4652cf..119055c64c7 100644 --- a/tests/integration/security/test_files_batches_access.py +++ b/tests/integration/security/test_files_batches_access.py @@ -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: diff --git a/tests/integration/spend/test_background_response_billed_once.py b/tests/integration/spend/test_background_response_billed_once.py index 1b7fc975217..b207e46204a 100644 --- a/tests/integration/spend/test_background_response_billed_once.py +++ b/tests/integration/spend/test_background_response_billed_once.py @@ -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) diff --git a/tests/integration/spend/test_baggage_tag_mappings.py b/tests/integration/spend/test_baggage_tag_mappings.py index 1bf90269484..e08348a074c 100644 --- a/tests/integration/spend/test_baggage_tag_mappings.py +++ b/tests/integration/spend/test_baggage_tag_mappings.py @@ -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": diff --git a/tests/integration/spend/test_failed_dispatch_tokens.py b/tests/integration/spend/test_failed_dispatch_tokens.py index b92606ebc88..c9875be562c 100644 --- a/tests/integration/spend/test_failed_dispatch_tokens.py +++ b/tests/integration/spend/test_failed_dispatch_tokens.py @@ -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 diff --git a/tests/integration/spend/test_streamed_usage_estimate.py b/tests/integration/spend/test_streamed_usage_estimate.py new file mode 100644 index 00000000000..c7bcdf4401a --- /dev/null +++ b/tests/integration/spend/test_streamed_usage_estimate.py @@ -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) diff --git a/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py b/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py index 45697723af1..3f14098f6a9 100644 --- a/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py +++ b/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py @@ -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) diff --git a/tests/unit/litellm_core_utils/test_fallback_generalizations.py b/tests/unit/litellm_core_utils/test_fallback_generalizations.py index e59252e1571..99fa849879c 100644 --- a/tests/unit/litellm_core_utils/test_fallback_generalizations.py +++ b/tests/unit/litellm_core_utils/test_fallback_generalizations.py @@ -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.""" diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py index ef002a73e52..4d5d6cc17a9 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -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={}, diff --git a/tests/unit/llms/anthropic/test_count_tokens_oauth.py b/tests/unit/llms/anthropic/test_count_tokens_oauth.py index 16ec513916c..b8a7c001b37 100644 --- a/tests/unit/llms/anthropic/test_count_tokens_oauth.py +++ b/tests/unit/llms/anthropic/test_count_tokens_oauth.py @@ -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 diff --git a/tests/unit/llms/base_llm/responses/test_transformation.py b/tests/unit/llms/base_llm/responses/test_transformation.py index 8d21af83668..50d918fa5be 100644 --- a/tests/unit/llms/base_llm/responses/test_transformation.py +++ b/tests/unit/llms/base_llm/responses/test_transformation.py @@ -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 diff --git a/tests/unit/llms/base_llm/test_base_model_iterator.py b/tests/unit/llms/base_llm/test_base_model_iterator.py index 204b458cadc..2971405885d 100644 --- a/tests/unit/llms/base_llm/test_base_model_iterator.py +++ b/tests/unit/llms/base_llm/test_base_model_iterator.py @@ -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 diff --git a/tests/unit/llms/bedrock/chat/test_native_output_fallback.py b/tests/unit/llms/bedrock/chat/test_native_output_fallback.py index 301fdaeb245..9ee8c14f5d6 100644 --- a/tests/unit/llms/bedrock/chat/test_native_output_fallback.py +++ b/tests/unit/llms/bedrock/chat/test_native_output_fallback.py @@ -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) diff --git a/tests/unit/llms/bedrock/test_bedrock_common_utils.py b/tests/unit/llms/bedrock/test_bedrock_common_utils.py index 7bde1805e18..dd89d9c46b6 100644 --- a/tests/unit/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/unit/llms/bedrock/test_bedrock_common_utils.py @@ -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 diff --git a/tests/unit/llms/openai_like/test_aiand_provider.py b/tests/unit/llms/openai_like/test_aiand_provider.py new file mode 100644 index 00000000000..804f3a3b6eb --- /dev/null +++ b/tests/unit/llms/openai_like/test_aiand_provider.py @@ -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"] diff --git a/tests/unit/proxy/client/cli/test_agents.py b/tests/unit/proxy/client/cli/test_agents.py index 3481346076f..3c79b916202 100644 --- a/tests/unit/proxy/client/cli/test_agents.py +++ b/tests/unit/proxy/client/cli/test_agents.py @@ -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 = {} diff --git a/tests/unit/proxy/client/cli/test_config_commands.py b/tests/unit/proxy/client/cli/test_config_commands.py index 2cd7a76ab95..e99a4311caf 100644 --- a/tests/unit/proxy/client/cli/test_config_commands.py +++ b/tests/unit/proxy/client/cli/test_config_commands.py @@ -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: diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index cd0d786f841..e056e723a04 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -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") diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index 9ff2a6de9dd..c8f601de280 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -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( diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index a44373b87c8..a32c40e5a43 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -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, diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index fa0c9eda02d..ebe3c7407d8 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -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", diff --git a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py index 0eb1aa26a65..8d5befd9443 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py @@ -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, diff --git a/tests/unit/responses/mcp/test_mcp_streaming_iterator.py b/tests/unit/responses/mcp/test_mcp_streaming_iterator.py index 78b0e3d4e5b..f745b5bfadf 100644 --- a/tests/unit/responses/mcp/test_mcp_streaming_iterator.py +++ b/tests/unit/responses/mcp/test_mcp_streaming_iterator.py @@ -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"), diff --git a/tests/unit/responses/test_responses_api_bridge_flag.py b/tests/unit/responses/test_responses_api_bridge_flag.py index 278bd4c1e52..de96e995dd4 100644 --- a/tests/unit/responses/test_responses_api_bridge_flag.py +++ b/tests/unit/responses/test_responses_api_bridge_flag.py @@ -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}) diff --git a/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py b/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py index 8fa1c799863..f79c8a6fadb 100644 --- a/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py @@ -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 diff --git a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py index d8bc4c45ab8..fa75f980cc0 100644 --- a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -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]] diff --git a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index 87fbf4433a1..a3745c7b23f 100644 --- a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -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 diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 8d893160780..9487e4d6321 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -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(): """ diff --git a/tests/unit/router_utils/test_cooldown_handlers.py b/tests/unit/router_utils/test_cooldown_handlers.py index bab6d257579..4299b61c703 100644 --- a/tests/unit/router_utils/test_cooldown_handlers.py +++ b/tests/unit/router_utils/test_cooldown_handlers.py @@ -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"}, diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 70b268bf2a2..f12129af1fc 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -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 diff --git a/tests/unit/test_model_prices_and_context_window.py b/tests/unit/test_model_prices_and_context_window.py index 1f84be7710f..08a39b94c12 100644 --- a/tests/unit/test_model_prices_and_context_window.py +++ b/tests/unit/test_model_prices_and_context_window.py @@ -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}) diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index d9b655a0aff..c6d02dddae0 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -7039,8 +7039,7 @@ async def test_aresponses_fallback_attempt_does_not_replay_held_lifecycle_events @pytest.mark.asyncio async def test_aresponses_streaming_iterator_partial_content_injects_continuation(): - """Mid-stream error: input is rewritten to include user prompt + - developer instruction + prior assistant message with partial output.""" + """Mid-stream error: preserve partial assistant output and request continuation in a user turn.""" from litellm.exceptions import MidStreamFallbackError router = _make_router_with_fallback() @@ -7076,11 +7075,11 @@ async def test_aresponses_streaming_iterator_partial_content_injects_continuatio assert isinstance(new_input, list) assert new_input[0]["role"] == "user" assert new_input[0]["content"][0]["text"] == "What's the capital of France?" - assert new_input[1]["role"] == "developer" - assert "do not repeat" in new_input[1]["content"][0]["text"].lower() - assert new_input[2]["role"] == "assistant" - assert new_input[2]["content"][0]["type"] == "output_text" - assert new_input[2]["content"][0]["text"] == "The capital of France is" + assert new_input[1]["role"] == "assistant" + assert new_input[1]["content"][0]["type"] == "output_text" + assert new_input[1]["content"][0]["text"] == "The capital of France is" + assert new_input[2]["role"] == "user" + assert "do not repeat" in new_input[2]["content"][0]["text"].lower() @pytest.mark.asyncio @@ -21239,9 +21238,9 @@ def test_router_retry_skip_stamp_feeds_deployment_ids_to_skip_on_retry( ], ) def test_router_as_retry_skipped_deployment_ids_keeps_only_a_tuple_of_strings(value, expected): - from litellm.router import _as_retry_skipped_deployment_ids + from litellm.router_utils.common_utils import as_retry_skipped_deployment_ids - assert _as_retry_skipped_deployment_ids(value) == expected + assert as_retry_skipped_deployment_ids(value) == expected @pytest.mark.parametrize( diff --git a/tests/unit/test_router_weighted_failover.py b/tests/unit/test_router_weighted_failover.py index 9cb38b00e1e..8b7fc81dbde 100644 --- a/tests/unit/test_router_weighted_failover.py +++ b/tests/unit/test_router_weighted_failover.py @@ -8,13 +8,14 @@ cross-group fallback runs. """ from collections import Counter -from typing import Optional +from typing import Final, Optional from unittest.mock import AsyncMock, patch import pytest import litellm from litellm import Router +from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo from litellm.utils import get_excluded_filtered_deployments @@ -944,3 +945,77 @@ async def test_failover_falls_through_to_external_fallback_when_remaining_in_coo ) assert response._hidden_params["model_id"] == "fallback" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("pre_call_check", "metadata"), + [ + ("session_affinity", {"session_id": "sess-1"}), + ("deployment_affinity", {"user_api_key_hash": "hashed-key-1"}), + ], +) +async def test_failover_reaches_healthy_sibling_when_affinity_pins_failed_deployment( + pre_call_check: str, metadata: dict[str, str] +): + router: Final = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "bad", + "mock_response": Exception("A down"), + "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"}, + }, + { + "model_name": "fallback-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "good", + "mock_response": "ok from fallback", + }, + "model_info": {"id": "fallback"}, + }, + ], + routing_strategy="simple-shuffle", + num_retries=0, + disable_cooldowns=True, + enable_weighted_failover=True, + optional_pre_call_checks=[pre_call_check], + fallbacks=[{"test-model": ["fallback-model"]}], + ) + + while_a_is_down: Final = await router.acompletion( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + metadata=dict(metadata), + ) + router.upsert_deployment( + Deployment( + model_name="test-model", + litellm_params=LiteLLM_Params(model="gpt-4o", api_key="bad", mock_response="ok from A", weight=1), + model_info=ModelInfo(id="A"), + ) + ) + after_a_recovers: Final = await router.acompletion( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + metadata=dict(metadata), + ) + + assert while_a_is_down._hidden_params["model_id"] == "B" + assert while_a_is_down._hidden_params["additional_headers"]["x-litellm-attempted-fallbacks"] == 1 + assert after_a_recovers._hidden_params["model_id"] == "A" diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx index bb6204b5219..c353a3e56be 100644 --- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx @@ -7,6 +7,7 @@ import { KeysResponse, useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; import { useTeamMemberBudgets } from "@/app/(dashboard)/hooks/teams/useTeamMemberBudgets"; import { KeyResponse } from "../key_team_helpers/key_list"; import { Organization } from "../networking"; +import KeyInfoView from "../templates/key_info_view"; vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({ useKeys: vi.fn(), @@ -489,6 +490,22 @@ describe("TeamVirtualKeysTable", () => { }); }); + it("should open a key without its own organization without the team's organization", async () => { + const keyWithoutOrganization = { + keys: [createMockKey({ token: "sk-no-org", key_alias: "no_org_key", organization_id: null })], + total_count: 1, + current_page: 1, + total_pages: 1, + }; + mockUseKeys.mockReturnValue(createKeysQueryResult(keyWithoutOrganization)); + renderWithProviders(); + + await userEvent.click(await screen.findByRole("button", { name: /sk-no-org/ })); + + await screen.findByText("Key Info View"); + expect(vi.mocked(KeyInfoView).mock.lastCall?.[0].keyData).toMatchObject({ organization_id: null }); + }); + describe("entity links out of the key rows", () => { const renderRow = async (key: KeyResponse, organization: Organization | null = null) => { mockUseKeys.mockReturnValue({ diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx index 29a1cc232b6..e212eb090cf 100644 --- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx @@ -92,15 +92,7 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi const { data: keys, isPending: isLoading, isFetching, refetch } = useKeys(pageIndex + 1, pageSize, keyListOptions); - const displayKeys = useMemo(() => { - const kList = keys?.keys || []; - const orgId = organization?.organization_id; - if (!orgId) return kList; - return kList.map((k: KeyResponse) => ({ - ...k, - organization_id: (k.organization_id ?? k.org_id) || orgId, - })); - }, [keys?.keys, organization?.organization_id]); + const displayKeys = useMemo(() => keys?.keys ?? [], [keys?.keys]); const teamMemberBudgetTeamIds = useMemo( () => (displayKeys.some((key) => key.max_budget == null && Boolean(key.user_id)) ? [teamId] : []), @@ -182,7 +174,7 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi }, { id: "organization_id", - accessorKey: "organization_id", + accessorFn: (key) => (key.organization_id ?? key.org_id) || organization?.organization_id || null, header: "Organization ID", size: 140, enableSorting: false, @@ -461,7 +453,7 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi }, }, ], - [expandedAccordions, teamMemberBudgets, teamId], + [expandedAccordions, teamMemberBudgets, teamId, organization?.organization_id], ); const handleSortingChange = useCallback((updaterOrValue: React.SetStateAction) => { diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx index c82a4823dfd..9e48c4e3aba 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx @@ -1,6 +1,7 @@ import { fireEvent, render, screen, waitFor } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { toast } from "@/lib/toast"; +import type { KeyResponse } from "../key_team_helpers/key_list"; vi.mock("@/lib/toast", () => ({ toast: { error: vi.fn(), success: vi.fn() }, @@ -655,3 +656,60 @@ describe("KeyInfoView handleKeyUpdate soft_budget", () => { expect(keyUpdateCallMock).not.toHaveBeenCalled(); }); }); + +describe("KeyInfoView handleKeyUpdate organization", () => { + const keyListRow = { ...baseKeyData, org_id: "org-1" } as KeyResponse; + + const adminAuth = { + accessToken: "access_abc", + userId: "user_1", + userRole: "Admin", + premiumUser: true, + token: "token_123", + userEmail: "test@example.com", + disabledPersonalKeyCreation: false, + showSSOBanner: false, + }; + + const renderKeyListRowSettings = () => { + mockUseAuthorized.mockReturnValue(adminAuth); + render( + {}} keyData={keyListRow} onKeyDataUpdate={() => {}} teams={[]} />, + ); + fireEvent.click(screen.getByText("Settings")); + }; + + const submitEditForm = (formValues: Record) => { + fireEvent.click(screen.getByText("Edit Settings")); + Object.assign(globalThis, { __TEST_FORM_VALUES: formValues }); + fireEvent.click(screen.getByText("Mock Submit")); + }; + + const shownOrganization = () => screen.getByText("Organization", { selector: "p" }).nextElementSibling; + + it("should show no organization after a save clears the one a key list row carried as org_id", async () => { + keyUpdateCallMock.mockResolvedValueOnce({ ...baseKeyData, organization_id: null }); + renderKeyListRowSettings(); + expect(shownOrganization()).toHaveTextContent("org-1"); + + submitEditForm({ token: "tok_123", organization_id: null }); + + await waitFor(() => { + expect(shownOrganization()).toHaveTextContent("Not Set"); + }); + expect(keyUpdateCallMock.mock.calls[0][1]).toHaveProperty("organization_id", null); + }); + + it("should leave organization_id out of a save that keeps the key's organization", async () => { + keyUpdateCallMock.mockResolvedValueOnce({ ...baseKeyData, organization_id: "org-1" }); + renderKeyListRowSettings(); + + submitEditForm({ token: "tok_123", organization_id: "org-1" }); + + await waitFor(() => { + expect(keyUpdateCallMock).toHaveBeenCalledTimes(1); + }); + expect(keyUpdateCallMock.mock.calls[0][1]).not.toHaveProperty("organization_id"); + expect(shownOrganization()).toHaveTextContent("org-1"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/templates/keyEditFieldNormalizers.ts b/ui/litellm-dashboard/src/components/templates/keyEditFieldNormalizers.ts index 1cfd1ae35f7..4ddc2c60bd3 100644 --- a/ui/litellm-dashboard/src/components/templates/keyEditFieldNormalizers.ts +++ b/ui/litellm-dashboard/src/components/templates/keyEditFieldNormalizers.ts @@ -9,6 +9,9 @@ const WORD_FORM_BUDGET_DURATIONS: Record = { export const canonicalBudgetDuration = (duration: string | null | undefined): string | null => duration ? WORD_FORM_BUDGET_DURATIONS[duration] ?? duration : null; +export const keyOrganizationId = (key: { organization_id?: string | null; org_id?: string | null }): string | null => + key.organization_id || key.org_id || null; + // Determine the key_type display value from allowed_routes export const keyTypeFromRoutes = (allowedRoutes: string[] | null | undefined): string => { if (!allowedRoutes || allowedRoutes.length === 0) return "default"; diff --git a/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts b/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts index 9dd3601ef31..1459d4d7301 100644 --- a/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts +++ b/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts @@ -6,7 +6,7 @@ import { ToolPermissionsChangeOrigin } from "../mcp_server_management/MCPToolPer import { extractLoggingSettings, formatMetadataForDisplay, stripTagsFromMetadata } from "../key_info_utils"; import { mapInternalToDisplayNames } from "../callback_info_helpers"; import { estimateChecks, estimateFields } from "./estimatedOutputTokens"; -import { canonicalBudgetDuration } from "./keyEditFieldNormalizers"; +import { canonicalBudgetDuration, keyOrganizationId } from "./keyEditFieldNormalizers"; export interface McpServersAndGroups { servers: string[]; @@ -109,7 +109,7 @@ export const toKeyEditFormValues = (keyData: KeyResponse): KeyEditFormValues => accessGroups: keyData.object_permission?.agent_access_groups || [], }, skills: keyData.object_permission?.skills || [], - organization_id: keyData.organization_id, + organization_id: keyOrganizationId(keyData), team_id: keyData.team_id, project_id: keyData.project_id, logging_settings: extractLoggingSettings(keyData.metadata), diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx index b6319b3713e..dce14adc2e5 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx @@ -1671,6 +1671,148 @@ describe("KeyEditView", () => { expect(await screen.findByText("Engineering")).toBeInTheDocument(); }); + it("should initialize organization from a key list row that only carries org_id", async () => { + const keyFromList = { + ...MOCK_KEY_DATA, + organization_id: null, + org_id: "org-1", + }; + + renderWithProviders( + {}} + onSubmit={async () => {}} + accessToken="" + userID="" + userRole="Admin" + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByLabelText("Organization")).toHaveValue("Engineering"); + }); + + await userEvent.click(screen.getByRole("combobox", { name: "Team" })); + + expect(await screen.findByRole("option", { name: /Alpha/ })).toBeInTheDocument(); + expect(screen.queryByRole("option", { name: /Beta/ })).not.toBeInTheDocument(); + }); + + it("should let a non-admin move a key to their team in another organization", async () => { + const onSubmitMock = vi.fn().mockResolvedValue(undefined); + + renderWithProviders( + {}} + onSubmit={onSubmitMock} + accessToken="" + userID="" + userRole="Internal User" + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByLabelText("Organization")).toHaveValue("Engineering"); + }); + + await userEvent.click(screen.getByRole("combobox", { name: "Team" })); + await userEvent.click(await screen.findByRole("option", { name: /Beta/ })); + + await waitFor(() => { + expect(screen.getByLabelText("Organization")).toHaveValue("Sales"); + }); + + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmitMock).toHaveBeenCalled(); + }); + expect(onSubmitMock.mock.calls[0][0]).toMatchObject({ organization_id: "org-2", team_id: "team-b" }); + }); + + it("should re-scope teams when the edited key switches to another organization", async () => { + const sharedProps = { + teams: [ + { team_id: "team-a", team_alias: "Alpha", organization_id: "org-1" }, + { team_id: "team-b", team_alias: "Beta", organization_id: "org-2" }, + ], + onCancel: () => {}, + onSubmit: async () => {}, + accessToken: "", + userID: "", + userRole: "Admin", + premiumUser: false, + }; + + const { rerender } = renderWithProviders( + , + ); + + await waitFor(() => { + expect(screen.getByLabelText("Organization")).toHaveValue("Engineering"); + }); + + rerender(); + + await waitFor(() => { + expect(screen.getByLabelText("Organization")).toHaveValue("Sales"); + }); + + await userEvent.click(screen.getByRole("combobox", { name: "Team" })); + + expect(await screen.findByRole("option", { name: /Beta/ })).toBeInTheDocument(); + expect(screen.queryByRole("option", { name: /Alpha/ })).not.toBeInTheDocument(); + }); + + it("should adopt the organization of a team picked in the form", async () => { + const onSubmitMock = vi.fn().mockResolvedValue(undefined); + + renderWithProviders( + {}} + onSubmit={onSubmitMock} + accessToken="" + userID="" + userRole="Admin" + premiumUser={false} + />, + ); + + await screen.findByRole("button", { name: /save changes/i }); + expect(screen.getByLabelText("Organization")).toHaveValue(""); + + await userEvent.click(screen.getByRole("combobox", { name: "Team" })); + await userEvent.click(await screen.findByRole("option", { name: /Beta/ })); + + await waitFor(() => { + expect(screen.getByLabelText("Organization")).toHaveValue("Sales"); + }); + + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmitMock).toHaveBeenCalled(); + }); + expect(onSubmitMock.mock.calls[0][0].organization_id).toBe("org-2"); + }); + it("should initialize organization from keyData", async () => { const keyWithOrg = { ...MOCK_KEY_DATA, diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index 44e9af0d316..3e844d38d68 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -123,7 +123,6 @@ function KeyEditSession({ ? mapInternalToDisplayNames(keyData.metadata.litellm_disabled_callbacks) : [], ); - const [selectedOrganizationId, setSelectedOrganizationId] = useState(keyData.organization_id || null); const [autoRotationEnabled, setAutoRotationEnabled] = useState(keyData.auto_rotate || false); const [rotationInterval, setRotationInterval] = useState(keyData.rotation_interval || ""); const [neverExpire, setNeverExpire] = useState(!keyData.expires); @@ -147,9 +146,11 @@ function KeyEditSession({ const { data: uiSettingsData } = useUISettings(); const enableProjectsUI = Boolean(uiSettingsData?.values?.enable_projects_ui); const hasProject = Boolean(keyData.project_id); + const canChangeOrganization = userRole === "Admin" && !hasProject; const detachProject = hasProject && form.watch("project_id") === null; const canDetachProject = canDetachKeyProject(team, organizations, userID, userRole); + const selectedOrganizationId = form.watch("organization_id") ?? null; const allowedRoutesValue = form.watch("allowed_routes"); const selectedModels = (form.watch("models") as string[] | undefined) ?? []; const allowedRoutes = parseAllowedRoutes(allowedRoutesValue); @@ -327,7 +328,6 @@ function KeyEditSession({ const handleOrganizationChange = (setField: (value: string | null) => void, orgId: string | null) => { setField(orgId); - setSelectedOrganizationId(orgId); form.setValue("team_id", null, { shouldDirty: true }); }; @@ -336,7 +336,6 @@ function KeyEditSession({ clearMcpGrants(form.setValue); const organizationId = teams?.find((t) => t.team_id === teamId)?.organization_id; if (!organizationId) return; - setSelectedOrganizationId(organizationId); form.setValue("organization_id", organizationId, { shouldDirty: true }); }; @@ -354,9 +353,10 @@ function KeyEditSession({ })), ]; - const visibleTeams = selectedOrganizationId - ? teams?.filter((t) => t.organization_id === selectedOrganizationId) - : teams; + const visibleTeams = + canChangeOrganization && selectedOrganizationId + ? teams?.filter((t) => t.organization_id === selectedOrganizationId) + : teams; const teamOptionsFor = (selectedTeamId: string | null): SearchSelectOption[] => { const selectedTeamOutsideFilter = teams?.find((t) => t.team_id === selectedTeamId && !visibleTeams?.includes(t)); @@ -797,7 +797,7 @@ function KeyEditSession({ value={value} organizations={organizations} loading={isOrganizationsLoading} - disabled={userRole !== "Admin" || hasProject} + disabled={!canChangeOrganization} onChange={(orgId) => handleOrganizationChange(onChange, orgId)} /> )} diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index 4680368943a..1189e4a8c7d 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -56,6 +56,7 @@ import { teamMemberBudgetGate, } from "../shared/InheritedBudgetHint"; import { KeyEditView } from "./key_edit_view"; +import { keyOrganizationId } from "./keyEditFieldNormalizers"; import { isTeamAdminEditingMemberKey, teamAdminMemberKeyPayload } from "./teamAdminMemberKeyPayload"; import { BudgetPeriodTooltip } from "@/components/shared/BudgetPeriodTooltip"; @@ -240,6 +241,10 @@ export default function KeyInfoView({ delete formValues.disable_global_guardrails; } + if ((formValues.organization_id || null) === keyOrganizationId(currentKeyData)) { + delete formValues.organization_id; + } + // Handle max budget empty string formValues.max_budget = mapEmptyStringToNull(formValues.max_budget); @@ -398,7 +403,9 @@ export default function KeyInfoView({ const newKeyValues = await keyUpdateCall(accessToken, formValues); - setCurrentKeyData((prevData) => (prevData ? { ...prevData, ...newKeyValues } : undefined)); + setCurrentKeyData((prevData) => + prevData ? { ...prevData, ...newKeyValues, org_id: newKeyValues.organization_id } : undefined, + ); void queryClient.invalidateQueries({ queryKey: keyKeys.all }); if (onKeyDataUpdate) { @@ -546,7 +553,7 @@ export default function KeyInfoView({ const lastConfiguredAt = currentKeyData.settings_updated_at || currentKeyData.created_at; const parentTeam = currentKeyData.team_id ? teamsData?.find((team) => team.team_id === currentKeyData.team_id) : null; - const orgId = currentKeyData.organization_id || currentKeyData.org_id || parentTeam?.organization_id || ""; + const orgId = keyOrganizationId(currentKeyData) || parentTeam?.organization_id || ""; const parentOrg = orgId ? organizations?.find((org) => org.organization_id === orgId) : null; const hasOwnBudget = currentKeyData.max_budget !== null; @@ -979,7 +986,7 @@ export default function KeyInfoView({

Organization

-

{(currentKeyData.organization_id ?? currentKeyData.org_id) || "Not Set"}

+

{keyOrganizationId(currentKeyData) || "Not Set"}