mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-06 07:48:06 +00:00
Merge remote-tracking branch 'origin/main' into litellm_openai_sdk3_docker
# Conflicts: # pyproject.toml
This commit is contained in:
mateo-berri
commit
b913a0fdbf
172 changed files with 11252 additions and 1271 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
115
litellm-rust/Cargo.lock
generated
115
litellm-rust/Cargo.lock
generated
|
|
@ -2342,7 +2342,7 @@ dependencies = [
|
|||
"hmac 0.13.0",
|
||||
"http 1.4.2",
|
||||
"jiff",
|
||||
"reqwest 0.13.5",
|
||||
"reqwest",
|
||||
"rustc_version",
|
||||
"rustls 0.23.42",
|
||||
"rustls-pki-types",
|
||||
|
|
@ -2400,7 +2400,7 @@ dependencies = [
|
|||
"pin-project",
|
||||
"prost",
|
||||
"prost-types",
|
||||
"reqwest 0.13.5",
|
||||
"reqwest",
|
||||
"rustc_version",
|
||||
"serde",
|
||||
"serde_json",
|
||||
|
|
@ -2896,7 +2896,6 @@ dependencies = [
|
|||
"tokio",
|
||||
"tokio-rustls 0.26.4",
|
||||
"tower-service",
|
||||
"webpki-roots",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -3389,7 +3388,7 @@ dependencies = [
|
|||
"litellm-auth-types",
|
||||
"litellm-http",
|
||||
"moka",
|
||||
"reqwest 0.12.28",
|
||||
"reqwest",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
|
|
@ -3473,7 +3472,7 @@ dependencies = [
|
|||
"litellm-cache-response",
|
||||
"litellm-cache-testing",
|
||||
"litellm-http",
|
||||
"reqwest 0.12.28",
|
||||
"reqwest",
|
||||
"rstest",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
|
|
@ -3508,7 +3507,7 @@ dependencies = [
|
|||
"litellm-cache-testing",
|
||||
"litellm-http",
|
||||
"percent-encoding",
|
||||
"reqwest 0.12.28",
|
||||
"reqwest",
|
||||
"rstest",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
|
|
@ -3535,7 +3534,7 @@ dependencies = [
|
|||
"litellm-cache-testing",
|
||||
"litellm-http",
|
||||
"qdrant-client",
|
||||
"reqwest 0.12.28",
|
||||
"reqwest",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
|
|
@ -3611,7 +3610,7 @@ dependencies = [
|
|||
"litellm-cache",
|
||||
"litellm-cache-testing",
|
||||
"litellm-http",
|
||||
"reqwest 0.12.28",
|
||||
"reqwest",
|
||||
"rstest",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
|
|
@ -3978,9 +3977,7 @@ dependencies = [
|
|||
"hyper-util",
|
||||
"litellm-core-utils",
|
||||
"rcgen",
|
||||
"reqwest 0.12.28",
|
||||
"reqwest 0.13.5",
|
||||
"rmcp",
|
||||
"reqwest",
|
||||
"rstest",
|
||||
"rustls 0.23.42",
|
||||
"rustls-native-certs",
|
||||
|
|
@ -3990,6 +3987,7 @@ dependencies = [
|
|||
"tempfile",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"tokio-rustls 0.26.4",
|
||||
"tokio-tungstenite",
|
||||
"tracing",
|
||||
"veil",
|
||||
|
|
@ -4013,7 +4011,7 @@ dependencies = [
|
|||
"litellm-llms",
|
||||
"litellm-secrets",
|
||||
"litellm-tracing",
|
||||
"reqwest 0.12.28",
|
||||
"reqwest",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
|
|
@ -4072,7 +4070,7 @@ dependencies = [
|
|||
"litellm-router-types",
|
||||
"litellm-secrets",
|
||||
"litellm-tracing",
|
||||
"reqwest 0.12.28",
|
||||
"reqwest",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
|
|
@ -4133,7 +4131,7 @@ dependencies = [
|
|||
"litellm-llms-types",
|
||||
"litellm-secrets",
|
||||
"litellm-tracing",
|
||||
"reqwest 0.12.28",
|
||||
"reqwest",
|
||||
"rstest",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
|
|
@ -4196,7 +4194,7 @@ dependencies = [
|
|||
"litellm-python-compat",
|
||||
"litellm-router-types",
|
||||
"litellm-secrets",
|
||||
"reqwest 0.12.28",
|
||||
"reqwest",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
|
|
@ -4278,7 +4276,7 @@ dependencies = [
|
|||
"pyo3",
|
||||
"pyo3-async-runtimes",
|
||||
"redis",
|
||||
"reqwest 0.12.28",
|
||||
"reqwest",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
|
|
@ -4371,7 +4369,7 @@ dependencies = [
|
|||
"litellm-secrets-hashicorp",
|
||||
"litellm-secrets-types",
|
||||
"moka",
|
||||
"reqwest 0.12.28",
|
||||
"reqwest",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
|
|
@ -4413,7 +4411,7 @@ dependencies = [
|
|||
"litellm-http",
|
||||
"litellm-secrets-types",
|
||||
"percent-encoding",
|
||||
"reqwest 0.12.28",
|
||||
"reqwest",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
|
|
@ -4436,7 +4434,7 @@ dependencies = [
|
|||
"moka",
|
||||
"percent-encoding",
|
||||
"rcgen",
|
||||
"reqwest 0.12.28",
|
||||
"reqwest",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
|
|
@ -4463,7 +4461,7 @@ dependencies = [
|
|||
"litellm-secrets-types",
|
||||
"moka",
|
||||
"percent-encoding",
|
||||
"reqwest 0.12.28",
|
||||
"reqwest",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
|
|
@ -4551,7 +4549,7 @@ version = "0.1.0"
|
|||
dependencies = [
|
||||
"flate2",
|
||||
"futures-util",
|
||||
"reqwest 0.12.28",
|
||||
"reqwest",
|
||||
"rstest",
|
||||
"semver",
|
||||
"serde",
|
||||
|
|
@ -5936,49 +5934,6 @@ version = "1.9.3"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ba39f3699c378cd8970968dcbff9c43159ea4cfbd88d43c00b22f2ef10a435d2"
|
||||
|
||||
[[package]]
|
||||
name = "reqwest"
|
||||
version = "0.12.28"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"futures-core",
|
||||
"futures-util",
|
||||
"h2 0.4.15",
|
||||
"http 1.4.2",
|
||||
"http-body 1.1.0",
|
||||
"http-body-util",
|
||||
"hyper 1.10.1",
|
||||
"hyper-rustls 0.27.9",
|
||||
"hyper-util",
|
||||
"js-sys",
|
||||
"log",
|
||||
"mime_guess",
|
||||
"percent-encoding",
|
||||
"pin-project-lite",
|
||||
"quinn",
|
||||
"rustls 0.23.42",
|
||||
"rustls-pki-types",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_urlencoded",
|
||||
"sync_wrapper",
|
||||
"tokio",
|
||||
"tokio-rustls 0.26.4",
|
||||
"tokio-util",
|
||||
"tower",
|
||||
"tower-http",
|
||||
"tower-service",
|
||||
"url",
|
||||
"wasm-bindgen",
|
||||
"wasm-bindgen-futures",
|
||||
"wasm-streams 0.4.2",
|
||||
"web-sys",
|
||||
"webpki-roots",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "reqwest"
|
||||
version = "0.13.5"
|
||||
|
|
@ -6000,6 +5955,7 @@ dependencies = [
|
|||
"js-sys",
|
||||
"log",
|
||||
"mime",
|
||||
"mime_guess",
|
||||
"percent-encoding",
|
||||
"pin-project-lite",
|
||||
"quinn",
|
||||
|
|
@ -6019,7 +5975,7 @@ dependencies = [
|
|||
"url",
|
||||
"wasm-bindgen",
|
||||
"wasm-bindgen-futures",
|
||||
"wasm-streams 0.5.0",
|
||||
"wasm-streams",
|
||||
"web-sys",
|
||||
]
|
||||
|
||||
|
|
@ -6066,7 +6022,7 @@ dependencies = [
|
|||
"pin-project-lite",
|
||||
"process-wrap",
|
||||
"rand 0.10.2",
|
||||
"reqwest 0.13.5",
|
||||
"reqwest",
|
||||
"schemars 1.2.2",
|
||||
"serde",
|
||||
"serde_json",
|
||||
|
|
@ -6189,7 +6145,7 @@ dependencies = [
|
|||
"async-trait",
|
||||
"bytes",
|
||||
"http 1.4.2",
|
||||
"reqwest 0.13.5",
|
||||
"reqwest",
|
||||
"rustify_derive",
|
||||
"serde",
|
||||
"serde_json",
|
||||
|
|
@ -7221,7 +7177,7 @@ dependencies = [
|
|||
"memchr",
|
||||
"parse-display",
|
||||
"pin-project-lite",
|
||||
"reqwest 0.13.5",
|
||||
"reqwest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_with",
|
||||
|
|
@ -7470,11 +7426,7 @@ checksum = "edc5f74e248dc973e0dbb7b74c7e0d6fcc301c694ff50049504004ef4d0cdcd9"
|
|||
dependencies = [
|
||||
"futures-util",
|
||||
"log",
|
||||
"rustls 0.23.42",
|
||||
"rustls-native-certs",
|
||||
"rustls-pki-types",
|
||||
"tokio",
|
||||
"tokio-rustls 0.26.4",
|
||||
"tungstenite",
|
||||
]
|
||||
|
||||
|
|
@ -7819,8 +7771,6 @@ dependencies = [
|
|||
"httparse",
|
||||
"log",
|
||||
"rand 0.8.7",
|
||||
"rustls 0.23.42",
|
||||
"rustls-pki-types",
|
||||
"sha1 0.10.7",
|
||||
"thiserror 1.0.69",
|
||||
"utf-8",
|
||||
|
|
@ -7860,7 +7810,7 @@ dependencies = [
|
|||
"futures",
|
||||
"pin-project",
|
||||
"rand 0.10.2",
|
||||
"reqwest 0.13.5",
|
||||
"reqwest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"time",
|
||||
|
|
@ -8074,7 +8024,7 @@ dependencies = [
|
|||
"async-trait",
|
||||
"derive_builder",
|
||||
"http 1.4.2",
|
||||
"reqwest 0.13.5",
|
||||
"reqwest",
|
||||
"rustify",
|
||||
"rustify_derive",
|
||||
"serde",
|
||||
|
|
@ -8221,19 +8171,6 @@ dependencies = [
|
|||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wasm-streams"
|
||||
version = "0.4.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65"
|
||||
dependencies = [
|
||||
"futures-util",
|
||||
"js-sys",
|
||||
"wasm-bindgen",
|
||||
"wasm-bindgen-futures",
|
||||
"web-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wasm-streams"
|
||||
version = "0.5.0"
|
||||
|
|
|
|||
|
|
@ -90,7 +90,7 @@ pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
|
|||
rand = "0.8"
|
||||
macro_rules_attribute = "0.2.3"
|
||||
schemars = "1"
|
||||
reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls", "http2", "stream"] }
|
||||
reqwest = { version = "0.13", default-features = false, features = ["json", "multipart", "rustls", "http2", "stream", "query"] }
|
||||
redis = "1.8.0"
|
||||
redis-test = "1.0.4"
|
||||
testcontainers-modules = "0.15.0"
|
||||
|
|
@ -100,10 +100,10 @@ rstest = "0.26.1"
|
|||
wiremock = "0.6.5"
|
||||
rstest_reuse = "0.7.0"
|
||||
rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] }
|
||||
rustls-native-certs = "0.8"
|
||||
rustify = "=0.7.0"
|
||||
rustify_derive = "=0.5.5"
|
||||
vaultrs = { version = "=0.8.0", default-features = false, features = ["rustls"] }
|
||||
rustls-native-certs = "0.8"
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
serde_json = { version = "1.0", features = ["float_roundtrip"] }
|
||||
serde_with = { version = "=3.21.0", default-features = false, features = ["std", "macros"] }
|
||||
|
|
@ -118,7 +118,7 @@ thiserror = "2.0"
|
|||
tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] }
|
||||
tiktoken-rs = "0.12.0"
|
||||
tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"] }
|
||||
tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] }
|
||||
tokio-tungstenite = { version = "0.24", default-features = false, features = ["handshake"] }
|
||||
futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] }
|
||||
base64 = "0.22"
|
||||
flate2 = "1"
|
||||
|
|
|
|||
|
|
@ -10,8 +10,11 @@ disallowed-methods = [
|
|||
{ path = "reqwest::Client::new", reason = "take litellm_http::Client from HttpClientPool" },
|
||||
{ path = "reqwest::Client::builder", reason = "HttpClientConfig owns client construction" },
|
||||
{ path = "reqwest::ClientBuilder::danger_accept_invalid_certs", reason = "set HttpClientConfig::verify instead" },
|
||||
{ path = "reqwest::ClientBuilder::tls_danger_accept_invalid_certs", reason = "set HttpClientConfig::verify instead" },
|
||||
{ path = "reqwest::ClientBuilder::tls_danger_accept_invalid_hostnames", reason = "set HttpClientConfig::verify instead" },
|
||||
{ path = "reqwest::ClientBuilder::identity", reason = "set HttpClientConfig::client_certificate instead" },
|
||||
{ path = "reqwest::ClientBuilder::use_preconfigured_tls", reason = "HttpClientConfig owns the TLS configuration" },
|
||||
{ path = "reqwest::ClientBuilder::tls_backend_preconfigured", reason = "HttpClientConfig owns the TLS configuration" },
|
||||
{ path = "sqlx::query", reason = "use sqlx::query! or query_file! so the SQL is checked against the migrated schema" },
|
||||
{ path = "sqlx::query_as", reason = "use sqlx::query_as! or query_file_as! so the SQL is checked against the migrated schema" },
|
||||
{ path = "sqlx::query_scalar", reason = "use sqlx::query_scalar! so the SQL is checked against the migrated schema" },
|
||||
|
|
|
|||
|
|
@ -1080,7 +1080,7 @@ mod tests {
|
|||
&no_env,
|
||||
)
|
||||
.await?;
|
||||
let client = litellm_http::Client::plain_for_test();
|
||||
let client = litellm_http::Client::for_test();
|
||||
let mut failures = Vec::new();
|
||||
|
||||
for region in ["us-west-2", "us-east-1"] {
|
||||
|
|
|
|||
|
|
@ -31,7 +31,7 @@ async fn connect(server: &MockServer) -> AzureBlobCache<JsonCodec<Value>> {
|
|||
None,
|
||||
ClientOptions {
|
||||
transport: Some(Transport::new(Arc::new(ReqwestTransport(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
)))),
|
||||
..ClientOptions::default()
|
||||
},
|
||||
|
|
|
|||
|
|
@ -97,7 +97,7 @@ async fn cache_exposes_its_configuration(#[future(awt)] server: MockServer) {
|
|||
path_service_account: Some("/secrets/sa.json".into()),
|
||||
..support::config(&server, Some("folder"))
|
||||
},
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
litellm_cache::JsonCodec::<Value>::new(),
|
||||
);
|
||||
assert_eq!(cache.bucket_name(), "bucket");
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@ pub fn cache_with_token(
|
|||
) -> JsonGcsCache {
|
||||
GcsCache::with_token_source(
|
||||
config(server, gcs_path),
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
JsonCodec::new(),
|
||||
token,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -108,7 +108,7 @@ fn config(base: String, timeout: Option<Duration>) -> OpenAiEmbedderConfig {
|
|||
async fn posts_embeddings_request_and_parses_vector() {
|
||||
let server = TestHttpServer::response("200 OK", r#"{"data":[{"embedding":[0.1,0.2]}]}"#).await;
|
||||
let embedder = OpenAiEmbedder::new(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
config(
|
||||
format!("{}/", server.base_url()),
|
||||
Some(Duration::from_secs(1)),
|
||||
|
|
@ -161,7 +161,7 @@ async fn status_timeout_and_body_errors_are_unavailable(
|
|||
let server =
|
||||
TestHttpServer::response_after(status, body, Duration::from_millis(delay_ms)).await;
|
||||
let embedder = OpenAiEmbedder::new(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
config(server.base_url(), timeout),
|
||||
);
|
||||
assert_eq!(embedder.async_embed("hello", None).await, expected);
|
||||
|
|
@ -170,7 +170,7 @@ async fn status_timeout_and_body_errors_are_unavailable(
|
|||
#[rstest]
|
||||
fn sync_embedding_is_unsupported() {
|
||||
let embedder = OpenAiEmbedder::new(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
config("http://127.0.0.1:9".to_owned(), None),
|
||||
);
|
||||
assert_eq!(
|
||||
|
|
|
|||
|
|
@ -768,7 +768,7 @@ async fn gcs() -> (wiremock::MockServer, Gcs) {
|
|||
path_service_account: None,
|
||||
endpoint: server.uri(),
|
||||
},
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
litellm_cache_response::ResponseCacheCodec,
|
||||
Arc::new(litellm_cache_gcs::StaticTokenSource("token".into())),
|
||||
)));
|
||||
|
|
|
|||
|
|
@ -34,7 +34,7 @@ pub fn config(endpoint: &str) -> S3CacheConfig {
|
|||
pub fn cache_with(config: S3CacheConfig, runtime: Handle) -> JsonS3Cache {
|
||||
S3Cache::new(
|
||||
config,
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
JsonCodec::new(),
|
||||
runtime,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -19,7 +19,7 @@ url.workspace = true
|
|||
moka.workspace = true
|
||||
sha2.workspace = true
|
||||
tower = { version = "0.5", features = ["util"] }
|
||||
rmcp = { version = "=3.4.1", default-features = false, features = ["server", "client", "elicitation", "transport-streamable-http-server", "transport-streamable-http-client", "transport-child-process"] }
|
||||
rmcp = { version = "=3.4.1", default-features = false, features = ["server", "client", "elicitation", "transport-streamable-http-server", "transport-streamable-http-client-reqwest", "transport-child-process"] }
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
|
|
|
|||
|
|
@ -21,7 +21,7 @@ litellm-auth-types.workspace = true
|
|||
litellm-config.workspace = true
|
||||
litellm-circuit-breaker.workspace = true
|
||||
litellm-redis-client.workspace = true
|
||||
litellm-http = { workspace = true, features = ["mcp"] }
|
||||
litellm-http.workspace = true
|
||||
litellm-llms.workspace = true
|
||||
litellm-secrets.workspace = true
|
||||
litellm-tracing.workspace = true
|
||||
|
|
|
|||
|
|
@ -114,11 +114,11 @@ pub async fn build_mcp(
|
|||
Auth::from_config(config, secrets.clone())
|
||||
.validate()
|
||||
.await?;
|
||||
let client = pool.mcp_client(http)?;
|
||||
let client = pool.client(http, ClientVariant::NoRedirect)?;
|
||||
Ok(Some(
|
||||
litellm_gateway_mcp::ConfiguredGateway::connect(
|
||||
&config.mcp_servers,
|
||||
client,
|
||||
(*client).clone(),
|
||||
secrets.as_ref(),
|
||||
shutdown,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -78,7 +78,9 @@ async fn upstream() -> Fixture {
|
|||
let transport = StreamableHttpService::new(
|
||||
|| Ok(Upstream),
|
||||
Arc::new(LocalSessionManager::default()),
|
||||
StreamableHttpServerConfig::default().with_cancellation_token(shutdown.clone()),
|
||||
StreamableHttpServerConfig::default()
|
||||
.with_allowed_hosts(["127.0.0.1", "upstream.invalid"])
|
||||
.with_cancellation_token(shutdown.clone()),
|
||||
);
|
||||
let app = Router::new().route_service("/mcp", transport);
|
||||
let stop = shutdown.clone();
|
||||
|
|
@ -222,6 +224,88 @@ mcp_servers:
|
|||
.unwrap();
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn mcp_tools_are_discovered_through_the_configured_proxy(#[future] upstream: Fixture) {
|
||||
let upstream = upstream.await;
|
||||
let config = Config::from_yaml(
|
||||
"general_settings: {master_key: gateway-key}\nmcp_servers: {docs: {url: 'http://upstream.invalid/mcp', timeout: 2}}",
|
||||
)
|
||||
.unwrap();
|
||||
let inference = litellm_gateway::build_inference(&config).unwrap();
|
||||
let proxy = upstream.url.trim_end_matches("/mcp");
|
||||
let http = litellm_http::HttpClientConfig {
|
||||
proxies: litellm_http::EnvironmentProxies::from_environment(&|name: &str| {
|
||||
(name == "HTTP_PROXY").then(|| proxy.to_owned())
|
||||
}),
|
||||
..inference.http.clone()
|
||||
};
|
||||
let shutdown = CancellationToken::new();
|
||||
let mcp = litellm_gateway::build_mcp(
|
||||
&config,
|
||||
inference.secrets.clone(),
|
||||
shutdown.clone(),
|
||||
&inference.resources.pool,
|
||||
&http,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let app = litellm_gateway_mcp::router(mcp.operations.clone(), HttpConfig::default());
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::get("/mcp-rest/tools/list?server_id=docs")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), 200);
|
||||
let body: Value =
|
||||
serde_json::from_slice(&to_bytes(response.into_body(), 65536).await.unwrap()).unwrap();
|
||||
assert_eq!(body["tools"][0]["name"], "echo");
|
||||
shutdown.cancel();
|
||||
tokio::time::timeout(Duration::from_secs(5), mcp.close())
|
||||
.await
|
||||
.unwrap();
|
||||
upstream.shutdown.cancel();
|
||||
tokio::time::timeout(Duration::from_secs(5), upstream.task)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn mcp_uses_the_configured_ca_bundle() {
|
||||
let directory = tempfile::tempdir().unwrap();
|
||||
let missing = directory.path().join("missing-ca.pem");
|
||||
let config = Config::from_yaml(
|
||||
"general_settings: {master_key: gateway-key}\nmcp_servers: {docs: {url: 'http://127.0.0.1:1/mcp'}}",
|
||||
)
|
||||
.unwrap();
|
||||
let inference = litellm_gateway::build_inference(&config).unwrap();
|
||||
let http = litellm_http::HttpClientConfig {
|
||||
verify: litellm_http::Verify::CaBundle(missing.clone()),
|
||||
..inference.http.clone()
|
||||
};
|
||||
let error = litellm_gateway::build_mcp(
|
||||
&config,
|
||||
inference.secrets.clone(),
|
||||
CancellationToken::new(),
|
||||
&inference.resources.pool,
|
||||
&http,
|
||||
)
|
||||
.await
|
||||
.err()
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
error,
|
||||
litellm_gateway::Error::Http(litellm_http::Error::Read { path, .. })
|
||||
if path == missing
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::unsupported_auth("auth_type: oauth2", "auth mode")]
|
||||
#[case::unsupported_policy("allowed_params: {echo: [safe]}", "allowed_params")]
|
||||
|
|
|
|||
|
|
@ -7,11 +7,8 @@ repository.workspace = true
|
|||
|
||||
[features]
|
||||
test-support = []
|
||||
mcp = ["dep:rmcp", "dep:reqwest-mcp"]
|
||||
|
||||
[dependencies]
|
||||
rmcp = { version = "=3.4.1", default-features = false, features = ["client", "transport-streamable-http-client-reqwest"], optional = true }
|
||||
reqwest-mcp = { package = "reqwest", version = "0.13.2", default-features = false, features = ["rustls-no-provider", "json", "stream", "http2"], optional = true }
|
||||
http.workspace = true
|
||||
litellm-core-utils.workspace = true
|
||||
hyper-util.workspace = true
|
||||
|
|
@ -32,5 +29,6 @@ webpki-roots.workspace = true
|
|||
futures-util.workspace = true
|
||||
rcgen = "0.14.10"
|
||||
tempfile.workspace = true
|
||||
tokio-rustls = { version = "0.26", default-features = false, features = ["ring"] }
|
||||
rstest.workspace = true
|
||||
tokio.workspace = true
|
||||
|
|
|
|||
|
|
@ -9,17 +9,7 @@ impl Client {
|
|||
}
|
||||
|
||||
#[cfg(any(test, feature = "test-support"))]
|
||||
pub fn for_test(client: reqwest::Client) -> Self {
|
||||
Self(client)
|
||||
}
|
||||
|
||||
#[cfg(any(test, feature = "test-support"))]
|
||||
pub fn plain_for_test() -> Self {
|
||||
Self(reqwest::Client::new())
|
||||
}
|
||||
|
||||
#[cfg(any(test, feature = "test-support"))]
|
||||
pub fn no_redirect_for_test() -> Self {
|
||||
pub fn for_test() -> Self {
|
||||
Self(
|
||||
reqwest::Client::builder()
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
|
|
@ -27,6 +17,11 @@ impl Client {
|
|||
.expect("a client without TLS or proxy settings builds"),
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(any(test, feature = "test-support"))]
|
||||
pub fn wrap_for_test(client: reqwest::Client) -> Self {
|
||||
Self(client)
|
||||
}
|
||||
}
|
||||
|
||||
impl Deref for Client {
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ pub enum Verify {
|
|||
Disabled,
|
||||
CaBundle(PathBuf),
|
||||
BuiltInRoots,
|
||||
NativeRoots,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
|
||||
|
|
@ -98,7 +99,7 @@ impl TryFrom<&HttpClientConfig> for reqwest::ClientBuilder {
|
|||
|
||||
fn try_from(config: &HttpClientConfig) -> Result<Self, Self::Error> {
|
||||
let base = reqwest::Client::builder()
|
||||
.use_preconfigured_tls(rustls::ClientConfig::try_from(config)?)
|
||||
.tls_backend_preconfigured(rustls::ClientConfig::try_from(config)?)
|
||||
.connect_timeout(config.connect_timeout)
|
||||
.pool_idle_timeout(config.pool_idle_timeout);
|
||||
let with_keepalive = match config.tcp_keepalive {
|
||||
|
|
|
|||
|
|
@ -7,8 +7,6 @@
|
|||
mod client;
|
||||
mod config;
|
||||
mod error;
|
||||
#[cfg(feature = "mcp")]
|
||||
mod mcp;
|
||||
pub mod media;
|
||||
pub mod outbound;
|
||||
mod pool;
|
||||
|
|
|
|||
|
|
@ -1,76 +0,0 @@
|
|||
use std::{
|
||||
collections::HashMap,
|
||||
net::{IpAddr, Ipv4Addr},
|
||||
sync::{Mutex, PoisonError},
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
|
||||
use crate::{Error, HttpClientConfig};
|
||||
|
||||
struct Entry {
|
||||
client: reqwest_mcp::Client,
|
||||
built_at: Instant,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub(super) struct Pool(Mutex<HashMap<HttpClientConfig, Entry>>);
|
||||
|
||||
impl Pool {
|
||||
pub fn client(
|
||||
&self,
|
||||
config: &HttpClientConfig,
|
||||
ttl: Duration,
|
||||
) -> Result<reqwest_mcp::Client, Error> {
|
||||
let mut clients = self.0.lock().unwrap_or_else(PoisonError::into_inner);
|
||||
if let Some(entry) = clients.get(config)
|
||||
&& entry.built_at.elapsed() < ttl
|
||||
{
|
||||
return Ok(entry.client.clone());
|
||||
}
|
||||
let client = build(config)?;
|
||||
clients.insert(
|
||||
config.clone(),
|
||||
Entry {
|
||||
client: client.clone(),
|
||||
built_at: Instant::now(),
|
||||
},
|
||||
);
|
||||
Ok(client)
|
||||
}
|
||||
}
|
||||
|
||||
fn build(config: &HttpClientConfig) -> Result<reqwest_mcp::Client, Error> {
|
||||
let base = reqwest_mcp::Client::builder()
|
||||
.tls_backend_preconfigured(rustls::ClientConfig::try_from(config)?)
|
||||
.connect_timeout(config.connect_timeout)
|
||||
.pool_idle_timeout(config.pool_idle_timeout)
|
||||
.redirect(reqwest_mcp::redirect::Policy::none());
|
||||
let keepalive = match config.tcp_keepalive {
|
||||
Some(value) => base
|
||||
.tcp_keepalive(value.idle)
|
||||
.tcp_keepalive_interval(value.interval)
|
||||
.tcp_keepalive_retries(value.retries),
|
||||
None => base,
|
||||
};
|
||||
let address = if config.force_ipv4 {
|
||||
keepalive.local_address(IpAddr::V4(Ipv4Addr::UNSPECIFIED))
|
||||
} else {
|
||||
keepalive
|
||||
};
|
||||
let protocol = if config.http2 {
|
||||
address
|
||||
} else {
|
||||
address.http1_only()
|
||||
};
|
||||
let agent = match &config.user_agent {
|
||||
Some(value) => protocol.user_agent(value),
|
||||
None => protocol,
|
||||
};
|
||||
config
|
||||
.proxies
|
||||
.mcp_proxies()
|
||||
.into_iter()
|
||||
.fold(agent.no_proxy(), reqwest_mcp::ClientBuilder::proxy)
|
||||
.build()
|
||||
.map_err(|error| Error::Client(error.without_url().to_string()))
|
||||
}
|
||||
|
|
@ -522,7 +522,7 @@ mod tests {
|
|||
b"HTTP/1.1 200 OK\r\nContent-Type: application/pdf; charset=binary\r\nContent-Length: 3\r\nConnection: close\r\n\r\nabc",
|
||||
)
|
||||
.await;
|
||||
let client = Client::no_redirect_for_test();
|
||||
let client = Client::for_test();
|
||||
let media = MediaFetcher::for_test(client)
|
||||
.fetch(url, policy(3, 0))
|
||||
.await
|
||||
|
|
@ -538,7 +538,7 @@ mod tests {
|
|||
b"HTTP/1.1 200 OK\r\nContent-Type: application/pdf\r\nContent-Length: 3\r\nConnection: close\r\n\r\nabc",
|
||||
)
|
||||
.await;
|
||||
let client = Client::no_redirect_for_test();
|
||||
let client = Client::for_test();
|
||||
let error = MediaFetcher::for_test(client)
|
||||
.fetch(url, policy(2, 0))
|
||||
.await
|
||||
|
|
@ -553,7 +553,7 @@ mod tests {
|
|||
b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n2\r\nab\r\n2\r\ncd\r\n0\r\n\r\n",
|
||||
)
|
||||
.await;
|
||||
let client = Client::no_redirect_for_test();
|
||||
let client = Client::for_test();
|
||||
let error = MediaFetcher::for_test(client)
|
||||
.fetch(url, policy(3, 0))
|
||||
.await
|
||||
|
|
|
|||
|
|
@ -6,7 +6,12 @@ use std::{
|
|||
|
||||
use reqwest::dns::Resolve;
|
||||
|
||||
use crate::{client::Client, config::HttpClientConfig, error::Error, proxy::EnvironmentProxies};
|
||||
use crate::{
|
||||
client::Client,
|
||||
config::{HttpClientConfig, Verify},
|
||||
error::Error,
|
||||
proxy::EnvironmentProxies,
|
||||
};
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
|
||||
pub enum ClientVariant {
|
||||
|
|
@ -14,6 +19,7 @@ pub enum ClientVariant {
|
|||
NoRedirect,
|
||||
Media,
|
||||
UnpinnedMedia,
|
||||
WebSocket,
|
||||
}
|
||||
|
||||
const CLIENT_TTL: Duration = Duration::from_secs(3600);
|
||||
|
|
@ -29,8 +35,6 @@ pub struct HttpClientPool {
|
|||
media_resolver: Arc<dyn Resolve>,
|
||||
ttl: Duration,
|
||||
clients: Mutex<Clients>,
|
||||
#[cfg(feature = "mcp")]
|
||||
mcp: crate::mcp::Pool,
|
||||
}
|
||||
|
||||
impl HttpClientPool {
|
||||
|
|
@ -43,19 +47,9 @@ impl HttpClientPool {
|
|||
media_resolver,
|
||||
ttl,
|
||||
clients: Mutex::default(),
|
||||
#[cfg(feature = "mcp")]
|
||||
mcp: crate::mcp::Pool::default(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "mcp")]
|
||||
pub fn mcp_client(
|
||||
&self,
|
||||
config: &HttpClientConfig,
|
||||
) -> Result<impl rmcp::transport::streamable_http_client::StreamableHttpClient, Error> {
|
||||
self.mcp.client(config, self.ttl)
|
||||
}
|
||||
|
||||
pub fn client(
|
||||
&self,
|
||||
config: &HttpClientConfig,
|
||||
|
|
@ -71,6 +65,14 @@ impl HttpClientPool {
|
|||
client_certificate: None,
|
||||
..config.clone()
|
||||
},
|
||||
ClientVariant::WebSocket => HttpClientConfig {
|
||||
http2: false,
|
||||
verify: match &config.verify {
|
||||
Verify::BuiltInRoots => Verify::NativeRoots,
|
||||
verify => verify.clone(),
|
||||
},
|
||||
..config.clone()
|
||||
},
|
||||
ClientVariant::Provider | ClientVariant::NoRedirect => config.clone(),
|
||||
};
|
||||
let key = (effective, variant);
|
||||
|
|
@ -103,12 +105,12 @@ impl HttpClientPool {
|
|||
) -> reqwest::ClientBuilder {
|
||||
match variant {
|
||||
ClientVariant::Provider => builder,
|
||||
ClientVariant::NoRedirect | ClientVariant::UnpinnedMedia => {
|
||||
ClientVariant::NoRedirect | ClientVariant::UnpinnedMedia | ClientVariant::WebSocket => {
|
||||
builder.redirect(reqwest::redirect::Policy::none())
|
||||
}
|
||||
ClientVariant::Media => builder
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.dns_resolver2(Arc::clone(&self.media_resolver)),
|
||||
.dns_resolver(Arc::clone(&self.media_resolver)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -390,56 +392,4 @@ mod tests {
|
|||
.is_err()
|
||||
);
|
||||
}
|
||||
#[cfg(feature = "mcp")]
|
||||
#[rstest::rstest]
|
||||
#[tokio::test]
|
||||
async fn mcp_clients_reuse_connections_and_apply_proxy_headers_without_redirects() {
|
||||
let (address, connections, requests) = serve("HTTP/1.1 302 Found").await;
|
||||
let pool = pool();
|
||||
let config = HttpClientConfig {
|
||||
proxies: proxied_through(&format!("http://{address}")),
|
||||
..config("mcp-pool-test")
|
||||
};
|
||||
for _ in 0..2 {
|
||||
let response = pool
|
||||
.mcp
|
||||
.client(&config, pool.ttl)
|
||||
.unwrap()
|
||||
.get("http://upstream.invalid/mcp")
|
||||
.timeout(Duration::from_secs(2))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), 302);
|
||||
response.bytes().await.unwrap();
|
||||
}
|
||||
assert_eq!(connections.load(Ordering::SeqCst), 1);
|
||||
let requests = requests.lock().unwrap();
|
||||
assert_eq!(requests.len(), 2);
|
||||
assert!(
|
||||
requests
|
||||
.iter()
|
||||
.all(|request| request.starts_with("GET http://upstream.invalid/mcp HTTP/1.1"))
|
||||
);
|
||||
assert!(
|
||||
requests
|
||||
.iter()
|
||||
.all(|request| request.contains("user-agent: mcp-pool-test"))
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(feature = "mcp")]
|
||||
#[rstest::rstest]
|
||||
#[tokio::test]
|
||||
async fn mcp_client_uses_host_tls_configuration() {
|
||||
let directory = tempfile::tempdir().unwrap();
|
||||
let config = HttpClientConfig {
|
||||
verify: Verify::CaBundle(directory.path().join("missing.pem")),
|
||||
..config("mcp-test")
|
||||
};
|
||||
assert!(matches!(
|
||||
pool().mcp_client(&config),
|
||||
Err(Error::Read { .. })
|
||||
));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -63,19 +63,6 @@ impl EnvironmentProxies {
|
|||
.map(|proxy| proxy.no_proxy(no_proxy.clone()))
|
||||
.collect()
|
||||
}
|
||||
#[cfg(feature = "mcp")]
|
||||
pub(crate) fn mcp_proxies(&self) -> Vec<reqwest_mcp::Proxy> {
|
||||
let no_proxy = reqwest_mcp::NoProxy::from_string(&self.no);
|
||||
[
|
||||
reqwest_mcp::Proxy::http(self.http.as_str()),
|
||||
reqwest_mcp::Proxy::https(self.https.as_str()),
|
||||
reqwest_mcp::Proxy::all(self.all.as_str()),
|
||||
]
|
||||
.into_iter()
|
||||
.filter_map(Result::ok)
|
||||
.map(|proxy| proxy.no_proxy(no_proxy.clone()))
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
|
|
|||
|
|
@ -195,6 +195,12 @@ impl TryFrom<&HttpClientConfig> for ClientConfig {
|
|||
Verify::BuiltInRoots => builder.with_root_certificates(RootCertStore {
|
||||
roots: webpki_roots::TLS_SERVER_ROOTS.to_vec(),
|
||||
}),
|
||||
Verify::NativeRoots => {
|
||||
let native = rustls_native_certs::load_native_certs();
|
||||
let mut roots = RootCertStore::empty();
|
||||
roots.add_parsable_certificates(native.certs);
|
||||
builder.with_root_certificates(roots)
|
||||
}
|
||||
Verify::CaBundle(path) => {
|
||||
builder.with_root_certificates(bundle_roots(path, TlsSource::CaBundle)?)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,62 +1,135 @@
|
|||
use std::{
|
||||
io,
|
||||
sync::{Arc, OnceLock},
|
||||
use http::{
|
||||
HeaderMap, HeaderValue, StatusCode,
|
||||
header::{
|
||||
CONNECTION, SEC_WEBSOCKET_ACCEPT, SEC_WEBSOCKET_KEY, SEC_WEBSOCKET_PROTOCOL,
|
||||
SEC_WEBSOCKET_VERSION, UPGRADE,
|
||||
},
|
||||
};
|
||||
|
||||
use rustls::{ClientConfig, RootCertStore};
|
||||
use tokio::net::TcpStream;
|
||||
use reqwest::Url;
|
||||
use tokio_tungstenite::{
|
||||
Connector, MaybeTlsStream, WebSocketStream, connect_async_tls_with_config,
|
||||
tungstenite::{client::IntoClientRequest, error::TlsError, handshake::client::Response},
|
||||
WebSocketStream,
|
||||
tungstenite::{
|
||||
handshake::{client::generate_key, derive_accept_key},
|
||||
protocol::Role,
|
||||
},
|
||||
};
|
||||
|
||||
pub type UpstreamWebSocket = WebSocketStream<MaybeTlsStream<TcpStream>>;
|
||||
use crate::{client::Client, transport::Error};
|
||||
|
||||
static TLS_CONFIG: OnceLock<Arc<ClientConfig>> = OnceLock::new();
|
||||
|
||||
fn build_tls_config() -> Result<ClientConfig, Box<tokio_tungstenite::tungstenite::Error>> {
|
||||
let native = rustls_native_certs::load_native_certs();
|
||||
let mut store = RootCertStore::empty();
|
||||
let (added, _ignored) = store.add_parsable_certificates(native.certs);
|
||||
if added == 0 {
|
||||
return Err(Box::new(tokio_tungstenite::tungstenite::Error::Io(
|
||||
io::Error::other(format!(
|
||||
"no usable native root certificates: {:?}",
|
||||
native.errors
|
||||
)),
|
||||
)));
|
||||
}
|
||||
ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider()))
|
||||
.with_safe_default_protocol_versions()
|
||||
.map(|builder| builder.with_root_certificates(store).with_no_client_auth())
|
||||
.map_err(|error| {
|
||||
Box::new(tokio_tungstenite::tungstenite::Error::Tls(
|
||||
TlsError::Rustls(error),
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
fn tls_config() -> Result<Arc<ClientConfig>, Box<tokio_tungstenite::tungstenite::Error>> {
|
||||
if let Some(config) = TLS_CONFIG.get() {
|
||||
return Ok(Arc::clone(config));
|
||||
}
|
||||
let built = Arc::new(build_tls_config()?);
|
||||
Ok(Arc::clone(TLS_CONFIG.get_or_init(|| built)))
|
||||
}
|
||||
pub type UpstreamWebSocket = WebSocketStream<reqwest::Upgraded>;
|
||||
|
||||
#[tracing::instrument(name = "litellm.websocket.handshake", level = "debug", skip_all)]
|
||||
pub async fn connect_upstream<R>(
|
||||
request: R,
|
||||
) -> Result<(UpstreamWebSocket, Response), Box<tokio_tungstenite::tungstenite::Error>>
|
||||
where
|
||||
R: IntoClientRequest + Unpin,
|
||||
{
|
||||
let request = request.into_client_request().map_err(Box::new)?;
|
||||
let connector = match request.uri().scheme_str() {
|
||||
Some("wss") => Some(Connector::Rustls(tls_config()?)),
|
||||
_ => None,
|
||||
};
|
||||
connect_async_tls_with_config(request, None, false, connector)
|
||||
pub async fn connect(
|
||||
client: &Client,
|
||||
url: &str,
|
||||
headers: HeaderMap,
|
||||
) -> Result<UpstreamWebSocket, Error> {
|
||||
let offered_protocols = headers.get(SEC_WEBSOCKET_PROTOCOL).cloned();
|
||||
let key = generate_key();
|
||||
let response = client
|
||||
.get(http_url(url)?)
|
||||
.headers(handshake_headers(headers, &key)?)
|
||||
.send()
|
||||
.await
|
||||
.map_err(Box::new)
|
||||
.map_err(Error::from_reqwest_before_dispatch)?;
|
||||
let status = response.status();
|
||||
if status != StatusCode::SWITCHING_PROTOCOLS {
|
||||
return Err(Error::Http {
|
||||
status: status.as_u16(),
|
||||
body: response.text().await.unwrap_or_default(),
|
||||
});
|
||||
}
|
||||
let expected = derive_accept_key(key.as_bytes());
|
||||
if response
|
||||
.headers()
|
||||
.get(SEC_WEBSOCKET_ACCEPT)
|
||||
.map(HeaderValue::as_bytes)
|
||||
!= Some(expected.as_bytes())
|
||||
{
|
||||
return Err(Error::Network(
|
||||
"WebSocket handshake returned a wrong Sec-WebSocket-Accept".into(),
|
||||
));
|
||||
}
|
||||
validate_response_headers(response.headers(), offered_protocols.as_ref())?;
|
||||
let upgraded = response.upgrade().await?;
|
||||
Ok(WebSocketStream::from_raw_socket(upgraded, Role::Client, None).await)
|
||||
}
|
||||
|
||||
fn validate_response_headers(
|
||||
headers: &HeaderMap,
|
||||
offered_protocols: Option<&HeaderValue>,
|
||||
) -> Result<(), Error> {
|
||||
if !headers
|
||||
.get(UPGRADE)
|
||||
.is_some_and(|value| value.as_bytes().eq_ignore_ascii_case(b"websocket"))
|
||||
{
|
||||
return Err(Error::Network(
|
||||
"WebSocket handshake returned an invalid Upgrade header".into(),
|
||||
));
|
||||
}
|
||||
if !headers.get_all(CONNECTION).iter().any(|value| {
|
||||
value.to_str().is_ok_and(|value| {
|
||||
value
|
||||
.split(',')
|
||||
.any(|token| token.trim().eq_ignore_ascii_case("upgrade"))
|
||||
})
|
||||
}) {
|
||||
return Err(Error::Network(
|
||||
"WebSocket handshake returned an invalid Connection header".into(),
|
||||
));
|
||||
}
|
||||
let protocol_matches = match (offered_protocols, headers.get(SEC_WEBSOCKET_PROTOCOL)) {
|
||||
(_, None) => true,
|
||||
(Some(offered), Some(selected)) => offered
|
||||
.to_str()
|
||||
.ok()
|
||||
.zip(selected.to_str().ok())
|
||||
.is_some_and(|(offered, selected)| {
|
||||
offered
|
||||
.split(',')
|
||||
.any(|protocol| protocol.trim() == selected)
|
||||
}),
|
||||
_ => false,
|
||||
};
|
||||
if !protocol_matches {
|
||||
return Err(Error::Network(
|
||||
"WebSocket handshake returned an invalid subprotocol".into(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn http_url(url: &str) -> Result<Url, Error> {
|
||||
let mut url = Url::parse(url).map_err(not_dispatched)?;
|
||||
let scheme = match url.scheme() {
|
||||
"ws" | "http" => "http",
|
||||
"wss" | "https" => "https",
|
||||
other => {
|
||||
return Err(not_dispatched(format!(
|
||||
"unsupported WebSocket URL scheme {other:?}"
|
||||
)));
|
||||
}
|
||||
};
|
||||
url.set_scheme(scheme)
|
||||
.map_err(|()| not_dispatched("WebSocket URL cannot be sent over HTTP"))?;
|
||||
Ok(url)
|
||||
}
|
||||
|
||||
fn handshake_headers(headers: HeaderMap, key: &str) -> Result<HeaderMap, Error> {
|
||||
let key = HeaderValue::from_str(key).map_err(not_dispatched)?;
|
||||
Ok([
|
||||
(CONNECTION, HeaderValue::from_static("Upgrade")),
|
||||
(UPGRADE, HeaderValue::from_static("websocket")),
|
||||
(SEC_WEBSOCKET_VERSION, HeaderValue::from_static("13")),
|
||||
(SEC_WEBSOCKET_KEY, key),
|
||||
]
|
||||
.into_iter()
|
||||
.fold(headers, |mut headers, (name, value)| {
|
||||
headers.insert(name, value);
|
||||
headers
|
||||
}))
|
||||
}
|
||||
|
||||
fn not_dispatched(error: impl std::fmt::Display) -> Error {
|
||||
Error::Connect(error.to_string())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,59 +1,539 @@
|
|||
use std::{
|
||||
net::SocketAddr,
|
||||
path::PathBuf,
|
||||
sync::{Arc, Mutex},
|
||||
};
|
||||
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use litellm_http::websocket::connect_upstream;
|
||||
use rstest::rstest;
|
||||
use tokio::net::TcpListener;
|
||||
use http::{HeaderMap, HeaderValue};
|
||||
use litellm_http::{
|
||||
ClientIdentity, ClientVariant, EnvironmentProxies, HttpClientConfig, HttpClientPool,
|
||||
HttpSettings, Resolution, Verify, media::PublicDnsResolver, transport::Error,
|
||||
websocket::connect,
|
||||
};
|
||||
use rcgen::{BasicConstraints, CertificateParams, CertifiedIssuer, IsCa, KeyPair};
|
||||
use rstest::{Context, rstest};
|
||||
use rustls::{
|
||||
RootCertStore, ServerConfig,
|
||||
crypto::ring,
|
||||
pki_types::{CertificateDer, PrivateKeyDer},
|
||||
server::WebPkiClientVerifier,
|
||||
};
|
||||
use tokio::{
|
||||
io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt},
|
||||
net::{TcpListener, TcpStream},
|
||||
};
|
||||
use tokio_rustls::TlsAcceptor;
|
||||
use tokio_tungstenite::{
|
||||
accept_hdr_async,
|
||||
tungstenite::{
|
||||
Message,
|
||||
client::IntoClientRequest,
|
||||
handshake::server::{Request, Response},
|
||||
},
|
||||
};
|
||||
|
||||
struct Pki {
|
||||
directory: tempfile::TempDir,
|
||||
ca: CertifiedIssuer<'static, KeyPair>,
|
||||
server_chain: Vec<CertificateDer<'static>>,
|
||||
server_key: Vec<u8>,
|
||||
}
|
||||
|
||||
impl Pki {
|
||||
fn new() -> Self {
|
||||
let mut ca_params = CertificateParams::new(Vec::<String>::new()).unwrap();
|
||||
ca_params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained);
|
||||
let ca = CertifiedIssuer::self_signed(ca_params, KeyPair::generate().unwrap()).unwrap();
|
||||
let server_key = KeyPair::generate().unwrap();
|
||||
let server = CertificateParams::new(vec!["localhost".into()])
|
||||
.unwrap()
|
||||
.signed_by(&server_key, &ca)
|
||||
.unwrap();
|
||||
let client_key = KeyPair::generate().unwrap();
|
||||
let client = CertificateParams::new(vec!["client".into()])
|
||||
.unwrap()
|
||||
.signed_by(&client_key, &ca)
|
||||
.unwrap();
|
||||
let directory = tempfile::tempdir().unwrap();
|
||||
std::fs::write(directory.path().join("ca.pem"), ca.pem()).unwrap();
|
||||
std::fs::write(
|
||||
directory.path().join("client.pem"),
|
||||
format!("{}{}", client.pem(), client_key.serialize_pem()),
|
||||
)
|
||||
.unwrap();
|
||||
Self {
|
||||
directory,
|
||||
server_chain: vec![server.der().clone()],
|
||||
server_key: server_key.serialize_der(),
|
||||
ca,
|
||||
}
|
||||
}
|
||||
|
||||
fn path(&self, name: &str) -> PathBuf {
|
||||
self.directory.path().join(name)
|
||||
}
|
||||
|
||||
fn acceptor(&self, require_client_certificate: bool) -> TlsAcceptor {
|
||||
let provider = Arc::new(ring::default_provider());
|
||||
let builder = ServerConfig::builder_with_provider(Arc::clone(&provider))
|
||||
.with_safe_default_protocol_versions()
|
||||
.unwrap();
|
||||
let builder = if require_client_certificate {
|
||||
let mut roots = RootCertStore::empty();
|
||||
roots.add(self.ca.der().clone()).unwrap();
|
||||
builder.with_client_cert_verifier(
|
||||
WebPkiClientVerifier::builder_with_provider(Arc::new(roots), provider)
|
||||
.build()
|
||||
.unwrap(),
|
||||
)
|
||||
} else {
|
||||
builder.with_no_client_auth()
|
||||
};
|
||||
let mut config = builder
|
||||
.with_single_cert(
|
||||
self.server_chain.clone(),
|
||||
PrivateKeyDer::try_from(self.server_key.clone()).unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
|
||||
TlsAcceptor::from(Arc::new(config))
|
||||
}
|
||||
}
|
||||
|
||||
fn config(verify: Verify) -> HttpClientConfig {
|
||||
HttpClientConfig {
|
||||
verify,
|
||||
http2: true,
|
||||
..Resolution::from(&HttpSettings::default()).config
|
||||
}
|
||||
}
|
||||
|
||||
fn proxied_through(proxy: SocketAddr) -> EnvironmentProxies {
|
||||
EnvironmentProxies::from_environment(&move |name: &str| {
|
||||
(name == "HTTPS_PROXY").then(|| format!("http://{proxy}"))
|
||||
})
|
||||
}
|
||||
|
||||
async fn open(
|
||||
config: &HttpClientConfig,
|
||||
url: &str,
|
||||
headers: HeaderMap,
|
||||
) -> Result<litellm_http::websocket::UpstreamWebSocket, Error> {
|
||||
let client = HttpClientPool::new(Arc::new(PublicDnsResolver))
|
||||
.client(config, ClientVariant::WebSocket)
|
||||
.unwrap();
|
||||
connect(&client, url, headers).await
|
||||
}
|
||||
|
||||
async fn echo_once<S: AsyncRead + AsyncWrite + Unpin>(
|
||||
stream: S,
|
||||
seen: Arc<Mutex<Option<Request>>>,
|
||||
) {
|
||||
#[expect(
|
||||
clippy::result_large_err,
|
||||
reason = "tungstenite requires an unboxed handshake error response"
|
||||
)]
|
||||
let Ok(mut socket) = accept_hdr_async(stream, |request: &Request, response: Response| {
|
||||
*seen.lock().unwrap() = Some(request.clone());
|
||||
Ok(response)
|
||||
})
|
||||
.await
|
||||
else {
|
||||
return;
|
||||
};
|
||||
if let Some(Ok(message)) = socket.next().await {
|
||||
let _ = socket.send(message).await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn serve_tls(
|
||||
pki: &Pki,
|
||||
require_client_certificate: bool,
|
||||
) -> (SocketAddr, Arc<Mutex<Option<Request>>>) {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let acceptor = pki.acceptor(require_client_certificate);
|
||||
let seen = Arc::new(Mutex::new(None));
|
||||
let recorded = Arc::clone(&seen);
|
||||
tokio::spawn(async move {
|
||||
while let Ok((stream, _)) = listener.accept().await {
|
||||
if let Ok(stream) = acceptor.accept(stream).await {
|
||||
echo_once(stream, Arc::clone(&recorded)).await;
|
||||
}
|
||||
}
|
||||
});
|
||||
(address, seen)
|
||||
}
|
||||
|
||||
async fn read_head(stream: &mut TcpStream) -> String {
|
||||
let mut head = Vec::new();
|
||||
while !head.ends_with(b"\r\n\r\n") {
|
||||
let mut byte = [0u8; 1];
|
||||
if stream.read(&mut byte).await.unwrap() == 0 {
|
||||
break;
|
||||
}
|
||||
head.push(byte[0]);
|
||||
}
|
||||
String::from_utf8(head).unwrap()
|
||||
}
|
||||
|
||||
async fn serve_tunnel_proxy(target: SocketAddr) -> (SocketAddr, Arc<Mutex<Vec<String>>>) {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let requests = Arc::new(Mutex::new(Vec::new()));
|
||||
let seen = Arc::clone(&requests);
|
||||
tokio::spawn(async move {
|
||||
while let Ok((mut client, _)) = listener.accept().await {
|
||||
let head = read_head(&mut client).await;
|
||||
seen.lock()
|
||||
.unwrap()
|
||||
.push(head.lines().next().unwrap_or_default().to_owned());
|
||||
let mut upstream = TcpStream::connect(target).await.unwrap();
|
||||
client
|
||||
.write_all(b"HTTP/1.1 200 Connection established\r\n\r\n")
|
||||
.await
|
||||
.unwrap();
|
||||
let _ = tokio::io::copy_bidirectional(&mut client, &mut upstream).await;
|
||||
}
|
||||
});
|
||||
(address, requests)
|
||||
}
|
||||
|
||||
async fn serve_raw(response: &'static str) -> SocketAddr {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
tokio::spawn(async move {
|
||||
let (mut stream, _) = listener.accept().await.unwrap();
|
||||
read_head(&mut stream).await;
|
||||
stream.write_all(response.as_bytes()).await.unwrap();
|
||||
let _ = stream.read(&mut [0u8; 1]).await;
|
||||
});
|
||||
address
|
||||
}
|
||||
|
||||
async fn round_trip(socket: &mut litellm_http::websocket::UpstreamWebSocket) {
|
||||
socket.send(Message::Text("ping".into())).await.unwrap();
|
||||
assert_eq!(
|
||||
socket.next().await.unwrap().unwrap(),
|
||||
Message::Text("ping".into())
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::without_query("/responses")]
|
||||
#[case::with_query("/responses?model=test-model")]
|
||||
#[tokio::test]
|
||||
async fn connects_with_caller_headers_and_exchanges_frames(#[case] path: &'static str) {
|
||||
async fn plain_websocket_sends_caller_headers_and_path(#[case] path: &'static str) {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let server = tokio::spawn(async move {
|
||||
let seen = Arc::new(Mutex::new(None));
|
||||
let recorded = Arc::clone(&seen);
|
||||
tokio::spawn(async move {
|
||||
let (stream, _) = listener.accept().await.unwrap();
|
||||
echo_once(stream, recorded).await;
|
||||
});
|
||||
let headers = HeaderMap::from_iter([(
|
||||
http::header::AUTHORIZATION,
|
||||
HeaderValue::from_static("Bearer test-key"),
|
||||
)]);
|
||||
|
||||
let mut socket = open(
|
||||
&config(Verify::BuiltInRoots),
|
||||
&format!("ws://{address}{path}"),
|
||||
headers,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
round_trip(&mut socket).await;
|
||||
let request = seen.lock().unwrap().take().unwrap();
|
||||
assert_eq!(request.uri().path_and_query().unwrap().as_str(), path);
|
||||
assert_eq!(request.headers()["authorization"], "Bearer test-key");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::ca_bundle_trusts_the_private_ca(|pki: &Pki| Verify::CaBundle(pki.path("ca.pem")), true)]
|
||||
#[case::ssl_verify_off_skips_verification(|_: &Pki| Verify::Disabled, true)]
|
||||
#[case::built_in_roots_reject_the_private_ca(|_: &Pki| Verify::BuiltInRoots, false)]
|
||||
#[tokio::test]
|
||||
async fn wss_honors_ssl_verify(#[case] verify: fn(&Pki) -> Verify, #[case] trusted: bool) {
|
||||
let pki = Pki::new();
|
||||
let (address, _) = serve_tls(&pki, false).await;
|
||||
|
||||
let result = open(
|
||||
&config(verify(&pki)),
|
||||
&format!("wss://localhost:{}/responses", address.port()),
|
||||
HeaderMap::new(),
|
||||
)
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(mut socket) if trusted => round_trip(&mut socket).await,
|
||||
Err(Error::Connect(_)) if !trusted => {}
|
||||
other => panic!("unexpected handshake outcome: {:?}", other.map(drop)),
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::plain_without_roots(false, false, true)]
|
||||
#[case::tls_with_roots(true, true, true)]
|
||||
#[case::tls_without_roots(true, false, false)]
|
||||
#[tokio::test]
|
||||
async fn websocket_uses_system_roots(
|
||||
#[context] context: Context,
|
||||
#[case] tls: bool,
|
||||
#[case] trusted_roots: bool,
|
||||
#[case] connects: bool,
|
||||
) {
|
||||
const CHILD_URL: &str = "LITELLM_TEST_SYSTEM_ROOTS_WSS_URL";
|
||||
if let Ok(url) = std::env::var(CHILD_URL) {
|
||||
let result = open(&config(Verify::BuiltInRoots), &url, HeaderMap::new()).await;
|
||||
match result {
|
||||
Ok(mut socket) if connects => round_trip(&mut socket).await,
|
||||
Err(Error::Connect(_)) if !connects => {}
|
||||
other => panic!("unexpected handshake outcome: {:?}", other.map(drop)),
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
let pki = Pki::new();
|
||||
let url = if tls {
|
||||
let (address, _) = serve_tls(&pki, false).await;
|
||||
format!("wss://localhost:{}/responses", address.port())
|
||||
} else {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
tokio::spawn(async move {
|
||||
let (stream, _) = listener.accept().await.unwrap();
|
||||
echo_once(stream, Arc::new(Mutex::new(None))).await;
|
||||
});
|
||||
format!("ws://{address}/responses")
|
||||
};
|
||||
let ca = if trusted_roots {
|
||||
pki.path("ca.pem")
|
||||
} else {
|
||||
let empty = pki.path("empty.pem");
|
||||
std::fs::write(&empty, "").unwrap();
|
||||
empty
|
||||
};
|
||||
let test = format!(
|
||||
"{}::case_{}_{}",
|
||||
context.name,
|
||||
context.case.unwrap() + 1,
|
||||
context.description.unwrap()
|
||||
);
|
||||
let output = tokio::task::spawn_blocking(move || {
|
||||
std::process::Command::new(std::env::current_exe().unwrap())
|
||||
.args(["--exact", &test, "--nocapture"])
|
||||
.env(CHILD_URL, url)
|
||||
.env("SSL_CERT_FILE", ca)
|
||||
.env("SSL_CERT_DIR", "")
|
||||
.output()
|
||||
.unwrap()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
let stdout = String::from_utf8_lossy(&output.stdout);
|
||||
assert!(
|
||||
output.status.success() && stdout.contains("1 passed"),
|
||||
"system-root WebSocket child failed:\n{stdout}\n{}",
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::with_certificate(true)]
|
||||
#[case::without_certificate(false)]
|
||||
#[tokio::test]
|
||||
async fn wss_presents_the_configured_client_certificate(#[case] configured: bool) {
|
||||
let pki = Pki::new();
|
||||
let (address, _) = serve_tls(&pki, true).await;
|
||||
let config = HttpClientConfig {
|
||||
client_certificate: configured.then(|| ClientIdentity::Pem(pki.path("client.pem"))),
|
||||
..config(Verify::CaBundle(pki.path("ca.pem")))
|
||||
};
|
||||
|
||||
let result = open(
|
||||
&config,
|
||||
&format!("wss://localhost:{}/responses", address.port()),
|
||||
HeaderMap::new(),
|
||||
)
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(mut socket) if configured => round_trip(&mut socket).await,
|
||||
Err(_) if !configured => {}
|
||||
other => panic!("unexpected handshake outcome: {:?}", other.map(drop)),
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn wss_tunnels_through_the_https_proxy() {
|
||||
let pki = Pki::new();
|
||||
let (target, _) = serve_tls(&pki, false).await;
|
||||
let (proxy, requests) = serve_tunnel_proxy(target).await;
|
||||
let config = HttpClientConfig {
|
||||
proxies: proxied_through(proxy),
|
||||
..config(Verify::CaBundle(pki.path("ca.pem")))
|
||||
};
|
||||
|
||||
let mut socket = open(
|
||||
&config,
|
||||
&format!("wss://localhost:{}/responses", target.port()),
|
||||
HeaderMap::new(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
round_trip(&mut socket).await;
|
||||
assert_eq!(
|
||||
*requests.lock().unwrap(),
|
||||
[format!("CONNECT localhost:{} HTTP/1.1", target.port())]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn refused_upgrade_reports_the_status_and_body() {
|
||||
let address =
|
||||
serve_raw("HTTP/1.1 403 Forbidden\r\nContent-Length: 6\r\nConnection: close\r\n\r\ndenied")
|
||||
.await;
|
||||
|
||||
let result = open(
|
||||
&config(Verify::BuiltInRoots),
|
||||
&format!("ws://{address}/responses"),
|
||||
HeaderMap::new(),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(
|
||||
result.map(drop),
|
||||
Err(Error::Http {
|
||||
status: 403,
|
||||
body: "denied".into()
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn upgrade_with_a_wrong_accept_key_is_rejected() {
|
||||
let address = serve_raw(
|
||||
"HTTP/1.1 101 Switching Protocols\r\nConnection: Upgrade\r\nUpgrade: websocket\r\nSec-WebSocket-Accept: wrong\r\n\r\n",
|
||||
)
|
||||
.await;
|
||||
|
||||
let result = open(
|
||||
&config(Verify::BuiltInRoots),
|
||||
&format!("ws://{address}/responses"),
|
||||
HeaderMap::new(),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(matches!(result, Err(Error::Network(_))));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::non_websocket_scheme("ftp://localhost/responses")]
|
||||
#[case::unparseable("not a url")]
|
||||
#[tokio::test]
|
||||
async fn invalid_url_is_rejected_before_dispatch(#[case] url: &str) {
|
||||
let result = open(&config(Verify::BuiltInRoots), url, HeaderMap::new()).await;
|
||||
|
||||
assert!(matches!(result, Err(Error::Connect(_))));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::missing_upgrade(None, Some("Upgrade"), None, None, false)]
|
||||
#[case::wrong_upgrade(Some("http"), Some("Upgrade"), None, None, false)]
|
||||
#[case::missing_connection(Some("websocket"), None, None, None, false)]
|
||||
#[case::wrong_connection(Some("websocket"), Some("keep-alive"), None, None, false)]
|
||||
#[case::unsolicited_protocol(Some("websocket"), Some("Upgrade"), None, Some("other"), false)]
|
||||
#[case::unoffered_protocol(
|
||||
Some("websocket"),
|
||||
Some("Upgrade"),
|
||||
Some("events"),
|
||||
Some("other"),
|
||||
false
|
||||
)]
|
||||
#[case::declined_optional_protocol(Some("websocket"), Some("Upgrade"), Some("events"), None, true)]
|
||||
#[case::matching_protocol(
|
||||
Some("websocket"),
|
||||
Some("Upgrade"),
|
||||
Some("events"),
|
||||
Some("events"),
|
||||
true
|
||||
)]
|
||||
#[case::multiple_offers(
|
||||
Some("websocket"),
|
||||
Some("Upgrade"),
|
||||
Some("events, realtime"),
|
||||
Some("realtime"),
|
||||
true
|
||||
)]
|
||||
#[case::case_insensitive_upgrade(Some("WebSocket"), Some("uPgRaDe"), None, None, true)]
|
||||
#[case::connection_tokens(Some("websocket"), Some("keep-alive, Upgrade"), None, None, true)]
|
||||
#[tokio::test]
|
||||
async fn upgrade_validates_headers_and_negotiated_protocol(
|
||||
#[case] upgrade: Option<&'static str>,
|
||||
#[case] connection: Option<&'static str>,
|
||||
#[case] offered: Option<&'static str>,
|
||||
#[case] selected: Option<&'static str>,
|
||||
#[case] accepted: bool,
|
||||
) {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
tokio::spawn(async move {
|
||||
let (stream, _) = listener.accept().await.unwrap();
|
||||
#[expect(
|
||||
clippy::result_large_err,
|
||||
reason = "tungstenite requires an unboxed handshake error response"
|
||||
)]
|
||||
let mut socket = accept_hdr_async(stream, move |request: &Request, response: Response| {
|
||||
assert_eq!(request.uri().path_and_query().unwrap().as_str(), path);
|
||||
assert_eq!(request.headers()["authorization"], "Bearer test-key");
|
||||
let result = accept_hdr_async(stream, move |_: &Request, mut response: Response| {
|
||||
response.headers_mut().remove("upgrade");
|
||||
response.headers_mut().remove("connection");
|
||||
if let Some(value) = upgrade {
|
||||
response
|
||||
.headers_mut()
|
||||
.insert("upgrade", HeaderValue::from_static(value));
|
||||
}
|
||||
if let Some(value) = connection {
|
||||
response
|
||||
.headers_mut()
|
||||
.insert("connection", HeaderValue::from_static(value));
|
||||
}
|
||||
if let Some(value) = selected {
|
||||
response
|
||||
.headers_mut()
|
||||
.insert("sec-websocket-protocol", HeaderValue::from_static(value));
|
||||
}
|
||||
Ok(response)
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
socket.next().await.unwrap().unwrap(),
|
||||
Message::Text("request".into())
|
||||
);
|
||||
socket.send(Message::Text("response".into())).await.unwrap();
|
||||
socket.close(None).await.unwrap();
|
||||
.await;
|
||||
let Ok(mut socket) = result else {
|
||||
return;
|
||||
};
|
||||
if let Some(Ok(message)) = socket.next().await {
|
||||
let _ = socket.send(message).await;
|
||||
}
|
||||
});
|
||||
let mut request = format!("ws://{address}{path}")
|
||||
.into_client_request()
|
||||
.unwrap();
|
||||
request
|
||||
.headers_mut()
|
||||
.insert("authorization", "Bearer test-key".parse().unwrap());
|
||||
let (mut socket, response) = connect_upstream(request).await.unwrap();
|
||||
assert_eq!(response.status(), 101);
|
||||
socket.send(Message::Text("request".into())).await.unwrap();
|
||||
assert_eq!(
|
||||
socket.next().await.unwrap().unwrap(),
|
||||
Message::Text("response".into())
|
||||
);
|
||||
assert!(matches!(
|
||||
socket.next().await.unwrap().unwrap(),
|
||||
Message::Close(_)
|
||||
));
|
||||
server.await.unwrap();
|
||||
let headers = offered.map_or_else(HeaderMap::new, |value| {
|
||||
HeaderMap::from_iter([(
|
||||
http::header::SEC_WEBSOCKET_PROTOCOL,
|
||||
HeaderValue::from_static(value),
|
||||
)])
|
||||
});
|
||||
let result = open(
|
||||
&config(Verify::BuiltInRoots),
|
||||
&format!("ws://{address}/responses"),
|
||||
headers,
|
||||
)
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(mut socket) if accepted => round_trip(&mut socket).await,
|
||||
Err(Error::Network(_)) if !accepted => {}
|
||||
other => panic!("unexpected handshake outcome: {:?}", other.map(drop)),
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -283,7 +283,7 @@ mod tests {
|
|||
|
||||
let (provider, config, request) = prepared(&upstream.uri());
|
||||
execute(
|
||||
&Client::plain_for_test(),
|
||||
&Client::for_test(),
|
||||
&AuthServices::default(),
|
||||
config,
|
||||
provider,
|
||||
|
|
@ -328,7 +328,7 @@ mod tests {
|
|||
|
||||
let (provider, config, request) = prepared(&upstream.uri());
|
||||
let error = execute(
|
||||
&Client::plain_for_test(),
|
||||
&Client::for_test(),
|
||||
&AuthServices::default(),
|
||||
config,
|
||||
provider,
|
||||
|
|
|
|||
|
|
@ -239,7 +239,7 @@ async fn chat_cache_identity_follows_resolved_configuration_and_request_callback
|
|||
.store(usize::from(call >= 2), Ordering::SeqCst);
|
||||
let cache = ScopedCache::new(cache.clone(), CacheScope::Shared);
|
||||
ChatCompletionsRoute::new(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
Arc::new(Default::default()),
|
||||
secrets.clone(),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -888,7 +888,7 @@ mod tests {
|
|||
.respond_with(ResponseTemplate::new(200).set_body_raw(body, "text/event-stream"))
|
||||
.mount(&upstream)
|
||||
.await;
|
||||
let response = litellm_http::Client::plain_for_test()
|
||||
let response = litellm_http::Client::for_test()
|
||||
.get(upstream.uri())
|
||||
.send()
|
||||
.await
|
||||
|
|
|
|||
|
|
@ -121,11 +121,7 @@ mod tests {
|
|||
};
|
||||
|
||||
fn client() -> OcrClient {
|
||||
OcrClient::for_test(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::no_redirect_for_test(),
|
||||
)
|
||||
.with_secrets(litellm_inference_testing::no_secrets())
|
||||
OcrClient::for_test().with_secrets(litellm_inference_testing::no_secrets())
|
||||
}
|
||||
|
||||
fn request(model: &str, base: &str, document: Value, options: Value) -> LiteLLMOcrRequest {
|
||||
|
|
|
|||
|
|
@ -25,6 +25,7 @@ tokio-tungstenite.workspace = true
|
|||
tracing.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-http = { workspace = true, features = ["test-support"] }
|
||||
litellm-cache.workspace = true
|
||||
litellm-cache-memory.workspace = true
|
||||
litellm-host-native.workspace = true
|
||||
|
|
|
|||
|
|
@ -1,12 +1,14 @@
|
|||
use std::{collections::HashMap, sync::Arc, time::Duration};
|
||||
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use litellm_http::websocket::{UpstreamWebSocket, connect_upstream};
|
||||
use litellm_http::{
|
||||
Client,
|
||||
websocket::{self, UpstreamWebSocket},
|
||||
};
|
||||
use tokio::sync::Mutex;
|
||||
use tokio_tungstenite::tungstenite::{
|
||||
Message,
|
||||
client::IntoClientRequest,
|
||||
http::{HeaderName, HeaderValue},
|
||||
http::{HeaderMap, HeaderName, HeaderValue},
|
||||
};
|
||||
|
||||
use super::Error;
|
||||
|
|
@ -24,42 +26,33 @@ impl ResponsesWebSocketConnection {
|
|||
fields(outcome)
|
||||
)]
|
||||
pub async fn connect_url(
|
||||
client: &Client,
|
||||
url: &str,
|
||||
headers: &HashMap<String, String>,
|
||||
timeout: Option<Duration>,
|
||||
) -> Result<Self, Error> {
|
||||
litellm_inference::diagnostic::operation("litellm.websocket.connect_url", async {
|
||||
let mut request = url.into_client_request().map_err(|error| {
|
||||
Error::Transport(litellm_http::transport::Error::Network(error.to_string()))
|
||||
})?;
|
||||
for (name, value) in headers {
|
||||
let header_name = name
|
||||
.parse::<HeaderName>()
|
||||
.map_err(|error| Error::InvalidRequest(error.to_string().into()))?;
|
||||
let header_value = HeaderValue::from_str(value)
|
||||
.map_err(|error| Error::InvalidRequest(error.to_string().into()))?;
|
||||
request.headers_mut().insert(header_name, header_value);
|
||||
}
|
||||
let connect = connect_upstream(request);
|
||||
let result = match timeout {
|
||||
let headers = headers
|
||||
.iter()
|
||||
.map(|(name, value)| {
|
||||
Ok((
|
||||
name.parse::<HeaderName>()
|
||||
.map_err(|error| Error::InvalidRequest(error.to_string().into()))?,
|
||||
HeaderValue::from_str(value)
|
||||
.map_err(|error| Error::InvalidRequest(error.to_string().into()))?,
|
||||
))
|
||||
})
|
||||
.collect::<Result<HeaderMap, Error>>()?;
|
||||
let connect = websocket::connect(client, url, headers);
|
||||
let socket = match timeout {
|
||||
Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| {
|
||||
Error::Transport(litellm_http::transport::Error::Network(
|
||||
"Responses WebSocket connection timed out".into(),
|
||||
))
|
||||
})?,
|
||||
None => connect.await,
|
||||
};
|
||||
let (socket, _) = result.map_err(|error| match *error {
|
||||
tokio_tungstenite::tungstenite::Error::Http(response) => {
|
||||
Error::Transport(litellm_http::transport::Error::Http {
|
||||
status: response.status().as_u16(),
|
||||
body: String::new(),
|
||||
})
|
||||
}
|
||||
other => {
|
||||
Error::Transport(litellm_http::transport::Error::Network(other.to_string()))
|
||||
}
|
||||
})?;
|
||||
}
|
||||
.map_err(Error::Transport)?;
|
||||
Ok(Self {
|
||||
socket: Arc::new(Mutex::new(Some(socket))),
|
||||
})
|
||||
|
|
|
|||
|
|
@ -538,6 +538,7 @@ async fn websocket_operations_trace_outcomes_without_capturing_frames_or_credent
|
|||
.logger()
|
||||
.instrument(async {
|
||||
let connection = ResponsesWebSocketConnection::connect_url(
|
||||
&litellm_http::Client::for_test(),
|
||||
&format!("ws://{address}/responses"),
|
||||
&std::collections::HashMap::from([(
|
||||
"authorization".into(),
|
||||
|
|
|
|||
|
|
@ -190,16 +190,14 @@ mod tests {
|
|||
clippy::disallowed_types,
|
||||
reason = "the pool has no default-header setting to stand in for provider credentials"
|
||||
)]
|
||||
let provider_http = litellm_http::Client::for_test(
|
||||
let provider_http = litellm_http::Client::wrap_for_test(
|
||||
reqwest::Client::builder()
|
||||
.default_headers(provider_headers)
|
||||
.build()
|
||||
.unwrap(),
|
||||
);
|
||||
let client = crate::base_llm::ocr::handler::OcrClient::for_test(
|
||||
provider_http,
|
||||
litellm_http::Client::no_redirect_for_test(),
|
||||
);
|
||||
let client =
|
||||
crate::base_llm::ocr::handler::OcrClient::for_test().with_provider_http(provider_http);
|
||||
let converted = inline_remote_document(
|
||||
client.document_fetcher(),
|
||||
OcrDocument::ImageUrl {
|
||||
|
|
|
|||
|
|
@ -86,21 +86,27 @@ impl OcrClient {
|
|||
}
|
||||
|
||||
#[cfg(any(test, feature = "test-support"))]
|
||||
pub fn for_test(provider_http: Client, no_redirect_http: Client) -> Self {
|
||||
pub fn for_test() -> Self {
|
||||
Self {
|
||||
secrets: Arc::new(
|
||||
litellm_secrets::source::EnvironmentSecrets::python_compatible(
|
||||
provider_http.clone(),
|
||||
),
|
||||
litellm_secrets::source::EnvironmentSecrets::python_compatible(Client::for_test()),
|
||||
),
|
||||
provider_http,
|
||||
polling_http: no_redirect_http.clone(),
|
||||
document_fetcher: MediaFetcher::for_test(no_redirect_http),
|
||||
provider_http: Client::for_test(),
|
||||
polling_http: Client::for_test(),
|
||||
document_fetcher: MediaFetcher::for_test(Client::for_test()),
|
||||
auth: Arc::new(AuthServices::default()),
|
||||
settings: OcrSettings::default(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(any(test, feature = "test-support"))]
|
||||
pub fn with_provider_http(self, provider_http: Client) -> Self {
|
||||
Self {
|
||||
provider_http,
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(any(test, feature = "test-support"))]
|
||||
pub fn with_settings(self, settings: OcrSettings) -> Self {
|
||||
Self { settings, ..self }
|
||||
|
|
@ -322,7 +328,7 @@ mod tests {
|
|||
let _connection = listener.accept().await.unwrap();
|
||||
tokio::time::sleep(Duration::from_secs(1)).await;
|
||||
});
|
||||
let error = litellm_http::Client::plain_for_test()
|
||||
let error = litellm_http::Client::for_test()
|
||||
.get(format!("http://{address}"))
|
||||
.timeout(Duration::from_millis(10))
|
||||
.send()
|
||||
|
|
@ -347,7 +353,7 @@ mod tests {
|
|||
socket.write_all(response.as_bytes()).await.unwrap();
|
||||
std::future::pending::<()>().await;
|
||||
});
|
||||
let response = litellm_http::Client::plain_for_test()
|
||||
let response = litellm_http::Client::for_test()
|
||||
.get(format!("http://{address}"))
|
||||
.send()
|
||||
.await
|
||||
|
|
|
|||
|
|
@ -566,10 +566,7 @@ mod tests {
|
|||
let params = ReductoParseV3Config
|
||||
.map_ocr_params(&overrides, "parse-v3")
|
||||
.unwrap();
|
||||
let client = OcrClient::for_test(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::no_redirect_for_test(),
|
||||
);
|
||||
let client = OcrClient::for_test();
|
||||
let connection = OcrConnection::default();
|
||||
let document = serde_json::from_value(
|
||||
json!({"type":"document_url","document_url":"reducto://ready.pdf"}),
|
||||
|
|
|
|||
|
|
@ -1,7 +1,9 @@
|
|||
use litellm_inference_responses::websocket::ResponsesWebSocketConnection as RustResponsesWebSocketConnection;
|
||||
use pyo3::prelude::*;
|
||||
use serde_json::Value;
|
||||
|
||||
use litellm_http::ClientVariant;
|
||||
use litellm_inference_responses::websocket::ResponsesWebSocketConnection as RustResponsesWebSocketConnection;
|
||||
|
||||
use crate::{
|
||||
errors::route_error_to_pyerr,
|
||||
marshal::{marshal_headers, optional_timeout},
|
||||
|
|
@ -25,10 +27,12 @@ impl ResponsesWebSocketConnection {
|
|||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let headers = marshal_headers(headers)?;
|
||||
let timeout = optional_timeout(timeout_seconds);
|
||||
let client = crate::http::host_client(py, ClientVariant::WebSocket)?;
|
||||
crate::execution::run_async_value(py, async move {
|
||||
let inner = RustResponsesWebSocketConnection::connect_url(&url, &headers, timeout)
|
||||
.await
|
||||
.map_err(route_error_to_pyerr)?;
|
||||
let inner =
|
||||
RustResponsesWebSocketConnection::connect_url(&client, &url, &headers, timeout)
|
||||
.await
|
||||
.map_err(route_error_to_pyerr)?;
|
||||
Ok(ResponsesWebSocketConnection { inner })
|
||||
})
|
||||
}
|
||||
|
|
|
|||
|
|
@ -192,7 +192,7 @@ handler.get_secret_from_manager = get_secret_from_manager
|
|||
KeyManagementSettings::default(),
|
||||
)),
|
||||
Arc::new(move |_: &str| fallback.map(str::to_owned)),
|
||||
OidcResolver::new(litellm_http::Client::plain_for_test()),
|
||||
OidcResolver::new(litellm_http::Client::for_test()),
|
||||
)
|
||||
.with_failure_policy(FailurePolicy::EnvironmentFallback);
|
||||
(resolver, locals, handler)
|
||||
|
|
|
|||
|
|
@ -84,7 +84,7 @@ mod tests {
|
|||
}
|
||||
|
||||
async fn resolve(state: Arc<SecretManagerState>, name: &'static str) -> Option<String> {
|
||||
ResolvedSecrets::from_state(state, litellm_http::Client::plain_for_test())
|
||||
ResolvedSecrets::from_state(state, litellm_http::Client::for_test())
|
||||
.resolve(&[name])
|
||||
.await
|
||||
.unwrap()
|
||||
|
|
@ -182,7 +182,7 @@ mod tests {
|
|||
.await;
|
||||
let source = ResolvedSecrets::from_state(
|
||||
state(&server, KeyManagementSettings::default()),
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
);
|
||||
let snapshot = source.resolve(&[declared]).await.unwrap();
|
||||
assert_eq!(snapshot.get(undeclared), None);
|
||||
|
|
@ -248,7 +248,7 @@ mod tests {
|
|||
async fn oidc_failures_are_not_converted_to_missing_secrets() {
|
||||
let result = ResolvedSecrets::from_state(
|
||||
Arc::new(SecretManagerState::default()),
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
)
|
||||
.resolve(&["oidc/"])
|
||||
.await;
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ use wiremock::{
|
|||
|
||||
fn manager(server: &MockServer) -> AzureKeyVault {
|
||||
AzureKeyVault::with_client(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
server.uri().parse().unwrap(),
|
||||
Arc::new(|name: &str| (name == "AZURE_AD_TOKEN").then(|| "fake".to_owned())),
|
||||
)
|
||||
|
|
@ -131,7 +131,7 @@ fn new_validates_vault_environment(
|
|||
#[case] missing_environment: bool,
|
||||
) {
|
||||
let result = AzureKeyVault::new(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
Arc::new(move |name: &str| {
|
||||
(name == "AZURE_KEY_VAULT_URI")
|
||||
.then(|| uri.map(str::to_owned))
|
||||
|
|
@ -158,7 +158,7 @@ fn new_validates_vault_environment(
|
|||
#[case::local("http://localhost:8080", "https://localhost/.default")]
|
||||
fn derives_scope_from_vault_host(#[case] uri: &str, #[case] expected: &str) {
|
||||
let manager = AzureKeyVault::with_client(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
uri.parse().unwrap(),
|
||||
Arc::new(|_: &str| None),
|
||||
)
|
||||
|
|
@ -187,7 +187,7 @@ async fn missing_credentials_do_not_request_vault() {
|
|||
|
||||
fn manager_without_credentials(server: &MockServer) -> AzureKeyVault {
|
||||
AzureKeyVault::with_client(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
server.uri().parse().unwrap(),
|
||||
Arc::new(|name: &str| {
|
||||
(name == "AZURE_CREDENTIAL").then(|| "ClientSecretCredential".to_owned())
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ use rstest::rstest;
|
|||
#[ignore]
|
||||
async fn reads_a_real_secret() {
|
||||
let environment = Arc::new(ProcessEnvironment);
|
||||
let manager = AzureKeyVault::new(litellm_http::Client::plain_for_test(), environment).unwrap();
|
||||
let manager = AzureKeyVault::new(litellm_http::Client::for_test(), environment).unwrap();
|
||||
let name = std::env::var("AZURE_KEY_VAULT_LIVE_SECRET_NAME").unwrap();
|
||||
let secret = manager.get_secret(&name).await.unwrap().unwrap();
|
||||
assert!(matches!(&secret, Secret::String(_)));
|
||||
|
|
|
|||
|
|
@ -77,7 +77,7 @@ async fn authentication_encodes_login(#[case] username: &str, #[case] expected_p
|
|||
.mount(&server)
|
||||
.await;
|
||||
let manager = CyberArkSecretManager::with_client(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
server.uri().parse().unwrap(),
|
||||
"acct".into(),
|
||||
username.into(),
|
||||
|
|
@ -433,7 +433,7 @@ async fn trailing_slash_endpoint_preserves_base_path() {
|
|||
.await;
|
||||
let endpoint = format!("{}/prefix/", server.uri()).parse().unwrap();
|
||||
let manager = CyberArkSecretManager::with_client(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
endpoint,
|
||||
"acct".into(),
|
||||
"admin".into(),
|
||||
|
|
|
|||
|
|
@ -62,7 +62,7 @@ pub(super) fn from_environment(
|
|||
|
||||
pub(super) fn manager(server: &MockServer, ttl: Duration) -> CyberArkSecretManager {
|
||||
CyberArkSecretManager::with_client(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
server.uri().parse().unwrap(),
|
||||
"acct".into(),
|
||||
"admin".into(),
|
||||
|
|
|
|||
|
|
@ -233,7 +233,7 @@ async fn writes_match_python_parity_fixture(parity_fixture: ParityFixture) {
|
|||
.mount(&server)
|
||||
.await;
|
||||
let manager = CyberArkSecretManager::with_client(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
server.uri().parse().unwrap(),
|
||||
parity_fixture.account,
|
||||
parity_fixture.username,
|
||||
|
|
@ -297,7 +297,7 @@ async fn live_conjur_round_trip() {
|
|||
.as_nanos()
|
||||
);
|
||||
let manager = CyberArkSecretManager::with_client(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
endpoint.clone(),
|
||||
account.clone(),
|
||||
username.clone(),
|
||||
|
|
@ -312,7 +312,7 @@ async fn live_conjur_round_trip() {
|
|||
.await
|
||||
.unwrap();
|
||||
let verifier = CyberArkSecretManager::with_client(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
endpoint.clone(),
|
||||
account.clone(),
|
||||
username.clone(),
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ use wiremock::{
|
|||
|
||||
fn manager(server: &MockServer, always_read: bool, ttl: Duration) -> GoogleSecretManager {
|
||||
GoogleSecretManager::with_client(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
server.uri().parse().unwrap(),
|
||||
"project".into(),
|
||||
Arc::new(|name: &str| (name == "VERTEX_AI_API_KEY").then(|| "token".into())),
|
||||
|
|
@ -215,7 +215,7 @@ async fn always_read_and_expired_cache_fetch_again(
|
|||
fn google_manager_requires_host_license_and_project_configuration() {
|
||||
assert!(matches!(
|
||||
GoogleSecretManager::new(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
Arc::new(|_: &str| None),
|
||||
false
|
||||
),
|
||||
|
|
@ -223,7 +223,7 @@ fn google_manager_requires_host_license_and_project_configuration() {
|
|||
));
|
||||
assert!(matches!(
|
||||
GoogleSecretManager::new(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
Arc::new(|_: &str| None),
|
||||
true
|
||||
),
|
||||
|
|
@ -244,7 +244,7 @@ fn google_manager_rejects_invalid_refresh_intervals(#[case] variable: &'static s
|
|||
});
|
||||
|
||||
assert!(matches!(
|
||||
GoogleSecretManager::new(litellm_http::Client::plain_for_test(), environment, true),
|
||||
GoogleSecretManager::new(litellm_http::Client::for_test(), environment, true),
|
||||
Err(Error::RefreshInterval)
|
||||
));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -53,7 +53,7 @@ async fn read_results_follow_the_selected_failure_policy(
|
|||
},
|
||||
)),
|
||||
Arc::new(move |_: &str| environment.map(str::to_owned)),
|
||||
OidcResolver::new(litellm_http::Client::plain_for_test()),
|
||||
OidcResolver::new(litellm_http::Client::for_test()),
|
||||
)
|
||||
.with_failure_policy(policy);
|
||||
let result = resolver
|
||||
|
|
@ -97,7 +97,7 @@ async fn primary_secret_values_other_than_strings_resolve_to_none(
|
|||
let resolver = SecretResolver::new_python_compatible(
|
||||
Arc::new(state(&server, settings)),
|
||||
Arc::new(|_: &str| Some("fallback".into())),
|
||||
OidcResolver::new(litellm_http::Client::plain_for_test()),
|
||||
OidcResolver::new(litellm_http::Client::for_test()),
|
||||
);
|
||||
let text = value.as_str();
|
||||
assert_eq!(
|
||||
|
|
@ -157,7 +157,7 @@ async fn gating_prediction_matches_actual_lookup(
|
|||
let resolver = SecretResolver::new_python_compatible(
|
||||
Arc::new(state),
|
||||
Arc::new(|_: &str| Some("environment".into())),
|
||||
OidcResolver::new(litellm_http::Client::plain_for_test()),
|
||||
OidcResolver::new(litellm_http::Client::for_test()),
|
||||
);
|
||||
assert_eq!(
|
||||
resolver
|
||||
|
|
|
|||
|
|
@ -22,7 +22,7 @@ async fn azure_handler_reads_missing_and_failed_secrets() {
|
|||
.await;
|
||||
let manager = SecretManager::AzureKeyVault(
|
||||
AzureKeyVault::with_client(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
server.uri().parse().unwrap(),
|
||||
std::sync::Arc::new(|name: &str| (name == "AZURE_AD_TOKEN").then(|| "fake".to_owned())),
|
||||
)
|
||||
|
|
@ -81,7 +81,7 @@ async fn successful_azure_responses_do_not_fall_back_when_the_value_is_empty_or_
|
|||
.mount(&server)
|
||||
.await;
|
||||
let manager = AzureKeyVault::with_client(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
server.uri().parse().unwrap(),
|
||||
Arc::new(|name: &str| (name == "AZURE_AD_TOKEN").then(|| "token".into())),
|
||||
)
|
||||
|
|
@ -92,7 +92,7 @@ async fn successful_azure_responses_do_not_fall_back_when_the_value_is_empty_or_
|
|||
Default::default(),
|
||||
)),
|
||||
Arc::new(|_: &str| Some("environment".into())),
|
||||
OidcResolver::new(litellm_http::Client::plain_for_test()),
|
||||
OidcResolver::new(litellm_http::Client::for_test()),
|
||||
);
|
||||
assert_eq!(
|
||||
resolver
|
||||
|
|
|
|||
|
|
@ -59,7 +59,7 @@ fn manager(provider: Provider, server: &MockServer) -> SecretManager {
|
|||
),
|
||||
Provider::Azure => SecretManager::AzureKeyVault(
|
||||
AzureKeyVault::with_client(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
server.uri().parse().unwrap(),
|
||||
environment,
|
||||
)
|
||||
|
|
@ -67,7 +67,7 @@ fn manager(provider: Provider, server: &MockServer) -> SecretManager {
|
|||
),
|
||||
Provider::Google => SecretManager::GoogleSecretManager(
|
||||
GoogleSecretManager::with_client(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
server.uri().parse().unwrap(),
|
||||
"project".into(),
|
||||
environment,
|
||||
|
|
@ -84,7 +84,7 @@ fn manager(provider: Provider, server: &MockServer) -> SecretManager {
|
|||
.unwrap(),
|
||||
),
|
||||
Provider::Cyberark => SecretManager::Cyberark(CyberArkSecretManager::with_client(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
server.uri().parse().unwrap(),
|
||||
"acct".into(),
|
||||
"admin".into(),
|
||||
|
|
@ -240,7 +240,7 @@ async fn python_read_failures_preserve_provider_fallback_rules(
|
|||
KeyManagementSettings::default(),
|
||||
)),
|
||||
Arc::new(move |_: &str| environment_value.map(str::to_owned)),
|
||||
OidcResolver::new(litellm_http::Client::plain_for_test()),
|
||||
OidcResolver::new(litellm_http::Client::for_test()),
|
||||
);
|
||||
let expected = if matches!(provider, Provider::Aws) {
|
||||
None
|
||||
|
|
|
|||
|
|
@ -24,7 +24,7 @@ async fn cyberark_handler_reads_values_and_surfaces_errors() {
|
|||
.mount(&server)
|
||||
.await;
|
||||
let manager = SecretManager::Cyberark(CyberArkSecretManager::with_client(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
server.uri().parse().unwrap(),
|
||||
"acct".into(),
|
||||
"admin".into(),
|
||||
|
|
|
|||
|
|
@ -25,7 +25,7 @@ async fn google_resolver_distinguishes_absence_from_failure(#[case] status: u16)
|
|||
_ => None,
|
||||
});
|
||||
let manager = GoogleSecretManager::with_client(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::for_test(),
|
||||
server.uri().parse().unwrap(),
|
||||
"project".into(),
|
||||
environment.clone(),
|
||||
|
|
@ -40,7 +40,7 @@ async fn google_resolver_distinguishes_absence_from_failure(#[case] status: u16)
|
|||
let resolver = SecretResolver::new_python_compatible(
|
||||
Arc::new(state),
|
||||
environment,
|
||||
OidcResolver::new(litellm_http::Client::plain_for_test()),
|
||||
OidcResolver::new(litellm_http::Client::for_test()),
|
||||
)
|
||||
.with_failure_policy(FailurePolicy::Propagate);
|
||||
let result = resolver.get_secret_str("KEY", None).await;
|
||||
|
|
|
|||
|
|
@ -56,7 +56,7 @@ async fn hashicorp_handler_resolves_found_missing_and_failed_values() {
|
|||
},
|
||||
)),
|
||||
Arc::new(|_: &str| None),
|
||||
litellm_secrets::OidcResolver::new(litellm_http::Client::plain_for_test()),
|
||||
litellm_secrets::OidcResolver::new(litellm_http::Client::for_test()),
|
||||
);
|
||||
assert_eq!(
|
||||
found_resolver
|
||||
|
|
@ -131,7 +131,7 @@ async fn hashicorp_handler_resolves_found_missing_and_failed_values() {
|
|||
let failed_resolver = SecretResolver::new_python_compatible(
|
||||
Arc::new(failed_state),
|
||||
Arc::new(|_: &str| None),
|
||||
litellm_secrets::OidcResolver::new(litellm_http::Client::plain_for_test()),
|
||||
litellm_secrets::OidcResolver::new(litellm_http::Client::for_test()),
|
||||
)
|
||||
.with_failure_policy(FailurePolicy::Propagate);
|
||||
assert!(matches!(
|
||||
|
|
|
|||
|
|
@ -30,7 +30,7 @@ async fn environment_sources_resolve_expected_value(
|
|||
("CIRCLE_OIDC_TOKEN_V2", "circle-v2"),
|
||||
]);
|
||||
assert_eq!(
|
||||
OidcResolver::new(litellm_http::Client::plain_for_test())
|
||||
OidcResolver::new(litellm_http::Client::for_test())
|
||||
.resolve(reference, env.as_ref())
|
||||
.await
|
||||
.unwrap()
|
||||
|
|
@ -43,7 +43,7 @@ async fn environment_sources_resolve_expected_value(
|
|||
#[tokio::test]
|
||||
async fn environment_sources_bypass_boolean_conversion_and_defaults() {
|
||||
let env = environment(&[("TOKEN", "true")]);
|
||||
let oidc = OidcResolver::new(litellm_http::Client::plain_for_test());
|
||||
let oidc = OidcResolver::new(litellm_http::Client::for_test());
|
||||
let resolver = SecretResolver::new(Arc::new(SecretManagerState::default()), env, oidc);
|
||||
assert_eq!(
|
||||
resolver
|
||||
|
|
@ -94,7 +94,7 @@ async fn github_requests_are_authenticated_cached_and_revalidate_environment() {
|
|||
),
|
||||
("ACTIONS_ID_TOKEN_REQUEST_TOKEN", "request-token"),
|
||||
]);
|
||||
let oidc = OidcResolver::new(litellm_http::Client::plain_for_test());
|
||||
let oidc = OidcResolver::new(litellm_http::Client::for_test());
|
||||
for _ in 0..2 {
|
||||
assert_eq!(
|
||||
oidc.resolve("oidc/github/https://service/oidc/path", env.as_ref())
|
||||
|
|
@ -131,7 +131,7 @@ async fn file_allowlist_resolves_symlinks_while_environment_paths_remain_explici
|
|||
("PATH_TOKEN", private.to_str().unwrap()),
|
||||
("AZURE_FEDERATED_TOKEN_FILE", token.to_str().unwrap()),
|
||||
]);
|
||||
let oidc = OidcResolver::new(litellm_http::Client::plain_for_test());
|
||||
let oidc = OidcResolver::new(litellm_http::Client::for_test());
|
||||
assert_eq!(
|
||||
oidc.resolve(&format!("oidc/file/{}", token.display()), env.as_ref())
|
||||
.await
|
||||
|
|
@ -213,7 +213,7 @@ async fn google_expiry_caps_cache_and_preserves_audience(
|
|||
.expect(calls)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let oidc = OidcResolver::new(litellm_http::Client::plain_for_test())
|
||||
let oidc = OidcResolver::new(litellm_http::Client::for_test())
|
||||
.with_google_identity_endpoint(server.uri().parse().unwrap())
|
||||
.with_clock(now);
|
||||
for _ in 0..2 {
|
||||
|
|
@ -235,7 +235,7 @@ async fn google_expiry_caps_cache_and_preserves_audience(
|
|||
#[tokio::test]
|
||||
async fn google_oidc_requires_its_build_feature() {
|
||||
assert!(matches!(
|
||||
OidcResolver::new(litellm_http::Client::plain_for_test())
|
||||
OidcResolver::new(litellm_http::Client::for_test())
|
||||
.resolve("oidc/google/audience", environment(&[]).as_ref())
|
||||
.await,
|
||||
Err(Error::UnsupportedOidc)
|
||||
|
|
@ -246,7 +246,7 @@ async fn google_oidc_requires_its_build_feature() {
|
|||
#[tokio::test]
|
||||
async fn azure_oidc_without_a_token_file_requires_its_build_feature() {
|
||||
assert!(matches!(
|
||||
OidcResolver::new(litellm_http::Client::plain_for_test())
|
||||
OidcResolver::new(litellm_http::Client::for_test())
|
||||
.resolve("oidc/azure/scope", environment(&[]).as_ref())
|
||||
.await,
|
||||
Err(Error::UnsupportedOidc)
|
||||
|
|
@ -262,7 +262,7 @@ async fn invalid_references_fail_before_environment_lookup(
|
|||
#[case] reference: &str,
|
||||
#[case] unsupported: bool,
|
||||
) {
|
||||
let error = OidcResolver::new(litellm_http::Client::plain_for_test())
|
||||
let error = OidcResolver::new(litellm_http::Client::for_test())
|
||||
.resolve(reference, &|_: &str| {
|
||||
panic!("invalid reference reached environment lookup")
|
||||
})
|
||||
|
|
@ -284,7 +284,7 @@ async fn unreadable_expiry_keeps_python_cache_fallback(#[case] token: &str) {
|
|||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let resolver = OidcResolver::new(litellm_http::Client::plain_for_test())
|
||||
let resolver = OidcResolver::new(litellm_http::Client::for_test())
|
||||
.with_google_identity_endpoint(server.uri().parse().unwrap());
|
||||
for _ in 0..2 {
|
||||
assert_eq!(
|
||||
|
|
@ -336,7 +336,7 @@ async fn azure_oidc_acquires_the_requested_scope_and_preserves_failures(#[case]
|
|||
})
|
||||
}
|
||||
}
|
||||
let oidc = OidcResolver::new(litellm_http::Client::plain_for_test())
|
||||
let oidc = OidcResolver::new(litellm_http::Client::for_test())
|
||||
.with_azure_token_provider(Arc::new(Provider(failed)));
|
||||
let resolver = SecretResolver::new_python_compatible(
|
||||
Arc::new(SecretManagerState::default()),
|
||||
|
|
@ -364,7 +364,7 @@ async fn azure_oidc_acquires_the_requested_scope_and_preserves_failures(#[case]
|
|||
#[tokio::test]
|
||||
async fn missing_oidc_environment_is_an_error(#[case] reference: &str) {
|
||||
assert!(matches!(
|
||||
OidcResolver::new(litellm_http::Client::plain_for_test())
|
||||
OidcResolver::new(litellm_http::Client::for_test())
|
||||
.resolve(reference, environment(&[]).as_ref())
|
||||
.await,
|
||||
Err(Error::MissingEnvironment)
|
||||
|
|
@ -383,7 +383,7 @@ async fn google_oidc_failures_are_not_cached_or_hidden_by_defaults() {
|
|||
let resolver = SecretResolver::new_python_compatible(
|
||||
Arc::new(SecretManagerState::default()),
|
||||
environment(&[]),
|
||||
OidcResolver::new(litellm_http::Client::plain_for_test())
|
||||
OidcResolver::new(litellm_http::Client::for_test())
|
||||
.with_google_identity_endpoint(server.uri().parse().unwrap()),
|
||||
);
|
||||
for _ in 0..2 {
|
||||
|
|
@ -424,7 +424,7 @@ async fn google_tokens_expire_at_the_python_cache_deadline(
|
|||
.expect(2)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let resolver = OidcResolver::new(litellm_http::Client::plain_for_test())
|
||||
let resolver = OidcResolver::new(litellm_http::Client::for_test())
|
||||
.with_google_identity_endpoint(server.uri().parse().unwrap())
|
||||
.with_clock(|| UNIX_EPOCH + Duration::from_secs(1000));
|
||||
assert_eq!(
|
||||
|
|
@ -477,7 +477,7 @@ async fn google_cache_uses_payload_expiry_without_requiring_a_jwt_header() {
|
|||
.expect(2)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let resolver = OidcResolver::new(litellm_http::Client::plain_for_test())
|
||||
let resolver = OidcResolver::new(litellm_http::Client::for_test())
|
||||
.with_google_identity_endpoint(server.uri().parse().unwrap())
|
||||
.with_clock(|| UNIX_EPOCH + Duration::from_secs(1000));
|
||||
for _ in 0..2 {
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ fn resolver(value: Option<&str>) -> SecretResolver {
|
|||
SecretResolver::new_python_compatible(
|
||||
Arc::new(SecretManagerState::default()),
|
||||
Arc::new(move |_: &str| value.clone()),
|
||||
OidcResolver::new(litellm_http::Client::plain_for_test()),
|
||||
OidcResolver::new(litellm_http::Client::for_test()),
|
||||
)
|
||||
}
|
||||
|
||||
|
|
@ -35,7 +35,7 @@ async fn native_reads_preserve_strings_and_report_conversion_errors(#[case] mana
|
|||
let resolver = SecretResolver::new(
|
||||
Arc::new(state),
|
||||
Arc::new(move |_: &str| Some(raw.to_owned())),
|
||||
OidcResolver::new(litellm_http::Client::plain_for_test()),
|
||||
OidcResolver::new(litellm_http::Client::for_test()),
|
||||
);
|
||||
assert_eq!(
|
||||
resolver
|
||||
|
|
@ -75,7 +75,7 @@ async fn native_defaults_apply_to_absence_but_never_hide_provider_failures() {
|
|||
KeyManagementSettings::default(),
|
||||
)),
|
||||
Arc::new(|_: &str| None),
|
||||
OidcResolver::new(litellm_http::Client::plain_for_test()),
|
||||
OidcResolver::new(litellm_http::Client::for_test()),
|
||||
);
|
||||
let result = resolver
|
||||
.get_secret_str("key", Some(SecretValue::new("default")))
|
||||
|
|
@ -189,7 +189,7 @@ fn managed(reply: Result<Option<Secret>, ()>, environment: Option<&'static str>)
|
|||
KeyManagementSettings::default(),
|
||||
)),
|
||||
Arc::new(move |_: &str| environment.map(str::to_owned)),
|
||||
OidcResolver::new(litellm_http::Client::plain_for_test()),
|
||||
OidcResolver::new(litellm_http::Client::for_test()),
|
||||
)
|
||||
.with_failure_policy(FailurePolicy::EnvironmentFallback)
|
||||
}
|
||||
|
|
@ -281,7 +281,7 @@ async fn prefix_is_removed_once_and_resolved_from_environment() {
|
|||
let resolver = SecretResolver::new_python_compatible(
|
||||
Arc::new(state),
|
||||
Arc::new(|name: &str| (name == "os.environ/KEY").then(|| "value".into())),
|
||||
OidcResolver::new(litellm_http::Client::plain_for_test()),
|
||||
OidcResolver::new(litellm_http::Client::for_test()),
|
||||
);
|
||||
assert_eq!(
|
||||
resolver
|
||||
|
|
@ -340,7 +340,7 @@ async fn excluded_hosted_keys_keep_the_python_manager_conversion_path(
|
|||
},
|
||||
)),
|
||||
Arc::new(move |_: &str| Some(raw.to_owned())),
|
||||
OidcResolver::new(litellm_http::Client::plain_for_test()),
|
||||
OidcResolver::new(litellm_http::Client::for_test()),
|
||||
);
|
||||
assert_eq!(
|
||||
resolver.get_secret("KEY", None).await.unwrap(),
|
||||
|
|
@ -373,7 +373,7 @@ async fn azure_callback_absence_preserves_none_but_errors_fall_back(
|
|||
KeyManagementSettings::default(),
|
||||
)),
|
||||
Arc::new(|_: &str| Some("environment".into())),
|
||||
OidcResolver::new(litellm_http::Client::plain_for_test()),
|
||||
OidcResolver::new(litellm_http::Client::for_test()),
|
||||
);
|
||||
assert_eq!(
|
||||
resolver
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ mod tests {
|
|||
#[case] expected: Option<&str>,
|
||||
) {
|
||||
unsafe { std::env::set_var(name, value) };
|
||||
let secret = EnvironmentSecrets::python_compatible(litellm_http::Client::plain_for_test())
|
||||
let secret = EnvironmentSecrets::python_compatible(litellm_http::Client::for_test())
|
||||
.resolve(&[name])
|
||||
.await
|
||||
.unwrap()
|
||||
|
|
@ -42,7 +42,7 @@ async fn dynamic_names_use_the_same_resolver_and_snapshots_never_do_fresh_lookup
|
|||
reads.fetch_add(1, Ordering::SeqCst);
|
||||
(name != "missing").then(|| name.to_owned())
|
||||
}),
|
||||
OidcResolver::new(litellm_http::Client::plain_for_test()),
|
||||
OidcResolver::new(litellm_http::Client::for_test()),
|
||||
);
|
||||
let snapshot = source.resolve(&["declared", "missing"]).await.unwrap();
|
||||
let name = format!("runtime-{}", "key");
|
||||
|
|
|
|||
|
|
@ -77,7 +77,7 @@ async fn database() -> Database {
|
|||
7,
|
||||
)
|
||||
.unwrap(),
|
||||
client: Client::no_redirect_for_test(),
|
||||
client: Client::for_test(),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@ async fn spend_insert_controls_target_and_retry_settings() {
|
|||
.mount(&server)
|
||||
.await;
|
||||
let config = Config::new("spend_test".into(), &format!("{}?query=DROP&async_insert=0&wait_for_async_insert=0&input_format_skip_unknown_fields=1&database=wrong&readonly=1", server.uri()), 7).unwrap();
|
||||
let client = Client::no_redirect_for_test();
|
||||
let client = Client::for_test();
|
||||
let row = BTreeMap::from([
|
||||
("request_id".into(), json!("request")),
|
||||
("start_time".into(), json!(1234)),
|
||||
|
|
@ -77,7 +77,7 @@ async fn spend_insert_preserves_storage_failure(#[case] status: u16) {
|
|||
.await;
|
||||
let config = Config::new("spend_test".into(), &server.uri(), 7).unwrap();
|
||||
let error = insert_rows(
|
||||
&Client::no_redirect_for_test(),
|
||||
&Client::for_test(),
|
||||
config.storage().writer(),
|
||||
"spend_test",
|
||||
vec![BTreeMap::from([("request_id".into(), json!("request"))])],
|
||||
|
|
@ -95,7 +95,7 @@ async fn empty_spend_insert_does_not_contact_storage() {
|
|||
let server = MockServer::start().await;
|
||||
let config = Config::new("spend_test".into(), &server.uri(), 7).unwrap();
|
||||
insert_rows(
|
||||
&Client::no_redirect_for_test(),
|
||||
&Client::for_test(),
|
||||
config.storage().writer(),
|
||||
"spend_test",
|
||||
vec![],
|
||||
|
|
@ -116,7 +116,7 @@ async fn failed_schema_setup_is_reported() {
|
|||
.await;
|
||||
let config = Config::new("spend_test".into(), &server.uri(), 7).unwrap();
|
||||
assert!(matches!(
|
||||
ensure_schema(&Client::no_redirect_for_test(), &config).await,
|
||||
ensure_schema(&Client::for_test(), &config).await,
|
||||
Err(Error::Storage(
|
||||
litellm_storage_clickhouse::Error::SchemaFailed(403)
|
||||
))
|
||||
|
|
|
|||
|
|
@ -48,7 +48,7 @@ async fn database() -> TestResult<ClickHouseDatabase> {
|
|||
Ok(ClickHouseDatabase {
|
||||
_container: container,
|
||||
url,
|
||||
client: Client::no_redirect_for_test(),
|
||||
client: Client::for_test(),
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -244,7 +244,7 @@ async fn failed_migration_is_not_recorded_and_retains_storage_error(
|
|||
|
||||
#[rstest]
|
||||
fn invalid_database_identifier_is_rejected() {
|
||||
let client = Client::no_redirect_for_test();
|
||||
let client = Client::for_test();
|
||||
let connection = Connection::writer("http://127.0.0.1:1").expect("valid URL");
|
||||
assert!(matches!(
|
||||
ClickHouseMigrate::new(
|
||||
|
|
@ -327,7 +327,7 @@ async fn duplicate_ledger_rows_are_tolerated(
|
|||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn revert_reports_forward_only_error() {
|
||||
let client = Client::no_redirect_for_test();
|
||||
let client = Client::for_test();
|
||||
let connection = Connection::writer("http://127.0.0.1:1").expect("valid URL");
|
||||
let migration = migration(
|
||||
1,
|
||||
|
|
@ -369,7 +369,7 @@ async fn applied_int64_versions_accept_numeric_and_quoted_json(
|
|||
)
|
||||
.mount(&mock_server)
|
||||
.await;
|
||||
let client = Client::no_redirect_for_test();
|
||||
let client = Client::for_test();
|
||||
let connection = Connection::writer(&mock_server.uri()).expect("valid URL");
|
||||
let migrator = migrator(vec![]);
|
||||
let mut adapter = ClickHouseMigrate::new(
|
||||
|
|
@ -406,7 +406,7 @@ async fn schema_requests_override_unsafe_connection_settings(
|
|||
mock_server.uri()
|
||||
);
|
||||
execute_statement(
|
||||
&Client::no_redirect_for_test(),
|
||||
&Client::for_test(),
|
||||
&Connection::writer(&url).expect("valid URL"),
|
||||
"CREATE DATABASE IF NOT EXISTS trace_test",
|
||||
REQUEST_TIMEOUT,
|
||||
|
|
@ -445,7 +445,7 @@ async fn oversized_ledger_response_is_rejected_before_migrations(
|
|||
)
|
||||
.mount(&mock_server)
|
||||
.await;
|
||||
let client = Client::no_redirect_for_test();
|
||||
let client = Client::for_test();
|
||||
let connection = Connection::writer(&mock_server.uri()).expect("valid URL");
|
||||
let migrator = migrator(vec![]);
|
||||
let mut adapter = ClickHouseMigrate::new(
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ async fn insert_rejects_invalid_identifiers(
|
|||
#[case] table: &str,
|
||||
#[case] invalid_database: bool,
|
||||
) {
|
||||
let client = Client::no_redirect_for_test();
|
||||
let client = Client::for_test();
|
||||
let connection = Connection::writer("http://localhost:8123").expect("valid URL");
|
||||
let result = insert_encoded_rows(&client, &connection, database, table, "token", "{}").await;
|
||||
|
||||
|
|
@ -24,7 +24,7 @@ async fn insert_rejects_invalid_identifiers(
|
|||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn read_rejects_empty_sql() {
|
||||
let client = Client::no_redirect_for_test();
|
||||
let client = Client::for_test();
|
||||
let connection = Connection::reader("http://localhost:8123", "litellm").expect("valid URL");
|
||||
|
||||
assert!(matches!(
|
||||
|
|
@ -87,7 +87,7 @@ async fn typed_fetch_encodes_parameters_and_validates_rows(
|
|||
.expect(2)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let client = Client::no_redirect_for_test();
|
||||
let client = Client::for_test();
|
||||
let connection = Connection::parse(&server.uri()).unwrap();
|
||||
let params = QueryParams {
|
||||
signed: i64::MIN,
|
||||
|
|
@ -131,7 +131,7 @@ async fn server_result_limits_allow_smaller_pages_without_retrying_other_failure
|
|||
.await;
|
||||
let connection = Connection::parse(&server.uri()).unwrap();
|
||||
let error = execute_read(
|
||||
&Client::no_redirect_for_test(),
|
||||
&Client::for_test(),
|
||||
&connection,
|
||||
"SELECT 1",
|
||||
&BTreeMap::new(),
|
||||
|
|
@ -174,7 +174,7 @@ async fn insert_timeout_environment_child() {
|
|||
.mount(&server)
|
||||
.await;
|
||||
let result = insert_encoded_rows(
|
||||
&Client::no_redirect_for_test(),
|
||||
&Client::for_test(),
|
||||
&Connection::parse(&server.uri()).unwrap(),
|
||||
"traces",
|
||||
"otel_traces",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()}}:
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -1,5 +1,7 @@
|
|||
# What is this?
|
||||
## Handler file for calling claude-3 on vertex ai
|
||||
import base64
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
|
|
@ -19,6 +21,19 @@ if TYPE_CHECKING:
|
|||
from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
|
||||
|
||||
_NON_DEFAULT_PARAMS: Final = TypeAdapter(Mapping[str, object])
|
||||
_VERTEX_LABELS_HEADER: Final = "X-Vertex-AI-Labels"
|
||||
_LABELS: Final = TypeAdapter(Mapping[object, object])
|
||||
|
||||
|
||||
def _vertex_labels_header(labels: object, header_names: Mapping[str, object]) -> str | None:
|
||||
if not isinstance(labels, dict) or any(name.lower() == _VERTEX_LABELS_HEADER.lower() for name in header_names):
|
||||
return None
|
||||
string_labels: Final = {
|
||||
key: value
|
||||
for key, value in _LABELS.validate_python(labels).items()
|
||||
if isinstance(key, str) and isinstance(value, str)
|
||||
}
|
||||
return base64.b64encode(json.dumps(string_labels).encode()).decode() if string_labels else None
|
||||
|
||||
|
||||
class VertexAIError(Exception):
|
||||
|
|
@ -102,6 +117,12 @@ class VertexAIAnthropicConfig(AnthropicConfig):
|
|||
litellm_params: dict,
|
||||
headers: dict,
|
||||
) -> dict:
|
||||
params: Final = _NON_DEFAULT_PARAMS.validate_python(optional_params)
|
||||
labels_header: Final = _vertex_labels_header(params.get("labels"), _NON_DEFAULT_PARAMS.validate_python(headers))
|
||||
if "labels" in params:
|
||||
del optional_params["labels"] # rebind-ok: labels travel in the X-Vertex-AI-Labels header, not the body
|
||||
if labels_header is not None:
|
||||
headers[_VERTEX_LABELS_HEADER] = labels_header # rebind-ok: outbound headers are returned by mutation
|
||||
data: Final = super().transform_request(
|
||||
model=model,
|
||||
messages=messages,
|
||||
|
|
|
|||
|
|
@ -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/", "")
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -34,7 +34,7 @@ from .claude_settings import (
|
|||
load_json_or_empty,
|
||||
)
|
||||
from .cmd_quoting import quote_for_cmd
|
||||
from .config import claude_subscription_default
|
||||
from .config import subscription_default
|
||||
from .pi import (
|
||||
LITELLM_PROXY_API_KEY_ENV,
|
||||
PI_PROVIDER_NAME,
|
||||
|
|
@ -83,9 +83,11 @@ _INSTALL_DOCS: Final[dict[str, str]] = {
|
|||
}
|
||||
|
||||
_HIDDEN_AGENTS: Final = frozenset({"pi"})
|
||||
_SUBSCRIPTION_AGENTS: Final = frozenset({"claude"})
|
||||
_SUBSCRIPTION_LOGINS: Final = MappingProxyType({"claude": "claude.ai", "codex": "ChatGPT"})
|
||||
_SUBSCRIPTION_AGENTS: Final = frozenset(_SUBSCRIPTION_LOGINS)
|
||||
|
||||
CODEX_PROXY_PROVIDER: Final = "litellm"
|
||||
CODEX_SUBSCRIPTION_PROVIDER: Final = "lite_subscription"
|
||||
CODEX_HOME_ENV: Final = "CODEX_HOME"
|
||||
CODEX_MODEL_CATALOG_FILENAME: Final = "litellm-models.json"
|
||||
_CODEX_PREFLIGHT_TIMEOUT_SECONDS: Final = 10.0
|
||||
|
|
@ -143,6 +145,10 @@ def build_agent_env(
|
|||
travels in the x-litellm-api-key header through ANTHROPIC_CUSTOM_HEADERS, and no
|
||||
ANTHROPIC_AUTH_TOKEN or ANTHROPIC_API_KEY is left in the environment, since either one makes
|
||||
Claude Code stop using the subscription (and its connectors).
|
||||
With subscription=True the OpenAI profile keeps Codex on its ChatGPT login: the key goes to
|
||||
LITELLM_PROXY_API_KEY, which Codex sends in the x-litellm-api-key header, and OPENAI_API_KEY and
|
||||
OPENAI_BASE_URL are left exactly as the caller had them, so the key is never offered to Codex as an
|
||||
OpenAI credential and a child process never pairs the caller's own OpenAI key with the proxy URL.
|
||||
"""
|
||||
env: Final = dict(base_env)
|
||||
root: Final = base_url.rstrip("/")
|
||||
|
|
@ -161,21 +167,24 @@ def build_agent_env(
|
|||
if ENABLE_GATEWAY_MODEL_DISCOVERY_ENV not in env:
|
||||
env[ENABLE_GATEWAY_MODEL_DISCOVERY_ENV] = ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE
|
||||
if PROFILE_OPENAI in profiles:
|
||||
env[OPENAI_BASE_URL_ENV] = root + "/v1"
|
||||
env[OPENAI_API_KEY_ENV] = api_key
|
||||
if subscription:
|
||||
env[LITELLM_PROXY_API_KEY_ENV] = api_key
|
||||
else:
|
||||
env[OPENAI_BASE_URL_ENV] = root + "/v1"
|
||||
env[OPENAI_API_KEY_ENV] = api_key
|
||||
if PROFILE_LITELLM in profiles:
|
||||
env[LITELLM_PROXY_API_KEY_ENV] = api_key
|
||||
return env
|
||||
|
||||
|
||||
def codex_proxy_provider(base_url: str) -> Mapping[str, str | bool]:
|
||||
def codex_proxy_provider(base_url: str, *, subscription: bool = False) -> Mapping[str, str | bool]:
|
||||
return MappingProxyType(
|
||||
{
|
||||
"name": "LiteLLM proxy",
|
||||
"base_url": base_url.rstrip("/") + "/v1",
|
||||
"wire_api": "responses",
|
||||
"supports_websockets": False,
|
||||
"requires_openai_auth": False,
|
||||
"requires_openai_auth": subscription,
|
||||
}
|
||||
)
|
||||
|
||||
|
|
@ -304,7 +313,7 @@ def subscription_launch_args(
|
|||
return (_SETTINGS_FLAG, str(overlay)), rest
|
||||
|
||||
|
||||
def _codex_proxy_args(base_url: str) -> list[str]:
|
||||
def _codex_proxy_args(base_url: str, subscription: bool) -> list[str]:
|
||||
"""Codex `-c` overrides that point it at the proxy.
|
||||
|
||||
Codex ignores OPENAI_BASE_URL (it always dials api.openai.com), so the env
|
||||
|
|
@ -312,24 +321,35 @@ def _codex_proxy_args(base_url: str) -> list[str]:
|
|||
inline; supports_websockets=false forces the HTTP/SSE Responses transport
|
||||
because the proxy does not speak the Responses WebSocket protocol. The key is
|
||||
read from OPENAI_API_KEY, which build_agent_env already exports.
|
||||
With subscription=True the provider is `lite_subscription` with requires_openai_auth, so Codex
|
||||
signs in with its own ChatGPT login, and the key travels in the x-litellm-api-key header from
|
||||
LITELLM_PROXY_API_KEY. It sets no env_key, which would win over the login. A `-c` override cannot
|
||||
remove keys, so the id is one no config.toml defines: `lite configure codex` writes `litellm`, and
|
||||
the cookbook tells users to write `litellm_subscription`, whose env_key or headers would otherwise apply.
|
||||
"""
|
||||
provider: Final = f"model_providers.{CODEX_PROXY_PROVIDER}"
|
||||
provider_id: Final = CODEX_SUBSCRIPTION_PROVIDER if subscription else CODEX_PROXY_PROVIDER
|
||||
provider: Final = f"model_providers.{provider_id}"
|
||||
credential: Final = (
|
||||
f"{provider}.env_http_headers={{{json.dumps(LITELLM_KEY_HEADER)}={json.dumps(LITELLM_PROXY_API_KEY_ENV)}}}"
|
||||
if subscription
|
||||
else f'{provider}.env_key="{OPENAI_API_KEY_ENV}"'
|
||||
)
|
||||
return [
|
||||
"-c",
|
||||
f'model_provider="{CODEX_PROXY_PROVIDER}"',
|
||||
f'model_provider="{provider_id}"',
|
||||
*(
|
||||
argument
|
||||
for key, value in codex_proxy_provider(base_url).items()
|
||||
for key, value in codex_proxy_provider(base_url, subscription=subscription).items()
|
||||
for argument in ("-c", f"{provider}.{key}={json.dumps(value)}")
|
||||
),
|
||||
"-c",
|
||||
f'{provider}.env_key="{OPENAI_API_KEY_ENV}"',
|
||||
credential,
|
||||
"-c",
|
||||
f"{provider}.http_headers={{}}",
|
||||
]
|
||||
|
||||
|
||||
_PROXY_ARGS: Final[dict[str, Callable[[str], list[str]]]] = {
|
||||
_PROXY_ARGS: Final[dict[str, Callable[[str, bool], list[str]]]] = {
|
||||
"codex": _codex_proxy_args,
|
||||
}
|
||||
|
||||
|
|
@ -410,14 +430,14 @@ _Preparer: TypeAlias = Callable[[str, str, Mapping[str, str]], Sequence[str]]
|
|||
_PREPARERS: Final[Mapping[str, _Preparer]] = MappingProxyType({"pi": prepare_pi, "codex": prepare_codex})
|
||||
|
||||
|
||||
def agent_launch_args(command: str, base_url: str) -> list[str]:
|
||||
def agent_launch_args(command: str, base_url: str, *, subscription: bool = False) -> list[str]:
|
||||
"""Extra CLI args an agent needs to actually honor the proxy.
|
||||
|
||||
Claude Code and OpenCode respect the exported env vars, so they get nothing
|
||||
here; Codex needs its provider pointed via config overrides.
|
||||
"""
|
||||
builder: Final = _PROXY_ARGS.get(os.path.basename(command))
|
||||
return builder(base_url) if builder else []
|
||||
return builder(base_url, subscription) if builder else []
|
||||
|
||||
|
||||
class ListedModel(LiteLLMBaseModel):
|
||||
|
|
@ -889,11 +909,13 @@ def run_agent(
|
|||
(verify_subscription if subscription else verify)(base_url, api_key)
|
||||
|
||||
env_before_sync: Final = base_env if base_env is not None else os.environ
|
||||
synced: Final = sync_models(binary, env_before_sync, base_url, api_key, skip_verify)
|
||||
synced: Final = (
|
||||
_NO_EXTRA_ENV if subscription else sync_models(binary, env_before_sync, base_url, api_key, skip_verify)
|
||||
)
|
||||
if isinstance(synced, ModelSyncSkipped):
|
||||
warn(f"litellm: not syncing {display_name} models from the proxy: {synced.reason}")
|
||||
|
||||
prepare: Final = preparers.get(os.path.basename(command[0]))
|
||||
prepare: Final = None if subscription else preparers.get(os.path.basename(command[0]))
|
||||
prepared_args: Final = tuple(prepare(base_url, api_key, env_before_sync)) if prepare is not None else ()
|
||||
|
||||
env: Final = MappingProxyType(
|
||||
|
|
@ -903,7 +925,11 @@ def run_agent(
|
|||
}
|
||||
)
|
||||
synced_args: Final = synced.args if isinstance(synced, ModelSyncArgs) else ()
|
||||
extra_args: Final = (*agent_launch_args(command[0], base_url), *synced_args, *prepared_args)
|
||||
extra_args: Final = (
|
||||
*agent_launch_args(command[0], base_url, subscription=subscription),
|
||||
*synced_args,
|
||||
*prepared_args,
|
||||
)
|
||||
settings_args, forwarded = (
|
||||
subscription_launch_args(
|
||||
command,
|
||||
|
|
@ -914,7 +940,7 @@ def run_agent(
|
|||
),
|
||||
write_overlay,
|
||||
)
|
||||
if subscription and os.path.basename(command[0]) in _SUBSCRIPTION_AGENTS
|
||||
if subscription and os.path.basename(command[0]) == "claude"
|
||||
else ((), tuple(command[1:]))
|
||||
)
|
||||
if reattach_terminal is not None:
|
||||
|
|
@ -948,10 +974,19 @@ def resolve_api_key(ctx: click.Context) -> str:
|
|||
|
||||
|
||||
_SKIP_VERIFY_HELP: Final = "Skip the pre-launch key check against the proxy."
|
||||
_SUBSCRIPTION_HELP: Final = (
|
||||
"Keep using your claude.ai login (and its connectors). The LiteLLM key is sent in the "
|
||||
"x-litellm-api-key header instead of replacing Claude Code's own credential. "
|
||||
"Make it the default with `lite config set claude_subscription true`."
|
||||
_SUBSCRIPTION_HELP: Final = MappingProxyType(
|
||||
{
|
||||
"claude": (
|
||||
"Keep using your claude.ai login (and its connectors). The LiteLLM key is sent in the "
|
||||
"x-litellm-api-key header instead of replacing Claude Code's own credential. "
|
||||
"Make it the default with `lite config set claude_subscription true`."
|
||||
),
|
||||
"codex": (
|
||||
"Keep using your ChatGPT login. The LiteLLM key is sent in the "
|
||||
"x-litellm-api-key header instead of replacing Codex's own credential. "
|
||||
"Make it the default with `lite config set codex_subscription true`."
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -964,10 +999,8 @@ def _launch(
|
|||
api_key: Final = resolve_api_key(ctx)
|
||||
|
||||
display_name, _profiles = agent_profile(binary)
|
||||
use_subscription: Final = (
|
||||
subscription if subscription is not None else binary in _SUBSCRIPTION_AGENTS and claude_subscription_default()
|
||||
)
|
||||
login_note: Final = " (keeping your claude.ai login)" if use_subscription else ""
|
||||
use_subscription: Final = subscription if subscription is not None else subscription_default(binary)
|
||||
login_note: Final = f" (keeping your {_SUBSCRIPTION_LOGINS[binary]} login)" if use_subscription else ""
|
||||
click.echo(f"litellm: routing {display_name} through proxy at {base_url.rstrip('/')}{login_note}")
|
||||
|
||||
try:
|
||||
|
|
@ -989,7 +1022,7 @@ _Decorated: Final = TypeVar("_Decorated", bound=Callable[..., object] | click.Co
|
|||
def _subscription_option(binary: str) -> Callable[[_Decorated], _Decorated]:
|
||||
if binary not in _SUBSCRIPTION_AGENTS:
|
||||
return lambda func: func
|
||||
return click.option("--subscription/--no-subscription", default=None, help=_SUBSCRIPTION_HELP)
|
||||
return click.option("--subscription/--no-subscription", default=None, help=_SUBSCRIPTION_HELP[binary])
|
||||
|
||||
|
||||
def _make_agent_command(binary: str, display_name: str) -> click.Command:
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ from litellm.litellm_core_utils.private_json import ensure_private_dir, write_pr
|
|||
|
||||
HIDDEN_COMMANDS_KEY: Final = "hidden_commands"
|
||||
CLAUDE_SUBSCRIPTION_KEY: Final = "claude_subscription"
|
||||
CODEX_SUBSCRIPTION_KEY: Final = "codex_subscription"
|
||||
PKCE_KEY: Final = "pkce"
|
||||
|
||||
_config_adapter: Final[TypeAdapter[Mapping[str, str]]] = TypeAdapter(Mapping[str, str])
|
||||
|
|
@ -114,9 +115,15 @@ def _normalize_boolean(key: str) -> Callable[[str], str]:
|
|||
return normalize
|
||||
|
||||
|
||||
def claude_subscription_default() -> bool:
|
||||
"""Whether `lite claude` keeps the claude.ai login unless told otherwise."""
|
||||
return get_config_value(CLAUDE_SUBSCRIPTION_KEY) == "true"
|
||||
_SUBSCRIPTION_KEYS: Final[Mapping[str, str]] = MappingProxyType(
|
||||
{"claude": CLAUDE_SUBSCRIPTION_KEY, "codex": CODEX_SUBSCRIPTION_KEY}
|
||||
)
|
||||
|
||||
|
||||
def subscription_default(agent: str) -> bool:
|
||||
"""Whether `lite <agent>` keeps the agent's own subscription login unless told otherwise."""
|
||||
key: Final = _SUBSCRIPTION_KEYS.get(agent)
|
||||
return key is not None and get_config_value(key) == "true"
|
||||
|
||||
|
||||
def pkce_default() -> bool:
|
||||
|
|
@ -129,6 +136,7 @@ _NORMALIZERS: Final[Mapping[str, Callable[[str], str]]] = MappingProxyType(
|
|||
"base_url": normalize_base_url,
|
||||
HIDDEN_COMMANDS_KEY: _normalize_hidden_commands,
|
||||
CLAUDE_SUBSCRIPTION_KEY: _normalize_boolean(CLAUDE_SUBSCRIPTION_KEY),
|
||||
CODEX_SUBSCRIPTION_KEY: _normalize_boolean(CODEX_SUBSCRIPTION_KEY),
|
||||
PKCE_KEY: _normalize_boolean(PKCE_KEY),
|
||||
}
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ model/{model_id}/update - PATCH endpoint for model update.
|
|||
|
||||
import asyncio
|
||||
import datetime
|
||||
import itertools
|
||||
import json
|
||||
from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence
|
||||
from contextlib import AbstractAsyncContextManager, asynccontextmanager, suppress
|
||||
|
|
@ -42,6 +43,7 @@ from litellm.litellm_core_utils.ptu_pricing import (
|
|||
ptu_config_error,
|
||||
)
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload
|
||||
from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping
|
||||
from litellm.proxy._types import (
|
||||
BlockModelRequest,
|
||||
CommonProxyErrors,
|
||||
|
|
@ -102,6 +104,18 @@ from litellm.proxy.management_helpers.auto_router_permissions import (
|
|||
authorize_member_auto_router_write,
|
||||
)
|
||||
from litellm.proxy.management_helpers.model_allowlist_rename_sync import sync_model_allowlists_for_renamed_model
|
||||
from litellm.proxy.management_helpers.model_quota_policy import (
|
||||
ENFORCE_PROJECT_MODEL_QUOTA_SETTING,
|
||||
ModelQuotaPolicy,
|
||||
QuotaLimit,
|
||||
deployment_route,
|
||||
exempts_any_route,
|
||||
limits_added_by_team_write,
|
||||
parse_model_quota_setting,
|
||||
required_limits,
|
||||
team_write_serves,
|
||||
unmet_limits,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import (
|
||||
PTU_COST_ATTRIBUTION_ENV_VAR,
|
||||
is_ptu_cost_attribution_enabled,
|
||||
|
|
@ -111,6 +125,7 @@ from litellm.repositories.chunked_in import find_many_in
|
|||
from litellm.repositories.credentials_repository import CredentialsRepository
|
||||
from litellm.repositories.model_repository import ModelRepository
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.repositories.project_repository import ProjectRepository
|
||||
from litellm.repositories.table_repositories import ModelTableRepository
|
||||
from litellm.repositories.team_repository import TeamRepository
|
||||
from litellm.router import Router
|
||||
|
|
@ -152,6 +167,7 @@ from litellm.types.router import (
|
|||
SPECIAL_MODEL_INFO_PARAMS,
|
||||
Deployment,
|
||||
GenericLiteLLMParams,
|
||||
LiteLLM_Params,
|
||||
ModelInfo,
|
||||
updateDeployment,
|
||||
updateLiteLLMParams,
|
||||
|
|
@ -762,26 +778,93 @@ async def _auto_router_capability_slot(
|
|||
ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING: Final = "enforce_rpm_tpm_on_model_add"
|
||||
|
||||
|
||||
def _raise_if_rate_limits_required_but_missing(*, litellm_params: GenericLiteLLMParams, enforced: bool) -> None:
|
||||
"""Require both rpm and tpm (each a positive value) when the operator opts in via config.yaml.
|
||||
_ROUTING_KEYS: Final = ("model", "custom_llm_provider", "api_base")
|
||||
_RATE_LIMIT_KEYS: Final = ("rpm", "tpm")
|
||||
_NOTHING_SENT: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
Off by default, so deployments keep adding models without limits. When
|
||||
``enforce_rpm_tpm_on_model_add: true`` is set under general_settings, a model added
|
||||
without both rpm and tpm set to a positive value is rejected rather than stored
|
||||
unbounded (or effectively excluded from routing by a zero/negative limit).
|
||||
"""
|
||||
if not enforced:
|
||||
|
||||
def _merged_litellm_params(
|
||||
incoming: GenericLiteLLMParams | None, existing: GenericLiteLLMParams | None, keys: tuple[str, ...]
|
||||
) -> Mapping[str, object]:
|
||||
sent: Final = _NOTHING_SENT if incoming is None else _dump_without_none(incoming)
|
||||
stored: Final = _NOTHING_SENT if existing is None else _decrypted_litellm_params(existing)
|
||||
return MappingProxyType({key: sent.get(key, stored.get(key)) for key in keys})
|
||||
|
||||
|
||||
def _written_routing_params(
|
||||
incoming: GenericLiteLLMParams | None, existing: GenericLiteLLMParams | None
|
||||
) -> Mapping[str, object]:
|
||||
return _merged_litellm_params(incoming, existing, _ROUTING_KEYS)
|
||||
|
||||
|
||||
def _quota_policy(general_settings: Mapping[str, object], setting: str) -> ModelQuotaPolicy | None:
|
||||
try:
|
||||
return parse_model_quota_setting(general_settings.get(setting))
|
||||
except ValidationError as e:
|
||||
raise ProxyException(
|
||||
message=f"general_settings.{setting} is invalid: {e}",
|
||||
type=ProxyErrorTypes.internal_server_error.value,
|
||||
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
param=None,
|
||||
) from e
|
||||
|
||||
|
||||
def _raise_if_rate_limits_required_but_missing(
|
||||
*, model_name: str, litellm_params: LiteLLM_Params, team_model: bool, policy: ModelQuotaPolicy | None
|
||||
) -> None:
|
||||
"""Off by default, so deployments keep adding models without limits. Once the operator opts in under
|
||||
general_settings, the policy decides per model name and provider which of rpm and tpm a new model needs,
|
||||
and a model missing one, or carrying a zero or negative one, is rejected rather than stored unbounded
|
||||
(or effectively excluded from routing)."""
|
||||
if policy is None:
|
||||
return
|
||||
missing: Final = tuple(
|
||||
field
|
||||
for field, value in (("rpm", litellm_params.rpm), ("tpm", litellm_params.tpm))
|
||||
if value is None or value <= 0
|
||||
_raise_on_missing_rate_limits(
|
||||
model_name,
|
||||
unmet_limits(
|
||||
required_limits(
|
||||
policy, model_name, deployment_route(_written_routing_params(litellm_params, None), team_model)
|
||||
),
|
||||
MappingProxyType({"rpm": litellm_params.rpm, "tpm": litellm_params.tpm}),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _raise_if_model_update_adds_missing_rate_limits(
|
||||
*,
|
||||
existing: Deployment,
|
||||
written_name: str,
|
||||
written_team_model: bool,
|
||||
incoming_params: GenericLiteLLMParams | None,
|
||||
remaining_params: GenericLiteLLMParams,
|
||||
policy: ModelQuotaPolicy | None,
|
||||
) -> None:
|
||||
if policy is None:
|
||||
return
|
||||
before: Final = required_limits(
|
||||
policy,
|
||||
existing.model_info.team_public_model_name or existing.model_name,
|
||||
deployment_route(
|
||||
_written_routing_params(None, existing.litellm_params), existing.model_info.team_id is not None
|
||||
),
|
||||
)
|
||||
after: Final = required_limits(
|
||||
policy,
|
||||
written_name,
|
||||
deployment_route(_written_routing_params(incoming_params, remaining_params), written_team_model),
|
||||
)
|
||||
written_limits: Final = _merged_litellm_params(incoming_params, remaining_params, _RATE_LIMIT_KEYS)
|
||||
checked: Final = after - before
|
||||
_raise_on_missing_rate_limits(
|
||||
written_name, unmet_limits(checked, MappingProxyType({limit: written_limits.get(limit) for limit in checked}))
|
||||
)
|
||||
|
||||
|
||||
def _raise_on_missing_rate_limits(model_name: str, missing: tuple[QuotaLimit, ...]) -> None:
|
||||
if not missing:
|
||||
return
|
||||
raise ProxyException(
|
||||
message=(
|
||||
f"{' and '.join(missing)} must be set to a positive value when "
|
||||
f"{' and '.join(missing)} must be set to a positive value for '{model_name}' when "
|
||||
f"'{ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING}' is enabled in general_settings"
|
||||
),
|
||||
type=ProxyErrorTypes.validation_error.value,
|
||||
|
|
@ -790,6 +873,113 @@ def _raise_if_rate_limits_required_but_missing(*, litellm_params: GenericLiteLLM
|
|||
)
|
||||
|
||||
|
||||
def _project_limit(metadata: object, limits_key: str, model: str) -> object:
|
||||
limits: Final = metadata.get(limits_key) if is_str_keyed_mapping(metadata) else None
|
||||
return limits.get(model) if is_str_keyed_mapping(limits) else None
|
||||
|
||||
|
||||
def _projects_short_on(
|
||||
projects: Sequence["prisma_models.LiteLLM_ProjectTable"], model: str, added: frozenset[QuotaLimit]
|
||||
) -> tuple[str, ...]:
|
||||
return tuple(
|
||||
project.project_alias or project.project_id
|
||||
for project in projects
|
||||
if model in project.models
|
||||
and unmet_limits(
|
||||
added,
|
||||
MappingProxyType(
|
||||
{
|
||||
"rpm": _project_limit(project.metadata, "model_rpm_limit", model),
|
||||
"tpm": _project_limit(project.metadata, "model_tpm_limit", model),
|
||||
}
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def _raise_if_team_write_leaves_projects_short(
|
||||
*,
|
||||
prisma_client: PrismaClient,
|
||||
general_settings: Mapping[str, object],
|
||||
llm_router: Router | None,
|
||||
team_id: str | None,
|
||||
public_name: str,
|
||||
replaced_id: str | None,
|
||||
written_params: Mapping[str, object] | None,
|
||||
) -> None:
|
||||
if team_id is None:
|
||||
return
|
||||
policy: Final = _quota_policy(general_settings, ENFORCE_PROJECT_MODEL_QUOTA_SETTING)
|
||||
if policy is None or not exempts_any_route(policy):
|
||||
return
|
||||
projects: Final = await ProjectRepository(prisma_client).table.find_many(where={"team_id": team_id})
|
||||
listed: Final = frozenset(itertools.chain.from_iterable(project.models for project in projects))
|
||||
served: Final = sorted(model for model in listed if team_write_serves(llm_router, team_id, public_name, model))
|
||||
shortfall: Final = next(
|
||||
(
|
||||
(model, added, short)
|
||||
for model in served
|
||||
if (
|
||||
added := limits_added_by_team_write(
|
||||
policy, llm_router, team_id, public_name, replaced_id, written_params, model
|
||||
)
|
||||
)
|
||||
and (short := _projects_short_on(projects, model, added))
|
||||
),
|
||||
None,
|
||||
)
|
||||
if shortfall is None:
|
||||
return
|
||||
short_model, short_limits, short_projects = shortfall
|
||||
raise ProxyException(
|
||||
message=(
|
||||
f"projects {', '.join(short_projects)} list '{short_model}' without a positive "
|
||||
f"{' and '.join(sorted(short_limits))} "
|
||||
f"quota, which '{ENFORCE_PROJECT_MODEL_QUOTA_SETTING}' requires once this team model change applies. "
|
||||
"Set those limits on the projects first"
|
||||
),
|
||||
type=ProxyErrorTypes.validation_error.value,
|
||||
code=status.HTTP_400_BAD_REQUEST,
|
||||
param=None,
|
||||
)
|
||||
|
||||
|
||||
async def _raise_if_team_update_leaves_projects_short(
|
||||
*,
|
||||
prisma_client: PrismaClient,
|
||||
general_settings: Mapping[str, object],
|
||||
llm_router: Router | None,
|
||||
existing: Deployment,
|
||||
written_team_id: str | None,
|
||||
written_name: str,
|
||||
incoming_params: GenericLiteLLMParams | None,
|
||||
remaining_params: GenericLiteLLMParams,
|
||||
) -> None:
|
||||
if written_team_id is None and existing.model_info.team_id is None:
|
||||
return
|
||||
stored_name: Final = existing.model_info.team_public_model_name or existing.model_name
|
||||
await _raise_if_team_write_leaves_projects_short(
|
||||
prisma_client=prisma_client,
|
||||
general_settings=general_settings,
|
||||
llm_router=llm_router,
|
||||
team_id=written_team_id,
|
||||
public_name=written_name,
|
||||
replaced_id=existing.model_info.id,
|
||||
written_params=_written_routing_params(incoming_params, remaining_params),
|
||||
)
|
||||
if (written_team_id, written_name) == (existing.model_info.team_id, stored_name):
|
||||
return
|
||||
await _raise_if_team_write_leaves_projects_short(
|
||||
prisma_client=prisma_client,
|
||||
general_settings=general_settings,
|
||||
llm_router=llm_router,
|
||||
team_id=existing.model_info.team_id,
|
||||
public_name=stored_name,
|
||||
replaced_id=existing.model_info.id,
|
||||
written_params=None,
|
||||
)
|
||||
|
||||
|
||||
_PTU_PRICED_PAIR: Final = frozenset({"ptu_count", "cost_per_ptu_per_hour"})
|
||||
|
||||
|
||||
|
|
@ -1308,6 +1498,7 @@ async def patch_model(
|
|||
ProxyException: For various error conditions including authentication and database errors
|
||||
"""
|
||||
from litellm.proxy.proxy_server import (
|
||||
general_settings_view,
|
||||
litellm_proxy_admin_name,
|
||||
llm_router,
|
||||
premium_user,
|
||||
|
|
@ -1434,6 +1625,29 @@ async def patch_model(
|
|||
existing_params=remaining_params,
|
||||
)
|
||||
|
||||
written_team_id: Final = (
|
||||
patch_data.model_info.team_id if patch_data.model_info else None
|
||||
) or db_model.model_info.team_id
|
||||
written_name: Final = _get_public_model_name(effective_patch, db_model)
|
||||
_raise_if_model_update_adds_missing_rate_limits(
|
||||
existing=db_model,
|
||||
written_name=written_name,
|
||||
written_team_model=written_team_id is not None,
|
||||
incoming_params=incoming_params,
|
||||
remaining_params=remaining_params,
|
||||
policy=_quota_policy(general_settings_view(), ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING),
|
||||
)
|
||||
await _raise_if_team_update_leaves_projects_short(
|
||||
prisma_client=prisma_client,
|
||||
general_settings=general_settings_view(),
|
||||
llm_router=llm_router,
|
||||
existing=db_model,
|
||||
written_team_id=written_team_id,
|
||||
written_name=written_name,
|
||||
incoming_params=incoming_params,
|
||||
remaining_params=remaining_params,
|
||||
)
|
||||
|
||||
effective_params: Final = _effective_complexity_router_params(incoming_params, remaining_params)
|
||||
requested_model_name: Final = patch_data.model_name
|
||||
stored_model_name: str | None = None
|
||||
|
|
@ -2509,6 +2723,7 @@ async def delete_model(
|
|||
|
||||
from litellm.proxy.proxy_server import (
|
||||
MODEL_RECONCILE_LOCK,
|
||||
general_settings_view,
|
||||
llm_router,
|
||||
premium_user,
|
||||
prisma_client,
|
||||
|
|
@ -2542,6 +2757,15 @@ async def delete_model(
|
|||
incoming_params=None,
|
||||
allow_missing_team=True,
|
||||
)
|
||||
await _raise_if_team_write_leaves_projects_short(
|
||||
prisma_client=prisma_client,
|
||||
general_settings=general_settings_view(),
|
||||
llm_router=llm_router,
|
||||
team_id=model_params.model_info.team_id,
|
||||
public_name=model_params.model_info.team_public_model_name or model_params.model_name,
|
||||
replaced_id=model_params.model_info.id,
|
||||
written_params=None,
|
||||
)
|
||||
|
||||
# update DB
|
||||
if store_model_in_db is True:
|
||||
|
|
@ -2708,6 +2932,8 @@ async def add_new_model(
|
|||
"""
|
||||
from litellm.proxy.proxy_server import (
|
||||
general_settings,
|
||||
general_settings_view,
|
||||
llm_router,
|
||||
premium_user,
|
||||
prisma_client,
|
||||
proxy_config,
|
||||
|
|
@ -2766,8 +2992,19 @@ async def add_new_model(
|
|||
)
|
||||
|
||||
_raise_if_rate_limits_required_but_missing(
|
||||
model_name=model_params.model_name,
|
||||
litellm_params=model_params.litellm_params,
|
||||
enforced=bool(general_settings.get(ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING, False)),
|
||||
team_model=model_params.model_info.team_id is not None,
|
||||
policy=_quota_policy(general_settings_view(), ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING),
|
||||
)
|
||||
await _raise_if_team_write_leaves_projects_short(
|
||||
prisma_client=prisma_client,
|
||||
general_settings=general_settings_view(),
|
||||
llm_router=llm_router,
|
||||
team_id=model_params.model_info.team_id,
|
||||
public_name=model_params.model_name,
|
||||
replaced_id=None,
|
||||
written_params=_written_routing_params(model_params.litellm_params, None),
|
||||
)
|
||||
|
||||
clean_model_info: Final = ModelInfo.model_validate(
|
||||
|
|
@ -2906,6 +3143,7 @@ async def update_model(
|
|||
"""
|
||||
from litellm.proxy.proxy_server import (
|
||||
LITELLM_PROXY_ADMIN_NAME,
|
||||
general_settings_view,
|
||||
llm_router,
|
||||
premium_user,
|
||||
prisma_client,
|
||||
|
|
@ -2970,6 +3208,30 @@ async def update_model(
|
|||
incoming_params=model_params.litellm_params,
|
||||
existing_params=deployment.litellm_params,
|
||||
)
|
||||
renamed_to: Final = (
|
||||
model_params.model_name
|
||||
if model_params.model_name not in (None, deployment.model_name) and deployment.model_info.team_id is None
|
||||
else None
|
||||
)
|
||||
written_name: Final = renamed_to or deployment.model_info.team_public_model_name or deployment.model_name
|
||||
_raise_if_model_update_adds_missing_rate_limits(
|
||||
existing=deployment,
|
||||
written_name=written_name,
|
||||
written_team_model=deployment.model_info.team_id is not None,
|
||||
incoming_params=model_params.litellm_params,
|
||||
remaining_params=deployment.litellm_params,
|
||||
policy=_quota_policy(general_settings_view(), ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING),
|
||||
)
|
||||
await _raise_if_team_update_leaves_projects_short(
|
||||
prisma_client=prisma_client,
|
||||
general_settings=general_settings_view(),
|
||||
llm_router=llm_router,
|
||||
existing=deployment,
|
||||
written_team_id=deployment.model_info.team_id,
|
||||
written_name=written_name,
|
||||
incoming_params=model_params.litellm_params,
|
||||
remaining_params=deployment.litellm_params,
|
||||
)
|
||||
effective_params: Final = _effective_complexity_router_params(
|
||||
model_params.litellm_params, deployment.litellm_params
|
||||
)
|
||||
|
|
@ -3009,12 +3271,6 @@ async def update_model(
|
|||
if value is not None or _existing_litellm_params_dict.get(key) is not None
|
||||
}
|
||||
|
||||
renamed_to: Final = (
|
||||
model_params.model_name
|
||||
if model_params.model_name not in (None, deployment.model_name)
|
||||
and deployment.model_info.team_id is None
|
||||
else None
|
||||
)
|
||||
base_update: Final[PrismaCompatibleUpdateDBModel] = {
|
||||
"litellm_params": json.dumps(merged_dictionary),
|
||||
"updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME,
|
||||
|
|
|
|||
327
litellm/proxy/management_helpers/model_quota_policy.py
Normal file
327
litellm/proxy/management_helpers/model_quota_policy.py
Normal file
|
|
@ -0,0 +1,327 @@
|
|||
import re
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Literal, TypeAlias
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator
|
||||
|
||||
import litellm
|
||||
from litellm.exceptions import BadRequestError
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider, is_openai_compatible_endpoint
|
||||
from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping
|
||||
from litellm.llms.openai_like.json_loader import JSONProviderRegistry
|
||||
from litellm.router_utils.common_utils import provider_for_generic_call
|
||||
from litellm.router_utils.pattern_match_deployments import PatternMatchRouter, PatternUtils
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm import Router
|
||||
|
||||
QuotaLimit: TypeAlias = Literal["rpm", "tpm"]
|
||||
|
||||
ENFORCE_PROJECT_MODEL_QUOTA_SETTING: Final = "enforce_project_model_quota"
|
||||
|
||||
_LIMITS: Final[tuple[QuotaLimit, ...]] = ("rpm", "tpm")
|
||||
_ROUTER_PATTERNS: Final = PatternMatchRouter()
|
||||
|
||||
|
||||
def _is_known_provider(provider: str) -> bool:
|
||||
return provider in litellm.provider_list or JSONProviderRegistry.exists(provider)
|
||||
|
||||
|
||||
class ModelQuotaPolicy(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
default: frozenset[QuotaLimit] = frozenset(_LIMITS)
|
||||
providers: Mapping[str, frozenset[QuotaLimit]] = Field(default_factory=dict)
|
||||
models: Mapping[str, frozenset[QuotaLimit]] = Field(default_factory=dict)
|
||||
|
||||
@field_validator("providers")
|
||||
@classmethod
|
||||
def _known_providers(cls, providers: Mapping[str, frozenset[QuotaLimit]]) -> Mapping[str, frozenset[QuotaLimit]]:
|
||||
unknown: Final = sorted(provider for provider in providers if not _is_known_provider(provider))
|
||||
if unknown:
|
||||
raise ValueError(f"unknown provider ids {unknown}")
|
||||
return providers
|
||||
|
||||
|
||||
REQUIRE_BOTH_LIMITS: Final = ModelQuotaPolicy()
|
||||
_POLICY_KEYS: Final = frozenset(ModelQuotaPolicy.model_fields)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TeamDeployment:
|
||||
provider: str | None
|
||||
called_model: str | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ModelRoute:
|
||||
providers: frozenset[str] = frozenset()
|
||||
team_deployments: frozenset[TeamDeployment] = frozenset()
|
||||
|
||||
|
||||
RouteOf: TypeAlias = Callable[[str], ModelRoute]
|
||||
NO_ROUTE: Final = ModelRoute()
|
||||
|
||||
|
||||
def parse_model_quota_setting(setting: object) -> ModelQuotaPolicy | None:
|
||||
"""Any falsy value is off and any other value requires both limits, as the boolean setting always
|
||||
behaved, so only a mapping that names a policy key is read as a per-provider and per-model policy."""
|
||||
if not setting:
|
||||
return None
|
||||
if is_str_keyed_mapping(setting) and not _POLICY_KEYS.isdisjoint(setting):
|
||||
return ModelQuotaPolicy.model_validate(setting)
|
||||
return REQUIRE_BOTH_LIMITS
|
||||
|
||||
|
||||
def _wildcard_matches(pattern: str, model: str) -> bool:
|
||||
return re.match(_ROUTER_PATTERNS.pattern_to_regex(pattern), model) is not None
|
||||
|
||||
|
||||
def _router_specificity(pattern: str) -> tuple[int, int]:
|
||||
return PatternUtils.calculate_pattern_specificity(_ROUTER_PATTERNS.pattern_to_regex(pattern))
|
||||
|
||||
|
||||
def _most_specific_model_pattern(model: str, entries: Mapping[str, frozenset[QuotaLimit]]) -> str | None:
|
||||
matching: Final = tuple(pattern for pattern in entries if "*" in pattern and _wildcard_matches(pattern, model))
|
||||
return max(matching, key=_router_specificity, default=None)
|
||||
|
||||
|
||||
def _model_entry(entries: Mapping[str, frozenset[QuotaLimit]], model: str) -> frozenset[QuotaLimit] | None:
|
||||
exact: Final = entries.get(model)
|
||||
if exact is not None:
|
||||
return exact
|
||||
pattern: Final = _most_specific_model_pattern(model, entries)
|
||||
return None if pattern is None else entries[pattern]
|
||||
|
||||
|
||||
def _provider_limits(policy: ModelQuotaPolicy, provider: str | None) -> frozenset[QuotaLimit]:
|
||||
return policy.default if provider is None else policy.providers.get(provider, policy.default)
|
||||
|
||||
|
||||
def _team_deployment_limits(policy: ModelQuotaPolicy, deployment: TeamDeployment) -> frozenset[QuotaLimit]:
|
||||
entry: Final = None if deployment.called_model is None else _model_entry(policy.models, deployment.called_model)
|
||||
return _provider_limits(policy, deployment.provider) if entry is None else entry
|
||||
|
||||
|
||||
def required_limits(policy: ModelQuotaPolicy, model: str, route: ModelRoute) -> frozenset[QuotaLimit]:
|
||||
"""A model entry (exact name, then the most specific `*` pattern) wins over the providers of the
|
||||
model's deployments, which win over the default. A model served by several providers needs every
|
||||
limit any of them requires. A team picks its model's public name, so a team deployment matches
|
||||
model entries by the model it calls instead."""
|
||||
entry: Final = None if route.team_deployments else _model_entry(policy.models, model)
|
||||
if entry is not None:
|
||||
return entry
|
||||
if not route.providers and not route.team_deployments:
|
||||
return policy.default
|
||||
return frozenset[QuotaLimit]().union(
|
||||
*(_team_deployment_limits(policy, deployment) for deployment in route.team_deployments),
|
||||
*(_provider_limits(policy, provider) for provider in route.providers),
|
||||
)
|
||||
|
||||
|
||||
def _is_positive(value: object) -> bool:
|
||||
return isinstance(value, int) and not isinstance(value, bool) and value > 0
|
||||
|
||||
|
||||
def unmet_limits(required: frozenset[QuotaLimit], values: Mapping[QuotaLimit, object]) -> tuple[QuotaLimit, ...]:
|
||||
return tuple(
|
||||
limit
|
||||
for limit in _LIMITS
|
||||
if (limit in required or values.get(limit) is not None) and not _is_positive(values.get(limit))
|
||||
)
|
||||
|
||||
|
||||
def no_route(_model: str) -> ModelRoute:
|
||||
return NO_ROUTE
|
||||
|
||||
|
||||
def _resolved_api_base(api_base: object) -> object:
|
||||
if isinstance(api_base, str) and api_base.startswith("os.environ/"):
|
||||
return get_secret_str(api_base)
|
||||
return api_base
|
||||
|
||||
|
||||
def _provider_behind_api_base(litellm_params: Mapping[str, object]) -> str | None:
|
||||
model: Final = litellm_params.get("model")
|
||||
api_base: Final = _resolved_api_base(litellm_params.get("api_base"))
|
||||
if litellm_params.get("custom_llm_provider") or not isinstance(model, str) or not isinstance(api_base, str):
|
||||
return None
|
||||
if model.split("/", 1)[0] in litellm.provider_list or not is_openai_compatible_endpoint(api_base):
|
||||
return None
|
||||
try:
|
||||
_, provider, _, _ = get_llm_provider(model=model, api_base=api_base)
|
||||
except BadRequestError:
|
||||
return None
|
||||
return provider
|
||||
|
||||
|
||||
def deployment_provider(litellm_params: object) -> str | None:
|
||||
if not is_str_keyed_mapping(litellm_params):
|
||||
return None
|
||||
return _provider_behind_api_base(litellm_params) or provider_for_generic_call(litellm_params)
|
||||
|
||||
|
||||
def _called_model(litellm_params: Mapping[str, object]) -> str | None:
|
||||
model: Final = litellm_params.get("model")
|
||||
if not isinstance(model, str):
|
||||
return None
|
||||
prefix, _, rest = model.partition("/")
|
||||
return rest if rest and _is_known_provider(prefix) else model
|
||||
|
||||
|
||||
def team_deployment(litellm_params: object) -> TeamDeployment:
|
||||
if not is_str_keyed_mapping(litellm_params):
|
||||
return TeamDeployment(provider=None, called_model=None)
|
||||
return TeamDeployment(provider=deployment_provider(litellm_params), called_model=_called_model(litellm_params))
|
||||
|
||||
|
||||
def deployment_route(litellm_params: Mapping[str, object], team_model: bool) -> ModelRoute:
|
||||
if team_model:
|
||||
return ModelRoute(team_deployments=frozenset((team_deployment(litellm_params),)))
|
||||
provider: Final = deployment_provider(litellm_params)
|
||||
return ModelRoute(providers=frozenset() if provider is None else frozenset((provider,)))
|
||||
|
||||
|
||||
def _named_deployments(llm_router: "Router", model: str, team_id: str | None) -> tuple[Mapping[str, object], ...]:
|
||||
return (
|
||||
*llm_router.get_all_deployments(model_name=model, team_id=team_id),
|
||||
*llm_router.get_model_list_from_routing_groups(model_name=model),
|
||||
)
|
||||
|
||||
|
||||
def _routed_deployments(llm_router: "Router", model: str, team_id: str | None) -> Sequence[Mapping[str, object]]:
|
||||
target: Final = llm_router.routable_model_group(model)
|
||||
return _named_deployments(llm_router, target, team_id) or llm_router.wildcard_route_deployments(target, team_id)
|
||||
|
||||
|
||||
def _is_team_deployment(deployment: Mapping[str, object]) -> bool:
|
||||
model_info: Final = deployment.get("model_info")
|
||||
return is_str_keyed_mapping(model_info) and model_info.get("team_id") is not None
|
||||
|
||||
|
||||
def router_model_routes(llm_router: "Router | None", team_id: str | None) -> RouteOf:
|
||||
if llm_router is None:
|
||||
return no_route
|
||||
|
||||
def route_of(model: str) -> ModelRoute:
|
||||
deployments: Final = _routed_deployments(llm_router, model, team_id)
|
||||
shared: Final = tuple(deployment for deployment in deployments if not _is_team_deployment(deployment))
|
||||
return ModelRoute(
|
||||
providers=frozenset(
|
||||
provider
|
||||
for deployment in shared
|
||||
if (provider := deployment_provider(deployment.get("litellm_params"))) is not None
|
||||
),
|
||||
team_deployments=frozenset(
|
||||
team_deployment(deployment.get("litellm_params"))
|
||||
for deployment in deployments
|
||||
if _is_team_deployment(deployment)
|
||||
),
|
||||
)
|
||||
|
||||
return route_of
|
||||
|
||||
|
||||
def _deployment_id(deployment: Mapping[str, object]) -> object:
|
||||
model_info: Final = deployment.get("model_info")
|
||||
return model_info.get("id") if is_str_keyed_mapping(model_info) else None
|
||||
|
||||
|
||||
def _routable_name(llm_router: "Router | None", model: str) -> str:
|
||||
return model if llm_router is None else llm_router.routable_model_group(model)
|
||||
|
||||
|
||||
def _patterns_tried_before(patterns: Sequence[str], regex: str) -> Sequence[str]:
|
||||
if regex in patterns:
|
||||
return patterns[: patterns.index(regex)]
|
||||
specificity: Final = PatternUtils.calculate_pattern_specificity(regex)
|
||||
return tuple(rival for rival in patterns if PatternUtils.calculate_pattern_specificity(rival) >= specificity)
|
||||
|
||||
|
||||
def _team_pattern_serves(llm_router: "Router", team_id: str, pattern: str, model: str) -> bool:
|
||||
regex: Final = _ROUTER_PATTERNS.pattern_to_regex(pattern)
|
||||
if re.match(regex, model) is None or _named_deployments(llm_router, model, team_id):
|
||||
return False
|
||||
team_router: Final = llm_router.team_pattern_routers.get(team_id)
|
||||
patterns: Final = () if team_router is None else tuple(team_router.patterns)
|
||||
return not any(re.match(rival, model) is not None for rival in _patterns_tried_before(patterns, regex))
|
||||
|
||||
|
||||
def team_write_serves(llm_router: "Router | None", team_id: str, public_name: str, model: str) -> bool:
|
||||
"""Whether a request from the team for ``model`` reaches the team route named ``public_name``: the name
|
||||
itself or an alias of it, or, for a `*` name, a model no named deployment shadows and no team pattern the
|
||||
router tries first takes."""
|
||||
target: Final = _routable_name(llm_router, model)
|
||||
if "*" not in public_name or llm_router is None:
|
||||
return target == public_name
|
||||
return _team_pattern_serves(llm_router, team_id, public_name, target)
|
||||
|
||||
|
||||
def _served_params(litellm_params: object, public_name: str, model: str) -> object:
|
||||
if "*" not in public_name or not is_str_keyed_mapping(litellm_params):
|
||||
return litellm_params
|
||||
matched: Final = re.match(_ROUTER_PATTERNS.pattern_to_regex(public_name), model)
|
||||
called: Final = litellm_params.get("model")
|
||||
if matched is None or not isinstance(called, str):
|
||||
return litellm_params
|
||||
return MappingProxyType(
|
||||
{
|
||||
**litellm_params,
|
||||
"model": PatternMatchRouter.set_deployment_model_name(
|
||||
matched_pattern=matched, litellm_deployment_litellm_model=called
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def team_route_after_write(
|
||||
llm_router: "Router | None",
|
||||
team_id: str,
|
||||
public_name: str,
|
||||
replaced_id: str | None,
|
||||
written_params: Mapping[str, object] | None,
|
||||
model: str,
|
||||
) -> ModelRoute | None:
|
||||
deployments: Final = () if llm_router is None else llm_router.get_all_deployments(public_name, team_id=team_id)
|
||||
kept: Final = tuple(
|
||||
deployment.get("litellm_params")
|
||||
for deployment in deployments
|
||||
if _is_team_deployment(deployment) and _deployment_id(deployment) != replaced_id
|
||||
)
|
||||
team_params: Final = kept if written_params is None else (*kept, written_params)
|
||||
target: Final = _routable_name(llm_router, model)
|
||||
if team_params:
|
||||
return ModelRoute(
|
||||
team_deployments=frozenset(
|
||||
team_deployment(_served_params(litellm_params, public_name, target)) for litellm_params in team_params
|
||||
)
|
||||
)
|
||||
if llm_router is None or not _routed_deployments(llm_router, model, None):
|
||||
return None
|
||||
return router_model_routes(llm_router, None)(model)
|
||||
|
||||
|
||||
def limits_added_by_team_write(
|
||||
policy: ModelQuotaPolicy,
|
||||
llm_router: "Router | None",
|
||||
team_id: str,
|
||||
public_name: str,
|
||||
replaced_id: str | None,
|
||||
written_params: Mapping[str, object] | None,
|
||||
model: str,
|
||||
) -> frozenset[QuotaLimit]:
|
||||
after: Final = team_route_after_write(llm_router, team_id, public_name, replaced_id, written_params, model)
|
||||
if after is None:
|
||||
return frozenset()
|
||||
now: Final = router_model_routes(llm_router, team_id)(model)
|
||||
return required_limits(policy, model, after) - required_limits(policy, model, now)
|
||||
|
||||
|
||||
def exempts_any_route(policy: ModelQuotaPolicy) -> bool:
|
||||
return any(
|
||||
limits != REQUIRE_BOTH_LIMITS.default
|
||||
for limits in (policy.default, *policy.providers.values(), *policy.models.values())
|
||||
)
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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"))
|
||||
|
|
|
|||
|
|
@ -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] = ""
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -368,6 +368,7 @@ class BaseResponsesAPIStreamingIterator:
|
|||
"""
|
||||
|
||||
_pending_logging_tasks: tuple[asyncio.Task[object], ...] = ()
|
||||
_started_tool_call: bool = False
|
||||
_earlier_rounds_spend: BilledRoundsSpend | None = None
|
||||
|
||||
def __init__(
|
||||
|
|
@ -491,12 +492,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,
|
||||
|
|
@ -1047,6 +1052,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,
|
||||
|
|
@ -1693,13 +1703,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(
|
||||
|
|
|
|||
|
|
@ -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 = (
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -4573,6 +4573,7 @@ class LlmProviders(str, Enum):
|
|||
POE = "poe"
|
||||
CHUTES = "chutes"
|
||||
NEOSANTARA = "neosantara"
|
||||
AIAND = "aiand"
|
||||
PARASAIL = "parasail"
|
||||
XIAOMI_MIMO = "xiaomi_mimo"
|
||||
TENSORMESH = "tensormesh"
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -296,7 +296,7 @@ ci = [
|
|||
"anthropic==0.84.0",
|
||||
"blockbuster==1.5.26",
|
||||
"beautifulsoup4==4.14.3",
|
||||
"pylint==4.0.5",
|
||||
"pylint==4.0.9",
|
||||
"langchain-openai==1.7.2",
|
||||
"langgraph>=1.2.4,<1.3.0",
|
||||
"langgraph-prebuilt>=1.1.0,<1.3.0",
|
||||
|
|
|
|||
|
|
@ -171,6 +171,8 @@ ignored_function_names = [
|
|||
"_deployment_field", # Read by is_batch_capable, tested through the mixed-group tests in test_router.py
|
||||
"_deployment_provider", # Read by without_deployments_that_cannot_batch, tested through the mixed-group tests in test_router.py
|
||||
"_cannot_run_batches", # Tested through the Azure and OpenAI mixed-group acreate_batch tests in test_router.py
|
||||
"_responses_continuation", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py
|
||||
"_responses_continuation_ceilings", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py
|
||||
"_team_deployments_named", # Tested through the team lookups in the concurrent delete tests in test_router.py
|
||||
]
|
||||
|
||||
|
|
|
|||
143
tests/integration/_support/riva.py
Normal file
143
tests/integration/_support/riva.py
Normal file
|
|
@ -0,0 +1,143 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import threading
|
||||
import wave
|
||||
from array import array
|
||||
from collections.abc import Iterable, Iterator
|
||||
from concurrent import futures
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from queue import SimpleQueue
|
||||
from typing import Final
|
||||
|
||||
import grpc
|
||||
from riva.client.proto import riva_asr_pb2, riva_asr_pb2_grpc
|
||||
|
||||
SAMPLE_RATE_HZ: Final = 16000
|
||||
OFFLINE_ONLY_REFUSAL: Final = (
|
||||
"Error: Unavailable model requested given these parameters: language_code=en-US; sample_rate=16000; type=online;"
|
||||
)
|
||||
_STALL_LIMIT_SECONDS: Final = 120
|
||||
|
||||
|
||||
class Rpc(Enum):
|
||||
RECOGNIZE = "Recognize"
|
||||
STREAMING = "StreamingRecognize"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Word:
|
||||
word: str
|
||||
start_ms: int
|
||||
end_ms: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RivaCall:
|
||||
rpc: Rpc
|
||||
audio: bytes
|
||||
cancelled: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RivaServer:
|
||||
address: str
|
||||
calls: SimpleQueue[RivaCall]
|
||||
|
||||
def drain(self) -> tuple[RivaCall, ...]:
|
||||
return tuple(self.calls.get_nowait() for _ in range(self.calls.qsize()))
|
||||
|
||||
|
||||
def _clip_samples(seconds: float) -> array[int]:
|
||||
return array("h", (((index * 37) % 2001 - 1000) * 16 for index in range(int(seconds * SAMPLE_RATE_HZ))))
|
||||
|
||||
|
||||
def pcm_wav(seconds: float) -> bytes:
|
||||
buffer: Final = io.BytesIO()
|
||||
with wave.open(buffer, "wb") as writer:
|
||||
writer.setnchannels(1)
|
||||
writer.setsampwidth(2)
|
||||
writer.setframerate(SAMPLE_RATE_HZ)
|
||||
writer.writeframes(_clip_samples(seconds).tobytes())
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
def pcm_bytes(seconds: float) -> int:
|
||||
return int(seconds * SAMPLE_RATE_HZ) * 2
|
||||
|
||||
|
||||
def carries_clip(audio: bytes, seconds: float) -> bool:
|
||||
sent: Final = _clip_samples(seconds)
|
||||
if len(audio) != len(sent) * 2:
|
||||
return False
|
||||
received: Final = array("h", audio)
|
||||
return all(abs(left - right) <= 1 for left, right in zip(sent, received, strict=True))
|
||||
|
||||
|
||||
class _Asr(riva_asr_pb2_grpc.RivaSpeechRecognitionServicer):
|
||||
def __init__(self, words: tuple[Word, ...], offline_only: bool, stall: bool, calls: SimpleQueue[RivaCall]) -> None:
|
||||
self._words: Final = words
|
||||
self._offline_only: Final = offline_only
|
||||
self._stall: Final = stall
|
||||
self._calls: Final = calls
|
||||
|
||||
def _alternative(self, config: riva_asr_pb2.RecognitionConfig) -> riva_asr_pb2.SpeechRecognitionAlternative:
|
||||
timed: Final = self._words if config.enable_word_time_offsets else ()
|
||||
return riva_asr_pb2.SpeechRecognitionAlternative(
|
||||
transcript=" ".join(word.word for word in self._words),
|
||||
confidence=0.9,
|
||||
words=[
|
||||
riva_asr_pb2.WordInfo(word=word.word, start_time=word.start_ms, end_time=word.end_ms) for word in timed
|
||||
],
|
||||
)
|
||||
|
||||
def Recognize(
|
||||
self, request: riva_asr_pb2.RecognizeRequest, context: grpc.ServicerContext
|
||||
) -> riva_asr_pb2.RecognizeResponse:
|
||||
if self._stall:
|
||||
released: Final = threading.Event()
|
||||
context.add_callback(released.set)
|
||||
released.wait(timeout=_STALL_LIMIT_SECONDS)
|
||||
self._calls.put(RivaCall(Rpc.RECOGNIZE, request.audio, cancelled=not context.is_active()))
|
||||
return riva_asr_pb2.RecognizeResponse()
|
||||
self._calls.put(RivaCall(Rpc.RECOGNIZE, request.audio, cancelled=False))
|
||||
return riva_asr_pb2.RecognizeResponse(
|
||||
results=[riva_asr_pb2.SpeechRecognitionResult(alternatives=[self._alternative(request.config)])]
|
||||
)
|
||||
|
||||
def StreamingRecognize(
|
||||
self, request_iterator: Iterable[riva_asr_pb2.StreamingRecognizeRequest], context: grpc.ServicerContext
|
||||
) -> Iterator[riva_asr_pb2.StreamingRecognizeResponse]:
|
||||
if self._offline_only:
|
||||
self._calls.put(RivaCall(Rpc.STREAMING, b"", cancelled=False))
|
||||
context.abort(grpc.StatusCode.INVALID_ARGUMENT, OFFLINE_ONLY_REFUSAL)
|
||||
requests: Final = tuple(request_iterator)
|
||||
config: Final = requests[0].streaming_config.config if requests else riva_asr_pb2.RecognitionConfig()
|
||||
self._calls.put(
|
||||
RivaCall(Rpc.STREAMING, b"".join(request.audio_content for request in requests), cancelled=False)
|
||||
)
|
||||
yield riva_asr_pb2.StreamingRecognizeResponse(
|
||||
results=[riva_asr_pb2.StreamingRecognitionResult(alternatives=[self._alternative(config)], is_final=True)]
|
||||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def riva_server(
|
||||
words: tuple[Word, ...],
|
||||
*,
|
||||
offline_only: bool = False,
|
||||
stall: bool = False,
|
||||
max_receive_bytes: int | None = None,
|
||||
) -> Iterator[RivaServer]:
|
||||
calls: Final[SimpleQueue[RivaCall]] = SimpleQueue()
|
||||
options: Final = () if max_receive_bytes is None else (("grpc.max_receive_message_length", max_receive_bytes),)
|
||||
server: Final = grpc.server(futures.ThreadPoolExecutor(max_workers=8), options=options)
|
||||
riva_asr_pb2_grpc.add_RivaSpeechRecognitionServicer_to_server(_Asr(words, offline_only, stall, calls), server)
|
||||
port: Final = server.add_insecure_port("127.0.0.1:0")
|
||||
server.start()
|
||||
try:
|
||||
yield RivaServer(f"127.0.0.1:{port}", calls)
|
||||
finally:
|
||||
server.stop(grace=None).wait()
|
||||
|
|
@ -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
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue